job-search/internal/websocket/tcpserver/server_test.go

package tcpserver

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"io"
	"log/slog"
	"net"
	"strings"
	"syscall"
	"testing"
	"time"

	"git.avsdev.uk/Go/job-search/internal/config"
	"git.avsdev.uk/Go/job-search/internal/websocket"
)

var discardLog = slog.New(slog.DiscardHandler)

// testConf is a server configuration for tests, changed by mutate if it is
// not nil.
func testConf(mutate func(*config.JSCServer)) *config.JSCServer {
	c := &config.JSCServer{
		Host: "127.0.0.1",
		Port: 0,
		WebSocket: config.JSCServerWebSocket{
			WriteTimeout: time.Second,
		},
		Timeouts: config.JSCServerTimeouts{
			ShutdownTimeout: time.Second,
		},
	}
	if mutate != nil {
		mutate(c)
	}
	return c
}

// newServer returns a server for conf with its own hub.
func newServer(conf *config.JSCServer) (*Server, *websocket.Hub) {
	h := websocket.NewHub(conf.WebSocket, websocket.WithLogger(discardLog))
	return New(conf, h, WithLogger(discardLog)), h
}

// startServer serves on a local port, returning the server's hub, its
// address, a function to cancel it and a channel with serve's result.
func startServer(
	t *testing.T, conf *config.JSCServer,
) (*websocket.Hub, string, context.CancelFunc, chan error) {
	t.Helper()
	s, h := newServer(conf)
	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)
	done := make(chan error, 1)
	go func() {
		done <- s.serve(ctx, ln)
	}()

	return h, ln.Addr().String(), cancel, done
}

// tcpClient is a plain TCP test client reading newline-delimited envelopes.
type tcpClient struct {
	t    *testing.T
	conn net.Conn
	r    *bufio.Reader
}

func dial(t *testing.T, addr string) *tcpClient {
	t.Helper()
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() {
		conn.Close()
	})
	return &tcpClient{
		t:    t,
		conn: conn,
		r:    bufio.NewReader(conn),
	}
}

// read returns the next envelope, failing the test after a second.
func (c *tcpClient) read() websocket.Envelope {
	c.t.Helper()
	c.conn.SetReadDeadline(time.Now().Add(time.Second))
	line, err := c.r.ReadString('\n')
	if err != nil {
		c.t.Fatalf("read: %v (got %q)", err, line)
	}

	var env websocket.Envelope
	if err := json.Unmarshal([]byte(line), &env); err != nil {
		c.t.Fatalf("bad envelope %q: %v", line, err)
	}

	return env
}

// waitClosed checks the server closes the connection, skipping any lines
// sent first. A reset counts as closed.
func (c *tcpClient) waitClosed() {
	c.t.Helper()
	c.conn.SetReadDeadline(time.Now().Add(time.Second))
	if _, err := io.ReadAll(c.r); err != nil &&
		!errors.Is(err, syscall.ECONNRESET) {
		c.t.Fatalf("connection not closed: %v", err)
	}
}

// Clients get pushed messages; cancelling the server's context tells them it
// is shutting down, closes them, stops accepting and returns nil.
func TestServerShutdown(t *testing.T) {
	h, addr, cancel, done := startServer(t, testConf(nil))
	c := dial(t, addr)
	if e := c.read(); e.Type != websocket.TypeWelcome {
		t.Fatalf("want welcome, got %+v", e)
	}
	for h.Count() != 1 {
		time.Sleep(time.Millisecond)
	}
	h.Push("item.created", map[string]string{"id": "x"})
	if e := c.read(); e.Type != "item.created" {
		t.Fatalf("want item.created, got %+v", e)
	}

	cancel()
	if e := c.read(); e.Type != websocket.TypeError ||
		e.Error != "server shutting down" {
		t.Fatalf("want shutdown error, got %+v", e)
	}
	c.waitClosed()
	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("serve = %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("serve did not return")
	}
	if conn, err := net.Dial("tcp", addr); err == nil {
		conn.Close()
		t.Error("listener still accepting")
	}
}

