Fix ipv6 stuffs 3

pkg/http/client.go

@@ -118,25 +118,6 @@ func NewHTTPClient(token string, maxRetries int, timeoutSecs int, restrictToHost
 			if err != nil {
 				return nil, err
 			}
-			if len(restrictToHosts) > 0 {
-				found := false
-				for _, h := range restrictToHosts {
-					if h == host {
-						found = true
-						break
-					}
-				}
-				if !found {
-					log.Warnf("Host %s is not an IPv6 host (ensure you have preferred_hosts properly set in your config.yml, if unset, run `zurg network-test -t ipv6`)", host)
-					// replace with a random ipv6 host
-					restrictToHostsLen := len(restrictToHosts)
-					randomHost := restrictToHosts[rand.Intn(restrictToHostsLen)]
-					host = randomHost
-					log.Warnf("Replacing with a random IPv6 host: %s", host)
-					address = net.JoinHostPort(host, port)
-				}
-			}
-			// todo: replace .com with .cloud if needed
 			ips, err := net.DefaultResolver.LookupIPAddr(ctx, host)
 			if err != nil {
 				return nil, err
@@ -164,6 +145,7 @@ func (r *HTTPClient) Do(req *http.Request) (*http.Response, error) {
 	if r.bearerToken != "" {
 		req.Header.Set("Authorization", "Bearer "+r.bearerToken)
 	}
+	r.replaceHostIfNeeded(req)
 	// check if Range header is set
 	hasRangeHeader := req.Header.Get("Range") != ""
 
@@ -203,3 +185,40 @@ func (r *HTTPClient) Do(req *http.Request) (*http.Response, error) {
 	}
 	return resp, err
 }
+
+func (r *HTTPClient) replaceHostIfNeeded(req *http.Request) {
+	if !r.cfg.ShouldForceIPv6() {
+		return
+	}
+	if len(r.restrictToHosts) == 0 {
+		// replace .com with .cloud
+		host := req.URL.Host
+		if strings.HasSuffix(host, ".com") {
+			newHost := strings.Replace(host, ".com", ".cloud", 1)
+			newURL := *req.URL
+			newURL.Host = newHost
+			req.URL = &newURL
+		}
+		return
+	}
+	host, port, err := net.SplitHostPort(req.URL.Host)
+	if err != nil {
+		host = req.URL.Host // Use the host without port
+		port = "443"        // Default HTTPS port
+	}
+	found := false
+	for _, h := range r.restrictToHosts {
+		if h == host {
+			found = true
+			break
+		}
+	}
+	if !found {
+		r.log.Warnf("Host %s is not an IPv6 host (ensure you have preferred_hosts properly set in your config.yml, if unset, run `zurg network-test -t ipv6`)", host)
+		randomHost := r.restrictToHosts[rand.Intn(len(r.restrictToHosts))]
+		newHost := net.JoinHostPort(randomHost, port)
+		newURL := *req.URL
+		newURL.Host = newHost
+		req.URL = &newURL
+	}
+}