Fixed configuration approach for JwksKeyProvider and JwtTokenResolverImpl

Author: artem-v

.gitignore

@@ -8,6 +8,7 @@
 **/target/
 *.iml
 **/logs/*.log
+**/logs/*.log.*
 *.db
 *.csv
 *.log

requirements.txt

@@ -1,7 +1,5 @@
 package io.scalecube.security.tokens.jwt;
 
-import static io.scalecube.security.tokens.jwt.Utils.toRsaPublicKey;
-
 import com.fasterxml.jackson.annotation.JsonAutoDetect;
 import com.fasterxml.jackson.annotation.JsonInclude;
 import com.fasterxml.jackson.annotation.PropertyAccessor;
@@ -11,63 +9,84 @@
 import java.io.BufferedInputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.math.BigInteger;
 import java.net.HttpURLConnection;
 import java.net.URL;
 import java.security.Key;
+import java.security.KeyFactory;
+import java.security.spec.KeySpec;
+import java.security.spec.RSAPublicKeySpec;
 import java.time.Duration;
+import java.util.Base64;
+import java.util.Base64.Decoder;
 import java.util.Optional;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import reactor.core.Exceptions;
 import reactor.core.publisher.Mono;
-import reactor.core.scheduler.Scheduler;
 import reactor.core.scheduler.Schedulers;
 
 public final class JwksKeyProvider implements KeyProvider {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(JwksKeyProvider.class);
 
-  private static final Duration CONNECT_TIMEOUT = Duration.ofSeconds(10);
-  private static final Duration READ_TIMEOUT = Duration.ofSeconds(10);
-
   private static final ObjectMapper OBJECT_MAPPER = newObjectMapper();
 
-  private final Scheduler scheduler;
-  private final String jwksUri;
-  private final long connectTimeoutMillis;
-  private final long readTimeoutMillis;
+  private String jwksUri;
+  private Duration connectTimeout = Duration.ofSeconds(10);
+  private Duration readTimeout = Duration.ofSeconds(10);
+
+  public JwksKeyProvider() {}
+
+  private JwksKeyProvider(JwksKeyProvider other) {
+    this.jwksUri = other.jwksUri;
+    this.connectTimeout = other.connectTimeout;
+    this.readTimeout = other.readTimeout;
+  }
 
   /**
-   * Constructor.
+   * Setter for jwksUri.
    *
    * @param jwksUri jwksUri
+   * @return new instance with applied setting
    */
