Merge pull request #6564 from vespa-engine/bjorncs/sia-backed-http-client

Bjorncs/sia backed http client

2 files changed, 146 insertions, 94 deletions

diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zts/DefaultZtsClient.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zts/DefaultZtsClient.java
index 8a94518cee7..951794798bf 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zts/DefaultZtsClient.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zts/DefaultZtsClient.java
@@ -22,11 +22,11 @@ import com.yahoo.vespa.athenz.client.zts.bindings.RoleTokenResponseEntity;
 import com.yahoo.vespa.athenz.client.zts.bindings.TenantDomainsResponseEntity;
 import com.yahoo.vespa.athenz.client.zts.utils.IdentityCsrGenerator;
 import com.yahoo.vespa.athenz.identity.ServiceIdentityProvider;
+import com.yahoo.vespa.athenz.identity.SiaBackedApacheHttpClient;
 import com.yahoo.vespa.athenz.tls.Pkcs10Csr;
 import com.yahoo.vespa.athenz.tls.Pkcs10CsrBuilder;
 import org.apache.http.HttpResponse;
 import org.apache.http.client.config.RequestConfig;
-import org.apache.http.client.methods.CloseableHttpResponse;
 import org.apache.http.client.methods.HttpUriRequest;
 import org.apache.http.client.methods.RequestBuilder;
 import org.apache.http.entity.ContentType;
@@ -45,9 +45,7 @@ import java.security.KeyPair;
 import java.security.cert.X509Certificate;
 import java.time.Duration;
 import java.util.List;
-import java.util.concurrent.locks.Lock;
-import java.util.concurrent.locks.ReadWriteLock;
-import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.function.Supplier;
 
 import static com.yahoo.vespa.athenz.tls.SignatureAlgorithm.SHA256_WITH_RSA;
 import static com.yahoo.vespa.athenz.tls.SubjectAlternativeName.Type.DNS_NAME;
