fix(workloadapi): handle rejected JWT bundle updates in JWT sources

java-spiffe-core/src/main/java/io/spiffe/bundle/jwtbundle/JwtBundle.java

@@ -87,7 +87,8 @@ public static JwtBundle load(TrustDomain trustDomain,
      * @param trustDomain a {@link TrustDomain}
      * @param bundleBytes an array of bytes representing the JWT bundle.
      * @return an instance of a {@link JwtBundle}
-     * @throws JwtBundleException if there is an error reading or parsing the file, or if a keyId is empty
+     * @throws JwtBundleException if there is an error reading or parsing the file, if a keyId is empty,
+     *                            or if the bundle contains a key type that is not supported
      */
     public static JwtBundle parse(
             TrustDomain trustDomain,
@@ -99,7 +100,7 @@ public static JwtBundle parse(
         try {
             JWKSet jwkSet = JWKSet.parse(new String(bundleBytes));
             return toJwtBundle(trustDomain, jwkSet);
-        } catch (ParseException | JOSEException | KeyException e) {
+        } catch (IllegalArgumentException | ParseException | JOSEException | KeyException e) {
             throw new JwtBundleException("Could not parse bundle from bytes", e);
         }
     }

java-spiffe-core/src/main/java/io/spiffe/workloadapi/CachedJwtSource.java

@@ -19,6 +19,7 @@
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
@@ -284,7 +285,8 @@ private boolean isTokenPastHalfLifetime(JwtSvid jwtSvid) {
 
     private void init(final Duration timeout) throws TimeoutException {
         CountDownLatch done = new CountDownLatch(1);
-        setJwtBundlesWatcher(done);
+        AtomicReference<Throwable> initError = new AtomicReference<>();
+        setJwtBundlesWatcher(done, initError);
 
         boolean success;
         if (timeout.isZero()) {
@@ -296,9 +298,13 @@ private void init(final Duration timeout) throws TimeoutException {
         if (!success) {
             throw new TimeoutException("Timeout waiting for JWT bundles update");
         }
+        Throwable error = initError.get();
+        if (error != null) {
+            throw new WatcherException("Error fetching JwtBundleSet", error);
+        }
     }
 
-    private void setJwtBundlesWatcher(final CountDownLatch done) {
+    private void setJwtBundlesWatcher(final CountDownLatch done, final AtomicReference<Throwable> initError) {
         workloadApiClient.watchJwtBundles(new Watcher<JwtBundleSet>() {
             @Override
             public void onUpdate(final JwtBundleSet update) {
@@ -310,8 +316,13 @@ public void onUpdate(final JwtBundleSet update) {
             @Override
             public void onError(final Throwable error) {
                 log.log(Level.SEVERE, "Error in JwtBundleSet watcher", error);
-                done.countDown();
-                throw new WatcherException("Error fetching JwtBundleSet", error);
+                // An error before the first update fails the initialization. After that, the current bundles
+                // are kept: throwing here would escape the gRPC stream observer and cancel the watch, so later
+                // bundle updates would never be applied.
+                if (done.getCount() > 0) {
+                    initError.compareAndSet(null, error);
+                    done.countDown();
+                }
             }
         });
     }

java-spiffe-core/src/main/java/io/spiffe/workloadapi/DefaultJwtSource.java

@@ -15,6 +15,7 @@
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
@@ -188,7 +189,8 @@ public void close() {
 
     private void init(final Duration timeout) throws TimeoutException {
         CountDownLatch done = new CountDownLatch(1);
-        setJwtBundlesWatcher(done);
+        AtomicReference<Throwable> initError = new AtomicReference<>();
+        setJwtBundlesWatcher(done, initError);
 
         boolean success;
         if (timeout.isZero()) {
@@ -200,9 +202,13 @@ private void init(final Duration timeout) throws TimeoutException {
         if (!success) {
             throw new TimeoutException("Timeout waiting for JWT bundles update");
         }
+        Throwable error = initError.get();
+        if (error != null) {
+            throw new WatcherException("Error fetching JwtBundleSet", error);
+        }
     }
 
-    private void setJwtBundlesWatcher(final CountDownLatch done) {
+    private void setJwtBundlesWatcher(final CountDownLatch done, final AtomicReference<Throwable> initError) {
         workloadApiClient.watchJwtBundles(new Watcher<JwtBundleSet>() {
             @Override
             public void onUpdate(final JwtBundleSet update) {
@@ -214,8 +220,13 @@ public void onUpdate(final JwtBundleSet update) {
             @Override
             public void onError(final Throwable error) {
                 log.log(Level.SEVERE, "Error in JwtBundleSet watcher", error);
-                done.countDown();
-                throw new WatcherException("Error fetching JwtBundleSet", error);
+                // An error before the first update fails the initialization. After that, the current bundles
+                // are kept: throwing here would escape the gRPC stream observer and cancel the watch, so later
+                // bundle updates would never be applied.
+                if (done.getCount() > 0) {
+                    initError.compareAndSet(null, error);
+                    done.countDown();
+                }
             }
         });
     }

java-spiffe-core/src/test/java/io/spiffe/bundle/jwtbundle/JwtBundleTest.java

@@ -253,6 +253,21 @@ void testParse_MissingKid_Fails() throws URISyntaxException, IOException {
         }
     }
 
+    @Test
+    void testParse_InvalidKeyType_ThrowsJwtBundleException() throws URISyntaxException, IOException {
+        Path path = Paths.get(toUri("testdata/jwtbundle/jwks_invalid_keytype.json"));
+        byte[] bundleBytes = Files.readAllBytes(path);
+        TrustDomain trustDomain = TrustDomain.parse("domain.test");
+
+        try {
+            JwtBundle.parse(trustDomain, bundleBytes);
+            fail("should have thrown exception");
+        } catch (JwtBundleException e) {
+            assertEquals("Could not parse bundle from bytes", e.getMessage());
+            assertEquals("Unsupported JWT family algorithm: OKP", e.getCause().getMessage());
+        }
+    }
+
     @Test
     void testParseInvalidJson() throws KeyException {
         try {

java-spiffe-core/src/test/java/io/spiffe/workloadapi/CachedJwtSourceTest.java

@@ -4,7 +4,9 @@
 import com.nimbusds.jose.jwk.Curve;
 import com.nimbusds.jwt.JWTClaimsSet;
 import io.spiffe.bundle.jwtbundle.JwtBundle;
+import io.spiffe.bundle.jwtbundle.JwtBundleSet;
 import io.spiffe.exception.BundleNotFoundException;
+import io.spiffe.exception.JwtBundleException;
 import io.spiffe.exception.JwtSourceException;
 import io.spiffe.exception.JwtSvidException;
 import io.spiffe.exception.SocketEndpointAddressException;
@@ -100,6 +102,19 @@ void testGetBundleForTrustDomain_SourceIsClosed_ThrowsIllegalStateException() th
         }
     }
 
+    @Test
+    void testJwtBundlesWatcherError_afterInit_keepsBundlesAndAppliesLaterUpdates() throws BundleNotFoundException {
+        TrustDomain trustDomain = TrustDomain.parse("example.org");
+        JwtBundle initialBundle = jwtSource.getBundleForTrustDomain(trustDomain);
+
+        workloadApiClient.jwtBundlesWatcher.onError(new JwtBundleException("Error processing JWT bundles update"));
+        assertSame(initialBundle, jwtSource.getBundleForTrustDomain(trustDomain));
+
+        JwtBundle updatedBundle = new JwtBundle(trustDomain);
+        workloadApiClient.jwtBundlesWatcher.onUpdate(JwtBundleSet.of(Collections.singleton(updatedBundle)));
+        assertSame(updatedBundle, jwtSource.getBundleForTrustDomain(trustDomain));
+    }
+
     @Test
     void testFetchJwtSvidWithSubject() {
         try {

java-spiffe-core/src/test/java/io/spiffe/workloadapi/DefaultJwtSourceTest.java

@@ -1,25 +1,37 @@
 package io.spiffe.workloadapi;
 
 import com.google.common.collect.Sets;
+import com.google.protobuf.ByteString;
+import io.grpc.stub.StreamObserver;
+import io.grpc.testing.GrpcCleanupRule;
 import io.spiffe.bundle.jwtbundle.JwtBundle;
+import io.spiffe.bundle.jwtbundle.JwtBundleSet;
 import io.spiffe.exception.BundleNotFoundException;
+import io.spiffe.exception.JwtBundleException;
 import io.spiffe.exception.JwtSourceException;
 import io.spiffe.exception.JwtSvidException;
 import io.spiffe.exception.SocketEndpointAddressException;
 import io.spiffe.spiffeid.SpiffeId;
 import io.spiffe.spiffeid.TrustDomain;
 import io.spiffe.svid.jwtsvid.JwtSvid;
+import io.spiffe.workloadapi.grpc.SpiffeWorkloadAPIGrpc;
+import io.spiffe.workloadapi.grpc.Workload;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import uk.org.webcompere.systemstubs.environment.EnvironmentVariables;
 
 import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Paths;
 import java.time.Duration;
+import java.util.Collections;
 import java.util.List;
 
+import static io.spiffe.utils.TestUtils.toUri;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
 
@@ -80,6 +92,40 @@ void testGetBundleForTrustDomain_SourceIsClosed_ThrowsIllegalStateException() th
         }
     }
 
+    @Test
+    void testJwtBundlesWatcherError_afterInit_keepsBundlesAndAppliesLaterUpdates() throws BundleNotFoundException {
+        TrustDomain trustDomain = TrustDomain.parse("example.org");
+        JwtBundle initialBundle = jwtSource.getBundleForTrustDomain(trustDomain);
+
+        workloadApiClient.jwtBundlesWatcher.onError(new JwtBundleException("Error processing JWT bundles update"));
+        assertSame(initialBundle, jwtSource.getBundleForTrustDomain(trustDomain));
+
+        JwtBundle updatedBundle = new JwtBundle(trustDomain);
+        workloadApiClient.jwtBundlesWatcher.onUpdate(JwtBundleSet.of(Collections.singleton(updatedBundle)));
+        assertSame(updatedBundle, jwtSource.getBundleForTrustDomain(trustDomain));
+    }
+
+    @Test
+    void testWatchJwtBundles_unsupportedKeyTypeUpdate_laterUpdatesAreApplied() throws Exception {
+        byte[] validBundle = Files.readAllBytes(Paths.get(toUri("testdata/workloadapi/bundle.json")));
+        byte[] unsupportedKeyTypeBundle = Files.readAllBytes(Paths.get(toUri("testdata/jwtbundle/jwks_invalid_keytype.json")));
+        SpiffeWorkloadAPIGrpc.SpiffeWorkloadAPIImplBase fakeWorkloadApi = new SpiffeWorkloadAPIGrpc.SpiffeWorkloadAPIImplBase() {
+            @Override
+            public void fetchJWTBundles(Workload.JWTBundlesRequest request, StreamObserver<Workload.JWTBundlesResponse> responseObserver) {
+                responseObserver.onNext(newJwtBundlesResponse("example.org", validBundle));
+                responseObserver.onNext(newJwtBundlesResponse("example.org", unsupportedKeyTypeBundle));
+                responseObserver.onNext(newJwtBundlesResponse("domain.test", validBundle));
+            }
+        };
+        DefaultWorkloadApiClient client = WorkloadApiClientTestUtil.create(fakeWorkloadApi, new GrpcCleanupRule());
+        JwtSourceOptions options = JwtSourceOptions.builder().workloadApiClient(client).build();
+
+        try (JwtSource source = DefaultJwtSource.newSource(options)) {
+            JwtBundle bundle = source.getBundleForTrustDomain(TrustDomain.parse("domain.test"));
+            assertEquals(TrustDomain.parse("domain.test"), bundle.getTrustDomain());
+        }
+    }
+
     @Test
     void testFetchJwtSvidWithSubject() {
         try {
@@ -281,4 +327,11 @@ void newSource_noSocketAddress() throws Exception {
             }
         });
     }
-}
+
+    private static Workload.JWTBundlesResponse newJwtBundlesResponse(String trustDomain, byte[] bundle) {
+        return Workload.JWTBundlesResponse
+                .newBuilder()
+                .putBundles(trustDomain, ByteString.copyFrom(bundle))
+                .build();
+    }
+}

java-spiffe-core/src/test/java/io/spiffe/workloadapi/GrpcConversionUtilsTest.java

@@ -53,6 +53,22 @@ void test_toJwtBundleSet_emtpyResponse() {
         }
     }
 
+    @Test
+    void test_toJwtBundleSet_unsupportedKeyType() throws URISyntaxException, IOException {
+        byte[] bundleBytes = Files.readAllBytes(Paths.get(toUri("testdata/jwtbundle/jwks_invalid_keytype.json")));
+        Workload.JWTBundlesResponse response = Workload.JWTBundlesResponse
+                .newBuilder()
+                .putBundles("example.org", ByteString.copyFrom(bundleBytes))
+                .build();
+
+        try {
+            GrpcConversionUtils.toJwtBundleSet(response);
+            fail("should have thrown exception");
+        } catch (JwtBundleException e) {
+            assertEquals("Could not parse bundle from bytes", e.getMessage());
+        }
+    }
+
     @Test
     void test_parseX509Bundle_corruptedBytes() {
         try {

java-spiffe-core/src/test/java/io/spiffe/workloadapi/WorkloadApiClientStub.java

@@ -39,6 +39,8 @@ public class WorkloadApiClientStub implements WorkloadApiClient {
 
     int fetchJwtSvidCallCount = 0;
 
+    Watcher<JwtBundleSet> jwtBundlesWatcher;
+
     boolean closed;
 
     Clock clock = Clock.systemDefaultZone();
@@ -106,6 +108,7 @@ public JwtSvid validateJwtSvid(String token, String audience) throws JwtSvidExce
 
     @Override
     public void watchJwtBundles(Watcher<JwtBundleSet> watcher) {
+        jwtBundlesWatcher = watcher;
         JwtBundleSet jwtBundleSet = generateJwtBundleSet();
         watcher.onUpdate(jwtBundleSet);
     }