job-search/internal/server/run_test.go

package server

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"io"
	"log/slog"
	"net"
	"net/http"
	"net/http/cookiejar"
	"net/http/httptest"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/coder/websocket"

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

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

func startRun(
	t *testing.T, mutate func(*config.JSCServer),
) (addr string, stop func() error) {
	t.Helper()
	cfg := defaultConf(t)
	addr = freeAddr(t)
	host, port, _ := net.SplitHostPort(addr)
	cfg.Host = host
	cfg.Port, _ = strconv.Atoi(port)
	cfg.Timeouts.ShutdownTimeout = 3 * time.Second
	if mutate != nil {
		mutate(cfg)
	}
	s, err := New(cfg, WithLogger(slog.New(slog.DiscardHandler)))
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() {
		done <- s.Run(ctx)
	}()

	deadline := time.Now().Add(3 * time.Second)
	for {
		c, err := net.Dial("tcp", addr)
		if err == nil {
			c.Close()
			break
		}
		if time.Now().After(deadline) {
			t.Fatal("server did not start")
		}
		time.Sleep(20 * time.Millisecond)
	}
	return addr, func() error {
		cancel()
		select {
		case err := <-done:
			return err
		case <-time.After(5 * time.Second):
			return errors.New("shutdown timed out")
		}
	}
}

// A PROXY v1 header from a trusted source sets the client IP.
func TestProxyProtocol(t *testing.T) {
	addr, stop := startRun(t, func(c *config.JSCServer) {
		c.Proxy.ProxyProtocol.Enabled = true
	})
	defer stop()

	conn, err := net.Dial("tcp", addr)
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	io.WriteString(conn, "PROXY TCP4 198.51.100.7 10.0.0.1 51234 8080\r\n"+
		"GET /api/health HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n")
	res, err := http.ReadResponse(bufio.NewReader(conn), nil)
	if err != nil {
		t.Fatal(err)
	}
	defer res.Body.Close()
	var out map[string]map[string]any
	json.NewDecoder(res.Body).Decode(&out)
	if out["client"]["ip"] != "198.51.100.7" {
		t.Fatalf("PROXY header not applied: %v", out["client"])
	}

	// Plain HTTP (no PROXY header) still works when not required.
	res2, err := http.Get("http://" + addr + "/api/health")
	if err != nil {
		t.Fatal(err)
	}
	res2.Body.Close()
	if res2.StatusCode != 200 {
		t.Fatalf("plain request: %d", res2.StatusCode)
	}
}

// Shutdown closes live WebSockets with 1001 and Run returns cleanly.
func TestGracefulShutdownClosesWebSockets(t *testing.T) {
	addr, stop := startRun(t, nil)
	jar, _ := cookiejar.New(nil)
	hc := &http.Client{
		Jar: jar,
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	url := "ws://" + addr + "/ws"
	conn, _, err := websocket.Dial(ctx, url, &websocket.DialOptions{
		HTTPClient: hc,
	})
	if err != nil {
		t.Fatal(err)
	}
	if _, _, err := conn.Read(ctx); err != nil { // welcome
		t.Fatal(err)
	}

	readErr := make(chan error, 1)
	go func() {
		_, _, err := conn.Read(ctx)
		readErr <- err
	}()

	if err := stop(); err != nil {
		t.Fatalf("Run returned error: %v", err)
	}
	err = <-readErr
	if websocket.CloseStatus(err) != websocket.StatusGoingAway {
		t.Fatalf("want close 1001, got %v", err)
	}
	if !strings.Contains(err.Error(), "shutting down") {
		t.Fatalf("unexpected close reason: %v", err)
	}
}

// The HTTP listener redirects to the same host and path over HTTPS, on the
// TLS listener's port (left out when it is 443).
func TestHTTPSRedirect(t *testing.T) {
	tests := []struct {
		tlsAddr, host, want string
	}{
		{"127.0.0.1:9393", "example.com:8080", "https://example.com:9393/x?y=1"},
		{"127.0.0.1:9393", "example.com", "https://example.com:9393/x?y=1"},
		{"127.0.0.1:443", "example.com:80", "https://example.com/x?y=1"},
		{"[::1]:9393", "[2001:db8::1]:8080", "https://[2001:db8::1]:9393/x?y=1"},
		{"[::1]:443", "[2001:db8::1]:80", "https://[2001:db8::1]/x?y=1"},
	}
	for _, tt := range tests {
		req := httptest.NewRequest("GET", "http://"+tt.host+"/x?y=1", nil)
		rec := httptest.NewRecorder()
		httpsRedirect(tt.tlsAddr).ServeHTTP(rec, req)
		location := rec.Header().Get("Location")
		if rec.Code != http.StatusPermanentRedirect || location != tt.want {
			t.Errorf(
				"%s via %s: %d %q; want 308 %q",
				tt.host,
				tt.tlsAddr,
				rec.Code,
				location,
				tt.want,
			)
		}
	}
}