Fix multiple issues in TLS session caching implementation

This commit addresses several issues identified in the TLS session caching feature:

1. SNI endpoint cache key collision: Added tls_session_cache_key property
   to EndPoint classes. SniEndPoint now includes server_name in the cache
   key to prevent collisions when multiple SNI endpoints use the same proxy.

2. TLS session caching for eventlet/twisted reactors: Added PyOpenSSL-based
   session caching support in EventletConnection and TwistedConnection
   (via _SSLCreator) to enable TLS session resumption for these reactors.

3. Missing hasattr check: Added hasattr check for session_reused attribute
   before accessing it, as not all SSL socket implementations have it.

4. Automatic cleanup of expired sessions: Added opportunistic cleanup in
   set_session() that runs every 100 operations to prevent memory
   accumulation from expired sessions.

5. Log message formatting: Changed log messages to use endpoint's __str__
   method instead of manually formatting address:port, which handles
   UnixSocketEndPoint (port=None) correctly.

6. Custom TLSSessionCache support: Updated Cluster to accept either
   TLSSessionCacheOptions or a TLSSessionCache instance directly,
   allowing users to provide custom cache implementations.

Author: dkropachev

cassandra/cluster.py

@@ -908,16 +908,19 @@ def default_retry_policy(self, policy):
 
     tls_session_cache_options = None
     """
-    Advanced TLS session cache configuration. Set to an instance of
-    :class:`~cassandra.tls.TLSSessionCacheOptions` for fine-grained control over
-    session caching behavior (e.g., cache_by_host_only option).
+    Advanced TLS session cache configuration. Can be set to:
+
+    - An instance of :class:`~cassandra.tls.TLSSessionCacheOptions` for
+      fine-grained control over session caching behavior (e.g., cache_by_host_only option).
+    - An instance of :class:`~cassandra.tls.TLSSessionCache` (or a custom subclass)
+      for complete control over session caching implementation.
 
     If None (default), a cache is created using :attr:`~.tls_session_cache_size`
     and :attr:`~.tls_session_cache_ttl` when SSL/TLS is enabled.
 
     This option takes precedence over the individual tls_session_cache_* parameters.
 
-    Example::
+    Example with options::
 
         from cassandra.tls import TLSSessionCacheOptions
 
@@ -929,6 +932,16 @@ def default_retry_policy(self, policy):
         )
         cluster = Cluster(ssl_context=ssl_context, tls_session_cache_options=options)
 
+    Example with custom cache::
+
+        from cassandra.tls import TLSSessionCache
+
+        class MyCustomCache(TLSSessionCache):
+            # Custom implementation
+            pass
+
+        cluster = Cluster(ssl_context=ssl_context, tls_session_cache_options=MyCustomCache())
+
     .. versionadded:: 3.30.0
     """
 
