RUST-1384 Implement auto encryption (#717)

Author: abr-egn

.evergreen/feature-combinations.sh

@@ -6,5 +6,5 @@
 export FEATURE_COMBINATIONS=(
     '' # default features
     '--no-default-features --features async-std-runtime,sync' # features that conflict w/ default features
-    '--features tokio-sync,zstd-compression,snappy-compression,zlib-compression,openssl-tls,aws-auth' # additive features
+    '--features tokio-sync,zstd-compression,snappy-compression,zlib-compression,openssl-tls,aws-auth,csfle' # additive features
 )

Cargo.toml

@@ -66,7 +66,7 @@ zlib-compression = ["flate2"]
 snappy-compression = ["snap"]
 
 # DO NOT USE; see https://jira.mongodb.org/browse/RUST-569 for the status of CSFLE support in the Rust driver.
-csfle = ["mongocrypt", "which"]
+csfle = ["mongocrypt", "which", "rayon"]
 
 [dependencies]
 async-trait = "0.1.42"
@@ -89,6 +89,7 @@ openssl-probe = { version = "0.1.5", optional = true }
 os_info = { version = "3.0.1", default-features = false }
 percent-encoding = "2.0.0"
 rand = { version = "0.8.3", features = ["small_rng"] }
+rayon = { version = "1.5.3", optional = true }
 rustc_version_runtime = "0.2.1"
 rustls-pemfile = "0.3.0"
 serde_with = "1.3.1"

src/client/csfle.rs

@@ -1,4 +1,5 @@
 pub mod options;
+mod state_machine;
 
 use std::{
     path::Path,
@@ -7,6 +8,7 @@ use std::{
 
 use derivative::Derivative;
 use mongocrypt::Crypt;
+use rayon::ThreadPool;
 
 use crate::{
     error::{Error, Result},
@@ -30,34 +32,37 @@ use super::WeakClient;
 #[derivative(Debug)]
 pub(super) struct ClientState {
     #[derivative(Debug = "ignore")]
-    #[allow(dead_code)]
-    crypt: Crypt,
+    pub(crate) crypt: Crypt,
     mongocryptd_client: Option<Client>,
     aux_clients: AuxClients,
     opts: AutoEncryptionOptions,
+    crypto_threads: ThreadPool,
 }
 
 #[derive(Debug)]
 struct AuxClients {
-    #[allow(dead_code)]
     key_vault_client: WeakClient,
-    #[allow(dead_code)]
     metadata_client: Option<WeakClient>,
-    #[allow(dead_code)]
-    internal_client: Option<Client>,
+    _internal_client: Option<Client>,
 }
 
 impl ClientState {
     pub(super) async fn new(client: &Client, opts: AutoEncryptionOptions) -> Result<Self> {
         let crypt = Self::make_crypt(&opts)?;
         let mongocryptd_client = Self::spawn_mongocryptd_if_needed(&opts, &crypt).await?;
         let aux_clients = Self::make_aux_clients(client, &opts)?;
+        let num_cpus = std::thread::available_parallelism()?.get();
+        let crypto_threads = rayon::ThreadPoolBuilder::new()
+            .num_threads(num_cpus)
+            .build()
+            .map_err(|e| Error::internal(format!("could not initialize thread pool: {}", e)))?;
 
         Ok(Self {
             crypt,
             mongocryptd_client,
             aux_clients,
             opts,
+            crypto_threads,
         })
     }
 
@@ -171,7 +176,7 @@ impl ClientState {
         Ok(AuxClients {
             key_vault_client,
             metadata_client,
-            internal_client,
+            _internal_client: internal_client,
         })
     }
 }

src/client/csfle/state_machine.rs

@@ -0,0 +1,171 @@
+use std::convert::TryInto;
+
+use bson::{Document, RawDocument, RawDocumentBuf};
+use futures_util::{stream, TryStreamExt};
+use mongocrypt::ctx::{Ctx, State};
+use tokio::{
+    io::{AsyncReadExt, AsyncWriteExt},
+    sync::oneshot,
+};
+
+use crate::{
+    client::options::ServerAddress,
+    cmap::options::StreamOptions,
+    error::{Error, Result},
+    operation::{RawOutput, RunCommand},
+    runtime::AsyncStream,
+    Client,
+};
+
+impl Client {
+    pub(crate) async fn run_mongocrypt_ctx(
+        &self,
+        ctx: Ctx,
+        db: Option<&str>,
+    ) -> Result<RawDocumentBuf> {
+        let guard = self.inner.csfle.read().await;
+        let csfle = match guard.as_ref() {
+            Some(csfle) => csfle,
+            None => return Err(Error::internal("no csfle state for mongocrypt ctx")),
+        };
+        let mut result = None;
+        // This needs to be a `Result` so that the `Ctx` can be temporarily owned by the processing
+        // thread for crypto finalization.  An `Option` would also work here, but `Result` means we
+        // can return a helpful error if things get into a broken state rather than panicing.
+        let mut ctx = Ok(ctx);
+        loop {
+            let state = result_ref(&ctx)?.state()?;
+            match state {
+                State::NeedMongoCollinfo => {
+                    let ctx = result_mut(&mut ctx)?;
+                    let filter = raw_to_doc(ctx.mongo_op()?)?;
+                    let metadata_client = csfle
+                        .aux_clients
+                        .metadata_client
+                        .as_ref()
+                        .and_then(|w| w.upgrade())
+                        .ok_or_else(|| {
+                            Error::internal("metadata_client required for NeedMongoCollinfo state")
+                        })?;
+                    let db = metadata_client.database(db.as_ref().ok_or_else(|| {
+                        Error::internal("db required for NeedMongoCollinfo state")
+                    })?);
+                    let mut cursor = db.list_collections(filter, None).await?;
+                    if cursor.advance().await? {
+                        ctx.mongo_feed(cursor.current())?;
+                    }
+                    ctx.mongo_done()?;
+                }
+                State::NeedMongoMarkings => {
+                    let ctx = result_mut(&mut ctx)?;
+                    let command = ctx.mongo_op()?.to_raw_document_buf();
+                    let db = db.as_ref().ok_or_else(|| {
+                        Error::internal("db required for NeedMongoMarkings state")
+                    })?;
+                    let op = RawOutput(RunCommand::new_raw(db.to_string(), command, None, None)?);
+                    let mongocryptd_client = csfle
+                        .mongocryptd_client
+                        .as_ref()
+                        .ok_or_else(|| Error::internal("mongocryptd client not found"))?;
+                    let response = mongocryptd_client.execute_operation(op, None).await?;
+                    ctx.mongo_feed(response.raw_body())?;
+                    ctx.mongo_done()?;
+                }
+                State::NeedMongoKeys => {
+                    let ctx = result_mut(&mut ctx)?;
+                    let filter = raw_to_doc(ctx.mongo_op()?)?;
+                    let kv_ns = &csfle.opts.key_vault_namespace;
+                    let kv_client = csfle
+                        .aux_clients
+                        .key_vault_client
+                        .upgrade()
+                        .ok_or_else(|| Error::internal("key vault client dropped"))?;
+                    let kv_coll = kv_client
+                        .database(&kv_ns.db)
+                        .collection::<RawDocumentBuf>(&kv_ns.coll);
+                    let mut cursor = kv_coll.find(filter, None).await?;
+                    while cursor.advance().await? {
+                        ctx.mongo_feed(cursor.current())?;
+                    }
+                    ctx.mongo_done()?;
+                }
+                State::NeedKms => {
+                    let ctx = result_mut(&mut ctx)?;
+                    let scope = ctx.kms_scope();
+                    let mut kms_ctxen: Vec<Result<_>> = vec![];
+                    while let Some(kms_ctx) = scope.next_kms_ctx() {
+                        kms_ctxen.push(Ok(kms_ctx));
+                    }
+                    stream::iter(kms_ctxen)
+                        .try_for_each_concurrent(None, |mut kms_ctx| async move {
+                            let endpoint = kms_ctx.endpoint()?;
+                            let addr = ServerAddress::parse(endpoint)?;
+                            let provider = kms_ctx.kms_provider()?;
+                            let tls_options = csfle
+                                .opts()
+                                .tls_options
+                                .as_ref()
+                                .and_then(|tls| tls.get(&provider))
+                                .cloned()
+                                .unwrap_or_default();
+                            let mut stream = AsyncStream::connect(
+                                StreamOptions::builder()
+                                    .address(addr)
+                                    .tls_options(tls_options)
+                                    .build(),
+                            )
+                            .await?;
+                            stream.write_all(kms_ctx.message()?).await?;
+                            let mut buf = vec![];
+                            while kms_ctx.bytes_needed() > 0 {
+                                let buf_size = kms_ctx.bytes_needed().try_into().map_err(|e| {
+                                    Error::internal(format!("buffer size overflow: {}", e))
+                                })?;
+                                buf.resize(buf_size, 0);
+                                let count = stream.read(&mut buf).await?;
+                                kms_ctx.feed(&buf[0..count])?;
+                            }
+                            Ok(())
+                        })
+                        .await?;
+                }
+                State::NeedKmsCredentials => todo!("RUST-1314"),
+                State::Ready => {
+                    let (tx, rx) = oneshot::channel();
+                    let mut thread_ctx = std::mem::replace(
+                        &mut ctx,
+                        Err(Error::internal("crypto context not present")),
+                    )?;
+                    csfle.crypto_threads.spawn(move || {
+                        let result = thread_ctx.finalize().map(|doc| doc.to_owned());
+                        let _ = tx.send((thread_ctx, result));
+                    });
+                    let (ctx_again, output) = rx
+                        .await
+                        .map_err(|_| Error::internal("crypto thread dropped"))?;
+                    ctx = Ok(ctx_again);
+                    result = Some(output?);
+                }
+                State::Done => break,
+                s => return Err(Error::internal(format!("unhandled state {:?}", s))),
+            }
+        }
+        match result {
+            Some(doc) => Ok(doc),
+            None => Err(Error::internal("libmongocrypt terminated without output")),
+        }
+    }
+}
+
+fn result_ref<T>(r: &Result<T>) -> Result<&T> {
+    r.as_ref().map_err(Error::clone)
+}
+
+fn result_mut<T>(r: &mut Result<T>) -> Result<&mut T> {
+    r.as_mut().map_err(|e| e.clone())
+}
+
+fn raw_to_doc(raw: &RawDocument) -> Result<Document> {
+    raw.try_into()
+        .map_err(|e| Error::internal(format!("could not parse raw document: {}", e)))
+}

src/client/executor.rs

@@ -1,4 +1,8 @@
+#[cfg(feature = "csfle")]
+use bson::RawDocumentBuf;
 use bson::{doc, RawBsonRef, RawDocument, Timestamp};
+#[cfg(feature = "csfle")]
+use futures_core::future::BoxFuture;
 use lazy_static::lazy_static;
 use serde::de::DeserializeOwned;
 
@@ -598,6 +602,21 @@ impl Client {
         let target_db = cmd.target_db.clone();
 
         let serialized = op.serialize_command(cmd)?;
+        #[cfg(feature = "csfle")]
+        let serialized = {
+            let guard = self.inner.csfle.read().await;
+            if let Some(ref csfle) = *guard {
+                if csfle.opts().bypass_auto_encryption != Some(true) {
+                    self.auto_encrypt(csfle, RawDocument::from_bytes(&serialized)?, &target_db)
+                        .await?
+                        .into_bytes()
+                } else {
+                    serialized
+                }
+            } else {
+                serialized
+            }
+        };
         let raw_cmd = RawCommand {
             name: cmd_name.clone(),
             target_db,
@@ -750,6 +769,21 @@ impl Client {
                     handler.handle_command_succeeded_event(command_succeeded_event);
                 });
 
+                #[cfg(feature = "csfle")]
+                let response = {
+                    let guard = self.inner.csfle.read().await;
+                    if let Some(ref csfle) = *guard {
+                        if csfle.opts().bypass_auto_encryption != Some(true) {
+                            let new_body = self.auto_decrypt(csfle, response.raw_body()).await?;
+                            RawCommandResponse::new_raw(response.source, new_body)
+                        } else {
+                            response
+                        }
+                    } else {
+                        response
+                    }
+                };
+
                 match op.handle_response(response, connection.stream_description()?) {
                     Ok(response) => Ok(response),
                     Err(mut err) => {
@@ -765,6 +799,34 @@ impl Client {
         }
     }
 
+    #[cfg(feature = "csfle")]
+    fn auto_encrypt<'a>(
+        &'a self,
+        csfle: &'a super::csfle::ClientState,
+        command: &'a RawDocument,
+        target_db: &'a str,
+    ) -> BoxFuture<'a, Result<RawDocumentBuf>> {
+        Box::pin(async move {
+            let ctx = csfle
+                .crypt
+                .ctx_builder()
+                .build_encrypt(target_db, command)?;
+            self.run_mongocrypt_ctx(ctx, Some(target_db)).await
+        })
+    }
+
+    #[cfg(feature = "csfle")]
+    fn auto_decrypt<'a>(
+        &'a self,
+        csfle: &'a super::csfle::ClientState,
+        response: &'a RawDocument,
+    ) -> BoxFuture<'a, Result<RawDocumentBuf>> {
+        Box::pin(async move {
+            let ctx = csfle.crypt.ctx_builder().build_decrypt(response)?;
+            self.run_mongocrypt_ctx(ctx, None).await
+        })
+    }
+
     /// Start an implicit session if the operation and write concern are compatible with sessions.
     async fn start_implicit_session<T: Operation>(&self, op: &T) -> Result<Option<ClientSession>> {
         match self.get_session_support_status().await? {

src/cmap/conn/command.rs

@@ -203,10 +203,11 @@ impl RawCommandResponse {
 
     pub(crate) fn new(source: ServerAddress, message: Message) -> Result<Self> {
         let raw = message.single_document_response()?;
-        Ok(Self {
-            source,
-            raw: RawDocumentBuf::from_bytes(raw)?,
-        })
+        Ok(Self::new_raw(source, RawDocumentBuf::from_bytes(raw)?))
+    }
+
+    pub(crate) fn new_raw(source: ServerAddress, raw: RawDocumentBuf) -> Self {
+        Self { source, raw }
     }
 
     pub(crate) fn body<'a, T: Deserialize<'a>>(&'a self) -> Result<T> {

src/coll/mod.rs

@@ -1189,6 +1189,10 @@ where
         }
 
         let ordered = options.as_ref().and_then(|o| o.ordered).unwrap_or(true);
+        #[cfg(feature = "csfle")]
+        let encrypted = self.client().auto_encryption_opts().await.is_some();
+        #[cfg(not(feature = "csfle"))]
+        let encrypted = false;
 
         let mut cumulative_failure: Option<BulkWriteFailure> = None;
         let mut error_labels: HashSet<String> = Default::default();
@@ -1198,7 +1202,7 @@ where
 
         while n_attempted < ds.len() {
             let docs: Vec<&T> = ds.iter().skip(n_attempted).map(Borrow::borrow).collect();
-            let insert = Insert::new(self.namespace(), docs, options.clone());
+            let insert = Insert::new_encrypted(self.namespace(), docs, options.clone(), encrypted);
 
             match self
                 .client()