feat(jans-auth-server): provided convenient method to add claim to AT…

… as JWT in modifyAccessToken() method #3579 (#3629)
JanssenProject · Jan 16, 2023 · cf0a824 · cf0a824
1 parent 8d4783b
commit cf0a824
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 17 deletions.
diff --git a/jans-auth-server/server/src/main/java/io/jans/as/server/model/common/AuthorizationGrant.java b/jans-auth-server/server/src/main/java/io/jans/as/server/model/common/AuthorizationGrant.java
@@ -191,16 +191,26 @@ public AccessToken createAccessToken(ExecutionContext context) {
                 log.trace("Failed to create access token with negative expiration time");
                 return null;
             }
+
+            JwtSigner jwtSigner = null;
             if (getClient().isAccessTokenAsJwt()) {
-                accessToken.setCode(createAccessTokenAsJwt(accessToken, context));
+                jwtSigner = createAccessTokenAsJwt(accessToken, context);
             }
 
-            boolean externalOk = externalUpdateTokenService.modifyAccessToken(accessToken, ExternalUpdateTokenContext.of(context));
+            boolean externalOk = externalUpdateTokenService.modifyAccessToken(accessToken, ExternalUpdateTokenContext.of(context, jwtSigner));
             if (!externalOk) {
                 log.trace("External script forbids access token creation.");
                 return null;
             }
 
+            if (getClient().isAccessTokenAsJwt() && jwtSigner != null) {
+                final String accessTokenCode = jwtSigner.sign().toString();
+                if (log.isTraceEnabled())
+                    log.trace("Created access token JWT: {}", accessTokenCode + ", claims: " + jwtSigner.getJwt().getClaims().toJsonString());
+
+                accessToken.setCode(accessTokenCode);
+            }
+
             final TokenEntity tokenEntity = asToken(accessToken);
             context.setAccessTokenEntity(tokenEntity);
 
@@ -218,7 +228,7 @@ public AccessToken createAccessToken(ExecutionContext context) {
         }
     }
 
-    public String createAccessTokenAsJwt(AccessToken accessToken, ExecutionContext context) throws Exception {
+    public JwtSigner createAccessTokenAsJwt(AccessToken accessToken, ExecutionContext context) throws Exception {
         final User user = getUser();
         final Client client = getClient();
 
@@ -257,11 +267,7 @@ public String createAccessTokenAsJwt(AccessToken accessToken, ExecutionContext c
             runIntrospectionScriptAndInjectValuesIntoJwt(jwt, context);
         }
 
-        final String accessTokenCode = jwtSigner.sign().toString();
-        if (log.isTraceEnabled())
-            log.trace("Created access token JWT: {}", accessTokenCode + ", claims: " + jwt.getClaims().toJsonString());
-
-        return accessTokenCode;
+        return jwtSigner;
     }
 
     private void runIntrospectionScriptAndInjectValuesIntoJwt(Jwt jwt, ExecutionContext executionContext) {

diff --git a/.../src/main/java/io/jans/as/server/service/external/context/ExternalUpdateTokenContext.java b/.../src/main/java/io/jans/as/server/service/external/context/ExternalUpdateTokenContext.java
@@ -6,18 +6,20 @@
 
 package io.jans.as.server.service.external.context;
 
+import com.google.common.collect.Lists;
 import io.jans.as.common.model.registration.Client;
 import io.jans.as.common.service.AttributeService;
 import io.jans.as.model.common.GrantType;
 import io.jans.as.model.configuration.AppConfiguration;
 import io.jans.as.model.jwt.Jwt;
+import io.jans.as.model.jwt.JwtClaims;
 import io.jans.as.server.model.common.AccessToken;
 import io.jans.as.server.model.common.AuthorizationGrant;
 import io.jans.as.server.model.common.ExecutionContext;
+import io.jans.as.server.model.token.JwtSigner;
 import io.jans.model.custom.script.conf.CustomScriptConfiguration;
-import org.jetbrains.annotations.Nullable;
-
 import jakarta.servlet.http.HttpServletRequest;
+import org.jetbrains.annotations.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -39,6 +41,7 @@ public class ExternalUpdateTokenContext extends ExternalScriptContext {
     private CustomScriptConfiguration script;
     @Nullable
     private ExecutionContext executionContext;
+    private JwtSigner jwtSigner;
 
     public ExternalUpdateTokenContext(HttpServletRequest httpRequest, AuthorizationGrant grant,
                                       Client client, AppConfiguration appConfiguration, AttributeService attributeService) {
@@ -50,8 +53,13 @@ public ExternalUpdateTokenContext(HttpServletRequest httpRequest, AuthorizationG
     }
 
     public static ExternalUpdateTokenContext of(ExecutionContext executionContext) {
+        return of(executionContext, null);
+    }
+
+    public static ExternalUpdateTokenContext of(ExecutionContext executionContext, JwtSigner jwtSigner) {
         ExternalUpdateTokenContext context = new ExternalUpdateTokenContext(executionContext.getHttpRequest(), executionContext.getGrant(), executionContext.getClient(), executionContext.getAppConfiguration(), executionContext.getAttributeService());
         context.setExecutionContext(executionContext);
+        context.setJwtSigner(jwtSigner);
         return context;
     }
 
@@ -72,6 +80,23 @@ private ExecutionContext createExecutionContext() {
         return result;
     }
 
+    public JwtClaims getClaims() {
+        Jwt jwt = getJwt();
+        return jwt != null ? jwt.getClaims() : null;
+    }
+
+    public Jwt getJwt() {
+        return jwtSigner != null ? jwtSigner.getJwt() : null;
+    }
+
+    public JwtSigner getJwtSigner() {
+        return jwtSigner;
+    }
+
+    public void setJwtSigner(JwtSigner jwtSigner) {
+        this.jwtSigner = jwtSigner;
+    }
+
     public CustomScriptConfiguration getScript() {
         return script;
     }
@@ -117,13 +142,9 @@ public void overwriteAccessTokenScopes(AccessToken accessToken, Set<String> newS
 
         grant.setScopes(newScopes);
 
-        // re-generate access token jwt to put new scopes into jwt
-        if (isValidJwt(accessToken.getCode())) {
-            try {
-                accessToken.setCode(grant.createAccessTokenAsJwt(accessToken, executionContext));
-            } catch (Exception e) {
-                log.error("Failed to generate access token jwt", e);
-            }
+        final Jwt jwt = getJwt();
+        if (jwt != null) {
+            jwt.getClaims().setClaim("scope", Lists.newArrayList(newScopes));
         }
     }