From b4382a679338712846c5fb04daeeab5ff5162631 Mon Sep 17 00:00:00 2001
From: Alexander Wang <alex@terrastruct.com>
Date: Mon, 18 Nov 2024 08:00:57 -0800
Subject: [PATCH] markdown: sanitize links

---
 ci/release/changelogs/next.md                 |   1 +
 d2compiler/compile_test.go                    |   7 ++
 lib/textmeasure/links.go                      |  26 ++++
 lib/textmeasure/markdown.go                   |   6 +-
 .../TestCompile/markdown_ampersand.exp.json   | 114 ++++++++++++++++++
 5 files changed, 153 insertions(+), 1 deletion(-)
 create mode 100644 lib/textmeasure/links.go
 create mode 100644 testdata/d2compiler/TestCompile/markdown_ampersand.exp.json

diff --git a/ci/release/changelogs/next.md b/ci/release/changelogs/next.md
index d4b30fab63..56f152973d 100644
--- a/ci/release/changelogs/next.md
+++ b/ci/release/changelogs/next.md
@@ -9,3 +9,4 @@
 #### Bugfixes ⛑️
 
 - Imports: fixes using substitutions in `icon` values [#2207](https://github.com/terrastruct/d2/pull/2207)
+- Markdown: fixes ampersands in URLs in markdown [#2219](https://github.com/terrastruct/d2/pull/2219)
diff --git a/d2compiler/compile_test.go b/d2compiler/compile_test.go
index 6870938d01..3c7c024ebd 100644
--- a/d2compiler/compile_test.go
+++ b/d2compiler/compile_test.go
@@ -935,6 +935,13 @@ b.(x -> y)[0]: two
 				}
 			},
 		},
+		{
+			name: "markdown_ampersand",
+			text: `memo: |md
+  <a href="https://www.google.com/search?q=d2&newwindow=1&amp;bar">d2</a>
+|
+`,
+		},
 		{
 			name: "unsemantic_markdown",
 
diff --git a/lib/textmeasure/links.go b/lib/textmeasure/links.go
new file mode 100644
index 0000000000..d1d07a652c
--- /dev/null
+++ b/lib/textmeasure/links.go
@@ -0,0 +1,26 @@
+package textmeasure
+
+import (
+	"fmt"
+	"regexp"
+	"strings"
+)
+
+func sanitizeLinks(input string) (string, error) {
+	re := regexp.MustCompile(`href="([^"]*)"`)
+
+	return re.ReplaceAllStringFunc(input, func(href string) string {
+		matches := re.FindStringSubmatch(href)
+		if len(matches) < 2 {
+			return href
+		}
+
+		value := matches[1]
+
+		value = strings.ReplaceAll(value, "&amp;", "TEMP_AMP")
+		value = strings.ReplaceAll(value, "&", "&amp;")
+		value = strings.ReplaceAll(value, "TEMP_AMP", "&amp;")
+
+		return fmt.Sprintf(`href="%s"`, value)
+	}), nil
+}
diff --git a/lib/textmeasure/markdown.go b/lib/textmeasure/markdown.go
index d2903343c0..41e2ea1351 100644
--- a/lib/textmeasure/markdown.go
+++ b/lib/textmeasure/markdown.go
@@ -83,7 +83,11 @@ func RenderMarkdown(m string) (string, error) {
 	if err := markdownRenderer.Convert([]byte(m), &output); err != nil {
 		return "", err
 	}
-	return output.String(), nil
+	sanitized, err := sanitizeLinks(output.String())
+	if err != nil {
+		return "", err
+	}
+	return sanitized, nil
 }
 
 func init() {
diff --git a/testdata/d2compiler/TestCompile/markdown_ampersand.exp.json b/testdata/d2compiler/TestCompile/markdown_ampersand.exp.json
new file mode 100644
index 0000000000..981b7f2479
--- /dev/null
+++ b/testdata/d2compiler/TestCompile/markdown_ampersand.exp.json
@@ -0,0 +1,114 @@
+{
+  "graph": {
+    "name": "",
+    "isFolderOnly": false,
+    "ast": {
+      "range": "d2/testdata/d2compiler/TestCompile/markdown_ampersand.d2,0:0:0-3:0:86",
+      "nodes": [
+        {
+          "map_key": {
+            "range": "d2/testdata/d2compiler/TestCompile/markdown_ampersand.d2,0:0:0-2:1:85",
+            "key": {
+              "range": "d2/testdata/d2compiler/TestCompile/markdown_ampersand.d2,0:0:0-0:4:4",
+              "path": [
+                {
+                  "unquoted_string": {
+                    "range": "d2/testdata/d2compiler/TestCompile/markdown_ampersand.d2,0:0:0-0:4:4",
+                    "value": [
+                      {
+                        "string": "memo",
+                        "raw_string": "memo"
+                      }
+                    ]
+                  }
+                }
+              ]
+            },
+            "primary": {},
+            "value": {
+              "block_string": {
+                "range": "d2/testdata/d2compiler/TestCompile/markdown_ampersand.d2,0:6:6-2:1:85",
+                "quote": "",
+                "tag": "md",
+                "value": "<a href=\"https://www.google.com/search?q=d2&newwindow=1&amp;bar\">d2</a>"
+              }
+            }
+          }
+        }
+      ]
+    },
+    "root": {
+      "id": "",
+      "id_val": "",
+      "attributes": {
+        "label": {
+          "value": ""
+        },
+        "labelDimensions": {
+          "width": 0,
+          "height": 0
+        },
+        "style": {},
+        "near_key": null,
+        "shape": {
+          "value": ""
+        },
+        "direction": {
+          "value": ""
+        },
+        "constraint": null
+      },
+      "zIndex": 0
+    },
+    "edges": null,
+    "objects": [
+      {
+        "id": "memo",
+        "id_val": "memo",
+        "references": [
+          {
+            "key": {
+              "range": "d2/testdata/d2compiler/TestCompile/markdown_ampersand.d2,0:0:0-0:4:4",
+              "path": [
+                {
+                  "unquoted_string": {
+                    "range": "d2/testdata/d2compiler/TestCompile/markdown_ampersand.d2,0:0:0-0:4:4",
+                    "value": [
+                      {
+                        "string": "memo",
+                        "raw_string": "memo"
+                      }
+                    ]
+                  }
+                }
+              ]
+            },
+            "key_path_index": 0,
+            "map_key_edge_index": -1
+          }
+        ],
+        "attributes": {
+          "label": {
+            "value": "<a href=\"https://www.google.com/search?q=d2&newwindow=1&amp;bar\">d2</a>"
+          },
+          "labelDimensions": {
+            "width": 0,
+            "height": 0
+          },
+          "style": {},
+          "near_key": null,
+          "language": "markdown",
+          "shape": {
+            "value": "text"
+          },
+          "direction": {
+            "value": ""
+          },
+          "constraint": null
+        },
+        "zIndex": 0
+      }
+    ]
+  },
+  "err": null
+}