nigig-org/crates/apps/pdf/pdf-document/tests/encryption_write.rs

//! Encrypted-write round trip: Phase 6's first exit criterion.
//!
//! A writer can only be tested by reading back what it wrote, and an
//! *encryption* writer has a second obligation on top of that: the
//! plaintext must not be in the file. Both are asserted here, and the
//! second one is asserted against the raw bytes rather than through any
//! API, because an API that decrypts is exactly the thing that cannot see
//! a leak.
//!
//! See `REVIEWS/adr/0024-pdf-encryption-on-save.md`.

use nigig_pdf_cos::encrypt_write::{EncryptionAlgorithm, EncryptionSettings};
use nigig_pdf_cos::object::{PdfDict, PdfObj};
use nigig_pdf_cos::writer::PdfDocBuilder;
use nigig_pdf_cos::Permissions;
use nigig_pdf_document::PdfDocument;

const SECRET: &[u8] = b"CONFIDENTIALPAYLOAD";

fn helvetica() -> PdfObj {
    let mut f = PdfDict::new();
    f.set("Type", PdfObj::Name("Font".to_string()));
    f.set("Subtype", PdfObj::Name("Type1".to_string()));
    f.set("BaseFont", PdfObj::Name("Helvetica".to_string()));
    PdfObj::Dict(f)
}

/// A one-page document whose content and title both contain `SECRET`.
fn encrypted_document(settings: Option<EncryptionSettings>) -> Vec<u8> {
    let mut builder = PdfDocBuilder::new();
    builder.add_font("F1", helvetica());
    builder.add_page(
        300.0,
        400.0,
        b"BT /F1 14 Tf 20 200 Td (CONFIDENTIALPAYLOAD) Tj ET",
    );
    builder.set_metadata(nigig_pdf_cos::writer::DocumentMetadata {
        title: Some("CONFIDENTIALPAYLOAD".to_string()),
        ..Default::default()
    });
    if let Some(settings) = settings {
        builder.set_encryption(settings);
    }
    builder.finish()
}

fn contains(haystack: &[u8], needle: &[u8]) -> bool {
    haystack.windows(needle.len()).any(|w| w == needle)
}

fn both_algorithms() -> [EncryptionAlgorithm; 2] {
    [EncryptionAlgorithm::Aes256, EncryptionAlgorithm::Aes128]
}

fn settings_for(algorithm: EncryptionAlgorithm, password: &str) -> EncryptionSettings {
    EncryptionSettings {
        algorithm,
        user_password: password.as_bytes().to_vec(),
        ..Default::default()
    }
}

// ------------------------------------------------------- the round trip

#[test]
fn an_encrypted_document_opens_with_its_password() {
    for algorithm in both_algorithms() {
        let pdf = encrypted_document(Some(settings_for(algorithm, "hunter2")));
        let mut doc = PdfDocument::parse_with_password(&pdf, b"hunter2")
            .unwrap_or_else(|e| panic!("{algorithm:?} should open: {e}"));

        assert_eq!(doc.page_count(), 1);
        let page = doc.page(0).expect("page 0");
        assert_eq!(page.media_box, [0.0, 0.0, 300.0, 400.0]);
        assert!(
            contains(&page.content_data, SECRET),
            "{algorithm:?}: content did not decrypt, got {:?}",
            String::from_utf8_lossy(&page.content_data)
        );
    }
}

#[test]
fn the_plaintext_is_not_in_the_file() {
    // The assertion that matters. Everything else can be right and this
    // still wrong: a stream written before the encryptor was attached, or
    // a string the recursive walk missed, leaves readable content in a
    // file every reader will happily call encrypted.
    for algorithm in both_algorithms() {
        let pdf = encrypted_document(Some(settings_for(algorithm, "hunter2")));
        assert!(
            !contains(&pdf, SECRET),
            "{algorithm:?}: plaintext found in the encrypted file"
        );
    }

    // Control: the same document unencrypted *does* contain it, so the
    // test above is not passing because the fixture lost its payload.
    let plain = encrypted_document(None);
    assert!(
        contains(&plain, SECRET),
        "the unencrypted control must contain the payload, or the leak \
         test proves nothing"
    );
}

