introduce a remote allocator option: NoModifyURL

allocate.go

@@ -95,7 +95,7 @@ func NewExecAllocator(parent context.Context, opts ...ExecAllocatorOption) (cont
 	return ctx, cancelWait
 }
 
-// ExecAllocatorOption is a exec allocator option.
+// ExecAllocatorOption is an exec allocator option.
 type ExecAllocatorOption = func(*ExecAllocator)
 
 // ExecAllocator is an Allocator which starts new browser processes on the host
@@ -510,6 +510,7 @@ func WSURLReadTimeout(t time.Duration) ExecAllocatorOption {
 // NewRemoteAllocator creates a new context set up with a RemoteAllocator,
 // suitable for use with NewContext. The url should point to the browser's
 // websocket address, such as "ws://127.0.0.1:$PORT/devtools/browser/...".
+//
 // If the url does not contain "/devtools/browser/", it will try to detect
 // the correct one by sending a request to "http://$HOST:$PORT/json/version".
 //
@@ -518,21 +519,34 @@ func WSURLReadTimeout(t time.Duration) ExecAllocatorOption {
 // * http://127.0.0.1:9222/
 //
 // But "ws://127.0.0.1:9222/devtools/browser/" are not accepted.
-// Because it contains "/devtools/browser/" and will be considered
-// as a valid websocket debugger URL.
-func NewRemoteAllocator(parent context.Context, url string) (context.Context, context.CancelFunc) {
+// 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) {
+	a := &RemoteAllocator{
+		wsURL: url,
+		modifyURLFunc: func(ctx context.Context, wsURL string) (string, error) {
+			return modifyURL(ctx, wsURL)
+		},
+	}
+	for _, o := range opts {
+		o(a)
+	}
+	c := &Context{Allocator: a}
+
 	ctx, cancel := context.WithCancel(parent)
-	c := &Context{Allocator: &RemoteAllocator{
-		wsURL: detectURL(url),
-	}}
 	ctx = context.WithValue(ctx, contextKey{}, c)
 	return ctx, cancel
 }
 
+// RemoteAllocatorOption is a remote allocator option.
+type RemoteAllocatorOption = func(*RemoteAllocator)
+
 // RemoteAllocator is an Allocator which connects to an already running Chrome
 // process via a websocket URL.
 type RemoteAllocator struct {
-	wsURL string
+	wsURL         string
+	modifyURLFunc func(ctx context.Context, wsURL string) (string, error)
 
 	wg sync.WaitGroup
 }
@@ -544,6 +558,15 @@ func (a *RemoteAllocator) Allocate(ctx context.Context, opts ...BrowserOption) (
 		return nil, ErrInvalidContext
 	}
 
+	wsURL := a.wsURL
+	var err error
+	if a.modifyURLFunc != nil {
+		wsURL, err = a.modifyURLFunc(ctx, wsURL)
+		if err != nil {
+			return nil, fmt.Errorf("failed to modify wsURL: %w", err)
+		}
+	}
+
 	// Use a different context for the websocket, so we can have a chance at
 	// closing the relevant pages before closing the websocket connection.
 	wctx, cancel := context.WithCancel(context.Background())
@@ -556,7 +579,8 @@ func (a *RemoteAllocator) Allocate(ctx context.Context, opts ...BrowserOption) (
 		cancel()    // close the websocket connection
 		a.wg.Done()
 	}()
-	browser, err := NewBrowser(wctx, a.wsURL, opts...)
+
+	browser, err := NewBrowser(wctx, wsURL, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -577,3 +601,9 @@ func (a *RemoteAllocator) Allocate(ctx context.Context, opts ...BrowserOption) (
 func (a *RemoteAllocator) Wait() {
 	a.wg.Wait()
 }
+
+// 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
+}

allocate_test.go

@@ -127,6 +127,8 @@ func TestRemoteAllocator(t *testing.T) {
 	tests := []struct {
 		name      string
 		modifyURL func(wsURL string) string
+		opts      []RemoteAllocatorOption
+		wantErr   string
 	}{
 		{
 			name:      "original wsURL",
@@ -164,15 +166,24 @@ func TestRemoteAllocator(t *testing.T) {
 				return u.String()
 			},
 		},
+		{
+			name: "NoModifyURL",
+			modifyURL: func(wsURL string) string {
+				return wsURL[0:strings.Index(wsURL, "devtools")]
+			},
+			opts:    []RemoteAllocatorOption{NoModifyURL},
+			wantErr: "could not dial",
+		},
 	}
 	for _, tt := range tests {
+		tt := tt
 		t.Run(tt.name, func(t *testing.T) {
-			testRemoteAllocator(t, tt.modifyURL)
+			testRemoteAllocator(t, tt.modifyURL, tt.wantErr, tt.opts)
 		})
 	}
 }
 
-func testRemoteAllocator(t *testing.T, modifyURL func(wsURL string) string) {
+func testRemoteAllocator(t *testing.T, modifyURL func(wsURL string) string, wantErr string, opts []RemoteAllocatorOption) {
 	tempDir := t.TempDir()
 
 	procCtx, cancel := context.WithCancel(context.Background())
@@ -204,7 +215,7 @@ func testRemoteAllocator(t *testing.T, modifyURL func(wsURL string) string) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	allocCtx, allocCancel := NewRemoteAllocator(context.Background(), modifyURL(wsURL))
+	allocCtx, allocCancel := NewRemoteAllocator(context.Background(), modifyURL(wsURL), opts...)
 	defer allocCancel()
 
 	taskCtx, taskCancel := NewContext(allocCtx,
@@ -214,6 +225,12 @@ func testRemoteAllocator(t *testing.T, modifyURL func(wsURL string) string) {
 
 	{
 		infos, err := Targets(taskCtx)
+		if len(wantErr) > 0 {
+			if err == nil || !strings.Contains(err.Error(), wantErr) {
+				t.Fatalf("\ngot error:\n\t%v\nwant error contains:\n\t%s", err, wantErr)
+			}
+			return
+		}
 		if err != nil {
 			t.Fatal(err)
 		}

browser.go

@@ -109,7 +109,6 @@ func NewBrowser(ctx context.Context, urlstr string, opts ...BrowserOption) (*Bro
 	}
 
 	var err error
-	urlstr = forceIP(urlstr)
 	b.conn, err = DialContext(dialCtx, urlstr, WithConnDebugf(b.dbgf))
 	if err != nil {
 		return nil, fmt.Errorf("could not dial %q: %w", urlstr, err)

util.go

@@ -1,11 +1,13 @@
 package chromedp
 
 import (
+	"context"
 	"encoding/json"
 	"net"
 	"net/http"
 	"net/url"
 	"strings"
+	"time"
 
 	"github.com/chromedp/cdproto"
 	"github.com/chromedp/cdproto/cdp"
@@ -15,62 +17,106 @@ import (
 //
 // Since Chrome 66+, Chrome DevTools Protocol clients connecting to a browser
 // must send the "Host:" header as either an IP address, or "localhost".
-func forceIP(urlstr string) string {
+// See https://github.com/chromium/chromium/commit/0e914b95f7cae6e8238e4e9075f248f801c686e6.
+func forceIP(ctx context.Context, urlstr string) (string, error) {
 	u, err := url.Parse(urlstr)
 	if err != nil {
-		return urlstr
+		return "", err
 	}
 	host, port, err := net.SplitHostPort(u.Host)
 	if err != nil {
-		return urlstr
+		return "", err
 	}
-	addr, err := net.ResolveIPAddr("ip", host)
+	host, err = resolveHost(ctx, host)
 	if err != nil {
-		return urlstr
+		return "", err
 	}
-	u.Host = net.JoinHostPort(addr.IP.String(), port)
-	return u.String()
+	u.Host = net.JoinHostPort(host, port)
+	return u.String(), nil
 }
 
-// detectURL detects the websocket debugger URL if the provided URL is not a
+// resolveHost tries to resolve a host to be an IP address. If the host is
+// an IP address or "localhost", it returns the host directly.
+func resolveHost(ctx context.Context, host string) (string, error) {
+	if host == "localhost" {
+		return host, nil
+	}
+	ip := net.ParseIP(host)
+	if ip != nil {
+		return host, nil
+	}
+
+	addrs, err := net.DefaultResolver.LookupIPAddr(ctx, host)
+	if err != nil {
+		return "", err
+	}
+
+	return addrs[0].IP.String(), nil
+}
+
+// modifyURL modifies the websocket debugger URL if the provided URL is not a
 // valid websocket debugger URL.
 //
-// A valid websocket debugger URL is something like:
-// ws://127.0.0.1:9222/devtools/browser/...
-// The original URL with the following formats are accepted:
-// * ws://127.0.0.1:9222/
-// * http://127.0.0.1:9222/
-func detectURL(urlstr string) string {
+// A websocket debugger URL containing "/devtools/browser/" are considered
+// valid. In this case, urlstr will only be modified by forceIP.
+//
+// Otherwise, it will construct a URL like http://[host]:[port]/json/version
+// and query the valid websocket debugger URL from this endpoint. The [host]
+// and [port] are parsed from the urlstr. If the host component is not an IP,
+// it will be resolved to an IP first. Example parameters:
+//   - 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) {
+	lctx, cancel := context.WithTimeout(ctx, 20*time.Second)
+	defer cancel()
+
 	if strings.Contains(urlstr, "/devtools/browser/") {
-		return urlstr
+		return forceIP(lctx, urlstr)
 	}
 
-	// replace the scheme and path to construct the URL like:
+	// replace the scheme and path to construct a URL like:
 	// http://127.0.0.1:9222/json/version
 	u, err := url.Parse(urlstr)
 	if err != nil {
-		return urlstr
+		return "", err
 	}
 	u.Scheme = "http"
+	host, port, err := net.SplitHostPort(u.Host)
+	if err != nil {
+		return "", err
+	}
+	host, err = resolveHost(ctx, host)
+	if err != nil {
+		return "", err
+	}
+	u.Host = net.JoinHostPort(host, port)
 	u.Path = "/json/version"
 
 	// to get "webSocketDebuggerUrl" in the response
-	resp, err := http.Get(forceIP(u.String()))
+	req, err := http.NewRequestWithContext(lctx, "GET", u.String(), nil)
+	if err != nil {
+		return "", err
+	}
+	resp, err := http.DefaultClient.Do(req)
 	if err != nil {
-		return urlstr
+		return "", err
 	}
 	defer resp.Body.Close()
 
 	var result map[string]interface{}
 	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
-		return urlstr
+		return "", err
 	}
-	// the browser will construct the debugger URL using the "host" header of the /json/version request.
-	// for example, run headless-shell in a container: docker run -d -p 9000:9222 chromedp/headless-shell:latest
-	// then: curl http://127.0.0.1:9000/json/version
-	// and the debugger URL will be something like: ws://127.0.0.1:9000/devtools/browser/...
+	// the browser will construct the debugger URL using the "host" header of
+	// the /json/version request. For example, run headless-shell in a container:
+	//     docker run -d -p 9000:9222 chromedp/headless-shell:latest
+	// then:
+	//     curl http://127.0.0.1:9000/json/version
+	// and the websocket debugger URL will be something like:
+	// ws://127.0.0.1:9000/devtools/browser/...
 	wsURL := result["webSocketDebuggerUrl"].(string)
-	return wsURL
+	return wsURL, nil
 }
 
 func runListeners(list []cancelableListener, ev interface{}) []cancelableListener {