Megalodon/internal/rcon/rcon.go

// Package rcon is a client for the Source RCON protocol as ARK: Survival Evolved speaks it. One
// Client holds one TCP connection, signs in on first use, reconnects when the server has dropped
// it, and runs one command at a time.
package rcon

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

const (
	typeResponse int32 = 0
	typeExec     int32 = 2 // a command going out, and the sign-in answer coming back
	typeAuth     int32 = 3

	// A response body this long may continue in the next packet.
	fullBody = 4000
	// ARK never sends anything near this; it only guards against a garbled length.
	maxPacket = 1 << 20
	// How long to wait for a continuation packet before deciding the response is complete.
	continuation = 250 * time.Millisecond
)

// ErrAuth is returned when the server rejects the password.
var ErrAuth = errors.New("rcon: wrong password")

// Client talks to one server. It is safe for concurrent use; commands are sent one at a time.
type Client struct {
	Addr     string        // host:port
	Password string        //
	Timeout  time.Duration // per dial and per command, default 5s

	mu   sync.Mutex
	conn net.Conn
	id   int32
}

// New returns a client that has not connected yet.
func New(addr, password string) *Client {
	return &Client{Addr: addr, Password: password, Timeout: 5 * time.Second}
}

// Close drops the connection. The client can be used again afterwards.
func (c *Client) Close() error {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.closeLocked()
}

func (c *Client) closeLocked() error {
	if c.conn == nil {
		return nil
	}
	err := c.conn.Close()
	c.conn = nil
	return err
}

// Exec runs one command and returns the server's reply. A connection that went stale while idle
// (the server restarted, a NAT forgot it) is replaced once, transparently.
func (c *Client) Exec(cmd string) (string, error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	fresh := false
	if c.conn == nil {
		if err := c.connect(); err != nil {
			return "", err
		}
		fresh = true
	}
	out, got, err := c.exec(cmd)
	if err == nil {
		return out, nil
	}
	c.closeLocked()
	if fresh || got {
		return "", err
	}
	// Nothing came back at all, so this was most likely a dead connection: try once on a new one.
	if err := c.connect(); err != nil {
		return "", err
	}
	out, _, err = c.exec(cmd)
	if err != nil {
		c.closeLocked()
		return "", err
	}
	return out, nil
}

func (c *Client) timeout() time.Duration {
	if c.Timeout > 0 {
		return c.Timeout
	}
	return 5 * time.Second
}

func (c *Client) connect() error {
	conn, err := net.DialTimeout("tcp", c.Addr, c.timeout())
	if err != nil {
		return err
	}
	conn.SetDeadline(time.Now().Add(c.timeout()))
	if err := writePacket(conn, 1, typeAuth, c.Password); err != nil {
		conn.Close()
		return err
	}
	// Servers may send an empty response before the sign-in answer.
	for i := 0; i < 4; i++ {
		p, err := readPacket(conn)
		if err != nil {
			conn.Close()
			return fmt.Errorf("rcon: sign-in: %w", err)
		}
		if p.typ != typeExec {
			continue
		}
		if p.id == -1 {
			conn.Close()
			return ErrAuth
		}
		c.conn = conn
		return nil
	}
	conn.Close()
	return errors.New("rcon: the server never answered the sign-in")
}

// exec reports whether any packet arrived, so Exec knows if a retry is safe.
func (c *Client) exec(cmd string) (body string, got bool, err error) {
	c.id++
	if c.id < 2 {
		c.id = 2
	}
	id := c.id
	c.conn.SetDeadline(time.Now().Add(c.timeout()))
	if err := writePacket(c.conn, id, typeExec, cmd); err != nil {
		return "", false, err
	}
	var sb strings.Builder
	for {
		p, err := readPacket(c.conn)
		if err != nil {
			return "", got, err
		}
		got = true
		if p.id != id || p.typ != typeResponse {
			continue // left over from an earlier command
		}
		sb.WriteString(p.body)
		if len(p.body) < fullBody {
			return sb.String(), true, nil
		}
		break
	}
	// A full-size packet: the reply may go on. There is no end marker, so wait briefly for more.
	for {
		c.conn.SetReadDeadline(time.Now().Add(continuation))
		p, err := readPacket(c.conn)
		if err != nil {
			var ne net.Error
			if errors.As(err, &ne) && ne.Timeout() {
				break
			}
			return "", true, err
		}
		if p.id != id || p.typ != typeResponse {
			continue
		}
		sb.WriteString(p.body)
		if len(p.body) < fullBody {
			break
		}
	}
	return sb.String(), true, nil
}

type packet struct {
	id, typ int32
	body    string
}

func writePacket(w io.Writer, id, typ int32, body string) error {
	buf := make([]byte, 4+4+4+len(body)+2)
	binary.LittleEndian.PutUint32(buf[0:], uint32(4+4+len(body)+2))
	binary.LittleEndian.PutUint32(buf[4:], uint32(id))
	binary.LittleEndian.PutUint32(buf[8:], uint32(typ))
	copy(buf[12:], body)
	_, err := w.Write(buf)
	return err
}

func readPacket(r io.Reader) (packet, error) {
	var hdr [4]byte
	if _, err := io.ReadFull(r, hdr[:]); err != nil {
		return packet{}, err
	}
	size := int(int32(binary.LittleEndian.Uint32(hdr[:])))
	if size < 10 || size > maxPacket {
		return packet{}, fmt.Errorf("rcon: bad packet length %d", size)
	}
	buf := make([]byte, size)
	if _, err := io.ReadFull(r, buf); err != nil {
		return packet{}, err
	}
	return packet{
		id:   int32(binary.LittleEndian.Uint32(buf[0:])),
		typ:  int32(binary.LittleEndian.Uint32(buf[4:])),
		body: strings.ToValidUTF8(string(buf[8:size-2]), "�"),
	}, nil
}