From 5d716553c6eb409c2f264864da8b65c0a0bcbe81 Mon Sep 17 00:00:00 2001
From: YuriyZ <yzabrovarniy@gmail.com>
Date: Thu, 23 Jun 2022 22:55:29 +0300
Subject: [PATCH] feat(jans-auth-server): added convenient method for
 up-scoping or down-scoping AT scopes #1218

https://github.com/JanssenProject/jans/issues/1218
---
 .../model/common/AuthorizationGrant.java      | 14 ++++-----
 .../context/ExternalUpdateTokenContext.java   | 30 +++++++++++++++++++
 2 files changed, 37 insertions(+), 7 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
index 13f2be3bc10..9e58193edfb 100644
--- 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
@@ -187,16 +187,13 @@ private void initTokenFromGrant(TokenEntity token) {
     public AccessToken createAccessToken(ExecutionContext context) {
         try {
             final AccessToken accessToken = super.createAccessToken(context);
-            if (getClient().isAccessTokenAsJwt()) {
-                accessToken.setCode(createAccessTokenAsJwt(accessToken, context));
-            }
             if (accessToken.getExpiresIn() < 0) {
                 log.trace("Failed to create access token with negative expiration time");
                 return null;
             }
-
-            final TokenEntity tokenEntity = asToken(accessToken);
-            context.setAccessTokenEntity(tokenEntity);
+            if (getClient().isAccessTokenAsJwt()) {
+                accessToken.setCode(createAccessTokenAsJwt(accessToken, context));
+            }
 
             boolean externalOk = externalUpdateTokenService.modifyAccessToken(accessToken, ExternalUpdateTokenContext.of(context));
             if (!externalOk) {
@@ -204,6 +201,9 @@ public AccessToken createAccessToken(ExecutionContext context) {
                 return null;
             }
 
+            final TokenEntity tokenEntity = asToken(accessToken);
+            context.setAccessTokenEntity(tokenEntity);
+
             persist(tokenEntity);
             statService.reportAccessToken(getGrantType());
             metricService.incCounter(MetricType.TOKEN_ACCESS_TOKEN_COUNT);
@@ -218,7 +218,7 @@ public AccessToken createAccessToken(ExecutionContext context) {
         }
     }
 
-    private String createAccessTokenAsJwt(AccessToken accessToken, ExecutionContext context) throws Exception {
+    public String createAccessTokenAsJwt(AccessToken accessToken, ExecutionContext context) throws Exception {
         final User user = getUser();
         final Client client = getClient();
 
diff --git a/jans-auth-server/server/src/main/java/io/jans/as/server/service/external/context/ExternalUpdateTokenContext.java b/jans-auth-server/server/src/main/java/io/jans/as/server/service/external/context/ExternalUpdateTokenContext.java
index ca2644d1455..6aba967e9b3 100644
--- a/jans-auth-server/server/src/main/java/io/jans/as/server/service/external/context/ExternalUpdateTokenContext.java
+++ b/jans-auth-server/server/src/main/java/io/jans/as/server/service/external/context/ExternalUpdateTokenContext.java
@@ -10,18 +10,26 @@
 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.server.model.common.AccessToken;
 import io.jans.as.server.model.common.AuthorizationGrant;
 import io.jans.as.server.model.common.ExecutionContext;
 import io.jans.model.custom.script.conf.CustomScriptConfiguration;
 import org.jetbrains.annotations.Nullable;
 
 import jakarta.servlet.http.HttpServletRequest;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Set;
 
 /**
  * @author Yuriy Movchan
  */
 public class ExternalUpdateTokenContext extends ExternalScriptContext {
 
+    private static final Logger log = LoggerFactory.getLogger(ExternalUpdateTokenContext.class);
+
     private final Client client;
     private final AuthorizationGrant grant;
 
@@ -100,4 +108,26 @@ public ExecutionContext getExecutionContext() {
     public void setExecutionContext(@Nullable ExecutionContext executionContext) {
         this.executionContext = executionContext;
     }
+
+    // Usually expected to be called in : "def modifyAccessToken(self, accessToken, context):"
+    public void overwriteAccessTokenScopes(AccessToken accessToken, Set<String> newScopes) {
+        if (grant == null) {
+            return;
+        }
+
+        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);
+            }
+        }
+    }
+
+    private boolean isValidJwt(String jwt) {
+        return Jwt.parseSilently(jwt) != null;
+    }
 }