fix: improve env variable substitution logic, only replace variables wrapped in curly braces (${var}) (#894)

This commit is contained in:
EquentR 2024-12-02 17:08:37 +08:00 committed by GitHub
parent d621053eaf
commit 3f253b327e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 20 additions and 3 deletions

View File

@ -4,11 +4,11 @@ import (
"bytes"
"fmt"
"os"
"regexp"
"strings"
"RedisShake/internal/log"
"github.com/a8m/envsubst"
"github.com/mcuadros/go-defaults"
"github.com/rs/zerolog"
"github.com/spf13/viper"
@ -101,13 +101,14 @@ func LoadConfig() *viper.Viper {
if len(os.Args) == 2 {
logger.Info().Msgf("load config from file: %s", os.Args[1])
configFile := os.Args[1]
buf, err := envsubst.ReadFile(configFile)
file, err := os.ReadFile(configFile)
if err != nil {
logger.Error().Msgf("failed to read config file: %v", err)
os.Exit(1)
}
fallback := envWithFallback(string(file))
v.SetConfigType("toml")
err = v.ReadConfig(bytes.NewReader(buf))
err = v.ReadConfig(bytes.NewReader([]byte(fallback)))
if err != nil {
logger.Error().Msgf("failed to read config file: %v", err)
os.Exit(1)
@ -124,3 +125,19 @@ func LoadConfig() *viper.Viper {
}
return v
}
// Custom substitution function that returns as is if the environment variable is empty
func envWithFallback(input string) string {
re := regexp.MustCompile(`\$\{([A-Za-z_][A-Za-z0-9_]*?)}`)
return re.ReplaceAllStringFunc(input, func(match string) string {
varName := match[1:]
if strings.HasPrefix(varName, "{") && strings.HasSuffix(varName, "}") {
varName = varName[1 : len(varName)-1] // Remove { and }
}
value := os.Getenv(varName)
if value == "" {
return match // if the environment variable is empty return as is
}
return value
})
}