From bbf4dde43f3ae5f93f250039925a4163cff3d12a Mon Sep 17 00:00:00 2001
From: Lea Anthony <lea.anthony@gmail.com>
Date: Fri, 12 Mar 2021 23:41:13 +1100
Subject: [PATCH] Support upserting environment variables

---
 v2/pkg/commands/build/base.go      | 33 ++++++++++++++++++++++++++----
 v2/pkg/commands/build/base_test.go | 31 ++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 4 deletions(-)
 create mode 100644 v2/pkg/commands/build/base_test.go

diff --git a/v2/pkg/commands/build/base.go b/v2/pkg/commands/build/base.go
index 9b38ffc5..7c82917a 100644
--- a/v2/pkg/commands/build/base.go
+++ b/v2/pkg/commands/build/base.go
@@ -223,9 +223,6 @@ func (b *BaseBuilder) CompileProject(options *Options) error {
 	// Set the directory
 	cmd.Dir = b.projectData.Path
 
-	// Set GO111MODULE environment variable
-	cmd.Env = append(os.Environ(), "GO111MODULE=on")
-
 	// Add CGO flags
 	// We use the project/build dir as a temporary place for our generated c headers
 	buildBaseDir, err := fs.RelativeToCwd("build")
@@ -233,7 +230,16 @@ func (b *BaseBuilder) CompileProject(options *Options) error {
 		return err
 	}
 
-	cmd.Env = append(os.Environ(), fmt.Sprintf("CGO_CFLAGS=-I%s", buildBaseDir))
+	cmd.Env = os.Environ() // inherit env
+
+	// Use upsertEnv so we don't overwrite user's CGO_CFLAGS
+	cmd.Env = upsertEnv(cmd.Env, "CGO_CFLAGS", func(v string) string {
+		if v != "" {
+			v += " "
+		}
+		v += "-I" + buildBaseDir
+		return v
+	})
 
 	// Setup buffers
 	var stdo, stde bytes.Buffer
@@ -391,3 +397,22 @@ func (b *BaseBuilder) ExtractAssets() (*html.AssetBundle, error) {
 	// Read in html
 	return html.NewAssetBundle(b.projectData.HTML)
 }
+
+func upsertEnv(env []string, key string, update func(v string) string) []string {
+	newEnv := make([]string, len(env), len(env)+1)
+	found := false
+	for i := range env {
+		if strings.HasPrefix(env[i], key+"=") {
+			eqIndex := strings.Index(env[i], "=")
+			val := env[i][eqIndex+1:]
+			newEnv[i] = fmt.Sprintf("%s=%v", key, update(val))
+			found = true
+			continue
+		}
+		newEnv[i] = env[i]
+	}
+	if !found {
+		newEnv = append(newEnv, fmt.Sprintf("%s=%v", key, update("")))
+	}
+	return newEnv
+}
diff --git a/v2/pkg/commands/build/base_test.go b/v2/pkg/commands/build/base_test.go
new file mode 100644
index 00000000..1f3c84cd
--- /dev/null
+++ b/v2/pkg/commands/build/base_test.go
@@ -0,0 +1,31 @@
+package build
+
+import "testing"
+
+func TestUpdateEnv(t *testing.T) {
+
+	env := []string{"one=1", "two=a=b", "three="}
+	newEnv := upsertEnv(env, "two", func(v string) string {
+		return v + "+added"
+	})
+	newEnv = upsertEnv(newEnv, "newVar", func(v string) string {
+		return "added"
+	})
+	newEnv = upsertEnv(newEnv, "three", func(v string) string {
+		return "3"
+	})
+
+	if len(newEnv) != 4 {
+		t.Errorf("expected: 4, got: %d", len(newEnv))
+	}
+	if newEnv[1] != "two=a=b+added" {
+		t.Errorf("expected: \"two=a=b+added\", got: %q", newEnv[1])
+	}
+	if newEnv[2] != "three=3" {
+		t.Errorf("expected: \"three=3\", got: %q", newEnv[2])
+	}
+	if newEnv[3] != "newVar=added" {
+		t.Errorf("expected: \"newVar=added\", got: %q", newEnv[3])
+	}
+
+}