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

268 lines
6.1 KiB
Go
Raw Normal View History

2020-01-25 04:47:37 +13:00
package websocket
import (
"fmt"
"net/http"
"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-04-06 10:34:51 +12:00
func New(sessions types.SessionManager, remote types.RemoteManager, webrtc types.WebRTCManager, conf *config.WebSocket) *WebSocketHandler {
2020-01-25 04:47:37 +13:00
logger := log.With().Str("module", "websocket").Logger()
return &WebSocketHandler{
logger: logger,
conf: conf,
sessions: sessions,
2020-04-06 10:34:51 +12:00
remote: remote,
2020-01-25 04:47:37 +13:00
upgrader: websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool {
return true
},
},
handler: &MessageHandler{
logger: logger.With().Str("subsystem", "handler").Logger(),
2020-04-06 10:34:51 +12:00
remote: remote,
2020-01-25 04:47:37 +13:00
sessions: sessions,
webrtc: webrtc,
banned: make(map[string]bool),
locked: false,
},
}
}
// 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
2020-04-06 10:34:51 +12:00
remote types.RemoteManager
2020-01-25 04:47:37 +13:00
conf *config.WebSocket
handler *MessageHandler
shutdown chan bool
}
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,
}
ok, reason, err := ws.handler.Connected(id, socket)
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")
defer func() {
ws.logger.
Debug().
Str("session", id).
Str("address", connection.RemoteAddr().String()).
Msg("session ended")
}()
ws.handle(connection, id)
return nil
}
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
}
if passwords[0] == ws.conf.AdminPassword {
2020-04-05 15:49:43 +12:00
return id, ip, true, nil
2020-01-25 04:47:37 +13:00
}
if passwords[0] == ws.conf.Password {
2020-04-05 15:49:43 +12:00
return id, ip, false, nil
2020-01-25 04:47:37 +13:00
}
2020-04-05 15:49:43 +12:00
return "", ip, false, fmt.Errorf("invalid password: %s", 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)
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
}
}
}
}