#[test]
fn strings_are_encrypted_too_not_just_streams() {
    // The title goes into /Info as a string. Encrypting streams and
    // forgetting strings is the easy half-implementation, and it leaks
    // document titles, form values and annotation contents.
    for algorithm in both_algorithms() {
        let pdf = encrypted_document(Some(settings_for(algorithm, "pw")));
        assert!(
            !contains(&pdf, b"CONFIDENTIALPAYLOAD"),
            "{algorithm:?}: the /Info title was written in the clear"
        );
        let doc = PdfDocument::parse_with_password(&pdf, b"pw").expect("opens");
        assert_eq!(
            doc.info().title.as_deref(),
            Some("CONFIDENTIALPAYLOAD"),
            "{algorithm:?}: the title did not decrypt"
        );
    }
}

#[test]
fn the_wrong_password_is_refused() {
    for algorithm in both_algorithms() {
        let pdf = encrypted_document(Some(settings_for(algorithm, "correct")));
        match PdfDocument::parse_with_password(&pdf, b"wrong") {
            Err(_) => {}
            Ok(mut doc) => {
                // If a document is somehow produced, it must not yield the
                // plaintext. Garbage that looks like content is the
                // failure mode to rule out.
                if let Ok(page) = doc.page(0) {
                    assert!(
                        !contains(&page.content_data, SECRET),
                        "{algorithm:?}: the wrong password recovered the content"
                    );
                }
            }
        }
    }
}

#[test]
fn an_empty_user_password_still_encrypts() {
    // Encrypting with an empty user password is the common "permissions
    // only" case: any reader can open it, but the bytes on disk are still
    // ciphertext and the permissions are declared.
    for algorithm in both_algorithms() {
        let pdf = encrypted_document(Some(settings_for(algorithm, "")));
        assert!(!contains(&pdf, SECRET), "{algorithm:?}: not encrypted");
        let mut doc = PdfDocument::parse(&pdf).expect("opens with no password");
        assert!(doc.is_encrypted());
        assert!(contains(&doc.page(0).unwrap().content_data, SECRET));
    }
}

#[test]
fn an_owner_password_also_opens_the_document() {
    for algorithm in both_algorithms() {
        let settings = EncryptionSettings {
            algorithm,
            user_password: b"user".to_vec(),
            owner_password: b"owner".to_vec(),
            ..Default::default()
        };
        let pdf = encrypted_document(Some(settings));
        for password in [b"user".as_slice(), b"owner".as_slice()] {
            let mut doc = PdfDocument::parse_with_password(&pdf, password).unwrap_or_else(|e| {
                panic!(
                    "{algorithm:?}: {:?} should open the document: {e}",
                    String::from_utf8_lossy(password)
                )
            });
            assert!(contains(&doc.page(0).unwrap().content_data, SECRET));
        }
    }
}

// ------------------------------------------------------------ structure

#[test]
fn the_encrypt_dictionary_is_written_unencrypted() {
    // It carries the salts and validation hashes a reader needs *before*
    // it has a key. Encrypting it makes the document permanently
    // unopenable — by us and by everyone else.
    for algorithm in both_algorithms() {
        let pdf = encrypted_document(Some(settings_for(algorithm, "pw")));
        let text = String::from_utf8_lossy(&pdf);
        assert!(text.contains("/Filter /Standard"), "{algorithm:?}");
        assert!(
            text.contains("/Encrypt "),
            "{algorithm:?}: no trailer entry"
        );
    }
}