-  public JwksKeyProvider(String jwksUri) {
-    this(jwksUri, newScheduler(), CONNECT_TIMEOUT, READ_TIMEOUT);
+  public JwksKeyProvider jwksUri(String jwksUri) {
+    final JwksKeyProvider c = copy();
+    c.jwksUri = jwksUri;
+    return c;
   }
 
   /**
-   * Constructor.
+   * Setter for connectTimeout.
    *
-   * @param jwksUri jwksUri
-   * @param scheduler scheduler
    * @param connectTimeout connectTimeout
+   * @return new instance with applied setting
+   */
+  public JwksKeyProvider connectTimeout(Duration connectTimeout) {
+    final JwksKeyProvider c = copy();
+    c.connectTimeout = connectTimeout;
+    return c;
+  }
+
+  /**
+   * Setter for readTimeout.
+   *
    * @param readTimeout readTimeout
+   * @return new instance with applied setting
    */
-  public JwksKeyProvider(
-      String jwksUri, Scheduler scheduler, Duration connectTimeout, Duration readTimeout) {
-    this.jwksUri = jwksUri;
-    this.scheduler = scheduler;
-    this.connectTimeoutMillis = connectTimeout.toMillis();
-    this.readTimeoutMillis = readTimeout.toMillis();
+  public JwksKeyProvider readTimeout(Duration readTimeout) {
+    final JwksKeyProvider c = copy();
+    c.readTimeout = readTimeout;
+    return c;
   }
 
   @Override
   public Mono<Key> findKey(String kid) {
     return computeKey(kid)
         .switchIfEmpty(Mono.error(new KeyNotFoundException("Key was not found, kid: " + kid)))
         .doOnSubscribe(s -> LOGGER.debug("[findKey] Looking up key in jwks, kid: {}", kid))
-        .subscribeOn(scheduler);
+        .subscribeOn(Schedulers.boundedElastic())
+        .publishOn(Schedulers.boundedElastic());
   }
 
   private Mono<Key> computeKey(String kid) {
@@ -78,8 +97,8 @@ private Mono<Key> computeKey(String kid) {
 
   private JwkInfoList computeKeyList() throws IOException {
     HttpURLConnection httpClient = (HttpURLConnection) new URL(jwksUri).openConnection();
-    httpClient.setConnectTimeout((int) connectTimeoutMillis);
-    httpClient.setReadTimeout((int) readTimeoutMillis);
+    httpClient.setConnectTimeout((int) connectTimeout.toMillis());
+    httpClient.setReadTimeout((int) readTimeout.toMillis());
 
     int responseCode = httpClient.getResponseCode();
     if (responseCode != 200) {
@@ -106,6 +125,18 @@ private Optional<Key> findRsaKey(JwkInfoList list, String kid) {
         .map(info -> toRsaPublicKey(info.modulus(), info.exponent()));
   }
 
+  static Key toRsaPublicKey(String n, String e) {
+    Decoder b64Decoder = Base64.getUrlDecoder();
+    BigInteger modulus = new BigInteger(1, b64Decoder.decode(n));
+    BigInteger exponent = new BigInteger(1, b64Decoder.decode(e));
+    KeySpec keySpec = new RSAPublicKeySpec(modulus, exponent);
+    try {
+      return KeyFactory.getInstance("RSA").generatePublic(keySpec);
+    } catch (Exception ex) {
+      throw Exceptions.propagate(ex);
+    }
+  }
+
   private static ObjectMapper newObjectMapper() {
     ObjectMapper mapper = new ObjectMapper();
     mapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
@@ -117,7 +148,7 @@ private static ObjectMapper newObjectMapper() {
     return mapper;
   }
 
-  private static Scheduler newScheduler() {
-    return Schedulers.newElastic("jwks-key-provider", 60, true);
+  private JwksKeyProvider copy() {
+    return new JwksKeyProvider(this);
   }
 }

tokens/src/main/java/io/scalecube/security/tokens/jwt/JwksKeyProvider.java

@@ -1,6 +1,5 @@
 package io.scalecube.security.tokens.jwt;
 
-import io.scalecube.security.tokens.jwt.jsonwebtoken.JsonwebtokenParserFactory;
 import java.security.Key;
 import java.time.Duration;
 import java.util.Map;
@@ -18,41 +17,68 @@ public final class JwtTokenResolverImpl implements JwtTokenResolver {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(JwtTokenResolver.class);
 
-  private static final Duration CLEANUP_INTERVAL = Duration.ofSeconds(60);
-
-  private final KeyProvider keyProvider;
-  private final JwtTokenParserFactory tokenParserFactory;
-  private final Scheduler scheduler;
-  private final Duration cleanupInterval;
+  private KeyProvider keyProvider;
+  private JwtTokenParserFactory tokenParserFactory;
+  private Scheduler scheduler = Schedulers.boundedElastic();
+  private Duration cleanupInterval = Duration.ofSeconds(60);
 
   private final Map<String, Mono<Key>> keyResolutions = new ConcurrentHashMap<>();
 
+  public JwtTokenResolverImpl() {}
+
+  private JwtTokenResolverImpl(JwtTokenResolverImpl other) {
+    this.keyProvider = other.keyProvider;
+    this.tokenParserFactory = other.tokenParserFactory;
+    this.scheduler = other.scheduler;
+    this.cleanupInterval = other.cleanupInterval;
+  }
+
+  /**
+   * Setter for keyProvider.
+   *
+   * @param keyProvider keyProvider
+   * @return new instance with applied setting
+   */
+  public JwtTokenResolverImpl keyProvider(KeyProvider keyProvider) {
+    final JwtTokenResolverImpl c = copy();
+    c.keyProvider = keyProvider;
+    return c;
+  }
+
   /**
-   * Constructor.
+   * Setter for tokenParserFactory.
    *
-   * @param keyProvider key provider
+   * @param tokenParserFactory tokenParserFactory
+   * @return new instance with applied setting
    */
-  public JwtTokenResolverImpl(KeyProvider keyProvider) {
-    this(keyProvider, new JsonwebtokenParserFactory(), newScheduler(), CLEANUP_INTERVAL);
+  public JwtTokenResolverImpl tokenParserFactory(JwtTokenParserFactory tokenParserFactory) {
+    final JwtTokenResolverImpl c = copy();
+    c.tokenParserFactory = tokenParserFactory;
+    return c;
   }
 
   /**
-   * Constructor.
+   * Setter for scheduler.
    *
-   * @param keyProvider key provider
-   * @param tokenParserFactory token parser factoty
-   * @param scheduler cleanup scheduler
-   * @param cleanupInterval cleanup interval for resolved cached keys
+   * @param scheduler scheduler
+   * @return new instance with applied setting
    */
-  public JwtTokenResolverImpl(
-      KeyProvider keyProvider,
-      JwtTokenParserFactory tokenParserFactory,
-      Scheduler scheduler,
-      Duration cleanupInterval) {
-    this.keyProvider = keyProvider;
-    this.tokenParserFactory = tokenParserFactory;
-    this.scheduler = scheduler;
-    this.cleanupInterval = cleanupInterval;
+  public JwtTokenResolverImpl scheduler(Scheduler scheduler) {
+    final JwtTokenResolverImpl c = copy();
+    c.scheduler = scheduler;
+    return c;
+  }
+
+  /**
+   * Setter for cleanupInterval.
+   *
+   * @param cleanupInterval cleanupInterval
+   * @return new instance with applied setting
+   */
+  public JwtTokenResolverImpl cleanupInterval(Duration cleanupInterval) {
+    final JwtTokenResolverImpl c = copy();
+    c.cleanupInterval = cleanupInterval;
+    return c;
   }
 
   @Override
@@ -69,8 +95,7 @@ public Mono<Map<String, Object>> resolve(String token) {
           Map<String, Object> body = jwtToken.body();
           String aud = (String) body.get("aud"); // optional
 
-          LOGGER.debug(
-              "[resolveToken][aud:{}][kid:{}] Resolving token {}", aud, kid, Utils.mask(token));
+          LOGGER.debug("[resolveToken][aud:{}][kid:{}] Resolving token {}", aud, kid, mask(token));
 
           // workaround to remove safely on errors
           AtomicReference<Mono<Key>> computedValueHolder = new AtomicReference<>();
@@ -84,15 +109,15 @@ public Mono<Map<String, Object>> resolve(String token) {
                           "[resolveToken][aud:{}][kid:{}][{}] Exception occurred: {}",
                           aud,
                           kid,
-                          Utils.mask(token),
+                          mask(token),
                           throwable.toString()))
               .doOnSuccess(
                   s ->
                       LOGGER.debug(
                           "[resolveToken][aud:{}][kid:{}] Resolved token {}",
                           aud,
                           kid,
-                          Utils.mask(token)));
+                          mask(token)));
         });
   }
 
@@ -122,7 +147,14 @@ private void cleanup(String kid, AtomicReference<Mono<Key>> computedValueHolder)
     }
   }
 
-  private static Scheduler newScheduler() {
-    return Schedulers.newElastic("token-resolver-cleaner", 60, true);
+  private static String mask(String data) {
+    if (data == null || data.isEmpty() || data.length() < 5) {
+      return "*****";
+    }
+    return data.replace(data.substring(2, data.length() - 2), "***");
+  }
+
+  private JwtTokenResolverImpl copy() {
+    return new JwtTokenResolverImpl(this);
   }
 }