From 895614df529b13c5fd5dd3a6adb440b7d7114d55 Mon Sep 17 00:00:00 2001
From: Brian H <brian.hardock@fermyon.com>
Date: Mon, 27 Jan 2025 15:27:48 -0700
Subject: [PATCH] feat(push): allow pushing composed components to registry

Signed-off-by: Brian H <brian.hardock@fermyon.com>
---
 Cargo.lock               |   3 +
 crates/oci/Cargo.toml    |   3 +
 crates/oci/src/client.rs | 146 ++++++++++++++++++++++++++++++---------
 src/commands/registry.rs |  13 +++-
 4 files changed, 130 insertions(+), 35 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 318106c7eb..1d4da00a9f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -8226,6 +8226,7 @@ dependencies = [
  "anyhow",
  "async-compression",
  "async-tar",
+ "async-trait",
  "base64 0.22.1",
  "chrono",
  "dirs 5.0.1",
@@ -8238,6 +8239,8 @@ dependencies = [
  "serde",
  "serde_json",
  "spin-common",
+ "spin-componentize",
+ "spin-compose",
  "spin-loader",
  "spin-locked-app",
  "tempfile",
diff --git a/crates/oci/Cargo.toml b/crates/oci/Cargo.toml
index 8e17099af7..6f28147920 100644
--- a/crates/oci/Cargo.toml
+++ b/crates/oci/Cargo.toml
@@ -8,6 +8,7 @@ edition = { workspace = true }
 anyhow = { workspace = true }
 async-compression = { version = "0.4", features = ["gzip", "tokio"] }
 async-tar = "0.5"
+async-trait = { workspace = true }
 base64 = "0.22"
 chrono = "0.4"
 # Fork with updated auth to support ACR login
@@ -22,6 +23,8 @@ reqwest = "0.12"
 serde = { workspace = true }
 serde_json = { workspace = true }
 spin-common = { path = "../common" }
+spin-componentize = { path = "../componentize" }
+spin-compose = { path = "../compose" }
 spin-loader = { path = "../loader" }
 spin-locked-app = { path = "../locked-app" }
 tempfile = { workspace = true }
diff --git a/crates/oci/src/client.rs b/crates/oci/src/client.rs
index 459ee4a5c0..26b7ba12bf 100644
--- a/crates/oci/src/client.rs
+++ b/crates/oci/src/client.rs
@@ -4,6 +4,7 @@ use std::collections::{BTreeMap, HashMap};
 use std::path::{Path, PathBuf};
 
 use anyhow::{bail, Context, Result};
+use async_trait::async_trait;
 use docker_credential::DockerCredential;
 use futures_util::future;
 use futures_util::stream::{self, StreamExt, TryStreamExt};
@@ -18,7 +19,7 @@ use spin_common::ui::quoted_path;
 use spin_common::url::parse_file_url;
 use spin_loader::cache::Cache;
 use spin_loader::FilesMountStrategy;
-use spin_locked_app::locked::{ContentPath, ContentRef, LockedApp};
+use spin_locked_app::locked::{ContentPath, ContentRef, LockedApp, LockedComponentSource};
 use tokio::fs;
 use walkdir::WalkDir;
 
@@ -67,6 +68,16 @@ enum AssemblyMode {
     Archive,
 }
 
+/// Indicates whether to compose the components of a Spin application when pushing an image.
+#[derive(Default)]
+enum ComposeMode {
+    /// Compose components before pushing the image.
+    #[default]
+    All,
+    /// Skip composing components before pushing the image.
+    Skip,
+}
+
 /// Client for interacting with an OCI registry for Spin applications.
 pub struct Client {
     /// Global cache for the metadata, Wasm modules, and static assets pulled from OCI registries.
@@ -119,6 +130,7 @@ impl Client {
         reference: impl AsRef<str>,
         annotations: Option<BTreeMap<String, String>>,
         infer_annotations: InferPredefinedAnnotations,
+        compose_mode: ComposeMode,
     ) -> Result<Option<String>> {
         let reference: Reference = reference
             .as_ref()
@@ -137,8 +149,15 @@ impl Client {
         )
         .await?;
 
-        self.push_locked_core(locked, auth, reference, annotations, infer_annotations)
-            .await
+        self.push_locked_core(
+            locked,
+            auth,
+            reference,
+            annotations,
+            infer_annotations,
+            compose_mode,
+        )
+        .await
     }
 
     /// Push a Spin application to an OCI registry and return the digest (or None
@@ -149,6 +168,7 @@ impl Client {
         reference: impl AsRef<str>,
         annotations: Option<BTreeMap<String, String>>,
         infer_annotations: InferPredefinedAnnotations,
+        compose_mode: ComposeMode,
     ) -> Result<Option<String>> {
         let reference: Reference = reference
             .as_ref()
@@ -156,8 +176,15 @@ impl Client {
             .with_context(|| format!("cannot parse reference {}", reference.as_ref()))?;
         let auth = Self::auth(&reference).await?;
 
-        self.push_locked_core(locked, auth, reference, annotations, infer_annotations)
-            .await
+        self.push_locked_core(
+            locked,
+            auth,
+            reference,
+            annotations,
+            infer_annotations,
+            compose_mode,
+        )
+        .await
     }
 
     /// Push a Spin application to an OCI registry and return the digest (or None
@@ -169,10 +196,11 @@ impl Client {
         reference: Reference,
         annotations: Option<BTreeMap<String, String>>,
         infer_annotations: InferPredefinedAnnotations,
+        compose_mode: ComposeMode,
     ) -> Result<Option<String>> {
         let mut locked_app = locked.clone();
         let mut layers = self
-            .assemble_layers(&mut locked_app, AssemblyMode::Simple)
+            .assemble_layers(&mut locked_app, AssemblyMode::Simple, compose_mode)
             .await
             .context("could not assemble layers for locked application")?;
 
@@ -183,7 +211,7 @@ impl Client {
         {
             locked_app = locked.clone();
             layers = self
-                .assemble_layers(&mut locked_app, AssemblyMode::Archive)
+                .assemble_layers(&mut locked_app, AssemblyMode::Archive, compose_mode)
                 .await
                 .context("could not assemble archive layers for locked application")?;
         }
@@ -246,43 +274,59 @@ impl Client {
         &mut self,
         locked: &mut LockedApp,
         assembly_mode: AssemblyMode,
+        compose_mode: ComposeMode,
     ) -> Result<Vec<ImageLayer>> {
         let mut layers = Vec::new();
         let mut components = Vec::new();
         for mut c in locked.clone().components {
-            // Add the wasm module for the component as layers.
-            let source = c
-                .clone()
-                .source
-                .content
-                .source
-                .context("component loaded from disk should contain a file source")?;
+            match compose_mode {
+                ComposeMode::All => {
+                    let composed = spin_compose::compose(&ComponentSourceLoader, &c)
+                        .await
+                        .with_context(|| {
+                            format!("failed to resolve dependencies for component {:?}", c.id)
+                        })?;
+
+                    let layer = ImageLayer::new(composed, WASM_LAYER_MEDIA_TYPE.to_string(), None);
+                    c.source.content = self.content_ref_for_layer(&layer);
+                    c.dependencies.clear();
+                    layers.push(layer);
+                }
+                ComposeMode::Skip => {
+                    // Add the wasm module for the component as layers.
+                    let source = c
+                        .clone()
+                        .source
+                        .content
+                        .source
+                        .context("component loaded from disk should contain a file source")?;
 
-            let source = parse_file_url(source.as_str())?;
-            let layer = Self::wasm_layer(&source).await?;
+                    let source = parse_file_url(source.as_str())?;
+                    let layer = Self::wasm_layer(&source).await?;
 
-            // Update the module source with the content ref of the layer.
-            c.source.content = self.content_ref_for_layer(&layer);
+                    // Update the module source with the content ref of the layer.
+                    c.source.content = self.content_ref_for_layer(&layer);
 
-            layers.push(layer);
+                    layers.push(layer);
 
-            let mut deps = BTreeMap::default();
-            for (dep_name, mut dep) in c.dependencies {
-                let source = dep
-                    .source
-                    .content
-                    .source
-                    .context("dependency loaded from disk should contain a file source")?;
-                let source = parse_file_url(source.as_str())?;
+                    let mut deps = BTreeMap::default();
+                    for (dep_name, mut dep) in c.dependencies {
+                        let source =
+                            dep.source.content.source.context(
+                                "dependency loaded from disk should contain a file source",
+                            )?;
+                        let source = parse_file_url(source.as_str())?;
 
-                let layer = Self::wasm_layer(&source).await?;
+                        let layer = Self::wasm_layer(&source).await?;
 
-                dep.source.content = self.content_ref_for_layer(&layer);
-                deps.insert(dep_name, dep);
+                        dep.source.content = self.content_ref_for_layer(&layer);
+                        deps.insert(dep_name, dep);
 
-                layers.push(layer);
+                        layers.push(layer);
+                    }
+                    c.dependencies = deps;
+                }
             }
-            c.dependencies = deps;
 
             let mut files = Vec::new();
             for f in c.files {
@@ -669,6 +713,32 @@ impl Client {
     }
 }
 
+struct ComponentSourceLoader;
+
+#[async_trait]
+impl spin_compose::ComponentSourceLoader for ComponentSourceLoader {
+    async fn load_component_source(
+        &self,
+        source: &LockedComponentSource,
+    ) -> anyhow::Result<Vec<u8>> {
+        let source = source
+            .content
+            .source
+            .as_ref()
+            .context("component loaded from disk should contain a file source")?;
+
+        let source = parse_file_url(source.as_str())?;
+
+        let bytes = fs::read(&source)
+            .await
+            .with_context(|| format!("cannot read wasm module {}", quoted_path(source)))?;
+
+        let component = spin_componentize::componentize_if_necessary(&bytes)?;
+
+        Ok(component.into())
+    }
+}
+
 /// Unpack contents of the provided archive layer, represented by bytes and its
 /// corresponding digest, into the provided cache.
 /// A temporary staging directory is created via tempfile::tempdir() to store
@@ -946,6 +1016,7 @@ mod test {
             locked_components: Vec<LockedComponent>,
             expected_layer_count: usize,
             expected_error: Option<&'static str>,
+            compose_mode: ComposeMode,
         }
 
         let tests: Vec<TestCase> = [
@@ -968,6 +1039,7 @@ mod test {
                 }}]),
                 expected_layer_count: 2,
                 expected_error: None,
+                compose_mode: ComposeMode::Skip,
             },
             TestCase {
                 name: "One component layer and two file layers",
@@ -992,6 +1064,7 @@ mod test {
                 }]),
                 expected_layer_count: 3,
                 expected_error: None,
+                compose_mode: ComposeMode::Skip,
             },
             TestCase {
                 name: "One component layer and one file with inlined content",
@@ -1012,6 +1085,7 @@ mod test {
                 }]),
                 expected_layer_count: 1,
                 expected_error: None,
+                compose_mode: ComposeMode::Skip,
             },
             TestCase {
                 name: "One component layer and one dependency component layer",
@@ -1036,6 +1110,7 @@ mod test {
                 }]),
                 expected_layer_count: 2,
                 expected_error: None,
+                compose_mode: ComposeMode::Skip,
             },
             TestCase {
                 name: "Component has no source",
@@ -1050,6 +1125,7 @@ mod test {
                 }]),
                 expected_layer_count: 0,
                 expected_error: Some("Invalid URL: \"\""),
