Make safe HTTP client options composable

fetch/fetcher.go

@@ -12,7 +12,6 @@ import (
 	"net"
 	"net/http"
 	"strconv"
-	"strings"
 	"time"
 
 	"github.com/rs/dnscache"
@@ -66,14 +65,17 @@ type Fetcher struct {
 	maxRetries   int
 	baseDelay    time.Duration
 	authFn       func(url string) (headerName, headerValue string)
-	allowPrivate map[string]bool
+	safeHTTPOpts safehttp.Options
+	ipChecker    *safehttp.HostIPChecker
 	stop         chan struct{}
 }
 
 // Option configures a Fetcher.
 type Option func(*Fetcher)
 
-// WithHTTPClient sets a custom HTTP client.
+// WithHTTPClient sets a custom HTTP client. To allow private hosts with a
+// custom client, construct it with safehttp.New and the desired options
+// before applying transport wrappers and passing it here.
 func WithHTTPClient(c *http.Client) Option {
 	return func(f *Fetcher) {
 		f.client = c
@@ -114,34 +116,8 @@ func WithAuthFunc(fn func(url string) (headerName, headerValue string)) Option {
 // Loopback and link-local addresses remain blocked.
 func WithAllowPrivateHosts(hosts ...string) Option {
 	return func(f *Fetcher) {
-		if f.allowPrivate == nil {
-			f.allowPrivate = make(map[string]bool, len(hosts))
-		}
-		for _, h := range hosts {
-			if h = normalizeHost(h); h != "" {
-				f.allowPrivate[h] = true
-			}
-		}
-	}
-}
-
-func normalizeHost(host string) string {
-	host = strings.TrimSpace(host)
-	if h, _, err := net.SplitHostPort(host); err == nil {
-		host = h
-	} else if strings.HasPrefix(host, "[") && strings.HasSuffix(host, "]") {
-		host = host[1 : len(host)-1]
-	}
-	return strings.ToLower(strings.TrimSuffix(host, "."))
-}
-
-// gateOptions returns the safehttp options for a dial to host.
-// Zero-value strict gate unless the host was whitelisted via WithAllowPrivateHosts.
-func (f *Fetcher) gateOptions(host string) safehttp.Options {
-	if f.allowPrivate[normalizeHost(host)] {
-		return safehttp.Options{AllowPrivate: true}
+		f.safeHTTPOpts.AllowPrivateHosts = append(f.safeHTTPOpts.AllowPrivateHosts, hosts...)
 	}
-	return safehttp.Options{}
 }
 
 // NewFetcher creates a new Fetcher with the given options.
@@ -190,7 +166,7 @@ func NewFetcher(opts ...Option) *Fetcher {
 					var lastErr error
 					for _, ip := range ips {
 						if parsed := net.ParseIP(ip); parsed != nil {
-							if err := safehttp.CheckIP(parsed, f.gateOptions(host)); err != nil {
+							if err := f.ipChecker.Check(host, parsed); err != nil {
 								lastErr = err
 								continue
 							}
@@ -222,6 +198,7 @@ func NewFetcher(opts ...Option) *Fetcher {
 	for _, opt := range opts {
 		opt(f)
 	}
+	f.ipChecker = safehttp.NewHostIPChecker(f.safeHTTPOpts)
 	return f
 }
 

fetch/fetcher_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"io"
+	"net"
 	"net/http"
 	"net/http/httptest"
 	"reflect"
@@ -411,39 +412,26 @@ func TestFetcherCloseStopsGoroutine(t *testing.T) {
 }
 
 func TestWithAllowPrivateHosts(t *testing.T) {
-	f := NewFetcher(WithAllowPrivateHosts(
+	want := []string{
 		" Registry.Internal.svc ",
 		"registry-with-port.internal:8080",
 		"registry-with-dot.internal.",
 		"[fd00::1]",
 		"",
-	))
+	}
+	f := NewFetcher(WithAllowPrivateHosts(want[:2]...), WithAllowPrivateHosts(want[2:]...))
 	defer func() { _ = f.Close() }()
 
-	for _, host := range []string{
-		"registry.internal.svc",
-		"REGISTRY-WITH-PORT.INTERNAL",
-		"registry-with-dot.internal",
-		"fd00::1",
-	} {
-		opts := f.gateOptions(host)
-		if !opts.AllowPrivate {
-			t.Errorf("gateOptions(%q).AllowPrivate = false, want true", host)
-		}
-		if opts.AllowLoopback {
-			t.Errorf("gateOptions(%q).AllowLoopback = true, want false", host)
-		}
+	if !reflect.DeepEqual(f.safeHTTPOpts.AllowPrivateHosts, want) {
+		t.Errorf("AllowPrivateHosts = %q; want %q", f.safeHTTPOpts.AllowPrivateHosts, want)
 	}
-
-	opts := f.gateOptions("other.example.com")
-	if opts.AllowPrivate || opts.AllowLoopback {
-		t.Errorf("non-whitelisted host exempted: %+v", opts)
+	if err := f.ipChecker.Check("registry.internal.svc", net.ParseIP("10.0.0.1")); err != nil {
+		t.Errorf("allowlisted private host rejected: %v", err)
 	}
 
 	strict := NewFetcher()
 	defer func() { _ = strict.Close() }()
-	opts = strict.gateOptions("registry.internal.svc")
-	if opts.AllowPrivate || opts.AllowLoopback {
-		t.Errorf("default fetcher not strict: %+v", opts)
+	if len(strict.safeHTTPOpts.AllowPrivateHosts) != 0 {
+		t.Errorf("default AllowPrivateHosts = %q; want empty", strict.safeHTTPOpts.AllowPrivateHosts)
 	}
 }

registries.go

@@ -134,6 +134,15 @@ var WithTimeout = client.WithTimeout
 // WithMaxRetries sets the maximum number of retries.
 var WithMaxRetries = client.WithMaxRetries
 
+// WithHTTPClient swaps in a caller-supplied HTTP client.
+var WithHTTPClient = client.WithHTTPClient
+
+// WithTransport replaces the underlying HTTP transport.
+var WithTransport = client.WithTransport
+
+// WithSafeHTTP applies protections for requests to untrusted hosts.
+var WithSafeHTTP = client.WithSafeHTTP
+
 // SupportedEcosystems returns all registered ecosystem types.
 // Note: ecosystems must be imported to be registered.
 func SupportedEcosystems() []string {

registries_test.go

@@ -6,7 +6,9 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"sort"
+	"strings"
 	"testing"
+	"time"
 
 	"github.com/git-pkgs/registries"
 	_ "github.com/git-pkgs/registries/all"
@@ -205,3 +207,32 @@ func TestConstants(t *testing.T) {
 		t.Errorf("StatusYanked constant mismatch")
 	}
 }
+
+func TestRootClientOptions(t *testing.T) {
+	custom := &http.Client{Timeout: 7 * time.Second}
+	c := registries.NewClient(registries.WithHTTPClient(custom))
+	if c.HTTPClient != custom {
+		t.Errorf("WithHTTPClient set HTTPClient to %p; want %p", c.HTTPClient, custom)
+	}
+
+	transport := http.DefaultTransport
+	c = registries.NewClient(registries.WithTransport(transport))
+	if c.HTTPClient.Transport != transport {
+		t.Errorf("WithTransport set Transport to %v; want %v", c.HTTPClient.Transport, transport)
+	}
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer server.Close()
+
+	c = registries.NewClient(registries.WithSafeHTTP())
+	resp, err := c.HTTPClient.Get(server.URL)
+	if err == nil {
+		_ = resp.Body.Close()
+		t.Fatal("WithSafeHTTP should refuse a loopback address")
+	}
+	if !strings.Contains(err.Error(), "loopback") {
+		t.Errorf("WithSafeHTTP error %v should mention loopback", err)
+	}
+}

safehttp/safehttp.go

@@ -29,6 +29,7 @@ import (
 	"net"
 	"net/http"
 	"net/url"
+	"strings"
 	"time"
 )
 
@@ -51,6 +52,12 @@ type Options struct {
 	// AllowPrivate disables the RFC1918 / ULA / CGNAT checks.
 	// Only set for tests or explicit operator allowlists.
 	AllowPrivate bool
+
+	// AllowPrivateHosts permits named hosts and IP literals to resolve or
+	// connect to RFC1918, ULA, or CGNAT addresses. Loopback and link-local
+	// addresses remain blocked. Host matching is case-insensitive and ignores
+	// ports and trailing dots.
+	AllowPrivateHosts []string
 }
 
 // testInsecure flips both AllowLoopback and AllowPrivate on at the
@@ -110,13 +117,50 @@ func CheckIP(ip net.IP, opts Options) error {
 	return newGate(opts).check(ip)
 }
 
+// CheckHostIP reports whether an IP is acceptable for the named host
+// under the supplied options. It applies AllowPrivateHosts in addition
+// to the checks performed by CheckIP. Use NewHostIPChecker when checking
+// multiple addresses with the same options.
+func CheckHostIP(host string, ip net.IP, opts Options) error {
+	return NewHostIPChecker(opts).Check(host, ip)
+}
+
+// HostIPChecker checks resolved IP addresses against a normalized host
+// allowlist. It can be reused across dials without rebuilding the allowlist.
+type HostIPChecker struct {
+	gate *ipGate
+}
+
+// NewHostIPChecker creates a reusable checker from opts.
+func NewHostIPChecker(opts Options) *HostIPChecker {
+	return &HostIPChecker{gate: newGate(opts)}
+}
+
+// Check reports whether ip is acceptable for host.
+func (c *HostIPChecker) Check(host string, ip net.IP) error {
+	return c.gate.checkHost(host, ip)
+}
+
 type ipGate struct {
-	allowLoopback bool
-	allowPrivate  bool
+	allowLoopback     bool
+	allowPrivate      bool
+	allowPrivateHosts map[string]bool
 }
 
 func newGate(opts Options) *ipGate {
-	return &ipGate{allowLoopback: opts.AllowLoopback, allowPrivate: opts.AllowPrivate}
+	g := &ipGate{
+		allowLoopback: opts.AllowLoopback,
+		allowPrivate:  opts.AllowPrivate,
+	}
+	if len(opts.AllowPrivateHosts) > 0 {
+		g.allowPrivateHosts = make(map[string]bool, len(opts.AllowPrivateHosts))
+		for _, host := range opts.AllowPrivateHosts {
+			if host = normalizeHost(host); host != "" {
+				g.allowPrivateHosts[host] = true
+			}
+		}
+	}
+	return g
 }
 
 func (g *ipGate) dial(ctx context.Context, network, addr string, dial func(context.Context, string, string) (net.Conn, error)) (net.Conn, error) {
@@ -126,7 +170,7 @@ func (g *ipGate) dial(ctx context.Context, network, addr string, dial func(conte
 	}
 
 	if ip := net.ParseIP(host); ip != nil {
-		if err := g.check(ip); err != nil {
+		if err := g.checkHost(host, ip); err != nil {
 			return nil, err
 		}
 		return dial(ctx, network, addr)
@@ -138,7 +182,7 @@ func (g *ipGate) dial(ctx context.Context, network, addr string, dial func(conte
 	}
 	var lastErr error
 	for _, ip := range ips {
-		if err := g.check(ip.IP); err != nil {
+		if err := g.checkHost(host, ip.IP); err != nil {
 			lastErr = err
 			continue
 		}
@@ -165,8 +209,12 @@ func mustCIDR(s string) *net.IPNet {
 }
 
 func (g *ipGate) check(ip net.IP) error {
+	return g.checkHost("", ip)
+}
+
+func (g *ipGate) checkHost(host string, ip net.IP) error {
 	allowLoopback := g.allowLoopback || testInsecure
-	allowPrivate := g.allowPrivate || testInsecure
+	allowPrivate := g.allowPrivate || g.allowPrivateHosts[normalizeHost(host)] || testInsecure
 
 	if ip.IsUnspecified() {
 		return blockedErr(ip, "unspecified")
@@ -191,6 +239,16 @@ func (g *ipGate) check(ip net.IP) error {
 	return nil
 }
 
+func normalizeHost(host string) string {
+	host = strings.TrimSpace(host)
+	if h, _, err := net.SplitHostPort(host); err == nil {
+		host = h
+	} else if strings.HasPrefix(host, "[") && strings.HasSuffix(host, "]") {
+		host = host[1 : len(host)-1]
+	}
+	return strings.ToLower(strings.TrimSuffix(host, "."))
+}
+
 func blockedErr(ip net.IP, kind string) error {
 	return fmt.Errorf("safehttp: refusing to connect to %s (%s)", ip, kind)
 }