Merge pull request #2888 from Shell32-Natsu/envs

explicitly specify envs to be exported in function

kyaml/fn/runtime/container/container.go

@@ -5,8 +5,6 @@ package container
 
 import (
 	"fmt"
-	"os"
-	"strings"
 
 	runtimeexec "sigs.k8s.io/kustomize/kyaml/fn/runtime/exec"
 	"sigs.k8s.io/kustomize/kyaml/fn/runtime/runtimeutil"
@@ -154,8 +152,6 @@ func (c *Filter) setupExec() {
 	c.Exec.Args = args
 }
 
-var tmpDirEnvKey string = "TMPDIR"
-
 // getArgs returns the command + args to run to spawn the container
 func (c *Filter) getCommand() (string, []string) {
 	// run the container using docker.  this is simpler than using the docker
@@ -177,19 +173,7 @@ func (c *Filter) getCommand() (string, []string) {
 		args = append(args, "--mount", storageMount.String())
 	}
 
-	// TODO: put these env processes into a separate function and call it in the outside of
-	// getCommand
-	os.Setenv("LOG_TO_STDERR", "true")
-	os.Setenv("STRUCTURED_RESULTS", "true")
-
-	// export the local environment vars to the container
-	for _, pair := range os.Environ() {
-		items := strings.Split(pair, "=")
-		if items[0] == "" || items[1] == "" || items[0] == tmpDirEnvKey {
-			continue
-		}
-		args = append(args, "-e", items[0])
-	}
+	args = append(args, runtimeutil.NewContainerEnvFromStringSlice(c.Env).GetDockerFlags()...)
 	a := append(args, c.Image)
 	return "docker", a
 }

kyaml/fn/runtime/container/container_test.go

@@ -6,8 +6,6 @@ package container
 import (
 	"bytes"
 	"fmt"
-	"os"
-	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -136,19 +134,11 @@ metadata:
 				t.FailNow()
 			}
 			tt.instance.Exec.FunctionConfig = cfg
-
-			os.Setenv("KYAML_TEST", "FOO")
+			tt.instance.Env = append(tt.instance.Env, "KYAML_TEST=FOO")
 			tt.instance.setupExec()
 
-			// configure expected env
-			for _, e := range os.Environ() {
-				// the process env
-				parts := strings.Split(e, "=")
-				if parts[0] == "" || parts[1] == "" || parts[0] == tmpDirEnvKey {
-					continue
-				}
-				tt.expectedArgs = append(tt.expectedArgs, "-e", parts[0])
-			}
+			tt.expectedArgs = append(tt.expectedArgs,
+				runtimeutil.NewContainerEnvFromStringSlice(tt.instance.Env).GetDockerFlags()...)
 			tt.expectedArgs = append(tt.expectedArgs, tt.instance.Image)
 
 			if !assert.Equal(t, "docker", tt.instance.Exec.Path) {
@@ -248,15 +238,3 @@ func TestFilter_ExitCode(t *testing.T) {
 		t.FailNow()
 	}
 }
-
-func TestIgnoreEnv(t *testing.T) {
-	os.Setenv(tmpDirEnvKey, "")
-
-	fltr := Filter{ContainerSpec: runtimeutil.ContainerSpec{Image: "example.com:version"}}
-	_, args := fltr.getCommand()
-	for _, arg := range args {
-		if arg == tmpDirEnvKey {
-			t.Fatalf("%s should not be exported to container", tmpDirEnvKey)
-		}
-	}
-}

kyaml/fn/runtime/runtimeutil/functiontypes.go

@@ -6,6 +6,7 @@ package runtimeutil
 import (
 	"fmt"
 	"os"
+	"sort"
 	"strings"
 
 	"sigs.k8s.io/kustomize/kyaml/yaml"
@@ -40,6 +41,103 @@ const (
 	NetworkNameNone  ContainerNetworkName = "none"
 	NetworkNameEmpty ContainerNetworkName = ""
 )
+const defaultEnvValue string = "true"
+
+// ContainerEnv defines the environment present in a container.
+type ContainerEnv struct {
+	// EnvVars is a key-value map that will be set as env in container
+	EnvVars map[string]string
+
+	// VarsToExport are only env key. Value will be the value in the host system
+	VarsToExport []string
+}
+
+// GetDockerFlags returns docker run style env flags
+func (ce *ContainerEnv) GetDockerFlags() []string {
+	envs := ce.EnvVars
+	if envs == nil {
+		envs = make(map[string]string)
+	}
+
+	flags := []string{}
+	// return in order to keep consistent among different runs
+	keys := []string{}
+	for k := range envs {
+		keys = append(keys, k)
+	}
+	sort.Strings(keys)
+	for _, key := range keys {
+		flags = append(flags, "-e", key+"="+envs[key])
+	}
+
+	for _, key := range ce.VarsToExport {
+		flags = append(flags, "-e", key)
+	}
+
+	return flags
+}
+
+// AddKeyValue adds a key-value pair into the envs
+func (ce *ContainerEnv) AddKeyValue(key, value string) {
+	if ce.EnvVars == nil {
+		ce.EnvVars = make(map[string]string)
+	}
+	ce.EnvVars[key] = value
+}
+
+// HasExportedKey returns true if the key is a exported key
+func (ce *ContainerEnv) HasExportedKey(key string) bool {
+	for _, k := range ce.VarsToExport {
+		if k == key {
+			return true
+		}
+	}
+	return false
+}
+
+// AddKey adds a key into the envs
+func (ce *ContainerEnv) AddKey(key string) {
+	if !ce.HasExportedKey(key) {
+		ce.VarsToExport = append(ce.VarsToExport, key)
+	}
+}
+
+// Raw returns a slice of string which represents the envs.
+// Example: [foo=bar, baz]
+func (ce *ContainerEnv) Raw() []string {
+	var ret []string
+	for k, v := range ce.EnvVars {
+		ret = append(ret, k+"="+v)
+	}
+
+	ret = append(ret, ce.VarsToExport...)
+	return ret
+}
+
+// NewContainerEnv returns a pointer to a new ContainerEnv
+func NewContainerEnv() *ContainerEnv {
+	var ce ContainerEnv
+	ce.EnvVars = make(map[string]string)
+	// default envs
+	ce.EnvVars["LOG_TO_STDERR"] = defaultEnvValue
+	ce.EnvVars["STRUCTURED_RESULTS"] = defaultEnvValue
+	return &ce
+}
+
+// NewContainerEnvFromStringSlice returns a new ContainerEnv pointer with parsing
+// input envStr. envStr example: ["foo=bar", "baz"]
+func NewContainerEnvFromStringSlice(envStr []string) *ContainerEnv {
+	ce := NewContainerEnv()
+	for _, e := range envStr {
+		parts := strings.SplitN(e, "=", 2)
+		if len(parts) == 1 {
+			ce.AddKey(e)
+		} else {
+			ce.AddKeyValue(parts[0], parts[1])
+		}
+	}
+	return ce
+}
 
 // FunctionSpec defines a spec for running a function
 type FunctionSpec struct {
@@ -75,6 +173,9 @@ type ContainerSpec struct {
 
 	// User is the username/uid that application runs as in continer
 	User ContainerUser `json:"user,omitempty" yaml:"user,omitempty"`
+
+	// Env is a slice of env string that will be exposed to container
+	Env []string `json:"envs,omitempty" yaml:"envs,omitempty"`
 }
 
 // ContainerNetwork

kyaml/fn/runtime/runtimeutil/runtimeutil_test.go

@@ -1428,3 +1428,92 @@ func Test_StringToStorageMount(t *testing.T) {
 		assert.Equal(t, tc.expectedOut, (&s).String())
 	}
 }
+
+func TestContainerEnvGetDockerFlags(t *testing.T) {
+	tests := []struct {
+		input  *ContainerEnv
+		output []string
+	}{
+		{
+			input:  NewContainerEnvFromStringSlice([]string{"foo=bar"}),
+			output: []string{"-e", "LOG_TO_STDERR=true", "-e", "STRUCTURED_RESULTS=true", "-e", "foo=bar"},
+		},
+		{
+			input:  NewContainerEnvFromStringSlice([]string{"foo"}),
+			output: []string{"-e", "LOG_TO_STDERR=true", "-e", "STRUCTURED_RESULTS=true", "-e", "foo"},
+		},
+		{
+			input:  NewContainerEnvFromStringSlice([]string{"foo=bar", "baz"}),
+			output: []string{"-e", "LOG_TO_STDERR=true", "-e", "STRUCTURED_RESULTS=true", "-e", "foo=bar", "-e", "baz"},
+		},
+		{
+			input:  NewContainerEnv(),
+			output: []string{"-e", "LOG_TO_STDERR=true", "-e", "STRUCTURED_RESULTS=true"},
+		},
+	}
+
+	for _, tc := range tests {
+		flags := tc.input.GetDockerFlags()
+		assert.Equal(t, tc.output, flags)
+	}
+}
+
+func TestGetContainerEnv(t *testing.T) {
+	tests := []struct {
+		input    string
+		expected ContainerEnv
+	}{
+		{
+			input: `
+apiVersion: v1
+kind: Foo
+metadata:
+  name: foo
+  configFn:
+    container:
+      image: gcr.io/kustomize-functions/example-tshirt:v0.1.0
+      envs:
+      - foo=bar
+`,
+			expected: *NewContainerEnvFromStringSlice([]string{"foo=bar"}),
+		},
+		{
+			input: `
+apiVersion: v1
+kind: Foo
+metadata:
+  name: foo
+  configFn:
+    container:
+      image: gcr.io/kustomize-functions/example-tshirt:v0.1.0
+      envs:
+      - foo=bar
+      - baz
+`,
+			expected: *NewContainerEnvFromStringSlice([]string{"foo=bar", "baz"}),
+		},
+		{
+			input: `
+apiVersion: v1
+kind: Foo
+metadata:
+  name: foo
+  configFn:
+    container:
+      image: gcr.io/kustomize-functions/example-tshirt:v0.1.0
+      envs:
+      - KUBECONFIG
+`,
+			expected: *NewContainerEnvFromStringSlice([]string{"KUBECONFIG"}),
+		},
+	}
+
+	for _, tc := range tests {
+		cfg, err := yaml.Parse(tc.input)
+		if !assert.NoError(t, err) {
+			return
+		}
+		fn := GetFunctionSpec(cfg)
+		assert.Equal(t, tc.expected, *NewContainerEnvFromStringSlice(fn.Container.Env))
+	}
+}

kyaml/runfn/runfn.go

@@ -90,6 +90,9 @@ type RunFns struct {
 
 	// User username used to run the application in container,
 	User runtimeutil.ContainerUser
+
+	// Env contains environment variables that will be exported to container
+	Env []string
 }
 
 // Execute runs the command
@@ -269,6 +272,22 @@ func (r RunFns) getFunctionsFromFunctions() ([]kio.Filter, error) {
 	return r.getFunctionFilters(true, r.Functions...)
 }
 
+// mergeContainerEnv will merge the envs specified by command line (imperative) and config
+// file (declarative). If they have same key, the imperative value will be respected.
+func (r RunFns) mergeContainerEnv(envs []string) []string {
+	imperative := runtimeutil.NewContainerEnvFromStringSlice(r.Env)
+	declarative := runtimeutil.NewContainerEnvFromStringSlice(envs)
+	for key, value := range imperative.EnvVars {
+		declarative.AddKeyValue(key, value)
+	}
+
+	for _, key := range imperative.VarsToExport {
+		declarative.AddKey(key)
+	}
+
+	return declarative.Raw()
+}
+
 func (r RunFns) getFunctionFilters(global bool, fns ...*yaml.RNode) (
 	[]kio.Filter, error) {
 	var fltrs []kio.Filter
@@ -282,10 +301,11 @@ func (r RunFns) getFunctionFilters(global bool, fns ...*yaml.RNode) (
 			}
 			spec.Container.Network.Name = runtimeutil.ContainerNetworkName(r.NetworkName)
 		}
-		// command line username has higher priority
-		if r.User != "" {
+		// command line username and envs has higher priority
+		if !r.User.IsEmpty() {
 			spec.Container.User = r.User
 		}
+		spec.Container.Env = r.mergeContainerEnv(spec.Container.Env)
 
 		c, err := r.functionFilterProvider(*spec, api)
 		if err != nil {
@@ -394,6 +414,7 @@ func (r *RunFns) ffp(spec runtimeutil.FunctionSpec, api *yaml.RNode) (kio.Filter
 			Network:       spec.Container.Network,
 			StorageMounts: r.StorageMounts,
 			User:          spec.Container.User,
+			Env:           spec.Container.Env,
 		})
 		cf := &c
 		cf.Exec.FunctionConfig = api

kyaml/runfn/runfn_test.go

@@ -986,3 +986,55 @@ func getFilterProvider(t *testing.T) func(runtimeutil.FunctionSpec, *yaml.RNode)
 		}, nil
 	}
 }
