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

295 lines
6.7 KiB
Go
Raw Normal View History

2020-01-25 04:47:37 +13:00
package websocket
import (
"fmt"
"net/http"
2021-03-20 09:33:49 +13:00
"sync/atomic"
2020-01-25 04:47:37 +13:00
"time"
"github.com/gorilla/websocket"
"github.com/rs/zerolog"
"github.com/rs/zerolog/log"
"n.eko.moe/neko/internal/types"
"n.eko.moe/neko/internal/types/config"
2020-01-25 04:47:37 +13:00
"n.eko.moe/neko/internal/types/event"
"n.eko.moe/neko/internal/types/message"
"n.eko.moe/neko/internal/utils"
)
2020-09-27 11:10:34 +13:00
func New(sessions types.SessionManager, remote types.RemoteManager, broadcast types.BroadcastManager, webrtc types.WebRTCManager, conf *config.WebSocket) *WebSocketHandler {
2020-01-25 04:47:37 +13:00
logger := log.With().Str("module", "websocket").Logger()
return &WebSocketHandler{
2020-09-27 11:10:34 +13:00
logger: logger,
conf: conf,
sessions: sessions,
remote: remote,
upgrader: websocket.Upgrader{
2020-01-25 04:47:37 +13:00
CheckOrigin: func(r *http.Request) bool {
return true
},
},
handler: &MessageHandler{
2020-09-27 11:10:34 +13:00
logger: logger.With().Str("subsystem", "handler").Logger(),
remote: remote,
broadcast: broadcast,
sessions: sessions,
webrtc: webrtc,
banned: make(map[string]bool),
locked: false,
2020-01-25 04:47:37 +13:00
},
2021-03-20 09:33:49 +13:00
conns: 0,
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 {
2020-09-27 11:10:34 +13:00
logger zerolog.Logger
upgrader websocket.Upgrader
sessions types.SessionManager
remote types.RemoteManager
conf *config.WebSocket
handler *MessageHandler
2021-03-20 09:33:49 +13:00
conns uint32
2020-09-27 11:10:34 +13:00
shutdown chan bool
2020-01-25 04:47:37 +13:00
}
func (ws *WebSocketHandler) Start() error {
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")
}
})
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")
}
})
2020-01-26 03:29:52 +13:00
go func() {
defer func() {
ws.logger.Info().Msg("shutdown")
}()
2020-04-06 10:34:51 +12:00
current := ws.remote.ReadClipboard()
2020-01-26 03:29:52 +13:00
for {
select {
case <-ws.shutdown:
return
default:
if ws.sessions.HasHost() {
2020-04-06 10:34:51 +12:00
text := ws.remote.ReadClipboard()
2020-02-03 11:48:23 +13:00
if text != current {
2020-01-26 03:29:52 +13:00
session, ok := ws.sessions.GetHost()
if ok {
session.Send(message.Clipboard{
Event: event.CONTROL_CLIPBOARD,
2020-02-03 11:48:23 +13:00
Text: text,
2020-01-26 03:29:52 +13:00
})
}
2020-02-03 11:48:23 +13:00
current = text
2020-01-26 03:29:52 +13:00
}
}
time.Sleep(100 * time.Millisecond)
}
}
}()
2020-01-25 04:47:37 +13:00
return nil
}
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
}
2020-04-05 15:49:43 +12:00
id, ip, 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
if err = connection.WriteJSON(message.Disconnect{
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,
2020-04-05 15:49:43 +12:00
address: ip,
2020-01-25 04:47:37 +13:00
connection: connection,
}
2021-01-15 09:41:00 +13:00
ok, reason, err := ws.handler.Connected(admin, socket)
2020-01-25 04:47:37 +13:00
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")
}
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-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")
}
2020-04-05 15:49:43 +12:00
func (ws *WebSocketHandler) authenticate(r *http.Request) (string, string, bool, error) {
2020-04-07 08:14:30 +12:00
ip := r.RemoteAddr
if ws.conf.Proxy {
ip = utils.ReadUserIP(r)
}
2020-04-05 15:49:43 +12:00
2020-01-25 04:47:37 +13:00
id, err := utils.NewUID(32)
if err != nil {
2020-04-05 15:49:43 +12:00
return "", ip, false, err
2020-01-25 04:47:37 +13:00
}
passwords, ok := r.URL.Query()["password"]
if !ok || len(passwords[0]) < 1 {
2020-04-05 15:49:43 +12:00
return "", ip, false, fmt.Errorf("no password provided")
2020-01-25 04:47:37 +13:00
}
2021-03-20 09:33:49 +13:00
isAdmin, err := ws.IsAdmin(passwords[0])
return id, ip, isAdmin, err
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)
go func() {
defer func() {
ticker.Stop()
ws.logger.Debug().Str("address", connection.RemoteAddr().String()).Msg("handle socket ending")
ws.handler.Disconnected(id)
}()
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")
}
case <-cancel:
return
case _ = <-ticker.C:
if err := connection.WriteMessage(websocket.PingMessage, nil); err != nil {
return
}
}
}
}