From 3856c05f99260a778852dc112cdfc3ea3f0e1a9e Mon Sep 17 00:00:00 2001
From: Trevor Foster <trevor.foster@hotmail.ca>
Date: Sat, 9 May 2020 05:38:39 -0400
Subject: [PATCH] Fix: Getting the value of a StringToString pflag (#874)

* add parsing for stringToString flags

* add logic to return flags default if not val set, add a test

* extract parsing into single func

* add a few more cases

* return nil if unable to parse instead of panicing

* return map[string]interface in order to work with cast.ToStringMap

* mostly copy pflags implementation of the conversion to a stringtostring
---
 viper.go      | 28 ++++++++++++++++++++++++++++
 viper_test.go | 47 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 75 insertions(+)

diff --git a/viper.go b/viper.go
index 5226090..f61f4ed 100644
--- a/viper.go
+++ b/viper.go
@@ -1083,6 +1083,8 @@ func (v *Viper) find(lcaseKey string, flagDefault bool) interface{} {
 			s = strings.TrimSuffix(s, "]")
 			res, _ := readAsCSV(s)
 			return cast.ToIntSlice(res)
+		case "stringToString":
+			return stringToStringConv(flag.ValueString())
 		default:
 			return flag.ValueString()
 		}
@@ -1158,6 +1160,8 @@ func (v *Viper) find(lcaseKey string, flagDefault bool) interface{} {
 				s = strings.TrimSuffix(s, "]")
 				res, _ := readAsCSV(s)
 				return cast.ToIntSlice(res)
+			case "stringToString":
+				return stringToStringConv(flag.ValueString())
 			default:
 				return flag.ValueString()
 			}
@@ -1177,6 +1181,30 @@ func readAsCSV(val string) ([]string, error) {
 	return csvReader.Read()
 }
 
+// mostly copied from pflag's implementation of this operation here https://github.com/spf13/pflag/blob/master/string_to_string.go#L79
+// alterations are: errors are swallowed, map[string]interface{} is returned in order to enable cast.ToStringMap
+func stringToStringConv(val string) interface{} {
+	val = strings.Trim(val, "[]")
+	// An empty string would cause an empty map
+	if len(val) == 0 {
+		return map[string]interface{}{}
+	}
+	r := csv.NewReader(strings.NewReader(val))
+	ss, err := r.Read()
+	if err != nil {
+		return nil
+	}
+	out := make(map[string]interface{}, len(ss))
+	for _, pair := range ss {
+		kv := strings.SplitN(pair, "=", 2)
+		if len(kv) != 2 {
+			return nil
+		}
+		out[kv[0]] = kv[1]
+	}
+	return out
+}
+
 // IsSet checks to see if the key has been set in any of the data locations.
 // IsSet is case-insensitive for a key.
 func IsSet(key string) bool { return v.IsSet(key) }
diff --git a/viper_test.go b/viper_test.go
index b8ceccb..fe942de 100644
--- a/viper_test.go
+++ b/viper_test.go
@@ -970,6 +970,53 @@ func TestBindPFlag(t *testing.T) {
 	assert.Equal(t, "testing_mutate", Get("testvalue"))
 }
 
+func TestBindPFlagStringToString(t *testing.T) {
+	tests := []struct {
+		Expected map[string]string
+		Value    string
+	}{
+		{map[string]string{}, ""},
+		{map[string]string{"yo": "hi"}, "yo=hi"},
+		{map[string]string{"yo": "hi", "oh": "hi=there"}, "yo=hi,oh=hi=there"},
+		{map[string]string{"yo": ""}, "yo="},
+		{map[string]string{"yo": "", "oh": "hi=there"}, "yo=,oh=hi=there"},
+	}
+
+	v := New() // create independent Viper object
+	defaultVal := map[string]string{}
+	v.SetDefault("stringtostring", defaultVal)
+
+	for _, testValue := range tests {
+		flagSet := pflag.NewFlagSet("test", pflag.ContinueOnError)
+		flagSet.StringToString("stringtostring", testValue.Expected, "test")
+
+		for _, changed := range []bool{true, false} {
+			flagSet.VisitAll(func(f *pflag.Flag) {
+				f.Value.Set(testValue.Value)
+				f.Changed = changed
+			})
+
+			err := v.BindPFlags(flagSet)
+			if err != nil {
+				t.Fatalf("error binding flag set, %v", err)
+			}
+
+			type TestMap struct {
+				StringToString map[string]string
+			}
+			val := &TestMap{}
+			if err := v.Unmarshal(val); err != nil {
+				t.Fatalf("%+#v cannot unmarshal: %s", testValue.Value, err)
+			}
+			if changed {
+				assert.Equal(t, testValue.Expected, val.StringToString)
+			} else {
+				assert.Equal(t, defaultVal, val.StringToString)
+			}
+		}
+	}
+}
+
 func TestBoundCaseSensitivity(t *testing.T) {
 	assert.Equal(t, "brown", Get("eyes"))