From 49e5061141885b9793c061419e7ba8644bd6db89 Mon Sep 17 00:00:00 2001 From: TheLastFlame Date: Sat, 10 Oct 2026 13:53:23 +1000 Subject: [PATCH 1/5] feat(mtproto): persist foreign dc auth keys in client session --- native/mtproto-rs/src/client/mod.rs | 29 ++++++++++++++++++++++++- native/mtproto-rs/src/client/persist.rs | 2 ++ native/mtproto-rs/src/session_file.rs | 4 ++++ 3 files changed, 34 insertions(+), 1 deletion(-) diff --git a/native/mtproto-rs/src/client/mod.rs b/native/mtproto-rs/src/client/mod.rs index 04206858..6ed351fb 100644 --- a/native/mtproto-rs/src/client/mod.rs +++ b/native/mtproto-rs/src/client/mod.rs @@ -97,6 +97,10 @@ pub(crate) struct ClientData { pub(crate) home_time_offset: i64, pub(crate) test_dc: bool, pub(crate) last_inline: Option, + /// Master auth keys for foreign datacenters cached across session runs. + pub(crate) dc_auth_keys: HashMap>, + /// Ephemeral bound temporary auth keys for datacenters (dc_id -> DcTempAuth). + pub(crate) dc_temp_auth: HashMap, pub(crate) persist_epoch: u64, pub(crate) persisted_epoch: u64, /// In-memory catch-up edge. Not written into the session file. @@ -108,6 +112,13 @@ pub(crate) struct ClientData { pub(crate) lazy_retry_at: u64, } +#[derive(Clone, Debug)] +pub(crate) struct DcTempAuth { + pub(crate) snapshot: Snapshot, + pub(crate) temp_expires_at: i32, + pub(crate) perm_key: Vec, +} + #[derive(Clone)] pub(crate) struct LastInlineQuery { pub(crate) chat_id: i64, @@ -163,6 +174,8 @@ pub(crate) struct ClientState { pub(crate) perm_auth_key: Option>, pub(crate) perm_salt: i64, pub(crate) perm_session_id: i64, + /// Master auth keys for foreign datacenters. + pub(crate) dc_auth_keys: HashMap>, } pub(crate) fn prefer_newer_cursor( @@ -230,6 +243,10 @@ pub fn create_client(api_id: i32, api_hash: String, session_path: String) -> u64 .as_ref() .map(|s| s.channel_recovery.clone()) .unwrap_or_default(); + let dc_auth_keys = loaded + .as_ref() + .map(|s| s.dc_auth_keys.clone()) + .unwrap_or_default(); let ( snapshot, user_id, @@ -325,6 +342,8 @@ pub fn create_client(api_id: i32, api_hash: String, session_path: String) -> u64 home_time_offset: snapshot.time_offset_micros, test_dc, last_inline: None, + dc_auth_keys, + dc_temp_auth: HashMap::new(), persist_epoch: 0, persisted_epoch: 0, is_syncing: false, @@ -536,6 +555,7 @@ pub(crate) fn with_client_mut( .flatten(), perm_salt: d.home_salt, perm_session_id: d.home_session_id, + dc_auth_keys: d.dc_auth_keys.clone(), } }; crate::rpc::set_use_test_dc(state.test_dc); @@ -548,6 +568,7 @@ pub(crate) fn with_client_mut( let before_new_session = state.new_session.clone(); let before_dead = state.session_dead; let before_test_dc = state.test_dc; + let before_dc_auth_keys = state.dc_auth_keys.clone(); let mut slot = io.transport.take(); let before_snapshot = io.snapshot.clone(); let result = with_client_transport(&client, &mut slot, || f(&mut state)); @@ -583,10 +604,14 @@ pub(crate) fn with_client_mut( let session_meta_changed = before_new_session != state.new_session || before_dead != state.session_dead || before_test_dc != state.test_dc; + let dc_keys_changed = before_dc_auth_keys != state.dc_auth_keys; let persist_needed; { let mut d = client.data.lock(); d.user_id = state.user_id; + if dc_keys_changed { + d.dc_auth_keys = state.dc_auth_keys; + } let maps_changed = if identity_changed { d.peers = state.peers; d.media = state.media; @@ -628,7 +653,8 @@ pub(crate) fn with_client_mut( || snapshot_changed || maps_changed || tokens_changed - || session_meta_changed); + || session_meta_changed + || dc_keys_changed); if persist_needed { d.persist_epoch = d.persist_epoch.saturating_add(1); } @@ -734,6 +760,7 @@ pub(crate) fn with_read_lane( .flatten(), perm_salt: d.home_salt, perm_session_id: d.home_session_id, + dc_auth_keys: d.dc_auth_keys.clone(), }) } }; diff --git a/native/mtproto-rs/src/client/persist.rs b/native/mtproto-rs/src/client/persist.rs index 85105491..faf27b83 100644 --- a/native/mtproto-rs/src/client/persist.rs +++ b/native/mtproto-rs/src/client/persist.rs @@ -79,6 +79,7 @@ pub(crate) fn persist(state: &ClientState) -> Result<(), MtprotoError> { new_session: state.new_session.clone(), session_dead: state.session_dead, test_dc: state.test_dc, + dc_auth_keys: state.dc_auth_keys.clone(), }; let clock = persist_generation(&state.session_path); let _write = PERSIST_WRITE.lock(); @@ -128,6 +129,7 @@ pub(crate) fn persist_updates_data(client: &Client, session_id: i64) -> Result<( new_session: d.new_session.clone(), session_dead: d.session_dead, test_dc: d.test_dc, + dc_auth_keys: d.dc_auth_keys.clone(), }; let path = d.session_path.clone(); let epoch = d.persist_epoch; diff --git a/native/mtproto-rs/src/session_file.rs b/native/mtproto-rs/src/session_file.rs index c5026887..85b9c069 100644 --- a/native/mtproto-rs/src/session_file.rs +++ b/native/mtproto-rs/src/session_file.rs @@ -53,6 +53,9 @@ pub struct ClientSession { /// Connected to Telegram test DCs (`99966XYYYY`). #[serde(default)] pub test_dc: bool, + /// Persistent master auth keys for foreign datacenters (dc_id -> 256-byte auth key). + #[serde(default)] + pub dc_auth_keys: HashMap>, } pub fn media_key(chat_id: i64, message_id: i32) -> String { @@ -131,6 +134,7 @@ impl FileSessionStore { logout_tokens: Vec::new(), new_session: None, test_dc: false, + dc_auth_keys: HashMap::new(), }; if !encrypted && key.is_some() { self.save(&session)?; From e4b9a4e0eaca420874ce4e68ba8e9ce81f13e677 Mon Sep 17 00:00:00 2001 From: TheLastFlame Date: Sat, 10 Oct 2026 13:55:43 +1000 Subject: [PATCH 2/5] feat(mtproto): cache imported dc authorizations across sessions --- .../mtproto-rs/src/client/media_download.rs | 2 +- native/mtproto-rs/src/client/session.rs | 25 +++++++++++++------ 2 files changed, 18 insertions(+), 9 deletions(-) diff --git a/native/mtproto-rs/src/client/media_download.rs b/native/mtproto-rs/src/client/media_download.rs index 5ab45977..30c12971 100644 --- a/native/mtproto-rs/src/client/media_download.rs +++ b/native/mtproto-rs/src/client/media_download.rs @@ -1027,7 +1027,7 @@ pub(crate) fn download_media_range_on_lane_dc( if !home_lease_ok(state, home.session_id) { return Err(expired_session_lease()); } - copy_authorization_to_dc(&mut state.snapshot, api_id, target_dc) + copy_authorization_to_dc(state, target_dc) })?; io.transport = None; } diff --git a/native/mtproto-rs/src/client/session.rs b/native/mtproto-rs/src/client/session.rs index d6a8c1fd..32a4b611 100644 --- a/native/mtproto-rs/src/client/session.rs +++ b/native/mtproto-rs/src/client/session.rs @@ -80,21 +80,27 @@ pub(crate) fn ensure_auth_key(state: &mut ClientState) -> Result<(), MtprotoErro } pub(crate) fn copy_authorization_to_dc( - home: &mut Snapshot, - api_id: i32, + state: &mut ClientState, target_dc: i32, ) -> Result { use tellers_mtproto::latest::api::{ AuthAuthorization, AuthExportAuthorizationRequest, AuthExportedAuthorization, AuthImportAuthorizationRequest, }; + if let Some(cached_key) = state.dc_auth_keys.get(&target_dc) { + let mut dest = fork_session(&state.snapshot); + dest.dc_id = target_dc; + dest.auth_key = Some(cached_key.clone()); + dest.server_salt = 0; + return Ok(dest); + } let exported: AuthExportedAuthorization = crate::api_invoke::invoke_api( - home, - api_id, + &mut state.snapshot, + state.api_id, AuthExportAuthorizationRequest { dc_id: target_dc }, )?; let AuthExportedAuthorization::AuthExportedAuthorization(exp) = exported; - let mut dest = fork_session(home); + let mut dest = fork_session(&state.snapshot); dest.dc_id = target_dc; dest.auth_key = None; dest.server_salt = 0; @@ -102,7 +108,7 @@ pub(crate) fn copy_authorization_to_dc( ensure_auth_key_on(&mut dest)?; let _: AuthAuthorization = crate::api_invoke::invoke_api( &mut dest, - api_id, + state.api_id, AuthImportAuthorizationRequest { id: exp.id, bytes: exp.bytes, @@ -110,6 +116,10 @@ pub(crate) fn copy_authorization_to_dc( )?; // The destination snapshot leaves this lane; never retain its TCP here. crate::rpc::drop_live_transport(); + if let Some(imported_key) = dest.auth_key.clone() { + state.dc_auth_keys.insert(target_dc, imported_key); + let _ = persist(state); + } Ok(dest) } @@ -199,8 +209,7 @@ pub(crate) fn ensure_ready_after_migrate( // exportAuthorization/importAuthorization. Dropping the // home auth key here turns a normal USER/NETWORK_MIGRATE // into an unnecessary logout and breaks all update lanes. - let mut home = state.snapshot.clone(); - let destination = copy_authorization_to_dc(&mut home, state.api_id, dc)?; + let destination = copy_authorization_to_dc(state, dc)?; state.snapshot = destination; // Authorization moved, but these are still the same // account's applied update sequences. From 4dbc40184b880848e178b776a9a1c3e5169726b1 Mon Sep 17 00:00:00 2001 From: TheLastFlame Date: Sat, 10 Oct 2026 14:03:47 +1000 Subject: [PATCH 3/5] perf(mtproto): cache temp pfs keys globally and reuse across media lane preparations --- .../mtproto-rs/src/client/media_download.rs | 84 ++++++++++++++----- .../tests/unit/client_mgr_channel_tests.rs | 1 + .../tests/unit/client_mgr_logout_tests.rs | 1 + native/mtproto-rs/tests/unit/lib_tests.rs | 1 + 4 files changed, 67 insertions(+), 20 deletions(-) diff --git a/native/mtproto-rs/src/client/media_download.rs b/native/mtproto-rs/src/client/media_download.rs index 30c12971..63acb417 100644 --- a/native/mtproto-rs/src/client/media_download.rs +++ b/native/mtproto-rs/src/client/media_download.rs @@ -1021,27 +1021,49 @@ pub(crate) fn download_media_range_on_lane_dc( && io.snapshot.dc_id == target_dc && io.snapshot.auth_key.as_ref() != home.auth_key.as_ref(); if !already_temp { - match prepare_media_lane_snapshot(&mut io.snapshot, home, target_dc) { - MediaLanePrep::NeedExport => { - io.snapshot = with_client_mut(handle, |state| { - if !home_lease_ok(state, home.session_id) { - return Err(expired_session_lease()); - } - copy_authorization_to_dc(state, target_dc) - })?; - io.transport = None; + let cached_temp = { + let d = client.data.lock(); + d.dc_temp_auth + .get(&target_dc) + .filter(|t| t.temp_expires_at > now) + .cloned() + }; + if let Some(cached) = cached_temp { + io.snapshot = fork_session(&cached.snapshot); + io.temp_expires_at = cached.temp_expires_at; + io.perm_key = Some(cached.perm_key); + io.transport = None; + } else { + match prepare_media_lane_snapshot(&mut io.snapshot, home, target_dc) { + MediaLanePrep::NeedExport => { + io.snapshot = with_client_mut(handle, |state| { + if !home_lease_ok(state, home.session_id) { + return Err(expired_session_lease()); + } + copy_authorization_to_dc(state, target_dc) + })?; + io.transport = None; + } + MediaLanePrep::Replaced => io.transport = None, + MediaLanePrep::Reuse => {} } - MediaLanePrep::Replaced => io.transport = None, - MediaLanePrep::Reuse => {} - } - if let Some(key) = io.snapshot.auth_key.clone() { - io.perm_key = Some(key.clone()); - let offset = io.snapshot.time_offset_micros; - let dc = io.snapshot.dc_id; - if let Ok(bound) = crate::auth::bind_temp::negotiate_bound_temp(&key, dc, offset) { - io.snapshot = bound.snapshot; - io.temp_expires_at = bound.expires_at; - io.transport = None; + if let Some(key) = io.snapshot.auth_key.clone() { + io.perm_key = Some(key.clone()); + let offset = io.snapshot.time_offset_micros; + let dc = io.snapshot.dc_id; + if let Ok(bound) = crate::auth::bind_temp::negotiate_bound_temp(&key, dc, offset) { + client.data.lock().dc_temp_auth.insert( + target_dc, + DcTempAuth { + snapshot: bound.snapshot.clone(), + temp_expires_at: bound.expires_at, + perm_key: key, + }, + ); + io.snapshot = bound.snapshot; + io.temp_expires_at = bound.expires_at; + io.transport = None; + } } } } @@ -1091,9 +1113,31 @@ pub(crate) fn download_media_range_on_lane_dc( let Some(mut second) = try_lock_second_media_lane(&worker_client) else { return false; }; if !session_lease_valid(&worker_client.data.lock(), expected_session) { return false; } if download_temp_ready(&second, target_dc, home_key.as_deref()) { return true; } + let cached_temp = { + let d = worker_client.data.lock(); + d.dc_temp_auth + .get(&target_dc) + .filter(|t| t.temp_expires_at > crate::auth::bind_temp::unix_now()) + .cloned() + }; + if let Some(cached) = cached_temp { + second.perm_key = Some(cached.perm_key); + second.snapshot = fork_session(&cached.snapshot); + second.temp_expires_at = cached.temp_expires_at; + second.transport = None; + return true; + } let Ok(bound) = crate::auth::bind_temp::negotiate_bound_temp(&key, target_dc, time_offset) else { return false; }; if crate::request_control::check().is_err() || !session_lease_valid(&worker_client.data.lock(), expected_session) { return false; } + worker_client.data.lock().dc_temp_auth.insert( + target_dc, + DcTempAuth { + snapshot: bound.snapshot.clone(), + temp_expires_at: bound.expires_at, + perm_key: key.clone(), + }, + ); second.perm_key = Some(key); second.snapshot = bound.snapshot; second.temp_expires_at = bound.expires_at; diff --git a/native/mtproto-rs/tests/unit/client_mgr_channel_tests.rs b/native/mtproto-rs/tests/unit/client_mgr_channel_tests.rs index 98b9f444..1ad0eeef 100644 --- a/native/mtproto-rs/tests/unit/client_mgr_channel_tests.rs +++ b/native/mtproto-rs/tests/unit/client_mgr_channel_tests.rs @@ -82,6 +82,7 @@ fn begin_updates_preserves_the_unapplied_cursor_in_persistence() { perm_auth_key: None, perm_salt: 0, perm_session_id: 0, + dc_auth_keys: HashMap::new(), }; begin_updates(&mut state); diff --git a/native/mtproto-rs/tests/unit/client_mgr_logout_tests.rs b/native/mtproto-rs/tests/unit/client_mgr_logout_tests.rs index 8c995f2c..9e66997e 100644 --- a/native/mtproto-rs/tests/unit/client_mgr_logout_tests.rs +++ b/native/mtproto-rs/tests/unit/client_mgr_logout_tests.rs @@ -31,6 +31,7 @@ fn logout_clears_unauthorized_local_session_without_network() { perm_auth_key: None, perm_salt: 0, perm_session_id: 0, + dc_auth_keys: HashMap::new(), }; persist(&state).expect("seed session"); assert!(path.exists()); diff --git a/native/mtproto-rs/tests/unit/lib_tests.rs b/native/mtproto-rs/tests/unit/lib_tests.rs index c051ff5f..26302fca 100644 --- a/native/mtproto-rs/tests/unit/lib_tests.rs +++ b/native/mtproto-rs/tests/unit/lib_tests.rs @@ -33,6 +33,7 @@ fn session_file_persists_user_id() { new_session: None, session_dead: true, test_dc: false, + dc_auth_keys: Default::default(), }) .expect("save"); let raw = std::fs::read(&path).expect("read compact"); From 132052c1a4a174acbfddb870d726f48746b2ca2b Mon Sep 17 00:00:00 2001 From: TheLastFlame Date: Sat, 10 Oct 2026 20:08:29 +1000 Subject: [PATCH 4/5] feat(media): implement resumable downloads matching Nekogram .temp.parts format and uncap large download queue --- native/mtproto-rs/src/ffi.rs | 4 +- native/mtproto-rs/src/media/download.rs | 182 ++++++++++++++---- native/mtproto-rs/src/transfer_policy.rs | 2 +- .../tests/unit/media_download_tests.rs | 46 ++++- 4 files changed, 194 insertions(+), 40 deletions(-) diff --git a/native/mtproto-rs/src/ffi.rs b/native/mtproto-rs/src/ffi.rs index 1f878c67..e9611ed9 100644 --- a/native/mtproto-rs/src/ffi.rs +++ b/native/mtproto-rs/src/ffi.rs @@ -101,8 +101,10 @@ pub fn set_dispatch_class(class: i32) { /// "Faster downloads" setting: media lanes (1..=8) and parts in flight (1..=16). #[uniffi::export] pub fn set_download_concurrency(lanes: i32, parts: i32) { - scheduler::set_active_media_lanes(lanes.max(1) as usize); + let lanes_cnt = lanes.max(1) as usize; + scheduler::set_active_media_lanes(lanes_cnt); client::set_pipeline_parts(parts.max(1) as usize); + *crate::transfer_policy::current().queue_limits.lock() = (lanes_cnt, lanes_cnt); } #[uniffi::export] diff --git a/native/mtproto-rs/src/media/download.rs b/native/mtproto-rs/src/media/download.rs index 5199da9c..e14b89a8 100644 --- a/native/mtproto-rs/src/media/download.rs +++ b/native/mtproto-rs/src/media/download.rs @@ -576,55 +576,109 @@ pub(crate) fn download_cancelled(dest_path: &Path) -> bool { cancel_marker(dest_path).exists() } -/// The caller publishes under its session lease; dropping any failed or stale -/// download removes only its own staging file, never a previously cached file. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct ByteRange { + pub start: i64, + pub end: i64, +} + +#[derive(Debug, Clone, Default)] +pub(crate) struct FilePartsMeta { + pub not_loaded: Vec, +} + +impl FilePartsMeta { + pub(crate) fn load(path: &Path) -> Option { + let bytes = fs::read(path).ok()?; + if bytes.len() < 4 { + return None; + } + let count = i32::from_be_bytes(bytes[0..4].try_into().ok()?) as usize; + if bytes.len() != 4 + count * 16 { + return None; + } + let mut not_loaded = Vec::with_capacity(count); + for i in 0..count { + let offset = 4 + i * 16; + let start = i64::from_be_bytes(bytes[offset..offset + 8].try_into().ok()?); + let end = i64::from_be_bytes(bytes[offset + 8..offset + 16].try_into().ok()?); + if start < 0 || end < start { + return None; + } + not_loaded.push(ByteRange { start, end }); + } + Some(Self { not_loaded }) + } + + pub(crate) fn save(&self, path: &Path) -> std::io::Result<()> { + let count = self.not_loaded.len() as i32; + let mut buf = Vec::with_capacity(4 + self.not_loaded.len() * 16); + buf.extend_from_slice(&count.to_be_bytes()); + for r in &self.not_loaded { + buf.extend_from_slice(&r.start.to_be_bytes()); + buf.extend_from_slice(&r.end.to_be_bytes()); + } + fs::write(path, buf) + } +} + +/// The caller publishes under its session lease; staging files (.temp and .temp.parts) +/// persist across retries/network drops and are only wiped on explicit cancellation. pub(crate) struct StagedDownload { path: PathBuf, + parts_path: PathBuf, + destination: PathBuf, } impl StagedDownload { pub(crate) fn new(destination: &Path) -> Result { - use std::sync::atomic::{AtomicU64, Ordering}; - static NEXT: AtomicU64 = AtomicU64::new(1); if let Some(parent) = destination.parent() { fs::create_dir_all(parent).map_err(|e| MtprotoError::Message(e.to_string()))?; } - loop { - let mut name = destination.as_os_str().to_os_string(); - name.push(format!( - ".{}.{}.part", - std::process::id(), - NEXT.fetch_add(1, Ordering::Relaxed) - )); - let path = PathBuf::from(name); - match fs::OpenOptions::new() - .write(true) - .create_new(true) - .open(&path) - { - Ok(_) => return Ok(Self { path }), - Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => continue, - Err(e) => return Err(MtprotoError::Message(e.to_string())), - } - } + let mut temp_name = destination.as_os_str().to_os_string(); + temp_name.push(".temp"); + let path = PathBuf::from(temp_name); + + let mut parts_name = destination.as_os_str().to_os_string(); + parts_name.push(".temp.parts"); + let parts_path = PathBuf::from(parts_name); + + Ok(Self { + path, + parts_path, + destination: destination.to_path_buf(), + }) } pub(crate) fn path(&self) -> &Path { &self.path } + pub(crate) fn parts_path(&self) -> &Path { + &self.parts_path + } + pub(crate) fn publish(self, destination: &Path) -> Result { if download_cancelled(destination) { + let _ = fs::remove_file(&self.path); + let _ = fs::remove_file(&self.parts_path); return Err(MtprotoError::Message("cancelled".into())); } + if destination.exists() { + let _ = fs::remove_file(destination); + } fs::rename(&self.path, destination).map_err(|e| MtprotoError::Message(e.to_string()))?; + let _ = fs::remove_file(&self.parts_path); Ok(destination.display().to_string()) } } impl Drop for StagedDownload { fn drop(&mut self) { - let _ = fs::remove_file(&self.path); + if download_cancelled(&self.destination) { + let _ = fs::remove_file(&self.path); + let _ = fs::remove_file(&self.parts_path); + } } } @@ -745,34 +799,70 @@ pub(crate) fn download_media_range_batched_streaming_capped( fs::create_dir_all(parent).map_err(|e| MtprotoError::Message(e.to_string()))?; } let result = (|| -> Result<(), MtprotoError> { - let mut out = - fs::File::create(dest_path).map_err(|e| MtprotoError::Message(e.to_string()))?; + let parts_path = PathBuf::from(format!("{}.temp.parts", cancellation_path.display())); + let mut out = fs::OpenOptions::new() + .read(true) + .write(true) + .create(true) + .open(dest_path) + .map_err(|e| MtprotoError::Message(e.to_string()))?; + let stream = offset.is_some(); + let mut initial_offset = offset.unwrap_or(0); + if !stream { + if let Some(meta) = FilePartsMeta::load(&parts_path) { + if let Some(first_missing) = meta.not_loaded.first() { + let file_len = out.metadata().map(|m| m.len() as i64).unwrap_or(0); + let safe_offset = first_missing.start.min(file_len); + if safe_offset > 0 { + let safe_aligned = safe_offset - (safe_offset % i64::from(DEFAULT_CHUNK)); + if safe_aligned > 0 { + initial_offset = safe_aligned; + } + } + } + } + } + + if initial_offset > 0 { + use std::io::{Seek, SeekFrom}; + let _ = out.set_len(initial_offset as u64); + let _ = out.seek(SeekFrom::Start(initial_offset as u64)); + } else if !stream { + use std::io::{Seek, SeekFrom}; + let _ = out.set_len(0); + let _ = out.seek(SeekFrom::Start(0)); + } + let window = crate::transfer_policy::current().floods.lock().window( dc_id, TransferClass::Download, std::time::Instant::now(), DownloadWindow { - chunk: resolve_chunk(offset), + chunk: resolve_chunk(Some(initial_offset).filter(|_| initial_offset > 0).or(offset)), in_flight: parts_in_flight.max(1), }, ); FORCED_CHUNK.with(|cell| cell.set(window.chunk)); let mut width = window.in_flight; - let stream = offset.is_some(); - let mut next = offset.unwrap_or(0); + let mut next = initial_offset; let mut retried_limit = false; let mut timeout_retries = 0u8; - let mut written_end = 0i64; - let mut downloaded_bytes = offset.unwrap_or(0); + let mut written_end = initial_offset; + let mut downloaded_bytes = initial_offset; let mut counted_offsets = HashSet::new(); let mut first_batch = true; + let mut last_saved_offset = initial_offset; let target_str = cancellation_path.display().to_string(); + + if initial_offset > 0 { + notify_progress(&target_str, downloaded_bytes, media.file_size.unwrap_or(0)); + } 'download: loop { width = crate::transfer_policy::current().floods.lock().window( dc_id, TransferClass::Download, std::time::Instant::now(), - DownloadWindow { chunk: resolve_chunk(offset), in_flight: parts_in_flight.max(1) }, + DownloadWindow { chunk: resolve_chunk(Some(next).filter(|_| next > 0).or(offset)), in_flight: parts_in_flight.max(1) }, ).in_flight.min(if retried_limit { width } else { parts_in_flight.max(1) }); let chunk = resolve_chunk(Some(next).filter(|_| next > 0).or(offset)); let batch = if stream { @@ -914,7 +1004,7 @@ pub(crate) fn download_media_range_batched_streaming_capped( timeout_retries += 1; next = first_missing_offset( &counted_offsets, - offset.unwrap_or(0), + initial_offset, next_to_request, chunk, ); @@ -998,7 +1088,7 @@ pub(crate) fn download_media_range_batched_streaming_capped( timeout_retries += 1; next = first_missing_offset( &counted_offsets, - offset.unwrap_or(0), + initial_offset, next_to_request, chunk, ); @@ -1015,10 +1105,27 @@ pub(crate) fn download_media_range_batched_streaming_capped( // Continue at the first hole instead of the cursor those refills advanced. next = first_missing_offset( &counted_offsets, - offset.unwrap_or(0), + initial_offset, next_to_request, chunk, ); + let contiguous = first_missing_offset( + &counted_offsets, + initial_offset, + written_end, + chunk, + ); + if contiguous > last_saved_offset { + let total = media.file_size.unwrap_or(0); + let meta = FilePartsMeta { + not_loaded: vec![ByteRange { + start: contiguous, + end: total, + }], + }; + let _ = meta.save(&parts_path); + last_saved_offset = contiguous; + } } } // 'download if !stream { @@ -1034,12 +1141,17 @@ pub(crate) fn download_media_range_batched_streaming_capped( .map_err(|e| MtprotoError::Message(e.to_string()))?; out.flush() .map_err(|e| MtprotoError::Message(e.to_string()))?; + let _ = fs::remove_file(&parts_path); Ok(()) })(); match result { Ok(()) => Ok(dest_path.display().to_string()), Err(err) => { - let _ = fs::remove_file(dest_path); + if download_cancelled(cancellation_path) { + let _ = fs::remove_file(dest_path); + let parts_path = PathBuf::from(format!("{}.temp.parts", cancellation_path.display())); + let _ = fs::remove_file(parts_path); + } Err(err) } } diff --git a/native/mtproto-rs/src/transfer_policy.rs b/native/mtproto-rs/src/transfer_policy.rs index be89ecd1..800673f6 100644 --- a/native/mtproto-rs/src/transfer_policy.rs +++ b/native/mtproto-rs/src/transfer_policy.rs @@ -99,7 +99,7 @@ impl TransferPolicy { return Err(crate::MtprotoError::Message(error.to_string())); } let limits = *self.queue_limits.lock(); - let limit = if large { limits.1.min(2) } else { limits.0 }; + let limit = if large { limits.1 } else { limits.0 }; let count = active.counts.get(&(dc, large)).copied().unwrap_or(0); let best = active .waiters diff --git a/native/mtproto-rs/tests/unit/media_download_tests.rs b/native/mtproto-rs/tests/unit/media_download_tests.rs index ea14d259..7428cc8b 100644 --- a/native/mtproto-rs/tests/unit/media_download_tests.rs +++ b/native/mtproto-rs/tests/unit/media_download_tests.rs @@ -16,16 +16,28 @@ use tellers_mtproto::latest::api::{ }; #[test] -fn stale_staging_preserves_existing_destination_and_removes_partial_bytes() { +fn stale_staging_preserves_partial_bytes_for_resuming_and_cleans_on_cancel() { let dest = std::env::temp_dir().join(format!("monogram-staging-{}", std::process::id())); fs::write(&dest, b"previous").unwrap(); let staged = StagedDownload::new(&dest).unwrap(); let path = staged.path().to_owned(); + let parts_path = staged.parts_path().to_owned(); fs::write(&path, b"uncommitted").unwrap(); + fs::write(&parts_path, b"meta").unwrap(); drop(staged); - assert!(!path.exists()); + // Preserved for resume on normal uncommitted drop + assert!(path.exists()); + assert!(parts_path.exists()); assert_eq!(fs::read(&dest).unwrap(), b"previous"); - fs::remove_file(dest).unwrap(); + + // Cleans up when explicitly cancelled + fs::write(cancel_marker(&dest), []).unwrap(); + let staged_cancel = StagedDownload::new(&dest).unwrap(); + drop(staged_cancel); + assert!(!path.exists()); + assert!(!parts_path.exists()); + let _ = fs::remove_file(cancel_marker(&dest)); + let _ = fs::remove_file(dest); } #[test] @@ -53,6 +65,34 @@ fn successful_staging_publishes_complete_file() { fs::remove_file(dest).unwrap(); } +#[test] +fn parts_meta_roundtrip_matches_nekogram_format() { + let parts_path = std::env::temp_dir().join(format!("test-parts-{}", std::process::id())); + let meta = FilePartsMeta { + not_loaded: vec![ + ByteRange { start: 1024, end: 4096 }, + ByteRange { start: 8192, end: 16384 }, + ], + }; + meta.save(&parts_path).unwrap(); + let loaded = FilePartsMeta::load(&parts_path).unwrap(); + assert_eq!(loaded.not_loaded.len(), 2); + assert_eq!(loaded.not_loaded[0], ByteRange { start: 1024, end: 4096 }); + assert_eq!(loaded.not_loaded[1], ByteRange { start: 8192, end: 16384 }); + + // Verify raw binary matches Nekogram big-endian layout + let raw = fs::read(&parts_path).unwrap(); + assert_eq!(raw.len(), 4 + 2 * 16); + let count = i32::from_be_bytes(raw[0..4].try_into().unwrap()); + assert_eq!(count, 2); + let start0 = i64::from_be_bytes(raw[4..12].try_into().unwrap()); + let end0 = i64::from_be_bytes(raw[12..20].try_into().unwrap()); + assert_eq!(start0, 1024); + assert_eq!(end0, 4096); + + let _ = fs::remove_file(parts_path); +} + #[test] fn chunk_is_telegram_max_part() { assert_eq!(CHUNK, 128 * 1024); From 04b9500b5106bd4caefd80b34bcfd8303a87f702 Mon Sep 17 00:00:00 2001 From: TheLastFlame Date: Sat, 10 Oct 2026 20:29:49 +1000 Subject: [PATCH 5/5] fix(media): write parts metadata continuously during download and preserve partial files across cancel/drop --- native/mtproto-rs/src/media/download.rs | 59 +++++++++++++++---- .../tests/unit/media_download_tests.rs | 21 ++++--- 2 files changed, 62 insertions(+), 18 deletions(-) diff --git a/native/mtproto-rs/src/media/download.rs b/native/mtproto-rs/src/media/download.rs index e14b89a8..6bc5c07f 100644 --- a/native/mtproto-rs/src/media/download.rs +++ b/native/mtproto-rs/src/media/download.rs @@ -660,8 +660,6 @@ impl StagedDownload { pub(crate) fn publish(self, destination: &Path) -> Result { if download_cancelled(destination) { - let _ = fs::remove_file(&self.path); - let _ = fs::remove_file(&self.parts_path); return Err(MtprotoError::Message("cancelled".into())); } if destination.exists() { @@ -675,10 +673,9 @@ impl StagedDownload { impl Drop for StagedDownload { fn drop(&mut self) { - if download_cancelled(&self.destination) { - let _ = fs::remove_file(&self.path); - let _ = fs::remove_file(&self.parts_path); - } + // Incomplete downloads (.temp and .temp.parts) are preserved on disk + // so that pausing, stopping, network drops, and app restarts resume from + // the last contiguous byte without restarting from 0. } } @@ -820,6 +817,14 @@ pub(crate) fn download_media_range_batched_streaming_capped( } } } + } else { + let file_len = out.metadata().map(|m| m.len() as i64).unwrap_or(0); + if file_len > 0 { + let safe_aligned = file_len - (file_len % i64::from(DEFAULT_CHUNK)); + if safe_aligned > 0 { + initial_offset = safe_aligned; + } + } } } @@ -939,6 +944,24 @@ pub(crate) fn download_media_range_batched_streaming_capped( stop_refill = true; } notify_progress(&target_str, downloaded_bytes, 0); + + let contiguous = first_missing_offset( + &counted_offsets, + initial_offset, + written_end, + chunk, + ); + if contiguous > last_saved_offset { + let total = media.file_size.unwrap_or(0); + let meta = FilePartsMeta { + not_loaded: vec![ByteRange { + start: contiguous, + end: total, + }], + }; + let _ = meta.save(&parts_path); + last_saved_offset = contiguous; + } } if index < written_parts.len() { written_parts[index] = true; @@ -1056,6 +1079,24 @@ pub(crate) fn download_media_range_batched_streaming_capped( if counted_offsets.insert(part_offset) { downloaded_bytes += bytes.len() as i64; notify_progress(&target_str, downloaded_bytes, 0); + + let contiguous = first_missing_offset( + &counted_offsets, + initial_offset, + written_end, + chunk, + ); + if contiguous > last_saved_offset { + let total = media.file_size.unwrap_or(0); + let meta = FilePartsMeta { + not_loaded: vec![ByteRange { + start: contiguous, + end: total, + }], + }; + let _ = meta.save(&parts_path); + last_saved_offset = contiguous; + } } if bytes.len() < chunk as usize || stream { last_part = true; @@ -1147,11 +1188,7 @@ pub(crate) fn download_media_range_batched_streaming_capped( match result { Ok(()) => Ok(dest_path.display().to_string()), Err(err) => { - if download_cancelled(cancellation_path) { - let _ = fs::remove_file(dest_path); - let parts_path = PathBuf::from(format!("{}.temp.parts", cancellation_path.display())); - let _ = fs::remove_file(parts_path); - } + // Partial downloads (.temp and .temp.parts) are preserved on disk for resume. Err(err) } } diff --git a/native/mtproto-rs/tests/unit/media_download_tests.rs b/native/mtproto-rs/tests/unit/media_download_tests.rs index 7428cc8b..2683ffc4 100644 --- a/native/mtproto-rs/tests/unit/media_download_tests.rs +++ b/native/mtproto-rs/tests/unit/media_download_tests.rs @@ -16,7 +16,7 @@ use tellers_mtproto::latest::api::{ }; #[test] -fn stale_staging_preserves_partial_bytes_for_resuming_and_cleans_on_cancel() { +fn stale_staging_preserves_partial_bytes_for_resuming_across_cancellations() { let dest = std::env::temp_dir().join(format!("monogram-staging-{}", std::process::id())); fs::write(&dest, b"previous").unwrap(); let staged = StagedDownload::new(&dest).unwrap(); @@ -30,27 +30,34 @@ fn stale_staging_preserves_partial_bytes_for_resuming_and_cleans_on_cancel() { assert!(parts_path.exists()); assert_eq!(fs::read(&dest).unwrap(), b"previous"); - // Cleans up when explicitly cancelled + // Preserved across cancellations as well so downloads can resume fs::write(cancel_marker(&dest), []).unwrap(); let staged_cancel = StagedDownload::new(&dest).unwrap(); drop(staged_cancel); - assert!(!path.exists()); - assert!(!parts_path.exists()); + assert!(path.exists()); + assert!(parts_path.exists()); let _ = fs::remove_file(cancel_marker(&dest)); + let _ = fs::remove_file(&path); + let _ = fs::remove_file(&parts_path); let _ = fs::remove_file(dest); } #[test] -fn cancel_before_publication_cleans_staging_without_publishing() { +fn cancel_before_publication_leaves_staging_without_publishing() { let dest = std::env::temp_dir().join(format!("monogram-publish-cancel-{}", std::process::id())); let staged = StagedDownload::new(&dest).unwrap(); let path = staged.path().to_owned(); + let parts_path = staged.parts_path().to_owned(); fs::write(&path, b"complete").unwrap(); + fs::write(&parts_path, b"parts").unwrap(); fs::write(cancel_marker(&dest), []).unwrap(); assert!(staged.publish(&dest).is_err()); - assert!(!path.exists()); + assert!(path.exists()); + assert!(parts_path.exists()); assert!(!dest.exists()); - fs::remove_file(cancel_marker(&dest)).unwrap(); + let _ = fs::remove_file(cancel_marker(&dest)); + let _ = fs::remove_file(&path); + let _ = fs::remove_file(&parts_path); } #[test]