Authenticate with Microsoft Entra ID using Token Cache

thjarvin · thjarvin · commit ef60ff2bc83a · 2025-06-06T15:28:50.000+03:00
diff --git a/pom.xml b/pom.xml
@@ -78,10 +78,16 @@
             <version>${pulsar.version}</version>
         </dependency>
 
+        <dependency>
+            <groupId>com.azure</groupId>
+            <artifactId>azure-identity</artifactId>
+            <version>1.11.2</version>
+        </dependency>
+
         <dependency>
             <groupId>redis.clients</groupId>
             <artifactId>jedis</artifactId>
-            <version>4.4.3</version>
+            <version>5.1.0</version>
             <type>jar</type>
             <scope>compile</scope>
         </dependency>
diff --git a/src/main/java/fi/hsl/common/pulsar/PulsarApplication.java b/src/main/java/fi/hsl/common/pulsar/PulsarApplication.java
@@ -1,14 +1,20 @@
 package fi.hsl.common.pulsar;
 
+import com.azure.core.credential.AccessToken;
+import com.azure.core.credential.TokenRequestContext;
+import com.azure.identity.DefaultAzureCredential;
+import com.azure.identity.DefaultAzureCredentialBuilder;
 import com.typesafe.config.Config;
 import fi.hsl.common.health.HealthServer;
+import fi.hsl.common.redis.RedisUtils;
 import org.apache.pulsar.client.admin.PulsarAdmin;
 import org.apache.pulsar.client.api.*;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import redis.clients.jedis.Jedis;
+import redis.clients.jedis.exceptions.JedisException;
 
 import java.util.Arrays;
 import java.util.HashMap;
