Skip to content

Commit 78a0aea

Browse files
ggivoggivo
committed
Token based auth integration with core extension
Provide a way for lettuce clients to use token-based authentication. TOKENs come with a TTL. After a Redis client authenticates with a TOKEN, if they didn't renew their authentication we need to evict (close) them. The suggested approach is to leverage the existing CredentialsProvider and add support for streaming credentials to handle token refresh scenarios. Each time a new token is received connection is reauthenticated.
1 parent 91871b6 commit 78a0aea

File tree

6 files changed

+472
-2
lines changed

6 files changed

+472
-2
lines changed

‎pom.xml

Lines changed: 33 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,19 @@
8989
<tag>HEAD</tag>
9090
</scm>
9191

92+
<repositories>
93+
<repository>
94+
<id>sonatype-snapshots</id>
95+
<url>https://oss.sonatype.org/content/repositories/snapshots/</url>
96+
<releases>
97+
<enabled>false</enabled>
98+
</releases>
99+
<snapshots>
100+
<enabled>true</enabled>
101+
</snapshots>
102+
</repository>
103+
</repositories>
104+
92105
<distributionManagement>
93106
<snapshotRepository>
94107
<id>ossrh</id>
@@ -173,12 +186,30 @@
173186
<type>pom</type>
174187
<scope>import</scope>
175188
</dependency>
176-
189+
<dependency>
190+
<groupId>redis.clients.authentication</groupId>
191+
<artifactId>redis-authx-core</artifactId>
192+
<version>0.1.0-SNAPSHOT</version>
193+
</dependency>
194+
<dependency>
195+
<groupId>redis.clients.authentication</groupId>
196+
<artifactId>redis-authx-entraid</artifactId>
197+
<version>0.1.0-SNAPSHOT</version>
198+
<scope>test</scope>
199+
</dependency>
177200
</dependencies>
178201
</dependencyManagement>
179202

180203
<dependencies>
181-
204+
<dependency>
205+
<groupId>redis.clients.authentication</groupId>
206+
<artifactId>redis-authx-core</artifactId>
207+
</dependency>
208+
<dependency>
209+
<groupId>redis.clients.authentication</groupId>
210+
<artifactId>redis-authx-entraid</artifactId>
211+
<scope>test</scope>
212+
</dependency>
182213
<!-- Start of core dependencies -->
183214

184215
<dependency>

‎src/main/java/io/lettuce/authx/TokenBasedRedisCredentialsProvider.java

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
package io.lettuce.authx;
2+
3+
import io.lettuce.core.RedisCredentials;
4+
import io.lettuce.core.StreamingCredentialsProvider;
5+
import reactor.core.publisher.Flux;
6+
import reactor.core.publisher.Mono;
7+
import reactor.core.publisher.Sinks;
8+
import redis.clients.authentication.core.Token;
9+
import redis.clients.authentication.core.TokenAuthConfig;
10+
import redis.clients.authentication.core.TokenListener;
11+
import redis.clients.authentication.core.TokenManager;
12+
13+
public class TokenBasedRedisCredentialsProvider implements StreamingCredentialsProvider {
14+
15+
private final TokenManager tokenManager;
16+
17+
private final Sinks.Many<RedisCredentials> credentialsSink = Sinks.many().replay().latest();
18+
19+
public TokenBasedRedisCredentialsProvider(TokenAuthConfig tokenAuthConfig) {
20+
this(new TokenManager(tokenAuthConfig.getIdentityProviderConfig().getProvider(),
21+
tokenAuthConfig.getTokenManagerConfig()));
22+
23+
}
24+
25+
public TokenBasedRedisCredentialsProvider(TokenManager tokenManager) {
26+
this.tokenManager = tokenManager;
27+
initializeTokenManager();
28+
}
29+
30+
/**
31+
* Initialize the TokenManager and subscribe to token renewal events.
32+
*/
33+
private void initializeTokenManager() {
34+
TokenListener listener = new TokenListener() {
35+
36+
@Override
37+
public void onTokenRenewed(Token token) {
38+
String username = token.tryGet("oid");
39+
char[] pass = token.getValue().toCharArray();
40+
RedisCredentials credentials = RedisCredentials.just(username, pass);
41+
credentialsSink.tryEmitNext(credentials);
42+
}
43+
44+
@Override
45+
public void onError(Exception exception) {
46+
credentialsSink.tryEmitError(exception);
47+
}
48+
49+
};
50+
51+
try {
52+
tokenManager.start(listener, false);
53+
} catch (Exception e) {
54+
credentialsSink.tryEmitError(e);
55+
}
56+
}
57+
58+
/**
59+
* Resolve the latest available credentials as a Mono.
60+
* <p>
61+
* This method returns a Mono that emits the most recent set of Redis credentials. The Mono will complete once the
62+
* credentials are emitted. If no credentials are available at the time of subscription, the Mono will wait until
63+
* credentials are available.
64+
*
65+
* @return a Mono that emits the latest Redis credentials
66+
*/
67+
@Override
68+
public Mono<RedisCredentials> resolveCredentials() {
69+
70+
return credentialsSink.asFlux().next();
71+
}
72+
73+
/**
74+
* Expose the Flux for all credential updates.
75+
* <p>
76+
* This method returns a Flux that emits all updates to the Redis credentials. Subscribers will receive the latest
77+
* credentials whenever they are updated. The Flux will continue to emit updates until the provider is shut down.
78+
*
79+
* @return a Flux that emits all updates to the Redis credentials
80+
*/
81+
@Override
82+
public Flux<RedisCredentials> credentials() {
83+
84+
return credentialsSink.asFlux().onBackpressureLatest(); // Provide a continuous stream of credentials
85+
}
86+
87+
/**
88+
* Stop the credentials provider and clean up resources.
89+
* <p>
90+
* This method stops the TokenManager and completes the credentials sink, ensuring that all resources are properly released.
91+
* It should be called when the credentials provider is no longer needed.
92+
*/
93+
public void shutdown() {
94+
credentialsSink.tryEmitComplete();
95+
tokenManager.stop();
96+
}
97+
98+
}

