Repository: compose-spec/compose-go
Pull request #276 (2022-06-24 00:14:52)

Fix environment variable expansion


Patch:
diff --git a/ci/Dockerfile b/ci/Dockerfile
index ce352e1..8340233 100644
--- a/ci/Dockerfile
+++ b/ci/Dockerfile
@@ -16,7 +16,7 @@ FROM golang:1.18
 
 WORKDIR /go/src
 
-ARG GOLANGCILINT_VERSION=v1.44.1
+ARG GOLANGCILINT_VERSION=v1.46.2
 RUN curl -sSfL https://raw.githubusercontent.com/golangci/golangci-lint/master/install.sh | sh -s -- -b $(go env GOPATH)/bin ${GOLANGCILINT_VERSION}
 RUN go install github.com/kunalkushwaha/ltag@latest && rm -rf /go/src/github.com/kunalkushwaha
 
diff --git a/dotenv/godotenv.go b/dotenv/godotenv.go
index 479831a..98d6620 100644
--- a/dotenv/godotenv.go
+++ b/dotenv/godotenv.go
@@ -24,6 +24,8 @@ import (
 	"sort"
 	"strconv"
 	"strings"
+
+	"github.com/compose-spec/compose-go/template"
 )
 
 const doubleQuoteSpecialChars = "\\\n\r\"!$`"
@@ -322,42 +324,24 @@ func parseValue(value string, envMap map[string]string, lookupFn LookupFn) strin
 		}
 
 		if singleQuotes == nil {
-			value = expandVariables(value, envMap, lookupFn)
+			value, _ = expandVariables(value, envMap, lookupFn)
 		}
 	}
 
 	return value
 }
 
