nvrl/proton-sdk-rs
1
0 Fork 0
mirror of https://github.com/narrrl/proton-sdk-rs.git synced 2026-11-03 16:41:26 +01:00
Code Issues Packages Projects Releases Wiki Activity

revision handler, parallel block download + gzip

Browse Source
This commit is contained in:
nvrl committed 2026-07-17 15:57:00 +02:00
1 parent 0e335e35ec
commit 0a15a8c86e
9 files changed
+885 -87
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

Cargo.lock
Generated
+46 -2
View File
@@ -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",
Cargo.toml
+2 -2
View File
@@ -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"
crates/proton-drive-rs/src/client.rs
+271 -71
View File
@@ -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<dyn Telemetry>,
/// 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<Semaphore>,
/// 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<Bootstrap>,
}
/// 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<bool>,
/// Decrypted node key per folder, used as the parent key for its children.
folder_keys: HashMap<NodeUid, PrivateKey>,
/// 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<HashMap<NodeUid, ShareId>>,
}
/// 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<Semaphore> {
self.block_slots.clone()
}
/// Resolve (and cache) the user's "My Files" root folder.
pub async fn get_my_files_folder(&self) -> Result<Node> {
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)
self.write_content_blocks(&blocks, &content_key, &mut manifest, output)
.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);
}
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<Vec<u8>> {
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<RevisionReader> {
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;
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<Vec<u8>> {
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<Vec<NodeUid>> {
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<Vec<NodeUid>> {
let mut uids = Vec::new();
let mut shares = HashMap::new();
let mut anchor: Option<String> = 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<HashMap<NodeUid, ShareId>> {
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<W: std::io::Write>(
&self,
blocks: &[BlockDto],
content_key: &ContentKey,
manifest: &mut Vec<u8>,
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<Option<PrivateKey>> {
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)
self.write_content_blocks(&blocks, &content_key, &mut manifest, output)
.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);
}
verify_manifest(&self.account, &revision, &node_key, &manifest).await;
Ok(())
crates/proton-drive-rs/src/lib.rs
+4 -3
View File
@@ -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,
crates/proton-drive-rs/src/revision.rs
+306
View File
@@ -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<Vec<u8>> {
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<u8>, Vec<u8>)> {
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<T>(
handle: tokio::task::JoinHandle<std::result::Result<T, proton_sdk::crypto::CryptoError>>,
) -> Result<T> {
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<BlockDto>,
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<BlockTable>,
/// 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<u64>,
file_size: u64,
}
impl RevisionReader {
pub(crate) fn new(
client: ProtonDriveClient,
uid: NodeUid,
revision_id: String,
content_key: ContentKey,
blocks: Vec<BlockDto>,
block_sizes: Vec<u64>,
) -> 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<Vec<u8>> {
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<usize> = 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<u8>, 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)
)
}
crates/proton-drive-rs/tests/common/mod.rs
+39 -7
View File
@@ -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<LiveClient> {
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::<serde_json::Value>("core/v4/users")
.await
{
Ok(_) => {
.get::<serde_json::Value>("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<LiveClient> {
_session: session,
});
}
Err(e) => eprintln!("[auth] cached session invalid ({e}); logging in fresh"),
}
}
@@ -121,6 +138,21 @@ async fn build_live_client() -> Option<LiveClient> {
})
}
/// 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)]
crates/proton-drive-rs/tests/live_upload_download.rs
+135
View File
@@ -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<HashMap<&'static str, usize>>,
}
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;
}
crates/proton-drive-rs/tests/spawnable.rs
+67
View File
@@ -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<proton_sdk::error::Result<Vec<u8>>> {
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<proton_sdk::error::Result<()>> {
tokio::spawn(async move {
let mut sink: Vec<u8> = 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<proton_sdk::error::Result<Vec<u8>>> {
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<proton_sdk::error::Result<()>> {
tokio::spawn(async move {
let reader: Arc<RevisionReader> = 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>(())
})
}
crates/proton-sdk/src/http.rs
+15 -2
View File
@@ -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<Self> {
// `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<Vec<u8>> {
///
/// Returns [`Bytes`] rather than `Vec<u8>` 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<Bytes> {
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<B: Serialize, T: DeserializeOwned>(
) -> Result<T> {
let http = reqwest::Client::builder()
.timeout(config.request_timeout)
.gzip(true)
.build()?;
let base_url = ensure_trailing_slash(&config.base_url);