src/tools/build-manifest/src/checksum.rs

   1 use crate::manifest::{FileHash, Manifest};
   2 use rayon::prelude::*;
   3 use sha2::{Digest, Sha256};
   4 use std::collections::{HashMap, HashSet};
   5 use std::error::Error;
   6 use std::fs::File;
   7 use std::io::BufReader;
   8 use std::path::{Path, PathBuf};
   9 use std::sync::Mutex;
  10 use std::time::Instant;
  11
  12 pub(crate) struct Checksums {
  13     cache_path: Option<PathBuf>,
  14     collected: Mutex<HashMap<PathBuf, String>>,
  15 }
  16
  17 impl Checksums {
  18     pub(crate) fn new() -> Result<Self, Box<dyn Error>> {
  19         let cache_path = std::env::var_os("BUILD_MANIFEST_CHECKSUM_CACHE").map(PathBuf::from);
  20
  21         let mut collected = HashMap::new();
  22         if let Some(path) = &cache_path {
  23             if path.is_file() {
  24                 collected = serde_json::from_slice(&std::fs::read(path)?)?;
  25             }
  26         }
  27
  28         Ok(Checksums { cache_path, collected: Mutex::new(collected) })
  29     }
  30
  31     pub(crate) fn store_cache(&self) -> Result<(), Box<dyn Error>> {
  32         if let Some(path) = &self.cache_path {
  33             std::fs::write(path, &serde_json::to_vec(&self.collected)?)?;
  34         }
  35         Ok(())
  36     }
  37
  38     pub(crate) fn fill_missing_checksums(&mut self, manifest: &mut Manifest) {
  39         let need_checksums = self.find_missing_checksums(manifest);
  40         if !need_checksums.is_empty() {
  41             self.collect_checksums(&need_checksums);
  42         }
  43         self.replace_checksums(manifest);
  44     }
  45
  46     fn find_missing_checksums(&mut self, manifest: &mut Manifest) -> HashSet<PathBuf> {
  47         let collected = self.collected.lock().unwrap();
  48         let mut need_checksums = HashSet::new();
  49         crate::manifest::visit_file_hashes(manifest, |file_hash| {
  50             if let FileHash::Missing(path) = file_hash {
  51                 let path = std::fs::canonicalize(path).unwrap();
  52                 if !collected.contains_key(&path) {
  53                     need_checksums.insert(path);
  54                 }
  55             }
  56         });
  57         need_checksums
  58     }
  59
  60     fn replace_checksums(&mut self, manifest: &mut Manifest) {
  61         let collected = self.collected.lock().unwrap();
  62         crate::manifest::visit_file_hashes(manifest, |file_hash| {
  63             if let FileHash::Missing(path) = file_hash {
  64                 let path = std::fs::canonicalize(path).unwrap();
  65                 match collected.get(&path) {
  66                     Some(hash) => *file_hash = FileHash::Present(hash.clone()),
  67                     None => panic!("missing hash for file {}", path.display()),
  68                 }
  69             }
  70         });
  71     }
  72
  73     fn collect_checksums(&mut self, files: &HashSet<PathBuf>) {
  74         let collection_start = Instant::now();
  75         println!(
  76             "collecting hashes for {} tarballs across {} threads",
  77             files.len(),
  78             rayon::current_num_threads().min(files.len()),
  79         );
  80
  81         files.par_iter().for_each(|path| match hash(path) {
  82             Ok(hash) => {
  83                 self.collected.lock().unwrap().insert(path.clone(), hash);
  84             }
  85             Err(err) => eprintln!("error while fetching the hash for {}: {}", path.display(), err),
  86         });
  87
  88         println!("collected {} hashes in {:.2?}", files.len(), collection_start.elapsed());
  89     }
  90 }
  91
  92 fn hash(path: &Path) -> Result<String, Box<dyn Error>> {
  93     let mut file = BufReader::new(File::open(path)?);
  94     let mut sha256 = Sha256::default();
  95     std::io::copy(&mut file, &mut sha256)?;
  96     Ok(hex::encode(sha256.finalize()))
  97 }