pdf-odin-llm/objstm_test.odin

package pdf

import "core:crypto/legacy/md5"
import "core:testing"

// The sample-PDF builders live in xref_test.odin; an object stream is only
// reachable through a cross-reference stream, so the two are tested with the
// same tools.

@(private = "file")
Objstm_Opts :: struct {
    // What the dictionary declares, or -1 to declare the truth.
    declared_n:     int,
    declared_first: int,
    // Overrides the /Type of the container, for the wrong-object case.
    type_name:      string,
    // The /Index field of the type-2 entries for objects 2, 3 and 4.
    indices:        [3]u64,
    // Write the container as a plain dictionary rather than a stream.
    not_a_stream:   bool,
    // Point every type-2 entry at this container instead of object 1.
    container:      u64,
}

@(private = "file")
objstm_defaults :: proc() -> Objstm_Opts {
    return Objstm_Opts{
        declared_n     = -1,
        declared_first = -1,
        indices        = {0, 1, 2},
        container      = 1,
    }
}

// build_objstm_file writes a modern file: one object stream holding the
// catalog, the page tree and a page, and a cross-reference stream pointing at
// all three.
@(private = "file")
build_objstm_file :: proc(opts: Objstm_Opts) -> Pdf_Builder {
    pb := pb_make("%PDF-1.5\n")

    bodies := [3]string{
        "<< /Type /Catalog /Pages 3 0 R >>",
        "<< /Type /Pages /Kids [4 0 R] /Count 1 >>",
        "<< /Type /Page /Parent 3 0 R /MediaBox [0 0 250 250] >>",
    }
    nums := [3]u32{2, 3, 4}
    payload, first := objstm_payload(nums[:], bodies[:])
    defer delete(payload)

    ty := opts.type_name if opts.type_name != "" else "/ObjStm"

    pb_mark(&pb, 1)
    pb_str(&pb, "1 0 obj\n<< /Type ")
    pb_str(&pb, ty)
    pb_str(&pb, " /N ")
    pb_num(&pb, opts.declared_n >= 0 ? opts.declared_n : 3)
    pb_str(&pb, " /First ")
    pb_num(&pb, opts.declared_first >= 0 ? opts.declared_first : first)
    if opts.not_a_stream {
        pb_str(&pb, " >>\nendobj\n")
    } else {
        pb_str(&pb, " /Length ")
        pb_num(&pb, len(payload))
        pb_str(&pb, " >>\nstream\n")
        pb_raw(&pb, payload)
        pb_str(&pb, "\nendstream\nendobj\n")
    }

    at := pb_here(&pb)
    c := opts.container
    rows := [6]Xref_Row{
        {0, 0, 65535},
        {1, u64(pb.offs[1]), 0},
        {2, c, opts.indices[0]},
        {2, c, opts.indices[1]},
        {2, c, opts.indices[2]},
        {1, u64(at), 0},
    }
    raw := xref_rows_encode(rows[:], 1, 4, 2)
    defer delete(raw)
    pb_stream_obj(&pb, 5, "/Type /XRef /Size 6 /W [1 4 2] /Root 2 0 R", raw)
    pb_startxref(&pb, at)
    return pb
}

@(private = "file")
expect_objstm_page :: proc(t: ^testing.T, doc: ^Document, width: f64) {
    d := get_dict(doc, Ref{4, 0})
    if !testing.expect(t, d != nil, "the page did not come out of the object stream") {
        return
    }
    testing.expect(t, dict_name(doc, d, "Type") == "Page", "object 4 is not a /Page")
    r := dict_rect(doc, d, "MediaBox")
    testing.expectf(t, r.x1 == width, "MediaBox width %v, wanted %v", r.x1, width)
}

// ---------------------------------------------------------------------------
// The ordinary case
// ---------------------------------------------------------------------------

