package websocket import ( "fmt" "github.com/gofiber/contrib/v3/websocket" "github.com/google/uuid" "github.com/rs/zerolog/log" "github.com/samber/lo" pb "git.solsynth.dev/goatworks/turbine/pkg/shared/proto/gen" ) // WebSocketController handles the WebSocket endpoint. type WebSocketController struct { Manager *Manager } // NewWebSocketController creates a new WebSocketController. func NewWebSocketController(manager *Manager) *WebSocketController { return &WebSocketController{ Manager: manager, } } // HandleWebSocket is the main handler for the /ws endpoint. func (wc *WebSocketController) HandleWebSocket(c *websocket.Conn) { // Mock Authentication for now based on C# example // In a real scenario, this would involve JWT verification, session lookup, etc. // For demonstration, we'll assume a dummy user and session. // The C# code uses HttpContext.Items to get CurrentUser and CurrentSession. // In Go Fiber, we can pass this through Locals or use middleware to set it up. // For now, let's create a dummy user and session. // TODO: Replace with actual authentication logic // For now, assume a dummy account and session // Based on the C# code, CurrentUser and CurrentSession are pb.Account and pb.AuthSession dummyAccount := &pb.Account{ Id: uuid.New().String(), Name: "dummy_user", Nick: "Dummy User", } dummySession := &pb.AuthSession{ ClientId: lo.ToPtr(uuid.New().String()), // This is used as deviceId if not present } // Device ID handling deviceAlt := c.Query("deviceAlt") if deviceAlt != "" { allowedDeviceAlternative := []string{"watch"} // Hardcoded for now found := false for _, alt := range allowedDeviceAlternative { if deviceAlt == alt { found = true break } } if !found { log.Warn().Msgf("Unsupported device alternative: %s", deviceAlt) bytes, _ := newErrorPacket("Unsupported device alternative: " + deviceAlt).ToBytes() c.WriteMessage(websocket.BinaryMessage, bytes) c.Close() return } } accountID := uuid.MustParse(dummyAccount.Id) deviceIDStr := "" if dummySession.ClientId == nil { deviceIDStr = uuid.New().String() } else { deviceIDStr = *dummySession.ClientId } if deviceAlt != "" { deviceIDStr = fmt.Sprintf("%s+%s", deviceIDStr, deviceAlt) } // Setup connection context for cancellation cancel := func() {} // Placeholder connectionKey := ConnectionKey{AccountID: accountID, DeviceID: deviceIDStr} // Add connection to manager if !wc.Manager.TryAdd(connectionKey, c, cancel) { bytes, _ := newErrorPacket("Too many connections from the same device and account.").ToBytes() c.WriteMessage(websocket.BinaryMessage, bytes) c.Close() return } log.Debug().Msgf("Connection established with user @%s#%s and device #%s", dummyAccount.Name, dummyAccount.Id, deviceIDStr) // Publish WebSocket connected event wc.Manager.PublishWebSocketConnectedEvent(accountID, deviceIDStr, false) // isOffline is false on connection defer func() { wc.Manager.Disconnect(connectionKey, "Client disconnected.") // Publish WebSocket disconnected event isOffline := !wc.Manager.GetAccountIsConnected(accountID) // Check if account is completely offline wc.Manager.PublishWebSocketDisconnectedEvent(accountID, deviceIDStr, isOffline) log.Debug().Msgf("Connection disconnected with user @%s#%s and device #%s", dummyAccount.Name, dummyAccount.Id, deviceIDStr) }() // Main event loop for { mt, msg, err := c.ReadMessage() if err != nil { log.Error().Err(err).Msg("WebSocket read error") break } if mt == websocket.CloseMessage { log.Info().Msg("Received close message from client") break } if mt != websocket.BinaryMessage { log.Warn().Msgf("Received non-binary message type: %d", mt) continue } packet, err := FromBytes(msg) if err != nil { log.Error().Err(err).Msg("Failed to deserialize WebSocket packet") bytes, _ := newErrorPacket("Failed to deserialize packet").ToBytes() c.WriteMessage(websocket.BinaryMessage, bytes) continue } if err := wc.Manager.HandlePacket(dummyAccount, deviceIDStr, packet, c); err != nil { log.Error().Err(err).Msg("Failed to handle incoming WebSocket packet") bytes, _ := newErrorPacket("Failed to process packet").ToBytes() c.WriteMessage(websocket.BinaryMessage, bytes) // Depending on error, might want to close connection } } } // newErrorPacket creates a new WebSocketPacket with an error type. func newErrorPacket(message string) *WebSocketPacket { return &WebSocketPacket{ Type: WebSocketPacketTypeError, ErrorMessage: message, } }