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))
     }
 }

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

@@ -8,6 +8,7 @@
 
 use std::io::Cursor;
 
+use bytes::Bytes;
 use pgp::armor::Dearmor;
 use pgp::composed::{MessageBuilder, PlainSessionKey, RawSessionKey};
 use pgp::crypto::aead::{AeadAlgorithm, ChunkSize};
@@ -347,7 +348,7 @@ impl ContentKey {
     /// rejects a message that *starts* with a SEIPD packet, so we parse packets
     /// directly, locate the SEIPD (tolerating a leading PKESK, e.g. a full
     /// round-trip message), and decrypt it with the supplied session key.
-    pub fn decrypt_block(&self, ciphertext: &[u8]) -> Result<Vec<u8>, CryptoError> {
+    pub fn decrypt_block(&self, ciphertext: &[u8]) -> Result<Bytes, CryptoError> {
         let inner = self.decrypt_seipd(ciphertext)?;
 
         // A content block's payload is a single literal-data packet.
@@ -357,7 +358,11 @@ impl ContentKey {
             .ok_or_else(|| CryptoError::Parse("block payload is empty".into()))?
             .map_err(|e| CryptoError::Parse(format!("block payload packet: {e}")))?;
         match literal {
-            Packet::LiteralData(data) => Ok(data.data().to_vec()),
+            // `into_bytes`, not `data().to_vec()`: the literal packet already
+            // owns its payload as `Bytes`, and a block is 4 MiB — copying it out
+            // doubled the peak footprint of every block delivered, times the ten
+            // a download has in flight.
+            Packet::LiteralData(data) => Ok(data.into_bytes()),
             other => Err(CryptoError::Parse(format!(
                 "block payload is not literal data: {:?}",
                 other.tag()

crates/proton-sdk/src/http.rs

@@ -68,7 +68,19 @@ struct Inner {
     base_url: String,
     config: ProtonClientConfiguration,
     session_id: SessionId,
-    tokens: Mutex<Tokens>,
+    /// The current tokens, read on every request and replaced by a refresh.
+    ///
+    /// A plain `std::sync::RwLock` over an `Arc` snapshot, deliberately *not*
+    /// the tokio mutex it used to be: reading the access token is what every
+    /// ordinary request does, and holding an async mutex for it meant queueing
+    /// behind whoever was refreshing — a refresh that stalls on the network (up
+    /// to the retry policy's whole budget) stalled every other API call with it.
+    /// The lock here is only ever held for a clone of an `Arc`, never across an
+    /// await; the mutual exclusion a refresh needs lives in [`Inner::refresh`].
+    tokens: std::sync::RwLock<Arc<Tokens>>,
+    /// Serializes token refreshes. Held across the refresh network call, so it
+    /// must be the async mutex; readers of `tokens` never touch it.
+    refresh: Mutex<()>,
     /// Telemetry sink for per-request events. Interior-mutable because the
     /// client is already shared (cloned into the Drive client) by the time a
     /// caller attaches a sink via [`ApiHttpClient::set_telemetry`]. Defaults to
@@ -103,7 +115,8 @@ impl ApiHttpClient {
                 base_url,
                 config,
                 session_id,
-                tokens: Mutex::new(tokens),
+                tokens: std::sync::RwLock::new(Arc::new(tokens)),
+                refresh: Mutex::new(()),
                 telemetry: std::sync::Mutex::new(NoopTelemetry::shared()),
                 on_tokens_refreshed: std::sync::Mutex::new(None),
             }),
@@ -144,7 +157,16 @@ impl ApiHttpClient {
 
     /// Snapshot the current tokens (e.g. to persist for a later `resume`).
     pub async fn current_tokens(&self) -> Tokens {
-        self.inner.tokens.lock().await.clone()
+        (*self.tokens()).clone()
+    }
+
+    /// The current tokens, as a cheap `Arc` snapshot. Never held across an await.
+    fn tokens(&self) -> Arc<Tokens> {
+        self.inner
+            .tokens
+            .read()
+            .expect("tokens rwlock poisoned")
+            .clone()
     }
 
     /// Attach a telemetry sink to receive a per-request
@@ -313,7 +335,7 @@ impl ApiHttpClient {
         let metadata = serde_json::to_vec(metadata)?;
         let mut timer = self.telemetry().start("http_request");
         timer.attr("method", "POST");
-        let rejected_token = self.inner.tokens.lock().await.access_token.clone();
+        let rejected_token = self.tokens().access_token.clone();
         let response = self
             .send_multipart_with_token(path, &metadata, binary_parts, &rejected_token)
             .await?;
@@ -371,7 +393,7 @@ impl ApiHttpClient {
         let mut timer = self.telemetry().start("http_request");
         timer.attr("method", method.as_str());
 
-        let access_token = self.inner.tokens.lock().await.access_token.clone();
+        let access_token = self.tokens().access_token.clone();
 
         // An early `?` here records the op as a failure (OpTimer defaults to it).
         let response = self
@@ -501,14 +523,19 @@ impl ApiHttpClient {
     /// in-memory access token already differs from the rejected one, another
     /// task refreshed first and we reuse its result.
     async fn refresh_access_token(&self, rejected_access_token: &str) -> Result<String> {
-        let mut guard = self.inner.tokens.lock().await;
+        // The refresh lock, not the token lock: ordinary requests read the
+        // tokens without ever waiting on this.
+        let _guard = self.inner.refresh.lock().await;
 
-        if guard.access_token != rejected_access_token {
-            return Ok(guard.access_token.clone());
+        // Re-read under the refresh lock — whoever held it before us may have
+        // already replaced the token we were handed a 401 for.
+        let current = self.tokens();
+        if current.access_token != rejected_access_token {
+            return Ok(current.access_token.clone());
         }
 
-        let refreshed = self.request_refresh(&guard.refresh_token).await?;
-        *guard = refreshed.clone();
+        let refreshed = self.request_refresh(&current.refresh_token).await?;
+        *self.inner.tokens.write().expect("tokens rwlock poisoned") = Arc::new(refreshed.clone());
 
         // Notify callback
         if let Some(ref cb) = *self
@@ -1027,4 +1054,98 @@ mod tests {
             "an ordinary session keeps refreshing across a route change"
         );
     }
+
+    /// An ordinary request reads the access token without queueing behind an
+    /// in-flight refresh.
+    ///
+    /// This is the property the token lock exists to have: the tokens were once
+    /// behind a tokio `Mutex` that the refresh path held across its network
+    /// call, so one slow `auth/v4/refresh` — up to the whole retry budget —
+    /// stalled every other API call on the session. The refresh here is made
+    /// deliberately slow; the concurrent GET has to finish long before it.
+    #[tokio::test]
+    async fn an_ordinary_request_does_not_wait_for_an_in_flight_refresh() {
+        use std::sync::atomic::{AtomicUsize, Ordering};
+        use std::time::Instant;
+        use tokio::io::{AsyncReadExt, AsyncWriteExt};
+
+        const REFRESH_DELAY: Duration = Duration::from_millis(1500);
+
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let addr = listener.local_addr().unwrap();
+
+        // `/needs-refresh` 401s once, which sends the caller to
+        // `auth/v4/refresh` — and that reply is held back. Everything else is
+        // answered immediately.
+        let unauthorized = Arc::new(AtomicUsize::new(0));
+        let seen = unauthorized.clone();
+        let server = tokio::spawn(async move {
+            loop {
+                let (mut sock, _) = listener.accept().await.unwrap();
+                let seen = seen.clone();
+                tokio::spawn(async move {
+                    let mut buf = [0u8; 2048];
+                    let n = sock.read(&mut buf).await.unwrap();
+                    let request = String::from_utf8_lossy(&buf[..n]).into_owned();
+
+                    let (status, body) = if request.contains("auth/v4/refresh") {
+                        tokio::time::sleep(REFRESH_DELAY).await;
+                        (
+                            "200 OK",
+                            r#"{"Code":1000,"AccessToken":"fresh","RefreshToken":"fresh-refresh","UID":"test-session"}"#.to_string(),
+                        )
+                    } else if request.contains("needs-refresh")
+                        && seen.fetch_add(1, Ordering::SeqCst) == 0
+                    {
+                        // Only the first attempt 401s; the retry after the
+                        // refresh succeeds, as it would against the real API.
+                        ("401 Unauthorized", r#"{"Code":401}"#.to_string())
+                    } else {
+                        ("200 OK", r#"{"Code":1000}"#.to_string())
+                    };
+
+                    let head = format!(
+                        "HTTP/1.1 {status}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
+                        body.len()
+                    );
+                    sock.write_all(head.as_bytes()).await.unwrap();
+                    sock.write_all(body.as_bytes()).await.unwrap();
+                    sock.flush().await.unwrap();
+                });
+            }
+        });
+
+        let config = ProtonClientConfiguration::new("test@1.0")
+            .with_base_url(format!("http://{addr}/"))
+            .with_retry_policy(RetryPolicy::disabled());
+        let client = ApiHttpClient::new(
+            config,
+            SessionId::from("test-session"),
+            Tokens {
+                access_token: "access".into(),
+                refresh_token: "refresh".into(),
+            },
+        )
+        .unwrap();
+
+        let refreshing = tokio::spawn({
+            let client = client.clone();
+            async move { client.get::<ApiResponse>("needs-refresh").await }
+        });
+
+        // Give the refresh time to be issued and be sitting on the network.
+        tokio::time::sleep(Duration::from_millis(200)).await;
+
+        let started = Instant::now();
+        let _: ApiResponse = client.get("ordinary").await.unwrap();
+        let waited = started.elapsed();
+
+        assert!(
+            waited < REFRESH_DELAY / 2,
+            "an ordinary request waited {waited:?} on a refresh still in flight"
+        );
+
+        refreshing.await.unwrap().unwrap();
+        server.abort();
+    }
 }