From 04e9e6b317e977d908bc596b1ee84bff82b297d4 Mon Sep 17 00:00:00 2001
From: Gus Caplan <me@gus.host>
Date: Tue, 24 Nov 2020 13:57:08 -0600
Subject: [PATCH] add SandboxBuilder image override

---
 src/cmd/mod.rs     |  8 +++++++
 src/cmd/sandbox.rs | 58 +++++++++++++++++++++++++++++++++++++++++++---
 src/crates/mod.rs  |  3 ++-
 src/workspace.rs   |  2 +-
 4 files changed, 66 insertions(+), 5 deletions(-)

diff --git a/src/cmd/mod.rs b/src/cmd/mod.rs
index d2af74e..051f42e 100644
--- a/src/cmd/mod.rs
+++ b/src/cmd/mod.rs
@@ -100,6 +100,14 @@ pub enum CommandError {
     #[error("invalid output of `docker inspect`: {0}")]
     InvalidDockerInspectOutput(#[source] serde_json::Error),
 
+    /// The data received from the `docker manifest inspect` command is not valid.
+    #[error("invalid output from `docker manifest inspect`: {0}")]
+    InvalidDockerManifestInspectOutput(#[source] serde_json::Error),
+
+    /// The remote image is larger than the specified size limit.
+    #[error("sandbox image is too large: {0}")]
+    SandboxImageTooLarge(usize),
+
     /// An I/O error occured while executing the command.
     #[error(transparent)]
     IO(#[from] std::io::Error),
diff --git a/src/cmd/sandbox.rs b/src/cmd/sandbox.rs
index 1f5fd33..367b3a4 100644
--- a/src/cmd/sandbox.rs
+++ b/src/cmd/sandbox.rs
@@ -12,6 +12,22 @@ pub struct SandboxImage {
     name: String,
 }
 
+#[derive(serde::Deserialize)]
+struct DockerManifest {
+    config: DockerManifestConfig,
+    layers: Vec<DockerManifestLayer>,
+}
+
+#[derive(serde::Deserialize)]
+struct DockerManifestConfig {
+    digest: String,
+}
+
+#[derive(serde::Deserialize)]
+struct DockerManifestLayer {
+    size: usize,
+}
+
 impl SandboxImage {
     /// Load a local image present in the host machine.
     ///
@@ -27,11 +43,33 @@ impl SandboxImage {
     ///
     /// This will access the network to download the image from the registry. If pulling fails an
     /// error will be returned instead.
-    pub fn remote(name: &str) -> Result<Self, CommandError> {
+    pub fn remote(name: &str, size_limit: Option<usize>) -> Result<Self, CommandError> {
         let mut image = SandboxImage { name: name.into() };
+        let digest = if let Some(size_limit) = size_limit {
+            let out = Command::new_workspaceless("docker")
+                .args(&["manifest", "inspect", name])
+                .run_capture()?
+                .stdout_lines()
+                .join("\n");
+            let m: DockerManifest = serde_json::from_str(&out)
+                .map_err(CommandError::InvalidDockerManifestInspectOutput)?;
+            let size = m.layers.iter().fold(0, |acc, l| acc + l.size);
+            if size > size_limit {
+                return Err(CommandError::SandboxImageTooLarge(size));
+            }
+            Some(m.config.digest)
+        } else {
+            None
+        };
         info!("pulling image {} from Docker Hub", name);
         Command::new_workspaceless("docker")
-            .args(&["pull", &name])
+            .args(&[
+                "pull",
+                &digest.map_or(name.to_string(), |digest| {
+                    let name = name.split(':').next().unwrap();
+                    format!("{}@{}", name, digest)
+                }),
+            ])
             .run()
             .map_err(|e| CommandError::SandboxImagePullFailed(Box::new(e)))?;
         if let Some(name_with_hash) = image.get_name_with_hash() {
@@ -146,6 +184,7 @@ pub struct SandboxBuilder {
     user: Option<String>,
     cmd: Vec<String>,
     enable_networking: bool,
+    image: Option<String>,
 }
 
 impl SandboxBuilder {
@@ -160,6 +199,7 @@ impl SandboxBuilder {
             user: None,
             cmd: Vec::new(),
             enable_networking: true,
+            image: None,
         }
     }
 
@@ -203,6 +243,14 @@ impl SandboxBuilder {
         self
     }
 
+    /// Override the image used for this sandbox.
+    ///
+    /// By default rustwide will use the image configured with [`WorkspaceBuilder::sandbox_image`].
+    pub fn image(mut self, image: SandboxImage) -> Self {
+        self.image = Some(image.name);
+        self
+    }
+
     pub(super) fn env<S1: Into<String>, S2: Into<String>>(mut self, key: S1, value: S2) -> Self {
         self.env.push((key.into(), value.into()));
         self
@@ -274,7 +322,11 @@ impl SandboxBuilder {
             args.push("--isolation=process".into());
         }
 
-        args.push(workspace.sandbox_image().name.clone());
+        if let Some(image) = self.image {
+            args.push(image);
+        } else {
+            args.push(workspace.sandbox_image().name.clone());
+        }
 
         for arg in self.cmd {
             args.push(arg);
diff --git a/src/crates/mod.rs b/src/crates/mod.rs
index 6e330c5..d9ce4a5 100644
--- a/src/crates/mod.rs
+++ b/src/crates/mod.rs
@@ -73,7 +73,8 @@ impl Crate {
         }
     }
 
-    pub(crate) fn copy_source_to(&self, workspace: &Workspace, dest: &Path) -> Result<(), Error> {
+    /// Copy this crate's source to `dest`. If `dest` already exists, it will be replaced.
+    pub fn copy_source_to(&self, workspace: &Workspace, dest: &Path) -> Result<(), Error> {
         if dest.exists() {
             info!(
                 "crate source directory {} already exists, cleaning it up",
diff --git a/src/workspace.rs b/src/workspace.rs
index daa9da1..2816610 100644
--- a/src/workspace.rs
+++ b/src/workspace.rs
@@ -146,7 +146,7 @@ impl WorkspaceBuilder {
             let sandbox_image = if let Some(img) = self.sandbox_image {
                 img
             } else {
-                SandboxImage::remote(DEFAULT_SANDBOX_IMAGE)?
+                SandboxImage::remote(DEFAULT_SANDBOX_IMAGE, None)?
             };
 
             let mut agent = attohttpc::Session::new();