[fix][broker] Let TokenAuthState update authenticationDataSource (apa…

…che#19282) (cherry picked from commit c875365)
michaeljmarshall · Apr 19, 2023 · 6b0a78e · 6b0a78e
1 parent 5e99e8a
commit 6b0a78e
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 0 deletions.
diff --git a/...on/src/main/java/org/apache/pulsar/broker/authentication/AuthenticationProviderToken.java b/...on/src/main/java/org/apache/pulsar/broker/authentication/AuthenticationProviderToken.java
@@ -331,6 +331,8 @@ private String getTokenAudience(ServiceConfiguration conf) throws IllegalArgumen
 
     private static final class TokenAuthenticationState implements AuthenticationState {
         private final AuthenticationProviderToken provider;
+        private final SocketAddress remoteAddress;
+        private final SSLSession sslSession;
         private AuthenticationDataSource authenticationDataSource;
         private Jwt<?, Claims> jwt;
         private long expiration;
@@ -344,6 +346,8 @@ private static final class TokenAuthenticationState implements AuthenticationSta
             String token = new String(authData.getBytes(), UTF_8);
             this.authenticationDataSource = new AuthenticationDataCommand(token, remoteAddress, sslSession);
             this.checkExpiration(token);
+            this.remoteAddress = remoteAddress;
+            this.sslSession = sslSession;
         }
 
         TokenAuthenticationState(
@@ -359,6 +363,10 @@ private static final class TokenAuthenticationState implements AuthenticationSta
             String token = httpHeaderValue.substring(HTTP_HEADER_VALUE_PREFIX.length());
             this.authenticationDataSource = new AuthenticationDataHttps(request);
             this.checkExpiration(token);
+
+            // These are not used when this constructor is invoked, set them to null.
+            this.sslSession = null;
+            this.remoteAddress = null;
         }
 
         @Override
@@ -375,6 +383,7 @@ public String getAuthRole() throws AuthenticationException {
         public AuthData authenticate(AuthData authData) throws AuthenticationException {
             String token = new String(authData.getBytes(), UTF_8);
             checkExpiration(token);
+            this.authenticationDataSource = new AuthenticationDataCommand(token, remoteAddress, sslSession);
             return null;
         }
 

diff --git a/...rc/test/java/org/apache/pulsar/broker/authentication/AuthenticationProviderTokenTest.java b/...rc/test/java/org/apache/pulsar/broker/authentication/AuthenticationProviderTokenTest.java
@@ -23,7 +23,9 @@
 import static org.mockito.Mockito.mock;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNotEquals;
 import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertNull;
 import static org.testng.Assert.assertTrue;
 import static org.testng.Assert.fail;
 import com.google.common.collect.Lists;
@@ -892,6 +894,40 @@ public void testTokenFromHttpHeaders() throws Exception {
         assertTrue(doFilter, "Authentication should have passed");
     }
 
+    @Test
+    public void testTokenStateUpdatesAuthenticationDataSource() throws Exception {
+        SecretKey secretKey = AuthTokenUtils.createSecretKey(SignatureAlgorithm.HS256);
+
+        @Cleanup
+        AuthenticationProviderToken provider = new AuthenticationProviderToken();
+
+        Properties properties = new Properties();
+        properties.setProperty(AuthenticationProviderToken.CONF_TOKEN_SECRET_KEY,
+                AuthTokenUtils.encodeKeyBase64(secretKey));
+
+        ServiceConfiguration conf = new ServiceConfiguration();
+        conf.setProperties(properties);
+        provider.initialize(conf);
+
+        String firstToken = AuthTokenUtils.createToken(secretKey, SUBJECT, Optional.empty());
+
+        AuthenticationState authState = provider.newAuthState(AuthData.of(firstToken.getBytes()),null, null);
+
+        AuthenticationDataSource firstAuthDataSource = authState.getAuthDataSource();
+        assertNotNull(firstAuthDataSource, "Should be initialized.");
+
+        String secondToken = AuthTokenUtils.createToken(secretKey, SUBJECT,
+                Optional.of(new Date(System.currentTimeMillis() + TimeUnit.SECONDS.toMillis(3))));
+
+        AuthData challenge = authState.authenticate(AuthData.of(secondToken.getBytes()));
+        AuthenticationDataSource secondAuthDataSource = authState.getAuthDataSource();
+
+        assertNull(challenge, "TokenAuth doesn't respond with challenges");
+        assertNotNull(secondAuthDataSource, "Created authDataSource");
+
+        assertNotEquals(firstAuthDataSource, secondAuthDataSource);
+    }
+
     private static String createTokenWithAudience(Key signingKey, String audienceClaim, List<String> audience) {
         JwtBuilder builder = Jwts.builder()
                 .setSubject(SUBJECT)