visionA/visionA-backend/internal/session/forwarder.go

// forwarder.go — api-server → remote-proxy 的 raw forward client。
//
// 這是雛形雙 binary 架構下「api-server 把前端 HTTP 請求轉發到 local agent」
// 的核心元件。
//
// 整條路徑：
//
//	browser ─HTTP─► api-server handler
//	                   │
//	                   │ Forwarder.ForwardHTTP / OpenStream
//	                   ▼
//	                raw TCP dial remote-proxy: POST /internal/forward/raw?token=...
//	                   │ (B3 Major-1 修復後新增的 hijack endpoint)
//	                   ▼
//	                remote-proxy hijack 自己的連線 → yamux.OpenStream → 雙向 io.Copy
//	                   │
//	                   ▼
//	                local agent (yamux client) 把 stream 上的 HTTP request
//	                   ▼
//	                轉到本地 127.0.0.1:3721（local-tool）回 response
//
// 對齊 `.autoflow/04-architecture/api/api-internal.md` §POST /internal/forward/raw
// 與 `.autoflow/04-architecture/tunnel.md` §3.3。

package session

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"
)

// defaultDialTimeout 是 raw TCP dial remote-proxy 的最大等待時間。
const defaultDialTimeout = 10 * time.Second

// defaultHandshakeTimeout 是讀取「HTTP/1.1 200 Connected」握手的最大等待時間。
const defaultHandshakeTimeout = 10 * time.Second

// Forwarder 把 api-server 的 HTTP 請求 forward 到 remote-proxy。
//
// 並發安全：本 struct 的方法不共享可變狀態，每個 OpenStream 走獨立 net.Conn；
// 多個 goroutine 可同時呼叫。
type Forwarder struct {
	// proxyHost 是從 baseURL 解析出來的 host:port，供 net.Dial 用。
	proxyHost string

	// dialer 用於 raw TCP dial。獨立成欄位以利測試 / 未來換成 TLS dial。
	dialer net.Dialer

	logger *slog.Logger
}

// NewForwarder 從 baseURL（例：http://localhost:3801）建立 Forwarder。
//
// baseURL 必須是 http:// 或 https:// 開頭；其他 scheme 視為錯誤但延遲到
// 第一次呼叫時才回（保持建構簽章簡單）。
//
// **注意**：雛形 internal port 是純 HTTP（network policy 阻擋外部存取，見
// api-internal.md §安全）。Phase 1 加 mTLS 時，本 Forwarder 需擴充支援 TLS。
func NewForwarder(baseURL string, logger *slog.Logger) *Forwarder {
	if logger == nil {
		logger = slog.Default()
	}
	host := parseHostFromBaseURL(baseURL)
	return &Forwarder{
		proxyHost: host,
		dialer:    net.Dialer{Timeout: defaultDialTimeout},
		logger:    logger,
	}
}

// parseHostFromBaseURL 從 baseURL 取出 host:port，失敗時回傳空字串
// （後續 OpenStream 會拒絕並回明確錯誤）。
func parseHostFromBaseURL(baseURL string) string {
	if baseURL == "" {
		return ""
	}
	u, err := url.Parse(baseURL)
	if err != nil {
		return ""
	}
	return u.Host
}

