[v1.x] For ECR, ensure we sanitize region input from environment vari…

…able (#19882) * Set default for cache_intermediate. * Make sure we sanitize region extracted from registry, since we pass it to os.system. Co-authored-by: Joe Evans <[email protected]>
apache · Feb 11, 2021 · 0138c29 · 0138c29
1 parent 9f3da90
commit 0138c29
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 12 deletions.
diff --git a/ci/build.py b/ci/build.py
@@ -97,7 +97,7 @@ def get_dockerfile(platform: str, path=get_dockerfiles_path()) -> str:
 
 
 def build_docker(platform: str, registry: str, num_retries: int, no_cache: bool,
-                 cache_intermediate: bool) -> str:
+                 cache_intermediate: bool=False) -> str:
     """
     Build a container for the given platform
     :param platform: Platform

diff --git a/ci/docker_cache.py b/ci/docker_cache.py
@@ -28,6 +28,7 @@
 import logging
 import os
 import subprocess
+import re
 import sys
 from typing import *
 
@@ -96,24 +97,30 @@ def _build_save_container(platform, registry, load_cache) -> Optional[str]:
         # Error handling is done by returning the errorous platform name. This is necessary due to
         # Parallel being unable to handle exceptions
 
+ECR_LOGGED_IN = False
 def _ecr_login(registry):
     """
     Use the AWS CLI to get credentials to login to ECR.
     """
     # extract region from registry
-    region = registry.split(".")[3]
+    global ECR_LOGGED_IN
+    if ECR_LOGGED_IN:
+        return
+    regionMatch = re.match(r'.*?\.dkr\.ecr\.([a-z]+\-[a-z]+\-\d+)\.amazonaws\.com', registry)
+    assert(regionMatch)
+    region = regionMatch.group(1)
     logging.info("Logging into ECR region %s using aws-cli..", region)
     os.system("$(aws ecr get-login --region "+region+" --no-include-email)")
+    ECR_LOGGED_IN = True
 
 def _upload_image(registry, docker_tag, image_id) -> None:
     """
-    Upload the passed image by id, tag it with docker tag and upload to S3 bucket
+    Upload the passed image by id, tag it with docker tag and upload to docker registry.
     :param registry: Docker registry name
     :param docker_tag: Docker tag
     :param image_id: Image id
     :return: None
     """
-
     if "dkr.ecr" in registry:
         _ecr_login(registry)
 
@@ -200,15 +207,19 @@ def script_name() -> str:
 
     platforms = build_util.get_platforms()
 
-    secret_name = os.environ['DOCKERHUB_SECRET_NAME']
-    endpoint_url = os.environ['DOCKERHUB_SECRET_ENDPOINT_URL']
-    region_name = os.environ['DOCKERHUB_SECRET_ENDPOINT_REGION']
-
-    try:
-        login_dockerhub(secret_name, endpoint_url, region_name)
+    if "dkr.ecr" in args.docker_registry:
+        _ecr_login(args.docker_registry)
         return build_save_containers(platforms=platforms, registry=args.docker_registry, load_cache=True)
-    finally:
-        logout_dockerhub()
+    else:
+        secret_name = os.environ['DOCKERHUB_SECRET_NAME']
+        endpoint_url = os.environ['DOCKERHUB_SECRET_ENDPOINT_URL']
+        region_name = os.environ['DOCKERHUB_SECRET_ENDPOINT_REGION']
+
+        try:
+            login_dockerhub(secret_name, endpoint_url, region_name)
+            return build_save_containers(platforms=platforms, registry=args.docker_registry, load_cache=True)
+        finally:
+            logout_dockerhub()
 
 
 if __name__ == '__main__':