From b31b2122579db25ef8ded25a15d622a3fa98b3d0 Mon Sep 17 00:00:00 2001
From: tangdegang <xiaodehappy@foxmail.com>
Date: Wed, 28 May 2025 11:42:52 +0800
Subject: [PATCH] fix: connect websocket for chromium transfer error with
 Sec-WebSocket-Key nil

---
 browser.go      | 6 ++++--
 browser_test.go | 4 ++--
 must.go         | 2 +-
 3 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/browser.go b/browser.go
index bd2045f5..4952b76e 100644
--- a/browser.go
+++ b/browser.go
@@ -9,6 +9,7 @@ package rod
 
 import (
 	"context"
+	"net/http"
 	"reflect"
 	"strings"
 	"sync"
@@ -145,7 +146,8 @@ func (b *Browser) NoDefaultDevice() *Browser {
 
 // Connect to the browser and start to control it.
 // If fails to connect, try to launch a local browser, if local browser not found try to download one.
-func (b *Browser) Connect() error {
+// If connect websocket not chromium，use header and set Sec-WebSocket-Key, for example websocket for transfer chromium.
+func (b *Browser) Connect(h http.Header) error {
 	if b.client == nil {
 		u := b.controlURL
 		if u == "" {
@@ -156,7 +158,7 @@ func (b *Browser) Connect() error {
 			}
 		}
 
-		c, err := cdp.StartWithURL(b.ctx, u, nil)
+		c, err := cdp.StartWithURL(b.ctx, u, h)
 		if err != nil {
 			return err
 		}
diff --git a/browser_test.go b/browser_test.go
index 419ebd3e..ce5e25da 100644
--- a/browser_test.go
+++ b/browser_test.go
@@ -434,7 +434,7 @@ func TestBrowserConnectFailure(t *testing.T) {
 
 	c := g.Context()
 	c.Cancel()
-	err := rod.New().Context(c).Connect()
+	err := rod.New().Context(c).Connect(nil)
 	if err == nil {
 		g.Fatal("expected an error on connect failure")
 	}
@@ -447,7 +447,7 @@ func TestBrowserPool(t *testing.T) {
 
 	b, err := pool.Get(func() (*rod.Browser, error) {
 		browser := rod.New()
-		return browser, browser.Connect()
+		return browser, browser.Connect(nil)
 	})
 	g.E(err)
 	pool.Put(b)
diff --git a/must.go b/must.go
index fa940210..82ea296b 100644
--- a/must.go
+++ b/must.go
@@ -48,7 +48,7 @@ func (b *Browser) WithPanic(fail func(interface{})) *Browser {
 
 // MustConnect is similar to [Browser.Connect].
 func (b *Browser) MustConnect() *Browser {
-	b.e(b.Connect())
+	b.e(b.Connect(nil))
 	return b
 }