job-search/internal/websocket/tcp_test.go

package websocket

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

	ws "github.com/coder/websocket"

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

var discardLog = slog.New(slog.DiscardHandler)

// tcpHub starts a hub serving plain TCP clients, via ServeTCP, on a local
// port and returns it with the address. It is shut down when the test ends.
func tcpHub(
	t *testing.T, mutate func(*config.JSCServerWebSocket),
) (*Hub, string) {
	t.Helper()
	conf := config.JSCServerWebSocket{
		MaxMessageSize: 1 << 10,
		WriteTimeout:   time.Second,
	}
	if mutate != nil {
		mutate(&conf)
	}

	h := NewHub(conf, WithLogger(discardLog))
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	go func() {
		for {
			conn, err := ln.Accept()
			if err != nil {
				return
			}
			go h.ServeTCP(ctx, conn)
		}
	}()

	t.Cleanup(func() {
		ln.Close()
		sctx, scancel := context.WithTimeout(context.Background(), time.Second)
		defer scancel()
		h.Shutdown(sctx)
		cancel()
	})

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

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

func dialTCP(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),
	}
}

func (c *tcpClient) send(line string) {
	c.t.Helper()
	if _, err := io.WriteString(c.conn, line); err != nil {
		c.t.Fatal(err)
	}
}

// read returns the next envelope, failing the test after a second.
func (c *tcpClient) read() 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 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: the server closing with unread data
// from the client resets the connection.
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)
	}
}

func waitCount(t *testing.T, h *Hub, want int) {
	t.Helper()
	for deadline := time.Now().Add(time.Second); h.Count() != want; {
		if time.Now().After(deadline) {
			t.Fatalf("Count = %d; want %d", h.Count(), want)
		}
		time.Sleep(5 * time.Millisecond)
	}
}

// ---------- ServeTCP ----------

func TestTCPMessages(t *testing.T) {
	h, addr := tcpHub(t, nil)
	c := dialTCP(t, addr)

	if e := c.read(); e.Type != TypeWelcome {
		t.Fatalf("want welcome, got %+v", e)
	}
	waitCount(t, h, 1)

	// Blank lines are skipped and a trailing \r is ignored.
	c.send("\n  \r\n{\"type\":\"ping\",\"id\":\"1\"}\r\n")
	if e := c.read(); e.Type != TypePong || e.ID != "1" {
		t.Fatalf("want pong 1, got %+v", e)
	}
	c.send(`{"type":"echo","id":"2","payload":{"a":1}}` + "\n")
	if e := c.read(); e.Type != TypeEcho || e.ID != "2" ||
		string(e.Payload) != `{"a":1}` {
		t.Fatalf("want echo 2, got %+v", e)
	}
	c.send("not json\n")
	if e := c.read(); e.Type != TypeError || e.Error != "invalid JSON envelope" {
		t.Fatalf("want invalid JSON error, got %+v", e)
	}
	c.send(`{"type":"bogus","id":"3"}` + "\n")
	if e := c.read(); e.Type != TypeError || e.ID != "3" {
		t.Fatalf("want unknown type error, got %+v", e)
	}

	// Pushed messages reach TCP clients.
	h.Push("item.created", map[string]string{"id": "x"})
	if e := c.read(); e.Type != "item.created" ||
		string(e.Payload) != `{"id":"x"}` {
		t.Fatalf("want item.created, got %+v", e)
	}
}

// With no write timeout, messages are still sent: a zero timeout must not
// mean one that has already expired.
func TestTCPNoWriteTimeout(t *testing.T) {
	_, addr := tcpHub(t, func(c *config.JSCServerWebSocket) {
		c.WriteTimeout = 0
	})

	c := dialTCP(t, addr)
	if e := c.read(); e.Type != TypeWelcome {
		t.Fatalf("want welcome, got %+v", e)
	}

	c.send(`{"type":"ping","id":"1"}` + "\n")
	if e := c.read(); e.Type != TypePong {
		t.Fatalf("want pong, got %+v", e)
	}
}

func TestTCPClientDisconnect(t *testing.T) {
	h, addr := tcpHub(t, nil)
	c := dialTCP(t, addr)
	c.read() // welcome
	waitCount(t, h, 1)

	c.conn.Close()
	waitCount(t, h, 0)
}

