From 7a767d8568772197503dca7f90b5a767f6f23572 Mon Sep 17 00:00:00 2001
From: BitMasher <61257372+BitMasher@users.noreply.github.com>
Date: Wed, 12 Mar 2025 07:00:29 -0500
Subject: [PATCH] feat: add CanGetPrivateClaimsFromRequest interface (#717)

---
 pkg/op/storage.go | 6 ++++++
 pkg/op/token.go   | 6 +++++-
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/pkg/op/storage.go b/pkg/op/storage.go
index 8488b28..a579810 100644
--- a/pkg/op/storage.go
+++ b/pkg/op/storage.go
@@ -144,6 +144,12 @@ type CanSetUserinfoFromRequest interface {
 	SetUserinfoFromRequest(ctx context.Context, userinfo *oidc.UserInfo, request IDTokenRequest, scopes []string) error
 }
 
+// CanGetPrivateClaimsFromRequest is an optional additional interface that may be implemented by
+// implementors of Storage. It allows setting the jwt token claims based on the request.
+type CanGetPrivateClaimsFromRequest interface {
+	GetPrivateClaimsFromRequest(ctx context.Context, request TokenRequest, restrictedScopes []string) (map[string]any, error)
+}
+
 // Storage is a required parameter for NewOpenIDProvider(). In addition to the
 // embedded interfaces below, if the passed Storage implements ClientCredentialsStorage
 // then the grant type "client_credentials" will be supported. In that case, the access
diff --git a/pkg/op/token.go b/pkg/op/token.go
index 04cd3cc..1df9cc2 100644
--- a/pkg/op/token.go
+++ b/pkg/op/token.go
@@ -147,7 +147,11 @@ func CreateJWT(ctx context.Context, issuer string, tokenRequest TokenRequest, ex
 				tokenExchangeRequest,
 			)
 		} else {
-			privateClaims, err = storage.GetPrivateClaimsFromScopes(ctx, tokenRequest.GetSubject(), client.GetID(), removeUserinfoScopes(restrictedScopes))
+			if fromRequest, ok := storage.(CanGetPrivateClaimsFromRequest); ok {
+				privateClaims, err = fromRequest.GetPrivateClaimsFromRequest(ctx, tokenRequest, removeUserinfoScopes(restrictedScopes))
+			} else {
+				privateClaims, err = storage.GetPrivateClaimsFromScopes(ctx, tokenRequest.GetSubject(), client.GetID(), removeUserinfoScopes(restrictedScopes))
+			}
 		}
 
 		if err != nil {