From b481eeac51083063653cadb194d26f60914fc204 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20Mart=C3=AD?= <mvdan@mvdan.cc>
Date: Mon, 8 Apr 2019 18:52:14 +0200
Subject: [PATCH] support fetching errors from cancellation

The API isn't very shiny, but it works. It doesn't matter that much, as
most users won't care about these errors.

Fixes #295.
---
 allocate.go      | 11 +++++++----
 chromedp.go      | 24 ++++++++++++++++++++----
 chromedp_test.go | 44 ++++++++++++++++++++++++++++++++++++++++++--
 3 files changed, 69 insertions(+), 10 deletions(-)

diff --git a/allocate.go b/allocate.go
index b99f78b..e0ed302 100644
--- a/allocate.go
+++ b/allocate.go
@@ -22,9 +22,6 @@ type Allocator interface {
 	// as temporary directories) will be freed.
 	Allocate(context.Context, ...BrowserOption) (*Browser, error)
 
-	// TODO: Wait should probably return an error, which can then be
-	// retrieved by the user if just calling cancel().
-
 	// Wait blocks until an allocator has freed all of its resources.
 	// Cancelling the context obtained via NewAllocator will already perform
 	// this operation, so normally there's no need to call Wait directly.
@@ -113,11 +110,17 @@ func (p *ExecAllocator) Allocate(ctx context.Context, opts ...BrowserOption) (*B
 		<-ctx.Done()
 		// First wait for the process to be finished.
 		if cmd != nil {
+			// TODO: do we care about this error in any scenario? if
+			// the user cancelled the context and killed chrome,
+			// this will most likely just be "signal: killed", which
+			// isn't interesting.
 			cmd.Wait()
 		}
 		// Then delete the temporary user data directory, if needed.
 		if removeDir {
-			os.RemoveAll(dataDir)
+			if err := os.RemoveAll(dataDir); c.cancelErr == nil {
+				c.cancelErr = err
+			}
 		}
 		p.wg.Done()
 		c.wg.Done()
diff --git a/chromedp.go b/chromedp.go
index d96ed73..3a56d0a 100644
--- a/chromedp.go
+++ b/chromedp.go
@@ -55,6 +55,11 @@ type Context struct {
 
 	// wg allows waiting for a target to be closed on cancellation.
 	wg sync.WaitGroup
+
+	// cancelErr is the first error encountered when cancelling this
+	// context, for example if a browser's temporary user data directory
+	// couldn't be deleted.
+	cancelErr error
 }
 
 // NewContext creates a browser context using the parent context.
@@ -98,14 +103,17 @@ func NewContext(parent context.Context, opts ...ContextOption) (context.Context,
 		defer cancel()
 		if id := c.Target.SessionID; id != "" {
 			action := target.DetachFromTarget().WithSessionID(id)
-			if err := action.Do(ctx, c.Browser); err != nil {
-				c.Browser.errf("%s", err)
+			if err := action.Do(ctx, c.Browser); c.cancelErr == nil {
+				c.cancelErr = err
 			}
 		}
 		if id := c.Target.TargetID; id != "" {
 			action := target.CloseTarget(id)
-			if _, err := action.Do(ctx, c.Browser); err != nil {
-				c.Browser.errf("%s", err)
+			if ok, err := action.Do(ctx, c.Browser); c.cancelErr == nil {
+				if !ok && err == nil {
+					err = fmt.Errorf("could not close target %q", id)
+				}
+				c.cancelErr = err
 			}
 		}
 		c.wg.Done()
@@ -125,6 +133,14 @@ func FromContext(ctx context.Context) *Context {
 	return c
 }
 
+func CancelError(ctx context.Context) error {
+	c := FromContext(ctx)
+	if c == nil {
+		return ErrInvalidContext
+	}
+	return c.cancelErr
+}
+
 // Run runs an action against the provided context. The provided context must
 // contain a valid Allocator; typically, that will be created via NewContext, or
 // via one of the allocator constructors like NewExecAllocator.
diff --git a/chromedp_test.go b/chromedp_test.go
index 9a1530c..b787e7a 100644
--- a/chromedp_test.go
+++ b/chromedp_test.go
@@ -28,7 +28,7 @@ var (
 func testAllocate(t *testing.T, path string) (_ context.Context, cancel func()) {
 	// Same browser, new tab; not needing to start new chrome browsers for
 	// each test gives a huge speed-up.
-	ctx, cancel := NewContext(browserCtx)
+	ctx, cancel_ := NewContext(browserCtx)
 
 	// Only navigate if we want a path, otherwise leave the blank page.
 	if path != "" {
@@ -37,7 +37,13 @@ func testAllocate(t *testing.T, path string) (_ context.Context, cancel func())
 		}
 	}
 
-	return ctx, cancel
+	cancelErr := func() {
+		cancel_() // not cancel(); that's the returned one
+		if err := CancelError(ctx); err != nil {
+			t.Error(err)
+		}
+	}
+	return ctx, cancelErr
 }
 
 func TestMain(m *testing.M) {
@@ -158,3 +164,37 @@ func TestBrowserQuit(t *testing.T) {
 		t.Fatalf("did not expect a standard context error: %v", err)
 	}
 }
+
+func TestCancelError(t *testing.T) {
+	t.Parallel()
+
+	ctx1, cancel1 := NewContext(context.Background())
+	defer cancel1()
+	if err := Run(ctx1); err != nil {
+		t.Fatal(err)
+	}
+
+	// Open and close a target normally; no error.
+	ctx2, cancel2 := NewContext(ctx1)
+	defer cancel2()
+	if err := Run(ctx2); err != nil {
+		t.Fatal(err)
+	}
+	cancel2()
+	if err := CancelError(ctx2); err != nil {
+		t.Fatalf("expected a nil error, got %v", err)
+	}
+
+	// Make "cancel" close the wrong target; error.
+	ctx3, cancel3 := NewContext(ctx1)
+	defer cancel3()
+	if err := Run(ctx3); err != nil {
+		t.Fatal(err)
+	}
+	FromContext(ctx3).Target.TargetID = "wrong"
+	cancel3()
+	time.Sleep(time.Second)
+	if err := CancelError(ctx3); err == nil {
+		t.Fatalf("expected a non-nil error, got %v", err)
+	}
+}