package config expands env like trpc_go.yaml

Author: liuzengh

config.go

@@ -25,6 +25,7 @@ import (
 	"time"
 
 	yaml "gopkg.in/yaml.v3"
+	"trpc.group/trpc-go/trpc-go/internal/expandenv"
 	trpcpb "trpc.group/trpc/trpc-protocol/pb/go/trpc"
 
 	"trpc.group/trpc-go/trpc-go/client"
@@ -608,11 +609,8 @@ func parseConfigFromFile(configPath string) (*Config, error) {
 	if err != nil {
 		return nil, err
 	}
-	// expand environment variables
-	buf = []byte(expandEnv(string(buf)))
-
 	cfg := defaultConfig()
-	if err := yaml.Unmarshal(buf, cfg); err != nil {
+	if err := yaml.Unmarshal(expandenv.ExpandEnv(buf), cfg); err != nil {
 		return nil, err
 	}
 	return cfg, nil

config/options.go

@@ -13,6 +13,8 @@
 
 package config
 
+import "trpc.group/trpc-go/trpc-go/internal/expandenv"
+
 // WithCodec returns an option which sets the codec's name.
 func WithCodec(name string) LoadOption {
 	return func(c *TrpcConfig) {
@@ -27,6 +29,14 @@ func WithProvider(name string) LoadOption {
 	}
 }
 
+// WithExpandEnv replaces ${var} in raw bytes with environment value of var.
+// Note, method TrpcConfig.Bytes will return the replaced bytes.
+func WithExpandEnv() LoadOption {
+	return func(c *TrpcConfig) {
+		c.expandEnv = expandenv.ExpandEnv
+	}
+}
+
 // options is config option.
 type options struct{}
 

config/trpc_config.go

@@ -162,6 +162,7 @@ type TrpcConfig struct {
 	path             string
 	decoder          Codec
 	rawData          []byte
+	expandEnv        func([]byte) []byte
 }
 
 func newTrpcConfig(path string) *TrpcConfig {
@@ -170,6 +171,7 @@ func newTrpcConfig(path string) *TrpcConfig {
 		unmarshalledData: make(map[string]interface{}),
 		path:             path,
 		decoder:          &YamlCodec{},
+		expandEnv:        func(bytes []byte) []byte { return bytes },
 	}
 }
 
@@ -189,7 +191,7 @@ func (c *TrpcConfig) Load() error {
 		return fmt.Errorf("trpc/config: failed to load %s: %s", c.path, err.Error())
 	}
 
-	c.rawData = data
+	c.rawData = c.expandEnv(data)
 	if err := c.decoder.Unmarshal(c.rawData, &c.unmarshalledData); err != nil {
 		return fmt.Errorf("trpc/config: failed to parse %s: %s", c.path, err.Error())
 	}
@@ -208,8 +210,8 @@ func (c *TrpcConfig) Reload() {
 		return
 	}
 
-	c.rawData = data
-	if err := c.decoder.Unmarshal(data, &c.unmarshalledData); err != nil {
+	c.rawData = c.expandEnv(data)
+	if err := c.decoder.Unmarshal(c.rawData, &c.unmarshalledData); err != nil {
 		log.Tracef("trpc/config: failed to parse %s: %v", c.path, err)
 		return
 	}

config/trpc_config_test.go

@@ -126,3 +126,43 @@ func TestYamlCodec_Unmarshal(t *testing.T) {
 		require.NotNil(t, GetCodec("yaml").Unmarshal([]byte("[1, 2]"), &tt))
 	})
 }
+
+func TestEnvExpanded(t *testing.T) {
+	RegisterProvider(NewEnvProvider(t.Name(), []byte(`
+password: ${pwd}
+`)))
+
+	t.Setenv("pwd", t.Name())
+	cfg, err := DefaultConfigLoader.Load(
+		t.Name(),
+		WithProvider(t.Name()),
+		WithExpandEnv())
+	require.Nil(t, err)
+
+	require.Equal(t, t.Name(), cfg.GetString("password", ""))
+	require.Contains(t, string(cfg.Bytes()), fmt.Sprintf("password: %s", t.Name()))
+}
+
+func NewEnvProvider(name string, data []byte) *EnvProvider {
+	return &EnvProvider{
+		name: name,
+		data: data,
+	}
+}
+
+type EnvProvider struct {
+	name string
+	data []byte
+}
+
+func (ep *EnvProvider) Name() string {
+	return ep.name
+}
+
+func (ep *EnvProvider) Read(string) ([]byte, error) {
+	return ep.data, nil
+}
+
+func (ep *EnvProvider) Watch(cb ProviderCallback) {
+	cb("", ep.data)
+}

internal/expandenv/expand_env.go

@@ -0,0 +1,55 @@
+// Package expandenv replaces ${key} in byte slices with the env value of key.
+package expandenv
+
+import (
+	"os"
+)
+
+// ExpandEnv looks for ${var} in s and replaces them with value of the corresponding environment variable.
+// It's not like os.ExpandEnv which handles both ${var} and $var.
+// $var is considered invalid, since configurations like password for redis/mysql may contain $.
+func ExpandEnv(s []byte) []byte {
+	var buf []byte
+	i := 0
+	for j := 0; j < len(s); j++ {
+		if s[j] == '$' && j+2 < len(s) && s[j+1] == '{' { // only ${var} instead of $var is valid
+			if buf == nil {
+				buf = make([]byte, 0, 2*len(s))
+			}
+			buf = append(buf, s[i:j]...)
+			name, w := getEnvName(s[j+1:])
+			if name == nil && w > 0 {
+				// invalid matching, remove the $
+			} else if name == nil {
+				buf = append(buf, s[j]) // keep the $
+			} else {
+				buf = append(buf, os.Getenv(string(name))...)
+			}
+			j += w
+			i = j + 1
+		}
+	}
+	if buf == nil {
+		return s
+	}
+	return append(buf, s[i:]...)
+}
+
+// getEnvName gets env name, that is, var from ${var}.
+// The env name and its len will be returned.
+func getEnvName(s []byte) ([]byte, int) {
+	// look for right curly bracket '}'
+	// it's guaranteed that the first char is '{' and the string has at least two char
+	for i := 1; i < len(s); i++ {
+		if s[i] == ' ' || s[i] == '\n' || s[i] == '"' { // "xx${xxx"
+			return nil, 0 // encounter invalid char, keep the $
+		}
+		if s[i] == '}' {
+			if i == 1 { // ${}
+				return nil, 2 // remove ${}
+			}
+			return s[1:i], i + 1
+		}
+	}
+	return nil, 0 // no }，keep the $
+}

internal/expandenv/expand_env_test.go

@@ -0,0 +1,47 @@
+package expandenv_test
+
+import (
+	"fmt"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	. "trpc.group/trpc-go/trpc-go/internal/expandenv"
+)
+
+func TestExpandEnv(t *testing.T) {
+	key := "env_key"
+	t.Run("no env", func(t *testing.T) {
+		require.Equal(t, []byte("abc"), ExpandEnv([]byte("abc")))
+	})
+	t.Run("${..} is expanded", func(t *testing.T) {
+		t.Setenv(key, t.Name())
+		require.Equal(t, fmt.Sprintf("head_%s_tail", t.Name()),
+			string(ExpandEnv([]byte(fmt.Sprintf("head_${%s}_tail", key)))))
+	})
+	t.Run("${ is not expanded", func(t *testing.T) {
+		require.Equal(t, "head_${_tail",
+			string(ExpandEnv([]byte(fmt.Sprintf("head_${_tail")))))
+	})
+	t.Run("${} is expanded as empty", func(t *testing.T) {
+		require.Equal(t, "head__tail",
+			string(ExpandEnv([]byte("head_${}_tail"))))
+	})
+	t.Run("${..} is not expanded if .. contains any space", func(t *testing.T) {
+		t.Setenv("key key", t.Name())
+		require.Equal(t, "head_${key key}_tail",
+			string(ExpandEnv([]byte("head_${key key}_tail"))))
+	})
+	t.Run("${..} is not expanded if .. contains any new line", func(t *testing.T) {
+		t.Setenv("key\nkey", t.Name())
+		require.Equal(t, t.Name(), os.Getenv("key\nkey"))
+		require.Equal(t, "head_${key\nkey}_tail",
+			string(ExpandEnv([]byte("head_${key\nkey}_tail"))))
+	})
+	t.Run(`${..} is not expanded if .. contains any "`, func(t *testing.T) {
+		t.Setenv(`key"key`, t.Name())
+		require.Equal(t, t.Name(), os.Getenv(`key"key`))
+		require.Equal(t, `head_${key"key}_tail`,
+			string(ExpandEnv([]byte(`head_${key"key}_tail`))))
+	})
+}

trpc_util.go

@@ -16,7 +16,6 @@ package trpc
 import (
 	"context"
 	"net"
-	"os"
 	"runtime"
 	"sync"
 	"time"
@@ -275,58 +274,6 @@ func Go(ctx context.Context, timeout time.Duration, handler func(context.Context
 	return DefaultGoer.Go(ctx, timeout, handler)
 }
 
-// expandEnv looks for ${var} in s and replaces them with value of the
-// corresponding environment variable.
-// $var is considered invalid.
-// It's not like os.ExpandEnv which will handle both ${var} and $var.
-// Since configurations like password for redis/mysql may contain $, this
-// method is needed.
-func expandEnv(s string) string {
-	var buf []byte
-	i := 0
-	for j := 0; j < len(s); j++ {
-		if s[j] == '$' && j+2 < len(s) && s[j+1] == '{' { // only ${var} instead of $var is valid
-			if buf == nil {
-				buf = make([]byte, 0, 2*len(s))
-			}
-			buf = append(buf, s[i:j]...)
-			name, w := getEnvName(s[j+1:])
-			if name == "" && w > 0 {
-				// invalid matching, remove the $
-			} else if name == "" {
-				buf = append(buf, s[j]) // keep the $
-			} else {
-				buf = append(buf, os.Getenv(name)...)
-			}
-			j += w
-			i = j + 1
-		}
-	}
-	if buf == nil {
-		return s
-	}
-	return string(buf) + s[i:]
-}
-
-// getEnvName gets env name, that is, var from ${var}.
-// And content of var and its len will be returned.
-func getEnvName(s string) (string, int) {
-	// look for right curly bracket '}'
-	// it's guaranteed that the first char is '{' and the string has at least two char
-	for i := 1; i < len(s); i++ {
-		if s[i] == ' ' || s[i] == '\n' || s[i] == '"' { // "xx${xxx"
-			return "", 0 // encounter invalid char, keep the $
-		}
-		if s[i] == '}' {
-			if i == 1 { // ${}
-				return "", 2 // remove ${}
-			}
-			return s[1:i], i + 1
-		}
-	}
-	return "", 0 // no }，keep the $
-}
-
 // --------------- the following code is IP Config related -----------------//
 
 // nicIP defines the parameters used to record the ip address (ipv4 & ipv6) of the nic.