From c2f7e0d82cdc134f584d6437e85dee51b3f5bc6e Mon Sep 17 00:00:00 2001
From: Vit Zikmund <vit.zikmund@themama.ai>
Date: Mon, 6 Feb 2023 12:36:12 +0000
Subject: [PATCH] Use a lazy-quoting @sh encoder

---
 pkg/yqlib/doc/operators/encode-decode.md   |  2 +-
 pkg/yqlib/encoder_sh.go                    | 42 ++++++++++++++++++----
 pkg/yqlib/operator_encoder_decoder_test.go |  2 +-
 3 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/pkg/yqlib/doc/operators/encode-decode.md b/pkg/yqlib/doc/operators/encode-decode.md
index 93b7abb4..b8cc1f75 100644
--- a/pkg/yqlib/doc/operators/encode-decode.md
+++ b/pkg/yqlib/doc/operators/encode-decode.md
@@ -478,7 +478,7 @@ yq '.coolData | @sh' sample.yml
 ```
 will output
 ```yaml
-'strings with spaces and a \'quote\''
+strings' with spaces and a '\'quote\'
 ```
 
 ## Decode a base64 encoded string
diff --git a/pkg/yqlib/encoder_sh.go b/pkg/yqlib/encoder_sh.go
index 53663158..a8272c90 100644
--- a/pkg/yqlib/encoder_sh.go
+++ b/pkg/yqlib/encoder_sh.go
@@ -9,7 +9,7 @@ import (
 	yaml "gopkg.in/yaml.v3"
 )
 
-var pattern = regexp.MustCompile(`[^\w@%+=:,./-]`)
+var unsafeChars = regexp.MustCompile(`[^\w@%+=:,./-]`)
 
 type shEncoder struct {
 }
@@ -36,9 +36,39 @@ func (e *shEncoder) Encode(writer io.Writer, originalNode *yaml.Node) error {
 		return fmt.Errorf("cannot encode %v as URI, can only operate on strings. Please first pipe through another encoding operator to convert the value to a string", node.Tag)
 	}
 
-	value := originalNode.Value
-	if pattern.MatchString(value) {
-		value = "'" + strings.ReplaceAll(value, "'", "\\'") + "'"
-	}
-	return writeString(writer, value)
+	return writeString(writer, e.encode(originalNode.Value))
+}
+
+// put any shell-unsafe characters into a single-quoted block, close the block lazily
+func (e *shEncoder) encode(input string) string {
+	const quote = '\''
+	var inQuoteBlock = false
+	var encoded strings.Builder
+	encoded.Grow(len(input))
+
+	for _, ir := range input {
+		// open or close a single-quote block
+		if ir == quote {
+			if inQuoteBlock {
+				// get out of a quote block for an input quote
+				encoded.WriteRune(quote)
+				inQuoteBlock = !inQuoteBlock
+			}
+			// escape the quote with a backslash
+			encoded.WriteRune('\\')
+		} else {
+			if unsafeChars.MatchString(string(ir)) && !inQuoteBlock {
+				// start a quote block on an unsafe character
+				encoded.WriteRune(quote)
+				inQuoteBlock = !inQuoteBlock
+			}
+		}
+		// pass on the input character
+		encoded.WriteRune(ir)
+	}
+	// close any pending quote block
+	if inQuoteBlock {
+		encoded.WriteRune(quote)
+	}
+	return encoded.String()
 }
diff --git a/pkg/yqlib/operator_encoder_decoder_test.go b/pkg/yqlib/operator_encoder_decoder_test.go
index 7801e6cf..6db682ec 100644
--- a/pkg/yqlib/operator_encoder_decoder_test.go
+++ b/pkg/yqlib/operator_encoder_decoder_test.go
@@ -263,7 +263,7 @@ var encoderDecoderOperatorScenarios = []expressionScenario{
 		document:       "coolData: strings with spaces and a 'quote'",
 		expression:     ".coolData | @sh",
 		expected: []string{
-			"D0, P[coolData], (!!str)::'strings with spaces and a \\'quote\\''\n",
+			"D0, P[coolData], (!!str)::strings' with spaces and a '\\'quote\\'\n",
 		},
 	},
 	{