Spinner/internal/proxy/proxy_test.go

package proxy

import (
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"io"
	"log/slog"
	"net"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"git.veillax.com/Veillax/spinner/internal/config"
	"git.veillax.com/Veillax/spinner/internal/health"
	"git.veillax.com/Veillax/spinner/internal/protocol"
)

var quiet = slog.New(slog.NewTextHandler(io.Discard, nil))

// backend is a fake Minecraft server: status pings get a canned response, and
// login connections are answered with "hello from <name>" plus the PROXY header
// it saw, if any.
func backend(t *testing.T, name string, expectProxy bool) string {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { ln.Close() })
	go func() {
		for {
			c, err := ln.Accept()
			if err != nil {
				return
			}
			go func() {
				defer c.Close()
				c.SetDeadline(time.Now().Add(5 * time.Second))
				seen := "none"
				if expectProxy {
					hdr := make([]byte, 16)
					if _, err := io.ReadFull(c, hdr); err != nil {
						return
					}
					rest := make([]byte, binary.BigEndian.Uint16(hdr[14:]))
					io.ReadFull(c, rest)
					if bytes.Equal(hdr[:12], v2sig) && hdr[12] == 0x21 && hdr[13] == 0x11 {
						seen = fmt.Sprintf("%d.%d.%d.%d", rest[0], rest[1], rest[2], rest[3])
					}
				}
				hs, _, err := protocol.ReadHandshake(c)
				if err != nil {
					return
				}
				if hs.IsStatus() {
					protocol.ServeStatus(c, `{"version":{"name":"fake","protocol":767},"players":{"online":2,"max":10},"description":{"text":"`+name+`"}}`)
					return
				}
				fmt.Fprintf(c, "hello from %s proxy=%s host=%s", name, seen, hs.Host)
			}()
		}
	}()
	return ln.Addr().String()
}

func start(t *testing.T, cfgJSON string) (addr string, p *Proxy, h *health.Tracker) {
	t.Helper()
	path := filepath.Join(t.TempDir(), "spinner.json")
	if err := os.WriteFile(path, []byte(cfgJSON), 0o600); err != nil {
		t.Fatal(err)
	}
	m, err := config.NewManager(path, quiet)
	if err != nil {
		t.Fatal(err)
	}
	h = health.New(quiet)
	p = New(m, h, quiet)
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)
	go p.Serve(ctx, ln)
	return ln.Addr().String(), p, h
}

func login(t *testing.T, addr, host string) string {
	t.Helper()
	c, err := net.Dial("tcp", addr)
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	c.SetDeadline(time.Now().Add(5 * time.Second))
	c.Write(protocol.BuildHandshake(767, host, 25565, 2))
	b, _ := io.ReadAll(c)
	return string(b)
}

func status(t *testing.T, addr, host string) string {
	t.Helper()
	c, err := net.Dial("tcp", addr)
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	js, err := protocol.QueryStatus(c, host, 25565, 5*time.Second)
	if err != nil {
		t.Fatal(err)
	}
	return js
}

func TestRoutesByHostname(t *testing.T) {
	a, b := backend(t, "alpha", false), backend(t, "beta", false)
	addr, _, _ := start(t, fmt.Sprintf(`{"servers":[
		{"name":"alpha","domains":["alpha.test"],"addresses":[%q]},
		{"name":"beta","domains":["*.beta.test"],"addresses":[%q]}]}`, a, b))

	if got := login(t, addr, "alpha.test"); !strings.HasPrefix(got, "hello from alpha") {
		t.Errorf("alpha: %q", got)
	}
	if got := login(t, addr, "x.beta.test"); !strings.HasPrefix(got, "hello from beta") || !strings.Contains(got, "host=x.beta.test") {
		t.Errorf("beta: %q", got)
	}
}

func TestProxyProtocolHeader(t *testing.T) {
	a := backend(t, "pp", true)
	addr, _, _ := start(t, fmt.Sprintf(`{"servers":[{"name":"pp","domains":["pp.test"],"addresses":[%q],"send_proxy_protocol":true}]}`, a))
	if got := login(t, addr, "pp.test"); !strings.Contains(got, "proxy=127.0.0.1") {
		t.Errorf("backend did not see client address in PROXY header: %q", got)
	}
}