@(test)
test_objstm_basic :: proc(t: ^testing.T) {
    pb := build_objstm_file(objstm_defaults())
    defer pb_destroy(&pb)

    doc, err := pb_open(&pb)
    defer close(doc)
    testing.expect(t, err == .None, "open failed")
    if doc == nil { return }

    testing.expect(t, !doc.is_damaged, "a well-formed object stream recorded a repair")
    testing.expect(t, doc.catalog != nil, "the catalog did not come out of the object stream")
    testing.expect(t, doc.xref[2].kind == .Compressed, "object 2 is not a compressed entry")
    testing.expect(t, dict_name(doc, get_dict(doc, Ref{3, 0}), "Type") == "Pages",
        "object 3 is not the page tree root")
    expect_objstm_page(t, doc, 250)
}

@(test)
test_objstm_header_is_cached :: proc(t: ^testing.T) {
    pb := build_objstm_file(objstm_defaults())
    defer pb_destroy(&pb)

    doc, err := pb_open(&pb)
    defer close(doc)
    testing.expect(t, err == .None, "open failed")
    if doc == nil { return }

    _ = get_object(doc, Ref{2, 0})
    stm, ok := doc.objstm_cache[1]
    testing.expect(t, ok && stm != nil, "the object stream header was not cached")
    if stm == nil { return }
    testing.expectf(t, len(stm.nums) == 3, "cached %v pairs, wanted 3", len(stm.nums))
    testing.expect(t, stm.nums[0] == 2 && stm.nums[1] == 3 && stm.nums[2] == 4,
        "the pair table was read in the wrong order")
    testing.expect(t, !stm.loading, "the loading flag was left set")
}

// ---------------------------------------------------------------------------
// A pair table that disagrees with the cross-reference table
// ---------------------------------------------------------------------------

@(test)
test_objstm_stale_index_is_searched_for :: proc(t: ^testing.T) {
    // Every type-2 entry claims index 0. The pair table is the thing that was
    // written next to the bytes, so it wins and the object is found by number.
    opts := objstm_defaults()
    opts.indices = {0, 0, 0}
    pb := build_objstm_file(opts)
    defer pb_destroy(&pb)

    doc, err := pb_open(&pb)
    defer close(doc)
    testing.expect(t, err == .None, "open failed")
    if doc == nil { return }

    expect_objstm_page(t, doc, 250)
    testing.expect(t, dict_name(doc, get_dict(doc, Ref{3, 0}), "Type") == "Pages",
        "object 3 was not found by number")
    testing.expect(t, doc_has_repair(doc, .Xref_Entry_Mismatch),
        "a stale index should be recorded as a repair")
}

@(test)
test_objstm_index_past_the_end :: proc(t: ^testing.T) {
    opts := objstm_defaults()
    opts.indices = {900, 901, 902}
    pb := build_objstm_file(opts)
    defer pb_destroy(&pb)

    doc, err := pb_open(&pb)
    defer close(doc)
    testing.expect(t, err == .None, "open failed")
    if doc == nil { return }
    expect_objstm_page(t, doc, 250)
}

@(test)
test_objstm_get_object_rejects_absent_number :: proc(t: ^testing.T) {
    pb := build_objstm_file(objstm_defaults())
    defer pb_destroy(&pb)

    doc, err := pb_open(&pb)
    defer close(doc)
    testing.expect(t, err == .None, "open failed")
    if doc == nil { return }

    obj, gerr := objstm_get_object(doc, 1, 0, 99)
    testing.expect(t, obj == nil && gerr == .Object_Not_Found,
        "an object number the stream does not hold must not resolve")

    _, zerr := objstm_get_object(doc, 1, 0, 0)
    testing.expect(t, zerr == .Invalid_Argument, "object 0 is not a legal request")
}

// ---------------------------------------------------------------------------
// /N and /First lying
// ---------------------------------------------------------------------------

