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