feat(mtproto): update layer to 230, update proxy logic and tests

Author: gdlbo

native/Cargo.lock

@@ -19,6 +19,10 @@ mod peer_avatar;
 
 use peer_avatar::reindex_peer_avatar;
 
+#[cfg(test)]
+#[path = "../../tests/unit/cdn_transfer_tests.rs"]
+mod cdn_transfer_tests;
+
 include!("media_lane.rs");
 
 
@@ -984,19 +988,15 @@ pub(crate) fn download_media_range_on_lane(
             Err(err) if media::is_cdn_redirect(&err) => {
                 let job = media::take_cdn_redirect().ok_or(err)?;
                 drop(admission);
-                let cdn_result = {
+                cdn_or_origin(|| {
                     let _cdn_admission = client.policy.admit_download(job.dc_id, large)?;
                     follow_cdn_redirect(handle, api_id, dest, &job)
-                };
-                match cdn_result {
-                    Ok(path) => Ok(path),
-                    Err(_) => media::with_cdn_supported(false, || {
-                        let _admission = client.policy.admit_download(target_dc, large)?;
-                        download_media_range_on_lane_dc(
-                            handle, client, home, api_id, media, dest, offset, target_dc,
-                        )
-                    }),
-                }
+                }, || {
+                    let _admission = client.policy.admit_download(target_dc, large)?;
+                    download_media_range_on_lane_dc(
+                        handle, client, home, api_id, media, dest, offset, target_dc,
+                    )
+                })
             }
             other => other,
         };

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

