From 877c1cadc1992cb9823504a6dfa105558d64d454 Mon Sep 17 00:00:00 2001
From: Lantao Liu <lantaol@google.com>
Date: Tue, 12 Feb 2019 10:29:45 -0800
Subject: [PATCH] Include default envs from containerd.

Signed-off-by: Lantao Liu <lantaol@google.com>
---
 pkg/server/helpers.go      |  9 ++++++++-
 pkg/server/helpers_test.go | 32 +++++++++++++++++++++++++++++++-
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/pkg/server/helpers.go b/pkg/server/helpers.go
index a1cbd4229..c4a0044ce 100644
--- a/pkg/server/helpers.go
+++ b/pkg/server/helpers.go
@@ -406,10 +406,17 @@ type generator struct {
 }
 
 func newCustomGenerator(g generate.Generator) generator {
-	return generator{
+	cg := generator{
 		Generator: g,
 		envCache:  make(map[string]int),
 	}
+	if g.Config != nil && g.Config.Process != nil {
+		for i, env := range g.Config.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
index 95e34a8dc..c0eb29839 100644
--- a/pkg/server/helpers_test.go
+++ b/pkg/server/helpers_test.go
@@ -25,6 +25,7 @@ import (
 	runcoptions "github.com/containerd/containerd/runtime/v2/runc/options"
 	"github.com/docker/distribution/reference"
 	imagedigest "github.com/opencontainers/go-digest"
+	runtimespec "github.com/opencontainers/runtime-spec/specs-go"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	runtime "k8s.io/kubernetes/pkg/kubelet/apis/cri/runtime/v1alpha2"
@@ -324,6 +325,7 @@ func TestRestrictOOMScoreAdj(t *testing.T) {
 
 func TestCustomGenerator(t *testing.T) {
 	for desc, test := range map[string]struct {
+		existing  []string
 		kv        [][2]string
 		expected  []string
 		expectNil bool
@@ -356,16 +358,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])
 		}