diff --git a/src/tools/genpolicy/Cargo.lock b/src/tools/genpolicy/Cargo.lock index 892d0d070e7c..4c18b6869a4a 100644 --- a/src/tools/genpolicy/Cargo.lock +++ b/src/tools/genpolicy/Cargo.lock @@ -359,6 +359,16 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "fs2" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9564fc758e15025b46aa6643b1b77d047d1a56a1aea6e01002ac0c7026876213" +dependencies = [ + "libc", + "winapi", +] + [[package]] name = "futures" version = "0.3.28" @@ -469,6 +479,7 @@ dependencies = [ "docker_credential", "env_logger", "flate2", + "fs2", "generic-array", "log", "oci", diff --git a/src/tools/genpolicy/Cargo.toml b/src/tools/genpolicy/Cargo.toml index e5fea58008d3..87429245b200 100644 --- a/src/tools/genpolicy/Cargo.toml +++ b/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" diff --git a/src/tools/genpolicy/src/registry.rs b/src/tools/genpolicy/src/registry.rs index fef7a5672ded..dd8763dec137 100644 --- a/src/tools/genpolicy/src/registry.rs +++ b/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 { - 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 = 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 { + // 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 = 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 { 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 {