Megalodon/internal/rcon/rcon_test.go

package rcon

import (
	"errors"
	"net"
	"strings"
	"sync/atomic"
	"testing"
)

// fakeServer speaks enough of the protocol to test the client: sign-in, then one reply per
// command, split into 4000-byte packets the way real servers do.
type fakeServer struct {
	ln       net.Listener
	password string
	handler  func(cmd string) string
	conns    atomic.Int32
	dropAt   int32 // close each connection after this many commands (0 = never)
}

func serve(t *testing.T, password string, h func(string) string) *fakeServer {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	f := &fakeServer{ln: ln, password: password, handler: h}
	t.Cleanup(func() { ln.Close() })
	go func() {
		for {
			c, err := ln.Accept()
			if err != nil {
				return
			}
			f.conns.Add(1)
			go f.handle(c)
		}
	}()
	return f
}

func (f *fakeServer) handle(c net.Conn) {
	defer c.Close()
	p, err := readPacket(c)
	if err != nil || p.typ != typeAuth {
		return
	}
	writePacket(c, p.id, typeResponse, "")
	if p.body != f.password {
		writePacket(c, -1, typeExec, "")
		return
	}
	writePacket(c, p.id, typeExec, "")
	var n int32
	for {
		p, err := readPacket(c)
		if err != nil {
			return
		}
		reply := f.handler(p.body)
		for {
			chunk := reply
			if len(chunk) > fullBody {
				chunk = reply[:fullBody]
			}
			writePacket(c, p.id, typeResponse, chunk)
			reply = reply[len(chunk):]
			if reply == "" {
				break
			}
		}
		n++
		if f.dropAt > 0 && n >= f.dropAt {
			return
		}
	}
}

func (f *fakeServer) addr() string { return f.ln.Addr().String() }

func TestExec(t *testing.T) {
	f := serve(t, "pw", func(cmd string) string { return "you said " + cmd })
	c := New(f.addr(), "pw")
	defer c.Close()
	for _, cmd := range []string{"ListPlayers", "GetChat"} {
		out, err := c.Exec(cmd)
		if err != nil || out != "you said "+cmd {
			t.Fatalf("%s: got %q, %v", cmd, out, err)
		}
	}
	if f.conns.Load() != 1 {
		t.Fatalf("the connection should be reused, saw %d", f.conns.Load())
	}
}

func TestWrongPassword(t *testing.T) {
	f := serve(t, "pw", func(string) string { return "" })
	c := New(f.addr(), "nope")
	if _, err := c.Exec("ListPlayers"); !errors.Is(err, ErrAuth) {
		t.Fatalf("want ErrAuth, got %v", err)
	}
}

func TestLongReplyArrivesWhole(t *testing.T) {
	long := strings.Repeat("0123456789", 1050) // 10500 bytes: three packets
	f := serve(t, "pw", func(string) string { return long })
	c := New(f.addr(), "pw")
	defer c.Close()
	out, err := c.Exec("GetChat")
	if err != nil || out != long {
		t.Fatalf("got %d bytes, %v", len(out), err)
	}
}

func TestReconnectsAfterDrop(t *testing.T) {
	f := serve(t, "pw", func(string) string { return "ok" })
	f.dropAt = 1
	c := New(f.addr(), "pw")
	defer c.Close()
	for i := 0; i < 3; i++ {
		if out, err := c.Exec("ListPlayers"); err != nil || out != "ok" {
			t.Fatalf("run %d: got %q, %v", i, out, err)
		}
	}
	if f.conns.Load() < 3 {
		t.Fatalf("expected a new connection per command, saw %d", f.conns.Load())
	}
}

func TestUnreachable(t *testing.T) {
	ln, _ := net.Listen("tcp", "127.0.0.1:0")
	addr := ln.Addr().String()
	ln.Close()
	if _, err := New(addr, "pw").Exec("x"); err == nil {
		t.Fatal("expected a connection error")
	}
}