From 61e3415d93e3055a3cd0513a7d405a19f1757e0c Mon Sep 17 00:00:00 2001
From: Bradley Kemp <bradley@bradleyjkemp.dev>
Date: Tue, 9 Apr 2024 17:42:24 +0100
Subject: [PATCH] add error handling for non-200 status codes

---
 urlscanio.go | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/urlscanio.go b/urlscanio.go
index 1d7e55b9..8b62b1cd 100644
--- a/urlscanio.go
+++ b/urlscanio.go
@@ -43,7 +43,10 @@ func InputFromURLScan(ctx context.Context, urlscanUUID string, client httpClient
 	g.Go(func() error {
 		domReq, _ := http.NewRequestWithContext(ctx, http.MethodGet, "https://urlscan.io/dom/"+result.Task.Uuid, nil)
 		domResp, err := client.Do(domReq)
-		if err != nil {
+		if err != nil || domResp.StatusCode != 200 {
+			if err == nil {
+				err = fmt.Errorf(domResp.Status)
+			}
 			return fmt.Errorf("failed to get result dom: %w", err)
 		}
 		defer domResp.Body.Close()
@@ -98,7 +101,10 @@ func InputFromURLScan(ctx context.Context, urlscanUUID string, client httpClient
 			// Fetch the response in parallel with other threads, only lock the mutex once we're modifying the Input{}
 			resourceReq, _ := http.NewRequestWithContext(ctx, http.MethodGet, "https://urlscan.io/responses/"+request.Response.Hash, nil)
 			resp, err := client.Do(resourceReq)
-			if err != nil {
+			if err != nil || resp.StatusCode != 200 && resp.StatusCode != 404 {
+				if err == nil {
+					err = fmt.Errorf(resp.Status)
+				}
 				return fmt.Errorf("failed to fetch resource %s %s: %w", request.Request.RequestId, request.Response.Hash, err)
 			}
 			resource, _ := io.ReadAll(resp.Body) // always read the body to completion to ensure proper connection re-use + caching