job-search/internal/websocket/http_test.go

package websocket

import (
	"context"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	ws "github.com/coder/websocket"
	"github.com/coder/websocket/wsjson"

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

// wsHub serves h's WebSocket clients on a test HTTP server and returns its
// ws:// URL.
func wsHub(t *testing.T, h *Hub) string {
	t.Helper()
	srv := httptest.NewServer(h)
	t.Cleanup(func() {
		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
		defer cancel()
		h.Shutdown(ctx)
		srv.Close()
	})
	return "ws" + strings.TrimPrefix(srv.URL, "http")
}

// dialWS connects a WebSocket client to url and reads its welcome.
func dialWS(t *testing.T, ctx context.Context, url string) *ws.Conn {
	t.Helper()
	conn, _, err := ws.Dial(ctx, url, nil)
	if err != nil {
		t.Fatalf("dial: %v", err)
	}
	t.Cleanup(func() {
		conn.CloseNow()
	})

	var e Envelope
	if err := wsjson.Read(ctx, conn, &e); err != nil || e.Type != TypeWelcome {
		t.Fatalf("want welcome, got %+v, %v", e, err)
	}

	return conn
}

// A hub without a CSRF handler accepts upgrades without a token, and the
// client gets replies. A max_message_size of 0 uses the default limit rather
// than refusing every message.
func TestHTTPNoCSRF(t *testing.T) {
	h := NewHub(
		config.JSCServerWebSocket{
			WriteTimeout: time.Second,
		},
		WithLogger(discardLog),
	)
	url := wsHub(t, h)
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	conn := dialWS(t, ctx, url)

	if err := wsjson.Write(ctx, conn, Envelope{
		Type: TypePing,
		ID:   "1",
	}); err != nil {
		t.Fatal(err)
	}
	var e Envelope
	if err := wsjson.Read(ctx, conn, &e); err != nil || e.Type != TypePong ||
		e.ID != "1" {
		t.Fatalf("want pong 1, got %+v, %v", e, err)
	}
}

// An upgrade is accepted without an Origin (a non-browser client), from the
// server's own origin or a trusted one, and refused from any other origin.
func TestHTTPOrigins(t *testing.T) {
	h := NewHub(
		config.JSCServerWebSocket{},
		WithLogger(discardLog),
		WithTrustedOrigins([]string{"https://admin.example.com"}),
	)
	url := wsHub(t, h)
	self := "http" + strings.TrimPrefix(url, "ws")

	tests := []struct {
		origin string
		ok     bool
	}{
		{"", true},
		{self, true},
		{"https://admin.example.com", true},
		{"http://admin.example.com", false},
		{"https://evil.example", false},
	}
	for _, tt := range tests {
		t.Run(tt.origin, func(t *testing.T) {
			ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
			defer cancel()
			header := http.Header{}
			if tt.origin != "" {
				header.Set("Origin", tt.origin)
			}
			conn, res, err := ws.Dial(ctx, url, &ws.DialOptions{
				HTTPHeader: header,
			})
			if tt.ok {
				if err != nil {
					t.Fatalf("dial from %q: %v", tt.origin, err)
				}
				conn.CloseNow()
				return
			}
			if err == nil || res == nil || res.StatusCode != http.StatusForbidden {
				t.Fatalf("dial from %q = %v (%v); want 403", tt.origin, err, res)
			}
		})
	}
}

// One Push reaches WebSocket and plain TCP clients of the same hub.
func TestPushReachesBothTransports(t *testing.T) {
	h, addr := tcpHub(t, nil)
	url := wsHub(t, h)

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	conn := dialWS(t, ctx, url)
	c := dialTCP(t, addr)
	c.read() // welcome
	waitCount(t, h, 2)

	h.Push("item.created", map[string]string{"id": "x"})
	var e Envelope
	if err := wsjson.Read(ctx, conn, &e); err != nil || e.Type != "item.created" {
		t.Fatalf("websocket: want item.created, got %+v, %v", e, err)
	}
	if e := c.read(); e.Type != "item.created" {
		t.Fatalf("tcp: want item.created, got %+v", e)
	}
}

// A payload that can't be marshalled is not sent.
func TestPushBadPayload(t *testing.T) {
	h, addr := tcpHub(t, nil)
	c := dialTCP(t, addr)
	c.read() // welcome
	waitCount(t, h, 1)

	h.Push("bad", func() {
	})
	h.Push("good", nil)
	if e := c.read(); e.Type != "good" || e.Payload != nil {
		t.Fatalf("want good with no payload, got %+v", e)
	}
}

// A binary frame closes the connection with 1003 (unsupported data).
func TestHTTPBinaryFrame(t *testing.T) {
	h := NewHub(
		config.JSCServerWebSocket{
			WriteTimeout: time.Second,
		},
		WithLogger(discardLog),
	)
	url := wsHub(t, h)
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	conn := dialWS(t, ctx, url)

	if err := conn.Write(ctx, ws.MessageBinary, []byte{1, 2, 3}); err != nil {
		t.Fatal(err)
	}
	_, _, err := conn.Read(ctx)
	if ws.CloseStatus(err) != ws.StatusUnsupportedData {
		t.Fatalf("read after binary frame = %v; want close 1003", err)
	}
}

// Messages over the message rate limit get a "rate limited" error and are not
// handled; the connection stays open.
func TestHTTPMessageRateLimit(t *testing.T) {
	conf := config.JSCServerWebSocket{
		WriteTimeout: time.Second,
		MessageRateLimit: config.JSCRateLimitRule{
			Enabled:   true,
			PerSecond: 1,
			Burst:     1,
		},
	}
	h := NewHub(conf, WithLogger(discardLog))
	url := wsHub(t, h)
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	conn := dialWS(t, ctx, url)

	for _, id := range []string{"1", "2"} {
		if err := wsjson.Write(ctx, conn, Envelope{
			Type: TypePing,
			ID:   id,
		}); err != nil {
			t.Fatal(err)
		}
	}
	var e Envelope
	if err := wsjson.Read(ctx, conn, &e); err != nil || e.Type != TypePong ||
		e.ID != "1" {
		t.Fatalf("first: want pong 1, got %+v, %v", e, err)
	}
	err := wsjson.Read(ctx, conn, &e)
	if err != nil || e.Type != TypeError || e.Error != "rate limited" {
		t.Fatalf("second: want rate limited, got %+v, %v", e, err)
	}
	if h.Count() != 1 {
		t.Errorf("Count = %d; want the client kept", h.Count())
	}
}

// A request that isn't a WebSocket upgrade is refused by the upgrade and no
// client is registered.
func TestHTTPNotAnUpgrade(t *testing.T) {
	h := NewHub(config.JSCServerWebSocket{}, WithLogger(discardLog))
	srv := httptest.NewServer(h)
	defer srv.Close()

	res, err := http.Get(srv.URL)
	if err != nil {
		t.Fatal(err)
	}
	res.Body.Close()

	if res.StatusCode < 400 {
		t.Errorf("status = %d; want an error", res.StatusCode)
	}
	if h.Count() != 0 {
		t.Errorf("Count = %d; want 0", h.Count())
	}
}

// After Shutdown, an upgrade succeeds but the client is closed at once with
// 1001 (going away).
func TestHTTPRefusedAfterShutdown(t *testing.T) {
	h := NewHub(config.JSCServerWebSocket{}, WithLogger(discardLog))
	url := wsHub(t, h)
	if err := h.Shutdown(context.Background()); err != nil {
		t.Fatal(err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	conn, _, err := ws.Dial(ctx, url, nil)
	if err != nil {
		t.Fatalf("dial: %v", err)
	}
	defer conn.CloseNow()

	if _, _, err := conn.Read(ctx); ws.CloseStatus(err) != ws.StatusGoingAway {
		t.Fatalf("read = %v; want close 1001", err)
	}
}