From de8eef1da506dd8c7d3278c5c64b03cab72dc47c Mon Sep 17 00:00:00 2001
From: Martin Atkins <mart@degeneration.co.uk>
Date: Wed, 15 Jun 2022 17:37:00 -0700
Subject: [PATCH] addrs: Format string instance keys in an HCL-compatible way

So far we've only ever needed to re-parse address strings that happen not
to contain instance keys and so we've gotten away with our serialization
of these not being quite right, but given how liberally we've expected to
be able to use address strings from this package for wire format
interchange it seems likely that this is going to surprise us eventually.

Now we'll use an escaping scheme compatible with HCL's parser rather
than Go's parser, and so we can safely rely on hclsyntax.ParseTraversal
as part of reversing this operation to transform an address string back
into an address equivalent to the value it was created from.
---
 internal/addrs/instance_key.go      | 62 ++++++++++++++++++++++--
 internal/addrs/instance_key_test.go | 75 +++++++++++++++++++++++++++++
 2 files changed, 134 insertions(+), 3 deletions(-)
 create mode 100644 internal/addrs/instance_key_test.go

diff --git a/internal/addrs/instance_key.go b/internal/addrs/instance_key.go
index ff128be5b4..2d46bfcbe0 100644
--- a/internal/addrs/instance_key.go
+++ b/internal/addrs/instance_key.go
@@ -2,6 +2,8 @@ package addrs
 
 import (
 	"fmt"
+	"strings"
+	"unicode"
 
 	"github.com/zclconf/go-cty/cty"
 	"github.com/zclconf/go-cty/cty/gocty"
@@ -72,9 +74,10 @@ func (k StringKey) instanceKeySigil() {
 }
 
 func (k StringKey) String() string {
-	// FIXME: This isn't _quite_ right because Go's quoted string syntax is
-	// slightly different than HCL's, but we'll accept it for now.
-	return fmt.Sprintf("[%q]", string(k))
+	// We use HCL's quoting syntax here so that we can in principle parse
+	// an address constructed by this package as if it were an HCL
+	// traversal, even if the string contains HCL's own metacharacters.
+	return fmt.Sprintf("[%s]", toHCLQuotedString(string(k)))
 }
 
 func (k StringKey) Value() cty.Value {
@@ -133,3 +136,56 @@ const (
 	IntKeyType    InstanceKeyType = 'I'
 	StringKeyType InstanceKeyType = 'S'
 )
+
+// toHCLQuotedString is a helper which formats the given string in a way that
+// HCL's expression parser would treat as a quoted string template.
+//
+// This includes:
+//   - Adding quote marks at the start and the end.
+//   - Using backslash escapes as needed for characters that cannot be represented directly.
+//   - Escaping anything that would be treated as a template interpolation or control sequence.
+func toHCLQuotedString(s string) string {
+	// This is an adaptation of a similar function inside the hclwrite package,
+	// inlined here because hclwrite's version generates HCL tokens but we
+	// only need normal strings.
+	if len(s) == 0 {
+		return `""`
+	}
+	var buf strings.Builder
+	buf.WriteByte('"')
+	for i, r := range s {
+		switch r {
+		case '\n':
+			buf.WriteString(`\n`)
+		case '\r':
+			buf.WriteString(`\r`)
+		case '\t':
+			buf.WriteString(`\t`)
+		case '"':
+			buf.WriteString(`\"`)
+		case '\\':
+			buf.WriteString(`\\`)
+		case '$', '%':
+			buf.WriteRune(r)
+			remain := s[i+1:]
+			if len(remain) > 0 && remain[0] == '{' {
+				// Double up our template introducer symbol to escape it.
+				buf.WriteRune(r)
+			}
+		default:
+			if !unicode.IsPrint(r) {
+				var fmted string
+				if r < 65536 {
+					fmted = fmt.Sprintf("\\u%04x", r)
+				} else {
+					fmted = fmt.Sprintf("\\U%08x", r)
+				}
+				buf.WriteString(fmted)
+			} else {
+				buf.WriteRune(r)
+			}
+		}
+	}
+	buf.WriteByte('"')
+	return buf.String()
+}
diff --git a/internal/addrs/instance_key_test.go b/internal/addrs/instance_key_test.go
new file mode 100644
index 0000000000..0d12888bf0
--- /dev/null
+++ b/internal/addrs/instance_key_test.go
@@ -0,0 +1,75 @@
+package addrs
+
+import (
+	"fmt"
+	"testing"
+)
+
+func TestInstanceKeyString(t *testing.T) {
+	tests := []struct {
+		Key  InstanceKey
+		Want string
+	}{
+		{
+			IntKey(0),
+			`[0]`,
+		},
+		{
+			IntKey(5),
+			`[5]`,
+		},
+		{
+			StringKey(""),
+			`[""]`,
+		},
+		{
+			StringKey("hi"),
+			`["hi"]`,
+		},
+		{
+			StringKey("0"),
+			`["0"]`, // intentionally distinct from IntKey(0)
+		},
+		{
+			// Quotes must be escaped
+			StringKey(`"`),
+			`["\""]`,
+		},
+		{
+			// Escape sequences must themselves be escaped
+			StringKey(`\r\n`),
+			`["\\r\\n"]`,
+		},
+		{
+			// Template interpolation sequences "${" must be escaped.
+			StringKey(`${hello}`),
+			`["$${hello}"]`,
+		},
+		{
+			// Template control sequences "%{" must be escaped.
+			StringKey(`%{ for something in something }%{ endfor }`),
+			`["%%{ for something in something }%%{ endfor }"]`,
+		},
+		{
+			// Dollar signs that aren't followed by { are not interpolation sequences
+			StringKey(`$hello`),
+			`["$hello"]`,
+		},
+		{
+			// Percent signs that aren't followed by { are not control sequences
+			StringKey(`%hello`),
+			`["%hello"]`,
+		},
+	}
+
+	for _, test := range tests {
+		testName := fmt.Sprintf("%#v", test.Key)
+		t.Run(testName, func(t *testing.T) {
+			got := test.Key.String()
+			want := test.Want
+			if got != want {
+				t.Errorf("wrong result\nreciever: %s\ngot:      %s\nwant:     %s", testName, got, want)
+			}
+		})
+	}
+}