fn_registry/functions/browser/cdp_conn.go

package browser

import (
	"bufio"
	"crypto/rand"
	"crypto/sha1"
	"encoding/base64"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

// cdpCmdTimeout es el tope que sendCDP espera por la respuesta a un comando antes
// de rendirse. Sin el, una respuesta que Chrome nunca envia (tab cerrada a media
// peticion, proceso colgado) bloquearia la goroutine del tool para siempre — el
// agente lo percibe como "lentitud infinita". Con el timeout, el tool falla limpio
// y el retry de withConn puede reconectar.
const cdpCmdTimeout = 30 * time.Second

// EventHandler es invocado cuando llega un evento CDP del metodo subscrito.
// El handler corre en la goroutine del readLoop — debe ser rapido o despachar
// a un canal/goroutine propio. params puede ser nil si Chrome no envia.
type EventHandler func(method string, params map[string]any)

// CDPConn es una conexion activa al Chrome DevTools Protocol.
// Gestiona el WebSocket raw y el protocolo JSON-RPC de CDP.
type CDPConn struct {
	conn     net.Conn
	reader   *bufio.Reader
	mu       sync.Mutex
	nextID   atomic.Int64
	port     int
	pid      int
	pending  map[int64]chan cdpResponse
	pendMu   sync.Mutex
	closed   bool
	handlers map[string][]EventHandler
	hMu      sync.Mutex

	// axEnabled/netEnabled/pageEnabled cachean si ya enviamos el enable de cada
	// dominio CDP en esta conexion. enable/disable es idempotente pero cuesta un
	// round-trip; en el hot path del agente (percibir->actuar repetido) re-enviar
	// Accessibility.enable / Network.enable en cada llamada duplica los RTT.
	// Habilitar una vez y cachear el flag elimina ese coste por percepcion/espera.
	axEnabled   atomic.Bool
	netEnabled  atomic.Bool
	pageEnabled atomic.Bool

	// frameCtx cachea el executionContextId del isolated world por frameID, para
	// que CdpEvalInFrame no cree un mundo aislado nuevo en cada llamada.
	// frameCtxMu protege solo el lazy-init del puntero (el cache tiene su mutex).
	frameCtx   *frameCtxCache
	frameCtxMu sync.Mutex
}

type cdpRequest struct {
	ID     int64          `json:"id"`
	Method string         `json:"method"`
	Params map[string]any `json:"params,omitempty"`
}

type cdpResponse struct {
	ID     int64          `json:"id"`
	Result map[string]any `json:"result"`
	Error  *cdpError      `json:"error"`
	Method string         `json:"method"` // para eventos
	Params map[string]any `json:"params"` // para eventos
}

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

// cdpVersionResponse es la respuesta de /json/version del endpoint CDP.
type cdpVersionResponse struct {
	WebSocketDebuggerURL string `json:"webSocketDebuggerUrl"`
	Browser              string `json:"Browser"`
}

// wsHandshake realiza el handshake WebSocket RFC 6455 sobre una conexion TCP ya abierta.
func wsHandshake(conn net.Conn, host, path string) (*bufio.Reader, error) {
	// Generar clave aleatoria de 16 bytes en base64
	keyBytes := make([]byte, 16)
	if _, err := rand.Read(keyBytes); err != nil {
		return nil, fmt.Errorf("ws handshake: generar clave: %w", err)
	}
	key := base64.StdEncoding.EncodeToString(keyBytes)

	// Enviar request HTTP upgrade
	req := fmt.Sprintf(
		"GET %s HTTP/1.1\r\nHost: %s\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Key: %s\r\nSec-WebSocket-Version: 13\r\n\r\n",
		path, host, key,
	)
	if _, err := fmt.Fprint(conn, req); err != nil {
		return nil, fmt.Errorf("ws handshake: enviar upgrade: %w", err)
	}

	// Leer respuesta HTTP
	reader := bufio.NewReaderSize(conn, 65536)
	status, err := reader.ReadString('\n')
	if err != nil {
		return nil, fmt.Errorf("ws handshake: leer status: %w", err)
	}
	if !strings.Contains(status, "101") {
		return nil, fmt.Errorf("ws handshake: status inesperado: %s", strings.TrimSpace(status))
	}

	// Consumir headers hasta linea vacia
	var acceptKey string
	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			return nil, fmt.Errorf("ws handshake: leer headers: %w", err)
		}
		line = strings.TrimRight(line, "\r\n")
		if line == "" {
			break
		}
		if strings.HasPrefix(strings.ToLower(line), "sec-websocket-accept:") {
			acceptKey = strings.TrimSpace(line[len("sec-websocket-accept:"):])
		}
	}

	// Verificar Sec-WebSocket-Accept
	magic := "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
	h := sha1.New()
	h.Write([]byte(key + magic))
	expected := base64.StdEncoding.EncodeToString(h.Sum(nil))
	if acceptKey != expected {
		return nil, fmt.Errorf("ws handshake: accept key invalida: got %q, want %q", acceptKey, expected)
	}

	return reader, nil
}

