package relay import ( "bufio" "context" "fmt" "io" "log/slog" "net" "net/http" "net/http/httptest" "net/url" "strings" "testing" "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "visiona-backend/internal/session" ) // TestEndToEnd_RawForward 驗證 B3 Review Major 1 修復後新增的 // `POST /internal/forward/raw` endpoint 能走完 raw TCP forwarding 路徑, // 並支援 streaming body(MJPEG / chunked)。 // // 路徑: // // fake api-server (raw TCP dial) // └─► POST /internal/forward/raw ──► remote-proxy internal server // └─► Hijack + OpenStream + 雙向 io.Copy // └─► fake tunnel client (yamux stream) // └─► fake local server(chunked response) // // 驗證重點: // 1. 「HTTP/1.1 200 Connected」握手成功 // 2. 完整 HTTP request 能寫進 hijacked 連線 → local server 收到 // 3. Response status / headers / body 能正確回來 // 4. Chunked streaming body 的 trailing chunks 也能收完(不像 JSON 版會一次收完) func TestEndToEnd_RawForward(t *testing.T) { // 1. Fake local server — 回 chunked streaming body 模擬 MJPEG / SSE localSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/event-stream") w.Header().Set("X-Test-Route", r.URL.Path) w.WriteHeader(http.StatusOK) flusher, _ := w.(http.Flusher) // 送 3 個 chunk 模擬 streaming for i := 0; i < 3; i++ { fmt.Fprintf(w, "data: chunk-%d\n\n", i) if flusher != nil { flusher.Flush() } time.Sleep(10 * time.Millisecond) } })) defer localSrv.Close() localAddr := strings.TrimPrefix(localSrv.URL, "http://") // 2. 起 remote-proxy(tunnel + internal server) store := session.NewInMemoryStore() relaySrv := NewServer(store, slog.Default(), Options{KeepAliveInterval: 500 * time.Millisecond}) internalSrv := NewInternalServer(store, slog.Default()) tunnelMux := http.NewServeMux() tunnelMux.HandleFunc("/tunnel/connect", relaySrv.HandleTunnelConnect) tunnelTS := httptest.NewServer(tunnelMux) defer tunnelTS.Close() internalMux := http.NewServeMux() internalSrv.Routes(internalMux) internalTS := httptest.NewServer(internalMux) defer internalTS.Close() // 3. Fake tunnel client — 把 stream 收到的 HTTP request 真 TCP 轉發給 localSrv const token = "vAc_cafecafecafecafecafecafecafecafe" stopTunnel := startTunnelClientForwardingTo(t, tunnelTS.URL, token, localAddr) defer stopTunnel() // 4. 等 session register require.Eventually(t, func() bool { ok, _ := store.Exists(context.Background(), token) return ok }, 2*time.Second, 20*time.Millisecond) // 5. 模擬 api-server 端:raw TCP dial → hijack 握手 → 送 HTTP request → 讀 response conn := dialRawForward(t, internalTS.URL, token) defer conn.Close() // 送一個真正的 HTTP GET / (走完整的 RFC 7230 格式,local agent 要會 parse) reqLine := "GET /api/stream HTTP/1.1\r\n" + "Host: 127.0.0.1\r\n" + "X-From-Api-Server: raw-test\r\n" + "Accept: text/event-stream\r\n" + "\r\n" _, werr := conn.Write([]byte(reqLine)) require.NoError(t, werr) // 讀 response — 用 http.ReadResponse 解析 chunked body reader := bufio.NewReader(conn) httpReq, _ := http.NewRequest(http.MethodGet, "/api/stream", nil) resp, rerr := http.ReadResponse(reader, httpReq) require.NoError(t, rerr) defer resp.Body.Close() assert.Equal(t, http.StatusOK, resp.StatusCode) assert.Equal(t, "/api/stream", resp.Header.Get("X-Test-Route"), "response header 應該被原封轉發") assert.Equal(t, "text/event-stream", resp.Header.Get("Content-Type")) // 讀 streaming body — 驗證三個 chunk 都收到 body, err := io.ReadAll(resp.Body) require.NoError(t, err) bodyStr := string(body) assert.Contains(t, bodyStr, "data: chunk-0") assert.Contains(t, bodyStr, "data: chunk-1") assert.Contains(t, bodyStr, "data: chunk-2") } // TestEndToEnd_RawForward_TunnelDisconnected 當 token 不存在時, // raw forward endpoint 應在 hijack 前回 502 JSON(可被一般 HTTP client 讀到)。 func TestEndToEnd_RawForward_TunnelDisconnected(t *testing.T) { store := session.NewInMemoryStore() internalSrv := NewInternalServer(store, slog.Default()) mux := http.NewServeMux() internalSrv.Routes(mux) ts := httptest.NewServer(mux) defer ts.Close() // 用一般 http client 打(沒 session 時還沒 hijack,會回一般 JSON response) resp, err := http.Post( ts.URL+"/internal/forward/raw?token=vAc_dddddddddddddddddddddddddddddddd", "application/octet-stream", nil, ) require.NoError(t, err) defer resp.Body.Close() assert.Equal(t, http.StatusBadGateway, resp.StatusCode) } // TestEndToEnd_RawForward_MissingToken 沒帶 token 應回 400。 func TestEndToEnd_RawForward_MissingToken(t *testing.T) { store := session.NewInMemoryStore() internalSrv := NewInternalServer(store, slog.Default()) mux := http.NewServeMux() internalSrv.Routes(mux) ts := httptest.NewServer(mux) defer ts.Close() resp, err := http.Post(ts.URL+"/internal/forward/raw", "application/octet-stream", nil) require.NoError(t, err) defer resp.Body.Close() assert.Equal(t, http.StatusBadRequest, resp.StatusCode) } // ---------------------------------------------------------------------- // Helpers // ---------------------------------------------------------------------- // dialRawForward 模擬 api-server 端:raw TCP dial remote-proxy, // 發一個帶 token 的 POST /internal/forward/raw 請求,讀取 "HTTP/1.1 200 Connected" // 握手回應,然後回傳這條已經「接管為 raw TCP」的連線供 caller 直接 io 使用。 // // 對齊 `HandleForwardRaw` 的協議。 func dialRawForward(t *testing.T, internalURL, token string) net.Conn { t.Helper() u, err := url.Parse(internalURL) require.NoError(t, err) // TCP dial conn, err := net.DialTimeout("tcp", u.Host, 5*time.Second) require.NoError(t, err) // 送 HTTP POST request(含 token query) reqLine := fmt.Sprintf( "POST /internal/forward/raw?token=%s HTTP/1.1\r\n"+ "Host: %s\r\n"+ "Content-Length: 0\r\n"+ "\r\n", token, u.Host, ) _, werr := conn.Write([]byte(reqLine)) require.NoError(t, werr) // 讀握手行 — 預期 "HTTP/1.1 200 Connected\r\n\r\n" reader := bufio.NewReader(conn) statusLine, err := reader.ReadString('\n') require.NoError(t, err, "failed to read status line") require.Contains(t, statusLine, "200 Connected", "expected 200 Connected, got: %q", statusLine) // 讀掉空白行(header 結束) for { line, err := reader.ReadString('\n') require.NoError(t, err) if line == "\r\n" || line == "\n" { break } } // Buffer 裡可能還有 reader 預讀的資料 — 不影響,因為後續我們會用 bufio.NewReader(conn) 再讀 // 但為了避免 reader 裡殘留的預讀資料被吃掉,caller 要自己管 bufio.NewReader // 這裡回傳 conn;caller 在 Write request 後,要用新的 bufio.NewReader(conn) 讀 response // // 注意:實務上 reader.Buffered() 應該是 0(server 還沒送 response body), // 所以直接回 conn 即可。 assert.Equal(t, 0, reader.Buffered(), "reader 不應該有預讀資料;若有則 caller 必須用此 reader 而非新建 bufio") return conn }