Allow flag or file based configuration

cmd/createtree/main.go

@@ -38,9 +38,11 @@ import (
 	"fmt"
 	"os"
 
+	"github.com/golang/glog"
 	"github.com/golang/protobuf/ptypes"
 	"github.com/golang/protobuf/ptypes/any"
 	"github.com/google/trillian"
+	"github.com/google/trillian/cmd"
 	"github.com/google/trillian/crypto/keyspb"
 	"github.com/google/trillian/crypto/sigpb"
 	"google.golang.org/grpc"
@@ -60,6 +62,8 @@ var (
 	privateKeyFormat = flag.String("private_key_format", "PEMKeyFile", "Type of private key to be used")
 	pemKeyPath       = flag.String("pem_key_path", "", "Path to the private key PEM file")
 	pemKeyPassword   = flag.String("pem_key_password", "", "Password of the private key PEM file")
+
+	configFile = flag.String("config", "", "Config file containing flags, file contents can be overridden by command line flags")
 )
 
 // createOpts contains all user-supplied options required to run the program.
@@ -175,6 +179,12 @@ func newOptsFromFlags() *createOpts {
 func main() {
 	flag.Parse()
 
+	if *configFile != "" {
+		if err := cmd.ParseFlagFile(*configFile); err != nil {
+			glog.Exitf("Failed to load flags from config file %q: %s", *configFile, err)
+		}
+	}
+
 	ctx := context.Background()
 	tree, err := createTree(ctx, newOptsFromFlags())
 	if err != nil {

cmd/flags.go

@@ -0,0 +1,56 @@
+// Copyright 2017 Google Inc. All Rights Reserved.
+//
+// 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.
+
+package cmd
+
+import (
+	"errors"
+	"flag"
+	"io/ioutil"
+	"os"
+
+	"bitbucket.org/creachadair/shell"
+)
+
+func parseFlags(file string) error {
+	args, valid := shell.Split(file)
+	if !valid {
+		return errors.New("flag file contains unclosed quotations")
+	}
+	// Expand any environment variables in the args
+	for i := range args {
+		args[i] = os.ExpandEnv(args[i])
+	}
+
+	if err := flag.CommandLine.Parse(args); err != nil {
+		return err
+	}
+
+	// Call flag.Parse() again so that command line flags
+	// can override flags provided in the provided flag file.
+	flag.Parse()
+	return nil
+}
+
+// ParseFlagFile parses a set of flags from a file at the provided
+// path. Re-calls flag.Parse() after parsing the flags in the file
+// so that flags provided on the command line take precedence over
+// flags provided in the file.
+func ParseFlagFile(path string) error {
+	file, err := ioutil.ReadFile(path)
+	if err != nil {
+		return err
+	}
+	return parseFlags(string(file))
+}

cmd/flags_test.go

@@ -0,0 +1,109 @@
+// Copyright 2017 Google Inc. All Rights Reserved.
+//
+// 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.
+
+package cmd
+
+import (
+	"flag"
+	"os"
+	"testing"
+)
+
+func TestParseFlags(t *testing.T) {
+	var a, b string
+	flag.StringVar(&a, "a", "", "")
+	flag.StringVar(&b, "b", "", "")
+
+	flag.CommandLine.Init(os.Args[0], flag.ContinueOnError)
+
+	tests := []struct {
+		name        string
+		contents    string
+		env         map[string]string
+		cliArgs     []string
+		expectedErr string
+		expectedA   string
+		expectedB   string
+	}{
+		{
+			name:      "two flags per line",
+			contents:  "-a one -b two",
+			expectedA: "one",
+			expectedB: "two",
+		},
+		{
+			name:      "one flag per line",
+			contents:  "-a one\n-b two",
+			expectedA: "one",
+			expectedB: "two",
+		},
+		{
+			name:      "one flag per line, with line continuation",
+			contents:  "-a one \\\n-b two",
+			expectedA: "one",
+			expectedB: "two",
+		},
+		{
+			name:      "one flag in file, one flag on command-line",
+			contents:  "-a one",
+			cliArgs:   []string{"-b", "two"},
+			expectedA: "one",
+			expectedB: "two",
+		},
+		{
+			name:      "two flags, one overridden by command-line",
+			contents:  "-a one\n-b two",
+			cliArgs:   []string{"-b", "three"},
+			expectedA: "one",
+			expectedB: "three",
+		},
+		{
+			name:      "two flags, one using an environment variable",
+			contents:  "-a one\n-b $TEST_VAR",
+			env:       map[string]string{"TEST_VAR": "from env"},
+			expectedA: "one",
+			expectedB: "from env",
+		},
+		{
+			name:        "three flags, one undefined",
+			contents:    "-a one -b two -c three",
+			expectedErr: "flag provided but not defined: -c",
+		},
+	}
+
+	initialArgs := os.Args[:]
+	for _, tc := range tests {
+		a, b = "", ""
+		os.Args = append(initialArgs, tc.cliArgs...)
+		for k, v := range tc.env {
+			if err := os.Setenv(k, v); err != nil {
+				t.Errorf("%v: os.SetEnv(%q, %q) = %q", tc.name, k, v, err)
+			}
+		}
+
+		if err := parseFlags(tc.contents); err != nil {
+			if err.Error() != tc.expectedErr {
+				t.Errorf("%v: parseFlags() = %q, want %q", tc.name, err, tc.expectedErr)
+			}
+			continue
+		}
+
+		if tc.expectedA != a {
+			t.Errorf("%v: flag 'a' not properly set: got %q, want %q", tc.name, a, tc.expectedA)
+		}
+		if tc.expectedB != b {
+			t.Errorf("%v: flag 'b' not properly set: got %q, want %q", tc.name, b, tc.expectedB)
+		}
+	}
+}

server/trillian_log_server/main.go

@@ -27,6 +27,7 @@ import (
 	etcdnaming "github.com/coreos/etcd/clientv3/naming"
 	"github.com/golang/glog"
 	"github.com/google/trillian"
+	"github.com/google/trillian/cmd"
 	"github.com/google/trillian/crypto/keys"
 	"github.com/google/trillian/extension"
 	"github.com/google/trillian/monitoring"
@@ -48,10 +49,19 @@ var (
 	etcdServers         = flag.String("etcd_servers", "", "A comma-separated list of etcd servers; no etcd registration if empty")
 	etcdService         = flag.String("etcd_service", "trillian-log", "Service name to announce ourselves under")
 	maxUnsequencedRows  = flag.Int("max_unsequenced_rows", mysqlq.DefaultMaxUnsequenced, "Max number of unsequenced rows before rate limiting kicks in")
+
+	configFile = flag.String("config", "", "Config file containing flags, file contents can be overridden by command line flags")
 )
 
 func main() {
 	flag.Parse()
+
+	if *configFile != "" {
+		if err := cmd.ParseFlagFile(*configFile); err != nil {
+			glog.Exitf("Failed to load flags from config file %q: %s", *configFile, err)
+		}
+	}
+
 	ctx := context.Background()
 
 	// First make sure we can access the database, quit if not

server/trillian_log_signer/main.go

@@ -23,6 +23,7 @@ import (
 	_ "github.com/go-sql-driver/mysql" // Load MySQL driver
 
 	"github.com/golang/glog"
+	"github.com/google/trillian/cmd"
 	"github.com/google/trillian/crypto/keys"
 	"github.com/google/trillian/extension"
 	"github.com/google/trillian/monitoring/metric"
@@ -50,10 +51,19 @@ var (
 	masterCheckInterval = flag.Duration("master_check_interval", 5*time.Second, "Interval between checking mastership still held")
 	masterHoldInterval  = flag.Duration("master_hold_interval", 60*time.Second, "Minimum interval to hold mastership for")
 	resignOdds          = flag.Int("resign_odds", 10, "Chance of resigning mastership after each check, the N in 1-in-N")
+
+	configFile = flag.String("config", "", "Config file containing flags, file contents can be overridden by command line flags")
 )
 
 func main() {
 	flag.Parse()
+
+	if *configFile != "" {
+		if err := cmd.ParseFlagFile(*configFile); err != nil {
+			glog.Exitf("Failed to load flags from config file %q: %s", *configFile, err)
+		}
+	}
+
 	glog.CopyStandardLogTo("WARNING")
 	glog.Info("**** Log Signer Starting ****")
 

server/vmap/trillian_map_server/main.go

@@ -23,6 +23,7 @@ import (
 
 	"github.com/golang/glog"
 	"github.com/google/trillian"
+	"github.com/google/trillian/cmd"
 	"github.com/google/trillian/crypto/keys"
 	"github.com/google/trillian/extension"
 	mysqlq "github.com/google/trillian/quota/mysql"
@@ -39,11 +40,19 @@ var (
 	rpcEndpoint        = flag.String("rpc_endpoint", "localhost:8090", "Endpoint for RPC requests (host:port)")
 	httpEndpoint       = flag.String("http_endpoint", "localhost:8091", "Endpoint for HTTP metrics and REST requests on (host:port, empty means disabled)")
 	maxUnsequencedRows = flag.Int("max_unsequenced_rows", mysqlq.DefaultMaxUnsequenced, "Max number of unsequenced rows before rate limiting kicks in")
+
+	configFile = flag.String("config", "", "Config file containing flags, file contents can be overridden by command line flags")
 )
 
 func main() {
 	flag.Parse()
 
+	if *configFile != "" {
+		if err := cmd.ParseFlagFile(*configFile); err != nil {
+			glog.Exitf("Failed to load flags from config file %q: %s", *configFile, err)
+		}
+	}
+
 	db, err := mysql.OpenDB(*mySQLURI)
 	if err != nil {
 		glog.Exitf("Failed to open database: %v", err)