Add support for registration of listeners on ServiceIdentityProvider

4 files changed, 74 insertions, 12 deletions

diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identity/ServiceIdentityProvider.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identity/ServiceIdentityProvider.java
index 6b318fb16be..f945783cf8a 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identity/ServiceIdentityProvider.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identity/ServiceIdentityProvider.java
@@ -15,4 +15,10 @@ import javax.net.ssl.SSLContext;
 public interface ServiceIdentityProvider {
     AthenzService identity();
     SSLContext getIdentitySslContext();
+    void addIdentityListener(Listener listener);
+    void removeIdentityListener(Listener listener);
+
+    interface Listener {
+        void onCredentialsUpdate(SSLContext sslContext, AthenzService identity);
+    }
 }
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identity/ServiceIdentityProviderListenerHelper.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identity/ServiceIdentityProviderListenerHelper.java
new file mode 100644
index 00000000000..836e46201ee
--- /dev/null
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identity/ServiceIdentityProviderListenerHelper.java
@@ -0,0 +1,40 @@
+// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.athenz.identity;
+
+import com.yahoo.vespa.athenz.api.AthenzService;
+
+import javax.net.ssl.SSLContext;
+import java.util.Set;
+import java.util.concurrent.ConcurrentSkipListSet;
+
+/**
+ * A helper class managing {@link ServiceIdentityProvider.Listener} instances for implementations of {@link ServiceIdentityProvider}.
+ *
+ * @author bjorncs
+ */
+public class ServiceIdentityProviderListenerHelper {
+
+    private final Set<ServiceIdentityProvider.Listener> listeners = new ConcurrentSkipListSet<>();
+    private final AthenzService identity;
+
+    public ServiceIdentityProviderListenerHelper(AthenzService identity) {
+        this.identity = identity;
+    }
+
+    public void addIdentityListener(ServiceIdentityProvider.Listener listener) {
+        listeners.add(listener);
+    }
+
+    public void removeIdentityListener(ServiceIdentityProvider.Listener listener) {
+        listeners.remove(listener);
+    }
+
+    public void onCredentialsUpdate(SSLContext sslContext) {
+        listeners.forEach(l -> l.onCredentialsUpdate(sslContext, identity));
+    }
+
+    public void clearListeners() {
+        listeners.clear();
+    }
+
+}
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identity/SiaIdentityProvider.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identity/SiaIdentityProvider.java
index 161438e2bbe..fb71ed65da1 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identity/SiaIdentityProvider.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identity/SiaIdentityProvider.java
@@ -13,16 +13,15 @@ import java.io.File;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.time.Duration;
-import java.util.Set;
-import java.util.concurrent.ConcurrentSkipListSet;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
-import java.util.function.Consumer;
 import java.util.logging.Logger;
 
 /**
+ * A {@link ServiceIdentityProvider} that provides the credentials stored on file system.
+ *
  * @author mortent
  * @author bjorncs
  */
@@ -38,7 +37,7 @@ public class SiaIdentityProvider extends AbstractComponent implements ServiceIde
     private final File certificateFile;
     private final File trustStoreFile;
     private final ScheduledExecutorService scheduler;
-    private final Set<Consumer<SSLContext>> listeners = new ConcurrentSkipListSet<>();
+    private final ServiceIdentityProviderListenerHelper listenerHelper;
 
     @Inject
     public SiaIdentityProvider(SiaProviderConfig config) {
@@ -70,6 +69,7 @@ public class SiaIdentityProvider extends AbstractComponent implements ServiceIde
         this.trustStoreFile = trustStoreFile;
         this.scheduler = scheduler;
         this.sslContext.set(createIdentitySslContext());
+        this.listenerHelper = new ServiceIdentityProviderListenerHelper(service);
         scheduler.scheduleAtFixedRate(this::reloadSslContext, REFRESH_INTERVAL.toMinutes(), REFRESH_INTERVAL.toMinutes(), TimeUnit.MINUTES);
     }
 
@@ -91,12 +91,14 @@ public class SiaIdentityProvider extends AbstractComponent implements ServiceIde
         return sslContext.get();
     }
 
