Use the correct keyId supporting key rotation (#91)

* Use the correct keyId supporting key rotation

* Add default region based on project id

examples/management-cli/pom.xml

@@ -19,7 +19,7 @@
         <dependency>
             <groupId>com.descope</groupId>
             <artifactId>java-sdk</artifactId>
-            <version>1.0.11</version>
+            <version>1.0.12</version>
         </dependency>
         <dependency>
             <groupId>info.picocli</groupId>

pom.xml

@@ -4,7 +4,7 @@
     <groupId>com.descope</groupId>
     <artifactId>java-sdk</artifactId>
     <modelVersion>4.0.0</modelVersion>
-    <version>1.0.11</version>
+    <version>1.0.12</version>
     <name>${project.groupId}:${project.artifactId}</name>
     <description>Java library used to integrate with Descope.</description>
     <url>https://github.com/descope/descope-java</url>

src/main/java/com/descope/client/Config.java

@@ -30,11 +30,10 @@ public class Config {
   // fail.
   private String managementKey;
 
-  // PublicKey (optional, "") - used to override or implicitly use a dedicated
-  // public key in order
-  // to decrypt and validate the JWT tokens during ValidateSessionRequest(). If
-  // empty, will attempt
-  // to fetch all public keys from the specified project id.
+  // PublicKey (optional, "") - used to override or implicitly use a dedicated public key in order
+  // to decrypt and validate the JWT tokens during ValidateSessionRequest().
+  // If empty, will attempt to fetch all public keys from the specified project id.
+  // Key should be a JSON in the format of com.descope.model.jwt.SigningKey
   private String publicKey;
 
   // DescopeBaseURL (optional, "https://api.descope.com") - override the default
@@ -47,20 +46,6 @@ public class Config {
   // with descope services.
   private Map<String, String> customDefaultHeaders;
 
-  // State whether session jwt should be sent to client in cookie or let the
-  // calling function handle
-  // the transfer of the jwt, defaults to leaving it for calling function, use
-  // cookie if session jwt
-  // will stay small (less than 1k) session cookie can grow bigger, in case of
-  // using authorization,
-  // or adding custom claims
-  private boolean sessionJWTViaCookie;
-
-  // When using cookies, set the cookie domain here. Alternatively this can be
-  // done via the Descope
-  // console.
-  private String sessionJWTCookieDomain;
-
   public String initializeProjectId() {
     if (StringUtils.isBlank(this.projectId)) {
       this.projectId = EnvironmentUtils.getProjectId();

src/main/java/com/descope/client/DescopeClient.java

@@ -3,18 +3,18 @@
 import com.descope.exception.ClientSetupException;
 import com.descope.exception.DescopeException;
 import com.descope.exception.ServerCommonException;
-import com.descope.model.auth.AuthParams;
 import com.descope.model.auth.AuthenticationServices;
 import com.descope.model.client.Client;
-import com.descope.model.client.ClientParams;
 import com.descope.model.client.SdkInfo;
-import com.descope.model.mgmt.ManagementParams;
+import com.descope.model.jwt.SigningKey;
 import com.descope.model.mgmt.ManagementServices;
 import com.descope.sdk.auth.impl.AuthenticationServiceBuilder;
+import com.descope.sdk.auth.impl.KeyProvider;
 import com.descope.sdk.mgmt.impl.ManagementServiceBuilder;
+import com.fasterxml.jackson.databind.DeserializationFeature;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import io.jsonwebtoken.lang.Collections;
 import java.util.HashMap;
-import java.util.Map;
 import lombok.Getter;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.lang3.StringUtils;
@@ -23,7 +23,8 @@
 @Getter
 public class DescopeClient {
 
-  private static final String DEFAULT_BASE_URL = "https://api.descope.com";
+  private static final String REGION_PLACEHOLDER = "<region>";
+  private static final String DEFAULT_BASE_URL = "https://api." + REGION_PLACEHOLDER + "descope.com";
 
   private final Config config;
   private final ManagementServices managementServices;
@@ -45,62 +46,45 @@ public DescopeClient(Config config) throws DescopeException {
 
     String publicKey = config.initializePublicKey();
     if (StringUtils.isNotBlank(publicKey)) {
-      log.info("Provided public key is set, forcing only provided public key validation");
+      log.debug("Provided public key is set, forcing only provided public key validation");
     }
     config.initializeManagementKey();
     config.initializeBaseURL();
 
     Client client = getClient(config);
-    this.authenticationServices = getAuthenticationServices(config, client);
-    this.managementServices = getManagementServices(config, projectId, client);
+    this.authenticationServices = AuthenticationServiceBuilder.buildServices(client);
+    this.managementServices = ManagementServiceBuilder.buildServices(client);
     this.config = config;
   }
 
-  private static ManagementServices getManagementServices(
-      Config config, String projectId, Client client) {
-    ManagementParams managementParams = ManagementParams.builder()
-        .projectId(projectId)
-        .managementKey(config.getManagementKey())
-        .build();
-    return ManagementServiceBuilder.buildServices(client, managementParams);
-  }
-
-  private static AuthenticationServices getAuthenticationServices(Config config, Client client) {
-    AuthParams authParams = AuthParams.builder()
-        .projectId(config.getProjectId())
-        .publicKey(config.getPublicKey())
-        .sessionJwtViaCookie(config.isSessionJWTViaCookie())
-        .cookieDomain(config.getSessionJWTCookieDomain())
-        .build();
-    return AuthenticationServiceBuilder.buildServices(client, authParams);
-  }
-
   private static Client getClient(Config config) {
-    ClientParams clientParams = ClientParams.builder()
-        .projectId(config.getProjectId())
-        .baseUrl(config.getDescopeBaseUrl())
-        .customDefaultHeaders(config.getCustomDefaultHeaders())
-        .build();
-    return getClient(clientParams);
-  }
-
-  private static Client getClient(ClientParams params) {
-    Map<String, String> customDefaultHeaders = params.getCustomDefaultHeaders();
-    Map<String, String> defaultHeaders = Collections.isEmpty(customDefaultHeaders)
-        ? new HashMap<>()
-        : new HashMap<>(customDefaultHeaders);
-
-    if (StringUtils.isBlank(params.getBaseUrl())) {
-      params.setBaseUrl(DEFAULT_BASE_URL);
+    final SdkInfo sdkInfo = getSdkInfo();
+    final String projectId = config.getProjectId();
+    if (projectId.length() < 28) {
+      throw ClientSetupException.invalidProjectId();
     }
-
-    SdkInfo sdkInfo = getSdkInfo();
-    return Client.builder()
-        .uri(params.getBaseUrl())
-        .params(params)
-        .headers(defaultHeaders)
+    final String region = projectId.substring(1, projectId.length() - 27);
+    final String baseUrl = DEFAULT_BASE_URL.replace(REGION_PLACEHOLDER, region.length() > 0 ? region + "." : "");
+    Client c = Client.builder()
+        .uri(StringUtils.isBlank(config.getDescopeBaseUrl()) ? baseUrl : config.getDescopeBaseUrl())
+        .projectId(projectId)
+        .managementKey(config.getManagementKey())
+        .headers(
+            Collections.isEmpty(config.getCustomDefaultHeaders())
+              ? new HashMap<>() : new HashMap<>(config.getCustomDefaultHeaders()))
         .sdkInfo(sdkInfo)
         .build();
+    if (StringUtils.isNotBlank(config.getPublicKey())) {
+      final ObjectMapper objectMapper = new ObjectMapper()
+          .configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
+      try {
+        SigningKey sk = objectMapper.readValue(config.getPublicKey(), SigningKey.class);
+        c.setProvidedKey(KeyProvider.getPublicKey(sk));
+      } catch (Exception e) {
+        throw ServerCommonException.invalidSigningKey(e.getMessage());
+      }
+    }
+    return c;
   }
 
   private static SdkInfo getSdkInfo() {

src/main/java/com/descope/exception/ClientSetupException.java

@@ -1,5 +1,6 @@
 package com.descope.exception;
 
+import static com.descope.exception.ErrorCode.INVALID_PROJECT_ID;
 import static com.descope.exception.ErrorCode.MISSING_PROJECT_ID;
 
 public class ClientSetupException extends DescopeException {
@@ -13,4 +14,9 @@ public static ClientSetupException missingProjectId() {
     String message = "Missing project ID";
     return new ClientSetupException(message, MISSING_PROJECT_ID);
   }
+
+  public static ClientSetupException invalidProjectId() {
+    String message = "Invalid project ID - must be over 27 characters long";
+    return new ClientSetupException(message, INVALID_PROJECT_ID);
+  }
 }

src/main/java/com/descope/exception/ErrorCode.java

@@ -13,6 +13,7 @@ public class ErrorCode {
 
   // client setup
   public static final String MISSING_PROJECT_ID = "G010001";
+  public static final String INVALID_PROJECT_ID = "G010002";
 
   // client functional errors
   public static final String INVALID_TOKEN = "G030002";

src/main/java/com/descope/model/client/Client.java

@@ -1,9 +1,10 @@
 package com.descope.model.client;
 
-import com.descope.model.jwt.Provider;
+import com.descope.sdk.auth.impl.KeyProvider;
 import java.security.Key;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
 import lombok.AllArgsConstructor;
 import lombok.Builder;
 import lombok.Data;
@@ -15,17 +16,24 @@
 @AllArgsConstructor
 public class Client {
   private String uri;
+  private String projectId;
+  private String managementKey;
   private Map<String, String> headers;
-  private ClientParams params;
   private SdkInfo sdkInfo;
+  private Key providedKey;
   @Builder.Default
-  private Provider provider = Provider.builder().keyMap(new HashMap<>()).build();
+  private AtomicReference<Map<String, Key>> keys = new AtomicReference<>(new HashMap<>());
 
-  public synchronized Key getProvidedKey() {
-    return provider.getProvidedKey();
-  }
-
-  public synchronized void setProvidedKey(Key key) {
-    provider.setProvidedKey(key);
+  public Key getKey(String keyId) {
+    if (providedKey != null) {
+      return providedKey;
+    }
+    Key k = keys.get().get(keyId);
+    // If key is not found, try to refresh key cache
+    if (k == null) {
+      keys.set(KeyProvider.getKeys(projectId, uri, sdkInfo));
+      k = keys.get().get(keyId);
+    }
+    return k;
   }
 }

src/main/java/com/descope/sdk/SdkServicesBase.java

@@ -6,12 +6,10 @@
 import com.descope.exception.ClientFunctionalException;
 import com.descope.model.client.Client;
 import com.descope.model.jwt.Token;
-import com.descope.sdk.auth.impl.KeyProvider;
 import com.descope.utils.JwtUtils;
 import com.descope.utils.UriUtils;
 import java.net.URI;
 import java.net.URLEncoder;
-import java.security.Key;
 import java.util.Map;
 import java.util.Map.Entry;
 import lombok.SneakyThrows;
@@ -21,7 +19,7 @@
 public abstract class SdkServicesBase {
   protected final Client client;
 
-  protected SdkServicesBase(Client client, String projectId) {
+  protected SdkServicesBase(Client client) {
     this.client = client;
   }
 
@@ -39,7 +37,7 @@ protected URI getQueryParamUri(String path, Map<String, String> params) {
     if (isNotEmpty(params)) {
       StringBuilder sb = new StringBuilder("?");
       for (Entry<String, String> e : params.entrySet()) {
-        if (sb.length() > 0) {
+        if (sb.length() > 1) {
           sb.append('&');
         }
         sb.append(
@@ -51,23 +49,11 @@ protected URI getQueryParamUri(String path, Map<String, String> params) {
     return UriUtils.getUri(client.getUri(), path);
   }
 
-  @SneakyThrows
-  protected Key requestKeys() {
-    Key key = client.getProvidedKey();
-    if (key != null) {
-      return client.getProvidedKey();
-    }
-
-    key = KeyProvider.getKey(client.getParams().getProjectId(), client.getUri(), client.getSdkInfo());
-    client.setProvidedKey(key);
-    return key;
-  }
-
   protected Token validateAndCreateToken(String jwt) {
     if (StringUtils.isBlank(jwt)) {
       throw ClientFunctionalException.invalidToken();
     }
-    return JwtUtils.getToken(jwt, requestKeys());
+    return JwtUtils.getToken(jwt, client);
   }
 
 }

src/main/java/com/descope/sdk/auth/impl/AuthenticationServiceBuilder.java

@@ -1,23 +1,22 @@
 package com.descope.sdk.auth.impl;
 
-import com.descope.model.auth.AuthParams;
 import com.descope.model.auth.AuthenticationServices;
 import com.descope.model.client.Client;
 import lombok.experimental.UtilityClass;
 
 @UtilityClass
 public class AuthenticationServiceBuilder {
-  public static AuthenticationServices buildServices(Client client, AuthParams authParams) {
+  public static AuthenticationServices buildServices(Client client) {
     return AuthenticationServices.builder()
-        .authService(new AuthenticationServiceImpl(client, authParams))
-        .otpService(new OTPServiceImpl(client, authParams))
-        .samlService(new SAMLServiceImpl(client, authParams))
-        .totpService(new TOTPServiceImpl(client, authParams))
-        .oauthService(new OAuthServiceImpl(client, authParams))
-        .passwordService(new PasswordServiceImpl(client, authParams))
-        .magicLinkService(new MagicLinkServiceImpl(client, authParams))
-        .enchantedLinkService(new EnchantedLinkServiceImpl(client, authParams))
-        .webAuthnService(new WebAuthnServiceImpl(client, authParams))
+        .authService(new AuthenticationServiceImpl(client))
+        .otpService(new OTPServiceImpl(client))
+        .samlService(new SAMLServiceImpl(client))
+        .totpService(new TOTPServiceImpl(client))
+        .oauthService(new OAuthServiceImpl(client))
+        .passwordService(new PasswordServiceImpl(client))
+        .magicLinkService(new MagicLinkServiceImpl(client))
+        .enchantedLinkService(new EnchantedLinkServiceImpl(client))
+        .webAuthnService(new WebAuthnServiceImpl(client))
         .build();
   }
 }