+                compose_mode: ComposeMode::Skip,
             },
             TestCase {
                 name: "Duplicate component sources",
@@ -1070,6 +1146,7 @@ mod test {
                 }}]),
                 expected_layer_count: 1,
                 expected_error: None,
+                compose_mode: ComposeMode::Skip,
             },
             TestCase {
                 name: "Duplicate file paths",
@@ -1107,6 +1184,7 @@ mod test {
                 }]),
                 expected_layer_count: 4,
                 expected_error: None,
+                compose_mode: ComposeMode::Skip,
             },
         ]
         .to_vec();
@@ -1137,7 +1215,7 @@ mod test {
                     assert_eq!(
                         e,
                         client
-                            .assemble_layers(&mut locked, AssemblyMode::Simple)
+                            .assemble_layers(&mut locked, AssemblyMode::Simple, tc.compose_mode)
                             .await
                             .unwrap_err()
                             .to_string(),
@@ -1149,7 +1227,7 @@ mod test {
                     assert_eq!(
                         tc.expected_layer_count,
                         client
-                            .assemble_layers(&mut locked, AssemblyMode::Simple)
+                            .assemble_layers(&mut locked, AssemblyMode::Simple, tc.compose_mode)
                             .await
                             .unwrap()
                             .len(),
diff --git a/src/commands/registry.rs b/src/commands/registry.rs
index a044ba6526..20730f95a3 100644
--- a/src/commands/registry.rs
+++ b/src/commands/registry.rs
@@ -3,7 +3,7 @@ use anyhow::{Context, Result};
 use clap::{Parser, Subcommand};
 use indicatif::{ProgressBar, ProgressStyle};
 use spin_common::arg_parser::parse_kv;
-use spin_oci::{client::InferPredefinedAnnotations, Client};
+use spin_oci::{client::InferPredefinedAnnotations, Client, ComposeMode};
 use std::{io::Read, path::PathBuf, time::Duration};
 
 /// Commands for working with OCI registries to distribute applications.
@@ -49,6 +49,10 @@ pub struct Push {
     )]
     pub insecure: bool,
 
+    /// Skip composing the application's components before pushing it.
+    #[clap(long = "skip-compose")]
+    pub skip_compose: bool,
+
     /// Specifies to perform `spin build` before pushing the application.
     #[clap(long, takes_value = false, env = ALWAYS_BUILD_ENV)]
     pub build: bool,
@@ -88,12 +92,19 @@ impl Push {
 
         let _spinner = create_dotted_spinner(2000, "Pushing app to the Registry".to_owned());
 
+        let compose_mode = if self.skip_compose {
+            ComposeMode::Skip
+        } else {
+            ComposeMode::All
+        };
+
         let digest = client
             .push(
                 &app_file,
                 &self.reference,
                 annotations,
                 InferPredefinedAnnotations::All,
+                self.skip_compose,
             )
             .await?;
         match digest {