Archived
2
0
This repository has been archived on 2024-06-24. You can view files and clone it, but cannot push or open issues or pull requests.
neko-custom/server/internal/websocket/websocket.go

453 lines
11 KiB
Go
Raw Normal View History

2020-01-25 04:47:37 +13:00
package websocket
import (
"fmt"
"net/http"
2022-10-31 14:06:05 +13:00
"os"
2021-09-12 04:21:28 +12:00
"sync"
2021-03-20 09:33:49 +13:00
"sync/atomic"
2020-01-25 04:47:37 +13:00
"time"
2022-11-05 15:43:18 +13:00
"github.com/fsnotify/fsnotify"
2020-01-25 04:47:37 +13:00
"github.com/gorilla/websocket"
"github.com/rs/zerolog"
"github.com/rs/zerolog/log"
2022-09-13 08:18:18 +12:00
"m1k1o/neko/internal/config"
2021-10-06 09:38:24 +13:00
"m1k1o/neko/internal/types"
"m1k1o/neko/internal/types/event"
"m1k1o/neko/internal/types/message"
"m1k1o/neko/internal/utils"
2022-09-13 08:36:56 +12:00
"m1k1o/neko/internal/websocket/handler"
2022-09-14 06:04:43 +12:00
"m1k1o/neko/internal/websocket/state"
2020-01-25 04:47:37 +13:00
)
2021-11-18 10:10:55 +13:00
const CONTROL_PROTECTION_SESSION = "by_control_protection"
2021-11-17 11:48:40 +13:00
2022-09-17 22:43:17 +12:00
func New(sessions types.SessionManager, desktop types.DesktopManager, capture types.CaptureManager, webrtc types.WebRTCManager, conf *config.WebSocket) *WebSocketHandler {
2020-01-25 04:47:37 +13:00
logger := log.With().Str("module", "websocket").Logger()
2022-11-20 08:26:45 +13:00
state := state.New(conf.FileTransferEnabled, conf.FileTransferPath)
2021-11-18 10:10:55 +13:00
// if control protection is enabled
if conf.ControlProtection {
2022-09-14 06:04:43 +12:00
state.Lock("control", CONTROL_PROTECTION_SESSION)
2021-11-18 10:10:55 +13:00
logger.Info().Msgf("control locked on behalf of control protection")
}
2022-11-20 08:26:45 +13:00
// create file transfer directory if not exists
if conf.FileTransferEnabled {
if _, err := os.Stat(conf.FileTransferPath); os.IsNotExist(err) {
err = os.Mkdir(conf.FileTransferPath, os.ModePerm)
logger.Err(err).Msg("creating file transfer directory")
}
2022-10-31 14:06:05 +13:00
}
2021-11-18 10:10:55 +13:00
// apply default locks
2021-11-17 11:00:24 +13:00
for _, lock := range conf.Locks {
2022-09-14 06:04:43 +12:00
state.Lock(lock, "") // empty session ID
2021-11-17 11:00:24 +13:00
}
if len(conf.Locks) > 0 {
logger.Info().Msgf("locked resources: %+v", conf.Locks)
}
2022-09-13 08:36:56 +12:00
handler := handler.New(
sessions,
desktop,
capture,
webrtc,
2022-09-14 06:04:43 +12:00
state,
2022-09-13 08:36:56 +12:00
)
2020-01-25 04:47:37 +13:00
return &WebSocketHandler{
2021-04-05 08:37:33 +12:00
logger: logger,
2021-09-12 04:21:28 +12:00
shutdown: make(chan interface{}),
2021-04-05 08:37:33 +12:00
conf: conf,
sessions: sessions,
2022-09-13 08:12:47 +12:00
desktop: desktop,
2022-09-13 08:36:56 +12:00
webrtc: webrtc,
2022-09-14 06:04:43 +12:00
state: state,
2021-04-05 08:37:33 +12:00
upgrader: websocket.Upgrader{
2020-01-25 04:47:37 +13:00
CheckOrigin: func(r *http.Request) bool {
return true
},
},
2022-09-13 08:36:56 +12:00
handler: handler,
2021-11-18 10:10:55 +13:00
serverStartedAt: time.Now(),
2020-01-25 04:47:37 +13:00
}
}
// Send pings to peer with this period. Must be less than pongWait.
const pingPeriod = 60 * time.Second
type WebSocketHandler struct {
2021-04-05 08:37:33 +12:00
logger zerolog.Logger
2021-09-12 04:21:28 +12:00
wg sync.WaitGroup
shutdown chan interface{}
2021-04-05 08:37:33 +12:00
upgrader websocket.Upgrader
sessions types.SessionManager
2022-09-13 08:12:47 +12:00
desktop types.DesktopManager
2022-09-13 08:36:56 +12:00
webrtc types.WebRTCManager
2022-09-14 06:04:43 +12:00
state *state.State
2021-04-05 08:37:33 +12:00
conf *config.WebSocket
2022-09-13 08:36:56 +12:00
handler *handler.MessageHandler
2021-11-18 10:10:55 +13:00
// stats
conns uint32
serverStartedAt time.Time
lastAdminLeftAt *time.Time
lastUserLeftAt *time.Time
2020-01-25 04:47:37 +13:00
}
2021-10-06 10:10:10 +13:00
func (ws *WebSocketHandler) Start() {
2020-01-25 04:47:37 +13:00
ws.sessions.OnCreated(func(id string, session types.Session) {
if err := ws.handler.SessionCreated(id, session); err != nil {
ws.logger.Warn().Str("id", id).Err(err).Msg("session created with and error")
} else {
ws.logger.Debug().Str("id", id).Msg("session created")
}
})
ws.sessions.OnConnected(func(id string, session types.Session) {
if err := ws.handler.SessionConnected(id, session); err != nil {
ws.logger.Warn().Str("id", id).Err(err).Msg("session connected with and error")
} else {
ws.logger.Debug().Str("id", id).Msg("session connected")
}
2021-11-17 11:48:40 +13:00
// if control protection is enabled and at least one admin
// and if room was locked on behalf control protection, unlock
2022-09-14 06:04:43 +12:00
sess, ok := ws.state.GetLocked("control")
2021-11-17 11:48:40 +13:00
if ok && ws.conf.ControlProtection && sess == CONTROL_PROTECTION_SESSION && len(ws.sessions.Admins()) > 0 {
2022-09-14 06:04:43 +12:00
ws.state.Unlock("control")
ws.sessions.SetControlLocked(false) // TODO: Handle locks in sessions as flags.
2021-11-17 11:48:40 +13:00
ws.logger.Info().Msgf("control unlocked on behalf of control protection")
if err := ws.sessions.Broadcast(
message.AdminLock{
Event: event.ADMIN_UNLOCK,
ID: id,
Resource: "control",
}, nil); err != nil {
ws.logger.Warn().Err(err).Msgf("broadcasting event %s has failed", event.ADMIN_UNLOCK)
}
}
2021-11-18 10:10:55 +13:00
// remove outdated stats
if session.Admin() {
ws.lastAdminLeftAt = nil
} else {
ws.lastUserLeftAt = nil
}
2020-01-25 04:47:37 +13:00
})
2020-04-07 08:14:30 +12:00
ws.sessions.OnDestroy(func(id string, session types.Session) {
2020-01-25 04:47:37 +13:00
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")
}
2021-11-17 11:48:40 +13:00
2021-11-18 10:10:55 +13:00
membersCount := len(ws.sessions.Members())
adminCount := len(ws.sessions.Admins())
2021-11-17 11:48:40 +13:00
// if control protection is enabled and no admin
2021-11-18 10:10:55 +13:00
// and room is not locked, lock
2022-09-14 06:04:43 +12:00
ok := ws.state.IsLocked("control")
2021-11-18 10:10:55 +13:00
if !ok && ws.conf.ControlProtection && adminCount == 0 {
2022-09-14 06:04:43 +12:00
ws.state.Lock("control", CONTROL_PROTECTION_SESSION)
ws.sessions.SetControlLocked(true) // TODO: Handle locks in sessions as flags.
2021-11-17 11:48:40 +13:00
ws.logger.Info().Msgf("control locked and released on behalf of control protection")
2022-09-13 08:36:56 +12:00
ws.handler.AdminRelease(id, session)
2021-11-17 11:48:40 +13:00
if err := ws.sessions.Broadcast(
message.AdminLock{
Event: event.ADMIN_LOCK,
ID: id,
Resource: "control",
}, nil); err != nil {
ws.logger.Warn().Err(err).Msgf("broadcasting event %s has failed", event.ADMIN_LOCK)
}
}
2021-11-18 10:10:55 +13:00
// if this was the last admin
if session.Admin() && adminCount == 0 {
now := time.Now()
ws.lastAdminLeftAt = &now
}
// if this was the last user
if !session.Admin() && membersCount-adminCount == 0 {
now := time.Now()
ws.lastUserLeftAt = &now
}
2020-01-25 04:47:37 +13:00
})
2022-09-16 10:01:15 +12:00
ws.desktop.OnClipboardUpdated(func() {
session, ok := ws.sessions.GetHost()
if !ok {
return
2020-01-26 03:29:52 +13:00
}
2022-09-16 10:01:15 +12:00
err := session.Send(message.Clipboard{
Event: event.CONTROL_CLIPBOARD,
Text: ws.desktop.ReadClipboard(),
})
ws.logger.Err(err).Msg("sync clipboard")
})
2022-11-05 15:43:18 +13:00
2022-11-20 08:26:45 +13:00
// watch for file changes and send file list if file transfer is enabled
if ws.conf.FileTransferEnabled {
watcher, err := fsnotify.NewWatcher()
if err != nil {
ws.logger.Err(err).Msg("unable to start file transfer dir watcher")
return
}
2022-11-05 15:43:18 +13:00
2022-11-20 08:26:45 +13:00
go func() {
for {
select {
case e, ok := <-watcher.Events:
if !ok {
ws.logger.Info().Msg("file transfer dir watcher closed")
return
}
if e.Has(fsnotify.Create) || e.Has(fsnotify.Remove) || e.Has(fsnotify.Rename) {
ws.logger.Debug().Str("event", e.String()).Msg("file transfer dir watcher event")
ws.handler.FileTransferRefresh(nil)
}
case err := <-watcher.Errors:
ws.logger.Err(err).Msg("error in file transfer dir watcher")
2022-11-20 06:29:21 +13:00
}
2022-11-05 15:43:18 +13:00
}
2022-11-20 08:26:45 +13:00
}()
2022-11-20 08:26:45 +13:00
if err := watcher.Add(ws.conf.FileTransferPath); err != nil {
ws.logger.Err(err).Msg("unable to add file transfer path to watcher")
}
2022-11-05 15:43:18 +13:00
}
2020-01-25 04:47:37 +13:00
}
func (ws *WebSocketHandler) Shutdown() error {
2021-09-12 04:21:28 +12:00
close(ws.shutdown)
ws.wg.Wait()
2020-01-25 04:47:37 +13:00
return nil
}
func (ws *WebSocketHandler) Upgrade(w http.ResponseWriter, r *http.Request) error {
ws.logger.Debug().Msg("attempting to upgrade connection")
2021-11-18 06:00:27 +13:00
id, err := utils.NewUID(32)
if err != nil {
ws.logger.Error().Err(err).Msg("failed to generate user id")
return err
}
2020-01-25 04:47:37 +13:00
connection, err := ws.upgrader.Upgrade(w, r, nil)
if err != nil {
ws.logger.Error().Err(err).Msg("failed to upgrade connection")
return err
}
2021-11-18 06:00:27 +13:00
admin, err := ws.authenticate(r)
2020-01-25 04:47:37 +13:00
if err != nil {
ws.logger.Warn().Err(err).Msg("authentication failed")
2020-01-25 04:47:37 +13:00
2021-08-16 02:05:26 +12:00
if err = connection.WriteJSON(message.SystemMessage{
2020-01-25 04:47:37 +13:00
Event: event.SYSTEM_DISCONNECT,
Message: "invalid_password",
2020-01-25 04:47:37 +13:00
}); err != nil {
ws.logger.Error().Err(err).Msg("failed to send disconnect")
}
if err = connection.Close(); err != nil {
return err
}
return nil
}
socket := &WebSocket{
id: id,
ws: ws,
2021-11-18 06:00:27 +13:00
address: utils.GetHttpRequestIP(r, ws.conf.Proxy),
2020-01-25 04:47:37 +13:00
connection: connection,
}
2022-09-13 08:36:56 +12:00
ok, reason := ws.handler.Connected(admin, socket.Address())
2020-01-25 04:47:37 +13:00
if !ok {
2021-08-16 02:05:26 +12:00
if err = connection.WriteJSON(message.SystemMessage{
2020-01-25 04:47:37 +13:00
Event: event.SYSTEM_DISCONNECT,
Message: reason,
}); err != nil {
ws.logger.Error().Err(err).Msg("failed to send disconnect")
}
if err = connection.Close(); err != nil {
return err
}
return nil
}
ws.sessions.New(id, admin, socket)
ws.logger.
Debug().
Str("session", id).
Str("address", connection.RemoteAddr().String()).
Msg("new connection created")
2021-03-20 09:33:49 +13:00
atomic.AddUint32(&ws.conns, uint32(1))
2020-01-25 04:47:37 +13:00
defer func() {
ws.logger.
Debug().
Str("session", id).
Str("address", connection.RemoteAddr().String()).
Msg("session ended")
2021-03-20 09:33:49 +13:00
atomic.AddUint32(&ws.conns, ^uint32(0))
2020-01-25 04:47:37 +13:00
}()
ws.handle(connection, id)
return nil
}
2021-03-20 10:06:40 +13:00
func (ws *WebSocketHandler) Stats() types.Stats {
host := ""
session, ok := ws.sessions.GetHost()
if ok {
host = session.ID()
}
return types.Stats{
Connections: atomic.LoadUint32(&ws.conns),
Host: host,
Members: ws.sessions.Members(),
2021-11-18 10:10:55 +13:00
2022-09-14 06:04:43 +12:00
Banned: ws.state.AllBanned(),
Locked: ws.state.AllLocked(),
2021-11-18 10:10:55 +13:00
ServerStartedAt: ws.serverStartedAt,
LastAdminLeftAt: ws.lastAdminLeftAt,
LastUserLeftAt: ws.lastUserLeftAt,
ControlProtection: ws.conf.ControlProtection,
2022-09-13 08:36:56 +12:00
ImplicitControl: ws.webrtc.ImplicitControl(),
2021-03-20 10:06:40 +13:00
}
2021-03-20 09:33:49 +13:00
}
2022-09-18 04:17:04 +12:00
func (ws *WebSocketHandler) IsLocked(resource string) bool {
return ws.state.IsLocked(resource)
}
2021-03-20 09:33:49 +13:00
func (ws *WebSocketHandler) IsAdmin(password string) (bool, error) {
if password == ws.conf.AdminPassword {
return true, nil
}
if password == ws.conf.Password {
return false, nil
}
return false, fmt.Errorf("invalid password")
}
2021-11-18 06:00:27 +13:00
func (ws *WebSocketHandler) authenticate(r *http.Request) (bool, error) {
2020-01-25 04:47:37 +13:00
passwords, ok := r.URL.Query()["password"]
if !ok || len(passwords[0]) < 1 {
2021-11-18 06:00:27 +13:00
return false, fmt.Errorf("no password provided")
2020-01-25 04:47:37 +13:00
}
2021-11-18 06:00:27 +13:00
return ws.IsAdmin(passwords[0])
2020-01-25 04:47:37 +13:00
}
func (ws *WebSocketHandler) handle(connection *websocket.Conn, id string) {
bytes := make(chan []byte)
cancel := make(chan struct{})
ticker := time.NewTicker(pingPeriod)
2021-09-12 04:21:28 +12:00
ws.wg.Add(1)
2020-01-25 04:47:37 +13:00
go func() {
defer func() {
ticker.Stop()
ws.logger.Debug().Str("address", connection.RemoteAddr().String()).Msg("handle socket ending")
ws.handler.Disconnected(id)
2021-09-12 04:21:28 +12:00
ws.wg.Done()
2020-01-25 04:47:37 +13:00
}()
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).
2020-02-13 17:45:29 +13:00
Str("address", connection.RemoteAddr().String()).
2020-01-25 04:47:37 +13:00
Str("raw", string(raw)).
Msg("received message from client")
2020-01-25 04:47:37 +13:00
if err := ws.handler.Message(id, raw); err != nil {
ws.logger.Error().Err(err).Msg("message handler has failed")
}
2021-09-12 04:21:28 +12:00
case <-ws.shutdown:
if err := connection.WriteJSON(message.SystemMessage{
Event: event.SYSTEM_DISCONNECT,
Message: "server_shutdown",
}); err != nil {
ws.logger.Err(err).Msg("failed to send disconnect")
}
if err := connection.Close(); err != nil {
ws.logger.Err(err).Msg("connection closed with an error")
}
return
2020-01-25 04:47:37 +13:00
case <-cancel:
return
2021-07-23 06:58:15 +12:00
case <-ticker.C:
2020-01-25 04:47:37 +13:00
if err := connection.WriteMessage(websocket.PingMessage, nil); err != nil {
return
}
}
}
}
2022-11-20 08:26:45 +13:00
//
// File transfer
//
func (ws *WebSocketHandler) CanTransferFiles(password string) (bool, error) {
if !ws.conf.FileTransferEnabled {
return false, nil
}
isAdmin, err := ws.IsAdmin(password)
if err != nil {
return false, err
}
return isAdmin || !ws.state.IsLocked("file_transfer"), nil
}
func (ws *WebSocketHandler) FileTransferPath(filename string) string {
return ws.state.FileTransferPath(filename)
}
func (ws *WebSocketHandler) FileTransferEnabled() bool {
return ws.conf.FileTransferEnabled
}