ModelRouter/internal/lua/plugin_browser.go

package lua

import (
	"context"
	"fmt"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

// Package lua — 服务端浏览器会话（R4-f，VNC 登录流程）
//
// 用户点明的场景：**有些平台的登录会重定向到本机回环地址**（OAuth loopback /
// RFC 8252 原生应用流程）。而：
//
//	浏览器在用户的机器上 → 重定向打到 127.0.0.1 是**用户的**回环 → 服务端收不到 code
//	浏览器在服务端       → 重定向打到 127.0.0.1 是**服务端的**回环 → 插件直接收到 code
//
// 后者不需要"提取浏览器里的凭据"——授权码会自己流到监听方。所以对回环类登录，
// 服务端浏览器不是权宜之计，而是**唯一自动化路径**。已实测验证（浏览器在服务端
// 时 loopback 回调的 code 落到了服务端监听器）。
//
// # 为什么仍然默认关闭
//
// 这不是"多此一举"，而是这套东西的风险确实最高：在服务端跑一个真实浏览器，
// 让它登录**用户的第三方账号**。所以：
//
//   - 默认关闭，需运维在设置页显式开启
//   - 每个平台单独授权（不是"开了 VNC 就能登全部"）
//   - 浏览器与 X 服务跑在独立进程组，插件卸载/会话结束即回收
//   - **不把浏览器的 CDP 端口暴露给插件 Lua**：插件拿到的是"会话 id + 打开 URL"，
//     不是任意控制浏览器的能力。否则一个插件就能读用户所有标签页的 cookie。
//
// # 与 plan 的关系
//
// 原计划把 VNC 定为"逃生舱、默认关闭"。用户指出回环场景后，它仍是逃生舱
// （默认关闭），但**不再是"最后手段"**：对回环类它是首选且唯一可行的流程。

// browserSessionTTL bounds one browser session. A server-side browser holding a
// logged-in third-party session must not outlive the login attempt.
const browserSessionTTL = 20 * time.Minute

// BrowserSession is one server-side browser exposed for interactive login.
type BrowserSession struct {
	ID       string `json:"id"`
	Plugin   string `json:"plugin"`
	Platform string `json:"platform"`
	// ViewURL is where the operator opens the VNC viewer (noVNC). It is served
	// by the gateway reverse-proxying this session's websockify port, so it
	// inherits the gateway's auth and the VNC port is never exposed directly.
	ViewURL string `json:"view_url"`
	// VNCPort / Display are internal; exposed for diagnostics only.
	VNCPort int    `json:"vnc_port"`
	Display string `json:"display"`
	// WSPort is the loopback port where noVNC's static files AND its WebSocket
	// live (websockify serves both). The gateway reverse-proxies it so the
	// viewer is reachable through the gateway's own auth instead of being
	// exposed on the network.
	WSPort int `json:"-"`
	// CallbackPort is the loopback port the login flow will be redirected to.
	// The operator never needs it, but it is the whole point of running the
	// browser server-side, so it is surfaced for troubleshooting.
	CallbackPort int       `json:"callback_port"`
	StartedAt    time.Time `json:"started_at"`
	ExpiresAt    time.Time `json:"expires_at"`
	// CapturedCode is the authorization code that arrived on the loopback
	// listener (empty until the flow completes).
	CapturedCode  string `json:"captured_code,omitempty"`
	CapturedState string `json:"captured_state,omitempty"`
	Error         string `json:"error,omitempty"`
	done          chan struct{}
}

// browserProcs holds the OS processes for one session so teardown is one place.
type browserProcs struct {
	xvfb     *exec.Cmd
	chromium *exec.Cmd
	x11vnc   *exec.Cmd
	ws       *exec.Cmd
	cbSrv    *http.Server
	profDir  string
}

// browserSupervisor owns every live browser session.
type browserSupervisor struct {
	mu       sync.Mutex
	sessions map[string]*browserSession
}

type browserSession struct {
	info  BrowserSession
	procs browserProcs
	mu    sync.Mutex
	done  chan struct{}
}

func newBrowserSupervisor() *browserSupervisor {
	return &browserSupervisor{sessions: map[string]*browserSession{}}
}

func (ps *Plugins) browsers() *browserSupervisor {
	ps.browserOnce.Do(func() {
		if ps.browserStore == nil {
			ps.browserStore = newBrowserSupervisor()
		}
	})
	return ps.browserStore
}

// loopbackCallback is what the loopback listener captured.
type loopbackCallback struct {
	code  string
	state string
	err   string
}

// StartBrowserSession launches a server-side browser for one login.
//
// It is deliberately a *lot* of moving parts (Xvfb + chromium + VNC + noVNC +
// loopback listener) in one function rather than spread out, because the failure
// modes are order-dependent: a leftover Xvfb from a crashed session makes the
// next chromium fail in a way that looks like a browser bug.
func (ps *Plugins) StartBrowserSession(plugin, platform, startURL string, callbackPort int) (*BrowserSession, error) {
	if !ps.VNCEnabled() {
		return nil, fmt.Errorf("server-side browser login is disabled")
	}
	if ps.find(plugin) == nil {
		return nil, fmt.Errorf("plugin %s not found", plugin)
	}
	if startURL == "" {
		return nil, fmt.Errorf("the browser flow needs a start URL")
	}
	if callbackPort == 0 {
		// Pick a free loopback port for the OAuth redirect to land on.
		p, err := freeLoopbackPort()
		if err != nil {
			return nil, err
		}
		callbackPort = p
	}

	display := fmt.Sprintf(":%d", 90+len(ps.browsers().sessions)%10)
	vncPort, err := freeLoopbackPort()
	if err != nil {
		return nil, err
	}
	wsPort, err := freeLoopbackPort()
	if err != nil {
		return nil, err
	}
	profDir, err := os.MkdirTemp("", "llmsproxy-browser-*")
	if err != nil {
		return nil, err
	}
	if err := os.Chmod(profDir, 0o700); err != nil {
		_ = os.RemoveAll(profDir)
		return nil, err
	}

	sess := &browserSession{
		info: BrowserSession{
			ID:           newLoginID(), // same unguessable-id generator
			Plugin:       plugin,
			Platform:     platform,
			VNCPort:      vncPort,
			WSPort:       wsPort,
			Display:      display,
			CallbackPort: callbackPort,
			StartedAt:    time.Now(),
			ExpiresAt:    time.Now().Add(browserSessionTTL),
		},
		procs: browserProcs{profDir: profDir},
		done:  make(chan struct{}),
	}

	if err := ps.launchBrowserStack(sess, startURL, callbackPort, wsPort, vncPort, display, profDir); err != nil {
		sess.teardown()
		return nil, err
	}

	ps.browsers().mu.Lock()
	ps.browsers().sessions[sess.info.ID] = sess
	ps.browsers().mu.Unlock()

	// The viewer is mounted under the plugin's own namespace so it inherits the
	// gateway's authentication. Exposing VNC directly would hand anyone who can
	// reach the port a logged-in browser.
	// A relative query so noVNC's bundled JS connects its WebSocket back through
	// THIS path (the gateway proxy) rather than guessing a host:port.
	sess.info.ViewURL = fmt.Sprintf("/browser/%s/vnc.html?autoconnect=1&path=browser/%s/websockify",
		sess.info.ID, sess.info.ID)
	out := sess.info
	return &out, nil
}

// launchBrowserStack starts the five processes and wires the loopback capture.
func (ps *Plugins) launchBrowserStack(sess *browserSession, startURL string, callbackPort, wsPort, vncPort int, display, profDir string) error {
	// 1. Virtual display.
	sess.procs.xvfb = exec.Command("Xvfb", display, "-screen", "0", "1280x800x24", "-nolisten", "tcp")
	sess.procs.xvfb.Stdout, sess.procs.xvfb.Stderr = nil, nil
	if err := sess.procs.xvfb.Start(); err != nil {
		return fmt.Errorf("start Xvfb (is it installed?): %w", err)
	}
	// Xvfb needs a moment before the display exists.
	if err := waitForDisplay(display, 5*time.Second); err != nil {
		return err
	}

	// 2. Loopback callback listener — the reason this whole thing exists.
	cb := make(chan loopbackCallback, 1)
	mux := http.NewServeMux()
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		q := r.URL.Query()
		cb <- loopbackCallback{code: q.Get("code"), state: q.Get("state"), err: q.Get("error")}
		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		fmt.Fprint(w, "<h2>Login received. You can close this tab.</h2>")
	})
	srv := &http.Server{Addr: fmt.Sprintf("127.0.0.1:%d", callbackPort), Handler: mux}
	ln, err := net.Listen("tcp", srv.Addr)
	if err != nil {
		return fmt.Errorf("bind loopback callback port %d: %w", callbackPort, err)
	}
	sess.procs.cbSrv = srv
	go func() { _ = srv.Serve(ln) }()

	// Capture the callback into the session record.
	go func() {
		select {
		case c := <-cb:
			sess.mu.Lock()
			sess.info.CapturedCode, sess.info.CapturedState, sess.info.Error = c.code, c.state, c.err
			sess.mu.Unlock()
		case <-sess.done:
		}
	}()

	// 3. Real browser on the virtual display.
	args := []string{
		"--no-sandbox", // the container has no usable sandbox; see the security note
		"--disable-dev-shm-usage",
		"--user-data-dir=" + profDir,
		"--no-first-run",
		"--no-default-browser-check",
		"--window-size=1280,800",
		startURL,
	}
	chromium, err := findChromium()
	if err != nil {
		return err
	}
	sess.procs.chromium = exec.Command(chromium, args...)
	sess.procs.chromium.Env = append(os.Environ(), "DISPLAY="+display)
	if err := sess.procs.chromium.Start(); err != nil {
		return fmt.Errorf("start chromium: %w", err)
	}

	// 4. VNC server on the same display.
	sess.procs.x11vnc = exec.Command("x11vnc", "-display", display,
		"-rfbport", fmt.Sprint(vncPort), "-localhost", "-nopw", "-forever", "-shared", "-quiet")
	if err := sess.procs.x11vnc.Start(); err != nil {
		return fmt.Errorf("start x11vnc: %w", err)
	}

	// 5. websockify bridges the browser's WebSocket to raw VNC for noVNC.
	sess.procs.ws = exec.Command("websockify", "--web", "/usr/share/novnc",
		fmt.Sprint(wsPort), fmt.Sprintf("127.0.0.1:%d", vncPort))
	if err := sess.procs.ws.Start(); err != nil {
		return fmt.Errorf("start websockify: %w", err)
	}

	return nil
}