@@ -1489,19 +1502,24 @@ def __init__(self,
         # Initialize TLS session cache if SSL is enabled and caching is enabled
         self._tls_session_cache = None
         if (ssl_context or ssl_options) and tls_session_cache_enabled:
-            from cassandra.tls import TLSSessionCacheOptions
+            from cassandra.tls import TLSSessionCache, TLSSessionCacheOptions
 
-            # Use provided options object or create one from individual parameters
             if tls_session_cache_options is not None:
-                cache_options = tls_session_cache_options
+                # Check if it's a TLSSessionCache instance (use directly)
+                # or TLSSessionCacheOptions (use create_cache())
+                if isinstance(tls_session_cache_options, TLSSessionCache):
+                    self._tls_session_cache = tls_session_cache_options
+                else:
+                    # Assume it's TLSSessionCacheOptions
+                    self._tls_session_cache = tls_session_cache_options.create_cache()
             else:
+                # Create default cache from individual parameters
                 cache_options = TLSSessionCacheOptions(
                     max_size=tls_session_cache_size,
                     ttl=tls_session_cache_ttl,
                     cache_by_host_only=False
                 )
-
-            self._tls_session_cache = cache_options.create_cache()
+                self._tls_session_cache = cache_options.create_cache()
 
         self.sockopts = sockopts
         self.cql_version = cql_version

cassandra/connection.py

@@ -161,6 +161,15 @@ def socket_family(self):
         """
         return socket.AF_UNSPEC
 
+    @property
+    def tls_session_cache_key(self):
+        """
+        Returns the cache key components for TLS session caching.
+        This is a tuple that uniquely identifies this endpoint for TLS session purposes.
+        Subclasses may override this to include additional components (e.g., SNI server name).
+        """
+        return (self.address, self.port)
+
     def resolve(self):
         """
         Resolve the endpoint to an address/port. This is called
@@ -275,6 +284,14 @@ def port(self):
     def ssl_options(self):
         return self._ssl_options
 
+    @property
+    def tls_session_cache_key(self):
+        """
+        Returns the cache key including server_name for SNI endpoints.
+        This prevents cache collisions when multiple SNI endpoints use the same proxy.
+        """
+        return (self.address, self.port, self._server_name)
+
     def resolve(self):
         try:
             resolved_addresses = socket.getaddrinfo(self._proxy_address, self._port,
@@ -349,6 +366,14 @@ def port(self):
     def socket_family(self):
         return socket.AF_UNIX
 
+    @property
+    def tls_session_cache_key(self):
+        """
+        Returns the cache key for Unix socket endpoints.
+        Since Unix sockets don't have a port, only the path is used.
+        """
+        return (self._unix_socket_path,)
+
     def resolve(self):
         return self.address, None
 
@@ -922,8 +947,7 @@ def _wrap_socket_from_context(self):
             cached_session = self.tls_session_cache.get_session(self.endpoint)
             if cached_session:
                 opts['session'] = cached_session
-                log.debug("Using cached TLS session for %s:%s", 
-                         self.endpoint.address, self.endpoint.port)
+                log.debug("Using cached TLS session for %s", self.endpoint)
 
         ssl_socket = self.ssl_context.wrap_socket(self._socket, **opts)
 
@@ -991,9 +1015,8 @@ def _connect_socket(self):
                 if self.tls_session_cache and self.ssl_context and hasattr(self._socket, 'session'):
                     if self._socket.session:
                         self.tls_session_cache.set_session(self.endpoint, self._socket.session)
-                        if self._socket.session_reused:
-                            log.debug("TLS session was reused for %s:%s", 
-                                     self.endpoint.address, self.endpoint.port)
+                        if hasattr(self._socket, 'session_reused') and self._socket.session_reused:
+                            log.debug("TLS session was reused for %s", self.endpoint)
 
                 sockerr = None
                 break

cassandra/io/eventletreactor.py

@@ -109,13 +109,27 @@ def _wrap_socket_from_context(self):
             # This is necessary for SNI
             self._socket.set_tlsext_host_name(self.ssl_options['server_hostname'].encode('ascii'))
 
+        # Apply cached TLS session for resumption (PyOpenSSL)
+        if self.tls_session_cache:
+            cached_session = self.tls_session_cache.get_session(self.endpoint)
+            if cached_session:
+                self._socket.set_session(cached_session)
+                log.debug("Using cached TLS session for %s", self.endpoint)
+
     def _initiate_connection(self, sockaddr):
         if self.uses_legacy_ssl_options:
             super(EventletConnection, self)._initiate_connection(sockaddr)
         else:
             self._socket.connect(sockaddr)
             if self.ssl_context or self.ssl_options:
                 self._socket.do_handshake()
+                # Store TLS session after successful handshake (PyOpenSSL)
+                if self.tls_session_cache:
+                    session = self._socket.get_session()
+                    if session:
+                        self.tls_session_cache.set_session(self.endpoint, session)
+                        if self._socket.session_reused():
+                            log.debug("TLS session was reused for %s", self.endpoint)
 
     def _match_hostname(self):
         if self.uses_legacy_ssl_options:

cassandra/io/twistedreactor.py

@@ -139,11 +139,12 @@ def _on_loop_timer(self):
 
 @implementer(IOpenSSLClientConnectionCreator)
 class _SSLCreator(object):
-    def __init__(self, endpoint, ssl_context, ssl_options, check_hostname, timeout):
+    def __init__(self, endpoint, ssl_context, ssl_options, check_hostname, timeout, tls_session_cache=None):
         self.endpoint = endpoint
         self.ssl_options = ssl_options
         self.check_hostname = check_hostname
         self.timeout = timeout
+        self.tls_session_cache = tls_session_cache
 
         if ssl_context:
             self.context = ssl_context
@@ -171,11 +172,27 @@ def info_callback(self, connection, where, ret):
                 transport = connection.get_app_data()
                 transport.failVerification(Failure(ConnectionException("Hostname verification failed", self.endpoint)))
 
+            # Store TLS session after successful handshake (PyOpenSSL)
+            if self.tls_session_cache:
+                session = connection.get_session()
+                if session:
+                    self.tls_session_cache.set_session(self.endpoint, session)
+                    if connection.session_reused():
+                        log.debug("TLS session was reused for %s", self.endpoint)
+
     def clientConnectionForTLS(self, tlsProtocol):
         connection = SSL.Connection(self.context, None)
         connection.set_app_data(tlsProtocol)
         if self.ssl_options and "server_hostname" in self.ssl_options:
             connection.set_tlsext_host_name(self.ssl_options['server_hostname'].encode('ascii'))
+
+        # Apply cached TLS session for resumption (PyOpenSSL)
+        if self.tls_session_cache:
+            cached_session = self.tls_session_cache.get_session(self.endpoint)
+            if cached_session:
+                connection.set_session(cached_session)
+                log.debug("Using cached TLS session for %s", self.endpoint)
+
         return connection
 
 
@@ -241,6 +258,7 @@ def add_connection(self):
                 self.ssl_options,
                 self._check_hostname,
                 self.connect_timeout,
+                tls_session_cache=self.tls_session_cache,
             )
 
             endpoint = SSL4ClientEndpoint(

cassandra/tls.py

@@ -91,10 +91,13 @@ class DefaultTLSSessionCache(TLSSessionCache):
     version-specific checks are needed.
     """
 
+    # Cleanup expired sessions every N set_session calls
+    _EXPIRY_CLEANUP_INTERVAL = 100
+
     def __init__(self, max_size=100, ttl=3600, cache_by_host_only=False):
         """
         Initialize the TLS session cache.
-        
+
         Args:
             max_size: Maximum number of sessions to cache (default: 100)
             ttl: Time-to-live for cached sessions in seconds (default: 3600)
@@ -106,13 +109,22 @@ def __init__(self, max_size=100, ttl=3600, cache_by_host_only=False):
         self._max_size = max_size
         self._ttl = ttl
         self._cache_by_host_only = cache_by_host_only
+        self._operation_count = 0  # Counter for opportunistic cleanup
 
     def _make_key(self, endpoint):
-        """Create a cache key from endpoint."""
+        """
+        Create a cache key from endpoint.
+
+        Uses the endpoint's tls_session_cache_key property which returns
+        appropriate components for each endpoint type (e.g., includes
+        server_name for SNI endpoints to prevent cache collisions).
+        """
+        key = endpoint.tls_session_cache_key
         if self._cache_by_host_only:
-            return (endpoint.address,)
+            # When caching by host only, use just the first component (address/path)
+            return (key[0],)
         else:
-            return (endpoint.address, endpoint.port)
+            return key
 
     def get_session(self, endpoint):
         """
@@ -143,41 +155,52 @@ def get_session(self, endpoint):
     def set_session(self, endpoint, session):
         """
         Store a TLS session for the given endpoint.
-        
+
         Args:
             endpoint: The EndPoint object representing the connection target
             session: The ssl.SSLSession object to cache
         """
         if session is None:
             return
-        
+
         key = self._make_key(endpoint)
         current_time = time.time()
-        
+
         with self._lock:
+            # Opportunistically clean up expired sessions periodically
+            self._operation_count += 1
+            if self._operation_count >= self._EXPIRY_CLEANUP_INTERVAL:
+                self._operation_count = 0
+                self._clear_expired_unlocked(current_time)
+
             # If key already exists, just update it
             if key in self._sessions:
                 self._sessions[key] = _SessionCacheEntry(session, current_time)
                 self._sessions.move_to_end(key)
                 return
-            
+
             # If cache is at max size, remove least recently used entry (first item)
             if len(self._sessions) >= self._max_size:
                 self._sessions.popitem(last=False)
-            
+
             # Store session with creation time
             self._sessions[key] = _SessionCacheEntry(session, current_time)
 
+    def _clear_expired_unlocked(self, current_time=None):
+        """Remove all expired sessions (must be called with lock held)."""
+        if current_time is None:
+            current_time = time.time()
+        expired_keys = [
+            key for key, entry in self._sessions.items()
+            if current_time - entry.timestamp > self._ttl
+        ]
+        for key in expired_keys:
+            del self._sessions[key]
+
     def clear_expired(self):
         """Remove all expired sessions from the cache."""
-        current_time = time.time()
         with self._lock:
-            expired_keys = [
-                key for key, entry in self._sessions.items()
-                if current_time - entry.timestamp > self._ttl
-            ]
-            for key in expired_keys:
-                del self._sessions[key]
+            self._clear_expired_unlocked()
 
     def clear(self):
         """Clear all sessions from the cache."""