prose test, incorporate more API changes

Author: isabelatkinson

src/action/csfle/create_data_key.rs

@@ -8,10 +8,10 @@ impl ClientEncryption {
     /// `await` will return d[`Result<Binary>`] (subtype 0x04) with the _id of the created
     /// document as a UUID.
     #[deeplink]
-    pub fn create_data_key(&self, master_key: MasterKey) -> CreateDataKey {
+    pub fn create_data_key(&self, master_key: impl Into<MasterKey>) -> CreateDataKey {
         CreateDataKey {
             client_enc: self,
-            master_key,
+            master_key: master_key.into(),
             options: None,
             #[cfg(test)]
             test_kms_provider: None,

src/client/csfle/client_encryption.rs

@@ -5,6 +5,7 @@ mod encrypt;
 
 use mongocrypt::{ctx::KmsProvider, Crypt};
 use serde::{Deserialize, Serialize};
+use typed_builder::TypedBuilder;
 
 use crate::{
     bson::{doc, spec::BinarySubtype, Binary, RawBinaryRef, RawDocumentBuf},
@@ -193,57 +194,169 @@ impl ClientEncryption {
 #[non_exhaustive]
 #[allow(missing_docs)]
 pub enum MasterKey {
-    #[serde(rename_all = "camelCase")]
-    Aws {
-        region: String,
-        /// The Amazon Resource Name (ARN) to the AWS customer master key (CMK).
-        key: String,
-        /// An alternate host identifier to send KMS requests to. May include port number. Defaults
-        /// to "kms.REGION.amazonaws.com"
-        endpoint: Option<String>,
-    },
-    #[serde(rename_all = "camelCase")]
-    Azure {
-        /// Host with optional port. Example: "example.vault.azure.net".
-        key_vault_endpoint: String,
-        key_name: String,
-        /// A specific version of the named key, defaults to using the key's primary version.
-        key_version: Option<String>,
-    },
-    #[serde(rename_all = "camelCase")]
-    Gcp {
-        project_id: String,
-        location: String,
-        key_ring: String,
-        key_name: String,
-        /// A specific version of the named key, defaults to using the key's primary version.
-        key_version: Option<String>,
-        /// Host with optional port. Defaults to "cloudkms.googleapis.com".
-        endpoint: Option<String>,
-    },
-    /// Master keys are not applicable to `KmsProvider::Local`.
-    Local,
-    #[serde(rename_all = "camelCase")]
-    Kmip {
-        /// keyId is the KMIP Unique Identifier to a 96 byte KMIP Secret Data managed object.  If
-        /// keyId is omitted, the driver creates a random 96 byte KMIP Secret Data managed object.
-        key_id: Option<String>,
-        /// If true (recommended), the KMIP server must decrypt this key. Defaults to false.
-        delegated: Option<bool>,
-        /// Host with optional port.
-        endpoint: Option<String>,
-    },
+    Aws(AwsMasterKey),
+    Azure(AzureMasterKey),
+    Gcp(GcpMasterKey),
+    Kmip(KmipMasterKey),
+    Local(LocalMasterKey),
+}
+
+/// An AWS master key.
+#[serde_with::skip_serializing_none]
+#[derive(Debug, Clone, Serialize, Deserialize, TypedBuilder)]
+#[builder(field_defaults(default, setter(into)))]
+#[serde(rename_all = "camelCase")]
+#[non_exhaustive]
+pub struct AwsMasterKey {
+    /// The name for the key. The value for this field must be the same as the corresponding
+    /// [`KmsProvider`](mongocrypt::ctx::KmsProvider)'s name.
+    #[serde(skip)]
+    pub name: Option<String>,
+
+    /// The region.
+    pub region: String,
+
+    /// The Amazon Resource Name (ARN) to the AWS customer master key (CMK).
+    pub key: String,
+
+    /// An alternate host identifier to send KMS requests to. May include port number. Defaults to
+    /// "kms.<region>.amazonaws.com".
+    pub endpoint: Option<String>,
+}
+
+impl From<AwsMasterKey> for MasterKey {
+    fn from(aws_master_key: AwsMasterKey) -> Self {
+        Self::Aws(aws_master_key)
+    }
+}
+
+/// An Azure master key.
+#[serde_with::skip_serializing_none]
+#[derive(Debug, Clone, Serialize, Deserialize, TypedBuilder)]
+#[builder(field_defaults(default, setter(into)))]
+#[serde(rename_all = "camelCase")]
+#[non_exhaustive]
+pub struct AzureMasterKey {
+    /// The name for the key. The value for this field must be the same as the corresponding
+    /// [`KmsProvider`](mongocrypt::ctx::KmsProvider)'s name.
+    #[serde(skip)]
+    pub name: Option<String>,
+
+    /// Host with optional port. Example: "example.vault.azure.net".
+    pub key_vault_endpoint: String,
+
+    /// The key name.
+    pub key_name: String,
+
+    /// A specific version of the named key, defaults to using the key's primary version.
+    pub key_version: Option<String>,
+}
+
+impl From<AzureMasterKey> for MasterKey {
+    fn from(azure_master_key: AzureMasterKey) -> Self {
+        Self::Azure(azure_master_key)
+    }
+}
+
+/// A GCP master key.
+#[serde_with::skip_serializing_none]
+#[derive(Debug, Clone, Serialize, Deserialize, TypedBuilder)]
+#[builder(field_defaults(default, setter(into)))]
+#[serde(rename_all = "camelCase")]
+#[non_exhaustive]
+pub struct GcpMasterKey {
+    /// The name for the key. The value for this field must be the same as the corresponding
+    /// [`KmsProvider`](mongocrypt::ctx::KmsProvider)'s name.
+    #[serde(skip)]
+    pub name: Option<String>,
+
+    /// The project ID.
+    pub project_id: String,
+
+    /// The location.
+    pub location: String,
+
+    /// The key ring.
+    pub key_ring: String,
+
+    /// The key name.
+    pub key_name: String,
+
+    /// A specific version of the named key. Defaults to using the key's primary version.
+    pub key_version: Option<String>,
+
+    /// Host with optional port. Defaults to "cloudkms.googleapis.com".
+    pub endpoint: Option<String>,
+}
+
+impl From<GcpMasterKey> for MasterKey {
+    fn from(gcp_master_key: GcpMasterKey) -> Self {
+        Self::Gcp(gcp_master_key)
+    }
+}
+
+/// A local master key.
+#[serde_with::skip_serializing_none]
+#[derive(Debug, Clone, Serialize, Deserialize, TypedBuilder)]
+#[builder(field_defaults(default, setter(into)))]
+#[serde(rename_all = "camelCase")]
+#[non_exhaustive]
+pub struct LocalMasterKey {
+    /// The name for the key. The value for this field must be the same as the corresponding
+    /// [`KmsProvider`](mongocrypt::ctx::KmsProvider)'s name.
+    #[serde(skip)]
+    pub name: Option<String>,
+}
+
+impl From<LocalMasterKey> for MasterKey {
+    fn from(local_master_key: LocalMasterKey) -> Self {
+        Self::Local(local_master_key)
+    }
+}
+
+/// A KMIP master key.
+#[serde_with::skip_serializing_none]
+#[derive(Debug, Clone, Serialize, Deserialize, TypedBuilder)]
+#[builder(field_defaults(default, setter(into)))]
+#[serde(rename_all = "camelCase")]
+#[non_exhaustive]
+pub struct KmipMasterKey {
+    /// The name for the key. The value for this field must be the same as the corresponding
+    /// [`KmsProvider`](mongocrypt::ctx::KmsProvider)'s name.
+    #[serde(skip)]
+    pub name: Option<String>,
+
+    /// The KMIP Unique Identifier to a 96 byte KMIP Secret Data managed object. If this field is
+    /// not specified, the driver creates a random 96 byte KMIP Secret Data managed object.
+    pub key_id: Option<String>,
+
+    /// If true (recommended), the KMIP server must decrypt this key. Defaults to false.
+    pub delegated: Option<bool>,
+
+    /// Host with optional port.
+    pub endpoint: Option<String>,
+}
+
+impl From<KmipMasterKey> for MasterKey {
+    fn from(kmip_master_key: KmipMasterKey) -> Self {
+        Self::Kmip(kmip_master_key)
+    }
 }
 
 impl MasterKey {
     /// Returns the `KmsProvider` associated with this key.
     pub fn provider(&self) -> KmsProvider {
-        match self {
-            MasterKey::Aws { .. } => KmsProvider::Aws { name: None },
-            MasterKey::Azure { .. } => KmsProvider::Azure { name: None },
-            MasterKey::Gcp { .. } => KmsProvider::Gcp { name: None },
-            MasterKey::Kmip { .. } => KmsProvider::Kmip { name: None },
-            MasterKey::Local => KmsProvider::Local { name: None },
+        let (provider, name) = match self {
+            MasterKey::Aws(AwsMasterKey { name, .. }) => (KmsProvider::aws(), name.clone()),
+            MasterKey::Azure(AzureMasterKey { name, .. }) => (KmsProvider::azure(), name.clone()),
+            MasterKey::Gcp(GcpMasterKey { name, .. }) => (KmsProvider::gcp(), name.clone()),
+            MasterKey::Kmip(KmipMasterKey { name, .. }) => (KmsProvider::kmip(), name.clone()),
+            MasterKey::Local(LocalMasterKey { name, .. }) => (KmsProvider::local(), name.clone()),
+        };
+        if let Some(name) = name {
+            provider.with_name(name)
+        } else {
+            provider
         }
     }
 }

src/client/csfle/client_encryption/create_data_key.rs

@@ -42,8 +42,8 @@ impl ClientEncryption {
         opts: Option<DataKeyOptions>,
     ) -> Result<Ctx> {
         let mut builder = self.crypt.ctx_builder();
-        let mut key_doc = doc! { "provider": kms_provider.name() };
-        if !matches!(master_key, MasterKey::Local) {
+        let mut key_doc = doc! { "provider": kms_provider.as_string() };
+        if !matches!(master_key, MasterKey::Local { .. }) {
             let master_doc = bson::to_document(&master_key)?;
             key_doc.extend(master_doc);
         }

src/client/csfle/options.rs

@@ -132,28 +132,32 @@ impl KmsProviders {
 
     #[cfg(test)]
     pub(crate) fn set_test_options(&mut self) {
-        use crate::{bson::doc, test::KMS_PROVIDERS_MAP};
-
-        for (provider, credentials) in self.credentials.iter_mut().filter(|(p, _)| {
-            matches!(
-                p,
-                KmsProvider::Aws { .. }
-                    | KmsProvider::Azure { .. }
-                    | KmsProvider::Gcp { .. }
-                    | KmsProvider::Kmip { .. }
-            )
-        }) {
-            let (test_credentials, tls) = KMS_PROVIDERS_MAP.get(provider).unwrap().clone();
-            *credentials = test_credentials;
+        use mongocrypt::ctx::KmsProviderType;
+
+        use crate::{bson::doc, test::csfle::ALL_KMS_PROVIDERS};
+
+        let all_kms_providers = ALL_KMS_PROVIDERS.clone();
+        let mut aws_tls_options = None;
+        for (provider, test_credentials, tls_options) in all_kms_providers {
+            let Some(credentials) = self.credentials.get_mut(&provider) else {
+                continue;
+            };
+            if !matches!(provider.provider_type(), KmsProviderType::Local) {
+                *credentials = test_credentials;
+            }
+
+            if let Some(tls_options) = tls_options {
+                if matches!(provider.provider_type(), KmsProviderType::Aws) {
+                    aws_tls_options = Some(tls_options.clone());
+                }
 
-            if let Some(tls) = tls {
                 self.tls_options
                     .get_or_insert_with(KmsProvidersTlsOptions::new)
-                    .insert(provider.clone(), tls);
+                    .insert(provider.clone(), tls_options);
             }
         }
 
-        let aws_temp_provider = KmsProvider::Other("awsTemporary".to_string());
+        let aws_temp_provider = KmsProvider::other("awsTemporary".to_string());
         if self.credentials.contains_key(&aws_temp_provider) {
             let aws_credentials = doc! {
                 "accessKeyId": std::env::var("CSFLE_AWS_TEMP_ACCESS_KEY_ID").unwrap(),
@@ -162,20 +166,16 @@ impl KmsProviders {
             };
             self.credentials.insert(KmsProvider::aws(), aws_credentials);
 
-            let aws_tls = KMS_PROVIDERS_MAP
-                .get(&KmsProvider::aws())
-                .and_then(|(_, t)| t.as_ref());
-            if let Some(tls) = aws_tls {
+            if let Some(ref aws_tls_options) = aws_tls_options {
                 self.tls_options
                     .get_or_insert_with(KmsProvidersTlsOptions::new)
-                    .insert(KmsProvider::aws(), tls.clone());
+                    .insert(KmsProvider::aws(), aws_tls_options.clone());
             }
 
             self.clear(&aws_temp_provider);
         }
 
-        let aws_temp_no_session_token_provider =
-            KmsProvider::Other("awsTemporaryNoSessionToken".to_string());
+        let aws_temp_no_session_token_provider = KmsProvider::other("awsTemporaryNoSessionToken");
         if self
             .credentials
             .contains_key(&aws_temp_no_session_token_provider)
@@ -186,14 +186,12 @@ impl KmsProviders {
             };
             self.credentials.insert(KmsProvider::aws(), aws_credentials);
 
-            let aws_tls = KMS_PROVIDERS_MAP
-                .get(&KmsProvider::aws())
-                .and_then(|(_, t)| t.as_ref());
-            if let Some(tls) = aws_tls {
+            if let Some(aws_tls_options) = aws_tls_options {
                 self.tls_options
                     .get_or_insert_with(KmsProvidersTlsOptions::new)
-                    .insert(KmsProvider::aws(), tls.clone());
+                    .insert(KmsProvider::aws(), aws_tls_options);
             }
+
             self.clear(&aws_temp_no_session_token_provider);
         }
     }

src/client/csfle/state_machine.rs

@@ -6,7 +6,7 @@ use std::{
 
 use bson::{rawdoc, Document, RawDocument, RawDocumentBuf};
 use futures_util::{stream, TryStreamExt};
-use mongocrypt::ctx::{Ctx, KmsProvider, State};
+use mongocrypt::ctx::{Ctx, KmsProviderType, State};
 use rayon::ThreadPool;
 use tokio::{
     io::{AsyncReadExt, AsyncWriteExt},
@@ -216,8 +216,8 @@ impl CryptExecutor {
                             continue;
                         }
 
-                        match provider {
-                            KmsProvider::Aws { .. } => {
+                        match provider.provider_type() {
+                            KmsProviderType::Aws => {
                                 #[cfg(feature = "aws-auth")]
                                 {
                                     use crate::{
@@ -237,7 +237,7 @@ impl CryptExecutor {
                                     if let Some(token) = aws_creds.session_token() {
                                         creds.append("sessionToken", token);
                                     }
-                                    kms_providers.append(provider.name(), creds);
+                                    kms_providers.append(provider.as_string(), creds);
                                 }
                                 #[cfg(not(feature = "aws-auth"))]
                                 {
@@ -247,11 +247,13 @@ impl CryptExecutor {
                                     ));
                                 }
                             }
-                            KmsProvider::Azure { .. } => {
+                            KmsProviderType::Azure => {
                                 #[cfg(feature = "azure-kms")]
                                 {
-                                    kms_providers
-                                        .append(provider.name(), self.azure.get_token().await?);
+                                    kms_providers.append(
+                                        provider.as_string(),
+                                        self.azure.get_token().await?,
+                                    );
                                 }
                                 #[cfg(not(feature = "azure-kms"))]
                                 {
@@ -261,7 +263,7 @@ impl CryptExecutor {
                                     ));
                                 }
                             }
-                            KmsProvider::Gcp { .. } => {
+                            KmsProviderType::Gcp => {
                                 #[cfg(feature = "gcp-kms")]
                                 {
                                     use crate::runtime::HttpClient;