From 6bc0b0f12420ad87398d3fbd398427d88b6b6e8c Mon Sep 17 00:00:00 2001
From: Mathieu Tortuyaux <mtortuyaux@microsoft.com>
Date: Thu, 22 Feb 2024 13:49:06 +0100
Subject: [PATCH] config: normalize keys of the config before any manipulation

Signed-off-by: Mathieu Tortuyaux <mtortuyaux@microsoft.com>
---
 config/config.go            | 36 ++++++++++++++++++++++++++++++++++++
 config/validate/validate.go |  5 ++++-
 2 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/config/config.go b/config/config.go
index f5a6605..bfdb37b 100644
--- a/config/config.go
+++ b/config/config.go
@@ -17,6 +17,7 @@ package config
 import (
 	"bufio"
 	"fmt"
+	"io"
 	"net/textproto"
 	"reflect"
 	"regexp"
@@ -70,10 +71,45 @@ func IsMultipartMime(userdata string) bool {
 	return strings.Contains(contentType, "multipart/mixed")
 }
 
+// Normalize transform the "-" in "_" for the keys
+// in the YAML configuration.
+// e.g reboot-strategy -> reboot_strategy
+func Normalize(r io.Reader) string {
+	var (
+		l string
+		w strings.Builder
+	)
+
+	sc := bufio.NewScanner(r)
+	for sc.Scan() {
+		l = sc.Text()
+		// sl = write_files:
+		sl := strings.Split(l, ":")
+		// sl = ["write_files", " "]
+		if len(sl) == 2 {
+			ssl := strings.Split(sl[0], "-")
+			// Handle case where it starts with -
+			// e.g:
+			// ---
+			// files:
+			//   - permissions: 0644
+			if len(ssl) == 2 && !strings.HasPrefix(ssl[1], " ") {
+				sl[0] = strings.Replace(sl[0], "-", "_", -1)
+			}
+			l = strings.Join(sl, ":")
+		}
+		w.WriteString(l + "\n")
+	}
+
+	return w.String()
+}
+
 // NewCloudConfig instantiates a new CloudConfig from the given contents (a
 // string of YAML), returning any error encountered. It will ignore unknown
 // fields but log encountering them.
 func NewCloudConfig(contents string) (*CloudConfig, error) {
+	contents = Normalize(strings.NewReader(contents))
+
 	var cfg CloudConfig
 	err := yaml.Unmarshal([]byte(contents), &cfg)
 	return &cfg, err
diff --git a/config/validate/validate.go b/config/validate/validate.go
index e9496ca..27d8a63 100644
--- a/config/validate/validate.go
+++ b/config/validate/validate.go
@@ -15,6 +15,7 @@
 package validate
 
 import (
+	"bytes"
 	"errors"
 	"fmt"
 	"regexp"
@@ -77,10 +78,12 @@ func validateCloudConfig(config []byte, rules []rule) (report Report, err error)
 // any parsing issues into the provided report. Unrecoverable errors are
 // returned as an error.
 func parseCloudConfig(cfg []byte, report *Report) (node, error) {
+	cfg = []byte(config.Normalize(bytes.NewReader(cfg)))
+
 	// unmarshal the config into an implicitly-typed form. The yaml library
 	// will implicitly convert types into their normalized form
 	// (e.g. 0744 -> 484, off -> false).
-	var weak map[interface{}]interface{}
+	var weak map[string]interface{}
 	if err := yaml.Unmarshal(cfg, &weak); err != nil {
 		matches := yamlLineError.FindStringSubmatch(err.Error())
 		if len(matches) == 3 {