enable using arbitrary containers for rbe_autoconfig (#264)

* enable using arbitrary containers for rbe_autoconfig

* address review comments

* fix issue

* reflect changes in sample bazelrc file

* fixes

* changing order of checks

* fix printing of error message

* addressing more review comments

* more review comments
diff --git a/examples/remotebuildexecution/rbe_system_check/.bazelrc.sample b/examples/remotebuildexecution/rbe_system_check/.bazelrc.sample
index 4c5a42a..02e10ed 100644
--- a/examples/remotebuildexecution/rbe_system_check/.bazelrc.sample
+++ b/examples/remotebuildexecution/rbe_system_check/.bazelrc.sample
@@ -31,12 +31,12 @@
 # Toolchain/platform flags that work with rbe_autoconfig rule.
 # Note target name must be set to "rbe_default" to use these flags as is.
 build:remote --crosstool_top=@rbe_default//rbe_config_cc:toolchain
-build:remote --host_javabase=@rbe_default//platforms:jdk8
-build:remote --javabase=@rbe_default//platforms:jdk8
+build:remote --host_javabase=@rbe_default//config:jdk8
+build:remote --javabase=@rbe_default//config:jdk8
 build:remote --host_java_toolchain=@bazel_tools//tools/jdk:toolchain_hostjdk8
 build:remote --java_toolchain=@bazel_tools//tools/jdk:toolchain_hostjdk8
-build:remote --extra_execution_platforms=@rbe_default//platforms:rbe_ubuntu1604
-build:remote --host_platform=@rbe_default//platforms:rbe_ubuntu1604
-build:remote --platforms=@rbe_default//platforms:rbe_ubuntu1604
-build:remote --extra_toolchains=@rbe_default//platforms:cc-toolchain
+build:remote --extra_execution_platforms=@rbe_default//config:platform
+build:remote --host_platform=@rbe_default//config:platform
+build:remote --platforms=@rbe_default//config:platform
+build:remote --extra_toolchains=@rbe_default//config:cc-toolchain
 
diff --git a/rules/BUILD.platform.tpl b/rules/BUILD.platform.tpl
index 44829f1..7daa34c 100644
--- a/rules/BUILD.platform.tpl
+++ b/rules/BUILD.platform.tpl
@@ -38,34 +38,24 @@
 toolchain(
     name = "cc-toolchain",
     exec_compatible_with = [
-        "@bazel_tools//platforms:linux",
-        "@bazel_tools//platforms:x86_64",
-        "@bazel_tools//tools/cpp:clang",
+        %{exec_compatible_with}
     ],
     target_compatible_with = [
-        "@bazel_tools//platforms:linux",
-        "@bazel_tools//platforms:x86_64",
+        %{target_compatible_with}
     ],
     toolchain = "%{toolchain}:cc-compiler-k8",
     toolchain_type = "@bazel_tools//tools/cpp:toolchain_type",
 )
 
-alias(
-    name = "rbe_ubuntu1604",
-    actual = ":rbe_ubuntu1604_%{revision}",
-)
-
 platform(
-    name = "rbe_ubuntu1604_%{revision}",
+    name = "platform",
     constraint_values = [
-        "@bazel_tools//platforms:x86_64",
-        "@bazel_tools//platforms:linux",
-        "@bazel_tools//tools/cpp:clang",
+        %{exec_compatible_with}
     ],
     remote_execution_properties = """
         properties: {
           name: "container-image"
-          value:"docker://gcr.io/cloud-marketplace/google/rbe-ubuntu16-04@%{rbe_ubuntu16_04_sha256}"
+          value:"docker://%{image_name}"
         }
         """,
 )
diff --git a/rules/extract.sh.tpl b/rules/extract.sh.tpl
index 20b5027..bd0c642 100644
--- a/rules/extract.sh.tpl
+++ b/rules/extract.sh.tpl
@@ -19,7 +19,7 @@
 # This is a generated file that runs a docker container, waits for it to
 # finish running and copies a file to an output location.
 
-id=$(docker run -d %{docker_run_flags} %{image_id} %{commands})
+id=$(docker run -d %{docker_run_flags} %{image_name} %{commands})
 
 docker wait $id
 docker cp $id:%{extract_file} %{output}
diff --git a/rules/rbe_repo.bzl b/rules/rbe_repo.bzl
index 98a1a04..285b433 100644
--- a/rules/rbe_repo.bzl
+++ b/rules/rbe_repo.bzl
@@ -12,18 +12,25 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Repository Rules to generate toolchain configs for an RBE-ubuntu16-04 container.
+"""Repository Rules to generate toolchain configs for a container image.
+
+The toolchain configs (+ platform) produced by this rule can be used to, e.g.,
+run a remote build in which remote actions will run inside a container image.
 
 Exposes the rbe_autoconfig macro that does the following:
-- Pulls an rbe-ubuntu 16_04 image (using 'docker pull')
-- Starts up a container using the rbe-ubuntu 16_04 image mounting the current project
+- Pulls an rbe-ubuntu 16_04 image (using 'docker pull'). Image to pull can be overriden.
+- Starts up a container using the pulled image, mounting either a small sample
+  project or the current project (if output_base is set).
 - Installs the current version of Bazel (one currently running) on the container
-  (or the one passed in with optional attr).
+  (or the one passed in with optional attr). Container must have tools required to install
+  and run Bazel (i.e., a jdk, a C/C++ compiler, a python interpreter).
 - Runs a bazel command to build the local_config_cc remote repository inside the container.
 - Extracts local_config_cc produced files (inside the container) to the produced
   remote repository.
+- Produces a default BUILD file with platform and toolchain targets to use the container
+  in a remote build.
 - Optionally copies the local_config_cc produced files to the project srcs under the
-  given output_base directory
+  given output_base directory.
 
 Add to your WORKSPACE file the following:
 
@@ -48,12 +55,30 @@
 
   load("@bazel_toolchains//rules:rbe_repo.bzl", "rbe_autoconfig")
 
+  # This is the simplest form of calling rbe_autoconfig.
+  # See below other examples, but your WORKSPACE should likely
+  # only need one single rbe_autoconfig target
   rbe_autoconfig(
     name = "rbe_default",
-    # Optional. See below.
+  )
+
+  # You can pass an optional output_base if you want the produced
+  # toolchain configs to be copied to your source tree (recommended).
+  rbe_autoconfig(
+    name = "rbe_default_with_output_base",
     output_base = "rbe-configs"
   )
 
+  # If you are using a custom container for remote builds, just
+  # set some extra args (see also about env variables)
+  rbe_autoconfig(
+    name = "rbe_my_custom_container",
+    registry = "gcr.io",
+    registry = "my-project/my-base",
+    # tag is not supported, always use a digest
+    digest = "sha256:deadbeef",
+  )
+
 For values of <latest_release> and other placeholders above, please see
 the WORKSPACE file in this repo.
 
@@ -72,20 +97,21 @@
 
       bazel build ... \
                 --crosstool_top=//rbe-configs/bazel_{bazel_version}:toolchain \
-                --host_javabase=//rbe-configs/bazel_{bazel_version}/platforms:jdk8 \
-                --javabase=//rbe-configs/bazel_{bazel_version}/platforms:jdk8 \
+                --host_javabase=//rbe-configs/bazel_{bazel_version}/config:jdk8 \
+                --javabase=//rbe-configs/bazel_{bazel_version}/config:jdk8 \
                 --host_java_toolchain=@bazel_tools//tools/jdk:toolchain_hostjdk8 \
                 --java_toolchain=@bazel_tools//tools/jdk:toolchain_hostjdk8 \
-                --extra_execution_platforms=/rbe-configs/bazel_{bazel_version}/platforms:rbe_ubuntu1604 \
-                --host_platform=/rbe-configs/bazel_{bazel_version}/platforms:rbe_ubuntu1604 \
-                --platforms=/rbe-configs/bazel_{bazel_version}/platforms:rbe_ubuntu1604 \
-                --extra_toolchains=/rbe-configs/bazel_{bazel_version}/platforms:cc-toolchain \
+                --extra_execution_platforms=/rbe-configs/bazel_{bazel_version}/config:platform \
+                --host_platform=/rbe-configs/bazel_{bazel_version}/config:platform \
+                --platforms=/rbe-configs/bazel_{bazel_version}/config:platform \
+                --extra_toolchains=/rbe-configs/bazel_{bazel_version}/config:cc-toolchain \
                 ... <other rbe flags> <build targets>
 
     We recommend you check in the code in //rbe-configs/bazel_{bazel_version}
-    so that users typically do not need to run this repo rule in order to do a
-    remote build (i.e., once files are checked in, you do not need to run this
-    rule until there is a new version of Bazel you want to support running with).
+    so that most devs/your CI typically do not need to run this repo rule
+    in order to do a remote build (i.e., once files are checked in,
+    you do not need to run this rule until there is a new version of Bazel
+    you want to support running with, or you need to update your container).
 
   2 - When output_base is not set (env var "RBE_AUTOCONF_ROOT" is not required),
     running this rule will create targets in the
@@ -94,14 +120,14 @@
 
       bazel build ... \
                 --crosstool_top=@rbe_default//rbe_config_cc:toolchain \
-                --host_javabase=@rbe_default//platforms:jdk8 \
-                --javabase=@rbe_default//platforms:jdk8 \
+                --host_javabase=@rbe_default//config:jdk8 \
+                --javabase=@rbe_default//config:jdk8 \
                 --host_java_toolchain=@bazel_tools//tools/jdk:toolchain_hostjdk8 \
                 --java_toolchain=@bazel_tools//tools/jdk:toolchain_hostjdk8 \
-                --extra_execution_platforms=@rbe_default//platforms:rbe_ubuntu1604 \
-                --host_platform=@rbe_default//platforms:rbe_ubuntu1604 \
-                --platforms=@rbe_default//platforms:rbe_ubuntu1604 \
-                --extra_toolchains=@rbe_default//platforms:cc-toolchain \
+                --extra_execution_platforms=@rbe_default//config:platform \
+                --host_platform=@rbe_default//config:platform \
+                --platforms=@rbe_default//config:platform \
+                --extra_toolchains=@rbe_default//config:cc-toolchain \
 
     Note running bazel clean --expunge_async, or otherwise modifying attrs or
     env variables used by this rule will trigger it to re-execute. Running this
@@ -127,6 +153,8 @@
   - docker
   - tar
   - bash utilities (e.g., cp, mv, rm, etc)
+  - docker authentication to pull the desired container should be set up
+    (rbe-ubuntu16-04 does not require any auth setup currently).
 
 Known limitations:
   - This rule cannot be executed inside a docker container.
@@ -151,7 +179,7 @@
 
 _BAZEL_CONFIG_DIR = "/bazel-config"
 _CONFIG_REPOS = ["local_config_cc"]
-_PLATFORM_DIR = "platforms"
+_PLATFORM_DIR = "config"
 _PROJECT_REPO_DIR = "project_src"
 _OUTPUT_DIR = _BAZEL_CONFIG_DIR + "/autoconf_out"
 _REPO_DIR = _BAZEL_CONFIG_DIR + "/" + _PROJECT_REPO_DIR
@@ -159,7 +187,17 @@
 _RBE_CONFIG_DIR = "rbe_config_cc"
 
 # We use 'l.gcr.io' to not require users to do gcloud login
-_RBE_UBUNTU_GCR = "l.gcr.io/google/rbe-ubuntu16-04@"
+_RBE_UBUNTU_REPO = "google/rbe-ubuntu16-04"
+_RBE_UBUNTU_REGISTRY = "l.gcr.io"
+_RBE_UBUNTU_EXEC_COMPAT_WITH = [
+    "@bazel_tools//platforms:x86_64",
+    "@bazel_tools//platforms:linux",
+    "@bazel_tools//tools/cpp:clang",
+]
+_RBE_UBUNTU_TARGET_COMPAT_WITH = [
+    "@bazel_tools//platforms:linux",
+    "@bazel_tools//platforms:x86_64",
+]
 _VERBOSE = False
 
 def _impl(ctx):
@@ -170,7 +208,7 @@
     project_root = ctx.os.environ.get(_RBE_AUTOCONF_ROOT, None)
     use_default_project = False
     if not project_root:
-        if ctx.attr.output_base != "":
+        if ctx.attr.output_base:
             fail(("%s env variable must be set for rbe_autoconfig" +
                   " to function properly when output_base is set") % _RBE_AUTOCONF_ROOT)
 
@@ -180,17 +218,18 @@
         # expected path.
         project_root = ctx.path(".").dirname.get_child("bazel_toolchains").get_child("rules").get_child("cc-sample-project")
         if not project_root.exists:
-            fail("Could not find default autoconf project in %s, please make sure " +
-                 "the bazel-toolchains repo is properly imported in your workspace")
+            fail(("Could not find default autoconf project in %s, please make sure " +
+                  "the bazel-toolchains repo is properly imported in your workspace") % str(project_root))
         project_root = str(project_root)
         use_default_project = True
 
     name = ctx.attr.name
     outputs_tar = ctx.attr.name + "_out.tar"
 
+    image_name = ctx.attr.registry + "/" + ctx.attr.repository + "@" + ctx.attr.digest
+
     # Pull the image using 'docker pull'
-    _pull_image(ctx)
-    image_id = _RBE_UBUNTU_GCR + ctx.attr.digest
+    _pull_image(ctx, image_name)
 
     bazel_version = None
     bazel_rc_version = None
@@ -207,7 +246,7 @@
         ctx,
         bazel_version = bazel_version,
         bazel_rc_version = bazel_rc_version,
-        image_id = image_id,
+        image_name = image_name,
         outputs_tar = outputs_tar,
         project_root = project_root,
         use_default_project = use_default_project,
@@ -215,7 +254,12 @@
 
     # Create a default BUILD file with the platform + toolchain targets that
     # will work with RBE with the produced toolchain
-    _create_platform(ctx, bazel_version = bazel_version, name = name)
+    _create_platform(
+        ctx,
+        bazel_version = bazel_version,
+        image_name = image_name,
+        name = name,
+    )
 
     # Expand outputs to project dir if user requested it
     _expand_outputs(
@@ -226,12 +270,12 @@
 
 # Convenience method to print results of execute (and fail on errors if needed).
 # Verbose logging is enabled via a global var in this bzl file.
-def _print_exec_results(prefix, exec_result, fail = False, args = None):
+def _print_exec_results(prefix, exec_result, fail_on_error = False, args = None):
     if _VERBOSE and exec_result.return_code != 0:
         print(prefix + "::error::" + exec_result.stderr)
     elif _VERBOSE:
         print(prefix + "::success::" + exec_result.stdout)
-    if fail and exec_result.return_code != 0:
+    if fail_on_error and exec_result.return_code != 0:
         if _VERBOSE and args:
             print("failed to run execute with the following args:" + str(args))
         fail("Failed to run:" + prefix + ":" + exec_result.stderr)
@@ -249,10 +293,10 @@
         fail("Cannot run rbe_autoconfig as 'tar' was not found on the path.")
 
 # Pulls an image using 'docker pull'.
-def _pull_image(ctx):
+def _pull_image(ctx, image_name):
     print("Pulling image.")
-    result = ctx.execute(["docker", "pull", _RBE_UBUNTU_GCR + ctx.attr.digest])
-    _print_exec_results("pull image", result, fail = True)
+    result = ctx.execute(["docker", "pull", image_name])
+    _print_exec_results("pull image", result, fail_on_error = True)
     print("Image pulled.")
 
 # Creates file "container/run_in_container.sh" which can be mounted onto container
@@ -281,6 +325,17 @@
     install_bazel_cmd += ["/tmp/bazel-installer.sh"]
     install_bazel_cmd += ["rm -f /tmp/bazel-installer.sh"]
 
+    # Command to recursively convert soft links to hard links in the config_repos
+    # Needed because some outputs of local_cc_config (e.g., dummy_toolchain.bzl)
+    # could be symlinks.
+    deref_symlinks_cmd = []
+    for config_repo in _CONFIG_REPOS:
+        symlinks_cmd = ("find $(bazel info output_base)/" +
+                        _EXTERNAL_FOLDER_PREFIX + config_repo +
+                        " -type l -exec bash -c 'ln -f \"$(readlink -m \"$0\")\" \"$0\"' {} \;")
+        deref_symlinks_cmd.append(symlinks_cmd)
+    deref_symlinks_cmd = " && ".join(deref_symlinks_cmd)
+
     # Command to copy produced toolchain configs to a tar at the root
     # of the container.
     copy_cmd = ["mkdir " + _OUTPUT_DIR]
@@ -321,6 +376,7 @@
     docker_cmd += setup_default_project_cmd
     docker_cmd += [
         bazel_cmd,
+        deref_symlinks_cmd,
         output_copy_cmd,
         clean_cmd,
     ]
@@ -332,7 +388,7 @@
         ctx,
         bazel_version,
         bazel_rc_version,
-        image_id,
+        image_name,
         outputs_tar,
         project_root,
         use_default_project):
@@ -367,7 +423,7 @@
         {
             "%{docker_run_flags}": " ".join(docker_run_flags),
             "%{commands}": "/container/run_in_container.sh",
-            "%{image_id}": image_id,
+            "%{image_name}": image_name,
             "%{extract_file}": "/" + outputs_tar,
             "%{output}": str(ctx.path(".")) + "/output.tar",
         },
@@ -390,19 +446,30 @@
     _print_exec_results("clean tools", result)
 
 # Creates a BUILD file with the java and cc toolchain + platform targets
-def _create_platform(ctx, bazel_version, name):
+def _create_platform(
+        ctx,
+        bazel_version,
+        image_name,
+        name):
     toolchain_target = "@" + name + "//" + _RBE_CONFIG_DIR
-    if ctx.attr.output_base != "":
+    if ctx.attr.output_base:
         toolchain_target = "//" + ctx.attr.output_base + "/bazel_" + bazel_version
-        if ctx.attr.config_dir != "":
+        if ctx.attr.config_dir:
             toolchain_target += ctx.attr.config_dir
     template = ctx.path(Label("@bazel_toolchains//rules:BUILD.platform.tpl"))
+    exec_compatible_with = ("\"" +
+                            ("\",\n        \"").join(ctx.attr.exec_compatible_with) +
+                            "\",")
+    target_compatible_with = ("\"" +
+                              ("\",\n        \"").join(ctx.attr.target_compatible_with) +
+                              "\",")
     ctx.template(
-        "platforms/BUILD",
+        _PLATFORM_DIR + "/BUILD",
         template,
         {
-            "%{revision}": ctx.attr.revision,
-            "%{rbe_ubuntu16_04_sha256}": ctx.attr.digest,
+            "%{exec_compatible_with}": exec_compatible_with,
+            "%{image_name}": image_name,
+            "%{target_compatible_with}": target_compatible_with,
             "%{toolchain}": toolchain_target,
         },
         False,
@@ -411,10 +478,10 @@
 # Copies all outputs of the autoconfig rule to a directory in the project
 # sources
 def _expand_outputs(ctx, bazel_version, project_root):
-    if ctx.attr.output_base != "":
+    if ctx.attr.output_base:
         print("Copying outputs to project directory")
         dest = project_root + "/" + ctx.attr.output_base + "/bazel_" + bazel_version + "/"
-        if ctx.attr.config_dir != "":
+        if ctx.attr.config_dir:
             dest += ctx.attr.config_dir + "/"
         platform_dest = dest + _PLATFORM_DIR + "/"
 
@@ -433,8 +500,8 @@
         result = ctx.execute(args)
         _print_exec_results("copy outputs", result, True, args)
 
-        # Copy the dest/platforms/BUILD file
-        result = ctx.execute("cp", str(ctx.path("platforms/BUILD")), platform_dest)
+        # Copy the dest/{_PLATFORM_DIR}/BUILD file
+        result = ctx.execute("cp", str(ctx.path(_PLATFORM_DIR + "/BUILD")), platform_dest)
 
 # Private declaration of _rbe_autoconfig repository rule. Do not use this
 # rule directly, use rbe_autoconfig macro declared below.
@@ -456,8 +523,9 @@
         ),
         "digest": attr.string(
             mandatory = True,
-            doc = ("The digest (sha256 sum) of the rbe-ubuntu16-04 " +
-                   "container to pull."),
+            doc = ("The digest (sha256 sum) of the image to pull. For example, " +
+                   "sha256:f1330b2f02714d3a3e98c5b1f6524fbb9c15154e44a31fb3caecb7a6ad4e8445" +
+                   ", note the digest includes 'sha256:'"),
         ),
         "env": attr.string_dict(
             doc = ("Optional. Dictionary from strings to strings. Additional env " +
@@ -470,12 +538,10 @@
                    "be used when generating the toolchain configs."),
         ),
         "output_base": attr.string(
-            default = "",
             doc = ("Optional. The directory (under the project root) where the " +
                    "produced toolchain configs will be copied to."),
         ),
         "config_dir": attr.string(
-            default = "",
             doc = ("Optional. Use only if output_base is defined. If you want to " +
                    "create multiple toolchain configs (for the same version of Bazel) " +
                    "you can use this attr to indicate a type of config (e.g., default, " +
@@ -488,9 +554,35 @@
                    "(inside the container) before running bazel to generate the " +
                    "toolchain configs"),
         ),
+        "registry": attr.string(
+            default = _RBE_UBUNTU_REGISTRY,
+            doc = ("The registry to pull the container from. For example, " +
+                   "l.gcr.io or marketplace.gcr.io. The default is the " +
+                   "value for rbe-ubuntu16-04 image."),
+        ),
+        "repository": attr.string(
+            default = _RBE_UBUNTU_REPO,
+            doc = ("The repository to pull the container from. For example," +
+                   " google/ubuntu. The default is the " +
+                   " value for the rbe-ubuntu16-04 image."),
+        ),
         "revision": attr.string(
-            mandatory = True,
-            doc = ("The revision of the rbe-ubuntu16-04 container"),
+            doc = ("The revision of the rbe-ubuntu16-04 container."),
+        ),
+        "exec_compatible_with": attr.string_list(
+            default = _RBE_UBUNTU_EXEC_COMPAT_WITH,
+            doc = ("The list of constraints that will be added to the " +
+                   "toolchain in its exec_compatible_with attribute (and to " +
+                   "the platform in its constraint_values attr). For " +
+                   "example, [\"@bazel_tools//platforms:linux\"]. Default " +
+                   " is set to values for rbe-ubuntu16-04 container."),
+        ),
+        "target_compatible_with": attr.string_list(
+            default = _RBE_UBUNTU_TARGET_COMPAT_WITH,
+            doc = ("The list of constraints that will be added to the " +
+                   "toolchain in its target_compatible_with attribute. For " +
+                   "example, [\"@bazel_tools//platforms:linux\"]. Default " +
+                   " is set to values for rbe-ubuntu16-04 container."),
         ),
     },
     environ = [
@@ -505,11 +597,16 @@
         name,
         bazel_version = None,
         bazel_rc = None,
-        output_base = "",
-        config_dir = "",
-        revision = "latest",
-        env = clang_env()):
-    """ Creates a repository with toolchain configs generated for an rbe-ubuntu container.
+        config_dir = None,
+        digest = None,
+        env = clang_env(),
+        exec_compatible_with = None,
+        output_base = None,
+        revision = None,
+        registry = None,
+        repository = None,
+        target_compatible_with = None):
+    """ Creates a repository with toolchain configs generated for a container image.
 
     This macro wraps (and simplifies) invocation of _rbe_autoconfig rule.
     Use this macro in your WORKSPACE.
@@ -523,22 +620,51 @@
           in attrs of _rbe_autoconfig for current latest).
       bazel_rc: The rc (for the given version of Bazel) to use. Must be published
           in https://releases.bazel.build
+      exec_compatible_with: Optional. List of constraints to add to the produced
+          toolchain/platform targets (e.g., ["@bazel_tools//platforms:linux"] in the
+          exec_compatible_with/constraint_values attrs, respectively.
+      digest: Optional. The digest of the image to pull. Should only be set if
+          a custom container is required.
+          Must be set together with registry and repository.
       output_base: Optional. The directory (under the project root) where the
           produced toolchain configs will be copied to.
       config_dir: Optional. Subdirectory where configs will be copied to.
           Use only if output_base is defined.
-      revision: a revision of an rbe-ubuntu16-04 container.
+      registry: Optional. The registry from which to pull the base image.
+          Should only be set if a custom container is required.
+          Must be set together with digest and repository.
+      repository: Optional. he `repository` of images to pull from.
+          Should only be set if a custom container is required.
+          Must be set together with registry and digest.
+      revision: Optional. A revision of an rbe-ubuntu16-04 container to use.
+          Should not be set if repository, registry and digest are used.
           See gcr.io/cloud-marketplace/google/rbe-ubuntu16-04
-      env: dict. Additional env variables that will be set when running the
-          Bazel command to generate the toolchain configs.
+      target_compatible_with: List of constraints to add to the produced
+          toolchain target (e.g., ["@bazel_tools//platforms:linux"]) in the
+          target_compatible_with attr.
+      env: dict. Optional. Additional env variables that will be set when
+          running the Bazel command to generate the toolchain configs.
+          Set to values for rbe-ubuntu16-04 container.
+          Does not need to be set if your custom container extends
+          an rbe-ubuntu16-04 container.
+          Should be overriden if a custom container does not extend
+          rbe-ubuntu16-04.
     """
