Upgrade ARG default values used in Docker FROM instructions

src/main/java/org/openrewrite/java/migrate/UpgradeDockerImageVersion.java

@@ -17,18 +17,29 @@
 
 import lombok.EqualsAndHashCode;
 import lombok.Value;
+import org.jspecify.annotations.Nullable;
 import org.openrewrite.ExecutionContext;
 import org.openrewrite.Option;
 import org.openrewrite.Recipe;
 import org.openrewrite.TreeVisitor;
+import org.openrewrite.docker.DockerIsoVisitor;
 import org.openrewrite.docker.trait.DockerFrom;
+import org.openrewrite.docker.tree.Docker;
+import org.openrewrite.internal.ListUtils;
 
+import java.util.HashMap;
 import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
 import java.util.Set;
+import java.util.UUID;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
 import static java.util.Arrays.asList;
+import static java.util.Collections.emptyMap;
+import static java.util.Collections.singletonList;
+import static java.util.Objects.requireNonNull;
 
 @EqualsAndHashCode(callSuper = false)
 @Value
@@ -50,48 +61,284 @@ public class UpgradeDockerImageVersion extends Recipe {
     private static final int OLDEST_VERSION = 8;
     private static final Pattern VERSIONED_TAG = Pattern.compile("(\\d{1,3})(\\D.*)?");
 
+    private static final String FROM_REPLACEMENTS = "fromReplacements";
+
     String displayName = "Upgrade Docker image Java version";
     String description = "Upgrade Docker image tags to use the specified Java version. " +
             "Updates common Java Docker images including eclipse-temurin, amazoncorretto, azul/zulu-openjdk, " +
             "and others. Also migrates deprecated images (openjdk, adoptopenjdk) to eclipse-temurin, " +
-            "preserving any tag suffix such as `-jre-alpine`. Image references built from build arguments or " +
-            "environment variables are left untouched, as their value can not be determined statically. A digest " +
-            "pin is dropped when the tag is upgraded, as the stale digest would otherwise keep resolving to the " +
-            "old image.";
+            "preserving any tag suffix such as `-jre-alpine`. When a `FROM` is built from a build argument, the " +
+            "default value of the corresponding global `ARG` is upgraded instead, such that `ARG java_version=17` " +
+            "used as `FROM eclipse-temurin:${java_version}` becomes `ARG java_version=25`. Image references built " +
+            "from arguments without a default value are left untouched, as their value can not be determined " +
+            "statically. A digest pin is dropped when the tag is upgraded, as the stale digest would otherwise " +
+            "keep resolving to the old image.";
 
     @Override
     public TreeVisitor<?, ExecutionContext> getVisitor() {
         if (version == null) {
             return TreeVisitor.noop();
         }
-        return new DockerFrom.Matcher().asVisitor((image, ctx) -> {
-            String imageName = image.getImageName().orElse("");
-            String tag = image.getTag().orElse("");
-            if (containsVariable(imageName) || containsVariable(tag)) {
-                return image.getTree();
+        return new DockerIsoVisitor<ExecutionContext>() {
+
+            @Override
+            public Docker.File visitFile(Docker.File file, ExecutionContext ctx) {
+                Map<String, String> defaults = new HashMap<>();
+                for (Docker.Arg arg : file.getGlobalArgs()) {
+                    Docker.Literal literalDefault = literalDefault(arg);
+                    if (literalDefault != null) {
+                        defaults.put(arg.getName().getText(), QuotedText.of(literalDefault.getText()).getText());
+                    }
+                }
+
+                ArgPlan plan = planUpgrades(file, defaults);
+                getCursor().putMessage(FROM_REPLACEMENTS, plan.getFromReplacements());
+                Docker.File f = super.visitFile(file, ctx);
+
+                Map<String, String> upgrades = plan.getArgUpgrades();
+                if (upgrades.isEmpty()) {
+                    return f;
+                }
+                return f.withGlobalArgs(ListUtils.map(f.getGlobalArgs(), arg -> {
+                    String upgraded = upgrades.get(arg.getName().getText());
+                    Docker.Literal literalDefault = literalDefault(arg);
+                    if (upgraded == null || literalDefault == null) {
+                        return arg;
+                    }
+                    String requoted = QuotedText.of(literalDefault.getText()).requote(upgraded);
+                    return arg.withValue(requireNonNull(arg.getValue())
+                            .withContents(singletonList(literalDefault.withText(requoted))));
+                }));
             }
 
-            Matcher matcher = VERSIONED_TAG.matcher(tag);
-            if (!matcher.matches()) {
-                return image.getTree();
+            @Override
+            public Docker.From visitFrom(Docker.From from, ExecutionContext ctx) {
+                if (containsVariable(from.getImageName()) || containsVariable(from.getTag())) {
+                    Map<UUID, Docker.From> replacements = getCursor().getNearestMessage(FROM_REPLACEMENTS, emptyMap());
+                    return replacements.getOrDefault(from.getId(), from);
+                }
+
+                DockerFrom image = new DockerFrom(getCursor());
+                String newTag = upgradedTag(image.getTag().orElse(""));
+                if (newTag == null) {
+                    return from;
+                }
+                String imageName = image.getImageName().orElse("");
+                if (DEPRECATED_IMAGES.contains(imageName)) {
+                    return image.withImageReference(NEW_IMAGE + ":" + newTag);
+                }
+                if (CURRENT_IMAGES.contains(imageName)) {
+                    return image.withTag(newTag).withDigest(null);
+                }
+                return from;
             }
-            int currentVersion = Integer.parseInt(matcher.group(1));
-            if (currentVersion < OLDEST_VERSION || version <= currentVersion) {
-                return image.getTree();
+        };
+    }
+
+    /**
+     * Withholding an argument can rule out the images that depend on it, which in turn can withhold further arguments,
+     * so the whole file is planned over and over until the set of withheld arguments stops growing. Only then is it
+     * known which `FROM` instructions may be rewritten, as an image may not be moved to a tag that is never written.
+     */
+    private ArgPlan planUpgrades(Docker.File file, Map<String, String> defaults) {
+        Map<String, String> upgrades = new HashMap<>();
+        Map<UUID, Docker.From> replacements = new HashMap<>();
+        Set<String> blocked = new HashSet<>();
+        int blockedCount;
+        do {
+            blockedCount = blocked.size();
+            upgrades.clear();
+            replacements.clear();
+            for (Docker.Stage stage : file.getStages()) {
+                Docker.From from = stage.getFrom();
+                if (containsVariable(from.getImageName()) || containsVariable(from.getTag())) {
+                    Docker.From planned = planFrom(from, defaults, upgrades, blocked);
+                    if (planned != from) {
+                        replacements.put(from.getId(), planned);
+                    }
+                }
             }
+        } while (blockedCount < blocked.size());
+        return new ArgPlan(upgrades, replacements);
+    }
 
-            String newTag = version + (matcher.group(2) == null ? "" : matcher.group(2));
-            if (DEPRECATED_IMAGES.contains(imageName)) {
-                return image.withImageReference(NEW_IMAGE + ":" + newTag);
+    private Docker.From planFrom(Docker.From from, Map<String, String> defaults, Map<String, String> upgrades, Set<String> blocked) {
+        String imageVariable = soleVariable(from.getImageName());
+        String tagVariable = from.getTag() == null ? null : leadingVariable(from.getTag());
+        String imageName = imageVariable == null ?
+                literalText(from.getImageName()) :
+                defaultValue(defaults, blocked, imageVariable);
+        if (imageName == null) {
+            return block(from, blocked, imageVariable, tagVariable);
+        }
+
+        String tag;
+        boolean wholeReference = from.getTag() == null;
+        if (wholeReference) {
+            // A single argument holding the whole reference, as in `FROM ${BASE_IMAGE}`
+            String[] reference = imageVariable == null ? null : splitReference(imageName);
+            if (reference == null) {
+                return block(from, blocked, imageVariable, null);
             }
-            if (CURRENT_IMAGES.contains(imageName)) {
-                return image.withTag(newTag).withDigest(null);
+            imageName = reference[0];
+            tag = reference[1];
+            tagVariable = imageVariable;
+        } else {
+            tag = tagVariable == null ? literalText(from.getTag()) : defaultValue(defaults, blocked, tagVariable);
+        }
+        if (tag == null) {
+            return block(from, blocked, imageVariable, tagVariable);
+        }
+
+        String newImageName = upgradedImageName(imageName);
+        String newTag = upgradedTag(tag);
+        if (newImageName == null || newTag == null) {
+            return block(from, blocked, imageVariable, tagVariable);
+        }
+
+        if (wholeReference) {
+            upgrades.put(requireNonNull(imageVariable), newImageName + ":" + newTag);
+            return from.withDigest(null);
+        }
+        if (tagVariable == null) {
+            from = from.withTag(withText(requireNonNull(from.getTag()), newTag));
+        } else {
+            upgrades.put(tagVariable, newTag);
+        }
+        if (!newImageName.equals(imageName)) {
+            if (imageVariable == null) {
+                from = from.withImageName(withText(from.getImageName(), newImageName));
+            } else {
+                upgrades.put(imageVariable, newImageName);
             }
-            return image.getTree();
-        });
+        }
+        return from.withDigest(null);
+    }
+
+    private @Nullable String upgradedImageName(String imageName) {
+        if (DEPRECATED_IMAGES.contains(imageName)) {
+            return NEW_IMAGE;
+        }
+        return CURRENT_IMAGES.contains(imageName) ? imageName : null;
+    }
+
+    private @Nullable String upgradedTag(String tag) {
+        Matcher matcher = VERSIONED_TAG.matcher(tag);
+        if (!matcher.matches()) {
+            return null;
+        }
+        int currentVersion = Integer.parseInt(matcher.group(1));
+        if (currentVersion < OLDEST_VERSION || version <= currentVersion) {
+            return null;
+        }
+        return version + (matcher.group(2) == null ? "" : matcher.group(2));
+    }
+
+    /**
+     * Withhold arguments feeding an image we leave alone, as a shared argument can not be bumped for one image alone.
+     */
+    private static Docker.From block(Docker.From from, Set<String> blocked, @Nullable String imageVariable, @Nullable String tagVariable) {
+        if (imageVariable != null) {
+            blocked.add(imageVariable);
+        }
+        if (tagVariable != null) {
+            blocked.add(tagVariable);
+        }
+        return from;
+    }
+
+    private static @Nullable String defaultValue(Map<String, String> defaults, Set<String> blocked, String variable) {
+        return blocked.contains(variable) ? null : defaults.get(variable);
     }
 
-    private static boolean containsVariable(String imageReferencePart) {
-        return imageReferencePart.indexOf('$') != -1;
+    private static boolean containsVariable(Docker.@Nullable Argument argument) {
+        if (argument != null) {
+            for (Docker.ArgumentContent content : argument.getContents()) {
+                if (content instanceof Docker.EnvironmentVariable) {
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
+
+    private static Docker.@Nullable Literal literalDefault(Docker.Arg arg) {
+        Docker.Argument value = arg.getValue();
+        return value == null ? null : sole(value.getContents(), Docker.Literal.class);
+    }
+
+    private static @Nullable String literalText(Docker.Argument argument) {
+        Docker.Literal literal = sole(argument.getContents(), Docker.Literal.class);
+        return literal == null ? null : literal.getText();
+    }
+
+    private static @Nullable String soleVariable(Docker.Argument argument) {
+        Docker.EnvironmentVariable variable = sole(argument.getContents(), Docker.EnvironmentVariable.class);
+        return variable == null ? null : variable.getName();
+    }
+
+    private static @Nullable String leadingVariable(Docker.Argument argument) {
+        List<Docker.ArgumentContent> contents = argument.getContents();
+        if (contents.isEmpty() || !(contents.get(0) instanceof Docker.EnvironmentVariable)) {
+            return null;
+        }
+        for (int i = 1; i < contents.size(); i++) {
+            if (!(contents.get(i) instanceof Docker.Literal)) {
+                return null;
+            }
+        }
+        return ((Docker.EnvironmentVariable) contents.get(0)).getName();
+    }
+
+    private static <T> @Nullable T sole(List<? extends Docker.ArgumentContent> contents, Class<T> type) {
+        if (contents.size() == 1 && type.isInstance(contents.get(0))) {
+            return type.cast(contents.get(0));
+        }
+        return null;
+    }
+
+    private static String @Nullable [] splitReference(String reference) {
+        int at = reference.indexOf('@');
+        String withoutDigest = at == -1 ? reference : reference.substring(0, at);
+        int colon = withoutDigest.indexOf(':', withoutDigest.lastIndexOf('/') + 1);
+        if (colon == -1) {
+            return null;
+        }
+        return new String[]{withoutDigest.substring(0, colon), withoutDigest.substring(colon + 1)};
+    }
+
+    private static Docker.Argument withText(Docker.Argument argument, String text) {
+        Docker.Literal literal = requireNonNull(sole(argument.getContents(), Docker.Literal.class));
+        return argument.withContents(singletonList(literal.withText(text)));
+    }
+
+    @Value
+    private static class ArgPlan {
+        Map<String, String> argUpgrades;
+        Map<UUID, Docker.From> fromReplacements;
+    }
+
+    /**
+     * The parser keeps any quotes around an {@code ARG} default value as part of the literal text, so they have to be
+     * taken off before matching a version, and put back on when writing the upgraded value.
+     */
+    @Value
+    private static class QuotedText {
+        String quote;
+        String text;
+
+        static QuotedText of(String source) {
+            if (source.length() > 1) {
+                char first = source.charAt(0);
+                if ((first == '"' || first == '\'') && source.charAt(source.length() - 1) == first) {
+                    return new QuotedText(String.valueOf(first), source.substring(1, source.length() - 1));
+                }
+            }
+            return new QuotedText("", source);
+        }
+
+        String requote(String replacement) {
+            return quote + replacement + quote;
+        }
     }
 }