simpleadmin-web 0.1.0: web panel for CS2-SimpleAdmin

internal/rcon/rcon.go

@@ -0,0 +1,230 @@
+// 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
+}