From 02fa534d713b5ecad8ea86ad5b93b2bc56f698f0 Mon Sep 17 00:00:00 2001
From: Evan Hazlett <ejhazlett@gmail.com>
Date: Wed, 24 May 2017 12:39:34 -0400
Subject: [PATCH] ctr: enable specifying additional environment variables

Signed-off-by: Evan Hazlett <ejhazlett@gmail.com>
---
 cmd/ctr/run.go         |  4 ++++
 cmd/ctr/run_unix.go    | 10 +++++++---
 cmd/ctr/run_windows.go |  4 +++-
 cmd/ctr/utils.go       | 38 ++++++++++++++++++++++++++++++++++++++
 4 files changed, 52 insertions(+), 4 deletions(-)

diff --git a/cmd/ctr/run.go b/cmd/ctr/run.go
index cacd83e27..f47317001 100644
--- a/cmd/ctr/run.go
+++ b/cmd/ctr/run.go
@@ -60,6 +60,10 @@ var runCommand = cli.Command{
 			Name:  "mount",
 			Usage: "specify additional container mount (ex: type=bind,src=/tmp,dest=/host,options=rbind:ro)",
 		},
+		cli.StringSliceFlag{
+			Name:  "env",
+			Usage: "specify additional container environment variables (i.e. FOO=bar)",
+		},
 		cli.BoolFlag{
 			Name:  "rm",
 			Usage: "remove the container after running",
diff --git a/cmd/ctr/run_unix.go b/cmd/ctr/run_unix.go
index 2b3db7f14..23f5c8712 100644
--- a/cmd/ctr/run_unix.go
+++ b/cmd/ctr/run_unix.go
@@ -52,10 +52,10 @@ var capabilities = []string{
 }
 
 func spec(id string, config *ocispec.ImageConfig, context *cli.Context, rootfs string) (*specs.Spec, error) {
-	env := []string{
+	defaultEnv := []string{
 		"PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin",
 	}
-	env = append(env, config.Env...)
+	defaultEnv = append(defaultEnv, config.Env...)
 	cmd := config.Cmd
 	if v := context.Args().Tail(); len(v) > 0 {
 		cmd = v
@@ -90,8 +90,12 @@ func spec(id string, config *ocispec.ImageConfig, context *cli.Context, rootfs s
 		}
 	}
 	if tty {
-		env = append(env, "TERM=xterm")
+		defaultEnv = append(defaultEnv, "TERM=xterm")
 	}
+
+	// additional environment vars
+	env := replaceOrAppendEnvValues(defaultEnv, context.StringSlice("env"))
+
 	cwd := config.WorkingDir
 	if cwd == "" {
 		cwd = "/"
diff --git a/cmd/ctr/run_windows.go b/cmd/ctr/run_windows.go
index bf804ee5d..411209348 100644
--- a/cmd/ctr/run_windows.go
+++ b/cmd/ctr/run_windows.go
@@ -61,6 +61,8 @@ func spec(id string, config *ocispec.ImageConfig, context *cli.Context) *specs.S
 		}
 	}
 
+	env := replaceOrAppendEnvValues(config.Env, context.StringSlice("env"))
+
 	return &specs.Spec{
 		Version: specs.Version,
 		Platform: specs.Platform{
@@ -74,7 +76,7 @@ func spec(id string, config *ocispec.ImageConfig, context *cli.Context) *specs.S
 			Args:     args,
 			Terminal: tty,
 			Cwd:      cwd,
-			Env:      config.Env,
+			Env:      env,
 			User: specs.User{
 				Username: config.User,
 			},
diff --git a/cmd/ctr/utils.go b/cmd/ctr/utils.go
index 3173ed496..3dfabeeff 100644
--- a/cmd/ctr/utils.go
+++ b/cmd/ctr/utils.go
@@ -198,3 +198,41 @@ func parseMountFlag(m string) (specs.Mount, error) {
 
 	return mount, nil
 }
+
+// replaceOrAppendEnvValues returns the defaults with the overrides either
+// replaced by env key or appended to the list
+func replaceOrAppendEnvValues(defaults, overrides []string) []string {
+	cache := make(map[string]int, len(defaults))
+	for i, e := range defaults {
+		parts := strings.SplitN(e, "=", 2)
+		cache[parts[0]] = i
+	}
+
+	for _, value := range overrides {
+		// Values w/o = means they want this env to be removed/unset.
+		if !strings.Contains(value, "=") {
+			if i, exists := cache[value]; exists {
+				defaults[i] = "" // Used to indicate it should be removed
+			}
+			continue
+		}
+
+		// Just do a normal set/update
+		parts := strings.SplitN(value, "=", 2)
+		if i, exists := cache[parts[0]]; exists {
+			defaults[i] = value
+		} else {
+			defaults = append(defaults, value)
+		}
+	}
+
+	// Now remove all entries that we want to "unset"
+	for i := 0; i < len(defaults); i++ {
+		if defaults[i] == "" {
+			defaults = append(defaults[:i], defaults[i+1:]...)
+			i--
+		}
+	}
+
+	return defaults
+}