-    public void addReloadListener(Consumer<SSLContext> listener) {
-        listeners.add(listener);
+    @Override
+    public void addIdentityListener(Listener listener) {
+        listenerHelper.addIdentityListener(listener);
     }
 
-    public void removeReloadListener(Consumer<SSLContext> listener) {
-        listeners.remove(listener);
+    @Override
+    public void removeIdentityListener(Listener listener) {
+        listenerHelper.removeIdentityListener(listener);
     }
 
     private SSLContext createIdentitySslContext() {
@@ -111,7 +113,7 @@ public class SiaIdentityProvider extends AbstractComponent implements ServiceIde
         try {
             SSLContext sslContext = createIdentitySslContext();
             this.sslContext.set(sslContext);
-            listeners.forEach(listener -> listener.accept(sslContext));
+            listenerHelper.onCredentialsUpdate(sslContext);
         } catch (Exception e) {
             log.log(LogLevel.SEVERE, "Failed to update SSLContext: " + e.getMessage(), e);
         }
@@ -130,7 +132,7 @@ public class SiaIdentityProvider extends AbstractComponent implements ServiceIde
         try {
             scheduler.shutdownNow();
             scheduler.awaitTermination(90, TimeUnit.SECONDS);
-            listeners.clear();
+            listenerHelper.clearListeners();
         } catch (InterruptedException e) {
             throw new RuntimeException(e);
         }
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzIdentityProviderImpl.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzIdentityProviderImpl.java
index eb0ae89fdcf..3817774cf06 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzIdentityProviderImpl.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identityprovider/client/AthenzIdentityProviderImpl.java
@@ -11,6 +11,7 @@ import com.yahoo.log.LogLevel;
 import com.yahoo.vespa.athenz.api.AthenzDomain;
 import com.yahoo.vespa.athenz.api.AthenzService;
 import com.yahoo.vespa.athenz.identity.ServiceIdentityProvider;
+import com.yahoo.vespa.athenz.identity.ServiceIdentityProviderListenerHelper;
 import com.yahoo.vespa.athenz.tls.KeyStoreType;
 import com.yahoo.vespa.athenz.tls.SslContextBuilder;
 import com.yahoo.vespa.defaults.Defaults;
@@ -49,6 +50,7 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
     private final ScheduledExecutorService scheduler;
     private final Clock clock;
     private final AthenzService identity;
+    private final ServiceIdentityProviderListenerHelper listenerHelper;
 
     // TODO IdentityConfig should contain ZTS uri and dns suffix
     @Inject
@@ -74,6 +76,7 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
         this.scheduler = scheduler;
         this.clock = clock;
         this.identity = new AthenzService(config.domain(), config.service());
+        this.listenerHelper = new ServiceIdentityProviderListenerHelper(this.identity);
         registerInstance();
     }
 
@@ -108,6 +111,16 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
     }
 
     @Override
+    public void addIdentityListener(Listener listener) {
+        listenerHelper.addIdentityListener(listener);
+    }
+
+    @Override
+    public void removeIdentityListener(Listener listener) {
+        listenerHelper.removeIdentityListener(listener);
+    }
+
+    @Override
     public SSLContext getRoleSslContext(String domain, String role) {
         // This ssl context should ideally be cached as it is quite expensive to create.
         PrivateKey privateKey = credentials.getKeyPair().getPrivate();
@@ -151,6 +164,7 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
         try {
             scheduler.shutdownNow();
             scheduler.awaitTermination(AWAIT_TERMINTATION_TIMEOUT.getSeconds(), TimeUnit.SECONDS);
+            listenerHelper.clearListeners();
         } catch (InterruptedException e) {
             throw new RuntimeException(e);
         }
@@ -170,10 +184,10 @@ public final class AthenzIdentityProviderImpl extends AbstractComponent implemen
 
     void refreshCertificate() {
         try {
-            AthenzCredentials newCredentials = isExpired(credentials)
+            credentials = isExpired(credentials)
                     ? athenzCredentialsService.registerInstance()
                     : athenzCredentialsService.updateCredentials(credentials.getIdentityDocument(), credentials.getIdentitySslContext());
-            credentials = newCredentials;
+            listenerHelper.onCredentialsUpdate(credentials.getIdentitySslContext());
         } catch (Throwable t) {
             log.log(LogLevel.WARNING, "Failed to update credentials: " + t.getMessage(), t);
         }