From 313d6c1c41dca4e238f9179bfafdcc3378f6f81e Mon Sep 17 00:00:00 2001
From: Mevan <mevan.karu@gmail.com>
Date: Thu, 3 Oct 2024 17:20:41 +0530
Subject: [PATCH] Add support to send PAT JWT to backend

---
 .../choreo/connect/enforcer/security/jwt/APIKeyUtils.java | 5 ++---
 .../connect/enforcer/security/jwt/JWTAuthenticator.java   | 8 +++++---
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/enforcer-parent/enforcer/src/main/java/org/wso2/choreo/connect/enforcer/security/jwt/APIKeyUtils.java b/enforcer-parent/enforcer/src/main/java/org/wso2/choreo/connect/enforcer/security/jwt/APIKeyUtils.java
index f6155bf2da..55e1c021ca 100644
--- a/enforcer-parent/enforcer/src/main/java/org/wso2/choreo/connect/enforcer/security/jwt/APIKeyUtils.java
+++ b/enforcer-parent/enforcer/src/main/java/org/wso2/choreo/connect/enforcer/security/jwt/APIKeyUtils.java
@@ -97,10 +97,10 @@ public static String generateAPIKeyHash(String apiKey) {
     /**
      * This function exchanges a given API key to an JWT token.
      *
-     * @param pat    PAT
+     * @param keyHash    Key Hash
      * @return JWT corresponding to given PAT.
      */
-    public static Optional<String> exchangePATToJWT(String pat) {
+    public static Optional<String> exchangePATToJWT(String keyHash) {
 
         URL url = null;
         try {
@@ -115,7 +115,6 @@ public static Optional<String> exchangePATToJWT(String pat) {
             // Create a request to exchange API key to JWT.
             HttpPost exchangeRequest = new HttpPost(url.toURI());
             exchangeRequest.addHeader("Content-Type", ContentType.APPLICATION_JSON.toString());
-            String keyHash = generateAPIKeyHash(pat);
             exchangeRequest.setEntity(new StringEntity(createPATExchangeRequest(keyHash)));
             try (CloseableHttpResponse response = httpClient.execute(exchangeRequest)) {
                 if (response.getStatusLine().getStatusCode() == 200) {
diff --git a/enforcer-parent/enforcer/src/main/java/org/wso2/choreo/connect/enforcer/security/jwt/JWTAuthenticator.java b/enforcer-parent/enforcer/src/main/java/org/wso2/choreo/connect/enforcer/security/jwt/JWTAuthenticator.java
index 6e9b8c3d09..ea467e6ee7 100644
--- a/enforcer-parent/enforcer/src/main/java/org/wso2/choreo/connect/enforcer/security/jwt/JWTAuthenticator.java
+++ b/enforcer-parent/enforcer/src/main/java/org/wso2/choreo/connect/enforcer/security/jwt/JWTAuthenticator.java
@@ -193,7 +193,7 @@ public AuthenticationContext authenticate(RequestContext requestContext) throws
             }
             // Handle PAT logic
             if (isPATEnabled && token.startsWith(APIKeyConstants.PAT_PREFIX)) {
-                token = exchangeJWTForPAT(token);
+                token = exchangeJWTForPAT(requestContext, token);
             }
             String context = requestContext.getMatchedAPI().getBasePath();
             String name = requestContext.getMatchedAPI().getName();
@@ -806,7 +806,7 @@ private String getJWTTokenIdentifier(SignedJWTInfo signedJWTInfo) {
         return signedJWTInfo.getSignedJWT().getSignature().toString();
     }
 
-    private String exchangeJWTForPAT(String pat) throws APISecurityException {
+    private String exchangeJWTForPAT(RequestContext requestContext, String pat) throws APISecurityException {
         if (!APIKeyUtils.isValidAPIKey(pat)) {
             throw new APISecurityException(APIConstants.StatusCodes.UNAUTHENTICATED.getCode(),
                     APISecurityConstants.API_AUTH_INVALID_CREDENTIALS,
@@ -820,13 +820,15 @@ private String exchangeJWTForPAT(String pat) throws APISecurityException {
             }
             return (String) cachedJWT;
         }
-        Optional<String> jwt = APIKeyUtils.exchangePATToJWT(pat);
+        Optional<String> jwt = APIKeyUtils.exchangePATToJWT(keyHash);
         if (jwt.isEmpty()) {
             throw new APISecurityException(APIConstants.StatusCodes.UNAUTHENTICATED.getCode(),
                     APISecurityConstants.API_AUTH_INVALID_CREDENTIALS,
                     APISecurityConstants.API_AUTH_INVALID_CREDENTIALS_MESSAGE);
         }
         CacheProvider.getGatewayAPIKeyJWTCache().put(keyHash, jwt.get());
+        // Add jwt to x-forwarded-authorization header.
+        requestContext.addOrModifyHeaders("x-forwarded-authorization", jwt.get());
         return jwt.get();
     }