pool: honor explicit SSL config for shard-aware ports

Problem: shard-aware endpoint selection treated legacy ssl_options as SSL-enabled only when the dict was truthy. An explicit empty ssl_options={} was therefore handled like plaintext and could select the non-SSL shard-aware port, diverging from the cluster-level SSL-enabled check.

Fix: treat SSL as enabled when ssl_context is set or ssl_options is not None. SSL-enabled configurations now use the SSL shard-aware port when advertised and otherwise fall back to regular non-shard-aware connections instead of the plaintext shard-aware port. Unit coverage now includes ssl_context, non-empty ssl_options, and empty ssl_options.

Author: dkropachev

cassandra/pool.py

@@ -675,15 +675,26 @@ def disable_advanced_shard_aware(self, secs):
         self.advanced_shardaware_block_until = max(time.time() + secs, self.advanced_shardaware_block_until)
 
     def _get_shard_aware_endpoint(self):
+        """
+        Return an endpoint for the advertised shard-aware port, if usable.
+
+        Plaintext clusters use shard_aware_port. SSL-enabled clusters use only
+        shard_aware_port_ssl; if it is absent, return None so the pool opens a
+        regular SSL connection instead of falling back to the plaintext port.
+        Explicit ssl_options={}, like ssl_context, marks the cluster SSL-enabled.
+        """
         if (self.advanced_shardaware_block_until and self.advanced_shardaware_block_until > time.time()) or \
            self._session.cluster.shard_aware_options.disable_shardaware_port:
             return None
 
+        cluster = self._session.cluster
+        ssl_enabled = cluster.ssl_context is not None or cluster.ssl_options is not None
+
         endpoint = None
-        if self._session.cluster.ssl_options and self.host.sharding_info.shard_aware_port_ssl:
+        if ssl_enabled and self.host.sharding_info.shard_aware_port_ssl:
             endpoint = copy.copy(self.host.endpoint)
             endpoint._port = self.host.sharding_info.shard_aware_port_ssl
-        elif self.host.sharding_info.shard_aware_port:
+        elif not ssl_enabled and self.host.sharding_info.shard_aware_port:
             endpoint = copy.copy(self.host.endpoint)
             endpoint._port = self.host.sharding_info.shard_aware_port
 
@@ -918,5 +929,3 @@ def open_count(self):
     @property
     def _excess_connection_limit(self):
         return self.host.sharding_info.shards_count * self.max_excess_connections_per_shard_multiplier
-
-

tests/unit/test_shard_aware.py

@@ -32,19 +32,19 @@ class MockSession(MagicMock):
     is_shutdown = False
     keyspace = "ks1"
 
-    def __init__(self, is_ssl=False, *args, **kwargs):
+    def __init__(self, ssl_options=None, ssl_context=None, sharding_info=None,
+                 *args, **kwargs):
         super(MockSession, self).__init__(*args, **kwargs)
         self.cluster = MagicMock()
-        if is_ssl:
-            self.cluster.ssl_options = {'some_ssl_options': True}
-        else:
-            self.cluster.ssl_options = None
+        self.cluster.ssl_options = ssl_options
+        self.cluster.ssl_context = ssl_context
         self.cluster.shard_aware_options = ShardAwareOptions()
         self.cluster.executor = ThreadPoolExecutor(max_workers=2)
         self.cluster.signal_connection_failure = lambda *args, **kwargs: False
         self.cluster.connection_factory = self.mock_connection_factory
         self.connection_counter = 0
         self.futures = []
+        self.sharding_info = sharding_info
 
     def submit(self, fn, *args, **kwargs):
         logging.info("Scheduling %s with args: %s, kwargs: %s", fn, args, kwargs)
@@ -60,8 +60,13 @@ def mock_connection_factory(self, *args, **kwargs):
         connection.is_closed = False
         connection.orphaned_threshold_reached = False
         connection.endpoint = args[0]
-        sharding_info = ShardingInfo(shard_id=1, shards_count=4, partitioner="", sharding_algorithm="", sharding_ignore_msb=0, shard_aware_port=19042, shard_aware_port_ssl=19045)
-        connection.features = ProtocolFeatures(shard_id=kwargs.get('shard_id', self.connection_counter), sharding_info=sharding_info)
+        sharding_info = self.sharding_info or ShardingInfo(
+            shard_id=1, shards_count=4, partitioner="",
+            sharding_algorithm="", sharding_ignore_msb=0,
+            shard_aware_port=19042, shard_aware_port_ssl=19045)
+        connection.features = ProtocolFeatures(
+            shard_id=kwargs.get('shard_id', self.connection_counter),
+            sharding_info=sharding_info)
         self.connection_counter += 1
 
         return connection
@@ -98,8 +103,12 @@ def test_advanced_shard_aware_port(self):
         host = MagicMock()
         host.endpoint = DefaultEndPoint("1.2.3.4")
 
-        for port, is_ssl in [(19042, False), (19045, True)]:
-            session = MockSession(is_ssl=is_ssl)
+        for port, ssl_options, ssl_context in [
+                (19042, None, None),
+                (19045, {'some_ssl_options': True}, None),
+                (19045, {}, None),
+                (19045, None, object())]:
+            session = MockSession(ssl_options=ssl_options, ssl_context=ssl_context)
             pool = HostConnection(host=host, host_distance=HostDistance.REMOTE, session=session)
             try:
                 for f in session.futures:
@@ -114,6 +123,36 @@ def test_advanced_shard_aware_port(self):
             finally:
                 session.cluster.executor.shutdown(wait=True)
 
+    def test_ssl_advanced_shard_aware_port_requires_ssl_port(self):
+        """
+        Test that SSL connections do not fall back to the plaintext
+        shard-aware port when the SSL shard-aware port is unavailable.
+        """
+        host = MagicMock()
+        host.endpoint = DefaultEndPoint("1.2.3.4")
+        sharding_info = ShardingInfo(
+            shard_id=1, shards_count=4, partitioner="", sharding_algorithm="",
+            sharding_ignore_msb=0, shard_aware_port=19042,
+            shard_aware_port_ssl=None)
+        for label, ssl_options, ssl_context in [
+                ('ssl_options', {'some_ssl_options': True}, None),
+                ('empty_ssl_options', {}, None),
+                ('ssl_context', None, object())]:
+            with self.subTest(label=label):
+                session = MockSession(
+                    ssl_options=ssl_options,
+                    ssl_context=ssl_context,
+                    sharding_info=sharding_info)
+                pool = HostConnection(host=host, host_distance=HostDistance.REMOTE, session=session)
+
+                try:
+                    for f in session.futures:
+                        f.result()
+
+                    assert pool._get_shard_aware_endpoint() is None
+                finally:
+                    session.cluster.executor.shutdown(wait=True)
+
     def test_advanced_shard_aware_cooldown(self):
         """
         `disable_advanced_shard_aware` must suppress the shard-aware endpoint for
@@ -123,7 +162,7 @@ def test_advanced_shard_aware_cooldown(self):
         """
         host = MagicMock()
         host.endpoint = DefaultEndPoint("1.2.3.4")
-        session = MockSession(is_ssl=False)
+        session = MockSession()
 
         pool = HostConnection(host=host, host_distance=HostDistance.REMOTE, session=session)
         for f in session.futures: