From 1e8d250d9f4414eb180b97908c717626ae33a122 Mon Sep 17 00:00:00 2001
From: Francesco Guardiani <me@slinkydeveloper.com>
Date: Mon, 3 Jun 2024 15:40:27 +0200
Subject: [PATCH] Allow more than one key (#340)

---
 .../RestateRequestIdentityVerifier.java       | 35 ++++++++++++++-----
 .../RestateRequestIdentityVerifierTest.java   |  2 +-
 2 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/sdk-request-identity/src/main/java/dev/restate/sdk/auth/signing/RestateRequestIdentityVerifier.java b/sdk-request-identity/src/main/java/dev/restate/sdk/auth/signing/RestateRequestIdentityVerifier.java
index deb56566..d7d7ea08 100644
--- a/sdk-request-identity/src/main/java/dev/restate/sdk/auth/signing/RestateRequestIdentityVerifier.java
+++ b/sdk-request-identity/src/main/java/dev/restate/sdk/auth/signing/RestateRequestIdentityVerifier.java
@@ -16,6 +16,9 @@
 import com.nimbusds.jose.util.Base64URL;
 import com.nimbusds.jwt.SignedJWT;
 import dev.restate.sdk.auth.RequestIdentityVerifier;
+import java.util.Arrays;
+import java.util.List;
+import java.util.stream.Collectors;
 
 public class RestateRequestIdentityVerifier implements RequestIdentityVerifier {
   private static final String SIGNATURE_SCHEME_HEADER = "x-restate-signature-scheme";
@@ -24,10 +27,10 @@ public class RestateRequestIdentityVerifier implements RequestIdentityVerifier {
   private static final String JWT_HEADER = "x-restate-jwt-v1";
   private static final String IDENTITY_V1_PREFIX = "publickeyv1_";
 
-  private final JWSVerifier verifier;
+  private final JWSVerifier[] verifiers;
 
-  private RestateRequestIdentityVerifier(JWSVerifier verifier) {
-    this.verifier = verifier;
+  private RestateRequestIdentityVerifier(List<JWSVerifier> verifier) {
+    this.verifiers = verifier.toArray(JWSVerifier[]::new);
   }
 
   @Override
@@ -37,10 +40,12 @@ public void verifyRequest(Headers headers) throws Exception {
       case SIGNATURE_SCHEME_V1:
         String jwtHeader = expectHeader(headers, JWT_HEADER);
         SignedJWT signedJWT = SignedJWT.parse(jwtHeader);
-        if (!signedJWT.verify(verifier)) {
-          throw new IllegalStateException("Verification of JWT token failed");
+        for (JWSVerifier verifier : verifiers) {
+          if (signedJWT.verify(verifier)) {
+            return;
+          }
         }
-        break;
+        throw new IllegalStateException("Verification of JWT token failed");
       case SIGNATURE_SCHEME_UNSIGNED:
         throw new IllegalStateException("Request has no identity, but one was expected");
       default:
@@ -58,6 +63,21 @@ private String expectHeader(Headers headers, String key) {
 
   /** Create the {@link RequestIdentityVerifier} from key strings. */
   public static RequestIdentityVerifier fromKey(String key) {
+    return fromKeys(key);
+  }
+
+  /** Create the {@link RequestIdentityVerifier} from key strings. */
+  public static RequestIdentityVerifier fromKeys(String... keys) {
+    if (keys.length == 0) {
+      throw new IllegalArgumentException("You must provide at least one key");
+    }
+    return new RestateRequestIdentityVerifier(
+        Arrays.stream(keys)
+            .map(RestateRequestIdentityVerifier::parseKey)
+            .collect(Collectors.toList()));
+  }
+
+  private static JWSVerifier parseKey(String key) {
     if (!key.startsWith(IDENTITY_V1_PREFIX)) {
       throw new IllegalArgumentException(
           "Identity v1 jwt public keys are expected to start with " + IDENTITY_V1_PREFIX);
@@ -77,7 +97,6 @@ public static RequestIdentityVerifier fromKey(String key) {
     } catch (JOSEException e) {
       throw new RuntimeException("Cannot create the verifier", e);
     }
-
-    return new RestateRequestIdentityVerifier(verifier);
+    return verifier;
   }
 }
diff --git a/sdk-request-identity/src/test/java/RestateRequestIdentityVerifierTest.java b/sdk-request-identity/src/test/java/RestateRequestIdentityVerifierTest.java
index 818e8b6a..3284a97f 100644
--- a/sdk-request-identity/src/test/java/RestateRequestIdentityVerifierTest.java
+++ b/sdk-request-identity/src/test/java/RestateRequestIdentityVerifierTest.java
@@ -13,7 +13,7 @@ public class RestateRequestIdentityVerifierTest {
 
   @Test
   void parseKey() {
-    RestateRequestIdentityVerifier.fromKey(
+    RestateRequestIdentityVerifier.fromKeys(
         "publickeyv1_ChjENKeMvCtRnqG2mrBK1HmPKufgFUc98K8B3ononQvp");
   }
 }