// OpenStream 對 remote-proxy 開一條 raw TCP 連線，完成 hijack 握手，並回傳
// 一條可以直接用 net.Conn 語意操作的連線（底層是 yamux stream）。
//
// 用法（典型 api-server handler）：
//
//	conn, err := forwarder.OpenStream(ctx, sessionToken)
//	if err != nil { ... }
//	defer conn.Close()
//
//	httpReq.Write(conn)                                  // 送 HTTP request
//	resp, _ := http.ReadResponse(bufio.NewReader(conn), httpReq)
//	io.Copy(browserResponseWriter, resp.Body)            // streaming friendly
//
// 失敗回傳的 error：
//   - ErrSessionNotFound：remote-proxy 在 hijack 前回 502 JSON
//   - 其他 wrapped error：dial / 握手 / 解析錯誤
//
// 注意：caller 拿到 conn 後**必須自己負責 Close**；本函式內部不會 set deadline，
// 因為 streaming 場景（MJPEG / SSE）需要無限長的存活時間。
func (f *Forwarder) OpenStream(ctx context.Context, sessionToken string) (net.Conn, error) {
	if f.proxyHost == "" {
		return nil, errors.New("session: forwarder has no proxy host (check VISIONA_PROXY_INTERNAL_URL)")
	}
	if sessionToken == "" {
		return nil, errors.New("session: forwarder.OpenStream requires non-empty sessionToken")
	}

	// 1. raw TCP dial
	conn, err := f.dialer.DialContext(ctx, "tcp", f.proxyHost)
	if err != nil {
		return nil, fmt.Errorf("session: dial remote-proxy %s: %w", f.proxyHost, err)
	}

	// 2. 寫 POST /internal/forward/raw?token=...
	// 仿 dialRawForward 測試 helper 的格式（見 internal/relay/integration_raw_test.go）。
	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",
		url.QueryEscape(sessionToken), f.proxyHost,
	)
	// 設一個短的握手 deadline，避免 remote-proxy 假死時 hang 住。
	if err := conn.SetWriteDeadline(time.Now().Add(defaultHandshakeTimeout)); err != nil {
		_ = conn.Close()
		return nil, fmt.Errorf("session: set write deadline: %w", err)
	}
	if _, err := conn.Write([]byte(reqLine)); err != nil {
		_ = conn.Close()
		return nil, fmt.Errorf("session: write forward request: %w", err)
	}

	// 3. 讀握手 — 預期 "HTTP/1.1 200 Connected\r\n\r\n"
	if err := conn.SetReadDeadline(time.Now().Add(defaultHandshakeTimeout)); err != nil {
		_ = conn.Close()
		return nil, fmt.Errorf("session: set read deadline: %w", err)
	}
	reader := bufio.NewReader(conn)
	statusLine, err := reader.ReadString('\n')
	if err != nil {
		_ = conn.Close()
		return nil, fmt.Errorf("session: read handshake status: %w", err)
	}
	statusLine = strings.TrimRight(statusLine, "\r\n")

	// 解析 status code
	// 格式：HTTP/1.1 200 Connected   或   HTTP/1.1 502 Bad Gateway
	if !strings.HasPrefix(statusLine, "HTTP/1.1 200") {
		// 非 200 → 把 body 讀出來幫 debug；常見：502 = TUNNEL_DISCONNECTED
		bodyHint := drainAndPeek(reader)
		_ = conn.Close()

		// session 不存在的明確錯誤對應 ErrSessionNotFound
		if strings.Contains(statusLine, "502") {
			return nil, fmt.Errorf("%w: remote-proxy responded %q (body hint: %s)",
				ErrSessionNotFound, statusLine, bodyHint)
		}
		return nil, fmt.Errorf("session: forward handshake failed: %q (body hint: %s)",
			statusLine, bodyHint)
	}

	// 4. 把握手後的 header 讀完（一直讀到空行）
	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			_ = conn.Close()
			return nil, fmt.Errorf("session: read handshake headers: %w", err)
		}
		if line == "\r\n" || line == "\n" {
			break
		}
	}

	// 5. 清掉 deadline，因為後續 streaming 場景不該再 timeout
	if err := conn.SetDeadline(time.Time{}); err != nil {
		_ = conn.Close()
		return nil, fmt.Errorf("session: clear deadline: %w", err)
	}

	// 6. 如果 reader 裡還有預讀資料（bufio.NewReader 可能讀超過一行），
	//    回傳一個包裝 conn 把預讀的 byte 接回 stream。
	//    這個情境在 raw forward 上理論上不會發生（remote-proxy 在發出
	//    "200 Connected\r\n\r\n" 之後不會主動寫資料 — 它要等 caller 寫
	//    request 才會從 yamux stream 收 response）；但保險起見處理。
	if buffered := reader.Buffered(); buffered > 0 {
		peek, _ := reader.Peek(buffered)
		f.logger.Warn("forwarder: unexpected bytes after handshake; wrapping conn",
			"bytes", buffered)
		return newPrefixConn(conn, append([]byte(nil), peek...)), nil
	}

	return conn, nil
}

