codebook_downloader/

lib.rs

use anyhow::Result;
use chrono::{DateTime, Utc};
use log::info;
use reqwest::blocking::{Client, Response};
use reqwest::header::{IF_MODIFIED_SINCE, LAST_MODIFIED};
use rustls::ClientConfig;
use rustls_platform_verifier::BuilderVerifierExt;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use std::fs::{self, File};
use std::io::{BufReader, Read};
use std::path::{Path, PathBuf};
use std::sync::{OnceLock, RwLock};
use tempfile::NamedTempFile;

const METADATA_FILE: &str = "_metadata.json";
const TWO_WEEKS: u64 = 14 * 24 * 3600;

#[derive(Debug, Serialize, Deserialize, Default)]
struct Metadata {
    files: HashMap<String, FileEntry>,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
struct FileEntry {
    path: PathBuf,
    last_checked: DateTime<Utc>,
    last_modified: Option<DateTime<Utc>>,
    content_hash: String,
}

pub struct Downloader {
    cache_dir: PathBuf,
    metadata_path: PathBuf,
    metadata: OnceLock<RwLock<Metadata>>,
    _client: OnceLock<Client>,
}

impl Downloader {
    pub fn new(cache_dir: impl AsRef<Path>) -> Self {
        let cache_dir = cache_dir.as_ref().to_path_buf();
        info!("Cache folder at: {cache_dir:?}");

        let metadata_path = cache_dir.join(METADATA_FILE);

        Self {
            cache_dir,
            metadata_path,
            metadata: OnceLock::new(),
            _client: OnceLock::new(),
        }
    }

    fn client(&self) -> &Client {
        self._client.get_or_init(|| {
            // Set up rustls_platform_verifier to use OS cert chains (proxy support)
            let arc_crypto_provider =
                std::sync::Arc::new(rustls::crypto::aws_lc_rs::default_provider());
            let config = ClientConfig::builder_with_provider(arc_crypto_provider)
                .with_safe_default_protocol_versions()
                .unwrap()
                .with_platform_verifier()
                .unwrap()
                .with_no_client_auth();
            reqwest::blocking::Client::builder()
                .use_preconfigured_tls(config)
                .build()
                .expect("Failed to build http client")
        })
    }

    fn metadata(&self) -> &RwLock<Metadata> {
        let metadata_path = self.metadata_path.clone();
        let cache_dir = self.cache_dir.clone();
        self.metadata.get_or_init(move || {
            fs::create_dir_all(&cache_dir)
                .expect("Failed to create cache directory: {cache_dir:?}");
            let metadata = Self::load_metadata(&metadata_path);
            RwLock::new(metadata)
        })
    }

    fn load_metadata(metadata_path: &Path) -> Metadata {
        match File::open(metadata_path) {
            Ok(file) => match serde_json::from_reader(file) {
                Ok(metadata) => metadata,
                Err(err) => {
                    log::warn!("Failed to parse metadata file {metadata_path:?}: {err}");
                    Metadata::default()
                }
            },
            Err(err) => {
                if err.kind() != std::io::ErrorKind::NotFound {
                    log::warn!("Failed to open metadata file {metadata_path:?}: {err}");
                }
                Metadata::default()
            }
        }
    }

    fn persist_metadata(&self, metadata: &Metadata) -> Result<()> {
        let file = File::create(&self.metadata_path)?;
        serde_json::to_writer_pretty(file, metadata)?;
        Ok(())
    }

    fn purge_stale_entry(&self, url: &str, stale_path: &Path) {
        let mut metadata = self.metadata().write().unwrap();
        if metadata
            .files
            .get(url)
            .map(|entry| entry.path == stale_path)
            .unwrap_or(false)
        {
            metadata.files.remove(url);
            if let Err(err) = self.persist_metadata(&metadata) {
                log::error!(
                    "Failed to persist metadata after removing stale entry for {url}: {err}"
                );
            }
        }
    }

