1package web
  2
  3import (
  4	"context"
  5	"fmt"
  6	"log/slog"
  7	"net/http"
  8	"sync/atomic"
  9
 10	"github.com/coder/websocket"
 11
 12	"goVNC/pkg/proxy"
 13	"goVNC/pkg/transport"
 14)
 15
 16// WebsockifyHandler handles WebSocket proxy connections.
 17type WebsockifyHandler struct {
 18	mux        *proxy.Multiplexer
 19	maxClients int
 20	clientSeq  uint64
 21}
 22
 23// NewWebsockifyHandler creates a new WebSocket proxy handler.
 24func NewWebsockifyHandler(mux *proxy.Multiplexer) *WebsockifyHandler {
 25	return &WebsockifyHandler{
 26		mux:        mux,
 27		maxClients: 10, // default max clients
 28	}
 29}
 30
 31// SetMaxClients sets the maximum number of concurrent clients.
 32func (h *WebsockifyHandler) SetMaxClients(max int) {
 33	h.maxClients = max
 34}
 35
 36// ServeHTTP handles WebSocket upgrade and proxying.
 37func (h *WebsockifyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 38	// Enforce max clients limit
 39	if h.maxClients > 0 && h.mux.ClientCount() >= h.maxClients {
 40		http.Error(w, "too many clients", http.StatusServiceUnavailable)
 41		return
 42	}
 43
 44	// Accept WebSocket connection
 45	conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{
 46		Subprotocols:   []string{"binary"},
 47		OriginPatterns: []string{"*"}, // Allow all origins for now
 48	})
 49	if err != nil {
 50		slog.Error("accepting websocket", slog.String("error", err.Error()))
 51		return
 52	}
 53
 54	// Generate client ID
 55	seq := atomic.AddUint64(&h.clientSeq, 1)
 56	clientID := fmt.Sprintf("client-%d", seq)
 57
 58	// Wrap in transport
 59	clientTransport := &wsTransport{conn: conn}
 60
 61	// Register with multiplexer
 62	client := h.mux.AddClient(clientID, clientTransport)
 63	defer h.mux.RemoveClient(clientID)
 64
 65	// Set up context
 66	ctx := r.Context()
 67
 68	// Get current session info to send initial state
 69	session := h.mux.VNCClient().Session()
 70	if session != nil {
 71		// Send ServerInit-like message so client knows dimensions
 72		// noVNC expects to receive the full handshake, but since we've
 73		// already done that, we need to replay it
 74		slog.Info("client connected to session",
 75			slog.String("client", clientID),
 76			slog.String("session", session.Name))
 77	}
 78
 79	// Run reader and writer concurrently
 80	errCh := make(chan error, 2)
 81
 82	go func() {
 83		errCh <- h.mux.RunClientWriter(ctx, client)
 84	}()
 85
 86	go func() {
 87		errCh <- h.mux.RunClientReader(ctx, client)
 88	}()
 89
 90	// Wait for either to complete
 91	err = <-errCh
 92	if err != nil && err != context.Canceled {
 93		slog.Debug("client connection ended",
 94			slog.String("client", clientID),
 95			slog.String("error", err.Error()))
 96	}
 97
 98	// Close the connection
 99	conn.Close(websocket.StatusNormalClosure, "")
100}
101
102// wsTransport wraps a WebSocket connection as a Transport.
103type wsTransport struct {
104	conn *websocket.Conn
105}
106
107func (t *wsTransport) Read(ctx context.Context) ([]byte, error) {
108	_, data, err := t.conn.Read(ctx)
109	return data, err
110}
111
112func (t *wsTransport) Write(ctx context.Context, data []byte) error {
113	return t.conn.Write(ctx, websocket.MessageBinary, data)
114}
115
116func (t *wsTransport) Close() error {
117	return t.conn.CloseNow()
118}
119
120func (t *wsTransport) SetReadLimit(limit int64) {
121	t.conn.SetReadLimit(limit)
122}
123
124var _ transport.Transport = (*wsTransport)(nil)