#[test]
fn an_encrypted_file_carries_an_id() {
    // Revisions 2-4 mix /ID into key derivation, so a missing or changed
    // /ID makes the file unopenable. Written for both revisions because
    // every other producer writes it.
    for algorithm in both_algorithms() {
        let pdf = encrypted_document(Some(settings_for(algorithm, "pw")));
        assert!(
            String::from_utf8_lossy(&pdf).contains("/ID ["),
            "{algorithm:?}: no /ID in the trailer"
        );
    }
}

#[test]
fn permissions_are_carried_into_the_saved_file() {
    // All bits set except print (bit 3) and copy (bit 5); the spec
    // numbers permission bits from 1, so those are shifts of 2 and 4.
    let restricted = Permissions(!(1 << 2) & !(1 << 4));
    for algorithm in both_algorithms() {
        let settings = EncryptionSettings {
            algorithm,
            user_password: b"pw".to_vec(),
            permissions: restricted,
            ..Default::default()
        };
        let pdf = encrypted_document(Some(settings));
        let doc = PdfDocument::parse_with_password(&pdf, b"pw").expect("opens");
        let permissions = doc.permissions().expect("permissions present");
        assert!(!permissions.can_print(), "{algorithm:?}");
        assert!(!permissions.can_copy(), "{algorithm:?}");
        assert!(permissions.can_annotate(), "{algorithm:?}");
    }
}

#[test]
fn two_saves_of_one_document_are_not_byte_identical() {
    // Encryption must not be deterministic: the file key, the salts and
    // every IV are fresh each time. Identical output would mean a fixed
    // IV, which leaks whether two documents share content.
    //
    // This is the one place the reproducible-build property of ADR 0023 is
    // deliberately given up, and it is given up for a reason.
    let a = encrypted_document(Some(settings_for(EncryptionAlgorithm::Aes256, "pw")));
    let b = encrypted_document(Some(settings_for(EncryptionAlgorithm::Aes256, "pw")));
    assert_ne!(a, b, "encryption is deterministic, so the IV is fixed");

    // Both still open, so the difference is randomness and not corruption.
    for pdf in [a, b] {
        let mut doc = PdfDocument::parse_with_password(&pdf, b"pw").expect("opens");
        assert!(contains(&doc.page(0).unwrap().content_data, SECRET));
    }
}

#[test]
fn one_encryptor_gives_two_identical_payloads_different_ciphertext() {
    // `two_saves_of_one_document_are_not_byte_identical` cannot see a
    // fixed IV: the AES-256 file key is fresh per save, so the output
    // differs even with a constant IV. Removing the IV randomisation
    // survived that test.
    //
    // Holding one encryptor fixed isolates the IV, which is the property
    // that actually matters: in CBC, encrypting identical plaintext under
    // one key with a repeated IV produces identical ciphertext and leaks
    // that the two are the same.
    use nigig_pdf_cos::encrypt_write::Encryptor;

    for algorithm in both_algorithms() {
        let enc =
            Encryptor::new(&settings_for(algorithm, "pw"), b"fixed-file-id").expect("encryptor");
        let a = enc.encrypt_stream(1, 0, b"identical payload").expect("a");
        let b = enc.encrypt_stream(1, 0, b"identical payload").expect("b");
        assert_ne!(
            a, b,
            "{algorithm:?}: same key, same object, same plaintext gave the \
             same bytes, so the IV is not random"
        );
        // And the IV really is the first 16 bytes, not something else that
        // happens to differ.
        assert_ne!(a[..16], b[..16], "{algorithm:?}: the IV block is constant");
    }
}

