From c44452648d5212b47db57e4ecb18734071df65b7 Mon Sep 17 00:00:00 2001
From: strNophix <nvdpoel01@gmail.com>
Date: Sat, 25 Jun 2022 20:56:28 +0200
Subject: [PATCH] Fetch an image from the Docker Registry

---
 src/main.rs     | 17 +++++----
 src/registry.rs | 92 +++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 102 insertions(+), 7 deletions(-)
 create mode 100644 src/registry.rs

diff --git a/src/main.rs b/src/main.rs
index 647aaa5..a7a13e6 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,18 +1,24 @@
 use std::{self, ffi::CString, path};
 
+mod registry;
+
 // Usage: your_docker.sh run <image> <command> <arg1> <arg2> ...
 fn main() {
     let args: Vec<_> = std::env::args().collect();
+    let image_name = &args[2];
     let command = &args[3];
     let command_args = &args[4..];
 
-    let base = "/tmp/mydocker";
-    let base_path = path::Path::new(base);
+    let base_path = std::env::temp_dir().join("mydocker");
 
     // prevent cryptic "no such file or directory" error inside chroot
     std::fs::create_dir_all(base_path.join("dev")).unwrap();
     std::fs::File::create(base_path.join("dev/null")).unwrap();
 
+    let image = registry::ImageIdentifier::from_string(image_name);
+    let mut reg = registry::Registry::default();
+    reg.pull(&image, base_path.to_str().unwrap());
+
     // copy over binary into chroot
     let command_path = path::Path::new(command).strip_prefix("/").unwrap();
     std::fs::create_dir_all(base_path.join(command_path.parent().unwrap()))
@@ -21,7 +27,7 @@ fn main() {
         .expect("Failed copying executed binary to chroot directory");
 
     // create and change into chroot
-    let cbase_path = CString::new(base.to_owned()).unwrap();
+    let cbase_path = CString::new(base_path.to_str().unwrap().to_owned()).unwrap();
     unsafe {
         libc::chroot(cbase_path.as_ptr());
     }
@@ -43,8 +49,5 @@ fn main() {
     let std_err = std::str::from_utf8(&output.stderr).unwrap();
     eprint!("{}", std_err);
 
-    match output.status.code() {
-        Some(code) => std::process::exit(code),
-        None => std::process::exit(1),
-    }
+    std::process::exit(output.status.code().unwrap_or(1));
 }
diff --git a/src/registry.rs b/src/registry.rs
new file mode 100644
index 0000000..83b7cea
--- /dev/null
+++ b/src/registry.rs
@@ -0,0 +1,92 @@
+// TODO: Enable derive feature
+
+pub struct Registry {
+    http_client: reqwest::blocking::Client,
+}
+
+impl Default for Registry {
+    fn default() -> Self {
+        return Registry {
+            http_client: reqwest::blocking::Client::new(),
+        };
+    }
+}
+
+impl Registry {
+    pub fn pull(&mut self, image: &ImageIdentifier, destination: &str) {
+        // Perform the little auth dance
+        let auth_url = format!("https://auth.docker.io/token?service=registry.docker.io&scope=repository%3A{}%2F{}%3Apull", image.author, image.name);
+        let auth: serde_json::Value = self
+            .http_client
+            .get(&auth_url)
+            .send()
+            .unwrap()
+            .json()
+            .unwrap();
+        let access_token = auth["token"].as_str().unwrap();
+
+        // Download the image manifest
+        let auth_header = format!("Bearer {}", access_token);
+        let image_url = format!(
+            "https://registry.hub.docker.com/v2/{}/{}/manifests/{}",
+            image.author, image.name, image.tag
+        );
+        let image_manifest: serde_json::Value = self
+            .http_client
+            .get(&image_url)
+            .header(reqwest::header::AUTHORIZATION, auth_header.to_owned())
+            .send()
+            .unwrap()
+            .json()
+            .unwrap();
+
+        // Download the image layers and extracts them
+        let temp_path = std::env::temp_dir();
+        for layer in image_manifest["fsLayers"].as_array().unwrap() {
+            let digest = layer["blobSum"].as_str().unwrap();
+            let blob_url = format!(
+                "https://registry.hub.docker.com/v2/{}/{}/blobs/{}",
+                image.author, image.name, digest
+            );
+            let blob = self
+                .http_client
+                .get(&blob_url)
+                .header(reqwest::header::AUTHORIZATION, auth_header.to_owned())
+                .send()
+                .unwrap()
+                .bytes()
+                .unwrap();
+
+            let layer_path = temp_path.join(digest);
+            std::fs::write(layer_path.to_owned(), blob).unwrap();
+
+            // TODO: handle exit code of untar
+            std::process::Command::new("tar")
+                .args(["-xf", layer_path.to_str().unwrap(), "-C", destination])
+                .output()
+                .unwrap();
+
+            std::fs::remove_file(layer_path).unwrap();
+        }
+    }
+}
+
+pub struct ImageIdentifier {
+    author: String,
+    name: String,
+    tag: String,
+}
+
+impl ImageIdentifier {
+    pub fn from_string(image: &String) -> Self {
+        let mut iter = image.rsplitn(2, ':');
+        let tag = iter.next().unwrap_or("latest").to_string();
+        let mut loc_iter = iter.next().unwrap().split('/').rev();
+        let name = loc_iter
+            .next()
+            .expect("No image name was supplied")
+            .to_string();
+        let author = loc_iter.next().unwrap_or("library").to_string();
+        return ImageIdentifier { author, name, tag };
+    }
+}