@@ -464,6 +464,11 @@ fn follow_cdn_redirect(
     dest: &Path,
     job: &media::CdnRedirect,
 ) -> Result<String, MtprotoError> {
+    if job.offset != 0 {
+        return Err(MtprotoError::Message("CDN resume requires origin download".into()));
+    }
+    let client = get_client(handle)?;
+    let session_id = client.data.lock().home_session_id;
     let pem = with_client_mut(handle, |state| {
         let cfg: CdnConfig =
             api_invoke::invoke_api(&mut state.snapshot, api_id, media::get_cdn_config_request())?;
@@ -503,63 +508,85 @@ fn follow_cdn_redirect(
         last_io: std::time::Instant::now(),
     });
     let staged = media::StagedDownload::new(dest)?;
-    let path = crate::rpc::with_live_transport(&mut slot, || {
+    crate::rpc::with_live_transport(&mut slot, || {
         crate::rpc::with_rpc_timeout_secs(45, || {
-            use std::io::{Seek, SeekFrom, Write};
-            let mut out = std::fs::File::create(staged.path())
-                .map_err(|e| MtprotoError::Message(e.to_string()))?;
-            let mut offset = job.offset;
-            loop {
-                if media::download_cancelled(dest) || crate::request_control::check().is_err() {
-                    return Err(MtprotoError::Message("cancelled".into()));
-                }
-                let mut part = loop {
-                    let response: UploadCdnFile = api_invoke::invoke_api_without_updates(
-                        &mut snap,
-                        api_id,
-                        media::get_cdn_file_request(job.file_token.clone(), offset, media::CHUNK),
-                    )?;
-                    match media::cdn_file_bytes(response)? {
-                        Ok(bytes) => break bytes,
-                        Err(request_token) => {
-                            let _ = with_client_mut(handle, |state| {
-                                api_invoke::invoke_api::<_, Vector<Box<FileHash>>>(
-                                    &mut state.snapshot,
-                                    api_id,
-                                    media::reupload_cdn_file_request(
-                                        job.file_token.clone(),
-                                        request_token,
-                                    ),
-                                )
-                            })?;
-                        }
-                    }
-                };
-                media::decrypt_cdn_part(
-                    &job.encryption_key,
-                    &job.encryption_iv,
-                    offset,
-                    &mut part,
+            receive_cdn_parts(job, dest, &staged, |offset| {
+                let response: UploadCdnFile = api_invoke::invoke_api_without_updates(
+                    &mut snap, api_id,
+                    media::get_cdn_file_request(job.file_token.clone(), offset, media::CHUNK),
                 )?;
-                out.seek(SeekFrom::Start(offset as u64))
-                    .map_err(|e| MtprotoError::Message(e.to_string()))?;
-                out.write_all(&part)
-                    .map_err(|e| MtprotoError::Message(e.to_string()))?;
-                if (part.len() as i32) < media::CHUNK {
-                    break;
-                }
-                offset += part.len() as i64;
-                media::notify_progress(&dest.display().to_string(), offset, 0);
-            }
-            out.flush()
-                .map_err(|e| MtprotoError::Message(e.to_string()))?;
-            Ok(staged.path().display().to_string())
+                media::cdn_file_bytes(response)
+            }, |request_token| {
+                with_client_mut(handle, |state| {
+                    api_invoke::invoke_api::<_, Vector<Box<FileHash>>>(
+                        &mut state.snapshot, api_id,
+                        media::reupload_cdn_file_request(job.file_token.clone(), request_token),
+                    )
+                }).map(|_| ())
+            })
         })
     })?;
-    let client = get_client(handle)?;
-    let session_id = client.data.lock().home_session_id;
     publish_media(&client, session_id, staged, dest)?;
-    let _ = path;
     Ok(dest.display().to_string())
 }
 
+fn cdn_or_origin<T>(
+    cdn: impl FnOnce() -> Result<T, MtprotoError>,
+    origin: impl FnOnce() -> Result<T, MtprotoError>,
+) -> Result<T, MtprotoError> {
+    match cdn() {
+        Ok(result) => Ok(result),
+        Err(_) => media::with_cdn_supported(false, origin),
+    }
+}
+
+fn receive_cdn_parts(
+    job: &media::CdnRedirect,
+    dest: &Path,
+    staged: &media::StagedDownload,
+    mut fetch: impl FnMut(i64) -> Result<Result<Vec<u8>, Vec<u8>>, MtprotoError>,
+    mut reupload: impl FnMut(Vec<u8>) -> Result<(), MtprotoError>,
+) -> Result<(), MtprotoError> {
+    use std::io::{Seek, SeekFrom, Write};
+    let mut out = std::fs::File::create(staged.path())
+        .map_err(|e| MtprotoError::Message(e.to_string()))?;
+    let mut offset = job.offset;
+    loop {
+        if media::download_cancelled(dest) || crate::request_control::check().is_err() {
+            return Err(MtprotoError::Message("cancelled".into()));
+        }
+        let mut reuploads = 0;
+        let mut part = loop {
+            if media::download_cancelled(dest) || crate::request_control::check().is_err() {
+                return Err(MtprotoError::Message("cancelled".into()));
+            }
+            match fetch(offset)? {
+                Ok(bytes) => break bytes,
+                Err(token) => {
+                    if reuploads >= 3 {
+                        return Err(MtprotoError::Message("CDN reupload limit exceeded".into()));
+                    }
+                    reuploads += 1;
+                    reupload(token)?;
+                }
+            }
+        };
+        if media::download_cancelled(dest) || crate::request_control::check().is_err() {
+            return Err(MtprotoError::Message("cancelled".into()));
+        }
+        media::decrypt_cdn_part(&job.encryption_key, &job.encryption_iv, offset, &mut part)?;
+        if part.len() > media::CHUNK as usize
+            || !media::verify_cdn_part(&job.file_hashes, offset, &part)
+        {
+            return Err(MtprotoError::Message("CDN hash verification failed".into()));
+        }
+        out.seek(SeekFrom::Start(offset as u64))
+            .map_err(|e| MtprotoError::Message(e.to_string()))?;
+        out.write_all(&part).map_err(|e| MtprotoError::Message(e.to_string()))?;
+        if part.len() < media::CHUNK as usize { break; }
+        offset += part.len() as i64;
+        media::notify_progress(&dest.display().to_string(), offset, 0);
+    }
+    out.flush().map_err(|e| MtprotoError::Message(e.to_string()))
+}
+

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

@@ -72,6 +72,10 @@ pub(crate) fn append_rich(text: &RichText, out: &mut media::FormattedText) {
             append_rich(&t.text, out);
             mark(out, start, "code", None);
         }
+        RichText::TextTonAddress(t) => {
+            append_rich(&t.text, out);
+            mark(out, start, "code", None);
+        }
         RichText::TextBankCard(t) => {
             append_rich(&t.text, out);
             mark(out, start, "code", None);

native/mtproto-rs/src/instant_view/rich.rs

@@ -28,6 +28,7 @@ pub struct CdnRedirect {
     pub encryption_key: Vec<u8>,
     pub encryption_iv: Vec<u8>,
     pub offset: i64,
+    pub file_hashes: Vec<FileHash>,
 }
 
 pub fn cdn_supported() -> bool {
@@ -59,6 +60,9 @@ pub fn store_cdn_redirect(redirect: &UploadFileCdnRedirectConstructor, offset: i
             encryption_key: redirect.encryption_key.clone(),
             encryption_iv: redirect.encryption_iv.clone(),
             offset,
+            file_hashes: crate::peers::vector_boxed_items(redirect.file_hashes.as_ref())
+                .cloned()
+                .collect(),
         });
     });
 }
@@ -152,8 +156,29 @@ pub fn hash_matches(file_hash: &FileHash, offset: i64, data: &[u8]) -> bool {
     if hash.offset != offset {
         return false;
     }
-    let take = hash.limit.max(0) as usize;
-    verify_cdn_hash(&data[..data.len().min(take)], &hash.hash)
+    let Ok(take) = usize::try_from(hash.limit) else {
+        return false;
+    };
+    take > 0 && data.len() >= take && verify_cdn_hash(&data[..take], &hash.hash)
+}
+
+pub(crate) fn verify_cdn_part(hashes: &[FileHash], offset: i64, data: &[u8]) -> bool {
+    if offset < 0 || data.is_empty() {
+        return false;
+    }
+    let mut consumed = 0;
+    while consumed < data.len() {
+        let Some(at) = offset.checked_add(consumed as i64) else { return false };
+        let Some(hash) = hashes.iter().find(|entry| matches!(entry,
+            FileHash::FileHash(hash) if hash.offset == at
+        )) else { return false };
+        if !hash_matches(hash, at, &data[consumed..]) {
+            return false;
+        }
+        let FileHash::FileHash(hash) = hash else { return false };
+        consumed += hash.limit as usize;
+    }
+    true
 }
 
 #[cfg(test)]

native/mtproto-rs/src/media/cdn.rs

@@ -522,6 +522,10 @@ fn push_rich(text: &RichText, out: &mut FormattedText) {
             push_rich(&t.text, out);
             mark(out, start, "code", None);
         }
+        RichText::TextTonAddress(t) => {
+            push_rich(&t.text, out);
+            mark(out, start, "code", None);
+        }
         RichText::TextBankCard(t) => {
             push_rich(&t.text, out);
             mark(out, start, "code", None);

native/mtproto-rs/src/media/page_plain.rs

@@ -97,6 +97,10 @@ mod gzip_tests;
 #[cfg(test)]
 #[path = "../tests/unit/rpc_push_tests.rs"]
 mod push_tests;
+#[cfg(test)]
+#[path = "../tests/unit/rpc_wire_tests.rs"]
+mod wire_tests;
+
 #[cfg(test)]
 #[path = "../tests/unit/rpc_session_tests.rs"]
 mod session_tests;