1package transport
  2
  3import (
  4	"context"
  5	"fmt"
  6	"net/http"
  7	"net/url"
  8
  9	"github.com/coder/websocket"
 10)
 11
 12// WebSocketTransport implements Transport over WebSocket connections.
 13// This is used for noVNC-style VNC connections where each RFB message
 14// is sent as a discrete WebSocket binary message.
 15type WebSocketTransport struct {
 16	conn *websocket.Conn
 17}
 18
 19// WebSocketDialer implements Dialer for WebSocket connections.
 20type WebSocketDialer struct{}
 21
 22// NewWebSocketDialer creates a new WebSocket dialer.
 23func NewWebSocketDialer() *WebSocketDialer {
 24	return &WebSocketDialer{}
 25}
 26
 27// Dial establishes a WebSocket connection to the target URL.
 28func (d *WebSocketDialer) Dial(ctx context.Context, target string, opts *DialOptions) (Transport, error) {
 29	wsURL, err := url.Parse(target)
 30	if err != nil {
 31		return nil, fmt.Errorf("parsing url=%q: %w", target, err)
 32	}
 33
 34	dialOpts := &websocket.DialOptions{}
 35
 36	if opts != nil {
 37		// Build HTTP headers
 38		if opts.Headers != nil {
 39			dialOpts.HTTPHeader = opts.Headers.Clone()
 40		} else {
 41			dialOpts.HTTPHeader = make(http.Header)
 42		}
 43
 44		// Add User-Agent if specified
 45		if opts.UserAgent != "" {
 46			dialOpts.HTTPHeader.Set("User-Agent", opts.UserAgent)
 47		}
 48
 49		// Add subprotocols if specified
 50		if len(opts.Subprotocols) > 0 {
 51			dialOpts.Subprotocols = opts.Subprotocols
 52		}
 53
 54		// TODO: Support proxy via custom HTTP client
 55		// TODO: Support custom TLS config
 56	}
 57
 58	// Apply timeout to context if specified
 59	if opts != nil && opts.Timeout > 0 {
 60		var cancel context.CancelFunc
 61		ctx, cancel = context.WithTimeout(ctx, opts.Timeout)
 62		defer cancel()
 63	}
 64
 65	conn, _, err := websocket.Dial(ctx, wsURL.String(), dialOpts)
 66	if err != nil {
 67		return nil, fmt.Errorf("connecting to url=%q: %w", wsURL.String(), err)
 68	}
 69
 70	return &WebSocketTransport{conn: conn}, nil
 71}
 72
 73// Read reads a single message from the WebSocket connection.
 74func (t *WebSocketTransport) Read(ctx context.Context) ([]byte, error) {
 75	_, data, err := t.conn.Read(ctx)
 76	if err != nil {
 77		return nil, err
 78	}
 79	return data, nil
 80}
 81
 82// Write writes a single message to the WebSocket connection.
 83func (t *WebSocketTransport) Write(ctx context.Context, data []byte) error {
 84	return t.conn.Write(ctx, websocket.MessageBinary, data)
 85}
 86
 87// Close closes the WebSocket connection.
 88func (t *WebSocketTransport) Close() error {
 89	return t.conn.CloseNow()
 90}
 91
 92// SetReadLimit sets the maximum message size for reads.
 93func (t *WebSocketTransport) SetReadLimit(limit int64) {
 94	t.conn.SetReadLimit(limit)
 95}
 96
 97// Conn returns the underlying WebSocket connection.
 98// This is useful for advanced operations not exposed by the Transport interface.
 99func (t *WebSocketTransport) Conn() *websocket.Conn {
100	return t.conn
101}
102
103// Compile-time interface checks
104var (
105	_ Transport = (*WebSocketTransport)(nil)
106	_ Dialer    = (*WebSocketDialer)(nil)
107)