From 2f577c07a2140853b8b5ce7873c8df13dc6cb882 Mon Sep 17 00:00:00 2001
From: Jacob Su <suzp1984@gmail.com>
Date: Tue, 19 May 2026 16:05:52 +0800
Subject: [PATCH] test: Add unit tests to verify proxy header copy and URL fix

1. Fix test expectation: change && to & in m3u8 rewrite test
2. Add TestHLSPlayStream_ServeByBackend_HeadersCopiedFromBackend to
   verify backend headers reach the client for m3u8 responses
3. Add TestHLSPlayStream_ServeByBackend_TSHeadersCopiedFromBackend to
   verify header copy for .ts file responses
4. Add TestHTTPFlvTsConn_ServeByBackend_HeadersCopiedFromBackend to
   verify header copy for FLV/TS streaming responses

These tests protect against regression where calling WriteHeader() before
Header.Add() causes headers to be silently discarded by Go's http.ResponseWriter.
---
 internal/proxy/http_test.go | 87 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 86 insertions(+), 1 deletion(-)

diff --git a/internal/proxy/http_test.go b/internal/proxy/http_test.go
index fa64225c4..0025daeb8 100644
--- a/internal/proxy/http_test.go
+++ b/internal/proxy/http_test.go
@@ -370,7 +370,7 @@ func TestHLSPlayStream_ServeByBackend_M3U8RewritesTSWithQuery(t *testing.T) {
 		&lb.OriginServer{IP: host, HTTP: []string{port}}); err != nil {
 		t.Fatalf("unexpected err: %v", err)
 	}
-	if want := "live-0.ts?spbhid=ABC&&token=foo"; !strings.Contains(rec.Body.String(), want) {
+	if want := "live-0.ts?spbhid=ABC&token=foo"; !strings.Contains(rec.Body.String(), want) {
 		t.Fatalf("missing %q in body: %q", want, rec.Body.String())
 	}
 }
@@ -396,6 +396,61 @@ func TestHLSPlayStream_ServeByBackend_AppendsRawQueryOnTS(t *testing.T) {
 	}
 }
 
+func TestHLSPlayStream_ServeByBackend_HeadersCopiedFromBackend(t *testing.T) {
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/vnd.apple.mpegurl")
+		w.Header().Set("Cache-Control", "no-cache")
+		w.Header().Set("X-Custom-Header", "custom-value")
+		_, _ = io.WriteString(w, "#EXTM3U\nlive-0.ts\n")
+	}))
+	defer ts.Close()
+	host, port := httptestHostPort(t, ts)
+
+	v := newHLSPlayStream()
+	req := httptest.NewRequest(http.MethodGet, "http://example.com/live.m3u8", nil)
+	rec := httptest.NewRecorder()
+	if err := v.serveByBackend(context.Background(), rec, req,
+		&lb.OriginServer{IP: host, HTTP: []string{port}}); err != nil {
+		t.Fatalf("unexpected err: %v", err)
+	}
+
+	// Verify headers are properly copied (not lost due to WriteHeader order)
+	if got := rec.Header().Get("Content-Type"); got != "application/vnd.apple.mpegurl" {
+		t.Errorf("Content-Type = %q, want application/vnd.apple.mpegurl", got)
+	}
+	if got := rec.Header().Get("Cache-Control"); got != "no-cache" {
+		t.Errorf("Cache-Control = %q, want no-cache", got)
+	}
+	if got := rec.Header().Get("X-Custom-Header"); got != "custom-value" {
+		t.Errorf("X-Custom-Header = %q, want custom-value", got)
+	}
+}
+
+func TestHLSPlayStream_ServeByBackend_TSHeadersCopiedFromBackend(t *testing.T) {
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "video/mp2t")
+		w.Header().Set("Cache-Control", "max-age=3600")
+		_, _ = w.Write([]byte{0x47, 0x00, 0x01})
+	}))
+	defer ts.Close()
+	host, port := httptestHostPort(t, ts)
+
+	v := newHLSPlayStream()
+	req := httptest.NewRequest(http.MethodGet, "http://example.com/live.ts", nil)
+	rec := httptest.NewRecorder()
+	if err := v.serveByBackend(context.Background(), rec, req,
+		&lb.OriginServer{IP: host, HTTP: []string{port}}); err != nil {
+		t.Fatalf("unexpected err: %v", err)
+	}
+
+	if got := rec.Header().Get("Content-Type"); got != "video/mp2t" {
+		t.Errorf("Content-Type = %q, want video/mp2t", got)
+	}
+	if got := rec.Header().Get("Cache-Control"); got != "max-age=3600" {
+		t.Errorf("Cache-Control = %q, want max-age=3600", got)
+	}
+}
+
 // =============================================================================
 // httpFlvTsConnection
 // =============================================================================
@@ -666,6 +721,36 @@ func TestHTTPFlvTsConn_ServeByBackend_PreservesMethod(t *testing.T) {
 	}
 }
 
+func TestHTTPFlvTsConn_ServeByBackend_HeadersCopiedFromBackend(t *testing.T) {
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "video/x-flv")
+		w.Header().Set("Cache-Control", "no-store")
+		w.Header().Set("X-Custom-Header", "flv-value")
+		_, _ = w.Write([]byte("FLV\x01\x05\x00\x00\x00\x09"))
+	}))
+	defer ts.Close()
+	host, port := httptestHostPort(t, ts)
+
+	v := newHTTPFlvTsConnection()
+	req := httptest.NewRequest(http.MethodGet, "http://example.com/live.flv", nil)
+	rec := httptest.NewRecorder()
+	if err := v.serveByBackend(context.Background(), rec, req,
+		&lb.OriginServer{IP: host, HTTP: []string{port}}); err != nil {
+		t.Fatalf("unexpected err: %v", err)
+	}
+
+	// Verify headers are properly copied (not lost due to WriteHeader order)
+	if got := rec.Header().Get("Content-Type"); got != "video/x-flv" {
+		t.Errorf("Content-Type = %q, want video/x-flv", got)
+	}
+	if got := rec.Header().Get("Cache-Control"); got != "no-store" {
+		t.Errorf("Cache-Control = %q, want no-store", got)
+	}
+	if got := rec.Header().Get("X-Custom-Header"); got != "flv-value" {
+		t.Errorf("X-Custom-Header = %q, want flv-value", got)
+	}
+}
+
 // =============================================================================
 // httpStreamProxyServer
 // =============================================================================