Change 'TlsContext' interface to return `X509SslContext'

4 files changed, 62 insertions, 52 deletions

diff --git a/security-utils/src/main/java/com/yahoo/security/tls/ConfigFileBasedTlsContext.java b/security-utils/src/main/java/com/yahoo/security/tls/ConfigFileBasedTlsContext.java
index ef1762ea7cd..176c6f95749 100644
--- a/security-utils/src/main/java/com/yahoo/security/tls/ConfigFileBasedTlsContext.java
+++ b/security-utils/src/main/java/com/yahoo/security/tls/ConfigFileBasedTlsContext.java
@@ -8,9 +8,8 @@ import com.yahoo.security.MutableX509KeyManager;
 import com.yahoo.security.MutableX509TrustManager;
 import com.yahoo.security.SslContextBuilder;
 import com.yahoo.security.X509CertificateUtils;
+import com.yahoo.security.X509SslContext;
 
-import javax.net.ssl.SSLContext;
-import javax.net.ssl.SSLEngine;
 import javax.net.ssl.SSLParameters;
 import java.io.IOException;
 import java.io.UncheckedIOException;
@@ -113,22 +112,20 @@ public class ConfigFileBasedTlsContext implements TlsContext {
         HostnameVerification hostnameVerification = options.isHostnameValidationDisabled() ? HostnameVerification.DISABLED : HostnameVerification.ENABLED;
         PeerAuthorizerTrustManager authorizerTrustManager =
                 new PeerAuthorizerTrustManager(options.getAuthorizedPeers(), mode, hostnameVerification, mutableTrustManager);
-        SSLContext sslContext = new SslContextBuilder()
+        var sslContext = new SslContextBuilder()
                 .withKeyManager(mutableKeyManager)
                 .withTrustManager(authorizerTrustManager)
-                .build();
+                .buildContext();
         List<String> acceptedCiphers = options.getAcceptedCiphers();
         Set<String> ciphers = acceptedCiphers.isEmpty() ? TlsContext.ALLOWED_CIPHER_SUITES : new HashSet<>(acceptedCiphers);
         List<String> acceptedProtocols = options.getAcceptedProtocols();
         Set<String> protocols = acceptedProtocols.isEmpty() ? TlsContext.ALLOWED_PROTOCOLS : new HashSet<>(acceptedProtocols);
-        return new DefaultTlsContext(sslContext, ciphers, protocols, peerAuthentication);
+        return DefaultTlsContext.of(sslContext, ciphers, protocols, peerAuthentication);
     }
 
     // Wrapped methods from TlsContext
-    @Override public SSLContext context() { return tlsContext.context(); }
+    @Override public X509SslContext sslContext() { return tlsContext.sslContext(); }
     @Override public SSLParameters parameters() { return tlsContext.parameters(); }
-    @Override public SSLEngine createSslEngine() { return tlsContext.createSslEngine(); }
-    @Override public SSLEngine createSslEngine(String peerHost, int peerPort) { return tlsContext.createSslEngine(peerHost, peerPort); }
 
     @Override
     public void close() {
diff --git a/security-utils/src/main/java/com/yahoo/security/tls/DefaultTlsContext.java b/security-utils/src/main/java/com/yahoo/security/tls/DefaultTlsContext.java
index 8f4838c9940..4e810c2d304 100644
--- a/security-utils/src/main/java/com/yahoo/security/tls/DefaultTlsContext.java
+++ b/security-utils/src/main/java/com/yahoo/security/tls/DefaultTlsContext.java
@@ -2,9 +2,9 @@
 package com.yahoo.security.tls;
 
 import com.yahoo.security.SslContextBuilder;
+import com.yahoo.security.X509SslContext;
 
 import javax.net.ssl.SSLContext;
-import javax.net.ssl.SSLEngine;
 import javax.net.ssl.SSLParameters;
 import java.security.PrivateKey;
 import java.security.cert.X509Certificate;
@@ -23,30 +23,35 @@ public class DefaultTlsContext implements TlsContext {
 
     private static final Logger log = Logger.getLogger(DefaultTlsContext.class.getName());
 
-    private final SSLContext sslContext;
+    private final X509SslContext sslContext;
     private final String[] validCiphers;
     private final String[] validProtocols;
     private final PeerAuthentication peerAuthentication;
 
-    public DefaultTlsContext(List<X509Certificate> certificates,
-                             PrivateKey privateKey,
-                             List<X509Certificate> caCertificates,
-                             AuthorizedPeers authorizedPeers,
-                             AuthorizationMode mode,
-                             PeerAuthentication peerAuthentication,
-                             HostnameVerification hostnameVerification) {
-        this(createSslContext(certificates, privateKey, caCertificates, authorizedPeers, mode, hostnameVerification), peerAuthentication);
+    public static DefaultTlsContext of(X509SslContext sslContext, PeerAuthentication peerAuthentication) {
+        return new DefaultTlsContext(sslContext, TlsContext.ALLOWED_CIPHER_SUITES, TlsContext.ALLOWED_PROTOCOLS, peerAuthentication);
     }
 
-    public DefaultTlsContext(SSLContext sslContext, PeerAuthentication peerAuthentication) {
-        this(sslContext, TlsContext.ALLOWED_CIPHER_SUITES, TlsContext.ALLOWED_PROTOCOLS, peerAuthentication);
+    public static DefaultTlsContext of(
+            List<X509Certificate> certificates, PrivateKey privateKey, List<X509Certificate> caCertificates,
+            AuthorizedPeers authorizedPeers, AuthorizationMode mode, PeerAuthentication peerAuthentication,
+            HostnameVerification hostnameVerification) {
+        var ctx = createSslContext(certificates, privateKey, caCertificates, authorizedPeers, mode, hostnameVerification);
+        return of(ctx, peerAuthentication);
     }
 
-    DefaultTlsContext(SSLContext sslContext, Set<String> acceptedCiphers, Set<String> acceptedProtocols, PeerAuthentication peerAuthentication) {
+    public static DefaultTlsContext of(
+            X509SslContext sslContext, Set<String> acceptedCiphers, Set<String> acceptedProtocols,
+            PeerAuthentication peerAuthentication) {
+        return new DefaultTlsContext(sslContext, acceptedCiphers, acceptedProtocols, peerAuthentication);
+    }
+
+    private DefaultTlsContext(X509SslContext sslContext, Set<String> acceptedCiphers, Set<String> acceptedProtocols,
+                              PeerAuthentication peerAuthentication) {
         this.sslContext = sslContext;
         this.peerAuthentication = peerAuthentication;
-        this.validCiphers = getAllowedCiphers(sslContext, acceptedCiphers);
-        this.validProtocols = getAllowedProtocols(sslContext, acceptedProtocols);
+        this.validCiphers = getAllowedCiphers(sslContext.context(), acceptedCiphers);
+        this.validProtocols = getAllowedProtocols(sslContext.context(), acceptedProtocols);
     }
 
     private static String[] getAllowedCiphers(SSLContext sslContext, Set<String> acceptedCiphers) {
@@ -78,7 +83,7 @@ public class DefaultTlsContext implements TlsContext {
     }
 
     @Override
-    public SSLContext context() {
+    public X509SslContext sslContext() {
         return sslContext;
     }
 
@@ -87,22 +92,8 @@ public class DefaultTlsContext implements TlsContext {
         return createSslParameters();
     }
 
-    @Override
-    public SSLEngine createSslEngine() {
-        SSLEngine sslEngine = sslContext.createSSLEngine();
-        sslEngine.setSSLParameters(createSslParameters());
-        return sslEngine;
-    }
-
-    @Override
-    public SSLEngine createSslEngine(String peerHost, int peerPort) {
-        SSLEngine sslEngine = sslContext.createSSLEngine(peerHost, peerPort);
-        sslEngine.setSSLParameters(createSslParameters());
-        return sslEngine;
-    }
-
     private SSLParameters createSslParameters() {
-        SSLParameters newParameters = sslContext.getDefaultSSLParameters();
+        SSLParameters newParameters = sslContext.context().getDefaultSSLParameters();
         newParameters.setCipherSuites(validCiphers);
         newParameters.setProtocols(validProtocols);
         switch (peerAuthentication) {
@@ -120,12 +111,9 @@ public class DefaultTlsContext implements TlsContext {
         return newParameters;
     }
 
-    private static SSLContext createSslContext(List<X509Certificate> certificates,
-                                               PrivateKey privateKey,
-                                               List<X509Certificate> caCertificates,
-                                               AuthorizedPeers authorizedPeers,
-                                               AuthorizationMode mode,
-                                               HostnameVerification hostnameVerification) {
+    private static X509SslContext createSslContext(
+            List<X509Certificate> certificates, PrivateKey privateKey, List<X509Certificate> caCertificates,
+            AuthorizedPeers authorizedPeers, AuthorizationMode mode, HostnameVerification hostnameVerification) {
         SslContextBuilder builder = new SslContextBuilder();
         if (!certificates.isEmpty()) {
             builder.withKeyStore(privateKey, certificates);
@@ -135,7 +123,7 @@ public class DefaultTlsContext implements TlsContext {
         }
         return builder.withTrustManagerFactory(truststore ->
                         new PeerAuthorizerTrustManager(authorizedPeers, mode, hostnameVerification, truststore))
-                .build();
+                .buildContext();
     }
 
 }
diff --git a/security-utils/src/main/java/com/yahoo/security/tls/TlsContext.java b/security-utils/src/main/java/com/yahoo/security/tls/TlsContext.java
index fff942ba6ab..6a530718363 100644
--- a/security-utils/src/main/java/com/yahoo/security/tls/TlsContext.java
+++ b/security-utils/src/main/java/com/yahoo/security/tls/TlsContext.java
@@ -1,9 +1,14 @@
 // Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.security.tls;
 
+import com.yahoo.security.X509SslContext;
+
 import javax.net.ssl.SSLContext;
 import javax.net.ssl.SSLEngine;
 import javax.net.ssl.SSLParameters;
+import javax.net.ssl.SSLServerSocket;
+import javax.net.ssl.SSLSocket;
+import java.io.IOException;
 import java.security.KeyManagementException;
 import java.security.NoSuchAlgorithmException;
 import java.util.Arrays;
@@ -92,13 +97,32 @@ public interface TlsContext extends AutoCloseable {
         } catch (KeyManagementException e) { throw new IllegalStateException(e); }
     }
 
-    SSLContext context();
-
+    X509SslContext sslContext();
     SSLParameters parameters();
 
-    SSLEngine createSslEngine();
+    default SSLEngine createSslEngine() {
+        SSLEngine sslEngine = sslContext().context().createSSLEngine();
+        sslEngine.setSSLParameters(parameters());
+        return sslEngine;
+    }
+
+    default SSLEngine createSslEngine(String peerHost, int peerPort) {
+        SSLEngine sslEngine = sslContext().context().createSSLEngine(peerHost, peerPort);
+        sslEngine.setSSLParameters(parameters());
+        return sslEngine;
+    }
+
+    default SSLSocket createClientSslSocket() throws IOException {
+        var socket = (SSLSocket) sslContext().context().getSocketFactory().createSocket();
+        socket.setSSLParameters(parameters());
+        return socket;
+    }
 
-    SSLEngine createSslEngine(String peerHost, int peerPort);
+    default SSLServerSocket createServerSslSocket() throws IOException {
+        var socket = (SSLServerSocket) sslContext().context().getServerSocketFactory().createServerSocket();
+        socket.setSSLParameters(parameters());
+        return socket;
+    }
 
     @Override default void close() {}
 
diff --git a/security-utils/src/test/java/com/yahoo/security/tls/DefaultTlsContextTest.java b/security-utils/src/test/java/com/yahoo/security/tls/DefaultTlsContextTest.java
index a8012f52e5c..ec7d5b8ca05 100644
--- a/security-utils/src/test/java/com/yahoo/security/tls/DefaultTlsContextTest.java
+++ b/security-utils/src/test/java/com/yahoo/security/tls/DefaultTlsContextTest.java
@@ -10,6 +10,7 @@ import javax.security.auth.x500.X500Principal;
 import java.security.KeyPair;
 import java.security.cert.X509Certificate;
 import java.time.Instant;
+import java.util.List;
 
 import static com.yahoo.security.KeyAlgorithm.EC;
 import static com.yahoo.security.SignatureAlgorithm.SHA256_WITH_ECDSA;
@@ -40,8 +41,8 @@ public class DefaultTlsContextTest {
                                 singletonList(RequiredPeerCredential.of(RequiredPeerCredential.Field.CN, "dummy")))));
 
         DefaultTlsContext tlsContext =
-                new DefaultTlsContext(
-                        singletonList(certificate), keyPair.getPrivate(), singletonList(certificate), authorizedPeers,
+                DefaultTlsContext.of(
+                        List.of(certificate), keyPair.getPrivate(), List.of(certificate), authorizedPeers,
                         AuthorizationMode.ENFORCE, PeerAuthentication.NEED, HostnameVerification.ENABLED);
 
         SSLEngine sslEngine = tlsContext.createSslEngine();