Feat/resumable downloads and dc auth cache

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(&mut state.snapshot, api_id, 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;

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<LastInlineQuery>,
+    /// Master auth keys for foreign datacenters cached across session runs.
+    pub(crate) dc_auth_keys: HashMap<i32, Vec<u8>>,
+    /// Ephemeral bound temporary auth keys for datacenters (dc_id -> DcTempAuth).
+    pub(crate) dc_temp_auth: HashMap<i32, DcTempAuth>,
     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<u8>,
+}
+
 #[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<Vec<u8>>,
     pub(crate) perm_salt: i64,
     pub(crate) perm_session_id: i64,
+    /// Master auth keys for foreign datacenters.
+    pub(crate) dc_auth_keys: HashMap<i32, Vec<u8>>,
 }
 
 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<T>(
                 .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<T>(
     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<T>(
     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<T>(
                 || 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<T>(
                     .flatten(),
                 perm_salt: d.home_salt,
                 perm_session_id: d.home_session_id,
+                dc_auth_keys: d.dc_auth_keys.clone(),
             })
         }
     };

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;

native/mtproto-rs/src/client/session.rs

@@ -80,36 +80,46 @@ 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<Snapshot, MtprotoError> {
     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;
     crate::rpc::drop_live_transport();
     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,
         },
     )?;
     // 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.

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]