fix(jans-auth-server): UpdateToken script is not invoked during Implicit Flow #6561

docs/admin/developer/scripts/update-token.md

@@ -204,13 +204,13 @@ Pseudocode and example - Issue Access token only if account balance is greater t
     def modifyAccessToken(self, accessToken, context):
 
 	    # header claims
-	    accessToken.getHeader().setClaim("header_name", "header_value")
+	    context.getHeader().setClaim("header_name", "header_value")
 
 	    #custom claims
-	    accessToken.getClaims().setClaim("claim_name", "claimValue")
+	    context.getClaims().setClaim("claim_name", "claimValue")
 
 	    #regular claims        
-	    accessToken.getClaims().setClaim("sub", claimValue)
+	    context.getClaims().setClaim("sub", claimValue)
 
 	    return True
 

jans-auth-server/server/src/main/java/io/jans/as/server/model/common/AuthorizationGrant.java

@@ -109,6 +109,8 @@ public void init(User user, AuthorizationGrantType authorizationGrantType, Clien
     }
 
     private IdToken createIdTokenInternal(AuthorizationCode authorizationCode, AccessToken accessToken, RefreshToken refreshToken, ExecutionContext executionContext) throws Exception {
+        executionContext.initFromGrantIfNeeded(this);
+
         JsonWebResponse jwr = idTokenFactory.createJwr(this, authorizationCode, accessToken, refreshToken, executionContext);
         final IdToken idToken = new IdToken(jwr.toString(), jwr.getClaims().getClaimAsDate(JwtClaimName.ISSUED_AT),
                 jwr.getClaims().getClaimAsDate(JwtClaimName.EXPIRATION_TIME));
@@ -189,6 +191,8 @@ private void initTokenFromGrant(TokenEntity token) {
     @Override
     public AccessToken createAccessToken(ExecutionContext context) {
         try {
+            context.initFromGrantIfNeeded(this);
+
             final AccessToken accessToken = super.createAccessToken(context);
             if (accessToken.getExpiresIn() < 0) {
                 log.trace("Failed to create access token with negative expiration time");
@@ -237,6 +241,8 @@ public JwtSigner createAccessTokenAsJwt(AccessToken accessToken, ExecutionContex
         final User user = getUser();
         final Client client = getClient();
 
+        context.initFromGrantIfNeeded(this);
+
         SignatureAlgorithm signatureAlgorithm = SignatureAlgorithm
                 .fromString(appConfiguration.getDefaultSignatureAlgorithm());
         if (client.getAccessTokenSigningAlg() != null
@@ -278,6 +284,8 @@ public JwtSigner createAccessTokenAsJwt(AccessToken accessToken, ExecutionContex
     }
 
     private void runIntrospectionScriptAndInjectValuesIntoJwt(Jwt jwt, ExecutionContext executionContext) {
+        executionContext.initFromGrantIfNeeded(this);
+
         JSONObject responseAsJsonObject = new JSONObject();
 
         ExternalIntrospectionContext context = new ExternalIntrospectionContext(this, executionContext.getHttpRequest(), executionContext.getHttpResponse(), appConfiguration, attributeService);
@@ -295,6 +303,8 @@ private void runIntrospectionScriptAndInjectValuesIntoJwt(Jwt jwt, ExecutionCont
 
     private RefreshToken saveRefreshToken(RefreshToken refreshToken, ExecutionContext executionContext) {
         try {
+            executionContext.initFromGrantIfNeeded(this);
+
             if (refreshToken.getExpiresIn() > 0) {
                 final TokenEntity entity = asToken(refreshToken);
                 executionContext.setRefreshTokenEntity(entity);
@@ -339,11 +349,13 @@ private RefreshToken saveRefreshToken(Supplier<RefreshToken> supplier, Execution
 
     @Override
     public RefreshToken createRefreshToken(ExecutionContext context) {
+        context.initFromGrantIfNeeded(this);
         return saveRefreshToken(() -> super.createRefreshToken(context), context);
     }
 
     @Override
     public RefreshToken createRefreshToken(ExecutionContext context, int lifetime) {
+        context.initFromGrantIfNeeded(this);
         return saveRefreshToken(() -> super.createRefreshToken(context, lifetime), context);
     }
 
@@ -361,6 +373,7 @@ public IdToken createIdToken(
             String nonce, AuthorizationCode authorizationCode, AccessToken accessToken, RefreshToken refreshToken,
             String state, ExecutionContext executionContext) {
         try {
+            executionContext.initFromGrantIfNeeded(this);
             executionContext.setScopes(getScopes());
             executionContext.setClaimsAsString(getClaims());
             executionContext.setNonce(nonce);

jans-auth-server/server/src/main/java/io/jans/as/server/model/common/ExecutionContext.java

@@ -334,4 +334,13 @@ public Response.ResponseBuilder getResponseBuilder() {
     public void setResponseBuilder(Response.ResponseBuilder responseBuilder) {
         this.responseBuilder = responseBuilder;
     }
+
+    public void initFromGrantIfNeeded(AuthorizationGrant authorizationGrant) {
+        if (client == null) {
+            client = authorizationGrant.getClient();
+        }
+        if (grant == null) {
+            grant = authorizationGrant;
+        }
+    }
 }

jans-auth-server/server/src/main/java/io/jans/as/server/service/external/ExternalUpdateTokenService.java

@@ -119,6 +119,7 @@ public int getRefreshTokenLifetimeInSeconds(ExternalUpdateTokenContext context)
     @NotNull
     private List<CustomScriptConfiguration> getScripts(@NotNull ExternalUpdateTokenContext context) {
         if (customScriptConfigurations == null || customScriptConfigurations.isEmpty() || context.getClient() == null) {
+            log.trace("No UpdateToken scripts or client is null.");
             return Lists.newArrayList();
         }