Skip to content

Commit

Permalink
Fix default and error variable expansions
Browse files Browse the repository at this point in the history
Signed-off-by: Ulysses Souza <[email protected]>
  • Loading branch information
ulyssessouza committed Jun 24, 2022
1 parent 3204ab7 commit fb94264
Show file tree
Hide file tree
Showing 4 changed files with 185 additions and 36 deletions.
40 changes: 12 additions & 28 deletions dotenv/godotenv.go
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,8 @@ import (
"sort"
"strconv"
"strings"

"github.com/compose-spec/compose-go/template"
)

const doubleQuoteSpecialChars = "\\\n\r\"!$`"
Expand Down Expand Up @@ -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 {
Expand Down
6 changes: 3 additions & 3 deletions dotenv/godotenv_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"},
},
Expand Down
154 changes: 154 additions & 0 deletions dotenv/godotenv_var_expansion_test.go
Original file line number Diff line number Diff line change
@@ -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)
})
}
}
21 changes: 16 additions & 5 deletions dotenv/parser.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -187,6 +196,8 @@ func expandEscapes(str string) string {
return "\n"
case "r":
return "\r"
case "$":
return "$$"
default:
return match
}
Expand Down

0 comments on commit fb94264

Please sign in to comment.