-    if output_base == "" and config_dir != "":
+    if not output_base and config_dir:
         fail("config_dir can only be used when output_base is set.")
-    if revision == "latest":
-        revision = RBE_UBUNTU16_04_LATEST
+    if revision and (digest or repository or registry):
+        fail("'revision' cannot be set if 'digest', 'repository' or " +
+             "'registry' are set.")
+    if not ((not digest and not repository and not registry) or
+            (digest and repository and registry)):
+        fail("All of 'digest', 'repository' and 'registry' or none of them " +
+             "must be set.")
     if bazel_rc and not bazel_version:
         fail("bazel_rc can only be used with bazel_version.")
-    digest = public_rbe_ubuntu16_04_sha256s().get(revision, None)
+    if not digest:
+        if not revision or revision == "latest":
+            revision = RBE_UBUNTU16_04_LATEST
+        digest = public_rbe_ubuntu16_04_sha256s().get(revision, None)
     if not digest:
         fail(("Could not find a valid digest for revision %s, " +
               "please make sure it is declared in " +
@@ -550,6 +676,10 @@
         config_dir = config_dir,
         digest = digest,
         env = env,
+        exec_compatible_with = exec_compatible_with,
         output_base = output_base,
+        registry = registry,
+        repository = repository,
         revision = revision,
+        target_compatible_with = target_compatible_with,
     )