Merge pull request #18 from bytecodealliance/wkg-get

Implement 'wkg get'

Cargo.toml

@@ -7,3 +7,8 @@ edition = "2021"
 version = "0.2.0"
 authors = ["The Wasmtime Project Developers"]
 license = "Apache-2.0 WITH LLVM-exception"
+
+[workspace.dependencies]
+tokio = "1.35.1"
+tracing = "0.1.40"
+wasm-pkg-loader = { version = "0.2.0", path = "crates/wasm-pkg-loader" }

crates/wasm-pkg-loader/src/package.rs

@@ -1,16 +1,19 @@
 use crate::{label::Label, Error};
 
+/// A package reference, consisting of kebab-case namespace and name, e.g. `wasm-pkg:loader`.
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
 pub struct PackageRef {
     namespace: Label,
     name: Label,
 }
 
 impl PackageRef {
+    /// Returns the namespace of the package.
     pub fn namespace(&self) -> &Label {
         &self.namespace
     }
 
+    /// Returns the name of the package.
     pub fn name(&self) -> &Label {
         &self.name
     }

crates/wkg/Cargo.toml

@@ -1,8 +1,18 @@
 [package]
 name = "wkg"
+description = "Wasm Package Tools CLI"
 edition.workspace = true
 version.workspace = true
 authors.workspace = true
 license.workspace = true
 
 [dependencies]
+anyhow = "1.0"
+clap = { version = "4.5.4", features = ["derive", "wrap_help"] }
+futures-util = { version = "0.3.29", features = ["io"] }
+tempfile = "3.10.1"
+tokio = { workspace = true, features = ["macros", "rt"] }
+tracing = { workspace = true }
+tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
+wasm-pkg-loader = { workspace = true }
+wit-component = "0.207"

crates/wkg/src/main.rs

@@ -1,3 +1,203 @@
-fn main() {
-    println!("wkg stub");
+mod package_spec;
+
+use std::{io::Seek, path::PathBuf};
+
+use anyhow::{ensure, Context};
+use clap::{Args, Parser, Subcommand, ValueEnum};
+use futures_util::TryStreamExt;
+use package_spec::PackageSpec;
+use tokio::io::AsyncWriteExt;
+use wasm_pkg_loader::ClientConfig;
+use wit_component::DecodedWasm;
+
+#[derive(Parser, Debug)]
+#[command(version)]
+struct Cli {
+    #[command(subcommand)]
+    command: Commands,
+}
+
+#[derive(Args, Debug)]
+struct RegistryArgs {
+    /// The registry domain to use. Overrides configuration file(s).
+    #[arg(long = "registry", value_name = "DOMAIN")]
+    domain: Option<String>,
+}
+
+#[derive(Subcommand, Debug)]
+enum Commands {
+    /// Get a package.
+    Get(GetCommand),
+}
+
+#[derive(Args, Debug)]
+struct GetCommand {
+    /// Output path. If this ends with a '/', a filename based on the package
+    /// name, version, and format will be appended, e.g.
+    /// `[email protected]``.
+    #[arg(long, short, default_value = "./")]
+    output: PathBuf,
+
+    /// Output format. The default of "auto" detects the format based on the
+    /// output filename or package contents.
+    #[arg(long, value_enum, default_value = "auto")]
+    format: Format,
+
+    /// Overwrite any existing output file.
+    #[arg(long)]
+    overwrite: bool,
+
+    /// The package to get, specified as <namespace>:<name> plus optional
+    /// @<version>, e.g. "wasi:cli" or "wasi:[email protected]".
+    package_spec: PackageSpec,
+
+    #[command(flatten)]
+    registry: RegistryArgs,
+}
+
+#[derive(ValueEnum, Clone, Debug, PartialEq)]
+enum Format {
+    Auto,
+    Wasm,
+    Wit,
+}
+
+impl GetCommand {
+    pub async fn run(self) -> anyhow::Result<()> {
+        let PackageSpec { package, version } = self.package_spec;
+
+        let mut client = {
+            let mut config = ClientConfig::default();
+            config.set_default_registry("bytecodealliance.org");
+            if let Some(file_config) = ClientConfig::from_default_file()? {
+                config.merge_config(file_config);
+            }
+            if let Some(registry) = self.registry.domain {
+                let namespace = package.namespace().to_string();
+                tracing::debug!(namespace, registry, "overriding namespace registry");
+                config.set_namespace_registry(namespace, registry);
+            }
+            config.to_client()
+        };
+
+        let version = match version {
+            Some(ver) => ver,
+            None => {
+                println!("No version specified; fetching version list...");
+                let versions = client.list_all_versions(&package).await?;
+                tracing::trace!(?versions);
+                versions
+                    .into_iter()
+                    .filter_map(|vi| (!vi.yanked).then_some(vi.version))
+                    .max()
+                    .context("No releases found")?
+            }
+        };
+
+        println!("Getting {package}@{version}...");
+        let release = client
+            .get_release(&package, &version)
+            .await
+            .context("Failed to get release details")?;
+        tracing::debug!(?release);
+
+        let output_trailing_slash = self.output.as_os_str().to_string_lossy().ends_with('/');
+        let parent_dir = if output_trailing_slash {
+            self.output.as_path()
+        } else {
+            self.output
+                .parent()
+                .context("Failed to resolve output parent dir")?
+        };
+
+        let (tmp_file, tmp_path) =
+            tempfile::NamedTempFile::with_prefix_in(".wkg-get", parent_dir)?.into_parts();
+        tracing::debug!(?tmp_path);
+
+        let mut content_stream = client.stream_content(&package, &release).await?;
+
+        let mut file = tokio::fs::File::from_std(tmp_file);
+        while let Some(chunk) = content_stream.try_next().await? {
+            file.write_all(&chunk).await?;
+        }
+
+        let mut format = self.format;
+        if let (Format::Auto, Some(ext)) = (&format, self.output.extension()) {
+            tracing::debug!("Inferring output format from file extension {ext:?}");
+            format = match ext.to_string_lossy().as_ref() {
+                "wasm" => Format::Wasm,
+                "wit" => Format::Wit,
+                _ => {
+                    println!(
+                        "Couldn't infer output format from file name {:?}",
+                        self.output.file_name().unwrap_or_default()
+                    );
+                    Format::Auto
+                }
+            }
+        }
+
+        let wit = if format == Format::Wasm {
+            None
+        } else {
+            let mut file = file.into_std().await;
+            file.rewind()?;
+            match wit_component::decode_reader(&mut file) {
+                Ok(DecodedWasm::WitPackage(resolve, pkg)) => {
+                    tracing::debug!(?pkg, "decoded WIT package");
+                    Some(wit_component::WitPrinter::default().print(&resolve, pkg)?)
+                }
+                Ok(_) => None,
+                Err(err) => {
+                    tracing::debug!(?err);
+                    if format == Format::Wit {
+                        return Err(err);
+                    }
+                    println!("Failed to detect package content type: {err:#}");
+                    None
+                }
+            }
+        };
+
+        let output_path = if output_trailing_slash {
+            let ext = if wit.is_some() { "wit" } else { "wasm" };
+            self.output.join(format!(
+                "{namespace}_{name}@{version}.{ext}",
+                namespace = package.namespace(),
+                name = package.name(),
+            ))
+        } else {
+            self.output
+        };
+        ensure!(
+            self.overwrite || !output_path.exists(),
+            "{output_path:?} already exists; you can use '--overwrite' to overwrite it"
+        );
+
+        if let Some(wit) = wit {
+            std::fs::write(&output_path, wit)
+                .with_context(|| format!("Failed to write WIT to {output_path:?}"))?
+        } else {
+            tmp_path
+                .persist(&output_path)
+                .with_context(|| format!("Failed to persist WASM to {output_path:?}"))?
+        }
+        println!("Wrote '{}'", output_path.display());
+
+        Ok(())
+    }
+}
+
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+    tracing_subscriber::fmt()
+        .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
+        .init();
+
+    let cli = Cli::parse();
+    tracing::debug!(?cli);
+
+    match cli.command {
+        Commands::Get(cmd) => cmd.run().await,
+    }
 }

crates/wkg/src/package_spec.rs

@@ -0,0 +1,25 @@
+use std::str::FromStr;
+
+use wasm_pkg_loader::{PackageRef, Version};
+
+// TODO: move to some library crate
+#[derive(Clone, Debug)]
+pub struct PackageSpec {
+    pub package: PackageRef,
+    pub version: Option<Version>,
+}
+
+impl FromStr for PackageSpec {
+    type Err = anyhow::Error;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let (package, version) = s
+            .split_once('@')
+            .map(|(pkg, ver)| (pkg, Some(ver)))
+            .unwrap_or((s, None));
+        Ok(Self {
+            package: package.parse()?,
+            version: version.map(|ver| ver.parse()).transpose()?,
+        })
+    }
+}