neko/internal/websocket/manager.go

235 lines
5.5 KiB
Go
Raw Normal View History

2020-11-02 04:09:48 +13:00
package websocket
import (
"net/http"
"time"
"github.com/gorilla/websocket"
"github.com/rs/zerolog"
"github.com/rs/zerolog/log"
"demodesk/neko/internal/websocket/handler"
"demodesk/neko/internal/types/event"
"demodesk/neko/internal/types/message"
"demodesk/neko/internal/types"
)
func New(
sessions types.SessionManager,
desktop types.DesktopManager,
capture types.CaptureManager,
webrtc types.WebRTCManager,
) *WebSocketManagerCtx {
logger := log.With().Str("module", "websocket").Logger()
return &WebSocketManagerCtx{
logger: logger,
sessions: sessions,
desktop: desktop,
upgrader: websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool {
return true
},
},
handler: handler.New(sessions, desktop, capture, webrtc),
}
}
// Send pings to peer with this period. Must be less than pongWait.
const pingPeriod = 60 * time.Second
type WebSocketManagerCtx struct {
logger zerolog.Logger
upgrader websocket.Upgrader
sessions types.SessionManager
desktop types.DesktopManager
handler *handler.MessageHandlerCtx
shutdown chan bool
}
func (ws *WebSocketManagerCtx) Start() {
ws.sessions.OnCreated(func(session types.Session) {
if err := ws.handler.SessionCreated(session); err != nil {
2020-11-02 08:53:25 +13:00
ws.logger.Warn().Str("id", session.ID()).Err(err).Msg("session created with an error")
2020-11-02 04:09:48 +13:00
} else {
ws.logger.Debug().Str("id", session.ID()).Msg("session created")
}
})
ws.sessions.OnConnected(func(session types.Session) {
if err := ws.handler.SessionConnected(session); err != nil {
2020-11-02 08:53:25 +13:00
ws.logger.Warn().Str("id", session.ID()).Err(err).Msg("session connected with an error")
2020-11-02 04:09:48 +13:00
} else {
ws.logger.Debug().Str("id", session.ID()).Msg("session connected")
}
})
2020-11-02 08:53:25 +13:00
ws.sessions.OnDisconnected(func(session types.Session) {
if err := ws.handler.SessionDisconnected(session); err != nil {
ws.logger.Warn().Str("id", session.ID()).Err(err).Msg("session disconnected with an error")
2020-11-02 04:09:48 +13:00
} else {
2020-11-02 08:53:25 +13:00
ws.logger.Debug().Str("id", session.ID()).Msg("session disconnected")
2020-11-02 04:09:48 +13:00
}
})
go func() {
2020-11-02 06:39:12 +13:00
ws.logger.Info().Msg("clipboard loop started")
2020-11-02 04:09:48 +13:00
defer func() {
2020-11-02 06:39:12 +13:00
ws.logger.Info().Msg("clipboard loop stopped")
2020-11-02 04:09:48 +13:00
}()
current := ws.desktop.ReadClipboard()
for {
select {
case <-ws.shutdown:
return
default:
session := ws.sessions.GetHost()
2020-11-17 09:57:35 +13:00
if session == nil {
2020-11-02 04:09:48 +13:00
break
}
text := ws.desktop.ReadClipboard()
if text == current {
break
}
2020-11-19 08:30:33 +13:00
if err := session.Send(
message.Clipboard{
Event: event.CONTROL_CLIPBOARD,
Text: text,
}); err != nil {
2020-11-02 04:09:48 +13:00
ws.logger.Warn().Err(err).Msg("could not sync clipboard")
}
current = text
}
time.Sleep(100 * time.Millisecond)
}
}()
}
func (ws *WebSocketManagerCtx) Shutdown() error {
ws.shutdown <- true
return nil
}
func (ws *WebSocketManagerCtx) Upgrade(w http.ResponseWriter, r *http.Request) error {
ws.logger.Debug().Msg("attempting to upgrade connection")
connection, err := ws.upgrader.Upgrade(w, r, nil)
if err != nil {
ws.logger.Error().Err(err).Msg("failed to upgrade connection")
return err
}
2020-11-02 08:23:09 +13:00
session, err := ws.sessions.Authenticate(r)
2020-11-02 04:09:48 +13:00
if err != nil {
ws.logger.Warn().Err(err).Msg("authentication failed")
// TODO: Refactor
if err = connection.WriteJSON(message.Disconnect{
Event: event.SYSTEM_DISCONNECT,
2020-11-02 08:23:09 +13:00
Message: "authentication failed",
2020-11-02 04:09:48 +13:00
}); err != nil {
ws.logger.Error().Err(err).Msg("failed to send disconnect")
}
return connection.Close()
}
2020-11-26 06:36:33 +13:00
websocket_peer := &WebSocketPeerCtx{
2020-11-15 04:03:12 +13:00
session: session,
2020-11-02 04:09:48 +13:00
ws: ws,
connection: connection,
}
2020-11-26 06:36:33 +13:00
ok, reason := ws.handler.Connected(session, websocket_peer)
2020-11-02 04:09:48 +13:00
if !ok {
// TODO: Refactor
if err = connection.WriteJSON(message.Disconnect{
Event: event.SYSTEM_DISCONNECT,
Message: reason,
}); err != nil {
ws.logger.Error().Err(err).Msg("failed to send disconnect")
}
return connection.Close()
}
2020-11-26 06:36:33 +13:00
session.SetWebSocketPeer(websocket_peer)
2020-11-02 04:09:48 +13:00
ws.logger.
Debug().
2020-11-02 08:23:09 +13:00
Str("session", session.ID()).
2020-11-02 04:09:48 +13:00
Str("address", connection.RemoteAddr().String()).
2020-11-15 04:03:39 +13:00
Msg("connection started")
2020-11-02 04:09:48 +13:00
defer func() {
ws.logger.
Debug().
2020-11-02 08:23:09 +13:00
Str("session", session.ID()).
2020-11-02 04:09:48 +13:00
Str("address", connection.RemoteAddr().String()).
2020-11-15 04:03:39 +13:00
Msg("connection ended")
2020-11-02 04:09:48 +13:00
}()
2020-11-02 08:53:25 +13:00
ws.handle(connection, session)
2020-11-02 04:09:48 +13:00
return nil
}
2020-11-02 08:53:25 +13:00
func (ws *WebSocketManagerCtx) handle(connection *websocket.Conn, session types.Session) {
2020-11-02 04:09:48 +13:00
bytes := make(chan []byte)
cancel := make(chan struct{})
ticker := time.NewTicker(pingPeriod)
go func() {
defer func() {
ticker.Stop()
ws.logger.Debug().Str("address", connection.RemoteAddr().String()).Msg("handle socket ending")
2020-11-02 09:03:25 +13:00
session.SetConnected(false)
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) {
ws.logger.Warn().Err(err).Msg("read message error")
} else {
ws.logger.Debug().Err(err).Msg("read message error")
}
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:
ws.logger.Debug().
2020-11-02 08:53:25 +13:00
Str("session", session.ID()).
2020-11-02 04:09:48 +13:00
Str("address", connection.RemoteAddr().String()).
Str("raw", string(raw)).
Msg("received message from client")
2020-11-15 04:03:12 +13:00
if err := ws.handler.Message(session, raw); err != nil {
2020-11-02 04:09:48 +13:00
ws.logger.Error().Err(err).Msg("message handler has failed")
}
case <-cancel:
return
case <-ticker.C:
if err := connection.WriteMessage(websocket.PingMessage, nil); err != nil {
ws.logger.Error().Err(err).Msg("ping message has failed")
return
}
}
}
}