Add builder methods accepting X509Certificate and PrivateKey

3 files changed, 72 insertions, 2 deletions

diff --git a/vespa-feed-client/src/main/java/ai/vespa/feed/client/ApacheCluster.java b/vespa-feed-client/src/main/java/ai/vespa/feed/client/ApacheCluster.java
index 7939e884551..ae764bc9a3d 100644
--- a/vespa-feed-client/src/main/java/ai/vespa/feed/client/ApacheCluster.java
+++ b/vespa-feed-client/src/main/java/ai/vespa/feed/client/ApacheCluster.java
@@ -147,9 +147,13 @@ class ApacheCluster implements Cluster {
         SslContextBuilder sslContextBuilder = new SslContextBuilder();
         if (builder.certificateFile != null && builder.privateKeyFile != null) {
             sslContextBuilder.withCertificateAndKey(builder.certificateFile, builder.privateKeyFile);
+        } else if (builder.certificate != null && builder.privateKey != null) {
+            sslContextBuilder.withCertificateAndKey(builder.certificate, builder.privateKey);
         }
         if (builder.caCertificatesFile != null) {
             sslContextBuilder.withCaCertificates(builder.caCertificatesFile);
+        } else if (builder.caCertificates != null) {
+            sslContextBuilder.withCaCertificates(builder.caCertificates);
         }
         return sslContextBuilder.build();
     }
diff --git a/vespa-feed-client/src/main/java/ai/vespa/feed/client/FeedClientBuilder.java b/vespa-feed-client/src/main/java/ai/vespa/feed/client/FeedClientBuilder.java
index bc9fa6b7e36..2c8e62d81c6 100644
--- a/vespa-feed-client/src/main/java/ai/vespa/feed/client/FeedClientBuilder.java
+++ b/vespa-feed-client/src/main/java/ai/vespa/feed/client/FeedClientBuilder.java
@@ -7,8 +7,11 @@ import java.io.IOException;
 import java.io.UncheckedIOException;
 import java.net.URI;
 import java.nio.file.Path;
