genpolicy: allow specifying layer cache file

Add --layers-cache-file-path flag to allow the user to
specify where the cache file for the container layers
is saved. This allows e.g. to have one cache file
independent of the user's working directory.

Signed-off-by: Leonard Cohnen <lc@edgeless.systems>

src/tools/genpolicy/src/registry.rs

@@ -97,7 +97,7 @@ impl Container {
                 let config_layer: DockerConfigLayer =
                     serde_json::from_str(&config_layer_str).unwrap();
                 let image_layers = get_image_layers(
-                    config.use_cache,
+                    config.layers_cache_file_path.clone(),
                     &mut client,
                     &reference,
                     &manifest,
@@ -228,7 +228,7 @@ impl Container {
 }
 
 async fn get_image_layers(
-    use_cached_files: bool,
+    layers_cache_file_path: Option<String>,
     client: &mut Client,
     reference: &Reference,
     manifest: &manifest::OciImageManifest,
@@ -247,7 +247,7 @@ async fn get_image_layers(
                 layers.push(ImageLayer {
                     diff_id: config_layer.rootfs.diff_ids[layer_index].clone(),
                     verity_hash: get_verity_hash(
-                        use_cached_files,
+                        layers_cache_file_path.clone(),
                         client,
                         reference,
                         &layer.digest,
@@ -267,7 +267,7 @@ async fn get_image_layers(
 }
 
 async fn get_verity_hash(
-    use_cached_files: bool,
+    layers_cache_file_path: Option<String>,
     client: &mut Client,
     reference: &Reference,
     layer_digest: &str,
@@ -275,7 +275,6 @@ async fn get_verity_hash(
 ) -> Result<String> {
     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);
@@ -289,8 +288,8 @@ async fn get_verity_hash(
     let mut error = false;
 
     // get value from store and return if it exists
-    if use_cached_files {
-        verity_hash = read_verity_from_store(cache_file, diff_id)?;
+    if let Some(path) = layers_cache_file_path.as_ref() {
+        verity_hash = read_verity_from_store(path, diff_id)?;
         info!("Using cache file");
         info!("dm-verity root hash: {verity_hash}");
     }
@@ -318,8 +317,8 @@ async fn get_verity_hash(
                 }
                 Ok(v) => {
                     verity_hash = v;
-                    if use_cached_files {
-                        add_verity_to_store(cache_file, diff_id, &verity_hash)?;
+                    if let Some(path) = layers_cache_file_path.as_ref() {
+                        add_verity_to_store(path, diff_id, &verity_hash)?;
                     }
                     info!("dm-verity root hash: {verity_hash}");
                 }
@@ -330,8 +329,8 @@ async fn get_verity_hash(
     temp_dir.close()?;
     if error {
         // remove the cache file if we're using it
-        if use_cached_files {
-            std::fs::remove_file(cache_file)?;
+        if let Some(path) = layers_cache_file_path.as_ref() {
+            std::fs::remove_file(path)?;
         }
         warn!("{error_message}");
     }
@@ -458,7 +457,12 @@ pub fn get_verity_hash_value(path: &Path) -> Result<String> {
 
 pub async fn get_container(config: &Config, image: &str) -> Result<Container> {
     if let Some(socket_path) = &config.containerd_socket_path {
-        return Container::new_containerd_pull(config.use_cache, image, socket_path).await;
+        return Container::new_containerd_pull(
+            config.layers_cache_file_path.clone(),
+            image,
+            socket_path,
+        )
+        .await;
     }
     Container::new(config, image).await
 }

src/tools/genpolicy/src/registry_containerd.rs

@@ -28,7 +28,7 @@ use tower::service_fn;
 
 impl Container {
     pub async fn new_containerd_pull(
-        use_cached_files: bool,
+        layers_cache_file_path: Option<String>,
         image: &str,
         containerd_socket_path: &str,
     ) -> Result<Self> {
@@ -58,8 +58,13 @@ impl Container {
         let config_layer = get_config_layer(image_ref_str, k8_cri_image_client)
             .await
             .unwrap();
-        let image_layers =
-            get_image_layers(use_cached_files, &manifest, &config_layer, &ctrd_client).await?;
+        let image_layers = get_image_layers(
+            layers_cache_file_path,
+            &manifest,
+            &config_layer,
+            &ctrd_client,
+        )
+        .await?;
 
         Ok(Container {
             config_layer,
@@ -242,7 +247,7 @@ pub fn build_auth(reference: &Reference) -> Option<AuthConfig> {
 }
 
 pub async fn get_image_layers(
-    use_cached_files: bool,
+    layers_cache_file_path: Option<String>,
     manifest: &serde_json::Value,
     config_layer: &DockerConfigLayer,
     client: &containerd_client::Client,
@@ -261,7 +266,7 @@ pub async fn get_image_layers(
                 let imageLayer = ImageLayer {
                     diff_id: config_layer.rootfs.diff_ids[layer_index].clone(),
                     verity_hash: get_verity_hash(
-                        use_cached_files,
+                        layers_cache_file_path.clone(),
                         layer["digest"].as_str().unwrap(),
                         client,
                         &config_layer.rootfs.diff_ids[layer_index].clone(),
@@ -280,14 +285,13 @@ pub async fn get_image_layers(
 }
 
 async fn get_verity_hash(
-    use_cached_files: bool,
+    layers_cache_file_path: Option<String>,
     layer_digest: &str,
     client: &containerd_client::Client,
     diff_id: &str,
 ) -> Result<String> {
     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);
@@ -300,8 +304,8 @@ async fn get_verity_hash(
     let mut error_message = "".to_string();
     let mut error = false;
 
-    if use_cached_files {
-        verity_hash = read_verity_from_store(cache_file, diff_id)?;
+    if let Some(path) = layers_cache_file_path.as_ref() {
+        verity_hash = read_verity_from_store(path, diff_id)?;
         info!("Using cache file");
         info!("dm-verity root hash: {verity_hash}");
     }
@@ -328,8 +332,8 @@ async fn get_verity_hash(
                 }
                 Ok(v) => {
                     verity_hash = v;
-                    if use_cached_files {
-                        add_verity_to_store(cache_file, diff_id, &verity_hash)?;
+                    if let Some(path) = layers_cache_file_path.as_ref() {
+                        add_verity_to_store(path, diff_id, &verity_hash)?;
                     }
                     info!("dm-verity root hash: {verity_hash}");
                 }
@@ -339,8 +343,8 @@ async fn get_verity_hash(
     temp_dir.close()?;
     if error {
         // remove the cache file if we're using it
-        if use_cached_files {
-            std::fs::remove_file(cache_file)?;
+        if let Some(path) = layers_cache_file_path.as_ref() {
+            std::fs::remove_file(path)?;
         }
         warn!("{error_message}");
     }

src/tools/genpolicy/src/utils.rs

@@ -88,6 +88,14 @@ struct CommandLineOptions {
         help = "If specified, resources that have a runtimeClassName field defined will only receive a policy if the parameter is a prefix one of the given runtime class names."
     )]
     runtime_class_names: Vec<String>,
+
+    #[clap(
+        long,
+        help = "Path to the layers cache file. This file is used to store the layers cache information. The default value is ./layers-cache.json.",
+        default_missing_value = "./layers-cache.json",
+        require_equals = true
+    )]
+    layers_cache_file_path: Option<String>,
 }
 
 /// Application configuration, derived from on command line parameters.
@@ -106,6 +114,7 @@ pub struct Config {
     pub raw_out: bool,
     pub base64_out: bool,
     pub containerd_socket_path: Option<String>,
+    pub layers_cache_file_path: Option<String>,
 }
 
 impl Config {
@@ -123,6 +132,12 @@ impl Config {
             None
         };
 
+        let mut layers_cache_file_path = args.layers_cache_file_path;
+        // preserve backwards compatibility for only using the `use_cached_files` flag
+        if args.use_cached_files && layers_cache_file_path.is_none() {
+            layers_cache_file_path = Some(String::from("./layers-cache.json"));
+        }
+
         let settings = settings::Settings::new(&args.json_settings_path);
 
         Self {
@@ -137,6 +152,7 @@ impl Config {
             raw_out: args.raw_out,
             base64_out: args.base64_out,
             containerd_socket_path: args.containerd_socket_path,
+            layers_cache_file_path,
         }
     }
 }

tests/integration/kubernetes/k8s-policy-pod.bats

@@ -63,6 +63,20 @@ setup() {
 	kubectl wait --for=condition=Ready "--timeout=${timeout}" pod "${pod_name}"
 }
 
+@test "Successful pod with auto-generated policy and custom layers cache path" {
+	tmp_path=$(mktemp -d)
+
+	auto_generate_policy "${pod_config_dir}" "${testcase_pre_generate_pod_yaml}" "${testcase_pre_generate_configmap_yaml}" \
+		"--layers-cache-file-path=${tmp_path}/cache.json"
+
+	[ -f "${tmp_path}/cache.json" ]
+	rm -r "${tmp_path}"
+
+	kubectl create -f "${testcase_pre_generate_configmap_yaml}"
+	kubectl create -f "${testcase_pre_generate_pod_yaml}"
+	kubectl wait --for=condition=Ready "--timeout=${timeout}" pod "${pod_name}"
+}
+
 # Common function for several test cases from this bats script.
 test_pod_policy_error() {
 	kubectl create -f "${correct_configmap_yaml}"