Initial commit: Beszel fork with Domain Locker integration

internal/hub/transport/websocket.go

@@ -0,0 +1,74 @@
+package transport
+
+import (
+	"context"
+	"errors"
+
+	"github.com/fxamacker/cbor/v2"
+	"github.com/henrygd/beszel"
+	"github.com/henrygd/beszel/internal/common"
+	"github.com/henrygd/beszel/internal/hub/ws"
+)
+
+// ErrWebSocketNotConnected indicates a WebSocket transport is not currently connected.
+var ErrWebSocketNotConnected = errors.New("websocket not connected")
+
+// WebSocketTransport implements Transport over WebSocket connections.
+type WebSocketTransport struct {
+	wsConn *ws.WsConn
+}
+
+// NewWebSocketTransport creates a new WebSocket transport wrapper.
+func NewWebSocketTransport(wsConn *ws.WsConn) *WebSocketTransport {
+	return &WebSocketTransport{wsConn: wsConn}
+}
+
+// Request sends a request to the agent via WebSocket and unmarshals the response.
+func (t *WebSocketTransport) Request(ctx context.Context, action common.WebSocketAction, req any, dest any) error {
+	if !t.IsConnected() {
+		return ErrWebSocketNotConnected
+	}
+
+	pendingReq, err := t.wsConn.SendRequest(ctx, action, req)
+	if err != nil {
+		return err
+	}
+
+	// Wait for response
+	select {
+	case message := <-pendingReq.ResponseCh:
+		defer message.Close()
+		defer pendingReq.Cancel()
+
+		// Legacy agents (< MinVersionAgentResponse) respond with a raw payload instead of an AgentResponse wrapper.
+		if t.wsConn.AgentVersion().LT(beszel.MinVersionAgentResponse) {
+			return cbor.Unmarshal(message.Data.Bytes(), dest)
+		}
+
+		var agentResponse common.AgentResponse
+		if err := cbor.Unmarshal(message.Data.Bytes(), &agentResponse); err != nil {
+			return err
+		}
+
+		if agentResponse.Error != "" {
+			return errors.New(agentResponse.Error)
+		}
+
+		return UnmarshalResponse(agentResponse, action, dest)
+
+	case <-pendingReq.Context.Done():
+		return pendingReq.Context.Err()
+	}
+}
+
+// IsConnected returns true if the WebSocket connection is active.
+func (t *WebSocketTransport) IsConnected() bool {
+	return t.wsConn != nil && t.wsConn.IsConnected()
+}
+
+// Close terminates the WebSocket connection.
+func (t *WebSocketTransport) Close() {
+	if t.wsConn != nil {
+		t.wsConn.Close(nil)
+	}
+}