go-omni/ws/client_test.go

package ws

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/gorilla/websocket"
)

func clientTestConfig() ClientConfig {
	return ClientConfig{ProxyMode: "direct", Session: SessionConfig{CloseTimeout: 100 * time.Millisecond}}
}
func clientForTest(t *testing.T, cfg ClientConfig) *Client {
	t.Helper()
	c, err := NewClient(cfg)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() {
		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
		defer cancel()
		if err := c.Shutdown(ctx); err != nil {
			t.Errorf("client shutdown: %v", err)
		}
	})
	return c
}
func echoWebSocket(w http.ResponseWriter, r *http.Request) {
	conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
	if err != nil {
		return
	}
	defer conn.Close()
	for {
		kind, data, err := conn.ReadMessage()
		if err != nil {
			return
		}
		if err = conn.WriteMessage(kind, data); err != nil {
			return
		}
	}
}
func wsURL(s *httptest.Server) string { return "ws" + strings.TrimPrefix(s.URL, "http") }
func receiveClientMessage(t *testing.T, ch <-chan Message) Message {
	t.Helper()
	select {
	case m := <-ch:
		return m
	case <-time.After(2 * time.Second):
		t.Fatal("message timeout")
		return Message{}
	}
}

func TestClientHandshakeDeadlineDoesNotExpireSession(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(echoWebSocket))
	defer server.Close()
	cfg := clientTestConfig()
	cfg.HandshakeTimeout = 40 * time.Millisecond
	c := clientForTest(t, cfg)
	messages := make(chan Message, 2)
	s, err := c.Dial(context.Background(), wsURL(server), Handlers{OnMessage: func(_ context.Context, _ *Session, m Message) error { messages <- m; return nil }})
	if err != nil {
		t.Fatal(err)
	}
	timer := time.NewTimer(100 * time.Millisecond)
	<-timer.C
	for _, kind := range []MessageType{Text, Binary} {
		if err = s.Send(context.Background(), Message{Type: kind, Data: []byte("alive")}); err != nil {
			t.Fatal(err)
		}
		got := receiveClientMessage(t, messages)
		if got.Type != kind || string(got.Data) != "alive" {
			t.Fatalf("unexpected echo: %#v", got)
		}
	}
	_ = s.Close(1000, "")
}

func TestClientTLSAndMutualTLS(t *testing.T) {
	for _, mutual := range []bool{false, true} {
		t.Run(fmt.Sprint(mutual), func(t *testing.T) {
			server := httptest.NewUnstartedServer(http.HandlerFunc(echoWebSocket))
			if mutual {
				server.TLS = &tls.Config{ClientAuth: tls.RequireAnyClientCert}
			}
			server.StartTLS()
			defer server.Close()
			roots := x509.NewCertPool()
			roots.AddCert(server.Certificate())
			cfg := clientTestConfig()
			cfg.TLSConfig = &tls.Config{RootCAs: roots, MinVersion: tls.VersionTLS12}
			if mutual {
				cfg.TLSConfig.Certificates = append([]tls.Certificate(nil), server.TLS.Certificates...)
			}
			c := clientForTest(t, cfg)
			cfg.TLSConfig.ServerName = "wrong-after-construction.invalid"
			messages := make(chan Message, 1)
			s, err := c.Dial(context.Background(), wsURL(server), Handlers{OnMessage: func(_ context.Context, _ *Session, m Message) error { messages <- m; return nil }})
			if err != nil {
				t.Fatal(err)
			}
			if err = s.Send(context.Background(), Message{Type: Text, Data: []byte("tls")}); err != nil {
				t.Fatal(err)
			}
			if string(receiveClientMessage(t, messages).Data) != "tls" {
				t.Fatal("bad echo")
			}
			_ = s.Close(1000, "")
		})
	}
}

