Coerce starlark dict keys to strings (Velocidex#1156)

Co-authored-by: Mike Cohen <mike@velocidex.com>

vql/tools/fixtures/TestStarlark.golden

@@ -27,6 +27,19 @@
      1,
      2,
      3
+    ],
+    {
+     "A": 1,
+     "2": 3,
+     "1.100000": 1
+    },
+    null,
+    true,
+    1.1,
+    "string",
+    [
+     1,
+     2
     ]
    ]
   }

vql/tools/starlark.go

@@ -99,14 +99,36 @@ func starlarkValueAsInterface(value starlark.Value) (interface{}, error) {
 	case *starlark.Dict:
 		result := ordereddict.NewDict()
 		for _, item := range v.Items() {
-			key := item[0].String()
+			key := item[0]
 			value := item[1]
 
 			dictValueInterfaced, err := starlarkValueAsInterface(value)
 			if err != nil {
 				return nil, err
 			}
-			result.Set(key[1:len(key)-1], dictValueInterfaced)
+
+			dictKeyInterfaced, err := starlarkValueAsInterface(key)
+			if err != nil {
+				return nil, err
+			}
+
+			// JSON keys must be strings so we need to
+			// convert from the Starlark type to a string.
+			switch t := dictKeyInterfaced.(type) {
+			case string:
+				result.Set(t, dictValueInterfaced)
+
+			case fmt.Stringer:
+				result.Set(t.String(), dictValueInterfaced)
+
+			case float64:
+				result.Set(fmt.Sprintf("%f", t),
+					dictValueInterfaced)
+
+			case int64:
+				result.Set(fmt.Sprintf("%d", t),
+					dictValueInterfaced)
+			}
 		}
 
 		return result, nil

vql/tools/starlark_test.go

@@ -47,7 +47,13 @@ LET X = starl(code=StarCode)
 SELECT X.Foo(X=2, Y="String", Z=[2, 3]) FROM scope()
 `, `
 def Foo(X, Y, Z):
-   return X + 2, Y + "A", [1,] + Z
+   return (X + 2, Y + "A", [1,] + Z,
+
+           # Dict with weird keys should be converted to string.
+           {"A":1, 2:3, 1.1: 1},
+
+           # Other starlark types.
+           None, True, 1.1, "string", [1, 2])
 `},
 }