@(test)
test_objstm_first_past_end_is_repaired :: proc(t: ^testing.T) {
    // /First outside the decoded stream. The pair table itself says where the
    // objects start, so the file is still readable.
    opts := objstm_defaults()
    opts.declared_first = 9999
    pb := build_objstm_file(opts)
    defer pb_destroy(&pb)

    doc, err := pb_open(&pb)
    defer close(doc)
    testing.expect(t, err == .None, "open failed")
    if doc == nil { return }

    testing.expect(t, doc_has_repair(doc, .Bad_Stream_Length), "no repair was recorded")
    expect_objstm_page(t, doc, 250)

    stm, ok := doc.objstm_cache[1]
    if testing.expect(t, ok && stm != nil, "the object stream was not cached") {
        testing.expectf(t, stm.first > 0 && stm.first <= len(stm.data),
            "/First was repaired to %v, outside a %v byte stream", stm.first, len(stm.data))
    }
}

@(test)
test_objstm_first_inside_pair_table_is_repaired :: proc(t: ^testing.T) {
    // /First 1 puts the objects on top of the pair table, which cannot be.
    opts := objstm_defaults()
    opts.declared_first = 1
    pb := build_objstm_file(opts)
    defer pb_destroy(&pb)

    doc, err := pb_open(&pb)
    defer close(doc)
    testing.expect(t, err == .None, "open failed")
    if doc == nil { return }
    testing.expect(t, doc_has_repair(doc, .Bad_Stream_Length), "no repair was recorded")
    // Reading the pair table again without the bogus ceiling recovers the file
    // completely.
    expect_objstm_page(t, doc, 250)
}

@(test)
test_objstm_n_too_large_is_clamped :: proc(t: ^testing.T) {
    opts := objstm_defaults()
    opts.declared_n = 100_000_000
    pb := build_objstm_file(opts)
    defer pb_destroy(&pb)

    doc, err := pb_open(&pb)
    defer close(doc)
    testing.expect(t, err == .None, "open failed")
    if doc == nil { return }

    stm, ok := doc.objstm_cache[1]
    if testing.expect(t, ok && stm != nil, "the object stream was not read at all") {
        testing.expectf(t, len(stm.nums) <= 64, "read %v pairs from a 3-object stream",
            len(stm.nums))
    }
    expect_objstm_page(t, doc, 250)
}

@(test)
test_objstm_n_missing_is_derived :: proc(t: ^testing.T) {
    opts := objstm_defaults()
    opts.declared_n = 0
    pb := build_objstm_file(opts)
    defer pb_destroy(&pb)

    doc, err := pb_open(&pb)
    defer close(doc)
    testing.expect(t, err == .None, "open failed")
    if doc == nil { return }
    expect_objstm_page(t, doc, 250)
}

// ---------------------------------------------------------------------------
// Containers that are not object streams
// ---------------------------------------------------------------------------

// The cross-reference stream says objects 2, 3 and 4 live inside object 1, but
// object 1 was written as a plain dictionary. There are no bytes to take
// apart, so the lookup has to fail rather than parse whatever follows the
// dictionary in the file -- which is the cross-reference stream itself.
@(test)
test_objstm_container_is_not_a_stream :: proc(t: ^testing.T) {
    opts := objstm_defaults()
    opts.not_a_stream = true
    pb := build_objstm_file(opts)
    defer pb_destroy(&pb)

    doc, _ := pb_open(&pb)
    defer close(doc)
    if doc == nil { return }

    obj, gerr := objstm_get_object(doc, 1, 0, 2)
    testing.expect(t, obj == nil, "a dictionary was taken apart as an object stream")
    // Which refusal comes back depends on whether the container has been tried
    // before -- `find_catalog` already walked the file during open, and a
    // container that failed once is remembered as failed. Both are clean.
    testing.expectf(t, gerr == .Object_Not_Found || gerr == .Bad_Stream,
        "container that is not a stream: got %s", error_string(gerr))

    // The container is remembered as failed, so the second and third objects
    // do not each retry the decode.
    stm, cached := doc.objstm_cache[1]
    testing.expect(t, cached && stm != nil, "the failed container was not cached")
    if stm != nil {
        testing.expect(t, stm.failed, "the container was not marked failed")
        testing.expect(t, !stm.loading, "the loading flag was left set")
        testing.expect(t, len(stm.nums) == 0, "a failed container kept a pair table")
    }

    // Everything the container was supposed to hold reads as null. Whether the
    // xref still calls those entries compressed or the loader gave up on the
    // section and rebuilt -- which is what happens here, and leaves 2, 3 and 4
    // free, because the rebuild finds only the objects physically in the file
    // -- the answer a caller gets is the same.
    for num in 2 ..= 4 {
        testing.expectf(t, get_object(doc, Ref{u32(num), 0}) == nil,
            "object %v came back out of a container that is not a stream", num)
    }

    // Silence would be the real bug: a caller has no other way to tell an
    // empty document from one whose objects were lost.
    testing.expect(t, doc.is_damaged, "losing three objects was not recorded as damage")
    testing.expect(t, doc_has_repair(doc, .Missing_Object),
        "no Missing_Object repair for the objects the container was supposed to hold")

    // Object 1 itself is a perfectly good dictionary and still reads.
    d := get_dict(doc, Ref{1, 0})
    testing.expect(t, d != nil && dict_name(doc, d, "Type") == "ObjStm",
        "the container dictionary itself stopped resolving")
}

