1// Copyright 2015 The etcd Authors 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 flags implements command-line flag parsing. 16package flags 17 18import ( 19 "flag" 20 "fmt" 21 "os" 22 "strings" 23 24 "github.com/spf13/pflag" 25 "go.uber.org/zap" 26) 27 28// SetFlagsFromEnv parses all registered flags in the given flagset, 29// and if they are not already set it attempts to set their values from 30// environment variables. Environment variables take the name of the flag but 31// are UPPERCASE, have the given prefix and any dashes are replaced by 32// underscores - for example: some-flag => ETCD_SOME_FLAG 33func SetFlagsFromEnv(lg *zap.Logger, prefix string, fs *flag.FlagSet) error { 34 var err error 35 alreadySet := make(map[string]bool) 36 fs.Visit(func(f *flag.Flag) { 37 alreadySet[FlagToEnv(prefix, f.Name)] = true 38 }) 39 usedEnvKey := make(map[string]bool) 40 fs.VisitAll(func(f *flag.Flag) { 41 if serr := setFlagFromEnv(lg, fs, prefix, f.Name, usedEnvKey, alreadySet, true); serr != nil { 42 err = serr 43 } 44 }) 45 verifyEnv(lg, prefix, usedEnvKey, alreadySet) 46 return err 47} 48 49// SetPflagsFromEnv is similar to SetFlagsFromEnv. However, the accepted flagset type is pflag.FlagSet 50// and it does not do any logging. 51func SetPflagsFromEnv(lg *zap.Logger, prefix string, fs *pflag.FlagSet) error { 52 var err error 53 alreadySet := make(map[string]bool) 54 usedEnvKey := make(map[string]bool) 55 fs.VisitAll(func(f *pflag.Flag) { 56 if f.Changed { 57 alreadySet[FlagToEnv(prefix, f.Name)] = true 58 } 59 if serr := setFlagFromEnv(lg, fs, prefix, f.Name, usedEnvKey, alreadySet, false); serr != nil { 60 err = serr 61 } 62 }) 63 verifyEnv(lg, prefix, usedEnvKey, alreadySet) 64 return err 65} 66 67// FlagToEnv converts flag string to upper-case environment variable key string. 68func FlagToEnv(prefix, name string) string { 69 return prefix + "_" + strings.ToUpper(strings.Replace(name, "-", "_", -1)) 70} 71 72func verifyEnv(lg *zap.Logger, prefix string, usedEnvKey, alreadySet map[string]bool) { 73 for _, env := range os.Environ() { 74 kv := strings.SplitN(env, "=", 2) 75 if len(kv) != 2 { 76 if lg != nil { 77 lg.Warn("found invalid environment variable", zap.String("environment-variable", env)) 78 } 79 } 80 if usedEnvKey[kv[0]] { 81 continue 82 } 83 if alreadySet[kv[0]] { 84 if lg != nil { 85 lg.Fatal( 86 "conflicting environment variable is shadowed by corresponding command-line flag (either unset environment variable or disable flag))", 87 zap.String("environment-variable", kv[0]), 88 ) 89 } 90 } 91 if strings.HasPrefix(env, prefix+"_") { 92 if lg != nil { 93 lg.Warn("unrecognized environment variable", zap.String("environment-variable", env)) 94 } 95 } 96 } 97} 98 99type flagSetter interface { 100 Set(fk string, fv string) error 101} 102 103func setFlagFromEnv(lg *zap.Logger, fs flagSetter, prefix, fname string, usedEnvKey, alreadySet map[string]bool, log bool) error { 104 key := FlagToEnv(prefix, fname) 105 if !alreadySet[key] { 106 val := os.Getenv(key) 107 if val != "" { 108 usedEnvKey[key] = true 109 if serr := fs.Set(fname, val); serr != nil { 110 return fmt.Errorf("invalid value %q for %s: %v", val, key, serr) 111 } 112 if log && lg != nil { 113 lg.Info( 114 "recognized and used environment variable", 115 zap.String("variable-name", key), 116 zap.String("variable-value", val), 117 ) 118 } 119 } 120 } 121 return nil 122} 123 124func IsSet(fs *flag.FlagSet, name string) bool { 125 set := false 126 fs.Visit(func(f *flag.Flag) { 127 if f.Name == name { 128 set = true 129 } 130 }) 131 return set 132} 133