func TestClientHeadersRefreshReconnectAndFreshSessions(t *testing.T) {
	var accepted atomic.Int32
	firstOpen := make(chan struct{})
	var firstOpenOnce sync.Once
	headerValues := make(chan string, 4)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		headerValues <- r.Header.Get("Authorization")
		conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
		if err != nil {
			return
		}
		defer conn.Close()
		if accepted.Add(1) == 1 {
			select {
			case <-firstOpen:
			case <-r.Context().Done():
				return
			}
			_ = conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(1012, "restart"), time.Now().Add(time.Second))
			return
		}
		for {
			if _, _, err = conn.ReadMessage(); err != nil {
				return
			}
		}
	}))
	defer server.Close()
	var providerCalls atomic.Int32
	var prepareCalls atomic.Int32
	ids := make(chan string, 4)
	cfg := clientTestConfig()
	cfg.Reconnect = ReconnectConfig{Enabled: true, MinBackoff: time.Millisecond, MaxBackoff: 5 * time.Millisecond, StableAfter: time.Second, MaxAttempts: 3}
	cfg.HeaderProvider = func(context.Context) (http.Header, error) {
		n := providerCalls.Add(1)
		return http.Header{"Authorization": {fmt.Sprint(n)}}, nil
	}
	cfg.Prepare = func(ctx context.Context) (Handlers, error) {
		prepareCalls.Add(1)
		return Handlers{OnOpen: func(sessionCtx context.Context, s *Session) error {
			if ctx.Err() != nil || sessionCtx.Err() != nil {
				return errors.New("prepare context canceled early")
			}
			ids <- s.ID()
			firstOpenOnce.Do(func() { close(firstOpen) })
			return nil
		}}, nil
	}
	c := clientForTest(t, cfg)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	finished := make(chan error, 1)
	go func() { finished <- c.Run(ctx, wsURL(server), Handlers{}) }()
	var first, second string
	select {
	case first = <-ids:
	case <-time.After(2 * time.Second):
		t.Fatal("first open timeout")
	}
	select {
	case second = <-ids:
	case <-time.After(2 * time.Second):
		t.Fatal("reopen timeout")
	}
	if first == second {
		t.Fatal("reused physical session ID")
	}
	if <-headerValues != "1" || <-headerValues != "2" {
		t.Fatal("headers not refreshed")
	}
	if prepareCalls.Load() != 2 {
		t.Fatal("prepare not recreated")
	}
	cancel()
	select {
	case err := <-finished:
		if !errors.Is(err, context.Canceled) {
			t.Fatalf("run: %v", err)
		}
	case <-time.After(time.Second):
		t.Fatal("run cancel timeout")
	}
}

func TestClientDoesNotReconnectAuthenticationOrProtocolFailure(t *testing.T) {
	for _, status := range []int{400, 401, 403, 404} {
		t.Run(fmt.Sprint(status), func(t *testing.T) {
			var calls atomic.Int32
			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				calls.Add(1)
				w.WriteHeader(status)
				_, _ = w.Write([]byte("secret response"))
			}))
			defer server.Close()
			cfg := clientTestConfig()
			cfg.Reconnect.Enabled = true
			c := clientForTest(t, cfg)
			err := c.Run(context.Background(), wsURL(server)+"/?token=TOPSECRET", Handlers{})
			var handshake *HandshakeError
			if !errors.As(err, &handshake) || handshake.StatusCode != status {
				t.Fatalf("error: %v", err)
			}
			if calls.Load() != 1 || strings.Contains(err.Error(), "SECRET") || strings.Contains(err.Error(), "secret") {
				t.Fatalf("retry/leak: calls=%d error=%v", calls.Load(), err)
			}
		})
	}
}

func TestClientCloseCancelsHandshakeAndBackoff(t *testing.T) {
	t.Run("handshake", func(t *testing.T) {
		entered := make(chan struct{})
		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { close(entered); <-r.Context().Done() }))
		defer server.Close()
		c := clientForTest(t, clientTestConfig())
		finished := make(chan error, 1)
		go func() { _, err := c.Dial(context.Background(), wsURL(server), Handlers{}); finished <- err }()
		<-entered
		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
		defer cancel()
		if err := c.Shutdown(ctx); err != nil {
			t.Fatal(err)
		}
		if err := <-finished; !errors.Is(err, ErrClosed) {
			t.Fatalf("dial: %v", err)
		}
	})
	t.Run("backoff", func(t *testing.T) {
		entered := make(chan struct{})
		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { close(entered); w.WriteHeader(503) }))
		defer server.Close()
		cfg := clientTestConfig()
		cfg.Reconnect = ReconnectConfig{Enabled: true, MinBackoff: time.Hour, MaxBackoff: time.Hour}
		c := clientForTest(t, cfg)
		finished := make(chan error, 1)
		go func() { finished <- c.Run(context.Background(), wsURL(server), Handlers{}) }()
		<-entered
		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
		defer cancel()
		if err := c.Shutdown(ctx); err != nil {
			t.Fatal(err)
		}
		if err := <-finished; err != nil {
			t.Fatal(err)
		}
	})
}

