YARN-11364. Docker Container to accept docker Image name with sha256 digest (#5092)

Co-authored-by: Ashutosh Gupta <ashugpt@amazon.com>
Reviewed-by: slfan1989 <55643692+slfan1989@users.noreply.github.com>
Signed-off-by: Chris Nauroth <cnauroth@apache.org>
(cherry picked from commit 83acb559817a97c14c4e3fd846dcc16ab615093e)

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/main/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/linux/runtime/DockerLinuxContainerRuntime.java

@@ -206,6 +206,8 @@ public class DockerLinuxContainerRuntime extends OCIContainerRuntime {
   private static final Pattern dockerImagePattern =
       Pattern.compile(DOCKER_IMAGE_PATTERN);
 
+  private static final Pattern DOCKER_DIGEST_PATTERN = Pattern.compile("^sha256:[a-z0-9]{12,64}$");
+
   private static final String DEFAULT_PROCFS = "/proc";
 
   @InterfaceAudience.Private
@@ -1178,9 +1180,17 @@ public class DockerLinuxContainerRuntime extends OCIContainerRuntime {
       throw new ContainerExecutionException(
           ENV_DOCKER_CONTAINER_IMAGE + " not set!");
     }
-    if (!dockerImagePattern.matcher(imageName).matches()) {
-      throw new ContainerExecutionException("Image name '" + imageName
-          + "' doesn't match docker image name pattern");
+    // check if digest is part of imageName, extract and validate it.
+    String digest = null;
+    if (imageName.contains("@sha256")) {
+      String[] digestParts = imageName.split("@");
+      digest = digestParts[1];
+      imageName = digestParts[0];
+    }
+    if (!dockerImagePattern.matcher(imageName).matches() || (digest != null
+        && !DOCKER_DIGEST_PATTERN.matcher(digest).matches())) {
+      throw new ContainerExecutionException(
+          "Image name '" + imageName + "' doesn't match docker image name pattern");
     }
   }
 

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/linux/runtime/TestDockerContainerRuntime.java

@@ -2034,19 +2034,27 @@ public class TestDockerContainerRuntime {
 
   @Test
   public void testDockerImageNamePattern() throws Exception {
-    String[] validNames =
-        { "ubuntu", "fedora/httpd:version1.0",
-            "fedora/httpd:version1.0.test",
-            "fedora/httpd:version1.0.TEST",
-            "myregistryhost:5000/ubuntu",
-            "myregistryhost:5000/fedora/httpd:version1.0",
-            "myregistryhost:5000/fedora/httpd:version1.0.test",
-            "myregistryhost:5000/fedora/httpd:version1.0.TEST"};
-
-    String[] invalidNames = { "Ubuntu", "ubuntu || fedora", "ubuntu#",
-        "myregistryhost:50AB0/ubuntu", "myregistry#host:50AB0/ubuntu",
-        ":8080/ubuntu"
-    };
+    String[] validNames = {"ubuntu", "fedora/httpd:version1.0", "fedora/httpd:version1.0.test",
+        "fedora/httpd:version1.0.TEST", "myregistryhost:5000/ubuntu",
+        "myregistryhost:5000/fedora/httpd:version1.0",
+        "myregistryhost:5000/fedora/httpd:version1.0.test",
+        "myregistryhost:5000/fedora/httpd:version1.0.TEST",
+        "123456789123.dkr.ecr.us-east-1.amazonaws.com/emr-docker-examples:pyspark-example"
+            + "@sha256:f1d4ae3f7261a72e98c6ebefe9985cf10a0ea5bd762585a43e0700ed99863807"};
+
+    String[] invalidNames = {"Ubuntu", "ubuntu || fedora", "ubuntu#", "myregistryhost:50AB0/ubuntu",
+        "myregistry#host:50AB0/ubuntu", ":8080/ubuntu",
+
+        // Invalid: contains "@sha256" but doesn't really contain a digest.
+        "123456789123.dkr.ecr.us-east-1.amazonaws.com/emr-docker-examples:pyspark-example@sha256",
+
+        // Invalid: digest is too short.
+        "123456789123.dkr.ecr.us-east-1.amazonaws.com/emr-docker-examples:pyspark-example"
+            + "@sha256:f1d4",
+
+        // Invalid: digest is too long
+        "123456789123.dkr.ecr.us-east-1.amazonaws.com/emr-docker-examples:pyspark-example"
+            + "@sha256:f1d4ae3f7261a72e98c6ebefe9985cf10a0ea5bd762585a43e0700ed99863807f"};
 
     for (String name : validNames) {
       DockerLinuxContainerRuntime.validateImageName(name);