- support full customization of different MSAL application types and advanced configurations with EntraIDTokenAuthConfigBuilder

atakavci · atakavci · commit ee60a45c25ac · 2024-12-01T19:11:34.000+03:00
- add more unit tests
diff --git a/entraid/src/main/java/redis/clients/authentication/entraid/EntraIDIdentityProvider.java b/entraid/src/main/java/redis/clients/authentication/entraid/EntraIDIdentityProvider.java
@@ -48,6 +48,10 @@ public EntraIDIdentityProvider(ManagedIdentityInfo info, Set<String> scopes) {
         resultSupplier = () -> supplierForManagedIdentityApp(app, params);
     }
 
+    public EntraIDIdentityProvider(Supplier<IAuthenticationResult> customEntraIdAppSupplier) {
+        this.resultSupplier = customEntraIdAppSupplier;
+    }
+
     private IClientCredential getClientCredential(ServicePrincipalInfo servicePrincipalInfo) {
         switch (servicePrincipalInfo.getAccessWith()) {
         case WithSecret:
diff --git a/entraid/src/main/java/redis/clients/authentication/entraid/EntraIDIdentityProviderConfig.java b/entraid/src/main/java/redis/clients/authentication/entraid/EntraIDIdentityProviderConfig.java
@@ -1,31 +1,33 @@
 package redis.clients.authentication.entraid;
 
 import java.util.Set;
+import java.util.function.Supplier;
+
+import com.microsoft.aad.msal4j.IAuthenticationResult;
 
 import redis.clients.authentication.core.IdentityProvider;
 import redis.clients.authentication.core.IdentityProviderConfig;
 
 public final class EntraIDIdentityProviderConfig implements IdentityProviderConfig, AutoCloseable {
 
-    private ServicePrincipalInfo servicePrincipalInfo;
-    private Set<String> scopes;
-    private ManagedIdentityInfo managedIdentityInfo;
+    private Supplier<IdentityProvider> providerSupplier;
+
+    public EntraIDIdentityProviderConfig(ServicePrincipalInfo info, Set<String> scopes) {
+        providerSupplier = () -> new EntraIDIdentityProvider(info, scopes);
+    }
+
+    public EntraIDIdentityProviderConfig(ManagedIdentityInfo info, Set<String> scopes) {
+        providerSupplier = () -> new EntraIDIdentityProvider(info, scopes);
+    }
 
-    public EntraIDIdentityProviderConfig(ServicePrincipalInfo servicePrincipalInfo,
-            ManagedIdentityInfo info, Set<String> scopes) {
-        this.servicePrincipalInfo = servicePrincipalInfo;
-        this.scopes = scopes;
-        this.managedIdentityInfo = info;
+    public EntraIDIdentityProviderConfig(
+            Supplier<IAuthenticationResult> customEntraIdAppSupplier) {
+        providerSupplier = () -> new EntraIDIdentityProvider(customEntraIdAppSupplier);
     }
 
     @Override
     public IdentityProvider getProvider() {
-        IdentityProvider identityProvider = null;
-        if (managedIdentityInfo != null) {
-            identityProvider = new EntraIDIdentityProvider(managedIdentityInfo, scopes);
-        } else {
-            identityProvider = new EntraIDIdentityProvider(servicePrincipalInfo, scopes);
-        }
+        IdentityProvider identityProvider = providerSupplier.get();
         clear();
         return identityProvider;
     }
@@ -36,8 +38,6 @@ public void close() throws Exception {
     }
 
     private void clear() {
-        servicePrincipalInfo = null;
-        managedIdentityInfo = null;
-        scopes = null;
+        providerSupplier = null;
     }
 }
diff --git a/entraid/src/main/java/redis/clients/authentication/entraid/EntraIDTokenAuthConfigBuilder.java b/entraid/src/main/java/redis/clients/authentication/entraid/EntraIDTokenAuthConfigBuilder.java
@@ -3,6 +3,9 @@
 import java.security.PrivateKey;
 import java.security.cert.X509Certificate;
 import java.util.Set;
+import java.util.function.Supplier;
+
+import com.microsoft.aad.msal4j.IAuthenticationResult;
 
 import redis.clients.authentication.core.TokenAuthConfig;
 import redis.clients.authentication.entraid.ManagedIdentityInfo.UserManagedIdentityType;
@@ -24,6 +27,7 @@ public class EntraIDTokenAuthConfigBuilder extends TokenAuthConfig.Builder
     private Set<String> scopes;
     private ServicePrincipalAccess accessWith;
     private ManagedIdentityInfo mii;
+    Supplier<IAuthenticationResult> customEntraIdAuthenticationSupplier;
 
     public EntraIDTokenAuthConfigBuilder() {
         this.expirationRefreshRatio(DEFAULT_EXPIRATION_REFRESH_RATIO)
@@ -66,6 +70,12 @@ public EntraIDTokenAuthConfigBuilder userAssignedManagedIdentity(
         return this;
     }
 
+    public EntraIDTokenAuthConfigBuilder customEntraIdAuthenticationSupplier(
+            Supplier<IAuthenticationResult> customEntraIdAuthenticationSupplier) {
+
+        return this;
+    }
+
     public EntraIDTokenAuthConfigBuilder scopes(Set<String> scopes) {
         this.scopes = scopes;
         return this;
@@ -85,11 +95,22 @@ public TokenAuthConfig build() {
         }
         if (spi != null && mii != null) {
             throw new RedisEntraIDException(
-                    "Cannot have both ServicePrincipal and ManagedIdentity");
+                    "Cannot have both ServicePrincipal and ManagedIdentity!");
+        }
+        if (this.customEntraIdAuthenticationSupplier != null && (spi != null || mii != null)) {
+            throw new RedisEntraIDException(
+                    "Cannot have both customEntraIdAuthenticationSupplier and ServicePrincipal/ManagedIdentity!");
+        }
+        if (spi != null) {
+            super.identityProviderConfig(new EntraIDIdentityProviderConfig(spi, scopes));
+        }
+        if (mii != null) {
+            super.identityProviderConfig(new EntraIDIdentityProviderConfig(mii, scopes));
+        }
+        if (customEntraIdAuthenticationSupplier != null) {
+            super.identityProviderConfig(
+                new EntraIDIdentityProviderConfig(customEntraIdAuthenticationSupplier));
         }
-        EntraIDIdentityProviderConfig idProviderConfig = new EntraIDIdentityProviderConfig(spi, mii,
-                scopes);
-        super.identityProviderConfig(idProviderConfig);
         return super.build();
     }
 
@@ -101,6 +122,7 @@ public void close() throws Exception {
         cert = null;
         authority = null;
         scopes = null;
+        customEntraIdAuthenticationSupplier = null;
     }
 
     public static EntraIDTokenAuthConfigBuilder builder() {
diff --git a/entraid/src/test/java/redis/clients/authentication/RedisEntraIDUnitTests.java b/entraid/src/test/java/redis/clients/authentication/RedisEntraIDUnitTests.java
@@ -16,8 +16,11 @@
 import static org.hamcrest.Matchers.both;
 import static org.hamcrest.MatcherAssert.assertThat;
 
+import java.time.Duration;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Date;
+import java.util.List;
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeoutException;
@@ -31,6 +34,7 @@
 
 import com.auth0.jwt.JWT;
 import com.auth0.jwt.algorithms.Algorithm;
+
 import redis.clients.authentication.core.IdentityProvider;
 import redis.clients.authentication.core.IdentityProviderConfig;
 import redis.clients.authentication.core.SimpleToken;
@@ -280,6 +284,16 @@ public void onError(Exception e) {
                 .until(() -> numberOfTokens.get(), is(2));
     }
 
+    // T.2.2
+
+    // Test that the Redis client is not blocked/interrupted during token renewal.
+    @Test
+    public void renewalDuringOperationsTest() {
+        // set the stage with consecutive get/set operations with unique keys which takes at least for 2000 ms with a jedispooled instace, 
+        // configure token manager to renew token every 500ms
+        // wait till all operations are completed and verify that token was renewed at least 3 times after initial token acquisition 
+    }
+
     // T.2.2
     // Ensure the system propagates error during renewal back to the user
     @Test
@@ -306,6 +320,7 @@ public void onError(Exception e) {
 
         Awaitility.await().pollInterval(ONE_HUNDRED_MILLISECONDS).atMost(TWO_SECONDS)
                 .until(() -> numberOfErrors.get(), is(1));
+
     }
 
     // T.2.3
@@ -330,6 +345,7 @@ public void onTokenRenewed(Token token) {
                     timeDiff.set((int) (token.getExpiresAt() - lastToken.getExpiresAt()));
                 }
                 lastToken = token;
+
             }
 
             @Override
@@ -350,7 +366,89 @@ public void onError(Exception e) {
     // T.2.3
     // Verify behavior with edge case renewal timing configurations (e.g., very low or high percentages).
     @Test
-    public void edgeCaseRenewalTimingTest() {
+    public void highPercentage_edgeCaseRenewalTimingTest() {
+        List<Token> tokens = new ArrayList<Token>();
+        int validDurationInMs = 1000;
+
+        IdentityProvider identityProvider = () -> new SimpleToken(TOKEN_VALUE,
+                System.currentTimeMillis() + validDurationInMs, System.currentTimeMillis(),
+                Collections.singletonMap("oid", TOKEN_OID));
+
+        TokenManagerConfig tokenManagerConfig = new TokenManagerConfig(0.99F, 0,
+                TOKEN_REQUEST_EXEC_TIMEOUT,
+                new TokenManagerConfig.RetryPolicy(RETRY_POLICY_MAX_ATTEMPTS, RETRY_POLICY_DELAY));
+
+        TokenManager tokenManager = new TokenManager(identityProvider, tokenManagerConfig);
+        TokenListener listener = new TokenListener() {
+
+            @Override
+            public void onTokenRenewed(Token token) {
+                tokens.add(token);
+            }
+
+            @Override
+            public void onError(Exception e) {
+            }
+        };
+
+        tokenManager.start(listener, false);
+
+        Awaitility.await().pollInterval(Duration.ofMillis(10)).atMost(Durations.TWO_SECONDS)
+                .until(() -> tokens.size(), is(2));
+
+        Token initialToken = tokens.get(0);
+        Token secondToken = tokens.get(1);
+        Long renewalWindowStart = initialToken.getReceivedAt()
+                + (long) (validDurationInMs * tokenManagerConfig.getExpirationRefreshRatio());
+        Long renewalWindowEnd = initialToken.getExpiresAt();
+        assertThat((Long) secondToken.getReceivedAt(),
+            both(greaterThanOrEqualTo(renewalWindowStart))
+                    .and(lessThanOrEqualTo(renewalWindowEnd)));
+
+    }
+
+        // T.2.3
+    // Verify behavior with edge case renewal timing configurations (e.g., very low or high percentages).
+    @Test
+    public void lowPercentage_edgeCaseRenewalTimingTest() {
+        List<Token> tokens = new ArrayList<Token>();
+        int validDurationInMs = 1000;
+
+        IdentityProvider identityProvider = () -> new SimpleToken(TOKEN_VALUE,
+                System.currentTimeMillis() + validDurationInMs, System.currentTimeMillis(),
+                Collections.singletonMap("oid", TOKEN_OID));
+
+        TokenManagerConfig tokenManagerConfig = new TokenManagerConfig(0.01F, 0,
+                TOKEN_REQUEST_EXEC_TIMEOUT,
+                new TokenManagerConfig.RetryPolicy(RETRY_POLICY_MAX_ATTEMPTS, RETRY_POLICY_DELAY));
+
+        TokenManager tokenManager = new TokenManager(identityProvider, tokenManagerConfig);
+        TokenListener listener = new TokenListener() {
+
+            @Override
+            public void onTokenRenewed(Token token) {
+                tokens.add(token);
+            }
+
+            @Override
+            public void onError(Exception e) {
+            }
+        };
+
+        tokenManager.start(listener, false);
+
+        Awaitility.await().pollInterval(ONE_MILLISECOND).atMost(Durations.TWO_SECONDS)
+                .until(() -> tokens.size(), is(2));
+
+        Token initialToken = tokens.get(0);
+        Token secondToken = tokens.get(1);
+        Long renewalWindowStart = initialToken.getReceivedAt()
+                + (long) (validDurationInMs * tokenManagerConfig.getExpirationRefreshRatio());
+        Long renewalWindowEnd = initialToken.getExpiresAt();
+        assertThat((Long) secondToken.getReceivedAt(),
+            both(greaterThanOrEqualTo(renewalWindowStart))
+                    .and(lessThanOrEqualTo(renewalWindowEnd)));
+
     }
 
     // T.2.4
@@ -363,6 +461,7 @@ public void expiredTokenCheckTest() {
 
         token = JWT.create().withExpiresAt(new Date(System.currentTimeMillis() + 1000))
                 .withClaim("oid", "user1").sign(Algorithm.none());
+
         assertFalse(new JWToken(token).isExpired());
     }
 
@@ -382,6 +481,13 @@ public void tokenParserTest() {
             lessThanOrEqualTo((Long) 10L));
     }
 
+    // T.2.5
+    // Ensure that token objects are immutable and cannot be modified after creation.
+    @Test
+    public void tokenImmutabilityTest() {
+        // ???
+    }
+
     // T.3.1
     // Verify that the most recent valid token is correctly cached and that the cache is initially empty
     @Test
@@ -399,6 +505,82 @@ public void tokenCachingTest() {
         assertNotNull(tokenManager.getCurrentToken());
     }
 
+    // T.3.1
+    // Ensure the token cache is updated when a new token is acquired or renewed.
+    @Test
+    public void cacheUpdateOnRenewalTest() {
+
+        AtomicInteger numberOfTokens = new AtomicInteger(0);
+        IdentityProvider identityProvider = () -> {
+            return new SimpleToken("" + numberOfTokens.incrementAndGet(),
+                    System.currentTimeMillis() + 500, System.currentTimeMillis(),
+                    Collections.singletonMap("oid", "user1"));
+        };
+        TokenManager tokenManager = new TokenManager(identityProvider, tokenManagerConfig);
+        assertNull(tokenManager.getCurrentToken());
+        tokenManager.start(mock(TokenListener.class), true);
+        assertNotNull(tokenManager.getCurrentToken());
+        assertEquals("1", tokenManager.getCurrentToken().getValue());
+        Awaitility.await().pollInterval(ONE_HUNDRED_MILLISECONDS).atMost(TWO_SECONDS)
+                .until(() -> tokenManager.getCurrentToken().getValue(), is("2"));
+
+    }
+
+    // T.3.2
+    // Verify that all existing connections can be re-authenticated when a new token is received.
+    @Test
+    public void allConnectionsReauthTest() {
+
+    }
+
+    // T.3.2
+    // Test system behavior when some connections fail to re-authenticate during bulk authentication. e.g when a network partition occurs for 1 or more of them
+    @Test
+    public void partialReauthFailureTest() {
+
+    }
+
+    // T.3.3
+    // Test authentication of a single connection using the current valid token.
+    @Test
+    public void singleConnectionAuthTest() {
+
+    }
+
+    // T.3.3
+    // Verify behavior when attempting to authenticate a single connection with an expired token.
+    @Test
+    public void connectionAuthWithExpiredTokenTest() {
+
+    }
+
+    // T.3.4
+    // Verify handling of reconnection and re-authentication after a network partition. (use cached token)
+    @Test
+    public void networkPartitionEvictionTest() {
+
+    }
+
+    // T.4.1
+    // Verify that token renewal timing can be configured correctly.
+    @Test
+    public void renewalTimingConfigTest() {
+
+    }
+
+    // T.4.2
+    // Verify that Azure AD-specific parameters can be configured correctly.
+    @Test
+    public void azureADConfigTest() {
+
+    }
+
+    // T.4.2
+    // Test configuration of custom identity provider parameters.
+    @Test
+    public void customProviderConfigTest() {
+    }
+
     private void delay(long durationInMs) {
         try {
             Thread.sleep(durationInMs);
