opsxapi/server/internal/proxy/ws_tunnel.go

package proxy

import (
	"crypto/rand"
	"encoding/hex"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

// WSTunnel 让浏览器借道 Gin 连接带自定义 Header 的 WebSocket：
//
//	POST /v1/proxy/ws/connect {url, headers, protocols} → {sessionId}
//	GET  /v1/proxy/ws/attach?sessionId=…  （浏览器 WebSocket 双向桥）
//	POST /v1/proxy/ws/close  {sessionId}
type wsTunnelSpec struct {
	URL       string            `json:"url" binding:"required"`
	Headers   map[string]string `json:"headers,omitempty"`
	Protocols []string          `json:"protocols,omitempty"`
}

type wsTunnelSession struct {
	id        string
	target    *websocket.Conn
	attach    chan *websocket.Conn
	createdAt time.Time
}

type WSTunnel struct {
	mu       sync.Mutex
	sessions map[string]*wsTunnelSession
	upgrader websocket.Upgrader
}

func NewWSTunnel() *WSTunnel {
	t := &WSTunnel{
		sessions: map[string]*wsTunnelSession{},
		upgrader: websocket.Upgrader{
			ReadBufferSize:  32 * 1024,
			WriteBufferSize: 32 * 1024,
			CheckOrigin: func(r *http.Request) bool {
				// 仅本地 UI 使用；Origin 校验交给 CORS 中间件
				return true
			},
		},
	}
	go t.gc()
	return t
}

func (t *WSTunnel) gc() {
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()
	for range ticker.C {
		t.mu.Lock()
		for id, s := range t.sessions {
			// 未 attach 或空闲超过 10 分钟的会话回收
			if time.Since(s.createdAt) > 10*time.Minute {
				_ = s.target.Close()
				delete(t.sessions, id)
			}
		}
		t.mu.Unlock()
	}
}

// ValidateWSURL 校验 ws/wss 目标（http/https 自动换算）
func ValidateWSURL(raw string) (*url.URL, error) {
	u, err := url.Parse(raw)
	if err != nil {
		return nil, err
	}
	switch u.Scheme {
	case "ws", "http":
		u.Scheme = "ws"
	case "wss", "https":
		u.Scheme = "wss"
	default:
		return nil, errOnlyWS
	}
	if u.Hostname() == "" {
		return nil, errEmptyHost
	}
	if _, blocked := blockedHosts[strings.ToLower(u.Hostname())]; blocked {
		return nil, errBlockedHost
	}
	return u, nil
}

var (
	errOnlyWS     = &wsTunnelError{"only ws/wss schemes are allowed"}
	errEmptyHost  = &wsTunnelError{"url host is required"}
	errBlockedHost = &wsTunnelError{"blocked metadata host"}
)

type wsTunnelError struct{ msg string }

func (e *wsTunnelError) Error() string { return e.msg }

func (t *WSTunnel) Connect(c *gin.Context) {
	var spec wsTunnelSpec
	if err := c.ShouldBindJSON(&spec); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid spec", "detail": err.Error()})
		return
	}
	u, err := ValidateWSURL(spec.URL)
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	dialer := websocket.Dialer{HandshakeTimeout: 15 * time.Second}
	dialHeader := http.Header{}
	for key, value := range spec.Headers {
		if !isHopByHop(key) {
			dialHeader.Set(key, value)
		}
	}
	target, _, err := dialer.Dial(u.String(), dialHeader)
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": "dial failed", "detail": err.Error()})
		return
	}

	id := randomSessionID()
	session := &wsTunnelSession{id: id, target: target, attach: make(chan *websocket.Conn, 1), createdAt: time.Now()}
	t.mu.Lock()
	t.sessions[id] = session
	t.mu.Unlock()

	c.JSON(http.StatusOK, gin.H{"sessionId": id})

	// 30 秒内没有浏览器 attach 就回收
	go func() {
		select {
		case client := <-session.attach:
			t.bridge(client, session)
		case <-time.After(30 * time.Second):
			t.mu.Lock()
			if s, ok := t.sessions[id]; ok {
				_ = s.target.Close()
				delete(t.sessions, id)
			}
			t.mu.Unlock()
		}
	}()
}

func (t *WSTunnel) Attach(c *gin.Context) {
	id := c.Query("sessionId")
	t.mu.Lock()
	session, ok := t.sessions[id]
	t.mu.Unlock()
	if !ok {
		c.JSON(http.StatusNotFound, gin.H{"error": "session not found"})
		return
	}
	client, err := t.upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		return
	}
	select {
	case session.attach <- client:
		// bridge goroutine 已在 Connect 中启动
	default:
		_ = client.Close()
	}
}

func (t *WSTunnel) Close(c *gin.Context) {
	var body struct {
		SessionID string `json:"sessionId"`
	}
	_ = c.ShouldBindJSON(&body)
	t.mu.Lock()
	session, ok := t.sessions[body.SessionID]
	if ok {
		delete(t.sessions, body.SessionID)
	}
	t.mu.Unlock()
	if ok {
		_ = session.target.Close()
	}
	c.JSON(http.StatusOK, gin.H{"ok": true})
}

// bridge 双向转发：浏览器 → 目标；目标 → 浏览器
func (t *WSTunnel) bridge(client *websocket.Conn, session *wsTunnelSession) {
	defer func() {
		_ = client.Close()
		_ = session.target.Close()
		t.mu.Lock()
		delete(t.sessions, session.id)
		t.mu.Unlock()
	}()

	done := make(chan struct{})
	// 浏览器 → 目标
	go func() {
		defer close(done)
		for {
			msgType, data, err := client.ReadMessage()
			if err != nil {
				return
			}
			if err := session.target.WriteMessage(msgType, data); err != nil {
				return
			}
		}
	}()
	// 目标 → 浏览器（bridge 协议：控制帧 + 数据帧）
	_ = client.WriteMessage(websocket.TextMessage, []byte(`{"type":"open"}`))
	for {
		msgType, data, err := session.target.ReadMessage()
		if err != nil {
			_ = client.WriteMessage(websocket.TextMessage, []byte(`{"type":"closed","reason":"target closed"}`))
			return
		}
		if err := client.WriteMessage(msgType, data); err != nil {
			return
		}
	}
}

func randomSessionID() string {
	b := make([]byte, 16)
	if _, err := rand.Read(b); err != nil {
		return time.Now().Format("20060102150405.000000000")
	}
	return hex.EncodeToString(b)
}