socket_connection.go

package velaros

import (
	"context"

	"github.com/coder/websocket"
)

// WebSocketConnection is a SocketConnection implementation that wraps
// github.com/coder/websocket.Conn. This is the default connection type used by the
// router when handling HTTP WebSocket upgrades.
type WebSocketConnection struct {
	webSocketConnection *websocket.Conn
}

var _ SocketConnection = &WebSocketConnection{}

// NewWebSocketConnection creates a WebSocketConnection from a github.com/coder/websocket.Conn.
// This is used internally by the router. Most applications don't need to call this directly
// unless implementing custom connection handling with Router.HandleConnection.
func NewWebSocketConnection(websocketConnection *websocket.Conn) *WebSocketConnection {
	return &WebSocketConnection{
		webSocketConnection: websocketConnection,
	}
}

// Read reads the next message from the WebSocket connection. Blocks until a message
// arrives or an error occurs. Implements SocketConnection.Read.
func (c *WebSocketConnection) Read(ctx context.Context) (*SocketMessage, error) {
	messageType, data, err := c.webSocketConnection.Read(ctx)
	if err != nil {
		return nil, err
	}

	return &SocketMessage{
		Type:    messageType,
		RawData: data,
	}, nil
}

// Write sends a message to the WebSocket connection. Implements SocketConnection.Write.
func (c *WebSocketConnection) Write(ctx context.Context, msg *SocketMessage) error {
	return c.webSocketConnection.Write(ctx, msg.Type, msg.Data)
}

// Close closes the WebSocket connection with the given status code and reason.
// Implements SocketConnection.Close.
func (c *WebSocketConnection) Close(status Status, reason string) error {
	return c.webSocketConnection.Close(websocket.StatusCode(status), reason)
}

// Ping sends a WebSocket-protocol PING frame and blocks until the peer's
// PONG arrives or ctx expires. Implements Pinger so the router can run
// application-level liveness checks (see Router.SetPingInterval).
//
// Underlying library auto-filters PING/PONG frames, so callers of Read never
// see them.
func (c *WebSocketConnection) Ping(ctx context.Context) error {
	return c.webSocketConnection.Ping(ctx)
}

// Pinger is an optional extension to SocketConnection. Connection
// implementations that can send a ping frame and observe the peer's pong
// should implement it to opt in to Router-driven liveness checks.
//
// The standard WebSocketConnection implements this. Custom SocketConnection
// implementations may omit it (in which case the router skips the ping loop
// for those connections).
type Pinger interface {
	Ping(ctx context.Context) error
}