nautilus_testkit/
files.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16use std::{
17    cmp,
18    ffi::OsString,
19    fmt::Display,
20    fs::{File, OpenOptions, remove_file},
21    io::{BufReader, BufWriter, Read, copy},
22    path::{Path, PathBuf},
23    sync::OnceLock,
24    thread::sleep,
25    time::{Duration, Instant},
26};
27
28use aws_lc_rs::digest::{self, Context};
29use nautilus_core::hex;
30use nautilus_network::retry::RetryConfig;
31use parking_lot::Mutex;
32use rand::{RngExt, rng};
33use reqwest::blocking::Client;
34use serde_json::Value;
35
36static LARGE_CHECKSUMS_LOCK: OnceLock<Mutex<()>> = OnceLock::new();
37
38fn lock_large_checksums() -> parking_lot::MutexGuard<'static, ()> {
39    LARGE_CHECKSUMS_LOCK.get_or_init(|| Mutex::new(())).lock()
40}
41
42#[derive(Debug)]
43enum DownloadError {
44    Retryable(String),
45    NonRetryable(String),
46}
47
48impl Display for DownloadError {
49    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
50        match self {
51            Self::Retryable(msg) => write!(f, "Retryable error: {msg}"),
52            Self::NonRetryable(msg) => write!(f, "Non-retryable error: {msg}"),
53        }
54    }
55}
56
57impl std::error::Error for DownloadError {}
58
59fn execute_with_retry_blocking<T, E, F>(
60    config: &RetryConfig,
61    mut op: F,
62    should_retry: impl Fn(&E) -> bool,
63) -> Result<T, E>
64where
65    E: std::error::Error,
66    F: FnMut() -> Result<T, E>,
67{
68    let start = Instant::now();
69    let mut delay = Duration::from_millis(config.initial_delay_ms);
70
71    for attempt in 0..=config.max_retries {
72        if attempt > 0 && !config.immediate_first {
73            let jitter = rng().random_range(0..=config.jitter_ms);
74            let sleep_for = delay + Duration::from_millis(jitter);
75            sleep(sleep_for);
76            delay = cmp::min(
77                next_retry_delay(delay, config.backoff_factor),
78                Duration::from_millis(config.max_delay_ms),
79            );
80        }
81
82        if let Some(max_total) = config.max_elapsed_ms
83            && start.elapsed() >= Duration::from_millis(max_total)
84        {
85            break;
86        }
87
88        match op() {
89            Ok(v) => return Ok(v),
90            Err(e) if attempt < config.max_retries && should_retry(&e) => {}
91            Err(e) => return Err(e),
92        }
93    }
94
95    op()
96}
97
98fn next_retry_delay(delay: Duration, backoff_factor: f64) -> Duration {
99    let next = delay.as_secs_f64() * backoff_factor;
100    if next.is_nan() || next.is_sign_negative() {
101        return Duration::ZERO;
102    }
103
104    Duration::try_from_secs_f64(next).unwrap_or(Duration::MAX)
105}
106
107/// Ensures that a file exists at the specified path by downloading it if necessary.
108///
109/// If the file already exists, it checks the integrity of the file using a SHA-256 checksum
110/// from the optional `checksums` file. If the checksum is valid, the function exits early. If
111/// the checksum is invalid or missing, the function updates the checksums file with the correct
112/// hash for the existing file without redownloading it.
113///
114/// If the file does not exist, it downloads the file from the specified `url` and updates the
115/// checksums file (if provided) with the calculated SHA-256 checksum of the downloaded file.
116///
117/// The `timeout_secs` parameter specifies the timeout in seconds for the HTTP request.
118/// If `None` is provided, a default timeout of 30 seconds will be used.
119///
120/// # Errors
121///
122/// Returns an error if:
123/// - The HTTP request cannot be sent or returns a non-success status code.
124/// - Any I/O operation fails during file creation, reading, or writing.
125/// - Checksum verification or JSON parsing fails.
126pub fn ensure_file_exists_or_download_http(
127    filepath: &Path,
128    url: &str,
129    checksums: Option<&Path>,
130    timeout_secs: Option<u64>,
131) -> anyhow::Result<()> {
132    ensure_file_exists_or_download_http_with_config(
133        filepath,
134        url,
135        checksums,
136        timeout_secs.unwrap_or(30),
137        None,
138        None,
139    )
140}
141
142/// Ensures that a file exists at the specified path by downloading it if necessary, with a custom timeout.
143///
144/// # Errors
145///
146/// Returns an error if:
147/// - The HTTP request cannot be sent or returns a non-success status code after retries.
148/// - Any I/O operation fails during file creation, reading, or writing.
149/// - Checksum verification or JSON parsing fails.
150pub fn ensure_file_exists_or_download_http_with_timeout(
151    filepath: &Path,
152    url: &str,
153    checksums: Option<&Path>,
154    timeout_secs: u64,
155) -> anyhow::Result<()> {
156    ensure_file_exists_or_download_http_with_config(
157        filepath,
158        url,
159        checksums,
160        timeout_secs,
161        None,
162        None,
163    )
164}
165
166/// Ensures that a file exists at the specified path by downloading it if necessary,
167/// with custom timeout, retry config, and initial jitter delay.
168///
169/// # Parameters
170///
171/// - `filepath`: The path where the file should exist.
172/// - `url`: The URL to download from if the file doesn't exist.
173/// - `checksums`: Optional path to checksums file for verification.
174/// - `timeout_secs`: Timeout in seconds for HTTP requests.
175/// - `retry_config`: Optional custom retry configuration (uses sensible defaults if None).
176/// - `initial_jitter_ms`: Optional initial jitter delay in milliseconds before download (defaults to 100-600ms if None).
177///
178/// # Errors
179///
180/// Returns an error if:
181/// - The HTTP request cannot be sent or returns a non-success status code after retries.
182/// - Any I/O operation fails during file creation, reading, or writing.
183/// - Checksum verification or JSON parsing fails.
184pub fn ensure_file_exists_or_download_http_with_config(
185    filepath: &Path,
186    url: &str,
187    checksums: Option<&Path>,
188    timeout_secs: u64,
189    retry_config: Option<RetryConfig>,
190    initial_jitter_ms: Option<u64>,
191) -> anyhow::Result<()> {
192    // Local/cached file path: accept the file if it exists, updating the
193    // checksum record when it differs (e.g. after local regeneration).
194    // This is intentionally lenient - download verification below is strict.
195    if filepath.exists() {
196        println!("File already exists (local/cached): {}", filepath.display());
197
198        if let Some(checksums_file) = checksums {
199            let _guard = lock_large_checksums();
200
201            if verify_sha256_checksum(filepath, checksums_file)? {
202                println!("Checksum verified");
203                return Ok(());
204            }
205
206            let new_checksum = calculate_sha256(filepath)?;
207            println!("Updating checksum for local file: {new_checksum}");
208            update_sha256_checksums(filepath, checksums_file, &new_checksum)?;
209            return Ok(());
210        }
211        return Ok(());
212    }
213
214    // Add a small random delay to avoid bursting the remote server when
215    // many downloads start concurrently. Can be disabled by passing Some(0).
216    if let Some(jitter_ms) = initial_jitter_ms {
217        if jitter_ms > 0 {
218            sleep(Duration::from_millis(jitter_ms));
219        }
220    } else {
221        let jitter_delay = {
222            let mut r = rng();
223            Duration::from_millis(r.random_range(100..=600))
224        };
225        sleep(jitter_delay);
226    }
227
228    download_file(filepath, url, timeout_secs, retry_config.clone())?;
229
230    // Verify checksum after download, retry once on mismatch (corrupt download)
231    if let Some(checksums_file) = checksums {
232        let guard = lock_large_checksums();
233
234        if !verify_sha256_checksum(filepath, checksums_file)? {
235            let actual = calculate_sha256(filepath)?;
236            println!("Checksum mismatch after download (calculated {actual}), retrying...");
237            remove_file(filepath)?;
238            drop(guard);
239
240            download_file(filepath, url, timeout_secs, retry_config)?;
241
242            let _guard = lock_large_checksums();
243
244            if !verify_sha256_checksum(filepath, checksums_file)? {
245                let actual = calculate_sha256(filepath)?;
246                remove_file(filepath)?;
247                anyhow::bail!(
248                    "Checksum mismatch after retry for {} (calculated {actual})",
249                    filepath.file_name().unwrap_or_default().display(),
250                );
251            }
252        }
253    }
254
255    Ok(())
256}
257
258fn download_file(
259    filepath: &Path,
260    url: &str,
261    timeout_secs: u64,
262    retry_config: Option<RetryConfig>,
263) -> anyhow::Result<()> {
264    // Validate HTTPS for security in production builds,
265    // HTTP is intentionally allowed in test builds for local test servers (127.0.0.1),
266    // CodeQL flags this as "non-https-url" but it's a deliberate design choice for testkit.
267    #[cfg(not(test))]
268    if !url.starts_with("https://") {
269        anyhow::bail!("URL must use HTTPS protocol for security: {url}");
270    }
271
272    println!("Downloading file from {url} to {}", filepath.display());
273
274    if let Some(parent) = filepath.parent() {
275        std::fs::create_dir_all(parent)?;
276    }
277
278    let client = Client::builder()
279        .timeout(Duration::from_secs(timeout_secs))
280        .build()?;
281
282    let cfg = if let Some(config) = retry_config {
283        config
284    } else {
285        // Default production config
286        let max_retries = 5u32;
287        let op_timeout_ms = timeout_secs.saturating_mul(1000);
288        // Make the provided timeout a hard ceiling for total elapsed time.
289        // Split it across attempts (at least 1000 ms per attempt) and cap total at op_timeout_ms.
290        let per_attempt_ms = std::cmp::max(1000u64, op_timeout_ms / (u64::from(max_retries) + 1));
291        RetryConfig {
292            max_retries,
293            initial_delay_ms: 1_000,
294            max_delay_ms: 10_000,
295            backoff_factor: 2.0,
296            jitter_ms: 1_000,
297            operation_timeout_ms: Some(per_attempt_ms),
298            immediate_first: false,
299            max_elapsed_ms: Some(op_timeout_ms),
300        }
301    };
302
303    let partial_path = partial_path_for(filepath);
304
305    let op = || -> Result<(), DownloadError> {
306        // Discard any leftover partial from a prior attempt before writing
307        let _ = remove_file(&partial_path);
308
309        match client.get(url).send() {
310            Ok(mut response) => {
311                let status = response.status();
312                if status.is_success() {
313                    let mut out = File::create(&partial_path)
314                        .map_err(|e| DownloadError::NonRetryable(e.to_string()))?;
315                    // Stream body to a sibling .partial path so a truncated copy never reaches the final filepath,
316                    // body-stream errors (TCP reset, chunked-encoding decode, premature EOF) are typically transient,
317                    // so surface them as Retryable.
318                    if let Err(e) = copy(&mut response, &mut out) {
319                        drop(out);
320                        let _ = remove_file(&partial_path);
321                        return Err(DownloadError::Retryable(format!("body stream error: {e}")));
322                    }
323                    drop(out);
324
325                    if let Err(e) = std::fs::rename(&partial_path, filepath) {
326                        let _ = remove_file(&partial_path);
327                        return Err(DownloadError::NonRetryable(format!(
328                            "rename {} -> {} failed: {e}",
329                            partial_path.display(),
330                            filepath.display(),
331                        )));
332                    }
333                    println!("File downloaded to {}", filepath.display());
334                    Ok(())
335                } else if status.is_server_error()
336                    || status.as_u16() == 429
337                    || status.as_u16() == 408
338                {
339                    println!("HTTP error {status}, retrying...");
340                    Err(DownloadError::Retryable(format!("HTTP {status}")))
341                } else {
342                    // Preserve existing error text used by tests
343                    Err(DownloadError::NonRetryable(format!(
344                        "Client error: HTTP {status}"
345                    )))
346                }
347            }
348            Err(e) => {
349                println!("Request failed: {e}");
350                Err(DownloadError::Retryable(e.to_string()))
351            }
352        }
353    };
354
355    let should_retry = |e: &DownloadError| matches!(e, DownloadError::Retryable(_));
356
357    execute_with_retry_blocking(&cfg, op, should_retry).map_err(|e| anyhow::anyhow!(e.to_string()))
358}
359
360fn partial_path_for(filepath: &Path) -> PathBuf {
361    // Suffix with pid + nanos so the staging path is unique per invocation,
362    // preventing collision with a real user file or another concurrent downloader.
363    let nanos = std::time::SystemTime::now()
364        .duration_since(std::time::UNIX_EPOCH)
365        .map_or(0, |d| d.as_nanos());
366    let mut p: OsString = filepath.as_os_str().to_owned();
367    p.push(format!(".partial.{}.{}", std::process::id(), nanos));
368    PathBuf::from(p)
369}
370
371fn calculate_sha256(filepath: &Path) -> anyhow::Result<String> {
372    let mut file = File::open(filepath)?;
373    let mut ctx = Context::new(&digest::SHA256);
374    let mut buffer = [0u8; 4096];
375
376    loop {
377        let count = file.read(&mut buffer)?;
378        if count == 0 {
379            break;
380        }
381        ctx.update(&buffer[..count]);
382    }
383
384    let digest = ctx.finish();
385    Ok(hex::encode(digest.as_ref()))
386}
387
388fn verify_sha256_checksum(filepath: &Path, checksums: &Path) -> anyhow::Result<bool> {
389    let file = File::open(checksums)?;
390    let reader = BufReader::new(file);
391    let checksums: Value = serde_json::from_reader(reader)?;
392
393    let filename = filepath.file_name().unwrap().to_str().unwrap();
394    if let Some(expected_checksum) = checksums.get(filename) {
395        let expected_checksum_str = expected_checksum.as_str().unwrap();
396        let expected_hash = expected_checksum_str
397            .strip_prefix("sha256:")
398            .unwrap_or(expected_checksum_str);
399        let calculated_checksum = calculate_sha256(filepath)?;
400        if expected_hash == calculated_checksum {
401            return Ok(true);
402        }
403    }
404
405    Ok(false)
406}
407
408fn update_sha256_checksums(
409    filepath: &Path,
410    checksums_file: &Path,
411    new_checksum: &str,
412) -> anyhow::Result<()> {
413    let checksums: Value = if checksums_file.exists() {
414        let file = File::open(checksums_file)?;
415        let reader = BufReader::new(file);
416        serde_json::from_reader(reader)?
417    } else {
418        serde_json::json!({})
419    };
420
421    let mut checksums_map = checksums.as_object().unwrap().clone();
422
423    // Add or update the checksum
424    let filename = filepath.file_name().unwrap().to_str().unwrap().to_string();
425    let prefixed_checksum = format!("sha256:{new_checksum}");
426    checksums_map.insert(filename, Value::String(prefixed_checksum));
427
428    let file = OpenOptions::new()
429        .write(true)
430        .create(true)
431        .truncate(true)
432        .open(checksums_file)?;
433    let writer = BufWriter::new(file);
434    serde_json::to_writer_pretty(writer, &serde_json::Value::Object(checksums_map))?;
435
436    Ok(())
437}
438
439#[cfg(test)]
440mod tests {
441    use std::{
442        fs,
443        io::{BufWriter, Write},
444        net::SocketAddr,
445        sync::{
446            Arc,
447            atomic::{AtomicUsize, Ordering},
448        },
449    };
450
451    use axum::{Router, http::StatusCode, routing::get, serve};
452    use rstest::*;
453    use serde_json::{json, to_writer};
454    use tempfile::TempDir;
455    use tokio::{
456        net::TcpListener,
457        task,
458        time::{Duration, sleep},
459    };
460
461    use super::*;
462
463    /// Creates a fast, deterministic retry config for tests.
464    /// Uses very short delays to make tests run quickly without introducing flakiness.
465    fn test_retry_config() -> RetryConfig {
466        RetryConfig {
467            max_retries: 5,
468            initial_delay_ms: 10,
469            max_delay_ms: 50,
470            backoff_factor: 2.0,
471            jitter_ms: 5,
472            operation_timeout_ms: Some(500),
473            immediate_first: false,
474            max_elapsed_ms: Some(2000),
475        }
476    }
477
478    #[rstest]
479    #[case::nan(f64::NAN, Duration::ZERO)]
480    #[case::negative(-1.0, Duration::ZERO)]
481    #[case::infinite(f64::INFINITY, Duration::MAX)]
482    #[case::normal(2.0, Duration::from_millis(20))]
483    fn next_retry_delay_handles_edge_factors(
484        #[case] backoff_factor: f64,
485        #[case] expected: Duration,
486    ) {
487        let delay = Duration::from_millis(10);
488
489        assert_eq!(next_retry_delay(delay, backoff_factor), expected);
490    }
491
492    async fn setup_test_server(
493        server_content: Option<String>,
494        status_code: StatusCode,
495    ) -> SocketAddr {
496        let server_content = Arc::new(server_content);
497        let server_content_clone = server_content.clone();
498        let app = Router::new().route(
499            "/testfile.txt",
500            get(move || {
501                let server_content = server_content_clone.clone();
502                async move {
503                    let response_body = match &*server_content {
504                        Some(content) => content.clone(),
505                        None => "File not found".to_string(),
506                    };
507                    (status_code, response_body)
508                }
509            }),
510        );
511
512        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
513        let addr = listener.local_addr().unwrap();
514        let server = serve(listener, app);
515
516        task::spawn(async move {
517            if let Err(e) = server.await {
518                eprintln!("server error: {e}");
519            }
520        });
521
522        sleep(Duration::from_millis(100)).await;
523
524        addr
525    }
526
527    #[tokio::test]
528    async fn test_file_already_exists() {
529        let temp_dir = TempDir::new().unwrap();
530        let file_path = temp_dir.path().join("testfile.txt");
531        fs::write(&file_path, "Existing file content").unwrap();
532
533        let url = "http://example.com/testfile.txt".to_string();
534        let result = ensure_file_exists_or_download_http(&file_path, &url, None, Some(5));
535
536        assert!(result.is_ok());
537        let content = fs::read_to_string(&file_path).unwrap();
538        assert_eq!(content, "Existing file content");
539    }
540
541    #[tokio::test]
542    async fn test_download_file_success() {
543        let temp_dir = TempDir::new().unwrap();
544        let filepath = temp_dir.path().join("testfile.txt");
545        let filepath_clone = filepath.clone();
546
547        let server_content = "Server file content".to_string();
548        let status_code = StatusCode::OK;
549        let addr = setup_test_server(Some(server_content.clone()), status_code).await;
550        let url = format!("http://{addr}/testfile.txt");
551
552        let result = tokio::task::spawn_blocking(move || {
553            ensure_file_exists_or_download_http_with_config(
554                &filepath_clone,
555                &url,
556                None,
557                5,
558                Some(test_retry_config()),
559                Some(0),
560            )
561        })
562        .await
563        .unwrap();
564
565        assert!(result.is_ok());
566        let content = fs::read_to_string(&filepath).unwrap();
567        assert_eq!(content, server_content);
568    }
569
570    #[tokio::test]
571    async fn test_download_file_not_found() {
572        let temp_dir = TempDir::new().unwrap();
573        let file_path = temp_dir.path().join("testfile.txt");
574
575        let server_content = None;
576        let status_code = StatusCode::NOT_FOUND;
577        let addr = setup_test_server(server_content, status_code).await;
578        let url = format!("http://{addr}/testfile.txt");
579
580        let result = tokio::task::spawn_blocking(move || {
581            ensure_file_exists_or_download_http_with_config(
582                &file_path,
583                &url,
584                None,
585                1,
586                Some(test_retry_config()),
587                Some(0),
588            )
589        })
590        .await
591        .unwrap();
592
593        assert!(result.is_err());
594        let err_msg = format!("{}", result.unwrap_err());
595        assert!(
596            err_msg.contains("Client error: HTTP"),
597            "Unexpected error message: {err_msg}"
598        );
599    }
600
601    #[tokio::test]
602    async fn test_network_error() {
603        let temp_dir = TempDir::new().unwrap();
604        let file_path = temp_dir.path().join("testfile.txt");
605
606        // Use an unreachable address to simulate a network error
607        let url = "http://127.0.0.1:0/testfile.txt".to_string();
608
609        let result = tokio::task::spawn_blocking(move || {
610            ensure_file_exists_or_download_http_with_config(
611                &file_path,
612                &url,
613                None,
614                2,
615                Some(test_retry_config()),
616                Some(0),
617            )
618        })
619        .await
620        .unwrap();
621
622        assert!(result.is_err());
623        let err_msg = format!("{}", result.unwrap_err());
624        assert!(
625            err_msg.contains("error"),
626            "Unexpected error message: {err_msg}"
627        );
628    }
629
630    #[tokio::test]
631    async fn test_retry_then_success_on_500() {
632        let temp_dir = TempDir::new().unwrap();
633        let filepath = temp_dir.path().join("testfile.txt");
634        let filepath_clone = filepath.clone();
635
636        let counter = Arc::new(AtomicUsize::new(0));
637        let counter_clone = counter.clone();
638
639        let app = Router::new().route(
640            "/testfile.txt",
641            get(move || {
642                let c = counter_clone.clone();
643                async move {
644                    let n = c.fetch_add(1, Ordering::SeqCst);
645                    if n < 2 {
646                        (StatusCode::INTERNAL_SERVER_ERROR, "temporary error")
647                    } else {
648                        (StatusCode::OK, "eventual success")
649                    }
650                }
651            }),
652        );
653
654        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
655        let addr = listener.local_addr().unwrap();
656        let server = serve(listener, app);
657        task::spawn(async move {
658            let _ = server.await;
659        });
660        sleep(Duration::from_millis(100)).await;
661
662        let url = format!("http://{addr}/testfile.txt");
663
664        let result = tokio::task::spawn_blocking(move || {
665            ensure_file_exists_or_download_http_with_config(
666                &filepath_clone,
667                &url,
668                None,
669                5,
670                Some(test_retry_config()),
671                Some(0),
672            )
673        })
674        .await
675        .unwrap();
676
677        assert!(result.is_ok());
678        let content = std::fs::read_to_string(&filepath).unwrap();
679        assert_eq!(content, "eventual success");
680        assert!(counter.load(Ordering::SeqCst) >= 2);
681    }
682
683    #[tokio::test]
684    async fn test_retry_then_success_on_429() {
685        let temp_dir = TempDir::new().unwrap();
686        let filepath = temp_dir.path().join("testfile.txt");
687        let filepath_clone = filepath.clone();
688
689        let counter = Arc::new(AtomicUsize::new(0));
690        let counter_clone = counter.clone();
691
692        let app = Router::new().route(
693            "/testfile.txt",
694            get(move || {
695                let c = counter_clone.clone();
696                async move {
697                    let n = c.fetch_add(1, Ordering::SeqCst);
698                    if n < 1 {
699                        (StatusCode::TOO_MANY_REQUESTS, "rate limited")
700                    } else {
701                        (StatusCode::OK, "ok after retry")
702                    }
703                }
704            }),
705        );
706
707        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
708        let addr = listener.local_addr().unwrap();
709        let server = serve(listener, app);
710        task::spawn(async move {
711            let _ = server.await;
712        });
713        sleep(Duration::from_millis(100)).await;
714
715        let url = format!("http://{addr}/testfile.txt");
716
717        let result = tokio::task::spawn_blocking(move || {
718            ensure_file_exists_or_download_http_with_config(
719                &filepath_clone,
720                &url,
721                None,
722                5,
723                Some(test_retry_config()),
724                Some(0),
725            )
726        })
727        .await
728        .unwrap();
729
730        assert!(result.is_ok());
731        let content = std::fs::read_to_string(&filepath).unwrap();
732        assert_eq!(content, "ok after retry");
733        assert!(counter.load(Ordering::SeqCst) >= 2);
734    }
735
736    #[tokio::test]
737    async fn test_no_retry_on_404() {
738        let temp_dir = TempDir::new().unwrap();
739        let filepath = temp_dir.path().join("testfile.txt");
740        let filepath_clone = filepath.clone();
741
742        let counter = Arc::new(AtomicUsize::new(0));
743        let counter_clone = counter.clone();
744
745        let app = Router::new().route(
746            "/testfile.txt",
747            get(move || {
748                let c = counter_clone.clone();
749                async move {
750                    c.fetch_add(1, Ordering::SeqCst);
751                    (StatusCode::NOT_FOUND, "missing")
752                }
753            }),
754        );
755
756        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
757        let addr = listener.local_addr().unwrap();
758        let server = serve(listener, app);
759        task::spawn(async move {
760            let _ = server.await;
761        });
762        sleep(Duration::from_millis(100)).await;
763
764        let url = format!("http://{addr}/testfile.txt");
765
766        let result = tokio::task::spawn_blocking(move || {
767            ensure_file_exists_or_download_http_with_config(
768                &filepath_clone,
769                &url,
770                None,
771                5,
772                Some(test_retry_config()),
773                Some(0),
774            )
775        })
776        .await
777        .unwrap();
778
779        assert!(result.is_err());
780        assert_eq!(counter.load(Ordering::SeqCst), 1, "should not retry on 404");
781    }
782
783    #[tokio::test]
784    async fn test_checksum_mismatch_retry_then_success() {
785        let temp_dir = TempDir::new().unwrap();
786        let filepath = temp_dir.path().join("testfile.txt");
787        let filepath_clone = filepath.clone();
788
789        let good_content = "correct content";
790        let good_checksum = calculate_sha256_bytes(good_content.as_bytes());
791
792        let checksums_path = temp_dir.path().join("checksums.json");
793        let checksums_data = json!({
794            "testfile.txt": format!("sha256:{good_checksum}")
795        });
796        let checksums_file = File::create(&checksums_path).unwrap();
797        to_writer(BufWriter::new(checksums_file), &checksums_data).unwrap();
798        let checksums_clone = checksums_path.clone();
799
800        // First request returns corrupt data, second returns correct data
801        let counter = Arc::new(AtomicUsize::new(0));
802        let counter_clone = counter.clone();
803
804        let app = Router::new().route(
805            "/testfile.txt",
806            get(move || {
807                let c = counter_clone.clone();
808                async move {
809                    let n = c.fetch_add(1, Ordering::SeqCst);
810                    if n == 0 {
811                        (StatusCode::OK, "corrupt data")
812                    } else {
813                        (StatusCode::OK, "correct content")
814                    }
815                }
816            }),
817        );
818
819        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
820        let addr = listener.local_addr().unwrap();
821        let server = serve(listener, app);
822        task::spawn(async move {
823            let _ = server.await;
824        });
825        sleep(Duration::from_millis(100)).await;
826
827        let url = format!("http://{addr}/testfile.txt");
828
829        let result = tokio::task::spawn_blocking(move || {
830            ensure_file_exists_or_download_http_with_config(
831                &filepath_clone,
832                &url,
833                Some(&checksums_clone),
834                5,
835                Some(test_retry_config()),
836                Some(0),
837            )
838        })
839        .await
840        .unwrap();
841
842        assert!(result.is_ok());
843        let content = fs::read_to_string(&filepath).unwrap();
844        assert_eq!(content, good_content);
845        assert_eq!(counter.load(Ordering::SeqCst), 2);
846    }
847
848    #[tokio::test]
849    async fn test_checksum_mismatch_retry_then_fail() {
850        let temp_dir = TempDir::new().unwrap();
851        let filepath = temp_dir.path().join("testfile.txt");
852        let filepath_clone = filepath.clone();
853
854        // Checksum for content that the server will never return
855        let checksums_path = temp_dir.path().join("checksums.json");
856        let checksums_data = json!({
857            "testfile.txt": "sha256:0000000000000000000000000000000000000000000000000000000000000000"
858        });
859        let checksums_file = File::create(&checksums_path).unwrap();
860        to_writer(BufWriter::new(checksums_file), &checksums_data).unwrap();
861        let checksums_clone = checksums_path.clone();
862
863        let counter = Arc::new(AtomicUsize::new(0));
864        let counter_clone = counter.clone();
865
866        let app = Router::new().route(
867            "/testfile.txt",
868            get(move || {
869                let c = counter_clone.clone();
870                async move {
871                    c.fetch_add(1, Ordering::SeqCst);
872                    (StatusCode::OK, "always wrong content")
873                }
874            }),
875        );
876
877        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
878        let addr = listener.local_addr().unwrap();
879        let server = serve(listener, app);
880        task::spawn(async move {
881            let _ = server.await;
882        });
883        sleep(Duration::from_millis(100)).await;
884
885        let url = format!("http://{addr}/testfile.txt");
886
887        let result = tokio::task::spawn_blocking(move || {
888            ensure_file_exists_or_download_http_with_config(
889                &filepath_clone,
890                &url,
891                Some(&checksums_clone),
892                5,
893                Some(test_retry_config()),
894                Some(0),
895            )
896        })
897        .await
898        .unwrap();
899
900        assert!(result.is_err());
901        let err_msg = format!("{}", result.unwrap_err());
902        assert!(err_msg.contains("Checksum mismatch after retry"));
903        assert_eq!(
904            counter.load(Ordering::SeqCst),
905            2,
906            "should download exactly twice"
907        );
908        assert!(!filepath.exists(), "corrupt file should be cleaned up");
909    }
910
911    /// First call overstates Content-Length to fail reqwest mid-stream; subsequent calls succeed.
912    async fn truncated_then_full_server(good_body: &'static str) -> (SocketAddr, Arc<AtomicUsize>) {
913        use tokio::io::{AsyncReadExt, AsyncWriteExt};
914
915        let counter = Arc::new(AtomicUsize::new(0));
916        let counter_clone = counter.clone();
917
918        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
919        let addr = listener.local_addr().unwrap();
920
921        task::spawn(async move {
922            loop {
923                let Ok((mut sock, _)) = listener.accept().await else {
924                    break;
925                };
926                let counter = counter_clone.clone();
927
928                task::spawn(async move {
929                    let mut buf = [0u8; 1024];
930                    let _ = sock.read(&mut buf).await;
931                    let n = counter.fetch_add(1, Ordering::SeqCst);
932
933                    if n == 0 {
934                        let resp = b"HTTP/1.1 200 OK\r\nContent-Length: 100\r\nConnection: close\r\n\r\nshort";
935                        let _ = sock.write_all(resp).await;
936                    } else {
937                        let resp = format!(
938                            "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
939                            good_body.len(),
940                            good_body,
941                        );
942                        let _ = sock.write_all(resp.as_bytes()).await;
943                    }
944                    let _ = sock.shutdown().await;
945                });
946            }
947        });
948
949        sleep(Duration::from_millis(100)).await;
950        (addr, counter)
951    }
952
953    fn count_partial_siblings(filepath: &Path) -> usize {
954        let parent = filepath.parent().unwrap();
955        let stem = filepath.file_name().unwrap().to_string_lossy().into_owned();
956        let prefix = format!("{stem}.partial.");
957        fs::read_dir(parent)
958            .unwrap()
959            .filter_map(Result::ok)
960            .filter(|e| e.file_name().to_string_lossy().starts_with(&prefix))
961            .count()
962    }
963
964    #[tokio::test]
965    async fn test_body_truncation_retries_and_recovers() {
966        let temp_dir = TempDir::new().unwrap();
967        let filepath = temp_dir.path().join("testfile.txt");
968        let filepath_clone = filepath.clone();
969
970        let (addr, counter) = truncated_then_full_server("complete payload").await;
971        let url = format!("http://{addr}/testfile.txt");
972
973        let result = tokio::task::spawn_blocking(move || {
974            ensure_file_exists_or_download_http_with_config(
975                &filepath_clone,
976                &url,
977                None,
978                5,
979                Some(test_retry_config()),
980                Some(0),
981            )
982        })
983        .await
984        .unwrap();
985
986        assert!(result.is_ok(), "should retry past the truncated response");
987        assert_eq!(counter.load(Ordering::SeqCst), 2);
988        assert_eq!(
989            count_partial_siblings(&filepath),
990            0,
991            "no .partial siblings must remain after success",
992        );
993        let content = fs::read_to_string(&filepath).unwrap();
994        assert_eq!(content, "complete payload");
995    }
996
997    #[tokio::test]
998    async fn test_body_truncation_exhausts_retries_leaves_no_corrupt_file() {
999        use tokio::io::{AsyncReadExt, AsyncWriteExt};
1000
1001        let temp_dir = TempDir::new().unwrap();
1002        let filepath = temp_dir.path().join("testfile.txt");
1003        let filepath_clone = filepath.clone();
1004
1005        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
1006        let addr = listener.local_addr().unwrap();
1007
1008        task::spawn(async move {
1009            loop {
1010                let Ok((mut sock, _)) = listener.accept().await else {
1011                    break;
1012                };
1013
1014                task::spawn(async move {
1015                    let mut buf = [0u8; 1024];
1016                    let _ = sock.read(&mut buf).await;
1017                    let resp =
1018                        b"HTTP/1.1 200 OK\r\nContent-Length: 100\r\nConnection: close\r\n\r\nshort";
1019                    let _ = sock.write_all(resp).await;
1020                    let _ = sock.shutdown().await;
1021                });
1022            }
1023        });
1024
1025        sleep(Duration::from_millis(100)).await;
1026        let url = format!("http://{addr}/testfile.txt");
1027
1028        let result = tokio::task::spawn_blocking(move || {
1029            ensure_file_exists_or_download_http_with_config(
1030                &filepath_clone,
1031                &url,
1032                None,
1033                5,
1034                Some(test_retry_config()),
1035                Some(0),
1036            )
1037        })
1038        .await
1039        .unwrap();
1040
1041        assert!(result.is_err(), "all retries should fail");
1042        assert!(!filepath.exists(), "no corrupt file at final path");
1043        assert_eq!(
1044            count_partial_siblings(&filepath),
1045            0,
1046            "no .partial siblings may leak after exhausted retries",
1047        );
1048    }
1049
1050    #[rstest]
1051    fn test_partial_path_for_is_unique_and_marked() {
1052        let target = Path::new("/tmp/data.parquet");
1053        let stem = target.file_name().unwrap().to_string_lossy().into_owned();
1054        let expected_prefix = format!("{stem}.partial.");
1055
1056        let a = partial_path_for(target);
1057        let b = partial_path_for(target);
1058        assert_ne!(a, b, "partial paths must be unique per call");
1059        for p in [&a, &b] {
1060            let name = p.file_name().unwrap().to_string_lossy().into_owned();
1061            assert!(name.starts_with(&expected_prefix), "got {name}");
1062        }
1063    }
1064
1065    #[tokio::test]
1066    async fn test_unrelated_partial_sibling_not_clobbered() {
1067        let temp_dir = TempDir::new().unwrap();
1068        let filepath = temp_dir.path().join("testfile.txt");
1069        let filepath_clone = filepath.clone();
1070        // A user-owned file whose name happens to start with `<target>.partial`
1071        let bystander = temp_dir.path().join("testfile.txt.partial");
1072        fs::write(&bystander, b"do not touch").unwrap();
1073
1074        let server_content = "downloaded".to_string();
1075        let addr = setup_test_server(Some(server_content.clone()), StatusCode::OK).await;
1076        let url = format!("http://{addr}/testfile.txt");
1077
1078        let result = tokio::task::spawn_blocking(move || {
1079            ensure_file_exists_or_download_http_with_config(
1080                &filepath_clone,
1081                &url,
1082                None,
1083                5,
1084                Some(test_retry_config()),
1085                Some(0),
1086            )
1087        })
1088        .await
1089        .unwrap();
1090
1091        assert!(result.is_ok());
1092        assert_eq!(fs::read_to_string(&filepath).unwrap(), server_content);
1093        assert_eq!(
1094            fs::read_to_string(&bystander).unwrap(),
1095            "do not touch",
1096            "unrelated sibling file must be preserved",
1097        );
1098    }
1099
1100    fn calculate_sha256_bytes(data: &[u8]) -> String {
1101        let mut ctx = digest::Context::new(&digest::SHA256);
1102        ctx.update(data);
1103        hex::encode(ctx.finish().as_ref())
1104    }
1105
1106    #[rstest]
1107    #[expect(clippy::panic_in_result_fn)]
1108    fn test_calculate_sha256() -> anyhow::Result<()> {
1109        let temp_dir = TempDir::new()?;
1110        let test_file_path = temp_dir.path().join("test_file.txt");
1111        let mut test_file = File::create(&test_file_path)?;
1112        let content = b"Hello, world!";
1113        test_file.write_all(content)?;
1114
1115        let expected_hash = "315f5bdb76d078c43b8ac0064e4a0164612b1fce77c869345bfc94c75894edd3";
1116        let calculated_hash = calculate_sha256(&test_file_path)?;
1117
1118        assert_eq!(calculated_hash, expected_hash);
1119        Ok(())
1120    }
1121
1122    #[rstest]
1123    #[expect(clippy::panic_in_result_fn)]
1124    fn test_verify_sha256_checksum() -> anyhow::Result<()> {
1125        let temp_dir = TempDir::new()?;
1126        let test_file_path = temp_dir.path().join("test_file.txt");
1127        let mut test_file = File::create(&test_file_path)?;
1128        let content = b"Hello, world!";
1129        test_file.write_all(content)?;
1130
1131        let calculated_checksum = calculate_sha256(&test_file_path)?;
1132
1133        // Create checksums.json containing the checksum
1134        let checksums_path = temp_dir.path().join("checksums.json");
1135        let checksums_data = json!({
1136            "test_file.txt": format!("sha256:{}", calculated_checksum)
1137        });
1138        let checksums_file = File::create(&checksums_path)?;
1139        let writer = BufWriter::new(checksums_file);
1140        to_writer(writer, &checksums_data)?;
1141
1142        let is_valid = verify_sha256_checksum(&test_file_path, &checksums_path)?;
1143        assert!(is_valid, "The checksum should be valid");
1144        Ok(())
1145    }
1146}
nautilus_testkit/files.rs

nautilus_testkit/
files.rs