@@ -19,6 +25,9 @@
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 
+import static fi.hsl.common.redis.RedisUtils.createJedisClient;
+import static fi.hsl.common.redis.RedisUtils.extractUsernameFromToken;
+
 public class PulsarApplication implements AutoCloseable {
     private static final Logger log = LoggerFactory.getLogger(PulsarApplication.class);
 
@@ -152,9 +161,54 @@ public PulsarApplicationContext initialize(@NotNull Config config) throws Except
     @NotNull
     protected Jedis createRedisClient(@NotNull String redisHost, int port, int connTimeOutSecs) {
         log.info("Connecting to Redis at " + redisHost + ":" + port + " with connection timeout of (s): "+ connTimeOutSecs);
-        int timeOutMs = connTimeOutSecs * 1000;
-        Jedis jedis = new Jedis(redisHost, port, timeOutMs);
-        jedis.connect();
+        
+        //Construct a Token Credential from Identity library, e.g. DefaultAzureCredential / ClientSecretCredential / Client CertificateCredential / ManagedIdentityCredential etc.
+        DefaultAzureCredential defaultAzureCredential = new DefaultAzureCredentialBuilder().build();
+
+        // Fetch a Microsoft Entra token to be used for authentication. This token will be used as the password.
+        TokenRequestContext trc = new TokenRequestContext().addScopes("https://redis.azure.com/.default");
+        RedisUtils.TokenRefreshCache tokenRefreshCache = new RedisUtils.TokenRefreshCache(defaultAzureCredential, trc);
+        AccessToken accessToken = tokenRefreshCache.getAccessToken();
+
+        // SSL connection is required.
+        boolean useSsl = true;
+        String username = extractUsernameFromToken(accessToken.getToken());
+        
+        // Create Jedis client and connect to the Azure Cache for Redis over the TLS/SSL port using the access token as password.
+        // Note: Cache Host Name, Port, Microsoft Entra access token and SSL connections are required below.
+        Jedis jedis = createJedisClient(redisHost, port, username, accessToken, useSsl);
+
+        // Configure the jedis instance for proactive authentication before token expires.
+        tokenRefreshCache.setJedisInstanceToAuthenticate(jedis);
+        
+        int maxTries = 3;
+        int i = 0;
+        
+        while (i < maxTries) {
+            try {
+                // Set a value against your key in the Redis cache.
+                jedis.set("Az:key", "testValue");
+                System.out.println(jedis.get("Az:key"));
+                break;
+            } catch (JedisException e) {
+                // Handle The Exception as required in your application.
+                e.printStackTrace();
+                
+                // For Exceptions containing Invalid Username Password / Permissions not granted error messages, look at troubleshooting section at the end of document.
+                
+                // Check if the client is broken, if it is then close and recreate it to create a new healthy connection.
+                if (jedis.isBroken()) {
+                    jedis.close();
+                    accessToken = tokenRefreshCache.getAccessToken();
+                    jedis = createJedisClient(redisHost, port, username, accessToken, useSsl);
+                    
+                    // Configure the jedis instance for proactive authentication before token expires.
+                    tokenRefreshCache.setJedisInstanceToAuthenticate(jedis);
+                }
+            }
+            i++;
+        }
+        
         log.info("Redis connected: " + jedis.isConnected());
         return jedis;
     }
diff --git a/src/main/java/fi/hsl/common/redis/RedisUtils.java b/src/main/java/fi/hsl/common/redis/RedisUtils.java
@@ -1,5 +1,11 @@
 package fi.hsl.common.redis;
 
+import com.azure.core.credential.AccessToken;
+import com.azure.core.credential.TokenCredential;
+import com.azure.core.credential.TokenRequestContext;
+import com.azure.core.util.CoreUtils;
+import com.google.gson.JsonObject;
+import com.google.gson.JsonParser;
 import fi.hsl.common.pulsar.PulsarApplicationContext;
 import fi.hsl.common.transitdata.TransitdataProperties;
 import org.jetbrains.annotations.NotNull;
@@ -10,9 +16,12 @@
 import redis.clients.jedis.params.ScanParams;
 import redis.clients.jedis.resps.ScanResult;
 
+import java.nio.charset.StandardCharsets;
+import java.time.Duration;
 import java.time.OffsetDateTime;
 import java.time.format.DateTimeFormatter;
 import java.util.*;
+import java.util.concurrent.ThreadLocalRandom;
 
 public class RedisUtils {
     private static final Logger log = LoggerFactory.getLogger(RedisUtils.class);
@@ -227,4 +236,105 @@ public boolean checkResponse(@Nullable final String response) {
     public boolean checkResponse(@Nullable final Long response) {
         return response != null && response == 1;
     }
+    
+    // Azure Cache for Redis helper code
+    public static Jedis createJedisClient(String cacheHostname, int port, String username, AccessToken accessToken, boolean useSsl) {
+        return new Jedis(cacheHostname, port, DefaultJedisClientConfig.builder()
+                .password(accessToken.getToken())
+                .user(username)
+                .ssl(useSsl)
+                .build());
+    }
+    
+    public static String extractUsernameFromToken(String token) {
+        String[] parts = token.split("\\.");
+        String base64 = parts[1];
+        
+        switch (base64.length() % 4) {
+            case 2:
+                base64 += "==";
+                break;
+            case 3:
+                base64 += "=";
+                break;
+        }
+        
+        byte[] jsonBytes = Base64.getDecoder().decode(base64);
+        String json = new String(jsonBytes, StandardCharsets.UTF_8);
+        JsonObject jwt = JsonParser.parseString(json).getAsJsonObject();
+        
+        return jwt.get("oid").getAsString();
+    }
+    
+    /**
+     * The token cache to store and proactively refresh the access token.
+     */
+    public static class TokenRefreshCache {
+        private final TokenCredential tokenCredential;
+        private final TokenRequestContext tokenRequestContext;
+        private final Timer timer;
+        private volatile AccessToken accessToken;
+        private final Duration maxRefreshOffset = Duration.ofMinutes(5);
+        private final Duration baseRefreshOffset = Duration.ofMinutes(2);
+        private Jedis jedisInstanceToAuthenticate;
+        private String username;
+        
+        /**
+         * Creates an instance of TokenRefreshCache
+         * @param tokenCredential the token credential to be used for authentication.
+         * @param tokenRequestContext the token request context to be used for authentication.
+         */
+        public TokenRefreshCache(TokenCredential tokenCredential, TokenRequestContext tokenRequestContext) {
+            this.tokenCredential = tokenCredential;
+            this.tokenRequestContext = tokenRequestContext;
+            this.timer = new Timer();
+        }
+        
+        /**
+         * Gets the cached access token.
+         * @return the AccessToken
+         */
+        public AccessToken getAccessToken() {
+            if (accessToken != null) {
+                return  accessToken;
+            } else {
+                TokenRefreshTask tokenRefreshTask = new TokenRefreshTask();
+                accessToken = tokenCredential.getToken(tokenRequestContext).block();
+                timer.schedule(tokenRefreshTask, getTokenRefreshDelay());
+                return accessToken;
+            }
+        }
+        
+        private class TokenRefreshTask extends TimerTask {
+            // Add your task here
+            public void run() {
+                accessToken = tokenCredential.getToken(tokenRequestContext).block();
+                username = extractUsernameFromToken(accessToken.getToken());
+                System.out.println("Refreshed Token with Expiry: " + accessToken.getExpiresAt().toEpochSecond());
+                
+                if (jedisInstanceToAuthenticate != null && !CoreUtils.isNullOrEmpty(username)) {
+                    jedisInstanceToAuthenticate.auth(username, accessToken.getToken());
+                    System.out.println("Refreshed Jedis Connection with fresh access token, token expires at : "
+                            + accessToken.getExpiresAt().toEpochSecond());
+                }
+                timer.schedule(new TokenRefreshTask(), getTokenRefreshDelay());
+            }
+        }
+        
+        private long getTokenRefreshDelay() {
+            return ((accessToken.getExpiresAt()
+                    .minusSeconds(ThreadLocalRandom.current().nextLong(baseRefreshOffset.getSeconds(), maxRefreshOffset.getSeconds()))
+                    .toEpochSecond() - OffsetDateTime.now().toEpochSecond()) * 1000);
+        }
+        
+        /**
+         * Sets the Jedis to proactively authenticate before token expiry.
+         * @param jedisInstanceToAuthenticate the instance to authenticate
+         * @return the updated instance
+         */
+        public TokenRefreshCache setJedisInstanceToAuthenticate(Jedis jedisInstanceToAuthenticate) {
+            this.jedisInstanceToAuthenticate = jedisInstanceToAuthenticate;
+            return this;
+        }
+    }
 }
