Merge pull request #48 from osusec/dr/image-tag-format

detjensrobert · web-flow · commit 45c7ba7524bd · 2025-04-28T22:53:11.000-07:00
Add config option for custom image tag format
diff --git a/src/access_handlers/docker.rs b/src/access_handlers/docker.rs
@@ -6,10 +6,11 @@ use bollard::{
 };
 use futures::{StreamExt, TryStreamExt};
 use itertools::Itertools;
+use minijinja;
 use tokio;
 use tracing::{debug, error, info, trace, warn};
 
-use crate::clients::docker;
+use crate::clients::{docker, render_strict};
 use crate::configparser::{get_config, get_profile_config};
 
 /// container registry / daemon access checks
@@ -26,7 +27,16 @@ pub async fn check(profile_name: &str) -> Result<()> {
 
     // build test image string
     // registry.example.com/somerepo/testimage:pleaseignore
-    let test_image = format!("{}/credstestimage", registry_config.domain);
+    let test_image = render_strict(
+        &registry_config.tag_format,
+        minijinja::context! {
+        domain => registry_config.domain,
+        challenge => "accesscheck",
+        container => "testimage",
+        profile => profile_name
+        },
+    )
+    .context("could not render tag format template")?;
     debug!("will push test image to {}", test_image);
 
     // push alpine image with build credentials
@@ -66,16 +76,16 @@ async fn check_build_credentials(client: &Docker, test_image: &str) -> Result<()
 
     let registry_config = &get_config()?.registry;
 
-    // rename alpine image as test image
-    let tag_opts = TagImageOptions {
-        repo: test_image,
-        tag: "latest",
-    };
+    // rename alpine image as test imag
+    let (repo, tag) = test_image
+        .rsplit_once(':')
+        .unwrap_or((test_image, "latest"));
+    let tag_opts = TagImageOptions { repo, tag };
     client.tag_image("alpine", Some(tag_opts)).await?;
 
     // now push test iamge to configured repo
-    debug!("pushing alpine to target registry");
-    let options = PushImageOptions { tag: "latest" };
+    debug!("pushing alpine to target registry as {}:{}", repo, tag);
+    let options = PushImageOptions { tag };
     let build_creds = DockerCredentials {
         username: Some(registry_config.build.user.clone()),
         password: Some(registry_config.build.pass.clone()),
@@ -84,7 +94,7 @@ async fn check_build_credentials(client: &Docker, test_image: &str) -> Result<()
     };
 
     client
-        .push_image(test_image, Some(options), Some(build_creds))
+        .push_image(repo, Some(options), Some(build_creds))
         .try_collect::<Vec<_>>()
         .await?;
 
@@ -100,8 +110,11 @@ async fn check_cluster_credentials(client: &Docker, test_image: &str) -> Result<
     let registry_config = &get_config()?.registry;
 
     // pull just-pushed alpine image from repo
+    let (repo, tag) = test_image
+        .rsplit_once(':')
+        .unwrap_or((test_image, "latest"));
     let alpine_test_image = CreateImageOptions {
-        from_image: test_image,
+        from_image: [repo, tag].join(":"),
         ..Default::default()
     };
     let cluster_creds = DockerCredentials {
diff --git a/src/builder/mod.rs b/src/builder/mod.rs
@@ -20,14 +20,6 @@ use crate::utils::TryJoinAll;
 pub mod artifacts;
 pub mod docker;
 
-// define tag format as reusable macro
-macro_rules! image_tag_str {
-    () => {
-        "{registry}/{challenge}-{container}:{profile}"
-    };
-}
-pub(super) use image_tag_str;
-
 /// Information about all of a challenge's build artifacts.
 #[derive(Debug)]
 pub struct BuildResult {
diff --git a/src/clients.rs b/src/clients.rs
@@ -345,3 +345,19 @@ pub async fn wait_for_status(client: &kube::Client, object: &DynamicObject) -> R
 
     Ok(())
 }
+
+//
+// Minijinja strict rendering with error
+//
+
+/// Similar to minijinja.render!(), but return Error if any undefined values.
+pub fn render_strict(template: &str, context: minijinja::Value) -> Result<String> {
+    let mut strict_env = minijinja::Environment::new();
+    // error on any undefined template variables
+    strict_env.set_undefined_behavior(minijinja::UndefinedBehavior::Strict);
+
+    let r = strict_env
+        .render_str(template, context)
+        .context(format!("could not render template {:?}", template))?;
+    Ok(r)
+}
diff --git a/src/configparser/challenge.rs b/src/configparser/challenge.rs
@@ -12,7 +12,7 @@ use std::str::FromStr;
 use tracing::{debug, error, info, trace, warn};
 use void::Void;
 
-use crate::builder::image_tag_str;
+use crate::clients::render_strict;
 use crate::configparser::config::Resource;
 use crate::configparser::field_coersion::string_or_struct;
 use crate::configparser::get_config;
@@ -153,13 +153,17 @@ impl ChallengeConfig {
 
         match &pod.image_source {
             ImageSource::Image(t) => Ok(t.to_string()),
-            ImageSource::Build(b) => Ok(format!(
-                image_tag_str!(),
-                registry = config.registry.domain,
-                challenge = self.name,
-                container = pod.name,
-                profile = profile_name
-            )),
+            // render image tag template from config
+            ImageSource::Build(b) => render_strict(
+                &get_config()?.registry.tag_format,
+                minijinja::context! {
+                    domain => config.registry.domain,
+                    challenge => self.slugify(),
+                    container => pod.name,
+                    profile => profile_name
+                },
+            )
+            .context("error rendering challenge image tag template"),
         }
     }
 
diff --git a/src/configparser/config.rs b/src/configparser/config.rs
@@ -18,6 +18,7 @@ pub fn parse() -> Result<RcdsConfig> {
         // keys by undoing the s/_/./ that the figment::split() did.
         var.to_string()
             .to_lowercase()
+            .replace("registry.tag.format", "registry.tag_format")
             .replace("frontend.", "frontend_")
             .replace("challenges.", "challenges_")
             .replace("s3.access.", "s3.access_")
@@ -60,10 +61,48 @@ struct RcdsConfig {
 #[derive(Debug, PartialEq, Serialize, Deserialize)]
 #[fully_pub]
 struct Registry {
+    /// Registry base url used to build part of the full image string.
+    ///
+    /// Example: `domain: "registry.io/myctf"`
     domain: String,
+
+    /// Container image tag format for challenge images.
+    ///
+    /// Format:
+    /// Jinja-style double-braces around field name (`{{ field_name }}`)
+    ///
+    /// Default, works for most registries (self-hosted, GCP, DigitalOcean, ...):
+    /// `"{{domain}}/{{challenge}}-{{container}}:{{profile}}"`
+    ///
+    /// For registries like AWS that make it hard to create individual repositories,
+    /// keep all the challenge info in the tag:
+    /// `"{{domain}}:{{challenge}}-{{container}}-{{profile}}"`
+    ///
+    /// Available fields:
+    /// - `domain`: the domain config field above; the repository base URL
+    /// - `challenge`: challenge name, slugified
+    /// - `container`: name of the specific pod in the challenge this image is for
+    /// - `profile`: the current deployment profile, for isolating images between environments
+    ///
+    /// Example:
+    ///
+    /// For challenge `pwn/notsh`, chal pod container `main`, profile `prod`, and the example domain:
+    /// ```py
+    /// the default --> "registry.io/myctf/pwn-notsh-main:prod"
+    ///
+    /// "{{domain}}:{{challenge}}-{{container}}" --> "registry.io/myctf:pwn-notsh-main"
+    /// ```
+    #[serde(default = "default_tag_format")]
+    tag_format: String,
+
+    /// Container registry login for pushing images during build/deploy
     build: UserPass,
+    /// Container registry login for pulling images in cluster. Can and should be read-only.
     cluster: UserPass,
 }
+fn default_tag_format() -> String {
+    "{{domain}}/{{challenge}}-{{container}}:{{profile}}".to_string()
+}
 
 #[derive(Debug, PartialEq, Serialize, Deserialize, Clone)]
 #[fully_pub]
diff --git a/src/tests/parsing/config.rs b/src/tests/parsing/config.rs
@@ -68,6 +68,131 @@ fn all_yaml() {
             flag_regex: "test{[a-zA-Z_]+}".to_string(),
             registry: Registry {
                 domain: "registry.example/test".to_string(),
+                tag_format: "{{domain}}/{{challenge}}-{{container}}:{{profile}}".to_string(),
+                build: UserPass {
+                    user: "admin".to_string(),
+                    pass: "notrealcreds".to_string(),
+                },
+                cluster: UserPass {
+                    user: "cluster".to_string(),
+                    pass: "alsofake".to_string(),
+                },
+            },
+            defaults: Defaults {
+                difficulty: 1,
+                resources: Resource {
+                    cpu: 1,
+                    memory: "500M".to_string(),
+                },
+            },
+            points: vec![ChallengePoints {
+                difficulty: 1,
+                min: 0,
+                max: 1337,
+            }],
+
+            deploy: HashMap::from([(
+                "testing".to_string(),
+                ProfileDeploy {
+                    challenges: HashMap::from([
+                        ("web/bar".to_string(), false),
+                        ("misc/foo".to_string(), true),
+                    ]),
+                },
+            )]),
+            profiles: HashMap::from([(
+                "testing".to_string(),
+                ProfileConfig {
+                    frontend_url: "https://frontend.example".to_string(),
+                    frontend_token: "secretsecretsecret".to_string(),
+                    challenges_domain: "chals.frontend.example".to_string(),
+                    kubeconfig: None,
+                    kubecontext: "testcluster".to_string(),
+                    s3: S3Config {
+                        bucket_name: "asset_testing".to_string(),
+                        endpoint: "s3.example".to_string(),
+                        region: "us-fake-1".to_string(),
+                        access_key: "accesskey".to_string(),
+                        secret_key: "secretkey".to_string(),
+                    },
+                    dns: serde_yml::to_value(HashMap::from([
+                        ("provider", "somebody"),
+                        ("thing", "whatever"),
+                    ]))
+                    .unwrap(),
+                },
+            )]),
+        };
+
+        assert_eq!(config, expected);
+
+        Ok(())
+    });
+}
+
+#[test]
+/// Test parsing RCDS config where all fields are specified in the yaml
+fn registry_tag_format() {
+    figment::Jail::expect_with(|jail| {
+        jail.clear_env();
+        jail.create_file(
+            "rcds.yaml",
+            r#"
+                flag_regex: test{[a-zA-Z_]+}
+
+                registry:
+                    domain: registry.example/test
+                    tag_format: "{{domain}}:{{challenge}}.{{container}}"
+                    build:
+                        user: admin
+                        pass: notrealcreds
+                    cluster:
+                        user: cluster
+                        pass: alsofake
+
+                defaults:
+                    difficulty: 1
+                    resources: { cpu: 1, memory: 500M }
+
+                points:
+                  - difficulty: 1
+                    min: 0
+                    max: 1337
+
+                deploy:
+                    testing:
+                        misc/foo: true
+                        web/bar: false
+
+                profiles:
+                    testing:
+                        frontend_url: https://frontend.example
+                        frontend_token: secretsecretsecret
+                        challenges_domain: chals.frontend.example
+                        kubecontext: testcluster
+                        s3:
+                            bucket_name: asset_testing
+                            endpoint: s3.example
+                            region: us-fake-1
+                            access_key: accesskey
+                            secret_key: secretkey
+                        dns:
+                            provider: somebody
+                            thing: whatever
+            "#,
+        )?;
+
+        let config = match parse() {
+            Ok(c) => Ok(c),
+            // figment::Error cannot coerce from anyhow::Error natively
+            Err(e) => Err(figment::Error::from(format!("{:?}", e))),
+        }?;
+
+        let expected = RcdsConfig {
+            flag_regex: "test{[a-zA-Z_]+}".to_string(),
+            registry: Registry {
+                domain: "registry.example/test".to_string(),
+                tag_format: "{{domain}}:{{challenge}}.{{container}}".to_string(),
                 build: UserPass {
                     user: "admin".to_string(),
                     pass: "notrealcreds".to_string(),
diff --git a/tests/repo/rcds.yaml b/tests/repo/rcds.yaml
@@ -2,6 +2,7 @@ flag_regex: dam{[a-zA-Z...]}
 
 registry:
   domain: localhost:5000/damctf
+  tag_format: "{{domain}}/{{challenge}}/{{container}}:{{profile}}"
   # or environment vars e.g. BEAVERCDS_REGISTRY_BUILD_USER / etc
   build:
     user: admin
