[lxc-devel] [distrobuilder/master] shared, chroot: Set environment variables

[monstermunchkin on Github]

A non-text attachment was scrubbed...
Name: not available
Type: text/x-mailbox
Size: 363 bytes
Desc: not available
URL: <http://lists.linuxcontainers.org/pipermail/lxc-devel/attachments/20180308/50eb361a/attachment.bin>
-------------- next part --------------
From 4981ed1fddd050c4000b7ce6717bbd1e3331975c Mon Sep 17 00:00:00 2001
From: Thomas Hipp <thomas.hipp at canonical.com>
Date: Thu, 8 Mar 2018 10:57:58 +0100
Subject: [PATCH] shared,chroot: Set environment variables

Signed-off-by: Thomas Hipp <thomas.hipp at canonical.com>
---
 distrobuilder/chroot.go | 12 ++++++++++++
 shared/util.go          | 45 +++++++++++++++++++++++++++++++--------------
 shared/util_test.go     | 33 +++++++++++++++++++++++++++++++++
 3 files changed, 76 insertions(+), 14 deletions(-)

diff --git a/distrobuilder/chroot.go b/distrobuilder/chroot.go
index e8259f1..c6fc395 100644
--- a/distrobuilder/chroot.go
+++ b/distrobuilder/chroot.go
@@ -198,9 +198,21 @@ func setupChroot(rootfs string) (func() error, error) {
 		return nil, err
 	}
 
+	// Set environment variables
+	oldEnvVariables := shared.SetEnvVariables(
+		[]shared.EnvVariable{
+			{"PATH", "/sbin:/bin:/usr/sbin:/usr/bin:/usr/local/sbin:/usr/local/bin", true},
+			{"SHELL", "/bin/sh", true},
+			{"TERM", "xterm", true},
+			{"DEBIAN_FRONTEND", "noninteractive", true},
+		})
+
 	return func() error {
 		defer root.Close()
 
+		// Reset old environment variables
+		shared.SetEnvVariables(oldEnvVariables)
+
 		// Switch back to the host rootfs
 		err = root.Chdir()
 		if err != nil {
diff --git a/shared/util.go b/shared/util.go
index fc112e5..72c4488 100644
--- a/shared/util.go
+++ b/shared/util.go
@@ -17,6 +17,12 @@ import (
 	"gopkg.in/flosch/pongo2.v3"
 )
 
+type EnvVariable struct {
+	Key   string
+	Value string
+	Set   bool
+}
+
 // Copy copies a file.
 func Copy(src, dest string) error {
 	var err error
@@ -47,13 +53,6 @@ func Copy(src, dest string) error {
 func RunCommand(name string, arg ...string) error {
 	cmd := exec.Command(name, arg...)
 
-	cmd.Env = []string{
-		"PATH=/sbin:/bin:/usr/sbin:/usr/bin:/usr/local/sbin:/usr/local/bin",
-		"SHELL=/bin/sh",
-		"TERM=xterm",
-		"DEBIAN_FRONTEND=noninteractive",
-	}
-
 	cmd.Stdin = os.Stdin
 	cmd.Stdout = os.Stdout
 	cmd.Stderr = os.Stderr
@@ -67,13 +66,7 @@ func RunCommand(name string, arg ...string) error {
 func RunScript(content string) error {
 	cmd := exec.Command("sh")
 
-	cmd.Env = []string{
-		"PATH=/sbin:/bin:/usr/sbin:/usr/bin:/usr/local/sbin:/usr/local/bin",
-		"SHELL=/bin/sh",
-		"TERM=xterm"}
-
-	buffer := bytes.NewBufferString(content)
-	cmd.Stdin = buffer
+	cmd.Stdin = bytes.NewBufferString(content)
 	cmd.Stdout = os.Stdout
 	cmd.Stderr = os.Stderr
 
@@ -233,3 +226,27 @@ func RenderTemplate(template string, ctx pongo2.Context) (string, error) {
 
 	return ret, err
 }
+
+// SetEnvVariables sets the provided environment variables and returns the
+// old ones.
+func SetEnvVariables(env []EnvVariable) []EnvVariable {
+	oldEnv := make([]EnvVariable, len(env))
+
+	for i := 0; i < len(env); i++ {
+		// Check whether the env variables are set at the moment
+		oldVal, set := os.LookupEnv(env[i].Key)
+
+		// Store old env variables
+		oldEnv[i].Key = env[i].Key
+		oldEnv[i].Value = oldVal
+		oldEnv[i].Set = set
+
+		if env[i].Set {
+			os.Setenv(env[i].Key, env[i].Value)
+		} else {
+			os.Unsetenv(env[i].Key)
+		}
+	}
+
+	return oldEnv
+}
diff --git a/shared/util_test.go b/shared/util_test.go
index eaad590..6b9c52f 100644
--- a/shared/util_test.go
+++ b/shared/util_test.go
@@ -171,3 +171,36 @@ func TestRenderTemplate(t *testing.T) {
 		}
 	}
 }
+
+func TestSetEnvVariables(t *testing.T) {
+	// Initial variables
+	os.Setenv("FOO", "bar")
+
+	env := []envVariable{
+		{"FOO", "bla", true},
+		{"BAR", "blub", true},
+	}
+
+	// Set new env variables
+	oldEnv := setEnvVariables(env)
+
+	for _, e := range env {
+		v, set := os.LookupEnv(e.key)
+		if !set || e.value != v {
+			t.Fatalf("Expected %s to be '%s', got '%s'", e.key, e.value, v)
+		}
+	}
+
+	// Reset env variables
+	setEnvVariables(oldEnv)
+
+	val, set := os.LookupEnv("FOO")
+	if !set || val != "bar" {
+		t.Fatalf("Expected %s to be '%s', got '%s'", "FOO", "bar", val)
+	}
+
+	val, set = os.LookupEnv("BAR")
+	if set {
+		t.Fatalf("Expected %s to be unset", "BAR")
+	}
+}