From 140d929add8bde65c5453184a3c706b9945eb820 Mon Sep 17 00:00:00 2001
From: Jakob Scheithe <jakob@contextflow.com>
Date: Wed, 14 Jun 2023 13:45:10 +0200
Subject: [PATCH 1/8] add DialOption for adding a header

---
 browser.go |  8 +++++++-
 conn.go    | 28 ++++++++++++++++++++++++++--
 2 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/browser.go b/browser.go
index 9b857d94..471603b6 100644
--- a/browser.go
+++ b/browser.go
@@ -42,6 +42,8 @@ type Browser struct {
 
 	dialTimeout time.Duration
 
+	dialOptions []DialOption
+
 	// pages keeps track of the attached targets, indexed by each's session
 	// ID. The only reason this is a field is so that the tests can check the
 	// map once a browser is closed.
@@ -109,7 +111,7 @@ func NewBrowser(ctx context.Context, urlstr string, opts ...BrowserOption) (*Bro
 	}
 
 	var err error
-	b.conn, err = DialContext(dialCtx, urlstr, WithConnDebugf(b.dbgf))
+	b.conn, err = DialContext(dialCtx, urlstr, append(b.dialOptions, WithConnDebugf(b.dbgf))...)
 	if err != nil {
 		return nil, fmt.Errorf("could not dial %q: %w", urlstr, err)
 	}
@@ -358,3 +360,7 @@ func WithConsolef(f func(string, ...interface{})) BrowserOption {
 func WithDialTimeout(d time.Duration) BrowserOption {
 	return func(b *Browser) { b.dialTimeout = d }
 }
+
+func WithDialOptions(dialOpts ...DialOption) BrowserOption {
+	return func(b *Browser) { b.dialOptions = append(b.dialOptions, dialOpts...) }
+}
diff --git a/conn.go b/conn.go
index edcb3af4..8c26f064 100644
--- a/conn.go
+++ b/conn.go
@@ -5,6 +5,7 @@ import (
 	"context"
 	"io"
 	"net"
+	"net/http"
 
 	"github.com/gobwas/ws"
 	"github.com/gobwas/ws/wsutil"
@@ -33,6 +34,8 @@ type Conn struct {
 	reader wsutil.Reader
 	writer wsutil.Writer
 
+	header http.Header
+
 	// reuse the easyjson structs to avoid allocs per Read/Write.
 	decoder jlexer.Lexer
 	encoder jwriter.Writer
@@ -42,8 +45,18 @@ type Conn struct {
 
 // DialContext dials the specified websocket URL using gobwas/ws.
 func DialContext(ctx context.Context, urlstr string, opts ...DialOption) (*Conn, error) {
+	// we apply options prematurely just to get the header information which we need before the connection is established
+	c := &Conn{}
+	for _, o := range opts {
+		o(c)
+	}
+
 	// connect
-	conn, br, _, err := ws.Dial(ctx, urlstr)
+	dialer := ws.Dialer{
+		Header: ws.HandshakeHeaderHTTP{c.header},
+	}
+
+	conn, br, _, err := dialer.Dial(ctx, urlstr)
 	if err != nil {
 		return nil, err
 	}
@@ -52,7 +65,7 @@ func DialContext(ctx context.Context, urlstr string, opts ...DialOption) (*Conn,
 	}
 
 	// apply opts
-	c := &Conn{
+	c = &Conn{
 		conn: conn,
 		// pass 0 to use the default initial buffer size (4KiB).
 		// github.com/gobwas/ws will grow the buffer size if needed.
@@ -146,3 +159,14 @@ func WithConnDebugf(f func(string, ...interface{})) DialOption {
 		c.dbgf = f
 	}
 }
+
+func WithDialHeader(header http.Header) DialOption {
+	return func(c *Conn) {
+		if c.header == nil {
+			c.header = make(http.Header)
+		}
+		for k, v := range header {
+			c.header[k] = v
+		}
+	}
+}

From 78f4fc3bd72f708e285ee98e40232eec55097989 Mon Sep 17 00:00:00 2001
From: Jakob Scheithe <jakob@contextflow.com>
Date: Wed, 14 Jun 2023 13:57:50 +0200
Subject: [PATCH 2/8] fix type for ws.Dialer

---
 conn.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/conn.go b/conn.go
index 8c26f064..4e72de80 100644
--- a/conn.go
+++ b/conn.go
@@ -53,7 +53,7 @@ func DialContext(ctx context.Context, urlstr string, opts ...DialOption) (*Conn,
 
 	// connect
 	dialer := ws.Dialer{
-		Header: ws.HandshakeHeaderHTTP{c.header},
+		Header: ws.HandshakeHeaderHTTP(c.header),
 	}
 
 	conn, br, _, err := dialer.Dial(ctx, urlstr)

From 42f1ca1fd3c97bd4968e03b0b2a85b08556573d6 Mon Sep 17 00:00:00 2001
From: Jakob Scheithe <jakob@contextflow.com>
Date: Wed, 14 Jun 2023 14:05:46 +0200
Subject: [PATCH 3/8] set dialHeader as extra argument to DialContext, instead
 of using a DialOption

---
 browser.go | 16 ++++++++++++----
 conn.go    | 26 +++-----------------------
 2 files changed, 15 insertions(+), 27 deletions(-)

diff --git a/browser.go b/browser.go
index 471603b6..bcd7d2dd 100644
--- a/browser.go
+++ b/browser.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"log"
+	"net/http"
 	"os"
 	"sync"
 	"sync/atomic"
@@ -42,7 +43,7 @@ type Browser struct {
 
 	dialTimeout time.Duration
 
-	dialOptions []DialOption
+	dialHeader http.Header
 
 	// pages keeps track of the attached targets, indexed by each's session
 	// ID. The only reason this is a field is so that the tests can check the
@@ -111,7 +112,7 @@ func NewBrowser(ctx context.Context, urlstr string, opts ...BrowserOption) (*Bro
 	}
 
 	var err error
-	b.conn, err = DialContext(dialCtx, urlstr, append(b.dialOptions, WithConnDebugf(b.dbgf))...)
+	b.conn, err = DialContext(dialCtx, urlstr, b.dialHeader, WithConnDebugf(b.dbgf))
 	if err != nil {
 		return nil, fmt.Errorf("could not dial %q: %w", urlstr, err)
 	}
@@ -361,6 +362,13 @@ func WithDialTimeout(d time.Duration) BrowserOption {
 	return func(b *Browser) { b.dialTimeout = d }
 }
 
-func WithDialOptions(dialOpts ...DialOption) BrowserOption {
-	return func(b *Browser) { b.dialOptions = append(b.dialOptions, dialOpts...) }
+func WithDialHeader(header http.Header) BrowserOption {
+	return func(b *Browser) {
+		if b.dialHeader == nil {
+			b.dialHeader = make(http.Header)
+		}
+		for k, v := range header {
+			b.dialHeader[k] = v
+		}
+	}
 }
diff --git a/conn.go b/conn.go
index 4e72de80..c3931207 100644
--- a/conn.go
+++ b/conn.go
@@ -34,8 +34,6 @@ type Conn struct {
 	reader wsutil.Reader
 	writer wsutil.Writer
 
-	header http.Header
-
 	// reuse the easyjson structs to avoid allocs per Read/Write.
 	decoder jlexer.Lexer
 	encoder jwriter.Writer
@@ -44,18 +42,11 @@ type Conn struct {
 }
 
 // DialContext dials the specified websocket URL using gobwas/ws.
-func DialContext(ctx context.Context, urlstr string, opts ...DialOption) (*Conn, error) {
-	// we apply options prematurely just to get the header information which we need before the connection is established
-	c := &Conn{}
-	for _, o := range opts {
-		o(c)
-	}
-
+func DialContext(ctx context.Context, urlstr string, header http.Header, opts ...DialOption) (*Conn, error) {
 	// connect
 	dialer := ws.Dialer{
-		Header: ws.HandshakeHeaderHTTP(c.header),
+		Header: ws.HandshakeHeaderHTTP(header),
 	}
-
 	conn, br, _, err := dialer.Dial(ctx, urlstr)
 	if err != nil {
 		return nil, err
@@ -65,7 +56,7 @@ func DialContext(ctx context.Context, urlstr string, opts ...DialOption) (*Conn,
 	}
 
 	// apply opts
-	c = &Conn{
+	c := &Conn{
 		conn: conn,
 		// pass 0 to use the default initial buffer size (4KiB).
 		// github.com/gobwas/ws will grow the buffer size if needed.
@@ -159,14 +150,3 @@ func WithConnDebugf(f func(string, ...interface{})) DialOption {
 		c.dbgf = f
 	}
 }
-
-func WithDialHeader(header http.Header) DialOption {
-	return func(c *Conn) {
-		if c.header == nil {
-			c.header = make(http.Header)
-		}
-		for k, v := range header {
-			c.header[k] = v
-		}
-	}
-}

From 725962f478e408e67a183e48c6601f4842fa4e03 Mon Sep 17 00:00:00 2001
From: Jakob Scheithe <jakob@contextflow.com>
Date: Wed, 14 Jun 2023 16:54:59 +0200
Subject: [PATCH 4/8] move dialHeader to RemoteAllocator to make it work with
 modifyURL

---
 allocate.go | 25 +++++++++++++++++++++++++
 browser.go  |  8 +++++---
 conn.go     |  1 +
 util.go     |  2 ++
 4 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/allocate.go b/allocate.go
index df7ce451..e42c9925 100644
--- a/allocate.go
+++ b/allocate.go
@@ -7,6 +7,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"net/http"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -30,6 +31,8 @@ type Allocator interface {
 	// Cancelling the allocator context will already perform this operation,
 	// so normally there's no need to call Wait directly.
 	Wait()
+
+	getDialHeader() http.Header
 }
 
 // setupExecAllocator is similar to NewExecAllocator, but it allows NewContext
@@ -327,6 +330,10 @@ func (a *ExecAllocator) Wait() {
 	a.wg.Wait()
 }
 
+func (a *ExecAllocator) getDialHeader() http.Header {
+	return nil
+}
+
 // ExecPath returns an ExecAllocatorOption which uses the given path to execute
 // browser processes. The given path can be an absolute path to a binary, or
 // just the name of the program to find via exec.LookPath.
@@ -552,6 +559,8 @@ type RemoteAllocator struct {
 	modifyURLFunc func(ctx context.Context, wsURL string) (string, error)
 
 	wg sync.WaitGroup
+
+	dialHeader http.Header
 }
 
 // Allocate satisfies the Allocator interface.
@@ -583,6 +592,7 @@ func (a *RemoteAllocator) Allocate(ctx context.Context, opts ...BrowserOption) (
 		a.wg.Done()
 	}()
 
+	opts = append(opts, WithDialHeaderBrowser(FromContext(ctx).Allocator.getDialHeader()))
 	browser, err := NewBrowser(wctx, wsURL, opts...)
 	if err != nil {
 		return nil, err
@@ -605,8 +615,23 @@ func (a *RemoteAllocator) Wait() {
 	a.wg.Wait()
 }
 
+func (a *RemoteAllocator) getDialHeader() http.Header {
+	return a.dialHeader
+}
+
 // NoModifyURL is a RemoteAllocatorOption that prevents the remote allocator
 // from modifying the websocket debugger URL passed to it.
 func NoModifyURL(a *RemoteAllocator) {
 	a.modifyURLFunc = nil
 }
+
+func WithDialHeader(h http.Header) RemoteAllocatorOption {
+	return func(a *RemoteAllocator) {
+		if a.dialHeader == nil {
+			a.dialHeader = make(http.Header)
+		}
+		for k, v := range h {
+			a.dialHeader[k] = v
+		}
+	}
+}
diff --git a/browser.go b/browser.go
index bcd7d2dd..d98581cb 100644
--- a/browser.go
+++ b/browser.go
@@ -42,8 +42,7 @@ type Browser struct {
 	closingGracefully chan struct{}
 
 	dialTimeout time.Duration
-
-	dialHeader http.Header
+	dialHeader  http.Header
 
 	// pages keeps track of the attached targets, indexed by each's session
 	// ID. The only reason this is a field is so that the tests can check the
@@ -362,7 +361,10 @@ func WithDialTimeout(d time.Duration) BrowserOption {
 	return func(b *Browser) { b.dialTimeout = d }
 }
 
-func WithDialHeader(header http.Header) BrowserOption {
+func WithDialHeaderBrowser(header http.Header) BrowserOption {
+	if header == nil {
+		return func(b *Browser) {}
+	}
 	return func(b *Browser) {
 		if b.dialHeader == nil {
 			b.dialHeader = make(http.Header)
diff --git a/conn.go b/conn.go
index c3931207..7ce60265 100644
--- a/conn.go
+++ b/conn.go
@@ -44,6 +44,7 @@ type Conn struct {
 // DialContext dials the specified websocket URL using gobwas/ws.
 func DialContext(ctx context.Context, urlstr string, header http.Header, opts ...DialOption) (*Conn, error) {
 	// connect
+	// h := FromContext(ctx).Allocator.getDialHeader()
 	dialer := ws.Dialer{
 		Header: ws.HandshakeHeaderHTTP(header),
 	}
diff --git a/util.go b/util.go
index adcd59d9..b07b53dc 100644
--- a/util.go
+++ b/util.go
@@ -95,6 +95,8 @@ func modifyURL(ctx context.Context, urlstr string) (string, error) {
 
 	// to get "webSocketDebuggerUrl" in the response
 	req, err := http.NewRequestWithContext(lctx, "GET", u.String(), nil)
+	req.Header = FromContext(ctx).Allocator.getDialHeader()
+
 	if err != nil {
 		return "", err
 	}

From daf717fcdd962c5d3151a2f1d36b02f6440ba4a5 Mon Sep 17 00:00:00 2001
From: Jakob Scheithe <jakob@contextflow.com>
Date: Thu, 15 Jun 2023 13:14:46 +0200
Subject: [PATCH 5/8] revert to 6e7da0b

---
 allocate.go | 25 -------------------------
 browser.go  | 18 +-----------------
 conn.go     |  9 ++-------
 util.go     |  2 --
 4 files changed, 3 insertions(+), 51 deletions(-)

diff --git a/allocate.go b/allocate.go
index e42c9925..df7ce451 100644
--- a/allocate.go
+++ b/allocate.go
@@ -7,7 +7,6 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"net/http"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -31,8 +30,6 @@ type Allocator interface {
 	// Cancelling the allocator context will already perform this operation,
 	// so normally there's no need to call Wait directly.
 	Wait()
-
-	getDialHeader() http.Header
 }
 
 // setupExecAllocator is similar to NewExecAllocator, but it allows NewContext
@@ -330,10 +327,6 @@ func (a *ExecAllocator) Wait() {
 	a.wg.Wait()
 }
 
-func (a *ExecAllocator) getDialHeader() http.Header {
-	return nil
-}
-
 // ExecPath returns an ExecAllocatorOption which uses the given path to execute
 // browser processes. The given path can be an absolute path to a binary, or
 // just the name of the program to find via exec.LookPath.
@@ -559,8 +552,6 @@ type RemoteAllocator struct {
 	modifyURLFunc func(ctx context.Context, wsURL string) (string, error)
 
 	wg sync.WaitGroup
-
-	dialHeader http.Header
 }
 
 // Allocate satisfies the Allocator interface.
@@ -592,7 +583,6 @@ func (a *RemoteAllocator) Allocate(ctx context.Context, opts ...BrowserOption) (
 		a.wg.Done()
 	}()
 
-	opts = append(opts, WithDialHeaderBrowser(FromContext(ctx).Allocator.getDialHeader()))
 	browser, err := NewBrowser(wctx, wsURL, opts...)
 	if err != nil {
 		return nil, err
@@ -615,23 +605,8 @@ func (a *RemoteAllocator) Wait() {
 	a.wg.Wait()
 }
 
-func (a *RemoteAllocator) getDialHeader() http.Header {
-	return a.dialHeader
-}
-
 // NoModifyURL is a RemoteAllocatorOption that prevents the remote allocator
 // from modifying the websocket debugger URL passed to it.
 func NoModifyURL(a *RemoteAllocator) {
 	a.modifyURLFunc = nil
 }
-
-func WithDialHeader(h http.Header) RemoteAllocatorOption {
-	return func(a *RemoteAllocator) {
-		if a.dialHeader == nil {
-			a.dialHeader = make(http.Header)
-		}
-		for k, v := range h {
-			a.dialHeader[k] = v
-		}
-	}
-}
diff --git a/browser.go b/browser.go
index d98581cb..9b857d94 100644
--- a/browser.go
+++ b/browser.go
@@ -5,7 +5,6 @@ import (
 	"errors"
 	"fmt"
 	"log"
-	"net/http"
 	"os"
 	"sync"
 	"sync/atomic"
@@ -42,7 +41,6 @@ type Browser struct {
 	closingGracefully chan struct{}
 
 	dialTimeout time.Duration
-	dialHeader  http.Header
 
 	// pages keeps track of the attached targets, indexed by each's session
 	// ID. The only reason this is a field is so that the tests can check the
@@ -111,7 +109,7 @@ func NewBrowser(ctx context.Context, urlstr string, opts ...BrowserOption) (*Bro
 	}
 
 	var err error
-	b.conn, err = DialContext(dialCtx, urlstr, b.dialHeader, WithConnDebugf(b.dbgf))
+	b.conn, err = DialContext(dialCtx, urlstr, WithConnDebugf(b.dbgf))
 	if err != nil {
 		return nil, fmt.Errorf("could not dial %q: %w", urlstr, err)
 	}
@@ -360,17 +358,3 @@ func WithConsolef(f func(string, ...interface{})) BrowserOption {
 func WithDialTimeout(d time.Duration) BrowserOption {
 	return func(b *Browser) { b.dialTimeout = d }
 }
-
-func WithDialHeaderBrowser(header http.Header) BrowserOption {
-	if header == nil {
-		return func(b *Browser) {}
-	}
-	return func(b *Browser) {
-		if b.dialHeader == nil {
-			b.dialHeader = make(http.Header)
-		}
-		for k, v := range header {
-			b.dialHeader[k] = v
-		}
-	}
-}
diff --git a/conn.go b/conn.go
index 7ce60265..edcb3af4 100644
--- a/conn.go
+++ b/conn.go
@@ -5,7 +5,6 @@ import (
 	"context"
 	"io"
 	"net"
-	"net/http"
 
 	"github.com/gobwas/ws"
 	"github.com/gobwas/ws/wsutil"
@@ -42,13 +41,9 @@ type Conn struct {
 }
 
 // DialContext dials the specified websocket URL using gobwas/ws.
-func DialContext(ctx context.Context, urlstr string, header http.Header, opts ...DialOption) (*Conn, error) {
+func DialContext(ctx context.Context, urlstr string, opts ...DialOption) (*Conn, error) {
 	// connect
-	// h := FromContext(ctx).Allocator.getDialHeader()
-	dialer := ws.Dialer{
-		Header: ws.HandshakeHeaderHTTP(header),
-	}
-	conn, br, _, err := dialer.Dial(ctx, urlstr)
+	conn, br, _, err := ws.Dial(ctx, urlstr)
 	if err != nil {
 		return nil, err
 	}
diff --git a/util.go b/util.go
index b07b53dc..adcd59d9 100644
--- a/util.go
+++ b/util.go
@@ -95,8 +95,6 @@ func modifyURL(ctx context.Context, urlstr string) (string, error) {
 
 	// to get "webSocketDebuggerUrl" in the response
 	req, err := http.NewRequestWithContext(lctx, "GET", u.String(), nil)
-	req.Header = FromContext(ctx).Allocator.getDialHeader()
-
 	if err != nil {
 		return "", err
 	}

From f1b0bc88a0d3f64e78735f6b6ceb323f6afb1000 Mon Sep 17 00:00:00 2001
From: Jakob Scheithe <jakob@contextflow.com>
Date: Thu, 15 Jun 2023 13:32:38 +0200
Subject: [PATCH 6/8] pass header as paramters and add as a field in
 RemoteAllocator

---
 allocate.go      | 21 ++++++++++++---------
 allocate_test.go |  2 +-
 browser.go       |  5 +++--
 chromedp_test.go |  8 ++++----
 conn.go          |  8 ++++++--
 util.go          |  3 ++-
 6 files changed, 28 insertions(+), 19 deletions(-)

diff --git a/allocate.go b/allocate.go
index df7ce451..889492cb 100644
--- a/allocate.go
+++ b/allocate.go
@@ -7,6 +7,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"net/http"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -256,7 +257,7 @@ func (a *ExecAllocator) Allocate(ctx context.Context, opts ...BrowserOption) (*B
 		return nil, err
 	}
 
-	browser, err := NewBrowser(ctx, wsURL, opts...)
+	browser, err := NewBrowser(ctx, wsURL, nil, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -525,11 +526,12 @@ func WSURLReadTimeout(t time.Duration) ExecAllocatorOption {
 // Because the allocator won't try to modify it and it's obviously invalid.
 //
 // Use chromedp.NoModifyURL to prevent it from modifying the url.
-func NewRemoteAllocator(parent context.Context, url string, opts ...RemoteAllocatorOption) (context.Context, context.CancelFunc) {
+func NewRemoteAllocator(parent context.Context, url string, header http.Header, opts ...RemoteAllocatorOption) (context.Context, context.CancelFunc) {
 	a := &RemoteAllocator{
-		wsURL: url,
-		modifyURLFunc: func(ctx context.Context, wsURL string) (string, error) {
-			return modifyURL(ctx, wsURL)
+		wsURL:    url,
+		wsHeader: header,
+		modifyURLFunc: func(ctx context.Context, wsURL string, wsHeader http.Header) (string, error) {
+			return modifyURL(ctx, wsURL, wsHeader)
 		},
 	}
 	for _, o := range opts {
@@ -549,7 +551,8 @@ type RemoteAllocatorOption = func(*RemoteAllocator)
 // process via a websocket URL.
 type RemoteAllocator struct {
 	wsURL         string
-	modifyURLFunc func(ctx context.Context, wsURL string) (string, error)
+	wsHeader      http.Header
+	modifyURLFunc func(ctx context.Context, wsURL string, wsHeader http.Header) (string, error)
 
 	wg sync.WaitGroup
 }
@@ -561,10 +564,10 @@ func (a *RemoteAllocator) Allocate(ctx context.Context, opts ...BrowserOption) (
 		return nil, ErrInvalidContext
 	}
 
-	wsURL := a.wsURL
+	wsURL, wsHeader := a.wsURL, a.wsHeader
 	var err error
 	if a.modifyURLFunc != nil {
-		wsURL, err = a.modifyURLFunc(ctx, wsURL)
+		wsURL, err = a.modifyURLFunc(ctx, wsURL, wsHeader)
 		if err != nil {
 			return nil, fmt.Errorf("failed to modify wsURL: %w", err)
 		}
@@ -583,7 +586,7 @@ func (a *RemoteAllocator) Allocate(ctx context.Context, opts ...BrowserOption) (
 		a.wg.Done()
 	}()
 
-	browser, err := NewBrowser(wctx, wsURL, opts...)
+	browser, err := NewBrowser(wctx, wsURL, wsHeader, opts...)
 	if err != nil {
 		return nil, err
 	}
diff --git a/allocate_test.go b/allocate_test.go
index 99098ec2..11030eb4 100644
--- a/allocate_test.go
+++ b/allocate_test.go
@@ -215,7 +215,7 @@ func testRemoteAllocator(t *testing.T, modifyURL func(wsURL string) string, want
 	if err != nil {
 		t.Fatal(err)
 	}
-	allocCtx, allocCancel := NewRemoteAllocator(context.Background(), modifyURL(wsURL), opts...)
+	allocCtx, allocCancel := NewRemoteAllocator(context.Background(), modifyURL(wsURL), nil, opts...)
 	defer allocCancel()
 
 	taskCtx, taskCancel := NewContext(allocCtx,
diff --git a/browser.go b/browser.go
index 9b857d94..dab60c6d 100644
--- a/browser.go
+++ b/browser.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"log"
+	"net/http"
 	"os"
 	"sync"
 	"sync/atomic"
@@ -78,7 +79,7 @@ type Browser struct {
 
 // NewBrowser creates a new browser. Typically, this function wouldn't be called
 // directly, as the Allocator interface takes care of it.
-func NewBrowser(ctx context.Context, urlstr string, opts ...BrowserOption) (*Browser, error) {
+func NewBrowser(ctx context.Context, urlstr string, header http.Header, opts ...BrowserOption) (*Browser, error) {
 	b := &Browser{
 		LostConnection:    make(chan struct{}),
 		closingGracefully: make(chan struct{}),
@@ -109,7 +110,7 @@ func NewBrowser(ctx context.Context, urlstr string, opts ...BrowserOption) (*Bro
 	}
 
 	var err error
-	b.conn, err = DialContext(dialCtx, urlstr, WithConnDebugf(b.dbgf))
+	b.conn, err = DialContext(dialCtx, urlstr, header, WithConnDebugf(b.dbgf))
 	if err != nil {
 		return nil, fmt.Errorf("could not dial %q: %w", urlstr, err)
 	}
diff --git a/chromedp_test.go b/chromedp_test.go
index 70fbaf7b..ffdfbe54 100644
--- a/chromedp_test.go
+++ b/chromedp_test.go
@@ -493,7 +493,7 @@ func TestDialTimeout(t *testing.T) {
 
 		ctx, cancel := context.WithCancel(context.Background())
 		defer cancel()
-		_, err = NewBrowser(ctx, url, WithDialTimeout(time.Microsecond))
+		_, err = NewBrowser(ctx, url, nil, WithDialTimeout(time.Microsecond))
 		got, want := fmt.Sprintf("%v", err), "i/o timeout"
 		if !strings.Contains(got, want) {
 			t.Fatalf("got %q, want %q", got, want)
@@ -516,7 +516,7 @@ func TestDialTimeout(t *testing.T) {
 
 		ctx, cancel := context.WithCancel(context.Background())
 		defer cancel()
-		_, err = NewBrowser(ctx, url, WithDialTimeout(0))
+		_, err = NewBrowser(ctx, url, nil, WithDialTimeout(0))
 		got := fmt.Sprintf("%v", err)
 		if !strings.Contains(got, "EOF") && !strings.Contains(got, "connection reset") {
 			t.Fatalf("got %q, want %q or %q", got, "EOF", "connection reset")
@@ -820,7 +820,7 @@ func TestBrowserContext(t *testing.T) {
 				if conn == nil {
 					t.Skip("skip when the remote debugging address is not available")
 				}
-				actx, _ := NewRemoteAllocator(context.Background(), "ws://"+conn.RemoteAddr().String())
+				actx, _ := NewRemoteAllocator(context.Background(), "ws://"+conn.RemoteAddr().String(), nil)
 				ctx, cancel := NewContext(actx, WithExistingBrowserContext(rootBrowserContextID1))
 				if err := Run(ctx); err != nil {
 					t.Fatal(err)
@@ -842,7 +842,7 @@ func TestBrowserContext(t *testing.T) {
 				if conn == nil {
 					t.Skip("skip when the remote debugging address is not available")
 				}
-				actx, _ := NewRemoteAllocator(context.Background(), "ws://"+conn.RemoteAddr().String())
+				actx, _ := NewRemoteAllocator(context.Background(), "ws://"+conn.RemoteAddr().String(), nil)
 				ctx, cancel := NewContext(actx, WithNewBrowserContext())
 				if err := Run(ctx); err != nil {
 					t.Fatal(err)
diff --git a/conn.go b/conn.go
index edcb3af4..c3931207 100644
--- a/conn.go
+++ b/conn.go
@@ -5,6 +5,7 @@ import (
 	"context"
 	"io"
 	"net"
+	"net/http"
 
 	"github.com/gobwas/ws"
 	"github.com/gobwas/ws/wsutil"
@@ -41,9 +42,12 @@ type Conn struct {
 }
 
 // DialContext dials the specified websocket URL using gobwas/ws.
-func DialContext(ctx context.Context, urlstr string, opts ...DialOption) (*Conn, error) {
+func DialContext(ctx context.Context, urlstr string, header http.Header, opts ...DialOption) (*Conn, error) {
 	// connect
-	conn, br, _, err := ws.Dial(ctx, urlstr)
+	dialer := ws.Dialer{
+		Header: ws.HandshakeHeaderHTTP(header),
+	}
+	conn, br, _, err := dialer.Dial(ctx, urlstr)
 	if err != nil {
 		return nil, err
 	}
diff --git a/util.go b/util.go
index adcd59d9..b6155b2f 100644
--- a/util.go
+++ b/util.go
@@ -67,7 +67,7 @@ func resolveHost(ctx context.Context, host string) (string, error) {
 //   - ws://127.0.0.1:9222/
 //   - http://127.0.0.1:9222/
 //   - http://container-name:9222/
-func modifyURL(ctx context.Context, urlstr string) (string, error) {
+func modifyURL(ctx context.Context, urlstr string, header http.Header) (string, error) {
 	lctx, cancel := context.WithTimeout(ctx, 20*time.Second)
 	defer cancel()
 
@@ -98,6 +98,7 @@ func modifyURL(ctx context.Context, urlstr string) (string, error) {
 	if err != nil {
 		return "", err
 	}
+	req.Header = header
 	resp, err := http.DefaultClient.Do(req)
 	if err != nil {
 		return "", err

From 2e53780d57358e53a3f11496df02830f9413fe2c Mon Sep 17 00:00:00 2001
From: Jakob Scheithe <jakob@contextflow.com>
Date: Mon, 24 Jul 2023 11:21:55 +0200
Subject: [PATCH 7/8] add ...WithHeader functions for NewBrowser and
 NewRemoteAllocator to avoid changing the public API

---
 allocate.go      | 10 +++++++---
 allocate_test.go |  2 +-
 browser.go       |  5 ++++-
 chromedp_test.go |  4 ++--
 4 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/allocate.go b/allocate.go
index 889492cb..7a922261 100644
--- a/allocate.go
+++ b/allocate.go
@@ -257,7 +257,7 @@ func (a *ExecAllocator) Allocate(ctx context.Context, opts ...BrowserOption) (*B
 		return nil, err
 	}
 
-	browser, err := NewBrowser(ctx, wsURL, nil, opts...)
+	browser, err := NewBrowser(ctx, wsURL, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -526,7 +526,11 @@ func WSURLReadTimeout(t time.Duration) ExecAllocatorOption {
 // Because the allocator won't try to modify it and it's obviously invalid.
 //
 // Use chromedp.NoModifyURL to prevent it from modifying the url.
-func NewRemoteAllocator(parent context.Context, url string, header http.Header, opts ...RemoteAllocatorOption) (context.Context, context.CancelFunc) {
+func NewRemoteAllocator(parent context.Context, url string, opts ...RemoteAllocatorOption) (context.Context, context.CancelFunc) {
+	return NewRemoteAllocatorWithHeader(parent, url, nil, opts...)
+}
+
+func NewRemoteAllocatorWithHeader(parent context.Context, url string, header http.Header, opts ...RemoteAllocatorOption) (context.Context, context.CancelFunc) {
 	a := &RemoteAllocator{
 		wsURL:    url,
 		wsHeader: header,
@@ -586,7 +590,7 @@ func (a *RemoteAllocator) Allocate(ctx context.Context, opts ...BrowserOption) (
 		a.wg.Done()
 	}()
 
-	browser, err := NewBrowser(wctx, wsURL, wsHeader, opts...)
+	browser, err := NewBrowserWithHeader(wctx, wsURL, wsHeader, opts...)
 	if err != nil {
 		return nil, err
 	}
diff --git a/allocate_test.go b/allocate_test.go
index 11030eb4..99098ec2 100644
--- a/allocate_test.go
+++ b/allocate_test.go
@@ -215,7 +215,7 @@ func testRemoteAllocator(t *testing.T, modifyURL func(wsURL string) string, want
 	if err != nil {
 		t.Fatal(err)
 	}
-	allocCtx, allocCancel := NewRemoteAllocator(context.Background(), modifyURL(wsURL), nil, opts...)
+	allocCtx, allocCancel := NewRemoteAllocator(context.Background(), modifyURL(wsURL), opts...)
 	defer allocCancel()
 
 	taskCtx, taskCancel := NewContext(allocCtx,
diff --git a/browser.go b/browser.go
index dab60c6d..e63cc5ec 100644
--- a/browser.go
+++ b/browser.go
@@ -79,7 +79,10 @@ type Browser struct {
 
 // NewBrowser creates a new browser. Typically, this function wouldn't be called
 // directly, as the Allocator interface takes care of it.
-func NewBrowser(ctx context.Context, urlstr string, header http.Header, opts ...BrowserOption) (*Browser, error) {
+func NewBrowser(ctx context.Context, urlstr string, opts ...BrowserOption) (*Browser, error) {
+	NewBrowserWithHeader(ctx, urlstr, nil, opts...)
+}
+func NewBrowserWithHeader(ctx context.Context, urlstr string, header http.Header, opts ...BrowserOption) (*Browser, error) {
 	b := &Browser{
 		LostConnection:    make(chan struct{}),
 		closingGracefully: make(chan struct{}),
diff --git a/chromedp_test.go b/chromedp_test.go
index ffdfbe54..888cbccc 100644
--- a/chromedp_test.go
+++ b/chromedp_test.go
@@ -493,7 +493,7 @@ func TestDialTimeout(t *testing.T) {
 
 		ctx, cancel := context.WithCancel(context.Background())
 		defer cancel()
-		_, err = NewBrowser(ctx, url, nil, WithDialTimeout(time.Microsecond))
+		_, err = NewBrowser(ctx, url, WithDialTimeout(time.Microsecond))
 		got, want := fmt.Sprintf("%v", err), "i/o timeout"
 		if !strings.Contains(got, want) {
 			t.Fatalf("got %q, want %q", got, want)
@@ -516,7 +516,7 @@ func TestDialTimeout(t *testing.T) {
 
 		ctx, cancel := context.WithCancel(context.Background())
 		defer cancel()
-		_, err = NewBrowser(ctx, url, nil, WithDialTimeout(0))
+		_, err = NewBrowser(ctx, url, WithDialTimeout(0))
 		got := fmt.Sprintf("%v", err)
 		if !strings.Contains(got, "EOF") && !strings.Contains(got, "connection reset") {
 			t.Fatalf("got %q, want %q or %q", got, "EOF", "connection reset")

From a08d36517d2f1dcee5bcb79a9d5c66d1e2ba8940 Mon Sep 17 00:00:00 2001
From: Jakob Scheithe <jakob@contextflow.com>
Date: Mon, 24 Jul 2023 13:13:46 +0200
Subject: [PATCH 8/8] add ...WithHeader function for DialContext

---
 browser.go       | 5 +++--
 chromedp_test.go | 4 ++--
 conn.go          | 4 ++++
 3 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/browser.go b/browser.go
index e63cc5ec..f9b91214 100644
--- a/browser.go
+++ b/browser.go
@@ -80,8 +80,9 @@ type Browser struct {
 // NewBrowser creates a new browser. Typically, this function wouldn't be called
 // directly, as the Allocator interface takes care of it.
 func NewBrowser(ctx context.Context, urlstr string, opts ...BrowserOption) (*Browser, error) {
-	NewBrowserWithHeader(ctx, urlstr, nil, opts...)
+	return NewBrowserWithHeader(ctx, urlstr, nil, opts...)
 }
+
 func NewBrowserWithHeader(ctx context.Context, urlstr string, header http.Header, opts ...BrowserOption) (*Browser, error) {
 	b := &Browser{
 		LostConnection:    make(chan struct{}),
@@ -113,7 +114,7 @@ func NewBrowserWithHeader(ctx context.Context, urlstr string, header http.Header
 	}
 
 	var err error
-	b.conn, err = DialContext(dialCtx, urlstr, header, WithConnDebugf(b.dbgf))
+	b.conn, err = DialContextWithHeader(dialCtx, urlstr, header, WithConnDebugf(b.dbgf))
 	if err != nil {
 		return nil, fmt.Errorf("could not dial %q: %w", urlstr, err)
 	}
diff --git a/chromedp_test.go b/chromedp_test.go
index 888cbccc..70fbaf7b 100644
--- a/chromedp_test.go
+++ b/chromedp_test.go
@@ -820,7 +820,7 @@ func TestBrowserContext(t *testing.T) {
 				if conn == nil {
 					t.Skip("skip when the remote debugging address is not available")
 				}
-				actx, _ := NewRemoteAllocator(context.Background(), "ws://"+conn.RemoteAddr().String(), nil)
+				actx, _ := NewRemoteAllocator(context.Background(), "ws://"+conn.RemoteAddr().String())
 				ctx, cancel := NewContext(actx, WithExistingBrowserContext(rootBrowserContextID1))
 				if err := Run(ctx); err != nil {
 					t.Fatal(err)
@@ -842,7 +842,7 @@ func TestBrowserContext(t *testing.T) {
 				if conn == nil {
 					t.Skip("skip when the remote debugging address is not available")
 				}
-				actx, _ := NewRemoteAllocator(context.Background(), "ws://"+conn.RemoteAddr().String(), nil)
+				actx, _ := NewRemoteAllocator(context.Background(), "ws://"+conn.RemoteAddr().String())
 				ctx, cancel := NewContext(actx, WithNewBrowserContext())
 				if err := Run(ctx); err != nil {
 					t.Fatal(err)
diff --git a/conn.go b/conn.go
index c3931207..0a9ffc53 100644
--- a/conn.go
+++ b/conn.go
@@ -43,6 +43,10 @@ type Conn struct {
 
 // DialContext dials the specified websocket URL using gobwas/ws.
 func DialContext(ctx context.Context, urlstr string, header http.Header, opts ...DialOption) (*Conn, error) {
+	return DialContextWithHeader(ctx, urlstr, nil, opts...)
+}
+
+func DialContextWithHeader(ctx context.Context, urlstr string, header http.Header, opts ...DialOption) (*Conn, error) {
 	// connect
 	dialer := ws.Dialer{
 		Header: ws.HandshakeHeaderHTTP(header),