+
+func TestRunfns_mergeContainerEnv(t *testing.T) {
+	testcases := []struct {
+		name      string
+		instance  RunFns
+		inputEnvs []string
+		expect    runtimeutil.ContainerEnv
+	}{
+		{
+			name:     "all empty",
+			instance: RunFns{},
+			expect:   *runtimeutil.NewContainerEnv(),
+		},
+		{
+			name:      "empty command line envs",
+			instance:  RunFns{},
+			inputEnvs: []string{"foo=bar"},
+			expect:    *runtimeutil.NewContainerEnvFromStringSlice([]string{"foo=bar"}),
+		},
+		{
+			name: "empty declarative envs",
+			instance: RunFns{
+				Env: []string{"foo=bar"},
+			},
+			expect: *runtimeutil.NewContainerEnvFromStringSlice([]string{"foo=bar"}),
+		},
+		{
+			name: "same key",
+			instance: RunFns{
+				Env: []string{"foo=bar", "foo"},
+			},
+			inputEnvs: []string{"foo=bar1", "bar"},
+			expect:    *runtimeutil.NewContainerEnvFromStringSlice([]string{"foo=bar", "bar", "foo"}),
+		},
+		{
+			name: "same exported key",
+			instance: RunFns{
+				Env: []string{"foo=bar", "foo"},
+			},
+			inputEnvs: []string{"foo1=bar1", "foo"},
+			expect:    *runtimeutil.NewContainerEnvFromStringSlice([]string{"foo=bar", "foo1=bar1", "foo"}),
+		},
+	}
+
+	for i := range testcases {
+		tc := testcases[i]
+		t.Run(tc.name, func(t *testing.T) {
+			envs := tc.instance.mergeContainerEnv(tc.inputEnvs)
+			assert.Equal(t, tc.expect.GetDockerFlags(), runtimeutil.NewContainerEnvFromStringSlice(envs).GetDockerFlags())
+		})
+	}
+}