iridium/neko
2
2
Fork 0
mirror of https://github.com/m1k1o/neko.git synced 2024-07-24 14:40:50 +12:00
Code Issues Packages Projects Releases Wiki Activity

major refactor.

Browse Source
This commit is contained in:
Miroslav Šedivý
2020-11-01 16:09:48 +01:00
parent 5c92b75cf7
commit 5d906e0a8b
53 changed files with 1189 additions and 1133 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
internal/websocket/broadcast/screen.go
View File

@ -1,17 +0,0 @@
package broadcast
import (
"demodesk/neko/internal/types"
"demodesk/neko/internal/types/event"
"demodesk/neko/internal/types/message"
)
func ScreenConfiguration(session types.SessionManager, id string, width int, height int, rate int) error {
return session.Broadcast(message.ScreenResolution{
Event: event.SCREEN_RESOLUTION,
ID: id,
Width: width,
Height: height,
Rate: rate,
}, nil)
}

20
internal/websocket/admin.go → internal/websocket/handler/admin.go
View File

@ -1,4 +1,4 @@
package websocket
package handler
import (
"strings"
@ -8,7 +8,7 @@ import (
"demodesk/neko/internal/types/message"
)
func (h *MessageHandler) adminLock(session types.Session) error {
func (h *MessageHandlerCtx) adminLock(session types.Session) error {
if !session.Admin() {
h.logger.Debug().Msg("user not admin")
return nil
@ -33,7 +33,7 @@ func (h *MessageHandler) adminLock(session types.Session) error {
return nil
}
func (h *MessageHandler) adminUnlock(session types.Session) error {
func (h *MessageHandlerCtx) adminUnlock(session types.Session) error {
if !session.Admin() {
h.logger.Debug().Msg("user not admin")
return nil
@ -58,7 +58,7 @@ func (h *MessageHandler) adminUnlock(session types.Session) error {
return nil
}
func (h *MessageHandler) adminControl(session types.Session) error {
func (h *MessageHandlerCtx) adminControl(session types.Session) error {
if !session.Admin() {
h.logger.Debug().Msg("user not admin")
return nil
@ -91,7 +91,7 @@ func (h *MessageHandler) adminControl(session types.Session) error {
return nil
}
func (h *MessageHandler) adminRelease(session types.Session) error {
func (h *MessageHandlerCtx) adminRelease(session types.Session) error {
if !session.Admin() {
h.logger.Debug().Msg("user not admin")
return nil
@ -124,7 +124,7 @@ func (h *MessageHandler) adminRelease(session types.Session) error {
return nil
}
func (h *MessageHandler) adminGive(session types.Session, payload *message.Admin) error {
func (h *MessageHandlerCtx) adminGive(session types.Session, payload *message.Admin) error {
if !session.Admin() {
h.logger.Debug().Msg("user not admin")
return nil
@ -153,7 +153,7 @@ func (h *MessageHandler) adminGive(session types.Session, payload *message.Admin
return nil
}
func (h *MessageHandler) adminMute(session types.Session, payload *message.Admin) error {
func (h *MessageHandlerCtx) adminMute(session types.Session, payload *message.Admin) error {
if !session.Admin() {
h.logger.Debug().Msg("user not admin")
return nil
@ -185,7 +185,7 @@ func (h *MessageHandler) adminMute(session types.Session, payload *message.Admin
return nil
}
func (h *MessageHandler) adminUnmute(session types.Session, payload *message.Admin) error {
func (h *MessageHandlerCtx) adminUnmute(session types.Session, payload *message.Admin) error {
if !session.Admin() {
h.logger.Debug().Msg("user not admin")
return nil
@ -212,7 +212,7 @@ func (h *MessageHandler) adminUnmute(session types.Session, payload *message.Adm
return nil
}
func (h *MessageHandler) adminKick(session types.Session, payload *message.Admin) error {
func (h *MessageHandlerCtx) adminKick(session types.Session, payload *message.Admin) error {
if !session.Admin() {
h.logger.Debug().Msg("user not admin")
return nil
@ -246,7 +246,7 @@ func (h *MessageHandler) adminKick(session types.Session, payload *message.Admin
return nil
}
func (h *MessageHandler) adminBan(session types.Session, payload *message.Admin) error {
func (h *MessageHandlerCtx) adminBan(session types.Session, payload *message.Admin) error {
if !session.Admin() {
h.logger.Debug().Msg("user not admin")
return nil

16
internal/websocket/broadcast.go → internal/websocket/handler/broadcast.go
View File

@ -1,4 +1,4 @@
package websocket
package handler
import (
"demodesk/neko/internal/types"
@ -6,13 +6,13 @@ import (
"demodesk/neko/internal/types/message"
)
func (h *MessageHandler) boradcastCreate(session types.Session, payload *message.BroadcastCreate) error {
func (h *MessageHandlerCtx) boradcastCreate(session types.Session, payload *message.BroadcastCreate) error {
if !session.Admin() {
h.logger.Debug().Msg("user not admin")
return nil
}
h.broadcast.Create(payload.URL)
h.capture.StartBroadcast(payload.URL)
if err := h.boradcastStatus(session); err != nil {
return err
@ -21,13 +21,13 @@ func (h *MessageHandler) boradcastCreate(session types.Session, payload *message
return nil
}
func (h *MessageHandler) boradcastDestroy(session types.Session) error {
func (h *MessageHandlerCtx) boradcastDestroy(session types.Session) error {
if !session.Admin() {
h.logger.Debug().Msg("user not admin")
return nil
}
h.broadcast.Destroy()
h.capture.StopBroadcast()
if err := h.boradcastStatus(session); err != nil {
return err
@ -36,7 +36,7 @@ func (h *MessageHandler) boradcastDestroy(session types.Session) error {
return nil
}
func (h *MessageHandler) boradcastStatus(session types.Session) error {
func (h *MessageHandlerCtx) boradcastStatus(session types.Session) error {
if !session.Admin() {
h.logger.Debug().Msg("user not admin")
return nil
@ -45,8 +45,8 @@ func (h *MessageHandler) boradcastStatus(session types.Session) error {
if err := session.Send(
message.BroadcastStatus{
Event: event.BORADCAST_STATUS,
IsActive: h.broadcast.IsActive(),
URL: h.broadcast.GetUrl(),
IsActive: h.capture.IsBoradcasting(),
URL: h.capture.BroadcastUrl(),
}); err != nil {
h.logger.Warn().Err(err).Msgf("sending event %s has failed", event.BORADCAST_STATUS)
return err

18
internal/websocket/control.go → internal/websocket/handler/control.go
View File

@ -1,4 +1,4 @@
package websocket
package handler
import (
"demodesk/neko/internal/types"
@ -6,7 +6,7 @@ import (
"demodesk/neko/internal/types/message"
)
func (h *MessageHandler) controlRelease(session types.Session) error {
func (h *MessageHandlerCtx) controlRelease(session types.Session) error {
// check if session is host
if !session.IsHost() {
h.logger.Debug().Str("id", session.ID()).Msg("is not the host")
@ -30,7 +30,7 @@ func (h *MessageHandler) controlRelease(session types.Session) error {
return nil
}
func (h *MessageHandler) controlRequest(session types.Session) error {
func (h *MessageHandlerCtx) controlRequest(session types.Session) error {
host := h.sessions.GetHost()
if host == nil {
@ -69,7 +69,7 @@ func (h *MessageHandler) controlRequest(session types.Session) error {
return nil
}
func (h *MessageHandler) controlGive(session types.Session, payload *message.Control) error {
func (h *MessageHandlerCtx) controlGive(session types.Session, payload *message.Control) error {
// check if session is host
if !session.IsHost() {
h.logger.Debug().Str("id", session.ID()).Msg("is not the host")
@ -99,18 +99,18 @@ func (h *MessageHandler) controlGive(session types.Session, payload *message.Con
return nil
}
func (h *MessageHandler) controlClipboard(session types.Session, payload *message.Clipboard) error {
func (h *MessageHandlerCtx) controlClipboard(session types.Session, payload *message.Clipboard) error {
// check if session is host
if !session.IsHost() {
h.logger.Debug().Str("id", session.ID()).Msg("is not the host")
return nil
}
h.remote.WriteClipboard(payload.Text)
h.desktop.WriteClipboard(payload.Text)
return nil
}
func (h *MessageHandler) controlKeyboard(session types.Session, payload *message.Keyboard) error {
func (h *MessageHandlerCtx) controlKeyboard(session types.Session, payload *message.Keyboard) error {
// check if session is host
if !session.IsHost() {
h.logger.Debug().Str("id", session.ID()).Msg("is not the host")
@ -119,7 +119,7 @@ func (h *MessageHandler) controlKeyboard(session types.Session, payload *message
// change layout
if payload.Layout != nil {
h.remote.SetKeyboardLayout(*payload.Layout)
h.desktop.SetKeyboardLayout(*payload.Layout)
}
// set num lock
@ -152,6 +152,6 @@ func (h *MessageHandler) controlKeyboard(session types.Session, payload *message
Int("ScrollLock", ScrollLock).
Msg("setting keyboard modifiers")
h.remote.SetKeyboardModifiers(NumLock, CapsLock, ScrollLock)
h.desktop.SetKeyboardModifiers(NumLock, CapsLock, ScrollLock)
return nil
}

47
internal/websocket/handler.go → internal/websocket/handler/handler.go
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"
"demodesk/neko/internal/types"
"demodesk/neko/internal/types/event"
@ -12,40 +13,60 @@ import (
"demodesk/neko/internal/utils"
)
type MessageHandler struct {
func New(
sessions types.SessionManager,
desktop types.DesktopManager,
capture types.CaptureManager,
webrtc types.WebRTCManager,
) *MessageHandlerCtx {
logger := log.With().Str("module", "handler").Logger()
return &MessageHandlerCtx{
logger: logger,
sessions: sessions,
desktop: desktop,
capture: capture,
webrtc: webrtc,
banned: make(map[string]bool),
locked: false,
}
}
type MessageHandlerCtx struct {
logger zerolog.Logger
sessions types.SessionManager
webrtc types.WebRTCManager
remote types.RemoteManager
broadcast types.BroadcastManager
desktop types.DesktopManager
capture types.CaptureManager
banned map[string]bool
locked bool
}
func (h *MessageHandler) Connected(id string, socket *WebSocket) (bool, string, error) {
func (h *MessageHandlerCtx) Connected(id string, socket types.WebSocket) (bool, string) {
address := socket.Address()
if address == "" {
h.logger.Debug().Msg("no remote address")
} else {
if address != "" {
ok, banned := h.banned[address]
if ok && banned {
h.logger.Debug().Str("address", address).Msg("banned")
return false, "banned", nil
return false, "banned"
}
} else {
h.logger.Debug().Msg("no remote address")
}
if h.locked {
session, ok := h.sessions.Get(id)
if !ok || !session.Admin() {
h.logger.Debug().Msg("server locked")
return false, "locked", nil
return false, "locked"
}
}
return true, "", nil
return true, ""
}
func (h *MessageHandler) Disconnected(id string) error {
func (h *MessageHandlerCtx) Disconnected(id string) error {
// TODO: Refactor.
if h.locked && len(h.sessions.Admins()) == 0 {
h.locked = false
}
@ -53,7 +74,7 @@ func (h *MessageHandler) Disconnected(id string) error {
return h.sessions.Destroy(id)
}
func (h *MessageHandler) Message(id string, raw []byte) error {
func (h *MessageHandlerCtx) Message(id string, raw []byte) error {
header := message.Message{}
if err := json.Unmarshal(raw, &header); err != nil {
return err

25
internal/websocket/screen.go → internal/websocket/handler/screen.go
View File

@ -1,33 +1,38 @@
package websocket
package handler
import (
"demodesk/neko/internal/types"
"demodesk/neko/internal/types/event"
"demodesk/neko/internal/types/message"
"demodesk/neko/internal/websocket/broadcast"
)
func (h *MessageHandler) screenSet(session types.Session, payload *message.ScreenResolution) error {
func (h *MessageHandlerCtx) screenSet(session types.Session, payload *message.ScreenResolution) error {
if !session.Admin() {
h.logger.Debug().Msg("user not admin")
return nil
}
if err := h.remote.ChangeResolution(payload.Width, payload.Height, payload.Rate); err != nil {
if err := h.capture.ChangeResolution(payload.Width, payload.Height, payload.Rate); err != nil {
h.logger.Warn().Err(err).Msgf("unable to change screen size")
return err
}
if err := broadcast.ScreenConfiguration(h.sessions, session.ID(), payload.Width, payload.Height, payload.Rate); err != nil {
h.logger.Warn().Err(err).Msgf("broadcasting event %s has failed", event.SCREEN_RESOLUTION)
if err := h.sessions.Broadcast(message.ScreenResolution{
Event: event.SCREEN_RESOLUTION,
ID: session.ID(),
Width: payload.Width,
Height: payload.Height,
Rate: payload.Rate,
}, nil); err != nil {
h.logger.Warn().Err(err).Msgf("sending event %s has failed", event.SCREEN_RESOLUTION)
return err
}
return nil
}
func (h *MessageHandler) screenResolution(session types.Session) error {
if size := h.remote.GetScreenSize(); size != nil {
func (h *MessageHandlerCtx) screenResolution(session types.Session) error {
if size := h.desktop.GetScreenSize(); size != nil {
if err := session.Send(message.ScreenResolution{
Event: event.SCREEN_RESOLUTION,
Width: size.Width,
@ -42,7 +47,7 @@ func (h *MessageHandler) screenResolution(session types.Session) error {
return nil
}
func (h *MessageHandler) screenConfigurations(session types.Session) error {
func (h *MessageHandlerCtx) screenConfigurations(session types.Session) error {
if !session.Admin() {
h.logger.Debug().Msg("user not admin")
return nil
@ -50,7 +55,7 @@ func (h *MessageHandler) screenConfigurations(session types.Session) error {
if err := session.Send(message.ScreenConfigurations{
Event: event.SCREEN_CONFIGURATIONS,
Configurations: h.remote.ScreenConfigurations(),
Configurations: h.desktop.ScreenConfigurations(),
}); err != nil {
h.logger.Warn().Err(err).Msgf("sending event %s has failed", event.SCREEN_CONFIGURATIONS)
return err

28
internal/websocket/session.go → internal/websocket/handler/session.go
View File

@ -1,4 +1,4 @@
package websocket
package handler
import (
"demodesk/neko/internal/types"
@ -6,7 +6,7 @@ import (
"demodesk/neko/internal/types/message"
)
func (h *MessageHandler) SessionCreated(session types.Session) error {
func (h *MessageHandlerCtx) SessionCreated(session types.Session) error {
// send sdp and id over to client
if err := h.signalProvide(session); err != nil {
return err
@ -27,11 +27,22 @@ func (h *MessageHandler) SessionCreated(session types.Session) error {
return nil
}
func (h *MessageHandler) SessionConnected(session types.Session) error {
func (h *MessageHandlerCtx) SessionConnected(session types.Session) error {
// TODO: Refactor.
members := []*message.MembersListEntry{}
for _, session := range h.sessions.Members() {
members = append(members, &message.MembersListEntry{
ID: session.ID(),
Name: session.Name(),
Admin: session.Admin(),
Muted: session.Muted(),
})
}
// send list of members to session
if err := session.Send(message.MembersList{
Event: event.MEMBER_LIST,
Memebers: h.sessions.Members(),
Memebers: members,
}); err != nil {
h.logger.Warn().Str("id", session.ID()).Err(err).Msgf("sending event %s has failed", event.MEMBER_LIST)
return err
@ -58,7 +69,12 @@ func (h *MessageHandler) SessionConnected(session types.Session) error {
if err := h.sessions.Broadcast(
message.Member{
Event: event.MEMBER_CONNECTED,
Member: session.Member(),
Member: &message.MembersListEntry{
ID: session.ID(),
Name: session.Name(),
Admin: session.Admin(),
Muted: session.Muted(),
},
}, nil); err != nil {
h.logger.Warn().Err(err).Msgf("broadcasting event %s has failed", event.CONTROL_RELEASE)
return err
@ -67,7 +83,7 @@ func (h *MessageHandler) SessionConnected(session types.Session) error {
return nil
}
func (h *MessageHandler) SessionDestroyed(id string) error {
func (h *MessageHandlerCtx) SessionDestroyed(id string) error {
// clear host if exists
host := h.sessions.GetHost()
if host != nil && host.ID() == id {

8
internal/websocket/signal.go → internal/websocket/handler/signal.go
View File

@ -1,4 +1,4 @@
package websocket
package handler
import (
"demodesk/neko/internal/types"
@ -6,8 +6,8 @@ import (
"demodesk/neko/internal/types/message"
)
func (h *MessageHandler) signalProvide(session types.Session) error {
sdp, lite, ice, err := h.webrtc.CreatePeer(session.ID(), session)
func (h *MessageHandlerCtx) signalProvide(session types.Session) error {
sdp, lite, ice, err := h.webrtc.CreatePeer(session)
if err != nil {
return err
}
@ -25,7 +25,7 @@ func (h *MessageHandler) signalProvide(session types.Session) error {
return nil
}
func (h *MessageHandler) signalAnswer(session types.Session, payload *message.SignalAnswer) error {
func (h *MessageHandlerCtx) signalAnswer(session types.Session, payload *message.SignalAnswer) error {
session.SetName(payload.DisplayName)
if err := session.SignalAnswer(payload.SDP); err != nil {

270
internal/websocket/manager.go Normal file
View File

@ -0,0 +1,270 @@
package websocket
import (
"fmt"
"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"
"demodesk/neko/internal/config"
"demodesk/neko/internal/utils"
)
func New(
sessions types.SessionManager,
desktop types.DesktopManager,
capture types.CaptureManager,
webrtc types.WebRTCManager,
conf *config.WebSocket,
) *WebSocketManagerCtx {
logger := log.With().Str("module", "websocket").Logger()
return &WebSocketManagerCtx{
logger: logger,
conf: conf,
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
conf *config.WebSocket
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 {
ws.logger.Warn().Str("id", session.ID()).Err(err).Msg("session created with and error")
} 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 {
ws.logger.Warn().Str("id", session.ID()).Err(err).Msg("session connected with and error")
} else {
ws.logger.Debug().Str("id", session.ID()).Msg("session connected")
}
})
ws.sessions.OnDestroy(func(id string) {
if err := ws.handler.SessionDestroyed(id); err != nil {
ws.logger.Warn().Str("id", id).Err(err).Msg("session destroyed with and error")
} else {
ws.logger.Debug().Str("id", id).Msg("session destroyed")
}
})
go func() {
defer func() {
ws.logger.Info().Msg("shutdown")
}()
current := ws.desktop.ReadClipboard()
for {
select {
case <-ws.shutdown:
return
default:
session := ws.sessions.GetHost()
if session != nil {
break
}
text := ws.desktop.ReadClipboard()
if text == current {
break
}
// TODO: Refactor
if err := session.Send(message.Clipboard{
Event: event.CONTROL_CLIPBOARD,
Text: text,
}); err != nil {
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
}
id, ip, admin, err := ws.authenticate(r)
if err != nil {
ws.logger.Warn().Err(err).Msg("authentication failed")
// TODO: Refactor
if err = connection.WriteJSON(message.Disconnect{
Event: event.SYSTEM_DISCONNECT,
Message: "invalid_password",
}); err != nil {
ws.logger.Error().Err(err).Msg("failed to send disconnect")
}
return connection.Close()
}
socket := &WebSocketCtx{
id: id,
ws: ws,
address: ip,
connection: connection,
}
ok, reason := ws.handler.Connected(id, socket)
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()
}
ws.sessions.New(id, admin, socket)
ws.logger.
Debug().
Str("session", id).
Str("address", connection.RemoteAddr().String()).
Msg("new connection created")
defer func() {
ws.logger.
Debug().
Str("session", id).
Str("address", connection.RemoteAddr().String()).
Msg("session ended")
}()
ws.handle(connection, id)
return nil
}
// TODO: Refactor
func (ws *WebSocketManagerCtx) authenticate(r *http.Request) (string, string, bool, error) {
ip := r.RemoteAddr
if ws.conf.Proxy {
ip = utils.ReadUserIP(r)
}
id, err := utils.NewUID(32)
if err != nil {
return "", ip, false, err
}
passwords, ok := r.URL.Query()["password"]
if !ok || len(passwords[0]) < 1 {
return "", ip, false, fmt.Errorf("no password provided")
}
if passwords[0] == ws.conf.AdminPassword {
return id, ip, true, nil
}
if passwords[0] == ws.conf.Password {
return id, ip, false, nil
}
return "", ip, false, fmt.Errorf("invalid password: %s", passwords[0])
}
func (ws *WebSocketManagerCtx) handle(connection *websocket.Conn, id string) {
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")
if err := ws.handler.Disconnected(id); err != nil {
ws.logger.Warn().Err(err).Msg("socket disconnected with error")
}
}()
for {
_, raw, err := connection.ReadMessage()
if err == nil {
bytes <- raw
continue
}
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)
}
}()
for {
select {
case raw := <-bytes:
ws.logger.Debug().
Str("session", id).
Str("address", connection.RemoteAddr().String()).
Str("raw", string(raw)).
Msg("received message from client")
if err := ws.handler.Message(id, raw); err != nil {
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
}
}
}
}

55
internal/websocket/socket.go
View File

@ -1,55 +0,0 @@
package websocket
import (
"encoding/json"
"strings"
"sync"
"github.com/gorilla/websocket"
)
type WebSocket struct {
id string
address string
ws *WebSocketHandler
connection *websocket.Conn
mu sync.Mutex
}
func (socket *WebSocket) Address() string {
//remote := socket.connection.RemoteAddr()
address := strings.SplitN(socket.address, ":", -1)
if len(address[0]) < 1 {
return socket.address
}
return address[0]
}
func (socket *WebSocket) Send(v interface{}) error {
socket.mu.Lock()
defer socket.mu.Unlock()
if socket.connection == nil {
return nil
}
raw, err := json.Marshal(v)
if err != nil {
return err
}
socket.ws.logger.Debug().
Str("session", socket.id).
Str("address", socket.connection.RemoteAddr().String()).
Str("raw", string(raw)).
Msg("sending message to client")
return socket.connection.WriteMessage(websocket.TextMessage, raw)
}
func (socket *WebSocket) Destroy() error {
if socket.connection == nil {
return nil
}
return socket.connection.Close()
}

275
internal/websocket/websocket.go
View File

@ -1,270 +1,57 @@
package websocket
import (
"fmt"
"net/http"
"time"
"encoding/json"
"strings"
"sync"
"github.com/gorilla/websocket"
"github.com/rs/zerolog"
"github.com/rs/zerolog/log"
"demodesk/neko/internal/types"
"demodesk/neko/internal/types/config"
"demodesk/neko/internal/types/event"
"demodesk/neko/internal/types/message"
"demodesk/neko/internal/utils"
)
func New(sessions types.SessionManager, remote types.RemoteManager, broadcast types.BroadcastManager, webrtc types.WebRTCManager, conf *config.WebSocket) *WebSocketHandler {
logger := log.With().Str("module", "websocket").Logger()
return &WebSocketHandler{
logger: logger,
conf: conf,
sessions: sessions,
remote: remote,
upgrader: websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool {
return true
},
},
handler: &MessageHandler{
logger: logger.With().Str("subsystem", "handler").Logger(),
remote: remote,
broadcast: broadcast,
sessions: sessions,
webrtc: webrtc,
banned: make(map[string]bool),
locked: false,
},
}
type WebSocketCtx struct {
id string
address string
ws *WebSocketManagerCtx
connection *websocket.Conn
mu sync.Mutex
}
// Send pings to peer with this period. Must be less than pongWait.
const pingPeriod = 60 * time.Second
type WebSocketHandler struct {
logger zerolog.Logger
upgrader websocket.Upgrader
sessions types.SessionManager
remote types.RemoteManager
conf *config.WebSocket
handler *MessageHandler
shutdown chan bool
}
func (ws *WebSocketHandler) Start() {
ws.sessions.OnCreated(func(session types.Session) {
if err := ws.handler.SessionCreated(session); err != nil {
ws.logger.Warn().Str("id", session.ID()).Err(err).Msg("session created with and error")
} 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 {
ws.logger.Warn().Str("id", session.ID()).Err(err).Msg("session connected with and error")
} else {
ws.logger.Debug().Str("id", session.ID()).Msg("session connected")
}
})
ws.sessions.OnDestroy(func(id string) {
if err := ws.handler.SessionDestroyed(id); err != nil {
ws.logger.Warn().Str("id", id).Err(err).Msg("session destroyed with and error")
} else {
ws.logger.Debug().Str("id", id).Msg("session destroyed")
}
})
go func() {
defer func() {
ws.logger.Info().Msg("shutdown")
}()
current := ws.remote.ReadClipboard()
for {
select {
case <-ws.shutdown:
return
default:
if ws.sessions.HasHost() {
text := ws.remote.ReadClipboard()
if text != current {
session := ws.sessions.GetHost()
if session != nil {
if err := session.Send(message.Clipboard{
Event: event.CONTROL_CLIPBOARD,
Text: text,
}); err != nil {
ws.logger.Warn().Err(err).Msg("could not sync clipboard")
}
}
current = text
}
}
time.Sleep(100 * time.Millisecond)
}
}
}()
}
func (ws *WebSocketHandler) Shutdown() error {
ws.shutdown <- true
return nil
}
func (ws *WebSocketHandler) 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
func (socket *WebSocketCtx) Address() string {
//remote := socket.connection.RemoteAddr()
address := strings.SplitN(socket.address, ":", -1)
if len(address[0]) < 1 {
return socket.address
}
id, ip, admin, err := ws.authenticate(r)
if err != nil {
ws.logger.Warn().Err(err).Msg("authentication failed")
return address[0]
}
if err = connection.WriteJSON(message.Disconnect{
Event: event.SYSTEM_DISCONNECT,
Message: "invalid_password",
}); err != nil {
ws.logger.Error().Err(err).Msg("failed to send disconnect")
}
func (socket *WebSocketCtx) Send(v interface{}) error {
socket.mu.Lock()
defer socket.mu.Unlock()
if err = connection.Close(); err != nil {
return err
}
if socket.connection == nil {
return nil
}
socket := &WebSocket{
id: id,
ws: ws,
address: ip,
connection: connection,
}
ok, reason, err := ws.handler.Connected(id, socket)
raw, err := json.Marshal(v)
if err != nil {
ws.logger.Error().Err(err).Msg("connection failed")
return err
}
if !ok {
if err = connection.WriteJSON(message.Disconnect{
Event: event.SYSTEM_DISCONNECT,
Message: reason,
}); err != nil {
ws.logger.Error().Err(err).Msg("failed to send disconnect")
}
socket.ws.logger.Debug().
Str("session", socket.id).
Str("address", socket.connection.RemoteAddr().String()).
Str("raw", string(raw)).
Msg("sending message to client")
if err = connection.Close(); err != nil {
return err
}
return socket.connection.WriteMessage(websocket.TextMessage, raw)
}
func (socket *WebSocketCtx) Destroy() error {
if socket.connection == nil {
return nil
}
ws.sessions.New(id, admin, socket)
ws.logger.
Debug().
Str("session", id).
Str("address", connection.RemoteAddr().String()).
Msg("new connection created")
defer func() {
ws.logger.
Debug().
Str("session", id).
Str("address", connection.RemoteAddr().String()).
Msg("session ended")
}()
ws.handle(connection, id)
return nil
}
func (ws *WebSocketHandler) authenticate(r *http.Request) (string, string, bool, error) {
ip := r.RemoteAddr
if ws.conf.Proxy {
ip = utils.ReadUserIP(r)
}
id, err := utils.NewUID(32)
if err != nil {
return "", ip, false, err
}
passwords, ok := r.URL.Query()["password"]
if !ok || len(passwords[0]) < 1 {
return "", ip, false, fmt.Errorf("no password provided")
}
if passwords[0] == ws.conf.AdminPassword {
return id, ip, true, nil
}
if passwords[0] == ws.conf.Password {
return id, ip, false, nil
}
return "", ip, false, fmt.Errorf("invalid password: %s", passwords[0])
}
func (ws *WebSocketHandler) handle(connection *websocket.Conn, id string) {
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")
if err := ws.handler.Disconnected(id); err != nil {
ws.logger.Warn().Err(err).Msg("socket disconnected with error")
}
}()
for {
_, raw, err := connection.ReadMessage()
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
}
bytes <- raw
}
}()
for {
select {
case raw := <-bytes:
ws.logger.Debug().
Str("session", id).
Str("address", connection.RemoteAddr().String()).
Str("raw", string(raw)).
Msg("received message from client")
if err := ws.handler.Message(id, raw); err != nil {
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 {
return
}
}
}
return socket.connection.Close()
}