+import java.security.PrivateKey;
+import java.security.cert.X509Certificate;
 import java.time.Duration;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -38,6 +41,9 @@ public class FeedClientBuilder {
     Path certificateFile;
     Path privateKeyFile;
     Path caCertificatesFile;
+    Collection<X509Certificate> certificate;
+    PrivateKey privateKey;
+    Collection<X509Certificate> caCertificates;
 
     public static FeedClientBuilder create(URI endpoint) { return new FeedClientBuilder(Collections.singletonList(endpoint)); }
 
@@ -115,11 +121,26 @@ public class FeedClientBuilder {
         return this;
     }
 
+    public FeedClientBuilder setCertificate(Collection<X509Certificate> certificate, PrivateKey privateKey) {
+        this.certificate = certificate;
+        this.privateKey = privateKey;
+        return this;
+    }
+
+    public FeedClientBuilder setCertificate(X509Certificate certificate, PrivateKey privateKey) {
+        return setCertificate(Collections.singletonList(certificate), privateKey);
+    }
+
     public FeedClientBuilder setCaCertificatesFile(Path caCertificatesFile) {
         this.caCertificatesFile = caCertificatesFile;
         return this;
     }
 
+    public FeedClientBuilder setCaCertificates(Collection<X509Certificate> caCertificates) {
+        this.caCertificates = caCertificates;
+        return this;
+    }
+
     public FeedClient build() {
         try {
             validateConfiguration();
@@ -130,9 +151,26 @@ public class FeedClientBuilder {
     }
 
     private void validateConfiguration() {
-        if (sslContext != null && (certificateFile != null || caCertificatesFile != null || privateKeyFile != null)) {
+        if (sslContext != null && (
+                certificateFile != null || caCertificatesFile != null || privateKeyFile != null ||
+                        certificate != null || caCertificates != null || privateKey != null)) {
             throw new IllegalArgumentException("Cannot set both SSLContext and certificate / CA certificates");
         }
+        if (certificate != null && certificateFile != null) {
+            throw new IllegalArgumentException("Cannot set both certificate directly and as file");
+        }
+        if (privateKey != null && privateKeyFile != null) {
+            throw new IllegalArgumentException("Cannot set both private key directly and as file");
+        }
+        if (caCertificates != null && caCertificatesFile != null) {
+            throw new IllegalArgumentException("Cannot set both CA certificates directly and as file");
+        }
+        if (certificate != null && certificate.isEmpty()) {
+            throw new IllegalArgumentException("Certificate cannot be empty");
+        }
+        if (caCertificates != null && caCertificates.isEmpty()) {
+            throw new IllegalArgumentException("CA certificates cannot be empty");
+        }
     }
 
 }
diff --git a/vespa-feed-client/src/main/java/ai/vespa/feed/client/SslContextBuilder.java b/vespa-feed-client/src/main/java/ai/vespa/feed/client/SslContextBuilder.java
index 7200d5fd943..9114e22f4a6 100644
--- a/vespa-feed-client/src/main/java/ai/vespa/feed/client/SslContextBuilder.java
+++ b/vespa-feed-client/src/main/java/ai/vespa/feed/client/SslContextBuilder.java
@@ -20,11 +20,14 @@ import java.nio.file.Path;
 import java.security.GeneralSecurityException;
 import java.security.KeyFactory;
 import java.security.KeyStore;
+import java.security.KeyStoreException;
 import java.security.PrivateKey;
 import java.security.cert.Certificate;
 import java.security.cert.X509Certificate;
 import java.security.spec.PKCS8EncodedKeySpec;
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
 import java.util.List;
 
 /**
@@ -39,6 +42,9 @@ class SslContextBuilder {
     private Path certificateFile;
     private Path privateKeyFile;
     private Path caCertificatesFile;
+    private Collection<X509Certificate> certificate;
+    private PrivateKey privateKey;
+    private Collection<X509Certificate> caCertificates;
 
     SslContextBuilder withCertificateAndKey(Path certificate, Path privateKey) {
         this.certificateFile = certificate;
@@ -46,20 +52,35 @@ class SslContextBuilder {
         return this;
     }
 
+    SslContextBuilder withCertificateAndKey(Collection<X509Certificate> certificate, PrivateKey privateKey) {
+        this.certificate = certificate;
+        this.privateKey = privateKey;
+        return this;
+    }
+
     SslContextBuilder withCaCertificates(Path caCertificates) {
         this.caCertificatesFile = caCertificates;
         return this;
     }
 
+    SslContextBuilder withCaCertificates(Collection<X509Certificate> caCertificates) {
+        this.caCertificates = caCertificates;
+        return this;
+    }
+
     SSLContext build() throws IOException {
         try {
             KeyStore keystore = KeyStore.getInstance("PKCS12");
             keystore.load(null);
             if (certificateFile != null && privateKeyFile != null) {
                 keystore.setKeyEntry("cert", privateKey(privateKeyFile), new char[0], certificates(certificateFile));
+            } else if (certificate != null && privateKey != null) {
+                keystore.setKeyEntry("cert", privateKey, new char[0], certificate.toArray(new Certificate[0]));
             }
             if (caCertificatesFile != null) {
-                keystore.setCertificateEntry("ca-cert", certificates(caCertificatesFile)[0]);
+                addCaCertificates(keystore, Arrays.asList(certificates(caCertificatesFile)));
+            } else if (caCertificates != null) {
+                addCaCertificates(keystore, caCertificates);
             }
             KeyManagerFactory kmf = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
             kmf.init(keystore, new char[0]);
@@ -73,6 +94,13 @@ class SslContextBuilder {
         }
     }
 
+    private static void addCaCertificates(KeyStore keystore, Collection<? extends Certificate> certificates) throws KeyStoreException {
+        int i = 0;
+        for (Certificate cert : certificates) {
+            keystore.setCertificateEntry("ca-cert-" + ++i, cert);
+        }
+    }
+
     private static Certificate[] certificates(Path file) throws IOException, GeneralSecurityException {
         try (PEMParser parser = new PEMParser(Files.newBufferedReader(file))) {
             List<X509Certificate> result = new ArrayList<>();