func TestClientConcurrentDialLimitAndClose(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(echoWebSocket))
	defer server.Close()
	cfg := clientTestConfig()
	cfg.MaxConnections = 2
	c := clientForTest(t, cfg)
	var wg sync.WaitGroup
	var successes atomic.Int32
	for i := 0; i < 20; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			_, err := c.Dial(context.Background(), wsURL(server), Handlers{})
			if err == nil {
				successes.Add(1)
			} else if !errors.Is(err, ErrConnectionLimit) {
				t.Errorf("dial: %v", err)
			}
		}()
	}
	wg.Wait()
	if successes.Load() != 2 {
		t.Fatalf("connections=%d", successes.Load())
	}
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	if err := c.Shutdown(ctx); err != nil {
		t.Fatal(err)
	}
	if _, err := c.Dial(ctx, wsURL(server), Handlers{}); !errors.Is(err, ErrClosed) {
		t.Fatalf("closed dial: %v", err)
	}
}

func TestClientRetryBudgetAndValidation(t *testing.T) {
	var calls atomic.Int32
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { calls.Add(1); w.WriteHeader(503) }))
	defer server.Close()
	cfg := clientTestConfig()
	cfg.Reconnect = ReconnectConfig{Enabled: true, MinBackoff: time.Millisecond, MaxBackoff: 2 * time.Millisecond, MaxAttempts: 2}
	c := clientForTest(t, cfg)
	if err := c.Run(context.Background(), wsURL(server), Handlers{}); err == nil {
		t.Fatal("wanted failure")
	}
	if calls.Load() != 3 {
		t.Fatalf("retry budget: %d", calls.Load())
	}
	for _, endpoint := range []string{"https://example.test", "ws://user:password@example.test/", "ws://", "ws://example.test/#fragment"} {
		if _, err := c.Dial(context.Background(), endpoint, Handlers{}); err == nil || strings.Contains(err.Error(), "password") {
			t.Fatalf("invalid endpoint: %v", err)
		}
	}
	for _, cfg := range []ClientConfig{{HandshakeTimeout: -1}, {ProxyMode: "bad"}, {ProxyMode: "direct", ProxyURL: "http://proxy"}, {Headers: http.Header{"Authorization": {"bad\r\nheader"}}}, {Headers: http.Header{"Sec-WebSocket-Key": {"fake"}}}, {RequireSubprotocol: true}, {Subprotocols: []string{"bad protocol"}}} {
		if _, err := NewClient(cfg); err == nil {
			t.Fatal("accepted invalid configuration")
		}
	}
	for i := 0; i < 100; i++ {
		d := retryBackoff(ReconnectConfig{MinBackoff: time.Second, MaxBackoff: 30 * time.Second}, i)
		if d <= 0 || d > 30*time.Second {
			t.Fatalf("unbounded delay %v", d)
		}
	}
}

func TestClientElapsedBudgetInterruptsStalledHandshake(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { <-r.Context().Done() }))
	defer server.Close()
	cfg := clientTestConfig()
	cfg.Reconnect = ReconnectConfig{Enabled: true, MaxElapsed: 50 * time.Millisecond}
	c := clientForTest(t, cfg)
	started := time.Now()
	err := c.Run(context.Background(), wsURL(server), Handlers{})
	if !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("budget error: %v", err)
	}
	if time.Since(started) > time.Second {
		t.Fatal("handshake exceeded reconnect budget")
	}
}

func TestClientSingleSlotRapidReconnect(t *testing.T) {
	var calls atomic.Int32
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
		if err != nil {
			return
		}
		defer conn.Close()
		code := 1012
		if calls.Add(1) >= 20 {
			code = 1000
		}
		_ = conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(code, ""), time.Now().Add(time.Second))
		_, _, _ = conn.ReadMessage()
	}))
	defer server.Close()
	cfg := clientTestConfig()
	cfg.MaxConnections = 1
	cfg.Reconnect = ReconnectConfig{Enabled: true, MinBackoff: time.Nanosecond, MaxBackoff: time.Nanosecond, MaxAttempts: 25}
	c := clientForTest(t, cfg)
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	if err := c.Run(ctx, wsURL(server), Handlers{}); err != nil {
		t.Fatalf("rapid reconnect: %v", err)
	}
	if calls.Load() != 20 {
		t.Fatalf("connections: %d", calls.Load())
	}
}

