fix(engine): keep earlier saves when undo crosses a security change (#3276)

From review. When undo restores security from before the last save, the
save now goes onto the newest file of this session that has the
snapshot's security and continues its source, not only the source
itself. An edit saved before the security change keeps its revision;
only revisions in the other security, all undone, are dropped.

save_bytes compares the saved file's security, recorded at mark_saved
as a SecurityIdentity (dictionary and, below revision 5, first /ID),
before opening anything. The chosen file is opened with the snapshot's
key, so a password that no longer opens the last save can't fail it.
One early return covers a snapshot that is its target unchanged.

overlay_changes_from keeps reading this file under its own /Encrypt
number and records the snapshot's only for writing, and refuses a
document with a pending security change.

crates/cos/src/document.rs

@@ -61,6 +61,16 @@ pub enum XrefEntry {
     InStream { stream: u32, index: u32 },
 }
 
+/// How a document is encrypted when written: its security dictionary and, below revision 5,
+/// where the file key is derived from it, the first file identifier. Revisions 5 and 6 wrap a
+/// key of their own in `/UE` and `/OE`. Equal identities encrypt with the same key, so an update
+/// written under one reads correctly under the other.
+#[derive(Clone, Debug, PartialEq)]
+pub struct SecurityIdentity {
+    dict: pdfcraft_crypt::EncryptDict,
+    first_id: Vec<u8>,
+}
+
 #[derive(Clone, Copy, Debug)]
 struct XrefRange {
     start: u32,
@@ -437,6 +447,9 @@ pub struct Document {
     /// The handler and `/Encrypt` object number that saves use, when encryption changed.
     out_security: Option<Arc<pdfcraft_crypt::SecurityHandler>>,
     out_encrypt_num: Option<u32>,
+    /// The `/Encrypt` number saves leave unencrypted when an overlay took it from another
+    /// document ([`Document::overlay_changes_from`]). Reading this file keeps `encrypt_num`.
+    written_encrypt_num: Option<Option<u32>>,
 }
 
 /// A temporary reader for whole-document work. Its caches never populate the editor or
@@ -565,6 +578,7 @@ impl Document {
             full_save: false,
             out_security: None,
             out_encrypt_num: None,
+            written_encrypt_num: None,
         };
         let mut pages = Dict::new();
         pages.set(b"Type".to_vec(), Object::name("Pages"));
@@ -648,6 +662,7 @@ impl Document {
             full_save: false,
             out_security: None,
             out_encrypt_num: None,
+            written_encrypt_num: None,
         };
         let mut log = Vec::new();
         if headerless {
@@ -792,7 +807,11 @@ impl Document {
     /// The handler used to write: the new one after `set_encryption` / `remove_encryption`,
     /// otherwise the one the document was opened with.
     pub(crate) fn output_security(&self) -> (Option<&pdfcraft_crypt::SecurityHandler>, Option<u32>) {
-        if self.encryption_changed { (self.out_security.as_deref(), self.out_encrypt_num) } else { (self.security.as_deref(), self.encrypt_num) }
+        if self.encryption_changed {
+            (self.out_security.as_deref(), self.out_encrypt_num)
+        } else {
+            (self.security.as_deref(), self.written_encrypt_num.unwrap_or(self.encrypt_num))
+        }
     }
 
     /// The security the next save writes: protection applied with `set_encryption`, none after
@@ -801,20 +820,21 @@ impl Document {
         self.output_security().0
     }
 
-    /// True when `self` and `other` are written with the same encryption: the same security
-    /// dictionary, and below revision 5, where the first file identifier is part of the key, the
-    /// same identifier. Together they determine the file key: revisions 2–4 derive it from the
-    /// password, `/O`, `/P` and `/ID`, and revisions 5–6 wrap it in `/UE` and `/OE`. Passwords
-    /// and authentication level are not an encryption identity: the same passwords can protect
-    /// a different key, algorithm or permissions, and user and owner openings share one. Where
-    /// the dictionary is stored (its object number, or inline) is not compared:
-    /// [`Document::overlay_changes_from`] takes that from the document it copies.
+    /// The encryption the next save writes ([`SecurityIdentity`]), or `None` when it writes none.
+    pub fn output_security_identity(&self) -> Option<SecurityIdentity> {
+        let dict = self.output_handler()?.dict().clone();
+        let first_id = if dict.r >= 5 { Vec::new() } else { self.first_id().to_vec() };
+        Some(SecurityIdentity { dict, first_id })
+    }
+
+    /// True when `self` and `other` are written with the same encryption
+    /// ([`Document::output_security_identity`]). Passwords and authentication level are not an
+    /// encryption identity: the same passwords can protect a different key, algorithm or
+    /// permissions, and user and owner openings share one. Where the dictionary is stored (its
+    /// object number, or inline) is not compared: [`Document::overlay_changes_from`] takes that
+    /// from the document it copies.
     pub fn has_same_output_security(&self, other: &Document) -> bool {
-        match (self.output_handler(), other.output_handler()) {
-            (None, None) => true,
-            (Some(a), Some(b)) => a.dict() == b.dict() && (a.dict().r >= 5 || self.first_id() == other.first_id()),
-            _ => false,
-        }
+        self.output_security_identity() == other.output_security_identity()
     }
 
     /// Protect the document with a password (§7.6.4). Takes effect on the next save, which is
@@ -1349,18 +1369,19 @@ impl Document {
     /// Record `current`'s object graph as unsaved edits of this document.
     ///
     /// Used when the working snapshot was restored from an earlier revision and the next
-    /// incremental save must still append to this file (the last saved revision), opened afresh:
-    /// which objects sit in object streams is read from the file. Objects that already match are
-    /// left untouched. Reads go through private readers so neither
-    /// document's editing cache is filled.
+    /// incremental save must still append to this file. Objects that already match are left
+    /// untouched, except ones packed in an object stream whose number `current` reuses or lacks:
+    /// those are written again on their own. Reads go through private readers, so neither
+    /// document's editing cache is filled. The `/Encrypt` dictionary comes with `current`'s objects and trailer, and so does
+    /// the number saves leave unencrypted, even where a full save had moved it; reading this
+    /// file still uses its own.
     ///
-    /// Returns `false`, changing nothing, when output security differs
-    /// ([`Document::has_same_output_security`]): an update appended to this file would be read
-    /// under other security than the objects before it, so append `current` to its own source. Otherwise the `/Encrypt` dictionary
-    /// comes with `current`'s objects and trailer, and so does the number the writer leaves
-    /// unencrypted, even where a full save had moved the dictionary.
+    /// Returns `Ok(false)`, having changed nothing, when output security differs
+    /// ([`Document::has_same_output_security`]) or this document has a pending security change.
+    /// An update appended to this file would then be read under other security than the
+    /// objects before it.
     pub fn overlay_changes_from(&mut self, current: &Document) -> Result<bool, CosError> {
-        if !self.has_same_output_security(current) {
+        if self.encryption_changed || !self.has_same_output_security(current) {
             return Ok(false);
         }
         let saved_nums = self.object_numbers();
@@ -1409,14 +1430,7 @@ impl Document {
             self.set(ObjRef::new(num, generation), (*object).clone());
         }
         *self.trailer_mut() = current.trailer().clone();
-        // Both numbers normally changed above (or the dictionary is inline), so nothing is read
-        // from this file under the other exemption.
-        let encrypt_num = current.output_security().1;
-        if self.encryption_changed {
-            self.out_encrypt_num = encrypt_num;
-        } else {
-            self.encrypt_num = encrypt_num;
-        }
+        self.written_encrypt_num = Some(current.output_security().1);
         Ok(true)
     }
 
@@ -2359,6 +2373,52 @@ const MAX_XREF_OBJECTS: usize = 8_388_607;
 mod tests {
     use super::*;
 
+    /// An overlay moves where saves leave `/Encrypt` unencrypted, never where reading this file
+    /// does: its own dictionary is still at its own number.
+    #[test]
+    fn an_overlay_moves_only_the_written_encrypt_exemption() {
+        fn protected(info: bool) -> Vec<u8> {
+            let objs = [
+                "<< /Type /Catalog /Pages 2 0 R >>",
+                "<< /Type /Pages /Kids [3 0 R] /Count 1 /MediaBox [0 0 300 200] >>",
+                "<< /Type /Page /Parent 2 0 R >>",
+                "<< /Title (t) >>",
+            ];
+            let mut out = b"%PDF-1.7\n".to_vec();
+            let mut offsets = Vec::new();
+            for (i, o) in objs.iter().enumerate() {
+                offsets.push(out.len());
+                out.extend_from_slice(format!("{} 0 obj\n{o}\nendobj\n", i + 1).as_bytes());
+            }
+            let xref = out.len();
+            out.extend_from_slice(format!("xref\n0 {}\n0000000000 65535 f \n", objs.len() + 1).as_bytes());
+            for o in offsets {
+                out.extend_from_slice(format!("{o:010} 00000 n \n").as_bytes());
+            }
+            let info = if info { " /Info 4 0 R" } else { "" };
+            out.extend_from_slice(format!("trailer\n<< /Size {} /Root 1 0 R{info} >>\nstartxref\n{xref}\n%%EOF\n", objs.len() + 1).as_bytes());
+            let mut doc = Document::open(Arc::new(out)).unwrap();
+            let params = crate::NewEncryption {
+                algorithm: crate::Algorithm::Aes256,
+                user_password: "u",
+                owner_password: "o",
+                permissions: -1,
+                encrypt_metadata: true,
+                seed: [1; 32],
+            };
+            doc.set_encryption(&params).unwrap();
+            crate::write_full(&doc, &crate::SaveOptions::default()).unwrap()
+        }
+        let open = |bytes: Vec<u8>| Document::open_with_password(Arc::new(bytes), Some("o")).unwrap();
+        let mut base = open(protected(true));
+        let current = open(protected(false));
+        let (own, theirs) = (base.encrypt_num, current.encrypt_num);
+        assert!(own.is_some() && theirs.is_some() && own != theirs);
+        assert!(base.overlay_changes_from(&current).unwrap());
+        assert_eq!(base.encrypt_num, own, "reading keeps this file's own dictionary");
+        assert_eq!(base.output_security().1, theirs, "saves leave the overlaid trailer's dictionary unencrypted");
+    }
+
     #[test]
     fn a_read_state_matches_a_clone_and_no_changed_document() {
         let mut doc = Document::new_empty();

crates/cos/src/lib.rs

@@ -24,7 +24,7 @@ mod writer;
 
 pub use budget::{DecodeBudget, MAX_PAGE_DECODED, MAX_PAGE_FORMS};
 pub use bytes::Bytes;
-pub use document::{Document, ReadState, Revision, SourceIdentity, XrefEntry};
+pub use document::{Document, ReadState, Revision, SecurityIdentity, SourceIdentity, XrefEntry};
 pub use object::{Dict, MAX_DECODED, Name, ObjRef, Object, PdfString, Stream};
 pub use parser::{Lexer, is_whitespace, parse_indirect};
 pub use pdfcraft_crypt::{

crates/cos/tests/it/encryption.rs

@@ -598,6 +598,17 @@ fn an_overlay_with_different_output_security_is_refused_before_any_change() {
         assert!(before == state(&target), "{case}: a refused overlay changed the target");
     }
 
+    // A target with a pending security change is refused even when the snapshot matches what
+    // that change will write.
+    let mut pending = Document::open(Arc::new(fixture())).unwrap();
+    rekey(&mut pending, Algorithm::Aes256, 9);
+    set_title(&mut pending, "Saved title");
+    let matching = Document::open_with_password(Arc::new(protect(Algorithm::Aes256, "user", "owner", -1)), Some("owner")).unwrap();
+    assert!(pending.has_same_output_security(&matching));
+    let before = state(&pending);
+    assert!(matches!(pending.overlay_changes_from(&matching), Ok(false)), "pending security change");
+    assert!(before == state(&pending), "a refused overlay changed the target");
+
     // Control: the same overlay with matching security does change the target.
     let mut target = encrypted_target();
     let mut current = open_owner();

crates/engine/src/history.rs

@@ -108,6 +108,9 @@ pub(crate) struct Editor {
     pub(crate) saved: Arc<Vec<u8>>,
     /// Passwords that open [`Self::saved`].
     pub(crate) saved_keys: Keys,
+    /// How [`Self::saved`] is encrypted, recorded with it: compared before deciding what a save
+    /// can append to.
+    pub(crate) saved_security: Option<pdfcraft_cos::SecurityIdentity>,
     /// [`Self::saved`] was never written where the document belongs: a new document (combine,
     /// extract, create) or one restored from recovery. It stays unsaved until a save, whatever
     /// undo brings back, or closing it would lose the whole document without asking.

crates/engine/src/open.rs

@@ -271,11 +271,13 @@ impl Session {
             Ok(Ok(cos)) => {
                 let keys = Keys { render: render_password.clone(), reopen: password.map(str::to_owned) };
                 let saved = cos.bytes().clone();
+                let saved_security = cos.output_security_identity();
                 (
                     Some(Editor {
                         cos,
                         saved,
                         saved_keys: keys.clone(),
+                        saved_security,
                         unsaved_file: false,
                         undo: Vec::new(),
                         redo: Vec::new(),

crates/engine/src/session.rs

@@ -8,7 +8,7 @@ use pdfcraft_render::{LayerOp, OpenError, ReadingColors, RenderConfig, RenderPoo
 #[cfg(test)]
 use crate::check_refresh_panic;
 use crate::edit::{EditCtx, check_permission, run_edit};
-use crate::history::{Keys, MAX_UNDO, Scope, keys_after, scope_of, uses_scripts};
+use crate::history::{Editor, Keys, MAX_UNDO, Scope, keys_after, scope_of, uses_scripts};
 use crate::{
     BookmarkPage, BookmarkRow, CombineSource, Combined, DataFormat, DocId, Document, Edit, EditError, Hidden, ImageResolution,
     LOW_RESOLUTION_PRINT_DPI, Label, ObjectKind, OptimizeReport, RecoverySnapshot, SignOptions, SourceKind, SplitPart, SummarySort, TrustStore,
@@ -665,38 +665,33 @@ impl Session {
     /// break its signatures, so with no edits it keeps its bytes, damage and all.
     ///
     /// Undo past a save that changed security restores a snapshot whose encryption differs from
-    /// the last saved file. An update appended to that file would be read with the wrong key and
-    /// keep the revision it was meant to replace, so the snapshot is appended to its own source
-    /// instead, whose encryption it shares (#3276). That keeps the source's revisions, object
-    /// numbers and signatures. Revisions saved after that source are dropped: undo went back
-    /// past all of them. An unedited snapshot is its source's bytes as they are.
+    /// the last saved file; see [`Self::save_target`] for where that is appended instead (#3276).
     pub fn save_bytes(&self, id: DocId) -> Result<Arc<Vec<u8>>, EditError> {
         let doc = self.get(id).ok_or(EditError::NoDocument)?;
         let Some(editor) = doc.editor.as_ref() else { return Ok(doc.bytes.clone()) };
-        let on_saved = Arc::ptr_eq(editor.cos.bytes(), &editor.saved);
-        if !editor.cos.is_modified() && on_saved && (!editor.cos.reconstructed() || doc.is_signed()) {
-            return Ok(editor.saved.clone());
+        let target = Self::save_target(editor);
+        let on_target = Arc::ptr_eq(editor.cos.bytes(), &target);
+        if !editor.cos.is_modified() && on_target && (!editor.cos.reconstructed() || doc.is_signed()) {
+            return Ok(target);
         }
         // What the user removed again before saving (a cleared redaction mark) stays out of the file.
         let opts = SaveOptions { mod_date: self.now().map(pdfcraft_cos::pdf_date), drop_unlinked: true, ..SaveOptions::default() };
-        // A restored snapshot can be based on an earlier revision. Append to the last saved
-        // file when it has the snapshot's security, so that revision stays in the chain. A full
-        // rewrite (redaction, new encryption, a reconstructed file) still writes the current
-        // graph on its own.
+        // A restored snapshot can be based on an earlier revision than the target. Append it to
+        // the target so that file's revisions stay in the chain. A full rewrite (redaction, new
+        // encryption, a reconstructed file) still writes the current graph on its own.
         let rewrite = editor.cos.full_save_required() || editor.cos.encryption_changed() || editor.cos.revisions().is_empty();
         let cos = writable(&editor.cos)?;
         guard(|| {
-            if on_saved || rewrite {
+            if on_target || rewrite {
                 return write_incremental(&cos, &opts).map(Arc::new);
             }
-            let mut base = pdfcraft_cos::Document::open_with_password(editor.saved.clone(), editor.saved_keys.reopen.as_deref())?;
+            // The target has the snapshot's security, so the snapshot's key opens it.
+            let mut base = editor.cos.open_with_same_credentials(target.clone())?;
             if base.overlay_changes_from(&cos)? {
                 write_incremental(&base, &opts).map(Arc::new)
-            } else if !cos.is_modified() {
-                Ok(cos.bytes().clone())
             } else {
-                // Undo restored security from before the last save: append the snapshot to its
-                // own source, which it hasn't changed security against.
+                // Not expected, as the target was chosen for this security: append to the
+                // snapshot's own source, which shares it.
                 write_incremental(&cos, &opts).map(Arc::new)
             }
         })
@@ -704,6 +699,32 @@ impl Session {
         .map_err(|e| EditError::Write(e.to_string()))
     }
 
+    /// The file a save appends to: the last saved file, unless undo restored security from before
+    /// it. An update appended there would be read with the wrong key and keep the revision it was
+    /// meant to replace (#3276). The snapshot then goes onto the newest file of this session that
+    /// has its security and continues its own source, which at worst is that source. Saves append
+    /// to the file they continue, so within one lineage the longest file is the newest. Revisions
+    /// saved after it have the other security and were all undone: they are dropped, as any
+    /// change of security drops the revisions before it.
+    fn save_target(editor: &Editor) -> Arc<Vec<u8>> {
+        let security = editor.cos.output_security_identity();
+        if editor.cos.encryption_changed() || security == editor.saved_security {
+            return editor.saved.clone();
+        }
+        let source = editor.cos.bytes();
+        editor
+            .undo
+            .iter()
+            .chain(&editor.redo)
+            .map(|(_, cos, _, _)| cos)
+            .filter(|cos| !cos.encryption_changed() && !cos.reconstructed() && cos.output_security_identity() == security)
+            .map(|cos| cos.bytes())
+            .filter(|bytes| !source.is_empty() && bytes.len() > source.len() && bytes.starts_with(source))
+            .max_by_key(|bytes| bytes.len())
+            .unwrap_or(source)
+            .clone()
+    }
+
     /// A compact, garbage-collected rewrite (Save As ▸ "Optimized" / Reduce File Size groundwork).
     pub fn save_full_bytes(&self, id: DocId) -> Result<Arc<Vec<u8>>, EditError> {
         let doc = self.get(id).ok_or(EditError::NoDocument)?;
@@ -745,6 +766,7 @@ impl Session {
             }
             editor.saved = bytes;
             editor.saved_keys = keys;
+            editor.saved_security = editor.cos.output_security_identity();
             editor.unsaved_file = false;
         }
         if let Some(p) = path {

crates/engine/src/tests/security_save.rs

@@ -323,6 +323,39 @@ fn undo_past_a_security_save_keeps_the_snapshot_edits() {
     }
 }
 
+/// Edit, save, change security, save, then undo past both and save: the edit's saved revision
+/// stays in the chain, in the original security; only the revision in the other security goes.
+#[test]
+fn undo_past_a_security_save_keeps_earlier_saves_of_the_same_security() {
+    for encrypted in [false, true] {
+        let (original, password) =
+            if encrypted { (encrypted_fixture(Algorithm::Aes256, false), Some("owner")) } else { (Arc::new(fixture(2)), None) };
+        let mut s = Session::new();
+        let id = s.open("a.pdf", None, original.clone(), password).unwrap();
+        s.apply(id, Edit::RotatePages { pages: vec![0], degrees: 90 }).unwrap();
+        let edited = s.save_bytes(id).unwrap();
+        assert!(edited.starts_with(original.as_slice()), "{encrypted}: the edit is an incremental save");
+        s.mark_saved(id, edited.clone(), None).unwrap();
+        let change = if encrypted { Edit::RemoveProtection } else { Edit::Protect(protection(Some("user"), Some("owner"))) };
+        s.apply(id, change).unwrap();
+        let changed = s.save_bytes(id).unwrap();
+        s.mark_saved(id, changed, None).unwrap();
+        s.undo(id).unwrap();
+        s.undo(id).unwrap();
+        let restored = s.save_bytes(id).unwrap();
+        assert!(restored.starts_with(edited.as_slice()) && restored.len() > edited.len(), "{encrypted}: appended to the edit's save");
+        let mut reopened = Session::new();
+        let rid = reopened.open("restored.pdf", None, restored.clone(), password).unwrap();
+        let doc = reopened.get(rid).unwrap();
+        assert_eq!((doc.info.encrypted, doc.info.pages[0].rotation), (encrypted, 0), "{encrypted}: the rotation is undone");
+        assert_eq!(page_texts(&reopened, rid), ["Page 1", "Page 2"], "{encrypted}");
+        assert_eq!(restored.windows(b"(Page 1)".len()).any(|w| w == b"(Page 1)"), !encrypted, "{encrypted}: plaintext only in a plain file");
+        s.mark_saved(id, restored.clone(), None).unwrap();
+        assert!(!s.get(id).unwrap().dirty, "{encrypted}");
+        assert_eq!(*s.save_bytes(id).unwrap(), *restored, "{encrypted}: an unchanged save is byte-stable");
+    }
+}
+
 /// An optimized save packs objects into object streams under new numbers. Undo past it restores
 /// a graph that puts other objects at those numbers: objects that look unchanged but sat in such
 /// a stream must still be readable from the next save.