#[test]
fn a_stream_length_describes_the_ciphertext_not_the_plaintext() {
    // /Length must count the encrypted bytes: the IV and PKCS#7 padding
    // make the stream longer than its plaintext. A reader that trusts a
    // short /Length truncates the stream and decrypts garbage — and the
    // round-trip tests miss it, because our own reader recovers by
    // scanning for `endstream` (a defect ADR 0023 had to fix).
    //
    // Worked on **bytes**, not on a lossy UTF-8 view: ciphertext is not
    // valid UTF-8, and `from_utf8_lossy` substitutes replacement
    // characters that shift every offset. The first draft of this test did
    // exactly that and reported a 156-byte stream declaring 80.
    let pdf = encrypted_document(Some(settings_for(EncryptionAlgorithm::Aes256, "pw")));

    let mut checked = 0usize;
    let mut cursor = 0usize;
    while let Some(rel) = pdf[cursor..].windows(7).position(|w| w == b"stream\n") {
        let at = cursor + rel;
        // `endstream` also contains "stream", so skip those matches.
        if at >= 3 && &pdf[at - 3..at] == b"end" {
            cursor = at + 7;
            continue;
        }
        let head = &pdf[..at];
        let Some(pos) = head.windows(8).rposition(|w| w == b"/Length ") else {
            cursor = at + 7;
            continue;
        };
        let declared: usize = head[pos + 8..]
            .iter()
            .take_while(|b| b.is_ascii_digit())
            .map(|b| (b - b'0') as usize)
            .fold(0, |acc, d| acc * 10 + d);

        let body_start = at + 7;
        let Some(end) = pdf[body_start..]
            .windows(10)
            .position(|w| w == b"\nendstream")
        else {
            cursor = body_start;
            continue;
        };
        assert_eq!(
            declared, end,
            "a stream declares /Length {declared} but holds {end} bytes"
        );
        checked += 1;
        cursor = body_start + end;
    }
    assert!(
        checked >= 1,
        "no streams were checked, so this proves nothing"
    );
}

#[test]
fn an_unencrypted_document_is_unchanged_by_this_work() {
    // The encryptor is opt-in. A document with no settings must be byte
    // reproducible, as ADR 0023 requires.
    let a = encrypted_document(None);
    let b = encrypted_document(None);
    assert_eq!(a, b, "unencrypted output stopped being reproducible");
    let mut doc = PdfDocument::parse(&a).expect("opens");
    assert!(!doc.is_encrypted());
    assert!(contains(&doc.page(0).unwrap().content_data, SECRET));
}

// -------------------------------------------------------------- content

#[test]
fn a_large_document_encrypts_every_page() {
    // Several pages, several streams. A per-object key derivation that is
    // wrong for object 12 shows up here and not in a one-page test.
    let mut builder = PdfDocBuilder::new();
    builder.add_font("F1", helvetica());
    for i in 0..12 {
        builder.add_page(
            200.0,
            200.0,
            format!("BT /F1 12 Tf 10 100 Td (PAGE{i:02}SECRET) Tj ET").as_bytes(),
        );
    }
    builder.set_encryption(settings_for(EncryptionAlgorithm::Aes256, "pw"));
    let pdf = builder.finish();

    for i in 0..12 {
        let marker = format!("PAGE{i:02}SECRET");
        assert!(
            !contains(&pdf, marker.as_bytes()),
            "page {i} was written in the clear"
        );
    }

    let mut doc = PdfDocument::parse_with_password(&pdf, b"pw").expect("opens");
    assert_eq!(doc.page_count(), 12);
    for i in 0..12 {
        let marker = format!("PAGE{i:02}SECRET");
        let page = doc.page(i).expect("page");
        assert!(
            contains(&page.content_data, marker.as_bytes()),
            "page {i} did not decrypt to its own content"
        );
    }
}

#[test]
fn binary_content_survives_encryption() {
    // A content stream is compressed by default, so its bytes are already
    // binary. This checks the padding and IV handling on data that is not
    // conveniently block-aligned text.
    let mut builder = PdfDocBuilder::new();
    let content: Vec<u8> = (0..=255u8).cycle().take(1000).collect();
    builder.set_compress(false);
    builder.add_page(100.0, 100.0, &content);
    builder.set_encryption(settings_for(EncryptionAlgorithm::Aes256, "pw"));
    let pdf = builder.finish();

    let mut doc = PdfDocument::parse_with_password(&pdf, b"pw").expect("opens");
    assert_eq!(doc.page(0).expect("page 0").content_data, content);
}