1313// SPDX-License-Identifier: Apache-2.0
1414
1515use assert_cmd:: prelude:: * ;
16- use file_diff:: diff_files;
16+ use futures:: StreamExt ;
17+ use md5:: Context ;
1718use rand:: distr:: { Alphanumeric , SampleString } ;
19+ use reqwest:: { Client , header} ;
1820use serde_json:: Value ;
19- use std:: fs:: File ;
20- use std:: io:: Cursor ;
21- use std:: io:: copy;
2221use std:: process:: Command ;
23- use tempfile:: Builder ;
22+ use std:: { error:: Error , path:: PathBuf } ;
23+ use tempfile:: { Builder , TempDir } ;
24+ use tokio:: { fs:: File , io:: AsyncReadExt , io:: AsyncWriteExt } ;
25+
26+ /// Downloads a file, saving it with the filename provided by the server.
27+ /// Returns (final_path, md5_checksum).
28+ pub async fn download_with_md5_and_filename (
29+ url : & str ,
30+ tmp_dir : & TempDir ,
31+ ) -> Result < ( PathBuf , String ) , Box < dyn Error > > {
32+ let client = Client :: new ( ) ;
33+ let response = client. get ( url) . send ( ) . await ?;
34+ response. error_for_status_ref ( ) ?; // fail fast on HTTP errors
35+
36+ // Try to extract filename from Content-Disposition or fallback to URL
37+ let filename = response
38+ . headers ( )
39+ . get ( header:: CONTENT_DISPOSITION )
40+ . and_then ( |val| val. to_str ( ) . ok ( ) )
41+ . and_then ( parse_filename_from_content_disposition)
42+ . or_else ( || extract_filename_from_url ( url) )
43+ . unwrap_or_else ( || "download.bin" . to_string ( ) ) ;
44+
45+ let path = tmp_dir. path ( ) . join ( & filename) ;
46+ let mut file = File :: create ( & path) . await ?;
47+ let mut context = Context :: new ( ) ;
48+
49+ let mut stream = response. bytes_stream ( ) ;
50+ while let Some ( chunk_result) = stream. next ( ) . await {
51+ let chunk = chunk_result?;
52+ file. write_all ( & chunk) . await ?;
53+ context. consume ( & chunk) ;
54+ }
55+
56+ file. flush ( ) . await ?;
57+ let digest = context. compute ( ) ;
58+ let checksum = format ! ( "{:x}" , digest) ;
59+
60+ Ok ( ( path, checksum) )
61+ }
62+
63+ /// Parse filename from a Content-Disposition header value.
64+ fn parse_filename_from_content_disposition ( header_value : & str ) -> Option < String > {
65+ // Simple extraction for headers like: attachment; filename="example.txt"
66+ header_value. split ( ';' ) . find_map ( |part| {
67+ let part = part. trim ( ) ;
68+ if part. starts_with ( "filename=" ) {
69+ Some (
70+ part. trim_start_matches ( "filename=" )
71+ . trim_matches ( '"' )
72+ . to_string ( ) ,
73+ )
74+ } else {
75+ None
76+ }
77+ } )
78+ }
79+
80+ /// Extracts filename from the URL path (fallback)
81+ fn extract_filename_from_url ( url : & str ) -> Option < String > {
82+ url. split ( '/' )
83+ . filter ( |s| !s. is_empty ( ) )
84+ . last ( )
85+ . map ( |s| s. to_string ( ) )
86+ }
2487
2588#[ tokio:: test]
2689async fn image_upload_download_roundtrip ( ) -> Result < ( ) , Box < dyn std:: error:: Error > > {
@@ -30,20 +93,13 @@ async fn image_upload_download_roundtrip() -> Result<(), Box<dyn std::error::Err
3093 "http://download.cirros-cloud.net/{ver}/cirros-{ver}-x86_64-disk.img" ,
3194 ver = cirros_ver
3295 ) ;
33- let response = reqwest:: get ( target) . await ?;
34- let ( mut img_data, fname) = {
35- let fname = response
36- . url ( )
37- . path_segments ( )
38- . and_then ( |segments| segments. last ( ) )
39- . and_then ( |name| if name. is_empty ( ) { None } else { Some ( name) } )
40- . unwrap_or ( "tmp.bin" ) ;
41-
42- let fname = tmp_dir. path ( ) . join ( fname) ;
43- ( File :: create ( fname. clone ( ) ) ?, fname)
44- } ;
45- let mut content = Cursor :: new ( response. bytes ( ) . await ?) ;
46- copy ( & mut content, & mut img_data) ?;
96+ let ( fname, checksum) = download_with_md5_and_filename ( & target, & tmp_dir)
97+ . await
98+ . expect ( "Download failed" ) ;
99+ assert_eq ! (
100+ "c8fc807773e5354afe61636071771906" , checksum,
101+ "Download checksum matches the expected"
102+ ) ;
47103
48104 let img_name = format ! (
49105 "test-rust-{}" ,
@@ -92,10 +148,18 @@ async fn image_upload_download_roundtrip() -> Result<(), Box<dyn std::error::Err
92148 . success ( ) ;
93149
94150 // Compare files
95- diff_files (
96- & mut img_data,
97- & mut File :: open ( & download_data_fname) . unwrap ( ) ,
98- ) ;
151+ let mut file = File :: open ( download_data_fname) . await ?;
152+ let mut buffer = [ 0u8 ; 8192 ] ;
153+ let mut context = Context :: new ( ) ;
154+
155+ loop {
156+ let n = file. read ( & mut buffer) . await ?;
157+ if n == 0 {
158+ break ;
159+ }
160+ context. consume ( & buffer[ ..n] ) ;
161+ }
162+ let download_digest = context. compute ( ) ;
99163
100164 // Delete image
101165 Command :: cargo_bin ( "osc" ) ?
@@ -106,5 +170,11 @@ async fn image_upload_download_roundtrip() -> Result<(), Box<dyn std::error::Err
106170 . assert ( )
107171 . success ( ) ;
108172
173+ assert_eq ! (
174+ format!( "{:x}" , download_digest) ,
175+ checksum,
176+ "Checksums match"
177+ ) ;
178+
109179 Ok ( ( ) )
110180}
0 commit comments