From 10e13c3716bbfab5fd224caae0a09ad6a16ee1e1 Mon Sep 17 00:00:00 2001
From: Alexander Wang <alex@terrastruct.com>
Date: Tue, 29 Nov 2022 13:19:03 -0800
Subject: [PATCH] fixes

---
 cmd/d2/main.go                    |  4 +--
 lib/imgbundler/imgbundler.go      | 20 ++++++++++-----
 lib/imgbundler/imgbundler_test.go | 42 +++++++++++++++++++++++++++++--
 3 files changed, 56 insertions(+), 10 deletions(-)

diff --git a/cmd/d2/main.go b/cmd/d2/main.go
index 9e8d290d1..1a7e68eac 100644
--- a/cmd/d2/main.go
+++ b/cmd/d2/main.go
@@ -205,7 +205,7 @@ func compile(ctx context.Context, ms *xmain.State, isWatching bool, plugin d2plu
 	if err != nil {
 		return nil, err
 	}
-	svg, err = imgbundler.InlineLocal(svg)
+	svg, err = imgbundler.InlineLocal(ms, svg)
 	if err != nil {
 		// Missing/broken images are fine during watch mode, as the user is likely building up a diagram.
 		// Otherwise, the assumption is that this diagram is building for production, and broken images are not okay.
@@ -217,7 +217,7 @@ func compile(ctx context.Context, ms *xmain.State, isWatching bool, plugin d2plu
 
 	out := svg
 	if filepath.Ext(outputPath) == ".png" {
-		svg, err = imgbundler.InlineRemote(svg)
+		svg, err = imgbundler.InlineRemote(ms, svg)
 		if err != nil {
 			if !isWatching {
 				return nil, err
diff --git a/lib/imgbundler/imgbundler.go b/lib/imgbundler/imgbundler.go
index 8782bea0e..43cfa5adf 100644
--- a/lib/imgbundler/imgbundler.go
+++ b/lib/imgbundler/imgbundler.go
@@ -16,6 +16,8 @@ import (
 
 	"go.uber.org/multierr"
 	"oss.terrastruct.com/xdefer"
+
+	"oss.terrastruct.com/d2/lib/xmain"
 )
 
 var imageRe = regexp.MustCompile(`<image href="([^"]+)"`)
@@ -26,15 +28,15 @@ type resp struct {
 	err    error
 }
 
-func InlineLocal(in []byte) ([]byte, error) {
-	return inline(in, false)
+func InlineLocal(ms *xmain.State, in []byte) ([]byte, error) {
+	return inline(ms, in, false)
 }
 
-func InlineRemote(in []byte) ([]byte, error) {
-	return inline(in, true)
+func InlineRemote(ms *xmain.State, in []byte) ([]byte, error) {
+	return inline(ms, in, true)
 }
 
-func inline(svg []byte, isRemote bool) (_ []byte, err error) {
+func inline(ms *xmain.State, svg []byte, isRemote bool) (_ []byte, err error) {
 	defer xdefer.Errorf(&err, "failed to bundle images")
 	imgs := imageRe.FindAllSubmatch(svg, -1)
 
@@ -93,10 +95,13 @@ func inline(svg []byte, isRemote bool) (_ []byte, err error) {
 	for {
 		select {
 		case <-ctx.Done():
+			ms.Log.Debug.Printf("there")
 			return nil, fmt.Errorf("failed waiting for imgbundler workers: %w", ctx.Err())
+		case <-time.After(time.Second * 5):
+			ms.Log.Info.Printf("fetching images...")
 		case resp, ok := <-respChan:
 			if !ok {
-				return svg, nil
+				return svg, err
 			}
 			if resp.err != nil {
 				err = multierr.Combine(err, resp.err)
@@ -124,6 +129,9 @@ func fetch(ctx context.Context, href string) (string, error) {
 		return "", err
 	}
 	defer imgResp.Body.Close()
+	if imgResp.StatusCode != 200 {
+		return "", fmt.Errorf("img %s returned status code %d", href, imgResp.StatusCode)
+	}
 	data, err := ioutil.ReadAll(imgResp.Body)
 	if err != nil {
 		return "", err
diff --git a/lib/imgbundler/imgbundler_test.go b/lib/imgbundler/imgbundler_test.go
index 4ccd69854..6ad608180 100644
--- a/lib/imgbundler/imgbundler_test.go
+++ b/lib/imgbundler/imgbundler_test.go
@@ -5,9 +5,15 @@ import (
 	"fmt"
 	"net/http"
 	"net/http/httptest"
+	"os"
 	"path/filepath"
 	"strings"
 	"testing"
+
+	"oss.terrastruct.com/cmdlog"
+	"oss.terrastruct.com/xos"
+
+	"oss.terrastruct.com/d2/lib/xmain"
 )
 
 //go:embed test_png.png
@@ -70,6 +76,17 @@ width="328" height="587" viewBox="-100 -131 328 587"><style type="text/css">
 }]]></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)
+
 	transport = roundTripFunc(func(req *http.Request) *http.Response {
 		respRecorder := httptest.NewRecorder()
 		switch req.URL.String() {
@@ -84,7 +101,7 @@ width="328" height="587" viewBox="-100 -131 328 587"><style type="text/css">
 		return respRecorder.Result()
 	})
 
-	out, err := InlineRemote([]byte(sampleSVG))
+	out, err := InlineRemote(ms, []byte(sampleSVG))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -97,6 +114,17 @@ width="328" height="587" viewBox="-100 -131 328 587"><style type="text/css">
 	if !strings.Contains(string(out), "image/png") {
 		t.Fatal("no png image inserted")
 	}
+
+	// Test error response
+	transport = roundTripFunc(func(req *http.Request) *http.Response {
+		respRecorder := httptest.NewRecorder()
+		respRecorder.WriteHeader(500)
+		return respRecorder.Result()
+	})
+	_, err = InlineRemote(ms, []byte(sampleSVG))
+	if err == nil {
+		t.Fatal("expected error")
+	}
 }
 
 func TestInlineLocal(t *testing.T) {
@@ -135,7 +163,17 @@ width="328" height="587" viewBox="-100 -131 328 587"><style type="text/css">
 }]]></style></svg>
 `, svgURL, pngURL)
 
-	out, err := InlineLocal([]byte(sampleSVG))
+	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)
 	}