Merge pull request #782 from docker/AIMI-4

Validate realm URL before token exchange

Author: ericcurtin

pkg/distribution/oci/remote/transport.go

@@ -2,9 +2,12 @@ package remote
 
 import (
 	"context"
+	"crypto/tls"
 	"encoding/json"
 	"fmt"
 	"io"
+	"log/slog"
+	"net"
 	"net/http"
 	"net/url"
 	"strings"
@@ -38,6 +41,153 @@ type Token struct {
 	ExpiresIn   int    `json:"expires_in"`
 }
 
+// privateOrLoopbackCIDRs lists IP ranges that must never be contacted as a
+// token-exchange realm. Allowing requests to these addresses would let a
+// malicious registry pivot Model Runner into an internal-service proxy
+// (SSRF), reaching endpoints that are not accessible from the public internet.
+var privateOrLoopbackCIDRs = func() []*net.IPNet {
+	cidrs := []string{
+		"127.0.0.0/8",    // loopback IPv4
+		"::1/128",        // loopback IPv6
+		"169.254.0.0/16", // link-local IPv4 / AWS EC2 instance-metadata
+		"fe80::/10",      // link-local IPv6
+		"10.0.0.0/8",     // RFC-1918 private
+		"172.16.0.0/12",  // RFC-1918 private
+		"192.168.0.0/16", // RFC-1918 private
+		"fc00::/7",       // IPv6 ULA
+	}
+	nets := make([]*net.IPNet, 0, len(cidrs))
+	for _, cidr := range cidrs {
+		_, n, err := net.ParseCIDR(cidr)
+		if err != nil {
+			// These are hardcoded compile-time constants; a parse failure
+			// indicates a programmer error (e.g. a typo). Panic immediately
+			// so the mistake is caught at startup rather than silently
+			// weakening the SSRF blocklist.
+			panic(fmt.Sprintf("internal error: failed to parse hardcoded CIDR %q: %v", cidr, err))
+		}
+		nets = append(nets, n)
+	}
+	return nets
+}()
+
+// internalHostnames lists hostnames that must never be used as a realm,
+// regardless of what IP address they resolve to.
+var internalHostnames = []string{
+	"localhost",
+	"host.docker.internal",
+	"model-runner.docker.internal",
+	"gateway.docker.internal",
+}
+
+// isDisallowedIP reports whether ip falls in any of the private/loopback/
+// link-local ranges that must not be contacted as a token-exchange realm.
+func isDisallowedIP(ip net.IP) bool {
+	for _, cidr := range privateOrLoopbackCIDRs {
+		if cidr.Contains(ip) {
+			return true
+		}
+	}
+	return false
+}
+
+// resolveAndValidateRealm parses the realm URL, validates the hostname against
+// a static blocklist, and resolves it to a dial address (ip:port) that is safe
+// to connect to. By returning the resolved IP, callers can use a custom
+// DialContext to connect to that exact address — preventing DNS-rebinding
+// attacks where a malicious DNS server could return different IPs for
+// successive lookups (TOCTOU).
+//
+// If hostname is a literal IP address it is validated directly without
+// triggering a DNS lookup.
+func resolveAndValidateRealm(rawURL string) (dialAddr, hostname string, err error) {
+	u, err := url.Parse(rawURL)
+	if err != nil {
+		return "", "", fmt.Errorf("invalid realm URL: %w", err)
+	}
+
+	hostname = u.Hostname()
+	port := u.Port()
+	if port == "" {
+		switch u.Scheme {
+		case "https":
+			port = "443"
+		default:
+			port = "80"
+		}
+	}
+
+	// Block well-known internal hostnames regardless of DNS resolution.
+	for _, internal := range internalHostnames {
+		if strings.EqualFold(hostname, internal) {
+			return "", "", fmt.Errorf("realm URL hostname %q is not allowed", hostname)
+		}
+	}
+
+	// If the hostname is a literal IP address, validate it directly without a
+	// DNS lookup — there is no DNS to rebind.
+	if ip := net.ParseIP(hostname); ip != nil {
+		if isDisallowedIP(ip) {
+			return "", "", fmt.Errorf("realm URL contains a disallowed IP address %s", hostname)
+		}
+		return net.JoinHostPort(hostname, port), hostname, nil
+	}
+
+	// Resolve the hostname and validate every returned address. Using the
+	// resolved IP as the dial address prevents DNS rebinding: the same IP that
+	// passed validation is the one that will be used for the connection.
+	ips, err := net.LookupHost(hostname)
+	if err != nil {
+		return "", "", fmt.Errorf("resolving realm hostname %q: %w", hostname, err)
+	}
+	if len(ips) == 0 {
+		return "", "", fmt.Errorf("realm hostname %q resolved to no addresses", hostname)
+	}
+	for _, ipStr := range ips {
+		ip := net.ParseIP(ipStr)
+		if ip == nil {
+			continue
+		}
+		if isDisallowedIP(ip) {
+			return "", "", fmt.Errorf("realm URL resolves to a disallowed address %s", ipStr)
+		}
+	}
+
+	// All resolved IPs passed validation. Use the first one as the dial
+	// address so the HTTP client never performs a second DNS lookup.
+	return net.JoinHostPort(ips[0], port), hostname, nil
+}
+
+// buildSafeTransport wraps base with a custom DialContext that connects
+// directly to dialAddr (a pre-validated "ip:port") instead of relying on DNS
+// resolution. This closes the TOCTOU window between realm URL validation and
+// the actual connection. For TLS connections, serverName is set as the SNI
+// value so that certificate validation still uses the original hostname.
+func buildSafeTransport(base http.RoundTripper, dialAddr, serverName string) (http.RoundTripper, error) {
+	if base == nil {
+		base = http.DefaultTransport
+	}
+
+	t, ok := base.(*http.Transport)
+	if !ok {
+		return nil, fmt.Errorf("cannot build safe transport from base of type %T; only *http.Transport is supported", base)
+	}
+
+	dial := func(ctx context.Context, network, _ string) (net.Conn, error) {
+		return (&net.Dialer{}).DialContext(ctx, network, dialAddr)
+	}
+
+	cloned := t.Clone()
+	cloned.DialContext = dial
+	if cloned.TLSClientConfig == nil {
+		cloned.TLSClientConfig = &tls.Config{MinVersion: tls.VersionTLS12} //nolint:gosec
+	} else {
+		cloned.TLSClientConfig = cloned.TLSClientConfig.Clone()
+	}
+	cloned.TLSClientConfig.ServerName = serverName
+	return cloned, nil
+}
+
 // Ping pings a registry and returns authentication information.
 func Ping(ctx context.Context, reg reference.Registry, transport http.RoundTripper) (*PingResponse, error) {
 	if transport == nil {
@@ -104,12 +254,21 @@ func parseWWWAuthenticate(header string) WWWAuthenticate {
 }
 
 // Exchange exchanges credentials for a bearer token.
-func Exchange(ctx context.Context, reg reference.Registry, auth authn.Authenticator, transport http.RoundTripper, scopes []string, pr *PingResponse) (*Token, error) {
+func Exchange(ctx context.Context, _ reference.Registry, auth authn.Authenticator, transport http.RoundTripper, scopes []string, pr *PingResponse) (*Token, error) {
 	if transport == nil {
 		transport = http.DefaultTransport
 	}
 
-	client := &http.Client{Transport: transport}
+	dialAddr, hostname, err := resolveAndValidateRealm(pr.WWWAuthenticate.Realm)
+	if err != nil {
+		return nil, fmt.Errorf("realm URL rejected: %w", err)
+	}
+
+	safeTransport, err := buildSafeTransport(transport, dialAddr, hostname)
+	if err != nil {
+		return nil, fmt.Errorf("failed to build safe transport: %w", err)
+	}
+	client := &http.Client{Transport: safeTransport}
 
 	// Build token request URL
 	tokenURL, err := url.Parse(pr.WWWAuthenticate.Realm)
@@ -150,7 +309,11 @@ func Exchange(ctx context.Context, reg reference.Registry, auth authn.Authentica
 
 	if resp.StatusCode != http.StatusOK {
 		body, _ := io.ReadAll(resp.Body)
-		return nil, fmt.Errorf("token request failed with status %d: %s", resp.StatusCode, string(body))
+		slog.DebugContext(ctx, "token request failed",
+			"status", resp.StatusCode,
+			"body", string(body),
+		)
+		return nil, fmt.Errorf("token request failed: unexpected status %d from token endpoint", resp.StatusCode)
 	}
 
 	var token Token

pkg/distribution/oci/remote/transport_test.go

@@ -0,0 +1,108 @@
+package remote_test
+
+import (
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync/atomic"
+	"testing"
+
+	"github.com/docker/model-runner/pkg/distribution/oci/reference"
+	"github.com/docker/model-runner/pkg/distribution/oci/remote"
+)
+
+// emptyRegistry is a zero-value reference.Registry.
+// reference.Registry is a concrete struct (not an interface), so nil cannot
+// be used; the zero value is a safe placeholder when registry-specific
+// behaviour is not exercised by the test.
+var emptyRegistry reference.Registry
+
+// pingResponseForRealm returns a *remote.PingResponse whose Realm field is
+// set to the given URL.  This simulates the result of calling Ping() against
+// a malicious registry that advertises an attacker-controlled realm.
+func pingResponseForRealm(realm string) *remote.PingResponse {
+	return &remote.PingResponse{
+		WWWAuthenticate: remote.WWWAuthenticate{
+			Realm:   realm,
+			Service: "evil-registry",
+			Scope:   "repository:evil/model:pull",
+		},
+	}
+}
+
+// TestExchangeSSRF_RequestSentToRealmURL verifies that prevents Exchange() from
+// contacting internal services via a malicious realm URL.
+func TestExchangeSSRF_RequestSentToRealmURL(t *testing.T) {
+	var hitCount atomic.Int32
+
+	// "Internal service" — simulates a host-local endpoint (127.0.0.1) that
+	// should never be reachable via a registry token-exchange flow.
+	internalService := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		hitCount.Add(1)
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusForbidden)
+		fmt.Fprintln(w, `{"error":"not a token endpoint"}`)
+	}))
+	defer internalService.Close()
+
+	// Build a PingResponse whose Realm points at the internal service —
+	// this is what a malicious registry would return in its 401 response.
+	pr := pingResponseForRealm(internalService.URL + "/internal/credentials")
+
+	_, err := remote.Exchange(t.Context(), emptyRegistry, nil, nil, []string{"repository:x:pull"}, pr)
+	if err == nil {
+		t.Fatal("Exchange() should have returned an error when the realm URL resolves to a loopback address")
+	}
+	if hitCount.Load() > 0 {
+		t.Errorf("SSRF not blocked: Exchange() sent %d request(s) to the internal service at %s — realm URL validation should have rejected 127.0.0.1 before making any HTTP request", hitCount.Load(), internalService.URL)
+	}
+	if !strings.Contains(err.Error(), "realm URL rejected") {
+		t.Errorf("expected error to mention realm URL rejection, got: %q", err.Error())
+	}
+}
+
+// TestExchangeSSRF_SensitiveBodyNotReflectedInError verifies that Exchange()
+// does NOT include a token-endpoint response body in the error it returns to
+// the caller.
+func TestExchangeSSRF_SensitiveBodyNotReflectedInError(t *testing.T) {
+	sensitiveData := map[string]string{
+		"db_password":      "s3cret_example",
+		"internal_api_key": "sk-example-key",
+		"proof":            "THIS_DATA_READ_VIA_SSRF",
+	}
+	sensitiveJSON, err := json.Marshal(sensitiveData)
+	if err != nil {
+		t.Fatalf("failed to marshal test data: %v", err)
+	}
+
+	// "Internal service" returns a non-200 response containing sensitive data.
+	internalService := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusUnauthorized)
+		w.Write(sensitiveJSON) //nolint:errcheck
+	}))
+	defer internalService.Close()
+
+	pr := pingResponseForRealm(internalService.URL + "/admin/api/keys")
+
+	_, err = remote.Exchange(t.Context(), emptyRegistry, nil, nil, []string{"repository:x:pull"}, pr)
+	if err == nil {
+		t.Fatal("expected an error from Exchange() when realm URL is blocked or token endpoint returns 401, got nil")
+	}
+
+	errMsg := err.Error()
+
+	// The error must indicate realm rejection, not a remote HTTP status code.
+	if !strings.Contains(errMsg, "realm URL rejected") {
+		t.Errorf("expected error to mention realm URL rejection (Fix 2a), got: %q", errMsg)
+	}
+
+	// the response body must never appear in the returned error
+	for _, sensitive := range []string{"s3cret_example", "sk-example-key", "THIS_DATA_READ_VIA_SSRF"} {
+		if strings.Contains(errMsg, sensitive) {
+			t.Errorf("BODY REFLECTION vulnerability: error contains sensitive value %q — the response body must not be included in errors returned to the caller (log at debug level instead)", sensitive)
+		}
+	}
+}