-var expandVarRegex = regexp.MustCompile(`(\\)?(\$)(\()?\{?([A-Z0-9_]+)?\}?`)
-
-func expandVariables(v string, envMap map[string]string, lookupFn LookupFn) string {
-	return expandVarRegex.ReplaceAllStringFunc(v, func(s string) string {
-		submatch := expandVarRegex.FindStringSubmatch(s)
-
-		if submatch == nil {
-			return s
+func expandVariables(value string, envMap map[string]string, lookupFn LookupFn) (string, error) {
+	retVal, err := template.Substitute(value, func(k string) (string, bool) {
+		if v, ok := envMap[k]; ok {
+			return v, ok
 		}
-		if submatch[1] == "\\" || submatch[2] == "(" {
-			return submatch[0][1:]
-		} else if submatch[4] != "" {
-			// first check if we have defined this already earlier
-			if envMap[submatch[4]] != "" {
-				return envMap[submatch[4]]
-			}
-			if lookupFn == nil {
-				return ""
-			}
-			// if we have not defined it, check the lookup function provided
-			// by the user
-			s2, ok := lookupFn(submatch[4])
-			if ok {
-				return s2
-			}
-			return ""
-		}
-		return s
+		return lookupFn(k)
 	})
+	if err != nil {
+		return value, err
+	}
+	return retVal, nil
 }
 
 func doubleQuoteEscape(line string) string {
diff --git a/dotenv/parser.go b/dotenv/parser.go
index 4ca6bd0..e049123 100644
--- a/dotenv/parser.go
+++ b/dotenv/parser.go
@@ -129,7 +129,7 @@ loop:
 }
 
 // extractVarValue extracts variable value and returns rest of slice
-func extractVarValue(src []byte, envMap map[string]string, lookupFn LookupFn) (value string, rest []byte, err error) {
+func extractVarValue(src []byte, envMap map[string]string, lookupFn LookupFn) (string, []byte, error) {
 	quote, isQuoted := hasQuotePrefix(src)
 	if !isQuoted {
 		// unquoted value - read until new line
@@ -138,13 +138,17 @@ func extractVarValue(src []byte, envMap map[string]string, lookupFn LookupFn) (v
 		if end < 0 {
 			value := strings.Split(string(src), "#")[0] // Remove inline comments on unquoted lines
 			value = strings.TrimRightFunc(value, unicode.IsSpace)
-			return expandVariables(value, envMap, lookupFn), nil, nil
+
+			retVal, err := expandVariables(value, envMap, lookupFn)
+			return retVal, nil, err
 		}
 
 		value := strings.Split(string(src[0:end]), "#")[0]
 		value = strings.TrimRightFunc(value, unicode.IsSpace)
 		rest = src[end:]
-		return expandVariables(value, envMap, lookupFn), rest, nil
+
+		retVal, err := expandVariables(value, envMap, lookupFn)
+		return retVal, rest, err
 	}
 
 	// lookup quoted string terminator
@@ -160,11 +164,16 @@ func extractVarValue(src []byte, envMap map[string]string, lookupFn LookupFn) (v
 
 		// trim quotes
 		trimFunc := isCharFunc(rune(quote))
-		value = string(bytes.TrimLeftFunc(bytes.TrimRightFunc(src[0:i], trimFunc), trimFunc))
+		value := string(bytes.TrimLeftFunc(bytes.TrimRightFunc(src[0:i], trimFunc), trimFunc))
 		if quote == prefixDoubleQuote {
 			// unescape newlines for double quote (this is compat feature)
 			// and expand environment variables
-			value = expandVariables(expandEscapes(value), envMap, lookupFn)
+
+			retVal, err := expandVariables(expandEscapes(value), envMap, lookupFn)
+			if err != nil {
+				return "", nil, err
+			}
+			value = retVal
 		}
 
 		return value, src[i+1:], nil
@@ -187,6 +196,8 @@ func expandEscapes(str string) string {
 			return "\n"
 		case "r":
 			return "\r"
+		case "$":
+			return "$$"
 		default:
 			return match
 		}
diff --git a/go.mod b/go.mod
index ab91790..ad2e472 100644
--- a/go.mod
+++ b/go.mod
@@ -13,6 +13,7 @@ require (
 	github.com/opencontainers/go-digest v1.0.0
 	github.com/pkg/errors v0.9.1
 	github.com/sirupsen/logrus v1.8.1
+	github.com/stretchr/testify v1.5.1
 	github.com/xeipuuv/gojsonschema v1.2.0
 	golang.org/x/sync v0.0.0-20210220032951-036812b2e83c
 	gopkg.in/yaml.v2 v2.4.0
@@ -20,6 +21,8 @@ require (
 )
 
 require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/xeipuuv/gojsonpointer v0.0.0-20180127040702-4e3ac2762d5f // indirect
 	github.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415 // indirect
 	golang.org/x/sys v0.0.0-20210119212857-b64e53b001e4 // indirect


Test patch:
diff --git a/dotenv/godotenv_test.go b/dotenv/godotenv_test.go
index efd6a13..1bc726f 100644
--- a/dotenv/godotenv_test.go
+++ b/dotenv/godotenv_test.go
@@ -257,17 +257,17 @@ func TestExpanding(t *testing.T) {
 			map[string]string{"BAR": "quote $FOO"},
 		},
 		{
-			"does not expand escaped variables",
+			"does not expand escaped variables 1",
 			`FOO="foo\$BAR"`,
 			map[string]string{"FOO": "foo$BAR"},
 		},
 		{
-			"does not expand escaped variables",
+			"does not expand escaped variables 2",
 			`FOO="foo\${BAR}"`,
 			map[string]string{"FOO": "foo${BAR}"},
 		},
 		{
-			"does not expand escaped variables",
+			"does not expand escaped variables 3",
 			"FOO=test\nBAR=\"foo\\${FOO} ${FOO}\"",
 			map[string]string{"FOO": "test", "BAR": "foo${FOO} test"},
 		},
diff --git a/dotenv/godotenv_var_expansion_test.go b/dotenv/godotenv_var_expansion_test.go
new file mode 100644
index 0000000..f8a2ff7
--- /dev/null
+++ b/dotenv/godotenv_var_expansion_test.go
@@ -0,0 +1,154 @@
+package dotenv
+
+import (
+	"testing"
+
+	"github.com/compose-spec/compose-go/template"
+	"github.com/stretchr/testify/assert"
+)
+
+var envMap = map[string]string{
+	// UNSET_VAR: <Cannot be here :D>
+	"EMPTY_VAR": "",
+	"TEST_VAR":  "Test Value",
+}
+
+var notFoundLookup = func(s string) (string, bool) {
+	return "", false
+}
+
+func TestExpandIfEmptyOrUnset(t *testing.T) {
+	templateResults := []struct {
+		name   string
+		input  string
+		result string
+	}{
+		{
+			"Expand if empty or unset: UNSET_VAR",
+			"RESULT=${UNSET_VAR:-Default Value}",
+			"RESULT=Default Value",
+		},
+		{
+			"Expand if empty or unset: EMPTY_VAR",
+			"RESULT=${EMPTY_VAR:-Default Value}",
+			"RESULT=Default Value",
+		},
+		{
+			"Expand if empty or unset: TEST_VAR",
+			"RESULT=${TEST_VAR:-Default Value}",
+			"RESULT=Test Value",
+		},
+	}
+
+	for _, expected := range templateResults {
+		t.Run(expected.name, func(t *testing.T) {
+			result, err := expandVariables(expected.input, envMap, notFoundLookup)
+			assert.Nil(t, err)
+			assert.Equal(t, result, expected.result)
+		})
+	}
+}
+
+func TestExpandIfUnset(t *testing.T) {
+	templateResults := []struct {
+		name   string
+		input  string
+		result string
+	}{
+		{
+			"Expand if unset: UNSET_VAR",
+			"RESULT=${UNSET_VAR-Default Value}",
+			"RESULT=Default Value",
+		},
+		{
+			"Expand if unset: EMPTY_VAR",
+			"RESULT=${EMPTY_VAR-Default Value}",
+			"RESULT=",
+		},
+		{
+			"Expand if unset: TEST_VAR",
+			"RESULT=${TEST_VAR-Default Value}",
+			"RESULT=Test Value",
+		},
+	}
+
+	for _, expected := range templateResults {
+		t.Run(expected.name, func(t *testing.T) {
+			result, err := expandVariables(expected.input, envMap, notFoundLookup)
+			assert.Nil(t, err)
+			assert.Equal(t, result, expected.result)
+		})
+	}
+}
+
+func TestErrorIfEmptyOrUnset(t *testing.T) {
+	templateResults := []struct {
+		name   string
+		input  string
+		result string
+		err    error
+	}{
+		{
+			"Error empty or unset: UNSET_VAR",
+			"RESULT=${UNSET_VAR:?Test error}",
+			"RESULT=${UNSET_VAR:?Test error}",
+			&template.InvalidTemplateError{Template: "required variable UNSET_VAR is missing a value: Test error"},
+		},
+		{
+			"Error empty or unset: EMPTY_VAR",
+			"RESULT=${EMPTY_VAR:?Test error}",
+			"RESULT=${EMPTY_VAR:?Test error}",
+			&template.InvalidTemplateError{Template: "required variable EMPTY_VAR is missing a value: Test error"},
+		},
+		{
+			"Error empty or unset: TEST_VAR",
+			"RESULT=${TEST_VAR:?Default Value}",
+			"RESULT=Test Value",
+			nil,
+		},
+	}
+
+	for _, expected := range templateResults {
+		t.Run(expected.name, func(t *testing.T) {
+			result, err := expandVariables(expected.input, envMap, notFoundLookup)
+			assert.Equal(t, expected.err, err)
+			assert.Equal(t, expected.result, result)
+		})
+	}
+}
+
+func TestErrorIfUnset(t *testing.T) {
+	templateResults := []struct {
+		name   string
+		input  string
+		result string
+		err    error
+	}{
+		{
+			"Error on unset: UNSET_VAR",
+			"RESULT=${UNSET_VAR?Test error}",
+			"RESULT=${UNSET_VAR?Test error}",
+			&template.InvalidTemplateError{Template: "required variable UNSET_VAR is missing a value: Test error"},
+		},
+		{
+			"Error on unset: EMPTY_VAR",
+			"RESULT=${EMPTY_VAR?Test error}",
+			"RESULT=",
+			nil,
+		},
+		{
+			"Error on unset: TEST_VAR",
+			"RESULT=${TEST_VAR?Default Value}",
+			"RESULT=Test Value",
+			nil,
+		},
+	}
+
+	for _, expected := range templateResults {
+		t.Run(expected.name, func(t *testing.T) {
+			result, err := expandVariables(expected.input, envMap, notFoundLookup)
+			assert.Equal(t, expected.err, err)
+			assert.Equal(t, expected.result, result)
+		})
+	}
+}