job-search/internal/cli/ws/dial_test.go

package ws

import (
	"bytes"
	"context"
	"errors"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"

	coder "github.com/coder/websocket"
	"github.com/spf13/cobra"

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

func newRoot() *cobra.Command {
	opts := &cli.Options{}
	return cmdtest.NewRoot(opts, NewCommands(opts)...)
}

// wsServer serves a hub that trusts https://admin.example.com on a test HTTP
// server, and returns the hub and its ws:// URL.
func wsServer(t *testing.T) (*websocket.Hub, string) {
	t.Helper()
	hub := websocket.NewHub(
		config.JSCServerWebSocket{
			WriteTimeout: time.Second,
		},
		websocket.WithLogger(slog.New(slog.DiscardHandler)),
		websocket.WithTrustedOrigins([]string{"https://admin.example.com"}),
	)
	server := httptest.NewServer(hub)
	t.Cleanup(func() {
		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
		defer cancel()
		hub.Shutdown(ctx)
		server.Close()
	})

	return hub, "ws" + strings.TrimPrefix(server.URL, "http")
}

// ping is a message the server answers with a pong.
const ping = `{"type":"ping","id":"1"}`

// With --message, dial prints the welcome, sends the message and prints the
// reply.
func TestDialPing(t *testing.T) {
	cmdtest.Setup(t, "")
	_, url := wsServer(t)

	out := cmdtest.MustRun(
		t, newRoot, "dial", "--http-upgrade", "--url", url, "--message", ping,
	)
	cmdtest.Contains(
		t,
		out,
		"connected (101 Switching Protocols)",
		`< {"type":"welcome"`,
		`> {"type":"ping","id":"1"}`,
		`< {"type":"pong","id":"1"`,
	)
}

// --ping sends a ping and waits for its pong, then closes the connection.
func TestDialPingFlag(t *testing.T) {
	cmdtest.Setup(t, "")
	_, url := wsServer(t)

	out := cmdtest.MustRun(
		t, newRoot, "dial", "--http-upgrade", "--url", url, "--ping",
	)
	cmdtest.Contains(
		t,
		out,
		`> {"type":"ping","id":"ping"}`,
		`< {"type":"pong","id":"ping"`,
	)
}

// --ping fails if no pong arrives within --timeout.
func TestDialPingTimeout(t *testing.T) {
	cmdtest.Setup(t, "")
	// A server that sends a welcome, then never replies.
	server := httptest.NewServer(http.HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			conn, err := coder.Accept(w, r, nil)
			if err != nil {
				return
			}
			defer conn.CloseNow()
			conn.Write(r.Context(), coder.MessageText, []byte(`{"type":"welcome"}`))
			for {
				if _, _, err := conn.Read(r.Context()); err != nil {
					return
				}
			}
		},
	))
	t.Cleanup(server.Close)
	url := "ws" + strings.TrimPrefix(server.URL, "http")

	_, err := cmdtest.Run(
		t,
		newRoot,
		"",
		"dial", "--http-upgrade", "--url", url, "--ping", "--timeout", "100ms",
	)
	cmdtest.ErrContains(t, err, "no pong within 100ms")
}

// Each --message is sent in turn, and the reply to each printed.
func TestDialMessages(t *testing.T) {
	cmdtest.Setup(t, "")
	_, url := wsServer(t)

	out := cmdtest.MustRun(
		t,
		newRoot,
		"dial", "--http-upgrade", "--url", url,
		"--message", `{"type":"echo","id":"7","payload":{"a":1}}`,
		"--message", `{"type":"nonsense","id":"8"}`,
	)
	cmdtest.Contains(
		t,
		out,
		`< {"type":"echo","id":"7","payload":{"a":1}}`,
		`< {"type":"error","id":"8"`,
	)
}

// An upgrade from a foreign origin is refused; a trusted one is accepted.
func TestDialOrigin(t *testing.T) {
	cmdtest.Setup(t, "")
	_, url := wsServer(t)

	_, err := cmdtest.Run(
		t,
		newRoot,
		"",
		"dial", "--http-upgrade", "--url", url, "--message", ping,
		"--origin", "https://evil.example",
	)
	cmdtest.ErrContains(t, err, "upgrade refused: 403 Forbidden")

	out := cmdtest.MustRun(
		t,
		newRoot,
		"dial", "--http-upgrade", "--url", url, "--message", ping,
		"--origin", "https://admin.example.com",
	)
	cmdtest.Contains(t, out, "connected (101 Switching Protocols)")
}

// With --listen, pushed messages are printed until the server closes the
// connection.
func TestDialListen(t *testing.T) {
	hub, url := wsServer(t)
	out := &syncBuffer{}

	done := make(chan error, 1)
	go func() {
		done <- dial(context.Background(), nil, out, dialSettings{
			url:      url,
			messages: []string{ping},
			listen:   true,
			timeout:  time.Second,
		})
	}()

	// Push only once the client has its pong, and shut down only once it has
	// the push: Shutdown closes clients without waiting for their queues.
	waitFor(t, out, `< {"type":"pong","id":"1"`)
	hub.Push("event", map[string]any{"n": 1})
	waitFor(t, out, `< {"type":"event","payload":{"n":1}}`)

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	hub.Shutdown(ctx)

	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("dial = %v; want nil when the server closes", err)
		}
	case <-time.After(time.Second):
		t.Fatal("dial did not return when the server closed")
	}
	waitFor(t, out, "closed by the server (1001)")
}

