RUST-2006 Add option to configure DEK cache lifetime (#1284)

Author: isabelatkinson

src/client/csfle.rs

@@ -122,6 +122,16 @@ impl ClientState {
         if opts.bypass_query_analysis == Some(true) {
             builder = builder.bypass_query_analysis();
         }
+        if let Some(key_cache_expiration) = opts.key_cache_expiration {
+            let expiration_ms: u64 = key_cache_expiration.as_millis().try_into().map_err(|_| {
+                Error::invalid_argument(format!(
+                    "key_cache_expiration must not exceed {} milliseconds, got {:?}",
+                    u64::MAX,
+                    key_cache_expiration
+                ))
+            })?;
+            builder = builder.key_cache_expiration(expiration_ms)?;
+        }
         let crypt = builder.build()?;
         if opts.extra_option(&EO_CRYPT_SHARED_REQUIRED)? == Some(true)
             && crypt.shared_lib_version().is_none()

src/client/csfle/client_builder.rs

@@ -1,3 +1,5 @@
+use std::time::Duration;
+
 use crate::{bson::Document, error::Result, options::ClientOptions, Client};
 
 use super::options::AutoEncryptionOptions;
@@ -101,6 +103,13 @@ impl EncryptedClientBuilder {
         self
     }
 
+    /// Set the duration of time after which the data encryption key cache should expire. Defaults
+    /// to 60 seconds if unset.
+    pub fn key_cache_expiration(mut self, expiration: impl Into<Option<Duration>>) -> Self {
+        self.enc_opts.key_cache_expiration = expiration.into();
+        self
+    }
+
     /// Constructs a new `Client` using automatic encryption.  May perform DNS lookups and/or spawn
     /// mongocryptd as part of `Client` initialization.
     pub async fn build(self) -> Result<Client> {

src/client/csfle/client_encryption.rs

@@ -3,6 +3,8 @@
 mod create_data_key;
 mod encrypt;
 
+use std::time::Duration;
+
 use mongocrypt::{ctx::KmsProvider, Crypt};
 use serde::{Deserialize, Serialize};
 use typed_builder::TypedBuilder;
@@ -61,32 +63,44 @@ impl ClientEncryption {
         key_vault_namespace: Namespace,
         kms_providers: impl IntoIterator<Item = (KmsProvider, bson::Document, Option<TlsOptions>)>,
     ) -> Result<Self> {
-        let kms_providers = KmsProviders::new(kms_providers)?;
-        let crypt = Crypt::builder()
-            .kms_providers(&kms_providers.credentials_doc()?)?
-            .use_need_kms_credentials_state()
-            .retry_kms(true)?
-            .use_range_v2()?
-            .build()?;
-        let exec = CryptExecutor::new_explicit(
-            key_vault_client.weak(),
-            key_vault_namespace.clone(),
-            kms_providers,
-        )?;
-        let key_vault = key_vault_client
-            .database(&key_vault_namespace.db)
-            .collection_with_options(
-                &key_vault_namespace.coll,
-                CollectionOptions::builder()
-                    .write_concern(WriteConcern::majority())
-                    .read_concern(ReadConcern::majority())
-                    .build(),
-            );
-        Ok(ClientEncryption {
-            crypt,
-            exec,
-            key_vault,
-        })
+        Self::builder(key_vault_client, key_vault_namespace, kms_providers).build()
+    }
+
+    /// Initialize a builder to construct a [`ClientEncryption`]. Methods on
+    /// [`ClientEncryptionBuilder`] can be chained to set options.
+    ///
+    /// ```no_run
+    /// # use bson::doc;
+    /// # use mongocrypt::ctx::KmsProvider;
+    /// # use mongodb::client_encryption::ClientEncryption;
+    /// # use mongodb::error::Result;
+    /// # fn func() -> Result<()> {
+    /// # let kv_client = todo!();
+    /// # let kv_namespace = todo!();
+    /// # let local_key = doc! { };
+    /// let enc = ClientEncryption::builder(
+    ///     kv_client,
+    ///     kv_namespace,
+    ///     [
+    ///         (KmsProvider::Local, doc! { "key": local_key }, None),
+    ///         (KmsProvider::Kmip, doc! { "endpoint": "localhost:5698" }, None),
+    ///     ]
+    /// )
+    /// .build()?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn builder(
+        key_vault_client: Client,
+        key_vault_namespace: Namespace,
+        kms_providers: impl IntoIterator<Item = (KmsProvider, bson::Document, Option<TlsOptions>)>,
+    ) -> ClientEncryptionBuilder {
+        ClientEncryptionBuilder {
+            key_vault_client,
+            key_vault_namespace,
+            kms_providers: kms_providers.into_iter().collect(),
+            key_cache_expiration: None,
+        }
     }
 
     // pub async fn rewrap_many_data_key(&self, _filter: Document, _opts: impl
@@ -189,6 +203,68 @@ impl ClientEncryption {
     }
 }
 
+/// Builder for constructing a [`ClientEncryption`]. Construct by calling
+/// [`ClientEncryption::builder`].
+pub struct ClientEncryptionBuilder {
+    key_vault_client: Client,
+    key_vault_namespace: Namespace,
+    kms_providers: Vec<(KmsProvider, bson::Document, Option<TlsOptions>)>,
+    key_cache_expiration: Option<Duration>,
+}
+
+impl ClientEncryptionBuilder {
+    /// Set the duration of time after which the data encryption key cache should expire. Defaults
+    /// to 60 seconds if unset.
+    pub fn key_cache_expiration(mut self, expiration: impl Into<Option<Duration>>) -> Self {
+        self.key_cache_expiration = expiration.into();
+        self
+    }
+
+    /// Build the [`ClientEncryption`].
+    pub fn build(self) -> Result<ClientEncryption> {
+        let kms_providers = KmsProviders::new(self.kms_providers)?;
+
+        let mut crypt_builder = Crypt::builder()
+            .kms_providers(&kms_providers.credentials_doc()?)?
+            .use_need_kms_credentials_state()
+            .use_range_v2()?
+            .retry_kms(true)?;
+        if let Some(key_cache_expiration) = self.key_cache_expiration {
+            let expiration_ms: u64 = key_cache_expiration.as_millis().try_into().map_err(|_| {
+                Error::invalid_argument(format!(
+                    "key_cache_expiration must not exceed {} milliseconds, got {:?}",
+                    u64::MAX,
+                    key_cache_expiration
+                ))
+            })?;
+            crypt_builder = crypt_builder.key_cache_expiration(expiration_ms)?;
+        }
+        let crypt = crypt_builder.build()?;
+
+        let exec = CryptExecutor::new_explicit(
+            self.key_vault_client.weak(),
+            self.key_vault_namespace.clone(),
+            kms_providers,
+        )?;
+        let key_vault = self
+            .key_vault_client
+            .database(&self.key_vault_namespace.db)
+            .collection_with_options(
+                &self.key_vault_namespace.coll,
+                CollectionOptions::builder()
+                    .write_concern(WriteConcern::majority())
+                    .read_concern(ReadConcern::majority())
+                    .build(),
+            );
+
+        Ok(ClientEncryption {
+            crypt,
+            exec,
+            key_vault,
+        })
+    }
+}
+
 /// A KMS-specific key used to encrypt data keys.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(untagged)]

src/client/csfle/options.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, time::Duration};
 
 use bson::Array;
 use mongocrypt::ctx::KmsProvider;