@@ -65,27 +63,21 @@ public class DefaultZtsClient implements ZtsClient {
     private static final ObjectMapper objectMapper = new ObjectMapper().registerModule(new JavaTimeModule());
 
     private final URI ztsUrl;
-    private final ReadWriteLock lock = new ReentrantReadWriteLock();
+    private final SiaBackedApacheHttpClient client;
     private final AthenzIdentity identity;
-    private final ServiceIdentityProvider identityProvider;
-    private final ServiceIdentityProviderListener identityListener;
-    private volatile CloseableHttpClient client;
 
     public DefaultZtsClient(URI ztsUrl, AthenzIdentity identity, SSLContext sslContext) {
-        this.ztsUrl = addTrailingSlash(ztsUrl);
-        this.client = createHttpClient(sslContext);
-        this.identity = identity;
-        this.identityProvider = null;
-        this.identityListener = null;
+        this(ztsUrl, identity, () -> sslContext);
     }
 
     public DefaultZtsClient(URI ztsUrl, ServiceIdentityProvider identityProvider) {
+        this(ztsUrl, identityProvider.identity(), identityProvider::getIdentitySslContext);
+    }
+
+    private DefaultZtsClient(URI ztsUrl, AthenzIdentity identity, Supplier<SSLContext> sslContextSupplier) {
         this.ztsUrl = addTrailingSlash(ztsUrl);
-        this.client = createHttpClient(identityProvider.getIdentitySslContext());
-        this.identity = identityProvider.identity();
-        this.identityProvider = identityProvider;
-        this.identityListener = new ServiceIdentityProviderListener();
-        identityProvider.addIdentityListener(this.identityListener);
+        this.identity = identity;
+        this.client = new SiaBackedApacheHttpClient(sslContextSupplier, DefaultZtsClient::createHttpClient);
     }
 
     @Override
@@ -101,11 +93,7 @@ public class DefaultZtsClient implements ZtsClient {
                 .setUri(ztsUrl.resolve("instance/"))
                 .setEntity(toJsonStringEntity(payload))
                 .build();
-        return withClient(client -> {
-            try (CloseableHttpResponse response = client.execute(request)) {
-                return getInstanceIdentity(response);
-            }
-        });
+        return client.execute(request, DefaultZtsClient::getInstanceIdentity);
     }
 
     @Override
@@ -125,11 +113,7 @@ public class DefaultZtsClient implements ZtsClient {
                 .setUri(uri)
                 .setEntity(toJsonStringEntity(payload))
                 .build();
-        return withClient(client -> {
-            try (CloseableHttpResponse response = client.execute(request)) {
-                return getInstanceIdentity(response);
-            }
-        });
+        return client.execute(request, DefaultZtsClient::getInstanceIdentity);
     }
 
     @Override
@@ -139,11 +123,9 @@ public class DefaultZtsClient implements ZtsClient {
                 .setUri(uri)
                 .setEntity(toJsonStringEntity(new IdentityRefreshRequestEntity(csr, keyId)))
                 .build();
-        return withClient(client -> {
-            try (CloseableHttpResponse response = client.execute(request)) {
-                IdentityResponseEntity entity = readEntity(response, IdentityResponseEntity.class);
-                return new Identity(entity.certificate(), entity.caCertificateBundle());
-            }
+        return client.execute(request, response -> {
+            IdentityResponseEntity entity = readEntity(response, IdentityResponseEntity.class);
+            return new Identity(entity.certificate(), entity.caCertificateBundle());
         });
     }
 
@@ -171,11 +153,9 @@ public class DefaultZtsClient implements ZtsClient {
             requestBuilder.addParameter("role", roleName);
         }
         HttpUriRequest request = requestBuilder.build();
-        return withClient(client -> {
-            try (CloseableHttpResponse response = client.execute(request)) {
-                RoleTokenResponseEntity roleTokenResponseEntity = readEntity(response, RoleTokenResponseEntity.class);
-                return roleTokenResponseEntity.token;
-            }
+        return client.execute(request, response -> {
+            RoleTokenResponseEntity roleTokenResponseEntity = readEntity(response, RoleTokenResponseEntity.class);
+            return roleTokenResponseEntity.token;
         });
     }
 
@@ -194,11 +174,9 @@ public class DefaultZtsClient implements ZtsClient {
         HttpUriRequest request = RequestBuilder.post(uri)
                 .setEntity(toJsonStringEntity(requestEntity))
                 .build();
-        return withClient(client -> {
-            try (CloseableHttpResponse response = client.execute(request)) {
-                RoleCertificateResponseEntity responseEntity = readEntity(response, RoleCertificateResponseEntity.class);
-                return responseEntity.certificate;
-            }
+        return client.execute(request, response -> {
+            RoleCertificateResponseEntity responseEntity = readEntity(response, RoleCertificateResponseEntity.class);
+            return responseEntity.certificate;
         });
     }
 
@@ -217,11 +195,9 @@ public class DefaultZtsClient implements ZtsClient {
                 .addParameter("roleName", roleName)
                 .addParameter("serviceName", providerIdentity.getName())
                 .build();
-        return withClient(client -> {
-            try (CloseableHttpResponse response = client.execute(request)) {
-                TenantDomainsResponseEntity entity = readEntity(response, TenantDomainsResponseEntity.class);
-                return entity.tenantDomainNames.stream().map(AthenzDomain::new).collect(toList());
-            }
+        return client.execute(request, response -> {
+            TenantDomainsResponseEntity entity = readEntity(response, TenantDomainsResponseEntity.class);
+            return entity.tenantDomainNames.stream().map(AthenzDomain::new).collect(toList());
         });
     }
 
@@ -256,37 +232,6 @@ public class DefaultZtsClient implements ZtsClient {
         }
     }
 
-    private <T> T withClient(RequestHandler<T> handler) {
-        Lock lock = this.lock.readLock();
-        lock.lock();
-        try {
-            return handler.doRequest(this.client);
-        } catch (IOException e) {
-            throw new UncheckedIOException(e);
-        } finally {
-            lock.unlock();
-        }
-    }
-
-    private void setClient(CloseableHttpClient newClient) {
-        Lock lock = this.lock.writeLock();
-        lock.lock();
-        CloseableHttpClient oldClient;
-        try {
-            oldClient = this.client;
-            this.client = newClient;
-        } finally {
-            lock.unlock();
-        }
-        if (oldClient != null) {
-            try {
-                oldClient.close();
-            } catch (IOException e) {
-                throw new UncheckedIOException(e);
-            }
-        }
-    }
-
     private static CloseableHttpClient createHttpClient(SSLContext sslContext) {
         return HttpClientBuilder.create()
                 .setRetryHandler(new DefaultHttpRequestRetryHandler(3, /*requestSentRetryEnabled*/true))
@@ -302,21 +247,7 @@ public class DefaultZtsClient implements ZtsClient {
 
     @Override
     public void close() {
-        if (identityProvider != null && identityListener != null) {
-            identityProvider.removeIdentityListener(identityListener);
-        }
-        setClient(null);
+        this.client.close();
     }
 
-    private class ServiceIdentityProviderListener implements ServiceIdentityProvider.Listener {
-        @Override
-        public void onCredentialsUpdate(SSLContext sslContext, AthenzService identity) {
-            setClient(createHttpClient(sslContext));
-        }
-    }
-
-    @FunctionalInterface
-    private interface RequestHandler<T> {
-        T doRequest(CloseableHttpClient client) throws IOException;
-    }
 }
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identity/SiaBackedApacheHttpClient.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identity/SiaBackedApacheHttpClient.java
new file mode 100644
index 00000000000..588ed5c7abd
--- /dev/null
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identity/SiaBackedApacheHttpClient.java
@@ -0,0 +1,121 @@
+// 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 org.apache.http.client.HttpClient;
+import org.apache.http.client.ResponseHandler;
+import org.apache.http.client.methods.HttpUriRequest;
+import org.apache.http.impl.client.CloseableHttpClient;
+import org.apache.http.impl.client.HttpClientBuilder;
+
+import javax.net.ssl.SSLContext;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.util.function.Supplier;
+
+/**
+ * A thread-safe http client wrapper on Apache's {@link HttpClient} that handles TLS client authentication using {@link ServiceIdentityProvider}.
+ *
+ * @author bjorncs
+ */
+public class SiaBackedApacheHttpClient implements AutoCloseable {
+
+    private final Object clientLock = new Object();
+    private final Supplier<SSLContext> sslContextSupplier;
+    private final HttpClientFactory httpClientFactory;
+    private HttpClientHolder client;
+    private boolean closed = false;
+
+    public SiaBackedApacheHttpClient(ServiceIdentityProvider identityProvider,
+                                     HttpClientFactory httpClientFactory) {
+        this(identityProvider::getIdentitySslContext, httpClientFactory);
+    }
+
+    public SiaBackedApacheHttpClient(SSLContext sslContext,
+                                     HttpClientFactory httpClientFactory) {
+        this(() -> sslContext, httpClientFactory);
+    }
+
+    public SiaBackedApacheHttpClient(Supplier<SSLContext> sslContextSupplier,
+                                     HttpClientFactory httpClientFactory) {
+        this.sslContextSupplier = sslContextSupplier;
+        this.httpClientFactory = httpClientFactory;
+        this.client = new HttpClientHolder(httpClientFactory, sslContextSupplier);
+    }
+
+    public <T> T execute(HttpUriRequest request, ResponseHandler<T> responseHandler) {
+        HttpClientHolder client = getClient();
+        try {
+            // Note: Apache http client will handle response/connection cleanup for this overload of execute().
+            return client.apacheClient.execute(request, responseHandler);
+        } catch (IOException e) {
+            throw new UncheckedIOException(e);
+        } finally {
+            synchronized (clientLock) {
+                client.release();
+            }
+        }
+    }
+
+    private HttpClientHolder getClient() {
+        synchronized (clientLock) {
+            if (closed) throw new IllegalStateException("Client has been closed!");
+            if (sslContextSupplier.get() != client.sslContext) {
+                client.release();
+                client = new HttpClientHolder(httpClientFactory, sslContextSupplier);
+            }
+            return client.refer();
+        }
+    }
+
+    /**
+     * Thread-safe and idempotent.
+     */
+    @Override
+    public void close() {
+        synchronized (clientLock) {
+            if (closed) return;
+            client.release();
+            closed = true;
+        }
+    }
+
+    private static class HttpClientHolder {
+        final Supplier<SSLContext> sslContextSupplier;
+        final CloseableHttpClient apacheClient;
+        final SSLContext sslContext;
+        int referenceCount = 1; // Owner's reference implicitly counted
+
+        HttpClientHolder(HttpClientFactory clientBuilder, Supplier<SSLContext> sslContextSupplier) {
+            SSLContext sslContext = sslContextSupplier.get();
+            this.sslContextSupplier = sslContextSupplier;
+            this.apacheClient = clientBuilder.createHttpClient(sslContext);
+            this.sslContext = sslContext;
+        }
+
+        HttpClientHolder refer() {
+            if (referenceCount == 0) throw new IllegalStateException("Client already closed!");
+            ++referenceCount;
+            return this;
+        }
+
+        void release() {
+            if (referenceCount == 0) throw new IllegalStateException("Client already closed!");
+            --referenceCount;
+            if (referenceCount == 0) {
+                try {
+                    apacheClient.close();
+                } catch (IOException e) {
+                    throw new UncheckedIOException(e);
+                }
+            }
+        }
+    }
+
+    /**
+     * A factory that returns a new instance of {@link CloseableHttpClient}. The implementor is responsible for configuring the {@link SSLContext}, e.g. using {@link HttpClientBuilder#setSSLContext(SSLContext)}.
+     */
+    @FunctionalInterface
+    public interface HttpClientFactory {
+        CloseableHttpClient createHttpClient(SSLContext sslContext);
+    }
+}