From 9ad44bdec052c143fb0d0affb91d8d4cfabedd79 Mon Sep 17 00:00:00 2001
From: Alexander Wang <alex@terrastruct.com>
Date: Mon, 28 Nov 2022 13:25:40 -0800
Subject: [PATCH] 2022-11-28 01:25:40PM

---
 lib/imgbundler/imgbundler.go      | 27 ++++++++++++++++++---------
 lib/imgbundler/imgbundler_test.go | 21 +++++++++++++++++++++
 2 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/lib/imgbundler/imgbundler.go b/lib/imgbundler/imgbundler.go
index 0cc88fc2b..5d705d60e 100644
--- a/lib/imgbundler/imgbundler.go
+++ b/lib/imgbundler/imgbundler.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io/ioutil"
 	"net/http"
+	"net/url"
 	"os"
 	"regexp"
 	"strings"
@@ -15,8 +16,7 @@ import (
 	"oss.terrastruct.com/d2/lib/xmain"
 )
 
-var uriImgRe = regexp.MustCompile(`<image href="(http[^"]+)"`)
-var localImgRe = regexp.MustCompile(`<image href="([^http][^"]+)"`)
+var imageRe = regexp.MustCompile(`<image href="([^"]+)"`)
 
 type resp struct {
 	srctxt string
@@ -25,26 +25,35 @@ type resp struct {
 }
 
 func InlineLocal(ms *xmain.State, in []byte) ([]byte, error) {
-	return inline(ms, localImgRe, in)
+	return inline(ms, in, false)
 }
 
 func InlineRemote(ms *xmain.State, in []byte) ([]byte, error) {
-	return inline(ms, uriImgRe, in)
+	return inline(ms, in, true)
 }
 
-func inline(ms *xmain.State, re *regexp.Regexp, in []byte) ([]byte, error) {
+func inline(ms *xmain.State, in []byte, isRemote bool) ([]byte, error) {
 	svg := string(in)
 
-	imgs := re.FindAllStringSubmatch(svg, -1)
+	imgs := imageRe.FindAllStringSubmatch(svg, -1)
+
+	var filtered [][]string
+	for _, img := range imgs {
+		u, err := url.Parse(img[1])
+		isRemoteImg := err == nil && strings.HasPrefix(u.Scheme, "http")
+		if isRemoteImg == isRemote {
+			filtered = append(filtered, img)
+		}
+	}
 
 	var wg sync.WaitGroup
 	respChan := make(chan resp)
 
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
-	wg.Add(len(imgs))
-	for _, img := range imgs {
-		if re == uriImgRe {
+	wg.Add(len(filtered))
+	for _, img := range filtered {
+		if isRemote {
 			go fetch(ctx, img[0], img[1], respChan)
 		} else {
 			go read(ctx, img[0], img[1], respChan)
diff --git a/lib/imgbundler/imgbundler_test.go b/lib/imgbundler/imgbundler_test.go
index de27ef019..6aa148207 100644
--- a/lib/imgbundler/imgbundler_test.go
+++ b/lib/imgbundler/imgbundler_test.go
@@ -26,6 +26,27 @@ func (f RoundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
 	return f(req), nil
 }
 
+func TestRegex(t *testing.T) {
+	urls := []string{
+		"https://icons.terrastruct.com/essentials/004-picture.svg",
+		"http://icons.terrastruct.com/essentials/004-picture.svg",
+	}
+
+	notURLs := []string{
+		"hi.png",
+		"./cat.png",
+		"/cat.png",
+	}
+
+	for _, href := range append(urls, notURLs...) {
+		str := fmt.Sprintf(`<image href="%s" />`, href)
+		matches := imageRe.FindAllStringSubmatch(str, -1)
+		if len(matches) != 1 {
+			t.Fatalf("uri regex didn't match %s", str)
+		}
+	}
+}
+
 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"