// A line longer than max_message_size drops the connection.
func TestTCPLineTooLong(t *testing.T) {
	h, addr := tcpHub(t, func(c *config.JSCServerWebSocket) {
		c.MaxMessageSize = 64
	})
	c := dialTCP(t, addr)
	c.read() // welcome

	c.send(`{"type":"echo","payload":"` + strings.Repeat("x", 100) + `"}` + "\n")
	c.waitClosed()
	waitCount(t, h, 0)
}

// Shutdown tells TCP clients why they are being closed, then closes them,
// and refuses new ones.
func TestTCPShutdown(t *testing.T) {
	h, addr := tcpHub(t, nil)
	c := dialTCP(t, addr)
	c.read() // welcome
	waitCount(t, h, 1)

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	if err := h.Shutdown(ctx); err != nil {
		t.Fatalf("Shutdown: %v", err)
	}
	if e := c.read(); e.Type != TypeError || e.Error != "server shutting down" {
		t.Fatalf("want shutdown error, got %+v", e)
	}
	c.waitClosed()

	late := dialTCP(t, addr)
	if e := late.read(); e.Type != TypeError ||
		e.Error != "server shutting down" {
		t.Fatalf("late client: want shutdown error, got %+v", e)
	}
	late.waitClosed()
}

// Ping writes a blank line to an idle client. The client's messages, which
// also count as activity, can arrive while pings run.
func TestTCPPing(t *testing.T) {
	_, addr := tcpHub(t, func(c *config.JSCServerWebSocket) {
		c.PingInterval = 20 * time.Millisecond
		c.PongTimeout = time.Second
	})
	c := dialTCP(t, addr)
	c.read() // welcome

	c.conn.SetReadDeadline(time.Now().Add(time.Second))
	line, err := c.r.ReadString('\n')
	if err != nil || line != "\n" {
		t.Fatalf("want a blank ping line, got %q, %v", line, err)
	}

	for i := range 5 {
		c.send(`{"type":"ping"}` + "\n")
		for {
			c.conn.SetReadDeadline(time.Now().Add(time.Second))
			line, err := c.r.ReadString('\n')
			if err != nil {
				t.Fatalf("message %d: %v", i, err)
			}
			if line != "\n" {
				break
			}
		}
		time.Sleep(10 * time.Millisecond)
	}
}

// A client that doesn't read is dropped as a slow consumer once its queue is
// full. net.Pipe has no buffer, so writes to it block until the peer reads.
func TestTCPSlowConsumer(t *testing.T) {
	h := NewHub(
		config.JSCServerWebSocket{
			WriteTimeout: 50 * time.Millisecond,
		},
		WithLogger(discardLog),
	)
	server, peer := net.Pipe()
	t.Cleanup(func() {
		server.Close()
		peer.Close()
	})
	done := make(chan struct{})
	go func() {
		h.ServeTCP(context.Background(), server)
		close(done)
	}()
	waitCount(t, h, 1)

	for range clientSendBuffer + 1 {
		h.Push("item.created", nil)
	}
	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("slow client not dropped")
	}
	waitCount(t, h, 0)
}

// A pong timeout of 0 means no timeout, not one that has already expired:
// idle clients are pinged and kept.
func TestTCPNoPongTimeout(t *testing.T) {
	h, addr := tcpHub(t, func(c *config.JSCServerWebSocket) {
		c.PingInterval = 10 * time.Millisecond
		c.PongTimeout = 0
	})
	c := dialTCP(t, addr)
	c.read() // welcome

	for range 3 {
		c.conn.SetReadDeadline(time.Now().Add(time.Second))
		if line, err := c.r.ReadString('\n'); err != nil || line != "\n" {
			t.Fatalf("want a blank ping line, got %q, %v", line, err)
		}
	}

	waitCount(t, h, 1)
}

// A max_message_size of 0 uses the default limit.
func TestTCPDefaultMessageSize(t *testing.T) {
	_, addr := tcpHub(t, func(c *config.JSCServerWebSocket) {
		c.MaxMessageSize = 0
	})
	c := dialTCP(t, addr)
	c.read() // welcome

	payload := strings.Repeat("x", 4096)
	c.send(`{"type":"echo","id":"1","payload":"` + payload + `"}` + "\n")
	if e := c.read(); e.Type != TypeEcho || e.ID != "1" {
		t.Fatalf("want echo, got %+v", e)
	}
}

// ---------- tcpConn ----------

