RUST-1383 Create or drop auxiliary collections when doing so for a collection with encrypted fields. (#710)

Author: abr-egn

src/client/csfle.rs

@@ -11,6 +11,7 @@ use mongocrypt::Crypt;
 use crate::{
     error::{Error, Result},
     Client,
+    Namespace,
 };
 
 use options::{
@@ -174,3 +175,21 @@ impl ClientState {
         })
     }
 }
+
+pub(crate) fn aux_collections(
+    base_ns: &Namespace,
+    enc_fields: &bson::Document,
+) -> Result<Vec<Namespace>> {
+    let mut out = vec![];
+    for &key in &["esc", "ecc", "ecoc"] {
+        let coll = match enc_fields.get_str(format!("{}Collection", key)) {
+            Ok(s) => s.to_string(),
+            Err(_) => format!("enxcol_.{}.{}", base_ns.coll, key),
+        };
+        out.push(Namespace {
+            coll,
+            ..base_ns.clone()
+        });
+    }
+    Ok(out)
+}

src/client/mod.rs

@@ -1,6 +1,6 @@
 pub mod auth;
 #[cfg(feature = "csfle")]
-mod csfle;
+pub(crate) mod csfle;
 mod executor;
 pub mod options;
 pub mod session;

src/coll/mod.rs

@@ -217,18 +217,25 @@ impl<T> Collection<T> {
         self.inner.write_concern.as_ref()
     }
 
+    #[allow(clippy::needless_option_as_deref)]
     async fn drop_common(
         &self,
         options: impl Into<Option<DropCollectionOptions>>,
         session: impl Into<Option<&mut ClientSession>>,
     ) -> Result<()> {
-        let session = session.into();
+        let mut session = session.into();
 
-        let mut options = options.into();
+        let mut options: Option<DropCollectionOptions> = options.into();
         resolve_options!(self, options, [write_concern]);
 
+        #[cfg(feature = "csfle")]
+        self.drop_aux_collections(options.as_ref(), session.as_deref_mut())
+            .await?;
+
         let drop = DropCollection::new(self.namespace(), options);
-        self.client().execute_operation(drop, session).await
+        self.client()
+            .execute_operation(drop, session.as_deref_mut())
+            .await
     }
 
     /// Drops the collection, deleting all data and indexes stored in it.
@@ -246,6 +253,67 @@ impl<T> Collection<T> {
         self.drop_common(options, session).await
     }
 
+    #[cfg(feature = "csfle")]
+    #[allow(clippy::needless_option_as_deref)]
+    async fn drop_aux_collections(
+        &self,
+        options: Option<&DropCollectionOptions>,
+        mut session: Option<&mut ClientSession>,
+    ) -> Result<()> {
+        // Find associated `encrypted_fields`:
+        // * from options to this call
+        let mut enc_fields = options.and_then(|o| o.encrypted_fields.as_ref());
+        let enc_opts = self.client().auto_encryption_opts().await;
+        // * from client-wide `encrypted_fields_map`:
+        let client_enc_fields = enc_opts
+            .as_ref()
+            .and_then(|eo| eo.encrypted_fields_map.as_ref());
+        if enc_fields.is_none() {
+            enc_fields =
+                client_enc_fields.and_then(|efm| efm.get(&format!("{}", self.namespace())));
+        }
+        // * from a `list_collections` call:
+        let found;
+        if enc_fields.is_none() && client_enc_fields.is_some() {
+            let filter = doc! { "name": self.name() };
+            let mut specs: Vec<_> = match session.as_deref_mut() {
+                Some(s) => {
+                    let mut cursor = self
+                        .inner
+                        .db
+                        .list_collections_with_session(filter, None, s)
+                        .await?;
+                    cursor.stream(s).try_collect().await?
+                }
+                None => {
+                    self.inner
+                        .db
+                        .list_collections(filter, None)
+                        .await?
+                        .try_collect()
+                        .await?
+                }
+            };
+            if let Some(spec) = specs.pop() {
+                if let Some(enc) = spec.options.encrypted_fields {
+                    found = enc;
+                    enc_fields = Some(&found);
+                }
+            }
+        }
+
+        // Drop the collections.
+        if let Some(enc_fields) = enc_fields {
+            for ns in crate::client::csfle::aux_collections(&self.namespace(), enc_fields)? {
+                let drop = DropCollection::new(ns, options.cloned());
+                self.client()
+                    .execute_operation(drop, session.as_deref_mut())
+                    .await?;
+            }
+        }
+        Ok(())
+    }
+
     /// Runs an aggregation operation.
     ///
     /// See the documentation [here](https://www.mongodb.com/docs/manual/aggregation/) for more