func TestFallsBackToSecondAddress(t *testing.T) {
	good := backend(t, "good", false)
	addr, _, _ := start(t, fmt.Sprintf(`{"servers":[{"name":"s","domains":["s.test"],"addresses":["127.0.0.1:1",%q]}]}`, good))
	if got := login(t, addr, "s.test"); !strings.HasPrefix(got, "hello from good") {
		t.Errorf("got %q", got)
	}
}

func TestOfflineBackend(t *testing.T) {
	addr, p, _ := start(t, `{"servers":[{"name":"dead","domains":["dead.test"],"addresses":["127.0.0.1:1"],"offline_motd":"napping"}]}`)
	if js := status(t, addr, "dead.test"); !strings.Contains(js, "napping") {
		t.Errorf("status: %s", js)
	}
	if got := login(t, addr, "dead.test"); !strings.Contains(got, "napping") {
		t.Errorf("login: %q", got)
	}
	if f := p.Stats.Snapshot().Servers[0].Failures; f < 2 {
		t.Errorf("failures = %d", f)
	}
}

func TestUnknownHost(t *testing.T) {
	addr, p, _ := start(t, `{"unknown_motd":"nobody home","servers":[]}`)
	if js := status(t, addr, "who.test"); !strings.Contains(js, "nobody home") {
		t.Errorf("status: %s", js)
	}
	if got := login(t, addr, "who.test"); !strings.Contains(got, "nobody home") {
		t.Errorf("login: %q", got)
	}
	if p.Stats.Snapshot().Unknown != 2 {
		t.Errorf("unknown = %d", p.Stats.Snapshot().Unknown)
	}
}

func TestStatusServedFromHealthCache(t *testing.T) {
	a := backend(t, "live", false)
	addr, p, h := start(t, fmt.Sprintf(`{"probe_interval_secs":60,"servers":[{"name":"live","domains":["live.test"],"addresses":[%q]}]}`, a))
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go h.Run(ctx, p.cfg)
	deadline := time.Now().Add(3 * time.Second)
	for {
		if _, ok := h.Cached("live", time.Minute); ok {
			break
		}
		if time.Now().After(deadline) {
			t.Fatal("probe never succeeded")
		}
		time.Sleep(20 * time.Millisecond)
	}
	if js := status(t, addr, "live.test"); !strings.Contains(js, `"online":2`) {
		t.Errorf("status: %s", js)
	}
	if snap := h.Snapshot(); len(snap) != 1 || !snap[0].Up || snap[0].Online != 2 || snap[0].Version != "fake" {
		t.Errorf("snapshot: %+v", snap)
	}
}

func TestPerIPLimit(t *testing.T) {
	addr, p, _ := start(t, `{"max_conns_per_ip":1,"servers":[]}`)
	hold, err := net.Dial("tcp", addr) // never sends a handshake, so it keeps its slot
	if err != nil {
		t.Fatal(err)
	}
	defer hold.Close()
	time.Sleep(100 * time.Millisecond)
	c, _ := net.Dial("tcp", addr)
	defer c.Close()
	c.SetDeadline(time.Now().Add(2 * time.Second))
	if _, err := io.ReadAll(c); err != nil && !strings.Contains(err.Error(), "reset") {
		t.Errorf("expected close, got %v", err)
	}
	if p.Stats.Snapshot().Rejected != 1 {
		t.Errorf("rejected = %d", p.Stats.Snapshot().Rejected)
	}
}

func TestConfigHotSwap(t *testing.T) {
	a := backend(t, "one", false)
	addr, p, _ := start(t, `{"servers":[]}`)
	if got := login(t, addr, "one.test"); strings.HasPrefix(got, "hello") {
		t.Fatalf("should be unknown: %q", got)
	}
	cfg, err := config.Parse([]byte(fmt.Sprintf(`{"servers":[{"name":"one","domains":["one.test"],"addresses":[%q]}]}`, a)))
	if err != nil {
		t.Fatal(err)
	}
	if err := p.cfg.Save(cfg); err != nil {
		t.Fatal(err)
	}
	if got := login(t, addr, "one.test"); !strings.HasPrefix(got, "hello from one") {
		t.Errorf("after save: %q", got)
	}
}