From c28968d2c57e8c343cc0c5b7930ae2ec60304729 Mon Sep 17 00:00:00 2001
From: David Sondermann <david.sondermann@hivemq.com>
Date: Thu, 3 Nov 2022 13:27:37 +0100
Subject: [PATCH] Add replacement of build args in Dockerfile FROM statements

---
 .../images/builder/ImageFromDockerfile.java   | 22 ++++++++++++++++---
 .../images/builder/DockerfileBuildTest.java   | 22 +++++++++++++++++--
 .../Dockerfile-from-buildarg                  |  9 ++++++++
 3 files changed, 48 insertions(+), 5 deletions(-)
 create mode 100644 core/src/test/resources/dockerfile-build-test/Dockerfile-from-buildarg

diff --git a/core/src/main/java/org/testcontainers/images/builder/ImageFromDockerfile.java b/core/src/main/java/org/testcontainers/images/builder/ImageFromDockerfile.java
index c453c24a779..35ab83039ae 100644
--- a/core/src/main/java/org/testcontainers/images/builder/ImageFromDockerfile.java
+++ b/core/src/main/java/org/testcontainers/images/builder/ImageFromDockerfile.java
@@ -35,6 +35,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.regex.Matcher;
 import java.util.zip.GZIPOutputStream;
 
 @Slf4j
@@ -182,24 +183,39 @@ protected void configure(BuildImageCmd buildImageCmd) {
 
     private void prePullDependencyImages(Set<String> imagesToPull) {
         imagesToPull.forEach(imageName -> {
+            String resolvedImageName = applyBuildArgsToImageName(imageName);
             try {
                 log.info(
                     "Pre-emptively checking local images for '{}', referenced via a Dockerfile. If not available, it will be pulled.",
-                    imageName
+                    resolvedImageName
                 );
-                new RemoteDockerImage(DockerImageName.parse(imageName))
+                new RemoteDockerImage(DockerImageName.parse(resolvedImageName))
                     .withImageNameSubstitutor(ImageNameSubstitutor.noop())
                     .get();
             } catch (Exception e) {
                 log.warn(
                     "Unable to pre-fetch an image ({}) depended upon by Dockerfile - image build will continue but may fail. Exception message was: {}",
-                    imageName,
+                    resolvedImageName,
                     e.getMessage()
                 );
             }
         });
     }
 
+    /**
+     * See {@code filterForEnvironmentVars()} in {@link com.github.dockerjava.core.dockerfile.DockerfileStatement}.
+     */
+    private String applyBuildArgsToImageName(String imageName) {
+        for (Map.Entry<String, String> entry : buildArgs.entrySet()) {
+            String value = Matcher.quoteReplacement(entry.getValue());
+            // handle: $VARIABLE case
+            imageName = imageName.replace("$" + entry.getKey(), value);
+            // handle ${VARIABLE} case
+            imageName = imageName.replace("${" + entry.getKey() + "}", value);
+        }
+        return imageName;
+    }
+
     public ImageFromDockerfile withBuildArg(final String key, final String value) {
         this.buildArgs.put(key, value);
         return this;
diff --git a/core/src/test/java/org/testcontainers/images/builder/DockerfileBuildTest.java b/core/src/test/java/org/testcontainers/images/builder/DockerfileBuildTest.java
index c45233ea390..10bf2e3d11c 100644
--- a/core/src/test/java/org/testcontainers/images/builder/DockerfileBuildTest.java
+++ b/core/src/test/java/org/testcontainers/images/builder/DockerfileBuildTest.java
@@ -8,6 +8,8 @@
 
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.HashMap;
+import java.util.Map;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -22,6 +24,13 @@ public class DockerfileBuildTest {
 
     @Parameterized.Parameters
     public static Object[][] parameters() {
+        Map<String, String> buildArgs = new HashMap<>(4);
+        buildArgs.put("BUILD_IMAGE", "alpine:3.16");
+        buildArgs.put("BASE_IMAGE", "alpine");
+        buildArgs.put("BASE_IMAGE_TAG", "3.12");
+        buildArgs.put("UNUSED", "ignored");
+
+        //noinspection deprecation
         return new Object[][] {
             // Dockerfile build without explicit per-file inclusion
             new Object[] {
@@ -38,7 +47,7 @@ public static Object[][] parameters() {
                 "test4567",
                 new ImageFromDockerfile().withFileFromPath(".", RESOURCE_PATH).withDockerfilePath("./Dockerfile-alt"),
             },
-            // Dockerfile build using build args
+            // Dockerfile build using withBuildArg()
             new Object[] {
                 "test7890",
                 new ImageFromDockerfile()
@@ -46,6 +55,14 @@ public static Object[][] parameters() {
                     .withDockerfilePath("./Dockerfile-buildarg")
                     .withBuildArg("CUSTOM_ARG", "test7890"),
             },
+            // Dockerfile build using withBuildArgs() with build args in FROM statement
+            new Object[] {
+                "test1234",
+                new ImageFromDockerfile()
+                    .withFileFromPath(".", RESOURCE_PATH)
+                    .withDockerfile(RESOURCE_PATH.resolve("Dockerfile-from-buildarg"))
+                    .withBuildArgs(buildArgs),
+            },
             // Dockerfile build using withDockerfile(File)
             new Object[] {
                 "test4567",
@@ -62,9 +79,10 @@ public DockerfileBuildTest(String expectedFileContent, ImageFromDockerfile image
     }
 
     @Test
+    @SuppressWarnings("resource")
     public void performTest() {
         try (
-            final GenericContainer container = new GenericContainer(image)
+            final GenericContainer<?> container = new GenericContainer<>(image)
                 .withStartupCheckStrategy(new OneShotStartupCheckStrategy())
                 .withCommand("cat", "/test.txt")
         ) {
diff --git a/core/src/test/resources/dockerfile-build-test/Dockerfile-from-buildarg b/core/src/test/resources/dockerfile-build-test/Dockerfile-from-buildarg
new file mode 100644
index 00000000000..16934b374e8
--- /dev/null
+++ b/core/src/test/resources/dockerfile-build-test/Dockerfile-from-buildarg
@@ -0,0 +1,9 @@
+ARG BUILD_IMAGE
+ARG BASE_IMAGE
+ARG BASE_IMAGE_TAG
+
+FROM ${BUILD_IMAGE} AS build
+COPY localfile.txt /test-build.txt
+
+FROM $BASE_IMAGE:${BASE_IMAGE_TAG} AS base
+COPY --from=build /test-build.txt /test.txt