110 lines
2.4 KiB
Go
110 lines
2.4 KiB
Go
package streaming
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"net/http"
|
|
|
|
"git.sleepycat.moe/sam/mercury/web/api"
|
|
"git.sleepycat.moe/sam/mercury/web/app"
|
|
"github.com/gorilla/websocket"
|
|
"github.com/rs/zerolog/log"
|
|
)
|
|
|
|
type App struct {
|
|
*app.App
|
|
}
|
|
|
|
func New(app *app.App) *App {
|
|
return &App{
|
|
App: app,
|
|
}
|
|
}
|
|
|
|
var upgrader = websocket.Upgrader{}
|
|
|
|
func (app *App) Streaming(w http.ResponseWriter, r *http.Request) error {
|
|
token, _ := app.TokenFromContext(r.Context())
|
|
conn, err := upgrader.Upgrade(w, r, nil)
|
|
if err != nil {
|
|
log.Err(err).Msg("performing websocket handshake")
|
|
hse := websocket.HandshakeError{}
|
|
if errors.As(err, &hse) {
|
|
return nil // the upgrader already sent a response for us
|
|
}
|
|
|
|
return err
|
|
}
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
socket, ok := SocketHolder.socketsFor(token.UserID).newSocket(ctx, cancel)
|
|
if !ok {
|
|
err := conn.WriteJSON(newEvent(EventTypeError, ErrorEvent{Code: api.ErrTooManyStreams, Message: "Too many streams open"}))
|
|
if err != nil {
|
|
log.Err(err).Msg("writing stream rejection message to socket")
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
go app.writeStream(conn, socket)
|
|
go app.readStream(conn, socket)
|
|
|
|
return nil
|
|
}
|
|
|
|
func (app *App) writeStream(conn *websocket.Conn, socket *socket) {
|
|
defer conn.Close()
|
|
|
|
for {
|
|
select {
|
|
case <-socket.ctx.Done():
|
|
return
|
|
case ev := <-socket.ch:
|
|
// at this point, the type should already have been filtered, so just send the event
|
|
err := conn.WriteJSON(ev)
|
|
if err != nil {
|
|
// write failed, bail and make client reconnect
|
|
log.Err(err).Msg("error writing JSON to socket")
|
|
socket.cancel()
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (app *App) readStream(conn *websocket.Conn, socket *socket) {
|
|
for {
|
|
select {
|
|
case <-socket.ctx.Done():
|
|
return
|
|
default:
|
|
var e IncomingEvent
|
|
err := conn.ReadJSON(&e)
|
|
if err != nil {
|
|
// read failed, bail and make client reconnect
|
|
log.Err(err).Msg("error reading JSON from socket")
|
|
socket.cancel()
|
|
return
|
|
}
|
|
|
|
switch e.Type {
|
|
case EventTypeSubscribe, EventTypeUnsubscribe:
|
|
var et EventType
|
|
err = json.Unmarshal(e.Data, &et)
|
|
if err != nil {
|
|
// invalid event type, log but don't disconnect
|
|
log.Err(err).Msg("reading event type to subscribe to")
|
|
continue
|
|
}
|
|
|
|
if !et.ValidReceive() {
|
|
// if it's not a valid event, ignore silently
|
|
continue
|
|
}
|
|
|
|
socket.setEvent(et, e.Type != EventTypeSubscribe)
|
|
}
|
|
}
|
|
}
|
|
}
|