diff --git a/allocate.go b/allocate.go index df7ce451..7a922261 100644 --- a/allocate.go +++ b/allocate.go @@ -7,6 +7,7 @@ import ( "errors" "fmt" "io" + "net/http" "os" "os/exec" "path/filepath" @@ -526,10 +527,15 @@ func WSURLReadTimeout(t time.Duration) ExecAllocatorOption { // // Use chromedp.NoModifyURL to prevent it from modifying the url. 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, - 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 +555,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 +568,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 +590,7 @@ func (a *RemoteAllocator) Allocate(ctx context.Context, opts ...BrowserOption) ( a.wg.Done() }() - browser, err := NewBrowser(wctx, wsURL, opts...) + browser, err := NewBrowserWithHeader(wctx, wsURL, wsHeader, opts...) if err != nil { return nil, err } diff --git a/browser.go b/browser.go index 9b857d94..f9b91214 100644 --- a/browser.go +++ b/browser.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "log" + "net/http" "os" "sync" "sync/atomic" @@ -79,6 +80,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, opts ...BrowserOption) (*Browser, error) { + 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{}), closingGracefully: make(chan struct{}), @@ -109,7 +114,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 = DialContextWithHeader(dialCtx, urlstr, header, WithConnDebugf(b.dbgf)) if err != nil { return nil, fmt.Errorf("could not dial %q: %w", urlstr, err) } diff --git a/conn.go b/conn.go index edcb3af4..0a9ffc53 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,16 @@ 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) { + return DialContextWithHeader(ctx, urlstr, nil, opts...) +} + +func DialContextWithHeader(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