RUST-1411 Add create_encrypted_collection helper (#853)

Author: abr-egn

src/client/csfle/state_machine.rs

@@ -14,12 +14,12 @@ use tokio::{
 };
 
 use crate::{
-    client::{auth::Credential, options::ServerAddress, WeakClient},
+    client::{options::ServerAddress, WeakClient},
     coll::options::FindOptions,
     error::{Error, Result},
     operation::{RawOutput, RunCommand},
     options::ReadConcern,
-    runtime::{AsyncStream, HttpClient, Process, TlsConfig},
+    runtime::{AsyncStream, Process, TlsConfig},
     Client,
     Namespace,
 };
@@ -209,6 +209,7 @@ impl CryptExecutor {
                 }
                 State::NeedKmsCredentials => {
                     let ctx = result_mut(&mut ctx)?;
+                    #[allow(unused_mut)]
                     let mut out = rawdoc! {};
                     if self
                         .kms_providers
@@ -219,8 +220,8 @@ impl CryptExecutor {
                         #[cfg(feature = "aws-auth")]
                         {
                             let aws_creds = crate::client::auth::aws::AwsCredential::get(
-                                &Credential::default(),
-                                &HttpClient::default(),
+                                &crate::client::auth::Credential::default(),
+                                &crate::runtime::HttpClient::default(),
                             )
                             .await?;
                             let mut creds = rawdoc! {

src/db/mod.rs

@@ -6,6 +6,8 @@ use std::{fmt::Debug, sync::Arc};
 use bson::doc;
 use futures_util::stream::TryStreamExt;
 
+#[cfg(feature = "in-use-encryption-unstable")]
+use crate::client_encryption::{ClientEncryption, MasterKey};
 use crate::{
     bson::{Bson, Document},
     change_stream::{
@@ -421,6 +423,56 @@ impl Database {
         self.create_collection_common(name, options, session).await
     }
 
+    /// Creates a new collection with encrypted fields, automatically creating new data encryption
+    /// keys when needed based on the configured [`CreateCollectionOptions::encrypted_fields`].
+    ///
+    /// Returns the potentially updated `encrypted_fields` along with status, as keys may have been
+    /// created even when a failure occurs.
+    ///
+    /// Does not affect any auto encryption settings on existing MongoClients that are already
+    /// configured with auto encryption.
+    #[cfg(feature = "in-use-encryption-unstable")]
+    pub async fn create_encrypted_collection(
+        &self,
+        ce: &ClientEncryption,
+        name: impl AsRef<str>,
+        options: impl Into<Option<CreateCollectionOptions>>,
+        master_key: MasterKey,
+    ) -> (Document, Result<()>) {
+        let options: Option<CreateCollectionOptions> = options.into();
+        let ef = match options.as_ref().and_then(|o| o.encrypted_fields.as_ref()) {
+            Some(ef) => ef,
+            None => {
+                return (
+                    doc! {},
+                    Err(Error::invalid_argument(
+                        "no encrypted_fields defined for collection",
+                    )),
+                );
+            }
+        };
+        let mut ef_prime = ef.clone();
+        if let Ok(fields) = ef_prime.get_array_mut("fields") {
+            for f in fields {
+                let f_doc = if let Some(d) = f.as_document_mut() {
+                    d
+                } else {
+                    continue;
+                };
+                if f_doc.get("keyId") == Some(&Bson::Null) {
+                    let d = match ce.create_data_key(master_key.clone()).run().await {
+                        Ok(v) => v,
+                        Err(e) => return (ef_prime, Err(e)),
+                    };
+                    f_doc.insert("keyId", d);
+                }
+            }
+        }
+        let mut opts_prime = options.unwrap().clone(); // safe unwrap: no options would be caught by the encrypted_fields check
+        opts_prime.encrypted_fields = Some(ef_prime.clone());
+        (ef_prime, self.create_collection(name, opts_prime).await)
+    }
+
     pub(crate) async fn run_command_common(
         &self,
         command: Document,

src/error.rs

@@ -158,15 +158,15 @@ impl Error {
     }
 
     pub(crate) fn is_max_time_ms_expired_error(&self) -> bool {
-        self.code() == Some(50)
+        self.sdam_code() == Some(50)
     }
 
     /// Whether a read operation should be retried if this error occurs.
     pub(crate) fn is_read_retryable(&self) -> bool {
         if self.is_network_error() {
             return true;
         }
-        match self.code() {
+        match self.sdam_code() {
             Some(code) => RETRYABLE_READ_CODES.contains(&code),
             None => false,
         }
@@ -187,7 +187,7 @@ impl Error {
         if self.is_network_error() {
             return true;
         }
-        match &self.code() {
+        match &self.sdam_code() {
             Some(code) => RETRYABLE_WRITE_CODES.contains(code),
             None => false,
         }
@@ -201,7 +201,7 @@ impl Error {
         {
             return true;
         }
-        match self.code() {
+        match self.sdam_code() {
             Some(code) => UNKNOWN_TRANSACTION_COMMIT_RESULT_LABEL_CODES.contains(&code),
             None => false,
         }
@@ -259,7 +259,7 @@ impl Error {
 
     /// Gets the code from this error for performing SDAM updates, if applicable.
     /// Any codes contained in WriteErrors are ignored.
-    pub(crate) fn code(&self) -> Option<i32> {
+    pub(crate) fn sdam_code(&self) -> Option<i32> {
         match self.kind.as_ref() {
             ErrorKind::Command(command_error) => Some(command_error.code),
             // According to SDAM spec, write concern error codes MUST also be checked, and
@@ -271,7 +271,22 @@ impl Error {
             ErrorKind::Write(WriteFailure::WriteConcernError(wc_error)) => Some(wc_error.code),
             _ => None,
         }
-        .or_else(|| self.source.as_ref().and_then(|s| s.code()))
+        .or_else(|| self.source.as_ref().and_then(|s| s.sdam_code()))
+    }
+
+    /// Gets the code from this error.
+    #[allow(unused)]
+    pub(crate) fn code(&self) -> Option<i32> {
+        match self.kind.as_ref() {
+            ErrorKind::Command(command_error) => Some(command_error.code),
+            ErrorKind::BulkWrite(BulkWriteFailure {
+                write_concern_error: Some(wc_error),
+                ..
+            }) => Some(wc_error.code),
+            ErrorKind::Write(e) => Some(e.code()),
+            _ => None,
+        }
+        .or_else(|| self.source.as_ref().and_then(|s| s.sdam_code()))
     }
 
     /// Gets the message for this error, if applicable, for use in testing.
@@ -333,21 +348,21 @@ impl Error {
 
     /// If this error corresponds to a "not writable primary" error as per the SDAM spec.
     pub(crate) fn is_notwritableprimary(&self) -> bool {
-        self.code()
+        self.sdam_code()
             .map(|code| NOTWRITABLEPRIMARY_CODES.contains(&code))
             .unwrap_or(false)
     }
 
     /// If this error corresponds to a "node is recovering" error as per the SDAM spec.
     pub(crate) fn is_recovering(&self) -> bool {
-        self.code()
+        self.sdam_code()
             .map(|code| RECOVERING_CODES.contains(&code))
             .unwrap_or(false)
     }
 
     /// If this error corresponds to a "node is shutting down" error as per the SDAM spec.
     pub(crate) fn is_shutting_down(&self) -> bool {
-        self.code()
+        self.sdam_code()
             .map(|code| SHUTTING_DOWN_CODES.contains(&code))
             .unwrap_or(false)
     }
@@ -361,7 +376,7 @@ impl Error {
         if !self.is_server_error() {
             return true;
         }
-        let code = self.code();
+        let code = self.sdam_code();
         if code == Some(43) {
             return true;
         }
@@ -388,6 +403,11 @@ impl Error {
         matches!(self.kind.as_ref(), ErrorKind::IncompatibleServer { .. })
     }
 
+    #[allow(unused)]
+    pub(crate) fn is_invalid_argument(&self) -> bool {
+        matches!(self.kind.as_ref(), ErrorKind::InvalidArgument { .. })
+    }
+
     pub(crate) fn with_source<E: Into<Option<Error>>>(mut self, source: E) -> Self {
         self.source = source.into().map(Box::new);
         self
@@ -825,6 +845,13 @@ impl WriteFailure {
             .into())
         }
     }
+
+    pub(crate) fn code(&self) -> i32 {
+        match self {
+            Self::WriteConcernError(e) => e.code,
+            Self::WriteError(e) => e.code,
+        }
+    }
 }
 
 /// An error that occurred during a GridFS operation.

src/gridfs/upload.rs

@@ -180,7 +180,7 @@ impl GridFsBucket {
             .build();
         // Ignore NamespaceExists errors if the collection has already been created.
         if let Err(error) = self.inner.db.create_collection(coll.name(), options).await {
-            if error.code() != Some(48) {
+            if error.sdam_code() != Some(48) {
                 return Err(error);
             }
         }

src/sync/db.rs

@@ -219,6 +219,28 @@ impl Database {
         ))
     }
 
+    /// Creates a new collection with encrypted fields, automatically creating new data encryption
+    /// keys when needed based on the configured [`CreateCollectionOptions::encrypted_fields`].
+    ///
+    /// Returns the potentially updated `encrypted_fields` along with status, as keys may have been
+    /// created even when a failure occurs.
+    ///
+    /// Does not affect any auto encryption settings on existing MongoClients that are already
+    /// configured with auto encryption.
+    #[cfg(feature = "in-use-encryption-unstable")]
+    pub fn create_encrypted_collection(
+        &self,
+        ce: &crate::client_encryption::ClientEncryption,
+        name: impl AsRef<str>,
+        options: impl Into<Option<CreateCollectionOptions>>,
+        master_key: crate::client_encryption::MasterKey,
+    ) -> (Document, Result<()>) {
+        runtime::block_on(
+            self.async_database
+                .create_encrypted_collection(ce, name, options, master_key),
+        )
+    }
+
     /// Runs a database-level command.
     ///
     /// Note that no inspection is done on `doc`, so the command will not use the database's default