// ForwardHTTP 是「給定 http.Request，回傳 *http.Response」的高階 helper。
//
// 內部實作：
//  1. OpenStream 拿 raw TCP（已 hijack）連線
//  2. req.Write(conn) 把完整 HTTP request 寫進去
//  3. http.ReadResponse 讀出 response（不消耗 body）
//
// 重要：response.Body **包住 conn 本身**（所以 caller 必須在用完後 Close
// response.Body）；這允許 streaming body（MJPEG / SSE / chunked）原樣轉發。
//
// req 的 URL.Host / Scheme 會被覆寫成 "127.0.0.1" / "http"，因為 local agent
// 收到的是「打到自己 localhost」的請求；caller 設定的 Host header 會被保留。
func (f *Forwarder) ForwardHTTP(ctx context.Context, sessionToken string, req *http.Request) (*http.Response, error) {
	if req == nil {
		return nil, errors.New("session: ForwardHTTP requires non-nil req")
	}

	conn, err := f.OpenStream(ctx, sessionToken)
	if err != nil {
		return nil, err
	}

	// 改寫 req 為「打給 local agent」格式：
	//   - URL.Scheme = http，URL.Host = 127.0.0.1 → req.Write 才不會報錯
	//   - RequestURI 必須清空（client 端不能設）
	//   - 不覆寫 req.Host：caller 自行決定要不要保留 browser 的 Host
	//
	// 注意：req 本身可能已被外部使用，這裡複製 URL 避免副作用。
	outReq := req.Clone(ctx)
	if outReq.URL == nil {
		outReq.URL = &url.URL{}
	}
	outReq.URL.Scheme = "http"
	outReq.URL.Host = "127.0.0.1"
	outReq.RequestURI = ""
	if outReq.Host == "" {
		outReq.Host = "127.0.0.1"
	}

	// 把 request 寫到 conn
	if err := outReq.Write(conn); err != nil {
		_ = conn.Close()
		return nil, fmt.Errorf("session: write request to forwarded conn: %w", err)
	}

	// 讀 response — 不可以 close conn，因為 response.Body 還會用到
	resp, err := http.ReadResponse(bufio.NewReader(conn), outReq)
	if err != nil {
		_ = conn.Close()
		return nil, fmt.Errorf("session: read response from forwarded conn: %w", err)
	}

	// 把 conn 包進 response.Body 的 close chain：caller close body 時連 conn 一起關
	resp.Body = &bodyWithConn{ReadCloser: resp.Body, conn: conn}
	return resp, nil
}

// ForwardWebSocket 預留 — B5 接前端 WS 時實作。
//
// 預期實作（草稿）：
//   - OpenStream 拿到 raw conn
//   - 把 WS upgrade request 透過 conn 寫過去
//   - 等 101 response 回來
//   - Hijack browser 端連線，與 conn 雙向 pipe
//
// 雛形先回 ErrNotImplemented，避免被誤用。
func (f *Forwarder) ForwardWebSocket(ctx context.Context, sessionToken string, req *http.Request) (net.Conn, error) {
	return nil, errors.New("session: ForwardWebSocket not implemented yet (TODO B5)")
}

// ----------------------------------------------------------------------
// Helpers
// ----------------------------------------------------------------------

// drainAndPeek 嘗試讀少量 byte 給 error message 加上 context；
// 不阻塞太久，最多 256 byte。
//
// 呼叫前提：caller 必須已經對 underlying conn 設過 ReadDeadline（這個函式只
// 在 OpenStream 握手失敗的 error path 被呼叫，該路徑已經 SetReadDeadline
// 到 defaultHandshakeTimeout），所以 Read 不會 hang 住；若 deadline 已過，
// Read 會立刻回 0 + deadline error，行為仍然是「不阻塞」。
func drainAndPeek(reader *bufio.Reader) string {
	buf := make([]byte, 256)
	n, _ := reader.Read(buf)
	return strings.TrimSpace(string(buf[:n]))
}

// bodyWithConn 把 ReadCloser 與底層 net.Conn 綁在一起，
// caller close body 時順便關 conn（避免 leak）。
type bodyWithConn struct {
	io.ReadCloser
	conn net.Conn
}

// Close 同時關閉 body 與底層 conn；以最後一個非 nil 的 error 回傳。
func (b *bodyWithConn) Close() error {
	bodyErr := b.ReadCloser.Close()
	connErr := b.conn.Close()
	if bodyErr != nil {
		return bodyErr
	}
	return connErr
}

// prefixConn 把預讀的 byte 接回 net.Conn 開頭，供 caller 透明使用。
//
// 並發說明：net.Conn 本身對單一 goroutine 讀 + 單一 goroutine 寫是安全的。
// prefixConn 只包裝 Read；prefix 的讀取不會跨 goroutine 共享（Read 慣例上
// 只由 reader goroutine 呼叫），所以這裡不需要額外的 mutex。
type prefixConn struct {
	net.Conn
	prefix []byte
}

func newPrefixConn(c net.Conn, prefix []byte) *prefixConn {
	return &prefixConn{Conn: c, prefix: prefix}
}

func (p *prefixConn) Read(b []byte) (int, error) {
	if len(p.prefix) > 0 {
		n := copy(b, p.prefix)
		p.prefix = p.prefix[n:]
		return n, nil
	}
	return p.Conn.Read(b)
}