websocket handler own submodule.

This commit is contained in:
Miroslav Šedivý 2022-09-12 22:36:56 +02:00
parent 777f7b4c37
commit 06e25df962
9 changed files with 70 additions and 46 deletions

View File

@ -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

View File

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

View File

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

View File

@ -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

View File

@ -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(

View File

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

View File

@ -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

View File

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

View File

@ -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(),
}
}