@(test)
test_objstm_container_has_wrong_type :: proc(t: ^testing.T) {
    opts := objstm_defaults()
    opts.type_name = "/XObject"
    pb := build_objstm_file(opts)
    defer pb_destroy(&pb)

    doc, _ := pb_open(&pb)
    defer close(doc)
    if doc == nil { return }

    obj, gerr := objstm_get_object(doc, 1, 0, 2)
    testing.expect(t, obj == nil && gerr == .Type_Mismatch,
        "a container that says it is something else must be refused")
}

// build_split_file writes a file whose catalog and page tree are ordinary
// top-level objects and whose page lives in an object stream. /Root therefore
// resolves without touching the object stream, so a broken entry for the page
// is a broken entry rather than a reason to rebuild the whole file -- which is
// what makes the degradation visible.
@(private = "file")
build_split_file :: proc(page_container: u64, objstm_self_hosted: bool) -> Pdf_Builder {
    pb := pb_make("%PDF-1.5\n")

    bodies := [1]string{"<< /Type /Page /Parent 3 0 R /MediaBox [0 0 250 250] >>"}
    nums := [1]u32{4}
    payload, first := objstm_payload(nums[:], bodies[:])
    defer delete(payload)

    pb_mark(&pb, 1)
    pb_str(&pb, "1 0 obj\n<< /Type /ObjStm /N 1 /First ")
    pb_num(&pb, first)
    pb_str(&pb, " /Length ")
    pb_num(&pb, len(payload))
    pb_str(&pb, " >>\nstream\n")
    pb_raw(&pb, payload)
    pb_str(&pb, "\nendstream\nendobj\n")

    pb_obj(&pb, 2, "<< /Type /Catalog /Pages 3 0 R >>")
    pb_obj(&pb, 3, "<< /Type /Pages /Kids [4 0 R] /Count 1 >>")

    at := pb_here(&pb)
    // Object 1 either sits in the file where it really is, or claims to live
    // inside itself.
    row1 := Xref_Row{1, u64(pb.offs[1]), 0}
    if objstm_self_hosted {
        row1 = Xref_Row{2, 1, 0}
    }
    rows := [6]Xref_Row{
        {0, 0, 65535},
        row1,
        {1, u64(pb.offs[2]), 0},
        {1, u64(pb.offs[3]), 0},
        {2, page_container, 0},
        {1, u64(at), 0},
    }
    raw := xref_rows_encode(rows[:], 1, 4, 2)
    defer delete(raw)
    pb_stream_obj(&pb, 5, "/Type /XRef /Size 6 /W [1 4 2] /Root 2 0 R", raw)
    pb_startxref(&pb, at)
    return pb
}

@(test)
test_objstm_split_file_baseline :: proc(t: ^testing.T) {
    pb := build_split_file(1, false)
    defer pb_destroy(&pb)

    doc, err := pb_open(&pb)
    defer close(doc)
    testing.expect(t, err == .None, "open failed")
    if doc == nil { return }
    testing.expect(t, !doc.is_damaged, "a well-formed split file recorded a repair")
    expect_objstm_page(t, doc, 250)
}