// pipeConn returns a tcpConn on one end of an in-memory pipe, and the other
// end. Writes on a pipe block until the other end reads.
func pipeConn(t *testing.T, closeTimeout time.Duration) (*tcpConn, net.Conn) {
	t.Helper()
	server, peer := net.Pipe()
	t.Cleanup(func() {
		server.Close()
		peer.Close()
	})
	return newTCPConn(server, 0, closeTimeout), peer
}

func TestTCPConnRead(t *testing.T) {
	tc, peer := pipeConn(t, time.Second)
	go func() {
		io.WriteString(peer, "\r\n{\"a\":1}\r\n")
		peer.Close()
	}()

	typ, data, err := tc.Read(context.Background())
	if err != nil || typ != ws.MessageText || string(data) != `{"a":1}` {
		t.Fatalf("Read = %v, %q, %v", typ, data, err)
	}

	// The peer closing is a normal closure.
	_, _, err = tc.Read(context.Background())
	if ws.CloseStatus(err) != ws.StatusNormalClosure {
		t.Fatalf("Read at EOF = %v; want a normal closure", err)
	}
}

func TestTCPConnReadCancelled(t *testing.T) {
	tc, _ := pipeConn(t, time.Second)
	ctx, cancel := context.WithCancel(context.Background())
	time.AfterFunc(10*time.Millisecond, cancel)
	if _, _, err := tc.Read(ctx); !errors.Is(err, context.Canceled) {
		t.Fatalf("Read = %v; want context.Canceled", err)
	}
}

func TestTCPConnWriteBinary(t *testing.T) {
	tc, _ := pipeConn(t, time.Second)
	err := tc.Write(context.Background(), ws.MessageBinary, []byte("x"))
	if !errors.Is(err, errUnsupportedData) {
		t.Fatalf("Write(binary) = %v; want errUnsupportedData", err)
	}
}

// Close sends its reason and later reads report its status; a second Close
// or a CloseNow does nothing more.
func TestTCPConnClose(t *testing.T) {
	tc, peer := pipeConn(t, time.Second)
	got := make(chan string, 1)
	go func() {
		line, _ := bufio.NewReader(peer).ReadString('\n')
		got <- line
	}()

	if err := tc.Close(ws.StatusGoingAway, "bye"); err != nil {
		t.Fatalf("Close: %v", err)
	}
	if line := <-got; line != `{"type":"error","error":"bye"}`+"\n" {
		t.Errorf("peer got %q", line)
	}

	_, _, err := tc.Read(context.Background())
	var ce ws.CloseError
	if !errors.As(err, &ce) || ce.Code != ws.StatusGoingAway ||
		ce.Reason != "bye" {
		t.Errorf("Read after Close = %v; want going away, bye", err)
	}

	if err := tc.Close(ws.StatusNormalClosure, "again"); err != nil {
		t.Errorf("second Close: %v", err)
	}
	if err := tc.CloseNow(); err != nil {
		t.Errorf("CloseNow after Close: %v", err)
	}
}

// CloseNow doesn't wait for a Close stuck sending its message to a peer that
// isn't reading, and cuts it short.
func TestTCPConnCloseNowDuringClose(t *testing.T) {
	tc, _ := pipeConn(t, 0) // no timeout: Close would wait for ever
	closed := make(chan error, 1)
	go func() {
		closed <- tc.Close(ws.StatusGoingAway, "bye")
	}()
	for !tc.isClosed() {
		time.Sleep(time.Millisecond)
	}

	done := make(chan error, 1)
	go func() {
		done <- tc.CloseNow()
	}()
	select {
	case err := <-done:
		if err != nil {
			t.Errorf("CloseNow: %v", err)
		}
	case <-time.After(time.Second):
		t.Fatal("CloseNow waited for Close")
	}
	select {
	case <-closed:
	case <-time.After(time.Second):
		t.Fatal("Close not cut short by CloseNow")
	}
}

// Close gives up sending its message after the close timeout.
func TestTCPConnCloseTimeout(t *testing.T) {
	tc, _ := pipeConn(t, 20*time.Millisecond)
	done := make(chan error, 1)
	go func() {
		done <- tc.Close(ws.StatusGoingAway, "bye")
	}()
	select {
	case err := <-done:
		if err != nil {
			t.Errorf("Close: %v", err)
		}
	case <-time.After(time.Second):
		t.Fatal("Close ignored its timeout")
	}
}