github.com/erda-project/erda-infra@v1.0.10-0.20240327085753-f3a249292aeb/pkg/config/config.go (about)

     1  // Copyright (c) 2021 Terminus, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package config
    16  
    17  import (
    18  	"bytes"
    19  	"encoding/json"
    20  	"fmt"
    21  	"io"
    22  	"io/ioutil"
    23  	"os"
    24  	"path/filepath"
    25  	"regexp"
    26  	"strconv"
    27  	"strings"
    28  
    29  	"github.com/hashicorp/hcl"
    30  	"github.com/magiconair/properties"
    31  	"github.com/mitchellh/mapstructure"
    32  	"github.com/pelletier/go-toml"
    33  	"github.com/recallsong/go-utils/reflectx"
    34  	"gopkg.in/ini.v1"
    35  	"gopkg.in/yaml.v2"
    36  )
    37  
    38  // TrimBOM .
    39  func TrimBOM(f []byte) []byte {
    40  	return bytes.TrimPrefix(f, []byte("\xef\xbb\xbf"))
    41  }
    42  
    43  var (
    44  	envVarRe      = regexp.MustCompile(`\$([\w\|]+)|\$\{([\w\|]+)(:[^}]*)?\}`)
    45  	envVarEscaper = strings.NewReplacer(
    46  		`"`, `\"`,
    47  		`\`, `\\`,
    48  	)
    49  )
    50  
    51  // EscapeEnv .
    52  func EscapeEnv(contents []byte) []byte {
    53  	params := envVarRe.FindAllSubmatch(contents, -1)
    54  	for _, param := range params {
    55  		if len(param) != 4 {
    56  			continue
    57  		}
    58  		var key, defval []byte
    59  		if len(param[1]) > 0 {
    60  			key = param[1]
    61  		} else if len(param[2]) > 0 {
    62  			key = param[2]
    63  		} else {
    64  			continue
    65  		}
    66  		if len(param[3]) > 0 {
    67  			defval = param[3][1:]
    68  		}
    69  		envKey := strings.TrimPrefix(reflectx.BytesToString(key), "$")
    70  		val, ok := os.LookupEnv(envKey)
    71  		if !ok && strings.Contains(envKey, "|") {
    72  			val, ok = lookupEnvWithBooleanExpression(envKey)
    73  		}
    74  		if !ok {
    75  			if len(param[1]) > 0 {
    76  				continue
    77  			}
    78  			val = string(defval)
    79  		}
    80  		val = envVarEscaper.Replace(val)
    81  		contents = bytes.Replace(contents, param[0], reflectx.StringToBytes(val), 1)
    82  	}
    83  	return contents
    84  }
    85  
    86  func lookupEnvWithBooleanExpression(envKey string) (string, bool) {
    87  	keys := strings.Split(envKey, "|")
    88  	if len(keys) == 1 {
    89  		// skip parse boolean value if doesn't contains "|"
    90  		return os.LookupEnv(envKey)
    91  	}
    92  
    93  	var vals []string
    94  	var allBool = true
    95  	var firstTrueVal string
    96  
    97  	for _, key := range keys {
    98  		val, ok := os.LookupEnv(key)
    99  		if !ok {
   100  			continue
   101  		}
   102  
   103  		vals = append(vals, val)
   104  		b, err := strconv.ParseBool(val)
   105  		if err != nil {
   106  			allBool = false
   107  			break
   108  		}
   109  		if b && len(firstTrueVal) == 0 {
   110  			firstTrueVal = val
   111  		}
   112  	}
   113  
   114  	if len(vals) == 0 {
   115  		return "", false
   116  	}
   117  
   118  	if allBool && len(firstTrueVal) > 0 {
   119  		return firstTrueVal, true
   120  	}
   121  	return vals[0], true
   122  }
   123  
   124  // ParseError denotes failing to parse configuration file.
   125  type ParseError struct {
   126  	err error
   127  }
   128  
   129  // Error returns the formatted configuration error.
   130  func (pe ParseError) Error() string {
   131  	return fmt.Sprintf("While parsing config: %s", pe.err.Error())
   132  }
   133  
   134  // UnmarshalToMap .
   135  func UnmarshalToMap(in io.Reader, typ string, c map[string]interface{}) (err error) {
   136  	buf := new(bytes.Buffer)
   137  	_, err = buf.ReadFrom(in)
   138  	if err != nil {
   139  		return err
   140  	}
   141  	err = polishBuffer(buf)
   142  	if err != nil {
   143  		return err
   144  	}
   145  	switch strings.ToLower(typ) {
   146  	case "yaml", "yml":
   147  		if err = yaml.Unmarshal(buf.Bytes(), &c); err != nil {
   148  			return ParseError{err}
   149  		}
   150  	case "json":
   151  		if err = json.Unmarshal(buf.Bytes(), &c); err != nil {
   152  			return ParseError{err}
   153  		}
   154  	case "hcl":
   155  		obj, err := hcl.Parse(reflectx.BytesToString(buf.Bytes()))
   156  		if err != nil {
   157  			return ParseError{err}
   158  		}
   159  		if err = hcl.DecodeObject(&c, obj); err != nil {
   160  			return ParseError{err}
   161  		}
   162  	case "toml":
   163  		tree, err := toml.LoadReader(buf)
   164  		if err != nil {
   165  			return ParseError{err}
   166  		}
   167  		tmap := tree.ToMap()
   168  		for k, v := range tmap {
   169  			c[k] = v
   170  		}
   171  	case "properties", "props", "prop":
   172  		props := properties.NewProperties()
   173  		var err error
   174  		if err = props.Load(buf.Bytes(), properties.UTF8); err != nil {
   175  			return ParseError{err}
   176  		}
   177  		for _, key := range props.Keys() {
   178  			value, _ := props.Get(key)
   179  			// recursively build nested maps
   180  			path := strings.Split(key, ".")
   181  			lastKey := strings.ToLower(path[len(path)-1])
   182  			deepestMap := deepSearch(c, path[0:len(path)-1])
   183  			// set innermost value
   184  			deepestMap[lastKey] = value
   185  		}
   186  	case "ini":
   187  		cfg := ini.Empty()
   188  		err = cfg.Append(buf.Bytes())
   189  		if err != nil {
   190  			return ParseError{err}
   191  		}
   192  		sections := cfg.Sections()
   193  		for i := 0; i < len(sections); i++ {
   194  			section := sections[i]
   195  			keys := section.Keys()
   196  			for j := 0; j < len(keys); j++ {
   197  				key := keys[j]
   198  				value := cfg.Section(section.Name()).Key(key.Name()).String()
   199  				c[section.Name()+"."+key.Name()] = value
   200  			}
   201  		}
   202  	}
   203  	toStringKeyMap(c)
   204  	return nil
   205  }
   206  
   207  func polishBuffer(buf *bytes.Buffer) error {
   208  	byts := buf.Bytes()
   209  	byts = TrimBOM(byts)
   210  	byts = EscapeEnv(byts)
   211  	buf.Reset()
   212  	_, err := buf.Write(byts)
   213  	return err
   214  }
   215  
   216  func toStringKeyMap(i interface{}) interface{} {
   217  	switch x := i.(type) {
   218  	case map[interface{}]interface{}:
   219  		m := map[string]interface{}{}
   220  		for k, v := range x {
   221  			m[fmt.Sprint(k)] = toStringKeyMap(v)
   222  		}
   223  		return m
   224  	case map[string]interface{}:
   225  		for k, v := range x {
   226  			x[k] = toStringKeyMap(v)
   227  		}
   228  	case []interface{}:
   229  		for i, v := range x {
   230  			x[i] = toStringKeyMap(v)
   231  		}
   232  	}
   233  	return i
   234  }
   235  
   236  // deepSearch scans deep maps, following the key indexes listed in the
   237  // sequence "path".
   238  // The last value is expected to be another map, and is returned.
   239  //
   240  // In case intermediate keys do not exist, or map to a non-map value,
   241  // a new map is created and inserted, and the search continues from there:
   242  // the initial map "m" may be modified!
   243  func deepSearch(m map[string]interface{}, path []string) map[string]interface{} {
   244  	for _, k := range path {
   245  		m2, ok := m[k]
   246  		if !ok {
   247  			// intermediate key does not exist
   248  			// => create it and continue from there
   249  			m3 := make(map[string]interface{})
   250  			m[k] = m3
   251  			m = m3
   252  			continue
   253  		}
   254  		m3, ok := m2.(map[string]interface{})
   255  		if !ok {
   256  			// intermediate key is a value
   257  			// => replace with a new map
   258  			m3 = make(map[string]interface{})
   259  			m[k] = m3
   260  		}
   261  		// continue search from here
   262  		m = m3
   263  	}
   264  	return m
   265  }
   266  
   267  // ConvertData .
   268  func ConvertData(input, output interface{}, tag string) error {
   269  	decoder, err := mapstructure.NewDecoder(&mapstructure.DecoderConfig{
   270  		Metadata:         nil,
   271  		Result:           output,
   272  		WeaklyTypedInput: true,
   273  		TagName:          tag,
   274  		DecodeHook: mapstructure.ComposeDecodeHookFunc(
   275  			mapstructure.StringToTimeDurationHookFunc(),
   276  			mapstructure.StringToSliceHookFunc(","),
   277  			mapstructure.StringToTimeHookFunc("2006-01-02 15:04:05"),
   278  		),
   279  	})
   280  	if err != nil {
   281  		return err
   282  	}
   283  	return decoder.Decode(input)
   284  }
   285  
   286  // LoadFile .
   287  func LoadFile(path string) ([]byte, error) {
   288  	byts, err := ioutil.ReadFile(path)
   289  	return byts, err
   290  }
   291  
   292  // LoadToMap .
   293  func LoadToMap(path string, c map[string]interface{}) error {
   294  	typ := filepath.Ext(path)
   295  	if len(typ) <= 0 {
   296  		return fmt.Errorf("%s unknown file extension", path)
   297  	}
   298  	byts, err := LoadFile(path)
   299  	if err != nil {
   300  		return err
   301  	}
   302  	return UnmarshalToMap(bytes.NewReader(byts), typ[1:], c)
   303  }
   304  
   305  // LoadEnvFileWithPath .
   306  func LoadEnvFileWithPath(path string, override bool) {
   307  	byts, err := ioutil.ReadFile(path)
   308  	if err != nil {
   309  		if os.IsNotExist(err) {
   310  			return
   311  		}
   312  		return
   313  	}
   314  	regex := regexp.MustCompile(`\s+\#`)
   315  	content := reflectx.BytesToString(byts)
   316  	for _, line := range strings.Split(content, "\n") {
   317  		if strings.HasPrefix(line, "#") {
   318  			continue
   319  		}
   320  		loc := regex.FindIndex(reflectx.StringToBytes(line))
   321  		if len(loc) > 0 {
   322  			line = line[0:loc[0]]
   323  		}
   324  		idx := strings.Index(line, "=")
   325  		if idx <= 0 {
   326  			continue
   327  		}
   328  		key := strings.TrimSpace(line[:idx])
   329  		if len(key) <= 0 {
   330  			continue
   331  		}
   332  		val := strings.TrimSpace(line[idx+1:])
   333  		if override {
   334  			os.Setenv(key, val)
   335  		} else {
   336  			_, ok := os.LookupEnv(key)
   337  			if !ok {
   338  				os.Setenv(key, val)
   339  			}
   340  		}
   341  
   342  	}
   343  }
   344  
   345  // LoadEnvFile .
   346  func LoadEnvFile() {
   347  	wd, err := os.Getwd()
   348  	if err != nil {
   349  		return
   350  	}
   351  	path := filepath.Join(wd, ".env")
   352  	LoadEnvFileWithPath(path, false)
   353  }