From 1e58c942631f43ceb35205eba060f42a4ff8d1ef Mon Sep 17 00:00:00 2001
From: Antoni Sawicki <tenox@google.com>
Date: Mon, 17 Jun 2019 00:38:52 -0700
Subject: [PATCH] fix for #40 hang on page with no nodes, added timeout for
 cdp.Nodes

---
 wrp.go | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/wrp.go b/wrp.go
index d250e25..88fbb0c 100644
--- a/wrp.go
+++ b/wrp.go
@@ -177,8 +177,7 @@ func capture(gourl string, w int64, h int64, s float64, co int, p int64, i bool,
 		chromedp.Evaluate(fmt.Sprintf("window.scrollTo(0, %d);", p*int64(float64(h)*float64(0.9))), &res),
 		chromedp.Sleep(time.Second*1),
 		chromedp.CaptureScreenshot(&pngbuf),
-		chromedp.Location(&loc),
-		chromedp.Nodes("a", &nodes, chromedp.ByQueryAll))
+		chromedp.Location(&loc))
 
 	if err != nil {
 		if err.Error() == "context canceled" {
@@ -216,6 +215,9 @@ func capture(gourl string, w int64, h int64, s float64, co int, p int64, i bool,
 	gifmap[imgpath] = gifbuf
 
 	// Process Nodes
+	tctx, cancel := context.WithTimeout(ctx, time.Second*2) // a context with timeout is needed for pages without nodes
+	defer cancel()
+	chromedp.Run(tctx, chromedp.Nodes("a", &nodes, chromedp.ByQueryAll))
 	base, _ := url.Parse(loc)
 	if i {
 		fmt.Fprintf(out, "<A HREF=\"%s\"><IMG SRC=\"%s\" ALT=\"wrp\" BORDER=\"0\" ISMAP></A>", mappath, imgpath)