fix(sdk) token mutex no longer spans the refresh

crates/proton-drive-rs/Cargo.toml

@@ -14,7 +14,7 @@ readme = "../../README.md"
 documentation = "https://docs.rs/proton-drive-rs"
 
 [dependencies]
-proton-sdk = { path = "../proton-sdk", version = "0.5.1" }
+proton-sdk = { path = "../proton-sdk", version = "0.6.0" }
 tokio.workspace = true
 reqwest.workspace = true
 serde.workspace = true

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

@@ -3426,21 +3426,41 @@ impl ProtonDriveClient {
             _ => generated.clone(),
         };
 
-        // Wrap the share session key with the SRP-salted passphrase, and store the
-        // password encrypted to the owner's address key.
-        let key_salt = generate_key_salt();
-        let salted_passphrase = derive_key_passphrase(full_password.as_bytes(), &key_salt)?;
-        let share_passphrase_key_packet =
-            base64_encode(share_session_key.encrypt_with_password(&salted_passphrase)?);
-        let armored_password = address_key.encrypt(full_password.as_bytes())?;
-
-        // SRP verifier over the link password, against a fresh signed modulus.
+        // The modulus is fetched first so that everything CPU-bound below can go
+        // to the blocking pool in one hop.
         let modulus = self.fetch_srp_modulus().await?;
-        let verifier = generate_verifier(
-            full_password.as_bytes(),
-            &modulus.modulus,
-            DEFAULT_BIT_LENGTH,
-        )?;
+
+        // Wrap the share session key with the SRP-salted passphrase, store the
+        // password encrypted to the owner's address key, and build the SRP
+        // verifier over the link password.
+        //
+        // All of that is pure CPU and none of it is cheap: `derive_key_passphrase`
+        // is a deliberately slow bcrypt, and `generate_verifier` a 2048-bit modular
+        // exponentiation. Inline, they stalled every other task on the reactor
+        // thread for the duration — so, like the rest of the crate's PGP work, they
+        // run on the blocking pool.
+        let key_salt = generate_key_salt();
+        let signed_modulus = modulus.modulus.clone();
+        let password_bytes = full_password.clone();
+        let salt_for_task = key_salt;
+        let (share_passphrase_key_packet, armored_password, verifier) =
+            tokio::task::spawn_blocking(move || {
+                let salted_passphrase =
+                    derive_key_passphrase(password_bytes.as_bytes(), &salt_for_task)?;
+                let key_packet =
+                    base64_encode(share_session_key.encrypt_with_password(&salted_passphrase)?);
+                let armored_password = address_key.encrypt(password_bytes.as_bytes())?;
+                let verifier = generate_verifier(
+                    password_bytes.as_bytes(),
+                    &signed_modulus,
+                    DEFAULT_BIT_LENGTH,
+                )?;
+                Ok::<_, ProtonError>((key_packet, armored_password, verifier))
+            })
+            .await
+            .map_err(|e| {
+                ProtonError::invalid_operation(format!("public link crypto task failed: {e}"))
+            })??;
 
         let flags = if custom_password.map(|p| !p.is_empty()).unwrap_or(false) {
             3 // random + custom password
@@ -3868,12 +3888,23 @@ impl ProtonDriveClient {
             .own_address_keys(&detail.invitation.invitee_email)
             .await
             .ok()?;
-        let passphrase = keys
-            .iter()
-            .find_map(|k| k.decrypt_armored_message(&detail.share.passphrase).ok())?;
-        let share_key = PrivateKey::from_armored(&detail.share.share_key, &passphrase).ok()?;
-        let name_bytes = share_key.decrypt_armored_message(&detail.link.name).ok()?;
-        String::from_utf8(name_bytes).ok()
+        // Unlocking the share key runs its S2K — tens of milliseconds of pure
+        // CPU, and this is called once per invitation in a listing. Off the
+        // reactor, like every other key unlock in the crate
+        // (`crypto::decrypt_link_verified`).
+        let armored_passphrase = detail.share.passphrase.clone();
+        let armored_share_key = detail.share.share_key.clone();
+        let armored_name = detail.link.name.clone();
+        tokio::task::spawn_blocking(move || {
+            let passphrase = keys
+                .iter()
+                .find_map(|k| k.decrypt_armored_message(&armored_passphrase).ok())?;
+            let share_key = PrivateKey::from_armored(&armored_share_key, &passphrase).ok()?;
+            let name_bytes = share_key.decrypt_armored_message(&armored_name).ok()?;
+            String::from_utf8(name_bytes).ok()
+        })
+        .await
+        .ok()?
     }
 
     /// The current user's own address keys for `email`: all its private keys (for

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

@@ -117,14 +117,26 @@ impl Credentials {
             ));
         }
 
