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/handler.go
2020-01-18 23:30:09 +00:00

206 lines
4.6 KiB
Go

package websocket
import (
"fmt"
"net/http"
"time"
"github.com/gorilla/websocket"
gonanoid "github.com/matoous/go-nanoid"
"github.com/rs/zerolog"
"github.com/rs/zerolog/log"
"n.eko.moe/neko/internal/config"
"n.eko.moe/neko/internal/session"
"n.eko.moe/neko/internal/webrtc"
)
func New(sessions *session.SessionManager, webrtc *webrtc.WebRTCManager, conf *config.WebSocket) *WebSocketHandler {
logger := log.With().Str("module", "websocket").Logger()
return &WebSocketHandler{
logger: logger,
conf: conf,
sessions: sessions,
upgrader: websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool {
return true
},
},
handler: &MessageHandler{
logger: logger.With().Str("subsystem", "handler").Logger(),
sessions: sessions,
webrtc: webrtc,
},
}
}
// Send pings to peer with this period. Must be less than pongWait.
const pingPeriod = 60 * time.Second
const alphabet = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
type WebSocketHandler struct {
logger zerolog.Logger
upgrader websocket.Upgrader
handler *MessageHandler
conf *config.WebSocket
sessions *session.SessionManager
shutdown chan bool
}
func (ws *WebSocketHandler) Start() error {
go func() {
defer func() {
ws.logger.Info().Msg("shutdown")
}()
for {
select {
case <-ws.shutdown:
return
}
}
}()
ws.sessions.OnCreated(func(id string, session *session.Session) {
if err := ws.handler.Created(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.OnDestroy(func(id string) {
if err := ws.handler.Destroyed(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")
}
})
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")
socket, err := ws.upgrader.Upgrade(w, r, nil)
if err != nil {
ws.logger.Error().Err(err).Msg("failed to upgrade connection")
return err
}
id, admin, err := ws.authenticate(r)
if err != nil {
ws.logger.Warn().Err(err).Msg("authenticatetion failed")
if err = socket.Close(); err != nil {
return err
}
return nil
}
ws.sessions.New(id, admin, socket)
ws.logger.
Debug().
Str("session", id).
Str("address", socket.RemoteAddr().String()).
Msg("new connection created")
defer func() {
ws.logger.
Debug().
Str("session", id).
Str("address", socket.RemoteAddr().String()).
Msg("session ended")
}()
if err = ws.handler.Connected(id, socket); err != nil {
ws.logger.Error().Err(err).Msg("connection failed")
if err = socket.Close(); err != nil {
return err
}
return nil
}
ws.handle(socket, id)
return nil
}
func (ws *WebSocketHandler) authenticate(r *http.Request) (string, bool, error) {
id, err := gonanoid.Generate(alphabet, 32)
if err != nil {
return "", false, err
}
passwords, ok := r.URL.Query()["password"]
if !ok || len(passwords[0]) < 1 {
return "", false, fmt.Errorf("no password provided")
}
if passwords[0] == ws.conf.AdminPassword {
return id, true, nil
}
if passwords[0] == ws.conf.Password {
return id, false, nil
}
return "", false, fmt.Errorf("invalid password: %s", passwords[0])
}
func (ws *WebSocketHandler) handle(socket *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", socket.RemoteAddr().String()).Msg("handle socket ending")
ws.handler.Disconnected(id)
}()
for {
_, raw, err := socket.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("raw", string(raw)).
Msg("recieved message from client")
if err := ws.handler.Message(id, raw); err != nil {
ws.logger.Error().Err(err).Msg("message handler has failed")
return
}
case <-cancel:
return
case _ = <-ticker.C:
if err := socket.WriteMessage(websocket.PingMessage, nil); err != nil {
return
}
}
}
}