fix: pr comment

supertokens · Sep 11, 2024 · 621befb · 621befb
1 parent ca95c13
commit 621befb
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 6 deletions.
diff --git a/src/main/java/io/supertokens/oauth/OAuth.java b/src/main/java/io/supertokens/oauth/OAuth.java
@@ -255,16 +255,16 @@ private static void checkNonSuccessResponse(HttpRequest.Response response) throw
         }
     }
 
-    public static JsonObject transformTokens(Main main, AppIdentifier appIdentifier, Storage storage, JsonObject jsonBody, String iss, boolean useDynamicKey) throws IOException, JWTException, InvalidKeyException, NoSuchAlgorithmException, StorageQueryException, StorageTransactionLogicException, UnsupportedJWTSigningAlgorithmException, TenantOrAppNotFoundException, InvalidKeySpecException, JWTCreationException, InvalidConfigException {
+    public static JsonObject transformTokens(Main main, AppIdentifier appIdentifier, Storage storage, JsonObject jsonBody, String iss, JsonObject accessTokenUpdate, JsonObject idTokenUpdate, boolean useDynamicKey) throws IOException, JWTException, InvalidKeyException, NoSuchAlgorithmException, StorageQueryException, StorageTransactionLogicException, UnsupportedJWTSigningAlgorithmException, TenantOrAppNotFoundException, InvalidKeySpecException, JWTCreationException, InvalidConfigException {
         if (jsonBody.has("access_token")) {
             String accessToken = jsonBody.get("access_token").getAsString();
-            accessToken = OAuthToken.reSignToken(appIdentifier, main, accessToken, iss, OAuthToken.TokenType.ACCESS_TOKEN, useDynamicKey, 0);
+            accessToken = OAuthToken.reSignToken(appIdentifier, main, accessToken, iss, accessTokenUpdate, OAuthToken.TokenType.ACCESS_TOKEN, useDynamicKey, 0);
             jsonBody.addProperty("access_token", accessToken);
         }
 
         if (jsonBody.has("id_token")) {
             String idToken = jsonBody.get("id_token").getAsString();
-            idToken = OAuthToken.reSignToken(appIdentifier, main, idToken, iss, OAuthToken.TokenType.ID_TOKEN, useDynamicKey, 0);
+            idToken = OAuthToken.reSignToken(appIdentifier, main, idToken, iss, idTokenUpdate, OAuthToken.TokenType.ID_TOKEN, useDynamicKey, 0);
             jsonBody.addProperty("id_token", idToken);
         }
 

diff --git a/src/main/java/io/supertokens/oauth/OAuthToken.java b/src/main/java/io/supertokens/oauth/OAuthToken.java
@@ -28,6 +28,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 public class OAuthToken {
     public enum TokenType {
@@ -45,6 +46,14 @@ public int getValue() {
         }
     }
 
+    private static Set<String> NON_OVERRIDABLE_TOKEN_PROPS = Set.of(
+        "kid", "typ", "alg", "aud",
+        "iss", "iat", "exp", "nbf", "jti", "ext",
+        "sid", "rat", "at_hash",
+        "client_id", "scp", "sub", "rsub",
+        "sessionHandle", "tId", "stt"
+    );
+
     public static JsonObject getPayloadFromJWTToken(AppIdentifier appIdentifier,
             @Nonnull Main main, @Nonnull String token)
             throws TenantOrAppNotFoundException, TryRefreshTokenException, StorageQueryException,
@@ -87,7 +96,7 @@ public static JsonObject getPayloadFromJWTToken(AppIdentifier appIdentifier,
         return jwtInfo.payload;
     }
 
-    public static String reSignToken(AppIdentifier appIdentifier, Main main, String token, String iss, TokenType tokenType, boolean useDynamicSigningKey, int retryCount) throws IOException, JWTException, InvalidKeyException, NoSuchAlgorithmException, StorageQueryException, StorageTransactionLogicException, UnsupportedJWTSigningAlgorithmException, TenantOrAppNotFoundException, InvalidKeySpecException,
+    public static String reSignToken(AppIdentifier appIdentifier, Main main, String token, String iss, JsonObject payloadUpdate, TokenType tokenType, boolean useDynamicSigningKey, int retryCount) throws IOException, JWTException, InvalidKeyException, NoSuchAlgorithmException, StorageQueryException, StorageTransactionLogicException, UnsupportedJWTSigningAlgorithmException, TenantOrAppNotFoundException, InvalidKeySpecException,
             JWTCreationException {
         // Load the JWKS from the specified URL
         JsonObject payload = JWT.getPayloadWithoutVerifying(token).payload;
@@ -103,6 +112,12 @@ public static String reSignToken(AppIdentifier appIdentifier, Main main, String
         payload.addProperty("iss", iss);
         payload.addProperty("stt", tokenType.getValue());
 
+        for (Map.Entry<String, JsonElement> entry : payloadUpdate.entrySet()) {
+            if (!NON_OVERRIDABLE_TOKEN_PROPS.contains(entry.getKey())) {
+                payload.add(entry.getKey(), entry.getValue());
+            }
+        }
+
         JWTSigningKeyInfo keyToUse;
         if (useDynamicSigningKey) {
             keyToUse = Utils.getJWTSigningKeyInfoFromKeyInfo(

diff --git a/src/main/java/io/supertokens/webserver/api/oauth/OAuthTokenAPI.java b/src/main/java/io/supertokens/webserver/api/oauth/OAuthTokenAPI.java
@@ -88,14 +88,17 @@ protected void doPost(HttpServletRequest req, HttpServletResponse resp) throws I
                         AppIdentifier appIdentifier = getAppIdentifier(req);
                         Storage storage = enforcePublicTenantAndGetPublicTenantStorage(req);
 
+                        JsonObject accessTokenUpdate = InputParser.parseJsonObjectOrThrowError(bodyFromSDK, "access_token", true);
+                        JsonObject idTokenUpdate = InputParser.parseJsonObjectOrThrowError(bodyFromSDK, "id_token", true);
+
                         // useStaticKeyInput defaults to true, so we check if it has been explicitly set to false
                         boolean useDynamicKey = Boolean.FALSE.equals(useStaticKeyInput);
-                        jsonBody = OAuth.transformTokens(super.main, appIdentifier, storage, jsonBody.getAsJsonObject(), iss, useDynamicKey);
+                        jsonBody = OAuth.transformTokens(super.main, appIdentifier, storage, jsonBody.getAsJsonObject(), iss, accessTokenUpdate, idTokenUpdate, useDynamicKey);
 
                     } catch (IOException | InvalidConfigException | TenantOrAppNotFoundException | BadPermissionException | StorageQueryException | InvalidKeyException | NoSuchAlgorithmException | InvalidKeySpecException | JWTCreationException | JWTException | StorageTransactionLogicException | UnsupportedJWTSigningAlgorithmException e) {
                         throw new ServletException(e);
                     }
-            
+
                     jsonBody.getAsJsonObject().addProperty("status", "OK");
                     super.sendJsonResponse(200, jsonBody, resp);
                 }