// teardown kills every process and removes the profile (which holds the session
// cookies for the logged-in account).
func (s *browserSession) teardown() {
	func() {
		defer func() { _ = recover() }()
		close(s.done)
	}()
	if s.procs.cbSrv != nil {
		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
		_ = s.procs.cbSrv.Shutdown(ctx)
		cancel()
	}
	for _, c := range []*exec.Cmd{s.procs.ws, s.procs.x11vnc, s.procs.chromium, s.procs.xvfb} {
		if c == nil || c.Process == nil {
			continue
		}
		_ = c.Process.Kill()
		_, _ = c.Process.Wait()
	}
	if s.procs.profDir != "" {
		// The profile holds live third-party session cookies: removing it is
		// part of the security contract, not housekeeping.
		_ = os.RemoveAll(s.procs.profDir)
	}
}

// GetBrowserSession returns a copy for the UI.
func (ps *Plugins) GetBrowserSession(id string) (*BrowserSession, bool) {
	bs := ps.browsers()
	bs.mu.Lock()
	s, ok := bs.sessions[id]
	bs.mu.Unlock()
	if !ok {
		return nil, false
	}
	s.mu.Lock()
	cp := s.info
	s.mu.Unlock()
	// Expiry applied on read, same as login sessions.
	if time.Now().After(cp.ExpiresAt) {
		ps.StopBrowserSession(id)
		cp.Error = "the browser session timed out"
	}
	return &cp, true
}