@@ -8,6 +8,7 @@ use crate::{
     bson::{Bson, Document},
     client::options::TlsOptions,
     error::{Error, Result},
+    serde_util,
     Namespace,
 };
 
@@ -59,6 +60,14 @@ pub(crate) struct AutoEncryptionOptions {
     #[cfg(test)]
     #[serde(skip)]
     pub(crate) disable_crypt_shared: Option<bool>,
+    /// The duration after which the data encryption key cache expires. Defaults to 60 seconds if
+    /// unset.
+    #[serde(
+        default,
+        rename = "keyExpirationMS",
+        deserialize_with = "serde_util::deserialize_duration_option_from_u64_millis"
+    )]
+    pub(crate) key_cache_expiration: Option<Duration>,
 }
 
 fn default_key_vault_namespace() -> Namespace {
@@ -81,6 +90,7 @@ impl AutoEncryptionOptions {
             bypass_query_analysis: None,
             #[cfg(test)]
             disable_crypt_shared: None,
+            key_cache_expiration: None,
         }
     }
 }

src/cmap/conn.rs

@@ -327,8 +327,8 @@ impl PinnedConnectionHandle {
         }
     }
 
-    /// Retrieve the pinned connection.  Will fail if the connection has been unpinned or is still in
-    /// use.
+    /// Retrieve the pinned connection.  Will fail if the connection has been unpinned or is still
+    /// in use.
     pub(crate) async fn take_connection(&self) -> Result<PooledConnection> {
         use tokio::sync::mpsc::error::TryRecvError;
         let mut receiver = self.receiver.lock().await;