// A pushed message larger than the WebSocket library's default read limit
// (32KiB), as a catch-up event can be, is printed in full.
func TestDialLargeMessage(t *testing.T) {
	hub, url := wsServer(t)
	out := &syncBuffer{}

	done := make(chan error, 1)
	go func() {
		done <- dial(context.Background(), nil, out, dialSettings{
			url:      url,
			messages: []string{ping},
			listen:   true,
			timeout:  time.Second,
		})
	}()

	waitFor(t, out, `< {"type":"pong","id":"1"`)
	large := strings.Repeat("x", 100_000)
	hub.Push("event", map[string]string{"large": large, "end": "here"})
	// encoding/json sorts a map's keys.
	waitFor(t, out, `{"end":"here","large":"`+large+`"}}`)

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	hub.Shutdown(ctx)
	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("dial = %v; want nil when the server closes", err)
		}
	case <-time.After(time.Second):
		t.Fatal("dial did not return when the server closed")
	}
}

// syncBuffer is a bytes.Buffer safe to write from one goroutine while
// another reads it.
type syncBuffer struct {
	mu     sync.Mutex
	buffer bytes.Buffer
}

func (b *syncBuffer) Write(data []byte) (int, error) {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.buffer.Write(data)
}

func (b *syncBuffer) String() string {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.buffer.String()
}

// waitFor fails the test unless out comes to contain every string in want
// within a second.
func waitFor(t *testing.T, out *syncBuffer, want ...string) {
	t.Helper()
	deadline := time.Now().Add(time.Second)
	for {
		missing := ""
		for _, wanted := range want {
			if !strings.Contains(out.String(), wanted) {
				missing = wanted
				break
			}
		}
		if missing == "" {
			return
		}
		if time.Now().After(deadline) {
			t.Fatalf("output does not contain %q:\n%s", missing, out)
		}
		time.Sleep(5 * time.Millisecond)
	}
}

// Without --message, each line of stdin is sent, skipping blank ones. The end
// of stdin doesn't end dial: pushes are still printed until Ctrl+C.
func TestDialStdin(t *testing.T) {
	hub, url := wsServer(t)
	stdin, input := io.Pipe()
	out := &syncBuffer{}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	done := make(chan error, 1)
	go func() {
		done <- dial(ctx, stdin, out, dialSettings{
			url:     url,
			timeout: time.Second,
		})
	}()

	io.WriteString(input, ping+"\n\n"+`{"type":"echo","id":"2"}`+"\n")
	input.Close()
	waitFor(
		t,
		out,
		`< {"type":"welcome"`,
		`< {"type":"pong","id":"1"`,
		`< {"type":"echo","id":"2"}`,
	)
	if strings.Contains(out.String(), `"type":"error"`) {
		t.Errorf("a blank line was sent:\n%s", out)
	}

	hub.Push("event", map[string]any{"n": 1})
	waitFor(t, out, `< {"type":"event","payload":{"n":1}}`)

	cancel() // Ctrl+C
	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("dial = %v; want nil after Ctrl+C", err)
		}
	case <-time.After(time.Second):
		t.Fatal("dial did not return after Ctrl+C")
	}
}

// Without --message, a failure reading stdin ends dial with an error, unlike
// the end of stdin.
func TestDialStdinError(t *testing.T) {
	_, url := wsServer(t)
	stdin, input := io.Pipe()
	out := &syncBuffer{}

	done := make(chan error, 1)
	go func() {
		done <- dial(context.Background(), stdin, out, dialSettings{
			url:     url,
			timeout: time.Second,
		})
	}()

	waitFor(t, out, `< {"type":"welcome"`)
	input.CloseWithError(errors.New("disk on fire"))

	select {
	case err := <-done:
		if err == nil ||
			!strings.Contains(err.Error(), "reading input: disk on fire") {
			t.Fatalf("dial = %v; want the input error", err)
		}
	case <-time.After(time.Second):
		t.Fatal("dial did not return after the input error")
	}
}

// Without --message, dial ends when the server closes the connection.
func TestDialStdinServerClose(t *testing.T) {
	hub, url := wsServer(t)
	stdin, _ := io.Pipe() // never written: stdin stays open
	out := &syncBuffer{}

	done := make(chan error, 1)
	go func() {
		done <- dial(context.Background(), stdin, out, dialSettings{
			url:     url,
			timeout: time.Second,
		})
	}()

	waitFor(t, out, `< {"type":"welcome"`)
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	hub.Shutdown(ctx)

	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("dial = %v; want nil when the server closes", err)
		}
	case <-time.After(time.Second):
		t.Fatal("dial did not return when the server closed")
	}
	waitFor(t, out, "closed by the server (1001)")
}

// Without --url, dial uses the configured host and port.
func TestDefaultURL(t *testing.T) {
	tests := []struct {
		name string
		conf config.JSCServer
		want string
	}{
		{"configured", config.JSCServer{
			Host: "10.0.0.1",
			Port: 8080,
		}, "ws://10.0.0.1:8080/ws"},
		{"every interface", config.JSCServer{
			Port: 9393,
		}, "ws://127.0.0.1:9393/ws"},
		{"IPv6", config.JSCServer{
			Host: "::1",
			Port: 9393,
		}, "ws://[::1]:9393/ws"},
		{"TLS", config.JSCServer{
			Host: "h",
			Port: 443,
			TLS: config.JSCServerTLS{
				Enabled: true,
			},
		}, "wss://h:443/ws"},
	}
	for _, tt := range tests {
		if got := defaultURL(tt.conf); got != tt.want {
			t.Errorf("%s: defaultURL = %q; want %q", tt.name, got, tt.want)
		}
	}
}