From bd8a389f953c79ec3c09962b15b607a0e6e4e8a4 Mon Sep 17 00:00:00 2001
From: Alexander Wang <alex@terrastruct.com>
Date: Sat, 26 Nov 2022 18:14:41 -0800
Subject: [PATCH] 2022-11-26 06:14:41PM

---
 cmd/d2/main.go                    |  6 ++-
 lib/imgbundler/imgbundler.go      | 40 +++++++++++++++++--
 lib/imgbundler/imgbundler_test.go | 66 ++++++++++++++++++++++++++++++-
 lib/imgbundler/test_svg.svg       |  7 ++++
 4 files changed, 112 insertions(+), 7 deletions(-)
 create mode 100644 lib/imgbundler/test_svg.svg

diff --git a/cmd/d2/main.go b/cmd/d2/main.go
index 71b49c9cb..04e37aa3e 100644
--- a/cmd/d2/main.go
+++ b/cmd/d2/main.go
@@ -205,10 +205,14 @@ func compile(ctx context.Context, ms *xmain.State, plugin d2plugin.Plugin, theme
 	if err != nil {
 		return nil, err
 	}
+	svg, err = imgbundler.InlineLocal(ms, svg)
+	if err != nil {
+		return nil, err
+	}
 
 	out := svg
 	if filepath.Ext(outputPath) == ".png" {
-		svg, err = imgbundler.Inline(ms, svg)
+		svg, err = imgbundler.InlineRemote(ms, svg)
 		if err != nil {
 			return nil, err
 		}
diff --git a/lib/imgbundler/imgbundler.go b/lib/imgbundler/imgbundler.go
index 58dddd435..0cc88fc2b 100644
--- a/lib/imgbundler/imgbundler.go
+++ b/lib/imgbundler/imgbundler.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io/ioutil"
 	"net/http"
+	"os"
 	"regexp"
 	"strings"
 	"sync"
@@ -14,7 +15,8 @@ import (
 	"oss.terrastruct.com/d2/lib/xmain"
 )
 
-var imgRe = regexp.MustCompile(`<image href="([^"]+)"`)
+var uriImgRe = regexp.MustCompile(`<image href="(http[^"]+)"`)
+var localImgRe = regexp.MustCompile(`<image href="([^http][^"]+)"`)
 
 type resp struct {
 	srctxt string
@@ -22,10 +24,18 @@ type resp struct {
 	err    error
 }
 
-func Inline(ms *xmain.State, in []byte) ([]byte, error) {
+func InlineLocal(ms *xmain.State, in []byte) ([]byte, error) {
+	return inline(ms, localImgRe, in)
+}
+
+func InlineRemote(ms *xmain.State, in []byte) ([]byte, error) {
+	return inline(ms, uriImgRe, in)
+}
+
+func inline(ms *xmain.State, re *regexp.Regexp, in []byte) ([]byte, error) {
 	svg := string(in)
 
-	imgs := imgRe.FindAllStringSubmatch(svg, -1)
+	imgs := re.FindAllStringSubmatch(svg, -1)
 
 	var wg sync.WaitGroup
 	respChan := make(chan resp)
@@ -34,7 +44,11 @@ func Inline(ms *xmain.State, in []byte) ([]byte, error) {
 	defer cancel()
 	wg.Add(len(imgs))
 	for _, img := range imgs {
-		go fetch(ctx, img[0], img[1], respChan)
+		if re == uriImgRe {
+			go fetch(ctx, img[0], img[1], respChan)
+		} else {
+			go read(ctx, img[0], img[1], respChan)
+		}
 	}
 
 	go func() {
@@ -93,3 +107,21 @@ func fetch(ctx context.Context, srctxt, href string, respChan chan resp) {
 		data:   fmt.Sprintf("data:%s;base64,%s", mimeType, enc),
 	}
 }
+
+func read(ctx context.Context, srctxt, href string, respChan chan resp) {
+	data, err := os.ReadFile(href)
+	if err != nil {
+		respChan <- resp{err: err}
+		return
+	}
+
+	mimeType := http.DetectContentType(data)
+	mimeType = strings.Replace(mimeType, "text/xml", "image/svg+xml", 1)
+
+	enc := base64.StdEncoding.EncodeToString(data)
+
+	respChan <- resp{
+		srctxt: srctxt,
+		data:   fmt.Sprintf("data:%s;base64,%s", mimeType, enc),
+	}
+}
diff --git a/lib/imgbundler/imgbundler_test.go b/lib/imgbundler/imgbundler_test.go
index 34a164c71..de27ef019 100644
--- a/lib/imgbundler/imgbundler_test.go
+++ b/lib/imgbundler/imgbundler_test.go
@@ -7,6 +7,7 @@ import (
 	"io/ioutil"
 	"net/http"
 	"os"
+	"path/filepath"
 	"strings"
 	"testing"
 
@@ -25,7 +26,7 @@ func (f RoundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
 	return f(req), nil
 }
 
-func TestInliner(t *testing.T) {
+func TestInlineRemote(t *testing.T) {
 	svgURL := "https://icons.terrastruct.com/essentials/004-picture.svg"
 	pngURL := "https://cdn4.iconfinder.com/data/icons/smart-phones-technologies/512/android-phone.png"
 
@@ -86,7 +87,7 @@ width="328" height="587" viewBox="-100 -131 328 587"><style type="text/css">
 		Env: xos.NewEnv(os.Environ()),
 	}
 	ms.Log = cmdlog.Log(ms.Env, os.Stderr)
-	out, err := Inline(ms, []byte(sampleSVG))
+	out, err := InlineRemote(ms, []byte(sampleSVG))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -100,3 +101,64 @@ width="328" height="587" viewBox="-100 -131 328 587"><style type="text/css">
 		t.Fatal("no png image inserted")
 	}
 }
+
+func TestInlineLocal(t *testing.T) {
+	svgURL, err := filepath.Abs("./test_svg.svg")
+	if err != nil {
+		t.Fatal(err)
+	}
+	pngURL, err := filepath.Abs("./test_png.png")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	sampleSVG := fmt.Sprintf(`<?xml version="1.0" encoding="utf-8"?>
+<svg
+style="background: white;"
+xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink"
+width="328" height="587" viewBox="-100 -131 328 587"><style type="text/css">
+<![CDATA[
+.shape {
+  shape-rendering: geometricPrecision;
+  stroke-linejoin: round;
+}
+.connection {
+  stroke-linecap: round;
+  stroke-linejoin: round;
+}
+
+]]>
+</style><g id="a"><g class="shape" ><image href="%s" x="0" y="0" width="128" height="128" style="fill:#FFFFFF;stroke:#0D32B2;opacity:1.000000;stroke-width:2;" /></g><text class="text-bold" x="64.000000" y="-15.000000" style="text-anchor:middle;font-size:16px;fill:#0A0F25">a</text></g><g id="b"><g class="shape" ><image href="%s" x="0" y="228" width="128" height="128" style="fill:#FFFFFF;stroke:#0D32B2;opacity:1.000000;stroke-width:2;" /></g><text class="text-bold" x="64.000000" y="213.000000" style="text-anchor:middle;font-size:16px;fill:#0A0F25">b</text></g><g id="(a -&gt; b)[0]"><marker id="mk-3990223579" markerWidth="10.000000" markerHeight="12.000000" refX="7.000000" refY="6.000000" viewBox="0.000000 0.000000 10.000000 12.000000" orient="auto" markerUnits="userSpaceOnUse"> <polygon class="connection" fill="#0D32B2" stroke-width="2" points="0.000000,0.000000 10.000000,6.000000 0.000000,12.000000" /> </marker><path d="M 64.000000 130.000000 C 64.000000 168.000000 64.000000 188.000000 64.000000 224.000000" class="connection" style="fill:none;stroke:#0D32B2;opacity:1.000000;stroke-width:2;" marker-end="url(#mk-3990223579)" /></g><style type="text/css"><![CDATA[
+.text-bold {
+	font-family: "font-bold";
+}
+@font-face {
+	font-family: font-bold;
+	src: url("REMOVED");
+}]]></style></svg>
+`, svgURL, pngURL)
+
+	ms := &xmain.State{
+		Name: "test",
+
+		Stdin:  os.Stdin,
+		Stdout: os.Stdout,
+		Stderr: os.Stderr,
+
+		Env: xos.NewEnv(os.Environ()),
+	}
+	ms.Log = cmdlog.Log(ms.Env, os.Stderr)
+	out, err := InlineLocal(ms, []byte(sampleSVG))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if strings.Contains(string(out), svgURL) {
+		t.Fatal("links still exist")
+	}
+	if !strings.Contains(string(out), "image/svg+xml") {
+		t.Fatal("no svg image inserted")
+	}
+	if !strings.Contains(string(out), "image/png") {
+		t.Fatal("no png image inserted")
+	}
+}
diff --git a/lib/imgbundler/test_svg.svg b/lib/imgbundler/test_svg.svg
new file mode 100644
index 000000000..803ac81da
--- /dev/null
+++ b/lib/imgbundler/test_svg.svg
@@ -0,0 +1,7 @@
+<?xml version="1.0" encoding="utf-8"?><svg width="165" height="58" viewBox="0 0 165 58" fill="none" xmlns="http://www.w3.org/2000/svg">
+<path d="M74.2187 57.213C73.5457 57.213 73 56.6673 73 55.9943V1.21875C73 0.545653 73.5457 0 74.2187 0H89.146C94.996 0 100.056 0.994499 104.327 2.9835C108.597 4.9725 111.903 8.073 114.243 12.285C116.641 16.4385 117.84 21.7912 117.84 28.3432C117.84 34.8952 116.67 40.3065 114.33 44.577C111.99 48.8475 108.714 52.0358 104.502 54.1418C100.349 56.1893 95.464 57.213 89.848 57.213H74.2187ZM85.8993 45.552C85.8993 46.2251 86.4449 46.7707 87.118 46.7707H88.3563C91.6323 46.7707 94.4695 46.1858 96.868 45.0158C99.325 43.8458 101.226 41.9152 102.572 39.2242C103.917 36.5332 104.59 32.9062 104.59 28.3432C104.59 23.7802 103.917 20.2117 102.572 17.6377C101.226 15.0052 99.325 13.1625 96.868 12.1095C94.4695 10.998 91.6323 10.4422 88.3563 10.4422H87.118C86.4449 10.4422 85.8993 10.9879 85.8993 11.661V45.552Z" fill="#0A0F25"/>
+<path d="M126.409 57.213C125.736 57.213 125.191 56.6673 125.191 55.9943V50.365C125.191 50.031 125.328 49.7116 125.57 49.4814C130.153 45.1203 134.209 41.1747 137.739 37.6448C141.424 33.9593 144.291 30.5955 146.338 27.5535C148.386 24.453 149.41 21.5865 149.41 18.954C149.41 16.2045 148.678 14.1277 147.216 12.7237C145.812 11.2612 143.911 10.53 141.512 10.53C139.523 10.53 137.71 11.115 136.072 12.285C134.816 13.1689 133.615 14.1822 132.47 15.3251C131.982 15.8119 131.19 15.8323 130.7 15.3481L125.423 10.137C124.952 9.67228 124.937 8.91536 125.403 8.44619C127.888 5.94519 130.42 4.03654 133 2.72025C135.867 1.19925 139.289 0.438748 143.267 0.438748C146.953 0.438748 150.17 1.17 152.92 2.6325C155.669 4.095 157.804 6.17175 159.325 8.86275C160.905 11.4952 161.695 14.5957 161.695 18.1642C161.695 21.2647 160.817 24.453 159.062 27.729C157.366 30.9465 155.172 34.1932 152.481 37.4692C149.79 40.6867 146.923 43.875 143.881 47.034C145.285 46.8585 146.865 46.7123 148.62 46.5953C150.375 46.4198 151.896 46.332 153.183 46.332H163.196C163.869 46.332 164.415 46.8777 164.415 47.5507V55.9943C164.415 56.6673 163.869 57.213 163.196 57.213H126.409Z" fill="#0A0F25"/>
+<path fill-rule="evenodd" clip-rule="evenodd" d="M40.5718 51.8158C40.5718 54.3741 40.5718 55.6533 39.6615 56.3965C38.7512 57.1397 37.5897 56.9026 35.2666 56.4282C27.1711 54.7753 19.4539 50.8089 13.1742 44.5291C7.07441 38.4294 3.15733 30.9735 1.42293 23.132C0.904136 20.7865 0.644741 19.6137 1.38829 18.6877C2.13183 17.7617 3.42795 17.7617 6.02019 17.7617L35.6186 17.7617C37.9536 17.7617 39.1211 17.7617 39.8465 18.4871C40.5718 19.2125 40.5718 20.38 40.5718 22.715L40.5718 51.8158Z" fill="#476CEF"/>
+<path fill-rule="evenodd" clip-rule="evenodd" d="M52.9077 38.7915C55.2308 39.2658 56.3923 39.503 57.3026 38.7598C58.2129 38.0165 58.213 36.7374 58.213 34.179L58.213 5.07828C58.213 2.74329 58.213 1.5758 57.4876 0.850412C56.7622 0.125025 55.5947 0.125026 53.2597 0.125025L23.6613 0.125025C21.0691 0.125026 19.7729 0.125026 19.0294 1.051C18.2859 1.97698 18.5452 3.14974 19.064 5.49527C20.7984 13.3368 24.7155 20.7927 30.8153 26.8924C37.0951 33.1722 44.8122 37.1386 52.9077 38.7915Z" fill="#94ABFF"/>
+<path fill-rule="evenodd" clip-rule="evenodd" d="M40.5698 34.2456L40.5698 22.715L40.5698 22.715C40.5698 20.38 40.5698 19.2125 39.8444 18.4871C39.119 17.7617 37.9515 17.7617 35.6165 17.7617L23.8125 17.7617C25.6894 21.0268 28.023 24.101 30.8134 26.8914C33.7817 29.8597 37.0711 32.3111 40.5698 34.2456Z" fill="#234CDA"/>
+</svg>