From 10c6d2af7146b4c64758faed65452fff4411e621 Mon Sep 17 00:00:00 2001
From: Anmol Sethi <hi@nhooyr.io>
Date: Tue, 6 Jun 2023 14:01:32 -0700
Subject: [PATCH] d2parser: Fix bad import key panic

---
 d2ir/compile_test.go                          |  8 ++++++-
 d2ir/import_test.go                           |  9 ++++++++
 d2parser/parse.go                             |  3 +++
 d2parser/parse_test.go                        | 12 ++++++++++
 .../d2parser/TestParse/import/#06.exp.json    | 22 +++++++++++++++++++
 5 files changed, 53 insertions(+), 1 deletion(-)
 create mode 100644 testdata/d2parser/TestParse/import/#06.exp.json

diff --git a/d2ir/compile_test.go b/d2ir/compile_test.go
index 2435d3d5f..516c42be4 100644
--- a/d2ir/compile_test.go
+++ b/d2ir/compile_test.go
@@ -54,10 +54,16 @@ func compileFS(t testing.TB, path string, mfs map[string]string) (*d2ir.Map, err
 	t.Helper()
 
 	ast, err := d2parser.Parse(path, strings.NewReader(mfs[path]), nil)
-	assert.Success(t, err)
+	if err != nil {
+		return nil, err
+	}
 
 	fs, err := mapfs.New(mfs)
 	assert.Success(t, err)
+	t.Cleanup(func() {
+		err = fs.Close()
+		assert.Success(t, err)
+	})
 	m, err := d2ir.Compile(ast, &d2ir.CompileOptions{
 		FS: fs,
 	})
diff --git a/d2ir/import_test.go b/d2ir/import_test.go
index 1fc5caffb..d6eff4809 100644
--- a/d2ir/import_test.go
+++ b/d2ir/import_test.go
@@ -98,6 +98,15 @@ label: meow`,
 					assert.ErrorString(t, err, `index.d2:1:1: failed to import "x.d2": open x.d2: no such file or directory`)
 				},
 			},
+			{
+				name: "escape",
+				run: func(t testing.TB) {
+					_, err := compileFS(t, "index.d2", map[string]string{
+						"index.d2": "...@'./../x.d2'",
+					})
+					assert.ErrorString(t, err, `index.d2:1:1: failed to import "../x.d2": stat ../x.d2: invalid argument`)
+				},
+			},
 			{
 				name: "absolute",
 				run: func(t testing.TB) {
diff --git a/d2parser/parse.go b/d2parser/parse.go
index efa1e2e06..18c4187d2 100644
--- a/d2parser/parse.go
+++ b/d2parser/parse.go
@@ -1710,6 +1710,9 @@ func (p *parser) parseImport(spread bool) *d2ast.Import {
 	}
 
 	k := p.parseKey()
+	if k == nil {
+		return imp
+	}
 	if k.Path[0].UnquotedString != nil && len(k.Path) > 1 && k.Path[1].UnquotedString != nil && k.Path[1].Unbox().ScalarString() == "d2" {
 		k.Path = append(k.Path[:1], k.Path[2:]...)
 	}
diff --git a/d2parser/parse_test.go b/d2parser/parse_test.go
index e9b1ea146..4591c729b 100644
--- a/d2parser/parse_test.go
+++ b/d2parser/parse_test.go
@@ -451,6 +451,18 @@ func testImport(t *testing.T) {
 				assert.Equal(t, "d2", ast.Nodes[0].Import.Path[1].Unbox().ScalarString())
 			},
 		},
+		{
+			text: "...@../file",
+			assert: func(t testing.TB, ast *d2ast.Map, err error) {
+				assert.ErrorString(t, err, "d2/testdata/d2parser/TestParse/import/#06.d2:1:5: unexpected text after import")
+			},
+		},
+		{
+			text: "@file",
+			assert: func(t testing.TB, ast *d2ast.Map, err error) {
+				assert.ErrorString(t, err, "d2/testdata/d2parser/TestParse/import/#06.d2:1:5: unexpected text after import")
+			},
+		},
 	}
 
 	runa(t, tca)
diff --git a/testdata/d2parser/TestParse/import/#06.exp.json b/testdata/d2parser/TestParse/import/#06.exp.json
new file mode 100644
index 000000000..8d0a4fa93
--- /dev/null
+++ b/testdata/d2parser/TestParse/import/#06.exp.json
@@ -0,0 +1,22 @@
+{
+  "ast": {
+    "range": "d2/testdata/d2parser/TestParse/import/#06.d2,0:0:0-0:11:11",
+    "nodes": [
+      {
+        "import": {
+          "range": "d2/testdata/d2parser/TestParse/import/#06.d2,0:0:0-0:4:4",
+          "spread": true,
+          "path": null
+        }
+      }
+    ]
+  },
+  "err": {
+    "errs": [
+      {
+        "range": "d2/testdata/d2parser/TestParse/import/#06.d2,0:4:4-0:11:11",
+        "errmsg": "d2/testdata/d2parser/TestParse/import/#06.d2:1:5: unexpected text after import"
+      }
+    ]
+  }
+}