// wsReadMessage lee un frame WebSocket y retorna el payload.
// Solo soporta frames de texto/binario no fragmentados (suficiente para CDP).
func wsReadMessage(reader *bufio.Reader) ([]byte, error) {
	// Leer primeros 2 bytes del frame
	header := make([]byte, 2)
	if _, err := io.ReadFull(reader, header); err != nil {
		return nil, fmt.Errorf("ws read: header: %w", err)
	}

	// fin := (header[0] & 0x80) != 0  // ignoramos fragmentacion
	opcode := header[0] & 0x0F
	masked := (header[1] & 0x80) != 0
	payloadLen := int64(header[1] & 0x7F)

	if opcode == 8 {
		return nil, fmt.Errorf("ws read: connection close frame")
	}

	// Leer longitud extendida
	switch payloadLen {
	case 126:
		var ext uint16
		if err := binary.Read(reader, binary.BigEndian, &ext); err != nil {
			return nil, fmt.Errorf("ws read: extended len 16: %w", err)
		}
		payloadLen = int64(ext)
	case 127:
		var ext uint64
		if err := binary.Read(reader, binary.BigEndian, &ext); err != nil {
			return nil, fmt.Errorf("ws read: extended len 64: %w", err)
		}
		payloadLen = int64(ext)
	}

	// Leer mascara si aplica (servidor->cliente normalmente no tiene mascara)
	var mask [4]byte
	if masked {
		if _, err := io.ReadFull(reader, mask[:]); err != nil {
			return nil, fmt.Errorf("ws read: mask: %w", err)
		}
	}

	// Leer payload
	payload := make([]byte, payloadLen)
	if _, err := io.ReadFull(reader, payload); err != nil {
		return nil, fmt.Errorf("ws read: payload: %w", err)
	}

	// Aplicar mascara si hay
	if masked {
		for i := range payload {
			payload[i] ^= mask[i%4]
		}
	}

	return payload, nil
}

// wsWriteMessage escribe un frame WebSocket enmascarado (cliente->servidor requiere mascara).
func wsWriteMessage(conn net.Conn, data []byte) error {
	// Generar mascara aleatoria
	var mask [4]byte
	if _, err := rand.Read(mask[:]); err != nil {
		return fmt.Errorf("ws write: generar mascara: %w", err)
	}

	// Construir frame
	payloadLen := len(data)
	var header []byte

	header = append(header, 0x81) // FIN + opcode text

	if payloadLen < 126 {
		header = append(header, byte(payloadLen)|0x80) // masked
	} else if payloadLen < 65536 {
		header = append(header, 126|0x80)
		header = append(header, byte(payloadLen>>8), byte(payloadLen))
	} else {
		header = append(header, 127|0x80)
		b := make([]byte, 8)
		binary.BigEndian.PutUint64(b, uint64(payloadLen))
		header = append(header, b...)
	}

	header = append(header, mask[:]...)

	// Aplicar mascara al payload
	masked := make([]byte, payloadLen)
	for i, b := range data {
		masked[i] = b ^ mask[i%4]
	}

	frame := append(header, masked...)
	if _, err := conn.Write(frame); err != nil {
		return fmt.Errorf("ws write: %w", err)
	}
	return nil
}

// sendCDP envia un comando CDP y espera la respuesta con el mismo ID.
func (c *CDPConn) sendCDP(method string, params map[string]any) (map[string]any, error) {
	id := c.nextID.Add(1)

	req := cdpRequest{ID: id, Method: method, Params: params}
	data, err := json.Marshal(req)
	if err != nil {
		return nil, fmt.Errorf("cdp send: marshal: %w", err)
	}

	// Registrar canal para la respuesta
	ch := make(chan cdpResponse, 1)
	c.pendMu.Lock()
	c.pending[id] = ch
	c.pendMu.Unlock()

	// Enviar frame WebSocket
	c.mu.Lock()
	err = wsWriteMessage(c.conn, data)
	c.mu.Unlock()
	if err != nil {
		c.pendMu.Lock()
		delete(c.pending, id)
		c.pendMu.Unlock()
		return nil, fmt.Errorf("cdp send %s: %w", method, err)
	}

	// Esperar respuesta (con timeout para no colgar el tool indefinidamente).
	select {
	case resp := <-ch:
		if resp.Error != nil {
			return nil, fmt.Errorf("cdp %s: error %d: %s", method, resp.Error.Code, resp.Error.Message)
		}
		return resp.Result, nil
	case <-time.After(cdpCmdTimeout):
		c.pendMu.Lock()
		delete(c.pending, id)
		c.pendMu.Unlock()
		return nil, fmt.Errorf("cdp %s: sin respuesta tras %s (conexion colgada?)", method, cdpCmdTimeout)
	}
}

