Ensure SSLEngine instances are purged on connection closed

2 files changed, 19 insertions, 12 deletions

diff --git a/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/JettyConnectionLogger.java b/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/JettyConnectionLogger.java
index d5cb9ee1e49..d337131b313 100644
--- a/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/JettyConnectionLogger.java
+++ b/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/JettyConnectionLogger.java
@@ -49,6 +49,8 @@ class JettyConnectionLogger extends AbstractLifeCycle implements Connection.List
     private static final Logger log = Logger.getLogger(JettyConnectionLogger.class.getName());
 
     private final SimpleConcurrentIdentityHashMap<SocketChannelEndPoint, ConnectionInfo> connectionInfos = new SimpleConcurrentIdentityHashMap<>();
+    private final SimpleConcurrentIdentityHashMap<SocketChannelEndPoint, SSLEngine> sslEngines = new SimpleConcurrentIdentityHashMap<>();
+    // Extra mapping as callbacks in SslHandshakeListener only provides SSLEngine (no connection reference) as argument
     private final SimpleConcurrentIdentityHashMap<SSLEngine, ConnectionInfo> sslToConnectionInfo = new SimpleConcurrentIdentityHashMap<>();
 
     private final boolean enabled;
@@ -91,7 +93,10 @@ class JettyConnectionLogger extends AbstractLifeCycle implements Connection.List
             String connectionClassName = connection.getClass().getSimpleName(); // For hidden implementations of Connection
             if (connection instanceof SslConnection) {
                 SSLEngine sslEngine = ((SslConnection) connection).getSSLEngine();
-                sslToConnectionInfo.put(sslEngine, info);
+                addReferenceToSslEngine(endpoint, info, sslEngine);
+            } else if (connection instanceof ALPNServerConnection) {
+                SSLEngine sslEngine = ((ALPNServerConnection) connection).getSSLEngine();
+                addReferenceToSslEngine(endpoint, info, sslEngine);
             } else if (connection instanceof HttpConnection) {
                 info.setHttpProtocol("HTTP/1.1");
             } else if (connection instanceof HTTP2ServerConnection) {
@@ -108,6 +113,14 @@ class JettyConnectionLogger extends AbstractLifeCycle implements Connection.List
         });
     }
 
+    private void addReferenceToSslEngine(SocketChannelEndPoint endpoint, ConnectionInfo info, SSLEngine sslEngine) {
+        if (sslEngine != null) {
+            sslEngines.put(endpoint, sslEngine)
+                    .ifPresent(sslToConnectionInfo::remove);
+            sslToConnectionInfo.put(sslEngine, info);
+        }
+    }
+
     @Override
     public void onClosed(Connection connection) {
         handleListenerInvocation("Connection.Listener", "onClosed", "%h", List.of(connection), () -> {
@@ -116,19 +129,13 @@ class JettyConnectionLogger extends AbstractLifeCycle implements Connection.List
             if (info == null) return; // Closed connection already handled
             if (connection instanceof HttpConnection) {
                 info.setHttpBytes(connection.getBytesIn(), connection.getBytesOut());
-            } else if (connection instanceof SslConnection) {
-                SSLEngine sslEngine = ((SslConnection) connection).getSSLEngine();
-                sslToConnectionInfo.remove(sslEngine);
-            } else if (connection instanceof ALPNServerConnection) {
-                SSLEngine sslEngine = ((ALPNServerConnection) connection).getSSLEngine();
-                if (sslEngine != null) {
-                    sslToConnectionInfo.remove(sslEngine);
-                }
             }
             if (!endpoint.isOpen()) {
                 info.setClosedAt(System.currentTimeMillis());
                 connectionLog.log(info.toLogEntry());
                 connectionInfos.remove(endpoint);
+                sslEngines.remove(endpoint)
+                        .ifPresent(sslToConnectionInfo::remove);
             }
         });
     }
@@ -169,7 +176,7 @@ class JettyConnectionLogger extends AbstractLifeCycle implements Connection.List
     public void handshakeSucceeded(Event event) {
         SSLEngine sslEngine = event.getSSLEngine();
         handleListenerInvocation("SslHandshakeListener", "handshakeSucceeded", "sslEngine=%h", List.of(sslEngine), () -> {
-            ConnectionInfo info = sslToConnectionInfo.remove(sslEngine).orElse(null);
+            ConnectionInfo info = sslToConnectionInfo.get(sslEngine).orElse(null);
             if (info == null) return;
             info.setSslSessionDetails(sslEngine.getSession());
         });
@@ -180,7 +187,7 @@ class JettyConnectionLogger extends AbstractLifeCycle implements Connection.List
         SSLEngine sslEngine = event.getSSLEngine();
         handleListenerInvocation("SslHandshakeListener", "handshakeFailed", "sslEngine=%h,failure=%s", List.of(sslEngine, failure), () -> {
             log.log(Level.FINE, failure, failure::toString);
-            ConnectionInfo info = sslToConnectionInfo.remove(sslEngine).orElse(null);
+            ConnectionInfo info = sslToConnectionInfo.get(sslEngine).orElse(null);
             if (info == null) return;
             info.setSslHandshakeFailure((SSLHandshakeException)failure);
         });
diff --git a/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/SimpleConcurrentIdentityHashMap.java b/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/SimpleConcurrentIdentityHashMap.java
index b2bfa2a5dda..59d606c640f 100644
--- a/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/SimpleConcurrentIdentityHashMap.java
+++ b/container-core/src/main/java/com/yahoo/jdisc/http/server/jetty/SimpleConcurrentIdentityHashMap.java
@@ -20,7 +20,7 @@ class SimpleConcurrentIdentityHashMap<K, V> {
 
     Optional<V> remove(K key) { return Optional.ofNullable(wrappedMap.remove(identityKey(key))); }
 
-    void put(K key, V value) { wrappedMap.put(identityKey(key), value); }
+    Optional<V> put(K key, V value) { return Optional.ofNullable(wrappedMap.put(identityKey(key), value)); }
 
     V computeIfAbsent(K key, Supplier<V> supplier) {
         return wrappedMap.computeIfAbsent(identityKey(key), ignored -> supplier.get());