fast node walk without decrypt + fixed crypto dub bug

Cargo.lock

@@ -1959,7 +1959,7 @@ dependencies = [
 
 [[package]]
 name = "proton-drive-rs"
-version = "0.1.6"
+version = "0.1.7"
 dependencies = [
  "async-trait",
  "base64",
@@ -1983,7 +1983,7 @@ dependencies = [
 
 [[package]]
 name = "proton-sdk"
-version = "0.1.6"
+version = "0.1.7"
 dependencies = [
  "aes-gcm 0.11.0",
  "async-trait",

Cargo.toml

@@ -3,7 +3,7 @@ resolver = "2"
 members = ["crates/proton-sdk", "crates/proton-drive-rs"]
 
 [workspace.package]
-version = "0.1.6"
+version = "0.1.7"
 edition = "2024"
 license = "MIT"
 repository = "https://github.com/narrrl/proton-sdk-rs"

crates/proton-drive-rs/src/client.rs

@@ -33,7 +33,7 @@ use sha1::Sha1;
 use crate::cache::DriveEntityCache;
 use crate::crypto::{
     decrypt_content_key_verified, decrypt_extended_attributes_verified, decrypt_link,
-    decrypt_link_verified, decrypt_share_key,
+    decrypt_link_name_verified, decrypt_link_verified, decrypt_share_key,
 };
 use crate::devices::{Device, DeviceMetadata, DeviceType};
 use crate::dtos::{
@@ -82,6 +82,23 @@ const TRASH_PAGE_SIZE: usize = 500;
 /// Photos returned per timeline page (C# `PhotosNodeOperations.TimelinePageSize`).
 const TIMELINE_PAGE_SIZE: usize = 500;
 
+/// How much of a node to decrypt when building it.
+///
+/// Unlocking a node's own key costs an S2K derivation (tens of milliseconds), so
+/// it is worth avoiding for callers that do not need what it protects.
+#[derive(Clone, Copy, PartialEq, Eq)]
+enum NodeDetail {
+    /// Decrypt everything: the node's key, and for a file the content key and
+    /// extended attributes it protects (so `claimed_size` and
+    /// `claimed_modification_time` are populated).
+    Full,
+    /// Decrypt only what the *parent* key can read: name, and signature
+    /// statuses. A file's key is left locked, leaving its claimed metadata and
+    /// content-key/xattr verification absent. Folders are unaffected — their key
+    /// is needed to reach their children.
+    Light,
+}
+
 /// Candidate name hashes checked per `checkAvailableHashes` request (C#
 /// `NodeOperations.GetAvailableNameAsync` `batchSize`).
 const NAME_AVAILABILITY_BATCH: usize = 10;
@@ -279,7 +296,7 @@ impl ProtonDriveClient {
             .resolve_parent_key(&uid.volume_id, &details.link)
             .await?;
         let node = self
-            .build_node(&uid.volume_id, &details, &parent_key)
+            .build_node(&uid.volume_id, &details, &parent_key, NodeDetail::Full)
             .await?;
         timer.success();
         Ok(Some(node))
@@ -343,23 +360,54 @@ impl ProtonDriveClient {
     /// to decrypt is logged and skipped (matching enumeration's partial-node
     /// behavior), so the result may be shorter than `uids`.
     pub async fn enumerate_nodes(&self, uids: &[NodeUid]) -> Result<Vec<Node>> {
+        self.enumerate_nodes_detail(uids, NodeDetail::Full).await
+    }
+
+    /// As [`enumerate_nodes`](Self::enumerate_nodes), but skips the parts of a
+    /// *file* that only its own node key can decrypt. The returned files carry
+    /// no [`claimed_size`](NodeKind::File::claimed_size) or
+    /// [`claimed_modification_time`](NodeKind::File::claimed_modification_time),
+    /// and no content-key or extended-attribute verification status. Everything
+    /// else — name, uid, parentage, timestamps, trashed/shared flags, and
+    /// `total_size_on_storage` — is identical. Folders are returned in full.
+    ///
+    /// Unlocking a file's node key costs an S2K derivation of tens of
+    /// milliseconds, which dominates a walk over a large tree. Use this to
+    /// enumerate cheaply and find what changed, then call
+    /// [`enumerate_nodes`](Self::enumerate_nodes) for the few nodes whose
+    /// claimed metadata you actually need.
+    pub async fn enumerate_nodes_light(&self, uids: &[NodeUid]) -> Result<Vec<Node>> {
+        self.enumerate_nodes_detail(uids, NodeDetail::Light).await
+    }
+
+    async fn enumerate_nodes_detail(
+        &self,
+        uids: &[NodeUid],
+        detail: NodeDetail,
+    ) -> Result<Vec<Node>> {
         let mut nodes = Vec::new();
 
         for (volume_id, link_ids) in group_by_volume(uids) {
             for chunk in link_ids.chunks(MAX_BATCH_COUNT) {
                 let details = self.get_link_details(&volume_id, chunk).await?;
-                for detail in &details.links {
-                    let parent_key = match self.resolve_parent_key(&volume_id, &detail.link).await {
+                for link_details in &details.links {
+                    let parent_key = match self
+                        .resolve_parent_key(&volume_id, &link_details.link)
+                        .await
+                    {
                         Ok(key) => key,
                         Err(e) => {
-                            tracing::warn!(link_id = %detail.link.id, error = %e, "skipping node: parent key unavailable");
+                            tracing::warn!(link_id = %link_details.link.id, error = %e, "skipping node: parent key unavailable");
                             continue;
                         }
                     };
-                    match self.build_node(&volume_id, detail, &parent_key).await {
+                    match self
+                        .build_node(&volume_id, link_details, &parent_key, detail)
+                        .await
+                    {
                         Ok(node) => nodes.push(node),
                         Err(e) => {
-                            tracing::warn!(link_id = %detail.link.id, error = %e, "skipping undecryptable node");
+                            tracing::warn!(link_id = %link_details.link.id, error = %e, "skipping undecryptable node");
                         }
                     }
                 }
@@ -3307,6 +3355,39 @@ impl ProtonDriveClient {
             .clone())
     }
 
+    /// The share key that unlocks a volume-*root* link. Normally that's the My
+    /// Files root (backed by the main share key), but a registered device's
+    /// root folder is also parentless and lives on the same volume while being
+    /// wrapped to that device's *own* share key. Match the root link id against
+    /// the device set and unlock the matching device share; fall back to the My
+    /// Files share for anything unrecognised (preserving the prior behaviour).
+    async fn root_link_share_key(
+        &self,
+        volume_id: &VolumeId,
+        root_id: &LinkId,
+    ) -> Result<PrivateKey> {
+        self.ensure_my_files().await?;
+        if self
+            .cache
+            .lock()
+            .await
+            .my_files_root
+            .as_ref()
+            .is_some_and(|uid| &uid.link_id == root_id)
+        {
+            return self.root_share_key().await;
+        }
+        if let Ok(devices) = self.device_metadata().await
+            && let Some(device) = devices.into_iter().find(|d| {
+                &d.root_folder_uid.link_id == root_id && &d.root_folder_uid.volume_id == volume_id
+            })
+        {
+            let (key, _address_id) = self.share_key_by_id(&device.share_id).await?;
+            return Ok(key);
+        }
+        self.root_share_key().await
+    }
+
     /// Resolve the Photos share + root folder, caching them.
     ///
     /// Mirrors C# `PhotosNodeOperations.GetFreshExistingPhotosFolderAsync`:
@@ -3408,7 +3489,13 @@ impl ProtonDriveClient {
             .resolve_parent_key_ctx(&uid.volume_id, &details.link, true)
             .await?;
         let node = self
-            .build_node_ctx(&uid.volume_id, &details, &parent_key, true)
+            .build_node_ctx(
+                &uid.volume_id,
+                &details,
+                &parent_key,
+                true,
+                NodeDetail::Full,
+            )
             .await?;
         Ok(Some(node))
     }
@@ -3433,7 +3520,7 @@ impl ProtonDriveClient {
                         }
                     };
                     match self
-                        .build_node_ctx(&volume_id, detail, &parent_key, true)
+                        .build_node_ctx(&volume_id, detail, &parent_key, true, NodeDetail::Full)
                         .await
                     {
                         Ok(node) => nodes.push(node),
@@ -3645,7 +3732,15 @@ impl ProtonDriveClient {
         let mut key = match base_key {
             Some(key) => key,
             None if for_photos => self.photos_share_key().await?,
-            None => self.root_share_key().await?,
+            None => {
+                // The topmost ancestor (or `link` itself when it is parentless)
+                // is a volume root. Usually that's My Files, but device roots
+                // live on the same volume with their *own* share key, so pick
+                // the share by matching the root link id rather than assuming
+                // My Files (which would fail to decrypt: "missing key").
+                let root_id = ancestry.last().map_or(&link.id, |a| &a.id);
+                self.root_link_share_key(volume_id, root_id).await?
+            }
         };
 
         for ancestor in ancestry.iter().rev() {
@@ -3693,8 +3788,9 @@ impl ProtonDriveClient {
         volume_id: &VolumeId,
         details: &LinkDetailsDto,
         parent_key: &PrivateKey,
+        detail: NodeDetail,
     ) -> Result<Node> {
-        self.build_node_ctx(volume_id, details, parent_key, false)
+        self.build_node_ctx(volume_id, details, parent_key, false, detail)
             .await
     }
 
@@ -3704,10 +3800,26 @@ impl ProtonDriveClient {
         details: &LinkDetailsDto,
         parent_key: &PrivateKey,
         for_photos: bool,
+        detail: NodeDetail,
     ) -> Result<Node> {
         let link = &details.link;
-        let (decrypted, mut verification) =
-            decrypt_link_verified(&self.account, parent_key, link).await?;
+
+        // A file's node key is only needed to read its contents (content key,
+        // extended attributes), so a `Light` build of a file skips unlocking it
+        // — that unlock is an S2K derivation, and it dominates a walk over many
+        // nodes. A folder's key is always unlocked: it is what its children are
+        // decrypted with, so the walk cannot proceed without it.
+        let want_node_key =
+            detail == NodeDetail::Full || !matches!(link.parsed_type(), LinkType::File);
+        let (name, node_key, mut verification) = if want_node_key {
+            let (decrypted, verification) =
+                decrypt_link_verified(&self.account, parent_key, link).await?;
+            (decrypted.name, Some(decrypted.node_key), verification)
+        } else {
+            let (name, verification) =
+                decrypt_link_name_verified(&self.account, parent_key, link).await?;
+            (name, None, verification)
+        };
 
         let uid = NodeUid::new(volume_id.clone(), link.id.clone());
         let parent_uid = link
@@ -3717,12 +3829,15 @@ impl ProtonDriveClient {
 
         let kind = match link.parsed_type() {
             LinkType::Folder | LinkType::Album => {
+                let node_key = node_key.as_ref().ok_or_else(|| {
+                    ProtonError::invalid_operation("folder node built without its key")
+                })?;
                 // Cache the folder's node key for later child enumeration.
                 self.cache
                     .lock()
                     .await
                     .folder_keys
-                    .insert(uid.clone(), decrypted.node_key.clone());
+                    .insert(uid.clone(), node_key.clone());
                 NodeKind::Folder
             }
             LinkType::File => {
@@ -3733,12 +3848,14 @@ impl ProtonDriveClient {
                 // Decrypt + verify the content key (`ContentKeyPacketSignature`).
                 // Best-effort: a decode/decrypt failure leaves the status absent
                 // rather than failing the whole node.
-                if let Some(packet_b64) = file.content_key_packet.as_deref() {
+                if let Some(node_key) = node_key.as_ref()
+                    && let Some(packet_b64) = file.content_key_packet.as_deref()
+                {
                     match BASE64.decode(packet_b64) {
                         Ok(packet) => {
                             match decrypt_content_key_verified(
                                 &self.account,
-                                &decrypted.node_key,
+                                node_key,
                                 link.signature_email.as_deref(),
                                 &packet,
                                 file.content_key_signature.as_deref(),
@@ -3764,12 +3881,13 @@ impl ProtonDriveClient {
                 // leaves the claimed metadata absent rather than failing the node.
                 let mut claimed_size = None;
                 let mut claimed_modification_time = None;
-                if let Some(rev) = file.active_revision.as_ref()
+                if let Some(node_key) = node_key.as_ref()
+                    && let Some(rev) = file.active_revision.as_ref()
                     && let Some(xattr) = rev.extended_attributes.as_deref()
                 {
                     match decrypt_extended_attributes_verified(
                         &self.account,
-                        &decrypted.node_key,
+                        node_key,
                         rev.signature_email.as_deref(),
                         xattr,
                     )
@@ -3812,7 +3930,7 @@ impl ProtonDriveClient {
             uid: uid.clone(),
             parent_uid,
             kind,
-            name: decrypted.name,
+            name,
             creation_time: link.creation_time,
             modification_time: link.modification_time,
             trashed: link.is_trashed(),

crates/proton-drive-rs/src/crypto.rs

@@ -107,6 +107,38 @@ pub async fn decrypt_link_verified(
     parent_key: &PrivateKey,
     link: &LinkDto,
 ) -> Result<(DecryptedLink, NodeVerification)> {
+    let (passphrase, name, verification) = decrypt_link_parts(account, parent_key, link).await?;
+    let node_key = PrivateKey::from_armored(&link.key, &passphrase)?;
+    Ok((DecryptedLink { node_key, name }, verification))
+}
+
+/// As [`decrypt_link_verified`], but stops short of unlocking the node's own
+/// key: it yields only what the parent key can read (the name, and the
+/// signature statuses).
+///
+/// Unlocking a node key runs its S2K, which costs tens of milliseconds — per
+/// node. A caller that only needs a link's name and identity (enumerating a tree
+/// to find what changed, say) pays that for nothing, so it can stop here. The
+/// node key is only needed to reach a node's *contents*: its content key, its
+/// extended attributes, or its children.
+pub async fn decrypt_link_name_verified(
+    account: &AccountClient,
+    parent_key: &PrivateKey,
+    link: &LinkDto,
+) -> Result<(String, NodeVerification)> {
+    let (_passphrase, name, verification) = decrypt_link_parts(account, parent_key, link).await?;
+    Ok((name, verification))
+}
+
+/// The shared body of [`decrypt_link_verified`] and
+/// [`decrypt_link_name_verified`]: everything the parent key alone can decrypt.
+/// Returns the node passphrase (for callers that go on to unlock the node key),
+/// the name, and the signature statuses.
+async fn decrypt_link_parts(
+    account: &AccountClient,
+    parent_key: &PrivateKey,
+    link: &LinkDto,
+) -> Result<(Vec<u8>, String, NodeVerification)> {
     let node_claim = AuthorshipClaim::create(account, link.signature_email.as_deref()).await;
     let name_claim = if link.name_signature_email == link.signature_email {
         node_claim.clone()
@@ -117,7 +149,6 @@ pub async fn decrypt_link_verified(
     // Passphrase: decrypted with the parent key, signed (detached) by the node
     // author; verify against the node claim (fallback = parent key).
     let passphrase = parent_key.decrypt_armored_message(&link.passphrase)?;
-    let node_key = PrivateKey::from_armored(&link.key, &passphrase)?;
     let passphrase_status = match &link.passphrase_signature {
         Some(sig) => verify_detached(sig, &passphrase, &node_claim.ring(parent_key)),
         None => VerificationStatus::NotSigned,
@@ -136,7 +167,7 @@ pub async fn decrypt_link_verified(
         content_key: None,
         extended_attributes: None,
     };
-    Ok((DecryptedLink { node_key, name }, verification))
+    Ok((passphrase, name, verification))
 }
 
 /// Decrypt a file's content key and verify its `ContentKeyPacketSignature`.

crates/proton-sdk/src/crypto/keys.rs

@@ -1,8 +1,10 @@
 //! PGP secret-key parsing and the layered Proton decryption chain.
 //!
-//! rPGP 0.16 decrypts directly with a `(key, password)` pair rather than a
-//! distinct pre-unlock step, so a [`PrivateKey`] simply pairs a parsed
-//! [`SignedSecretKey`] with the passphrase that unlocks it.
+//! rPGP decrypts with a `(key, password)` pair rather than a distinct pre-unlock
+//! step, so a [`PrivateKey`] pairs a parsed [`SignedSecretKey`] with the
+//! passphrase that unlocks it. The secret material is decrypted once at parse
+//! time (see [`PrivateKey::from_armored`]) so that each later operation does not
+//! re-run the key's S2K derivation.
 
 use pgp::composed::{Deserializable, DetachedSignature, SignedSecretKey};
 use pgp::types::Password;
@@ -13,7 +15,9 @@ use super::messages;
 /// A Proton private key: a parsed secret key plus the passphrase that unlocks it.
 ///
 /// `Password` is not `Clone`, so we retain the raw passphrase bytes and build a
-/// `Password` on demand for each operation.
+/// `Password` on demand for each operation. The passphrase is still needed after
+/// construction only for key material that could not be decrypted up front (see
+/// [`PrivateKey::from_armored`]); decrypted material ignores the password.
 #[derive(Clone)]
 pub struct PrivateKey {
     key: SignedSecretKey,
@@ -22,19 +26,40 @@ pub struct PrivateKey {
 
 impl PrivateKey {
     /// Parse an armored secret key and validate that `passphrase` unlocks it.
+    ///
+    /// The secret material is decrypted in place and kept that way. rPGP derives
+    /// the key's S2K on *every* operation that takes a locked key plus a
+    /// password, and Proton's S2K is deliberately expensive — so leaving the key
+    /// locked here costs one full derivation per decrypt/sign, which dominates
+    /// any walk over many nodes (each node key is unlocked to read its children).
+    /// Decrypting once turns that into a one-off cost per key.
+    ///
+    /// This keeps the plaintext secret material resident for the key's lifetime,
+    /// which is no weaker than the alternative: we hold `passphrase` alongside
+    /// the locked key regardless, so the material is recoverable from memory
+    /// either way. Only the in-memory form changes — the armored key uploaded to
+    /// the server is built separately and stays locked.
     pub fn from_armored(armored: &str, passphrase: &[u8]) -> Result<Self, CryptoError> {
-        let (key, _headers) = SignedSecretKey::from_string(armored)
+        let (mut key, _headers) = SignedSecretKey::from_string(armored)
             .map_err(|e| CryptoError::Parse(format!("secret key: {e}")))?;
 
-        // Confirm the passphrase is correct up front, mirroring the C# `Unlock`.
-        // `unlock` returns `Result<Result<T>>`: the outer layer reports whether
-        // the secret params could be decrypted (i.e. the passphrase fit), the
-        // inner layer is our (trivial) closure result.
+        // Decrypting the primary key doubles as the passphrase check the C#
+        // `Unlock` does: a wrong passphrase cannot decrypt the secret params.
+        // (A key whose material is already plain is left alone and accepted, as
+        // an `unlock` check would have been.)
         let password = password_from_bytes(passphrase);
-        key.unlock(&password, |_, _| Ok(()))
-            .map_err(|e| CryptoError::Unlock(e.to_string()))?
+        key.primary_key
+            .remove_password(&password)
             .map_err(|e| CryptoError::Unlock(e.to_string()))?;
 
+        // Subkeys are locked with the same passphrase, so this normally succeeds
+        // for all of them. A subkey that resists is left locked rather than
+        // failing the whole key: it still works via `password()` (just at full
+        // S2K cost per use), and the primary may be all the caller needs.
+        for subkey in &mut key.secret_subkeys {
+            let _ = subkey.key.remove_password(&password);
+        }
+
         Ok(Self {
             key,
             passphrase: passphrase.to_vec(),
@@ -127,3 +152,49 @@ pub fn decrypt_armored_with_keys(
 pub(crate) fn password_from_bytes(bytes: &[u8]) -> Password {
     Password::from(bytes)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::crypto::encrypt::generate_node_key;
+
+    #[test]
+    fn from_armored_rejects_a_wrong_passphrase() {
+        // `from_armored` doubles as the passphrase check, so a key that does not
+        // belong to the passphrase must not parse.
+        let node = generate_node_key().expect("generate node key");
+        // Not `expect_err`: `PrivateKey` deliberately has no `Debug` (it holds
+        // decrypted secret material), so match the result instead.
+        match PrivateKey::from_armored(&node.locked_armored, b"not the passphrase") {
+            Err(CryptoError::Unlock(_)) => {}
+            Err(e) => panic!("expected an unlock error, got {e:?}"),
+            Ok(_) => panic!("wrong passphrase must not unlock the key"),
+        }
+    }
+
+    #[test]
+    fn key_decrypted_at_parse_time_still_decrypts_and_signs() {
+        // The secret material is decrypted in place at parse time; every
+        // operation must keep working against a key re-parsed from its armored
+        // (locked) form.
+        let node = generate_node_key().expect("generate node key");
+        let reparsed = PrivateKey::from_armored(&node.locked_armored, &node.passphrase)
+            .expect("correct passphrase unlocks the key");
+
+        // Decrypt (exercises the X25519 subkey) — the message is encrypted to the
+        // original handle and must open with the re-parsed one.
+        let plaintext = b"node passphrase".to_vec();
+        let armored = node.key.encrypt(&plaintext).expect("encrypt");
+        assert_eq!(
+            reparsed.decrypt_armored_message(&armored).expect("decrypt"),
+            plaintext
+        );
+
+        // Sign (exercises the Ed25519 primary).
+        let data = b"manifest bytes";
+        let sig = reparsed.sign_detached(data).expect("detached sign");
+        reparsed
+            .verify_detached_signature(&sig, data)
+            .expect("verify detached signature");
+    }
+}