trpc.group/trpc-go/trpc-go@v1.0.3/internal/expandenv/expand_env.go (about)

     1  // Package expandenv replaces ${key} in byte slices with the env value of key.
     2  package expandenv
     3  
     4  import (
     5  	"os"
     6  )
     7  
     8  // ExpandEnv looks for ${var} in s and replaces them with value of the corresponding environment variable.
     9  // It's not like os.ExpandEnv which handles both ${var} and $var.
    10  // $var is considered invalid, since configurations like password for redis/mysql may contain $.
    11  func ExpandEnv(s []byte) []byte {
    12  	var buf []byte
    13  	i := 0
    14  	for j := 0; j < len(s); j++ {
    15  		if s[j] == '$' && j+2 < len(s) && s[j+1] == '{' { // only ${var} instead of $var is valid
    16  			if buf == nil {
    17  				buf = make([]byte, 0, 2*len(s))
    18  			}
    19  			buf = append(buf, s[i:j]...)
    20  			name, w := getEnvName(s[j+1:])
    21  			if name == nil && w > 0 {
    22  				// invalid matching, remove the $
    23  			} else if name == nil {
    24  				buf = append(buf, s[j]) // keep the $
    25  			} else {
    26  				buf = append(buf, os.Getenv(string(name))...)
    27  			}
    28  			j += w
    29  			i = j + 1
    30  		}
    31  	}
    32  	if buf == nil {
    33  		return s
    34  	}
    35  	return append(buf, s[i:]...)
    36  }
    37  
    38  // getEnvName gets env name, that is, var from ${var}.
    39  // The env name and its len will be returned.
    40  func getEnvName(s []byte) ([]byte, int) {
    41  	// look for right curly bracket '}'
    42  	// it's guaranteed that the first char is '{' and the string has at least two char
    43  	for i := 1; i < len(s); i++ {
    44  		if s[i] == ' ' || s[i] == '\n' || s[i] == '"' { // "xx${xxx"
    45  			return nil, 0 // encounter invalid char, keep the $
    46  		}
    47  		if s[i] == '}' {
    48  			if i == 1 { // ${}
    49  				return nil, 2 // remove ${}
    50  			}
    51  			return s[1:i], i + 1
    52  		}
    53  	}
    54  	return nil, 0 // no }，keep the $
    55  }