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)") } // 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) } } }