@(test)
test_objstm_container_does_not_exist :: proc(t: ^testing.T) {
    // The page claims to live inside object 90, which the file does not have.
    // /Root still resolves, so nothing triggers a rebuild and the reference
    // has to degrade to null on its own.
    pb := build_split_file(90, false)
    defer pb_destroy(&pb)

    doc, err := pb_open(&pb)
    defer close(doc)
    testing.expect(t, err == .None, "open failed")
    if doc == nil { return }

    testing.expect(t, get_object(doc, Ref{4, 0}) == nil,
        "an entry pointing at a container that is not there must read as null")
    testing.expect(t, doc.catalog != nil, "one bad entry must not cost the rest of the file")
    testing.expect(t, doc_has_repair(doc, .Missing_Object), "no repair was recorded")
}

// ---------------------------------------------------------------------------
// Cycles
// ---------------------------------------------------------------------------

@(test)
test_objstm_contains_itself :: proc(t: ^testing.T) {
    // The cross-reference table says the object stream itself is one of the
    // objects inside it. Reading it needs its own bytes, which needs reading
    // it: without a guard this is an unbounded recursion.
    pb := build_split_file(1, true)
    defer pb_destroy(&pb)

    doc, err := pb_open(&pb)
    defer close(doc)
    testing.expect(t, err == .None, "open failed")
    if doc == nil { return }

    testing.expect(t, get_object(doc, Ref{1, 0}) == nil,
        "a self-containing object stream must read as null, not recurse")
    // Everything the stream held goes with it, but the rest of the file does
    // not.
    testing.expect(t, doc.catalog != nil, "the catalog was lost with the object stream")
    testing.expect(t, doc.is_damaged, "the cycle should have been recorded")
}

@(test)
test_objstm_mutual_cycle :: proc(t: ^testing.T) {
    // Object 1 says it lives in object 2's stream, and object 2 says it lives
    // in object 1's.
    pb := pb_make("%PDF-1.5\n")
    defer pb_destroy(&pb)

    bodies := [1]string{"<< /Type /Catalog >>"}
    for num in ([2]u32{1, 2}) {
        nums := [1]u32{num == 1 ? 2 : 1}
        payload, first := objstm_payload(nums[:], bodies[:])
        defer delete(payload)
        pb_mark(&pb, num)
        pb_num(&pb, int(num))
        pb_str(&pb, " 0 obj\n<< /Type /ObjStm /N 1 /First ")
        pb_num(&pb, first)
        pb_str(&pb, " /Length ")
        pb_num(&pb, len(payload))
        pb_str(&pb, " >>\nstream\n")
        pb_raw(&pb, payload)
        pb_str(&pb, "\nendstream\nendobj\n")
    }

    at := pb_here(&pb)
    rows := [4]Xref_Row{
        {0, 0, 65535},
        {2, 2, 0},
        {2, 1, 0},
        {1, u64(at), 0},
    }
    raw := xref_rows_encode(rows[:], 1, 4, 2)
    defer delete(raw)
    pb_stream_obj(&pb, 3, "/Type /XRef /Size 4 /W [1 4 2] /Root 1 0 R", raw)
    pb_startxref(&pb, at)

    doc, _ := pb_open(&pb)
    defer close(doc)
    if doc == nil { return }
    _ = get_object(doc, Ref{1, 0})
    _ = get_object(doc, Ref{2, 0})
    testing.expect(t, true, "a mutual object-stream cycle must terminate")
}

// ---------------------------------------------------------------------------
// Object streams found by a rebuild
// ---------------------------------------------------------------------------

