From 0cf760362ae6a9dc06ca5f570d2cdfc03273a644 Mon Sep 17 00:00:00 2001 From: Nils Pukropp Date: Sun, 16 Aug 2026 18:53:29 +0200 Subject: [PATCH] fix(sdk) token mutex no longer spans the refresh --- Cargo.lock | 44 ++++++- Cargo.toml | 7 +- README.md | 51 ++++++-- crates/proton-drive-rs/Cargo.toml | 2 +- crates/proton-drive-rs/src/client.rs | 71 ++++++++--- crates/proton-drive-rs/src/public_link.rs | 64 +++++++--- crates/proton-drive-rs/src/revision.rs | 6 +- crates/proton-drive-rs/src/transport.rs | 57 ++++++--- crates/proton-sdk/src/crypto/content.rs | 9 +- crates/proton-sdk/src/http.rs | 141 ++++++++++++++++++++-- 10 files changed, 363 insertions(+), 89 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ed1d1d0..235fbc1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1160,6 +1160,25 @@ dependencies = [ "subtle", ] +[[package]] +name = "h2" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6cb093c84e8bd9b188d4c4a8cb6579fc016968d14c99882163cd3ff402a4f155" +dependencies = [ + "atomic-waker", + "bytes", + "fnv", + "futures-core", + "futures-sink", + "http", + "indexmap", + "slab", + "tokio", + "tokio-util", + "tracing", +] + [[package]] name = "hashbrown" version = "0.15.5" @@ -1171,6 +1190,12 @@ dependencies = [ "foldhash", ] +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + [[package]] name = "heck" version = "0.5.0" @@ -1277,6 +1302,7 @@ dependencies = [ "bytes", "futures-channel", "futures-core", + "h2", "http", "http-body", "httparse", @@ -1467,6 +1493,16 @@ dependencies = [ "icu_properties", ] +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + [[package]] name = "inout" version = "0.1.4" @@ -1630,7 +1666,7 @@ version = "0.12.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "234cf4f4a04dc1f57e24b96cc0cd600cf2af460d4161ac5ecdd0af8e1f3b2a38" dependencies = [ - "hashbrown", + "hashbrown 0.15.5", ] [[package]] @@ -2036,7 +2072,7 @@ dependencies = [ [[package]] name = "proton-drive-rs" -version = "0.5.1" +version = "0.6.0" dependencies = [ "async-trait", "base64 0.22.1", @@ -2061,7 +2097,7 @@ dependencies = [ [[package]] name = "proton-sdk" -version = "0.5.1" +version = "0.6.0" dependencies = [ "aes-gcm 0.11.0", "async-trait", @@ -2240,6 +2276,7 @@ dependencies = [ "bytes", "futures-core", "futures-util", + "h2", "http", "http-body", "http-body-util", @@ -2879,6 +2916,7 @@ dependencies = [ "bytes", "futures-core", "futures-sink", + "libc", "pin-project-lite", "tokio", ] diff --git a/Cargo.toml b/Cargo.toml index 404e403..c7afc2a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,7 +3,7 @@ resolver = "2" members = ["crates/proton-sdk", "crates/proton-drive-rs"] [workspace.package] -version = "0.5.1" +version = "0.6.0" edition = "2024" license = "MIT" repository = "https://github.com/narrrl/proton-sdk-rs" @@ -24,7 +24,10 @@ opt-level = 2 [workspace.dependencies] tokio = { version = "1", features = ["rt-multi-thread", "macros", "sync", "fs", "io-util", "time"] } -reqwest = { version = "0.13", default-features = false, features = ["json", "rustls", "webpki-roots", "stream", "multipart", "gzip"] } +# `http2` is in reqwest's default set, which `default-features = false` drops: +# without it every concurrent block download pays its own TCP+TLS handshake +# instead of multiplexing over one connection. +reqwest = { version = "0.13", default-features = false, features = ["json", "rustls", "webpki-roots", "stream", "multipart", "gzip", "http2"] } serde = { version = "1", features = ["derive"] } serde_json = "1" serde_repr = "0.1" diff --git a/README.md b/README.md index 08b2e3e..da042cd 100644 --- a/README.md +++ b/README.md @@ -239,6 +239,14 @@ Moving a file or folder within the same volume is an offline cryptographic opera `proton-sdk-rs` implements **Milestone 2 (Authenticated Read/Write)** and advanced desktop features. Parity with the official C# SDK (pinned at the upstream commit documented in [UPSTREAM_SYNC.md](file:///home/narl/dev/private/proton-sdk-rs/UPSTREAM_SYNC.md)) is actively maintained. +The C# SDK is the parity *reference*, not the ceiling. Several subsystems have no C# public API and were +ported from the official TypeScript SDK instead, and a few surfaces exist in neither upstream โ€” they were +built here because a real files-on-demand client (`proton-drive-linux`) needs them. Divergences are logged +per-sync in [UPSTREAM_SYNC.md](file:///home/narl/dev/private/proton-sdk-rs/UPSTREAM_SYNC.md). + +**Legend** โ€” โœ… ported from the C# SDK ยท ๐ŸŸฆ no C# public API, ported from the TypeScript SDK ยท โญ exists in +neither upstream SDK, original to `proton-sdk-rs` ยท โŒ not implemented. + ### Feature Matrix | Module | Feature | Status | Notes / Upstream Parity | @@ -247,40 +255,57 @@ Moving a file or folder within the same volume is an offline cryptographic opera | | TOTP 2FA | โœ… | Support for applying second-factor TOTP codes. | | | Token Refresh | โœ… | Scope refresh and transparent 401 token refresh. | | | Resume Session | โœ… | Restore clients instantly via serialized tokens. | +| **Account** | Addresses & Keys | โœ… | Address enumeration, address-key unlocking, public-key lookup by email (`AccountClient`). | +| | Key-salt Injection | โญ | `with_key_salts` lets a client hand in cached salts and chain an entity repository onto it, so a warm start skips the salt round-trip. | +| | Quota | โœ… | `quota()` โ€” used/total space with `available()` / `used_fraction()` helpers. | | **HTTP Client** | Retries & Jitter | โœ… | Exponential backoff for 408, 429, and 5xx, respecting `Retry-After`. | | | Telemetry | โœ… | Pluggable request tracking (`ITelemetry` analogue). | | **Drive / Volume** | Volume Resolution | โœ… | Auto-resolves volume lists; auto-creates the `my-files` root volume on first login. | | | Node Listing | โœ… | `get_node`, `enumerate_folder_children`, `enumerate_trash` with key resolution. | +| | Metadata-only Listing | โญ | `enumerate_nodes_light` skips content-key resolution for listings that only need names and sizes. | +| | Path Helpers | โญ | `get_node_by_path`, `get_node_path`, `get_node_hierarchy`, `create_folder_path`, `get_available_name` โ€” a filesystem front-end wants paths, not UID walks. | | | File Operations | โœ… | Rename, trash, restore, delete, empty trash โ€” batch calls report one outcome per node (`trash_nodes`), or stream them as each batch lands (`trash_nodes_streaming`). | -| | Move Operations | โœ… | Same-volume move. Batch moves are chunked automatically. | -| | Cross-volume Move | โŒ | Not supported (throwing `NotImplementedException` in the C# public API too). | +| | Move Operations | โœ… | Same-volume move. Batch moves are chunked automatically, one outcome per node. | +| | Cross-volume Move | โŒ | Not supported here, and not by `proton-drive-linux` either โ€” the C# public API throws `NotImplementedException`, and a cross-volume move means re-uploading content under the destination volume's keys. | +| **Events / Sync** | Event Enumeration | โœ… | `enumerate_events` + `latest_event_id` per volume scope; `invalidate_caches_for_event` keeps the entity cache coherent. | +| | Event Manager | ๐ŸŸฆ | Background poll loop (`EventManager`) with a broadcast channel and a foreground/background interval switch, modelled on the TS `eventScheduler`. | | **Uploads** | Block Uploading | โœ… | Encrypts and uploads chunks (4 MiB default) to block storage. | | | Streaming Upload | โœ… | Streams uploads from any type implementing `std::io::Read`. | -| | Revision Control | โœ… | Create drafts, upload blocks, and seal new file revisions. | +| | Revision Control | โœ… | Create drafts, upload blocks, and seal new file revisions; `upload_file_replacing_draft_from` reclaims an abandoned draft. | +| | Atomic Small Uploads | โœ… | Single-request upload for small buffered files, opt-in via `with_small_file_upload` (no remote feature-flag provider in this crate). | +| | Upload Concurrency | โญ | `with_max_inflight_blocks` bounds in-flight block uploads/downloads per client. | | | AEAD Block Support | โœ… | SEIPDv2 / AES-256-GCM AEAD encryption alongside legacy SEIPDv1 (AES-256-CFB). | | | Inline Thumbnails | โœ… | Generates and encrypts inline-signed thumbnails. | | **Downloads** | Block Downloading | โœ… | Streams block retrieval, decryption, and signature checks. | +| | Seekable Range Reads | โญ | `open_revision` โ†’ `RevisionReader::read_at` and `download_range` serve arbitrary byte ranges by fetching only the covering blocks โ€” what makes FUSE files-on-demand possible. Neither upstream SDK has it. | +| | Thumbnail Retrieval | โœ… | `download_thumbnail`, plus the plural `enumerate_thumbnails` batch API. | | | Content Integrity | โœ… | Verification of file manifests and block-level SHA-256 digests. | | | Signature Verification | โœ… | Detached signature verification (non-fatal, reports status). | +| **Revisions** | Revision History | ๐ŸŸฆ | `enumerate_revisions` / `get_revision`, plus `download_revision` and `download_revision_to` for any past revision. | +| | Restore & Delete | ๐ŸŸฆ | `restore_revision`, `delete_revision`. | | **Devices** | Device Registration | โœ… | Create and register desktop/sync clients (`create_device`). | | | Device Listing | โœ… | List registered devices and their sync root folders (`enumerate_devices`). | | | Device Operations | โœ… | Rename and delete registered devices. | -| **Sharing & Links** | Member Management | โœ… | Invite users (`invite_users`), list members, update roles, and remove members. | -| | Public Share Links | โœ… | Create public shareable links (`create_public_link`) with password protection and expiry. | +| **Sharing & Links** | Share Creation | ๐ŸŸฆ | `share_node` bootstraps a share on a node (no C# public API for it). | +| | Member Management | โœ… | Invite users (`invite_users`), list members, update roles, and remove members. | +| | Public Share Links | ๐ŸŸฆ | Create public shareable links (`create_public_link`) with password protection and expiry; `get_public_link`, `remove_public_link`. | | | External Invites | โœ… | Invite non-Proton users (`invite_external_users`) and track registration status. | -| | Bookmarks | โœ… | Save and list public link bookmarks (`create_bookmark`, `list_bookmarks`). | -| | Incoming Invites | โœ… | List, accept, or reject invitations shared *with* the current user. | -| | Public Link (consume) | โœ… | Open someone else's link with no Proton account (`ProtonDrivePublicLinkClient`): SRP handshake, custom passwords, browse and download the shared subtree. | -| | Public Link streaming | โœ… | Seekable range reads over a shared file (`open_revision` โ†’ `RevisionReader::read_at`), plus thumbnails and in-place session renewal โ€” a visitor can stream a large file without downloading it. | +| | Bookmarks | โœ… | Save, list, and delete public link bookmarks (`create_bookmark`, `list_bookmarks`, `delete_bookmark`). | +| | Incoming Invites | โœ… | List, accept, or reject invitations shared *with* the current user; `leave_shared_node`. | +| | Shared Listings | โœ… | `enumerate_shared_by_me_node_uids`, `enumerate_shared_with_me` (drive and photos volumes). | +| | Public Link (consume) | ๐ŸŸฆ | Open someone else's link with no Proton account (`ProtonDrivePublicLinkClient`): SRP handshake, custom passwords, browse and download the shared subtree. | +| | Public Link streaming | โญ | Seekable range reads over a shared file (`open_revision` โ†’ `RevisionReader::read_at`), plus thumbnails, batch enumeration and in-place session renewal โ€” a visitor streams a large file without downloading it. No TS or C# counterpart. | | **Photos** | Photos Timeline | โœ… | `ProtonPhotosClient` maps photostream, timeline enumeration, and photo downloads. | | | Photo Uploads | โœ… | Uploading photos with `PhotoUploadMetadata` (capture time, tags, grouping). | | | Albums (read) | โœ… | List albums (`enumerate_album_node_uids`) and their photos (`enumerate_album`); album nodes carry `Node::album`. | | | Photo Tags | โœ… | Add/remove classification tags (`update_photos`, or `update_photos_streaming` for per-photo outcomes as they finish); `Favorite` on photos in our own timeline. | -| | Shared Photos | โœ… | Photos/albums shared *with* us (`enumerate_shared_with_me_node_uids`, `enumerate_shared_with_me_album_uids`). | -| | Favorite Shared Photos | โŒ | Needs the photo re-encrypted for our timeline root; not ported. | -| | Album Writes | โŒ | Creating albums and adding/removing photos is not ported. | +| | Shared Photos | โœ… | Photos/albums shared *with* us (`enumerate_shared_with_me_node_uids`, `enumerate_shared_with_me_album_uids`) and by us (`enumerate_shared_node_uids`). | +| | Duplicate Detection | โœ… | `find_duplicates` matches candidates on name/content HMAC before re-uploading. | +| | Favorite Shared Photos | โŒ | Needs the photo re-encrypted for our timeline root; not ported (upstream `03b1cb7f`). | +| | Album Writes | โŒ | Creating albums and adding/removing photos โ€” no C# public API upstream either. | | | Photos Volume Create| โŒ | Volume creation is not yet ported. | -| **Caching** | Pluggable Cache | โœ… | Pluggable entity cache (`with_entity_cache`, or `with_entity_repository` to chain one onto a `with_key_salts` client); keys/secrets remain strictly in memory. | +| **Caching** | Pluggable Cache | โœ… | Pluggable entity cache (`with_entity_cache`, or `with_entity_repository` to chain one onto a `with_key_salts` client); keys/secrets remain strictly in memory. Retained deliberately โ€” upstream C# removed its entity cache. | +| **Validation** | Node Names | โœ… | Empty or >255-char names are rejected client-side before create/rename (counted in `char`s; C# counts UTF-16 units). | --- diff --git a/crates/proton-drive-rs/Cargo.toml b/crates/proton-drive-rs/Cargo.toml index 3cc922d..5758b37 100644 --- a/crates/proton-drive-rs/Cargo.toml +++ b/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 diff --git a/crates/proton-drive-rs/src/client.rs b/crates/proton-drive-rs/src/client.rs index 032c174..7c4e8ea 100644 --- a/crates/proton-drive-rs/src/client.rs +++ b/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 diff --git a/crates/proton-drive-rs/src/public_link.rs b/crates/proton-drive-rs/src/public_link.rs index 9569b0c..b995641 100644 --- a/crates/proton-drive-rs/src/public_link.rs +++ b/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 { + /// 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 { 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(work: F) -> Result +where + F: FnOnce() -> Result + 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 { diff --git a/crates/proton-drive-rs/src/revision.rs b/crates/proton-drive-rs/src/revision.rs index f799fc5..6f8a54f 100644 --- a/crates/proton-drive-rs/src/revision.rs +++ b/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> { +) -> Result { 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, Vec)> { +) -> Result<(Vec, 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, 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 diff --git a/crates/proton-drive-rs/src/transport.rs b/crates/proton-drive-rs/src/transport.rs index dd293d8..91dc6a8 100644 --- a/crates/proton-drive-rs/src/transport.rs +++ b/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)> { const PAGE_SIZE: i32 = 50; + const MAX_PAGE_WINDOW: usize = 6; - let mut blocks: Vec = Vec::new(); - let mut metadata: Option = 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::(&path).await } + }; - if metadata.is_none() { - metadata = Some(revision); - } - blocks.extend(page); + let mut first = page(0).await?.revision; + let mut blocks: Vec = 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 = (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)) } } diff --git a/crates/proton-sdk/src/crypto/content.rs b/crates/proton-sdk/src/crypto/content.rs index bd182b3..107bb42 100644 --- a/crates/proton-sdk/src/crypto/content.rs +++ b/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, CryptoError> { + pub fn decrypt_block(&self, ciphertext: &[u8]) -> Result { 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() diff --git a/crates/proton-sdk/src/http.rs b/crates/proton-sdk/src/http.rs index 7607ad1..6d206ca 100644 --- a/crates/proton-sdk/src/http.rs +++ b/crates/proton-sdk/src/http.rs @@ -68,7 +68,19 @@ struct Inner { base_url: String, config: ProtonClientConfiguration, session_id: SessionId, - tokens: Mutex, + /// 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>, + /// 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 { + 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 { - 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(¤t.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::("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(); + } }