Update caching mechanism to allow for concurrent genpolicy runs

src/tools/genpolicy/Cargo.toml

@@ -50,3 +50,4 @@ sha2 = "0.10.6"
 tarindex = { path = "../../tardev-snapshotter/tarindex" }
 tempfile = "3.5.0"
 zerocopy = "0.6.1"
+fs2 = "0.4.3"

src/tools/genpolicy/src/registry.rs

@@ -18,7 +18,11 @@ use oci_distribution::{manifest, secrets::RegistryAuth, Client, Reference};
 use serde::{Deserialize, Serialize};
 use sha2::{digest::typenum::Unsigned, digest::OutputSizeUser, Sha256};
 use std::{io, io::Seek, io::Write, path::Path};
-use tokio::{fs, io::AsyncWriteExt};
+use tokio::{io::AsyncWriteExt};
+use std::io::{BufWriter};
+use std::fs::OpenOptions;
+use fs2::FileExt;
+
 
 /// Container image properties obtained from an OCI repository.
 #[derive(Clone, Debug, Default)]
@@ -54,7 +58,7 @@ struct DockerRootfs {
 }
 
 /// This application's image layer properties.
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct ImageLayer {
     pub diff_id: String,
     pub verity_hash: String,
@@ -238,6 +242,7 @@ async fn get_image_layers(
                         client,
                         reference,
                         &layer.digest,
+                        &config_layer.rootfs.diff_ids[layer_index].clone(),
                     )
                     .await?,
                 });
@@ -252,130 +257,170 @@ async fn get_image_layers(
     Ok(layers)
 }
 