    pub fn get(&self, url: &str) -> Result<PathBuf> {
        let entry = {
            let metadata = self.metadata().read().unwrap();
            metadata.files.get(url).cloned()
        };

        let result = match entry {
            Some(entry) => {
                if !entry.path.exists() {
                    self.purge_stale_entry(url, &entry.path);
                    self.download_new(url)
                } else {
                    let needs_update =
                        entry.last_checked.timestamp() + TWO_WEEKS as i64 <= Utc::now().timestamp();
                    if needs_update {
                        self.try_update(url)
                    } else {
                        Ok(entry.path)
                    }
                }
            }
            None => self.download_new(url),
        };

        result.or_else(|e| {
            log::error!("Failed to update, using cached version: {e}");
            let entry = {
                let metadata = self.metadata().read().unwrap();
                metadata
                    .files
                    .get(url)
                    .map(|file_info| file_info.path.clone())
            };
            match entry {
                Some(path) => {
                    if path.exists() {
                        Ok(path)
                    } else {
                        self.purge_stale_entry(url, &path);
                        // If fallback path doesn't exist, try to download anyway
                        self.download_new(url)
                    }
                }
                None => {
                    // URL not found in the files hashmap, try to download it
                    log::warn!("URL not found in cache, attempting fresh download: {url}");
                    self.download_new(url)
                }
            }
        })
    }

    fn try_update(&self, url: &str) -> Result<PathBuf> {
        // Get last modified time with read lock
        let last_modified = {
            self.metadata()
                .read()
                .unwrap()
                .files
                .get(url)
                .and_then(|e| e.last_modified)
        };
        // log::info!("{:?}", last_modified);
        // log::info!("URL {:?}", url);

        let mut request = self.client().get(url);
        if let Some(lm) = last_modified {
            request = request.header(IF_MODIFIED_SINCE, lm.with_timezone(&Utc).to_rfc2822());
        }
        // log::info!("{:?}", request);

        let response = request.send()?;
        // log::info!("RESPONSE {:?}", response);

        match response.status().as_u16() {
            304 => self.update_check_time(url),
            200 => self.handle_updated_response(url, response),
            status => {
                let _ = self.update_check_time(url);
                Err(anyhow::anyhow!("Unexpected status code: {}", status))
            }
        }
    }

    fn handle_updated_response(&self, url: &str, response: Response) -> Result<PathBuf> {
        let last_modified = parse_last_modified(&response);
        let temp_file = self.download_to_temp(response)?;
        let new_hash = compute_file_hash(temp_file.path())?;
        let old_hash = {
            let metadata = self
                .metadata()
                .read()
                .map_err(|e| anyhow::anyhow!("Lock error: {}", e))?;
            metadata.files.get(url).unwrap().content_hash.clone()
        };
        if new_hash == old_hash {
            self.update_check_time(url)
        } else {
            self.replace_file(url, temp_file, last_modified, new_hash)
        }
    }

    fn download_new(&self, url: &str) -> Result<PathBuf> {
        let response = self.client().get(url).send()?;
        let last_modified = parse_last_modified(&response);
        let temp_file = self.download_to_temp(response)?;
        let new_hash = compute_file_hash(temp_file.path())?;
        self.store_new_file(url, temp_file, last_modified, new_hash)
    }

    fn download_to_temp(&self, mut response: Response) -> Result<NamedTempFile> {
        let mut temp_file = NamedTempFile::new_in(&self.cache_dir)?;
        std::io::copy(&mut response, &mut temp_file)?;
        Ok(temp_file)
    }

    fn store_new_file(
        &self,
        url: &str,
        temp_file: NamedTempFile,
        last_modified: Option<DateTime<Utc>>,
        content_hash: String,
    ) -> Result<PathBuf> {
        let filename = hash_url(url);
        let path = self.cache_dir.join(filename);
        temp_file.persist(&path)?;

        let entry = FileEntry {
            path: path.clone(),
            last_checked: Utc::now(),
            last_modified,
            content_hash,
        };
        {
            let mut metadata = self.metadata().write().unwrap();
            metadata.files.insert(url.to_string(), entry);
            self.persist_metadata(&metadata)?;
        }
        Ok(path)
    }

    fn replace_file(
        &self,
        url: &str,
        temp_file: NamedTempFile,
        last_modified: Option<DateTime<Utc>>,
        content_hash: String,
    ) -> Result<PathBuf> {
        let new_path: PathBuf;
        {
            let mut metadata = self.metadata().write().unwrap();
            let entry = metadata.files.get_mut(url).unwrap();
            let old_path = entry.path.clone();

            new_path = self.cache_dir.join(hash_url(url));
            temp_file.persist(&new_path)?;

            // Remove old file if it's different
            if old_path != new_path && old_path.exists() {
                fs::remove_file(old_path)?;
            }

            entry.path = new_path.clone();
            entry.last_checked = Utc::now();
            entry.last_modified = last_modified;
            entry.content_hash = content_hash;
            self.persist_metadata(&metadata)?;
        }

        Ok(new_path)
    }

