vitrine

test: prove websocket upgrade and SSE streaming through forwards

commit 06f4ceb908738d2986a46f68994a8399c1c0c3d1 · parent a052659317 · browse this tree
Ximianth, 2026-08-14 16:36:16 UTC

httputil.ReverseProxy documents both behaviors; these tests demand
them over real sockets rather than trusting the docs.

The websocket test hand-rolls the Upgrade handshake (the subject is
upgrade pass-through, not ws framing -- no ws library needed) and
proves bytes flow both directions on the hijacked connection.

The SSE test cannot false-pass: the backend holds the stream open
until the client confirms the event arrived, so a proxy that buffers
to EOF deadlocks into a damage-control timeout instead of passing.
No sleeps anywhere; all coordination is channel-gated.

diff --git a/forward_stream_test.go b/forward_stream_test.go
new file mode 100644
index 0000000..9b96e82
--- /dev/null
+++ b/forward_stream_test.go
@@ -0,0 +1,132 @@
+package main
+
+import (
+	"bufio"
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+)
+
+// The tests in this file prove streaming behavior through a forward:
+// websocket upgrades and SSE both demand more than request/response --
+// bytes must flow while the connection stays open.
+// httptest.ResponseRecorder can neither hijack nor observe mid-stream,
+// so these run vitrine in a real httptest.Server and speak over real sockets.
+
+// TestForwardWebsocketUpgrade proves an Upgrade round-trip through the proxy:
+// handshake, then bytes both ways on the hijacked connection.
+// The handshake is hand-rolled: the subject is upgrade *pass-through*,
+// not websocket framing, so no ws library is needed on either end.
+func TestForwardWebsocketUpgrade(t *testing.T) {
+	// The backend hijacks upgrade requests and echoes lines back prefixed.
+	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if !strings.EqualFold(r.Header.Get("Upgrade"), "websocket") {
+			t.Errorf("backend saw no Upgrade header; got %q", r.Header.Get("Upgrade"))
+			http.Error(w, "expected upgrade", http.StatusBadRequest)
+			return
+		}
+		conn, rw, err := w.(http.Hijacker).Hijack()
+		if err != nil {
+			t.Errorf("hijack: %v", err)
+			return
+		}
+		defer conn.Close()
+		rw.WriteString("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n")
+		rw.Flush()
+		for {
+			line, err := rw.ReadString('\n')
+			if err != nil {
+				return
+			}
+			rw.WriteString("echo:" + line)
+			rw.Flush()
+		}
+	}))
+	defer backend.Close()
+
+	vitrine := httptest.NewServer(forwardServer(t, map[string]string{"game": backend.URL}))
+	defer vitrine.Close()
+
+	// Dial vitrine raw and perform the upgrade handshake through the proxy.
+	conn, err := net.Dial("tcp", strings.TrimPrefix(vitrine.URL, "http://"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer conn.Close()
+	fmt.Fprintf(conn, "GET /game/ws HTTP/1.1\r\nHost: x\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n")
+	br := bufio.NewReader(conn)
+	resp, err := http.ReadResponse(br, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if resp.StatusCode != http.StatusSwitchingProtocols {
+		t.Fatalf("status = %d, want 101", resp.StatusCode)
+	}
+
+	// The connection is now a two-way pipe through the proxy; prove both directions.
+	for _, msg := range []string{"hello\n", "worlds\n"} {
+		if _, err := io.WriteString(conn, msg); err != nil {
+			t.Fatal(err)
+		}
+		line, err := br.ReadString('\n')
+		if err != nil {
+			t.Fatal(err)
+		}
+		if line != "echo:"+msg {
+			t.Fatalf("read %q, want %q", line, "echo:"+msg)
+		}
+	}
+}
+
+// TestForwardSSEStreamsMidResponse proves the proxy flushes streamed bytes
+// while the backend holds the connection open, rather than buffering to EOF.
+// The backend sends one event, then blocks until the client confirms receipt
+// (channel-gated, no sleeps): if the event arrives at all, it arrived mid-stream.
+func TestForwardSSEStreamsMidResponse(t *testing.T) {
+	received := make(chan struct{})
+	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "text/event-stream")
+		fmt.Fprint(w, "data: tick\n\n")
+		w.(http.Flusher).Flush()
+		// Hold the response open until the client has read the event.
+		// Only then may the stream end; a proxy that buffers to EOF deadlocks
+		// here instead of passing, which is exactly the point.
+		select {
+		case <-received:
+		case <-r.Context().Done():
+			t.Error("backend: client gave up before confirming the event")
+		}
+	}))
+	defer backend.Close()
+
+	vitrine := httptest.NewServer(forwardServer(t, map[string]string{"": backend.URL}))
+	defer vitrine.Close()
+
+	// The client timeout is damage control only, never load-bearing:
+	// a correctly-flushing proxy delivers the event immediately.
+	// Without it, a buffering proxy would deadlock the test instead of failing it
+	// (client blocked on read, backend blocked on the confirmation channel).
+	client := &http.Client{Timeout: 30 * time.Second}
+	resp, err := client.Get(vitrine.URL + "/events")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer resp.Body.Close()
+	if ct := resp.Header.Get("Content-Type"); ct != "text/event-stream" {
+		t.Fatalf("content-type = %q", ct)
+	}
+	br := bufio.NewReader(resp.Body)
+	line, err := br.ReadString('\n')
+	if err != nil {
+		t.Fatal(err)
+	}
+	if line != "data: tick\n" {
+		t.Fatalf("read %q, want the SSE event line", line)
+	}
+	close(received)
+}