geo/server/internal/tenant/transport/desktop_dispatch_handler.go

package transport

import (
	"encoding/json"
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"go.uber.org/zap"

	"github.com/geo-platform/tenant-api/internal/bootstrap"
	"github.com/geo-platform/tenant-api/internal/shared/stream"
)

const (
	dispatchWriteWait       = 10 * time.Second
	dispatchPongWait        = 60 * time.Second
	dispatchPingPeriod      = 25 * time.Second
	dispatchReadLimitBytes  = 2048
	dispatchCloseGracePause = 500 * time.Millisecond
)

type DesktopDispatchHandler struct {
	app      *bootstrap.App
	upgrader websocket.Upgrader
}

func NewDesktopDispatchHandler(a *bootstrap.App) *DesktopDispatchHandler {
	return &DesktopDispatchHandler{
		app: a,
		upgrader: websocket.Upgrader{
			ReadBufferSize:  4096,
			WriteBufferSize: 4096,
			CheckOrigin: func(_ *http.Request) bool {
				// Origin enforcement happens at the gateway / LB layer; this
				// endpoint is already gated by the bearer token middleware.
				return true
			},
		},
	}
}

// Dispatch upgrades the HTTP request into a WebSocket and streams dispatch
// events for the current desktop client. The contract intentionally mirrors
// the existing SSE events endpoint: connect with a bearer token, receive JSON
// frames whose shape matches DesktopDispatchEvent, and reconnect on drop.
func (h *DesktopDispatchHandler) Dispatch(c *gin.Context) {
	client := MustDesktopClient(c.Request.Context())
	hub := h.app.DesktopDispatch
	if hub == nil {
		c.AbortWithStatus(http.StatusServiceUnavailable)
		return
	}

	conn, err := h.upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		if h.app.Logger != nil {
			h.app.Logger.Warn("desktop dispatch upgrade failed",
				zap.String("client_id", client.ID.String()),
				zap.Error(err),
			)
		}
		return
	}

	subscriber, cancel := hub.Subscribe(client.ID.String())
	defer cancel()

	conn.SetReadLimit(dispatchReadLimitBytes)
	_ = conn.SetReadDeadline(time.Now().Add(dispatchPongWait))
	conn.SetPongHandler(func(string) error {
		_ = conn.SetReadDeadline(time.Now().Add(dispatchPongWait))
		return nil
	})

	// Reader goroutine: drains control frames and signals disconnect. We do
	// not expect application data from the client over this socket.
	readerDone := make(chan struct{})
	go func() {
		defer close(readerDone)
		for {
			if _, _, err := conn.ReadMessage(); err != nil {
				return
			}
		}
	}()

	connected := stream.DesktopDispatchEvent{
		Type:           "connected",
		TargetClientID: client.ID.String(),
		WorkspaceID:    client.WorkspaceID,
		UpdatedAt:      time.Now().UTC(),
	}
	if err := writeDispatchFrame(conn, connected); err != nil {
		_ = conn.Close()
		return
	}

	pingTicker := time.NewTicker(dispatchPingPeriod)
	defer pingTicker.Stop()

	defer func() {
		_ = conn.SetWriteDeadline(time.Now().Add(dispatchWriteWait))
		_ = conn.WriteMessage(websocket.CloseMessage,
			websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
		time.Sleep(dispatchCloseGracePause)
		_ = conn.Close()
	}()

	outbound := subscriber.Outbound()
	for {
		select {
		case <-c.Request.Context().Done():
			return
		case <-readerDone:
			return
		case event, ok := <-outbound:
			if !ok {
				return
			}
			if err := writeDispatchFrame(conn, event); err != nil {
				return
			}
		case <-pingTicker.C:
			_ = conn.SetWriteDeadline(time.Now().Add(dispatchWriteWait))
			if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}

func writeDispatchFrame(conn *websocket.Conn, event stream.DesktopDispatchEvent) error {
	body, err := json.Marshal(event)
	if err != nil {
		return err
	}
	_ = conn.SetWriteDeadline(time.Now().Add(dispatchWriteWait))
	return conn.WriteMessage(websocket.TextMessage, body)
}