    fn update_check_time(&self, url: &str) -> Result<PathBuf> {
        let path: PathBuf;
        {
            let mut metadata = self.metadata().write().unwrap();
            let entry = metadata.files.get_mut(url).unwrap();
            entry.last_checked = Utc::now();
            path = entry.path.clone();
            self.persist_metadata(&metadata)?;
        }
        Ok(path)
    }
}

fn hash_url(url: &str) -> String {
    let mut hasher = Sha256::new();
    hasher.update(url.as_bytes());
    format!("{:x}", hasher.finalize())
}

fn compute_file_hash(path: &Path) -> Result<String> {
    let file = File::open(path)?;
    let mut reader = BufReader::new(file);
    let mut hasher = Sha256::new();
    let mut buffer = [0; 1024];

    loop {
        let count = reader.read(&mut buffer)?;
        if count == 0 {
            break;
        }
        hasher.update(&buffer[..count]);
    }

    Ok(format!("{:x}", hasher.finalize()))
}

fn parse_last_modified(response: &Response) -> Option<DateTime<Utc>> {
    response
        .headers()
        .get(LAST_MODIFIED)
        .and_then(|hv| hv.to_str().ok())
        .and_then(|s| DateTime::parse_from_rfc2822(s).ok())
        .map(|dt| dt.with_timezone(&Utc))
}

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::Duration;
    use httpmock::MockServer;
    use tempfile::tempdir;

    #[test]
    fn test_download_new_file() {
        let server = MockServer::start();
        let mock = server.mock(|when, then| {
            when.method("GET").path("/test.txt");
            then.status(200)
                .body("test content")
                .header("Last-Modified", "Wed, 21 Oct 2023 07:28:00 GMT");
        });

        let temp_dir = tempdir().unwrap();
        let downloader = Downloader::new(temp_dir.path());
        let path = downloader.get(&server.url("/test.txt")).unwrap();

        mock.assert();
        assert!(path.exists());
        assert_eq!(std::fs::read_to_string(path).unwrap(), "test content");
        let metadata = downloader.metadata().read().unwrap();
        let entry = metadata.files.get(&server.url("/test.txt")).unwrap();
        assert_eq!(entry.content_hash, compute_file_hash(&entry.path).unwrap());
    }

    #[test]
    fn test_returns_cached_file_when_offline() {
        let server = MockServer::start();
        let temp_dir = tempdir().unwrap();

        // First download to cache
        let mock = server.mock(|when, then| {
            when.method("GET").path("/test.txt");
            then.status(200).body("cached content");
        });

        let downloader = Downloader::new(temp_dir.path());
        let path = downloader.get(&server.url("/test.txt")).unwrap();
        mock.assert();

        // Now simulate offline
        let downloader = Downloader::new(temp_dir.path());
        // server.stop(); // Make sure server isn't running
        let cached_path = downloader.get(&server.url("/test.txt")).unwrap();

        assert_eq!(path, cached_path);
        assert_eq!(
            std::fs::read_to_string(cached_path).unwrap(),
            "cached content"
        );
    }

    #[test]
    fn test_updates_file_when_modified() {
        let server = MockServer::start();
        let temp_dir = tempdir().unwrap();
        let test_path = server.url("/test.txt");

        // Initial download
        let initial_last_modified = "Wed, 21 Oct 2020 07:28:00 GMT";
        let mut mock1 = server.mock(|when, then| {
            when.method("GET").path("/test.txt");
            then.status(200)
                .body("v1")
                .header("Last-Modified", initial_last_modified);
        });

        let downloader = Downloader::new(temp_dir.path());
        let path_v1 = downloader.get(&test_path).unwrap();
        mock1.assert();
        mock1.delete();

        // Get stored metadata
        let stored_last_modified = {
            let metadata = downloader.metadata().read().unwrap();
            metadata.files[&test_path].last_modified
        };

        // Set last checked to 3 weeks ago
        {
            let mut metadata = downloader.metadata().write().unwrap();
            let entry = metadata.files.get_mut(&test_path).unwrap();
            entry.last_checked = stored_last_modified.unwrap() - Duration::weeks(3);
        }

        // Update mock with new content
        let mock2 = server.mock(|when, then| {
            when.method("GET").path("/test.txt").header(
                IF_MODIFIED_SINCE.as_str(),
                stored_last_modified.unwrap().to_rfc2822(),
            );
            then.status(200)
                .body("v2")
                .header("Last-Modified", "Fri, 23 Oct 2020 07:28:00 GMT");
        });

        let path_v2 = downloader.get(&test_path).unwrap();
        println!("path: {path_v2:?}");
        mock2.assert();

        assert_eq!(path_v1, path_v2);
        assert_eq!(std::fs::read_to_string(path_v2).unwrap(), "v2");
    }
    #[test]
    fn test_uses_stale_file_when_update_fails() {
        let server = MockServer::start();
        let temp_dir = tempdir().unwrap();

        // Initial download
        let mock1 = server.mock(|when, then| {
            when.method("GET").path("/test.txt");
            then.status(200).body("original");
        });

        let downloader = Downloader::new(temp_dir.path());
        let original_path = downloader.get(&server.url("/test.txt")).unwrap();
        mock1.assert();

        // Force update check time and break the server
        {
            let mut metadata = downloader.metadata().try_write().unwrap();
            if let Some(entry) = metadata.files.get_mut(&server.url("/test.txt")) {
                entry.last_checked = Utc::now() - Duration::seconds(TWO_WEEKS as i64 * 2);
            }
        }
        // server.stop();

        let cached_path = downloader.get(&server.url("/test.txt")).unwrap();
        assert_eq!(original_path, cached_path);
    }

