1 files changed, 182 insertions, 0 deletions
diff --git a/taskcluster/docker/image_builder/build-image/src/main.rs b/taskcluster/docker/image_builder/build-image/src/main.rs
new file mode 100644
index 0000000000..997617c84e
--- /dev/null
+++ b/taskcluster/docker/image_builder/build-image/src/main.rs
@@ -0,0 +1,182 @@
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+#![forbid(unsafe_code)]
+
+use std::collections::HashMap;
+use std::path::Path;
+use std::process::Command;
+
+use anyhow::{ensure, Context, Result};
+use fs_extra::dir::{move_dir, CopyOptions};
+use serde::Deserialize;
+
+mod config;
+mod taskcluster;
+
+use config::Config;
+
+fn log_step(msg: &str) {
+    println!("[build-image] {}", msg);
+}
+
+fn read_image_digest(path: &str) -> Result<String> {
+    let output = Command::new("/kaniko/skopeo")
+        .arg("inspect")
+        .arg(format!("docker-archive:{}", path))
+        .stdout(std::process::Stdio::piped())
+        .spawn()?
+        .wait_with_output()?;
+    ensure!(output.status.success(), "Could not inspect parent image.");
+
+    #[derive(Deserialize, Debug)]
+    #[serde(rename_all = "PascalCase")]
+    struct ImageInfo {
+        #[serde(skip_serializing_if = "Option::is_none")]
+        name: Option<String>,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        tag: Option<String>,
+        digest: String,
+        // ...
+    }
+
+    let image_info: ImageInfo = serde_json::from_slice(&output.stdout)
+        .with_context(|| format!("Could parse image info from {:?}", path))?;
+    Ok(image_info.digest)
+}
+
+fn download_parent_image(
+    cluster: &taskcluster::TaskCluster,
+    task_id: &str,
+    dest: &str,
+) -> Result<String> {
+    zstd::stream::copy_decode(
+        cluster.stream_artifact(&task_id, "public/image.tar.zst")?,
+        std::fs::File::create(dest)?,
+    )
+    .context("Could not download parent image.")?;
+
+    read_image_digest(dest)
+}
+
+fn build_image(
+    context_path: &str,
+    dest: &str,
+    debug: bool,
+    build_args: HashMap<String, String>,
+) -> Result<()> {
+    let mut command = Command::new("/kaniko/executor");
+    command
+        .stderr(std::process::Stdio::inherit())
+        .args(&["--context", &format!("tar://{}", context_path)])
+        .args(&["--destination", "image"])
+        .args(&["--dockerfile", "Dockerfile"])
+        .arg("--no-push")
+        .args(&["--cache-dir", "/workspace/cache"])
+        .arg("--single-snapshot")
+        // FIXME: Generating reproducible layers currently causes OOM.
+        // .arg("--reproducible")
+        .arg("--whitelist-var-run=false")
+        .args(&["--tarPath", dest]);
+    if debug {
+        command.args(&["-v", "debug"]);
+    }
+    for (key, value) in build_args {
+        command.args(&["--build-arg", &format!("{}={}", key, value)]);
+    }
+    let status = command.status()?;
+    ensure!(status.success(), "Could not build image.");
+    Ok(())
+}
+
+fn repack_image(source: &str, dest: &str, image_name: &str) -> Result<()> {
+    let status = Command::new("/kaniko/skopeo")
+        .arg("copy")
+        .arg(format!("docker-archive:{}", source))
+        .arg(format!("docker-archive:{}:{}", dest, image_name))
+        .stderr(std::process::Stdio::inherit())
+        .status()?;
+    ensure!(status.success(), "Could repack image.");
+    Ok(())
+}
+
+fn main() -> Result<()> {
+    // Kaniko expects everything to be in /kaniko, so if not running from there, move
+    // everything there.
+    if let Some(path) = std::env::current_exe()?.parent() {
+        if path != Path::new("/kaniko") {
+            let mut options = CopyOptions::new();
+            options.copy_inside = true;
+            move_dir(path, "/kaniko", &options)?;
+        }
+    }
+
+    let config = Config::from_env().context("Could not parse environment variables.")?;
+
+    let cluster = taskcluster::TaskCluster::from_env()?;
+
+    let mut build_args = config.docker_build_args;
+
+    build_args.insert("TASKCLUSTER_ROOT_URL".into(), cluster.root_url());
+
+    log_step("Downloading context.");
+
+    std::io::copy(
+        &mut cluster.stream_artifact(&config.context_task_id, &config.context_path)?,
+        &mut std::fs::File::create("/workspace/context.tar.gz")?,
+    )
+    .context("Could not download image context.")?;
+
+    if let Some(parent_task_id) = config.parent_task_id {
+        log_step("Downloading image.");
+        let digest = download_parent_image(&cluster, &parent_task_id, "/workspace/parent.tar")?;
+
+        log_step(&format!("Parent image digest {}", &digest));
+        std::fs::create_dir_all("/workspace/cache")?;
+        std::fs::rename(
+            "/workspace/parent.tar",
+            format!("/workspace/cache/{}", digest),
+        )?;
+
+        build_args.insert(
+            "DOCKER_IMAGE_PARENT".into(),
+            format!("parent:latest@{}", digest),
+        );
+    }
+
+    log_step("Building image.");
+    build_image(
+        "/workspace/context.tar.gz",
+        "/workspace/image-pre.tar",
+        config.debug,
+        build_args,
+    )?;
+    log_step("Repacking image.");
+    repack_image(
+        "/workspace/image-pre.tar",
+        "/workspace/image.tar",
+        &config.image_name,
+    )?;
+
+    log_step("Compressing image.");
+    compress_file(
+        "/workspace/image.tar",
+        "/workspace/image.tar.zst",
+        config.docker_image_zstd_level,
+    )?;
+
+    Ok(())
+}
+
+fn compress_file(
+    source: impl AsRef<std::path::Path>,
+    dest: impl AsRef<std::path::Path>,
+    zstd_level: i32,
+) -> Result<()> {
+    Ok(zstd::stream::copy_encode(
+        std::fs::File::open(source)?,
+        std::fs::File::create(dest)?,
+        zstd_level,
+    )?)
+}