From 1decbccd743f6e194d3d421876ea08d6628ef9e5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20Mart=C3=AD?= <mvdan@mvdan.cc>
Date: Mon, 1 Apr 2019 14:31:11 +0100
Subject: [PATCH] store a Target pointer directly in Context

That way, we avoid the racy map access via Browser.executorForTarget. If
a context is attached to a target, the Target field must be non-nil.

The Browser.pages map is still racy, since multiple tabs can be created
concurrently; we'll fix this other data race in another commit.
---
 browser.go      |  9 +++++----
 context.go      | 12 +++++-------
 example_test.go |  2 +-
 handler.go      |  4 ++--
 4 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/browser.go b/browser.go
index 1a8a268..2d57395 100644
--- a/browser.go
+++ b/browser.go
@@ -9,6 +9,7 @@ package chromedp
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"log"
 	"sync/atomic"
 
@@ -102,16 +103,16 @@ func (b *Browser) send(method cdproto.MethodType, params easyjson.RawMessage) er
 	return b.conn.Write(msg)
 }
 
-func (b *Browser) executorForTarget(ctx context.Context, sessionID target.SessionID) *Target {
+func (b *Browser) newExecutorForTarget(ctx context.Context, sessionID target.SessionID) *Target {
 	if sessionID == "" {
 		panic("empty session ID")
 	}
-	if t, ok := b.pages[sessionID]; ok {
-		return t
+	if _, ok := b.pages[sessionID]; ok {
+		panic(fmt.Sprintf("executor for %q already exists", sessionID))
 	}
 	t := &Target{
 		browser:   b,
-		sessionID: sessionID,
+		SessionID: sessionID,
 
 		eventQueue: make(chan *cdproto.Message, 1024),
 		waitQueue:  make(chan func(cur *cdp.Frame) bool, 1024),
diff --git a/context.go b/context.go
index 87f5251..e7520d1 100644
--- a/context.go
+++ b/context.go
@@ -19,7 +19,7 @@ type Context struct {
 
 	Browser *Browser
 
-	SessionID target.SessionID
+	Target *Target
 }
 
 // NewContext creates a browser context using the parent context.
@@ -72,12 +72,12 @@ func Run(ctx context.Context, action Action) error {
 		}
 		c.Browser = browser
 	}
-	if c.SessionID == "" {
+	if c.Target == nil {
 		if err := c.newSession(ctx); err != nil {
 			return err
 		}
 	}
-	return action.Do(ctx, c.Browser.executorForTarget(ctx, c.SessionID))
+	return action.Do(ctx, c.Target)
 }
 
 func (c *Context) newSession(ctx context.Context) error {
@@ -93,7 +93,7 @@ func (c *Context) newSession(ctx context.Context) error {
 		return err
 	}
 
-	target := c.Browser.executorForTarget(ctx, sessionID)
+	c.Target = c.Browser.newExecutorForTarget(ctx, sessionID)
 
 	// enable domains
 	for _, enable := range []Action{
@@ -105,12 +105,10 @@ func (c *Context) newSession(ctx context.Context) error {
 		dom.Enable(),
 		css.Enable(),
 	} {
-		if err := enable.Do(ctx, target); err != nil {
+		if err := enable.Do(ctx, c.Target); err != nil {
 			return fmt.Errorf("unable to execute %T: %v", enable, err)
 		}
 	}
-
-	c.SessionID = sessionID
 	return nil
 }
 
diff --git a/example_test.go b/example_test.go
index d9c7f39..efb9aad 100644
--- a/example_test.go
+++ b/example_test.go
@@ -99,7 +99,7 @@ func ExampleManyTabs() {
 	c2 := chromedp.FromContext(ctx2)
 
 	fmt.Printf("Same browser: %t\n", c1.Browser == c2.Browser)
-	fmt.Printf("Same tab: %t\n", c1.SessionID == c2.SessionID)
+	fmt.Printf("Same tab: %t\n", c1.Target == c2.Target)
 
 	// wait for the resources to be cleaned up
 	cancel()
diff --git a/handler.go b/handler.go
index 7cd16f6..b49c551 100644
--- a/handler.go
+++ b/handler.go
@@ -19,7 +19,7 @@ import (
 // Target manages a Chrome DevTools Protocol target.
 type Target struct {
 	browser   *Browser
-	sessionID target.SessionID
+	SessionID target.SessionID
 
 	waitQueue  chan func(cur *cdp.Frame) bool
 	eventQueue chan *cdproto.Message
@@ -87,7 +87,7 @@ func (t *Target) Execute(ctx context.Context, method string, params json.Marshal
 		return err
 	}
 	sendParams := target.SendMessageToTarget(string(msgJSON)).
-		WithSessionID(t.sessionID)
+		WithSessionID(t.SessionID)
 	sendParamsJSON, _ := json.Marshal(sendParams)
 
 	// We want to grab the response from the inner message.