fix(staged-dockerfile): correction for ENV and ARG instructions handling

1. ARG should affect all RUN instruction digests which defined after ARG.
2. ARG could be used in the ENV instruction expansion.
3. ARG should not be expanded in the RUN instruction statically, instead expand it in runtime.

Internally args should be passed as env vars when running RUN instructions.

Signed-off-by: Timofey Kirillov <timofey.kirillov@flant.com>

pkg/build/image/image.go

@@ -3,6 +3,7 @@ package image
 import (
 	"context"
 	"fmt"
+	"strings"
 
 	"github.com/gookit/color"
 
@@ -203,12 +204,22 @@ func (i *Image) GetRebuilt() bool {
 	return i.rebuilt
 }
 
+func (i *Image) ExpandDependencies(ctx context.Context, baseEnv map[string]string) error {
+	for _, stg := range i.stages {
+		if err := stg.ExpandDependencies(ctx, i.Conveyor, baseEnv); err != nil {
+			return fmt.Errorf("unable to expand dependencies for stage %q: %w", stg.Name(), err)
+		}
+	}
+	return nil
+}
+
 func (i *Image) SetupBaseImage(ctx context.Context, storageManager manager.StorageManagerInterface, storageOpts manager.StorageOptions) error {
 	switch i.baseImageType {
 	case StageAsBaseImage:
 		i.stageAsBaseImage = i.Conveyor.GetImage(i.baseImageName).GetLastNonEmptyStage()
 		i.baseImageReference = i.stageAsBaseImage.GetStageImage().Image.Name()
 		i.baseStageImage = i.stageAsBaseImage.GetStageImage()
+
 	case ImageFromRegistryAsBaseImage:
 		if i.IsDockerfileImage && i.dockerfileExpanderFactory != nil {
 			dependenciesArgs := stage.ResolveDependenciesArgs(i.DockerfileImageConfig.Dependencies, i.Conveyor)
@@ -218,27 +229,33 @@ func (i *Image) SetupBaseImage(ctx context.Context, storageManager manager.Stora
 			}
 			i.baseImageReference = ref
 		}
+
 		i.baseStageImage = i.Conveyor.GetOrCreateStageImage(i.baseImageReference, nil, nil, i)
+
+		info, err := storageManager.GetImageInfo(ctx, i.baseImageReference, storageOpts)
+		if err != nil {
+			return fmt.Errorf("unable to get base image %q manifest: %w", i.baseImageReference, err)
+		}
+
+		i.baseStageImage.Image.SetStageDescription(&image.StageDescription{
+			StageID: nil, // this is not a stage actually, TODO
+			Info:    info,
+		})
+
+		// for _, expression := range info.OnBuild {
+		// 	fmt.Printf(">> %q\n", expression)
+		// }
+
 	case NoBaseImage:
+
 	default:
 		panic(fmt.Sprintf("unknown base image type %q", i.baseImageType))
 	}
 
-	if i.IsDockerfileImage && i.DockerfileImageConfig.Staged {
-		switch i.baseImageType {
-		case StageAsBaseImage, ImageFromRegistryAsBaseImage:
-
-			fmt.Printf("-- %s SetupBaseImage %q\n", i.Name, i.baseImageReference)
-
-			info, err := storageManager.GetImageInfo(ctx, i.baseImageReference, storageOpts)
-			if err != nil {
-				return fmt.Errorf("unable to get base image %q manifest: %w", i.baseImageReference, err)
-			}
-
-			fmt.Printf("-- %s SetupBaseImage %q -> %#v\n", i.Name, i.baseImageReference, info)
-			for _, expression := range info.OnBuild {
-				fmt.Printf(">> %q\n", expression)
-			}
+	switch i.baseImageType {
+	case StageAsBaseImage, ImageFromRegistryAsBaseImage:
+		if err := i.ExpandDependencies(ctx, EnvToMap(i.baseStageImage.Image.GetStageDescription().Info.Env)); err != nil {
+			return err
 		}
 	}
 
@@ -257,6 +274,8 @@ func (i *Image) GetBaseImageReference() string {
 func (i *Image) FetchBaseImage(ctx context.Context) error {
 	switch i.baseImageType {
 	case ImageFromRegistryAsBaseImage:
+		// TODO: Refactor, move manifest fetching into SetupBaseImage, only pull image in FetchBaseImage method
+
 		if info, err := i.ContainerBackend.GetImageInfo(ctx, i.baseStageImage.Image.Name(), container_backend.GetImageInfoOpts{}); err != nil {
 			return fmt.Errorf("unable to inspect local image %s: %w", i.baseStageImage.Image.Name(), err)
 		} else if info != nil {
@@ -348,3 +367,22 @@ func (i *Image) getFromBaseImageIdFromRegistry(ctx context.Context, reference st
 
 	return i.baseImageRepoId, nil
 }
+
+func EnvToMap(env []string) map[string]string {
+	res := make(map[string]string)
+	for _, kv := range env {
+		k, v := parseKeyValue(kv)
+		res[k] = v
+	}
+	return res
+}
+
+func parseKeyValue(env string) (string, string) {
+	parts := strings.SplitN(env, "=", 2)
+	v := ""
+	if len(parts) > 1 {
+		v = parts[1]
+	}
+
+	return parts[0], v
+}

pkg/build/stage/base.go

@@ -144,6 +144,10 @@ func (s *BaseStage) Name() StageName {
 	panic("name must be defined!")
 }
 
+func (s *BaseStage) ExpandDependencies(ctx context.Context, c Conveyor, baseEnv map[string]string) error {
+	return nil
+}
+
 func (s *BaseStage) FetchDependencies(_ context.Context, _ Conveyor, _ container_backend.ContainerBackend, _ docker_registry.ApiInterface) error {
 	return nil
 }

pkg/build/stage/instruction/add.go

@@ -24,10 +24,7 @@ func NewAdd(i *dockerfile.DockerfileStageInstruction[*instructions.AddCommand],
 }
 
 func (stg *Add) GetDependencies(ctx context.Context, c stage.Conveyor, cb container_backend.ContainerBackend, prevImage, prevBuiltImage *stage.StageImage, buildContextArchive container_backend.BuildContextArchiver) (string, error) {
-	args, err := stg.getDependencies(ctx, c, cb, prevImage, prevBuiltImage, buildContextArchive, stg)
-	if err != nil {
-		return "", err
-	}
+	var args []string
 
 	args = append(args, append([]string{"Sources"}, stg.instruction.Data.Sources()...)...)
 	args = append(args, "Dest", stg.instruction.Data.Dest())

pkg/build/stage/instruction/base.go

@@ -17,6 +17,9 @@ type Base[T dockerfile.InstructionDataInterface, BT container_backend.Instructio
 	backendInstruction BT
 	dependencies       []*config.Dependency
 	hasPrevStage       bool
+
+	baseEnv     map[string]string
+	expandedEnv map[string]string
 }
 
 func NewBase[T dockerfile.InstructionDataInterface, BT container_backend.InstructionInterface](instruction *dockerfile.DockerfileStageInstruction[T], backendInstruction BT, dependencies []*config.Dependency, hasPrevStage bool, opts *stage.BaseStageOptions) *Base[T, BT] {
@@ -41,24 +44,65 @@ func (stg *Base[T, BT]) UsesBuildContext() bool {
 	return stg.backendInstruction.UsesBuildContext()
 }
 
-func (stg *Base[T, BT]) getDependencies(ctx context.Context, c stage.Conveyor, cb container_backend.ContainerBackend, prevImage, prevBuiltImage *stage.StageImage, buildContextArchive container_backend.BuildContextArchiver, expander InstructionExpander) ([]string, error) {
-	if err := expander.ExpandInstruction(ctx, c, cb, prevImage, prevBuiltImage, buildContextArchive); err != nil {
-		return nil, fmt.Errorf("unable to expand instruction %q: %w", stg.instruction.Data.Name(), err)
+func (stg *Base[T, BT]) ExpandDependencies(ctx context.Context, c stage.Conveyor, baseEnv map[string]string) error {
+	return stg.doExpandDependencies(ctx, c, baseEnv, stg)
+}
+
+func (stg *Base[T, BT]) doExpandDependencies(ctx context.Context, c stage.Conveyor, baseEnv map[string]string, expander InstructionExpander) error {
+	if err := stg.expandBaseEnv(baseEnv); err != nil {
+		return fmt.Errorf("2nd stage env expansion failed: %w", err)
 	}
-	return nil, nil
+	if err := expander.ExpandInstruction(c, stg.GetExpandedEnv(c)); err != nil {
+		return fmt.Errorf("2nd stage instruction expansion failed: %w", err)
+	}
+	return nil
 }
 
 func (stg *Base[T, BT]) PrepareImage(ctx context.Context, c stage.Conveyor, cb container_backend.ContainerBackend, prevBuiltImage, stageImage *stage.StageImage, buildContextArchive container_backend.BuildContextArchiver) error {
 	stageImage.Builder.DockerfileStageBuilder().AppendInstruction(stg.backendInstruction)
 	return nil
 }
 
-func (stg *Base[T, BT]) ExpandInstruction(ctx context.Context, c stage.Conveyor, cb container_backend.ContainerBackend, prevBuiltImage, stageImage *stage.StageImage, buildContextArchive container_backend.BuildContextArchiver) error {
-	dependenciesArgs := stage.ResolveDependenciesArgs(stg.dependencies, c)
-	// NOTE: do not skip unset envs during second stage expansion
-	return stg.instruction.Expand(dependenciesArgs, dockerfile.ExpandOptions{SkipUnsetEnv: false})
+func (stg *Base[T, BT]) expandBaseEnv(baseEnv map[string]string) error {
+	env := make(map[string]string)
+
+	// NOTE: default fallback builtin PATH
+	env["PATH"] = "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"
+	for k, v := range baseEnv {
+		env[k] = v
+	}
+
+	// NOTE: this is 2nd stage expansion (without skipping unset envs)
+	instructionEnv, err := stg.instruction.ExpandEnv(env, dockerfile.ExpandOptions{})
+	if err != nil {
+		return fmt.Errorf("error expanding env: %w", err)
+	}
+	for k, v := range instructionEnv {
+		env[k] = v
+	}
+
+	stg.baseEnv = baseEnv
+	stg.expandedEnv = env
+
+	return nil
+}
+
+func (stg *Base[T, BT]) GetExpandedEnv(c stage.Conveyor) map[string]string {
+	env := make(map[string]string)
+	for k, v := range stg.expandedEnv {
+		env[k] = v
+	}
+	for k, v := range stage.ResolveDependenciesArgs(stg.dependencies, c) {
+		env[k] = v
+	}
+	return env
+}
+
+func (stg *Base[T, BT]) ExpandInstruction(_ stage.Conveyor, env map[string]string) error {
+	// NOTE: this is 2nd stage expansion (without skipping unset envs)
+	return stg.instruction.Expand(env, dockerfile.ExpandOptions{})
 }
 
 type InstructionExpander interface {
-	ExpandInstruction(ctx context.Context, c stage.Conveyor, cb container_backend.ContainerBackend, prevBuiltImage, stageImage *stage.StageImage, buildContextArchive container_backend.BuildContextArchiver) error
+	ExpandInstruction(c stage.Conveyor, env map[string]string) error
 }

pkg/build/stage/instruction/cmd.go

@@ -23,10 +23,7 @@ func NewCmd(i *dockerfile.DockerfileStageInstruction[*instructions.CmdCommand],
 }
 
 func (stg *Cmd) GetDependencies(ctx context.Context, c stage.Conveyor, cb container_backend.ContainerBackend, prevImage, prevBuiltImage *stage.StageImage, buildContextArchive container_backend.BuildContextArchiver) (string, error) {
-	args, err := stg.getDependencies(ctx, c, cb, prevImage, prevBuiltImage, buildContextArchive, stg)
-	if err != nil {
-		return "", err
-	}
+	var args []string
 
 	args = append(args, append([]string{"Cmd"}, stg.instruction.Data.CmdLine...)...)
 	args = append(args, "PrependShell", fmt.Sprintf("%v", stg.instruction.Data.PrependShell))

pkg/build/stage/instruction/copy.go

@@ -22,8 +22,12 @@ func NewCopy(i *dockerfile.DockerfileStageInstruction[*instructions.CopyCommand]
 	return &Copy{Base: NewBase(i, backend_instruction.NewCopy(i.Data), dependencies, hasPrevStage, opts)}
 }
 
-func (stg *Copy) ExpandInstruction(ctx context.Context, c stage.Conveyor, cb container_backend.ContainerBackend, prevBuiltImage, stageImage *stage.StageImage, buildContextArchive container_backend.BuildContextArchiver) error {
-	if err := stg.Base.ExpandInstruction(ctx, c, cb, prevBuiltImage, stageImage, buildContextArchive); err != nil {
+func (stg *Copy) ExpandDependencies(ctx context.Context, c stage.Conveyor, baseEnv map[string]string) error {
+	return stg.doExpandDependencies(ctx, c, baseEnv, stg)
+}
+
+func (stg *Copy) ExpandInstruction(c stage.Conveyor, env map[string]string) error {
+	if err := stg.Base.ExpandInstruction(c, env); err != nil {
 		return err
 	}
 
@@ -38,10 +42,7 @@ func (stg *Copy) ExpandInstruction(ctx context.Context, c stage.Conveyor, cb con
 }
 
 func (stg *Copy) GetDependencies(ctx context.Context, c stage.Conveyor, cb container_backend.ContainerBackend, prevImage, prevBuiltImage *stage.StageImage, buildContextArchive container_backend.BuildContextArchiver) (string, error) {
-	args, err := stg.getDependencies(ctx, c, cb, prevImage, prevBuiltImage, buildContextArchive, stg)
-	if err != nil {
-		return "", err
-	}
+	var args []string
 
 	args = append(args, "From", stg.instruction.Data.From)
 	args = append(args, append([]string{"Sources"}, stg.instruction.Data.Sources()...)...)

pkg/build/stage/instruction/entrypoint.go

@@ -23,12 +23,10 @@ func NewEntrypoint(i *dockerfile.DockerfileStageInstruction[*instructions.Entryp
 }
 
 func (stg *Entrypoint) GetDependencies(ctx context.Context, c stage.Conveyor, cb container_backend.ContainerBackend, prevImage, prevBuiltImage *stage.StageImage, buildContextArchive container_backend.BuildContextArchiver) (string, error) {
-	args, err := stg.getDependencies(ctx, c, cb, prevImage, prevBuiltImage, buildContextArchive, stg)
-	if err != nil {
-		return "", err
-	}
+	var args []string
 
 	args = append(args, append([]string{"Entrypoint"}, stg.instruction.Data.CmdLine...)...)
 	args = append(args, "PrependShell", fmt.Sprintf("%v", stg.instruction.Data.PrependShell))
+
 	return util.Sha256Hash(args...), nil
 }

pkg/build/stage/instruction/env.go

@@ -22,16 +22,14 @@ func NewEnv(i *dockerfile.DockerfileStageInstruction[*instructions.EnvCommand],
 }
 
 func (stg *Env) GetDependencies(ctx context.Context, c stage.Conveyor, cb container_backend.ContainerBackend, prevImage, prevBuiltImage *stage.StageImage, buildContextArchive container_backend.BuildContextArchiver) (string, error) {
-	args, err := stg.getDependencies(ctx, c, cb, prevImage, prevBuiltImage, buildContextArchive, stg)
-	if err != nil {
-		return "", err
-	}
+	var args []string
 
 	if len(stg.instruction.Data.Env) > 0 {
 		args = append(args, "Env")
 		for _, item := range stg.instruction.Data.Env {
 			args = append(args, item.Key, item.Value)
 		}
 	}
+
 	return util.Sha256Hash(args...), nil
 }

pkg/build/stage/instruction/expose.go

@@ -22,11 +22,9 @@ func NewExpose(i *dockerfile.DockerfileStageInstruction[*instructions.ExposeComm
 }
 
 func (stg *Expose) GetDependencies(ctx context.Context, c stage.Conveyor, cb container_backend.ContainerBackend, prevImage, prevBuiltImage *stage.StageImage, buildContextArchive container_backend.BuildContextArchiver) (string, error) {
-	args, err := stg.getDependencies(ctx, c, cb, prevImage, prevBuiltImage, buildContextArchive, stg)
-	if err != nil {
-		return "", err
-	}
+	var args []string
 
 	args = append(args, append([]string{"Ports"}, stg.instruction.Data.Ports...)...)
+
 	return util.Sha256Hash(args...), nil
 }

pkg/build/stage/instruction/healthcheck.go

@@ -23,15 +23,13 @@ func NewHealthcheck(i *dockerfile.DockerfileStageInstruction[*instructions.Healt
 }
 
 func (stg *Healthcheck) GetDependencies(ctx context.Context, c stage.Conveyor, cb container_backend.ContainerBackend, prevImage, prevBuiltImage *stage.StageImage, buildContextArchive container_backend.BuildContextArchiver) (string, error) {
-	args, err := stg.getDependencies(ctx, c, cb, prevImage, prevBuiltImage, buildContextArchive, stg)
-	if err != nil {
-		return "", err
-	}
+	var args []string
 
 	args = append(args, append([]string{"Test"}, stg.instruction.Data.Health.Test...)...)
 	args = append(args, "Interval", stg.instruction.Data.Health.Interval.String())
 	args = append(args, "Timeout", stg.instruction.Data.Health.Timeout.String())
 	args = append(args, "StartPeriod", stg.instruction.Data.Health.StartPeriod.String())
 	args = append(args, "Retries", fmt.Sprintf("%d", stg.instruction.Data.Health.Retries))
+
 	return util.Sha256Hash(args...), nil
 }

pkg/build/stage/instruction/label.go

@@ -22,16 +22,14 @@ func NewLabel(i *dockerfile.DockerfileStageInstruction[*instructions.LabelComman
 }
 
 func (stg *Label) GetDependencies(ctx context.Context, c stage.Conveyor, cb container_backend.ContainerBackend, prevImage, prevBuiltImage *stage.StageImage, buildContextArchive container_backend.BuildContextArchiver) (string, error) {
-	args, err := stg.getDependencies(ctx, c, cb, prevImage, prevBuiltImage, buildContextArchive, stg)
-	if err != nil {
-		return "", err
-	}
+	var args []string
 
 	if len(stg.instruction.Data.Labels) > 0 {
 		args = append(args, "Labels")
 		for _, item := range stg.instruction.Data.Labels {
 			args = append(args, item.Key, item.Value)
 		}
 	}
+
 	return util.Sha256Hash(args...), nil
 }

pkg/build/stage/instruction/maintainer.go

@@ -22,11 +22,9 @@ func NewMaintainer(i *dockerfile.DockerfileStageInstruction[*instructions.Mainta
 }
 
 func (stg *Maintainer) GetDependencies(ctx context.Context, c stage.Conveyor, cb container_backend.ContainerBackend, prevImage, prevBuiltImage *stage.StageImage, buildContextArchive container_backend.BuildContextArchiver) (string, error) {
-	args, err := stg.getDependencies(ctx, c, cb, prevImage, prevBuiltImage, buildContextArchive, stg)
-	if err != nil {
-		return "", err
-	}
+	var args []string
 
 	args = append(args, "Maintainer", stg.instruction.Data.Maintainer)
+
 	return util.Sha256Hash(args...), nil
 }

pkg/build/stage/instruction/on_build.go

@@ -22,11 +22,9 @@ func NewOnBuild(i *dockerfile.DockerfileStageInstruction[*instructions.OnbuildCo
 }
 
 func (stg *OnBuild) GetDependencies(ctx context.Context, c stage.Conveyor, cb container_backend.ContainerBackend, prevImage, prevBuiltImage *stage.StageImage, buildContextArchive container_backend.BuildContextArchiver) (string, error) {
-	args, err := stg.getDependencies(ctx, c, cb, prevImage, prevBuiltImage, buildContextArchive, stg)
-	if err != nil {
-		return "", err
-	}
+	var args []string
 
 	args = append(args, "Expression", stg.instruction.Data.Expression)
+
 	return util.Sha256Hash(args...), nil
 }