From 70e7ad4346091d078914e1f000620333fe871cf5 Mon Sep 17 00:00:00 2001
From: Max Lambrecht <maxlambrecht@gmail.com>
Date: Tue, 18 Sep 2018 14:15:58 -0300
Subject: [PATCH] Adding CountDownLatch to make the SpiffeIdManager consumer
 wait until the SpiffeSVID has been initialized

Signed-off-by: Max Lambrecht <maxlambrecht@gmail.com>
---
 .../java/spiffe/provider/SpiffeIdManager.java | 32 ++++++++++++++++---
 1 file changed, 27 insertions(+), 5 deletions(-)

diff --git a/src/main/java/spiffe/provider/SpiffeIdManager.java b/src/main/java/spiffe/provider/SpiffeIdManager.java
index 2484353..f5f2c6d 100644
--- a/src/main/java/spiffe/provider/SpiffeIdManager.java
+++ b/src/main/java/spiffe/provider/SpiffeIdManager.java
@@ -2,13 +2,14 @@ package spiffe.provider;
 
 import spiffe.api.svid.Fetcher;
 import spiffe.api.svid.Workload;
-import spiffe.api.svid.Workload.X509SVID;
 import spiffe.api.svid.X509SVIDFetcher;
 
 import java.security.PrivateKey;
 import java.security.cert.X509Certificate;
-import java.util.List;
 import java.util.Set;
+import java.util.concurrent.CountDownLatch;
+import java.util.logging.Level;
+import java.util.logging.Logger;
 
 import static java.util.Collections.EMPTY_SET;
 
@@ -21,6 +22,7 @@ import static java.util.Collections.EMPTY_SET;
 public class SpiffeIdManager {
 
     private static final SpiffeIdManager INSTANCE = new SpiffeIdManager();
+    private static final Logger LOGGER = Logger.getLogger(SpiffeIdManager.class.getName());
 
     static SpiffeIdManager getInstance() {
         return INSTANCE;
@@ -36,6 +38,11 @@ public class SpiffeIdManager {
      */
     private final FunctionalReadWriteLock guard;
 
+    /**
+     * Used to make the getters wait until there's a spiffeSVID initialized
+     */
+    private final CountDownLatch completedSpiffeSVIDUpdate = new CountDownLatch(1);
+
     /**
      * Private Constructor
      *
@@ -49,6 +56,7 @@ public class SpiffeIdManager {
     }
 
     public SpiffeSVID getSpiffeSVID() {
+        awaitSpiffeSVID();
         return guard.read(() -> spiffeSVID);
     }
 
@@ -58,18 +66,32 @@ public class SpiffeIdManager {
      */
     private void updateSVID(Workload.X509SVIDResponse x509SVIDResponse) {
         guard.write(() -> spiffeSVID = new SpiffeSVID(x509SVIDResponse));
+        completedSpiffeSVIDUpdate.countDown();
+        LOGGER.log(Level.FINE, "Spiffe SVID has been updated ");
     }
 
-    X509Certificate getCertificate() {
+    public X509Certificate getCertificate() {
+        awaitSpiffeSVID();
         return guard.read(() -> spiffeSVID != null ? spiffeSVID.getCertificate() : null);
     }
 
-    PrivateKey getPrivateKey() {
+    public PrivateKey getPrivateKey() {
+        awaitSpiffeSVID();
         return guard.read(() -> spiffeSVID != null ? spiffeSVID.getPrivateKey() : null);
     }
 
     @SuppressWarnings("unchecked")
-    Set<X509Certificate> getTrustedCerts() {
+    public Set<X509Certificate> getTrustedCerts() {
+        awaitSpiffeSVID();
         return guard.read(() -> spiffeSVID != null ? spiffeSVID.getTrustedCerts() : EMPTY_SET);
     }
+
+    private void awaitSpiffeSVID() {
+        try {
+            completedSpiffeSVIDUpdate.await();
+        } catch (InterruptedException e) {
+            LOGGER.info("Interrupted " + e.getMessage());
+            Thread.currentThread().interrupt();
+        }
+    }
 }