feat: add option to follow only same-host redirects

client.go

@@ -124,11 +124,6 @@ func (c *Client) getConn(protocol, host string, options *Options) (Conn, error)
 }
 
 func (c *Client) do(method, url, uripath string, headers map[string][]string, body io.Reader, redirectstatus *RedirectStatus, options *Options) (*http.Response, error) {
-	protocol := "http"
-	if strings.HasPrefix(strings.ToLower(url), "https://") {
-		protocol = "https"
-	}
-
 	if headers == nil {
 		headers = make(map[string][]string)
 	}
@@ -137,19 +132,12 @@ func (c *Client) do(method, url, uripath string, headers map[string][]string, bo
 		return nil, err
 	}
 
-	host := u.Host
 	if options.AutomaticHostHeader {
 		// add automatic space
-		headers["Host"] = []string{fmt.Sprintf(" %s", host)}
+		headers["Host"] = []string{" " + u.Host}
 	}
 
-	if !strings.Contains(host, ":") {
-		if protocol == "https" {
-			host += ":443"
-		} else {
-			host += ":80"
-		}
-	}
+	protocol, host := dialTarget(url, u.Host)
 
 	// standard path
 	path := u.Path
@@ -164,10 +152,6 @@ func (c *Client) do(method, url, uripath string, headers map[string][]string, bo
 		path = uripath
 	}
 
-	if strings.HasPrefix(url, "https://") {
-		protocol = "https"
-	}
-
 	conn, err := c.getConn(protocol, host, options)
 	if err != nil {
 		return nil, err
@@ -198,23 +182,57 @@ func (c *Client) do(method, url, uripath string, headers map[string][]string, bo
 		return nil, err
 	}
 
-	if resp.Status.IsRedirect() && redirectstatus.FollowRedirects && redirectstatus.Current <= redirectstatus.MaxRedirects {
+	follow := redirectstatus.FollowRedirects || options.FollowHostRedirects
+	if resp.Status.IsRedirect() && follow && redirectstatus.Current <= redirectstatus.MaxRedirects {
+		loc := headerValue(r.Header, "Location")
+		if strings.HasPrefix(loc, "/") {
+			loc = fmt.Sprintf("%s://%s%s", protocol, host, loc)
+		}
+
+		if options.FollowHostRedirects && !dialsTo(loc, protocol, host) {
+			return r, nil
+		}
 		// consume the response body
 		_, err := io.Copy(io.Discard, r.Body)
 		if err := firstErr(err, r.Body.Close()); err != nil {
 			return nil, err
 		}
-		loc := headerValue(r.Header, "Location")
-		if strings.HasPrefix(loc, "/") {
-			loc = fmt.Sprintf("%s://%s%s", protocol, host, loc)
-		}
 		redirectstatus.Current++
 		return c.do(method, loc, uripath, headers, body, redirectstatus, options)
 	}
 
 	return r, err
 }
 
+// dialTarget returns the protocol and the host:port address that a request
+// to rawURL connects to, given host, the host parsed from rawURL.
+func dialTarget(rawURL, host string) (string, string) {
+	protocol, port := "http", ":80"
+	if strings.HasPrefix(strings.ToLower(rawURL), "https://") {
+		protocol, port = "https", ":443"
+	}
+
+	if strings.Contains(host, ":") {
+		return protocol, host
+	}
+
+	return protocol, host + port
+}
+
+// dialsTo reports whether a request to rawURL connects to addr over protocol.
+// The address must match exactly: case folding would also equate distinct
+// hostnames, such as "ſecret.com" and "secret.com".
+func dialsTo(rawURL, protocol, addr string) bool {
+	u, err := urlutil.ParseURL(rawURL, true)
+	if err != nil {
+		return false
+	}
+
+	p, a := dialTarget(rawURL, u.Host)
+
+	return p == protocol && a == addr
+}
+
 // RedirectStatus is the current redirect status for the request
 type RedirectStatus struct {
 	FollowRedirects bool

client_test.go

@@ -6,7 +6,9 @@ import (
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"slices"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -150,3 +152,141 @@ func TestResponseBodyOwnsConnection(t *testing.T) {
 		t.Errorf("Close called %d times, want 1", conn.closeCalls)
 	}
 }
+
+func TestFollowHostRedirects(t *testing.T) {
+	t.Parallel()
+
+	var (
+		mu   sync.Mutex
+		hits []string
+	)
+
+	record := func(server string, r *http.Request) {
+		mu.Lock()
+		defer mu.Unlock()
+
+		hits = append(hits, server+" "+r.URL.Path)
+	}
+
+	other := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
+		record("other", r)
+	}))
+	defer other.Close()
+
+	target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		record("target", r)
+
+		location, ok := map[string]string{
+			"/to-path":       "/landing",
+			"/to-self":       "http://" + r.Host + "/landing",
+			"/to-https":      "https://" + r.Host + "/landing",
+			"/to-other-port": other.URL + "/landing",
+			"/to-other-host": "http://" + strings.Replace(r.Host, "127.0.0.1", "localhost", 1) + "/landing",
+		}[r.URL.Path]
+		if ok {
+			w.Header().Set("Location", location)
+			w.WriteHeader(http.StatusFound)
+		}
+	}))
+	defer target.Close()
+
+	sameHost := Options{FollowRedirects: true, FollowHostRedirects: true, MaxRedirects: 10, AutomaticHostHeader: true, Timeout: 5 * time.Second}
+
+	tests := []struct {
+		name       string
+		path       string
+		options    Options
+		wantStatus int
+		wantHits   []string
+	}{
+		{
+			name: "relative location", path: "/to-path", options: sameHost,
+			wantStatus: http.StatusOK, wantHits: []string{"target /to-path", "target /landing"},
+		},
+		{
+			name: "absolute location on the same host", path: "/to-self", options: sameHost,
+			wantStatus: http.StatusOK, wantHits: []string{"target /to-self", "target /landing"},
+		},
+		{
+			name: "other scheme", path: "/to-https", options: sameHost,
+			wantStatus: http.StatusFound, wantHits: []string{"target /to-https"},
+		},
+		{
+			name: "other port", path: "/to-other-port", options: sameHost,
+			wantStatus: http.StatusFound, wantHits: []string{"target /to-other-port"},
+		},
+		{
+			name: "other hostname", path: "/to-other-host", options: sameHost,
+			wantStatus: http.StatusFound, wantHits: []string{"target /to-other-host"},
+		},
+		{
+			name: "FollowHostRedirects without FollowRedirects", path: "/to-self",
+			options:    Options{FollowHostRedirects: true, MaxRedirects: 10, AutomaticHostHeader: true, Timeout: 5 * time.Second},
+			wantStatus: http.StatusOK, wantHits: []string{"target /to-self", "target /landing"},
+		},
+		{
+			name: "FollowRedirects alone", path: "/to-other-port",
+			options:    Options{FollowRedirects: true, MaxRedirects: 10, AutomaticHostHeader: true, Timeout: 5 * time.Second},
+			wantStatus: http.StatusOK, wantHits: []string{"target /to-other-port", "other /landing"},
+		},
+	}
+	// The subtests share hits, so they run one at a time.
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			mu.Lock()
+			hits = nil
+			mu.Unlock()
+
+			c := &Client{dialer: new(dialer), Options: &tt.options}
+			headers := map[string][]string{"Authorization": {"Bearer secret"}}
+
+			resp, err := c.DoRawWithOptions("GET", target.URL+tt.path, "", headers, nil, &tt.options)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			if err := resp.Body.Close(); err != nil {
+				t.Fatal(err)
+			}
+
+			mu.Lock()
+			defer mu.Unlock()
+
+			if resp.StatusCode != tt.wantStatus || !slices.Equal(hits, tt.wantHits) {
+				t.Errorf("status = %d, requests = %q; want %d, %q", resp.StatusCode, hits, tt.wantStatus, tt.wantHits)
+			}
+		})
+	}
+}
+
+func TestDialsTo(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		url, protocol, addr string
+		want                bool
+	}{
+		{"http://example.com/x", "http", "example.com:80", true},
+		{"http://example.com:80/x", "http", "example.com:80", true},
+		{"https://example.com/x", "https", "example.com:443", true},
+		{"https://example.com:80/x", "http", "example.com:80", false},
+		// Without "//" the parser keeps Scheme "https" but no host; such a URL
+		// has always been dialed as plain HTTP.
+		{"https:example.com/x", "http", ":80", true},
+		{"http://example.com:8080/x", "http", "example.com:80", false},
+		{"http://example.org/x", "http", "example.com:80", false},
+		{"http://example.com:80@example.org/x", "http", "example.com:80", false},
+		// Case folding would map these hosts onto ASCII ones (U+017F to s,
+		// U+212A to k) while they resolve to different names.
+		{"http://ſecret.com/x", "http", "secret.com:80", false},
+		{"http://Key.com/x", "http", "key.com:80", false},
+		{"http://EXAMPLE.com/x", "http", "example.com:80", false},
+		{"landing", "http", "example.com:80", false},
+		{"", "http", "example.com:80", false},
+	}
+	for _, tt := range tests {
+		if got := dialsTo(tt.url, tt.protocol, tt.addr); got != tt.want {
+			t.Errorf("dialsTo(%q, %q, %q) = %v, want %v", tt.url, tt.protocol, tt.addr, got, tt.want)
+		}
+	}
+}

options.go

@@ -9,8 +9,14 @@ import (
 
 // Options contains configuration options for rawhttp client
 type Options struct {
-	Timeout                time.Duration
-	FollowRedirects        bool
+	Timeout         time.Duration
+	FollowRedirects bool
+	// FollowHostRedirects follows a redirect only when the next request
+	// connects to the same host and port as the current one, over the same
+	// protocol (HTTP or HTTPS). The host must match exactly, letter case
+	// included. Otherwise the redirect response is returned. It follows such
+	// redirects even when FollowRedirects is false.
+	FollowHostRedirects    bool
 	MaxRedirects           int
 	AutomaticHostHeader    bool
 	AutomaticContentLength bool