lupus-lite/routes/ws.go

package routes

import (
	"log"

	"github.com/aziis98/lupus-lite/events"
	"github.com/aziis98/lupus-lite/util"
	"github.com/gofiber/fiber/v2"
	"github.com/gofiber/websocket/v2"
)

type clientMessage struct {
	Type  string `json:"type"`
	Event string `json:"event"`
}

type serverMessage struct {
	Event string `json:"event"`
	Data  any    `json:"data"`
}

// BindWebsocketToEventBus creates a WebSocket handler that can subscribe the socket to some events generated on the server, for example a message from the client like { type: "subscribe", event: "foo" } binds the connection to the event "foo" and forwards events and the associated data to client as JSON.
func BindWebsocketToEventBus(eb *events.EventBus) fiber.Handler {
	return websocket.New(func(c *websocket.Conn) {
		clientId := util.GenerateRandomString(10)

		log.Printf("[%v] Connection started", clientId)
		defer log.Printf("[%v] Connection closed", clientId)

		listeners := []*events.Listener{}
		defer func() {
			for _, l := range listeners {
				eb.Unsubscribe(l)
			}
		}()

		done := make(chan error)
		for {
			var msg clientMessage
			if err := c.ReadJSON(&msg); err != nil {
				log.Printf(`[%v] Read error: %v`, clientId, err)
				return
			}

			log.Printf("[%v] Received: %v", clientId, msg)

			switch msg.Type {
			case "subscribe":
				newListener := eb.Subscribe(
					msg.Event,
					func(e interface{}) {
						log.Printf(`[%v] Sent message`, clientId)
						if err := c.WriteJSON(serverMessage{
							Event: msg.Event,
							Data:  e,
						}); err != nil {
							log.Printf(`[%v] Write error: %v`, clientId, err)
							done <- err
						}
					},
				)
				listeners = append(listeners, newListener)
			default:
				log.Printf("[%v] Unknown message type %q", clientId, msg.Type)
				return
			}

			select {
			case err := <-done:
				if err != nil {
					return
				}
			default:
			}
		}
	})
}