Improve FindDockerImageUses recipe with multistage files and gitlab-ci files

build.gradle.kts

@@ -10,5 +10,6 @@ val rewriteVersion = rewriteRecipe.rewriteVersion.get()
 dependencies {
     implementation(platform("org.openrewrite:rewrite-bom:$rewriteVersion"))
     implementation("org.openrewrite:rewrite-core")
+    implementation("org.openrewrite:rewrite-yaml")
     testImplementation("org.openrewrite:rewrite-test")
 }

src/main/java/org/openrewrite/docker/DockerImageVersion.java

@@ -22,11 +22,10 @@
 public class DockerImageVersion {
     String imageName;
 
-    @Nullable
     String version;
 
-    @Override
-    public String toString() {
-        return imageName + (version != null ? ":" + version : "");
+    public static DockerImageVersion of(String value) {
+        String[] imageVersionStr = value.split(":");
+        return new DockerImageVersion(imageVersionStr[0], imageVersionStr.length > 1 ? imageVersionStr[1].split(" ")[0] : "");
     }
 }

src/main/java/org/openrewrite/docker/search/FindDockerImageUses.java

@@ -15,17 +15,22 @@
  */
 package org.openrewrite.docker.search;
 
-import org.openrewrite.ExecutionContext;
-import org.openrewrite.Recipe;
-import org.openrewrite.TreeVisitor;
+import lombok.EqualsAndHashCode;
+import lombok.Value;
+import org.jspecify.annotations.Nullable;
+import org.openrewrite.*;
 import org.openrewrite.docker.DockerImageVersion;
 import org.openrewrite.docker.table.DockerBaseImages;
+import org.openrewrite.docker.trait.ImageMatcher;
 import org.openrewrite.marker.SearchResult;
+import org.openrewrite.text.Find;
+import org.openrewrite.text.PlainText;
+import org.openrewrite.trait.Reference;
 
-import java.util.List;
-import java.util.stream.Collectors;
+import java.nio.file.Path;
+import java.util.*;
 
-import static org.openrewrite.docker.trait.Traits.dockerfile;
+import static java.util.stream.Collectors.joining;
 
 public class FindDockerImageUses extends Recipe {
     transient DockerBaseImages dockerBaseImages = new DockerBaseImages(this);
@@ -37,25 +42,52 @@ public String getDisplayName() {
 
     @Override
     public String getDescription() {
-        return "Produce an impact analysis of base images used in Dockerfiles.";
+        return "Produce an impact analysis of base images used in Dockerfiles, .gitlab-ci files, Kubernetes Deployment file, etc.";
     }
 
     @Override
     public TreeVisitor<?, ExecutionContext> getVisitor() {
-        return dockerfile().asVisitor((docker, ctx) -> {
-            List<DockerImageVersion> froms = docker.getFroms();
-            if (!froms.isEmpty()) {
-                for (DockerImageVersion from : froms) {
-                    dockerBaseImages.insertRow(ctx, new DockerBaseImages.Row(
-                            from.getImageName(),
-                            from.getVersion() == null ? "" : from.getVersion()
-                    ));
+        return new TreeVisitor<Tree, ExecutionContext>() {
+
+            @Override
+            public @Nullable Tree visit(@Nullable Tree tree, ExecutionContext ctx) {
+                if (tree instanceof SourceFileWithReferences) {
+                    SourceFileWithReferences sourceFile = (SourceFileWithReferences) tree;
+                    Path sourcePath = sourceFile.getSourcePath();
+                    Collection<Reference> references = sourceFile.getReferences().findMatches(new ImageMatcher());
+                    Map<Tree, List<Reference>> matches = new HashMap<>();
+                    for (Reference ref : references) {
+                        DockerImageVersion from = DockerImageVersion.of(ref.getValue());
+                        dockerBaseImages.insertRow(ctx,
+                                new DockerBaseImages.Row(sourcePath.toString(), tree.getClass().getSimpleName(), from.getImageName(), from.getVersion())
+                        );
+                        matches.computeIfAbsent(ref.getTree(), t -> new ArrayList<>()).add(ref);
+                    }
+                    return new ReferenceFindSearchResultVisitor(matches).visit(tree, ctx, getCursor());
+                }
+                return tree;
+            }
+        };
+    }
+
+    @Value
+    @EqualsAndHashCode(callSuper = false)
+    private static class ReferenceFindSearchResultVisitor extends TreeVisitor<Tree, ExecutionContext> {
+        Map<Tree, List<Reference>> matches;
+
+        @Override
+        public @Nullable Tree postVisit(Tree tree, ExecutionContext ctx) {
+            List<Reference> references = matches.get(tree);
+            if (references != null) {
+                if (tree instanceof PlainText) {
+                    String find = references.stream().map(Reference::getValue).sorted().collect(joining("|"));
+                    return new Find(find, true, null, null, null, null, true)
+                            .getVisitor()
+                            .visitNonNull(tree, ctx);
                 }
-                return SearchResult.found(docker.getTree(),
-                        froms.stream().map(DockerImageVersion::toString)
-                                .collect(Collectors.joining(", ")));
+                return SearchResult.found(tree, references.get(0).getValue());
             }
-            return docker.getTree();
-        });
+            return tree;
+        }
     }
 }

src/main/java/org/openrewrite/docker/table/DockerBaseImages.java

@@ -30,6 +30,14 @@ public DockerBaseImages(Recipe recipe) {
 
     @Value
     public static class Row {
+        @Column(displayName = "Source path before the run",
+                description = "The source path of the file before the run.")
+        String sourcePath;
+
+        @Column(displayName = "LST type",
+                description = "The LST model type that the file is parsed as.")
+        String type;
+
         @Column(displayName = "Image name",
                 description = "The full name of the image.")
         String imageName;

src/main/java/org/openrewrite/docker/trait/DockerfileImageReference.java

@@ -0,0 +1,63 @@
+package org.openrewrite.docker.trait;
+
+import lombok.Value;
+import org.openrewrite.Cursor;
+import org.openrewrite.SourceFile;
+import org.openrewrite.internal.StringUtils;
+import org.openrewrite.text.PlainText;
+import org.openrewrite.trait.Reference;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.Set;
+
+@Value
+public class DockerfileImageReference implements Reference {
+    Cursor cursor;
+    String value;
+
+    @Override
+    public Kind getKind() {
+        return Kind.IMAGE;
+    }
+
+    public static class Provider implements Reference.Provider {
+        @Override
+        public boolean isAcceptable(SourceFile sourceFile) {
+            if (sourceFile instanceof PlainText) {
+                PlainText text = (PlainText) sourceFile;
+                String fileName = text.getSourcePath().toFile().getName();
+                return (fileName.endsWith("Dockerfile") || fileName.equals("Containerfile"))
+                        && (text.getText().contains("FROM") || text.getText().contains("from"));
+            }
+            return false;
+        }
+
+        @Override
+        public Set<Reference> getReferences(SourceFile sourceFile) {
+            Cursor c = new Cursor(new Cursor(null, Cursor.ROOT_VALUE), sourceFile);
+            String[] words = ((PlainText) sourceFile).getText()
+                    .replaceAll("\\s*#.*?\\n", "") // remove comments
+                    .replaceAll("\".*?\"", "") // remove string literals
+                    .split("\\s+");
+
+            Set<Reference> references = new HashSet<>();
+            ArrayList<String> imageVariables = new ArrayList<>();
+            for (int i = 0, wordsLength = words.length; i < wordsLength; i++) {
+                if ("from".equalsIgnoreCase(words[i])) {
+                    String image = words[i + 1].startsWith("--platform") ? words[i + 2] : words[i + 1];
+                    references.add(new DockerfileImageReference(c, image));
+                } else if ("as".equalsIgnoreCase(words[i])) {
+                    imageVariables.add(words[i + 1]);
+                } else if (words[i].startsWith("--from") && words[i].split("=").length == 2) {
+                    String image = words[i].split("=")[1];
+                    if (!imageVariables.contains(image) && !StringUtils.isNumeric(image)) {
+                        references.add(new DockerfileImageReference(c, image));
+                    }
+                }
+            }
+
+            return references;
+        }
+    }
+}

src/main/java/org/openrewrite/docker/trait/ImageMatcher.java

@@ -0,0 +1,16 @@
+package org.openrewrite.docker.trait;
+
+import org.openrewrite.trait.Reference;
+
+public class ImageMatcher implements Reference.Matcher {
+
+    @Override
+    public boolean matchesReference(Reference reference) {
+        return reference.getKind().equals(Reference.Kind.IMAGE);
+    }
+
+    @Override
+    public Reference.Renamer createRenamer(String newName) {
+        return reference -> newName;
+    }
+}

src/main/java/org/openrewrite/docker/trait/YamlImageReference.java

@@ -0,0 +1,48 @@
+package org.openrewrite.docker.trait;
+
+import lombok.EqualsAndHashCode;
+import lombok.Value;
+import org.jspecify.annotations.Nullable;
+import org.openrewrite.Cursor;
+import org.openrewrite.trait.Reference;
+import org.openrewrite.trait.SimpleTraitMatcher;
+import org.openrewrite.yaml.trait.YamlReference;
+import org.openrewrite.yaml.tree.Yaml;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+
+@Value
+@EqualsAndHashCode(callSuper = false)
+public class YamlImageReference extends YamlReference {
+    Cursor cursor;
+
+    @Override
+    public Reference.Kind getKind() {
+        return Reference.Kind.IMAGE;
+    }
+
+    public static class Provider extends YamlProvider {
+        private static final SimpleTraitMatcher<YamlReference> matcher = new SimpleTraitMatcher<YamlReference>() {
+            private final AtomicBoolean found = new AtomicBoolean(false);
+
+            @Override
+            protected @Nullable YamlReference test(Cursor cursor) {
+                Object value = cursor.getValue();
+                if (value instanceof Yaml.Scalar) {
+                    if (found.get()) {
+                        found.set(false);
+                        return new YamlImageReference(cursor);
+                    } else if ("image".equals(((Yaml.Scalar) value).getValue())) {
+                        found.set(true);
+                    }
+                }
+                return null;
+            }
+        };
+
+        @Override
+        public SimpleTraitMatcher<YamlReference> getMatcher() {
+            return matcher;
+        }
+    }
+}

src/main/resources/META-INF/services/org.openrewrite.trait.Reference$Provider

@@ -0,0 +1,2 @@
+org.openrewrite.docker.trait.YamlImageReference$Provider
+org.openrewrite.docker.trait.DockerfileImageReference$Provider