From 7cba21d3ea83b20f359516b7dc2e91424c8f48da Mon Sep 17 00:00:00 2001
From: Rudo Thomas <rudo.thomas@gmail.com>
Date: Fri, 22 Nov 2024 11:03:48 +0100
Subject: [PATCH] fix: handle quotes in --tla-str and --ext-str in "tk eval"
 (#1237)

* fix: Handle quotes in --tla-str and --ext-str in "tk eval".

Problem: If a string with a quote was passes to tk eval's --tla-str or
--ext-str, things would break in an unexpected fashion.

Fix: Properly JSON-quote any given string.

* test: Add test for the --tla-* and --ext-* flags parsing.
---
 cmd/tk/flags.go      | 10 ++++++++--
 cmd/tk/flags_test.go | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 2 deletions(-)
 create mode 100644 cmd/tk/flags_test.go

diff --git a/cmd/tk/flags.go b/cmd/tk/flags.go
index 79e436274..61618b637 100644
--- a/cmd/tk/flags.go
+++ b/cmd/tk/flags.go
@@ -1,7 +1,7 @@
 package main
 
 import (
-	"fmt"
+	"encoding/json"
 	"strings"
 
 	"github.com/rs/zerolog/log"
@@ -95,7 +95,13 @@ func cliCodeParser(fs *pflag.FlagSet) (func() map[string]string, func() map[stri
 				if len(split) != 2 {
 					log.Fatal().Msgf(kind+"-str argument has wrong format: `%s`. Expected `key=<value>`", s)
 				}
-				m[split[0]] = fmt.Sprintf(`"%s"`, split[1])
+				// Properly quote the string; note that fmt.Sprintf("%q",...) could
+				// produce \U escapes which are not valid Jsonnet.
+				js, err := json.Marshal(split[1])
+				if err != nil {
+					log.Fatal().Msgf("impossible: failed to convert string to JSON: %s", err)
+				}
+				m[split[0]] = string(js)
 			}
 			return m
 		}
diff --git a/cmd/tk/flags_test.go b/cmd/tk/flags_test.go
new file mode 100644
index 000000000..434da9937
--- /dev/null
+++ b/cmd/tk/flags_test.go
@@ -0,0 +1,34 @@
+package main
+
+import (
+	"testing"
+
+	"github.com/spf13/pflag"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCliCodeParser(t *testing.T) {
+	fs := pflag.NewFlagSet("test-cli-code-parser", pflag.ContinueOnError)
+	parseExt, parseTLA := cliCodeParser(fs)
+	err := fs.Parse([]string{
+		"--ext-str", "es=1a \" \U0001f605 ' b\nc\u010f",
+		"--tla-str", "ts=2a \" \U0001f605 ' b\nc\u010f",
+		"--ext-code", "ec=1+2",
+		"--tla-code", "tc=2+3",
+		"-A", "ts2=ts2", // tla-str
+		"-V", "es2=es2", // ext-str
+	})
+	assert.NoError(t, err)
+	ext := parseExt()
+	assert.Equal(t, map[string]string{
+		"es":  `"1a \" ` + "\U0001f605" + ` ' b\nc` + "\u010f" + `"`,
+		"ec":  "1+2",
+		"es2": `"es2"`,
+	}, ext)
+	tla := parseTLA()
+	assert.Equal(t, map[string]string{
+		"ts":  `"2a \" ` + "\U0001f605" + ` ' b\nc` + "\u010f" + `"`,
+		"tc":  "2+3",
+		"ts2": `"ts2"`,
+	}, tla)
+}