From 434e791036ca767694d8251bb002660547114fac Mon Sep 17 00:00:00 2001
From: Martin Hutchinson <mhutchinson@gmail.com>
Date: Wed, 31 Jan 2024 12:10:09 +0000
Subject: [PATCH] Allow the witness identities to be provided as flags (#77)

This will allow us to easily configure the CI and dev environments from CloudRun. The configuration file option cannot be used because of the nature of cloud run not having a file system. The other option would be to allow the config file to be read from a URL, or GCS, but this option is cleaner.

Also fixed a bug where it was always parsing the built-in configuration
instead of whatever was provided at runtime.
---
 cmd/main.go      | 39 +++++++++++++++++++++++++++++++++++----
 config/config.go |  2 +-
 2 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/cmd/main.go b/cmd/main.go
index c3f7a72..0c47463 100644
--- a/cmd/main.go
+++ b/cmd/main.go
@@ -24,6 +24,7 @@ import (
 	"net"
 	"net/http"
 	"os"
+	"strings"
 
 	"cloud.google.com/go/cloudsqlconn"
 	"github.com/golang/glog"
@@ -33,6 +34,7 @@ import (
 	"github.com/transparency-dev/distributor/config"
 	"golang.org/x/mod/sumdb/note"
 	"golang.org/x/sync/errgroup"
+	"gopkg.in/yaml.v3"
 
 	_ "embed"
 
@@ -45,10 +47,12 @@ var (
 	useCloudSql = flag.Bool("use_cloud_sql", false, "Set to true to set up the DB connection using cloudsql connection. This will ignore mysql_uri and generate it from env variables.")
 	mysqlURI    = flag.String("mysql_uri", "", "URI for MySQL DB")
 
-	witnessConfigFile = flag.String("witness_config_file", "", "Path to a file containing the public keys of allowed witnesses")
+	witnessConfigFile = flag.String("witness_config_file", "", "Path to a file containing the public keys of allowed witnesses. Mutually exclusive with witkey.")
+	witnessKeys       witFlags
 )
 
 func main() {
+	flag.Var(&witnessKeys, "witkey", "Provide one or more witness keys directly as flags (can specify multiple times). Mutually exclusive with witness_config_file.")
 	flag.Parse()
 	ctx := context.Background()
 
@@ -153,16 +157,31 @@ func getLogsOrDie() map[string]config.LogInfo {
 }
 
 func getWitnessesOrDie() map[string]note.Verifier {
-	cfg := config.WitnessesYAML
-	if *witnessConfigFile != "" {
+	var cfg []byte
+	if witFile, witFlags := *witnessConfigFile != "", len(witnessKeys) > 0; witFile && !witFlags {
 		c, err := os.ReadFile(*witnessConfigFile)
 		if err != nil {
 			glog.Exitf("Failed to read witness_config_file (%q): %v", *witnessConfigFile, err)
 		}
 		glog.Infof("Witness list read from %v", *witnessConfigFile)
 		cfg = c
+	} else if !witFile && witFlags {
+		// This is a bit messy to turn flags into yaml and then parse them again, but the cost
+		// is small, and the benefit is that we guarantee the same parsing & instantiation logic.
+		witCfg := struct {
+			Witnesses []string `yaml:"Witnesses"`
+		}{}
+		witCfg.Witnesses = witnessKeys
+		var err error
+		cfg, err = yaml.Marshal(witCfg)
+		if err != nil {
+			glog.Exitf("Failed to marshal witness config: %v", err)
+		}
+	} else if !witFile && !witFlags {
+		glog.Info("Flags witness_config_file nor witkey are specified; default witness list will be used")
+		cfg = config.WitnessesYAML
 	} else {
-		glog.Info("Flag witness_config_file not specified; default witness list will be used")
+		glog.Exitf("Only one of witness_config_file and witkey can be specified")
 	}
 	w, err := config.ParseWitnessesConfig(cfg)
 	if err != nil {
@@ -173,5 +192,17 @@ func getWitnessesOrDie() map[string]note.Verifier {
 	for _, v := range w {
 		r[v.Name()] = v
 	}
+	glog.Infof("Configured with %d witness keys: %s", len(r), r)
 	return r
 }
+
+type witFlags []string
+
+func (wf *witFlags) String() string {
+	return strings.Join(*wf, ",")
+}
+
+func (wf *witFlags) Set(w string) error {
+	*wf = append(*wf, w)
+	return nil
+}
diff --git a/config/config.go b/config/config.go
index bbf9a3c..700d9ac 100644
--- a/config/config.go
+++ b/config/config.go
@@ -85,7 +85,7 @@ func ParseWitnessesConfig(y []byte) (map[uint32]note.Verifier, error) {
 	witCfg := struct {
 		Witnesses []string `yaml:"Witnesses"`
 	}{}
-	if err := yaml.Unmarshal(WitnessesYAML, &witCfg); err != nil {
+	if err := yaml.Unmarshal(y, &witCfg); err != nil {
 		return nil, fmt.Errorf("failed to unmarshal witness config: %v", err)
 	}
 	ws := make(map[uint32]note.Verifier)