From bfab428e7e7c001bc587f29354e37147d735b6fc Mon Sep 17 00:00:00 2001
From: Marc Wickenden <marc@4armed.com>
Date: Wed, 17 Jan 2024 09:25:38 +0000
Subject: [PATCH] timeout URLs after 2 seconds

---
 internal/screenshot/screenshot.go | 23 +++++++++++++++--------
 1 file changed, 15 insertions(+), 8 deletions(-)

diff --git a/internal/screenshot/screenshot.go b/internal/screenshot/screenshot.go
index 92ed2c0..78f8aa4 100644
--- a/internal/screenshot/screenshot.go
+++ b/internal/screenshot/screenshot.go
@@ -1,9 +1,12 @@
 package screenshot
 
 import (
+	"context"
+	"errors"
 	"net/url"
 	"strings"
 	"sync"
+	"time"
 
 	"github.com/4armed/kubeshot/internal/config"
 	"github.com/go-rod/rod"
@@ -51,23 +54,27 @@ func screenshot(url string, c *config.Config) {
 		logger.Warning("getFilename returned error: %v", err)
 	}
 
-	browser := rod.New().MustConnect()
+	browser := rod.New().
+		MustConnect().
+		MustIgnoreCertErrors(true) // Ignore certificate warnings
 	defer browser.MustClose()
 
-	// Ignore certificate warnings
-	browser.MustIgnoreCertErrors(true)
+	page := browser.MustPage()
 
-	page, err := browser.Page(proto.TargetCreateTarget{
-		URL: url,
+	err = rod.Try(func() {
+		page.Timeout(2 * time.Second).MustNavigate(url)
 	})
-	if err != nil {
-		logger.Warning("could not load page %s: %v", url, err)
+	if errors.Is(err, context.DeadlineExceeded) {
+		logger.Warning("timeout loading URL %s: %v", url, err)
+		return
+	} else if err != nil {
+		logger.Warning("could not load URL %s: %v", url, err)
 		return
 	}
 
 	err = page.WaitLoad()
 	if err != nil {
-		logger.Warning("could not load URL %s: %v", url, err)
+		logger.Warning("error waiting for window.onload at %s: %v", url, err)
 		return
 	}