feat: avoid retrieving token eagerly

(cherry picked from commit 264d0f7)
camunda-community-hub · Dec 25, 2023 · 209e71f · 209e71f
1 parent f8060d6
commit 209e71f
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 41 deletions.
diff --git a/camunda-sdk-java/java-common/src/main/java/io/camunda/common/auth/SaaSAuthentication.java b/camunda-sdk-java/java-common/src/main/java/io/camunda/common/auth/SaaSAuthentication.java
@@ -22,14 +22,12 @@ public class SaaSAuthentication extends JwtAuthentication {
   private static final Logger LOG = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());
   private JwtConfig jwtConfig;
   private Map<Product, String> tokens;
-  private Map<Product, LocalDateTime> expirations;
 
   // TODO: have a single object mapper to be used all throughout the SDK, i.e.bean injection
   private JsonMapper jsonMapper = new SdkObjectMapper();
 
   public SaaSAuthentication() {
     tokens = new HashMap<>();
-    expirations = new HashMap<>();
   }
 
   public static SaaSAuthenticationBuilder builder() {
@@ -42,11 +40,10 @@ public void setJwtConfig(JwtConfig jwtConfig) {
 
   @Override
   public Authentication build() {
-    jwtConfig.getMap().forEach(this::retrieveToken);
     return this;
   }
 
-  private void retrieveToken(Product product, JwtCredential jwtCredential) {
+  private String retrieveToken(Product product, JwtCredential jwtCredential) {
     try {
       HttpPost httpPost = new HttpPost(jwtCredential.authUrl);
       httpPost.addHeader("Content-Type", "application/json");
@@ -57,29 +54,22 @@ private void retrieveToken(Product product, JwtCredential jwtCredential) {
       CloseableHttpResponse response = client.execute(httpPost);
       TokenResponse tokenResponse = jsonMapper.fromJson(EntityUtils.toString(response.getEntity()), TokenResponse.class);
       tokens.put(product, tokenResponse.getAccessToken());
-      expirations.put(product, LocalDateTime.now().plusSeconds(tokenResponse.getExpiresIn()));
     } catch (Exception e) {
       LOG.warn("Authenticating for " + product + " failed due to " + e);
       throw new RuntimeException("Unable to authenticate", e);
     }
-  }
-
-  private void retrieveToken(Product product) {
-    JwtCredential jwtCredential = jwtConfig.getMap().get(product);
-    retrieveToken(product, jwtCredential);
+    return tokens.get(product);
   }
 
   @Override
   public Map.Entry<String, String> getTokenHeader(Product product) {
-    refreshToken();
-    return new AbstractMap.SimpleEntry<>("Authorization", "Bearer " + tokens.get(product));
-  }
-
-  private void refreshToken() {
-    expirations.forEach((product, expiration) -> {
-      if (expiration.isAfter(LocalDateTime.now())) {
-        retrieveToken(product);
-      }
-    });
+    String token;
+    if (tokens.containsKey(product)) {
+      token = tokens.get(product);
+    } else {
+      JwtCredential jwtCredential = jwtConfig.getProduct(product);
+      token = retrieveToken(product, jwtCredential);
+    }
+    return new AbstractMap.SimpleEntry<>("Authorization", "Bearer " + token);
   }
 }
diff --git a/...-sdk-java/java-common/src/main/java/io/camunda/common/auth/SelfManagedAuthentication.java b/...-sdk-java/java-common/src/main/java/io/camunda/common/auth/SelfManagedAuthentication.java
@@ -31,14 +31,12 @@ public class SelfManagedAuthentication extends JwtAuthentication {
   private String keycloakUrl;
   private JwtConfig jwtConfig;
   private Map<Product, String> tokens;
-  private Map<Product, LocalDateTime> expirations;
 
   // TODO: have a single object mapper to be used all throughout the SDK, i.e.bean injection
   private JsonMapper jsonMapper = new SdkObjectMapper();
 
   public SelfManagedAuthentication() {
     tokens = new HashMap<>();
-    expirations = new HashMap<>();
   }
 
   public static SelfManagedAuthenticationBuilder builder() {
@@ -60,11 +58,10 @@ public void setJwtConfig(JwtConfig jwtConfig) {
   @Override
   public Authentication build() {
     authUrl = keycloakUrl+"/auth/realms/"+keycloakRealm+"/protocol/openid-connect/token";
-    jwtConfig.getMap().forEach(this::retrieveToken);
     return this;
   }
 
-  private void retrieveToken(Product product, JwtCredential jwtCredential) {
+  private String retrieveToken(Product product, JwtCredential jwtCredential) {
     try {
       HttpPost httpPost = new HttpPost(authUrl);
       httpPost.addHeader("Content-Type", "application/x-www-form-urlencoded");
@@ -89,31 +86,23 @@ private void retrieveToken(Product product, JwtCredential jwtCredential) {
       CloseableHttpClient client = HttpClient.getInstance();
       CloseableHttpResponse response = client.execute(httpPost);
       TokenResponse tokenResponse =  jsonMapper.fromJson(EntityUtils.toString(response.getEntity()), TokenResponse.class);
-      // TODO: verify JWT has the desired permission vs what the user requested
       tokens.put(product, tokenResponse.getAccessToken());
-      expirations.put(product, LocalDateTime.now().plusSeconds(tokenResponse.getExpiresIn()));
     } catch (Exception e) {
       LOG.warn("Authenticating for " + product + " failed due to " + e);
       throw new RuntimeException("Unable to authenticate", e);
     }
-  }
-
-  private void retrieveToken(Product product) {
-    JwtCredential jwtCredential = jwtConfig.getMap().get(product);
-    retrieveToken(product, jwtCredential);
+    return tokens.get(product);
   }
 
   @Override
   public Map.Entry<String, String> getTokenHeader(Product product) {
-    refreshToken();
-    return new AbstractMap.SimpleEntry<>("Authorization", "Bearer " + tokens.get(product));
-  }
-
-  private void refreshToken() {
-    expirations.forEach((product, expiration) -> {
-      if (expiration.isAfter(LocalDateTime.now())) {
-        retrieveToken(product);
-      }
-    });
+    String token;
+    if (tokens.containsKey(product)) {
+      token = tokens.get(product);
+    } else {
+      JwtCredential jwtCredential = jwtConfig.getProduct(product);
+      token = retrieveToken(product, jwtCredential);
+    }
+    return new AbstractMap.SimpleEntry<>("Authorization", "Bearer " + token);
   }
 }