From 5b61bcc4e03455e339ae1bfabd2ba0e22b1c74c5 Mon Sep 17 00:00:00 2001
From: Pedro Silva <pedro@acryl.io>
Date: Fri, 15 Jul 2022 11:58:59 +0100
Subject: [PATCH] fix(gms): Change MessageDigest to be thread safe (#5405)

---
 .../token/StatefulTokenService.java           | 33 ++++---------------
 1 file changed, 7 insertions(+), 26 deletions(-)

diff --git a/metadata-service/auth-impl/src/main/java/com/datahub/authentication/token/StatefulTokenService.java b/metadata-service/auth-impl/src/main/java/com/datahub/authentication/token/StatefulTokenService.java
index 19efd4f613..33d89e2dfa 100644
--- a/metadata-service/auth-impl/src/main/java/com/datahub/authentication/token/StatefulTokenService.java
+++ b/metadata-service/auth-impl/src/main/java/com/datahub/authentication/token/StatefulTokenService.java
@@ -16,13 +16,6 @@ import com.linkedin.metadata.resources.entity.AspectUtils;
 import com.linkedin.metadata.utils.AuditStampUtils;
 import com.linkedin.metadata.utils.GenericRecordUtils;
 import com.linkedin.mxe.MetadataChangeProposal;
-import lombok.extern.slf4j.Slf4j;
-import org.apache.commons.lang.ArrayUtils;
-
-import javax.annotation.Nonnull;
-import javax.annotation.Nullable;
-import java.security.MessageDigest;
-import java.security.NoSuchAlgorithmException;
 import java.util.Base64;
 import java.util.Date;
 import java.util.HashMap;
@@ -31,11 +24,13 @@ import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+import lombok.extern.slf4j.Slf4j;
+import org.apache.commons.codec.digest.DigestUtils;
+import org.apache.commons.lang.ArrayUtils;
 
-import static com.datahub.authentication.token.TokenClaims.ACTOR_ID_CLAIM_NAME;
-import static com.datahub.authentication.token.TokenClaims.ACTOR_TYPE_CLAIM_NAME;
-import static com.datahub.authentication.token.TokenClaims.TOKEN_TYPE_CLAIM_NAME;
-import static com.datahub.authentication.token.TokenClaims.TOKEN_VERSION_CLAIM_NAME;
+import static com.datahub.authentication.token.TokenClaims.*;
 
 
 /**
@@ -48,7 +43,6 @@ public class StatefulTokenService extends StatelessTokenService {
   private final EntityService _entityService;
   private final LoadingCache<String, Boolean> _revokedTokenCache;
   private final String salt;
-  private final MessageDigest sha256;
 
   public StatefulTokenService(@Nonnull final String signingKey, @Nonnull final String signingAlgorithm,
       @Nullable final String iss, @Nonnull final EntityService entityService, @Nonnull final String salt) {
@@ -65,11 +59,6 @@ public class StatefulTokenService extends StatelessTokenService {
           }
         });
     this.salt = salt;
-    try {
-      this.sha256 = MessageDigest.getInstance("SHA-256");
-    } catch (NoSuchAlgorithmException e) {
-      throw new RuntimeException("Unable to get SHA-256 algorithm.");
-    }
   }
 
   /**
@@ -179,14 +168,6 @@ public class StatefulTokenService extends StatelessTokenService {
     throw new TokenException("Access token no longer exists");
   }
 
-  public boolean isTokenRevoked(@Nonnull String hashToken) {
-    try {
-      return _revokedTokenCache.get(hashToken);
-    } catch (ExecutionException e) {
-      return false;
-    }
-  }
-
   /**
    * Hashes the input after salting it.
    */
@@ -194,7 +175,7 @@ public class StatefulTokenService extends StatelessTokenService {
     final byte[] saltingKeyBytes = this.salt.getBytes();
     final byte[] inputBytes = input.getBytes();
     final byte[] concatBytes = ArrayUtils.addAll(inputBytes, saltingKeyBytes);
-    final byte[] bytes = sha256.digest(concatBytes);
+    final byte[] bytes = DigestUtils.sha256(concatBytes);
     return Base64.getEncoder().encodeToString(bytes);
   }
 }