forward_stream_test.go

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