neko/internal/websocket/manager.go

345 lines
8.8 KiB
Go
Raw Normal View History

2020-11-02 04:09:48 +13:00
package websocket
import (
2021-08-30 09:00:51 +12:00
"encoding/json"
2020-11-02 04:09:48 +13:00
"net/http"
"time"
"github.com/gorilla/websocket"
"github.com/rs/zerolog"
"github.com/rs/zerolog/log"
2021-01-10 11:01:39 +13:00
"demodesk/neko/internal/types"
2020-11-02 04:09:48 +13:00
"demodesk/neko/internal/types/event"
"demodesk/neko/internal/types/message"
2021-02-15 02:40:17 +13:00
"demodesk/neko/internal/websocket/handler"
2020-11-02 04:09:48 +13:00
)
func New(
sessions types.SessionManager,
desktop types.DesktopManager,
capture types.CaptureManager,
webrtc types.WebRTCManager,
) *WebSocketManagerCtx {
logger := log.With().Str("module", "websocket").Logger()
return &WebSocketManagerCtx{
2021-02-15 02:40:17 +13:00
logger: logger,
sessions: sessions,
desktop: desktop,
handler: handler.New(sessions, desktop, capture, webrtc),
2021-08-30 09:01:51 +12:00
handlers: []types.WebSocketHandler{},
2020-11-02 04:09:48 +13:00
}
}
// Send pings to peer with this period. Must be less than pongWait.
2021-03-04 06:35:58 +13:00
const pingPeriod = 10 * time.Second
2020-11-02 04:09:48 +13:00
type WebSocketManagerCtx struct {
2021-02-15 02:40:17 +13:00
logger zerolog.Logger
sessions types.SessionManager
desktop types.DesktopManager
handler *handler.MessageHandlerCtx
2021-08-30 09:01:51 +12:00
handlers []types.WebSocketHandler
2020-11-02 04:09:48 +13:00
}
2021-02-15 05:11:21 +13:00
func (manager *WebSocketManagerCtx) Start() {
manager.sessions.OnCreated(func(session types.Session) {
2021-08-30 04:23:58 +12:00
logger := manager.logger.With().Str("session_id", session.ID()).Logger()
2021-02-15 05:11:21 +13:00
if err := manager.handler.SessionCreated(session); err != nil {
2021-08-30 04:23:58 +12:00
logger.Warn().Err(err).Msg("session created with an error")
2020-12-03 06:59:54 +13:00
} else {
2021-08-30 04:23:58 +12:00
logger.Debug().Msg("session created")
2020-12-03 06:59:54 +13:00
}
})
2021-02-15 05:11:21 +13:00
manager.sessions.OnDeleted(func(session types.Session) {
2021-08-30 04:23:58 +12:00
logger := manager.logger.With().Str("session_id", session.ID()).Logger()
2021-02-15 05:11:21 +13:00
if err := manager.handler.SessionDeleted(session); err != nil {
2021-08-30 04:23:58 +12:00
logger.Warn().Err(err).Msg("session deleted with an error")
2020-12-03 06:59:54 +13:00
} else {
2021-08-30 04:23:58 +12:00
logger.Debug().Msg("session deleted")
2020-12-03 06:59:54 +13:00
}
})
2021-02-15 05:11:21 +13:00
manager.sessions.OnConnected(func(session types.Session) {
2021-08-30 04:23:58 +12:00
logger := manager.logger.With().Str("session_id", session.ID()).Logger()
2021-02-15 05:11:21 +13:00
if err := manager.handler.SessionConnected(session); err != nil {
2021-08-30 04:23:58 +12:00
logger.Warn().Err(err).Msg("session connected with an error")
2020-11-02 04:09:48 +13:00
} else {
2021-08-30 04:23:58 +12:00
logger.Debug().Msg("session connected")
2020-11-02 04:09:48 +13:00
}
})
2021-02-15 05:11:21 +13:00
manager.sessions.OnDisconnected(func(session types.Session) {
2021-08-30 04:23:58 +12:00
logger := manager.logger.With().Str("session_id", session.ID()).Logger()
2021-02-15 05:11:21 +13:00
if err := manager.handler.SessionDisconnected(session); err != nil {
2021-08-30 04:23:58 +12:00
logger.Warn().Err(err).Msg("session disconnected with an error")
2020-11-02 04:09:48 +13:00
} else {
2021-08-30 04:23:58 +12:00
logger.Debug().Msg("session disconnected")
2020-11-02 04:09:48 +13:00
}
})
2021-02-15 05:11:21 +13:00
manager.sessions.OnProfileChanged(func(session types.Session) {
2021-08-30 04:23:58 +12:00
logger := manager.logger.With().Str("session_id", session.ID()).Logger()
2021-02-15 05:11:21 +13:00
if err := manager.handler.SessionProfileChanged(session); err != nil {
2021-08-30 04:23:58 +12:00
logger.Warn().Err(err).Msg("session profile changed with an error")
2020-12-03 06:59:54 +13:00
} else {
2021-08-30 09:03:04 +12:00
logger.Debug().Interface("profile", session.Profile()).Msg("session profile changed")
2020-12-03 06:59:54 +13:00
}
})
2021-02-15 05:11:21 +13:00
manager.sessions.OnStateChanged(func(session types.Session) {
2021-08-30 04:23:58 +12:00
logger := manager.logger.With().Str("session_id", session.ID()).Logger()
2021-02-15 05:11:21 +13:00
if err := manager.handler.SessionStateChanged(session); err != nil {
2021-08-30 04:23:58 +12:00
logger.Warn().Err(err).Msg("session state changed with an error")
2020-12-03 06:59:54 +13:00
} else {
2021-08-30 09:03:04 +12:00
logger.Debug().Interface("state", session.State()).Msg("session state changed")
2020-12-03 06:59:54 +13:00
}
})
2021-03-14 09:37:00 +13:00
manager.sessions.OnHostChanged(func(session types.Session) {
msg := message.ControlHost{
Event: event.CONTROL_HOST,
HasHost: session != nil,
}
if msg.HasHost {
msg.HostID = session.ID()
}
manager.sessions.Broadcast(msg, nil)
2021-08-30 04:23:58 +12:00
manager.logger.Debug().
Bool("has_host", msg.HasHost).
Str("host_id", msg.HostID).
Msg("session host changed")
2021-03-14 09:37:00 +13:00
})
2021-02-15 05:11:21 +13:00
manager.desktop.OnClipboardUpdated(func() {
session := manager.sessions.GetHost()
2021-03-14 12:45:51 +13:00
if session == nil || !session.Profile().CanAccessClipboard {
2021-01-12 03:30:53 +13:00
return
}
2021-02-15 05:11:21 +13:00
data, err := manager.desktop.ClipboardGetText()
2021-01-29 09:12:35 +13:00
if err != nil {
2021-02-15 05:11:21 +13:00
manager.logger.Warn().Err(err).Msg("could not get clipboard content")
2021-02-02 00:23:37 +13:00
return
2021-01-29 09:12:35 +13:00
}
if err := session.Send(message.ClipboardData{
2021-01-12 03:30:53 +13:00
Event: event.CLIPBOARD_UPDATED,
Text: data.Text,
// TODO: Send HTML?
2021-01-29 09:12:35 +13:00
}); err != nil {
2021-02-15 05:11:21 +13:00
manager.logger.Warn().Err(err).Msg("could not sync clipboard")
2021-08-30 04:23:58 +12:00
return
2020-11-02 04:09:48 +13:00
}
2021-08-30 04:23:58 +12:00
manager.logger.Debug().Msg("session sync clipboard")
2021-01-12 03:30:53 +13:00
})
2021-01-12 12:09:43 +13:00
2021-02-15 05:11:21 +13:00
manager.fileChooserDialogEvents()
2021-08-30 04:23:58 +12:00
manager.logger.Info().Msg("websocket starting")
2020-11-02 04:09:48 +13:00
}
2021-02-15 05:11:21 +13:00
func (manager *WebSocketManagerCtx) Shutdown() error {
2021-08-30 04:23:58 +12:00
manager.logger.Info().Msg("websocket shutdown")
2020-11-02 04:09:48 +13:00
return nil
}
2021-08-30 09:01:51 +12:00
func (manager *WebSocketManagerCtx) AddHandler(handler types.WebSocketHandler) {
2021-02-15 05:11:21 +13:00
manager.handlers = append(manager.handlers, handler)
2020-12-13 05:40:10 +13:00
}
2021-03-16 22:45:22 +13:00
func (manager *WebSocketManagerCtx) Upgrade(w http.ResponseWriter, r *http.Request, checkOrigin types.CheckOrigin) {
2021-08-30 09:03:04 +12:00
// add request data to logger context
logger := manager.logger.With().
Str("address", r.RemoteAddr).
Str("agent", r.UserAgent()).
Logger()
logger.Debug().Msg("attempting to upgrade connection")
2020-11-02 04:09:48 +13:00
2021-01-24 06:18:14 +13:00
upgrader := websocket.Upgrader{
CheckOrigin: checkOrigin,
}
connection, err := upgrader.Upgrade(w, r, nil)
2020-11-02 04:09:48 +13:00
if err != nil {
2021-08-30 09:03:04 +12:00
logger.Error().Err(err).Msg("failed to upgrade connection")
2021-03-16 22:45:22 +13:00
return
2020-11-02 04:09:48 +13:00
}
2021-03-14 08:42:56 +13:00
session, err := manager.sessions.Authenticate(r)
2020-11-02 04:09:48 +13:00
if err != nil {
2021-08-30 09:03:04 +12:00
logger.Debug().Err(err).Msg("authentication failed")
2020-11-02 04:09:48 +13:00
2020-12-07 06:49:08 +13:00
// TODO: Refactor, return error code.
2020-12-02 07:44:23 +13:00
if err = connection.WriteJSON(
message.SystemDisconnect{
Event: event.SYSTEM_DISCONNECT,
2020-12-07 06:49:08 +13:00
Message: err.Error(),
}); err != nil {
2021-08-30 09:03:04 +12:00
logger.Error().Err(err).Msg("failed to send disconnect event")
2020-12-07 06:49:08 +13:00
}
2021-03-16 22:45:22 +13:00
if err := connection.Close(); err != nil {
2021-08-30 09:03:04 +12:00
logger.Warn().Err(err).Msg("connection closed with an error")
2021-03-16 22:45:22 +13:00
}
return
2020-12-07 06:49:08 +13:00
}
2021-08-30 09:03:04 +12:00
// use session id with defeault logger context
logger = manager.logger.With().Str("session_id", session.ID()).Logger()
2021-08-30 04:23:58 +12:00
2021-03-14 12:45:51 +13:00
if !session.Profile().CanConnect {
2021-08-30 04:23:58 +12:00
logger.Debug().Msg("connection disabled")
2021-03-16 22:45:22 +13:00
2020-12-07 06:49:08 +13:00
// TODO: Refactor, return error code.
if err = connection.WriteJSON(
message.SystemDisconnect{
Event: event.SYSTEM_DISCONNECT,
Message: "connection disabled",
2020-12-02 07:44:23 +13:00
}); err != nil {
2021-08-30 04:23:58 +12:00
logger.Error().Err(err).Msg("failed to send disconnect event")
2021-03-16 22:45:22 +13:00
}
if err := connection.Close(); err != nil {
2021-08-30 04:23:58 +12:00
logger.Warn().Err(err).Msg("connection closed with an error")
2020-11-02 04:09:48 +13:00
}
2021-03-16 22:45:22 +13:00
return
2020-11-02 04:09:48 +13:00
}
2021-03-14 12:50:08 +13:00
if session.State().IsConnected {
2021-08-30 04:23:58 +12:00
logger.Warn().Msg("already connected")
2021-08-24 04:41:19 +12:00
if !manager.sessions.MercifulReconnect() {
// TODO: Refactor, return error code.
if err = connection.WriteJSON(
message.SystemDisconnect{
Event: event.SYSTEM_DISCONNECT,
Message: "already connected",
}); err != nil {
2021-08-30 04:23:58 +12:00
logger.Error().Err(err).Msg("failed to send disconnect event")
2021-08-24 04:41:19 +12:00
}
if err := connection.Close(); err != nil {
2021-08-30 04:23:58 +12:00
logger.Warn().Err(err).Msg("connection closed with an error")
2021-08-24 04:41:19 +12:00
}
2021-03-16 22:45:22 +13:00
2021-08-24 04:41:19 +12:00
return
2021-03-16 22:45:22 +13:00
}
2021-08-30 04:23:58 +12:00
logger.Info().Msg("replacing peer connection")
2021-08-30 05:17:10 +12:00
// destroy previous peer connection
session.GetWebSocketPeer().Destroy()
}
2021-03-24 03:38:37 +13:00
peer := &WebSocketPeerCtx{
2021-08-30 04:23:58 +12:00
logger: logger,
2020-11-15 04:03:12 +13:00
session: session,
2020-11-02 04:09:48 +13:00
connection: connection,
2021-03-24 03:38:37 +13:00
}
session.SetWebSocketPeer(peer)
2020-11-02 04:09:48 +13:00
2021-08-30 04:23:58 +12:00
logger.Info().
2020-11-02 04:09:48 +13:00
Str("address", connection.RemoteAddr().String()).
2021-08-30 09:03:04 +12:00
Str("agent", r.UserAgent()).
2020-11-15 04:03:39 +13:00
Msg("connection started")
2020-11-02 04:09:48 +13:00
2021-03-24 03:38:37 +13:00
session.SetWebSocketConnected(peer, true)
2020-11-26 08:26:11 +13:00
2020-11-02 04:09:48 +13:00
defer func() {
2021-08-30 04:23:58 +12:00
logger.Info().
2020-11-02 04:09:48 +13:00
Str("address", connection.RemoteAddr().String()).
2021-08-30 09:03:04 +12:00
Str("agent", r.UserAgent()).
2020-11-15 04:03:39 +13:00
Msg("connection ended")
2020-11-26 08:26:11 +13:00
2021-03-24 03:38:37 +13:00
session.SetWebSocketConnected(peer, false)
2020-11-02 04:09:48 +13:00
}()
2021-02-15 05:11:21 +13:00
manager.handle(connection, session)
2020-11-02 04:09:48 +13:00
}
2021-02-15 05:11:21 +13:00
func (manager *WebSocketManagerCtx) handle(connection *websocket.Conn, session types.Session) {
2021-08-30 04:23:58 +12:00
// add session id to logger context
logger := manager.logger.With().Str("session_id", session.ID()).Logger()
2020-11-02 04:09:48 +13:00
bytes := make(chan []byte)
cancel := make(chan struct{})
2021-01-24 06:18:14 +13:00
ticker := time.NewTicker(pingPeriod)
defer ticker.Stop()
2020-11-02 04:09:48 +13:00
2020-11-26 08:26:11 +13:00
go func() {
2020-11-02 04:09:48 +13:00
for {
_, raw, err := connection.ReadMessage()
2020-11-02 04:37:00 +13:00
if err != nil {
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
2021-08-30 04:23:58 +12:00
logger.Warn().Err(err).Msg("read message error")
2020-11-02 04:37:00 +13:00
} else {
2021-08-30 04:23:58 +12:00
logger.Debug().Err(err).Msg("read message error")
2020-11-02 04:37:00 +13:00
}
2021-02-03 06:28:32 +13:00
2020-11-02 04:37:00 +13:00
close(cancel)
break
2020-11-02 04:09:48 +13:00
}
2020-11-02 04:37:00 +13:00
bytes <- raw
2020-11-02 04:09:48 +13:00
}
}()
for {
select {
case raw := <-bytes:
2021-08-30 09:00:51 +12:00
data := types.WebSocketMessage{}
if err := json.Unmarshal(raw, &data); err != nil {
logger.Error().Err(err).Msg("message parsing has failed")
break
}
// TODO: Switch to payload based messages.
data.Payload = raw
2021-08-30 04:23:58 +12:00
logger.Debug().
2020-11-02 04:09:48 +13:00
Str("address", connection.RemoteAddr().String()).
2021-08-30 09:00:51 +12:00
Str("event", data.Event).
Str("payload", string(data.Payload)).
2020-11-02 04:09:48 +13:00
Msg("received message from client")
2021-08-30 09:00:51 +12:00
handled := manager.handler.Message(session, data)
2021-02-15 05:11:21 +13:00
for _, handler := range manager.handlers {
2020-12-13 05:40:10 +13:00
if handled {
break
}
2021-08-30 09:00:51 +12:00
handled = handler(session, data)
2020-12-13 05:40:10 +13:00
}
if !handled {
2021-08-30 09:00:51 +12:00
logger.Warn().Str("event", data.Event).Msg("unhandled message")
2020-11-02 04:09:48 +13:00
}
case <-cancel:
return
case <-ticker.C:
if err := connection.WriteMessage(websocket.PingMessage, nil); err != nil {
2021-08-30 04:23:58 +12:00
logger.Error().Err(err).Msg("ping message has failed")
2020-11-02 04:09:48 +13:00
return
}
}
}
}