@(test)
test_objstm_found_by_rebuild :: proc(t: ^testing.T) {
    // The cross-reference stream is unreachable, so every object in the file
    // has to be found by scanning -- and three of the four only exist inside
    // an object stream.
    pb := build_objstm_file(objstm_defaults())
    defer pb_destroy(&pb)

    // Point startxref at nothing.
    at := last_index_bytes(pb.buf[:], transmute([]byte)string("startxref"))
    testing.expect(t, at > 0, "the sample has no startxref")
    if at <= 0 { return }
    resize(&pb.buf, at)
    pb_str(&pb, "startxref\n99000000\n%%EOF\n")

    doc, err := pb_open(&pb)
    defer close(doc)
    testing.expect(t, err == .None, "a rebuild should have opened this file")
    if doc == nil { return }

    testing.expect(t, doc_has_repair(doc, .Rebuilt_Xref), "no rebuild was recorded")
    testing.expect(t, doc.xref[4].kind == .Compressed,
        "the rebuild did not register the object stream's contents")
    testing.expect(t, doc.catalog != nil, "the catalog inside the object stream was lost")
    expect_objstm_page(t, doc, 250)
}

@(test)
test_objstm_rebuild_prefers_later_top_level_definition :: proc(t: ^testing.T) {
    // The file was updated by appending an uncompressed object 4, which is
    // newer than the object stream and must win.
    pb := build_objstm_file(objstm_defaults())
    defer pb_destroy(&pb)

    at := last_index_bytes(pb.buf[:], transmute([]byte)string("startxref"))
    if at <= 0 { return }
    resize(&pb.buf, at)
    pb_obj(&pb, 4, "<< /Type /Page /Parent 3 0 R /MediaBox [0 0 700 700] >>")
    pb_str(&pb, "startxref\n99000000\n%%EOF\n")

    doc, err := pb_open(&pb)
    defer close(doc)
    testing.expect(t, err == .None, "a rebuild should have opened this file")
    if doc == nil { return }
    expect_objstm_page(t, doc, 700)
}

// ---------------------------------------------------------------------------
// Damaged payloads
// ---------------------------------------------------------------------------

@(test)
test_objstm_truncation_survives :: proc(t: ^testing.T) {
    pb := build_objstm_file(objstm_defaults())
    defer pb_destroy(&pb)

    for cut := 1; cut <= len(pb.buf); cut += 5 {
        doc, _ := open_bytes(pb.buf[:cut], Open_Options{skip_page_tree = true})
        if doc != nil {
            for num in 1 ..< object_count(doc) {
                _ = get_object(doc, Ref{u32(num), 0})
            }
            close(doc)
        }
    }
    testing.expect(t, true, "a truncated object stream must never crash or hang")
}

@(test)
test_objstm_garbage_payload_survives :: proc(t: ^testing.T) {
    cases := [?]string{
        "",
        "   ",
        "1",
        "1 ",
        "1 0",
        "1 999999999999 << >>",
        "-4 -4 << >>",
        "1 0 <<",
        "99999999999999999999 0 << >>",
    }
    for body in cases {
        pb := pb_make("%PDF-1.5\n")
        defer pb_destroy(&pb)

        payload := transmute([]byte)body
        pb_mark(&pb, 1)
        pb_str(&pb, "1 0 obj\n<< /Type /ObjStm /N 1 /First 4 /Length ")
        pb_num(&pb, len(payload))
        pb_str(&pb, " >>\nstream\n")
        pb_raw(&pb, payload)
        pb_str(&pb, "\nendstream\nendobj\n")

        at := pb_here(&pb)
        rows := [3]Xref_Row{
            {0, 0, 65535},
            {1, u64(pb.offs[1]), 0},
            {1, u64(at), 0},
        }
        raw := xref_rows_encode(rows[:], 1, 4, 2)
        defer delete(raw)
        pb_stream_obj(&pb, 2, "/Type /XRef /Size 3 /W [1 4 2] /Root 1 0 R", raw)
        pb_startxref(&pb, at)

        doc, _ := pb_open(&pb)
        if doc != nil {
            _, _ = objstm_get_object(doc, 1, 0, 7)
            close(doc)
        }
    }
    testing.expect(t, true, "a garbage object-stream payload must never crash")
}

