dockerTools.buildLayeredImage: support fromImage

It is now possible to pass a `fromImage` to `buildLayeredImage` and
`streamLayeredImage`, similar to what `buildImage` currently supports.

This will prepend the layers of the given base image to the resulting
image, while ensuring that at most `maxLayers` are used. It will also
ensure that environment variables from the base image are propagated
to the final image.

3 files changed, 178 insertions, 30 deletions

diff --git a/pkgs/build-support/docker/default.nix b/pkgs/build-support/docker/default.nix
index fec289f0ff1..a73737cb123 100644
--- a/pkgs/build-support/docker/default.nix
+++ b/pkgs/build-support/docker/default.nix
@@ -729,6 +729,8 @@ rec {
     name,
     # Image tag, the Nix's output hash will be used if null
     tag ? null,
+    # Parent image, to append to.
+    fromImage ? null,
     # Files to put on the image (a nix store path or list of paths).
     contents ? [],
     # Docker config; e.g. what command to run on the container.
@@ -791,7 +793,7 @@ rec {
       unnecessaryDrvs = [ baseJson overallClosure ];
 
       conf = runCommand "${baseName}-conf.json" {
-        inherit maxLayers created;
+        inherit fromImage maxLayers created;
         imageName = lib.toLower name;
         passthru.imageTag =
           if tag != null
@@ -821,6 +823,27 @@ rec {
                          unnecessaryDrvs}
         }
 
+        # Compute the number of layers that are already used by a potential
+        # 'fromImage' as well as the customization layer. Ensure that there is
+        # still at least one layer available to store the image contents.
+        usedLayers=0
+
+        # subtract number of base image layers
+        if [[ -n "$fromImage" ]]; then
+          (( usedLayers += $(tar -xOf "$fromImage" manifest.json | jq '.[0].Layers | length') ))
+        fi
+
+        # one layer will be taken up by the customisation layer
+        (( usedLayers += 1 ))
+
+        if ! (( $usedLayers < $maxLayers )); then
+          echo >&2 "Error: usedLayers $usedLayers layers to store 'fromImage' and" \
+                    "'extraCommands', but only maxLayers=$maxLayers were" \
+                    "allowed. At least 1 layer is required to store contents."
+          exit 1
+        fi
+        availableLayers=$(( maxLayers - usedLayers ))
+
         # Create $maxLayers worth of Docker Layers, one layer per store path
         # unless there are more paths than $maxLayers. In that case, create
         # $maxLayers-1 for the most popular layers, and smush the remainaing
@@ -838,18 +861,20 @@ rec {
                 | (.[:$maxLayers-1] | map([.])) + [ .[$maxLayers-1:] ]
                 | map(select(length > 0))
             ' \
-              --argjson maxLayers "$(( maxLayers - 1 ))" # one layer will be taken up by the customisation layer
+              --argjson maxLayers "$availableLayers"
         )"
 
         cat ${baseJson} | jq '
           . + {
             "store_dir": $store_dir,
+            "from_image": $from_image,
             "store_layers": $store_layers,
             "customisation_layer", $customisation_layer,
             "repo_tag": $repo_tag,
             "created": $created
           }
           ' --arg store_dir "${storeDir}" \
+            --argjson from_image ${if fromImage == null then "null" else "'\"${fromImage}\"'"} \
             --argjson store_layers "$store_layers" \
             --arg customisation_layer ${customisationLayer} \
             --arg repo_tag "$imageName:$imageTag" \
diff --git a/pkgs/build-support/docker/examples.nix b/pkgs/build-support/docker/examples.nix
index 9e33a42af23..9c7d4681214 100644
--- a/pkgs/build-support/docker/examples.nix
+++ b/pkgs/build-support/docker/examples.nix
@@ -188,7 +188,25 @@ rec {
     };
   };
 
-  # 12. example of running something as root on top of a parent image
+  # 12 Create a layered image on top of a layered image
+  layered-on-top-layered = pkgs.dockerTools.buildLayeredImage {
+    name = "layered-on-top-layered";
+    tag = "latest";
+    fromImage = layered-image;
+    extraCommands = ''
+      mkdir ./example-output
+      chmod 777 ./example-output
+    '';
+    config = {
+      Env = [ "PATH=${pkgs.coreutils}/bin/" ];
+      WorkingDir = "/example-output";
+      Cmd = [
+        "${pkgs.bash}/bin/bash" "-c" "echo hello > foo; cat foo"
+      ];
+    };
+  };
+
+  # 13. example of running something as root on top of a parent image
   # Regression test related to PR #52109
   runAsRootParentImage = buildImage {
     name = "runAsRootParentImage";
@@ -197,7 +215,7 @@ rec {
     fromImage = bash;
   };
 
-  # 13. example of 3 layers images This image is used to verify the
+  # 14. example of 3 layers images This image is used to verify the
   # order of layers is correct.
   # It allows to validate
   # - the layer of parent are below
@@ -235,23 +253,36 @@ rec {
     '';
   };
 
-  # 14. Environment variable inheritance.
+  # 15. Environment variable inheritance.
   # Child image should inherit parents environment variables,
   # optionally overriding them.
-  environmentVariables = let
-    parent = pkgs.dockerTools.buildImage {
-      name = "parent";
-      tag = "latest";
-      config = {
-        Env = [
-          "FROM_PARENT=true"
-          "LAST_LAYER=parent"
-        ];
-      };
+  environmentVariablesParent = pkgs.dockerTools.buildImage {
+    name = "parent";
+    tag = "latest";
+    config = {
+      Env = [
+        "FROM_PARENT=true"
+        "LAST_LAYER=parent"
+      ];
     };
-  in pkgs.dockerTools.buildImage {
+  };
+
+  environmentVariables = pkgs.dockerTools.buildImage {
+    name = "child";
+    fromImage = environmentVariablesParent;
+    tag = "latest";
+    contents = [ pkgs.coreutils ];
+    config = {
+      Env = [
+        "FROM_CHILD=true"
+        "LAST_LAYER=child"
+      ];
+    };
+  };
+
+  environmentVariablesLayered = pkgs.dockerTools.buildLayeredImage {
     name = "child";
-    fromImage = parent;
+    fromImage = environmentVariablesParent;
     tag = "latest";
     contents = [ pkgs.coreutils ];
     config = {
@@ -262,14 +293,14 @@ rec {
     };
   };
 
-  # 15. Create another layered image, for comparing layers with image 10.
+  # 16. Create another layered image, for comparing layers with image 10.
   another-layered-image = pkgs.dockerTools.buildLayeredImage {
     name = "another-layered-image";
     tag = "latest";
     config.Cmd = [ "${pkgs.hello}/bin/hello" ];
   };
 
-  # 16. Create a layered image with only 2 layers
+  # 17. Create a layered image with only 2 layers
   two-layered-image = pkgs.dockerTools.buildLayeredImage {
     name = "two-layered-image";
     tag = "latest";
@@ -278,7 +309,7 @@ rec {
     maxLayers = 2;
   };
 
-  # 17. Create a layered image with more packages than max layers.
+  # 18. Create a layered image with more packages than max layers.
   # coreutils and hello are part of the same layer
   bulk-layer = pkgs.dockerTools.buildLayeredImage {
     name = "bulk-layer";
@@ -289,7 +320,19 @@ rec {
     maxLayers = 2;
   };
 
-  # 18. Create a "layered" image without nix store layers. This is not
+  # 19. Create a layered image with a base image and more packages than max
+  # layers. coreutils and hello are part of the same layer
+  layered-bulk-layer = pkgs.dockerTools.buildLayeredImage {
+    name = "layered-bulk-layer";
+    tag = "latest";
+    fromImage = two-layered-image;
+    contents = with pkgs; [
+      coreutils hello
+    ];
+    maxLayers = 4;
+  };
+
+  # 20. Create a "layered" image without nix store layers. This is not
   # recommended, but can be useful for base images in rare cases.
   no-store-paths = pkgs.dockerTools.buildLayeredImage {
     name = "no-store-paths";
@@ -321,7 +364,7 @@ rec {
     };
   };
 
-  # 19. Support files in the store on buildLayeredImage
+  # 21. Support files in the store on buildLayeredImage
   # See: https://github.com/NixOS/nixpkgs/pull/91084#issuecomment-653496223
   filesInStore = pkgs.dockerTools.buildLayeredImageWithNixDb {
     name = "file-in-store";
@@ -341,7 +384,7 @@ rec {
     };
   };
 
-  # 20. Ensure that setting created to now results in a date which
+  # 22. Ensure that setting created to now results in a date which
   # isn't the epoch + 1 for layered images.
   unstableDateLayered = pkgs.dockerTools.buildLayeredImage {
     name = "unstable-date-layered";
diff --git a/pkgs/build-support/docker/stream_layered_image.py b/pkgs/build-support/docker/stream_layered_image.py
index 60d67442c16..3e5781ba1c8 100644
--- a/pkgs/build-support/docker/stream_layered_image.py
+++ b/pkgs/build-support/docker/stream_layered_image.py
@@ -33,6 +33,7 @@ function does all this.
 
 import io
 import os
+import re
 import sys
 import json
 import hashlib
@@ -126,10 +127,85 @@ class ExtractChecksum:
         return (self._digest.hexdigest(), self._size)
 
 
+FromImage = namedtuple("FromImage", ["tar", "manifest_json", "image_json"])
 # Some metadata for a layer
 LayerInfo = namedtuple("LayerInfo", ["size", "checksum", "path", "paths"])
 
 
+def load_from_image(from_image_str):
+    """
+    Loads the given base image, if any.
+
+    from_image_str: Path to the base image archive.
+
+    Returns: A 'FromImage' object with references to the loaded base image,
+             or 'None' if no base image was provided.
+    """
+    if from_image_str is None:
+        return None
+
+    base_tar = tarfile.open(from_image_str)
+
+    manifest_json_tarinfo = base_tar.getmember("manifest.json")
+    with base_tar.extractfile(manifest_json_tarinfo) as f:
+        manifest_json = json.load(f)
+
+    image_json_tarinfo = base_tar.getmember(manifest_json[0]["Config"])
+    with base_tar.extractfile(image_json_tarinfo) as f:
+        image_json = json.load(f)
+
+    return FromImage(base_tar, manifest_json, image_json)
+
+
+def add_base_layers(tar, from_image):
+    """
+    Adds the layers from the given base image to the final image.
+
+    tar: 'tarfile.TarFile' object for new layers to be added to.
+    from_image: 'FromImage' object with references to the loaded base image.
+    """
+    if from_image is None:
+        print("No 'fromImage' provided", file=sys.stderr)
+        return []
+
+    layers = from_image.manifest_json[0]["Layers"]
+    checksums = from_image.image_json["rootfs"]["diff_ids"]
+    layers_checksums = zip(layers, checksums)
+
+    for num, (layer, checksum) in enumerate(layers_checksums, start=1):
+        layer_tarinfo = from_image.tar.getmember(layer)
+        checksum = re.sub(r"^sha256:", "", checksum)
+
+        tar.addfile(layer_tarinfo, from_image.tar.extractfile(layer_tarinfo))
+        path = layer_tarinfo.path
+        size = layer_tarinfo.size
+
+        print("Adding base layer", num, "from", path, file=sys.stderr)
+        yield LayerInfo(size=size, checksum=checksum, path=path, paths=[path])
+
+    from_image.tar.close()
+
+
+def overlay_base_config(from_image, final_config):
+    """
+    Overlays the final image 'config' JSON on top of selected defaults from the
+    base image 'config' JSON.
+
+    from_image: 'FromImage' object with references to the loaded base image.
+    final_config: 'dict' object of the final image 'config' JSON.
+    """
+    if from_image is None:
+        return final_config
+
+    base_config = from_image.image_json["config"]
+
+    # Preserve environment from base image
+    final_env = base_config.get("Env", []) + final_config.get("Env", [])
+    if final_env:
+        final_config["Env"] = final_env
+    return final_config
+
+
 def add_layer_dir(tar, paths, store_dir, mtime):
     """
     Appends given store paths to a TarFile object as a new layer.
@@ -248,17 +324,21 @@ def main():
     mtime = int(created.timestamp())
     store_dir = conf["store_dir"]
 
+    from_image = load_from_image(conf["from_image"])
+
     with tarfile.open(mode="w|", fileobj=sys.stdout.buffer) as tar:
         layers = []
-        for num, store_layer in enumerate(conf["store_layers"]):
-            print(
-              "Creating layer", num,
-              "from paths:", store_layer,
-              file=sys.stderr)
+        layers.extend(add_base_layers(tar, from_image))
+
+        start = len(layers) + 1
+        for num, store_layer in enumerate(conf["store_layers"], start=start):
+            print("Creating layer", num, "from paths:", store_layer,
+                  file=sys.stderr)
             info = add_layer_dir(tar, store_layer, store_dir, mtime=mtime)
             layers.append(info)
 
-        print("Creating the customisation layer...", file=sys.stderr)
+        print("Creating layer", len(layers) + 1, "with customisation...",
+              file=sys.stderr)
         layers.append(
           add_customisation_layer(
             tar,
@@ -273,7 +353,7 @@ def main():
             "created": datetime.isoformat(created),
             "architecture": conf["architecture"],
             "os": "linux",
-            "config": conf["config"],
+            "config": overlay_base_config(from_image, conf["config"]),
             "rootfs": {
                 "diff_ids": [f"sha256:{layer.checksum}" for layer in layers],
                 "type": "layers",