From ec189e3c142b74745e19d24b2b538923a4823d91 Mon Sep 17 00:00:00 2001
From: "Timothy J. Raymond" <traymond@microsoft.com>
Date: Tue, 3 Sep 2024 11:52:07 -0400
Subject: [PATCH] Add ability to overlay config to GetConfig

In order to make it possible for specific fields to be overridden by
user-controllable config files, the GetConfig function must be modified
to support an additional configuration file that can be filtered to
specific fields. This modifies GetConfig in a backwards-compatible way
to support a FilteredConfig type, specifiying additional overlays with
allowable fields from those config overlays.
---
 pkg/config/config.go                 | 50 ++++++++++++++++++++--------
 pkg/config/config_test.go            |  7 +++-
 pkg/config/internal/internal.go      |  7 ++--
 pkg/config/internal/internal_test.go |  2 +-
 4 files changed, 47 insertions(+), 19 deletions(-)

diff --git a/pkg/config/config.go b/pkg/config/config.go
index 0efe098a28..89a7f48558 100644
--- a/pkg/config/config.go
+++ b/pkg/config/config.go
@@ -4,10 +4,12 @@ package config
 
 import (
 	"fmt"
+	"os"
 	"reflect"
 	"strings"
 	"time"
 
+	"github.com/microsoft/retina/pkg/config/internal"
 	"github.com/mitchellh/mapstructure"
 	"github.com/pkg/errors"
 	"github.com/spf13/viper"
@@ -67,26 +69,37 @@ type Config struct {
 	MonitorSockPath          string        `yaml:"monitorSockPath"`
 }
 
-func GetConfig(files ...string) (*Config, error) {
-	if len(files) == 0 {
-		viper.SetConfigName("config")
-		viper.AddConfigPath("/retina/config")
+type FilteredConfig struct {
+	Filename      string
+	AllowedFields []string
+}
+
+func mergeConfig(file FilteredConfig) error {
+	f, err := os.Open(file.Filename)
+	if err != nil {
+		return errors.Wrapf(err, "opening config file %q", file)
+	}
+	defer f.Close()
+
+	fy, err := internal.NewFilteredYAML(f, file.AllowedFields)
+	if err != nil {
+		return errors.Wrap(err, "creating FilteredYAML")
+	}
+
+	err = viper.MergeConfig(fy)
+	if err != nil {
+		return errors.Wrap(err, "merging config with viper")
 	}
+	return nil
+}
 
-	cfgFilename := files[0]
-	if cfgFilename != "" {
-		viper.SetConfigFile(cfgFilename)
+func GetConfig(primaryCfg string, overlays ...FilteredConfig) (*Config, error) {
+	if primaryCfg != "" {
+		viper.SetConfigFile(primaryCfg)
 	} else {
 		viper.SetConfigName("config")
 		viper.AddConfigPath("/retina/config")
 	}
-	for _, file := range files[1:] {
-		viper.SetConfigFile(file)
-		if err := viper.MergeInConfig(); err != nil {
-			return nil, errors.Wrapf(err, "loading config file %q", file)
-		}
-	}
-
 	viper.SetEnvPrefix("retina")
 	viper.AutomaticEnv()
 	// NOTE(mainred): RetinaEndpoint is currently the only supported solution to cache Pod, and before an alternative is implemented,
@@ -97,6 +110,15 @@ func GetConfig(files ...string) (*Config, error) {
 	if err != nil {
 		return nil, fmt.Errorf("fatal error config file: %s", err)
 	}
+
+	// apply overlay configs
+	for _, file := range overlays {
+		err := mergeConfig(file)
+		if err != nil {
+			return nil, errors.Wrapf(err, "merging config for %q", file)
+		}
+	}
+
 	var config Config
 	decoderConfigOption := func(dc *mapstructure.DecoderConfig) {
 		dc.DecodeHook = mapstructure.ComposeDecodeHookFunc(
diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go
index 7b686a55a8..cc9623b990 100644
--- a/pkg/config/config_test.go
+++ b/pkg/config/config_test.go
@@ -32,7 +32,12 @@ func TestGetConfig(t *testing.T) {
 }
 
 func TestGetConfigOverlay(t *testing.T) {
-	c, err := GetConfig("./testdata/config.yaml", "./testdata/overlay.yaml")
+	c, err := GetConfig("./testdata/config.yaml", FilteredConfig{
+		Filename: "./testdata/overlay.yaml",
+		AllowedFields: []string{
+			"logLevel",
+		},
+	})
 	if err != nil {
 		t.Fatal("err getting config: err:", err)
 	}
diff --git a/pkg/config/internal/internal.go b/pkg/config/internal/internal.go
index 0fecf037a4..70aff16113 100644
--- a/pkg/config/internal/internal.go
+++ b/pkg/config/internal/internal.go
@@ -8,7 +8,7 @@ import (
 	"gopkg.in/yaml.v2"
 )
 
-func NewFilteredYAML(source io.Reader, allowedFields []string) (*FilteredYAML, error) {
+func NewFilteredYAML(source io.ReadCloser, allowedFields []string) (*FilteredYAML, error) {
 	f := &FilteredYAML{
 		YAML:          source,
 		AllowedFields: allowedFields,
@@ -26,12 +26,13 @@ func NewFilteredYAML(source io.Reader, allowedFields []string) (*FilteredYAML, e
 // fields. Any additional fields found will be removed, such that the resulting
 // configuration is the subset of fields found in the allowlist.
 type FilteredYAML struct {
-	YAML          io.Reader // the input YAML
-	AllowedFields []string  // the set of allowed fields in the resulting YAML
+	YAML          io.ReadCloser // the input YAML
+	AllowedFields []string      // the set of allowed fields in the resulting YAML
 	buf           *bytes.Buffer
 }
 
 func (f *FilteredYAML) filter() error {
+	defer f.YAML.Close()
 	f.buf = bytes.NewBufferString("")
 
 	decoded := make(map[string]any)
diff --git a/pkg/config/internal/internal_test.go b/pkg/config/internal/internal_test.go
index de43f4009e..cdde62762b 100644
--- a/pkg/config/internal/internal_test.go
+++ b/pkg/config/internal/internal_test.go
@@ -42,7 +42,7 @@ func TestFilteredYAML(t *testing.T) {
 		t.Run(test.name, func(t *testing.T) {
 			t.Parallel()
 
-			fy, err := internal.NewFilteredYAML(strings.NewReader(test.in), test.allowed)
+			fy, err := internal.NewFilteredYAML(io.NopCloser(strings.NewReader(test.in)), test.allowed)
 			if err != nil {
 				t.Fatal("unexpected error creating filtered yaml: err:", err)
 			}