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    fmt::Display,
19    fs::{File, OpenOptions, remove_file},
20    io::{BufReader, BufWriter, Read, copy},
21    path::Path,
22    sync::{Mutex, OnceLock},
23    thread::sleep,
24    time::{Duration, Instant},
25};
26
27use aws_lc_rs::digest::{self, Context};
28use nautilus_core::hex;
29use nautilus_network::retry::RetryConfig;
30use rand::{RngExt, rng};
31use reqwest::blocking::Client;
32use serde_json::Value;
33
34static LARGE_CHECKSUMS_LOCK: OnceLock<Mutex<()>> = OnceLock::new();
35
36fn lock_large_checksums() -> anyhow::Result<std::sync::MutexGuard<'static, ()>> {
37    LARGE_CHECKSUMS_LOCK
38        .get_or_init(|| Mutex::new(()))
39        .lock()
40        .map_err(|e| anyhow::anyhow!("Failed to lock checksums file access: {e}"))
41}
42
43#[derive(Debug)]
44enum DownloadError {
45    Retryable(String),
46    NonRetryable(String),
47}
48
49impl Display for DownloadError {
50    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
51        match self {
52            Self::Retryable(msg) => write!(f, "Retryable error: {msg}"),
53            Self::NonRetryable(msg) => write!(f, "Non-retryable error: {msg}"),
54        }
55    }
56}
57
58impl std::error::Error for DownloadError {}
59
60fn execute_with_retry_blocking<T, E, F>(
61    config: &RetryConfig,
62    mut op: F,
63    should_retry: impl Fn(&E) -> bool,
64) -> Result<T, E>
65where
66    E: std::error::Error,
67    F: FnMut() -> Result<T, E>,
68{
69    let start = Instant::now();
70    let mut delay = Duration::from_millis(config.initial_delay_ms);
71
72    for attempt in 0..=config.max_retries {
73        if attempt > 0 && !config.immediate_first {
74            let jitter = rng().random_range(0..=config.jitter_ms);
75            let sleep_for = delay + Duration::from_millis(jitter);
76            sleep(sleep_for);
77            let next = (delay.as_millis() as f64 * config.backoff_factor) as u64;
78            delay = cmp::min(
79                Duration::from_millis(next),
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
100/// Ensures that a file exists at the specified path by downloading it if necessary.
101///
102/// If the file already exists, it checks the integrity of the file using a SHA-256 checksum
103/// from the optional `checksums` file. If the checksum is valid, the function exits early. If
104/// the checksum is invalid or missing, the function updates the checksums file with the correct
105/// hash for the existing file without redownloading it.
106///
107/// If the file does not exist, it downloads the file from the specified `url` and updates the
108/// checksums file (if provided) with the calculated SHA-256 checksum of the downloaded file.
109///
110/// The `timeout_secs` parameter specifies the timeout in seconds for the HTTP request.
111/// If `None` is provided, a default timeout of 30 seconds will be used.
112///
113/// # Errors
114///
115/// Returns an error if:
116/// - The HTTP request cannot be sent or returns a non-success status code.
117/// - Any I/O operation fails during file creation, reading, or writing.
118/// - Checksum verification or JSON parsing fails.
119pub fn ensure_file_exists_or_download_http(
120    filepath: &Path,
121    url: &str,
122    checksums: Option<&Path>,
123    timeout_secs: Option<u64>,
124) -> anyhow::Result<()> {
125    ensure_file_exists_or_download_http_with_config(
126        filepath,
127        url,
128        checksums,
129        timeout_secs.unwrap_or(30),
130        None,
131        None,
132    )
133}
134
135/// Ensures that a file exists at the specified path by downloading it if necessary, with a custom timeout.
136///
137/// # Errors
138///
139/// Returns an error if:
140/// - The HTTP request cannot be sent or returns a non-success status code after retries.
141/// - Any I/O operation fails during file creation, reading, or writing.
142/// - Checksum verification or JSON parsing fails.
143pub fn ensure_file_exists_or_download_http_with_timeout(
144    filepath: &Path,
145    url: &str,
146    checksums: Option<&Path>,
147    timeout_secs: u64,
148) -> anyhow::Result<()> {
149    ensure_file_exists_or_download_http_with_config(
150        filepath,
151        url,
152        checksums,
153        timeout_secs,
154        None,
155        None,
156    )
157}
158
159/// Ensures that a file exists at the specified path by downloading it if necessary,
160/// with custom timeout, retry config, and initial jitter delay.
161///
162/// # Parameters
163///
164/// - `filepath`: The path where the file should exist.
165/// - `url`: The URL to download from if the file doesn't exist.
166/// - `checksums`: Optional path to checksums file for verification.
167/// - `timeout_secs`: Timeout in seconds for HTTP requests.
168/// - `retry_config`: Optional custom retry configuration (uses sensible defaults if None).
169/// - `initial_jitter_ms`: Optional initial jitter delay in milliseconds before download (defaults to 100-600ms if None).
170///
171/// # Errors
172///
173/// Returns an error if:
174/// - The HTTP request cannot be sent or returns a non-success status code after retries.
175/// - Any I/O operation fails during file creation, reading, or writing.
176/// - Checksum verification or JSON parsing fails.
177pub fn ensure_file_exists_or_download_http_with_config(
178    filepath: &Path,
179    url: &str,
180    checksums: Option<&Path>,
181    timeout_secs: u64,
182    retry_config: Option<RetryConfig>,
183    initial_jitter_ms: Option<u64>,
184) -> anyhow::Result<()> {
185    // Local/cached file path: accept the file if it exists, updating the
186    // checksum record when it differs (e.g. after local regeneration).
187    // This is intentionally lenient — download verification below is strict.
188    if filepath.exists() {
189        println!("File already exists (local/cached): {}", filepath.display());
190
191        if let Some(checksums_file) = checksums {
192            let _guard = lock_large_checksums()?;
193
194            if verify_sha256_checksum(filepath, checksums_file)? {
195                println!("Checksum verified");
196                return Ok(());
197            } else {
198                let new_checksum = calculate_sha256(filepath)?;
199                println!("Updating checksum for local file: {new_checksum}");
200                update_sha256_checksums(filepath, checksums_file, &new_checksum)?;
201                return Ok(());
202            }
203        }
204        return Ok(());
205    }
206
207    // Add a small random delay to avoid bursting the remote server when
208    // many downloads start concurrently. Can be disabled by passing Some(0).
209    if let Some(jitter_ms) = initial_jitter_ms {
210        if jitter_ms > 0 {
211            sleep(Duration::from_millis(jitter_ms));
212        }
213    } else {
214        let jitter_delay = {
215            let mut r = rng();
216            Duration::from_millis(r.random_range(100..=600))
217        };
218        sleep(jitter_delay);
219    }
220
221    download_file(filepath, url, timeout_secs, retry_config.clone())?;
222
223    // Verify checksum after download, retry once on mismatch (corrupt download)
224    if let Some(checksums_file) = checksums {
225        let _guard = lock_large_checksums()?;
226
227        if !verify_sha256_checksum(filepath, checksums_file)? {
228            let actual = calculate_sha256(filepath)?;
229            println!("Checksum mismatch after download (calculated {actual}), retrying...");
230            remove_file(filepath)?;
231            drop(_guard);
232
233            download_file(filepath, url, timeout_secs, retry_config)?;
234
235            let _guard = lock_large_checksums()?;
236
237            if !verify_sha256_checksum(filepath, checksums_file)? {
238                let actual = calculate_sha256(filepath)?;
239                remove_file(filepath)?;
240                anyhow::bail!(
241                    "Checksum mismatch after retry for {} (calculated {actual})",
242                    filepath.file_name().unwrap_or_default().display(),
243                );
244            }
245        }
246    }
247
248    Ok(())
249}
250
251fn download_file(
252    filepath: &Path,
253    url: &str,
254    timeout_secs: u64,
255    retry_config: Option<RetryConfig>,
256) -> anyhow::Result<()> {
257    // Validate HTTPS for security in production builds,
258    // HTTP is intentionally allowed in test builds for local test servers (127.0.0.1),
259    // CodeQL flags this as "non-https-url" but it's a deliberate design choice for testkit.
260    #[cfg(not(test))]
261    if !url.starts_with("https://") {
262        anyhow::bail!("URL must use HTTPS protocol for security: {url}");
263    }
264
265    println!("Downloading file from {url} to {}", filepath.display());
266
267    if let Some(parent) = filepath.parent() {
268        std::fs::create_dir_all(parent)?;
269    }
270
271    let client = Client::builder()
272        .timeout(Duration::from_secs(timeout_secs))
273        .build()?;
274
275    let cfg = if let Some(config) = retry_config {
276        config
277    } else {
278        // Default production config
279        let max_retries = 5u32;
280        let op_timeout_ms = timeout_secs.saturating_mul(1000);
281        // Make the provided timeout a hard ceiling for total elapsed time.
282        // Split it across attempts (at least 1000 ms per attempt) and cap total at op_timeout_ms.
283        let per_attempt_ms = std::cmp::max(1000u64, op_timeout_ms / (max_retries as u64 + 1));
284        RetryConfig {
285            max_retries,
286            initial_delay_ms: 1_000,
287            max_delay_ms: 10_000,
288            backoff_factor: 2.0,
289            jitter_ms: 1_000,
290            operation_timeout_ms: Some(per_attempt_ms),
291            immediate_first: false,
292            max_elapsed_ms: Some(op_timeout_ms),
293        }
294    };
295
296    let op = || -> Result<(), DownloadError> {
297        match client.get(url).send() {
298            Ok(mut response) => {
299                let status = response.status();
300                if status.is_success() {
301                    let mut out = File::create(filepath)
302                        .map_err(|e| DownloadError::NonRetryable(e.to_string()))?;
303                    // Stream the response body directly to disk to avoid large allocations
304                    copy(&mut response, &mut out)
305                        .map_err(|e| DownloadError::NonRetryable(e.to_string()))?;
306                    println!("File downloaded to {}", filepath.display());
307                    Ok(())
308                } else if status.is_server_error()
309                    || status.as_u16() == 429
310                    || status.as_u16() == 408
311                {
312                    println!("HTTP error {status}, retrying...");
313                    Err(DownloadError::Retryable(format!("HTTP {status}")))
314                } else {
315                    // Preserve existing error text used by tests
316                    Err(DownloadError::NonRetryable(format!(
317                        "Client error: HTTP {status}"
318                    )))
319                }
320            }
321            Err(e) => {
322                println!("Request failed: {e}");
323                Err(DownloadError::Retryable(e.to_string()))
324            }
325        }
326    };
327
328    let should_retry = |e: &DownloadError| matches!(e, DownloadError::Retryable(_));
329
330    execute_with_retry_blocking(&cfg, op, should_retry).map_err(|e| anyhow::anyhow!(e.to_string()))
331}
332
333fn calculate_sha256(filepath: &Path) -> anyhow::Result<String> {
334    let mut file = File::open(filepath)?;
335    let mut ctx = Context::new(&digest::SHA256);
336    let mut buffer = [0u8; 4096];
337
338    loop {
339        let count = file.read(&mut buffer)?;
340        if count == 0 {
341            break;
342        }
343        ctx.update(&buffer[..count]);
344    }
345
346    let digest = ctx.finish();
347    Ok(hex::encode(digest.as_ref()))
348}
349
350fn verify_sha256_checksum(filepath: &Path, checksums: &Path) -> anyhow::Result<bool> {
351    let file = File::open(checksums)?;
352    let reader = BufReader::new(file);
353    let checksums: Value = serde_json::from_reader(reader)?;
354
355    let filename = filepath.file_name().unwrap().to_str().unwrap();
356    if let Some(expected_checksum) = checksums.get(filename) {
357        let expected_checksum_str = expected_checksum.as_str().unwrap();
358        let expected_hash = expected_checksum_str
359            .strip_prefix("sha256:")
360            .unwrap_or(expected_checksum_str);
361        let calculated_checksum = calculate_sha256(filepath)?;
362        if expected_hash == calculated_checksum {
363            return Ok(true);
364        }
365    }
366
367    Ok(false)
368}
369
370fn update_sha256_checksums(
371    filepath: &Path,
372    checksums_file: &Path,
373    new_checksum: &str,
374) -> anyhow::Result<()> {
375    let checksums: Value = if checksums_file.exists() {
376        let file = File::open(checksums_file)?;
377        let reader = BufReader::new(file);
378        serde_json::from_reader(reader)?
379    } else {
380        serde_json::json!({})
381    };
382
383    let mut checksums_map = checksums.as_object().unwrap().clone();
384
385    // Add or update the checksum
386    let filename = filepath.file_name().unwrap().to_str().unwrap().to_string();
387    let prefixed_checksum = format!("sha256:{new_checksum}");
388    checksums_map.insert(filename, Value::String(prefixed_checksum));
389
390    let file = OpenOptions::new()
391        .write(true)
392        .create(true)
393        .truncate(true)
394        .open(checksums_file)?;
395    let writer = BufWriter::new(file);
396    serde_json::to_writer_pretty(writer, &serde_json::Value::Object(checksums_map))?;
397
398    Ok(())
399}
400
401#[cfg(test)]
402mod tests {
403    use std::{
404        fs,
405        io::{BufWriter, Write},
406        net::SocketAddr,
407        sync::{
408            Arc,
409            atomic::{AtomicUsize, Ordering},
410        },
411    };
412
413    use axum::{Router, http::StatusCode, routing::get, serve};
414    use rstest::*;
415    use serde_json::{json, to_writer};
416    use tempfile::TempDir;
417    use tokio::{
418        net::TcpListener,
419        task,
420        time::{Duration, sleep},
421    };
422
423    use super::*;
424
425    /// Creates a fast, deterministic retry config for tests.
426    /// Uses very short delays to make tests run quickly without introducing flakiness.
427    fn test_retry_config() -> RetryConfig {
428        RetryConfig {
429            max_retries: 5,
430            initial_delay_ms: 10,
431            max_delay_ms: 50,
432            backoff_factor: 2.0,
433            jitter_ms: 5,
434            operation_timeout_ms: Some(500),
435            immediate_first: false,
436            max_elapsed_ms: Some(2000),
437        }
438    }
439
440    async fn setup_test_server(
441        server_content: Option<String>,
442        status_code: StatusCode,
443    ) -> SocketAddr {
444        let server_content = Arc::new(server_content);
445        let server_content_clone = server_content.clone();
446        let app = Router::new().route(
447            "/testfile.txt",
448            get(move || {
449                let server_content = server_content_clone.clone();
450                async move {
451                    let response_body = match &*server_content {
452                        Some(content) => content.clone(),
453                        None => "File not found".to_string(),
454                    };
455                    (status_code, response_body)
456                }
457            }),
458        );
459
460        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
461        let addr = listener.local_addr().unwrap();
462        let server = serve(listener, app);
463
464        task::spawn(async move {
465            if let Err(e) = server.await {
466                eprintln!("server error: {e}");
467            }
468        });
469
470        sleep(Duration::from_millis(100)).await;
471
472        addr
473    }
474
475    #[tokio::test]
476    async fn test_file_already_exists() {
477        let temp_dir = TempDir::new().unwrap();
478        let file_path = temp_dir.path().join("testfile.txt");
479        fs::write(&file_path, "Existing file content").unwrap();
480
481        let url = "http://example.com/testfile.txt".to_string();
482        let result = ensure_file_exists_or_download_http(&file_path, &url, None, Some(5));
483
484        assert!(result.is_ok());
485        let content = fs::read_to_string(&file_path).unwrap();
486        assert_eq!(content, "Existing file content");
487    }
488
489    #[tokio::test]
490    async fn test_download_file_success() {
491        let temp_dir = TempDir::new().unwrap();
492        let filepath = temp_dir.path().join("testfile.txt");
493        let filepath_clone = filepath.clone();
494
495        let server_content = "Server file content".to_string();
496        let status_code = StatusCode::OK;
497        let addr = setup_test_server(Some(server_content.clone()), status_code).await;
498        let url = format!("http://{addr}/testfile.txt");
499
500        let result = tokio::task::spawn_blocking(move || {
501            ensure_file_exists_or_download_http_with_config(
502                &filepath_clone,
503                &url,
504                None,
505                5,
506                Some(test_retry_config()),
507                Some(0),
508            )
509        })
510        .await
511        .unwrap();
512
513        assert!(result.is_ok());
514        let content = fs::read_to_string(&filepath).unwrap();
515        assert_eq!(content, server_content);
516    }
517
518    #[tokio::test]
519    async fn test_download_file_not_found() {
520        let temp_dir = TempDir::new().unwrap();
521        let file_path = temp_dir.path().join("testfile.txt");
522
523        let server_content = None;
524        let status_code = StatusCode::NOT_FOUND;
525        let addr = setup_test_server(server_content, status_code).await;
526        let url = format!("http://{addr}/testfile.txt");
527
528        let result = tokio::task::spawn_blocking(move || {
529            ensure_file_exists_or_download_http_with_config(
530                &file_path,
531                &url,
532                None,
533                1,
534                Some(test_retry_config()),
535                Some(0),
536            )
537        })
538        .await
539        .unwrap();
540
541        assert!(result.is_err());
542        let err_msg = format!("{}", result.unwrap_err());
543        assert!(
544            err_msg.contains("Client error: HTTP"),
545            "Unexpected error message: {err_msg}"
546        );
547    }
548
549    #[tokio::test]
550    async fn test_network_error() {
551        let temp_dir = TempDir::new().unwrap();
552        let file_path = temp_dir.path().join("testfile.txt");
553
554        // Use an unreachable address to simulate a network error
555        let url = "http://127.0.0.1:0/testfile.txt".to_string();
556
557        let result = tokio::task::spawn_blocking(move || {
558            ensure_file_exists_or_download_http_with_config(
559                &file_path,
560                &url,
561                None,
562                2,
563                Some(test_retry_config()),
564                Some(0),
565            )
566        })
567        .await
568        .unwrap();
569
570        assert!(result.is_err());
571        let err_msg = format!("{}", result.unwrap_err());
572        assert!(
573            err_msg.contains("error"),
574            "Unexpected error message: {err_msg}"
575        );
576    }
577
578    #[tokio::test]
579    async fn test_retry_then_success_on_500() {
580        let temp_dir = TempDir::new().unwrap();
581        let filepath = temp_dir.path().join("testfile.txt");
582        let filepath_clone = filepath.clone();
583
584        let counter = Arc::new(AtomicUsize::new(0));
585        let counter_clone = counter.clone();
586
587        let app = Router::new().route(
588            "/testfile.txt",
589            get(move || {
590                let c = counter_clone.clone();
591                async move {
592                    let n = c.fetch_add(1, Ordering::SeqCst);
593                    if n < 2 {
594                        (StatusCode::INTERNAL_SERVER_ERROR, "temporary error")
595                    } else {
596                        (StatusCode::OK, "eventual success")
597                    }
598                }
599            }),
600        );
601
602        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
603        let addr = listener.local_addr().unwrap();
604        let server = serve(listener, app);
605        task::spawn(async move {
606            let _ = server.await;
607        });
608        sleep(Duration::from_millis(100)).await;
609
610        let url = format!("http://{addr}/testfile.txt");
611
612        let result = tokio::task::spawn_blocking(move || {
613            ensure_file_exists_or_download_http_with_config(
614                &filepath_clone,
615                &url,
616                None,
617                5,
618                Some(test_retry_config()),
619                Some(0),
620            )
621        })
622        .await
623        .unwrap();
624
625        assert!(result.is_ok());
626        let content = std::fs::read_to_string(&filepath).unwrap();
627        assert_eq!(content, "eventual success");
628        assert!(counter.load(Ordering::SeqCst) >= 2);
629    }
630
631    #[tokio::test]
632    async fn test_retry_then_success_on_429() {
633        let temp_dir = TempDir::new().unwrap();
634        let filepath = temp_dir.path().join("testfile.txt");
635        let filepath_clone = filepath.clone();
636
637        let counter = Arc::new(AtomicUsize::new(0));
638        let counter_clone = counter.clone();
639
640        let app = Router::new().route(
641            "/testfile.txt",
642            get(move || {
643                let c = counter_clone.clone();
644                async move {
645                    let n = c.fetch_add(1, Ordering::SeqCst);
646                    if n < 1 {
647                        (StatusCode::TOO_MANY_REQUESTS, "rate limited")
648                    } else {
649                        (StatusCode::OK, "ok after retry")
650                    }
651                }
652            }),
653        );
654
655        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
656        let addr = listener.local_addr().unwrap();
657        let server = serve(listener, app);
658        task::spawn(async move {
659            let _ = server.await;
660        });
661        sleep(Duration::from_millis(100)).await;
662
663        let url = format!("http://{addr}/testfile.txt");
664
665        let result = tokio::task::spawn_blocking(move || {
666            ensure_file_exists_or_download_http_with_config(
667                &filepath_clone,
668                &url,
669                None,
670                5,
671                Some(test_retry_config()),
672                Some(0),
673            )
674        })
675        .await
676        .unwrap();
677
678        assert!(result.is_ok());
679        let content = std::fs::read_to_string(&filepath).unwrap();
680        assert_eq!(content, "ok after retry");
681        assert!(counter.load(Ordering::SeqCst) >= 2);
682    }
683
684    #[tokio::test]
685    async fn test_no_retry_on_404() {
686        let temp_dir = TempDir::new().unwrap();
687        let filepath = temp_dir.path().join("testfile.txt");
688        let filepath_clone = filepath.clone();
689
690        let counter = Arc::new(AtomicUsize::new(0));
691        let counter_clone = counter.clone();
692
693        let app = Router::new().route(
694            "/testfile.txt",
695            get(move || {
696                let c = counter_clone.clone();
697                async move {
698                    c.fetch_add(1, Ordering::SeqCst);
699                    (StatusCode::NOT_FOUND, "missing")
700                }
701            }),
702        );
703
704        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
705        let addr = listener.local_addr().unwrap();
706        let server = serve(listener, app);
707        task::spawn(async move {
708            let _ = server.await;
709        });
710        sleep(Duration::from_millis(100)).await;
711
712        let url = format!("http://{addr}/testfile.txt");
713
714        let result = tokio::task::spawn_blocking(move || {
715            ensure_file_exists_or_download_http_with_config(
716                &filepath_clone,
717                &url,
718                None,
719                5,
720                Some(test_retry_config()),
721                Some(0),
722            )
723        })
724        .await
725        .unwrap();
726
727        assert!(result.is_err());
728        assert_eq!(counter.load(Ordering::SeqCst), 1, "should not retry on 404");
729    }
730
731    #[tokio::test]
732    async fn test_checksum_mismatch_retry_then_success() {
733        let temp_dir = TempDir::new().unwrap();
734        let filepath = temp_dir.path().join("testfile.txt");
735        let filepath_clone = filepath.clone();
736
737        let good_content = "correct content";
738        let good_checksum = calculate_sha256_bytes(good_content.as_bytes());
739
740        let checksums_path = temp_dir.path().join("checksums.json");
741        let checksums_data = json!({
742            "testfile.txt": format!("sha256:{good_checksum}")
743        });
744        let checksums_file = File::create(&checksums_path).unwrap();
745        to_writer(BufWriter::new(checksums_file), &checksums_data).unwrap();
746        let checksums_clone = checksums_path.clone();
747
748        // First request returns corrupt data, second returns correct data
749        let counter = Arc::new(AtomicUsize::new(0));
750        let counter_clone = counter.clone();
751
752        let app = Router::new().route(
753            "/testfile.txt",
754            get(move || {
755                let c = counter_clone.clone();
756                async move {
757                    let n = c.fetch_add(1, Ordering::SeqCst);
758                    if n == 0 {
759                        (StatusCode::OK, "corrupt data")
760                    } else {
761                        (StatusCode::OK, "correct content")
762                    }
763                }
764            }),
765        );
766
767        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
768        let addr = listener.local_addr().unwrap();
769        let server = serve(listener, app);
770        task::spawn(async move {
771            let _ = server.await;
772        });
773        sleep(Duration::from_millis(100)).await;
774
775        let url = format!("http://{addr}/testfile.txt");
776
777        let result = tokio::task::spawn_blocking(move || {
778            ensure_file_exists_or_download_http_with_config(
779                &filepath_clone,
780                &url,
781                Some(&checksums_clone),
782                5,
783                Some(test_retry_config()),
784                Some(0),
785            )
786        })
787        .await
788        .unwrap();
789
790        assert!(result.is_ok());
791        let content = fs::read_to_string(&filepath).unwrap();
792        assert_eq!(content, good_content);
793        assert_eq!(counter.load(Ordering::SeqCst), 2);
794    }
795
796    #[tokio::test]
797    async fn test_checksum_mismatch_retry_then_fail() {
798        let temp_dir = TempDir::new().unwrap();
799        let filepath = temp_dir.path().join("testfile.txt");
800        let filepath_clone = filepath.clone();
801
802        // Checksum for content that the server will never return
803        let checksums_path = temp_dir.path().join("checksums.json");
804        let checksums_data = json!({
805            "testfile.txt": "sha256:0000000000000000000000000000000000000000000000000000000000000000"
806        });
807        let checksums_file = File::create(&checksums_path).unwrap();
808        to_writer(BufWriter::new(checksums_file), &checksums_data).unwrap();
809        let checksums_clone = checksums_path.clone();
810
811        let counter = Arc::new(AtomicUsize::new(0));
812        let counter_clone = counter.clone();
813
814        let app = Router::new().route(
815            "/testfile.txt",
816            get(move || {
817                let c = counter_clone.clone();
818                async move {
819                    c.fetch_add(1, Ordering::SeqCst);
820                    (StatusCode::OK, "always wrong content")
821                }
822            }),
823        );
824
825        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
826        let addr = listener.local_addr().unwrap();
827        let server = serve(listener, app);
828        task::spawn(async move {
829            let _ = server.await;
830        });
831        sleep(Duration::from_millis(100)).await;
832
833        let url = format!("http://{addr}/testfile.txt");
834
835        let result = tokio::task::spawn_blocking(move || {
836            ensure_file_exists_or_download_http_with_config(
837                &filepath_clone,
838                &url,
839                Some(&checksums_clone),
840                5,
841                Some(test_retry_config()),
842                Some(0),
843            )
844        })
845        .await
846        .unwrap();
847
848        assert!(result.is_err());
849        let err_msg = format!("{}", result.unwrap_err());
850        assert!(err_msg.contains("Checksum mismatch after retry"));
851        assert_eq!(
852            counter.load(Ordering::SeqCst),
853            2,
854            "should download exactly twice"
855        );
856        assert!(!filepath.exists(), "corrupt file should be cleaned up");
857    }
858
859    fn calculate_sha256_bytes(data: &[u8]) -> String {
860        let mut ctx = digest::Context::new(&digest::SHA256);
861        ctx.update(data);
862        hex::encode(ctx.finish().as_ref())
863    }
864
865    #[rstest]
866    #[expect(clippy::panic_in_result_fn)]
867    fn test_calculate_sha256() -> anyhow::Result<()> {
868        let temp_dir = TempDir::new()?;
869        let test_file_path = temp_dir.path().join("test_file.txt");
870        let mut test_file = File::create(&test_file_path)?;
871        let content = b"Hello, world!";
872        test_file.write_all(content)?;
873
874        let expected_hash = "315f5bdb76d078c43b8ac0064e4a0164612b1fce77c869345bfc94c75894edd3";
875        let calculated_hash = calculate_sha256(&test_file_path)?;
876
877        assert_eq!(calculated_hash, expected_hash);
878        Ok(())
879    }
880
881    #[rstest]
882    #[expect(clippy::panic_in_result_fn)]
883    fn test_verify_sha256_checksum() -> anyhow::Result<()> {
884        let temp_dir = TempDir::new()?;
885        let test_file_path = temp_dir.path().join("test_file.txt");
886        let mut test_file = File::create(&test_file_path)?;
887        let content = b"Hello, world!";
888        test_file.write_all(content)?;
889
890        let calculated_checksum = calculate_sha256(&test_file_path)?;
891
892        // Create checksums.json containing the checksum
893        let checksums_path = temp_dir.path().join("checksums.json");
894        let checksums_data = json!({
895            "test_file.txt": format!("sha256:{}", calculated_checksum)
896        });
897        let checksums_file = File::create(&checksums_path)?;
898        let writer = BufWriter::new(checksums_file);
899        to_writer(writer, &checksums_data)?;
900
901        let is_valid = verify_sha256_checksum(&test_file_path, &checksums_path)?;
902        assert!(is_valid, "The checksum should be valid");
903        Ok(())
904    }
905}
nautilus_testkit/files.rs

nautilus_testkit/
files.rs
