From 958dbdf5599bd11fa556b1c7baf23cbcb583490c Mon Sep 17 00:00:00 2001
From: Matthias Simon <5nord@users.noreply.github.com>
Date: Sun, 22 May 2022 02:02:33 +0200
Subject: [PATCH] Fix override for OverLoad and co. (#15)

This commit adds an override-flag to make local variables override
environment variables.

Closes #14.
---
 gotenv.go      | 25 ++++++++++++++-----------
 gotenv_test.go | 18 ++++++++++++++++++
 2 files changed, 32 insertions(+), 11 deletions(-)

diff --git a/gotenv.go b/gotenv.go
index 55c17ab..7ffc06b 100644
--- a/gotenv.go
+++ b/gotenv.go
@@ -84,7 +84,7 @@ func loadenv(override bool, filenames ...string) error {
 
 // parse and set :)
 func parset(r io.Reader, override bool) error {
-	env, err := StrictParse(r)
+	env, err := strictParse(r, override)
 	if err != nil {
 		return err
 	}
@@ -110,7 +110,7 @@ func setenv(key, val string, override bool) {
 // It expands the value of a variable from the environment variable but does not set the value to the environment itself.
 // This function is skipping any invalid lines and only processing the valid one.
 func Parse(r io.Reader) Env {
-	env, _ := StrictParse(r)
+	env, _ := strictParse(r, false)
 	return env
 }
 
@@ -118,6 +118,10 @@ func Parse(r io.Reader) Env {
 // It expands the value of a variable from the environment variable but does not set the value to the environment itself.
 // This function is returning an error if there are any invalid lines.
 func StrictParse(r io.Reader) (Env, error) {
+	return strictParse(r, false)
+}
+
+func strictParse(r io.Reader, override bool) (Env, error) {
 	env := make(Env)
 	scanner := bufio.NewScanner(r)
 
@@ -133,7 +137,7 @@ func StrictParse(r io.Reader) (Env, error) {
 
 		i++
 
-		err := parseLine(line, env)
+		err := parseLine(line, env, override)
 		if err != nil {
 			return env, err
 		}
@@ -141,8 +145,7 @@ func StrictParse(r io.Reader) (Env, error) {
 
 	return env, nil
 }
-
-func parseLine(s string, env Env) error {
+func parseLine(s string, env Env, override bool) error {
 	rl := regexp.MustCompile(linePattern)
 	rm := rl.FindStringSubmatch(s)
 
@@ -177,11 +180,11 @@ func parseLine(s string, env Env) error {
 
 	rv := regexp.MustCompile(variablePattern)
 	fv := func(s string) string {
-		return varReplacement(s, hsq, env)
+		return varReplacement(s, hsq, env, override)
 	}
 
 	val = rv.ReplaceAllStringFunc(val, fv)
-	val = parseVal(val, env, hdq)
+	val = parseVal(val, env, hdq, override)
 
 	env[key] = val
 	return nil
@@ -201,7 +204,7 @@ func parseExport(st string, env Env) error {
 	return nil
 }
 
-func varReplacement(s string, hsq bool, env Env) string {
+func varReplacement(s string, hsq bool, env Env, override bool) string {
 	if strings.HasPrefix(s, "\\") {
 		return strings.TrimPrefix(s, "\\")
 	}
@@ -220,7 +223,7 @@ func varReplacement(s string, hsq bool, env Env) string {
 
 	v := mn[3]
 
-	if replace, ok := os.LookupEnv(v); ok {
+	if replace, ok := os.LookupEnv(v); ok && !override {
 		return replace
 	}
 
@@ -246,7 +249,7 @@ func checkFormat(s string, env Env) error {
 	return fmt.Errorf("line `%s` doesn't match format", s)
 }
 
-func parseVal(val string, env Env, ignoreNewlines bool) string {
+func parseVal(val string, env Env, ignoreNewlines bool, override bool) string {
 	if strings.Contains(val, "=") && !ignoreNewlines {
 		kv := strings.Split(val, "\r")
 
@@ -254,7 +257,7 @@ func parseVal(val string, env Env, ignoreNewlines bool) string {
 			val = kv[0]
 
 			for i := 1; i < len(kv); i++ {
-				parseLine(kv[i], env)
+				parseLine(kv[i], env, override)
 			}
 		}
 	}
diff --git a/gotenv_test.go b/gotenv_test.go
index 490e86a..26a5eb4 100644
--- a/gotenv_test.go
+++ b/gotenv_test.go
@@ -342,6 +342,24 @@ func TestOverLoad_overriding(t *testing.T) {
 	os.Clearenv()
 }
 
+func TestOverLoad_overrideVars(t *testing.T) {
+	os.Setenv("A", "fromEnv")
+	err := gotenv.OverLoad("fixtures/vars.env")
+	assert.Nil(t, err)
+	assert.Equal(t, "fromFile", os.Getenv("B"))
+	os.Clearenv()
+}
+
+func TestOverLoad_overrideVars2(t *testing.T) {
+	os.Setenv("C", "fromEnv")
+	err := gotenv.OverLoad("fixtures/vars.env")
+	assert.Nil(t, err)
+	// The value for D is not "fromFile" because C is defined after the
+	// definition of D.
+	assert.Equal(t, "fromEnv", os.Getenv("D"), "C defined after usage")
+	os.Clearenv()
+}
+
 func TestMustOverLoad_nonExist(t *testing.T) {
 	assert.Panics(t, func() { gotenv.Must(gotenv.OverLoad, ".env.not.exist") }, "Caling gotenv.Must with Overgotenv.Load and non exist file SHOULD panic")
 }