| 1 | package cdp |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "errors" |
| 7 | "fmt" |
| 8 | "net/http" |
| 9 | "sync" |
| 10 | "time" |
| 11 | |
| 12 | "github.com/gorilla/websocket" |
| 13 | ) |
| 14 | |
| 15 | const ( |
| 16 | // Screenshots and accessibility trees travel inline on this socket, so the |
| 17 | // read limit sits well above any control frame. |
| 18 | maxFrameBytes = 96 << 20 |
| 19 | writeTimeout = 15 * time.Second |
| 20 | ) |
| 21 | |
| 22 | // errConnClosed reports that the DevTools socket is gone; callers translate it |
| 23 | // into an unknown outcome for writes and a plain failure for reads. |
| 24 | var errConnClosed = errors.New("cdp: devtools connection closed") |
| 25 | |
| 26 | // message is one DevTools frame in either direction. Commands carry ID and |
| 27 | // Method, replies carry ID with Result or Error, and events carry Method only. |
| 28 | type message struct { |
| 29 | ID int64 `json:"id,omitempty"` |
| 30 | Method string `json:"method,omitempty"` |
| 31 | Params json.RawMessage `json:"params,omitempty"` |
| 32 | Result json.RawMessage `json:"result,omitempty"` |
| 33 | Error *protocolError `json:"error,omitempty"` |
| 34 | SessionID string `json:"sessionId,omitempty"` |
| 35 | } |
| 36 | |
| 37 | // protocolError is Chrome's refusal of one command. |
| 38 | type protocolError struct { |
| 39 | Code int `json:"code"` |
| 40 | Message string `json:"message"` |
| 41 | Data string `json:"data,omitempty"` |
| 42 | } |
| 43 | |
| 44 | func (e *protocolError) Error() string { |
| 45 | if e.Data == "" { |
| 46 | return fmt.Sprintf("cdp: %s (%d)", e.Message, e.Code) |
| 47 | } |
| 48 | return fmt.Sprintf("cdp: %s: %s (%d)", e.Message, e.Data, e.Code) |
| 49 | } |
| 50 | |
| 51 | // handler observes events of one method on one session; a true return |
| 52 | // unsubscribes it. Handlers run on the read loop and must never call conn. |
| 53 | type handler struct { |
| 54 | session string |
| 55 | method string |
| 56 | fn func(params json.RawMessage) bool |
| 57 | } |
| 58 | |
| 59 | // conn multiplexes commands and events over one DevTools WebSocket. Flat |
| 60 | // sessions put every target on this socket, so IDs are unique across sessions |
| 61 | // and only the sessionId field says which target answered. |
| 62 | type conn struct { |
| 63 | ws *websocket.Conn |
| 64 | writeMu sync.Mutex |
| 65 | |
| 66 | mu sync.Mutex |
| 67 | nextID int64 |
| 68 | pending map[int64]chan message |
| 69 | handlers map[int]handler |
| 70 | nextHnd int |
| 71 | err error |
| 72 | |
| 73 | done chan struct{} |
| 74 | } |
| 75 | |
| 76 | // dialConn opens the DevTools socket at wsURL and starts its read loop. |
| 77 | func dialConn(ctx context.Context, wsURL string) (*conn, error) { |
| 78 | dialer := websocket.Dialer{HandshakeTimeout: 20 * time.Second, ReadBufferSize: 64 << 10, WriteBufferSize: 64 << 10} |
| 79 | ws, resp, err := dialer.DialContext(ctx, wsURL, http.Header{}) |
| 80 | if err != nil { |
| 81 | if resp != nil { |
| 82 | return nil, fmt.Errorf("cdp: dial %s: %w (status %d)", wsURL, err, resp.StatusCode) |
| 83 | } |
| 84 | return nil, fmt.Errorf("cdp: dial %s: %w", wsURL, err) |
| 85 | } |
| 86 | ws.SetReadLimit(maxFrameBytes) |
| 87 | c := &conn{ws: ws, pending: map[int64]chan message{}, handlers: map[int]handler{}, done: make(chan struct{})} |
| 88 | go c.readLoop() |
| 89 | return c, nil |
| 90 | } |
| 91 | |
| 92 | func (c *conn) readLoop() { |
| 93 | for { |
| 94 | _, data, err := c.ws.ReadMessage() |
| 95 | if err != nil { |
| 96 | c.shutdown(err) |
| 97 | return |
| 98 | } |
| 99 | var msg message |
| 100 | if err := json.Unmarshal(data, &msg); err != nil { |
| 101 | continue |
| 102 | } |
| 103 | if msg.ID != 0 { |
| 104 | c.deliver(msg) |
| 105 | continue |
| 106 | } |
| 107 | c.dispatch(msg) |
| 108 | } |
| 109 | } |
| 110 | |
| 111 | func (c *conn) deliver(msg message) { |
| 112 | c.mu.Lock() |
| 113 | ch, ok := c.pending[msg.ID] |
| 114 | delete(c.pending, msg.ID) |
| 115 | c.mu.Unlock() |
| 116 | if ok { |
| 117 | ch <- msg |
| 118 | } |
| 119 | } |
| 120 | |
| 121 | // dispatch fans one event out to its handlers. The snapshot is taken under the |
| 122 | // lock and the handlers run without it, so a handler may unsubscribe itself. |
| 123 | func (c *conn) dispatch(msg message) { |
| 124 | c.mu.Lock() |
| 125 | matched := make([]int, 0, 4) |
| 126 | for id, h := range c.handlers { |
| 127 | if h.method == msg.Method && (h.session == "" || h.session == msg.SessionID) { |
| 128 | matched = append(matched, id) |
| 129 | } |
| 130 | } |
| 131 | fns := make(map[int]func(json.RawMessage) bool, len(matched)) |
| 132 | for _, id := range matched { |
| 133 | fns[id] = c.handlers[id].fn |
| 134 | } |
| 135 | c.mu.Unlock() |
| 136 | for id, fn := range fns { |
| 137 | if fn(msg.Params) { |
| 138 | c.removeHandler(id) |
| 139 | } |
| 140 | } |
| 141 | } |
| 142 | |
| 143 | func (c *conn) shutdown(cause error) { |
| 144 | c.mu.Lock() |
| 145 | if c.err != nil { |
| 146 | c.mu.Unlock() |
| 147 | return |
| 148 | } |
| 149 | if cause == nil { |
| 150 | cause = errConnClosed |
| 151 | } |
| 152 | c.err = cause |
| 153 | pending := c.pending |
| 154 | c.pending = map[int64]chan message{} |
| 155 | c.mu.Unlock() |
| 156 | close(c.done) |
| 157 | for _, ch := range pending { |
| 158 | close(ch) |
| 159 | } |
| 160 | _ = c.ws.Close() |
| 161 | } |
| 162 | |
| 163 | // close tears the socket down; in-flight callers observe errConnClosed. |
| 164 | func (c *conn) close() { |
| 165 | c.shutdown(errConnClosed) |
| 166 | } |
| 167 | |
| 168 | func (c *conn) closed() bool { |
| 169 | select { |
| 170 | case <-c.done: |
| 171 | return true |
| 172 | default: |
| 173 | return false |
| 174 | } |
| 175 | } |
| 176 | |
| 177 | func (c *conn) addHandler(session, method string, fn func(json.RawMessage) bool) int { |
| 178 | c.mu.Lock() |
| 179 | defer c.mu.Unlock() |
| 180 | c.nextHnd++ |
| 181 | id := c.nextHnd |
| 182 | c.handlers[id] = handler{session: session, method: method, fn: fn} |
| 183 | return id |
| 184 | } |
| 185 | |
| 186 | func (c *conn) removeHandler(id int) { |
| 187 | c.mu.Lock() |
| 188 | delete(c.handlers, id) |
| 189 | c.mu.Unlock() |
| 190 | } |
| 191 | |
| 192 | // on subscribes to every event of method on session ("" for any session) and |
| 193 | // returns the unsubscribe function. |
| 194 | func (c *conn) on(session, method string, fn func(json.RawMessage)) func() { |
| 195 | id := c.addHandler(session, method, func(p json.RawMessage) bool { |
| 196 | fn(p) |
| 197 | return false |
| 198 | }) |
| 199 | return func() { c.removeHandler(id) } |
| 200 | } |
| 201 | |
| 202 | // once delivers the first matching event and unsubscribes. The returned |
| 203 | // cancel must run even when the event arrives, so callers defer it. |
| 204 | func (c *conn) once(session, method string) (<-chan json.RawMessage, func()) { |
| 205 | ch := make(chan json.RawMessage, 1) |
| 206 | id := c.addHandler(session, method, func(p json.RawMessage) bool { |
| 207 | ch <- p |
| 208 | return true |
| 209 | }) |
| 210 | return ch, func() { c.removeHandler(id) } |
| 211 | } |
| 212 | |
| 213 | // call sends one command and waits for its reply. A cancelled context or a |
| 214 | // dead socket returns without a verdict, which write paths translate into |
| 215 | // browser.ErrUnknownOutcome. |
| 216 | func (c *conn) call(ctx context.Context, session, method string, params, out any) error { |
| 217 | body, err := encodeParams(params) |
| 218 | if err != nil { |
| 219 | return fmt.Errorf("cdp: encode %s: %w", method, err) |
| 220 | } |
| 221 | reply := make(chan message, 1) |
| 222 | c.mu.Lock() |
| 223 | if c.err != nil { |
| 224 | c.mu.Unlock() |
| 225 | return c.err |
| 226 | } |
| 227 | c.nextID++ |
| 228 | id := c.nextID |
| 229 | c.pending[id] = reply |
| 230 | c.mu.Unlock() |
| 231 | |
| 232 | if err := c.write(message{ID: id, Method: method, Params: body, SessionID: session}); err != nil { |
| 233 | c.mu.Lock() |
| 234 | delete(c.pending, id) |
| 235 | c.mu.Unlock() |
| 236 | return err |
| 237 | } |
| 238 | select { |
| 239 | case <-ctx.Done(): |
| 240 | c.mu.Lock() |
| 241 | delete(c.pending, id) |
| 242 | c.mu.Unlock() |
| 243 | return ctx.Err() |
| 244 | case <-c.done: |
| 245 | return c.failure() |
| 246 | case msg, ok := <-reply: |
| 247 | if !ok { |
| 248 | return c.failure() |
| 249 | } |
| 250 | if msg.Error != nil { |
| 251 | return msg.Error |
| 252 | } |
| 253 | if out == nil { |
| 254 | return nil |
| 255 | } |
| 256 | if err := json.Unmarshal(msg.Result, out); err != nil { |
| 257 | return fmt.Errorf("cdp: decode %s reply: %w", method, err) |
| 258 | } |
| 259 | return nil |
| 260 | } |
| 261 | } |
| 262 | |
| 263 | func (c *conn) failure() error { |
| 264 | c.mu.Lock() |
| 265 | defer c.mu.Unlock() |
| 266 | if c.err != nil { |
| 267 | return c.err |
| 268 | } |
| 269 | return errConnClosed |
| 270 | } |
| 271 | |
| 272 | func (c *conn) write(msg message) error { |
| 273 | data, err := json.Marshal(msg) |
| 274 | if err != nil { |
| 275 | return fmt.Errorf("cdp: encode %s: %w", msg.Method, err) |
| 276 | } |
| 277 | c.writeMu.Lock() |
| 278 | defer c.writeMu.Unlock() |
| 279 | if err := c.ws.SetWriteDeadline(time.Now().Add(writeTimeout)); err != nil { |
| 280 | return fmt.Errorf("cdp: %s: %w", msg.Method, err) |
| 281 | } |
| 282 | if err := c.ws.WriteMessage(websocket.TextMessage, data); err != nil { |
| 283 | return fmt.Errorf("cdp: %s: %w", msg.Method, err) |
| 284 | } |
| 285 | return nil |
| 286 | } |
| 287 | |
| 288 | func encodeParams(params any) (json.RawMessage, error) { |
| 289 | if params == nil { |
| 290 | return json.RawMessage(`{}`), nil |
| 291 | } |
| 292 | return json.Marshal(params) |
| 293 | } |
| 294 |