dockerfile: update env replacement efficiency

Signed-off-by: Tonis Tiigi <[email protected]>

Author: tonistiigi

frontend/dockerfile/dockerfile2llb/convert.go

@@ -11,9 +11,11 @@ import (
 	"os"
 	"path"
 	"path/filepath"
+	"runtime"
 	"sort"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/containerd/containerd/platforms"
@@ -302,7 +304,7 @@ func toDispatchState(ctx context.Context, dt []byte, opt ConvertOpt) (*dispatchS
 			info := argInfo{definition: metaArg, location: cmd.Location()}
 			if v, ok := opt.BuildArgs[metaArg.Key]; !ok {
 				if metaArg.Value != nil {
-					result, err := shlex.ProcessWordWithMatches(*metaArg.Value, metaArgsToMap(optMetaArgs))
+					result, err := shlex.ProcessWordWithMatches(*metaArg.Value, metaArgsToEnvs(optMetaArgs))
 					if err != nil {
 						return nil, parser.WithLocation(err, cmd.Location())
 					}
@@ -329,9 +331,12 @@ func toDispatchState(ctx context.Context, dt []byte, opt ConvertOpt) (*dispatchS
 
 	// set base state for every image
 	for i, st := range stages {
-		nameMatch, err := shlex.ProcessWordWithMatches(st.BaseName, metaArgsToMap(optMetaArgs))
+		nameMatch, err := shlex.ProcessWordWithMatches(st.BaseName, metaArgsToEnvs(optMetaArgs))
 		reportUnusedFromArgs(metaArgsKeys(optMetaArgs), nameMatch.Unmatched, st.Location, lint)
 		used := nameMatch.Matched
+		if used == nil {
+			used = map[string]struct{}{}
+		}
 
 		if err != nil {
 			return nil, parser.WithLocation(err, st.Location)
@@ -353,7 +358,7 @@ func toDispatchState(ctx context.Context, dt []byte, opt ConvertOpt) (*dispatchS
 		}
 
 		if v := st.Platform; v != "" {
-			platMatch, err := shlex.ProcessWordWithMatches(v, metaArgsToMap(optMetaArgs))
+			platMatch, err := shlex.ProcessWordWithMatches(v, metaArgsToEnvs(optMetaArgs))
 			reportUnusedFromArgs(metaArgsKeys(optMetaArgs), platMatch.Unmatched, st.Location, lint)
 
 			if err != nil {
@@ -644,7 +649,7 @@ func toDispatchState(ctx context.Context, dt []byte, opt ConvertOpt) (*dispatchS
 		}
 
 		// make sure that PATH is always set
-		if _, ok := shell.BuildEnvs(d.image.Config.Env)["PATH"]; !ok {
+		if _, ok := shell.EnvsFromSlice(d.image.Config.Env).Get("PATH"); !ok {
 			var osName string
 			if d.platform != nil {
 				osName = d.platform.OS
@@ -770,14 +775,38 @@ func toDispatchState(ctx context.Context, dt []byte, opt ConvertOpt) (*dispatchS
 	return target, nil
 }
 
-func metaArgsToMap(metaArgs []instructions.KeyValuePairOptional) map[string]string {
-	m := map[string]string{}
+func metaArgsToEnvs(metaArgs []instructions.KeyValuePairOptional) shell.EnvGetter {
+	return &envsFromKeyValuePairs{in: metaArgs}
+}
 
-	for _, arg := range metaArgs {
-		m[arg.Key] = arg.ValueString()
+type envsFromKeyValuePairs struct {
+	in   []instructions.KeyValuePairOptional
+	once sync.Once
+	m    map[string]string
+}
+
+func (e *envsFromKeyValuePairs) init() {
+	if len(e.in) == 0 {
+		return
 	}
+	e.m = make(map[string]string, len(e.in))
+	for _, kv := range e.in {
+		e.m[kv.Key] = kv.ValueString()
+	}
+}
+
+func (e *envsFromKeyValuePairs) Get(key string) (string, bool) {
+	e.once.Do(e.init)
+	v, ok := e.m[key] // windows: case-insensitive
+	return v, ok
+}
 
-	return m
+func (e *envsFromKeyValuePairs) Keys() []string {
+	keys := make([]string, len(e.in))
+	for i, kp := range e.in {
+		keys[i] = kp.Key
+	}
+	return keys
 }
 
 func metaArgsKeys(metaArgs []instructions.KeyValuePairOptional) []string {
@@ -840,17 +869,41 @@ type dispatchOpt struct {
 	lint              *linter.Linter
 }
 
+func getEnv(state llb.State) shell.EnvGetter {
+	return &envsFromState{state: &state}
+}
+
+type envsFromState struct {
+	state *llb.State
+	once  sync.Once
+	env   shell.EnvGetter
+}
+
+func (e *envsFromState) init() {
+	env, err := e.state.Env(context.TODO())
+	if err != nil {
+		return
+	}
+	e.env = shell.EnvsFromSlice(env)
+}
+
+func (e *envsFromState) Get(key string) (string, bool) {
+	e.once.Do(e.init)
+	return e.env.Get(key)
+}
+
+func (e *envsFromState) Keys() []string {
+	e.once.Do(e.init)
+	return e.env.Keys()
+}
+
 func dispatch(d *dispatchState, cmd command, opt dispatchOpt) error {
 	var err error
 	// ARG command value could be ignored, so defer handling the expansion error
 	_, isArg := cmd.Command.(*instructions.ArgCommand)
 	if ex, ok := cmd.Command.(instructions.SupportsSingleWordExpansion); ok && !isArg {
 		err := ex.Expand(func(word string) (string, error) {
-			env, err := d.state.Env(context.TODO())
-			if err != nil {
-				return "", err
-			}
-
+			env := getEnv(d.state)
 			newword, unmatched, err := opt.shlex.ProcessWord(word, env)
 			reportUnmatchedVariables(cmd, d.buildArgs, env, unmatched, &opt)
 			return newword, err
@@ -861,12 +914,9 @@ func dispatch(d *dispatchState, cmd command, opt dispatchOpt) error {
 	}
 	if ex, ok := cmd.Command.(instructions.SupportsSingleWordExpansionRaw); ok {
 		err := ex.ExpandRaw(func(word string) (string, error) {
-			env, err := d.state.Env(context.TODO())
-			if err != nil {
-				return "", err
-			}
 			lex := shell.NewLex('\\')
 			lex.SkipProcessQuotes = true
+			env := getEnv(d.state)
 			newword, unmatched, err := lex.ProcessWord(word, env)
 			reportUnmatchedVariables(cmd, d.buildArgs, env, unmatched, &opt)
 			return newword, err
@@ -1185,10 +1235,6 @@ func dispatchRun(d *dispatchState, c *instructions.RunCommand, proxy *llb.ProxyE
 		args = withShell(d.image, args)
 	}
 
-	env, err := d.state.Env(context.TODO())
-	if err != nil {
-		return err
-	}
 	opt = append(opt, llb.Args(args), dfCmd(c), location(dopt.sourceMap, c.Location()))
 	if d.ignoreCache {
 		opt = append(opt, llb.IgnoreCache)
@@ -1233,6 +1279,7 @@ func dispatchRun(d *dispatchState, c *instructions.RunCommand, proxy *llb.ProxyE
 	if err != nil {
 		return err
 	}
+	env := getEnv(d.state)
 	opt = append(opt, llb.WithCustomName(prefixCommand(d, uppercaseCmd(processCmdEnv(&shlex, customname, env)), d.prefixPlatform, pl, env)))
 	for _, h := range dopt.extraHosts {
 		opt = append(opt, llb.AddExtraHost(h.Host, h.IP))
@@ -1251,7 +1298,7 @@ func dispatchRun(d *dispatchState, c *instructions.RunCommand, proxy *llb.ProxyE
 	}
 
 	d.state = d.state.Run(opt...).Root()
-	return commitToHistory(&d.image, "RUN "+runCommandString(args, d.buildArgs, shell.BuildEnvs(env)), true, &d.state, d.epoch)
+	return commitToHistory(&d.image, "RUN "+runCommandString(args, d.buildArgs, env), true, &d.state, d.epoch)
 }
 
 func dispatchWorkdir(d *dispatchState, c *instructions.WorkdirCommand, commit bool, opt *dispatchOpt) error {
@@ -1297,10 +1344,7 @@ func dispatchWorkdir(d *dispatchState, c *instructions.WorkdirCommand, commit bo
 			if d.platform != nil {
 				platform = *d.platform
 			}
-			env, err := d.state.Env(context.TODO())
-			if err != nil {
-				return err
-			}
+			env := getEnv(d.state)
 			d.state = d.state.File(llb.Mkdir(wd, 0755, mkdirOpt...),
 				llb.WithCustomName(prefixCommand(d, uppercaseCmd(processCmdEnv(opt.shlex, c.String(), env)), d.prefixPlatform, &platform, env)),
 				location(opt.sourceMap, c.Location()),
@@ -1375,11 +1419,7 @@ func dispatchCopy(d *dispatchState, cfg copyConfig) error {
 		platform = *d.platform
 	}
 
-	env, err := d.state.Env(context.TODO())
-	if err != nil {
-		return err
-	}
-
+	env := getEnv(d.state)
 	name := uppercaseCmd(processCmdEnv(cfg.opt.shlex, cfg.cmdToPrint.String(), env))
 	pgName := prefixCommand(d, name, d.prefixPlatform, &platform, env)
 
@@ -1636,10 +1676,7 @@ func dispatchHealthcheck(d *dispatchState, c *instructions.HealthCheckCommand, l
 
 func dispatchExpose(d *dispatchState, c *instructions.ExposeCommand, shlex *shell.Lex) error {
 	ports := []string{}
-	env, err := d.state.Env(context.TODO())
-	if err != nil {
-		return err
-	}
+	env := getEnv(d.state)
 	for _, p := range c.Ports {
 		ps, err := shlex.ProcessWords(p, env)
 		if err != nil {
@@ -1720,10 +1757,7 @@ func dispatchArg(d *dispatchState, c *instructions.ArgCommand, opt *dispatchOpt)
 			v := opt.buildArgValues[arg.Key]
 			arg.Value = &v
 		} else if hasDefault {
-			env, err := d.state.Env(context.TODO())
-			if err != nil {
-				return err
-			}
+			env := getEnv(d.state)
 			v, unmatched, err := opt.shlex.ProcessWord(*arg.Value, env)
 			reportUnmatchedVariables(c, d.buildArgs, env, unmatched, opt)
 			if err != nil {
@@ -1824,10 +1858,10 @@ func dfCmd(cmd interface{}) llb.ConstraintsOpt {
 	})
 }
 
-func runCommandString(args []string, buildArgs []instructions.KeyValuePairOptional, envMap map[string]string) string {
+func runCommandString(args []string, buildArgs []instructions.KeyValuePairOptional, env shell.EnvGetter) string {
 	var tmpBuildEnv []string
 	for _, arg := range buildArgs {
-		v, ok := envMap[arg.Key]
+		v, ok := env.Get(arg.Key)
 		if !ok {
 			v = arg.ValueString()
 		}
@@ -2006,15 +2040,15 @@ func uppercaseCmd(str string) string {
 	return strings.Join(p, " ")
 }
 
-func processCmdEnv(shlex *shell.Lex, cmd string, env []string) string {
+func processCmdEnv(shlex *shell.Lex, cmd string, env shell.EnvGetter) string {
 	w, _, err := shlex.ProcessWord(cmd, env)
 	if err != nil {
 		return cmd
 	}
 	return w
 }
 
-func prefixCommand(ds *dispatchState, str string, prefixPlatform bool, platform *ocispecs.Platform, env []string) string {
+func prefixCommand(ds *dispatchState, str string, prefixPlatform bool, platform *ocispecs.Platform, env shell.EnvGetter) string {
 	if ds.cmdTotal == 0 {
 		return str
 	}
@@ -2057,26 +2091,26 @@ func formatTargetPlatform(base ocispecs.Platform, target *ocispecs.Platform) str
 }
 
 // platformFromEnv returns defined platforms based on TARGET* environment variables
-func platformFromEnv(env []string) *ocispecs.Platform {
+func platformFromEnv(env shell.EnvGetter) *ocispecs.Platform {
 	var p ocispecs.Platform
 	var set bool
-	for _, v := range env {
-		parts := strings.SplitN(v, "=", 2)
-		switch parts[0] {
+	for _, key := range env.Keys() {
+		switch key {
 		case "TARGETPLATFORM":
-			p, err := platforms.Parse(parts[1])
+			v, _ := env.Get(key)
+			p, err := platforms.Parse(v)
 			if err != nil {
 				continue
 			}
 			return &p
 		case "TARGETOS":
-			p.OS = parts[1]
+			p.OS, _ = env.Get(key)
 			set = true
 		case "TARGETARCH":
-			p.Architecture = parts[1]
+			p.Architecture, _ = env.Get(key)
 			set = true
 		case "TARGETVARIANT":
-			p.Variant = parts[1]
+			p.Variant, _ = env.Get(key)
 			set = true
 		}
 	}
@@ -2199,7 +2233,7 @@ func validateStageNames(stages []instructions.Stage, lint *linter.Linter) {
 	}
 }
 
-func reportUnmatchedVariables(cmd instructions.Command, buildArgs []instructions.KeyValuePairOptional, env []string, unmatched map[string]struct{}, opt *dispatchOpt) {
+func reportUnmatchedVariables(cmd instructions.Command, buildArgs []instructions.KeyValuePairOptional, env shell.EnvGetter, unmatched map[string]struct{}, opt *dispatchOpt) {
 	if len(unmatched) == 0 {
 		return
 	}
@@ -2210,15 +2244,12 @@ func reportUnmatchedVariables(cmd instructions.Command, buildArgs []instructions
 		return
 	}
 	options := metaArgsKeys(opt.metaArgs)
-	for _, envVar := range env {
-		key, _ := parseKeyValue(envVar)
-		options = append(options, key)
-	}
+	options = append(options, env.Keys()...)
 	for cmdVar := range unmatched {
 		if _, nonEnvOk := nonEnvArgs[cmdVar]; nonEnvOk {
 			continue
 		}
-		match, _ := suggest.Search(cmdVar, options, true)
+		match, _ := suggest.Search(cmdVar, options, runtime.GOOS != "windows")
 		msg := linter.RuleUndefinedVar.Format(cmdVar, match)
 		opt.lint.Run(&linter.RuleUndefinedVar, cmd.Location(), msg)
 	}