refactor detectURL and rename it to modifyURL

Changes:

1. make it context aware.
2. make it return an error if it fails to modify the URL.
3. bypass resolution when the host is "localhost" or an IP address.
4. improve comments.

allocate.go

@@ -518,8 +518,7 @@ 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.
+// Because the allocator won't try to modify it and it's obviously invalid.
 func NewRemoteAllocator(parent context.Context, url string) (context.Context, context.CancelFunc) {
 	ctx, cancel := context.WithCancel(parent)
 	c := &Context{Allocator: &RemoteAllocator{
@@ -544,6 +543,11 @@ func (a *RemoteAllocator) Allocate(ctx context.Context, opts ...BrowserOption) (
 		return nil, ErrInvalidContext
 	}
 
+	wsURL, err := modifyURL(ctx, a.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,8 +560,7 @@ func (a *RemoteAllocator) Allocate(ctx context.Context, opts ...BrowserOption) (
 		cancel()    // close the websocket connection
 		a.wg.Done()
 	}()
-	wsURL := detectURL(a.wsURL)
-	wsURL = forceIP(wsURL)
+
 	browser, err := NewBrowser(wctx, wsURL, opts...)
 	if err != nil {
 		return nil, 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 {