@@ -387,7 +455,7 @@ impl<T> Collection<T> {
             .await
     }
 
-    async fn create_index_common(
+    pub(crate) async fn create_index_common(
         &self,
         index: IndexModel,
         options: impl Into<Option<CreateIndexOptions>>,
@@ -1379,6 +1447,21 @@ impl Namespace {
             coll: String::new(),
         }
     }
+
+    pub(crate) fn from_str(s: &str) -> Option<Self> {
+        let mut parts = s.split('.');
+
+        let db = parts.next();
+        let coll = parts.collect::<Vec<_>>().join(".");
+
+        match (db, coll) {
+            (Some(db), coll) if !coll.is_empty() => Some(Self {
+                db: db.to_string(),
+                coll,
+            }),
+            _ => None,
+        }
+    }
 }
 
 impl fmt::Display for Namespace {
@@ -1393,18 +1476,8 @@ impl<'de> Deserialize<'de> for Namespace {
         D: Deserializer<'de>,
     {
         let s: String = Deserialize::deserialize(deserializer)?;
-        let mut parts = s.split('.');
-
-        let db = parts.next();
-        let coll = parts.collect::<Vec<_>>().join(".");
-
-        match (db, coll) {
-            (Some(db), coll) if !coll.is_empty() => Ok(Self {
-                db: db.to_string(),
-                coll,
-            }),
-            _ => Err(D::Error::custom("Missing one or more fields in namespace")),
-        }
+        Self::from_str(&s)
+            .ok_or_else(|| D::Error::custom("Missing one or more fields in namespace"))
     }
 }
 

src/coll/options.rs

@@ -948,6 +948,11 @@ pub struct CreateIndexOptions {
 pub struct DropCollectionOptions {
     /// The write concern for the operation.
     pub write_concern: Option<WriteConcern>,
+
+    /// Map of encrypted fields for the collection.
+    #[cfg(feature = "csfle")]
+    #[serde(skip)]
+    pub encrypted_fields: Option<Document>,
 }
 
 /// Specifies the options to a

src/db/mod.rs

@@ -2,6 +2,8 @@ pub mod options;
 
 use std::{fmt::Debug, sync::Arc};
 
+#[cfg(feature = "csfle")]
+use bson::doc;
 use futures_util::stream::TryStreamExt;
 
 use crate::{
@@ -285,23 +287,93 @@ impl Database {
             .await
     }
 
+    #[allow(clippy::needless_option_as_deref)]
     async fn create_collection_common(
         &self,
         name: impl AsRef<str>,
         options: impl Into<Option<CreateCollectionOptions>>,
         session: impl Into<Option<&mut ClientSession>>,
     ) -> Result<()> {
-        let mut options = options.into();
+        let mut options: Option<CreateCollectionOptions> = options.into();
         resolve_options!(self, options, [write_concern]);
+        let mut session = session.into();
 
-        let create = Create::new(
-            Namespace {
-                db: self.name().to_string(),
-                coll: name.as_ref().to_string(),
-            },
-            options,
-        );
-        self.client().execute_operation(create, session).await
+        let ns = Namespace {
+            db: self.name().to_string(),
+            coll: name.as_ref().to_string(),
+        };
+
+        #[cfg(feature = "csfle")]
+        self.create_aux_collections(&ns, &mut options, session.as_deref_mut())
+            .await?;
+
+        let create = Create::new(ns.clone(), options);
+        self.client()
+            .execute_operation(create, session.as_deref_mut())
+            .await?;
+
+        #[cfg(feature = "csfle")]
+        {
+            let coll = self.collection::<Document>(&ns.coll);
+            coll.create_index_common(
+                crate::IndexModel {
+                    keys: doc! {"__safeContent__": 1},
+                    options: None,
+                },
+                None,
+                session.as_deref_mut(),
+            )
+            .await?;
+        }
+
+        Ok(())
+    }
+
+    #[cfg(feature = "csfle")]
+    #[allow(clippy::needless_option_as_deref)]
+    async fn create_aux_collections(
+        &self,
+        base_ns: &Namespace,
+        options: &mut Option<CreateCollectionOptions>,
+        mut session: Option<&mut ClientSession>,
+    ) -> Result<()> {
+        let has_encrypted_fields = options
+            .as_ref()
+            .and_then(|o| o.encrypted_fields.as_ref())
+            .is_some();
+        // If options does not have `associated_fields`, populate it from client-wide
+        // `encrypted_fields_map`:
+        if !has_encrypted_fields {
+            let enc_opts = self.client().auto_encryption_opts().await;
+            if let Some(enc_opts_fields) = enc_opts
+                .as_ref()
+                .and_then(|eo| eo.encrypted_fields_map.as_ref())
+                .and_then(|efm| efm.get(&format!("{}", &base_ns)))
+            {
+                options
+                    .get_or_insert_with(Default::default)
+                    .encrypted_fields = Some(enc_opts_fields.clone());
+            }
+        }
+
+        if let Some(opts) = options.as_ref() {
+            if let Some(enc_fields) = opts.encrypted_fields.as_ref() {
+                for ns in crate::client::csfle::aux_collections(base_ns, enc_fields)? {
+                    let mut sub_opts = opts.clone();
+                    sub_opts.clustered_index = Some(self::options::ClusteredIndex {
+                        key: doc! { "_id": 1 },
+                        unique: true,
+                        name: None,
+                        v: None,
+                    });
+                    let create = Create::new(ns, Some(sub_opts));
+                    self.client()
+                        .execute_operation(create, session.as_deref_mut())
+                        .await?;
+                }
+            }
+        }
+        Ok(())
     }
 
     /// Creates a new collection in the database with the given `name` and `options`.

src/db/options.rs

@@ -110,6 +110,10 @@ pub struct CreateCollectionOptions {
 
     /// Options for clustered collections.
     pub clustered_index: Option<ClusteredIndex>,
+
+    /// Map of encrypted fields for the created collection.
+    #[cfg(feature = "csfle")]
+    pub encrypted_fields: Option<Document>,
 }
 
 /// Specifies how strictly the database should apply validation rules to existing documents during

src/operation/drop_collection/test.rs

@@ -1,54 +1,9 @@
 use crate::{
     bson::doc,
-    cmap::StreamDescription,
-    concern::{Acknowledgment, WriteConcern},
     error::{ErrorKind, WriteFailure},
-    operation::{test::handle_response_test, DropCollection, Operation},
-    options::DropCollectionOptions,
-    Namespace,
+    operation::{test::handle_response_test, DropCollection},
 };
 
-#[test]
-fn build() {
-    let options = DropCollectionOptions {
-        write_concern: Some(WriteConcern {
-            w: Some(Acknowledgment::Custom("abc".to_string())),
-            ..Default::default()
-        }),
-    };
-
-    let ns = Namespace {
-        db: "test_db".to_string(),
-        coll: "test_coll".to_string(),
-    };
-
-    let mut op = DropCollection::new(ns.clone(), Some(options));
-
-    let description = StreamDescription::new_testing();
-    let cmd = op.build(&description).expect("build should succeed");
-
-    assert_eq!(cmd.name.as_str(), "drop");
-    assert_eq!(cmd.target_db.as_str(), "test_db");
-    assert_eq!(
-        cmd.body,
-        doc! {
-            "drop": "test_coll",
-            "writeConcern": { "w": "abc" }
-        }
-    );
-
-    let mut op = DropCollection::new(ns, None);
-    let cmd = op.build(&description).expect("build should succeed");
-    assert_eq!(cmd.name.as_str(), "drop");
-    assert_eq!(cmd.target_db.as_str(), "test_db");
-    assert_eq!(
-        cmd.body,
-        doc! {
-            "drop": "test_coll",
-        }
-    );
-}
-
 #[test]
 fn handle_success() {
     let op = DropCollection::empty();