-        let proofs = generate_proofs(
-            info.version,
-            self.password.as_bytes(),
-            &decode_base64(&info.url_password_salt, "UrlPasswordSalt")?,
-            &info.modulus,
-            &decode_base64(&info.server_ephemeral, "ServerEphemeral")?,
-            DEFAULT_BIT_LENGTH,
-        )?;
+        // SRP proofs are a bcrypt derivation plus 2048-bit modular
+        // exponentiations — tens of milliseconds of pure CPU. Inline they stalled
+        // the reactor thread; the crate offloads its crypto everywhere else
+        // (`crypto::decrypt_link_verified`) and this is no different.
+        let version = info.version;
+        let password = self.password.clone();
+        let url_password_salt = decode_base64(&info.url_password_salt, "UrlPasswordSalt")?;
+        let modulus = info.modulus.clone();
+        let server_ephemeral = decode_base64(&info.server_ephemeral, "ServerEphemeral")?;
+        let proofs = blocking(move || {
+            Ok(generate_proofs(
+                version,
+                password.as_bytes(),
+                &url_password_salt,
+                &modulus,
+                &server_ephemeral,
+                DEFAULT_BIT_LENGTH,
+            )?)
+        })
+        .await?;
 
         let request = PublicLinkAuthRequest {
             client_proof: BASE64.encode(&proofs.client_proof),
@@ -147,7 +159,7 @@ impl Credentials {
             ));
         }
 
-        let share_key = self.unlock_share_key(&response)?;
+        let share_key = self.unlock_share_key(&response).await?;
 
         // An absent AccessToken means the server accepted an existing Proton
         // session instead of minting an anonymous one. We have no such session
