From 6511f2f0fae936572626890efd486163f1931672 Mon Sep 17 00:00:00 2001
From: Alex Zenla <alex@edera.dev>
Date: Sat, 24 Feb 2024 22:05:06 +0000
Subject: [PATCH] controller: implement enhanced image name format support

---
 controller/src/image/fetch.rs |  7 ++++++-
 controller/src/image/name.rs  | 38 +++++++++++++++++++++++++++--------
 scripts/krata-debug-common.sh | 29 ++++++++++++++++++++++++++
 scripts/kratactl-debug.sh     | 14 ++++---------
 scripts/kratanet-debug.sh     | 12 ++++-------
 5 files changed, 73 insertions(+), 27 deletions(-)
 create mode 100644 scripts/krata-debug-common.sh

diff --git a/controller/src/image/fetch.rs b/controller/src/image/fetch.rs
index dadce06..fcd9b99 100644
--- a/controller/src/image/fetch.rs
+++ b/controller/src/image/fetch.rs
@@ -6,6 +6,9 @@ use std::ops::DerefMut;
 use ureq::{Agent, Request, Response};
 use url::Url;
 
+const MANIFEST_PICKER_PLATFORM: Os = Os::Linux;
+const MANIFEST_PICKER_ARCHITECTURE: Arch = Arch::Amd64;
+
 pub struct RegistryClient {
     agent: Agent,
     url: Url,
@@ -86,7 +89,9 @@ impl RegistryClient {
     fn pick_manifest(&mut self, index: ImageIndex) -> Option<Descriptor> {
         for item in index.manifests() {
             if let Some(platform) = item.platform() {
-                if *platform.os() == Os::Linux && *platform.architecture() == Arch::Amd64 {
+                if *platform.os() == MANIFEST_PICKER_PLATFORM
+                    && *platform.architecture() == MANIFEST_PICKER_ARCHITECTURE
+                {
                     return Some(item.clone());
                 }
             }
diff --git a/controller/src/image/name.rs b/controller/src/image/name.rs
index dd1cc83..a6b1d2c 100644
--- a/controller/src/image/name.rs
+++ b/controller/src/image/name.rs
@@ -2,6 +2,9 @@ use anyhow::Result;
 use std::fmt;
 use url::Url;
 
+const DOCKER_HUB_MIRROR: &str = "mirror.gcr.io";
+const DEFAULT_IMAGE_TAG: &str = "latest";
+
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct ImageName {
     pub hostname: String,
@@ -33,20 +36,39 @@ impl Default for ImageName {
 
 impl ImageName {
     pub fn parse(name: &str) -> Result<Self> {
-        let (hostname, name) = name
+        let full_name = name.to_string();
+        let name = full_name.clone();
+        let (mut hostname, mut name) = name
             .split_once('/')
-            .unwrap_or(("registry-1.docker.io", name));
-        let (hostname, port) = if let Some((hostname, port)) = hostname.split_once(':') {
-            (hostname, Some(str::parse(port)?))
+            .map(|x| (x.0.to_string(), x.1.to_string()))
+            .unwrap_or_else(|| (DOCKER_HUB_MIRROR.to_string(), format!("library/{}", name)));
+
+        // heuristic to find any docker hub image formats
+        // that may be in the hostname format. for example:
+        // abc/xyz:latest will trigger this if check, but abc.io/xyz:latest will not,
+        // and neither will abc/hello/xyz:latest
+        if !hostname.contains('.') && full_name.chars().filter(|x| *x == '/').count() == 1 {
+            name = format!("{}/{}", hostname, name);
+            hostname = DOCKER_HUB_MIRROR.to_string();
+        }
+
+        let (hostname, port) = if let Some((hostname, port)) = hostname
+            .split_once(':')
+            .map(|x| (x.0.to_string(), x.1.to_string()))
+        {
+            (hostname, Some(str::parse(&port)?))
         } else {
             (hostname, None)
         };
-        let (name, reference) = name.split_once(':').unwrap_or((name, "latest"));
+        let (name, reference) = name
+            .split_once(':')
+            .map(|x| (x.0.to_string(), x.1.to_string()))
+            .unwrap_or((name.to_string(), DEFAULT_IMAGE_TAG.to_string()));
         Ok(ImageName {
-            hostname: hostname.to_string(),
+            hostname,
             port,
-            name: name.to_string(),
-            reference: reference.to_string(),
+            name,
+            reference,
         })
     }
 
diff --git a/scripts/krata-debug-common.sh b/scripts/krata-debug-common.sh
new file mode 100644
index 0000000..8a92156
--- /dev/null
+++ b/scripts/krata-debug-common.sh
@@ -0,0 +1,29 @@
+#!/bin/sh
+set -e
+
+REAL_SCRIPT="$(realpath "${0}")"
+cd "$(dirname "${REAL_SCRIPT}")/.."
+
+if [ -z "${RUST_LOG}" ]
+then
+  RUST_LOG="INFO"
+fi
+
+CARGO_BUILD_FLAGS=""
+
+if [ "${KRATA_BUILD_QUIET}" = "1" ]
+then
+  CARGO_BUILD_FLAGS="-q"
+fi
+
+build_and_run() {
+  EXE_TARGET="${1}"
+  shift
+  if [ "${KRATA_BUILD_INITRD}" = "1" ]
+  then
+    ./initrd/build.sh -q
+    sudo cp "target/initrd/initrd" "/var/lib/krata/default/initrd"
+  fi
+  cargo build ${CARGO_BUILD_FLAGS} --target x86_64-unknown-linux-gnu --bin "${EXE_TARGET}"
+  exec sudo RUST_LOG="${RUST_LOG}" "target/x86_64-unknown-linux-gnu/debug/${EXE_TARGET}" "${@}"
+}
diff --git a/scripts/kratactl-debug.sh b/scripts/kratactl-debug.sh
index 1a22472..860cd8a 100755
--- a/scripts/kratactl-debug.sh
+++ b/scripts/kratactl-debug.sh
@@ -1,14 +1,8 @@
 #!/bin/sh
 set -e
 
-if [ -z "${RUST_LOG}" ]
-then
-  RUST_LOG="INFO"
-fi
-
 REAL_SCRIPT="$(realpath "${0}")"
-cd "$(dirname "${REAL_SCRIPT}")/.."
-./initrd/build.sh -q
-sudo cp "target/initrd/initrd" "/var/lib/krata/default/initrd"
-cargo build --target x86_64-unknown-linux-gnu --bin kratactl
-exec sudo RUST_LOG="${RUST_LOG}" target/x86_64-unknown-linux-gnu/debug/kratactl "${@}"
+# shellcheck source-path=krata-debug-common.sh
+. "$(dirname "${REAL_SCRIPT}")/krata-debug-common.sh"
+
+KRATA_BUILD_INITRD=1 build_and_run kratactl "${@}"
diff --git a/scripts/kratanet-debug.sh b/scripts/kratanet-debug.sh
index 82c14eb..90b6986 100755
--- a/scripts/kratanet-debug.sh
+++ b/scripts/kratanet-debug.sh
@@ -1,12 +1,8 @@
 #!/bin/sh
 set -e
 
-if [ -z "${RUST_LOG}" ]
-then
-  RUST_LOG="INFO"
-fi
-
 REAL_SCRIPT="$(realpath "${0}")"
-cd "$(dirname "${REAL_SCRIPT}")/.."
-cargo build --target x86_64-unknown-linux-gnu --bin kratanet
-exec sudo RUST_LOG="${RUST_LOG}" target/x86_64-unknown-linux-gnu/debug/kratanet "${@}"
+# shellcheck source-path=krata-debug-common.sh
+. "$(dirname "${REAL_SCRIPT}")/krata-debug-common.sh"
+
+build_and_run kratanet "${@}"