// ---------------------------------------------------------------------------
// Encrypted documents
//
// A modern encrypted file keeps its catalog inside an object stream, and that
// stream cannot be decoded until the file key exists -- which is after the
// cross-reference table has been read. Two things in xref.odin exist for this
// case and are only exercised here: /Root is not resolved as a validity check
// when /Encrypt is present, and every object and object stream read during the
// load is dropped afterwards, because all of it was read in the dark.
// ---------------------------------------------------------------------------

@(private = "file")
rc4_object_key :: proc(file_key: []byte, num: u32, gen: u16, out: []byte) -> int {
    ctx: md5.Context
    md5.init(&ctx)
    md5.update(&ctx, file_key)
    tail := [5]byte{
        byte(num), byte(num >> 8), byte(num >> 16),
        byte(gen), byte(gen >> 8),
    }
    md5.update(&ctx, tail[:])
    digest: [md5.DIGEST_SIZE]byte
    md5.final(&ctx, digest[:])
    n := min(len(file_key) + 5, 16)
    copy(out[:n], digest[:n])
    return n
}

@(private = "file")
pb_hex :: proc(pb: ^Pdf_Builder, data: []byte) {
    hex := transmute([]byte)string("0123456789abcdef")
    pb_str(pb, "<")
    for c in data {
        append(&pb.buf, hex[c >> 4])
        append(&pb.buf, hex[c & 0x0f])
    }
    pb_str(pb, ">")
}

// build_encrypted_objstm_file writes an RC4 40-bit, revision 2 document with an
// empty user password -- the oldest scheme, and the one whose key derivation is
// short enough to spell out here -- whose catalog, page tree and page all live
// inside an encrypted object stream.
@(private = "file")
build_encrypted_objstm_file :: proc() -> Pdf_Builder {
    id0 := [16]byte{
        0x2b, 0x1c, 0x7d, 0x0e, 0x55, 0x91, 0xa3, 0x44,
        0x10, 0xff, 0x6c, 0x22, 0x83, 0x39, 0xd7, 0x00,
    }
    permissions := i32(-1)

    // /O, Algorithm 3 with an empty owner password.
    padded: [32]byte
    crypt_pad_password(nil, padded[:])
    ctx: md5.Context
    md5.init(&ctx)
    md5.update(&ctx, padded[:])
    digest: [md5.DIGEST_SIZE]byte
    md5.final(&ctx, digest[:])
    o_value: [32]byte
    crypt_pad_password(nil, o_value[:])
    rc4_crypt(digest[:5], o_value[:], o_value[:])

    file_key: [5]byte
    crypt_key_r234(nil, o_value[:], permissions, id0[:], 2, 5, true, file_key[:])

    u_value: [32]byte
    crypt_user_value(file_key[:], 2, id0[:], u_value[:])

    bodies := [3]string{
        "<< /Type /Catalog /Pages 3 0 R >>",
        "<< /Type /Pages /Kids [4 0 R] /Count 1 >>",
        "<< /Type /Page /Parent 3 0 R /MediaBox [0 0 321 321] >>",
    }
    nums := [3]u32{2, 3, 4}
    plain, first := objstm_payload(nums[:], bodies[:])
    defer delete(plain)

    // The whole object stream is encrypted with object 1's key; the objects
    // inside it are then in the clear (ISO 32000-2, 7.5.7).
    cipher := make([]byte, len(plain))
    defer delete(cipher)
    okey: [16]byte
    n_key := rc4_object_key(file_key[:], 1, 0, okey[:])
    rc4_crypt(okey[:n_key], cipher, plain)

    pb := pb_make("%PDF-1.5\n")

    pb_mark(&pb, 1)
    pb_str(&pb, "1 0 obj\n<< /Type /ObjStm /N 3 /First ")
    pb_num(&pb, first)
    pb_str(&pb, " /Length ")
    pb_num(&pb, len(cipher))
    pb_str(&pb, " >>\nstream\n")
    pb_raw(&pb, cipher)
    pb_str(&pb, "\nendstream\nendobj\n")

    // The /Encrypt dictionary's own strings are never encrypted: they are what
    // the key is derived from.
    pb_mark(&pb, 5)
    pb_str(&pb, "5 0 obj\n<< /Filter /Standard /V 1 /R 2 /P -1 /O ")
    pb_hex(&pb, o_value[:])
    pb_str(&pb, " /U ")
    pb_hex(&pb, u_value[:])
    pb_str(&pb, " >>\nendobj\n")

    at := pb_here(&pb)
    rows := [7]Xref_Row{
        {0, 0, 65535},
        {1, u64(pb.offs[1]), 0},
        {2, 1, 0},
        {2, 1, 1},
        {2, 1, 2},
        {1, u64(pb.offs[5]), 0},
        {1, u64(at), 0},
    }
    raw := xref_rows_encode(rows[:], 1, 4, 2)
    defer delete(raw)

    // A cross-reference stream is never encrypted either, so it is written in
    // the clear and its /ID strings stay as they are.
    pb_mark(&pb, 6)
    pb_str(&pb, "6 0 obj\n<< /Type /XRef /Size 7 /W [1 4 2] /Root 2 0 R /Encrypt 5 0 R /ID [")
    pb_hex(&pb, id0[:])
    pb_str(&pb, " ")
    pb_hex(&pb, id0[:])
    pb_str(&pb, "] /Length ")
    pb_num(&pb, len(raw))
    pb_str(&pb, " >>\nstream\n")
    pb_raw(&pb, raw)
    pb_str(&pb, "\nendstream\nendobj\n")
    pb_startxref(&pb, at)
    return pb
}

@(test)
test_objstm_encrypted_catalog_in_object_stream :: proc(t: ^testing.T) {
    pb := build_encrypted_objstm_file()
    defer pb_destroy(&pb)

    doc, err := pb_open(&pb)
    defer close(doc)
    testing.expectf(t, err == .None, "open failed: %v", error_string(err))
    if doc == nil { return }

    testing.expect(t, document_is_encrypted(doc), "the document is not marked encrypted")
    testing.expect(t, !doc_has_repair(doc, .Rebuilt_Xref),
        "an encrypted catalog inside an object stream must not force a rebuild")
    testing.expect(t, doc.catalog != nil, "the catalog was not decrypted")
    testing.expect(t, dict_name(doc, doc.catalog, "Type") == "Catalog",
        "the catalog decoded to something else")
    expect_objstm_page(t, doc, 321)
}

@(test)
test_objstm_encrypted_rebuild_degrades_cleanly :: proc(t: ^testing.T) {
    // An encrypted file that also needs a rebuild is the one case where the
    // compressed objects genuinely cannot be recovered: the file key is not
    // derived until after the table is read, so the object streams decode to
    // ciphertext. What must not happen is that the garbage gets registered as
    // if it were a pair table -- the invented entries would overwrite the real
    // ones the scan found.
    pb := build_encrypted_objstm_file()
    defer pb_destroy(&pb)

    at := last_index_bytes(pb.buf[:], transmute([]byte)string("startxref"))
    testing.expect(t, at > 0, "the sample has no startxref")
    if at <= 0 { return }
    resize(&pb.buf, at)
    pb_str(&pb, "startxref\n99000000\n%%EOF\n")

    doc, err := pb_open(&pb)
    defer close(doc)
    testing.expectf(t, err == .None, "open failed: %v", error_string(err))
    if doc == nil { return }

    testing.expect(t, doc_has_repair(doc, .Rebuilt_Xref), "no rebuild was recorded")
    testing.expect(t, document_is_encrypted(doc),
        "the /Encrypt dictionary survived the rebuild, so the file is still encrypted")
    // The three objects in the stream are gone, but the objects that were
    // written at the top level are exactly where the scan found them.
    testing.expect(t, doc.xref[1].kind == .Normal, "the object stream itself was lost")
    testing.expect(t, doc.xref[5].kind == .Normal, "the /Encrypt dictionary was lost")
    for num in 1 ..< object_count(doc) {
        testing.expect(t, doc.xref[num].kind != .Compressed,
            "ciphertext was read as a pair table")
        _ = get_object(doc, Ref{u32(num), 0})
    }
}