From e998a95b96b9f04d803f6cd5a202985ab70c0215 Mon Sep 17 00:00:00 2001
From: Alexander Wang <alex@terrastruct.com>
Date: Wed, 31 May 2023 23:04:09 -0700
Subject: [PATCH] d2ir: prevent illegal non-tail keywords

---
 d2compiler/compile_test.go                           |  7 +++++++
 d2graph/d2graph.go                                   |  3 ++-
 d2ir/d2ir.go                                         |  7 +++----
 .../TestCompile/keyword-container.exp.json           | 12 ++++++++++++
 4 files changed, 24 insertions(+), 5 deletions(-)
 create mode 100644 testdata/d2compiler/TestCompile/keyword-container.exp.json

diff --git a/d2compiler/compile_test.go b/d2compiler/compile_test.go
index 657d09575..cbe248773 100644
--- a/d2compiler/compile_test.go
+++ b/d2compiler/compile_test.go
@@ -1722,6 +1722,13 @@ y -> x.style
 `,
 			expErr: `d2/testdata/d2compiler/TestCompile/edge_to_style.d2:2:8: reserved keywords are prohibited in edges`,
 		},
+		{
+			name: "keyword-container",
+
+			text: `a.near.b
+`,
+			expErr: `d2/testdata/d2compiler/TestCompile/keyword-container.d2:1:3: "near" must be the last part of the key`,
+		},
 		{
 			name: "escaped_id",
 
diff --git a/d2graph/d2graph.go b/d2graph/d2graph.go
index 301cc5bd8..6a0bfd51b 100644
--- a/d2graph/d2graph.go
+++ b/d2graph/d2graph.go
@@ -1645,11 +1645,12 @@ var SimpleReservedKeywords = map[string]struct{}{
 	"classes":        {},
 }
 
-// ReservedKeywordHolders are reserved keywords that are meaningless on its own and exist solely to hold a set of reserved keywords
+// ReservedKeywordHolders are reserved keywords that can hold composites
 var ReservedKeywordHolders = map[string]struct{}{
 	"style":            {},
 	"source-arrowhead": {},
 	"target-arrowhead": {},
+	"classes":          {},
 }
 
 // StyleKeywords are reserved keywords which cannot exist outside of the "style" keyword
diff --git a/d2ir/d2ir.go b/d2ir/d2ir.go
index 50fb8c104..6f0ad2c67 100644
--- a/d2ir/d2ir.go
+++ b/d2ir/d2ir.go
@@ -669,10 +669,9 @@ func (m *Map) ensureField(i int, kp *d2ast.KeyPath, refctx *RefContext) (*Field,
 
 	if _, ok := d2graph.ReservedKeywords[strings.ToLower(head)]; ok {
 		head = strings.ToLower(head)
-	}
-
-	if head == "class" && i < len(kp.Path)-1 {
-		return nil, d2parser.Errorf(kp.Path[i].Unbox(), `"class" must be the last part of the key`)
+		if _, ok := d2graph.ReservedKeywordHolders[head]; !ok && i < len(kp.Path)-1 {
+			return nil, d2parser.Errorf(kp.Path[i].Unbox(), fmt.Sprintf(`"%s" must be the last part of the key`, head))
+		}
 	}
 
 	if head == "_" {
diff --git a/testdata/d2compiler/TestCompile/keyword-container.exp.json b/testdata/d2compiler/TestCompile/keyword-container.exp.json
new file mode 100644
index 000000000..898e9071a
--- /dev/null
+++ b/testdata/d2compiler/TestCompile/keyword-container.exp.json
@@ -0,0 +1,12 @@
+{
+  "graph": null,
+  "err": {
+    "ioerr": null,
+    "errs": [
+      {
+        "range": "d2/testdata/d2compiler/TestCompile/keyword-container.d2,0:2:2-0:6:6",
+        "errmsg": "d2/testdata/d2compiler/TestCompile/keyword-container.d2:1:3: \"near\" must be the last part of the key"
+      }
+    ]
+  }
+}