AutoZero/cdp/cdp.go

// Package cdp 实现跨平台 browser 工具所需的 Chrome DevTools Protocol 客户端：
// WebSocket 传输（coder/websocket）、CDP 连接层、Chrome 启动发现、browser 工具。
package cdp

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/coder/websocket"
)

// maxWSMessage 限制单条 CDP 消息大小（截图 base64 可达数 MB，留足余量）。
const maxWSMessage = 64 << 20

// Conn 是一条到 CDP target 的 WebSocket 连接：请求/响应按 id 配对，
// 事件按方法名分发给先注册的等待者（必须先 Subscribe 再触发事件）。
type Conn struct {
	ws  *websocket.Conn
	ctx context.Context // 连接生命周期（readLoop 用）

	wmu sync.Mutex // coder/websocket 要求写方串行

	mu      sync.Mutex
	nextID  int64
	pending map[int64]chan cdpReply
	subs    map[string][]*EventWaiter
	closed  bool
}

type cdpReply struct {
	Result json.RawMessage
	Err    *cdpError
}

type cdpError struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
}

type cdpEvent struct {
	Method string
	Params json.RawMessage
}

type cdpWireMsg struct {
	ID     int64           `json:"id"`
	Result json.RawMessage `json:"result"`
	Error  *cdpError       `json:"error"`
	Method string          `json:"method"`
	Params json.RawMessage `json:"params"`
}

// Dial 连接一个 target 的 WebSocket 调试端点并启动读取循环。
func Dial(ctx context.Context, wsURL string) (*Conn, error) {
	dctx, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()
	ws, _, err := websocket.Dial(dctx, wsURL, &websocket.DialOptions{CompressionMode: websocket.CompressionDisabled})
	if err != nil {
		return nil, fmt.Errorf("ws dial %s: %w", wsURL, err)
	}
	ws.SetReadLimit(maxWSMessage)
	c := &Conn{
		ws:      ws,
		ctx:     ctx,
		pending: map[int64]chan cdpReply{},
		subs:    map[string][]*EventWaiter{},
	}
	go c.readLoop()
	return c, nil
}

// Call 发送一条 CDP 命令并等待同 id 响应。
func (c *Conn) Call(ctx context.Context, method string, params any) (json.RawMessage, error) {
	c.mu.Lock()
	if c.closed {
		c.mu.Unlock()
		return nil, errors.New("cdp connection closed")
	}
	c.nextID++
	id := c.nextID
	ch := make(chan cdpReply, 1)
	c.pending[id] = ch
	c.mu.Unlock()

	msg := map[string]any{"id": id, "method": method}
	if params != nil {
		msg["params"] = params
	}
	data, err := json.Marshal(msg)
	if err != nil {
		c.removePending(id)
		return nil, err
	}
	c.wmu.Lock()
	err = c.ws.Write(ctx, websocket.MessageText, data)
	c.wmu.Unlock()
	if err != nil {
		c.removePending(id)
		return nil, fmt.Errorf("%s: %w", method, err)
	}
	select {
	case r := <-ch:
		if r.Err != nil {
			return nil, fmt.Errorf("%s: %s (code %d)", method, r.Err.Message, r.Err.Code)
		}
		return r.Result, nil
	case <-ctx.Done():
		c.removePending(id)
		return nil, ctx.Err()
	}
}

// Subscribe 注册事件等待者（必须在事件触发前调用；缓冲 8，满则丢弃）。
func (c *Conn) Subscribe(event string) *EventWaiter {
	w := &EventWaiter{ch: make(chan cdpEvent, 8), method: event}
	c.mu.Lock()
	c.subs[event] = append(c.subs[event], w)
	c.mu.Unlock()
	return w
}

// EventWaiter 等待一个已订阅的 CDP 事件。
type EventWaiter struct {
	ch     chan cdpEvent
	method string
}

// Wait 等待事件到达（超时返回错误）。
func (w *EventWaiter) Wait(ctx context.Context, timeout time.Duration) (json.RawMessage, error) {
	timer := time.NewTimer(timeout)
	defer timer.Stop()
	select {
	case ev := <-w.ch:
		return ev.Params, nil
	case <-ctx.Done():
		return nil, ctx.Err()
	case <-timer.C:
		return nil, fmt.Errorf("timeout waiting for event %s", w.method)
	}
}

func (c *Conn) removePending(id int64) {
	c.mu.Lock()
	delete(c.pending, id)
	c.mu.Unlock()
}

// readLoop 按 id 分发响应、按方法名分发事件；连接断开时唤醒所有等待者。
func (c *Conn) readLoop() {
	for {
		_, data, err := c.ws.Read(c.ctx)
		if err != nil {
			c.failAll(err)
			return
		}
		var m cdpWireMsg
		if err := json.Unmarshal(data, &m); err != nil {
			continue
		}
		if m.ID != 0 {
			c.mu.Lock()
			ch := c.pending[m.ID]
			delete(c.pending, m.ID)
			c.mu.Unlock()
			if ch != nil {
				ch <- cdpReply{Result: m.Result, Err: m.Error}
			}
			continue
		}
		if m.Method != "" {
			c.mu.Lock()
			waiters := c.subs[m.Method]
			c.mu.Unlock()
			for _, w := range waiters {
				select {
				case w.ch <- cdpEvent{Method: m.Method, Params: m.Params}:
				default: // 满则丢弃，避免阻塞 readLoop
				}
			}
		}
	}
}

func (c *Conn) failAll(err error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.closed = true
	for id, ch := range c.pending {
		ch <- cdpReply{Err: &cdpError{Message: "connection lost: " + err.Error()}}
		delete(c.pending, id)
	}
}

// Close 关闭连接（幂等）。
func (c *Conn) Close() {
	c.mu.Lock()
	if c.closed {
		c.mu.Unlock()
		return
	}
	c.closed = true
	pending := c.pending
	subs := c.subs
	c.pending = map[int64]chan cdpReply{}
	c.subs = map[string][]*EventWaiter{}
	c.mu.Unlock()
	for id, ch := range pending {
		ch <- cdpReply{Err: &cdpError{Message: "connection closed"}}
		delete(pending, id)
	}
	for _, ws := range subs {
		for _, w := range ws {
			close(w.ch)
		}
	}
	_ = c.ws.Close(websocket.StatusNormalClosure, "")
}