@@ -197,18 +209,36 @@ impl Credentials {
     /// `DriveCrypto.decryptKeyWithSrpPassword`. The salt is a bcrypt salt, not
     /// the SRP one: the same `derive_key_passphrase` the create side used to
     /// wrap this passphrase.
-    fn unlock_share_key(&self, response: &PublicLinkAuthResponse) -> Result<PrivateKey> {
+    /// Runs on the blocking pool: the bcrypt derivation and the share key's S2K
+    /// are each tens of milliseconds of CPU.
+    async fn unlock_share_key(&self, response: &PublicLinkAuthResponse) -> Result<PrivateKey> {
         let salt = decode_base64(&response.share.share_password_salt, "SharePasswordSalt")?;
-        let key_password = derive_key_passphrase(self.password.as_bytes(), &salt)?;
-        let passphrase =
-            decrypt_armored_with_password(&response.share.share_passphrase, &key_password)?;
-        Ok(PrivateKey::from_armored(
-            &response.share.share_key,
-            &passphrase,
-        )?)
+        let password = self.password.clone();
+        let armored_passphrase = response.share.share_passphrase.clone();
+        let armored_share_key = response.share.share_key.clone();
+        blocking(move || {
+            let key_password = derive_key_passphrase(password.as_bytes(), &salt)?;
+            let passphrase = decrypt_armored_with_password(&armored_passphrase, &key_password)?;
+            Ok(PrivateKey::from_armored(&armored_share_key, &passphrase)?)
+        })
+        .await
     }
 }
 
+/// Run CPU-bound crypto off the runtime, mapping a join failure into a
+/// [`ProtonError`]. The public-link handshake is the one path in the crate whose
+/// crypto is not already behind `crypto::blocking` (that helper is about link
+/// decryption and takes an `AccountClient`, which a visitor has no equivalent of).
+async fn blocking<T, F>(work: F) -> Result<T>
+where
+    F: FnOnce() -> Result<T> + Send + 'static,
+    T: Send + 'static,
+{
+    tokio::task::spawn_blocking(work).await.map_err(|e| {
+        ProtonError::invalid_operation(format!("public link crypto task failed: {e}"))
+    })?
+}
+
 /// The client a public-link session is currently issuing requests through,
 /// tagged with the handshake that produced it.
 struct LiveSession {

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

@@ -38,7 +38,7 @@ pub(crate) const MAX_CONCURRENT_BLOCK_DOWNLOADS: usize = 10;
 pub(crate) async fn decrypt_block_blocking(
     content_key: ContentKey,
     ciphertext: Bytes,
-) -> Result<Vec<u8>> {
+) -> Result<Bytes> {
     join_decrypt(tokio::task::spawn_blocking(move || {
         content_key.decrypt_block(&ciphertext)
     }))
@@ -52,7 +52,7 @@ pub(crate) async fn decrypt_block_blocking(
 pub(crate) async fn digest_and_decrypt_block_blocking(
     content_key: ContentKey,
     ciphertext: Bytes,
-) -> Result<(Vec<u8>, Vec<u8>)> {
+) -> Result<(Vec<u8>, Bytes)> {
     let handle = tokio::task::spawn_blocking(move || {
         let digest = Sha256::digest(&ciphertext).to_vec();
         content_key
@@ -249,7 +249,7 @@ impl RevisionReader {
     /// in-flight block slot, and dropping it here rather than at the caller
     /// would let `buffered` pile up decrypted blocks that nothing is accounting
     /// for. The caller drops it once the bytes are consumed.
-    async fn block_plaintext(&self, index: usize) -> Result<(Vec<u8>, OwnedSemaphorePermit)> {
+    async fn block_plaintext(&self, index: usize) -> Result<(Bytes, OwnedSemaphorePermit)> {
         // Held across fetch *and* decrypt: both halves are resident 4 MiB
         // buffers, and the plaintext outlives the ciphertext.
         let permit = self

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

@@ -26,6 +26,7 @@
 use std::sync::Arc;
 
 use async_trait::async_trait;
+use futures::stream::{self, StreamExt, TryStreamExt};
 use proton_sdk::error::{ProtonError, Result};
 use proton_sdk::http::ApiHttpClient;
 use proton_sdk::ids::{LinkId, VolumeId};
@@ -143,6 +144,18 @@ impl RevisionTransport {
     }
 
     /// The revision's metadata and its full, ordered block table.
+    ///
+    /// Pages after the first are fetched several at a time rather than one per
+    /// round trip: a 4 GiB file is 21 pages, and walking them serially costs 21
+    /// RTTs before the first byte of content can be requested. The window is
+    /// speculative — a page's `FromBlockIndex` is derived from its position
+    /// rather than from the previous page's contents, which the contiguity check
+    /// below already requires to hold — so a short page inside a batch simply
+    /// ends the walk and the surplus responses are discarded.
+    ///
+    /// The window ramps `2 → 4 → …→ MAX_PAGE_WINDOW` so that a file only just
+    /// over a page boundary wastes one speculative request rather than five,
+    /// while a genuinely large table still reaches full width after two batches.
     pub(crate) async fn list_blocks(
         &self,
         volume_id: &VolumeId,
@@ -150,29 +163,39 @@ impl RevisionTransport {
         revision_id: &str,
     ) -> Result<(RevisionDto, Vec<BlockDto>)> {
         const PAGE_SIZE: i32 = 50;
+        const MAX_PAGE_WINDOW: usize = 6;
 
-        let mut blocks: Vec<BlockDto> = Vec::new();
-        let mut metadata: Option<RevisionDto> = None;
-        let mut from_index: i32 = 1;
-
-        loop {
+        let page = |page_index: i32| {
+            let from_index = page_index * PAGE_SIZE + 1;
             let path = format!(
                 "v2/volumes/{volume_id}/files/{link_id}/revisions/{revision_id}?FromBlockIndex={from_index}&PageSize={PAGE_SIZE}&NoBlockUrls=0"
             );
-            let response: RevisionResponse = self.get(&path).await?;
-            let mut revision = response.revision;
-            let page = std::mem::take(&mut revision.blocks);
-            let page_len = page.len();
+            async move { self.get::<RevisionResponse>(&path).await }
+        };
 
-            if metadata.is_none() {
-                metadata = Some(revision);
-            }
-            blocks.extend(page);
+        let mut first = page(0).await?.revision;
+        let mut blocks: Vec<BlockDto> = std::mem::take(&mut first.blocks);
+        let mut complete = blocks.len() < PAGE_SIZE as usize;
+        let metadata = first;
 
-            if page_len < PAGE_SIZE as usize {
-                break;
+        let mut next_page: i32 = 1;
+        let mut window: usize = 2;
+        while !complete {
+            let batch: Vec<i32> = (next_page..next_page + window as i32).collect();
+            next_page += window as i32;
+
+            let mut responses = stream::iter(batch.into_iter().map(page)).buffered(window);
+            window = (window * 2).min(MAX_PAGE_WINDOW);
+            while let Some(response) = responses.try_next().await? {
+                if complete {
+                    // A page in this batch already came up short; anything after
+                    // it is past the end of the block table.
+                    continue;
+                }
+                let page_blocks = response.revision.blocks;
+                complete = page_blocks.len() < PAGE_SIZE as usize;
+                blocks.extend(page_blocks);
             }
-            from_index = blocks.iter().map(|b| b.index).max().unwrap_or(from_index) + 1;
         }
 
         blocks.sort_by_key(|b| b.index);
@@ -184,8 +207,6 @@ impl RevisionTransport {
             }
         }
 
-        let metadata = metadata
-            .ok_or_else(|| ProtonError::invalid_operation("revision returned no metadata"))?;
         Ok((metadata, blocks))
     }
 }