dockerfile: bug fix for env parsing (#6422)

buildkit changed the AST format for the Env command

fixes #6421

Signed-off-by: Nick Santos <nick.santos@docker.com>

internal/dockerfile/ast.go

@@ -210,8 +210,8 @@ func (a AST) printNode(node *parser.Node, writer io.Writer) (int, error) {
 	// https://github.com/moby/buildkit/blob/2ec7d53b00f24624cda0adfbdceed982623a93b3/frontend/dockerfile/parser/parser.go#L152
 	case command.Cmd, command.Entrypoint, command.Run, command.Shell:
 		v = fmtCmd(node)
-	case command.Label:
-		v = fmtLabel(node)
+	case command.Label, command.Env:
+		v = fmtKeyValuePairs(node)
 	default:
 		v = fmtDefault(node)
 	}
@@ -286,7 +286,7 @@ func fmtDefault(node *parser.Node) string {
 	return appendHeredocs(node, strings.Join(cmd, " "))
 }
 
-func fmtLabel(node *parser.Node) string {
+func fmtKeyValuePairs(node *parser.Node) string {
 	cmd := getCmd(node)
 	assignments := []string{cmd[0]}
 	for i := 1; i < len(cmd); i += 3 {

internal/dockerfile/ast_test.go

@@ -125,6 +125,33 @@ FROM golang:10
 	assertPrint(t, orig, expected)
 }
 
+func TestPrintArgAndEnv(t *testing.T) {
+	assertPrintSame(t, `
+
+FROM busybox
+ARG APP_NAME=TheAppName
+
+ENV APP_NAME=$APP_NAME
+`)
+}
+
+func TestPrintMultipleEnv(t *testing.T) {
+	assertPrintSame(t, `
+FROM busybox
+ENV VAR_1=foo VAR_2=bar
+`)
+}
+
+func TestPrintReformatEnv(t *testing.T) {
+	assertPrint(t, `
+FROM busybox
+ENV APP_NAME TheAppName
+`, `
+FROM busybox
+ENV APP_NAME=TheAppName
+`)
+}
+
 // Convert the dockerfile into an AST, print it, and then
 // assert that the result is the same as the original.
 func assertPrintSame(t *testing.T, original string) {