feat(boil): Add image size command

rust/boil/src/cli/image.rs

@@ -1,6 +1,12 @@
 use clap::{Args, Subcommand, ValueEnum};
 
-use crate::{cli::Cli, core::image::ImageSelector};
+use crate::{
+    cli::Cli,
+    core::{
+        image::ImageSelector,
+        platform::{Architecture, TargetPlatform},
+    },
+};
 
 #[derive(Debug, Args)]
 pub struct ImageArguments {
@@ -17,6 +23,9 @@ pub enum ImageCommand {
     ///
     /// Access tokens must be provided with the following name: `BOIL_REGISTRY_TOKEN_<REGISTRY_URI>`.
     Check(ImageCheckArguments),
+
+    /// Calculates the size of images known by boil.
+    Size(ImageSizeArguments),
 }
 
 #[derive(Debug, Args)]
@@ -35,14 +44,46 @@ pub struct ImageCheckArguments {
     pub image: Vec<ImageSelector>,
 
     // NOTE (@Techassi): Should this maybe be renamed to vendor_version?
-    /// The image version being built.
+    /// The image version to check.
+    #[arg(
+        short, long,
+        value_parser = Cli::parse_image_version,
+        default_value_t = Cli::default_image_version(),
+        help_heading = "Image Options"
+    )]
+    pub image_version: String,
+}
+
+#[derive(Debug, Args)]
+pub struct ImageSizeArguments {
+    /// Optionally specify one or more images to check. Checks all images by default.
+    pub image: Vec<ImageSelector>,
+
+    // NOTE (@Techassi): Should this maybe be renamed to vendor_version?
+    /// The image version to use.
     #[arg(
         short, long,
         value_parser = Cli::parse_image_version,
         default_value_t = Cli::default_image_version(),
         help_heading = "Image Options"
     )]
     pub image_version: String,
+
+    /// Target platform of the image.
+    #[arg(
+        short, long,
+        short_alias = 'a', alias = "architecture",
+        default_value_t = Self::default_architecture(),
+        help_heading = "Image Options"
+    )]
+    pub target_platform: TargetPlatform,
+}
+
+impl ImageSizeArguments {
+    // TODO: Auto-detect this
+    fn default_architecture() -> TargetPlatform {
+        TargetPlatform::Linux(Architecture::Amd64)
+    }
 }
 
 // #[derive(Clone, Debug, Default, strum::Display, strum::EnumString)]

rust/boil/src/cmd/image.rs

@@ -1,14 +1,21 @@
-use std::{collections::BTreeMap, io::IsTerminal};
+use std::{
+    collections::{BTreeMap, HashMap},
+    io::IsTerminal,
+};
 
 use secrecy::{ExposeSecret, SecretString};
+use serde::Serialize;
 use snafu::{ResultExt, Snafu};
 
 use crate::{
-    cli::{ImageCheckArguments, ImageListArguments, Pretty},
+    cli::{ImageCheckArguments, ImageListArguments, ImageSizeArguments, Pretty},
     config::Config,
     core::bakefile::{self, Targets, TargetsOptions},
-    models::TagList,
-    utils::{format_image_index_manifest_tag, format_registry_token_env_var_name},
+    models::{Manifest, TagList},
+    utils::{
+        format_image_index_manifest_tag, format_image_manifest_tag,
+        format_registry_token_env_var_name,
+    },
 };
 
 #[derive(Debug, Snafu)]
