simpleadmin-web/internal/rcon/rcon.go

// Package rcon is a Source RCON client (the TCP protocol CS2 still speaks).
//
// One connection is kept open and shared; commands are serialised with a mutex. A response can
// span several packets, so every command is followed by an empty SERVERDATA_RESPONSE_VALUE packet:
// the server answers packets in order, so seeing that packet's echo means the command's output is
// complete. If the server never echoes it, reading stops after a short idle timeout instead.
package rcon

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"strings"
	"sync"
	"time"
)

const (
	typeResponseValue = 0
	typeExecCommand   = 2
	typeAuthResponse  = 2
	typeAuth          = 3

	maxPacket = 4096 + 14
)

// ErrAuth means the server rejected the RCON password.
var ErrAuth = errors.New("rcon: wrong password")

// ErrDown means a recent attempt to reach the server failed, so this one wasn't tried.
var ErrDown = errors.New("rcon: server unreachable, retrying shortly")

// backoff is how long Exec fails fast after the server couldn't be reached, so page loads and
// actions don't each wait out a connection timeout while the game server is down.
const backoff = 10 * time.Second

type Client struct {
	addr     string
	password string
	timeout  time.Duration

	mu        sync.Mutex
	conn      net.Conn
	nextID    int32
	downUntil time.Time
	lastErr   error
}

func New(addr, password string) *Client {
	return &Client{addr: addr, password: password, timeout: 5 * time.Second, nextID: 1}
}

// Exec runs a console command and returns everything it printed.
func (c *Client) Exec(cmd string) (string, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	if time.Now().Before(c.downUntil) {
		return "", fmt.Errorf("%w: %v", ErrDown, c.lastErr)
	}
	hadConn := c.conn != nil
	out, err := c.exec(cmd)
	if err != nil && hadConn && !errors.Is(err, ErrAuth) {
		// The kept-open connection may have gone stale (server restart, idle timeout): retry once
		// on a fresh one. A failed fresh connection isn't retried.
		c.close()
		out, err = c.exec(cmd)
	}
	if err != nil {
		c.close()
		c.downUntil, c.lastErr = time.Now().Add(backoff), err
	}
	return out, err
}

func (c *Client) exec(cmd string) (string, error) {
	if c.conn == nil {
		if err := c.connect(); err != nil {
			return "", err
		}
	}
	id := c.id()
	endID := c.id()
	if err := c.write(id, typeExecCommand, cmd); err != nil {
		return "", err
	}
	if err := c.write(endID, typeResponseValue, ""); err != nil {
		return "", err
	}

	var out strings.Builder
	deadline := time.Now().Add(c.timeout)
	idle := 400 * time.Millisecond
	got := false
	for {
		wait := deadline
		if got {
			// Once output has started, stop after a quiet gap even without the end marker.
			if d := time.Now().Add(idle); d.Before(wait) {
				wait = d
			}
		}
		_ = c.conn.SetReadDeadline(wait)
		pid, _, body, err := c.read()
		if err != nil {
			var ne net.Error
			if errors.As(err, &ne) && ne.Timeout() && got {
				return out.String(), nil
			}
			return "", err
		}
		switch pid {
		case id:
			out.WriteString(body)
			got = true
		case endID:
			// Some servers answer the empty packet with a second, odd-looking packet. Drain it.
			_ = c.conn.SetReadDeadline(time.Now().Add(50 * time.Millisecond))
			_, _, _, _ = c.read()
			return out.String(), nil
		}
	}
}

func (c *Client) connect() error {
	conn, err := net.DialTimeout("tcp", c.addr, 2*time.Second)
	if err != nil {
		return fmt.Errorf("rcon: connect %s: %w", c.addr, err)
	}
	c.conn = conn
	id := c.id()
	if err := c.write(id, typeAuth, c.password); err != nil {
		c.close()
		return err
	}
	_ = conn.SetReadDeadline(time.Now().Add(c.timeout))
	for {
		pid, ptype, _, err := c.read()
		if err != nil {
			c.close()
			return fmt.Errorf("rcon: auth: %w", err)
		}
		if ptype != typeAuthResponse {
			continue // an empty RESPONSE_VALUE precedes the auth response
		}
		if pid == -1 {
			c.close()
			return ErrAuth
		}
		return nil
	}
}

func (c *Client) id() int32 {
	c.nextID++
	if c.nextID > 1<<30 {
		c.nextID = 1
	}
	return c.nextID
}

func (c *Client) write(id int32, ptype int32, body string) error {
	var buf bytes.Buffer
	size := int32(len(body) + 10)
	_ = binary.Write(&buf, binary.LittleEndian, size)
	_ = binary.Write(&buf, binary.LittleEndian, id)
	_ = binary.Write(&buf, binary.LittleEndian, ptype)
	buf.WriteString(body)
	buf.Write([]byte{0, 0})
	_ = c.conn.SetWriteDeadline(time.Now().Add(c.timeout))
	_, err := c.conn.Write(buf.Bytes())
	return err
}

func (c *Client) read() (id int32, ptype int32, body string, err error) {
	var size int32
	if err = binary.Read(c.conn, binary.LittleEndian, &size); err != nil {
		return
	}
	if size < 10 || size > maxPacket {
		err = fmt.Errorf("rcon: bad packet size %d", size)
		return
	}
	data := make([]byte, size)
	if _, err = io.ReadFull(c.conn, data); err != nil {
		return
	}
	id = int32(binary.LittleEndian.Uint32(data[0:4]))
	ptype = int32(binary.LittleEndian.Uint32(data[4:8]))
	body = string(bytes.TrimRight(data[8:], "\x00"))
	return
}

func (c *Client) close() {
	if c.conn != nil {
		_ = c.conn.Close()
		c.conn = nil
	}
}

// Close drops the connection. The next Exec reconnects.
func (c *Client) Close() {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.close()
}

// Arg makes s safe to pass as one quoted console argument: the console splits commands on ';' and
// newlines and has no escape for '"', so those are removed along with other control characters.
func Arg(s string, max int) string {
	var b strings.Builder
	for _, r := range s {
		switch {
		case r == '"' || r == ';' || r == '\\':
			continue
		case r < 0x20 || r == 0x7f:
			b.WriteRune(' ')
		default:
			b.WriteRune(r)
		}
	}
	out := strings.TrimSpace(b.String())
	if max > 0 && len([]rune(out)) > max {
		out = string([]rune(out)[:max])
	}
	return out
}