use crate::client::Client;
use crate::http::{
HTTP_STATUS_GONE, HTTP_STATUS_NOT_FOUND, HTTP_STATUS_OK, HTTP_STATUS_REDIRECTION_START,
HttpClient, HttpStatusError,
};
use crate::mediaconn::{MEDIA_AUTH_REFRESH_RETRY_ATTEMPTS, MediaConn, is_media_auth_error};
use anyhow::{Result, anyhow};
use std::sync::Arc;
use wacore::runtime::Runtime;
pub use wacore::download::{
DEFAULT_MEDIA_HOSTS, DownloadUtils, DownloadWriter, Downloadable, MediaDecryption,
MediaDecryptionError, MediaHost, MediaRoute, MediaType,
};
const DOWNLOAD_PREALLOC_CAP: u64 = 64 * 1024 * 1024;
impl From<&MediaConn> for MediaRoute {
fn from(conn: &MediaConn) -> Self {
MediaRoute::authenticated(
conn.hosts
.iter()
.map(|h| MediaHost::new(h.hostname.clone()))
.collect(),
conn.auth.clone(),
)
}
}
pub struct DownloadParams {
pub direct_path: String,
pub media_key: Option<Vec<u8>>,
pub file_sha256: Vec<u8>,
pub file_enc_sha256: Option<Vec<u8>>,
pub file_length: u64,
pub media_type: MediaType,
}
impl DownloadParams {
pub fn encrypted(
direct_path: impl Into<String>,
media_key: &[u8],
file_sha256: &[u8],
file_enc_sha256: &[u8],
file_length: u64,
media_type: MediaType,
) -> Self {
Self {
direct_path: direct_path.into(),
media_key: Some(media_key.to_vec()),
file_sha256: file_sha256.to_vec(),
file_enc_sha256: Some(file_enc_sha256.to_vec()),
file_length,
media_type,
}
}
}
impl Downloadable for DownloadParams {
fn direct_path(&self) -> Option<&str> {
Some(&self.direct_path)
}
fn media_key(&self) -> Option<&[u8]> {
self.media_key.as_deref()
}
fn file_enc_sha256(&self) -> Option<&[u8]> {
self.file_enc_sha256.as_deref()
}
fn file_sha256(&self) -> Option<&[u8]> {
Some(&self.file_sha256)
}
fn file_length(&self) -> Option<u64> {
Some(self.file_length)
}
fn app_info(&self) -> MediaType {
self.media_type
}
}
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum MediaDownloadError {
#[error("the CDN rejected the media reference: {0}")]
ReferenceRejected(#[source] anyhow::Error),
#[error("every media host failed: {0}")]
HostsUnreachable(#[source] anyhow::Error),
#[error("the media route names no hosts")]
NoHosts,
#[error("{0}")]
Other(#[from] anyhow::Error),
}
impl From<DownloadRequestError> for MediaDownloadError {
fn from(err: DownloadRequestError) -> Self {
match err {
DownloadRequestError::Auth(e) | DownloadRequestError::NotFound(e) => {
Self::ReferenceRejected(e)
}
DownloadRequestError::Other(e) => Self::HostsUnreachable(e),
DownloadRequestError::Prepare(e) => Self::Other(e),
DownloadRequestError::NoHosts => Self::NoHosts,
}
}
}
#[derive(Debug)]
enum DownloadRequestError {
Auth(anyhow::Error),
NotFound(anyhow::Error),
Other(anyhow::Error),
Prepare(anyhow::Error),
NoHosts,
}
impl DownloadRequestError {
fn auth(status_code: u16) -> Self {
Self::Auth(Self::refused(
status_code,
format!("Download failed with status: {status_code}"),
))
}
fn not_found(status_code: u16) -> Self {
Self::NotFound(Self::refused(
status_code,
format!("Download media not found/expired with status: {status_code}"),
))
}
fn refused_status(status_code: u16) -> Self {
Self::Other(Self::refused(
status_code,
format!("Download failed with status: {status_code}"),
))
}
fn refused(status_code: u16, context: String) -> anyhow::Error {
HttpStatusError {
status: status_code,
}
.into_error(context)
}
fn other(err: impl Into<anyhow::Error>) -> Self {
Self::Other(err.into())
}
fn is_auth(&self) -> bool {
matches!(self, Self::Auth(_))
}
fn is_not_found(&self) -> bool {
matches!(self, Self::NotFound(_))
}
fn into_anyhow(self) -> anyhow::Error {
match self {
Self::Auth(err) | Self::NotFound(err) | Self::Other(err) | Self::Prepare(err) => err,
Self::NoHosts => anyhow!("Failed to download from all available media hosts"),
}
}
}
fn validate_download_status(status_code: u16) -> std::result::Result<(), DownloadRequestError> {
if status_code < HTTP_STATUS_REDIRECTION_START {
return Ok(());
}
let error = if is_media_auth_error(status_code) {
DownloadRequestError::auth(status_code)
} else if matches!(status_code, HTTP_STATUS_NOT_FOUND | HTTP_STATUS_GONE) {
DownloadRequestError::not_found(status_code)
} else {
DownloadRequestError::refused_status(status_code)
};
Err(error)
}
fn decrypt_or_validate_buffered_body(
body: &mut Vec<u8>,
decryption: &MediaDecryption,
) -> std::result::Result<(), DownloadRequestError> {
match decryption {
MediaDecryption::Encrypted {
media_key,
media_type,
} => DownloadUtils::verify_and_decrypt_in_place(body, media_key, *media_type)
.map_err(DownloadRequestError::other),
MediaDecryption::Plaintext { file_sha256 } => {
DownloadUtils::validate_plaintext_sha256(body, file_sha256)
.map_err(DownloadRequestError::other)
}
}
}
async fn download_media_with_retry<
PrepareRequests,
PrepareRequestsFut,
InvalidateMediaConn,
InvalidateMediaConnFut,
ExecuteRequest,
ExecuteRequestFut,
>(
max_refresh_attempts: usize,
mut prepare_requests: PrepareRequests,
mut invalidate_media_conn: InvalidateMediaConn,
mut execute_request: ExecuteRequest,
) -> std::result::Result<Vec<u8>, DownloadRequestError>
where
PrepareRequests: FnMut(bool) -> PrepareRequestsFut,
PrepareRequestsFut: Future<Output = Result<Vec<wacore::download::DownloadRequest>>>,
InvalidateMediaConn: FnMut() -> InvalidateMediaConnFut,
InvalidateMediaConnFut: Future<Output = ()>,
ExecuteRequest: FnMut(wacore::download::DownloadRequest) -> ExecuteRequestFut,
ExecuteRequestFut: Future<Output = std::result::Result<Vec<u8>, DownloadRequestError>>,
{
let mut force_refresh = false;
let mut last_err: Option<anyhow::Error> = None;
for attempt in 0..=max_refresh_attempts {
let requests = prepare_requests(force_refresh)
.await
.map_err(DownloadRequestError::Prepare)?;
let mut retry_with_fresh_auth = false;
for request in requests {
match execute_request(request.clone()).await {
Ok(data) => return Ok(data),
Err(err)
if (err.is_auth() || err.is_not_found()) && attempt < max_refresh_attempts =>
{
invalidate_media_conn().await;
force_refresh = true;
retry_with_fresh_auth = true;
break;
}
Err(err) if err.is_auth() || err.is_not_found() => return Err(err),
Err(err) => {
let err = err.into_anyhow();
log::warn!(
"Failed to download from URL {}: {:?}. Trying next host.",
request.url,
err
);
last_err = Some(err);
}
}
}
if !retry_with_fresh_auth {
break;
}
}
match last_err {
Some(err) => Err(DownloadRequestError::Other(err)),
None => Err(DownloadRequestError::NoHosts),
}
}
async fn download_to_writer_with_retry<
W,
PrepareRequests,
PrepareRequestsFut,
InvalidateMediaConn,
InvalidateMediaConnFut,
ExecuteRequest,
ExecuteRequestFut,
>(
max_refresh_attempts: usize,
runtime: &Arc<dyn Runtime>,
mut writer: W,
mut prepare_requests: PrepareRequests,
mut invalidate_media_conn: InvalidateMediaConn,
mut execute_request: ExecuteRequest,
) -> std::result::Result<W, DownloadRequestError>
where
W: DownloadWriter + Send + 'static,
PrepareRequests: FnMut(bool) -> PrepareRequestsFut,
PrepareRequestsFut: Future<Output = Result<Vec<wacore::download::DownloadRequest>>>,
InvalidateMediaConn: FnMut() -> InvalidateMediaConnFut,
InvalidateMediaConnFut: Future<Output = ()>,
ExecuteRequest: FnMut(wacore::download::DownloadRequest, W) -> ExecuteRequestFut,
ExecuteRequestFut: Future<Output = Result<(W, std::result::Result<(), DownloadRequestError>)>>,
{
let mut force_refresh = false;
let mut last_err: Option<anyhow::Error> = None;
for attempt in 0..=max_refresh_attempts {
let requests = match prepare_requests(force_refresh).await {
Ok(requests) => requests,
Err(err) => {
discard_failed_write(runtime, writer).await;
return Err(DownloadRequestError::Prepare(err));
}
};
let mut retry_with_fresh_auth = false;
for request in requests {
let (next_writer, result) = match execute_request(request.clone(), writer).await {
Ok(outcome) => outcome,
Err(err) => return Err(DownloadRequestError::Other(err)),
};
writer = next_writer;
match result {
Ok(()) => return Ok(writer),
Err(err)
if (err.is_auth() || err.is_not_found()) && attempt < max_refresh_attempts =>
{
invalidate_media_conn().await;
force_refresh = true;
retry_with_fresh_auth = true;
break;
}
Err(err) if err.is_auth() || err.is_not_found() => {
discard_failed_write(runtime, writer).await;
return Err(err);
}
Err(err) => {
let err = err.into_anyhow();
log::warn!(
"Failed to stream-download from URL {}: {:?}. Trying next host.",
request.url,
err
);
last_err = Some(err);
}
}
}
if !retry_with_fresh_auth {
break;
}
}
discard_failed_write(runtime, writer).await;
match last_err {
Some(err) => Err(DownloadRequestError::Other(err)),
None => Err(DownloadRequestError::NoHosts),
}
}
async fn execute_request_into_memory(
http_client: &Arc<dyn HttpClient>,
runtime: &Arc<dyn Runtime>,
request: &wacore::download::DownloadRequest,
capacity: usize,
) -> std::result::Result<Vec<u8>, DownloadRequestError> {
if http_client.supports_streaming() {
let writer = std::io::Cursor::new(Vec::with_capacity(capacity));
match streaming_download_and_decrypt(http_client, runtime, request, writer).await {
Ok((writer, Ok(()))) => Ok(writer.into_inner()),
Ok((_, Err(e))) => Err(e),
Err(e) => Err(DownloadRequestError::other(e)),
}
} else {
buffered_download_to_vec(http_client, runtime, request).await
}
}
fn download_capacity(downloadable: &dyn Downloadable) -> usize {
downloadable
.file_length()
.unwrap_or(0)
.min(DOWNLOAD_PREALLOC_CAP) as usize
}
pub struct MediaDownloader {
http_client: Arc<dyn HttpClient>,
runtime: Arc<dyn Runtime>,
route: MediaRoute,
}
const NO_MEDIA_CONN_REFRESH: usize = 0;
impl MediaDownloader {
pub fn new(
http_client: Arc<dyn HttpClient>,
runtime: Arc<dyn Runtime>,
route: MediaRoute,
) -> Self {
Self {
http_client,
runtime,
route,
}
}
pub fn with_default_hosts(http_client: Arc<dyn HttpClient>, runtime: Arc<dyn Runtime>) -> Self {
Self::new(http_client, runtime, MediaRoute::default_hosts())
}
pub fn route(&self) -> &MediaRoute {
&self.route
}
#[cfg_attr(
feature = "tracing",
tracing::instrument(
name = "wa.media.download_via_route",
level = "debug",
skip_all,
err(Debug)
)
)]
pub async fn download(
&self,
downloadable: &dyn Downloadable,
) -> std::result::Result<Vec<u8>, MediaDownloadError> {
let capacity = download_capacity(downloadable);
download_media_with_retry(
NO_MEDIA_CONN_REFRESH,
|_force| async { DownloadUtils::prepare_download_requests(downloadable, &self.route) },
|| async {},
|request| async move {
execute_request_into_memory(&self.http_client, &self.runtime, &request, capacity)
.await
},
)
.await
.map_err(MediaDownloadError::from)
}
#[cfg_attr(
feature = "tracing",
tracing::instrument(
name = "wa.media.download_via_route_to_writer",
level = "debug",
skip_all,
err(Debug)
)
)]
pub async fn download_to_writer<W: DownloadWriter + Send + 'static>(
&self,
downloadable: &dyn Downloadable,
writer: W,
) -> std::result::Result<W, MediaDownloadError> {
download_to_writer_with_retry(
NO_MEDIA_CONN_REFRESH,
&self.runtime,
writer,
|_force| async { DownloadUtils::prepare_download_requests(downloadable, &self.route) },
|| async {},
|request, writer| async move {
streaming_download_and_decrypt(&self.http_client, &self.runtime, &request, writer)
.await
},
)
.await
.map_err(MediaDownloadError::from)
}
}
impl Client {
#[cfg_attr(
feature = "tracing",
tracing::instrument(name = "wa.media.download", level = "debug", skip_all, err(Debug))
)]
pub async fn download(&self, downloadable: &dyn Downloadable) -> Result<Vec<u8>> {
let capacity = download_capacity(downloadable);
download_media_with_retry(
MEDIA_AUTH_REFRESH_RETRY_ATTEMPTS,
|force| self.prepare_requests(downloadable, force),
|| async { self.invalidate_media_conn().await },
|request| async move {
execute_request_into_memory(&self.http_client, &self.runtime, &request, capacity)
.await
},
)
.await
.map_err(DownloadRequestError::into_anyhow)
}
#[cfg_attr(
feature = "tracing",
tracing::instrument(
name = "wa.media.fetch_sticker_pack",
level = "debug",
skip_all,
err(Debug)
)
)]
pub async fn fetch_sticker_pack(
&self,
pack_id: &str,
locale: &str,
) -> Result<wacore::sticker_pack::StickerPack> {
let url = wacore::sticker_pack::sticker_pack_data_url(pack_id, locale);
let response = self
.http_client
.execute(crate::http::HttpRequest::get(&url))
.await
.map_err(|e| anyhow!("sticker pack request failed: {e}"))?;
if response.status_code != HTTP_STATUS_OK {
let status = response.status_code;
return Err(HttpStatusError { status }
.into_error(format!("sticker pack endpoint returned status {status}")));
}
wacore::sticker_pack::parse_sticker_pack_response(&response.body)
}
#[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.media.download_from_params", level = "debug", skip_all, fields(kind = ?params.media_type), err(Debug)))]
pub async fn download_from_params(&self, params: &DownloadParams) -> Result<Vec<u8>> {
self.download(params).await
}
async fn prepare_requests(
&self,
downloadable: &dyn Downloadable,
force_refresh: bool,
) -> Result<Vec<wacore::download::DownloadRequest>> {
let route = if downloadable.static_url().is_some() {
MediaRoute::unauthenticated(Vec::new())
} else {
MediaRoute::from(&self.refresh_media_conn(force_refresh).await?)
};
DownloadUtils::prepare_download_requests(downloadable, &route)
}
#[cfg_attr(
feature = "tracing",
tracing::instrument(
name = "wa.media.download_to_writer",
level = "debug",
skip_all,
err(Debug)
)
)]
pub async fn download_to_writer<W: DownloadWriter + Send + 'static>(
&self,
downloadable: &dyn Downloadable,
writer: W,
) -> Result<W> {
download_to_writer_with_retry(
MEDIA_AUTH_REFRESH_RETRY_ATTEMPTS,
&self.runtime,
writer,
|force| self.prepare_requests(downloadable, force),
|| async { self.invalidate_media_conn().await },
|request, writer| async move {
streaming_download_and_decrypt(&self.http_client, &self.runtime, &request, writer)
.await
},
)
.await
.map_err(DownloadRequestError::into_anyhow)
}
#[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.media.download_from_params_to_writer", level = "debug", skip_all, fields(kind = ?params.media_type), err(Debug)))]
pub async fn download_from_params_to_writer<W: DownloadWriter + Send + 'static>(
&self,
params: &DownloadParams,
writer: W,
) -> Result<W> {
self.download_to_writer(params, writer).await
}
}
fn clear_writer<W: DownloadWriter>(writer: &mut W) -> std::io::Result<()> {
writer.truncate(0)?;
writer.rewind()?;
Ok(())
}
fn finish_verified_write<W: DownloadWriter>(
writer: &mut W,
) -> std::result::Result<(), DownloadRequestError> {
writer.rewind().map_err(DownloadRequestError::other)
}
async fn discard_failed_write<W: DownloadWriter + Send + 'static>(
runtime: &Arc<dyn Runtime>,
writer: W,
) {
let mut writer = writer;
wacore::runtime::blocking(&**runtime, move || {
if let Err(e) = clear_writer(&mut writer) {
log::warn!(
"Failed to empty the writer after a failed media download: {e}. \
It may still hold unverified bytes."
);
}
})
.await
}
async fn streaming_download_and_decrypt<W: DownloadWriter + Send + 'static>(
http_client: &Arc<dyn HttpClient>,
runtime: &Arc<dyn Runtime>,
request: &wacore::download::DownloadRequest,
writer: W,
) -> Result<(W, std::result::Result<(), DownloadRequestError>)> {
if !http_client.supports_streaming() {
return buffered_download_and_decrypt(http_client, runtime, request, writer).await;
}
let http_client = http_client.clone();
let url = request.url.clone();
let decryption = request.decryption.clone();
Ok(wacore::runtime::blocking(&**runtime, move || {
let mut writer = writer;
if let Err(e) = clear_writer(&mut writer) {
return (writer, Err(DownloadRequestError::other(e)));
}
let result = (|| -> std::result::Result<(), DownloadRequestError> {
let http_request = crate::http::HttpRequest::get(url);
let resp = http_client
.execute_streaming(http_request)
.map_err(DownloadRequestError::other)?;
validate_download_status(resp.status_code)?;
match &decryption {
MediaDecryption::Encrypted {
media_key,
media_type,
} => {
DownloadUtils::decrypt_stream_to_writer(
resp.body,
media_key,
*media_type,
&mut writer,
)
.map_err(DownloadRequestError::other)?;
}
MediaDecryption::Plaintext { file_sha256 } => {
DownloadUtils::copy_and_validate_plaintext_to_writer(
resp.body,
file_sha256,
&mut writer,
)
.map_err(DownloadRequestError::other)?;
}
}
finish_verified_write(&mut writer)
})();
(writer, result)
})
.await)
}
async fn buffered_download_and_decrypt<W: DownloadWriter + Send + 'static>(
http_client: &Arc<dyn HttpClient>,
runtime: &Arc<dyn Runtime>,
request: &wacore::download::DownloadRequest,
writer: W,
) -> Result<(W, std::result::Result<(), DownloadRequestError>)> {
let mut body = match buffered_download_body(http_client, request).await {
Ok(body) => body,
Err(err) => return Ok((writer, Err(err))),
};
let decryption = request.decryption.clone();
Ok(wacore::runtime::blocking(&**runtime, move || {
let mut writer = writer;
let result = (|| {
decrypt_or_validate_buffered_body(&mut body, &decryption)?;
clear_writer(&mut writer).map_err(DownloadRequestError::other)?;
writer
.write_all(&body)
.map_err(DownloadRequestError::other)?;
finish_verified_write(&mut writer)
})();
(writer, result)
})
.await)
}
async fn buffered_download_body(
http_client: &Arc<dyn HttpClient>,
request: &wacore::download::DownloadRequest,
) -> std::result::Result<Vec<u8>, DownloadRequestError> {
let http_request = crate::http::HttpRequest::get(request.url.clone());
let response = http_client
.execute(http_request)
.await
.map_err(DownloadRequestError::other)?;
validate_download_status(response.status_code)?;
Ok(response.body)
}
async fn buffered_download_to_vec(
http_client: &Arc<dyn HttpClient>,
runtime: &Arc<dyn Runtime>,
request: &wacore::download::DownloadRequest,
) -> std::result::Result<Vec<u8>, DownloadRequestError> {
let mut body = buffered_download_body(http_client, request).await?;
let decryption = request.decryption.clone();
wacore::runtime::blocking(&**runtime, move || {
decrypt_or_validate_buffered_body(&mut body, &decryption)?;
Ok(body)
})
.await
}
#[cfg(test)]
mod tests {
use super::*;
use crate::error::ErrorChainExt;
use crate::mediaconn::{MediaConn, MediaConnHost};
use async_lock::Mutex;
use std::io::{Cursor, Seek, SeekFrom, Write};
use std::sync::Arc;
use wacore::time::Instant;
use waproto::whatsapp as wa;
struct PlaintextDownloadable {
direct_path: String,
file_sha256: Vec<u8>,
}
impl Downloadable for PlaintextDownloadable {
fn direct_path(&self) -> Option<&str> {
Some(&self.direct_path)
}
fn media_key(&self) -> Option<&[u8]> {
None
}
fn file_enc_sha256(&self) -> Option<&[u8]> {
None
}
fn file_sha256(&self) -> Option<&[u8]> {
Some(&self.file_sha256)
}
fn file_length(&self) -> Option<u64> {
None
}
fn app_info(&self) -> MediaType {
MediaType::Image
}
}
fn media_conn(auth: &str, hosts: &[&str]) -> MediaConn {
MediaConn {
auth: auth.to_string(),
ttl: 60,
auth_ttl: None,
hosts: hosts
.iter()
.map(|hostname| MediaConnHost::new((*hostname).to_string()))
.collect(),
fetched_at: Instant::now(),
}
}
fn plaintext_sha256(data: &[u8]) -> Vec<u8> {
wacore::upload::encrypt_media(data, MediaType::Image)
.expect("hash derivation should succeed")
.file_sha256
.to_vec()
}
#[cfg(feature = "ureq-client")]
fn spawn_cdn_server(status: u16, reason: &'static str, body: Vec<u8>) -> String {
use std::io::Read;
use std::net::TcpListener;
let listener = TcpListener::bind("127.0.0.1:0").expect("bind ephemeral port");
let addr = listener.local_addr().expect("local addr");
std::thread::spawn(move || {
let Ok((mut stream, _)) = listener.accept() else {
return;
};
let mut buf = Vec::new();
let mut tmp = [0u8; 1024];
while !buf.windows(4).any(|w| w == b"\r\n\r\n") {
match stream.read(&mut tmp) {
Ok(0) | Err(_) => return,
Ok(n) => buf.extend_from_slice(&tmp[..n]),
}
}
let header = format!(
"HTTP/1.1 {status} {reason}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
body.len()
);
let _ = stream.write_all(header.as_bytes());
let _ = stream.write_all(&body);
});
format!("http://{addr}")
}
#[cfg(feature = "ureq-client")]
fn spawn_cdn_status_server(status: u16, reason: &'static str) -> String {
spawn_cdn_server(status, reason, b"denied".to_vec())
}
#[cfg(feature = "ureq-client")]
fn plaintext_request(url: String) -> wacore::download::DownloadRequest {
wacore::download::DownloadRequest {
url,
decryption: MediaDecryption::Plaintext {
file_sha256: vec![0u8; 32],
},
}
}
#[cfg(feature = "ureq-client")]
async fn ureq_client() -> Arc<Client> {
crate::test_utils::create_test_client_with_http(
"cdn-status",
Arc::new(whatsapp_rust_ureq_http_client::UreqHttpClient::new()),
)
.await
}
#[cfg(feature = "ureq-client")]
#[tokio::test]
async fn cdn_auth_status_reaches_the_classifier_on_the_streaming_path() {
for status in [401u16, 403] {
let url = spawn_cdn_status_server(status, "Forbidden");
let client = ureq_client().await;
let (_writer, result) = streaming_download_and_decrypt(
&client.http_client,
&client.runtime,
&plaintext_request(url),
Cursor::new(Vec::new()),
)
.await
.expect("the request itself completes; the status is the failure");
let err = result.expect_err("a non-2xx CDN response must fail the download");
assert!(
err.is_auth(),
"{status} must classify as an auth error so the media conn is refreshed, got {err:?}"
);
}
}
#[cfg(feature = "ureq-client")]
#[tokio::test]
async fn cdn_expired_status_reaches_the_classifier_on_the_buffered_path() {
for status in [404u16, 410] {
let url = spawn_cdn_status_server(status, "Gone");
let client = ureq_client().await;
let err = buffered_download_body(&client.http_client, &plaintext_request(url))
.await
.expect_err("a non-2xx CDN response must fail the download");
assert!(
err.is_not_found(),
"{status} must classify as expired so the URL is re-derived, got {err:?}"
);
}
}
#[cfg(feature = "ureq-client")]
#[tokio::test]
async fn stale_auth_403_invalidates_the_media_conn_and_the_retry_recovers() {
let body = b"download me".to_vec();
let file_sha256 = plaintext_sha256(&body);
let stale_host = spawn_cdn_status_server(403, "Forbidden");
let fresh_host = spawn_cdn_server(200, "OK", body.clone());
let client = ureq_client().await;
let invalidations = Arc::new(Mutex::new(0usize));
let attempts = Arc::new(Mutex::new(Vec::new()));
let downloaded = download_media_with_retry(
MEDIA_AUTH_REFRESH_RETRY_ATTEMPTS,
{
let attempts = Arc::clone(&attempts);
move |force| {
let attempts = Arc::clone(&attempts);
let url = if force {
fresh_host.clone()
} else {
stale_host.clone()
};
let file_sha256 = file_sha256.clone();
async move {
attempts.lock().await.push(force);
Ok(vec![wacore::download::DownloadRequest {
url,
decryption: MediaDecryption::Plaintext { file_sha256 },
}])
}
}
},
{
let invalidations = Arc::clone(&invalidations);
move || {
let invalidations = Arc::clone(&invalidations);
async move {
*invalidations.lock().await += 1;
}
}
},
|request| {
let client = Arc::clone(&client);
async move {
match streaming_download_and_decrypt(
&client.http_client,
&client.runtime,
&request,
Cursor::new(Vec::new()),
)
.await
{
Ok((writer, Ok(()))) => Ok(writer.into_inner()),
Ok((_, Err(e))) => Err(e),
Err(e) => Err(DownloadRequestError::other(e)),
}
}
},
)
.await
.expect("the forced-refresh retry must recover the download");
assert_eq!(downloaded, body);
assert_eq!(
*invalidations.lock().await,
1,
"a 403 must invalidate the cached media conn"
);
assert_eq!(
*attempts.lock().await,
vec![false, true],
"the second attempt must ask for a refreshed media conn"
);
}
#[cfg(feature = "ureq-client")]
#[tokio::test]
async fn the_cdn_status_survives_to_the_public_error() {
for (status, reason) in [
(403u16, "Forbidden"),
(410, "Gone"),
(429, "Too Many Requests"),
] {
let first = spawn_cdn_status_server(status, reason);
let refreshed = spawn_cdn_status_server(status, reason);
let client = ureq_client().await;
let err = download_media_with_retry(
MEDIA_AUTH_REFRESH_RETRY_ATTEMPTS,
move |force| {
let url = if force {
refreshed.clone()
} else {
first.clone()
};
async move { Ok(vec![plaintext_request(url)]) }
},
|| async {},
|request| {
let client = Arc::clone(&client);
async move {
match streaming_download_and_decrypt(
&client.http_client,
&client.runtime,
&request,
Cursor::new(Vec::new()),
)
.await
{
Ok((writer, Ok(()))) => Ok(writer.into_inner()),
Ok((_, Err(e))) => Err(e),
Err(e) => Err(DownloadRequestError::other(e)),
}
}
},
)
.await
.map_err(DownloadRequestError::into_anyhow)
.expect_err("a non-2xx CDN response must fail the download");
let cause: &(dyn std::error::Error + 'static) = err.as_ref();
assert_eq!(
cause.http_status(),
Some(status),
"the consumer must recover {status} by type, got: {err:?}"
);
assert!(
format!("{err}").contains(&status.to_string()),
"the message should still name the status, got: {err}"
);
}
}
#[cfg(feature = "ureq-client")]
#[tokio::test]
async fn a_failure_with_no_exchange_reports_no_status() {
let client = ureq_client().await;
let request = plaintext_request("http://127.0.0.1:1".to_string());
let err = download_media_with_retry(
MEDIA_AUTH_REFRESH_RETRY_ATTEMPTS,
move |_force| {
let request = request.clone();
async move { Ok(vec![request]) }
},
|| async {},
|request| {
let client = Arc::clone(&client);
async move {
match streaming_download_and_decrypt(
&client.http_client,
&client.runtime,
&request,
Cursor::new(Vec::new()),
)
.await
{
Ok((writer, Ok(()))) => Ok(writer.into_inner()),
Ok((_, Err(e))) => Err(e),
Err(e) => Err(DownloadRequestError::other(e)),
}
}
},
)
.await
.map_err(DownloadRequestError::into_anyhow)
.expect_err("a refused connection must fail the download");
let cause: &(dyn std::error::Error + 'static) = err.as_ref();
assert_eq!(
cause.http_status(),
None,
"a transport failure must not be laundered into an upstream status, got: {err:?}"
);
}
#[test]
fn download_statuses_have_one_shared_classification() {
use crate::http::{HTTP_STATUS_FORBIDDEN, HTTP_STATUS_UNAUTHORIZED};
assert!(validate_download_status(HTTP_STATUS_OK).is_ok());
assert!(matches!(
validate_download_status(HTTP_STATUS_UNAUTHORIZED),
Err(DownloadRequestError::Auth(_))
));
assert!(matches!(
validate_download_status(HTTP_STATUS_FORBIDDEN),
Err(DownloadRequestError::Auth(_))
));
assert!(matches!(
validate_download_status(HTTP_STATUS_NOT_FOUND),
Err(DownloadRequestError::NotFound(_))
));
assert!(matches!(
validate_download_status(HTTP_STATUS_GONE),
Err(DownloadRequestError::NotFound(_))
));
assert!(matches!(
validate_download_status(HTTP_STATUS_REDIRECTION_START),
Err(DownloadRequestError::Other(_))
));
}
#[test]
fn process_downloaded_media_ok() {
let data = b"Hello media test";
let enc = wacore::upload::encrypt_media(data, MediaType::Image)
.expect("encryption should succeed");
let mut cursor = Cursor::new(Vec::<u8>::new());
let plaintext = DownloadUtils::verify_and_decrypt(
&enc.data_to_upload,
&enc.media_key,
MediaType::Image,
)
.expect("decryption should succeed");
cursor.write_all(&plaintext).expect("write should succeed");
assert_eq!(cursor.into_inner(), data);
}
#[test]
fn process_downloaded_media_bad_mac() {
let data = b"Tamper";
let mut enc = wacore::upload::encrypt_media(data, MediaType::Image)
.expect("encryption should succeed");
let last = enc.data_to_upload.len() - 1;
enc.data_to_upload[last] ^= 0x01;
let err = DownloadUtils::verify_and_decrypt(
&enc.data_to_upload,
&enc.media_key,
MediaType::Image,
)
.unwrap_err();
assert!(
matches!(&err, MediaDecryptionError::InvalidMac),
"Expected InvalidMac, got: {}",
err
);
}
#[tokio::test]
async fn download_retries_with_forced_media_conn_refresh_after_auth_error() {
let body = b"download me".to_vec();
let downloadable = PlaintextDownloadable {
direct_path: "/v/t62.7118-24/123".to_string(),
file_sha256: plaintext_sha256(&body),
};
let first_conn = media_conn("stale-auth", &["cdn1.example.com"]);
let refreshed_conn = media_conn("fresh-auth", &["cdn2.example.com"]);
let refresh_calls = Arc::new(Mutex::new(Vec::new()));
let invalidations = Arc::new(Mutex::new(0usize));
let seen_urls = Arc::new(Mutex::new(Vec::new()));
let downloaded = download_media_with_retry(
MEDIA_AUTH_REFRESH_RETRY_ATTEMPTS,
{
let refresh_calls = Arc::clone(&refresh_calls);
let downloadable = &downloadable;
move |force| {
let refresh_calls = Arc::clone(&refresh_calls);
let first_conn = first_conn.clone();
let refreshed_conn = refreshed_conn.clone();
async move {
refresh_calls.lock().await.push(force);
let media_conn = if force { refreshed_conn } else { first_conn };
DownloadUtils::prepare_download_requests(
downloadable,
&MediaRoute::from(&media_conn),
)
}
}
},
{
let invalidations = Arc::clone(&invalidations);
move || {
let invalidations = Arc::clone(&invalidations);
async move {
*invalidations.lock().await += 1;
}
}
},
{
let seen_urls = Arc::clone(&seen_urls);
let body = body.clone();
move |request| {
let seen_urls = Arc::clone(&seen_urls);
let body = body.clone();
let url = request.url.clone();
async move {
seen_urls.lock().await.push(url.clone());
if url.contains("stale-auth") {
Err(DownloadRequestError::auth(401))
} else {
Ok(body)
}
}
}
},
)
.await
.expect("download should succeed after refreshing media auth");
assert_eq!(downloaded, body);
assert_eq!(*refresh_calls.lock().await, vec![false, true]);
assert_eq!(*invalidations.lock().await, 1);
let seen_urls = seen_urls.lock().await.clone();
assert_eq!(seen_urls.len(), 2);
assert!(seen_urls[0].contains("auth=stale-auth"));
assert!(seen_urls[1].contains("auth=fresh-auth"));
}
#[tokio::test]
async fn download_fails_over_to_next_host_without_refresh() {
let body = b"failover me".to_vec();
let downloadable = PlaintextDownloadable {
direct_path: "/v/t62.7118-24/failover".to_string(),
file_sha256: plaintext_sha256(&body),
};
let conn = media_conn(
"auth-tok",
&["bad-host.example.com", "good-host.example.com"],
);
let refresh_calls = Arc::new(Mutex::new(Vec::new()));
let invalidations = Arc::new(Mutex::new(0usize));
let seen_urls = Arc::new(Mutex::new(Vec::new()));
let downloaded = download_media_with_retry(
MEDIA_AUTH_REFRESH_RETRY_ATTEMPTS,
{
let refresh_calls = Arc::clone(&refresh_calls);
let downloadable = &downloadable;
let conn = conn.clone();
move |force| {
let refresh_calls = Arc::clone(&refresh_calls);
let conn = conn.clone();
async move {
refresh_calls.lock().await.push(force);
DownloadUtils::prepare_download_requests(
downloadable,
&MediaRoute::from(&conn),
)
}
}
},
{
let invalidations = Arc::clone(&invalidations);
move || {
let invalidations = Arc::clone(&invalidations);
async move {
*invalidations.lock().await += 1;
}
}
},
{
let seen_urls = Arc::clone(&seen_urls);
let body = body.clone();
move |request| {
let seen_urls = Arc::clone(&seen_urls);
let body = body.clone();
let url = request.url.clone();
async move {
seen_urls.lock().await.push(url.clone());
if url.contains("bad-host") {
Err(DownloadRequestError::other(anyhow!("connection reset")))
} else {
Ok(body)
}
}
}
},
)
.await
.expect("download should fail over to the healthy host");
assert_eq!(downloaded, body);
assert_eq!(*refresh_calls.lock().await, vec![false]);
assert_eq!(*invalidations.lock().await, 0);
let seen_urls = seen_urls.lock().await.clone();
assert_eq!(seen_urls.len(), 2);
assert!(seen_urls[0].contains("bad-host"));
assert!(seen_urls[1].contains("good-host"));
}
#[tokio::test]
async fn download_propagates_last_error_when_all_hosts_fail() {
let body = b"never arrives".to_vec();
let downloadable = PlaintextDownloadable {
direct_path: "/v/t62.7118-24/allfail".to_string(),
file_sha256: plaintext_sha256(&body),
};
let conn = media_conn("auth-tok", &["host-a.example.com", "host-b.example.com"]);
let invalidations = Arc::new(Mutex::new(0usize));
let seen_urls = Arc::new(Mutex::new(Vec::new()));
let err = download_media_with_retry(
MEDIA_AUTH_REFRESH_RETRY_ATTEMPTS,
{
let downloadable = &downloadable;
let conn = conn.clone();
move |_force| {
let conn = conn.clone();
async move {
DownloadUtils::prepare_download_requests(
downloadable,
&MediaRoute::from(&conn),
)
}
}
},
{
let invalidations = Arc::clone(&invalidations);
move || {
let invalidations = Arc::clone(&invalidations);
async move {
*invalidations.lock().await += 1;
}
}
},
{
let seen_urls = Arc::clone(&seen_urls);
move |request| {
let seen_urls = Arc::clone(&seen_urls);
let url = request.url.clone();
async move {
seen_urls.lock().await.push(url.clone());
Err::<Vec<u8>, _>(DownloadRequestError::other(anyhow!("host {url} down")))
}
}
},
)
.await
.expect_err("all hosts failing must surface an error")
.into_anyhow();
assert!(
err.to_string().contains("down"),
"expected the propagated last_err, got: {err}"
);
assert_eq!(*invalidations.lock().await, 0);
assert_eq!(seen_urls.lock().await.len(), 2);
}
#[tokio::test]
async fn download_to_writer_retries_with_forced_media_conn_refresh_after_auth_error() {
let body = b"stream me".to_vec();
let downloadable = PlaintextDownloadable {
direct_path: "/v/t62.7118-24/stream".to_string(),
file_sha256: plaintext_sha256(&body),
};
let first_conn = media_conn("stale-auth", &["cdn1.example.com"]);
let refreshed_conn = media_conn("fresh-auth", &["cdn2.example.com"]);
let refresh_calls = Arc::new(Mutex::new(Vec::new()));
let invalidations = Arc::new(Mutex::new(0usize));
let seen_urls = Arc::new(Mutex::new(Vec::new()));
let runtime: Arc<dyn Runtime> = Arc::new(crate::TokioRuntime);
let writer = download_to_writer_with_retry(
MEDIA_AUTH_REFRESH_RETRY_ATTEMPTS,
&runtime,
Cursor::new(Vec::<u8>::new()),
{
let refresh_calls = Arc::clone(&refresh_calls);
let downloadable = &downloadable;
move |force| {
let refresh_calls = Arc::clone(&refresh_calls);
let first_conn = first_conn.clone();
let refreshed_conn = refreshed_conn.clone();
async move {
refresh_calls.lock().await.push(force);
let media_conn = if force { refreshed_conn } else { first_conn };
DownloadUtils::prepare_download_requests(
downloadable,
&MediaRoute::from(&media_conn),
)
}
}
},
{
let invalidations = Arc::clone(&invalidations);
move || {
let invalidations = Arc::clone(&invalidations);
async move {
*invalidations.lock().await += 1;
}
}
},
{
let seen_urls = Arc::clone(&seen_urls);
let body = body.clone();
move |request, mut writer| {
let seen_urls = Arc::clone(&seen_urls);
let body = body.clone();
let url = request.url.clone();
async move {
seen_urls.lock().await.push(url.clone());
writer.seek(SeekFrom::Start(0))?;
if url.contains("stale-auth") {
Ok((writer, Err(DownloadRequestError::auth(403))))
} else {
writer.write_all(&body)?;
writer.seek(SeekFrom::Start(0))?;
Ok((writer, Ok(())))
}
}
}
},
)
.await
.expect("streaming download should succeed after refreshing media auth");
assert_eq!(writer.into_inner(), body);
assert_eq!(*refresh_calls.lock().await, vec![false, true]);
assert_eq!(*invalidations.lock().await, 1);
let seen_urls = seen_urls.lock().await.clone();
assert_eq!(seen_urls.len(), 2);
assert!(seen_urls[0].contains("auth=stale-auth"));
assert!(seen_urls[1].contains("auth=fresh-auth"));
}
struct RoutedHttpClient {
routes: Vec<(&'static str, u16, Vec<u8>)>,
fallback: (u16, Vec<u8>),
streaming: bool,
seen_urls: std::sync::Mutex<Vec<String>>,
}
impl RoutedHttpClient {
fn new(routes: Vec<(&'static str, u16, Vec<u8>)>, fallback: (u16, Vec<u8>)) -> Arc<Self> {
Arc::new(Self {
routes,
fallback,
streaming: false,
seen_urls: std::sync::Mutex::new(Vec::new()),
})
}
fn streaming(
routes: Vec<(&'static str, u16, Vec<u8>)>,
fallback: (u16, Vec<u8>),
) -> Arc<Self> {
Arc::new(Self {
routes,
fallback,
streaming: true,
seen_urls: std::sync::Mutex::new(Vec::new()),
})
}
}
impl RoutedHttpClient {
fn record(&self, url: &str) {
self.seen_urls
.lock()
.expect("test mutex is never poisoned")
.push(url.to_string());
}
fn urls(&self) -> Vec<String> {
self.seen_urls
.lock()
.expect("test mutex is never poisoned")
.clone()
}
fn respond(&self, url: &str) -> (u16, Vec<u8>) {
self.routes
.iter()
.find(|(needle, _, _)| url.contains(needle))
.map(|(_, status, body)| (*status, body.clone()))
.unwrap_or_else(|| self.fallback.clone())
}
}
#[async_trait::async_trait]
impl HttpClient for RoutedHttpClient {
async fn execute(
&self,
request: crate::http::HttpRequest,
) -> Result<crate::http::HttpResponse> {
self.record(&request.url);
let (status_code, body) = self.respond(&request.url);
Ok(crate::http::HttpResponse { status_code, body })
}
fn supports_streaming(&self) -> bool {
self.streaming
}
fn execute_streaming(
&self,
request: crate::http::HttpRequest,
) -> Result<wacore::net::StreamingHttpResponse> {
self.record(&request.url);
let (status_code, body) = self.respond(&request.url);
Ok(wacore::net::StreamingHttpResponse {
status_code,
body: Box::new(Cursor::new(body)),
})
}
}
fn downloader(http: Arc<RoutedHttpClient>, hosts: &[&str]) -> MediaDownloader {
MediaDownloader::new(
http,
Arc::new(crate::TokioRuntime),
MediaRoute::unauthenticated(hosts.iter().copied().map(MediaHost::new).collect()),
)
}
fn encrypted_params(data: &[u8]) -> (DownloadParams, Vec<u8>) {
let enc = wacore::upload::encrypt_media(data, MediaType::Image)
.expect("encryption should succeed");
let params = DownloadParams::encrypted(
"/v/t62.7118-24/no-session",
&enc.media_key,
&enc.file_sha256,
&enc.file_enc_sha256,
data.len() as u64,
MediaType::Image,
);
(params, enc.data_to_upload)
}
fn forged_body(plaintext: &[u8], media_key: &[u8; 32]) -> Vec<u8> {
let mut body =
wacore::upload::encrypt_media_with_key(plaintext, MediaType::Image, Some(media_key))
.expect("encryption should succeed")
.data_to_upload;
let last = body.len() - 1;
body[last] ^= 1;
body
}
#[tokio::test]
async fn a_failed_host_leaves_no_tail_behind_a_shorter_successful_retry() {
let media_key = [0x5b; 32];
let original = b"short but verified".to_vec();
let good =
wacore::upload::encrypt_media_with_key(&original, MediaType::Image, Some(&media_key))
.expect("encryption should succeed");
let params = DownloadParams::encrypted(
"/v/t62.7118-24/tail",
&good.media_key,
&good.file_sha256,
&good.file_enc_sha256,
original.len() as u64,
MediaType::Image,
);
let forged = forged_body(&vec![0xAA; 64 * 1024], &media_key);
assert!(forged.len() > good.data_to_upload.len() * 10);
for streaming in [true, false] {
let routes = vec![
("forging-host", 200, forged.clone()),
("honest-host", 200, good.data_to_upload.clone()),
];
let http = if streaming {
RoutedHttpClient::streaming(routes, (500, Vec::new()))
} else {
RoutedHttpClient::new(routes, (500, Vec::new()))
};
let writer = downloader(
http.clone(),
&["forging-host.example.com", "honest-host.example.com"],
)
.download_to_writer(¶ms, Cursor::new(Vec::new()))
.await
.expect("the honest host must still satisfy the download");
assert_eq!(
http.urls().len(),
2,
"the forged host must have been tried first (streaming={streaming})"
);
assert_eq!(
writer.into_inner(),
original,
"the forged host's plaintext must not survive past the media (streaming={streaming})"
);
}
}
#[tokio::test]
async fn a_writer_that_arrives_with_content_still_ends_up_holding_only_the_media() {
let original = b"exactly this".to_vec();
let (params, encrypted) = encrypted_params(&original);
let http = RoutedHttpClient::streaming(Vec::new(), (200, encrypted));
let writer = downloader(http, &["cdn.example.com"])
.download_to_writer(¶ms, Cursor::new(vec![0xFF; 4096]))
.await
.expect("download should succeed");
assert_eq!(writer.into_inner(), original);
}
#[tokio::test]
async fn an_append_mode_file_still_ends_up_holding_only_the_media() {
let original = b"appended, yet exact".to_vec();
let (params, encrypted) = encrypted_params(&original);
let path = std::env::temp_dir().join(format!(
"wa-rust-append-{}-{:?}.bin",
std::process::id(),
std::thread::current().id()
));
std::fs::write(&path, b"stale bytes the caller left behind")
.expect("fixture write should succeed");
let file = std::fs::OpenOptions::new()
.append(true)
.open(&path)
.expect("append open should succeed");
let http = RoutedHttpClient::streaming(Vec::new(), (200, encrypted));
downloader(http, &["cdn.example.com"])
.download_to_writer(¶ms, file)
.await
.expect("download should succeed");
let written = std::fs::read(&path).expect("read back should succeed");
let _ = std::fs::remove_file(&path);
assert_eq!(written, original);
}
#[derive(Clone, Debug)]
struct SharedWriter(Arc<std::sync::Mutex<Cursor<Vec<u8>>>>);
impl SharedWriter {
fn new() -> Self {
Self(Arc::new(std::sync::Mutex::new(Cursor::new(Vec::new()))))
}
fn contents(&self) -> Vec<u8> {
self.with(|inner| inner.get_ref().clone())
}
fn with<T>(&self, f: impl FnOnce(&mut Cursor<Vec<u8>>) -> T) -> T {
f(&mut self.0.lock().expect("test mutex is never poisoned"))
}
}
impl Write for SharedWriter {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
self.with(|inner| inner.write(buf))
}
fn flush(&mut self) -> std::io::Result<()> {
self.with(|inner| inner.flush())
}
}
impl Seek for SharedWriter {
fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
self.with(|inner| inner.seek(pos))
}
}
impl DownloadWriter for SharedWriter {
fn truncate(&mut self, len: u64) -> std::io::Result<()> {
self.with(|inner| inner.truncate(len))
}
}
#[tokio::test]
async fn a_download_that_fails_everywhere_empties_the_writer() {
let media_key = [0x77; 32];
let reference = encrypted_params(b"unused").0;
let params = DownloadParams::encrypted(
"/v/t62.7118-24/doomed",
&media_key,
&reference.file_sha256,
reference.file_enc_sha256.as_deref().unwrap_or_default(),
12,
MediaType::Image,
);
let forged = forged_body(&vec![0x11; 32 * 1024], &media_key);
let http = RoutedHttpClient::streaming(vec![("only-host", 200, forged)], (500, Vec::new()));
let sink = SharedWriter::new();
let err = downloader(http, &["only-host.example.com"])
.download_to_writer(¶ms, sink.clone())
.await
.expect_err("a forged body must not be reported as a download");
assert!(matches!(err, MediaDownloadError::HostsUnreachable(_)));
assert!(
sink.contents().is_empty(),
"a failed download must not leave plaintext behind"
);
}
#[tokio::test]
async fn media_downloader_fetches_without_a_session_or_auth() {
let original = b"media that outlived its session".to_vec();
let (params, encrypted) = encrypted_params(&original);
let http = RoutedHttpClient::new(
vec![("good-host", 200, encrypted)],
(500, b"server error".to_vec()),
);
let downloaded = downloader(
http.clone(),
&["bad-host.example.com", "good-host.example.com"],
)
.download(¶ms)
.await
.expect("an injected host list is all a download needs");
assert_eq!(downloaded, original);
let seen = http.urls();
assert_eq!(seen.len(), 2, "the failing host must fail over to the next");
assert!(
seen[0].starts_with("https://bad-host.example.com/v/t62.7118-24/no-session?token=")
);
assert!(
seen[1].starts_with("https://good-host.example.com/v/t62.7118-24/no-session?token=")
);
assert!(seen.iter().all(|url| !url.contains("auth=")));
}
#[tokio::test]
async fn media_downloader_streams_to_a_writer_without_a_session() {
let original = b"streamed without a session".to_vec();
let (params, encrypted) = encrypted_params(&original);
let http = RoutedHttpClient::streaming(Vec::new(), (200, encrypted.clone()));
assert!(
http.supports_streaming(),
"this must not fall back to buffered"
);
let writer = downloader(http, &["cdn.example.com"])
.download_to_writer(¶ms, Cursor::new(Vec::new()))
.await
.expect("the streaming path must work without a session too");
assert_eq!(writer.into_inner(), original);
let buffered = RoutedHttpClient::new(Vec::new(), (200, encrypted));
let writer = downloader(buffered, &["cdn.example.com"])
.download_to_writer(¶ms, Cursor::new(Vec::new()))
.await
.expect("the buffered fallback must work too");
assert_eq!(writer.into_inner(), original);
}
#[tokio::test]
async fn media_downloader_separates_an_unbuildable_reference_from_a_dead_host() {
let params = DownloadParams {
direct_path: "/v/t62.7118-24/incomplete".to_string(),
media_key: Some(vec![1u8; 32]),
file_sha256: vec![2u8; 32],
file_enc_sha256: None,
file_length: 16,
media_type: MediaType::Image,
};
let http = RoutedHttpClient::new(Vec::new(), (200, Vec::new()));
let err = downloader(http.clone(), &["cdn1.example.com", "cdn2.example.com"])
.download(¶ms)
.await
.expect_err("an incomplete reference cannot succeed");
assert!(
matches!(err, MediaDownloadError::Other(_)),
"a reference that cannot build a URL is not a host failure, got {err:?}"
);
assert!(
err.to_string().contains("Missing file_enc_sha256"),
"the cause must survive the classification, got: {err}"
);
assert!(http.urls().is_empty());
}
#[tokio::test]
async fn media_downloader_separates_an_expired_reference_from_a_dead_host() {
let (params, _) = encrypted_params(b"gone");
let expired = RoutedHttpClient::new(Vec::new(), (410, b"gone".to_vec()));
let err = downloader(expired.clone(), &["cdn1.example.com", "cdn2.example.com"])
.download(¶ms)
.await
.expect_err("an expired reference must not read as success");
assert!(
matches!(err, MediaDownloadError::ReferenceRejected(_)),
"expected a rejected reference, got {err:?}"
);
assert_eq!(
expired.urls().len(),
1,
"no host rotation and no refresh: nothing here can re-sign the reference"
);
let dead = RoutedHttpClient::new(Vec::new(), (500, b"boom".to_vec()));
let err = downloader(dead.clone(), &["cdn1.example.com", "cdn2.example.com"])
.download(¶ms)
.await
.expect_err("every host failing must not read as success");
assert!(
matches!(err, MediaDownloadError::HostsUnreachable(_)),
"expected unreachable hosts, got {err:?}"
);
assert_eq!(dead.urls().len(), 2, "every host is tried");
}
#[tokio::test]
async fn media_downloader_defaults_to_the_known_cdn_hosts() {
let downloader = MediaDownloader::with_default_hosts(
RoutedHttpClient::new(Vec::new(), (200, Vec::new())),
Arc::new(crate::TokioRuntime),
);
assert!(downloader.route().auth.is_none());
assert_eq!(
downloader
.route()
.hosts
.iter()
.map(|h| h.hostname.as_str())
.collect::<Vec<_>>(),
DEFAULT_MEDIA_HOSTS.to_vec(),
);
}
#[tokio::test]
async fn media_downloader_without_hosts_reports_it() {
let (params, _) = encrypted_params(b"nowhere to go");
let http = RoutedHttpClient::new(Vec::new(), (200, Vec::new()));
let err = downloader(http.clone(), &[])
.download(¶ms)
.await
.expect_err("an empty route cannot succeed");
assert!(
matches!(err, MediaDownloadError::NoHosts),
"expected NoHosts, got {err:?}"
);
assert!(http.urls().is_empty());
}
#[tokio::test]
async fn static_url_download_asks_for_no_media_conn() {
let client = crate::test_utils::create_test_client_with_name("static_url_no_iq").await;
let with_static_url = wa::message::ImageMessage {
static_url: Some("https://static.cdn.example.com/media/abc123".to_string()),
direct_path: Some("/v/t62.7118-24/unused".to_string()),
file_sha256: Some(vec![7u8; 32]),
..Default::default()
};
let requests = client
.prepare_requests(&with_static_url, false)
.await
.expect("a static URL needs no hosts, so it must not need a session");
assert_eq!(requests.len(), 1);
assert_eq!(
requests[0].url,
"https://static.cdn.example.com/media/abc123"
);
let without_static_url = wa::message::ImageMessage {
direct_path: Some("/v/t62.7118-24/needs-hosts".to_string()),
file_sha256: Some(vec![7u8; 32]),
..Default::default()
};
let err = client
.prepare_requests(&without_static_url, false)
.await
.expect_err("host construction still needs a media conn");
assert!(
err.to_string().contains("not connected"),
"expected the media-conn IQ to be attempted, got: {err}"
);
}
struct CannedHttpClient {
status: u16,
body: Vec<u8>,
seen_url: Mutex<Option<String>>,
}
#[async_trait::async_trait]
impl HttpClient for CannedHttpClient {
async fn execute(
&self,
request: crate::http::HttpRequest,
) -> Result<crate::http::HttpResponse> {
*self.seen_url.lock().await = Some(request.url);
Ok(crate::http::HttpResponse {
status_code: self.status,
body: self.body.clone(),
})
}
}
#[tokio::test]
async fn fetch_sticker_pack_hits_cdn_and_parses() {
use base64::engine::{Engine, general_purpose::STANDARD};
let body = format!(
r#"[{{"sticker-pack-id":"P1","name":"Cats","stickers":[
{{"media-key":"{}","file-hash":"{}","enc-file-hash":"{}","direct-path":"/d","file-size":9}}
]}}]"#,
STANDARD.encode([1u8; 32]),
STANDARD.encode([2u8; 32]),
STANDARD.encode([3u8; 32]),
);
let http = Arc::new(CannedHttpClient {
status: 200,
body: body.into_bytes(),
seen_url: Mutex::new(None),
});
let client =
crate::test_utils::create_test_client_with_http("sticker_fetch", http.clone()).await;
let pack = client.fetch_sticker_pack("P1", "en").await.unwrap();
assert_eq!(pack.sticker_pack_id.as_deref(), Some("P1"));
assert_eq!(pack.stickers.len(), 1);
assert_eq!(pack.stickers[0].direct_path(), Some("/d"));
let url = http.seen_url.lock().await.clone().unwrap();
assert_eq!(
url,
"https://static.whatsapp.net/sticker?lottie=1&cat=sticker_pack_data&id=P1&lg=en"
);
}
#[tokio::test]
async fn fetch_sticker_pack_errors_on_non_200() {
let http = Arc::new(CannedHttpClient {
status: 404,
body: Vec::new(),
seen_url: Mutex::new(None),
});
let client = crate::test_utils::create_test_client_with_http("sticker_404", http).await;
let err = client
.fetch_sticker_pack("P1", "en")
.await
.expect_err("a non-200 sticker pack response must fail");
let cause: &(dyn std::error::Error + 'static) = err.as_ref();
assert_eq!(cause.http_status(), Some(404), "got: {err:?}");
}
}