// StopBrowserSession tears one session down and forgets it.
func (ps *Plugins) StopBrowserSession(id string) {
	bs := ps.browsers()
	bs.mu.Lock()
	s, ok := bs.sessions[id]
	if ok {
		delete(bs.sessions, id)
	}
	bs.mu.Unlock()
	if ok {
		s.teardown()
	}
}

// StopBrowserSessionsFor tears down every session a plugin owns (unload path).
func (ps *Plugins) StopBrowserSessionsFor(plugin string) int {
	bs := ps.browsers()
	bs.mu.Lock()
	var victims []*browserSession
	for id, s := range bs.sessions {
		if s.info.Plugin == plugin {
			victims = append(victims, s)
			delete(bs.sessions, id)
		}
	}
	bs.mu.Unlock()
	for _, s := range victims {
		s.teardown()
	}
	return len(victims)
}

// ---- helpers ----

func findChromium() (string, error) {
	for _, name := range []string{"chromium", "chromium-browser", "google-chrome"} {
		if p, err := exec.LookPath(name); err == nil {
			return p, nil
		}
	}
	return "", fmt.Errorf("no chromium found on PATH; the server-side browser flow needs one")
}

func freeLoopbackPort() (int, error) {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return 0, err
	}
	defer l.Close()
	return l.Addr().(*net.TCPAddr).Port, nil
}

func waitForDisplay(display string, timeout time.Duration) error {
	// Xvfb does not signal readiness; probing the unix socket is the reliable way.
	sock := filepath.Join("/tmp/.X11-unix", "X"+strings.TrimPrefix(display, ":"))
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		if _, err := os.Stat(sock); err == nil {
			return nil
		}
		time.Sleep(50 * time.Millisecond)
	}
	return fmt.Errorf("display %s did not come up within %s", display, timeout)
}