fix(workloadapi): keep X.509 context watch alive after rejected updates

java-spiffe-core/src/main/java/io/spiffe/workloadapi/DefaultX509Source.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.function.Function;
 import java.util.logging.Level;
 import java.util.logging.Logger;
@@ -201,7 +202,8 @@ private static WorkloadApiClient createClient(final X509SourceOptions options)
 
     private void init(final Duration timeout) throws TimeoutException {
         CountDownLatch done = new CountDownLatch(1);
-        setX509ContextWatcher(done);
+        AtomicReference<Throwable> initError = new AtomicReference<>();
+        setX509ContextWatcher(done, initError);
 
         final boolean success;
         if (timeout.isZero()) {
@@ -213,9 +215,13 @@ private void init(final Duration timeout) throws TimeoutException {
         if (!success) {
             throw new TimeoutException("Timeout waiting for X.509 Context update");
         }
+        Throwable error = initError.get();
+        if (error != null) {
+            throw new WatcherException("Error in X509Context watcher", error);
+        }
     }
 
-    private void setX509ContextWatcher(final CountDownLatch done) {
+    private void setX509ContextWatcher(final CountDownLatch done, final AtomicReference<Throwable> initError) {
         workloadApiClient.watchX509Context(new Watcher<X509Context>() {
             @Override
             public void onUpdate(final X509Context update) {
@@ -228,8 +234,13 @@ public void onUpdate(final X509Context update) {
             @Override
             public void onError(final Throwable error) {
                 log.log(Level.SEVERE, "Error in X509Context watcher", error);
-                done.countDown();
-                throw new WatcherException("Error in X509Context watcher", error);
+                // An error before the first update fails the initialization. After that, the current SVID and
+                // bundles are kept: throwing here would escape the gRPC stream observer and cancel the watch, so
+                // later X.509 context updates (SVID rotations) would never be applied.
+                if (done.getCount() > 0) {
+                    initError.compareAndSet(null, error);
+                    done.countDown();
+                }
             }
         });
     }

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

@@ -1,21 +1,33 @@
 package io.spiffe.workloadapi;
 
+import com.google.protobuf.ByteString;
+import io.grpc.stub.StreamObserver;
+import io.grpc.testing.GrpcCleanupRule;
 import io.spiffe.bundle.x509bundle.X509Bundle;
+import io.spiffe.bundle.x509bundle.X509BundleSet;
 import io.spiffe.exception.BundleNotFoundException;
 import io.spiffe.exception.SocketEndpointAddressException;
+import io.spiffe.exception.X509ContextException;
 import io.spiffe.exception.X509SourceException;
 import io.spiffe.spiffeid.SpiffeId;
 import io.spiffe.spiffeid.TrustDomain;
 import io.spiffe.svid.x509svid.X509Svid;
+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.nio.file.Files;
+import java.nio.file.Paths;
 import java.time.Duration;
+import java.util.Collections;
 
+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;
 
@@ -96,6 +108,44 @@ void testGetX509Svid_SourceIsClosed_ThrowsIllegalStateException() {
         }
     }
 
+    @Test
+    void testX509ContextWatcherError_afterInit_keepsX509ContextAndAppliesLaterUpdates() throws BundleNotFoundException {
+        TrustDomain trustDomain = TrustDomain.parse("example.org");
+        X509Svid initialSvid = x509Source.getX509Svid();
+        X509Bundle initialBundle = x509Source.getBundleForTrustDomain(trustDomain);
+
+        workloadApiClient.x509ContextWatcher.onError(new X509ContextException("Error processing X.509 Context update"));
+        assertSame(initialSvid, x509Source.getX509Svid());
+        assertSame(initialBundle, x509Source.getBundleForTrustDomain(trustDomain));
+
+        X509Bundle updatedBundle = new X509Bundle(trustDomain);
+        X509Context update = X509Context.of(Collections.singletonList(initialSvid), X509BundleSet.of(Collections.singleton(updatedBundle)));
+        workloadApiClient.x509ContextWatcher.onUpdate(update);
+        assertSame(updatedBundle, x509Source.getBundleForTrustDomain(trustDomain));
+    }
+
+    @Test
+    void testWatchX509Context_corruptedUpdate_laterUpdatesAreApplied() throws Exception {
+        byte[] svid = Files.readAllBytes(Paths.get(toUri("testdata/workloadapi/svid.der")));
+        byte[] key = Files.readAllBytes(Paths.get(toUri("testdata/workloadapi/svid.key.der")));
+        byte[] bundle = Files.readAllBytes(Paths.get(toUri("testdata/workloadapi/bundle.der")));
+        byte[] corrupted = Files.readAllBytes(Paths.get(toUri("testdata/workloadapi/corrupted")));
+        SpiffeWorkloadAPIGrpc.SpiffeWorkloadAPIImplBase fakeWorkloadApi = new SpiffeWorkloadAPIGrpc.SpiffeWorkloadAPIImplBase() {
+            @Override
+            public void fetchX509SVID(Workload.X509SVIDRequest request, StreamObserver<Workload.X509SVIDResponse> responseObserver) {
+                responseObserver.onNext(newX509SvidResponse(svid, key, bundle, "initial"));
+                responseObserver.onNext(newX509SvidResponse(corrupted, corrupted, corrupted, "corrupted"));
+                responseObserver.onNext(newX509SvidResponse(svid, key, bundle, "rotated"));
+            }
+        };
+        DefaultWorkloadApiClient client = WorkloadApiClientTestUtil.create(fakeWorkloadApi, new GrpcCleanupRule());
+        DefaultX509Source.X509SourceOptions options = DefaultX509Source.X509SourceOptions.builder().workloadApiClient(client).build();
+
+        try (DefaultX509Source source = DefaultX509Source.newSource(options)) {
+            assertEquals("rotated", source.getX509Svid().getHint());
+        }
+    }
+
     @Test
     void newSource_success() {
         DefaultX509Source.X509SourceOptions options = DefaultX509Source.X509SourceOptions
@@ -172,4 +222,19 @@ void newSource_noSocketAddress() throws Exception {
             }
         });
     }
-}
+
+    private static Workload.X509SVIDResponse newX509SvidResponse(byte[] svid, byte[] key, byte[] bundle, String hint) {
+        Workload.X509SVID x509Svid = Workload.X509SVID
+                .newBuilder()
+                .setSpiffeId("spiffe://example.org/workload-server")
+                .setX509Svid(ByteString.copyFrom(svid))
+                .setX509SvidKey(ByteString.copyFrom(key))
+                .setBundle(ByteString.copyFrom(bundle))
+                .setHint(hint)
+                .build();
+        return Workload.X509SVIDResponse
+                .newBuilder()
+                .addSvids(x509Svid)
+                .build();
+    }
+}

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

@@ -39,6 +39,8 @@ public class WorkloadApiClientStub implements WorkloadApiClient {
 
     int fetchJwtSvidCallCount = 0;
 
+    Watcher<X509Context> x509ContextWatcher;
+
     Watcher<JwtBundleSet> jwtBundlesWatcher;
 
     boolean closed;
@@ -52,6 +54,7 @@ public X509Context fetchX509Context() {
 
     @Override
     public void watchX509Context(Watcher<X509Context> watcher) {
+        x509ContextWatcher = watcher;
         X509Context update = generateX509Context();
         watcher.onUpdate(update);
     }