PYTHON-1884 Support auto encryption in cursors

Author: ShaneHarvey

pymongo/encryption.py

@@ -197,19 +197,27 @@ def __init__(self, io_callbacks, opts):
             opts._kms_providers, schema_map))
         self._bypass_auto_encryption = opts._bypass_auto_encryption
 
-    def encrypt(self, database, cmd):
+    def encrypt(self, database, cmd, check_keys, codec_options):
         """Encrypt a MongoDB command.
 
         :Parameters:
           - `database`: The database for this command.
-          - `cmd`: A command as BSON.
+          - `cmd`: A command document.
+          - `check_keys`: If True, check `cmd` for invalid keys.
+          - `codec_options`: The CodecOptions to use while encoding `cmd`.
 
         :Returns:
           The encrypted command to execute.
         """
-        encrypted_cmd = self._auto_encrypter.encrypt(database, cmd)
+        # Workaround for $clusterTime which is incompatible with check_keys.
+        cluster_time = check_keys and cmd.pop('$clusterTime', None)
+        encrypted_cmd = self._auto_encrypter.encrypt(
+            database, _dict_to_bson(cmd, check_keys, codec_options))
         # TODO: PYTHON-1922 avoid decoding the encrypted_cmd.
-        return _inflate_bson(encrypted_cmd, DEFAULT_RAW_BSON_OPTIONS)
+        encrypt_cmd = _inflate_bson(encrypted_cmd, DEFAULT_RAW_BSON_OPTIONS)
+        if cluster_time:
+            encrypt_cmd['$clusterTime'] = cluster_time
+        return encrypt_cmd
 
     def decrypt(self, response):
         """Decrypt a MongoDB command response.

pymongo/message.py

@@ -305,6 +305,12 @@ def as_command(self, sock_info):
                     'readConcern', {})[
                     'afterClusterTime'] = session.operation_time
         sock_info.send_cluster_time(cmd, session, self.client)
+        # Support auto encryption
+        client = self.client
+        if (client._encrypter and
+                not client._encrypter._bypass_auto_encryption):
+            cmd = client._encrypter.encrypt(
+                self.db, cmd, False, self.codec_options)
         self._as_command = cmd, self.db
         return self._as_command
 
@@ -393,6 +399,12 @@ def as_command(self, sock_info):
         if self.session:
             self.session._apply_to(cmd, False, self.read_preference)
         sock_info.send_cluster_time(cmd, self.session, self.client)
+        # Support auto encryption
+        client = self.client
+        if (client._encrypter and
+                not client._encrypter._bypass_auto_encryption):
+            cmd = client._encrypter.encrypt(
+                self.db, cmd, False, self.codec_options)
         self._as_command = cmd, self.db
         return self._as_command
 

pymongo/network.py

@@ -34,7 +34,7 @@
 except ImportError:
     _SELECT_ERROR = OSError
 
-from bson import _dict_to_bson, _decode_all_selective
+from bson import _decode_all_selective
 from bson.py3compat import PY3
 
 from pymongo import helpers, message
@@ -117,13 +117,8 @@ def command(sock, dbname, spec, slave_ok, is_mongos,
 
     if (client and client._encrypter and
             not client._encrypter._bypass_auto_encryption):
-        # Workaround for $clusterTime which is incompatible with check_keys.
-        if check_keys:
-            cluster_time = spec.pop('$clusterTime', None)
         spec = orig = client._encrypter.encrypt(
-            dbname, _dict_to_bson(spec, check_keys, codec_options))
-        if check_keys and cluster_time:
-            spec['$clusterTime'] = cluster_time
+            dbname, spec, check_keys, codec_options)
         # We already checked the keys, no need to do it again.
         check_keys = False
 

pymongo/server.py

@@ -16,6 +16,8 @@
 
 from datetime import datetime
 
+from bson import _decode_all_selective
+
 from pymongo.errors import NotMasterError, OperationFailure
 from pymongo.helpers import _check_command_response
 from pymongo.message import _convert_exception
@@ -164,6 +166,15 @@ def run_operation_with_response(
                 duration, res, operation.name, request_id,
                 sock_info.address)
 
+        # Decrypt response.
+        client = operation.client
+        if client and client._encrypter:
+            if use_cmd:
+                decrypted = client._encrypter.decrypt(
+                    reply.raw_command_response())
+                docs = _decode_all_selective(
+                    decrypted, operation.codec_options, user_fields)
+
         if exhaust:
             response = ExhaustResponse(
                 data=reply,

test/test_encryption.py

@@ -157,22 +157,38 @@ def _test_auto_encrypt(self, opts):
         self.addCleanup(key_vault.drop)
 
         # Collection.insert_one auto encrypts.
+        docs = [{'_id': 1, 'ssn': '123'},
+                {'_id': 2, 'ssn': '456'},
+                {'_id': 3, 'ssn': '789'}]
         encrypted_coll = client.pymongo_test.test
-        encrypted_coll.insert_one({'_id': 1, 'ssn': '123'})
+        for doc in docs:
+            encrypted_coll.insert_one(doc)
 
         # Database.command auto decrypts.
         res = client.pymongo_test.command(
             'find', 'test', filter={'ssn': '123'})
         decrypted_docs = res['cursor']['firstBatch']
         self.assertEqual(decrypted_docs, [{'_id': 1, 'ssn': '123'}])
 
+        # Collection.find auto decrypts.
+        decrypted_docs = list(encrypted_coll.find())
+        self.assertEqual(decrypted_docs, docs)
+
+        # Collection.find auto decrypts getMores.
+        decrypted_docs = list(encrypted_coll.find(batch_size=1))
+        self.assertEqual(decrypted_docs, docs)
+
         # Collection.aggregate auto decrypts.
         decrypted_docs = list(encrypted_coll.aggregate([]))
-        self.assertEqual(decrypted_docs, [{'_id': 1, 'ssn': '123'}])
+        self.assertEqual(decrypted_docs, docs)
+
+        # Collection.aggregate auto decrypts getMores.
+        decrypted_docs = list(encrypted_coll.aggregate([], batchSize=1))
+        self.assertEqual(decrypted_docs, docs)
 
         # Collection.distinct auto decrypts.
         decrypted_ssns = encrypted_coll.distinct('ssn')
-        self.assertEqual(decrypted_ssns, ['123'])
+        self.assertEqual(decrypted_ssns, ['123', '456', '789'])
 
         # Make sure the field is actually encrypted.
         encrypted_doc = self.db.test.find_one()