fix: change auth token provider to accept token strings instead of pr…

…incipals (#9255) * fix: change auth token provider to accept token strings instead of principals * update unit test * spotbugs
confluentinc · Jul 13, 2022 · b45841a · b45841a
1 parent 848e179
commit b45841a
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 26 deletions.
diff --git a/ksqldb-engine/src/main/java/io/confluent/ksql/security/KsqlAuthTokenProvider.java b/ksqldb-engine/src/main/java/io/confluent/ksql/security/KsqlAuthTokenProvider.java
@@ -15,9 +15,6 @@
 
 package io.confluent.ksql.security;
 
-import java.security.Principal;
-import java.util.Optional;
-
 /**
  * Interface to extract auth token information to ksqlDB
  */
@@ -26,8 +23,8 @@ public interface KsqlAuthTokenProvider {
   /**
    * Extract the lifetime of a token from the Principal.
    *
-   * @param principal The {@link Principal} that's carrying the auth token.
+   * @param token The auth token.
    * @return An {@Optional} containing the expiration time of the token in ms if there is one
    */
-  Optional<Long> getLifetimeMs(Principal principal);
+  long getLifetimeMs(String token);
 }
diff --git a/ksqldb-rest-app/src/main/java/io/confluent/ksql/rest/server/KsqlServerEndpoints.java b/ksqldb-rest-app/src/main/java/io/confluent/ksql/rest/server/KsqlServerEndpoints.java
@@ -312,7 +312,7 @@ public void executeWebsocketStream(final ServerWebSocket webSocket, final MultiM
             ksqlSecurityContext,
             context,
             new AuthenticationUtil(Clock.systemUTC())
-                .getTokenTimeout(apiSecurityContext.getPrincipal(), ksqlConfig, authTokenProvider)
+                .getTokenTimeout(apiSecurityContext.getAuthToken(), ksqlConfig, authTokenProvider)
         );
       } finally {
         ksqlSecurityContext.getServiceContext().close();

diff --git a/ksqldb-rest-app/src/main/java/io/confluent/ksql/rest/util/AuthenticationUtil.java b/ksqldb-rest-app/src/main/java/io/confluent/ksql/rest/util/AuthenticationUtil.java
@@ -16,7 +16,6 @@
 package io.confluent.ksql.rest.util;
 
 import io.confluent.ksql.security.KsqlAuthTokenProvider;
-import io.confluent.ksql.security.KsqlPrincipal;
 import io.confluent.ksql.util.KsqlConfig;
 import java.time.Clock;
 import java.util.Objects;
@@ -31,19 +30,16 @@ public AuthenticationUtil(final Clock clock) {
   }
 
   public Optional<Long> getTokenTimeout(
-      final Optional<KsqlPrincipal> principal,
+      final Optional<String> token,
       final KsqlConfig ksqlConfig,
       final Optional<KsqlAuthTokenProvider> authTokenProvider
   ) {
     final long maxTimeout =
         ksqlConfig.getLong(KsqlConfig.KSQL_WEBSOCKET_CONNECTION_MAX_TIMEOUT_MS);
     if (maxTimeout > 0) {
-      if (authTokenProvider.isPresent()
-          && principal.isPresent()
-          && authTokenProvider.get().getLifetimeMs(principal.get()).isPresent()
-      ) {
-        final long tokenTimeout = authTokenProvider.get().getLifetimeMs(principal.get()).get()
-            - clock.millis();
+      if (authTokenProvider.isPresent() && token.isPresent()) {
+        final long tokenTimeout = Math.max(authTokenProvider.get().getLifetimeMs(token.get())
+            - clock.millis(), 0);
         return Optional.of(Math.min(tokenTimeout, maxTimeout));
       } else {
         return Optional.of(maxTimeout);

diff --git a/ksqldb-rest-app/src/test/java/io/confluent/ksql/rest/util/AuthenticationUtilTest.java b/ksqldb-rest-app/src/test/java/io/confluent/ksql/rest/util/AuthenticationUtilTest.java
@@ -20,7 +20,6 @@
 import static org.mockito.Mockito.when;
 
 import io.confluent.ksql.security.KsqlAuthTokenProvider;
-import io.confluent.ksql.security.KsqlPrincipal;
 import io.confluent.ksql.util.KsqlConfig;
 import java.time.Clock;
 import java.time.Instant;
@@ -39,14 +38,13 @@ public class AuthenticationUtilTest {
   private KsqlConfig ksqlConfig;
   @Mock
   private KsqlAuthTokenProvider authTokenProvider;
-  @Mock
-  private KsqlPrincipal ksqlPrincipal;
+  private static final String TOKEN = "TOKEN";
   private final AuthenticationUtil authenticationUtil
       = new AuthenticationUtil(Clock.fixed(Instant.ofEpochMilli(0), ZoneId.of("UTC")));
 
   @Before
   public void init() {
-    when(authTokenProvider.getLifetimeMs(ksqlPrincipal)).thenReturn(Optional.of(50000L));
+    when(authTokenProvider.getLifetimeMs(TOKEN)).thenReturn(50000L);
     when(ksqlConfig.getLong(KsqlConfig.KSQL_WEBSOCKET_CONNECTION_MAX_TIMEOUT_MS)).thenReturn(60000L);
   }
 
@@ -56,7 +54,7 @@ public void shouldReturnEmptyWhenConfigSetToZero() {
     when(ksqlConfig.getLong(KsqlConfig.KSQL_WEBSOCKET_CONNECTION_MAX_TIMEOUT_MS)).thenReturn(0L);
 
     // Then:
-    assertThat(authenticationUtil.getTokenTimeout(Optional.of(ksqlPrincipal), ksqlConfig, Optional.of(authTokenProvider)), equalTo(Optional.empty()));
+    assertThat(authenticationUtil.getTokenTimeout(Optional.of(TOKEN), ksqlConfig, Optional.of(authTokenProvider)), equalTo(Optional.empty()));
   }
 
   @Test
@@ -66,29 +64,29 @@ public void shouldReturnDefaultWhenNoPrincipalPresent() {
 
   @Test
   public void shouldReturnDefaultWhenNoAuthTokenProviderPresent() {
-    assertThat(authenticationUtil.getTokenTimeout(Optional.of(ksqlPrincipal), ksqlConfig, Optional.empty()), equalTo(Optional.of(60000L)));
+    assertThat(authenticationUtil.getTokenTimeout(Optional.of(TOKEN), ksqlConfig, Optional.empty()), equalTo(Optional.of(60000L)));
   }
 
   @Test
-  public void shouldReturnDefaultWhenPrincipalHasNoExpiry() {
+  public void shouldReturnZeroWhenPrincipalHasTooLowExpiryTime() {
     // Given:
-    when(authTokenProvider.getLifetimeMs(ksqlPrincipal)).thenReturn(Optional.empty());
+    when(authTokenProvider.getLifetimeMs(TOKEN)).thenReturn(-10L);
 
     // Then:
-    assertThat(authenticationUtil.getTokenTimeout(Optional.of(ksqlPrincipal), ksqlConfig, Optional.of(authTokenProvider)), equalTo(Optional.of(60000L)));
+    assertThat(authenticationUtil.getTokenTimeout(Optional.of(TOKEN), ksqlConfig, Optional.of(authTokenProvider)), equalTo(Optional.of(0L)));
   }
 
   @Test
   public void shouldReturnTokenExpiryTime() {
-    assertThat(authenticationUtil.getTokenTimeout(Optional.of(ksqlPrincipal), ksqlConfig, Optional.of(authTokenProvider)), equalTo(Optional.of(50000L)));
+    assertThat(authenticationUtil.getTokenTimeout(Optional.of(TOKEN), ksqlConfig, Optional.of(authTokenProvider)), equalTo(Optional.of(50000L)));
   }
 
   @Test
   public void shouldReturnMaxTimeout() {
     // Given:
-    when(authTokenProvider.getLifetimeMs(ksqlPrincipal)).thenReturn(Optional.of(50000000L));
+    when(authTokenProvider.getLifetimeMs(TOKEN)).thenReturn(50000000L);
 
     // Then:
-    assertThat(authenticationUtil.getTokenTimeout(Optional.of(ksqlPrincipal), ksqlConfig, Optional.of(authTokenProvider)), equalTo(Optional.of(60000L)));
+    assertThat(authenticationUtil.getTokenTimeout(Optional.of(TOKEN), ksqlConfig, Optional.of(authTokenProvider)), equalTo(Optional.of(60000L)));
   }
 }