Verify MNIST download integrity before using cached file

src/data/mnist.rs

@@ -4,6 +4,7 @@
 use std::fs;
 use std::io::{Read, Write};
 use std::path::{Path, PathBuf};
+use std::sync::atomic::{AtomicU32, Ordering};
 
 use crate::error::{Error, Result};
 use crate::tensor::Tensor;
@@ -59,10 +60,8 @@ impl Mnist {
     /// Load a split from an explicit directory.
     pub fn load_from(cache: &Path, split: Split) -> Result<Mnist> {
         let (image_file, label_file) = split.files();
-        download_missing(cache, &[image_file, label_file])?;
-
-        let (images, image_count) = parse_images(&fs::read(cache.join(image_file))?)?;
-        let (labels, label_count) = parse_labels(&fs::read(cache.join(label_file))?)?;
+        let (images, image_count) = load_cached(cache, image_file, parse_images, fetch)?;
+        let (labels, label_count) = load_cached(cache, label_file, parse_labels, fetch)?;
         if image_count != label_count {
             return Err(Error::Dataset(format!(
                 "mnist: {image_count} images but {label_count} labels"
@@ -119,32 +118,78 @@ pub fn default_cache_dir() -> PathBuf {
         .join("mnist")
 }
 
-/// Download and decompress any of `files` not already cached.
-fn download_missing(cache: &Path, files: &[&str]) -> Result<()> {
-    fs::create_dir_all(cache)?;
-    for name in files {
-        let target = cache.join(name);
-        if target.exists() {
-            continue;
+/// Parse a cached file, downloading it first when it is missing.
+///
+/// A cached file that does not parse is deleted and fetched once more. Without
+/// that, a file truncated by an interrupted write would be trusted forever,
+/// since the cache check only asks whether the file exists.
+fn load_cached<T>(
+    cache: &Path,
+    name: &str,
+    parse: fn(&[u8]) -> Result<T>,
+    fetch: impl Fn(&Path, &str) -> Result<Vec<u8>>,
+) -> Result<T> {
+    let target = cache.join(name);
+    if target.exists() {
+        match parse(&fs::read(&target)?) {
+            Ok(parsed) => return Ok(parsed),
+            Err(e) => {
+                eprintln!(
+                    "mnist: {} is not usable ({e}), downloading it again",
+                    target.display()
+                );
+                let _ = fs::remove_file(&target);
+            }
         }
+    }
+    fs::create_dir_all(cache)?;
+    let bytes = fetch(cache, name)?;
+    parse(&bytes).map_err(|e| Error::Dataset(format!("mnist: {}: {e}", target.display())))
+}
 
-        let url = format!("{MIRROR}/{name}.gz");
-        eprintln!("mnist: downloading {url}");
-        let response = ureq::get(&url)
-            .call()
-            .map_err(|e| Error::Dataset(format!("could not fetch {url}: {e}")))?;
+/// Download `name` from the mirror, decompress it, and write it into the cache.
+fn fetch(cache: &Path, name: &str) -> Result<Vec<u8>> {
+    let url = format!("{MIRROR}/{name}.gz");
+    eprintln!("mnist: downloading {url}");
+    let response = ureq::get(&url)
+        .call()
+        .map_err(|e| Error::Dataset(format!("could not fetch {url}: {e}")))?;
 
-        let mut compressed = Vec::new();
-        response
-            .into_body()
-            .into_reader()
-            .read_to_end(&mut compressed)?;
-        let mut raw = Vec::new();
-        flate2::read::GzDecoder::new(&compressed[..]).read_to_end(&mut raw)?;
+    let mut compressed = Vec::new();
+    response
+        .into_body()
+        .into_reader()
+        .read_to_end(&mut compressed)?;
+    let mut raw = Vec::new();
+    flate2::read::GzDecoder::new(&compressed[..]).read_to_end(&mut raw)?;
 
-        fs::File::create(&target)?.write_all(&raw)?;
+    write_atomic(&cache.join(name), &raw)?;
+    Ok(raw)
+}
+
+/// Write `bytes` to `target` so a crash leaves it either absent or complete:
+/// a temporary file next to it, `sync_all`, then a rename into place.
+fn write_atomic(target: &Path, bytes: &[u8]) -> Result<()> {
+    // The pid and a counter keep concurrent writers, in and across processes,
+    // off each other's temporary file.
+    static SERIAL: AtomicU32 = AtomicU32::new(0);
+    let name = target.file_name().unwrap_or_default().to_string_lossy();
+    let partial = target.with_file_name(format!(
+        "{name}.{}.{}.partial",
+        std::process::id(),
+        SERIAL.fetch_add(1, Ordering::Relaxed)
+    ));
+
+    let written = (|| {
+        let mut file = fs::File::create(&partial)?;
+        file.write_all(bytes)?;
+        file.sync_all()?;
+        fs::rename(&partial, target)
+    })();
+    if written.is_err() {
+        let _ = fs::remove_file(&partial);
     }
-    Ok(())
+    written.map_err(Error::from)
 }
 
 /// IDX images: magic, count, rows, cols, then one byte per pixel.
@@ -204,3 +249,134 @@ fn read_header(bytes: &[u8], words: usize, magic: u32, what: &str) -> Result<Vec
     }
     Ok(values)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// A fresh directory under the system temp dir, removed on drop.
+    struct TempDir(PathBuf);
+
+    impl TempDir {
+        fn new(name: &str) -> TempDir {
+            let dir =
+                std::env::temp_dir().join(format!("fastnn-mnist-{}-{name}", std::process::id()));
+            let _ = fs::remove_dir_all(&dir);
+            fs::create_dir_all(&dir).unwrap();
+            TempDir(dir)
+        }
+    }
+
+    impl Drop for TempDir {
+        fn drop(&mut self) {
+            let _ = fs::remove_dir_all(&self.0);
+        }
+    }
+
+    /// A valid IDX image file holding `count` one-pixel images.
+    fn image_file(count: u32) -> Vec<u8> {
+        let mut bytes = Vec::new();
+        for word in [IMAGE_MAGIC, count, 1, 1] {
+            bytes.extend_from_slice(&word.to_be_bytes());
+        }
+        bytes.resize(IMAGE_HEADER + count as usize, 7);
+        bytes
+    }
+
+    fn partials_in(dir: &Path) -> usize {
+        fs::read_dir(dir)
+            .unwrap()
+            .filter(|entry| {
+                entry
+                    .as_ref()
+                    .unwrap()
+                    .file_name()
+                    .to_string_lossy()
+                    .ends_with(".partial")
+            })
+            .count()
+    }
+
+    #[test]
+    fn write_atomic_writes_the_bytes_and_leaves_no_partial() {
+        let dir = TempDir::new("write");
+        let target = dir.0.join("data.bin");
+        write_atomic(&target, b"hello").unwrap();
+        assert_eq!(fs::read(&target).unwrap(), b"hello");
+        assert_eq!(partials_in(&dir.0), 0);
+    }
+
+    #[test]
+    fn write_atomic_replaces_an_existing_file() {
+        let dir = TempDir::new("replace");
+        let target = dir.0.join("data.bin");
+        fs::write(&target, b"old").unwrap();
+        write_atomic(&target, b"new").unwrap();
+        assert_eq!(fs::read(&target).unwrap(), b"new");
+    }
+
+    #[test]
+    fn write_atomic_cleans_up_when_the_rename_fails() {
+        let dir = TempDir::new("rename-fails");
+        // A file cannot replace a directory, so the rename fails after the write.
+        let target = dir.0.join("taken");
+        fs::create_dir(&target).unwrap();
+        assert!(write_atomic(&target, b"x").is_err());
+        assert_eq!(partials_in(&dir.0), 0);
+    }
+
+    #[test]
+    fn a_valid_cached_file_is_used_without_fetching() {
+        let dir = TempDir::new("valid");
+        fs::write(dir.0.join("images"), image_file(3)).unwrap();
+        let (_, count) = load_cached(&dir.0, "images", parse_images, |_, _| {
+            panic!("a valid cached file must not be fetched")
+        })
+        .unwrap();
+        assert_eq!(count, 3);
+    }
+
+    #[test]
+    fn a_missing_file_is_fetched() {
+        let dir = TempDir::new("missing");
+        let (_, count) = load_cached(&dir.0, "images", parse_images, |cache, name| {
+            let bytes = image_file(2);
+            write_atomic(&cache.join(name), &bytes)?;
+            Ok(bytes)
+        })
+        .unwrap();
+        assert_eq!(count, 2);
+        assert!(parse_images(&fs::read(dir.0.join("images")).unwrap()).is_ok());
+    }
+
+    #[test]
+    fn a_truncated_cached_file_is_fetched_again() {
+        let dir = TempDir::new("truncated");
+        // Valid magic and header, body cut short: what an interrupted write leaves.
+        let mut truncated = image_file(5);
+        truncated.truncate(IMAGE_HEADER + 2);
+        fs::write(dir.0.join("images"), &truncated).unwrap();
+
+        let fetches = std::cell::Cell::new(0);
+        let (_, count) = load_cached(&dir.0, "images", parse_images, |cache, name| {
+            fetches.set(fetches.get() + 1);
+            let bytes = image_file(5);
+            write_atomic(&cache.join(name), &bytes)?;
+            Ok(bytes)
+        })
+        .unwrap();
+        assert_eq!(count, 5);
+        assert_eq!(fetches.get(), 1);
+    }
+
+    #[test]
+    fn a_bad_download_names_the_file() {
+        let dir = TempDir::new("bad-download");
+        fs::write(dir.0.join("images"), b"junk").unwrap();
+        let err = load_cached(&dir.0, "images", parse_images, |_, _| {
+            Ok(b"still junk".to_vec())
+        })
+        .unwrap_err();
+        assert!(err.to_string().contains("images"), "{err}");
+    }
+}