Sync with upstream up to 3.29.1

cassandra/cluster.py

@@ -2545,12 +2545,6 @@ def __init__(self, cluster, hosts, keyspace=None):
 
         self.encoder = Encoder()
 
-        if self.cluster.column_encryption_policy is not None:
-            try:
-                self.client_protocol_handler.column_encryption_policy = self.cluster.column_encryption_policy
-            except AttributeError:
-                log.info("Unable to set column encryption policy for session")
-
         # create connection pools in parallel
         self._initial_connect_futures = set()
         for host in hosts:
@@ -2571,6 +2565,15 @@ def __init__(self, cluster, hosts, keyspace=None):
         self.session_id = uuid.uuid4()
         self._graph_paging_available = self._check_graph_paging_available()
 
+        if self.cluster.column_encryption_policy is not None:
+            try:
+                self.client_protocol_handler = type(
+                    str(self.session_id) + "-ProtocolHandler",
+                    (ProtocolHandler,),
+                    {"column_encryption_policy": self.cluster.column_encryption_policy})
+            except AttributeError:
+                log.info("Unable to set column encryption policy for session")
+
         if self.cluster.monitor_reporting_enabled:
             cc_host = self.cluster.get_control_connection_host()
             valid_insights_version = (cc_host and version_supports_insights(cc_host.dse_version))

tests/integration/standard/column_encryption/test_policies.py

@@ -95,7 +95,11 @@ def test_end_to_end_simple(self):
         self.assertEquals(expected, encrypted)
         self.assertEquals(expected, unencrypted)
 
-    def test_end_to_end_different_cle_contexts(self):
+    def test_end_to_end_different_cle_contexts_different_ivs(self):
+        """
+        Test to validate PYTHON-1350.  We should be able to decode the data from two different contexts (with two different IVs)
+        since the IV used to decrypt the data is actually now stored with the data.
+        """
 
         expected = 2
 
@@ -133,3 +137,34 @@ def test_end_to_end_different_cle_contexts(self):
         (encrypted,unencrypted) = session2.execute("select encrypted, unencrypted from foo.bar where unencrypted = %s allow filtering", (expected,)).one()
         self.assertEquals(expected, encrypted)
         self.assertEquals(expected, unencrypted)
+
+    def test_end_to_end_different_cle_contexts_different_policies(self):
+        """
+        Test to validate PYTHON-1356.  Class variables used to pass CLE policy down to protocol handler shouldn't persist.
+        """
+
+        expected = 3
+
+        key = os.urandom(AES256_KEY_SIZE_BYTES)
+        (col_desc, cl_policy) = self._create_policy(key)
+        cluster = TestCluster(column_encryption_policy=cl_policy)
+        session = cluster.connect()
+        self._recreate_keyspace(session)
+
+        # Use encode_and_encrypt helper function to populate date
+        session.execute("insert into foo.bar (encrypted, unencrypted) values (%s,%s)",(cl_policy.encode_and_encrypt(col_desc, expected), expected))
+
+        # We now open a new session _without_ the CLE policy specified.  We should _not_ be able to read decrypted bits from this session.
+        cluster2 = TestCluster()
+        session2 = cluster2.connect()
+
+        # A straight select from the database will now return the decrypted bits.  We select both encrypted and unencrypted
+        # values here to confirm that we don't interfere with regular processing of unencrypted vals.
+        (encrypted,unencrypted) = session2.execute("select encrypted, unencrypted from foo.bar where unencrypted = %s allow filtering", (expected,)).one()
+        self.assertEquals(cl_policy.encode_and_encrypt(col_desc, expected), encrypted)
+        self.assertEquals(expected, unencrypted)
+
+        # Confirm the same behaviour from a subsequent prepared statement as well
+        prepared = session2.prepare("select encrypted, unencrypted from foo.bar where unencrypted = ? allow filtering")
+        (encrypted,unencrypted) = session2.execute(prepared, [expected]).one()
+        self.assertEquals(cl_policy.encode_and_encrypt(col_desc, expected), encrypted)