visionA/visionA-backend/internal/relay/integration_raw_test.go

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
}