From 06e25df962fcb3b36018d1d9d77491e7c337119d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Miroslav=20=C5=A0ediv=C3=BD?= Date: Mon, 12 Sep 2022 22:36:56 +0200 Subject: [PATCH] websocket handler own submodule. --- .../internal/websocket/{ => handler}/admin.go | 14 +++--- .../websocket/{ => handler}/broadcast.go | 2 +- .../internal/websocket/{ => handler}/chat.go | 2 +- .../websocket/{ => handler}/control.go | 6 +-- .../websocket/{ => handler}/handler.go | 38 +++++++++++---- .../websocket/{ => handler}/screen.go | 2 +- .../websocket/{ => handler}/session.go | 4 +- .../websocket/{ => handler}/signal.go | 2 +- server/internal/websocket/websocket.go | 46 ++++++++++--------- 9 files changed, 70 insertions(+), 46 deletions(-) rename server/internal/websocket/{ => handler}/admin.go (96%) rename server/internal/websocket/{ => handler}/broadcast.go (99%) rename server/internal/websocket/{ => handler}/chat.go (98%) rename server/internal/websocket/{ => handler}/control.go (98%) rename server/internal/websocket/{ => handler}/handler.go (85%) rename server/internal/websocket/{ => handler}/screen.go (98%) rename server/internal/websocket/{ => handler}/session.go (98%) rename server/internal/websocket/{ => handler}/signal.go (98%) diff --git a/server/internal/websocket/admin.go b/server/internal/websocket/handler/admin.go similarity index 96% rename from server/internal/websocket/admin.go rename to server/internal/websocket/handler/admin.go index 72e38a4..462cd33 100644 --- a/server/internal/websocket/admin.go +++ b/server/internal/websocket/handler/admin.go @@ -1,4 +1,4 @@ -package websocket +package handler import ( "strings" @@ -14,7 +14,7 @@ func (h *MessageHandler) adminLock(id string, session types.Session, payload *me return nil } - _, ok := h.locked[payload.Resource] + _, ok := h.Locked[payload.Resource] if ok { h.logger.Debug().Str("resource", payload.Resource).Msg("resource already locked...") return nil @@ -30,7 +30,7 @@ func (h *MessageHandler) adminLock(id string, session types.Session, payload *me h.sessions.SetControlLocked(true) } - h.locked[payload.Resource] = id + h.Locked[payload.Resource] = id if err := h.sessions.Broadcast( message.AdminLock{ @@ -51,7 +51,7 @@ func (h *MessageHandler) adminUnlock(id string, session types.Session, payload * return nil } - _, ok := h.locked[payload.Resource] + _, ok := h.Locked[payload.Resource] if !ok { h.logger.Debug().Str("resource", payload.Resource).Msg("resource not locked...") return nil @@ -62,7 +62,7 @@ func (h *MessageHandler) adminUnlock(id string, session types.Session, payload * h.sessions.SetControlLocked(false) } - delete(h.locked, payload.Resource) + delete(h.Locked, payload.Resource) if err := h.sessions.Broadcast( message.AdminLock{ @@ -114,7 +114,7 @@ func (h *MessageHandler) adminControl(id string, session types.Session) error { return nil } -func (h *MessageHandler) adminRelease(id string, session types.Session) error { +func (h *MessageHandler) AdminRelease(id string, session types.Session) error { if !session.Admin() { h.logger.Debug().Msg("user not admin") return nil @@ -302,7 +302,7 @@ func (h *MessageHandler) adminBan(id string, session types.Session, payload *mes } h.logger.Debug().Str("address", remote).Msg("adding address to banned") - h.banned[address[0]] = id + h.Banned[address[0]] = id if err := target.Kick("banned"); err != nil { return err diff --git a/server/internal/websocket/broadcast.go b/server/internal/websocket/handler/broadcast.go similarity index 99% rename from server/internal/websocket/broadcast.go rename to server/internal/websocket/handler/broadcast.go index 865c031..947bb53 100644 --- a/server/internal/websocket/broadcast.go +++ b/server/internal/websocket/handler/broadcast.go @@ -1,4 +1,4 @@ -package websocket +package handler import ( "m1k1o/neko/internal/types" diff --git a/server/internal/websocket/chat.go b/server/internal/websocket/handler/chat.go similarity index 98% rename from server/internal/websocket/chat.go rename to server/internal/websocket/handler/chat.go index 6b1392d..6b7ebf6 100644 --- a/server/internal/websocket/chat.go +++ b/server/internal/websocket/handler/chat.go @@ -1,4 +1,4 @@ -package websocket +package handler import ( "m1k1o/neko/internal/types" diff --git a/server/internal/websocket/control.go b/server/internal/websocket/handler/control.go similarity index 98% rename from server/internal/websocket/control.go rename to server/internal/websocket/handler/control.go index c24bc3a..4ac12eb 100644 --- a/server/internal/websocket/control.go +++ b/server/internal/websocket/handler/control.go @@ -1,4 +1,4 @@ -package websocket +package handler import ( "m1k1o/neko/internal/types" @@ -34,7 +34,7 @@ func (h *MessageHandler) controlRequest(id string, session types.Session) error // check for host if !h.sessions.HasHost() { // check if control is locked or user is admin - _, ok := h.locked["control"] + _, ok := h.Locked["control"] if ok && !session.Admin() { h.logger.Debug().Msg("control is locked") return nil @@ -98,7 +98,7 @@ func (h *MessageHandler) controlGive(id string, session types.Session, payload * } // check if control is locked or giver is admin - _, ok := h.locked["control"] + _, ok := h.Locked["control"] if ok && !session.Admin() { h.logger.Debug().Msg("control is locked") return nil diff --git a/server/internal/websocket/handler.go b/server/internal/websocket/handler/handler.go similarity index 85% rename from server/internal/websocket/handler.go rename to server/internal/websocket/handler/handler.go index e92832d..e45f432 100644 --- a/server/internal/websocket/handler.go +++ b/server/internal/websocket/handler/handler.go @@ -1,10 +1,11 @@ -package websocket +package handler import ( "encoding/json" "github.com/pkg/errors" "github.com/rs/zerolog" + "github.com/rs/zerolog/log" "m1k1o/neko/internal/types" "m1k1o/neko/internal/types/event" @@ -15,27 +16,46 @@ import ( type MessageHandler struct { logger zerolog.Logger sessions types.SessionManager - webrtc types.WebRTCManager desktop types.DesktopManager capture types.CaptureManager + webrtc types.WebRTCManager broadcast types.BroadcastManager - banned map[string]string // IP -> session ID (that banned it) - locked map[string]string // resource name -> session ID (that locked it) + + Banned map[string]string // IP -> session ID (that banned it) + Locked map[string]string // resource name -> session ID (that locked it) } -func (h *MessageHandler) Connected(admin bool, socket *WebSocket) (bool, string) { - address := socket.Address() +func New( + sessions types.SessionManager, + desktop types.DesktopManager, + capture types.CaptureManager, + webrtc types.WebRTCManager, + broadcast types.BroadcastManager, +) *MessageHandler { + return &MessageHandler{ + logger: log.With().Str("module", "websocket").Str("submodule", "handler").Logger(), + sessions: sessions, + desktop: desktop, + capture: capture, + webrtc: webrtc, + broadcast: broadcast, + Banned: make(map[string]string), + Locked: make(map[string]string), + } +} + +func (h *MessageHandler) Connected(admin bool, address string) (bool, string) { if address == "" { h.logger.Debug().Msg("no remote address") } else { - _, ok := h.banned[address] + _, ok := h.Banned[address] if ok { h.logger.Debug().Str("address", address).Msg("banned") return false, "banned" } } - _, ok := h.locked["login"] + _, ok := h.Locked["login"] if ok && !admin { h.logger.Debug().Msg("server locked") return false, "locked" @@ -150,7 +170,7 @@ func (h *MessageHandler) Message(id string, raw []byte) error { case event.ADMIN_CONTROL: return errors.Wrapf(h.adminControl(id, session), "%s failed", header.Event) case event.ADMIN_RELEASE: - return errors.Wrapf(h.adminRelease(id, session), "%s failed", header.Event) + return errors.Wrapf(h.AdminRelease(id, session), "%s failed", header.Event) case event.ADMIN_GIVE: payload := &message.Admin{} return errors.Wrapf( diff --git a/server/internal/websocket/screen.go b/server/internal/websocket/handler/screen.go similarity index 98% rename from server/internal/websocket/screen.go rename to server/internal/websocket/handler/screen.go index d1648f5..d59da88 100644 --- a/server/internal/websocket/screen.go +++ b/server/internal/websocket/handler/screen.go @@ -1,4 +1,4 @@ -package websocket +package handler import ( "m1k1o/neko/internal/types" diff --git a/server/internal/websocket/session.go b/server/internal/websocket/handler/session.go similarity index 98% rename from server/internal/websocket/session.go rename to server/internal/websocket/handler/session.go index 2f5fb65..100aff3 100644 --- a/server/internal/websocket/session.go +++ b/server/internal/websocket/handler/session.go @@ -1,4 +1,4 @@ -package websocket +package handler import ( "m1k1o/neko/internal/types" @@ -16,7 +16,7 @@ func (h *MessageHandler) SessionCreated(id string, session types.Session) error if err := session.Send(message.SystemInit{ Event: event.SYSTEM_INIT, ImplicitHosting: h.webrtc.ImplicitControl(), - Locks: h.locked, + Locks: h.Locked, }); err != nil { h.logger.Warn().Str("id", id).Err(err).Msgf("sending event %s has failed", event.SYSTEM_INIT) return err diff --git a/server/internal/websocket/signal.go b/server/internal/websocket/handler/signal.go similarity index 98% rename from server/internal/websocket/signal.go rename to server/internal/websocket/handler/signal.go index d416adb..409d3a7 100644 --- a/server/internal/websocket/signal.go +++ b/server/internal/websocket/handler/signal.go @@ -1,4 +1,4 @@ -package websocket +package handler import ( "m1k1o/neko/internal/types" diff --git a/server/internal/websocket/websocket.go b/server/internal/websocket/websocket.go index 29f908c..11e550c 100644 --- a/server/internal/websocket/websocket.go +++ b/server/internal/websocket/websocket.go @@ -16,6 +16,7 @@ import ( "m1k1o/neko/internal/types/event" "m1k1o/neko/internal/types/message" "m1k1o/neko/internal/utils" + "m1k1o/neko/internal/websocket/handler" ) const CONTROL_PROTECTION_SESSION = "by_control_protection" @@ -40,28 +41,30 @@ func New(sessions types.SessionManager, desktop types.DesktopManager, capture ty logger.Info().Msgf("locked resources: %+v", conf.Locks) } + handler := handler.New( + sessions, + desktop, + capture, + webrtc, + broadcast, + ) + + // set inital locks + handler.Locked = locks + return &WebSocketHandler{ logger: logger, shutdown: make(chan interface{}), conf: conf, sessions: sessions, desktop: desktop, - + webrtc: webrtc, upgrader: websocket.Upgrader{ CheckOrigin: func(r *http.Request) bool { return true }, }, - handler: &MessageHandler{ - logger: logger.With().Str("subsystem", "handler").Logger(), - desktop: desktop, - capture: capture, - broadcast: broadcast, - sessions: sessions, - webrtc: webrtc, - banned: make(map[string]string), - locked: locks, - }, + handler: handler, serverStartedAt: time.Now(), } } @@ -76,8 +79,9 @@ type WebSocketHandler struct { upgrader websocket.Upgrader sessions types.SessionManager desktop types.DesktopManager + webrtc types.WebRTCManager conf *config.WebSocket - handler *MessageHandler + handler *handler.MessageHandler // stats conns uint32 @@ -104,9 +108,9 @@ func (ws *WebSocketHandler) Start() { // if control protection is enabled and at least one admin // and if room was locked on behalf control protection, unlock - sess, ok := ws.handler.locked["control"] + sess, ok := ws.handler.Locked["control"] if ok && ws.conf.ControlProtection && sess == CONTROL_PROTECTION_SESSION && len(ws.sessions.Admins()) > 0 { - delete(ws.handler.locked, "control") + delete(ws.handler.Locked, "control") ws.sessions.SetControlLocked(false) // TODO: Handle locks in sessions as flags. ws.logger.Info().Msgf("control unlocked on behalf of control protection") @@ -140,12 +144,12 @@ func (ws *WebSocketHandler) Start() { // if control protection is enabled and no admin // and room is not locked, lock - _, ok := ws.handler.locked["control"] + _, ok := ws.handler.Locked["control"] if !ok && ws.conf.ControlProtection && adminCount == 0 { - ws.handler.locked["control"] = CONTROL_PROTECTION_SESSION + ws.handler.Locked["control"] = CONTROL_PROTECTION_SESSION ws.sessions.SetControlLocked(true) // TODO: Handle locks in sessions as flags. ws.logger.Info().Msgf("control locked and released on behalf of control protection") - ws.handler.adminRelease(id, session) + ws.handler.AdminRelease(id, session) if err := ws.sessions.Broadcast( message.AdminLock{ @@ -258,7 +262,7 @@ func (ws *WebSocketHandler) Upgrade(w http.ResponseWriter, r *http.Request) erro connection: connection, } - ok, reason := ws.handler.Connected(admin, socket) + ok, reason := ws.handler.Connected(admin, socket.Address()) if !ok { if err = connection.WriteJSON(message.SystemMessage{ Event: event.SYSTEM_DISCONNECT, @@ -310,15 +314,15 @@ func (ws *WebSocketHandler) Stats() types.Stats { Host: host, Members: ws.sessions.Members(), - Banned: ws.handler.banned, - Locked: ws.handler.locked, + Banned: ws.handler.Banned, + Locked: ws.handler.Locked, ServerStartedAt: ws.serverStartedAt, LastAdminLeftAt: ws.lastAdminLeftAt, LastUserLeftAt: ws.lastUserLeftAt, ControlProtection: ws.conf.ControlProtection, - ImplicitControl: ws.handler.webrtc.ImplicitControl(), + ImplicitControl: ws.webrtc.ImplicitControl(), } }