Support LABEL in docker_build rule

Allows users to associate custom metadata with docker images.

Discussion:
https://groups.google.com/d/msg/bazel-dev/FTQVg2U3CvQ/X-8RJ01_AgAJ

--
Change-Id: Ia9f6ceb1dd99aa91cf0c41f3d7afc447ab5792e5
Reviewed-on: https://bazel-review.googlesource.com/#/c/2350/
MOS_MIGRATED_REVID=110489115

7 files changed, 286 insertions, 43 deletions

diff --git a/tools/build_defs/docker/BUILD b/tools/build_defs/docker/BUILD
index d5d6f7b1f7..2832767596 100644
--- a/tools/build_defs/docker/BUILD
+++ b/tools/build_defs/docker/BUILD
@@ -22,6 +22,8 @@ TEST_TARGETS = [
     "generated_tarball",
     "with_env",
     "with_double_env",
+    "with_label",
+    "with_double_label",
     "workdir_with_tar_base",
     "link_with_files_base",
 ]
diff --git a/tools/build_defs/docker/README.md b/tools/build_defs/docker/README.md
index 420472f2d0..be0f120f97 100644
--- a/tools/build_defs/docker/README.md
+++ b/tools/build_defs/docker/README.md
@@ -199,7 +199,7 @@ repositories: pull and push docker image.
 ## docker_build
 
 ```python
-docker_build(name, base, data_path, directory, files, mode, tars, debs, symlinks, entrypoint, cmd, env, ports, volumes, workdir, repository)
+docker_build(name, base, data_path, directory, files, mode, tars, debs, symlinks, entrypoint, cmd, env, labels, ports, volumes, workdir, repository)
 ```
 
 <table class="table table-condensed table-bordered table-implicit">
