From fe75e7067b3564185ed3ff1999ec44b9a2d3e6f6 Mon Sep 17 00:00:00 2001
From: Alexander Wang <alex@terrastruct.com>
Date: Tue, 29 Nov 2022 13:30:22 -0800
Subject: [PATCH] 2022-11-29 01:30:22PM

---
 lib/imgbundler/imgbundler.go      |  6 +++++-
 lib/imgbundler/imgbundler_test.go | 29 +++++++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/lib/imgbundler/imgbundler.go b/lib/imgbundler/imgbundler.go
index 43cfa5adf..0b1addc22 100644
--- a/lib/imgbundler/imgbundler.go
+++ b/lib/imgbundler/imgbundler.go
@@ -20,6 +20,9 @@ import (
 	"oss.terrastruct.com/d2/lib/xmain"
 )
 
+// 32 MB
+var max_img_size int64 = 33_554_432
+
 var imageRe = regexp.MustCompile(`<image href="([^"]+)"`)
 
 type resp struct {
@@ -132,7 +135,8 @@ func fetch(ctx context.Context, href string) (string, error) {
 	if imgResp.StatusCode != 200 {
 		return "", fmt.Errorf("img %s returned status code %d", href, imgResp.StatusCode)
 	}
-	data, err := ioutil.ReadAll(imgResp.Body)
+	r := http.MaxBytesReader(nil, imgResp.Body, max_img_size)
+	data, err := ioutil.ReadAll(r)
 	if err != nil {
 		return "", err
 	}
diff --git a/lib/imgbundler/imgbundler_test.go b/lib/imgbundler/imgbundler_test.go
index 6ad608180..33196c194 100644
--- a/lib/imgbundler/imgbundler_test.go
+++ b/lib/imgbundler/imgbundler_test.go
@@ -1,6 +1,7 @@
 package imgbundler
 
 import (
+	"crypto/rand"
 	_ "embed"
 	"fmt"
 	"net/http"
@@ -115,6 +116,34 @@ width="328" height="587" viewBox="-100 -131 328 587"><style type="text/css">
 		t.Fatal("no png image inserted")
 	}
 
+	// Test almost too large response
+	transport = roundTripFunc(func(req *http.Request) *http.Response {
+		respRecorder := httptest.NewRecorder()
+		bytes := make([]byte, max_img_size-1)
+		rand.Read(bytes)
+		respRecorder.Write(bytes)
+		respRecorder.WriteHeader(200)
+		return respRecorder.Result()
+	})
+	_, err = InlineRemote(ms, []byte(sampleSVG))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Test too large response
+	transport = roundTripFunc(func(req *http.Request) *http.Response {
+		respRecorder := httptest.NewRecorder()
+		bytes := make([]byte, max_img_size+1)
+		rand.Read(bytes)
+		respRecorder.Write(bytes)
+		respRecorder.WriteHeader(200)
+		return respRecorder.Result()
+	})
+	_, err = InlineRemote(ms, []byte(sampleSVG))
+	if err == nil {
+		t.Fatal("expected error")
+	}
+
 	// Test error response
 	transport = roundTripFunc(func(req *http.Request) *http.Response {
 		respRecorder := httptest.NewRecorder()