// Over websocket.connection_limit, a client is told it is rate limited and
// closed; other clients are unaffected.
func TestServerConnectionLimit(t *testing.T) {
	conf := testConf(func(c *config.JSCServer) {
		c.WebSocket.ConnectionLimit = config.JSCRateLimitRule{
			Enabled:   true,
			PerSecond: 1,
			Burst:     1,
		}
		c.RateLimits.IdleTTL = time.Minute
	})
	_, addr, cancel, _ := startServer(t, conf)
	defer cancel()

	first := dial(t, addr)
	if e := first.read(); e.Type != websocket.TypeWelcome {
		t.Fatalf("first: want welcome, got %+v", e)
	}
	second := dial(t, addr)
	if e := second.read(); e.Type != websocket.TypeError ||
		e.Error != "rate limited" {
		t.Fatalf("second: want rate limited, got %+v", e)
	}
	second.waitClosed()

	first.conn.Write([]byte(`{"type":"ping","id":"1"}` + "\n"))
	if e := first.read(); e.Type != websocket.TypePong {
		t.Fatalf("first: want pong, got %+v", e)
	}
}

// With the connection limit off (or no rate), there is no limiter.
func TestServerConnectionLimitOff(t *testing.T) {
	for _, rule := range []config.JSCRateLimitRule{
		{Enabled: false, PerSecond: 1, Burst: 1},
		{Enabled: true, PerSecond: 0, Burst: 1},
	} {
		s, _ := newServer(testConf(func(c *config.JSCServer) {
			c.WebSocket.ConnectionLimit = rule
		}))
		if l := s.connectionLimiter(context.Background()); l != nil {
			t.Errorf("rule %+v: got a limiter", rule)
		}
	}
}

func TestServerRun(t *testing.T) {
	s, _ := newServer(testConf(nil))
	ctx, cancel := context.WithCancel(context.Background())
	time.AfterFunc(20*time.Millisecond, cancel)
	if err := s.Run(ctx); err != nil {
		t.Fatalf("Run = %v", err)
	}
}

func TestServerRunListenError(t *testing.T) {
	s, _ := newServer(testConf(func(c *config.JSCServer) {
		c.Port = -1
	}))
	if err := s.Run(context.Background()); err == nil ||
		!strings.Contains(err.Error(), "listen") {
		t.Fatalf("Run = %v; want a listen error", err)
	}
}

// flakyListener fails Accept with a temporary error a number of times, then
// as if closed.
type flakyListener struct {
	net.Listener
	fails int
}

func (l *flakyListener) Accept() (net.Conn, error) {
	if l.fails > 0 {
		l.fails--
		return nil, errors.New("too many open files")
	}
	return nil, net.ErrClosed
}

// Accept errors other than a closed listener are retried after a back off.
func TestServerAcceptRetries(t *testing.T) {
	s, _ := newServer(testConf(nil))
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer ln.Close()

	fl := &flakyListener{
		Listener: ln,
		fails:    3,
	}

	err = s.accept(context.Background(), context.Background(), fl)
	if !errors.Is(err, net.ErrClosed) || fl.fails != 0 {
		t.Fatalf(
			"accept = %v with %d failures left; want net.ErrClosed", err, fl.fails,
		)
	}
}

// Cancelling ctx stops accept while it is backing off.
func TestServerAcceptCancelledInBackOff(t *testing.T) {
	s, _ := newServer(testConf(nil))
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer ln.Close()

	ctx, cancel := context.WithCancel(context.Background())
	fl := &flakyListener{
		Listener: ln,
		fails:    1 << 30,
	}

	time.AfterFunc(20*time.Millisecond, cancel)
	if err := s.accept(ctx, ctx, fl); !errors.Is(err, context.Canceled) {
		t.Fatalf("accept = %v; want context.Canceled", err)
	}
}

// A listener closed from elsewhere is an error, after a graceful shutdown.
func TestServerListenerClosed(t *testing.T) {
	s, _ := newServer(testConf(nil))
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	ln.Close()
	if err := s.serve(context.Background(), ln); !errors.Is(err, net.ErrClosed) {
		t.Fatalf("serve = %v; want net.ErrClosed", err)
	}
}

func TestRemoteIP(t *testing.T) {
	a, b := net.Pipe()
	defer a.Close()
	defer b.Close()
	if ip := remoteIP(a); ip != "" {
		t.Errorf("remoteIP(pipe) = %q; want empty", ip)
	}
}