snow/handlers/socket.go

package handlers

import (
	"github.com/beevik/etree"
	"github.com/ectrc/snow/aid"
	"github.com/ectrc/snow/person"
	"github.com/gofiber/contrib/websocket"
	"github.com/gofiber/fiber/v2"
	"github.com/google/uuid"
)

type SocketType string
var (
	SocketTypeXmpp SocketType = "xmpp"
	SocketTypeUnknown SocketType = "unknown"
)

type Socket struct {
	ID string
	JID string
	Type SocketType
	Connection *websocket.Conn
	Person *person.Person
}

type MessageToWrite struct {
	Socket *Socket
	Message []byte
}

func (s *Socket) Write(message []byte) {
	socketWriteQueue <- MessageToWrite{
		Socket: s,
		Message: message,
	}
}

func (s *Socket) WriteTree(message *etree.Document) {
	bytes, err := message.WriteToBytes()
	if err != nil {
		return
	}

	socketWriteQueue <- MessageToWrite{
		Socket: s,
		Message: bytes,
	}
}

var (
	socketWriteQueue = make(chan MessageToWrite, 1000)
	socketHandlers = map[SocketType]func(string) {
		SocketTypeXmpp: handlePresenceSocket,
	}

	sockets = aid.GenericSyncMap[Socket]{}
)

func MiddlewareWebsocket(c *fiber.Ctx) error {
	if !websocket.IsWebSocketUpgrade(c) {
		return fiber.ErrUpgradeRequired
	}

	var protocol SocketType
	switch c.Get("Sec-WebSocket-Protocol") {
	case "xmpp":
		protocol = SocketTypeXmpp
	default:
		protocol = SocketTypeUnknown
	}

	c.Locals("uuid", uuid.New().String())
	c.Locals("protocol", protocol)
	return c.Next()
}

func WebsocketConnection(c *websocket.Conn) {
	protocol := c.Locals("protocol").(SocketType)
	uuid := c.Locals("uuid").(string)

	sockets.Add(uuid, &Socket{
		ID: uuid,
		Type: protocol,
		Connection: c,
	})
	defer func() {
		socket, ok := sockets.Get(uuid)
		if !ok {
			return
		}
		socket.Connection.Close()
		sockets.Delete(uuid)
		aid.Print("(xmpp) connection closed", uuid)
	}()

	if handle, ok := socketHandlers[protocol]; ok {
		handle(uuid)
	} 
}

func init() {
	go func() {
		for {
			if aid.Config != nil {
				break
			}
		}

		aid.Print("(socket) write queue started")

		for {
			message := <-socketWriteQueue
			aid.Print("(socket) message sent", message.Socket.ID, string(message.Message))
			message.Socket.Connection.WriteMessage(websocket.TextMessage, message.Message)
		}
	}()
}