func TestClientSubprotocol(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := (&websocket.Upgrader{Subprotocols: []string{"omni.v1"}}).Upgrade(w, r, nil)
		if err != nil {
			return
		}
		defer conn.Close()
		for {
			if _, _, err = conn.ReadMessage(); err != nil {
				return
			}
		}
	}))
	defer server.Close()
	cfg := clientTestConfig()
	cfg.Subprotocols = []string{"omni.v1"}
	cfg.RequireSubprotocol = true
	c := clientForTest(t, cfg)
	s, err := c.Dial(context.Background(), wsURL(server), Handlers{})
	if err != nil {
		t.Fatal(err)
	}
	if s.Subprotocol() != "omni.v1" {
		t.Fatal("wrong protocol")
	}
	_ = s.Close(1000, "")
	cfg.Subprotocols = []string{"missing"}
	other := clientForTest(t, cfg)
	if _, err = other.Dial(context.Background(), wsURL(server), Handlers{}); err == nil {
		t.Fatal("required protocol not enforced")
	}
}

func TestClientHTTPProxyAuthenticationAndTunnel(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Proxy-Authorization") != "" {
			t.Error("proxy credential leaked to target")
		}
		echoWebSocket(w, r)
	}))
	defer server.Close()
	var connected atomic.Int32
	proxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodConnect || r.Header.Get("Proxy-Authorization") != "Basic dXNlcjpwYXNz" {
			w.WriteHeader(http.StatusProxyAuthRequired)
			return
		}
		upstream, err := net.Dial("tcp", r.Host)
		if err != nil {
			t.Error(err)
			w.WriteHeader(502)
			return
		}
		conn, rw, err := w.(http.Hijacker).Hijack()
		if err != nil {
			upstream.Close()
			t.Error(err)
			return
		}
		connected.Add(1)
		_, _ = conn.Write([]byte("HTTP/1.1 200 Connection Established\r\n\r\n"))
		go func() { defer conn.Close(); defer upstream.Close(); _, _ = io.Copy(upstream, rw) }()
		go func() { defer conn.Close(); defer upstream.Close(); _, _ = io.Copy(conn, upstream) }()
	}))
	defer proxy.Close()
	cfg := clientTestConfig()
	cfg.ProxyMode = "url"
	cfg.ProxyURL = strings.Replace(proxy.URL, "http://", "http://user:pass@", 1)
	c := clientForTest(t, cfg)
	messages := make(chan Message, 1)
	s, err := c.Dial(context.Background(), wsURL(server), Handlers{OnMessage: func(_ context.Context, _ *Session, m Message) error { messages <- m; return nil }})
	if err != nil {
		t.Fatal(err)
	}
	if err = s.Send(context.Background(), Message{Type: Text, Data: []byte("proxied")}); err != nil {
		t.Fatal(err)
	}
	if string(receiveClientMessage(t, messages).Data) != "proxied" || connected.Load() != 1 {
		t.Fatal("proxy not used")
	}
	_ = s.Close(1000, "")
}

func BenchmarkClientRoundTrip(b *testing.B) {
	server := httptest.NewServer(http.HandlerFunc(echoWebSocket))
	defer server.Close()
	c, err := NewClient(clientTestConfig())
	if err != nil {
		b.Fatal(err)
	}
	defer c.Shutdown(context.Background())
	received := make(chan struct{}, 1)
	s, err := c.Dial(context.Background(), wsURL(server), Handlers{OnMessage: func(context.Context, *Session, Message) error { received <- struct{}{}; return nil }})
	if err != nil {
		b.Fatal(err)
	}
	message := Message{Type: Text, Data: []byte("benchmark")}
	b.ReportAllocs()
	b.SetBytes(int64(len(message.Data)))
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if err := s.Send(context.Background(), message); err != nil {
			b.Fatal(err)
		}
		<-received
	}
	b.StopTimer()
	_ = s.Close(1000, "")
}