@@ -379,6 +379,42 @@ docker_build(name, base, data_path, directory, files, mode, tars, debs, symlinks
       </td>
     </tr>
     <tr>
+      <td><code>env</code></td>
+      <td>
+        <code>Dictionary from strings to strings, optional</code>
+        <p><a href="https://docs.docker.com/reference/builder/#env">Dictionary
+               from environment variable names to their values when running the
+               docker image.</a></p>
+        <p>
+          <code>
+          env = {
+            "FOO": "bar",
+            ...
+          },
+          </code>
+        </p>
+      </td>
+    </tr>
+    <tr>
+      <td><code>labels</code></td>
+      <td>
+        <code>Dictionary from strings to strings, optional</code>
+        <p><a href="https://docs.docker.com/reference/builder/#label">Dictionary
+               from custom metadata names to their values. You can also put a
+               file name prefixed by '@' as a value. Then the value is replaced
+               with the contents of the file.
+        <p>
+          <code>
+          labels = {
+            "com.example.foo": "bar",
+            "com.example.baz": "@metadata.json",
+            ...
+          },
+          </code>
+        </p>
+      </td>
+    </tr>
+    <tr>
       <td><code>ports</code></td>
       <td>
         <code>String list, optional</code>
diff --git a/tools/build_defs/docker/build_test.sh b/tools/build_defs/docker/build_test.sh
index 7aea3fd4f7..046705bd7b 100755
--- a/tools/build_defs/docker/build_test.sh
+++ b/tools/build_defs/docker/build_test.sh
@@ -110,6 +110,13 @@ function check_env() {
   check_property Env "notop_${input}" "${@}"
 }
 
+function check_label() {
+  input="$1"
+  shift
+  check_property Label "${input}" "${@}"
+  check_property Label "notop_${input}" "${@}"
+}
+
 function check_workdir() {
   input="$1"
   shift
@@ -341,6 +348,27 @@ function test_with_double_env() {
     '["bar=blah blah blah", "baz=/asdf blah blah blah", "foo=/asdf"]'
 }
 
+function test_with_label() {
+  check_layers "with_label" \
+    "125e7cfb9d4a6d803a57b88bcdb05d9a6a47ac0d6312a8b4cff52a2685c5c858" \
+    "eba6abda3d259ab6ed5f4d48b76df72a5193fad894d4ae78fbf0a363d8f9e8fd"
+
+  check_label "with_label" \
+    "eba6abda3d259ab6ed5f4d48b76df72a5193fad894d4ae78fbf0a363d8f9e8fd" \
+    '["com.example.bar={\"name\": \"blah\"}", "com.example.baz=qux", "com.example.foo={\"name\": \"blah\"}"]'
+}
+
+function test_with_double_label() {
+  check_layers "with_double_label" \
+    "125e7cfb9d4a6d803a57b88bcdb05d9a6a47ac0d6312a8b4cff52a2685c5c858" \
+    "eba6abda3d259ab6ed5f4d48b76df72a5193fad894d4ae78fbf0a363d8f9e8fd" \
+    "bfe88fbb5e24fc5bff138f7a1923d53a2ee1bbc8e54b6f5d9c371d5f48b6b023" \
+
+  check_label "with_double_label" \
+    "bfe88fbb5e24fc5bff138f7a1923d53a2ee1bbc8e54b6f5d9c371d5f48b6b023" \
+    '["com.example.bar={\"name\": \"blah\"}", "com.example.baz=qux", "com.example.foo={\"name\": \"blah\"}", "com.example.qux={\"name\": \"blah-blah\"}"]'
+}
+
 function get_layer_listing() {
   local input=$1
   local layer=$2
diff --git a/tools/build_defs/docker/docker.bzl b/tools/build_defs/docker/docker.bzl
index 395f0277d0..eb454a2b63 100644
--- a/tools/build_defs/docker/docker.bzl
+++ b/tools/build_defs/docker/docker.bzl
@@ -133,23 +133,43 @@ def _get_base_artifact(ctx):
       fail("base attribute should be a single tar file.")
     return ctx.files.base[0]
 
+def _serialize_dict(dict_value):
+    return ",".join(["%s=%s" % (k, dict_value[k]) for k in dict_value])
+
 def _metadata_action(ctx, layer, name, output):
   """Generate the action to create the JSON metadata for the layer."""
   rewrite_tool = ctx.executable._rewrite_tool
-  env = ctx.attr.env
+
+  label_file_dict = dict()
+  for i in range(len(ctx.files.label_files)):
+    fname = ctx.attr.label_file_strings[i]
+    file = ctx.files.label_files[i]
+    label_file_dict[fname] = file
+
+  labels = dict()
+  for l in ctx.attr.labels:
+    fname = ctx.attr.labels[l]
+    if fname[0] == '@':
+      labels[l] = "@" + label_file_dict[fname[1:]].path
+    else:
+      labels[l] = fname
+
   args = [
       "--output=%s" % output.path,
       "--layer=%s" % layer.path,
       "--name=@%s" % name.path,
       "--entrypoint=%s" % ",".join(ctx.attr.entrypoint),
       "--command=%s" % ",".join(ctx.attr.cmd),
-      "--env=%s" % ",".join(["%s=%s" % (k, env[k]) for k in env]),
+      "--labels=%s" % _serialize_dict(labels),
+      "--env=%s" % _serialize_dict(ctx.attr.env),
       "--ports=%s" % ",".join(ctx.attr.ports),
       "--volumes=%s" % ",".join(ctx.attr.volumes)
       ]
   if ctx.attr.workdir:
     args += ["--workdir=" + ctx.attr.workdir]
   inputs = [layer, rewrite_tool, name]
+  if ctx.attr.label_files:
+    inputs += ctx.files.label_files
   base = _get_base_artifact(ctx)
   if base:
     args += ["--base=%s" % base.path]
@@ -302,11 +322,15 @@ docker_build_ = rule(
         "entrypoint": attr.string_list(),
         "cmd": attr.string_list(),
         "env": attr.string_dict(),
+        "labels": attr.string_dict(),
         "ports": attr.string_list(),  # Skylark doesn't support int_list...
         "volumes": attr.string_list(),
         "workdir": attr.string(),
         "repository": attr.string(default="bazel"),
         # Implicit dependencies.
+        "label_files": attr.label_list(
+            allow_files=True),
+        "label_file_strings": attr.string_list(),
         "_build_layer": attr.label(
             default=Label("//tools/build_defs/pkg:build_tar"),
             cfg=HOST_CFG,
@@ -447,6 +471,13 @@ def docker_build(**kwargs):
   """
   if "cmd" in kwargs:
     kwargs["cmd"] = _validate_command("cmd", kwargs["cmd"])
+  for reserved in ["label_files", "label_file_strings"]:
+    if reserved in kwargs:
+      fail("reserved for internal use by docker_build macro", attr=reserved)
+  if "labels" in kwargs:
+    files = sorted(set([v[1:] for v in kwargs["labels"].values() if v[0] == '@']))
+    kwargs["label_files"] = files
+    kwargs["label_file_strings"] = files
   if "entrypoint" in kwargs:
     kwargs["entrypoint"] = _validate_command("entrypoint", kwargs["entrypoint"])
   docker_build_(**kwargs)
diff --git a/tools/build_defs/docker/rewrite_json.py b/tools/build_defs/docker/rewrite_json.py
index fe62a91149..11933e5d51 100644
--- a/tools/build_defs/docker/rewrite_json.py
+++ b/tools/build_defs/docker/rewrite_json.py
@@ -42,6 +42,8 @@ gflags.DEFINE_list(
     'command', None,
     'Override the "Cmd" of the previous layer')
 
+gflags.DEFINE_list('labels', None, 'Augment the "Label" of the previous layer')
+
 gflags.DEFINE_list(
     'ports', None,
     'Augment the "ExposedPorts" of the previous layer')
@@ -60,23 +62,37 @@ gflags.DEFINE_list(
 
 FLAGS = gflags.FLAGS
 
-_MetadataOptionsT = namedtuple(
-    'MetadataOptionsT',
-    ['name', 'parent', 'size', 'entrypoint', 'cmd', 'env', 'ports', 'volumes',
-     'workdir'])
+_MetadataOptionsT = namedtuple('MetadataOptionsT',
+                               ['name', 'parent', 'size', 'entrypoint', 'cmd',
+                                'env', 'labels', 'ports', 'volumes', 'workdir'])
 
 
 class MetadataOptions(_MetadataOptionsT):
   """Docker image layer metadata options."""
 
-  def __new__(cls, name=None, parent=None, size=None,
-              entrypoint=None, cmd=None, env=None,
-              ports=None, volumes=None, workdir=None):
+  def __new__(cls,
+              name=None,
+              parent=None,
+              size=None,
+              entrypoint=None,
+              cmd=None,
+              labels=None,
+              env=None,
+              ports=None,
+              volumes=None,
+              workdir=None):
     """Constructor."""
-    return super(MetadataOptions, cls).__new__(
-        cls, name=name, parent=parent, size=size,
-        entrypoint=entrypoint, cmd=cmd, env=env,
-        ports=ports, volumes=volumes, workdir=workdir)
+    return super(MetadataOptions, cls).__new__(cls,
+                                               name=name,
+                                               parent=parent,
+                                               size=size,
+                                               entrypoint=entrypoint,
+                                               cmd=cmd,
+                                               labels=labels,
+                                               env=env,
+                                               ports=ports,
+                                               volumes=volumes,
+                                               workdir=workdir)
 
 
 _DOCKER_VERSION = '1.5.0'
@@ -104,6 +120,15 @@ def DeepCopySkipNull(data):
   return copy.deepcopy(data)
 
 
+def KeyValueToDict(pair):
+  """Converts an iterable object of key=value pairs to dictionary."""
+  d = dict()
+  for kv in pair:
+    (k, v) = kv.split('=', 1)
+    d[k] = v
+  return d
+
+
 def RewriteMetadata(data, options):
   """Rewrite and return a copy of the input data according to options.
 
@@ -147,20 +172,23 @@ def RewriteMetadata(data, options):
   output['architecture'] = _PROCESSOR_ARCHITECTURE
   output['os'] = _OPERATING_SYSTEM
 
+  def Dict2ConfigValue(d):
+    return ['%s=%s' % (k, d[k]) for k in sorted(d.keys())]
+
   if options.env:
-    environ_dict = {}
     # Build a dictionary of existing environment variables (used by Resolve).
-    for kv in output['config'].get('Env', []):
-      (k, v) = kv.split('=', 1)
-      environ_dict[k] = v
+    environ_dict = KeyValueToDict(output['config'].get('Env', []))
     # Merge in new environment variables, resolving references.
-    for kv in options.env:
-      (k, v) = kv.split('=', 1)
+    for k, v in options.env.iteritems():
       # Resolve handles scenarios like "PATH=$PATH:...".
-      v = Resolve(v, environ_dict)
-      environ_dict[k] = v
-    output['config']['Env'] = [
-        '%s=%s' % (k, environ_dict[k]) for k in sorted(environ_dict.keys())]
+      environ_dict[k] = Resolve(v, environ_dict)
+    output['config']['Env'] = Dict2ConfigValue(environ_dict)
+
+  if options.labels:
+    label_dict = KeyValueToDict(output['config'].get('Label', []))
+    for k, v in options.labels.iteritems():
+      label_dict[k] = v
+    output['config']['Label'] = Dict2ConfigValue(label_dict)
 
   if options.ports:
     if 'ExposedPorts' not in output['config']:
@@ -263,16 +291,23 @@ def main(unused_argv):
     with open(name[1:], 'r') as f:
       name = f.read()
 
-  output = RewriteMetadata(data, MetadataOptions(
-      name=name,
-      parent=parent,
-      size=os.path.getsize(FLAGS.layer),
-      entrypoint=FLAGS.entrypoint,
-      cmd=FLAGS.command,
-      env=FLAGS.env,
-      ports=FLAGS.ports,
-      volumes=FLAGS.volumes,
-      workdir=FLAGS.workdir))
+  labels = KeyValueToDict(FLAGS.labels)
+  for label, value in labels.iteritems():
+    if value.startswith('@'):
+      with open(value[1:], 'r') as f:
+        labels[label] = f.read()
+
+  output = RewriteMetadata(data,
+                           MetadataOptions(name=name,
+                                           parent=parent,
+                                           size=os.path.getsize(FLAGS.layer),
+                                           entrypoint=FLAGS.entrypoint,
+                                           cmd=FLAGS.command,
+                                           labels=labels,
+                                           env=KeyValueToDict(FLAGS.env),
+                                           ports=FLAGS.ports,
+                                           volumes=FLAGS.volumes,
+                                           workdir=FLAGS.workdir))
 
   with open(FLAGS.output, 'w') as fp:
     json.dump(output, fp, sort_keys=True)
diff --git a/tools/build_defs/docker/rewrite_json_test.py b/tools/build_defs/docker/rewrite_json_test.py
index 3d89d1dbf5..d751208065 100644
--- a/tools/build_defs/docker/rewrite_json_test.py
+++ b/tools/build_defs/docker/rewrite_json_test.py
@@ -556,17 +556,17 @@ class RewriteJsonTest(unittest.TestCase):
     }
     name = 'deadbeef'
     parent = 'blah'
-    env = [
-        'baz=blah',
-        'foo=bar',
-    ]
+    env = {'baz': 'blah', 'foo': 'bar',}
     expected = {
         'id': name,
         'parent': parent,
         'config': {
             'User': 'mattmoor',
             'WorkingDir': '/usr/home/mattmoor',
-            'Env': env,
+            'Env': [
+                'baz=blah',
+                'foo=bar',
+            ],
         },
         'docker_version': _DOCKER_VERSION,
         'architecture': _PROCESSOR_ARCHITECTURE,
@@ -591,10 +591,7 @@ class RewriteJsonTest(unittest.TestCase):
     }
     name = 'deadbeef'
     parent = 'blah'
-    env = [
-        'baz=replacement',
-        'foo=$foo:asdf',
-    ]
+    env = {'baz': 'replacement', 'foo': '$foo:asdf',}
     expected = {
         'id': name,
         'parent': parent,
@@ -616,6 +613,75 @@ class RewriteJsonTest(unittest.TestCase):
         name=name, env=env, parent=parent))
     self.assertEquals(expected, actual)
 
+  def testLabel(self):
+    in_data = {
+        'config': {
+            'User': 'mattmoor',
+            'WorkingDir': '/usr/home/mattmoor'
+        }
+    }
+    name = 'deadbeef'
+    parent = 'blah'
+    labels = {'baz': 'blah', 'foo': 'bar',}
+    expected = {
+        'id': name,
+        'parent': parent,
+        'config': {
+            'User': 'mattmoor',
+            'WorkingDir': '/usr/home/mattmoor',
+            'Label': [
+                'baz=blah',
+                'foo=bar',
+            ],
+        },
+        'docker_version': _DOCKER_VERSION,
+        'architecture': _PROCESSOR_ARCHITECTURE,
+        'os': _OPERATING_SYSTEM,
+    }
+
+    actual = RewriteMetadata(in_data,
+                             MetadataOptions(name=name,
+                                             labels=labels,
+                                             parent=parent))
+    self.assertEquals(expected, actual)
+
+  def testAugmentLabel(self):
+    in_data = {
+        'config': {
+            'User': 'mattmoor',
+            'WorkingDir': '/usr/home/mattmoor',
+            'Label': [
+                'baz=blah',
+                'blah=still around',
+            ],
+        }
+    }
+    name = 'deadbeef'
+    parent = 'blah'
+    labels = {'baz': 'replacement', 'foo': 'bar',}
+    expected = {
+        'id': name,
+        'parent': parent,
+        'config': {
+            'User': 'mattmoor',
+            'WorkingDir': '/usr/home/mattmoor',
+            'Label': [
+                'baz=replacement',
+                'blah=still around',
+                'foo=bar',
+            ],
+        },
+        'docker_version': _DOCKER_VERSION,
+        'architecture': _PROCESSOR_ARCHITECTURE,
+        'os': _OPERATING_SYSTEM,
+    }
+
+    actual = RewriteMetadata(in_data,
+                             MetadataOptions(name=name,
+                                             labels=labels,
+                                             parent=parent))
+    self.assertEquals(expected, actual)
+
   def testAugmentVolumeWithNullInput(self):
     in_data = {
         'config': {
diff --git a/tools/build_defs/docker/testdata/BUILD b/tools/build_defs/docker/testdata/BUILD
index c1d1896e50..f8b077f88a 100644
--- a/tools/build_defs/docker/testdata/BUILD
+++ b/tools/build_defs/docker/testdata/BUILD
@@ -195,6 +195,33 @@ docker_build(
 )
 
 docker_build(
+    name = "with_label",
+    base = ":base_with_volume",
+    labels = {
+        "com.example.foo": "@blah.json",
+        "com.example.bar": "@blah.json",
+        "com.example.baz": "qux",
+    },
+)
+
+docker_build(
+    name = "with_double_label",
+    base = ":with_label",
+    labels = {
+        "com.example.qux": "@blah-blah.json",
+    },
+)
+
+[genrule(
+    name = "label-" + n,
+    outs = ["%s.json" % n],
+    cmd = "echo -n '{\"name\": \"%s\"}' > $@" % n,
+) for n in [
+    "blah",
+    "blah-blah",
+]]
+
+docker_build(
     name = "link_with_files_base",
     base = ":files_base",
     symlinks = {
@@ -308,6 +335,24 @@ docker_build(
 )
 
 docker_build(
+    name = "notop_with_label",
+    base = ":notop_base_with_volume",
+    labels = {
+        "com.example.foo": "@blah.json",
+        "com.example.bar": "@blah.json",
+        "com.example.baz": "qux",
+    },
+)
+
+docker_build(
+    name = "notop_with_double_label",
+    base = ":notop_with_label",
+    labels = {
+        "com.example.qux": "@blah-blah.json",
+    },
+)
+
+docker_build(
     name = "notop_link_with_files_base",
     base = ":notop_files_base",
     symlinks = {