Repository: candlecorp/wick (language: rust)

Issue:
Panic on invalid OCI reference
**Describe the bug**
Doing `wick reg pull` or other ops that require pulling from the registry will panic on index out-of-bounds for common reference mistakes

**To Reproduce**
Run `wick reg pull candle.registry.dev/anything:1`

**Expected behavior**
An nice error telling me I don't have an org in my reference.

**Actual behavior**
An unhelpful panic message.

**Environment:**
 - OS: Mac
 - Architecture: m1
 - Build Version: head



Discussion / hints:

Fix patch:
diff --git a/crates/wick/wick-oci-utils/Cargo.toml b/crates/wick/wick-oci-utils/Cargo.toml
index f99e978af..8bd136685 100644
--- a/crates/wick/wick-oci-utils/Cargo.toml
+++ b/crates/wick/wick-oci-utils/Cargo.toml
@@ -35,3 +35,4 @@ clap = { workspace = true, features = ["derive", "env"] }
 async-recursion = { workspace = true }
 env_logger = { workspace = true }
 anyhow = { workspace = true }
+rstest = { workspace = true }
diff --git a/crates/wick/wick-oci-utils/src/error.rs b/crates/wick/wick-oci-utils/src/error.rs
index 6c72c82be..064bda3be 100644
--- a/crates/wick/wick-oci-utils/src/error.rs
+++ b/crates/wick/wick-oci-utils/src/error.rs
@@ -18,6 +18,8 @@ pub enum OciError {
   #[error("Invalid manifest found at {}. Try deleting your cache directory.",.0.display())]
   InvalidManifest(PathBuf),
 
+  #[error("Reference '{0}' did not contain a tag or digest")]
+  NoTagOrDigest(String),
   /// Error thrown when attempting to fetch an image with :latest when forbidden.
   #[error("Configuration disallows fetching artifacts with the :latest tag ({0})")]
   LatestDisallowed(String),
@@ -64,7 +66,7 @@ pub enum OciError {
 
   /// Failed to parse image reference location
   #[error("Failed to parse the image reference: {0}")]
-  InvalidReference(String),
+  InvalidReferenceFormat(String),
 
   /// Failed to push package
   #[error("Failed to pull the package: {0}")]
diff --git a/crates/wick/wick-oci-utils/src/utils.rs b/crates/wick/wick-oci-utils/src/utils.rs
index 861201986..8ec86353c 100644
--- a/crates/wick/wick-oci-utils/src/utils.rs
+++ b/crates/wick/wick-oci-utils/src/utils.rs
@@ -8,10 +8,56 @@ use tokio::fs;
 
 use crate::Error;
 
+// Lovingly borrowed from oci-distribution who doesn't export it
+pub static REFERENCE_REGEXP: &str = r"^((?:(?:[a-zA-Z0-9]|[a-zA-Z0-9][a-zA-Z0-9-]*[a-zA-Z0-9])(?:(?:\.(?:[a-zA-Z0-9]|[a-zA-Z0-9][a-zA-Z0-9-]*[a-zA-Z0-9]))+)?(?::[0-9]+)?/)?[a-z0-9]+(?:(?:(?:[._]|__|[-]*)[a-z0-9]+)+)?(?:(?:/[a-z0-9]+(?:(?:(?:[._]|__|[-]*)[a-z0-9]+)+)?)+)?)(?::([\w][\w.-]{0,127}))?(?:@([A-Za-z][A-Za-z0-9]*(?:[-_+.][A-Za-z][A-Za-z0-9]*)*[:][[:xdigit:]]{32,}))?$";
+
+pub const DEFAULT_REGISTRY: &str = "registry.candle.dev";
+
+static RE: Lazy<Regex> = Lazy::new(|| {
+  regex::RegexBuilder::new(REFERENCE_REGEXP)
+    .size_limit(10 * (1 << 21))
+    .build()
+    .unwrap()
+});
+
+fn normalize_reference(reference: &str) -> Result<String, Error> {
+  let captures = RE
+    .captures(reference)
+    .ok_or(Error::InvalidReferenceFormat(reference.to_owned()))?;
+  let name = &captures[1];
+  let tag = captures.get(2).map(|m| m.as_str().to_owned());
+  let digest = captures.get(3).map(|m| m.as_str().to_owned());
+  if tag.is_none() && digest.is_none() {
+    return Err(Error::NoTagOrDigest(reference.to_owned()));
+  }
+
+  let (registry, repository) = split_domain(name);
+  Ok(format!(
+    "{}/{}:{}",
+    registry,
+    repository,
+    tag.unwrap_or_else(|| digest.unwrap())
+  ))
+}
+
+// Also borrowed from oci-distribution who borrowed it from the go docker implementation.
+fn split_domain(name: &str) -> (String, String) {
+  match name.split_once('/') {
+    None => (DEFAULT_REGISTRY.to_owned(), name.to_owned()),
+    Some((left, right)) => {
+      if !(left.contains('.') || left.contains(':')) && left != "localhost" {
+        (DEFAULT_REGISTRY.to_owned(), name.to_owned())
+      } else {
+        (left.to_owned(), right.to_owned())
+      }
+    }
+  }
+}
+
 /// Parse a `&str` as a Reference.
 pub fn parse_reference(reference: &str) -> Result<Reference, Error> {
-  oci_distribution::Reference::from_str(reference)
-    .map_err(|e| Error::OCIParseError(reference.to_owned(), e.to_string()))
+  let reference = normalize_reference(reference)?;
+  oci_distribution::Reference::from_str(&reference).map_err(|e| Error::OCIParseError(reference, e.to_string()))
 }
 
 /// Parse a `&str` as a Reference and return the protocol to use.
@@ -19,8 +65,8 @@ pub fn parse_reference_and_protocol(
   reference: &str,
   allowed_insecure: &[String],
 ) -> Result<(Reference, oci_distribution::client::ClientProtocol), Error> {
-  let reference =
-    Reference::from_str(reference).map_err(|e| Error::OCIParseError(reference.to_owned(), e.to_string()))?;
+  let reference = parse_reference(reference)?;
+
   let insecure = allowed_insecure.contains(&reference.registry().to_owned());
   Ok((
     reference,
@@ -33,24 +79,18 @@ pub fn parse_reference_and_protocol(
 }
 
 pub(crate) fn get_cache_directory(input: &str, basedir: impl AsRef<Path>) -> Result<PathBuf, Error> {
-  let image_ref_result = Reference::from_str(input);
-  let image_ref = match image_ref_result {
-    Ok(image_ref) => image_ref,
-    Err(_) => {
-      return Err(Error::InvalidReference(input.to_owned()));
-    }
-  };
-
-  let registry = image_ref.registry().split(':').collect::<Vec<&str>>()[0];
-  let org = image_ref.repository().split('/').collect::<Vec<&str>>()[0];
-  let repo = image_ref.repository().split('/').collect::<Vec<&str>>()[1];
-  let version = image_ref.tag().ok_or(Error::NoName)?;
+  let image_ref = parse_reference(input)?;
 
-  let parts = vec![registry, org, repo, version];
+  let registry = image_ref
+    .registry()
+    .split_once(':')
+    .map_or(image_ref.registry(), |(reg, _port)| reg);
+  let (org, repo) = image_ref.repository().split_once('/').ok_or(Error::OCIParseError(
+    input.to_owned(),
+    "repository was not in org/repo format".to_owned(),
+  ))?;
 
-  if parts.len() != 4 {
-    return Err(Error::InvalidReference(input.to_owned()));
-  }
+  let version = image_ref.tag().ok_or(Error::NoName)?;
 
   // Create the wick_components directory if it doesn't exist
   let target_dir = basedir.as_ref().join(registry).join(org).join(repo).join(version);
@@ -81,19 +121,34 @@ mod tests {
 
   use super::*;
 
-  #[test]
-  fn test_directory_structure() {
-    let input = "localhost:5555/test/integration:0.0.3";
-
-    let expected_dir = Path::new("localhost/test/integration/0.0.3");
-    let result = get_cache_directory(input, "").unwrap();
-    assert_eq!(result, expected_dir);
-
-    let input = "example.com/myorg/myrepo:1.0.0";
-    let expected_dir = Path::new("/foo/bar/example.com/myorg/myrepo/1.0.0");
-    let result = get_cache_directory(input, "/foo/bar").unwrap();
+  #[rstest::rstest]
+  #[case("localhost:5555/test/integration:0.0.3", "", "localhost/test/integration/0.0.3")]
+  #[case(
+    "example.com/myorg/myrepo:1.0.0",
+    "/foo/bar",
+    "/foo/bar/example.com/myorg/myrepo/1.0.0"
+  )]
+  #[case("org/myrepo:1.0.1", "", "registry.candle.dev/org/myrepo/1.0.1")]
+  fn directory_structure_positive(#[case] input: &str, #[case] basedir: &str, #[case] expected: &str) {
+    let expected_dir = Path::new(expected);
+    let result = get_cache_directory(input, basedir).unwrap();
     assert_eq!(result, expected_dir);
   }
+  #[rstest::rstest]
+  #[case("example.com/myrepo:1.0.0")]
+  #[case("example.com/org/myrepo")]
+  #[case("example.com/myrepo")]
+  #[case("example.com:5000/myrepo:1.0.0")]
+  #[case("example.com:5000/org/myrepo")]
+  #[case("example.com:5000/myrepo")]
+  #[case("myrepo:1.0.0")]
+  #[case("org/myrepo")]
+  #[case("myrepo")]
+  fn directory_structure_negative(#[case] input: &str) {
+    let result = get_cache_directory(input, "");
+    println!("{:?}", result);
+    assert!(result.is_err());
+  }
 
   #[test]
   fn test_good_wickref() -> Result<()> {


Test patch: