add container run and commit layer command

contrib/BUILD

@@ -44,6 +44,13 @@ py_library(
     srcs_version = "PY2AND3",
 )
 
+py_binary(
+    name = "extract_last_layer",
+    srcs = [":extract_last_layer.py"],
+    legacy_create_init = False,
+    python_version = "PY3",
+)
+
 py_binary(
     name = "compare_ids_test",
     srcs = [":compare_ids_test.py"],

contrib/extract_last_layer.py

@@ -0,0 +1,95 @@
+# Copyright 2020 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Extracts the last layer of a docker image out of an image tarball
+
+Takes three arguments, the path to the tarball, and the output file for the layer, and the output file for the layer diffID
+"""
+
+
+from __future__ import print_function
+from json import JSONDecoder
+import hashlib
+import sys
+import tarfile
+
+
+def extract_last_layer(tar_path, layer_path, diffid_path):
+  """Extracts the last layer from a docker image from an image tarball
+
+  Args:
+    tar_path: str path to the tarball
+    layer_path: str path for the output layer
+    diffid_path: str path for the layer diff ID
+
+
+  Returns:
+    str the diff ID of the layer
+
+  """
+  tar = tarfile.open(tar_path, mode="r")
+
+  decoder = JSONDecoder()
+  try:
+    # Extracts it as a file object (not to the disk)
+    manifest = tar.extractfile("manifest.json").read().decode("utf-8")
+  except Exception as e:
+    print((
+        "Unable to extract manifest.json, make sure {} "
+        "is a valid docker image.\n").format(tar_path),
+          e,
+          file=sys.stderr)
+    exit(1)
+
+  # Get the manifest dictionary from JSON
+  manifest = decoder.decode(manifest)[0]
+
+  layers = manifest["Layers"]
+
+  last_layer_path = layers[-1]
+
+  layer_id = last_layer_path.split("/")[0]
+
+  diff_id = hashlib.sha256()
+
+  try:
+    last_layer = tar.extractfile(last_layer_path)
+    with open(layer_path, "wb") as f:
+      while True:
+        buf = last_layer.read(4096)
+        if buf:
+          diff_id.update(buf)
+          f.write(buf)
+        else:
+          break
+  except Exception as e:
+    print((
+        "Unable to extract last layer {} to {}, make sure {} "
+        "is a valid docker image and that the layer path is writable\n").format(layer_id, layer_path, tar_path),
+          e,
+          file=sys.stderr)
+    exit(1)
+
+  diff_id_digest = diff_id.hexdigest()
+  try:
+    with open(diffid_path, "w") as f:
+      f.write(diff_id_digest)
+  except Exception as e:
+    print("Unable to write layer Diff ID {} to {}, make sure the path is writeable\n".format(diff_id_digest, diffid_path), e, file=sys.stderr)
+    exit(1)
+
+  return layer_id
+
+
+if __name__ == "__main__":
+  print(extract_last_layer(sys.argv[1], sys.argv[2], sys.argv[3]))

docker/util/BUILD

@@ -34,6 +34,7 @@ exports_files([
     "commit.sh.tpl",
     "extract.sh.tpl",
     "image_util.sh.tpl",
+    "commit_layer.sh.tpl",
 ])
 
 bzl_library(

docker/util/commit_layer.sh.tpl

@@ -0,0 +1,40 @@
+#!/bin/bash
+
+set -ex
+
+# Load utils
+source %{util_script}
+
+# Resolve the docker tool path
+DOCKER="%{docker_tool_path}"
+DOCKER_FLAGS="%{docker_flags}"
+
+if [[ -z "$DOCKER" ]]; then
+    echo >&2 "error: docker not found; do you need to manually configure the docker toolchain?"
+    exit 1
+fi
+
+# Load the image and remember its name
+image_id=$(%{image_id_extractor_path} %{image_tar})
+$DOCKER $DOCKER_FLAGS load -i %{image_tar}
+
+id=$($DOCKER $DOCKER_FLAGS run -d %{docker_run_flags} $image_id %{commands})
+# Actually wait for the container to finish running its commands
+retcode=$($DOCKER $DOCKER_FLAGS wait $id)
+# Trigger a failure if the run had a non-zero exit status
+if [ $retcode != 0 ]; then
+  $DOCKER $DOCKER_FLAGS logs $id && false
+fi
+OUTPUT_IMAGE_TAR="%{output_layer_tar}.image.tar"
+reset_cmd $image_id $id %{output_image}
+$DOCKER $DOCKER_FLAGS save %{output_image} -o $OUTPUT_IMAGE_TAR
+$DOCKER $DOCKER_FLAGS rm $id
+$DOCKER $DOCKER_FLAGS rmi %{output_image}
+
+%{image_last_layer_extractor_path} $OUTPUT_IMAGE_TAR %{output_layer_tar} %{output_diff_id}
+
+rm $OUTPUT_IMAGE_TAR
+
+
+
+

docker/util/run.bzl

@@ -18,6 +18,14 @@ to new container image, or extract specified targets to a directory on
 the host machine.
 """
 
+load("@bazel_skylib//lib:dicts.bzl", "dicts")
+load(
+    "@bazel_tools//tools/build_defs/hash:hash.bzl",
+    _hash_tools = "tools",
+)
+load("@io_bazel_rules_docker//container:providers.bzl", "LayerInfo")
+load("@io_bazel_rules_docker//container:layer.bzl", "zip_layer")
+
 def _extract_impl(
         ctx,
         name = "",
@@ -277,6 +285,169 @@ commit = struct(
     implementation = _commit_impl,
 )
 
+
+def _commit_layer_impl(
+        ctx,
+        name = None,
+        image = None,
+        commands = None,
+        docker_run_flags = None,
+        compression = None,
+        compression_options = None,
+        output_layer_tar = None):
+    """Implementation for the container_run_and_commit_layer rule.
+
+    This rule runs a set of commands in a given image, waits for the commands
+    to finish, and then commits the container to a new image.
+
+    Args:
+        ctx: The bazel rule context
+        name: A unique name for this rule.
+        image: The input image tarball
+        commands: The commands to run in the input image container
+        docker_run_flags: String list, overrides ctx.attr.docker_run_flags
+        output_image_tar: The output image obtained as a result of running
+                          the commands on the input image
+    """
+
+    name = name or ctx.attr.name
+    image = image or ctx.file.image
+    commands = commands or ctx.attr.commands
+    docker_run_flags = docker_run_flags or ctx.attr.docker_run_flags
+    script = ctx.actions.declare_file(name + ".build")
+    compression = compression or ctx.attr.compression
+    compression_options = compression_options or ctx.attr.compression_options
+    output_layer_tar = output_layer_tar or ctx.outputs.layer
+
+    toolchain_info = ctx.toolchains["@io_bazel_rules_docker//toolchains/docker:toolchain_type"].info
+
+    # Generate a shell script to execute the reset cmd
+    image_utils = ctx.actions.declare_file("image_util.sh")
+    ctx.actions.expand_template(
+        template = ctx.file._image_utils_tpl,
+        output = image_utils,
+        substitutions = {
+            "%{docker_flags}": " ".join(toolchain_info.docker_flags),
+            "%{docker_tool_path}": toolchain_info.tool_path,
+        },
+        is_executable = True,
+    )
+
+    output_diff_id = ctx.actions.declare_file(output_layer_tar.basename + ".sha256")
+
+    # Generate a shell script to execute the run statement
+    ctx.actions.expand_template(
+        template = ctx.file._run_tpl,
+        output = script,
+        substitutions = {
+            "%{commands}": _process_commands(commands),
+            "%{docker_flags}": " ".join(toolchain_info.docker_flags),
+            "%{docker_run_flags}": " ".join(docker_run_flags),
+            "%{docker_tool_path}": toolchain_info.tool_path,
+            "%{image_last_layer_extractor_path}": ctx.executable._last_layer_extractor_tool.path,
+            "%{image_id_extractor_path}": ctx.executable._extract_image_id.path,
+            "%{image_tar}": image.path,
+            "%{output_image}": "bazel/%s:%s" % (
+                ctx.label.package or "default",
+                name,
+            ),
+            "%{output_layer_tar}": output_layer_tar.path,
+            "%{output_diff_id}": output_diff_id.path,
+            "%{util_script}": image_utils.path,
+        },
+        is_executable = True,
+    )
+
+    runfiles = [image, image_utils]
+
+    ctx.actions.run(
+        outputs = [output_layer_tar, output_diff_id],
+        inputs = runfiles,
+        executable = script,
+        tools = [ctx.executable._extract_image_id, ctx.executable._last_layer_extractor_tool],
+        use_default_shell_env = True,
+    )
+
+    zipped_layer, blob_sum = zip_layer(
+        ctx,
+        output_layer_tar,
+        compression = compression,
+        compression_options = compression_options,
+    )
+
+
+    return [
+        LayerInfo(
+            unzipped_layer = output_layer_tar,
+            diff_id = output_diff_id,
+            zipped_layer = zipped_layer,
+            blob_sum = blob_sum,
+            env = {}
+        )
+    ] 
+
+_commit_layer_attrs = dicts.add({
+    "commands": attr.string_list(
+        doc = "A list of commands to run (sequentially) in the container.",
+        mandatory = True,
+        allow_empty = False,
+    ),
+    "docker_run_flags": attr.string_list(
+        doc = "Extra flags to pass to the docker run command.",
+        mandatory = False,
+    ),
+    "image": attr.label(
+        doc = "The image to run the commands in.",
+        mandatory = True,
+        allow_single_file = True,
+        cfg = "target",
+    ),
+    "compression": attr.string(default = "gzip"),
+    "compression_options": attr.string_list(),
+    "_image_utils_tpl": attr.label(
+        default = "//docker/util:image_util.sh.tpl",
+        allow_single_file = True,
+    ),
+    "_run_tpl": attr.label(
+        default = Label("//docker/util:commit_layer.sh.tpl"),
+        allow_single_file = True,
+    ),
+    "_last_layer_extractor_tool": attr.label(
+        default = Label("//contrib:extract_last_layer"),
+        cfg = "host",
+        executable = True,
+        allow_files = True,
+    ),
+    "_extract_image_id": attr.label(
+        default = Label("//contrib:extract_image_id"),
+        cfg = "host",
+        executable = True,
+        allow_files = True,
+    ),
+}, _hash_tools)
+
+_commit_layer_outputs = {
+    "layer": "%{name}-layer.tar",
+}
+
+container_run_and_commit_layer = rule(
+    attrs = _commit_layer_attrs,
+    doc = ("This rule runs a set of commands in a given image, waits" +
+           "for the commands to finish, and then commits the" +
+           "container state to a new layer."),
+    executable = False,
+    outputs = _commit_layer_outputs,
+    implementation = _commit_layer_impl,
+    toolchains = ["@io_bazel_rules_docker//toolchains/docker:toolchain_type"],
+)
+
+# Export container_run_and_commit_layer rule for other bazel rules to depend on.
+commit_layer = struct(
+    attrs = _commit_layer_attrs,
+    outputs = _commit_layer_outputs,
+    implementation = _commit_layer_impl,
+)
+
 def _process_commands(command_list):
     # Use the $ to allow escape characters in string
     return 'sh -c $\"{0}\"'.format(" && ".join(command_list))