‎src/test/java/io/lettuce/authx/TokenBasedRedisCredentialsProviderTest.java

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
package io.lettuce.authx;
2+
3+
import io.lettuce.core.RedisCredentials;
4+
import io.lettuce.core.TestTokenManager;
5+
import org.junit.jupiter.api.BeforeEach;
6+
import org.junit.jupiter.api.Test;
7+
import reactor.core.Disposable;
8+
import reactor.core.publisher.Flux;
9+
import reactor.core.publisher.Mono;
10+
import reactor.test.StepVerifier;
11+
import redis.clients.authentication.core.SimpleToken;
12+
13+
import java.time.Duration;
14+
import java.util.Collections;
15+
16+
import static org.assertj.core.api.Assertions.assertThat;
17+
18+
public class TokenBasedRedisCredentialsProviderTest {
19+
20+
private TestTokenManager tokenManager;
21+
22+
private TokenBasedRedisCredentialsProvider credentialsProvider;
23+
24+
@BeforeEach
25+
public void setUp() {
26+
// Use TestToken manager to emit tokens/errors on request
27+
tokenManager = new TestTokenManager(null, null);
28+
credentialsProvider = new TokenBasedRedisCredentialsProvider(tokenManager);
29+
}
30+
31+
@Test
32+
public void shouldReturnPreviouslyEmittedTokenWhenResolved() {
33+
tokenManager.emitToken(testToken("test-user", "token-1"));
34+
35+
Mono<RedisCredentials> credentials = credentialsProvider.resolveCredentials();
36+
37+
StepVerifier.create(credentials).assertNext(actual -> {
38+
assertThat(actual.getUsername()).isEqualTo("test-user");
39+
assertThat(new String(actual.getPassword())).isEqualTo("token-1");
40+
}).verifyComplete();
41+
}
42+
43+
@Test
44+
public void shouldReturnLatestEmittedTokenWhenResolved() {
45+
tokenManager.emitToken(testToken("test-user", "token-2"));
46+
tokenManager.emitToken(testToken("test-user", "token-3")); // Latest token
47+
48+
Mono<RedisCredentials> credentials = credentialsProvider.resolveCredentials();
49+
50+
StepVerifier.create(credentials).assertNext(actual -> {
51+
assertThat(actual.getUsername()).isEqualTo("test-user");
52+
assertThat(new String(actual.getPassword())).isEqualTo("token-3");
53+
}).verifyComplete();
54+
}
55+
56+
@Test
57+
public void shouldReturnTokenEmittedBeforeSubscription() {
58+
59+
tokenManager.emitToken(testToken("test-user", "token-1"));
60+
61+
// Test resolveCredentials
62+
Mono<RedisCredentials> credentials1 = credentialsProvider.resolveCredentials();
63+
64+
StepVerifier.create(credentials1).assertNext(actual -> {
65+
assertThat(actual.getUsername()).isEqualTo("test-user");
66+
assertThat(new String(actual.getPassword())).isEqualTo("token-1");
67+
}).verifyComplete();
68+
69+
// Emit second token and subscribe another
70+
tokenManager.emitToken(testToken("test-user", "token-2"));
71+
tokenManager.emitToken(testToken("test-user", "token-3"));
72+
Mono<RedisCredentials> credentials2 = credentialsProvider.resolveCredentials();
73+
StepVerifier.create(credentials2).assertNext(actual -> {
74+
assertThat(actual.getUsername()).isEqualTo("test-user");
75+
assertThat(new String(actual.getPassword())).isEqualTo("token-3");
76+
}).verifyComplete();
77+
}
78+
79+
@Test
80+
public void shouldWaitForAndReturnTokenWhenEmittedLater() {
81+
Mono<RedisCredentials> result = credentialsProvider.resolveCredentials();
82+
83+
tokenManager.emitTokenWithDelay(testToken("test-user", "delayed-token"), 100); // Emit token after 100ms
84+
StepVerifier.create(result)
85+
.assertNext(credentials -> assertThat(String.valueOf(credentials.getPassword())).isEqualTo("delayed-token"))
86+
.verifyComplete();
87+
}
88+
89+
@Test
90+
public void shouldCompleteAllSubscribersOnStop() {
91+
Flux<RedisCredentials> credentialsFlux1 = credentialsProvider.credentials();
92+
Flux<RedisCredentials> credentialsFlux2 = credentialsProvider.credentials();
93+
94+
Disposable subscription1 = credentialsFlux1.subscribe();
95+
Disposable subscription2 = credentialsFlux2.subscribe();
96+
97+
tokenManager.emitToken(testToken("test-user", "token-1"));
98+
99+
new Thread(() -> {
100+
try {
101+
Thread.sleep(100); // Delay of 100 milliseconds
102+
} catch (InterruptedException e) {
103+
Thread.currentThread().interrupt();
104+
}
105+
credentialsProvider.shutdown();
106+
}).start();
107+
108+
StepVerifier.create(credentialsFlux1)
109+
.assertNext(credentials -> assertThat(String.valueOf(credentials.getPassword())).isEqualTo("token-1"))
110+
.verifyComplete();
111+
112+
StepVerifier.create(credentialsFlux2)
113+
.assertNext(credentials -> assertThat(String.valueOf(credentials.getPassword())).isEqualTo("token-1"))
114+
.verifyComplete();
115+
}
116+
117+
@Test
118+
public void shouldPropagateMultipleTokensOnStream() {
119+
120+
Flux<RedisCredentials> result = credentialsProvider.credentials();
121+
StepVerifier.create(result).then(() -> tokenManager.emitToken(testToken("test-user", "token1")))
122+
.then(() -> tokenManager.emitToken(testToken("test-user", "token2")))
123+
.assertNext(credentials -> assertThat(String.valueOf(credentials.getPassword())).isEqualTo("token1"))
124+
.assertNext(credentials -> assertThat(String.valueOf(credentials.getPassword())).isEqualTo("token2"))
125+
.thenCancel().verify(Duration.ofMillis(100));
126+
}
127+
128+
@Test
129+
public void shouldHandleTokenRequestErrorGracefully() {
130+
Exception simulatedError = new RuntimeException("Token request failed");
131+
tokenManager.emitError(simulatedError);
132+
133+
Flux<RedisCredentials> result = credentialsProvider.credentials();
134+
135+
StepVerifier.create(result).expectErrorMatches(
136+
throwable -> throwable instanceof RuntimeException && "Token request failed".equals(throwable.getMessage()))
137+
.verify();
138+
}
139+
140+
private SimpleToken testToken(String username, String value) {
141+
return new SimpleToken(value, System.currentTimeMillis() + 5000, // expires in 5 seconds
142+
System.currentTimeMillis(), Collections.singletonMap("oid", username));
143+
144+
}
145+
146+
}

