From 8adb2155228c48a36dfbbad4a7351ead5527a79f Mon Sep 17 00:00:00 2001
From: Kathryn Baldauf <k_baldauf@apple.com>
Date: Tue, 5 Aug 2025 14:34:34 -0700
Subject: [PATCH] Native Builder: Add parser support for CMD and LABEL
 instructions (#448)

This PR adds support for CMD and LABEL instructions in the native
builder's parser.

This also changes how options are tokenized. Options now include the raw
string so that when constructing a command for instructions like CMD and
RUN, we can use the exact user input without having to add logic in the
tokenizer to know when we're parsing a command verses other options,
etc.

Closes https://github.com/apple/container/issues/428 and
https://github.com/apple/container/issues/429

Signed-off-by: Kathryn Baldauf <k_baldauf@apple.com>
---
 .../ContainerBuildIR/Graph/GraphBuilder.swift |  22 ++-
 .../Docker/DockerInstructionVisitor.swift     |  12 +-
 .../Docker/DockerfileParser.swift             | 105 ++++++++---
 .../Docker/DockerfileTokenizer.swift          |   5 +-
 .../Docker/Instructions/CopyInstruction.swift |   2 +-
 .../Instructions/DockerInstruction.swift      |  22 ++-
 .../Docker/Instructions/RunInstruction.swift  |  11 +-
 .../ContainerBuildParser/Types.swift          |  14 +-
 .../ParserTests.swift                         | 133 +++++++++++++-
 .../TokenizerTests.swift                      | 163 +++++++++++++-----
 .../VisitorTests.swift                        |  80 +++++++--
 11 files changed, 477 insertions(+), 92 deletions(-)

diff --git a/Sources/NativeBuilder/ContainerBuildIR/Graph/GraphBuilder.swift b/Sources/NativeBuilder/ContainerBuildIR/Graph/GraphBuilder.swift
index a118c05c..f52d76a7 100644
--- a/Sources/NativeBuilder/ContainerBuildIR/Graph/GraphBuilder.swift
+++ b/Sources/NativeBuilder/ContainerBuildIR/Graph/GraphBuilder.swift
@@ -179,8 +179,20 @@ public final class GraphBuilder {
         network: NetworkMode = .default,
     ) throws -> Self {
         let cmd = shell ? Command.shell(command) : Command.exec(command.split(separator: " ").map(String.init))
-        let envVars = env.map { (key: $0.key, value: EnvironmentValue.literal($0.value)) }
+        return try runWithCmd(cmd, shell: shell, env: env, workdir: workdir, user: user, mounts: mounts, network: network)
+    }
 
+    @discardableResult
+    public func runWithCmd(
+        _ cmd: Command,
+        shell: Bool = true,
+        env: [String: String] = [:],
+        workdir: String? = nil,
+        user: User? = nil,
+        mounts: [Mount] = [],
+        network: NetworkMode = .default,
+    ) throws -> Self {
+        let envVars = env.map { (key: $0.key, value: EnvironmentValue.literal($0.value)) }
         let operation = ExecOperation(
             command: cmd,
             environment: Environment(envVars),
@@ -280,6 +292,14 @@ public final class GraphBuilder {
         return try add(operation)
     }
 
+    @discardableResult
+    public func labelBatch(labels: [String: String]) throws -> Self {
+        let operation = MetadataOperation(
+            action: .setLabelBatch(labels)
+        )
+        return try add(operation)
+    }
+
     /// Expose port
     @discardableResult
     public func expose(_ port: Int, protocolType: PortSpec.NetworkProtocol = .tcp) throws -> Self {
diff --git a/Sources/NativeBuilder/ContainerBuildParser/Docker/DockerInstructionVisitor.swift b/Sources/NativeBuilder/ContainerBuildParser/Docker/DockerInstructionVisitor.swift
index 6546cf62..114e58ac 100644
--- a/Sources/NativeBuilder/ContainerBuildParser/Docker/DockerInstructionVisitor.swift
+++ b/Sources/NativeBuilder/ContainerBuildParser/Docker/DockerInstructionVisitor.swift
@@ -20,6 +20,8 @@ protocol InstructionVisitor {
     func visit(_ from: FromInstruction) throws
     func visit(_ run: RunInstruction) throws
     func visit(_ copy: CopyInstruction) throws
+    func visit(_ cmd: CMDInstruction) throws
+    func visit(_ label: LabelInstruction) throws
 }
 
 /// DockerInstructionVisitor visits each provided DockerInstruction and builds a
@@ -114,7 +116,7 @@ extension DockerInstructionVisitor {
             mounts.append(graphMount)
         }
 
-        try graphBuilder.run(run.command, shell: run.shell, mounts: mounts)
+        try graphBuilder.runWithCmd(run.command, mounts: mounts)
     }
 
     func visit(_ copy: CopyInstruction) throws {
@@ -136,4 +138,12 @@ extension DockerInstructionVisitor {
         }
         try graphBuilder.copyFromContext(paths: copy.sources, to: copy.destination, chown: copy.chown, chmod: copy.chmod)
     }
+
+    func visit(_ cmd: CMDInstruction) throws {
+        try graphBuilder.cmd(cmd.command)
+    }
+
+    func visit(_ label: LabelInstruction) throws {
+        try graphBuilder.labelBatch(labels: label.labels)
+    }
 }
diff --git a/Sources/NativeBuilder/ContainerBuildParser/Docker/DockerfileParser.swift b/Sources/NativeBuilder/ContainerBuildParser/Docker/DockerfileParser.swift
index 09239fe7..96a9de7a 100644
--- a/Sources/NativeBuilder/ContainerBuildParser/Docker/DockerfileParser.swift
+++ b/Sources/NativeBuilder/ContainerBuildParser/Docker/DockerfileParser.swift
@@ -63,6 +63,12 @@ public struct DockerfileParser: BuildParser {
             return try tokensToFromInstruction(tokens: tokens)
         case .RUN:
             return try tokensToRunInstruction(tokens: tokens)
+        case .COPY:
+            return try tokensToCopyInstruction(tokens: tokens)
+        case .CMD:
+            return try tokensToCMDInstruction(tokens: tokens)
+        case .LABEL:
+            return try tokensToLabelInstruction(tokens: tokens)
         default:
             throw ParseError.invalidInstruction(value)
         }
@@ -78,13 +84,13 @@ public struct DockerfileParser: BuildParser {
 
         // Step 1: parse options
         while index < tokens.endIndex {
-            guard case .option(let key, let value) = tokens[index] else {
+            guard case .option(let option) = tokens[index] else {
                 break
             }
-            guard FromOptions(rawValue: key) == .platform else {
+            guard FromOptions(rawValue: option.key) == .platform else {
                 throw ParseError.unexpectedValue
             }
-            platform = value
+            platform = option.value
             index += 1
         }
 
@@ -133,29 +139,42 @@ public struct DockerfileParser: BuildParser {
 
         // Step 1: parse options
         while index < tokens.endIndex {
-            guard case .option(let key, let value) = tokens[index] else {
+            guard case .option(let option) = tokens[index] else {
                 break
             }
 
-            guard let option = RunOptions(rawValue: key) else {
+            guard let runOpt = RunOptions(rawValue: option.key) else {
                 throw ParseError.unexpectedValue
             }
 
-            switch option {
+            switch runOpt {
             case .mount:
-                rawMounts.append(value)
+                rawMounts.append(option.value)
             case .network:
-                network = value
+                network = option.value
             default:
                 throw ParseError.unexpectedValue
             }
             index += 1
         }
 
+        // Step 2: parse run command
+        let (newIndex, cmd) = getCommand(start: index, tokens: tokens)
+        index = newIndex
+
+        // check for extra tokens
+        if index < tokens.endIndex {
+            throw ParseError.unexpectedValue
+        }
+
+        return try RunInstruction(command: cmd, rawMounts: rawMounts, network: network)
+    }
+
+    private func getCommand(start: Int, tokens: [Token]) -> (index: Int, cmd: Command) {
         var command = [String]()
         var shell = true
 
-        // Step 2: parse run command and if we're using shell or exec form
+        var index = start
         while index < tokens.endIndex {
             if case .stringList(let value) = tokens[index], command.isEmpty {
                 // when using the exec form, there should only be a single list for the command
@@ -166,18 +185,16 @@ public struct DockerfileParser: BuildParser {
                 break
             } else if case .stringLiteral(let value) = tokens[index] {
                 command.append(value)
+            } else if case .option(let option) = tokens[index] {
+                command.append(option.raw)
             } else {
                 break
             }
             index += 1
         }
 
-        // check for extra tokens
-        if index < tokens.endIndex {
-            throw ParseError.unexpectedValue
-        }
-
-        return try RunInstruction(command: command, shell: shell, rawMounts: rawMounts, network: network)
+        let cmd = shell ? Command.shell(command.joined(separator: " ")) : Command.exec(command)
+        return (index, cmd)
     }
 
     internal func tokensToCopyInstruction(tokens: [Token]) throws -> CopyInstruction {
@@ -191,35 +208,35 @@ public struct DockerfileParser: BuildParser {
 
         // Step 1: parse options
         while index < tokens.endIndex {
-            guard case .option(let key, let value) = tokens[index] else {
+            guard case .option(let option) = tokens[index] else {
                 break
             }
 
-            guard let option = CopyOptions(rawValue: key) else {
+            guard let copyOpt = CopyOptions(rawValue: option.key) else {
                 throw ParseError.unexpectedValue
             }
 
-            switch option {
+            switch copyOpt {
             case .from:
                 if from != nil {
                     throw ParseError.duplicateOptionSet(CopyOptions.from.rawValue)
                 }
-                from = value
+                from = option.value
             case .chown:
                 if chown != nil {
                     throw ParseError.duplicateOptionSet(CopyOptions.chown.rawValue)
                 }
-                chown = value
+                chown = option.value
             case .chmod:
                 if chmod != nil {
                     throw ParseError.duplicateOptionSet(CopyOptions.chmod.rawValue)
                 }
-                chmod = value
+                chmod = option.value
             case .link:
                 if link != nil {
                     throw ParseError.duplicateOptionSet(CopyOptions.link.rawValue)
                 }
-                link = value
+                link = option.value
             default:
                 throw ParseError.unexpectedValue
             }
@@ -250,4 +267,48 @@ public struct DockerfileParser: BuildParser {
         return try CopyInstruction(sources: sources, destination: destination, from: from, ownership: chown, permissions: chmod)
     }
 
+    internal func tokensToCMDInstruction(tokens: [Token]) throws -> CMDInstruction {
+        var index = tokens.startIndex
+        index += 1
+
+        // get the command
+        let (newIndex, cmd) = getCommand(start: index, tokens: tokens)
+        index = newIndex
+
+        // check for extra tokens
+        if index < tokens.endIndex {
+            throw ParseError.unexpectedValue
+        }
+
+        return CMDInstruction(command: cmd)
+    }
+
+    internal func tokensToLabelInstruction(tokens: [Token]) throws -> LabelInstruction {
+        var index = tokens.startIndex
+        index += 1
+
+        var labels: [String: String] = [:]
+        while index < tokens.endIndex {
+            guard case .stringLiteral(let option) = tokens[index] else {
+                break
+            }
+            let components = option.split(separator: "=", maxSplits: 1)
+            guard components.count == 2 else {
+                throw ParseError.unexpectedValue
+            }
+            let key = String(components[0])
+            guard labels[key] == nil else {
+                throw ParseError.duplicateOptionSet(key)
+            }
+            labels[key] = String(components[1])
+            index += 1
+        }
+
+        // check for extra tokens
+        if index < tokens.endIndex {
+            throw ParseError.unexpectedValue
+        }
+
+        return LabelInstruction(labels: labels)
+    }
 }
diff --git a/Sources/NativeBuilder/ContainerBuildParser/Docker/DockerfileTokenizer.swift b/Sources/NativeBuilder/ContainerBuildParser/Docker/DockerfileTokenizer.swift
index 36e154cd..a3c13947 100644
--- a/Sources/NativeBuilder/ContainerBuildParser/Docker/DockerfileTokenizer.swift
+++ b/Sources/NativeBuilder/ContainerBuildParser/Docker/DockerfileTokenizer.swift
@@ -135,10 +135,11 @@ struct DockerfileTokenizer {
             let valueStart = position
             parseWord()
             let rawValue = input[valueStart..<position]
-            return .option(String(rawWord), String(rawValue))
+            let raw = input[wordStart..<position]
+            return .option(Option(key: String(rawWord), value: String(rawValue), raw: String(raw)))
         }
         // split by equal
         let optionComponents = rawWord.split(separator: "=", maxSplits: 1)
-        return .option(String(optionComponents[0]), String(optionComponents[1]))
+        return .option(Option(key: String(optionComponents[0]), value: String(optionComponents[1]), raw: String(rawWord)))
     }
 }
diff --git a/Sources/NativeBuilder/ContainerBuildParser/Docker/Instructions/CopyInstruction.swift b/Sources/NativeBuilder/ContainerBuildParser/Docker/Instructions/CopyInstruction.swift
index 3b87c1e5..9f863c75 100644
--- a/Sources/NativeBuilder/ContainerBuildParser/Docker/Instructions/CopyInstruction.swift
+++ b/Sources/NativeBuilder/ContainerBuildParser/Docker/Instructions/CopyInstruction.swift
@@ -23,7 +23,7 @@ enum CopyOptions: String {
     case link = "--link"
 }
 
-struct CopyInstruction: DockerInstruction, Equatable {
+struct CopyInstruction: DockerInstruction {
     let sources: [String]
     let destination: String
     let from: String?
diff --git a/Sources/NativeBuilder/ContainerBuildParser/Docker/Instructions/DockerInstruction.swift b/Sources/NativeBuilder/ContainerBuildParser/Docker/Instructions/DockerInstruction.swift
index 5d536e6c..2dcd58fd 100644
--- a/Sources/NativeBuilder/ContainerBuildParser/Docker/Instructions/DockerInstruction.swift
+++ b/Sources/NativeBuilder/ContainerBuildParser/Docker/Instructions/DockerInstruction.swift
@@ -19,7 +19,7 @@ import ContainerizationOCI
 
 /// DockerInstruction represents a single docker instruction with its given options
 /// and arguments. Instructions are "visited" to add to a build graph.
-protocol DockerInstruction {
+protocol DockerInstruction: Sendable, Equatable {
     func accept(_ visitor: DockerInstructionVisitor) throws
 }
 
@@ -27,7 +27,7 @@ enum FromOptions: String {
     case platform = "--platform"
 }
 
-struct FromInstruction: DockerInstruction, Equatable {
+struct FromInstruction: DockerInstruction {
     let image: ImageReference
     let platform: Platform?
     let stageName: String?
@@ -55,6 +55,9 @@ struct FromInstruction: DockerInstruction, Equatable {
 enum DockerInstructionName: String {
     case FROM = "from"
     case RUN = "run"
+    case COPY = "copy"
+    case CMD = "cmd"
+    case LABEL = "label"
 }
 
 /// DockerKeyword defines words that are used as keywords within a line of a dockerfile
@@ -62,3 +65,18 @@ enum DockerInstructionName: String {
 enum DockerKeyword: String {
     case AS = "as"
 }
+
+struct CMDInstruction: DockerInstruction {
+    let command: Command
+
+    func accept(_ visitor: DockerInstructionVisitor) throws {
+        try visitor.visit(self)
+    }
+}
+
+struct LabelInstruction: DockerInstruction {
+    let labels: [String: String]
+    func accept(_ visitor: DockerInstructionVisitor) throws {
+        try visitor.visit(self)
+    }
+}
diff --git a/Sources/NativeBuilder/ContainerBuildParser/Docker/Instructions/RunInstruction.swift b/Sources/NativeBuilder/ContainerBuildParser/Docker/Instructions/RunInstruction.swift
index 4047b1e0..29eb70ef 100644
--- a/Sources/NativeBuilder/ContainerBuildParser/Docker/Instructions/RunInstruction.swift
+++ b/Sources/NativeBuilder/ContainerBuildParser/Docker/Instructions/RunInstruction.swift
@@ -257,21 +257,18 @@ struct RunMountOptions: Equatable {
 
 /// RunInstruction represents a RUN instruction from a dockerfile
 struct RunInstruction: DockerInstruction, Equatable {
-    let command: String
-    let shell: Bool
+    let command: Command
     let mounts: [RunMount]
     let network: NetworkMode
 
     init() {
-        self.command = ""
-        self.shell = false
+        self.command = .shell("")
         self.mounts = []
         self.network = .default
     }
 
-    init(command: [String], shell: Bool, rawMounts: [String], network: String?) throws {
-        self.command = command.joined(separator: " ")
-        self.shell = shell
+    init(command: Command, rawMounts: [String], network: String?) throws {
+        self.command = command
         self.network = try RunInstruction.parseNetworkMode(mode: network)
         var parsedMounts: [RunMount] = []
         for m in rawMounts {
diff --git a/Sources/NativeBuilder/ContainerBuildParser/Types.swift b/Sources/NativeBuilder/ContainerBuildParser/Types.swift
index 48efb1cc..d2290f45 100644
--- a/Sources/NativeBuilder/ContainerBuildParser/Types.swift
+++ b/Sources/NativeBuilder/ContainerBuildParser/Types.swift
@@ -41,5 +41,17 @@ public enum ParseError: Error, Equatable {
 public enum Token: Sendable, Equatable {
     case stringLiteral(String)
     case stringList([String])
-    case option(String, String)
+    case option(Option)
+}
+
+public struct Option: Sendable, Equatable {
+    let key: String
+    let value: String
+    let raw: String
+
+    init(key: String, value: String, raw: String) {
+        self.key = key
+        self.value = value
+        self.raw = raw
+    }
 }
diff --git a/Tests/NativeBuilderTests/ContainerBuildParserTests/ParserTests.swift b/Tests/NativeBuilderTests/ContainerBuildParserTests/ParserTests.swift
index cedcd3df..8ab7f6fd 100644
--- a/Tests/NativeBuilderTests/ContainerBuildParserTests/ParserTests.swift
+++ b/Tests/NativeBuilderTests/ContainerBuildParserTests/ParserTests.swift
@@ -93,9 +93,140 @@ import Testing
             try parser.parse(dockerfile)
         }
     }
+
+    @Test func testSimpleDockerfileWithCopy() throws {
+        let imageRef = ImageReference(parsing: "alpine:latest")
+        #expect(imageRef != nil, "Failed to parse image reference")
+        let dockerfile =
+            #"""
+            FROM alpine AS build-context
+
+            FROM alpine AS stage-two
+            COPY --from=build-context /test /test
+            """#
+        let parser = DockerfileParser()
+        let actualGraph = try parser.parse(dockerfile)
+
+        #expect(!actualGraph.stages.isEmpty)
+        let stages = actualGraph.stages
+
+        #expect(stages.count == 2, "expected 2 stages, instead got \(stages.count)")
+        #expect(stages[0].name == "build-context", "expected stage name build-context, got \(stages[0].name)")
+        #expect(stages[1].name == "stage-two", "expected stage name stage-two, got \(stages[1].name)")
+        #expect(stages[1].nodes.count == 1, "expected 1 node, instead got \(stages[1].nodes.count)")
+
+        let node = stages[1].nodes[0]
+        #expect(node.operation is FilesystemOperation)
+
+        let copy = node.operation as! FilesystemOperation
+        #expect(copy.action == .copy)
+    }
+
+    @Test func testSimpleDockerfileRun() throws {
+        let imageRef = ImageReference(parsing: "alpine:latest")
+        #expect(imageRef != nil, "Failed to parse image reference")
+        let dockerfile =
+            #"""
+            FROM alpine:latest AS build
+
+            RUN ["ls", "-la"]
+            """#
+        let parser = DockerfileParser()
+        let actualGraph = try parser.parse(dockerfile)
+
+        #expect(!actualGraph.stages.isEmpty)
+
+        let stages = actualGraph.stages
+        #expect(stages.count == 1, "expected 1 stage, instead got \(actualGraph.stages.count)")
+
+        let stage = stages[0]
+        #expect(stage.nodes.count == 1, "expected 1 node, instead got \(stage.nodes.count)")
+
+        let run = stage.nodes[0].operation as! ExecOperation
+        #expect(run.command.displayString == "ls -la")
+    }
+
+    @Test func testSimpleDockerfileRunShell() throws {
+        let imageRef = ImageReference(parsing: "alpine:latest")
+        #expect(imageRef != nil, "Failed to parse image reference")
+        let dockerfile =
+            #"""
+            FROM alpine:latest AS build
+
+            RUN build.sh --verbose
+            """#
+        let parser = DockerfileParser()
+        let actualGraph = try parser.parse(dockerfile)
+
+        #expect(!actualGraph.stages.isEmpty)
+
+        let stages = actualGraph.stages
+        #expect(stages.count == 1, "expected 1 stage, instead got \(actualGraph.stages.count)")
+
+        let stage = stages[0]
+        #expect(stage.nodes.count == 1, "expected 1 node, instead got \(stage.nodes.count)")
+
+        let run = stage.nodes[0].operation as! ExecOperation
+        #expect(run.command.displayString == "build.sh --verbose")
+
+    }
+
+    @Test func testSimpleDockerfileLabel() throws {
+        let imageRef = ImageReference(parsing: "alpine:latest")
+        #expect(imageRef != nil, "Failed to parse image reference")
+        let dockerfile =
+            #"""
+            FROM alpine:latest AS build
+
+            LABEL label1=value1
+            """#
+        let parser = DockerfileParser()
+        let actualGraph = try parser.parse(dockerfile)
+
+        #expect(!actualGraph.stages.isEmpty)
+
+        let stages = actualGraph.stages
+        #expect(stages.count == 1, "expected 1 stage, instead got \(actualGraph.stages.count)")
+
+        let stage = stages[0]
+        #expect(stage.nodes.count == 1, "expected 1 node, instead got \(stage.nodes.count)")
+
+        let label = stage.nodes[0].operation as! MetadataOperation
+        #expect(label.action == .setLabelBatch(["label1": "value1"]))
+    }
+
+    @Test func testSimpleDockerfileCMD() throws {
+        let imageRef = ImageReference(parsing: "alpine:latest")
+        #expect(imageRef != nil, "Failed to parse image reference")
+        let dockerfile =
+            #"""
+            FROM alpine:latest AS build
+
+            CMD ["./test.sh", "--verbose"]
+            """#
+        let parser = DockerfileParser()
+        let actualGraph = try parser.parse(dockerfile)
+
+        #expect(!actualGraph.stages.isEmpty)
+
+        let stages = actualGraph.stages
+        #expect(stages.count == 1, "expected 1 stage, instead got \(actualGraph.stages.count)")
+
+        let stage = stages[0]
+        #expect(stage.nodes.count == 1, "expected 1 node, instead got \(stage.nodes.count)")
+
+        let cmd = stage.nodes[0].operation as! MetadataOperation
+        switch cmd.action {
+        case .setCmd(let command):
+            #expect(command.displayString == "./test.sh --verbose")
+        default:
+            Issue.record("expected .setCmd action type, instead got \(cmd.action)")
+            return
+        }
+    }
 }
 
-// tests for parsing RUN mount options
+// tests for parsing options for the different instructions
 extension ParserTest {
     struct RunMountTestCase {
         let rawMount: String
diff --git a/Tests/NativeBuilderTests/ContainerBuildParserTests/TokenizerTests.swift b/Tests/NativeBuilderTests/ContainerBuildParserTests/TokenizerTests.swift
index 00911755..386a823c 100644
--- a/Tests/NativeBuilderTests/ContainerBuildParserTests/TokenizerTests.swift
+++ b/Tests/NativeBuilderTests/ContainerBuildParserTests/TokenizerTests.swift
@@ -46,7 +46,7 @@ import Testing
             input: "RUN --mount=type=cache /app",
             expectedTokens: [
                 .stringLiteral("RUN"),
-                .option("--mount", "type=cache"),
+                .option(Option(key: "--mount", value: "type=cache", raw: "--mount=type=cache")),
                 .stringLiteral("/app"),
             ]
         ),
@@ -54,7 +54,7 @@ import Testing
             input: "RUN --network=default /app",
             expectedTokens: [
                 .stringLiteral("RUN"),
-                .option("--network", "default"),
+                .option(Option(key: "--network", value: "default", raw: "--network=default")),
                 .stringLiteral("/app"),
             ]
         ),
@@ -62,8 +62,8 @@ import Testing
             input: "RUN --mount=type=bind,target=/target --network=host build.sh",
             expectedTokens: [
                 .stringLiteral("RUN"),
-                .option("--mount", "type=bind,target=/target"),
-                .option("--network", "host"),
+                .option(Option(key: "--mount", value: "type=bind,target=/target", raw: "--mount=type=bind,target=/target")),
+                .option(Option(key: "--network", value: "host", raw: "--network=host")),
                 .stringLiteral("build.sh"),
             ]
         ),
@@ -71,7 +71,7 @@ import Testing
             input: "RUN --mount type=cache /app",
             expectedTokens: [
                 .stringLiteral("RUN"),
-                .option("--mount", "type=cache"),
+                .option(Option(key: "--mount", value: "type=cache", raw: "--mount type=cache")),
                 .stringLiteral("/app"),
             ]
         ),
@@ -82,9 +82,9 @@ import Testing
                 """,
             expectedTokens: [
                 .stringLiteral("RUN"),
-                .option("--mount", "type=cache"),
+                .option(Option(key: "--mount", value: "type=cache", raw: "--mount=type=cache")),
                 .stringLiteral("build.sh"),
-                .option("--input", "hello"),
+                .option(Option(key: "--input", value: "hello", raw: "--input hello")),
             ]
         ),
         tokenizerTestInput(
@@ -94,7 +94,7 @@ import Testing
                 """#,
             expectedTokens: [
                 .stringLiteral("RUN"),
-                .option("--mount", "type=cache"),
+                .option(Option(key: "--mount", value: "type=cache", raw: "--mount=type=cache")),
                 .stringList(["build.sh", "--input", "hello"]),
             ]
         ),
@@ -105,7 +105,7 @@ import Testing
                 """#,
             expectedTokens: [
                 .stringLiteral("RUN"),
-                .option("--mount", "type=cache"),
+                .option(Option(key: "--mount", value: "type=cache", raw: "--mount=type=cache")),
                 .stringList(["build.sh", "--input", "hello"]),
             ]
         ),
@@ -116,7 +116,7 @@ import Testing
                 """#,
             expectedTokens: [
                 .stringLiteral("RUN"),
-                .option("--mount", "type=cache"),
+                .option(Option(key: "--mount", value: "type=cache", raw: "--mount=type=cache")),
                 .stringLiteral("build.sh --input hello"),
             ]
         ),
@@ -128,7 +128,7 @@ import Testing
                 """#,
             expectedTokens: [
                 .stringLiteral("RUN"),
-                .option("--mount", "type=cache"),
+                .option(Option(key: "--mount", value: "type=cache", raw: "--mount=type=cache")),
                 .stringLiteral("build.sh --input hello"),
             ]
         ),
@@ -139,7 +139,7 @@ import Testing
                 """#,
             expectedTokens: [
                 .stringLiteral("RUN"),
-                .option("--mount", "type=cache"),
+                .option(Option(key: "--mount", value: "type=cache", raw: "--mount=type=cache")),
                 .stringLiteral("build.sh --input hello"),
             ]
         ),
@@ -147,7 +147,7 @@ import Testing
             input: "COPY --from=alpine src /dest",
             expectedTokens: [
                 .stringLiteral("COPY"),
-                .option("--from", "alpine"),
+                .option(Option(key: "--from", value: "alpine", raw: "--from=alpine")),
                 .stringLiteral("src"),
                 .stringLiteral("/dest"),
             ]
@@ -157,7 +157,7 @@ import Testing
             input: "COPY --from=alpine src src1 src2 src3 /dest",
             expectedTokens: [
                 .stringLiteral("COPY"),
-                .option("--from", "alpine"),
+                .option(Option(key: "--from", value: "alpine", raw: "--from=alpine")),
                 .stringLiteral("src"),
                 .stringLiteral("src1"),
                 .stringLiteral("src2"),
@@ -169,7 +169,7 @@ import Testing
             input: "COPY --chown=10:11 src /dest",
             expectedTokens: [
                 .stringLiteral("COPY"),
-                .option("--chown", "10:11"),
+                .option(Option(key: "--chown", value: "10:11", raw: "--chown=10:11")),
                 .stringLiteral("src"),
                 .stringLiteral("/dest"),
             ]
@@ -178,7 +178,7 @@ import Testing
             input: "COPY --chown=bin stuff.txt /stuffdest/",
             expectedTokens: [
                 .stringLiteral("COPY"),
-                .option("--chown", "bin"),
+                .option(Option(key: "--chown", value: "bin", raw: "--chown=bin")),
                 .stringLiteral("stuff.txt"),
                 .stringLiteral("/stuffdest/"),
             ]
@@ -187,7 +187,7 @@ import Testing
             input: "COPY --chown=1 source /destination",
             expectedTokens: [
                 .stringLiteral("COPY"),
-                .option("--chown", "1"),
+                .option(Option(key: "--chown", value: "1", raw: "--chown=1")),
                 .stringLiteral("source"),
                 .stringLiteral("/destination"),
             ]
@@ -196,7 +196,7 @@ import Testing
             input: "COPY --chmod=440 src /dest/",
             expectedTokens: [
                 .stringLiteral("COPY"),
-                .option("--chmod", "440"),
+                .option(Option(key: "--chmod", value: "440", raw: "--chmod=440")),
                 .stringLiteral("src"),
                 .stringLiteral("/dest/"),
             ]
@@ -205,11 +205,52 @@ import Testing
             input: "COPY --link=false src /dest/",
             expectedTokens: [
                 .stringLiteral("COPY"),
-                .option("--link", "false"),
+                .option(Option(key: "--link", value: "false", raw: "--link=false")),
                 .stringLiteral("src"),
                 .stringLiteral("/dest/"),
             ]
         ),
+        tokenizerTestInput(
+            input: "LABEL key=value",
+            expectedTokens: [
+                .stringLiteral("LABEL"),
+                .stringLiteral("key=value"),
+            ]
+        ),
+        tokenizerTestInput(
+            input:
+                #"""
+                LABEL key=value anotherkey=anothervalue quoted="quotelabel"
+                """#,
+            expectedTokens: [
+                .stringLiteral("LABEL"),
+                .stringLiteral("key=value"),
+                .stringLiteral("anotherkey=anothervalue"),
+                .stringLiteral("quoted=\"quotelabel\""),
+            ]
+        ),
+        tokenizerTestInput(
+            input:
+                #"""
+                CMD ["executable", "param1", "param2"]
+                """#,
+            expectedTokens: [
+                .stringLiteral("CMD"),
+                .stringList(["executable", "param1", "param2"]),
+            ]
+        ),
+        tokenizerTestInput(
+            input:
+                #"""
+                CMD command param1 param2
+                """#,
+            expectedTokens: [
+                .stringLiteral("CMD"),
+                .stringLiteral("command"),
+                .stringLiteral("param1"),
+                .stringLiteral("param2"),
+            ]
+        ),
     ]
 
     @Test func testTokenization() throws {
@@ -235,9 +276,9 @@ import Testing
         return true
     }
 
-    struct TokenTest {
+    struct TokenTest: Sendable {
         let tokens: [Token]
-        let expectedInstruction: DockerInstruction
+        let expectedInstruction: any DockerInstruction
     }
 
     @Test func tokenTranslationFrom() throws {
@@ -261,7 +302,7 @@ import Testing
             TokenTest(
                 tokens: [
                     .stringLiteral("FROM"),
-                    .option("--platform", "linux/arm64"),
+                    .option(Option(key: "--platform", value: "linux/arm64", raw: "--platform=linux/arm64")),
                     .stringLiteral("alpine"),
                 ],
                 expectedInstruction: try FromInstruction(image: "alpine", platform: "linux/arm64")
@@ -282,53 +323,45 @@ import Testing
     @Test func testTokensToRunWithShellCommand() throws {
         let tokens: [Token] = [
             .stringLiteral("RUN"),
-            .option("--mount", "type=cache,target=/cache"),
+            .option(Option(key: "--mount", value: "type=cache,target=/cache", raw: "--mount=type=cache,target=/cache")),
             .stringLiteral("build.sh --input hello"),
         ]
 
         let parser = DockerfileParser()
         let actual = try parser.tokensToRunInstruction(tokens: tokens)
 
-        #expect(actual.shell)
-        #expect(actual.command == "build.sh --input hello")
+        #expect(actual.command.displayString == "build.sh --input hello")
     }
 
     @Test func testTokensToRunWithoutShell() throws {
         let command = ["build.sh", "--input", "hello"]
         let tokens: [Token] = [
             .stringLiteral("RUN"),
-            .option("--mount", "type=cache,target=/mytarget"),
+            .option(Option(key: "--mount", value: "type=cache,target=/mytarget", raw: "--mount=type=cache,target=/mytarget")),
             .stringList(command),
         ]
 
         let parser = DockerfileParser()
         let actual = try parser.tokensToRunInstruction(tokens: tokens)
 
-        #expect(actual.shell == false)
-        #expect(actual.command == command.joined(separator: " "))
+        #expect(actual.command.displayString == command.joined(separator: " "))
     }
 
     static let extraTokensTests: [[Token]] = [
         [
             .stringLiteral("RUN"),
-            .option("--mount", "type=tmpfs,size=1000"),
+            .option(Option(key: "--mount", value: "type=tmpfs,size=1000", raw: "--mount=type=tmpfs,size=1000")),
             .stringList(["build.sh", "--input", "hello"]),
             .stringLiteral("extra"),
         ],
         [
             .stringLiteral("RUN"),
-            .option("--mount", "type=bind,target=/target"),
+            .option(Option(key: "--mount", value: "type=bind,target=/target", raw: "--mount=type=bind,target=/target")),
             .stringLiteral("build.sh"),
             .stringLiteral("--input"),
             .stringLiteral("hello"),
             .stringList(["extra"]),
         ],
-        [
-            .stringLiteral("RUN"),
-            .option("--mount", "type=bind,target=/target"),
-            .stringLiteral("build.sh"),
-            .option("key", "value"),
-        ],
     ]
 
     @Test("Parsing to run instruction fails when there's extra tokens", arguments: extraTokensTests)
@@ -344,7 +377,7 @@ import Testing
             TokenTest(
                 tokens: [
                     .stringLiteral("COPY"),
-                    .option("--link", "false"),
+                    .option(Option(key: "--link", value: "false", raw: "--link=false")),
                     .stringLiteral("src"),
                     .stringLiteral("/dest/"),
                 ],
@@ -356,7 +389,7 @@ import Testing
             TokenTest(
                 tokens: [
                     .stringLiteral("COPY"),
-                    .option("--chmod", "440"),
+                    .option(Option(key: "--chmod", value: "440", raw: "--chmod=440")),
                     .stringLiteral("src"),
                     .stringLiteral("/dest"),
                 ],
@@ -369,7 +402,7 @@ import Testing
             TokenTest(
                 tokens: [
                     .stringLiteral("COPY"),
-                    .option("--chown", "11:mygroup"),
+                    .option(Option(key: "--chown", value: "11:mygroup", raw: "--chown 11:mygroup")),
                     .stringLiteral("source"),
                     .stringLiteral("destination"),
                 ],
@@ -382,7 +415,7 @@ import Testing
             TokenTest(
                 tokens: [
                     .stringLiteral("COPY"),
-                    .option("--from", "alpine"),
+                    .option(Option(key: "--from", value: "alpine", raw: "--from=alpine")),
                     .stringLiteral("src"),
                     .stringLiteral("src1"),
                     .stringLiteral("src2"),
@@ -398,7 +431,7 @@ import Testing
             TokenTest(
                 tokens: [
                     .stringLiteral("COPY"),
-                    .option("--from", "base"),
+                    .option(Option(key: "--from", value: "base", raw: "--from base")),
                     .stringLiteral("Source"),
                     .stringLiteral("Dest"),
                 ],
@@ -430,7 +463,7 @@ import Testing
         [
             // no sources
             .stringLiteral("COPY"),
-            .option("--from", "alpine"),
+            .option(Option(key: "--from", value: "alpine", raw: "--from=alpine")),
         ],
     ]
 
@@ -441,4 +474,50 @@ import Testing
             let _ = try parser.tokensToCopyInstruction(tokens: tokens)
         }
     }
+
+    static let cmdTokenTests: [TokenTest] = [
+        TokenTest(
+            tokens: [
+                .stringLiteral("CMD"),
+                .stringList(["executable", "param1", "param2"]),
+            ],
+            expectedInstruction: CMDInstruction(command: .exec(["executable", "param1", "param2"]))
+        ),
+        TokenTest(
+            tokens: [
+                .stringLiteral("CMD"),
+                .stringLiteral("command"),
+                .stringLiteral("param1"),
+                .stringLiteral("param2"),
+            ], expectedInstruction: CMDInstruction(command: .shell("command param1 param2"))
+        ),
+    ]
+
+    @Test("Successful tokens to CMD Instruction conversion", arguments: cmdTokenTests)
+    func testTokensToCMDInstruction(_ testCase: TokenTest) throws {
+        let parser = DockerfileParser()
+        let actual = try parser.tokensToCMDInstruction(tokens: testCase.tokens)
+        guard let expected = testCase.expectedInstruction as? CMDInstruction else {
+            Issue.record("Instruction is not the correct type, \(testCase.expectedInstruction)")
+            return
+        }
+        #expect(actual == expected)
+    }
+
+    @Test func testTokensToLabelInstruction() throws {
+        let tokens: [Token] = [
+            .stringLiteral("LABEL"),
+            .stringLiteral("key=value"),
+            .stringLiteral("anotherkey=anothervalue"),
+            .stringLiteral("quoted=\"quotelabel\""),
+        ]
+        let expectedLabels: [String: String] = [
+            "key": "value",
+            "anotherkey": "anothervalue",
+            "quoted": "\"quotelabel\"",
+        ]
+        let expectedInstruction = LabelInstruction(labels: expectedLabels)
+        let actual = try DockerfileParser().tokensToLabelInstruction(tokens: tokens)
+        #expect(actual == expectedInstruction)
+    }
 }
diff --git a/Tests/NativeBuilderTests/ContainerBuildParserTests/VisitorTests.swift b/Tests/NativeBuilderTests/ContainerBuildParserTests/VisitorTests.swift
index 571f23e8..859caa40 100644
--- a/Tests/NativeBuilderTests/ContainerBuildParserTests/VisitorTests.swift
+++ b/Tests/NativeBuilderTests/ContainerBuildParserTests/VisitorTests.swift
@@ -20,6 +20,8 @@ import Testing
 
 @testable import ContainerBuildParser
 
+/// VisitorTest covers testing that visiting instructions is successful and graphs
+/// made from those instructions are structured as expected.
 @Suite class VisitorTest {
     @Test func simpleVisitFrom() throws {
         let imageName = "alpine"
@@ -42,11 +44,10 @@ import Testing
 
     @Test func simpleVisitRun() throws {
         let from = try FromInstruction(image: "scratch")
-        let command = ["sh", "-c", "top"]
+        let command: Command = .exec(["sh", "-c", "top"])
         let network = "default"
         let run = try RunInstruction(
             command: command,
-            shell: false,
             rawMounts: [],
             network: network
         )
@@ -65,12 +66,8 @@ import Testing
 
         #expect(node.operation is ExecOperation)
 
-        guard let exec = node.operation as? ExecOperation else {
-            Issue.record("expected ExecOperation, instead got \(node.operation)")
-            return
-        }
-
-        #expect(exec.command.arguments == command)
+        let exec = node.operation as! ExecOperation
+        #expect(exec.command == command)
         #expect(exec.mounts.isEmpty)
         #expect(exec.network == .default)
     }
@@ -99,11 +96,9 @@ import Testing
         #expect(stage.nodes.count == 1)
         let node = stage.nodes[0]
 
-        guard let copyNode = node.operation as? FilesystemOperation else {
-            Issue.record("expected FilesystemOperation, instead got \(node.operation)")
-            return
-        }
+        #expect(node.operation is FilesystemOperation)
 
+        let copyNode = node.operation as! FilesystemOperation
         #expect(copyNode.action == .copy)
 
         let expectedSource = ContextSource(
@@ -118,4 +113,65 @@ import Testing
         let expectedPerms: Permissions = .mode(777)
         #expect(copyNode.fileMetadata.permissions == expectedPerms)
     }
+
+    @Test func simpleVisitLabel() throws {
+        let from = try FromInstruction(image: "scratch")
+        let labels = [
+            "label1": "value1",
+            "label2": "label2",
+        ]
+        let labelInst = LabelInstruction(labels: labels)
+
+        let visitor = DockerInstructionVisitor()
+        try visitor.visit(from)
+        try visitor.visit(labelInst)
+
+        let graph = try visitor.graphBuilder.build()
+        #expect(graph.stages.count == 1)
+
+        let stage = graph.stages[0]
+        #expect(stage.nodes.count == 1)
+
+        let node = stage.nodes[0]
+        #expect(node.operation is MetadataOperation)
+
+        let meta = node.operation as! MetadataOperation
+        switch meta.action {
+        case .setLabelBatch(let batch):
+            #expect(batch == labels)
+        default:
+            Issue.record("expected .setLabelBatch action type, instead got \(meta.action)")
+            return
+        }
+    }
+
+    @Test func simpleVisitCMD() throws {
+        let from = try FromInstruction(image: "scratch")
+        let rawCommand = Command.shell("./test.sh")
+        let cmd = CMDInstruction(command: rawCommand)
+
+        let visitor = DockerInstructionVisitor()
+        try visitor.visit(from)
+        try visitor.visit(cmd)
+
+        let graph = try visitor.graphBuilder.build()
+        #expect(graph.stages.count == 1)
+
+        let stage = graph.stages[0]
+        #expect(stage.nodes.count == 1)
+
+        let node = stage.nodes[0]
+        #expect(node.operation is MetadataOperation)
+
+        let meta = node.operation as! MetadataOperation
+
+        switch meta.action {
+        case .setCmd(let command):
+            #expect(command == rawCommand)
+        default:
+            Issue.record("expected .setCmd action type, instead got \(meta.action)")
+            return
+        }
+    }
+
 }