    #[test]
    fn test_doesnt_check_within_two_weeks() {
        let server = MockServer::start();
        let temp_dir = tempdir().unwrap();

        // Initial download
        let mock = server.mock(|when, then| {
            when.method("GET").path("/test.txt");
            then.status(200).body("content");
        });

        let downloader = Downloader::new(temp_dir.path());
        downloader.get(&server.url("/test.txt")).unwrap();
        mock.assert_calls(1);

        // Subsequent call within two weeks
        let mock2 = server.mock(|when, then| {
            when.method("GET").path("/test.txt");
            then.status(200);
        });
        downloader.get(&server.url("/test.txt")).unwrap();
        mock2.assert_calls(0); // Should not make any new requests
    }

    #[test]
    fn test_handles_304_not_modified() {
        let server = MockServer::start();
        let temp_dir = tempdir().unwrap();

        // Initial download
        let mut mock1 = server.mock(|when, then| {
            when.method("GET").path("/test.txt");
            then.status(200)
                .body("content")
                .header("Last-Modified", "Wed, 21 Oct 2020 07:28:00 GMT");
        });

        let downloader = Downloader::new(temp_dir.path());
        let original_path = downloader.get(&server.url("/test.txt")).unwrap();
        mock1.assert();
        mock1.delete();

        // Force check time
        {
            let mut metadata = downloader.metadata().write().unwrap();
            if let Some(entry) = metadata.files.get_mut(&server.url("/test.txt")) {
                entry.last_checked = DateTime::parse_from_rfc2822("Wed, 21 Oct 2020 07:28:00 GMT")
                    .unwrap()
                    .with_timezone(&Utc);
            }
        }

        // 304 response
        let mock2 = server.mock(|when, then| {
            when.method("GET")
                .path("/test.txt")
                .header("If-Modified-Since", "Wed, 21 Oct 2020 07:28:00 +0000");
            then.status(304);
        });

        let cached_path = downloader.get(&server.url("/test.txt")).unwrap();
        mock2.assert();
        assert_eq!(original_path, cached_path);
        let metadata = downloader.metadata().read().unwrap();
        let entry = metadata.files.get(&server.url("/test.txt")).unwrap();
        assert!(entry.last_checked > Utc::now() - Duration::seconds(1));
    }

    #[test]
    fn test_file_hashing() {
        let url1 = "https://example.com/file1";
        let url2 = "https://example.com/file2";

        assert_ne!(hash_url(url1), hash_url(url2));

        let same_url = "https://example.com/same";
        assert_eq!(hash_url(same_url), hash_url(same_url));
    }

    #[test]
    fn test_redownloads_when_file_missing() {
        let server = MockServer::start();
        let temp_dir = tempdir().unwrap();

        // First download
        let mut mock1 = server.mock(|when, then| {
            when.method("GET").path("/test.txt");
            then.status(200).body("content");
        });

        let downloader = Downloader::new(temp_dir.path());
        let path = downloader.get(&server.url("/test.txt")).unwrap();
        mock1.assert();
        mock1.delete();

        // Simulate file deletion but keep metadata
        std::fs::remove_file(&path).unwrap();
        assert!(!path.exists());

        // Second request should redownload
        let mock2 = server.mock(|when, then| {
            when.method("GET").path("/test.txt");
            then.status(200).body("redownloaded content");
        });

        let new_path = downloader.get(&server.url("/test.txt")).unwrap();
        mock2.assert();

        assert!(new_path.exists());
        assert_eq!(
            std::fs::read_to_string(&new_path).unwrap(),
            "redownloaded content"
        );
    }
}