Fix default and error variable expansions

Signed-off-by: Ulysses Souza <[email protected]>
compose-spec · Jun 24, 2022 · fb94264 · fb94264
1 parent 3204ab7
commit fb94264
Show file tree

Hide file tree

Showing 4 changed files with 185 additions and 36 deletions.
diff --git 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/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
@@ -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
+		templ  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.templ, envMap, notFoundLookup)
+			assert.Equal(t, expected.err, err)
+			assert.Equal(t, expected.result, result)
+		})
+	}
+}
+
+func TestErrorIfUnset(t *testing.T) {
+	templateResults := []struct {
+		name   string
+		templ  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.templ, envMap, notFoundLookup)
+			assert.Equal(t, expected.err, err)
+			assert.Equal(t, expected.result, result)
+		})
+	}
+}
diff --git 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
 		}