Merge pull request #3 from avast/RotatingKeysSupport

Rotating Keycloak public keys support

Author: augi

README.md

@@ -48,7 +48,6 @@ YourService.newStub(aChannel).withCallCredentials(callCredentials);
 
 ### Server usage
 This ensures that only requests with valid `JWT` in `Authorization` header are processed.
- The `fromConfig` method automatically downloads the Keycloak public key before the instance is actually created.
 ```java
 import io.grpc.ServerServiceDefinition;
 import com.avast.grpc.jwt.keycloak.server.KeycloakJwtServerInterceptor;
@@ -78,4 +77,4 @@ On other hand, `Context key` is set of values that are available during request
 
 So when implementing interceptors, you must be sure that you read Context values from the right thread. It's actually no issue for us because:
 1. The right thread is automatically handled by gRPC-core when using`CallCredentials`. So you can call `applier.apply()` method on any thread.
-2. Our `ServerInterceptor` implementation is fully synchronous.
+2. Our `ServerInterceptor` implementation handles it correctly.

core/src/main/java/com/avast/grpc/jwt/server/DelayedServerCallListener.java

@@ -0,0 +1,64 @@
+package com.avast.grpc.jwt.server;
+
+import io.grpc.ServerCall;
+import java.util.ArrayList;
+import java.util.List;
+
+// https://stackoverflow.com/a/53656689/181796
+class DelayedServerCallListener<ReqT> extends ServerCall.Listener<ReqT> {
+  private ServerCall.Listener<ReqT> delegate;
+  private List<Runnable> events = new ArrayList<>();
+
+  @Override
+  public synchronized void onMessage(ReqT message) {
+    if (delegate == null) {
+      events.add(() -> delegate.onMessage(message));
+    } else {
+      delegate.onMessage(message);
+    }
+  }
+
+  @Override
+  public synchronized void onHalfClose() {
+    if (delegate == null) {
+      events.add(() -> delegate.onHalfClose());
+    } else {
+      delegate.onHalfClose();
+    }
+  }
+
+  @Override
+  public synchronized void onCancel() {
+    if (delegate == null) {
+      events.add(() -> delegate.onCancel());
+    } else {
+      delegate.onCancel();
+    }
+  }
+
+  @Override
+  public synchronized void onComplete() {
+    if (delegate == null) {
+      events.add(() -> delegate.onComplete());
+    } else {
+      delegate.onComplete();
+    }
+  }
+
+  @Override
+  public synchronized void onReady() {
+    if (delegate == null) {
+      events.add(() -> delegate.onReady());
+    } else {
+      delegate.onReady();
+    }
+  }
+
+  public synchronized void setDelegate(ServerCall.Listener<ReqT> delegate) {
+    this.delegate = delegate;
+    for (Runnable runnable : events) {
+      runnable.run();
+    }
+    events = null;
+  }
+}

core/src/main/java/com/avast/grpc/jwt/server/JwtServerInterceptor.java

@@ -37,19 +37,38 @@ public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
       call.close(Status.UNAUTHENTICATED.withDescription(msg), new Metadata());
       return new ServerCall.Listener<ReqT>() {};
     }
-    T token;
+    DelayedServerCallListener<ReqT> delayedListener = new DelayedServerCallListener<>();
+    Context context = Context.current(); // we must call this on the right thread
     try {
-      token = tokenParser.parseToValid(authHeader.substring(AUTH_HEADER_PREFIX.length()));
+      tokenParser
+          .parseToValid(authHeader.substring(AUTH_HEADER_PREFIX.length()))
+          .whenComplete(
+              (token, e) ->
+                  context.run(
+                      () -> {
+                        if (e == null) {
+                          delayedListener.setDelegate(
+                              Contexts.interceptCall(
+                                  Context.current().withValue(AccessTokenContextKey, token),
+                                  call,
+                                  headers,
+                                  next));
+                        } else {
+                          delayedListener.setDelegate(handleException(e, call));
+                        }
+                      }));
     } catch (Exception e) {
-      String msg =
-          Constants.AuthorizationMetadataKey.name()
-              + " header validation failed: "
-              + e.getMessage();
-      LOGGER.warn(msg, e);
-      call.close(Status.UNAUTHENTICATED.withDescription(msg).withCause(e), new Metadata());
-      return new ServerCall.Listener<ReqT>() {};
+      return handleException(e, call);
     }
-    return Contexts.interceptCall(
-        Context.current().withValue(AccessTokenContextKey, token), call, headers, next);
+    return delayedListener;
+  }
+
+  private <ReqT, RespT> ServerCall.Listener<ReqT> handleException(
+      Throwable e, ServerCall<ReqT, RespT> call) {
+    String msg =
+        Constants.AuthorizationMetadataKey.name() + " header validation failed: " + e.getMessage();
+    LOGGER.warn(msg, e);
+    call.close(Status.UNAUTHENTICATED.withDescription(msg).withCause(e), new Metadata());
+    return new ServerCall.Listener<ReqT>() {};
   }
 }

core/src/main/java/com/avast/grpc/jwt/server/JwtTokenParser.java

@@ -1,8 +1,10 @@
 package com.avast.grpc.jwt.server;
 
+import java.util.concurrent.CompletableFuture;
+
 @FunctionalInterface
 public interface JwtTokenParser<T> {
 
   /** Get valid JWT token, throws an exception otherwise. */
-  T parseToValid(String jwtToken) throws Exception;
+  CompletableFuture<T> parseToValid(String jwtToken);
 }

core/src/test/java/com/avast/grpc/jwt/server/JwtServerInterceptorTest.java

@@ -7,12 +7,21 @@
 import io.grpc.Metadata;
 import io.grpc.ServerCall;
 import io.grpc.ServerCallHandler;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicReference;
 import org.junit.Test;
 
 public class JwtServerInterceptorTest {
 
-  JwtTokenParser<String> jwtTokenParser = jwtToken -> jwtToken;
+  JwtTokenParser<String> jwtTokenParser =
+      jwtToken -> {
+        if (jwtToken.equals("Invalid Token")) {
+          CompletableFuture<String> res = new CompletableFuture<>();
+          res.completeExceptionally(new RuntimeException("invalid token"));
+          return res;
+        }
+        return CompletableFuture.completedFuture(jwtToken);
+      };
   JwtServerInterceptor<String> target = new JwtServerInterceptor<>(jwtTokenParser);
   ServerCall<Object, Object> serverCall = (ServerCall<Object, Object>) mock(ServerCall.class);
   ServerCallHandler<Object, Object> next =
@@ -34,6 +43,15 @@ public void closesCallOnInvalidHeader() {
     verify(next, never()).startCall(any(), any());
   }
 
+  @Test
+  public void closesCallOnInvalidToken() {
+    Metadata metadata = new Metadata();
+    metadata.put(Constants.AuthorizationMetadataKey, "Bearer Invalid Token");
+    target.interceptCall(serverCall, metadata, next);
+    verify(serverCall).close(any(), any());
+    verify(next, never()).startCall(any(), any());
+  }
+
   @Test
   public void callNextStageWithContextKeyOnValidHeader() {
     Metadata metadata = new Metadata();

keycloak/src/main/java/com/avast/grpc/jwt/keycloak/server/DefaultKeycloakPublicKeyProvider.java

@@ -0,0 +1,79 @@
+package com.avast.grpc.jwt.keycloak.server;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import java.io.IOException;
+import java.net.URL;
+import java.security.PublicKey;
+import java.time.Clock;
+import java.time.Duration;
+import java.time.Instant;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import org.keycloak.constants.ServiceUrlConstants;
+import org.keycloak.jose.jwk.JSONWebKeySet;
+import org.keycloak.jose.jwk.JWK;
+import org.keycloak.util.JWKSUtils;
+
+public class DefaultKeycloakPublicKeyProvider implements KeycloakPublicKeyProvider {
+
+  private final String serverUrl;
+  private final String realm;
+  private final Duration minTimeBetweenJwksRequests;
+  private final Duration publicKeyCacheTtl;
+  private final Clock clock;
+
+  private Map<String, PublicKey> currentKeys = new ConcurrentHashMap<>();
+  private volatile Instant lastRequestTime = Instant.MIN;
+
+  public DefaultKeycloakPublicKeyProvider(
+      String serverUrl,
+      String realm,
+      Duration minTimeBetweenJwksRequests,
+      Duration publicKeyCacheTtl,
+      Clock clock) {
+    this.serverUrl = serverUrl;
+    this.realm = realm;
+    this.minTimeBetweenJwksRequests = minTimeBetweenJwksRequests;
+    this.publicKeyCacheTtl = publicKeyCacheTtl;
+    this.clock = clock;
+  }
+
+  @Override
+  public PublicKey get(String keyId) {
+    if (lastRequestTime.plus(publicKeyCacheTtl).isBefore(clock.instant())) {
+      updateKeys();
+    }
+    PublicKey fromCache = currentKeys.get(keyId);
+    if (fromCache != null) {
+      return fromCache;
+    }
+    updateKeys();
+    PublicKey res = currentKeys.get(keyId);
+    if (res == null) {
+      throw new RuntimeException("Key with following ID not found: " + keyId);
+    }
+    return res;
+  }
+
+  protected void updateKeys() {
+    synchronized (this) {
+      if (clock.instant().isAfter(lastRequestTime.plus(minTimeBetweenJwksRequests))) {
+        Map<String, PublicKey> newKeys = fetchNewKeys();
+        currentKeys.clear();
+        currentKeys.putAll(newKeys);
+        lastRequestTime = clock.instant();
+      }
+    }
+  }
+
+  protected Map<String, PublicKey> fetchNewKeys() {
+    try {
+      ObjectMapper om = new ObjectMapper();
+      String jwksUrl = serverUrl + ServiceUrlConstants.JWKS_URL.replace("{realm-name}", realm);
+      JSONWebKeySet jwks = om.readValue(new URL(jwksUrl).openStream(), JSONWebKeySet.class);
+      return JWKSUtils.getKeysForUse(jwks, JWK.Use.SIG);
+    } catch (IOException e) {
+      throw new RuntimeException("Cannot fetch key from Keycloak server", e);
+    }
+  }
+}

keycloak/src/main/java/com/avast/grpc/jwt/keycloak/server/KeycloakJwtServerInterceptor.java

@@ -5,6 +5,7 @@
 import com.avast.grpc.jwt.server.JwtServerInterceptor;
 import com.avast.grpc.jwt.server.JwtTokenParser;
 import com.typesafe.config.Config;
+import java.time.Clock;
 import org.keycloak.representations.AccessToken;
 
 public class KeycloakJwtServerInterceptor extends JwtServerInterceptor<AccessToken> {
@@ -14,8 +15,16 @@ public KeycloakJwtServerInterceptor(JwtTokenParser<AccessToken> tokenParser) {
 
   public static KeycloakJwtServerInterceptor fromConfig(Config config) {
     Config fc = config.withFallback(DefaultConfig);
+    KeycloakPublicKeyProvider publicKeyProvider =
+        new DefaultKeycloakPublicKeyProvider(
+            fc.getString("serverUrl"),
+            fc.getString("realm"),
+            fc.getDuration("minTimeBetweenJwksRequests"),
+            fc.getDuration("publicKeyCacheTtl"),
+            Clock.systemUTC());
     KeycloakJwtTokenParser tokenParser =
-        KeycloakJwtTokenParser.create(fc.getString("serverUrl"), fc.getString("realm"));
+        new KeycloakJwtTokenParser(
+            fc.getString("serverUrl"), fc.getString("realm"), publicKeyProvider);
     tokenParser = tokenParser.withExpectedAudience(fc.getString("expectedAudience"));
     tokenParser = tokenParser.withExpectedIssuedFor(fc.getString("expectedIssuedFor"));
     return new KeycloakJwtServerInterceptor(tokenParser);

keycloak/src/main/java/com/avast/grpc/jwt/keycloak/server/KeycloakJwtTokenParser.java

@@ -1,28 +1,25 @@
 package com.avast.grpc.jwt.keycloak.server;
 
 import com.avast.grpc.jwt.server.JwtTokenParser;
-import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.base.Strings;
-import java.net.URL;
-import java.security.PublicKey;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
 import org.keycloak.TokenVerifier;
 import org.keycloak.common.VerificationException;
 import org.keycloak.constants.ServiceUrlConstants;
-import org.keycloak.jose.jwk.JSONWebKeySet;
-import org.keycloak.jose.jwk.JWK;
-import org.keycloak.jose.jwk.JWKParser;
 import org.keycloak.representations.AccessToken;
 import org.keycloak.util.TokenUtil;
 
 public class KeycloakJwtTokenParser implements JwtTokenParser<AccessToken> {
 
-  protected final PublicKey publicKey;
+  protected final KeycloakPublicKeyProvider publicKeyProvider;
   protected final TokenVerifier.Predicate<AccessToken>[] checks;
   protected String expectedAudience;
   protected String expectedIssuedFor;
 
-  protected KeycloakJwtTokenParser(String serverUrl, String realm, PublicKey publicKey) {
-    this.publicKey = publicKey;
+  public KeycloakJwtTokenParser(
+      String serverUrl, String realm, KeycloakPublicKeyProvider publicKeyProvider) {
+    this.publicKeyProvider = publicKeyProvider;
     String realmUrl =
         serverUrl + ServiceUrlConstants.REALM_INFO_PATH.replace("{realm-name}", realm);
     this.checks =
@@ -35,20 +32,36 @@ protected KeycloakJwtTokenParser(String serverUrl, String realm, PublicKey publi
   }
 
   @Override
-  public AccessToken parseToValid(String jwtToken) throws VerificationException {
-    TokenVerifier<AccessToken> verifier = createTokenVerifier(jwtToken);
-    return verifier.verify().getToken();
+  public CompletableFuture<AccessToken> parseToValid(String jwtToken) {
+    TokenVerifier<AccessToken> verifier;
+    try {
+      verifier = createTokenVerifier(jwtToken);
+    } catch (VerificationException e) {
+      CompletableFuture<AccessToken> r = new CompletableFuture<>();
+      r.completeExceptionally(e);
+      return r;
+    }
+    return CompletableFuture.supplyAsync(
+        () -> {
+          try {
+            return verifier.verify().getToken();
+          } catch (VerificationException e) {
+            throw new CompletionException(e);
+          }
+        });
   }
 
-  protected TokenVerifier<AccessToken> createTokenVerifier(String jwtToken) {
+  protected TokenVerifier<AccessToken> createTokenVerifier(String jwtToken)
+      throws VerificationException {
     TokenVerifier<AccessToken> verifier =
-        TokenVerifier.create(jwtToken, AccessToken.class).withChecks(checks).publicKey(publicKey);
+        TokenVerifier.create(jwtToken, AccessToken.class).withChecks(checks);
     if (!Strings.isNullOrEmpty(expectedAudience)) {
       verifier = verifier.audience(expectedAudience);
     }
     if (!Strings.isNullOrEmpty(expectedIssuedFor)) {
       verifier = verifier.issuedFor(expectedIssuedFor);
     }
+    verifier.publicKey(publicKeyProvider.get(verifier.getHeader().getKeyId()));
     return verifier;
   }
 
@@ -61,20 +74,4 @@ public KeycloakJwtTokenParser withExpectedIssuedFor(String expectedIssuedFor) {
     this.expectedIssuedFor = expectedIssuedFor;
     return this;
   }
-
-  public static KeycloakJwtTokenParser create(String serverUrl, String realm) {
-    try {
-      ObjectMapper om = new ObjectMapper();
-      String jwksUrl = serverUrl + ServiceUrlConstants.JWKS_URL.replace("{realm-name}", realm);
-      JSONWebKeySet jwks = om.readValue(new URL(jwksUrl).openStream(), JSONWebKeySet.class);
-      if (jwks.getKeys().length == 0) {
-        throw new RuntimeException("No keys found");
-      }
-      JWK jwk = jwks.getKeys()[0];
-      PublicKey publicKey = JWKParser.create(jwk).toPublicKey();
-      return new KeycloakJwtTokenParser(serverUrl, realm, publicKey);
-    } catch (Exception e) {
-      throw new RuntimeException("Exception when obtaining public key from " + serverUrl, e);
-    }
-  }
 }

keycloak/src/main/java/com/avast/grpc/jwt/keycloak/server/KeycloakPublicKeyProvider.java

@@ -0,0 +1,7 @@
+package com.avast.grpc.jwt.keycloak.server;
+
+import java.security.PublicKey;
+
+public interface KeycloakPublicKeyProvider {
+  PublicKey get(String keyId);
+}

keycloak/src/main/resources/reference.conf

@@ -10,4 +10,6 @@ keycloakDefaults {
   // for server
   expectedAudience = ""
   expectedIssuedFor = ""
+  minTimeBetweenJwksRequests = 10 seconds
+  publicKeyCacheTtl = 1 day
 }