Include default envs from containerd.

Random-Liu · Random-Liu · commit 5edec1d8a6ef · 2019-02-12T19:14:34.000-08:00
Signed-off-by: Lantao Liu &lt;lantaol@google.com&gt;
diff --git a/pkg/server/helpers.go b/pkg/server/helpers.go
@@ -446,10 +446,18 @@ type generator struct {
 }
 
 func newCustomGenerator(g generate.Generator) generator {
-	return generator{
+	cg := generator{
 		Generator: g,
 		envCache:  make(map[string]int),
 	}
+	spec := g.Spec()
+	if spec != nil && spec.Process != nil {
+		for i, env := range spec.Process.Env {
+			kv := strings.SplitN(env, "=", 2)
+			cg.envCache[kv[0]] = i
+		}
+	}
+	return cg
 }
 
 // AddProcessEnv overrides the original AddProcessEnv. It uses
diff --git a/pkg/server/helpers_test.go b/pkg/server/helpers_test.go
@@ -24,6 +24,7 @@ import (
 	"github.com/containerd/containerd/containers"
 	"github.com/containerd/containerd/linux/runctypes"
 	imagedigest "github.com/opencontainers/go-digest"
+	runtimespec "github.com/opencontainers/runtime-spec/specs-go"
 	"github.com/stretchr/testify/assert"
 	"golang.org/x/net/context"
 	runtime "k8s.io/kubernetes/pkg/kubelet/apis/cri/runtime/v1alpha2"
@@ -216,6 +217,7 @@ func TestOrderedMounts(t *testing.T) {
 
 func TestCustomGenerator(t *testing.T) {
 	for desc, test := range map[string]struct {
+		existing  []string
 		kv        [][2]string
 		expected  []string
 		expectNil bool
@@ -248,16 +250,44 @@ func TestCustomGenerator(t *testing.T) {
 				{"k3", "v3"},
 				{"k3", "v4"},
 				{"k1", "v5"},
+				{"k4", "v6"},
 			},
 			expected: []string{
 				"k1=v5",
 				"k2=v2",
 				"k3=v4",
+				"k4=v6",
+			},
+		},
+		"existing env": {
+			existing: []string{
+				"k1=v1",
+				"k2=v2",
+				"k3=v3",
+			},
+			kv: [][2]string{
+				{"k3", "v4"},
+				{"k2", "v5"},
+				{"k4", "v6"},
+			},
+			expected: []string{
+				"k1=v1",
+				"k2=v5",
+				"k3=v4",
+				"k4=v6",
 			},
 		},
 	} {
 		t.Logf("TestCase %q", desc)
-		g := newSpecGenerator(nil)
+		var spec *runtimespec.Spec
+		if len(test.existing) > 0 {
+			spec = &runtimespec.Spec{
+				Process: &runtimespec.Process{
+					Env: test.existing,
+				},
+			}
+		}
+		g := newSpecGenerator(spec)
 		for _, kv := range test.kv {
 			g.AddProcessEnv(kv[0], kv[1])
 		}

Original file line number	Diff line number	Diff line change
`@@ -446,10 +446,18 @@ type generator struct {`
`446`	`446`	`}`
`447`	`447`
`448`	`448`	`func newCustomGenerator(g generate.Generator) generator {`
`449`		`- return generator{`
	`449`	`+ cg := generator{`
`450`	`450`	`Generator: g,`
`451`	`451`	`envCache: make(map[string]int),`
`452`	`452`	`}`
	`453`	`+ spec := g.Spec()`
	`454`	`+ if spec != nil && spec.Process != nil {`
	`455`	`+ for i, env := range spec.Process.Env {`
	`456`	`+ kv := strings.SplitN(env, "=", 2)`
	`457`	`+ cg.envCache[kv[0]] = i`
	`458`	`+ }`
	`459`	`+ }`
	`460`	`+ return cg`
`453`	`461`	`}`
`454`	`462`
`455`	`463`	`// AddProcessEnv overrides the original AddProcessEnv. It uses`