@@ -22,8 +29,8 @@ pub enum Error {
     #[snafu(display("failed to build request client"))]
     BuildClient { source: reqwest::Error },
 
-    #[snafu(display("failed to send request"))]
-    SendRequest { source: reqwest::Error },
+    #[snafu(display("failed to send request ({url})"))]
+    SendRequest { source: reqwest::Error, url: String },
 
     #[snafu(display("failed to deserialize response"))]
     DeserializeResponse { source: reqwest::Error },
@@ -51,7 +58,7 @@ pub fn list_images(arguments: ImageListArguments) -> Result<(), Error> {
         .context(BuildTargetsSnafu)?
     };
 
-    let list = targets
+    let list: BTreeMap<String, Vec<String>> = targets
         .into_iter()
         .map(|(image_name, (image_config, _))| {
             let versions: Vec<_> = image_config
@@ -109,7 +116,7 @@ pub async fn check_images(arguments: ImageCheckArguments, config: Config) -> Res
                 "https://{registry}/v2/{registry_namespace}/{image_name}/tags/list",
                 registry_namespace = registry_options.namespace,
             );
-            let request = client.get(url);
+            let request = client.get(&url);
 
             let request = match &registry_token {
                 Some(registry_token) => request.bearer_auth(registry_token.expose_secret()),
@@ -119,7 +126,7 @@ pub async fn check_images(arguments: ImageCheckArguments, config: Config) -> Res
             let tag_list: TagList = request
                 .send()
                 .await
-                .context(SendRequestSnafu)?
+                .context(SendRequestSnafu { url })?
                 .json()
                 .await
                 .context(DeserializeResponseSnafu)?;
@@ -148,14 +155,108 @@ pub async fn check_images(arguments: ImageCheckArguments, config: Config) -> Res
     Ok(())
 }
 
-fn print_to_stdout(list: BTreeMap<String, Vec<String>>, pretty: Pretty) -> Result<(), Error> {
+pub async fn calculate_size(arguments: ImageSizeArguments, config: Config) -> Result<(), Error> {
+    let targets = if arguments.image.is_empty() {
+        Targets::all(TargetsOptions {
+            only_entry: true,
+            non_recursive: true,
+        })
+        .context(BuildTargetsSnafu)?
+    } else {
+        Targets::set(
+            &arguments.image,
+            TargetsOptions {
+                only_entry: true,
+                non_recursive: true,
+            },
+        )
+        .context(BuildTargetsSnafu)?
+    };
+
+    let mut registry_tokens = BTreeMap::new();
+    let client = reqwest::ClientBuilder::new()
+        .build()
+        .context(BuildClientSnafu)?;
+
+    #[derive(Serialize)]
+    struct SizeResult {
+        images: HashMap<String, u64>,
+        total: u64,
+    }
+
+    let mut result = SizeResult {
+        images: HashMap::new(),
+        total: 0,
+    };
+
+    for (image_name, (image_config, _)) in targets {
+        for (registry, registry_options) in image_config.metadata.registries {
+            // Add tokens to a map so that we don't need construct the key and retrieve the value
+            // over and over again.
+            let registry_token = registry_tokens.entry(registry.clone()).or_insert_with(|| {
+                let name = format_registry_token_env_var_name(&registry);
+                std::env::var(name).ok().map(SecretString::from)
+            });
+
+            for (image_version, _) in image_config.versions.iter() {
+                let image_index_manifest_tag = format_image_index_manifest_tag(
+                    image_version,
+                    &config.metadata.vendor_tag_prefix,
+                    &arguments.image_version,
+                );
+
+                let manifest_tag = format_image_manifest_tag(
+                    &image_index_manifest_tag,
+                    arguments.target_platform.architecture(),
+                    // Never strip the architecture, because we need to reference the exact manifest
+                    // to be able to calculate the sizes
+                    false,
+                );
+
+                let url = format!(
+                    "https://{registry}/v2/{registry_namespace}/{image_name}/manifests/{manifest_tag}",
+                    registry_namespace = registry_options.namespace,
+                );
+
+                let request = client.get(&url);
+
+                let request = match &registry_token {
+                    Some(registry_token) => request.bearer_auth(registry_token.expose_secret()),
+                    None => request,
+                };
+
+                let manifest: Manifest = request
+                    .send()
+                    .await
+                    .context(SendRequestSnafu { url })?
+                    .json()
+                    .await
+                    .context(DeserializeResponseSnafu)?;
+
+                let layer_size = manifest.layers.iter().fold(0u64, |acc, e| acc + e.size);
+
+                let size = result.images.entry(image_name.clone()).or_default();
+                *size += layer_size;
+
+                result.total += layer_size;
+            }
+        }
+    }
+
+    print_to_stdout(result, Pretty::Always)
+}
+
+fn print_to_stdout<T>(value: T, pretty: Pretty) -> Result<(), Error>
+where
+    T: Serialize,
+{
     let stdout = std::io::stdout();
 
     match pretty {
         Pretty::Always | Pretty::Auto if stdout.is_terminal() => {
-            serde_json::to_writer_pretty(stdout, &list)
+            serde_json::to_writer_pretty(stdout, &value)
         }
-        _ => serde_json::to_writer(stdout, &list),
+        _ => serde_json::to_writer(stdout, &value),
     }
     .context(SerializeListSnafu)
 }

rust/boil/src/main.rs

@@ -75,6 +75,12 @@ async fn main() -> Result<(), Error> {
                     .await
                     .context(ImageSnafu)
             }
+            ImageCommand::Size(arguments) => {
+                let config = Config::from_file(&cli.config_path).context(ReadConfigSnafu)?;
+                cmd::image::calculate_size(arguments, config)
+                    .await
+                    .context(ImageSnafu)
+            }
         },
         Command::Images(arguments) => cmd::image::list_images(arguments).context(ImageSnafu),
         Command::Completions(arguments) => {

rust/boil/src/models/mod.rs

@@ -6,3 +6,16 @@ pub struct TagList {
     // pub _name: String,
     pub tags: Vec<String>,
 }
+
+// TODO (@Techassi): We should eventually use the complete, upstream types from oci-spec
+/// A partial OCI manifest.
+#[derive(Debug, Deserialize)]
+pub struct Manifest {
+    pub layers: Vec<ManifestLayer>,
+}
+
+/// A partial OCI manifest layer.
+#[derive(Debug, Deserialize)]
+pub struct ManifestLayer {
+    pub size: u64,
+}