From 9d960d3fdceba7a28aa4bbec53177c20c308b10e Mon Sep 17 00:00:00 2001
From: maxime mouial <hush-hush@users.noreply.github.com>
Date: Thu, 28 Nov 2024 20:07:45 +0100
Subject: [PATCH] Implement AllKeyLowercase to be compatible with Viper
 (#31480)

---
 pkg/config/nodetreemodel/config.go  | 31 +++++++++++++++++++++++------
 pkg/config/nodetreemodel/helpers.go | 11 ++++++++++
 2 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/pkg/config/nodetreemodel/config.go b/pkg/config/nodetreemodel/config.go
index 57d8ee8244c48..9e293bb9a7592 100644
--- a/pkg/config/nodetreemodel/config.go
+++ b/pkg/config/nodetreemodel/config.go
@@ -19,7 +19,6 @@ import (
 
 	"github.com/DataDog/viper"
 	"go.uber.org/atomic"
-	"golang.org/x/exp/maps"
 	"golang.org/x/exp/slices"
 
 	"github.com/DataDog/datadog-agent/pkg/config/model"
@@ -100,6 +99,8 @@ type ntmConfig struct {
 	// keys that have been used but are unknown
 	// used to warn (a single time) on use
 	unknownKeys map[string]struct{}
+	// allSettings contains all settings that we have a value for in the default tree
+	allSettings []string
 
 	// extraConfigFilePaths represents additional configuration file paths that will be merged into the main configuration when ReadInConfig() is called.
 	extraConfigFilePaths []string
@@ -238,9 +239,9 @@ func (c *ntmConfig) SetKnown(key string) {
 	if c.isReady() {
 		panic("cannot SetKnown() once the config has been marked as ready for use")
 	}
+
 	key = strings.ToLower(key)
 	c.knownKeys[key] = struct{}{}
-	c.setDefault(key, nil)
 }
 
 // IsKnown returns whether a key is known
@@ -296,6 +297,24 @@ func (c *ntmConfig) mergeAllLayers() error {
 	return nil
 }
 
+func computeAllSettings(node InnerNode, path string) []string {
+	knownKeys := []string{}
+	for _, name := range node.ChildrenKeys() {
+		newPath := joinKey(path, name)
+
+		child, _ := node.GetChild(name)
+		if _, ok := child.(LeafNode); ok {
+			knownKeys = append(knownKeys, newPath)
+		} else if inner, ok := child.(InnerNode); ok {
+			knownKeys = append(knownKeys, computeAllSettings(inner, newPath)...)
+		} else {
+			log.Errorf("unknown node type in the tree: %T", child)
+		}
+	}
+	slices.Sort(knownKeys)
+	return knownKeys
+}
+
 // BuildSchema is called when Setup is complete, and the config is ready to be used
 func (c *ntmConfig) BuildSchema() {
 	c.Lock()
@@ -305,6 +324,7 @@ func (c *ntmConfig) BuildSchema() {
 	if err := c.mergeAllLayers(); err != nil {
 		c.warnings = append(c.warnings, err.Error())
 	}
+	c.allSettings = computeAllSettings(c.defaults, "")
 }
 
 // Stringify stringifies the config, but only with the test build tag
@@ -392,14 +412,12 @@ func (c *ntmConfig) IsSet(key string) bool {
 	return c.IsKnown(key)
 }
 
-// AllKeysLowercased returns all keys lower-cased
+// AllKeysLowercased returns all keys lower-cased from the default tree, but not keys that are merely marked as known
 func (c *ntmConfig) AllKeysLowercased() []string {
 	c.RLock()
 	defer c.RUnlock()
 
-	res := maps.Keys(c.knownKeys)
-	slices.Sort(res)
-	return res
+	return slices.Clone(c.allSettings)
 }
 
 func (c *ntmConfig) leafAtPath(key string) LeafNode {
@@ -686,6 +704,7 @@ func NewConfig(name string, envPrefix string, envKeyReplacer *strings.Replacer)
 		ready:              atomic.NewBool(false),
 		configEnvVars:      map[string]string{},
 		knownKeys:          map[string]struct{}{},
+		allSettings:        []string{},
 		unknownKeys:        map[string]struct{}{},
 		defaults:           newInnerNode(nil),
 		file:               newInnerNode(nil),
diff --git a/pkg/config/nodetreemodel/helpers.go b/pkg/config/nodetreemodel/helpers.go
index 33f96491b583a..9cbb3beb8e2d1 100644
--- a/pkg/config/nodetreemodel/helpers.go
+++ b/pkg/config/nodetreemodel/helpers.go
@@ -16,6 +16,17 @@ func splitKey(key string) []string {
 	return strings.Split(strings.ToLower(key), ".")
 }
 
+func joinKey(parts ...string) string {
+	nonEmptyParts := make([]string, 0, len(parts))
+	for idx := range parts {
+		if parts[idx] == "" {
+			continue
+		}
+		nonEmptyParts = append(nonEmptyParts, parts[idx])
+	}
+	return strings.Join(nonEmptyParts, ".")
+}
+
 func safeMul(a, b uint) uint {
 	c := a * b
 	// detect multiplication overflows