feat: allow overriding PKCE verifier validation via CodeVerifierStrategy

handler/pkce/handler.go

@@ -5,9 +5,6 @@ package pkce
 
 import (
 	"context"
-	"crypto/sha256"
-	"encoding/base64"
-	"regexp"
 
 	"github.com/ory/x/errorsx"
 
@@ -27,11 +24,24 @@ type Handler struct {
 		fosite.EnforcePKCEForPublicClientsProvider
 		fosite.EnablePKCEPlainChallengeMethodProvider
 	}
+
+	// Verifier validates the code_verifier's format and its match against a
+	// bound code_challenge. If nil, DefaultCodeVerifierStrategy is used, which
+	// enforces RFC 7636 as written. See CodeVerifierStrategy for when to
+	// override this.
+	Verifier CodeVerifierStrategy
 }
 
 var _ fosite.TokenEndpointHandler = (*Handler)(nil)
 
-var verifierWrongFormat = regexp.MustCompile("[^\\w\\.\\-~]")
+// codeVerifierStrategy returns Verifier, defaulting to
+// DefaultCodeVerifierStrategy when unset.
+func (c *Handler) codeVerifierStrategy() CodeVerifierStrategy {
+	if c.Verifier == nil {
+		return DefaultCodeVerifierStrategy{}
+	}
+	return c.Verifier
+}
 
 func (c *Handler) HandleAuthorizeEndpointRequest(ctx context.Context, ar fosite.AuthorizeRequester, resp fosite.AuthorizeResponder) error {
 	// This let's us define multiple response types, for example open id connect's id_token
@@ -146,10 +156,6 @@ func (c *Handler) HandleTokenEndpointRequest(ctx context.Context, request fosite
 		return errorsx.WithStack(fosite.ErrServerError.WithWrap(err).WithDebug(err.Error()))
 	}
 
-	if err := c.Storage.DeletePKCERequestSession(ctx, signature); err != nil {
-		return errorsx.WithStack(fosite.ErrServerError.WithWrap(err).WithDebug(err.Error()))
-	}
-
 	challenge := pkceRequest.GetRequestForm().Get("code_challenge")
 	method := pkceRequest.GetRequestForm().Get("code_challenge_method")
 	client := pkceRequest.GetClient()
@@ -160,72 +166,47 @@ func (c *Handler) HandleTokenEndpointRequest(ctx context.Context, request fosite
 	nc := len(challenge)
 
 	if !c.Config.GetEnforcePKCE(ctx) && nc == 0 && nv == 0 {
-		return nil
+		// No challenge was bound and none is required, so this is a valid
+		// non-PKCE exchange. Consume the session before allowing it through.
+		return c.consumePKCERequestSession(ctx, signature)
 	}
 
-	// NOTE: The code verifier SHOULD have enough entropy to make it
-	// 	impractical to guess the value.  It is RECOMMENDED that the output of
-	// 	a suitable random number generator be used to create a 32-octet
-	// 	sequence.  The octet sequence is then base64url-encoded to produce a
-	// 	43-octet URL safe string to use as the code verifier.
-
-	// Validation
-	if nv < 43 {
-		return errorsx.WithStack(fosite.ErrInvalidGrant.
-			WithHint("The PKCE code verifier must be at least 43 characters."))
-	} else if nv > 128 {
-		return errorsx.WithStack(fosite.ErrInvalidGrant.
-			WithHint("The PKCE code verifier can not be longer than 128 characters."))
-	} else if verifierWrongFormat.MatchString(verifier) {
-		return errorsx.WithStack(fosite.ErrInvalidGrant.
-			WithHint("The PKCE code verifier must only contain [a-Z], [0-9], '-', '.', '_', '~'."))
+	// Validation. See DefaultCodeVerifierStrategy for the RFC 7636 rules this
+	// applies by default, and CodeVerifierStrategy for how to change them.
+	if err := c.codeVerifierStrategy().ValidateVerifierFormat(ctx, verifier); err != nil {
+		return err
 	} else if nc == 0 {
+		// A verifier was presented against a session that never had a challenge
+		// bound to it. There is nothing here a downgrade could exploit, so the
+		// session is consumed before rejecting the request.
+		if err := c.consumePKCERequestSession(ctx, signature); err != nil {
+			return err
+		}
+
 		return errorsx.WithStack(fosite.ErrInvalidGrant.
 			WithHint("The PKCE code verifier was provided but the code challenge was absent from the authorization request."))
 	}
 
-	// Upon receipt of the request at the token endpoint, the server
-	// verifies it by calculating the code challenge from the received
-	// "code_verifier" and comparing it with the previously associated
-	// "code_challenge", after first transforming it according to the
-	// "code_challenge_method" method specified by the client.
-	//
-	// 	If the "code_challenge_method" from Section 4.3 was "S256", the
-	// received "code_verifier" is hashed by SHA-256, base64url-encoded, and
-	// then compared to the "code_challenge", i.e.:
-	//
-	// BASE64URL-ENCODE(SHA256(ASCII(code_verifier))) == code_challenge
-	//
-	// If the "code_challenge_method" from Section 4.3 was "plain", they are
-	// compared directly, i.e.:
-	//
-	// code_verifier == code_challenge.
-	//
-	// 	If the values are equal, the token endpoint MUST continue processing
-	// as normal (as defined by OAuth 2.0 [RFC6749]).  If the values are not
-	// equal, an error response indicating "invalid_grant" as described in
-	// Section 5.2 of [RFC6749] MUST be returned.
-	switch method {
-	case "S256":
-		hash := sha256.New()
-		if _, err := hash.Write([]byte(verifier)); err != nil {
-			return errorsx.WithStack(fosite.ErrServerError.WithWrap(err).WithDebug(err.Error()))
-		}
-
-		if base64.RawURLEncoding.EncodeToString(hash.Sum([]byte{})) != challenge {
-			return errorsx.WithStack(fosite.ErrInvalidGrant.
-				WithHint("The PKCE code challenge did not match the code verifier."))
-		}
-		break
-	case "plain":
-		fallthrough
-	default:
-		if verifier != challenge {
-			return errorsx.WithStack(fosite.ErrInvalidGrant.
-				WithHint("The PKCE code challenge did not match the code verifier."))
-		}
+	// The session is deleted only once the verifier is confirmed to match the
+	// bound challenge below -- never beforehand, and never on a failed match.
+	// Deleting it on a failed attempt would strip the challenge, after which
+	// the same code could be replayed with no verifier at all, downgrading the
+	// exchange to a non-PKCE one.
+	if err := c.codeVerifierStrategy().ValidateChallenge(ctx, method, challenge, verifier); err != nil {
+		return err
 	}
 
+	return c.consumePKCERequestSession(ctx, signature)
+}
+
+// consumePKCERequestSession deletes the PKCE request session tied to
+// signature. Call this only once a request has been fully validated, or
+// determined not to need PKCE at all -- see the downgrade note in
+// HandleTokenEndpointRequest.
+func (c *Handler) consumePKCERequestSession(ctx context.Context, signature string) error {
+	if err := c.Storage.DeletePKCERequestSession(ctx, signature); err != nil {
+		return errorsx.WithStack(fosite.ErrServerError.WithWrap(err).WithDebug(err.Error()))
+	}
 	return nil
 }
 

handler/pkce/handler_test.go

@@ -10,6 +10,8 @@ import (
 	"fmt"
 	"testing"
 
+	"github.com/ory/x/errorsx"
+
 	"github.com/pkg/errors"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -391,6 +393,186 @@ func TestPKCEHandleTokenEndpointRequest(t *testing.T) {
 	}
 }
 
+// sessionExists reports whether a PKCE request session is still stored under
+// signature.
+func sessionExists(t *testing.T, s PKCERequestStorage, signature string) bool {
+	t.Helper()
+	_, err := s.GetPKCERequestSession(context.Background(), signature, nil)
+	if err == nil {
+		return true
+	}
+	require.ErrorIs(t, err, fosite.ErrNotFound)
+	return false
+}
+
+// TestHandleTokenEndpointRequest_SessionLifecycle locks in when
+// HandleTokenEndpointRequest consumes (deletes) the PKCE request session tied
+// to a code, per the downgrade note on that method: never on a failed
+// verification, so a failed attempt can't strip a bound challenge and let the
+// same code be replayed with no verifier as a non-PKCE exchange.
+func TestHandleTokenEndpointRequest_SessionLifecycle(t *testing.T) {
+	s256verifier := "KGCt4m8AmjUvIR5ArTByrmehjtbxn1A49YpTZhsH8N7fhDr7LQayn9xx6mck"
+	hash := sha256.New()
+	hash.Write([]byte(s256verifier))
+	s256challenge := base64.RawURLEncoding.EncodeToString(hash.Sum([]byte{}))
+
+	for _, tc := range []struct {
+		d               string
+		force           bool
+		challenge       string
+		method          string
+		verifier        string
+		wantErr         bool
+		wantSessionLeft bool
+	}{
+		{
+			d:         "match: session is consumed",
+			challenge: s256challenge,
+			method:    "S256",
+			verifier:  s256verifier,
+		},
+		{
+			d:               "mismatch: session survives so a retry can't downgrade",
+			challenge:       s256challenge,
+			method:          "S256",
+			verifier:        "wrong-verifier-wrong-verifier-wrong-verifier",
+			wantErr:         true,
+			wantSessionLeft: true,
+		},
+		{
+			d:               "verifier missing when a challenge was bound: session survives",
+			challenge:       s256challenge,
+			method:          "S256",
+			wantErr:         true,
+			wantSessionLeft: true,
+		},
+		{
+			d:        "no challenge or verifier, not enforced: session is consumed",
+			verifier: "",
+		},
+		{
+			d:        "verifier presented against a session with no bound challenge: session is consumed",
+			method:   "S256",
+			verifier: s256verifier,
+			wantErr:  true,
+		},
+	} {
+		t.Run(tc.d, func(t *testing.T) {
+			s := storage.NewMemoryStore()
+			ms := &mockCodeStrategy{signature: "code-under-test"}
+			h := &Handler{
+				Storage:               s,
+				AuthorizeCodeStrategy: ms,
+				Config:                &fosite.Config{EnforcePKCE: tc.force},
+			}
+			client := &fosite.DefaultClient{}
+
+			ar := fosite.NewAuthorizeRequest()
+			ar.Client = client
+			if tc.challenge != "" {
+				ar.Form.Add("code_challenge", tc.challenge)
+			}
+			if tc.method != "" {
+				ar.Form.Add("code_challenge_method", tc.method)
+			}
+			require.NoError(t, s.CreatePKCERequestSession(context.Background(), ms.signature, ar))
+
+			r := fosite.NewAccessRequest(nil)
+			r.Client = client
+			r.GrantTypes = fosite.Arguments{"authorization_code"}
+			if tc.verifier != "" {
+				r.Form.Add("code_verifier", tc.verifier)
+			}
+
+			err := h.HandleTokenEndpointRequest(context.Background(), r)
+			if tc.wantErr {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+			}
+
+			assert.Equal(t, tc.wantSessionLeft, sessionExists(t, s, ms.signature))
+		})
+	}
+}
+
+// permissiveCodeVerifierStrategy is a CodeVerifierStrategy stub that skips the
+// RFC 7636 length and character-set checks DefaultCodeVerifierStrategy applies,
+// while still requiring an exact match between verifier and challenge. It
+// stands in for authorization servers migrating from a provider that never
+// enforced RFC 7636's minimum verifier length.
+type permissiveCodeVerifierStrategy struct{}
+
+func (permissiveCodeVerifierStrategy) ValidateVerifierFormat(_ context.Context, _ string) error {
+	return nil
+}
+
+func (permissiveCodeVerifierStrategy) ValidateChallenge(_ context.Context, _, challenge, verifier string) error {
+	if verifier != challenge {
+		return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint("stub: verifier does not match challenge"))
+	}
+	return nil
+}
+
+// TestHandlerHonorsCustomCodeVerifierStrategy proves Handler defers verifier
+// format and comparison entirely to a custom Verifier when one is set, rather
+// than only being able to relax those rules by forking HandleTokenEndpointRequest.
+func TestHandlerHonorsCustomCodeVerifierStrategy(t *testing.T) {
+	s := storage.NewMemoryStore()
+	ms := &mockCodeStrategy{signature: "short-verifier-code"}
+	h := &Handler{
+		Storage:               s,
+		AuthorizeCodeStrategy: ms,
+		Config:                &fosite.Config{EnablePKCEPlainChallengeMethod: true},
+		Verifier:              permissiveCodeVerifierStrategy{},
+	}
+	client := &fosite.DefaultClient{}
+
+	ar := fosite.NewAuthorizeRequest()
+	ar.Client = client
+	ar.Form.Add("code_challenge", "short")
+	ar.Form.Add("code_challenge_method", "plain")
+	require.NoError(t, s.CreatePKCERequestSession(context.Background(), ms.signature, ar))
+
+	r := fosite.NewAccessRequest(nil)
+	r.Client = client
+	r.GrantTypes = fosite.Arguments{"authorization_code"}
+	r.Form.Add("code_verifier", "short")
+
+	// DefaultCodeVerifierStrategy would reject "short" outright (below the
+	// 43-character RFC 7636 minimum); the custom strategy above has no such
+	// floor, so the exact match succeeds.
+	assert.NoError(t, h.HandleTokenEndpointRequest(context.Background(), r))
+}
+
+// TestCustomCodeVerifierStrategyStillRequiresRegisteredChallenge proves
+// Handler still rejects a code_verifier with no bound code_challenge before
+// ever consulting the strategy's ValidateChallenge, regardless of how
+// permissive the strategy's ValidateVerifierFormat is.
+func TestCustomCodeVerifierStrategyStillRequiresRegisteredChallenge(t *testing.T) {
+	s := storage.NewMemoryStore()
+	ms := &mockCodeStrategy{signature: "no-challenge-code"}
+	h := &Handler{
+		Storage:               s,
+		AuthorizeCodeStrategy: ms,
+		Config:                &fosite.Config{},
+		Verifier:              permissiveCodeVerifierStrategy{},
+	}
+	client := &fosite.DefaultClient{}
+
+	ar := fosite.NewAuthorizeRequest()
+	ar.Client = client
+	require.NoError(t, s.CreatePKCERequestSession(context.Background(), ms.signature, ar))
+
+	r := fosite.NewAccessRequest(nil)
+	r.Client = client
+	r.GrantTypes = fosite.Arguments{"authorization_code"}
+	r.Form.Add("code_verifier", "short")
+
+	err := h.HandleTokenEndpointRequest(context.Background(), r)
+	assert.EqualError(t, newtesterr(err), "The provided authorization grant (e.g., authorization code, resource owner credentials) or refresh token is invalid, expired, revoked, does not match the redirection URI used in the authorization request, or was issued to another client. The PKCE code verifier was provided but the code challenge was absent from the authorization request.")
+}
+
 func newtesterr(err error) error {
 	if err == nil {
 		return nil