From 24c72f33565648b2188c4d14cbb870b60349a971 Mon Sep 17 00:00:00 2001
From: James Elliott <james-d-elliott@users.noreply.github.com>
Date: Fri, 21 Jul 2023 16:14:03 +1000
Subject: [PATCH] fix(pkce): session generated needlessly

---
 handler/pkce/handler.go      | 43 +++++++++++++++++++++++-------------
 handler/pkce/handler_test.go |  1 +
 2 files changed, 29 insertions(+), 15 deletions(-)

diff --git a/handler/pkce/handler.go b/handler/pkce/handler.go
index f457b8bea..86fe99ac1 100644
--- a/handler/pkce/handler.go
+++ b/handler/pkce/handler.go
@@ -47,6 +47,11 @@ func (c *Handler) HandleAuthorizeEndpointRequest(ctx context.Context, ar fosite.
 		return err
 	}
 
+	// We don't need a session if it's not enforced and the PKCE parameters are not provided by the client.
+	if challenge == "" && method == "" {
+		return nil
+	}
+
 	code := resp.GetCode()
 	if len(code) == 0 {
 		return errorsx.WithStack(fosite.ErrServerError.WithDebug("The PKCE handler must be loaded after the authorize code handler."))
@@ -71,17 +76,7 @@ func (c *Handler) validate(ctx context.Context, challenge, method string, client
 		// error response with the "error" value set to "invalid_request".  The
 		// "error_description" or the response of "error_uri" SHOULD explain the
 		// nature of error, e.g., code challenge required.
-		if c.Config.GetEnforcePKCE(ctx) {
-			return errorsx.WithStack(fosite.ErrInvalidRequest.
-				WithHint("Clients must include a code_challenge when performing the authorize code flow, but it is missing.").
-				WithDebug("The server is configured in a way that enforces PKCE for clients."))
-		}
-		if c.Config.GetEnforcePKCEForPublicClients(ctx) && client.IsPublic() {
-			return errorsx.WithStack(fosite.ErrInvalidRequest.
-				WithHint("This client must include a code_challenge when performing the authorize code flow, but it is missing.").
-				WithDebug("The server is configured in a way that enforces PKCE for this client."))
-		}
-		return nil
+		return c.validateNoPKCE(ctx, client)
 	}
 
 	// If the server supporting PKCE does not support the requested
@@ -108,6 +103,20 @@ func (c *Handler) validate(ctx context.Context, challenge, method string, client
 	return nil
 }
 
+func (c *Handler) validateNoPKCE(ctx context.Context, client fosite.Client) error {
+	if c.Config.GetEnforcePKCE(ctx) {
+		return errorsx.WithStack(fosite.ErrInvalidRequest.
+			WithHint("Clients must include a code_challenge when performing the authorize code flow, but it is missing.").
+			WithDebug("The server is configured in a way that enforces PKCE for clients."))
+	}
+	if c.Config.GetEnforcePKCEForPublicClients(ctx) && client.IsPublic() {
+		return errorsx.WithStack(fosite.ErrInvalidRequest.
+			WithHint("This client must include a code_challenge when performing the authorize code flow, but it is missing.").
+			WithDebug("The server is configured in a way that enforces PKCE for this client."))
+	}
+	return nil
+}
+
 func (c *Handler) HandleTokenEndpointRequest(ctx context.Context, request fosite.AccessRequester) error {
 	if !c.CanHandleTokenEndpointRequest(ctx, request) {
 		return errorsx.WithStack(fosite.ErrUnknownRequest)
@@ -123,8 +132,12 @@ func (c *Handler) HandleTokenEndpointRequest(ctx context.Context, request fosite
 
 	code := request.GetRequestForm().Get("code")
 	signature := c.AuthorizeCodeStrategy.AuthorizeCodeSignature(ctx, code)
-	authorizeRequest, err := c.Storage.GetPKCERequestSession(ctx, signature, request.GetSession())
+	pkceRequest, err := c.Storage.GetPKCERequestSession(ctx, signature, request.GetSession())
 	if errors.Is(err, fosite.ErrNotFound) {
+		if verifier == "" {
+			return c.validateNoPKCE(ctx, request.GetClient())
+		}
+
 		return errorsx.WithStack(fosite.ErrInvalidGrant.WithHint("Unable to find initial PKCE data tied to this request").WithWrap(err).WithDebug(err.Error()))
 	} else if err != nil {
 		return errorsx.WithStack(fosite.ErrServerError.WithWrap(err).WithDebug(err.Error()))
@@ -134,9 +147,9 @@ func (c *Handler) HandleTokenEndpointRequest(ctx context.Context, request fosite
 		return errorsx.WithStack(fosite.ErrServerError.WithWrap(err).WithDebug(err.Error()))
 	}
 
-	challenge := authorizeRequest.GetRequestForm().Get("code_challenge")
-	method := authorizeRequest.GetRequestForm().Get("code_challenge_method")
-	client := authorizeRequest.GetClient()
+	challenge := pkceRequest.GetRequestForm().Get("code_challenge")
+	method := pkceRequest.GetRequestForm().Get("code_challenge_method")
+	client := pkceRequest.GetClient()
 	if err := c.validate(ctx, challenge, method, client); err != nil {
 		return err
 	}
diff --git a/handler/pkce/handler_test.go b/handler/pkce/handler_test.go
index 5058cc7c2..d7dca4968 100644
--- a/handler/pkce/handler_test.go
+++ b/handler/pkce/handler_test.go
@@ -126,6 +126,7 @@ func TestPKCEHandlerValidate(t *testing.T) {
 			expectErr: fosite.ErrInvalidGrant,
 			client:    pc,
 			code:      "invalid-code-2",
+			verifier:  "foofoofoofoofoofoofoofoofoofoofoofoofoofoofoofoofoofoofoofoofoofoofoofoofoofoofoofoofoofoofoofoofoofoofoo",
 		},
 		{
 			d:      "passes because auth code flow but pkce is not forced and no challenge given",