panel/agent/internal/rcon/websocket.go

package rcon

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/url"
	"sync"
	"sync/atomic"
	"time"

	"golang.org/x/net/websocket"
)

// WebSocketDialer speaks Facepunch's WebSocket RCON — the protocol used
// by Rust (app 258550) since they removed Source RCON years ago. A small
// subset of games have adopted the same pattern (a few Oxide-compatible
// community servers, some Unity-based survival games).
//
// Wire protocol (JSON text frames over a plain websocket):
//
//	client → server:  {"Identifier": N, "Message": "command", "Name": "panel"}
//	server → client:  {"Identifier": N, "Message": "result text",
//	                    "Type": "Generic", "Stacktrace": ""}
//
// The password is passed as the URL path, not a header:
//
//	ws://<host>:<port>/<password>
//
// Connections are long-lived — Rust's RCON doesn't close after each command
// the way Conan's does — but we still redial transparently on EOF/reset
// via the tracker's existing logic.
type WebSocketDialer struct{}

// Name identifies this adapter in the module manifest.
func (WebSocketDialer) Name() string { return "websocket_rcon" }

// Dial opens a websocket to ws://host:port/password. Returns a client that
// serializes Exec calls; Rust sometimes interleaves async log spam into the
// same connection as command replies, so we match responses back by id.
func (WebSocketDialer) Dial(ctx context.Context, opts DialOptions) (Client, error) {
	if opts.Host == "" || opts.Port == 0 {
		return nil, errors.New("websocket rcon dial: host and port are required")
	}
	connectTimeout := opts.ConnectTimeout
	if connectTimeout == 0 {
		connectTimeout = 5 * time.Second
	}

	wsURL := fmt.Sprintf("ws://%s:%d/%s", opts.Host, opts.Port, url.PathEscape(opts.Password))
	origin := fmt.Sprintf("http://%s:%d/", opts.Host, opts.Port)
	cfg, err := websocket.NewConfig(wsURL, origin)
	if err != nil {
		return nil, fmt.Errorf("websocket rcon: build config: %w", err)
	}

	// x/net/websocket.Dial doesn't honor a context; run it in a goroutine
	// so we can bail out when ctx.Done fires. Rust's authentication is
	// implicit — bad password = server closes the handshake with a 401.
	type dialResult struct {
		ws  *websocket.Conn
		err error
	}
	ch := make(chan dialResult, 1)
	go func() {
		ws, err := websocket.DialConfig(cfg)
		ch <- dialResult{ws, err}
	}()
	select {
	case r := <-ch:
		if r.err != nil {
			return nil, fmt.Errorf("websocket rcon dial %s:%d: %w", opts.Host, opts.Port, r.err)
		}
		return newWebSocketClient(r.ws), nil
	case <-time.After(connectTimeout):
		return nil, fmt.Errorf("websocket rcon dial %s:%d: connect timeout", opts.Host, opts.Port)
	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

type wsRequest struct {
	Identifier int32  `json:"Identifier"`
	Message    string `json:"Message"`
	Name       string `json:"Name"`
}

type wsResponse struct {
	Message    string `json:"Message"`
	Identifier int32  `json:"Identifier"`
	Type       string `json:"Type"`
	Stacktrace string `json:"Stacktrace"`
}

type wsClient struct {
	ws      *websocket.Conn
	nextID  atomic.Int32
	execMu  sync.Mutex // serializes Exec calls
	readMu  sync.Mutex // guards ws reads; Close may race
	closed  atomic.Bool
	unsolic chan wsResponse // buffers server-pushed log events we don't consume
}

func newWebSocketClient(ws *websocket.Conn) *wsClient {
	c := &wsClient{ws: ws, unsolic: make(chan wsResponse, 64)}
	c.nextID.Store(1)
	return c
}

// Exec runs one command and returns the matching response body. Rust
// echoes our request's Identifier back, so we can reject unrelated log
// messages the server pushes on the same connection.
func (c *wsClient) Exec(ctx context.Context, cmd string) (string, error) {
	if c.closed.Load() {
		return "", errors.New("websocket rcon: client closed")
	}
	c.execMu.Lock()
	defer c.execMu.Unlock()

	id := c.nextID.Add(1)
	req := wsRequest{Identifier: id, Message: cmd, Name: "panel"}
	payload, err := json.Marshal(req)
	if err != nil {
		return "", fmt.Errorf("marshal: %w", err)
	}

	deadline, ok := ctx.Deadline()
	if !ok {
		deadline = time.Now().Add(10 * time.Second)
	}
	if err := c.ws.SetDeadline(deadline); err != nil {
		return "", fmt.Errorf("set deadline: %w", err)
	}
	defer c.ws.SetDeadline(time.Time{}) //nolint:errcheck

	if err := websocket.Message.Send(c.ws, string(payload)); err != nil {
		return "", fmt.Errorf("send: %w", err)
	}

	// Read frames until one matches our Identifier. Rust pipes chat, log,
	// and player-event messages through the same socket; if their id is
	// 0 (or doesn't match ours) we drop them into unsolic for someone
	// else to consume. 200ms allowance for interleaved frames.
	for {
		c.readMu.Lock()
		var raw string
		rerr := websocket.Message.Receive(c.ws, &raw)
		c.readMu.Unlock()
		if rerr != nil {
			return "", fmt.Errorf("recv: %w", rerr)
		}
		var resp wsResponse
		if err := json.Unmarshal([]byte(raw), &resp); err != nil {
			// Non-JSON frame — Rust shouldn't send those but be forgiving.
			continue
		}
		if resp.Identifier != id {
			// Async server push (log line, chat, etc.). Buffer non-blockingly.
			select {
			case c.unsolic <- resp:
			default:
			}
			continue
		}
		if resp.Stacktrace != "" {
			return resp.Message, fmt.Errorf("rust rcon error: %s", resp.Stacktrace)
		}
		return resp.Message, nil
	}
}

// Close tears down the underlying websocket connection.
func (c *wsClient) Close() error {
	if !c.closed.CompareAndSwap(false, true) {
		return nil
	}
	return c.ws.Close()
}