From 4392c87cde108b0a9ca7569576aa5b55aa847d3b Mon Sep 17 00:00:00 2001
From: Bashir Sadjad <bashir@google.com>
Date: Thu, 31 Aug 2023 14:17:54 -0400
Subject: [PATCH] Fixed thread-safety and memory issues around JWT verifiers
 (#185)

---
 .../google/fhir/gateway/TokenVerifier.java    | 40 ++++++++++++-------
 1 file changed, 25 insertions(+), 15 deletions(-)

diff --git a/server/src/main/java/com/google/fhir/gateway/TokenVerifier.java b/server/src/main/java/com/google/fhir/gateway/TokenVerifier.java
index c88e6a30..5837174f 100644
--- a/server/src/main/java/com/google/fhir/gateway/TokenVerifier.java
+++ b/server/src/main/java/com/google/fhir/gateway/TokenVerifier.java
@@ -38,6 +38,8 @@
 import java.security.spec.InvalidKeySpecException;
 import java.security.spec.X509EncodedKeySpec;
 import java.util.Base64;
+import java.util.HashMap;
+import java.util.Map;
 import org.apache.http.HttpResponse;
 import org.apache.http.util.EntityUtils;
 import org.slf4j.Logger;
@@ -56,7 +58,12 @@ public class TokenVerifier {
   private static final String SIGN_ALGORITHM = "RS256";
 
   private final String tokenIssuer;
+  // Note the Verification class is _not_ thread-safe but the JWTVerifier instances created by its
+  // `build()` are thread-safe and reusable. It is important to reuse those instances, otherwise
+  // we may end up with a memory leak; details: https://github.com/auth0/java-jwt/issues/592
+  // Access to `jwtVerifierConfig` and `verifierForIssuer` should be non-concurrent.
   private final Verification jwtVerifierConfig;
+  private final Map<String, JWTVerifier> verifierForIssuer;
   private final HttpUtil httpUtil;
   private final String configJson;
 
@@ -68,6 +75,7 @@ public class TokenVerifier {
     RSAPublicKey issuerPublicKey = fetchAndDecodePublicKey();
     jwtVerifierConfig = JWT.require(Algorithm.RSA256(issuerPublicKey, null));
     this.configJson = httpUtil.fetchWellKnownConfig(tokenIssuer, wellKnownEndpoint);
+    this.verifierForIssuer = new HashMap<>();
   }
 
   public static TokenVerifier createFromEnvVars() throws IOException {
@@ -126,21 +134,23 @@ private RSAPublicKey fetchAndDecodePublicKey() throws IOException {
     return null;
   }
 
-  private JWTVerifier buildJwtVerifier(String issuer) {
-
-    if (tokenIssuer.equals(issuer)) {
-      return jwtVerifierConfig.withIssuer(tokenIssuer).build();
-    } else if (FhirProxyServer.isDevMode()) {
-      // If server is in DEV mode, set issuer to one from request
-      logger.warn("Server run in DEV mode. Setting issuer to issuer from request.");
-      return jwtVerifierConfig.withIssuer(issuer).build();
-    } else {
-      ExceptionUtil.throwRuntimeExceptionAndLog(
-          logger,
-          String.format("The token issuer %s does not match the expected token issuer", issuer),
-          AuthenticationException.class);
-      return null;
+  private synchronized JWTVerifier getJwtVerifier(String issuer) {
+    if (!tokenIssuer.equals(issuer)) {
+      if (FhirProxyServer.isDevMode()) {
+        // If server is in DEV mode, set issuer to one from request
+        logger.warn("Server run in DEV mode. Setting issuer to issuer from request.");
+      } else {
+        ExceptionUtil.throwRuntimeExceptionAndLog(
+            logger,
+            String.format("The token issuer %s does not match the expected token issuer", issuer),
+            AuthenticationException.class);
+        return null;
+      }
+    }
+    if (!verifierForIssuer.containsKey(issuer)) {
+      verifierForIssuer.put(issuer, jwtVerifierConfig.withIssuer(issuer).build());
     }
+    return verifierForIssuer.get(issuer);
   }
 
   @VisibleForTesting
@@ -161,7 +171,7 @@ DecodedJWT decodeAndVerifyBearerToken(String authHeader) {
     }
     String issuer = jwt.getIssuer();
     String algorithm = jwt.getAlgorithm();
-    JWTVerifier jwtVerifier = buildJwtVerifier(issuer);
+    JWTVerifier jwtVerifier = getJwtVerifier(issuer);
     logger.info(
         String.format(
             "JWT issuer is %s, audience is %s, and algorithm is %s",