Docker: a COPY --from carries an image reference in the grammar

rewrite-docker/src/main/antlr/DockerLexer.g4

@@ -17,6 +17,31 @@ import java.util.Queue;}
     private boolean atLineStart = true;
     // Track if we're after HEALTHCHECK to recognize CMD/NONE as keywords
     private boolean afterHealthcheck = false;
+    // Track if we're in the flag section of a COPY or ADD, where --from carries an image reference
+    private boolean copyAddFlags = false;
+
+    // Every flag that is scoped to one logical line is cleared here, so that a mode of its own does
+    // not have to remember which of them the newline it matches instead of NEWLINE should reset
+    private void resetLine() {
+        atLineStart = true;
+        afterHealthcheck = false;
+        copyAddFlags = false;
+    }
+
+    // Whether what follows the '--' that both flag rules begin with is the '--from=' of a COPY or ADD
+    private boolean atFromFlag() {
+        if (!copyAddFlags) {
+            return false;
+        }
+        String fromFlag = "from=";
+        for (int i = 0; i < fromFlag.length(); i++) {
+            int c = _input.LA(i + 1);
+            if (c == -1 || Character.toLowerCase(c) != fromFlag.charAt(i)) {
+                return false;
+            }
+        }
+        return true;
+    }
 }
 
 options {
@@ -42,8 +67,8 @@ NONE       : 'NONE'       { if (!afterHealthcheck) setType(UNQUOTED_TEXT); atLin
 LABEL      : 'LABEL'      { if (!atLineStart) setType(UNQUOTED_TEXT); atLineStart = false; };
 EXPOSE     : 'EXPOSE'     { if (!atLineStart) setType(UNQUOTED_TEXT); atLineStart = false; };
 ENV        : 'ENV'        { if (!atLineStart) setType(UNQUOTED_TEXT); atLineStart = false; };
-ADD        : 'ADD'        { if (!atLineStart) setType(UNQUOTED_TEXT); atLineStart = false; };
-COPY       : 'COPY'       { if (!atLineStart) setType(UNQUOTED_TEXT); atLineStart = false; };
+ADD        : 'ADD'        { if (!atLineStart) setType(UNQUOTED_TEXT); else copyAddFlags = true; atLineStart = false; };
+COPY       : 'COPY'       { if (!atLineStart) setType(UNQUOTED_TEXT); else copyAddFlags = true; atLineStart = false; };
 ENTRYPOINT : 'ENTRYPOINT' { if (!atLineStart) setType(UNQUOTED_TEXT); atLineStart = false; };
 VOLUME     : 'VOLUME'     { if (!atLineStart) setType(UNQUOTED_TEXT); atLineStart = false; };
 USER       : 'USER'       { if (!atLineStart) setType(UNQUOTED_TEXT); else pushMode(USER_SPEC); atLineStart = false; };
@@ -86,9 +111,16 @@ EQUALS     : '=' { if (!afterHealthcheck) atLineStart = false; };
 // Captures the entire flag as a single token, stopping at whitespace
 // This avoids the greedy flagValue+ parsing issue while keeping shell commands working
 // Flag values can contain quoted strings (which may include spaces)
-FLAG : FLAG_TOKEN { if (!afterHealthcheck) atLineStart = false; };
+// The predicate excludes this rule where FROM_FLAG applies. Lexing is maximal munch, so without it
+// this longer token would always win and the image reference of a --from would stay unsplit. It sits
+// after the '--' because a predicate reachable without consuming anything stops the lexer caching the
+// start state of the mode that holds it, which costs every token in that mode a closure computation.
+FLAG : '--' {!atFromFlag()}? FLAG_BODY { if (!afterHealthcheck) atLineStart = false; };
 
-fragment FLAG_TOKEN : '--' [a-z] [a-z0-9_-]* ('=' FLAG_VALUE_PART+)?;
+// The --from of a COPY or ADD names an image, so its value is lexed as an image reference
+FROM_FLAG : '--' {atFromFlag()}? 'from=' { atLineStart = false; } -> pushMode(FLAG_IMAGE_REF);
+
+fragment FLAG_BODY : [a-z] [a-z0-9_-]* ('=' FLAG_VALUE_PART+)?;
 fragment FLAG_VALUE_PART
     : '"' ( '\\' ~[\r\n] | ~["\\\r\n] )* '"'   // Double-quoted string (with escapes)
     | '\'' ~['\r\n]* '\''                        // Single-quoted string (literal)
@@ -164,7 +196,7 @@ UNQUOTED_TEXT
     | '<' ~[< \t\r\n\\"'$[\]=] ( UNQUOTED_CHAR | ESCAPED_CHAR )*  // Single < followed by non-<
     | '<'  // Just a < by itself
     | ESCAPED_CHAR ( UNQUOTED_CHAR | ESCAPED_CHAR )*  // Start with escaped char (e.g., \; in find -exec)
-    ) { if (!afterHealthcheck) atLineStart = false; }
+    ) { if (!afterHealthcheck) atLineStart = false; copyAddFlags = false; }
     ;
 
 // Whitespace - HIDDEN in main mode
@@ -173,7 +205,7 @@ WS : WS_CHAR+ -> channel(HIDDEN);
 fragment WS_CHAR : [ \t];
 
 // Newlines - HIDDEN in main mode, reset state for next line
-NEWLINE : NEWLINE_CHAR+ { atLineStart = true; afterHealthcheck = false; } -> channel(HIDDEN);
+NEWLINE : NEWLINE_CHAR+ { resetLine(); } -> channel(HIDDEN);
 
 fragment NEWLINE_CHAR : [\r\n];
 
@@ -188,15 +220,15 @@ mode IMAGE_REF;
 IR_WS                : WS_CHAR+      -> type(WS), channel(HIDDEN);
 IR_LINE_CONTINUATION : LINE_CONT     -> type(LINE_CONTINUATION), channel(HIDDEN);
 IR_COMMENT           : '#' ~[\r\n]*  -> type(COMMENT), channel(HIDDEN);
-IR_NEWLINE           : NEWLINE_CHAR+ { atLineStart = true; afterHealthcheck = false; } -> type(NEWLINE), channel(HIDDEN), popMode;
+IR_NEWLINE           : NEWLINE_CHAR+ { resetLine(); } -> type(NEWLINE), channel(HIDDEN), popMode;
 
 COLON : ':';
 AT    : '@';
 
 // AS ends the reference; the stage name that follows it is ordinary text
 AS : 'AS' -> popMode;
 
-IR_FLAG                 : FLAG_TOKEN      -> type(FLAG);
+IR_FLAG                 : '--' FLAG_BODY  -> type(FLAG);
 IR_DOUBLE_QUOTED_STRING : DQ_STRING       -> type(DOUBLE_QUOTED_STRING);
 IR_SINGLE_QUOTED_STRING : SQ_STRING       -> type(SINGLE_QUOTED_STRING);
 IR_ENV_VAR              : VAR_REF         -> type(ENV_VAR);
@@ -205,11 +237,40 @@ IR_DOLLAR               : '$'             -> type(DOLLAR);
 
 // Text of one part of the reference. A colon that a '/' follows belongs to a registry port rather
 // than to a tag ('host:5000/img:tag'), so it stays inside the token.
-IR_UNQUOTED_TEXT : ( IR_TEXT_CHAR | ESCAPED_CHAR | IR_PORT_COLON )+ -> type(UNQUOTED_TEXT);
+IR_UNQUOTED_TEXT : IR_TEXT -> type(UNQUOTED_TEXT);
 
+// Shared with FLAG_IMAGE_REF, which reads the same reference. ESCAPE_SEQUENCE rather than
+// ESCAPED_CHAR because it stops before a newline, which lexing longest-match-first would otherwise
+// take into the token and so hide the line continuation that ends the reference.
+fragment IR_TEXT       : ( IR_TEXT_CHAR | ESCAPE_SEQUENCE | IR_PORT_COLON )+;
 fragment IR_TEXT_CHAR  : ~[:@ \t\r\n\\"'$];
 fragment IR_PORT_COLON : ':' ( IR_TEXT_CHAR | ':' )* '/';
 
+// ----------------------------------------------------------------------------------------------
+// FLAG_IMAGE_REF mode - the image reference carried by the --from flag of a COPY or ADD
+// As IMAGE_REF, except that the reference ends at the whitespace before the paths that follow it
+// rather than at the end of the line, and AS is not a keyword because no stage alias can appear here.
+// ----------------------------------------------------------------------------------------------
+mode FLAG_IMAGE_REF;
+
+// The end of the reference is a token of its own rather than hidden whitespace: popping a mode does
+// not bound a parser rule, so without a token to name the `imageName` of `COPY --from=build --link .`
+// would carry on into the flag that follows it.
+FLAG_END : ( WS_CHAR+ | LINE_CONT ) -> popMode;
+
+FIR_NEWLINE : NEWLINE_CHAR+ { resetLine(); } -> type(NEWLINE), channel(HIDDEN), popMode;
+
+FIR_COLON : ':' -> type(COLON);
+FIR_AT    : '@' -> type(AT);
+
+FIR_DOUBLE_QUOTED_STRING : DQ_STRING       -> type(DOUBLE_QUOTED_STRING);
+FIR_SINGLE_QUOTED_STRING : SQ_STRING       -> type(SINGLE_QUOTED_STRING);
+FIR_ENV_VAR              : VAR_REF         -> type(ENV_VAR);
+FIR_SPECIAL_VAR          : SPECIAL_VAR_REF -> type(SPECIAL_VAR);
+FIR_DOLLAR               : '$'             -> type(DOLLAR);
+
+FIR_UNQUOTED_TEXT : IR_TEXT -> type(UNQUOTED_TEXT);
+
 // ----------------------------------------------------------------------------------------------
 // USER_SPEC mode - the user:group of a USER instruction
 // Entered from the USER keyword and left at the end of the line. As IMAGE_REF, minus the '@' and the
@@ -221,7 +282,7 @@ mode USER_SPEC;
 US_WS                : WS_CHAR+     -> type(WS), channel(HIDDEN);
 US_LINE_CONTINUATION : LINE_CONT    -> type(LINE_CONTINUATION), channel(HIDDEN);
 US_COMMENT           : '#' ~[\r\n]* -> type(COMMENT), channel(HIDDEN);
-US_NEWLINE           : NEWLINE_CHAR+ { atLineStart = true; afterHealthcheck = false; } -> type(NEWLINE), channel(HIDDEN), popMode;
+US_NEWLINE           : NEWLINE_CHAR+ { resetLine(); } -> type(NEWLINE), channel(HIDDEN), popMode;
 
 US_COLON : ':' -> type(COLON);
 

rewrite-docker/src/main/antlr/DockerParser.g4

@@ -149,7 +149,7 @@ maintainerInstruction
 
 // Common elements
 flags
-    : flag+
+    : ( flag | fromFlag )+
     ;
 
 // Flag token captures entire flag: --name or --name=value
@@ -158,6 +158,13 @@ flag
     : FLAG
     ;
 
+// The --from of a COPY or ADD holds the same name:tag@digest reference a FROM does, split by the
+// same rule. FLAG_END is the whitespace that ends the reference, without which this rule would
+// carry on into the flags and paths that follow it.
+fromFlag
+    : FROM_FLAG imageReference? FLAG_END?
+    ;
+
 execForm
     : jsonArray
     ;

rewrite-docker/src/main/java/org/openrewrite/docker/internal/DockerParserVisitor.java

@@ -133,24 +133,35 @@ public Docker visitFromInstruction(DockerParser.FromInstructionContext ctx) {
 
         List<Docker.Flag> flags = ctx.flags() != null ? convertFlags(ctx.flags()) : null;
 
-        DockerParser.ImageReferenceContext reference = ctx.imageReference();
-        TerminalNode colon = reference.COLON();
-        TerminalNode at = reference.AT();
-        Docker.Argument imageName = separatedPart(reference.imageName(), colon != null ? colon : at);
+        Docker.@Nullable Argument[] reference = imageReferenceParts(ctx.imageReference());
+
+        Docker.From.As as = ctx.AS() != null ? visitFromAs(ctx) : null;
+
+        return new Docker.From(randomId(), prefix, Markers.EMPTY, fromKeyword, flags, reference[0], reference[1], reference[2], as);
+    }
+
+    /// The `{imageName, tag, digest}` a reference splits into, or an empty name where the reference is
+    /// absent, as it is in the `--from=` of a `COPY` that names nothing. Shared by `FROM` and by the
+    /// `--from` of a `COPY`/`ADD`, which the grammar splits by the same `imageReference` rule.
+    private Docker.@Nullable Argument[] imageReferenceParts(DockerParser.@Nullable ImageReferenceContext ctx) {
+        if (ctx == null) {
+            return new Docker.@Nullable Argument[]{
+                    new Docker.Argument(randomId(), Space.EMPTY, Markers.EMPTY, emptyList()), null, null};
+        }
+        TerminalNode colon = ctx.COLON();
+        TerminalNode at = ctx.AT();
+        Docker.Argument imageName = separatedPart(ctx.imageName(), colon != null ? colon : at);
         Docker.Argument tag = null;
         if (colon != null) {
             skip(colon.getSymbol());
-            tag = separatedPart(reference.tag(), at);
+            tag = separatedPart(ctx.tag(), at);
         }
         Docker.Argument digest = null;
         if (at != null) {
             skip(at.getSymbol());
-            digest = separatedPart(reference.digest(), null);
+            digest = separatedPart(ctx.digest(), null);
         }
-
-        Docker.From.As as = ctx.AS() != null ? visitFromAs(ctx) : null;
-
-        return new Docker.From(randomId(), prefix, Markers.EMPTY, fromKeyword, flags, imageName, tag, digest, as);
+        return new Docker.@Nullable Argument[]{imageName, tag, digest};
     }
 
     private Docker.From.As visitFromAs(DockerParser.FromInstructionContext ctx) {
@@ -934,12 +945,31 @@ private Docker.CommandForm visitCommandFormForEntrypoint(DockerParser.Entrypoint
 
     private List<Docker.Flag> convertFlags(DockerParser.FlagsContext ctx) {
         List<Docker.Flag> flags = new ArrayList<>();
-        for (DockerParser.FlagContext flagCtx : ctx.flag()) {
-            flags.add(parseFlag(flagCtx.FLAG().getSymbol()));
+        for (ParseTree child : ctx.children) {
+            if (child instanceof DockerParser.FlagContext) {
+                flags.add(parseFlag(((DockerParser.FlagContext) child).FLAG().getSymbol()));
+            } else if (child instanceof DockerParser.FromFlagContext) {
+                flags.add(parseFromFlag((DockerParser.FromFlagContext) child));
+            }
         }
         return flags;
     }
 
+    /// The `--from` of a `COPY` or `ADD`, whose value the `FLAG_IMAGE_REF` lexer mode splits by the
+    /// same rule that splits the reference of a `FROM`. A flag holds one value, so the parts are
+    /// flattened back into it with their separators, the form [ImageReferences] reads them from.
+    private Docker.Flag parseFromFlag(DockerParser.FromFlagContext ctx) {
+        Token token = ctx.FROM_FLAG().getSymbol();
+        Space flagPrefix = prefix(token);
+        String tokenText = token.getText();
+        skip(token);
+
+        String flagName = tokenText.substring(2, tokenText.indexOf('='));
+        Docker.Argument value = new Docker.Argument(randomId(), Space.EMPTY, Markers.EMPTY,
+                ImageReferences.contents(imageReferenceParts(ctx.imageReference())));
+        return new Docker.Flag(randomId(), flagPrefix, Markers.EMPTY, flagName, value);
+    }
+
     /**
      * Parse a FLAG token into a Flag AST node.
      * FLAG token format: --name or --name=value