websocket handler own submodule.

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

server/internal/websocket/handler/broadcast.go

@@ -1,4 +1,4 @@
-package websocket
+package handler
 
 import (
 	"m1k1o/neko/internal/types"

server/internal/websocket/handler/chat.go

@@ -1,4 +1,4 @@
-package websocket
+package handler
 
 import (
 	"m1k1o/neko/internal/types"

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

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) {
+func New(
-	address := socket.Address()
+	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(

server/internal/websocket/handler/screen.go

@@ -1,4 +1,4 @@
-package websocket
+package handler
 
 import (
 	"m1k1o/neko/internal/types"

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

server/internal/websocket/handler/signal.go

@@ -1,4 +1,4 @@
-package websocket
+package handler
 
 import (
 	"m1k1o/neko/internal/types"

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{
+		handler:         handler,
-			logger:    logger.With().Str("subsystem", "handler").Logger(),
-			desktop:   desktop,
-			capture:   capture,
-			broadcast: broadcast,
-			sessions:  sessions,
-			webrtc:    webrtc,
-			banned:    make(map[string]string),
-			locked:    locks,
-		},
 		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,
+		Banned: ws.handler.Banned,
-		Locked: ws.handler.locked,
+		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(),
 	}
 }