spark-server/internal/transport/tcp/server_test.go

package tcp

import (
	"context"
	"encoding/binary"
	"io"
	"log/slog"
	"net"
	"testing"
	"time"

	"spark-server/internal/infrastructure/id"
	"spark-server/internal/module/auth"
	"spark-server/internal/transport/tcp/protocol"
	"spark-server/internal/transport/tcp/session"
)

func TestHeartbeatWithoutAuth(t *testing.T) {
	srv := startTCP(t, &Options{})
	c := dial(t, srv.addr)
	defer c.Close()
	writeFrame(t, c, protocol.Frame{Version: protocol.CurrentVersion, Command: protocol.CmdHeartbeat, SequenceID: 3})
	got := readFrame(t, c)
	if got.Command != protocol.CmdHeartbeat || got.SequenceID != 3 {
		t.Fatalf("heartbeat reply %+v", got)
	}
}

func TestAuthGateAndGenerateMsgID(t *testing.T) {
	ids, err := id.New(1)
	if err != nil {
		t.Fatal(err)
	}
	srv := startTCP(t, &Options{
		Auth: auth.Static{Token: "secret", UserID: 42, DeviceID: 7},
		IDs:  ids,
	})
	c := dial(t, srv.addr)
	defer c.Close()

	writeFrame(t, c, protocol.Frame{
		Version: protocol.CurrentVersion, Command: protocol.CmdGenerateMsgID, RequestID: 9, SequenceID: 1,
	})
	denied := readFrame(t, c)
	eb, err := protocol.DecodeError(denied.Body)
	if err != nil {
		t.Fatal(err)
	}
	if eb.Code != protocol.CodeNotAuth {
		t.Fatalf("want not-auth, got %+v", eb)
	}

	writeFrame(t, c, protocol.Frame{
		Version: protocol.CurrentVersion, Command: protocol.CmdAuth, RequestID: 2, SequenceID: 2,
		Body: protocol.EncodeAuthRequest(protocol.AuthRequest{Token: "secret", DeviceID: 7}),
	})
	authed := readFrame(t, c)
	if authed.Command != protocol.CmdAuth || len(authed.Body) == 0 {
		t.Fatalf("auth reply %+v", authed)
	}

	writeFrame(t, c, protocol.Frame{
		Version: protocol.CurrentVersion, Command: protocol.CmdGenerateMsgID, RequestID: 3, SequenceID: 3,
	})
	msgid := readFrame(t, c)
	if msgid.Command != protocol.CmdGenerateMsgID || len(msgid.Body) == 0 {
		t.Fatalf("msgid reply %+v", msgid)
	}
}

func TestProxyProtocolThenHeartbeat(t *testing.T) {
	ids, err := id.New(1)
	if err != nil {
		t.Fatal(err)
	}
	srv := startTCP(t, &Options{
		ProxyProtocol: true,
		Auth:          auth.Static{Token: "secret", UserID: 42, DeviceID: 7},
		IDs:           ids,
	})
	c := dial(t, srv.addr)
	defer c.Close()
	if _, err := c.Write([]byte("PROXY TCP4 203.0.113.10 192.0.2.1 443 9000\r\n")); err != nil {
		t.Fatal(err)
	}
	writeFrame(t, c, protocol.Frame{
		Version: protocol.CurrentVersion, Command: protocol.CmdAuth, SequenceID: 1,
		Body: protocol.EncodeAuthRequest(protocol.AuthRequest{Token: "secret", DeviceID: 7}),
	})
	authed := readFrame(t, c)
	if authed.Command != protocol.CmdAuth {
		t.Fatalf("auth after proxy %+v", authed)
	}
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		sessions := srv.sessions.GetUserSessions(42)
		if len(sessions) == 1 && sessions[0].RealIP() == "203.0.113.10:443" {
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatalf("real ip not set, sessions=%v", dumpRealIP(srv.sessions.GetUserSessions(42)))
}

func dumpRealIP(ss []session.Session) []string {
	out := make([]string, 0, len(ss))
	for _, s := range ss {
		out = append(out, s.RealIP())
	}
	return out
}

func startTCP(t *testing.T, o *Options) *Server {
	t.Helper()
	if o == nil {
		o = &Options{}
	}
	if o.Addr == "" {
		o.Addr = freeAddr(t)
	}
	if o.IdleTimeout == 0 {
		o.IdleTimeout = time.Hour
	}
	o.KeepAlive = 30 * time.Second
	log := slog.New(slog.NewTextHandler(io.Discard, nil))
	pool := NewTaskPool(4, 32)
	t.Cleanup(pool.Close)
	s := NewServer(log, pool, session.NewManager(), o)
	errCh := make(chan error, 1)
	go func() { errCh <- s.Start() }()
	select {
	case <-s.ready:
	case err := <-errCh:
		t.Fatalf("start: %v", err)
	case <-time.After(3 * time.Second):
		t.Fatal("boot timeout")
	}
	t.Cleanup(func() {
		ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
		defer cancel()
		_ = s.Stop(ctx)
		select {
		case <-errCh:
		case <-time.After(3 * time.Second):
		}
	})
	return s
}

func freeAddr(t *testing.T) string {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	addr := ln.Addr().String()
	_ = ln.Close()
	return addr
}

func dial(t *testing.T, addr string) net.Conn {
	t.Helper()
	var last error
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		c, err := net.DialTimeout("tcp", addr, 200*time.Millisecond)
		if err == nil {
			_ = c.SetDeadline(time.Now().Add(3 * time.Second))
			return c
		}
		last = err
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatalf("dial %s: %v", addr, last)
	return nil
}

func writeFrame(t *testing.T, c net.Conn, f protocol.Frame) {
	t.Helper()
	raw, err := protocol.Encode(f)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := c.Write(raw); err != nil {
		t.Fatal(err)
	}
}

func readFrame(t *testing.T, c net.Conn) protocol.Frame {
	t.Helper()
	header := make([]byte, protocol.HeaderSize)
	if _, err := io.ReadFull(c, header); err != nil {
		t.Fatal(err)
	}
	bodyLen := int(binary.BigEndian.Uint32(header[19:23]))
	raw := make([]byte, protocol.HeaderSize+bodyLen)
	copy(raw, header)
	if bodyLen > 0 {
		if _, err := io.ReadFull(c, raw[protocol.HeaderSize:]); err != nil {
			t.Fatal(err)
		}
	}
	f, err := protocol.Decode(raw, 65536)
	if err != nil {
		t.Fatal(err)
	}
	return f
}

func TestLimiterAllow(t *testing.T) {
	l := newLimiter(1)
	if !l.Allow() {
		t.Fatal("first token")
	}
	if l.Allow() {
		t.Fatal("burst exhausted")
	}
}