From f606ca9e73822457f634cd5468d09896dffa4ab9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20Mart=C3=AD?= <mvdan@mvdan.cc>
Date: Thu, 18 Apr 2019 15:07:33 +0900
Subject: [PATCH] fix a couple of crashes with premature cancels

The first would fail as we'd decrement the WaitGroup before adding to
it, and the escond would fail with a nil pointer dereference panic.

Add tests to reproduce both, and fix them.
---
 chromedp.go      | 11 +++++++++--
 chromedp_test.go | 27 +++++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/chromedp.go b/chromedp.go
index 650701f..51f9577 100644
--- a/chromedp.go
+++ b/chromedp.go
@@ -80,7 +80,7 @@ func NewContext(parent context.Context, opts ...ContextOption) (context.Context,
 	if pc := FromContext(parent); pc != nil {
 		c.Allocator = pc.Allocator
 		c.Browser = pc.Browser
-		// don't inherit SessionID, so that NewContext can be used to
+		// don't inherit Target, so that NewContext can be used to
 		// create a new tab on the same browser.
 
 		c.first = c.Browser == nil
@@ -94,6 +94,7 @@ func NewContext(parent context.Context, opts ...ContextOption) (context.Context,
 	}
 
 	ctx = context.WithValue(ctx, contextKey{}, c)
+	c.wg.Add(1)
 	go func() {
 		<-ctx.Done()
 		if c.first {
@@ -103,6 +104,13 @@ func NewContext(parent context.Context, opts ...ContextOption) (context.Context,
 			return
 		}
 
+		if c.Target == nil {
+			// This is a new tab, but we didn't create it and attach
+			// to it yet. Nothing to do.
+			c.wg.Done()
+			return
+		}
+
 		// Not the original browser tab; simply detach and close it.
 		// We need a new context, as ctx is cancelled; use a 1s timeout.
 		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
@@ -215,7 +223,6 @@ func (c *Context) newSession(ctx context.Context) error {
 	if err != nil {
 		return err
 	}
-	c.wg.Add(1)
 
 	c.Target = c.Browser.newExecutorForTarget(ctx, targetID, sessionID)
 
diff --git a/chromedp_test.go b/chromedp_test.go
index abef5e1..4e17e41 100644
--- a/chromedp_test.go
+++ b/chromedp_test.go
@@ -198,3 +198,30 @@ func TestCancelError(t *testing.T) {
 		t.Fatalf("expected a non-nil error, got %v", err)
 	}
 }
+
+func TestPrematureCancel(t *testing.T) {
+	t.Parallel()
+
+	// Cancel before the browser is allocated.
+	ctx, cancel := NewContext(context.Background())
+	cancel()
+	if err := Run(ctx); err != context.Canceled {
+		t.Fatalf("wanted canceled context error, got %v", err)
+	}
+}
+
+func TestPrematureCancelTab(t *testing.T) {
+	t.Parallel()
+
+	ctx1, cancel := NewContext(context.Background())
+	defer cancel()
+	if err := Run(ctx1); err != nil {
+		t.Fatal(err)
+	}
+
+	// Cancel after the browser is allocated, but before we've created a new
+	// tab.
+	ctx2, cancel := NewContext(ctx1)
+	cancel()
+	Run(ctx2)
+}