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

nautilus_testkit/
files.rs
