From b6e71ef9909d13163b2b404968670633a560a123 Mon Sep 17 00:00:00 2001
From: Terry Cain <terry@terrys-home.co.uk>
Date: Fri, 1 Apr 2022 23:02:06 +0100
Subject: [PATCH] Added support for repeated extra arguments

Problem:
Specifying extra arguments for the API server for example is not supported as
the arguments get stored in a map before being passed to the API server.

Solution:
Updated the GetArgs function to store the arguments in a map that can have
multiple values. Some more logic is added so that repeated extra arguments
retain their order when sorted whilst overall the arguments can still be
sorted for improved readability when logged.

Support has been added for prefixing and suffixing default argument values
by using -= and += when specifying extra arguments.

Signed-off-by: Terry Cain <terry@terrys-home.co.uk>
---
 pkg/daemons/config/types.go      | 80 ++++++++++++++++++++++++++------
 pkg/daemons/config/types_test.go | 61 ++++++++++++++++++++++++
 2 files changed, 128 insertions(+), 13 deletions(-)

diff --git a/pkg/daemons/config/types.go b/pkg/daemons/config/types.go
index 0ee3861b22..f40648bd91 100644
--- a/pkg/daemons/config/types.go
+++ b/pkg/daemons/config/types.go
@@ -283,24 +283,78 @@ func (a ArgString) String() string {
 	return b.String()
 }
 
-// GetArgs appends extra arguments to existing arguments overriding any default options.
-func GetArgs(argsMap map[string]string, extraArgs []string) []string {
+// GetArgs appends extra arguments to existing arguments with logic to override any default
+// arguments whilst also allowing to prefix and suffix default string slice arguments.
+func GetArgs(initialArgs map[string]string, extraArgs []string) []string {
 	const hyphens = "--"
 
-	// add extra args to args map to override any default option
-	for _, arg := range extraArgs {
-		splitArg := strings.SplitN(strings.TrimPrefix(arg, hyphens), "=", 2)
-		if len(splitArg) < 2 {
-			argsMap[splitArg[0]] = "true"
-			continue
+	multiArgs := make(map[string][]string)
+
+	for _, unsplitArg := range extraArgs {
+		splitArg := strings.SplitN(strings.TrimPrefix(unsplitArg, hyphens), "=", 2)
+		arg := splitArg[0]
+		value := "true"
+		if len(splitArg) > 1 {
+			value = splitArg[1]
 		}
-		argsMap[splitArg[0]] = splitArg[1]
+
+		// After the first iteration, initial args will be empty when handling
+		// duplicate arguments as they will form part of existingValues
+		cleanedArg := strings.TrimRight(arg, "-+")
+		initialValue, initialValueExists := initialArgs[cleanedArg]
+		existingValues, existingValuesFound := multiArgs[cleanedArg]
+
+		newValues := make([]string, 0)
+		if strings.HasSuffix(arg, "+") { // Append value to initial args
+			if initialValueExists {
+				newValues = append(newValues, initialValue)
+			}
+			if existingValuesFound {
+				newValues = append(newValues, existingValues...)
+			}
+			newValues = append(newValues, value)
+
+		} else if strings.HasSuffix(arg, "-") { // Prepend value to initial args
+			newValues = append(newValues, value)
+			if initialValueExists {
+				newValues = append(newValues, initialValue)
+			}
+			if existingValuesFound {
+				newValues = append(newValues, existingValues...)
+			}
+		} else { // Append value ignoring initial args
+			if existingValuesFound {
+				newValues = append(newValues, existingValues...)
+			}
+			newValues = append(newValues, value)
+		}
+
+		delete(initialArgs, cleanedArg)
+		multiArgs[cleanedArg] = newValues
+
 	}
+
+	// Add any remaining initial args to the map
+	for arg, value := range initialArgs {
+		multiArgs[arg] = []string{value}
+	}
+
+	// Get args so we can output them sorted whilst preserving the order of
+	// repeated keys
+	var keys []string
+	for arg := range multiArgs {
+		keys = append(keys, arg)
+	}
+	sort.Strings(keys)
+
 	var args []string
-	for arg, value := range argsMap {
-		cmd := fmt.Sprintf("%s%s=%s", hyphens, strings.TrimPrefix(arg, hyphens), value)
-		args = append(args, cmd)
+	for _, arg := range keys {
+		values := multiArgs[arg]
+		for _, value := range values {
+			cmd := fmt.Sprintf("%s%s=%s", hyphens, strings.TrimPrefix(arg, hyphens), value)
+			args = append(args, cmd)
+		}
 	}
-	sort.Strings(args)
+
 	return args
 }
diff --git a/pkg/daemons/config/types_test.go b/pkg/daemons/config/types_test.go
index f94d6093b7..9744abe2b7 100644
--- a/pkg/daemons/config/types_test.go
+++ b/pkg/daemons/config/types_test.go
@@ -79,6 +79,67 @@ func Test_UnitGetArgs(t *testing.T) {
 				"--iii=II",
 			},
 		},
+		{
+			name: "Multiple args with defaults Test",
+			args: args{
+				argsMap: map[string]string{
+					"aaa": "A",
+					"bbb": "B",
+				},
+				extraArgs: []string{
+					"--ccc=C",
+					"--bbb=DD",
+					"--bbb=AA",
+				},
+			},
+
+			want: []string{
+				"--aaa=A",
+				"--bbb=DD",
+				"--bbb=AA",
+				"--ccc=C",
+			},
+		},
+		{
+			name: "Multiple args with defaults and prefix",
+			args: args{
+				argsMap: map[string]string{
+					"aaa": "A",
+					"bbb": "B",
+				},
+				extraArgs: []string{
+					"--ccc=C",
+					"--bbb-=DD",
+				},
+			},
+
+			want: []string{
+				"--aaa=A",
+				"--bbb=DD",
+				"--bbb=B",
+				"--ccc=C",
+			},
+		},
+		{
+			name: "Multiple args with defaults and suffix",
+			args: args{
+				argsMap: map[string]string{
+					"aaa": "A",
+					"bbb": "B",
+				},
+				extraArgs: []string{
+					"--ccc=C",
+					"--bbb+=DD",
+				},
+			},
+
+			want: []string{
+				"--aaa=A",
+				"--bbb=B",
+				"--bbb=DD",
+				"--ccc=C",
+			},
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {