Harden CIMD client resolution and caching

Tighten redirect URI validation for metadata documents, honor
Cache-Control no-store when caching fetched documents, and resolve
clients on the same transaction as authorization. Load
external_client_id from the database and parse unbounded max-stale
directives in cachecontrol.

Signed-off-by: Bryan Frimin <bryan@probo.com>

pkg/cachecontrol/cachecontrol.go

@@ -137,6 +137,11 @@ func ParseRequest(header string) (*RequestDirective, error) {
 
 			dir.maxAge = &seconds
 		case MaxStale:
+			if token.Value == "" {
+				dir.maxStaleUnbounded = true
+				break
+			}
+
 			seconds, err := parseDeltaSeconds(token.Value)
 			if err != nil {
 				return nil, fmt.Errorf("cannot parse max-stale: %w", err)
@@ -292,6 +297,10 @@ func parseDirectives(header string, parse func(string) (*TokenPair, error)) ([]*
 		tokens = append(tokens, token)
 	}
 
+	if err := scanner.Err(); err != nil {
+		return nil, fmt.Errorf("cannot scan cache-control directives: %w", err)
+	}
+
 	return tokens, nil
 }
 
@@ -370,6 +379,7 @@ func scanCommaSeparatedWords(data []byte, atEOF bool) (advance int, token []byte
 
 	for width := 0; start < len(data); start += width {
 		var r rune
+
 		r, width = utf8.DecodeRune(data[start:])
 		if !isSpace(r) {
 			break
@@ -377,10 +387,12 @@ func scanCommaSeparatedWords(data []byte, atEOF bool) (advance int, token []byte
 	}
 
 	var ws int
+
 	inQuotes := false
 
 	for width, i := 0, start; i < len(data); i += width {
 		var r rune
+
 		r, width = utf8.DecodeRune(data[i:])
 
 		switch {

pkg/cachecontrol/cachecontrol_test.go

@@ -172,6 +172,37 @@ func TestParseRequest(t *testing.T) {
 			require.Error(t, err)
 		},
 	)
+
+	t.Run(
+		"max-stale without value",
+		func(t *testing.T) {
+			t.Parallel()
+
+			dir, err := cachecontrol.ParseRequest("max-stale")
+			require.NoError(t, err)
+			assert.True(t, dir.MaxStaleUnbounded())
+
+			_, bounded, ok := dir.MaxStale()
+			require.True(t, ok)
+			assert.False(t, bounded)
+		},
+	)
+
+	t.Run(
+		"max-stale with value",
+		func(t *testing.T) {
+			t.Parallel()
+
+			dir, err := cachecontrol.ParseRequest("max-stale=120")
+			require.NoError(t, err)
+			assert.False(t, dir.MaxStaleUnbounded())
+
+			seconds, bounded, ok := dir.MaxStale()
+			require.True(t, ok)
+			assert.True(t, bounded)
+			assert.Equal(t, uint64(120), seconds)
+		},
+	)
 }
 
 func TestParseResponse(t *testing.T) {
@@ -378,10 +409,12 @@ func TestResponseMaxAgeDuration(t *testing.T) {
 						_, gotOK := dir.MaxAgeDuration()
 						assert.False(t, gotOK)
 					}
+
 					return
 				}
 
 				require.NoError(t, err)
+
 				gotAge, gotOK := dir.MaxAgeDuration()
 				assert.True(t, gotOK)
 				assert.Equal(t, tt.wantAge, gotAge)

pkg/cachecontrol/directives.go

@@ -20,6 +20,7 @@ type (
 	RequestDirective struct {
 		maxAge            *uint64
 		maxStale          *uint64
+		maxStaleUnbounded bool
 		minFresh          *uint64
 		noCache           bool
 		noStore           bool
@@ -50,12 +51,20 @@ func (d *RequestDirective) MaxAge() (uint64, bool) {
 	return 0, false
 }
 
-func (d *RequestDirective) MaxStale() (uint64, bool) {
+func (d *RequestDirective) MaxStale() (seconds uint64, bounded bool, ok bool) {
-	if v := d.maxStale; v != nil {
+	if d.maxStaleUnbounded {
-		return *v, true
+		return 0, false, true
 	}
 
-	return 0, false
+	if v := d.maxStale; v != nil {
+		return *v, true, true
+	}
+
+	return 0, false, false
+}
+
+func (d *RequestDirective) MaxStaleUnbounded() bool {
+	return d.maxStaleUnbounded
 }
 
 func (d *RequestDirective) MinFresh() (uint64, bool) {

pkg/coredata/oauth2_client.go

@@ -136,6 +136,7 @@ func (c *OAuth2Client) LoadByID(
 	q := `
 SELECT
 	id,
+	COALESCE(external_client_id, '') AS external_client_id,
 	organization_id,
 	client_secret_hash,
 	client_name,
@@ -190,6 +191,7 @@ func (c *OAuth2Client) LoadByExternalClientID(
 	q := `
 SELECT
 	id,
+	COALESCE(external_client_id, '') AS external_client_id,
 	organization_id,
 	client_secret_hash,
 	client_name,
@@ -242,6 +244,7 @@ func (c *OAuth2Clients) LoadByOrganizationID(
 	q := `
 SELECT
 	id,
+	COALESCE(external_client_id, '') AS external_client_id,
 	organization_id,
 	client_secret_hash,
 	client_name,

pkg/iam/oauth2/cimd.go

@@ -17,6 +17,7 @@ package oauth2
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
@@ -196,12 +197,8 @@ func validateClientMetadataDocument(clientIDURL string, doc *ClientMetadataDocum
 	}
 
 	for _, redirectURI := range doc.RedirectURIs {
-		parsed, err := url.Parse(redirectURI)
+		if err := validateCIMDRedirectURI(redirectURI); err != nil {
-		if err != nil || parsed.Scheme == "" || parsed.Host == "" {
+			return err
-			return NewError(
-				ErrInvalidClient,
-				WithDescription("client metadata document contains invalid redirect_uri"),
-			)
 		}
 	}
 
@@ -223,6 +220,41 @@ func validateClientMetadataDocument(clientIDURL string, doc *ClientMetadataDocum
 	return nil
 }
 
+func validateCIMDRedirectURI(redirectURI string) error {
+	parsed, err := url.Parse(redirectURI)
+	if err != nil || parsed.Scheme == "" || parsed.Host == "" {
+		return NewError(
+			ErrInvalidClient,
+			WithDescription("client metadata document contains invalid redirect_uri"),
+		)
+	}
+
+	if parsed.User != nil || parsed.Fragment != "" {
+		return NewError(
+			ErrInvalidClient,
+			WithDescription("client metadata document contains invalid redirect_uri"),
+		)
+	}
+
+	switch parsed.Scheme {
+	case "https":
+	case "http":
+		if !net.IsLoopback(parsed.Hostname()) {
+			return NewError(
+				ErrInvalidClient,
+				WithDescription("client metadata document contains invalid redirect_uri"),
+			)
+		}
+	default:
+		return NewError(
+			ErrInvalidClient,
+			WithDescription("client metadata document contains invalid redirect_uri"),
+		)
+	}
+
+	return nil
+}
+
 func cimdRedirectURIAllowed(doc *ClientMetadataDocument, redirectURI string) bool {
 	for _, allowed := range doc.RedirectURIs {
 		if redirectURI == allowed {
@@ -280,8 +312,14 @@ func (f *cimdFetcher) loadCache(clientIDURL string) (*ClientMetadataDocument, bo
 }
 
 func (f *cimdFetcher) storeCache(clientIDURL string, doc *ClientMetadataDocument, cacheControl string) {
+	dir, err := cachecontrol.ParseResponse(cacheControl)
+	if err == nil && dir.NoStore() {
+		return
+	}
+
 	ttl := cimdDefaultCacheTTL
-	if dir, err := cachecontrol.ParseResponse(cacheControl); err == nil {
+
+	if err == nil {
 		if maxAge, ok := dir.MaxAgeDuration(); ok {
 			ttl = min(ttl, maxAge)
 		}
@@ -300,8 +338,30 @@ func (s *Service) ResolveClient(
 	ctx context.Context,
 	clientIDRaw string,
 	redirectURI string,
+) (*coredata.OAuth2Client, error) {
+	return s.resolveClient(ctx, nil, clientIDRaw, redirectURI)
+}
+
+func (s *Service) resolveClient(
+	ctx context.Context,
+	tx pg.Tx,
+	clientIDRaw string,
+	redirectURI string,
 ) (*coredata.OAuth2Client, error) {
 	if clientID, err := gid.ParseGID(clientIDRaw); err == nil {
+		if tx != nil {
+			client := coredata.OAuth2Client{}
+			if err := client.LoadByID(ctx, tx, coredata.NewNoScope(), clientID); err != nil {
+				if errors.Is(err, coredata.ErrResourceNotFound) {
+					return nil, NewError(ErrInvalidClient, WithDescription("client not found"))
+				}
+
+				return nil, fmt.Errorf("cannot load oauth2 client: %w", err)
+			}
+
+			return &client, nil
+		}
+
 		return s.GetClientByID(ctx, clientID)
 	}
 
@@ -325,7 +385,7 @@ func (s *Service) ResolveClient(
 		return nil, ErrInvalidRedirectURI
 	}
 
-	client, err := s.upsertCIMDClient(ctx, clientIDRaw, doc)
+	client, err := s.upsertCIMDClient(ctx, tx, clientIDRaw, doc)
 	if err != nil {
 		return nil, err
 	}
@@ -335,6 +395,7 @@ func (s *Service) ResolveClient(
 
 func (s *Service) upsertCIMDClient(
 	ctx context.Context,
+	tx pg.Tx,
 	externalClientID string,
 	doc *ClientMetadataDocument,
 ) (*coredata.OAuth2Client, error) {
@@ -358,6 +419,7 @@ func (s *Service) upsertCIMDClient(
 	)
 
 	now := time.Now()
+
 	candidate, err := coredata.NewCIMDClient(
 		externalClientID,
 		doc.ClientName,
@@ -373,16 +435,28 @@ func (s *Service) upsertCIMDClient(
 
 	var client coredata.OAuth2Client
 
-	err = s.pg.WithTx(
+	upsert := func(ctx context.Context, conn pg.Tx) error {
-		ctx,
-		func(ctx context.Context, tx pg.Tx) error {
 		client = *candidate
 
-			if err := client.UpsertCIMD(ctx, tx); err != nil {
+		if err := client.UpsertCIMD(ctx, conn); err != nil {
 			return fmt.Errorf("cannot upsert cimd oauth2 client: %w", err)
 		}
 
 		return nil
+	}
+
+	if tx != nil {
+		if err := upsert(ctx, tx); err != nil {
+			return nil, err
+		}
+
+		return &client, nil
+	}
+
+	err = s.pg.WithTx(
+		ctx,
+		func(ctx context.Context, conn pg.Tx) error {
+			return upsert(ctx, conn)
 		},
 	)
 	if err != nil {

pkg/iam/oauth2/cimd_test.go

@@ -102,6 +102,31 @@ func TestValidateClientMetadataDocument(t *testing.T) {
 
 	require.NoError(t, validateClientMetadataDocument(clientID, &doc))
 
+	t.Run(
+		"http redirect on non-loopback rejected",
+		func(t *testing.T) {
+			t.Parallel()
+
+			bad := doc
+			bad.RedirectURIs = []string{"http://example.com/callback"}
+
+			err := validateClientMetadataDocument(clientID, &bad)
+			require.Error(t, err)
+		},
+	)
+
+	t.Run(
+		"http loopback redirect allowed",
+		func(t *testing.T) {
+			t.Parallel()
+
+			loopback := doc
+			loopback.RedirectURIs = []string{"http://127.0.0.1:3000/callback"}
+
+			require.NoError(t, validateClientMetadataDocument(clientID, &loopback))
+		},
+	)
+
 	t.Run(
 		"mismatched client_id",
 		func(t *testing.T) {
@@ -143,6 +168,11 @@ func TestCIMDRedirectURIAllowed(t *testing.T) {
 func TestCIMDFetcherFetch(t *testing.T) {
 	t.Parallel()
 
+	t.Run(
+		"caches response with max-age",
+		func(t *testing.T) {
+			t.Parallel()
+
 			doc := ClientMetadataDocument{
 				ClientName:              "Test MCP Client",
 				RedirectURIs:            []string{"http://127.0.0.1:3000/callback"},
@@ -174,4 +204,46 @@ func TestCIMDFetcherFetch(t *testing.T) {
 			cached, err := fetcher.fetch(t.Context(), clientID)
 			require.NoError(t, err)
 			require.Equal(t, fetched.ClientName, cached.ClientName)
+		},
+	)
+
+	t.Run(
+		"no-store response is not cached",
+		func(t *testing.T) {
+			t.Parallel()
+
+			doc := ClientMetadataDocument{
+				ClientName:              "Test MCP Client",
+				RedirectURIs:            []string{"http://127.0.0.1:3000/callback"},
+				TokenEndpointAuthMethod: "none",
+			}
+
+			requestCount := 0
+			server := httptest.NewTLSServer(
+				http.HandlerFunc(
+					func(w http.ResponseWriter, _ *http.Request) {
+						requestCount++
+
+						w.Header().Set("Cache-Control", "no-store")
+						_ = json.NewEncoder(w).Encode(doc)
+					},
+				),
+			)
+			t.Cleanup(server.Close)
+
+			clientID := server.URL + "/oauth/client.json"
+			doc.ClientID = clientID
+
+			fetcher := &cimdFetcher{
+				httpClient: server.Client(),
+				logger:     log.NewLogger(),
+			}
+
+			_, err := fetcher.fetch(t.Context(), clientID)
+			require.NoError(t, err)
+			_, err = fetcher.fetch(t.Context(), clientID)
+			require.NoError(t, err)
+			assert.Equal(t, 2, requestCount)
+		},
+	)
 }

pkg/iam/oauth2/service.go

@@ -1433,7 +1433,7 @@ func (s *Service) Authorize(
 	if err := s.pg.WithTx(
 		ctx,
 		func(ctx context.Context, tx pg.Tx) error {
-			client, err := s.ResolveClient(ctx, req.ClientIDRaw, req.RedirectURI)
+			client, err := s.resolveClient(ctx, tx, req.ClientIDRaw, req.RedirectURI)
 			if err != nil {
 				return err
 			}