// ensureAX habilita el dominio Accessibility una sola vez por conexion (necesario
// antes de Accessibility.getFullAXTree). Idempotente y cacheado: la segunda y
// sucesivas llamadas son no-op, evitando un round-trip por percepcion.
func (c *CDPConn) ensureAX() error {
	if c.axEnabled.Load() {
		return nil
	}
	if _, err := c.sendCDP("Accessibility.enable", nil); err != nil {
		return err
	}
	c.axEnabled.Store(true)
	return nil
}

// ensureNetwork habilita el dominio Network una sola vez por conexion. Cacheado:
// no lo deshabilitamos al terminar una espera (eso borraria el estado y forzaria
// el enable de nuevo); los handlers de eventos se desregistran por su cancel().
func (c *CDPConn) ensureNetwork() error {
	if c.netEnabled.Load() {
		return nil
	}
	if _, err := c.sendCDP("Network.enable", nil); err != nil {
		return err
	}
	c.netEnabled.Store(true)
	return nil
}

// ensurePage habilita el dominio Page una sola vez por conexion (necesario para
// recibir Page.loadEventFired y demas eventos de ciclo de vida de la pagina).
func (c *CDPConn) ensurePage() error {
	if c.pageEnabled.Load() {
		return nil
	}
	if _, err := c.sendCDP("Page.enable", nil); err != nil {
		return err
	}
	c.pageEnabled.Store(true)
	return nil
}

// readLoop lee mensajes del WebSocket y los enruta a los canales pendientes
// (respuestas a comandos) o a los handlers registrados (eventos CDP).
// Debe ejecutarse en una goroutine.
func (c *CDPConn) readLoop() {
	for {
		data, err := wsReadMessage(c.reader)
		if err != nil {
			// Conexion cerrada o error — notificar a todos los pendientes
			c.pendMu.Lock()
			for _, ch := range c.pending {
				ch <- cdpResponse{Error: &cdpError{Message: err.Error()}}
			}
			c.pending = map[int64]chan cdpResponse{}
			c.pendMu.Unlock()
			return
		}

		var resp cdpResponse
		if err := json.Unmarshal(data, &resp); err != nil {
			continue
		}

		// Si tiene ID, es respuesta a un comando
		if resp.ID > 0 {
			c.pendMu.Lock()
			ch, ok := c.pending[resp.ID]
			if ok {
				delete(c.pending, resp.ID)
			}
			c.pendMu.Unlock()
			if ok {
				ch <- resp
			}
			continue
		}

		// Sin ID = evento CDP. Llamar handlers registrados para ese metodo.
		if resp.Method != "" {
			c.hMu.Lock()
			hs := append([]EventHandler(nil), c.handlers[resp.Method]...)
			c.hMu.Unlock()
			for _, h := range hs {
				// Aislamos panics de handlers ajenos para que un handler
				// roto no mate la conexion entera.
				func(h EventHandler) {
					defer func() { _ = recover() }()
					h(resp.Method, resp.Params)
				}(h)
			}
		}
	}
}

// OnEvent registra un handler para un metodo CDP (ej "Network.requestWillBeSent").
// Devuelve una funcion `cancel` que des-registra el handler. Multiples handlers
// para el mismo metodo se invocan en orden de registro.
//
// El handler corre en la goroutine de lectura — mantenlo rapido. Para trabajo
// pesado, despacha a un canal/goroutine propios.
func (c *CDPConn) OnEvent(method string, h EventHandler) (cancel func()) {
	if c == nil || h == nil || method == "" {
		return func() {}
	}
	c.hMu.Lock()
	if c.handlers == nil {
		c.handlers = make(map[string][]EventHandler)
	}
	c.handlers[method] = append(c.handlers[method], h)
	idx := len(c.handlers[method]) - 1
	c.hMu.Unlock()

	return func() {
		c.hMu.Lock()
		defer c.hMu.Unlock()
		hs := c.handlers[method]
		if idx < len(hs) {
			c.handlers[method] = append(hs[:idx], hs[idx+1:]...)
		}
	}
}

// cdpGetWSURL obtiene el webSocketDebuggerUrl del endpoint HTTP de CDP.
func cdpGetWSURL(port int) (string, error) {
	resp, err := http.Get(fmt.Sprintf("http://localhost:%d/json/version", port))
	if err != nil {
		return "", fmt.Errorf("cdp version: %w", err)
	}
	defer resp.Body.Close()

	var info cdpVersionResponse
	if err := json.NewDecoder(resp.Body).Decode(&info); err != nil {
		return "", fmt.Errorf("cdp version: decode: %w", err)
	}
	if info.WebSocketDebuggerURL == "" {
		return "", fmt.Errorf("cdp version: webSocketDebuggerUrl vacio")
	}
	return info.WebSocketDebuggerURL, nil
}