diff --git a/Cargo.lock b/Cargo.lock index 82d1363..f4ff7de 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -110,6 +110,18 @@ dependencies = [ "zeroize", ] +[[package]] +name = "async-compression" +version = "0.4.42" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e79b3f8a79cccc2898f31920fc69f304859b3bd567490f75ebf51ae1c792a9ac" +dependencies = [ + "compression-codecs", + "compression-core", + "pin-project-lite", + "tokio", +] + [[package]] name = "async-trait" version = "0.1.89" @@ -456,6 +468,23 @@ dependencies = [ "memchr", ] +[[package]] +name = "compression-codecs" +version = "0.4.38" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce2548391e9c1929c21bf6aa2680af86fe4c1b33e6cea9ac1cfeec0bd11218cf" +dependencies = [ + "compression-core", + "flate2", + "memchr", +] + +[[package]] +name = "compression-core" +version = "0.4.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc14f565cf027a105f7a44ccf9e5b424348421a1d8952a8fc9d499d313107789" + [[package]] name = "const-oid" version = "0.9.6" @@ -523,6 +552,15 @@ version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fd121741cf3eb82c08dd3023eb55bf2665e5f60ec20f89760cf836ae4562e6a0" +[[package]] +name = "crc32fast" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + [[package]] name = "crypto-bigint" version = "0.5.5" @@ -909,6 +947,7 @@ version = "1.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" dependencies = [ + "crc32fast", "miniz_oxide", "zlib-rs", ] @@ -1959,7 +1998,7 @@ dependencies = [ [[package]] name = "proton-drive-rs" -version = "0.1.7" +version = "0.1.8" dependencies = [ "async-trait", "base64", @@ -1983,7 +2022,7 @@ dependencies = [ [[package]] name = "proton-sdk" -version = "0.1.7" +version = "0.1.8" dependencies = [ "aes-gcm 0.11.0", "async-trait", @@ -2839,12 +2878,17 @@ version = "0.6.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4cfcf7e2740e6fc6d4d688b4ef00650406bb94adf4731e43c096c3a19fe40840" dependencies = [ + "async-compression", "bitflags", "bytes", + "futures-core", "futures-util", "http", "http-body", + "http-body-util", "pin-project-lite", + "tokio", + "tokio-util", "tower", "tower-layer", "tower-service", diff --git a/Cargo.toml b/Cargo.toml index a6d2ed4..6c8f0c2 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.1.7" +version = "0.1.8" edition = "2024" license = "MIT" repository = "https://github.com/narrrl/proton-sdk-rs" @@ -15,7 +15,7 @@ categories = ["api-bindings", "cryptography", "filesystem"] [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"] } +reqwest = { version = "0.13", default-features = false, features = ["json", "rustls", "webpki-roots", "stream", "multipart", "gzip"] } serde = { version = "1", features = ["derive"] } serde_json = "1" serde_repr = "0.1" diff --git a/crates/proton-drive-rs/src/client.rs b/crates/proton-drive-rs/src/client.rs index fe179c3..79758f9 100644 --- a/crates/proton-drive-rs/src/client.rs +++ b/crates/proton-drive-rs/src/client.rs @@ -6,8 +6,9 @@ use std::sync::Arc; use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64; +use futures::stream::{self, StreamExt, TryStreamExt}; use sha2::{Digest, Sha256}; -use tokio::sync::Mutex; +use tokio::sync::{Mutex, Semaphore}; use proton_sdk::account::{AccountClient, KeySalt}; use proton_sdk::cache::{CacheRepository, InMemoryCacheRepository}; @@ -63,6 +64,9 @@ use crate::dtos::{ use crate::events::{DriveEvent, DriveEventScopeId}; use crate::node::{FileThumbnail, Node, NodeKind, RevisionState, Thumbnail, ThumbnailType}; use crate::photos::{PhotoUploadMetadata, PhotosTimelineItem}; +use crate::revision::{ + MAX_CONCURRENT_BLOCK_DOWNLOADS, RevisionReader, digest_and_decrypt_block_blocking, +}; use crate::sharing::{ Bookmark, ExternalInvitation, ExternalInvitationState, IncomingInvitation, MemberRole, PublicLink, ShareInvitation, ShareMember, @@ -75,6 +79,19 @@ const DEFAULT_BLOCK_SIZE: usize = 1 << 22; /// `NodeOperations.MaximumBatchCount`). const MAX_BATCH_COUNT: usize = 150; +/// Content blocks a client keeps in memory at once, across every download it is +/// running. +/// +/// [`MAX_CONCURRENT_BLOCK_DOWNLOADS`] bounds one file; nothing bounded the host, +/// so N concurrent downloads cost N times that — a mount pinning a handful of +/// large files could hold hundreds of MiB of block buffers. This is the missing +/// global ceiling (TypeScript caps concurrent *files* instead, in +/// `internal/download/queue.ts`; capping blocks bounds the memory directly). +/// +/// Sized a little above one file's concurrency so a lone download still +/// saturates its pipeline and a second one is not starved. +const DEFAULT_MAX_INFLIGHT_BLOCKS: usize = 12; + /// Trashed links requested per page when enumerating the trash (C# /// `VolumeOperations.TrashPageSize`). const TRASH_PAGE_SIZE: usize = 500; @@ -125,6 +142,31 @@ pub struct ProtonDriveClient { /// [`NoopTelemetry`]; supply one via /// [`ProtonDriveClient::with_telemetry`]. telemetry: Arc, + /// Global cap on content blocks resident in memory across *all* downloads + /// this client is running. See + /// [`with_max_inflight_blocks`](ProtonDriveClient::with_max_inflight_blocks). + /// Shared across clones, which is the point: the bound is per host, not per + /// download. + block_slots: Arc, + /// Serializes the My Files and Photos bootstraps against themselves. + /// + /// [`ensure_my_files`](ProtonDriveClient::ensure_my_files) and + /// [`ensure_photos`](ProtonDriveClient::ensure_photos) are check-then-act: + /// they read `cache`, drop its guard, fetch and decrypt a share over the + /// network, then re-acquire and store. Concurrent first calls would each run + /// the whole bootstrap. Holding this across the network is safe precisely + /// because it guards nothing else — `cache` stays free for every other + /// caller — and the loser re-checks and finds the work already done. + bootstrap: Arc, +} + +/// The one-time-per-account setup that [`ProtonDriveClient`] single-flights. +/// Separate locks: a Photos bootstrap has no reason to wait behind a My Files +/// one, and neither ever waits on the other. +#[derive(Default)] +struct Bootstrap { + my_files: Mutex<()>, + photos: Mutex<()>, } #[derive(Default)] @@ -142,6 +184,10 @@ struct DriveCache { photos_volume_exists: Option, /// Decrypted node key per folder, used as the parent key for its children. folder_keys: HashMap, + /// The membership share behind each node another user shares with us, keyed + /// by the shared node itself (which is that share's root link). `None` until + /// [`ProtonDriveClient::shared_with_me_shares`] pages `v2/sharedwithme`. + shared_with_me_shares: Option>, } /// An open revision draft ready to receive content blocks: the target revision @@ -208,6 +254,8 @@ impl ProtonDriveClient { cache: Arc::new(Mutex::new(DriveCache::default())), entities: DriveEntityCache::new(entity_repository), telemetry: NoopTelemetry::shared(), + block_slots: Arc::new(Semaphore::new(DEFAULT_MAX_INFLIGHT_BLOCKS)), + bootstrap: Arc::new(Bootstrap::default()), } } @@ -231,10 +279,19 @@ impl ProtonDriveClient { cache: Arc::new(Mutex::new(DriveCache::default())), entities: DriveEntityCache::new(InMemoryCacheRepository::shared()), telemetry: NoopTelemetry::shared(), + block_slots: Arc::new(Semaphore::new(DEFAULT_MAX_INFLIGHT_BLOCKS)), + bootstrap: Arc::new(Bootstrap::default()), } } /// The account client backing this Drive client's key chain. + /// The underlying HTTP client, for sibling modules that issue their own + /// requests against the same session (e.g. [`RevisionReader`] fetching + /// block bodies from storage). + pub(crate) fn http(&self) -> &ApiHttpClient { + &self.http + } + pub fn account(&self) -> &AccountClient { &self.account } @@ -263,6 +320,28 @@ impl ProtonDriveClient { self } + /// Cap the content blocks this client holds in memory at once, across every + /// download it is running. Defaults to [`DEFAULT_MAX_INFLIGHT_BLOCKS`]. + /// + /// A permit is held from the moment a block's fetch starts until its + /// plaintext is consumed, so `blocks * BLOCK_SIZE` is roughly the ceiling on + /// resident block memory — the knob to turn down on a memory-constrained + /// host, at the cost of download throughput. Must be non-zero; a zero cap + /// would deadlock every download, so it is clamped to 1. + /// + /// Applies to this client and every clone made from it *afterwards*; clones + /// already taken keep the previous cap (they share the previous semaphore). + pub fn with_max_inflight_blocks(mut self, blocks: usize) -> Self { + self.block_slots = Arc::new(Semaphore::new(blocks.max(1))); + self + } + + /// The global in-flight block permits, for [`RevisionReader`] and the + /// whole-file download path to acquire against. + pub(crate) fn block_slots(&self) -> Arc { + self.block_slots.clone() + } + /// Resolve (and cache) the user's "My Files" root folder. pub async fn get_my_files_folder(&self) -> Result { self.ensure_my_files().await?; @@ -610,43 +689,40 @@ impl ProtonDriveClient { } } - for block in &blocks { - let ciphertext = self - .http - .get_storage_blob(&block.bare_url, &block.token) - .await?; - let digest = Sha256::digest(&ciphertext); - let plaintext = content_key.decrypt_block(&ciphertext)?; - output - .write_all(&plaintext) - .map_err(|e| ProtonError::invalid_operation(format!("write block: {e}")))?; - manifest.extend_from_slice(&digest); - } + self.write_content_blocks(&blocks, &content_key, &mut manifest, output) + .await?; verify_manifest(&self.account, &revision, &node_key, &manifest).await; timer.success(); Ok(()) } - /// Download and decrypt only the plaintext byte range `[offset, offset + length)` - /// of a file's active revision. + /// Open a seekable reader on a file's active revision. /// /// Each content block decrypts independently under the revision's content - /// key ([`ContentKey::decrypt_block`]), so an on-demand reader can fetch - /// just the blocks that overlap the requested range instead of the whole - /// file — the basis for a FUSE/placeholder mount that hydrates on access. + /// key ([`ContentKey::decrypt_block`]), so a reader can fetch just the + /// blocks that overlap a requested range instead of the whole file — the + /// basis for a FUSE/placeholder mount that hydrates on access. + /// + /// This resolves everything a read needs — link details, the ancestor chain, + /// the node key (an S2K unlock), the content key, the block table and the + /// per-block plaintext sizes — **once**. Subsequent + /// [`RevisionReader::read_at`] calls cost only the block bodies they + /// overlap. A caller that reads a file more than once (any on-demand mount: + /// a `read(2)` is far smaller than a 4 MiB block) should hold the reader for + /// as long as the file is open rather than calling + /// [`download_range`](Self::download_range) per read. /// /// Block plaintext sizes come from the revision's extended attributes /// (`Common.BlockSizes`); absent that, blocks are assumed to be /// [`DEFAULT_BLOCK_SIZE`] with a possibly-shorter final block inferred from - /// the recorded total size. The range is clamped to the file's length, so a - /// read at or past EOF yields fewer bytes (or none). + /// the recorded total size. /// - /// Unlike [`download_file_to`](Self::download_file_to), a partial read - /// cannot recompute the full content manifest, so manifest-signature - /// verification is skipped. - pub async fn download_range(&self, uid: &NodeUid, offset: u64, length: u64) -> Result> { - let mut timer = self.telemetry.start("download_range"); + /// The reader is pinned to the revision that is active now; it does not + /// follow later revisions of the same file. + pub async fn open_revision(&self, uid: &NodeUid) -> Result { + let mut timer = self.telemetry.start("open_revision"); + let details = self .get_link_details(&uid.volume_id, std::slice::from_ref(&uid.link_id)) .await?; @@ -684,40 +760,38 @@ impl ProtonDriveClient { let block_sizes = self .resolve_block_sizes(&node_key, &revision, blocks.len()) .await; - let file_size: u64 = block_sizes.iter().sum(); - if length == 0 || offset >= file_size { - timer.success(); - return Ok(Vec::new()); - } - let end = offset.saturating_add(length).min(file_size); - - // Walk blocks in order, fetching and decrypting only those whose - // plaintext span overlaps `[offset, end)`. - let mut out = Vec::with_capacity((end - offset) as usize); - let mut block_start: u64 = 0; - for (block, &bsize) in blocks.iter().zip(block_sizes.iter()) { - let block_end = block_start + bsize; - if block_end <= offset { - block_start = block_end; - continue; - } - if block_start >= end { - break; - } - let ciphertext = self - .http - .get_storage_blob(&block.bare_url, &block.token) - .await?; - let plaintext = content_key.decrypt_block(&ciphertext)?; - let from = offset.saturating_sub(block_start) as usize; - let to = ((end - block_start) as usize).min(plaintext.len()); - if from < to { - out.extend_from_slice(&plaintext[from..to]); - } - block_start = block_end; - } + timer.success(); + Ok(RevisionReader::new( + self.clone(), + uid.clone(), + revision_id, + content_key, + blocks, + block_sizes, + )) + } + /// Download and decrypt only the plaintext byte range `[offset, offset + length)` + /// of a file's active revision. + /// + /// A one-shot convenience over [`open_revision`](Self::open_revision): it + /// resolves the revision's keys and block table, reads the range, and drops + /// them. Callers that issue more than one read against the same file should + /// hold a [`RevisionReader`] instead — this call repeats the whole + /// resolution (two API round-trips and an S2K node-key unlock) every time. + /// + /// The range is clamped to the file's length, so a read at or past EOF + /// yields fewer bytes (or none). + /// + /// Unlike [`download_file_to`](Self::download_file_to), a partial read + /// cannot recompute the full content manifest, so manifest-signature + /// verification is skipped. + pub async fn download_range(&self, uid: &NodeUid, offset: u64, length: u64) -> Result> { + let mut timer = self.telemetry.start("download_range"); + let reader = self.open_revision(uid).await?; + let out = reader.read_at(offset, length).await?; + timer.attr("byte_count", out.len()); timer.success(); Ok(out) } @@ -1639,7 +1713,20 @@ impl ProtonDriveClient { /// [`enumerate_nodes`](Self::enumerate_nodes). pub async fn enumerate_shared_with_me_node_uids(&self) -> Result> { let mut timer = self.telemetry.start("enumerate_shared_with_me_node_uids"); + let uids = self.page_shared_with_me().await?; + timer.success(); + Ok(uids) + } + + /// Page `v2/sharedwithme` in the order the API returns, recording each item's + /// membership share on the way through. Returns the uids; the share index is + /// left in the cache for [`shared_with_me_shares`](Self::shared_with_me_shares). + /// + /// Order is the API's and is preserved: it is what a front-end lists, and a + /// set's iteration order would reshuffle the page on every refresh. + async fn page_shared_with_me(&self) -> Result> { let mut uids = Vec::new(); + let mut shares = HashMap::new(); let mut anchor: Option = None; loop { @@ -1653,7 +1740,9 @@ impl ProtonDriveClient { let is_drive_item = ShareTargetType::from_raw(link.share_target_type) .is_some_and(ShareTargetType::is_drive_item); if is_drive_item { - uids.push(NodeUid::new(link.volume_id.clone(), link.link_id.clone())); + let uid = NodeUid::new(link.volume_id.clone(), link.link_id.clone()); + shares.insert(uid.clone(), link.share_id.clone()); + uids.push(uid); } } @@ -1663,10 +1752,31 @@ impl ProtonDriveClient { } } - timer.success(); + self.cache.lock().await.shared_with_me_shares = Some(shares); Ok(uids) } + /// The membership share id behind each node shared with us, keyed by the + /// shared node. Pages `v2/sharedwithme` on a miss; `refresh` re-pages even + /// when the index is already cached. + /// + /// A shared node is the root of the share its owner granted us, so it comes + /// back parentless on *their* volume and only that share's key unlocks it — + /// see [`root_link_share_key`](Self::root_link_share_key). + async fn shared_with_me_shares(&self, refresh: bool) -> Result> { + if !refresh && let Some(shares) = self.cache.lock().await.shared_with_me_shares.clone() { + return Ok(shares); + } + self.page_shared_with_me().await?; + Ok(self + .cache + .lock() + .await + .shared_with_me_shares + .clone() + .unwrap_or_default()) + } + /// The nodes I have shared with others, as [`NodeUid`]s. /// /// Ported from the TypeScript SDK's `iterateSharedNodeUids`: page @@ -3222,7 +3332,61 @@ impl ProtonDriveClient { /// Fetch every block of a revision (paginated), returning the revision /// metadata plus the contiguous, index-sorted block list. - async fn fetch_revision_blocks( + /// Fetch every content block, decrypt it, write the plaintext to `output` + /// and append the ciphertext digest to `manifest`. + /// + /// Blocks are fetched [`MAX_CONCURRENT_BLOCK_DOWNLOADS`] at a time and + /// decrypted off-runtime, but are consumed strictly in index order: both the + /// manifest (whose signature is checked over the concatenated digests) and + /// `output` depend on that ordering. + async fn write_content_blocks( + &self, + blocks: &[BlockDto], + content_key: &ContentKey, + manifest: &mut Vec, + output: &mut W, + ) -> Result<()> { + // Each fetch owns its url/token rather than borrowing `blocks`. Borrowing + // here makes the resulting future carry a higher-ranked lifetime that + // `tokio::spawn` rejects ("implementation of `FnOnce` is not general + // enough") in *callers* of `download_file_to` — a downstream break, not a + // local one, so it does not show up in this crate's own build. + let fetches: Vec<(String, String)> = blocks + .iter() + .map(|block| (block.bare_url.clone(), block.token.clone())) + .collect(); + + let mut decrypted = stream::iter(fetches.into_iter().map(|(url, token)| { + let http = self.http.clone(); + let content_key = content_key.clone(); + let slots = self.block_slots(); + async move { + // The client-wide in-flight block cap. Held across fetch and + // decrypt, and returned to the caller so it outlives the + // plaintext sitting in `buffered`'s queue — see + // `RevisionReader::block_plaintext`. + let permit = slots.acquire_owned().await.map_err(|e| { + ProtonError::invalid_operation(format!("block slots closed: {e}")) + })?; + let ciphertext = http.get_storage_blob(&url, &token).await?; + let (digest, plaintext) = + digest_and_decrypt_block_blocking(content_key, ciphertext).await?; + Ok::<_, ProtonError>((digest, plaintext, permit)) + } + })) + .buffered(MAX_CONCURRENT_BLOCK_DOWNLOADS); + + while let Some((digest, plaintext, _permit)) = decrypted.try_next().await? { + output + .write_all(&plaintext) + .map_err(|e| ProtonError::invalid_operation(format!("write block: {e}")))?; + manifest.extend_from_slice(&digest); + } + + Ok(()) + } + + pub(crate) async fn fetch_revision_blocks( &self, volume_id: &VolumeId, link_id: &LinkId, @@ -3273,6 +3437,15 @@ impl ProtonDriveClient { return Ok(()); } + // Only one of a burst of first calls does the bootstrap; the rest wait + // here and take the second check below. On failure the guard is dropped + // with nothing cached, so the next caller retries rather than inheriting + // the error. + let _guard = self.bootstrap.my_files.lock().await; + if self.cache.lock().await.my_files_share.is_some() { + return Ok(()); + } + // Mirrors C# `NodeOperations.GetOrCreateMyFilesFolderAsync`: a brand-new // account has no My Files volume yet, so the share lookup fails. Create // the volume, then re-read it through the normal path to populate caches. @@ -3358,9 +3531,11 @@ impl ProtonDriveClient { /// 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). + /// wrapped to that device's *own* share key, and a node another user shares + /// with us is parentless on *their* volume, wrapped to the membership share + /// we were granted. Match the root link id against the device set, then the + /// shared-with-me set, unlocking whichever share owns it; fall back to the + /// My Files share for anything unrecognised. async fn root_link_share_key( &self, volume_id: &VolumeId, @@ -3385,9 +3560,33 @@ impl ProtonDriveClient { let (key, _address_id) = self.share_key_by_id(&device.share_id).await?; return Ok(key); } + + let root_uid = NodeUid::new(volume_id.clone(), root_id.clone()); + if let Some(key) = self.shared_with_me_share_key(&root_uid).await? { + return Ok(key); + } + self.root_share_key().await } + /// The membership share key for a node shared with us, or `None` when the + /// node is not one. A cached listing that misses is re-paged once, so a + /// share joined since the last listing still resolves. + async fn shared_with_me_share_key(&self, uid: &NodeUid) -> Result> { + let was_cached = self.cache.lock().await.shared_with_me_shares.is_some(); + let mut share_id = self.shared_with_me_shares(false).await?.get(uid).cloned(); + if share_id.is_none() && was_cached { + share_id = self.shared_with_me_shares(true).await?.get(uid).cloned(); + } + match share_id { + Some(share_id) => { + let (key, _address_id) = self.share_key_by_id(&share_id).await?; + Ok(Some(key)) + } + None => Ok(None), + } + } + /// Resolve the Photos share + root folder, caching them. /// /// Mirrors C# `PhotosNodeOperations.GetFreshExistingPhotosFolderAsync`: @@ -3403,6 +3602,17 @@ impl ProtonDriveClient { } } + // Single-flighted for the same reason as `ensure_my_files`: a burst of + // first photos calls (a gallery opening pages several at once) would + // otherwise each fetch and decrypt the share. + let _guard = self.bootstrap.photos.lock().await; + { + let cache = self.cache.lock().await; + if let Some(exists) = cache.photos_volume_exists { + return Ok(exists); + } + } + let response: MyFilesShareResponse = match self.http.get("v2/shares/photos").await { Ok(response) => response, Err(ProtonError::Api(e)) if e.code == proton_sdk::api::ResponseCode::DoesNotExist => { @@ -3638,18 +3848,8 @@ impl ProtonDriveClient { } } - for block in &blocks { - let ciphertext = self - .http - .get_storage_blob(&block.bare_url, &block.token) - .await?; - let digest = Sha256::digest(&ciphertext); - let plaintext = content_key.decrypt_block(&ciphertext)?; - output - .write_all(&plaintext) - .map_err(|e| ProtonError::invalid_operation(format!("write block: {e}")))?; - manifest.extend_from_slice(&digest); - } + self.write_content_blocks(&blocks, &content_key, &mut manifest, output) + .await?; verify_manifest(&self.account, &revision, &node_key, &manifest).await; Ok(()) diff --git a/crates/proton-drive-rs/src/lib.rs b/crates/proton-drive-rs/src/lib.rs index 183f12f..d0ec7fa 100644 --- a/crates/proton-drive-rs/src/lib.rs +++ b/crates/proton-drive-rs/src/lib.rs @@ -1,8 +1,7 @@ //! High-level Proton Drive client (Rust port). //! -//! Built on top of [`proton_sdk`]. Milestone 1 covers session-backed read -//! operations: resolving the My Files folder, fetching node metadata and -//! enumerating folder children. +//! Built on top of [`proton_sdk`]. Implements full read/write, uploads/downloads, +//! sharing, bookmarks, device sync registration, events tracking, and photos timeline. //! //! ```no_run //! # async fn run() -> proton_sdk::error::Result<()> { @@ -41,6 +40,7 @@ mod dtos; mod events; mod node; mod photos; +mod revision; mod sharing; pub use cache::{CachedNodeInfo, DriveEntityCache}; @@ -51,6 +51,7 @@ pub use node::{Node, NodeKind, RevisionState, Thumbnail, ThumbnailType}; pub use photos::{PhotoTag, PhotoUploadMetadata, PhotosTimelineItem, ProtonPhotosClient}; pub use proton_sdk::account::KeySalt; pub use proton_sdk::cache::{CacheRepository, EncryptedCacheRepository, InMemoryCacheRepository}; +pub use revision::RevisionReader; pub use sharing::{ Bookmark, ExternalInvitation, ExternalInvitationState, IncomingInvitation, MemberRole, PublicLink, ShareInvitation, ShareMember, diff --git a/crates/proton-drive-rs/src/revision.rs b/crates/proton-drive-rs/src/revision.rs new file mode 100644 index 0000000..1a0e3f3 --- /dev/null +++ b/crates/proton-drive-rs/src/revision.rs @@ -0,0 +1,306 @@ +//! A seekable reader over a file revision's content blocks. +//! +//! Mirrors the TypeScript SDK's `internal/download/seekableStream.ts` + +//! `blockIndex.ts`: resolve a revision's keys and block table once, then serve +//! arbitrary plaintext ranges by fetching only the blocks that overlap them. + +use bytes::Bytes; +use futures::stream::{self, StreamExt, TryStreamExt}; +use proton_sdk::crypto::ContentKey; +use proton_sdk::error::{ProtonError, Result}; +use proton_sdk::ids::NodeUid; +use sha2::{Digest, Sha256}; +use tokio::sync::{Mutex, OwnedSemaphorePermit, RwLock}; + +use crate::client::ProtonDriveClient; +use crate::dtos::BlockDto; + +/// Content blocks fetched concurrently within a single file. +/// +/// Matches the TypeScript SDK's `MAX_DOWNLOAD_BLOCK_SIZE` +/// (`internal/download/fileDownloader.ts`). Each in-flight block holds its 4 MiB +/// ciphertext and then its plaintext, so this bounds one download at roughly +/// `2 * 4 MiB * N`. +/// +/// This is the *per-file* ceiling only. The host-wide one is the client's +/// in-flight block permits (`DEFAULT_MAX_INFLIGHT_BLOCKS`), without which N +/// concurrent downloads would cost N times this bound. TypeScript caps +/// concurrent files instead (`internal/download/queue.ts`). +pub(crate) const MAX_CONCURRENT_BLOCK_DOWNLOADS: usize = 10; + +/// Decrypt one content block off the async runtime. +/// +/// PGP block decryption is CPU-bound and takes long enough on a 4 MiB block to +/// stall every other task sharing the reactor thread — including the other block +/// fetches this download is running concurrently. Hand it to the blocking pool. +pub(crate) async fn decrypt_block_blocking( + content_key: ContentKey, + ciphertext: Bytes, +) -> Result> { + join_decrypt(tokio::task::spawn_blocking(move || { + content_key.decrypt_block(&ciphertext) + })) + .await +} + +/// As [`decrypt_block_blocking`], but also returns the SHA-256 of the +/// *ciphertext* — the download manifest's per-block input. Hashing 4 MiB is +/// itself CPU-bound, so it rides along in the same blocking task rather than +/// costing a second hop. +pub(crate) async fn digest_and_decrypt_block_blocking( + content_key: ContentKey, + ciphertext: Bytes, +) -> Result<(Vec, Vec)> { + let handle = tokio::task::spawn_blocking(move || { + let digest = Sha256::digest(&ciphertext).to_vec(); + content_key + .decrypt_block(&ciphertext) + .map(|pt| (digest, pt)) + }); + + join_decrypt(handle).await +} + +/// Await a blocking decrypt task, mapping both the join failure and the crypto +/// failure into a [`ProtonError`]. +async fn join_decrypt( + handle: tokio::task::JoinHandle>, +) -> Result { + handle + .await + .map_err(|e| ProtonError::invalid_operation(format!("block decrypt task failed: {e}")))? + .map_err(Into::into) +} + +/// The block table, versioned so a refresh can be deduplicated: a caller that +/// hit an expired URL records the generation it saw, and only re-lists the +/// revision if nobody else has already replaced that generation. +struct BlockTable { + blocks: Vec, + generation: u64, +} + +/// An open handle on a file revision: its content key plus the block table and +/// per-block plaintext sizes, resolved once by +/// [`ProtonDriveClient::open_revision`](crate::ProtonDriveClient::open_revision). +/// +/// Reading a range costs only the block bodies it overlaps — no link details, no +/// ancestor walk, no node-key unlock, no revision listing. That makes it the +/// right shape for an on-demand mount, which reads the same file many times at +/// a granularity far below the 4 MiB block size. +/// +/// A reader is pinned to the revision that was active when it was opened. If the +/// file gains a new revision, this reader keeps serving the old one; reopen to +/// follow the change. +pub struct RevisionReader { + client: ProtonDriveClient, + uid: NodeUid, + revision_id: String, + content_key: ContentKey, + blocks: RwLock, + /// Serializes block-table refreshes so a burst of expired-URL failures + /// triggers one re-listing rather than one per block. + refresh: Mutex<()>, + /// Plaintext size of each block, in block order. + block_sizes: Vec, + file_size: u64, +} + +impl RevisionReader { + pub(crate) fn new( + client: ProtonDriveClient, + uid: NodeUid, + revision_id: String, + content_key: ContentKey, + blocks: Vec, + block_sizes: Vec, + ) -> Self { + let file_size = block_sizes.iter().sum(); + Self { + client, + uid, + revision_id, + content_key, + blocks: RwLock::new(BlockTable { + blocks, + generation: 0, + }), + refresh: Mutex::new(()), + block_sizes, + file_size, + } + } + + /// The node this reader was opened on. + pub fn uid(&self) -> &NodeUid { + &self.uid + } + + /// The revision this reader is pinned to. + pub fn revision_id(&self) -> &str { + &self.revision_id + } + + /// Total plaintext size of the revision, summed from the block sizes + /// recorded in its extended attributes. + pub fn size(&self) -> u64 { + self.file_size + } + + /// Plaintext size of each content block, in block order. A caller that + /// aligns its reads to these boundaries fetches each block exactly once. + pub fn block_sizes(&self) -> &[u64] { + &self.block_sizes + } + + /// Read the plaintext byte range `[offset, offset + length)`. + /// + /// The range is clamped to the revision's length, so a read at or past EOF + /// yields fewer bytes (or none). Only the blocks overlapping the range are + /// fetched, and they are fetched up to + /// [`MAX_CONCURRENT_BLOCK_DOWNLOADS`] at a time. + /// + /// A partial read cannot recompute the content manifest, so — as with + /// [`ProtonDriveClient::download_range`](crate::ProtonDriveClient::download_range) + /// — manifest-signature verification is skipped. Use + /// [`ProtonDriveClient::download_file_to`](crate::ProtonDriveClient::download_file_to) + /// when the whole file is wanted and authenticity should be checked. + pub async fn read_at(&self, offset: u64, length: u64) -> Result> { + if length == 0 || offset >= self.file_size { + return Ok(Vec::new()); + } + let end = offset.saturating_add(length).min(self.file_size); + + // Which blocks overlap the range, and where each one starts in the + // plaintext — resolved up front so the fetches can run concurrently. + let mut wanted: Vec<(usize, u64)> = Vec::new(); + let mut block_start: u64 = 0; + for (index, &block_size) in self.block_sizes.iter().enumerate() { + if block_start >= end { + break; + } + let block_end = block_start + block_size; + if block_end > offset { + wanted.push((index, block_start)); + } + block_start = block_end; + } + + // `buffered` yields in input order, so the slices append in block order. + // The closure takes the index *by value*: a closure taking a reference + // would give the fetch future a higher-ranked lifetime that `tokio::spawn` + // rejects in callers ("implementation of `FnOnce` is not general enough"). + let indices: Vec = wanted.iter().map(|&(index, _)| index).collect(); + let mut blocks = stream::iter(indices.into_iter().map(|index| self.block_plaintext(index))) + .buffered(MAX_CONCURRENT_BLOCK_DOWNLOADS); + + let mut out = Vec::with_capacity((end - offset) as usize); + let mut next = 0_usize; + // `permit` is the in-flight block slot; it frees at the end of each + // iteration, once the bytes we want are copied into `out`. + while let Some((plaintext, _permit)) = blocks.try_next().await? { + let (_, block_start) = wanted[next]; + next += 1; + + let from = offset.saturating_sub(block_start) as usize; + let to = ((end - block_start) as usize).min(plaintext.len()); + if from < to { + out.extend_from_slice(&plaintext[from..to]); + } + } + + Ok(out) + } + + /// Fetch and decrypt one content block, re-listing the revision once if its + /// storage URL has expired. + /// + /// Returns the permit alongside the plaintext: it is the client-wide + /// 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)> { + // Held across fetch *and* decrypt: both halves are resident 4 MiB + // buffers, and the plaintext outlives the ciphertext. + let permit = self + .client + .block_slots() + .acquire_owned() + .await + .map_err(|e| ProtonError::invalid_operation(format!("block slots closed: {e}")))?; + + let (url, token, generation) = self.block_location(index).await?; + + let ciphertext = match self.client.http().get_storage_blob(&url, &token).await { + Ok(bytes) => bytes, + Err(e) if is_expired_block_url(&e) => { + self.refresh_blocks(generation).await?; + let (url, token, _) = self.block_location(index).await?; + self.client.http().get_storage_blob(&url, &token).await? + } + Err(e) => return Err(e), + }; + + let plaintext = decrypt_block_blocking(self.content_key.clone(), ciphertext).await?; + Ok((plaintext, permit)) + } + + /// The storage URL and token for a block, plus the generation of the table + /// they came from. + async fn block_location(&self, index: usize) -> Result<(String, String, u64)> { + let table = self.blocks.read().await; + let block = table.blocks.get(index).ok_or_else(|| { + ProtonError::invalid_operation(format!( + "block {index} is missing from revision {}", + self.revision_id + )) + })?; + Ok(( + block.bare_url.clone(), + block.token.clone(), + table.generation, + )) + } + + /// Re-list the revision to obtain fresh block URLs, unless another task has + /// already replaced the generation the caller saw. + async fn refresh_blocks(&self, seen_generation: u64) -> Result<()> { + let _guard = self.refresh.lock().await; + + if self.blocks.read().await.generation != seen_generation { + // Someone else refreshed while we waited; their table is at least as + // fresh as one we would fetch now. + return Ok(()); + } + + let (_, blocks) = self + .client + .fetch_revision_blocks(&self.uid.volume_id, &self.uid.link_id, &self.revision_id) + .await?; + + if blocks.len() != self.block_sizes.len() { + return Err(ProtonError::invalid_operation(format!( + "revision {} changed block count while open ({} -> {})", + self.revision_id, + self.block_sizes.len(), + blocks.len() + ))); + } + + let mut table = self.blocks.write().await; + table.blocks = blocks; + table.generation += 1; + Ok(()) + } +} + +/// Whether a storage fetch failed in a way that an expired block URL would +/// produce. Block URLs carry their own authorization and are handed out with a +/// server-side lifetime we neither control nor are told, so recovery is +/// reactive: on these statuses, re-list the revision and try once more. +fn is_expired_block_url(error: &ProtonError) -> bool { + matches!( + error, + ProtonError::Api(e) if matches!(e.http_status, 401 | 403 | 404) + ) +} diff --git a/crates/proton-drive-rs/tests/common/mod.rs b/crates/proton-drive-rs/tests/common/mod.rs index c27850d..be413fd 100644 --- a/crates/proton-drive-rs/tests/common/mod.rs +++ b/crates/proton-drive-rs/tests/common/mod.rs @@ -16,6 +16,7 @@ use std::time::{SystemTime, UNIX_EPOCH}; use hmac::{Hmac, KeyInit, Mac}; use proton_drive_rs::ProtonDriveClient; use proton_sdk::config::ProtonClientConfiguration; +use proton_sdk::error::ProtonError; use proton_sdk::session::{PasswordMode, ProtonApiSession, ResumeParameters}; use serde::{Deserialize, Serialize}; use sha1::Sha1; @@ -70,16 +71,33 @@ async fn build_live_client() -> Option { let config = ProtonClientConfiguration::new(APP_VERSION).with_user_agent(USER_AGENT); // Fast path: resume a cached session (no TOTP). Validate with one cheap - // authenticated call so a dead/revoked session falls through to a fresh login. + // authenticated call so a session the tests cannot use falls through to a + // fresh login. + // + // The probe is `core/v4/keys/salts` specifically. A resumed session keeps + // working for most routes — `core/v4/users`, `core/v4/addresses`, even + // `drive/v2/shares/my-files` all still answer — long after its *password* + // scope has lapsed. Only the key-salts route reflects that, with + // `InsufficientScope` (HTTP 403). Since the salts are step one of every + // decryption chain (salt → bcrypt → unlock user keys), a session that cannot + // read them is useless to these tests even though it looks alive. Probing a + // route that does not need the scope is what let a dead session through and + // failed the whole binary. if let Some(stored) = load_cached_session() { let session = ProtonApiSession::resume(config.clone(), stored.into_params()) .expect("resume cached session"); - match session + let probe = session .http() - .get::("core/v4/users") - .await - { - Ok(_) => { + .get::("core/v4/keys/salts") + .await; + + match probe { + Err(e) if needs_fresh_login(&e) => { + eprintln!("[auth] cached session lost its password scope ({e}); logging in fresh"); + } + // Anything else — success, or a transport/parse hiccup — leaves the + // session usable; don't burn a TOTP window over it. + Ok(_) | Err(_) => { eprintln!("[auth] resumed cached session (no TOTP)"); // Refresh may have rotated the tokens; re-persist the current set. save_session(&session).await; @@ -89,7 +107,6 @@ async fn build_live_client() -> Option { _session: session, }); } - Err(e) => eprintln!("[auth] cached session invalid ({e}); logging in fresh"), } } @@ -121,6 +138,21 @@ async fn build_live_client() -> Option { }) } +/// Whether an error means the cached session cannot be used and a fresh SRP+2FA +/// login is required. +/// +/// `InsufficientScope` is the important one: no amount of token refreshing fixes +/// it (the user must re-authenticate with their password), so the http client's +/// automatic 401-refresh path cannot recover and every later call needing that +/// scope fails the same way. A bare 401 that survived that refresh path means the +/// refresh token is dead too. +fn needs_fresh_login(error: &ProtonError) -> bool { + matches!( + error, + ProtonError::Api(e) if e.is_insufficient_scope() || e.is_unauthorized() + ) +} + /// Persisted session credentials. Mirrors [`ResumeParameters`], serialized to a /// gitignored cache file so reruns can skip SRP+2FA. #[derive(Serialize, Deserialize)] diff --git a/crates/proton-drive-rs/tests/live_upload_download.rs b/crates/proton-drive-rs/tests/live_upload_download.rs index e295b82..3854ae6 100644 --- a/crates/proton-drive-rs/tests/live_upload_download.rs +++ b/crates/proton-drive-rs/tests/live_upload_download.rs @@ -8,7 +8,44 @@ mod common; +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; + use proton_sdk::ids::NodeUid; +use proton_sdk::telemetry::{Telemetry, TelemetryEvent}; + +/// Counts telemetry events per operation, so a test can assert *which* work an +/// SDK call actually did (e.g. that a read issued no API requests). +#[derive(Default)] +struct OpCounter { + counts: Mutex>, +} + +impl Telemetry for OpCounter { + fn record(&self, event: &TelemetryEvent) { + *self + .counts + .lock() + .expect("counter poisoned") + .entry(event.operation) + .or_default() += 1; + } +} + +impl OpCounter { + fn count(&self, operation: &str) -> usize { + self.counts + .lock() + .expect("counter poisoned") + .get(operation) + .copied() + .unwrap_or(0) + } + + fn reset(&self) { + self.counts.lock().expect("counter poisoned").clear(); + } +} /// Trash then permanently delete the given nodes; best-effort, logs on failure. async fn cleanup(client: &proton_drive_rs::ProtonDriveClient, uids: &[NodeUid]) { @@ -260,3 +297,101 @@ async fn download_range_multi_block() { cleanup(client, &[uid]).await; } + +/// A held `RevisionReader` serves the same ranges as `download_range`, and — +/// the point of the handle — resolves the revision's metadata only once: the +/// reads themselves must issue no API requests at all, only block fetches. +#[tokio::test] +#[ignore = "live: needs test-account credentials"] +async fn open_revision_reader_reuses_resolved_metadata() { + let Some(live) = common::live_client().await else { + return; + }; + + let counter = Arc::new(OpCounter::default()); + // Clones share the session, caches and connection pool; only the telemetry + // sink differs, so the counts below are this test's own work. + let client = live.client.clone().with_telemetry(counter.clone()); + + let root = client + .get_my_files_folder() + .await + .expect("get my-files root"); + + // 10 MiB → blocks of 4 MiB, 4 MiB, 2 MiB. + let block = 4 * 1024 * 1024u64; + let size = (10 * 1024 * 1024) as usize; + let mut payload = vec![0u8; size]; + for (i, b) in payload.iter_mut().enumerate() { + *b = (i as u32).wrapping_mul(2_654_435_761) as u8; + } + + let name = format!("rt-reader-{}.bin", common::unique_suffix()); + let uid = client + .upload_file_from( + &root.uid, + &name, + "application/octet-stream", + std::io::Cursor::new(payload.clone()), + size as i64, + Vec::new(), + None, + false, + ) + .await + .expect("upload_file_from"); + + let reader = client.open_revision(&uid).await.expect("open_revision"); + + assert_eq!(reader.size(), size as u64, "reader size"); + assert_eq!( + reader.block_sizes(), + &[block, block, 2 * 1024 * 1024], + "reader block sizes" + ); + + // Everything the reader needs is now resolved. From here on, reads must + // only fetch block bodies. + counter.reset(); + + let total = size as u64; + let cases: &[(u64, u64)] = &[ + (0, 256), // start of block 1 + (block - 128, 256), // straddles block 1 → block 2 + (block, 4096), // exact start of block 2 + (2 * block - 1, 2), // straddles block 2 → block 3 (final short block) + (2 * block + 1000, 50_000), // interior of final short block + (total - 100, 500), // tail, length past EOF → clamps + (block - 10, block + 20), // spans a full block plus both neighbors + (0, total), // whole file + (total, 4096), // at EOF → empty + ]; + + for &(off, len) in cases { + let got = reader + .read_at(off, len) + .await + .unwrap_or_else(|e| panic!("read_at({off},{len}): {e}")); + let from = (off as usize).min(size); + let to = ((off + len) as usize).min(size); + let want = &payload[from..to]; + assert_eq!( + got.len(), + want.len(), + "read_at({off},{len}) length mismatch" + ); + assert_eq!(got, want, "read_at({off},{len}) byte mismatch"); + } + + assert_eq!( + counter.count("http_request"), + 0, + "reads must not re-resolve link details or the revision listing" + ); + assert!( + counter.count("storage_download") > 0, + "reads must fetch block bodies" + ); + + cleanup(&client, &[uid]).await; +} diff --git a/crates/proton-drive-rs/tests/spawnable.rs b/crates/proton-drive-rs/tests/spawnable.rs new file mode 100644 index 0000000..b30b3da --- /dev/null +++ b/crates/proton-drive-rs/tests/spawnable.rs @@ -0,0 +1,67 @@ +//! Compile-time guard: the download futures must survive `tokio::spawn`. +//! +//! `download_file`/`download_file_to` and `RevisionReader::read_at` internally +//! drive a `futures::stream::buffered` pipeline. If any stage of that pipeline +//! borrows from a closure argument, the resulting future picks up a +//! higher-ranked lifetime and callers get "implementation of `FnOnce` is not +//! general enough" the moment they `tokio::spawn` it. That error surfaces only +//! *downstream* — this crate builds fine either way — so it went unnoticed until +//! a consumer broke. +//! +//! These never run; they only have to compile. + +#![allow(dead_code)] + +use std::sync::Arc; + +use tokio::task::JoinHandle; + +use proton_drive_rs::{ProtonDriveClient, RevisionReader}; +use proton_sdk::ids::NodeUid; + +/// `download_file` inside a spawned task. +fn download_file_is_spawnable( + client: ProtonDriveClient, + uid: NodeUid, +) -> JoinHandle>> { + tokio::spawn(async move { client.download_file(&uid).await }) +} + +/// `download_file_to` inside a spawned task, writing into an owned sink. +fn download_file_to_is_spawnable( + client: ProtonDriveClient, + uid: NodeUid, +) -> JoinHandle> { + tokio::spawn(async move { + let mut sink: Vec = Vec::new(); + client.download_file_to(&uid, &mut sink).await + }) +} + +/// `download_range` inside a spawned task. +fn download_range_is_spawnable( + client: ProtonDriveClient, + uid: NodeUid, +) -> JoinHandle>> { + tokio::spawn(async move { client.download_range(&uid, 0, 4096).await }) +} + +/// `open_revision` + a shared reader's `read_at` inside a spawned task — the +/// shape an on-demand mount uses to fetch several blocks concurrently. +fn reader_read_at_is_spawnable( + client: ProtonDriveClient, + uid: NodeUid, +) -> JoinHandle> { + tokio::spawn(async move { + let reader: Arc = Arc::new(client.open_revision(&uid).await?); + let mut set = tokio::task::JoinSet::new(); + for i in 0..4u64 { + let reader = reader.clone(); + set.spawn(async move { reader.read_at(i * 4096, 4096).await }); + } + while let Some(joined) = set.join_next().await { + joined.expect("task panicked")?; + } + Ok::<_, proton_sdk::error::ProtonError>(()) + }) +} diff --git a/crates/proton-sdk/src/http.rs b/crates/proton-sdk/src/http.rs index bec4804..c5801f6 100644 --- a/crates/proton-sdk/src/http.rs +++ b/crates/proton-sdk/src/http.rs @@ -8,6 +8,7 @@ use std::sync::Arc; use std::time::Duration; +use bytes::Bytes; use rand::RngExt; use reqwest::{Method, StatusCode}; use serde::Serialize; @@ -72,8 +73,14 @@ impl ApiHttpClient { session_id: SessionId, tokens: Tokens, ) -> Result { + // `gzip` makes reqwest advertise Accept-Encoding and transparently decode + // the response. The Proton API honours it for the JSON envelope, which is + // the bulk of a metadata-heavy walk (link details, listings). Block bodies + // are already ciphertext and will not compress — the header costs nothing + // there, and storage responses are unaffected either way. let http = reqwest::Client::builder() .timeout(config.request_timeout) + .gzip(true) .build()?; let base_url = ensure_trailing_slash(&config.base_url); @@ -150,7 +157,12 @@ impl ApiHttpClient { /// rather than the session bearer. Mirrors C# `StorageApiClient /// .GetBlobStreamAsync`. A successful response is raw binary; an error /// response is the usual JSON envelope. - pub async fn get_storage_blob(&self, url: &str, token: &str) -> Result> { + /// + /// Returns [`Bytes`] rather than `Vec` so the 4 MiB body is not copied + /// on its way to the decryptor — the reference-counted buffer reqwest + /// already assembled is handed straight through, including into the + /// blocking decrypt task. + pub async fn get_storage_blob(&self, url: &str, token: &str) -> Result { let mut timer = self.telemetry().start("storage_download"); let response = send_retrying(&self.inner.config.retry_policy, || { let mut request = self.inner.http.get(url).header(STORAGE_TOKEN_HEADER, token); @@ -176,7 +188,7 @@ impl ApiHttpClient { } timer.success(); - Ok(bytes.to_vec()) + Ok(bytes) } /// `POST {url}` a block blob to storage as `multipart/form-data`. @@ -431,6 +443,7 @@ pub async fn post_unauthenticated( ) -> Result { let http = reqwest::Client::builder() .timeout(config.request_timeout) + .gzip(true) .build()?; let base_url = ensure_trailing_slash(&config.base_url);