-fn delete_files(decompressed_path: &Path, compressed_path: &Path, verity_path: &Path) {
-    let _ = fs::remove_file(&decompressed_path);
-    let _ = fs::remove_file(&compressed_path);
-    let _ = fs::remove_file(&verity_path);
-}
-
 async fn get_verity_hash(
     use_cached_files: bool,
     client: &mut Client,
     reference: &Reference,
     layer_digest: &str,
+    diff_id: &str,
 ) -> Result<String> {
-    let base_dir = std::path::Path::new("layers_cache");
-
+    let temp_dir = tempfile::tempdir_in(".")?;
+    let base_dir = temp_dir.path();
+    let cache_file = "layers-cache.json";
     // Use file names supported by both Linux and Windows.
     let file_name = str::replace(&layer_digest, ":", "-");
-
     let mut decompressed_path = base_dir.join(file_name);
     decompressed_path.set_extension("tar");
 
     let mut compressed_path = decompressed_path.clone();
     compressed_path.set_extension("gz");
 
-    let mut verity_path = decompressed_path.clone();
-    verity_path.set_extension("verity");
-
     let mut verity_hash = "".to_string();
     let mut error_message = "".to_string();
     let mut error = false;
 
-    if use_cached_files && verity_path.exists() {
-        info!("Using cached file {:?}", &verity_path);
-    } else if let Err(e) = create_verity_hash_file(
-        use_cached_files,
-        client,
-        reference,
-        layer_digest,
-        &base_dir,
-        &decompressed_path,
-        &compressed_path,
-        &verity_path,
-    )
-    .await
-    {
-        error = true;
-        error_message = format!("Failed to create verity hash for {layer_digest}, error {e}");
+    // get value from store and return if it exists
+    if use_cached_files {
+        verity_hash = read_verity_from_store(&cache_file, &diff_id)?;
+        info!("Using cache file");
+        info!("dm-verity root hash: {verity_hash}");
     }
 
-    if !error {
-        match std::fs::read_to_string(&verity_path) {
-            Err(e) => {
-                error = true;
-                error_message = format!("Failed to read {:?}, error {e}", &verity_path);
-            }
-            Ok(v) => {
-                verity_hash = v;
-                info!("dm-verity root hash: {verity_hash}");
+    // create the layer files
+    if verity_hash == "" {
+        if let Err(e) = create_decompressed_layer_file(
+            client,
+            reference,
+            layer_digest,
+            &decompressed_path,
+            &compressed_path,
+        )
+        .await
+        {
+            error_message = format!(
+                "Failed to create verity hash for {layer_digest}, error {e}"
+            );
+            error = true
+        };
+
+        if !error {
+            match get_verity_hash_value(&decompressed_path) {
+                Err(e) => {
+                    error_message = format!("Failed to get verity hash {e}");
+                    error = true;
+                }
+                Ok(v) => {
+                    verity_hash = v;
+                    if use_cached_files {
+                        add_verity_to_store(&cache_file, &diff_id, &verity_hash)?;
+                    }
+                    info!("dm-verity root hash: {verity_hash}");
+                }
             }
         }
     }
 
-    if !use_cached_files {
-        let _ = std::fs::remove_dir_all(&base_dir);
-    } else if error {
-        delete_files(&decompressed_path, &compressed_path, &verity_path);
-    }
-
+    temp_dir.close()?;
     if error {
-        panic!("{error_message}");
-    } else {
-        Ok(verity_hash)
+        // remove the cache file if we're using it
+        if use_cached_files {
+            std::fs::remove_file(&cache_file)?;
+        }
+        warn!("{error_message}");
     }
+    Ok(verity_hash)
 }
 
-async fn create_verity_hash_file(
-    use_cached_files: bool,
-    client: &mut Client,
-    reference: &Reference,
-    layer_digest: &str,
-    base_dir: &Path,
-    decompressed_path: &Path,
-    compressed_path: &Path,
-    verity_path: &Path,
+// the store is a json file that matches layer hashes to verity hashes
+fn add_verity_to_store(
+    cache_file: &str,
+    diff_id: &str,
+    verity_hash: &str,
 ) -> Result<()> {
-    if use_cached_files && decompressed_path.exists() {
-        info!("Using cached file {:?}", &decompressed_path);
+    // open the json file in read mode, create it if it doesn't exist
+    let read_file = OpenOptions::new()
+        .read(true)
+        .write(true)
+        .create(true)
+        .open(cache_file)?;
+
+    let mut data: Vec<ImageLayer> = if let Ok(vec) = serde_json::from_reader(read_file) {
+        vec
     } else {
-        std::fs::create_dir_all(&base_dir)?;
+        // Delete the malformed file here if it's present
+        Vec::new()
+    };
 
-        create_decompressed_layer_file(
-            use_cached_files,
-            client,
-            reference,
-            layer_digest,
-            &decompressed_path,
-            &compressed_path,
-        )
-        .await?;
+    // Add new data to the deserialized JSON
+    data.push(ImageLayer{
+        diff_id: diff_id.to_string(),
+        verity_hash: verity_hash.to_string(),
+    });
+
+    // Serialize in pretty format
+    let serialized = serde_json::to_string_pretty(&data)?;
+
+    // Open the JSON file to write
+    let file = OpenOptions::new()
+        .write(true)
+        .open(cache_file)?;
+
+    // try to lock the file, if it fails, get the error
+    let result = file.try_lock_exclusive();
+    if result.is_err() {
+        warn!("Waiting to lock file: {cache_file}");
+        file.lock_exclusive()?;
+    }
+    // Write the serialized JSON to the file
+    let mut writer = BufWriter::new(&file);
+    writeln!(writer, "{}", serialized)?;
+    writer.flush()?;
+    file.unlock()?;
+    Ok(())
+}
+
+// helper function to read the verity hash from the store
+// returns empty string if not found or file does not exist
+fn read_verity_from_store(cache_file: &str, diff_id: &str) -> Result<String> {
+    // see if file exists, return empty if not
+    if !Path::new(cache_file).exists() {
+        return Ok("".to_string());
+    }
+
+    let file = OpenOptions::new()
+        .read(true)
+        .open(cache_file)?;
+
+    // If the file is empty, return empty string
+    if file.metadata()?.len() == 0 {
+        return Ok("".to_string());
     }
 
-    do_create_verity_hash_file(decompressed_path, verity_path)
+    let data: Vec<ImageLayer> = serde_json::from_reader(file)?;
+    for layer in data {
+        if layer.diff_id == diff_id {
+            return Ok(layer.verity_hash);
+        }
+    }
+    Ok("".to_string())
 }
 
 async fn create_decompressed_layer_file(
-    use_cached_files: bool,
     client: &mut Client,
     reference: &Reference,
     layer_digest: &str,
     decompressed_path: &Path,
     compressed_path: &Path,
 ) -> Result<()> {
-    if use_cached_files && compressed_path.exists() {
-        info!("Using cached file {:?}", &compressed_path);
-    } else {
-        info!("Pulling layer {layer_digest}");
-        let mut file = tokio::fs::File::create(&compressed_path)
-            .await
-            .map_err(|e| anyhow!(e))?;
-        client
-            .pull_blob(&reference, layer_digest, &mut file)
-            .await
-            .map_err(|e| anyhow!(e))?;
-        file.flush().await.map_err(|e| anyhow!(e))?;
-    }
+    info!("Pulling layer {:?}", layer_digest);
+    let mut file = tokio::fs::File::create(&compressed_path)
+        .await
+        .map_err(|e| anyhow!(e))?;
+    client
+        .pull_blob(&reference, layer_digest, &mut file)
+        .await
+        .map_err(|e| anyhow!(e))?;
+    file.flush().await.map_err(|e| anyhow!(e))?;
 
     info!("Decompressing layer");
     let compressed_file = std::fs::File::open(&compressed_path).map_err(|e| anyhow!(e))?;
@@ -396,7 +441,7 @@ async fn create_decompressed_layer_file(
     Ok(())
 }
 
-fn do_create_verity_hash_file(path: &Path, verity_path: &Path) -> Result<()> {
+fn get_verity_hash_value(path: &Path) -> Result<String> {
     info!("Calculating dm-verity root hash");
     let mut file = std::fs::File::open(path)?;
     let size = file.seek(std::io::SeekFrom::End(0))?;
@@ -409,13 +454,7 @@ fn do_create_verity_hash_file(path: &Path, verity_path: &Path) -> Result<()> {
     let hash = verity::traverse_file(&mut file, 0, false, v, &mut verity::no_write)?;
     let result = format!("{:x}", hash);
 
-    let mut verity_file = std::fs::File::create(verity_path).map_err(|e| anyhow!(e))?;
-    verity_file
-        .write_all(result.as_bytes())
-        .map_err(|e| anyhow!(e))?;
-    verity_file.flush().map_err(|e| anyhow!(e))?;
-
-    Ok(())
+    Ok(result)
 }
 
 pub async fn get_container(use_cache: bool, image: &str) -> Result<Container> {