2016-05-13 06:48:53 +03:00
|
|
|
// Copyright 2015 The etcd Authors
|
2015-01-25 06:19:16 +03:00
|
|
|
//
|
|
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
|
|
// you may not use this file except in compliance with the License.
|
|
|
|
// You may obtain a copy of the License at
|
|
|
|
//
|
|
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
//
|
|
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
|
|
// See the License for the specific language governing permissions and
|
|
|
|
// limitations under the License.
|
2014-10-18 02:41:22 +04:00
|
|
|
|
2015-11-18 07:54:10 +03:00
|
|
|
// Package flags implements command-line flag parsing.
|
2014-10-27 19:46:37 +03:00
|
|
|
package flags
|
2014-09-25 03:06:28 +04:00
|
|
|
|
|
|
|
import (
|
|
|
|
"flag"
|
|
|
|
"fmt"
|
|
|
|
"os"
|
2014-09-25 03:29:47 +04:00
|
|
|
"strings"
|
2014-10-01 22:40:35 +04:00
|
|
|
|
2016-03-25 21:40:24 +03:00
|
|
|
"github.com/spf13/pflag"
|
2020-02-12 22:15:41 +03:00
|
|
|
"go.uber.org/zap"
|
2015-06-11 18:52:55 +03:00
|
|
|
)
|
|
|
|
|
2014-09-25 03:29:47 +04:00
|
|
|
// SetFlagsFromEnv parses all registered flags in the given flagset,
|
|
|
|
// and if they are not already set it attempts to set their values from
|
|
|
|
// environment variables. Environment variables take the name of the flag but
|
2016-03-25 21:40:24 +03:00
|
|
|
// are UPPERCASE, have the given prefix and any dashes are replaced by
|
2014-09-25 03:29:47 +04:00
|
|
|
// underscores - for example: some-flag => ETCD_SOME_FLAG
|
2020-02-12 22:15:41 +03:00
|
|
|
func SetFlagsFromEnv(lg *zap.Logger, prefix string, fs *flag.FlagSet) error {
|
2014-11-07 03:40:13 +03:00
|
|
|
var err error
|
2014-09-25 03:29:47 +04:00
|
|
|
alreadySet := make(map[string]bool)
|
|
|
|
fs.Visit(func(f *flag.Flag) {
|
2016-10-14 21:15:28 +03:00
|
|
|
alreadySet[FlagToEnv(prefix, f.Name)] = true
|
2014-09-25 03:29:47 +04:00
|
|
|
})
|
2015-08-27 22:08:31 +03:00
|
|
|
usedEnvKey := make(map[string]bool)
|
2014-09-25 03:29:47 +04:00
|
|
|
fs.VisitAll(func(f *flag.Flag) {
|
2020-02-12 22:15:41 +03:00
|
|
|
if serr := setFlagFromEnv(lg, fs, prefix, f.Name, usedEnvKey, alreadySet, true); serr != nil {
|
2017-11-27 17:41:43 +03:00
|
|
|
err = serr
|
|
|
|
}
|
2016-03-25 21:40:24 +03:00
|
|
|
})
|
2020-02-12 22:15:41 +03:00
|
|
|
verifyEnv(lg, prefix, usedEnvKey, alreadySet)
|
2016-03-25 21:40:24 +03:00
|
|
|
return err
|
|
|
|
}
|
|
|
|
|
|
|
|
// SetPflagsFromEnv is similar to SetFlagsFromEnv. However, the accepted flagset type is pflag.FlagSet
|
|
|
|
// and it does not do any logging.
|
2020-02-12 22:15:41 +03:00
|
|
|
func SetPflagsFromEnv(lg *zap.Logger, prefix string, fs *pflag.FlagSet) error {
|
2016-03-25 21:40:24 +03:00
|
|
|
var err error
|
|
|
|
alreadySet := make(map[string]bool)
|
|
|
|
usedEnvKey := make(map[string]bool)
|
|
|
|
fs.VisitAll(func(f *pflag.Flag) {
|
|
|
|
if f.Changed {
|
2016-10-14 21:15:28 +03:00
|
|
|
alreadySet[FlagToEnv(prefix, f.Name)] = true
|
2016-03-25 21:40:24 +03:00
|
|
|
}
|
2020-02-12 22:15:41 +03:00
|
|
|
if serr := setFlagFromEnv(lg, fs, prefix, f.Name, usedEnvKey, alreadySet, false); serr != nil {
|
2016-03-25 21:40:24 +03:00
|
|
|
err = serr
|
2014-09-25 03:29:47 +04:00
|
|
|
}
|
|
|
|
})
|
2020-02-12 22:15:41 +03:00
|
|
|
verifyEnv(lg, prefix, usedEnvKey, alreadySet)
|
2016-03-25 21:40:24 +03:00
|
|
|
return err
|
|
|
|
}
|
|
|
|
|
2016-10-14 21:15:28 +03:00
|
|
|
// FlagToEnv converts flag string to upper-case environment variable key string.
|
|
|
|
func FlagToEnv(prefix, name string) string {
|
2022-11-14 23:18:53 +03:00
|
|
|
return prefix + "_" + strings.ToUpper(strings.ReplaceAll(name, "-", "_"))
|
2016-03-25 21:40:24 +03:00
|
|
|
}
|
2015-08-27 22:08:31 +03:00
|
|
|
|
2020-02-12 22:15:41 +03:00
|
|
|
func verifyEnv(lg *zap.Logger, prefix string, usedEnvKey, alreadySet map[string]bool) {
|
2015-08-27 22:08:31 +03:00
|
|
|
for _, env := range os.Environ() {
|
|
|
|
kv := strings.SplitN(env, "=", 2)
|
|
|
|
if len(kv) != 2 {
|
2020-02-12 22:15:41 +03:00
|
|
|
if lg != nil {
|
|
|
|
lg.Warn("found invalid environment variable", zap.String("environment-variable", env))
|
|
|
|
}
|
2015-08-27 22:08:31 +03:00
|
|
|
}
|
|
|
|
if usedEnvKey[kv[0]] {
|
|
|
|
continue
|
|
|
|
}
|
|
|
|
if alreadySet[kv[0]] {
|
2020-02-12 22:15:41 +03:00
|
|
|
if lg != nil {
|
|
|
|
lg.Fatal(
|
|
|
|
"conflicting environment variable is shadowed by corresponding command-line flag (either unset environment variable or disable flag))",
|
|
|
|
zap.String("environment-variable", kv[0]),
|
|
|
|
)
|
|
|
|
}
|
2015-08-27 22:08:31 +03:00
|
|
|
}
|
2017-01-20 10:43:40 +03:00
|
|
|
if strings.HasPrefix(env, prefix+"_") {
|
2020-02-12 22:15:41 +03:00
|
|
|
if lg != nil {
|
|
|
|
lg.Warn("unrecognized environment variable", zap.String("environment-variable", env))
|
|
|
|
}
|
2015-08-27 22:08:31 +03:00
|
|
|
}
|
|
|
|
}
|
2016-03-25 21:40:24 +03:00
|
|
|
}
|
2015-08-27 22:08:31 +03:00
|
|
|
|
2016-03-25 21:40:24 +03:00
|
|
|
type flagSetter interface {
|
|
|
|
Set(fk string, fv string) error
|
2014-09-25 03:29:47 +04:00
|
|
|
}
|
2014-10-01 22:40:35 +04:00
|
|
|
|
2020-02-12 22:15:41 +03:00
|
|
|
func setFlagFromEnv(lg *zap.Logger, fs flagSetter, prefix, fname string, usedEnvKey, alreadySet map[string]bool, log bool) error {
|
2016-10-14 21:15:28 +03:00
|
|
|
key := FlagToEnv(prefix, fname)
|
2016-03-25 21:40:24 +03:00
|
|
|
if !alreadySet[key] {
|
|
|
|
val := os.Getenv(key)
|
|
|
|
if val != "" {
|
|
|
|
usedEnvKey[key] = true
|
|
|
|
if serr := fs.Set(fname, val); serr != nil {
|
|
|
|
return fmt.Errorf("invalid value %q for %s: %v", val, key, serr)
|
|
|
|
}
|
2020-02-12 22:15:41 +03:00
|
|
|
if log && lg != nil {
|
|
|
|
lg.Info(
|
|
|
|
"recognized and used environment variable",
|
|
|
|
zap.String("variable-name", key),
|
|
|
|
zap.String("variable-value", val),
|
|
|
|
)
|
2016-03-25 21:40:24 +03:00
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
return nil
|
2015-08-27 22:08:31 +03:00
|
|
|
}
|
|
|
|
|
2015-01-23 00:40:45 +03:00
|
|
|
func IsSet(fs *flag.FlagSet, name string) bool {
|
|
|
|
set := false
|
|
|
|
fs.Visit(func(f *flag.Flag) {
|
|
|
|
if f.Name == name {
|
|
|
|
set = true
|
|
|
|
}
|
|
|
|
})
|
|
|
|
return set
|
|
|
|
}
|