package wsapi import ( "encoding/json" "errors" "io" "strings" "dynatron.me/x/blasphem/pkg/auth" "dynatron.me/x/blasphem/pkg/blas/core" "github.com/gorilla/websocket" "github.com/labstack/echo/v4" "github.com/mitchellh/mapstructure" "github.com/rs/zerolog/log" ) var ( NoSuchHandlerErr = errors.New("bad websocket command") NoMessageIDErr = errors.New("no message ID") ) type Type string type MsgBase struct { Type Type `json:"type"` } type ( wsSession struct { *websocket.Conn b core.Blas ec echo.Context h phaseHandler user *auth.User refreshToken *auth.RefreshToken } phaseHandler interface { handleMsg(io.Reader) error } cmdHandler struct { *wsSession } wsEntry struct { dataNew core.NewData hnd core.Handler } wsRegistry map[string]wsEntry wsManager struct { r wsRegistry } ) func (wsm *wsManager) RegisterWSCommand(cmd string, hnd core.Handler, dataNew core.NewData) { wsm.r[cmd] = wsEntry{ dataNew: dataNew, hnd: hnd, } } func (wsm *wsManager) WSCommandHandler(cmd string, cmdSplit []string) (core.NewData, core.Handler, error) { if wse, ok := wsm.r[cmd]; ok { return wse.dataNew, wse.hnd, nil } if wse, ok := wsm.r[cmdSplit[0]]; ok { return wse.dataNew, wse.hnd, nil } return nil, nil, NoSuchHandlerErr } func NewManager() core.WebSocketManager { return &wsManager{ r: make(wsRegistry), } } func NewSession(s core.Blas, c echo.Context, conn *websocket.Conn) core.WebSocketSession { ws := &wsSession{ Conn: conn, b: s, ec: c, } ws.h = &authPhase{ws} return ws } func (ws *wsSession) WSConn() *websocket.Conn { return ws.Conn } func (ws *wsSession) Blas() core.Blas { return ws.b } func (ws *wsSession) Go() error { err := ws.sendAuthRequired() if err != nil { return err } for { _, rdr, err := ws.NextReader() switch { case err == nil: case websocket.IsCloseError(err, websocket.CloseGoingAway): return nil case err != nil: log.Error().Err(err).Str("remote", ws.ec.Request().RemoteAddr).Msg("websocket read fail") return err } err = ws.h.handleMsg(rdr) if err != nil { log.Error().Err(err).Msg("handleMsg") } } } type cmdMsg struct { } type MsgType string const ( ResultMsgType MsgType = "result" ) type Error struct { Code string `json:"code"` Message string `json:"message"` } type WSError struct { ID *int `json:"id,omitempty"` Type MsgType `json:"type"` Success bool `json:"success"` Error Error `json:"error"` } func (ws *cmdHandler) writeError(id int, err Error) error { return ws.WriteJSON(WSError{ ID: &id, Type: ResultMsgType, Success: false, Error: err, }) } func (ws *cmdHandler) handleMsg(r io.Reader) error { var msgMap map[string]interface{} err := json.NewDecoder(r).Decode(&msgMap) if err != nil { return err } msgType, ok := msgMap["type"].(string) if !ok { return NoSuchHandlerErr } idFl, ok := msgMap["id"].(float64) if !ok { return ws.WriteJSON( WSError{ Type: ResultMsgType, Success: false, Error: Error{ Code: "invalid_id", Message: "command has no ID", }, }) } id := int(idFl) cmd := strings.Split(msgType, "/") newData, hand, err := ws.b.WSCommandHandler(cmd[0], cmd) switch err { case nil: case NoSuchHandlerErr: return ws.writeError(id, Error{ Code: "invalid_type", Message: "no such command", }) default: return err } nd := newData(cmd) err = mapstructure.Decode(&msgMap, &nd) if err != nil { ws.writeError(id, Error{ Code: "invalid_format", Message: err.Error(), }) return nil } return hand(ws, id, cmd, nd) }