‎src/test/java/io/lettuce/core/AuthenticationIntegrationTests.java

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,11 @@
55

66
import javax.inject.Inject;
77

8+
import io.lettuce.authx.TokenBasedRedisCredentialsProvider;
89
import io.lettuce.core.event.command.CommandListener;
910
import io.lettuce.core.event.command.CommandSucceededEvent;
1011
import io.lettuce.core.protocol.RedisCommand;
12+
import io.lettuce.test.Delay;
1113
import org.awaitility.Awaitility;
1214
import org.junit.jupiter.api.BeforeEach;
1315
import org.junit.jupiter.api.Tag;
@@ -26,8 +28,10 @@
2628
import reactor.core.publisher.Flux;
2729
import reactor.core.publisher.Mono;
2830
import reactor.core.publisher.Sinks;
31+
import redis.clients.authentication.core.SimpleToken;
2932

3033
import java.time.Duration;
34+
import java.time.Instant;
3135
import java.util.ArrayList;
3236
import java.util.Collections;
3337
import java.util.List;
@@ -170,6 +174,59 @@ private boolean isAuthCommandWithCredentials(RedisCommand<?, ?, ?> command, Stri
170174
return false;
171175
}
172176

177+
@Test
178+
@Inject
179+
void tokenBasedCredentialProvider(RedisClient client) {
180+
181+
ClientOptions clientOptions = ClientOptions.builder()
182+
.disconnectedBehavior(ClientOptions.DisconnectedBehavior.REJECT_COMMANDS).build();
183+
client.setOptions(clientOptions);
184+
// Connection used to simulate test user credential rotation
185+
StatefulRedisConnection<String, String> defaultConnection = client.connect();
186+
187+
String testUser = "streaming_cred_test_user";
188+
char[] testPassword1 = "token_1".toCharArray();
189+
char[] testPassword2 = "token_2".toCharArray();
190+
191+
TestTokenManager tokenManager = new TestTokenManager(null, null);
192+
193+
// streaming credentials provider that emits redis credentials which will trigger connection re-authentication
194+
// token manager is used to emit updated credentials
195+
TokenBasedRedisCredentialsProvider credentialsProvider = new TokenBasedRedisCredentialsProvider(tokenManager);
196+
197+
RedisURI uri = RedisURI.builder().withTimeout(Duration.ofSeconds(1)).withClientName("streaming_cred_test")
198+
.withHost(TestSettings.host()).withPort(TestSettings.port()).withAuthentication(credentialsProvider).build();
199+
200+
// create test user with initial credentials set to 'testPassword1'
201+
createTestUser(defaultConnection, testUser, testPassword1);
202+
tokenManager.emitToken(testToken(testUser, testPassword1));
203+
204+
StatefulRedisConnection<String, String> connection = client.connect(StringCodec.UTF8, uri);
205+
assertThat(connection.sync().aclWhoami()).isEqualTo(testUser);
206+
207+
// update test user credentials in Redis server (password changed to testPassword2)
208+
// then emit updated credentials trough streaming credentials provider
209+
// and trigger re-connect to force re-authentication
210+
// updated credentials should be used for re-authentication
211+
updateTestUser(defaultConnection, testUser, testPassword2);
212+
tokenManager.emitToken(testToken(testUser, testPassword2));
213+
connection.sync().quit();
214+
215+
Delay.delay(Duration.ofMillis(100));
216+
assertThat(connection.sync().ping()).isEqualTo("PONG");
217+
218+
String res = connection.sync().aclWhoami();
219+
assertThat(res).isEqualTo(testUser);
220+
221+
defaultConnection.close();
222+
connection.close();
223+
}
224+
225+
private SimpleToken testToken(String username, char[] password) {
226+
return new SimpleToken(String.valueOf(password), Instant.now().plusMillis(500).toEpochMilli(),
227+
Instant.now().toEpochMilli(), Collections.singletonMap("oid", username));
228+
}
229+
173230
static class RenewableRedisCredentialsProvider implements StreamingCredentialsProvider {
174231

175232
private final Sinks.Many<RedisCredentials> credentialsSink = Sinks.many().replay().latest();

0 commit comments

Comments
 (0)