RUST-2161 Support auto encryption in unified tests (#1426)

Author: abr-egn

src/client/csfle.rs

@@ -230,3 +230,15 @@ pub(crate) fn aux_collections(
     }
     Ok(out)
 }
+
+impl Client {
+    pub(crate) async fn init_csfle(&self, opts: AutoEncryptionOptions) -> Result<()> {
+        let mut csfle_state = self.inner.csfle.write().await;
+        if csfle_state.is_some() {
+            return Err(Error::internal("double initialization of csfle state"));
+        }
+        *csfle_state = Some(ClientState::new(self, opts).await?);
+
+        Ok(())
+    }
+}

src/client/csfle/client_builder.rs

@@ -112,8 +112,7 @@ impl EncryptedClientBuilder {
     /// mongocryptd as part of `Client` initialization.
     pub async fn build(self) -> Result<Client> {
         let client = Client::with_options(self.client_options)?;
-        *client.inner.csfle.write().await =
-            Some(super::ClientState::new(&client, self.enc_opts).await?);
+        client.init_csfle(self.enc_opts).await?;
         Ok(client)
     }
 }

src/coll/action/drop.rs

@@ -52,7 +52,7 @@ where
         }
         // * from a `list_collections` call:
         let found;
-        if enc_fields.is_none() && client_enc_fields.is_some() {
+        if enc_fields.is_none() && enc_opts.is_some() {
             let filter = doc! { "name": self.name() };
             let mut specs: Vec<_> = match session.as_deref_mut() {
                 Some(s) => {

src/test/csfle.rs

@@ -44,6 +44,11 @@ pub(crate) type KmsProviderList = Vec<KmsInfo>;
 static CSFLE_LOCAL_KEY: Lazy<String> = Lazy::new(|| get_env_var("CSFLE_LOCAL_KEY"));
 static FLE_AWS_KEY: Lazy<String> = Lazy::new(|| get_env_var("FLE_AWS_KEY"));
 static FLE_AWS_SECRET: Lazy<String> = Lazy::new(|| get_env_var("FLE_AWS_SECRET"));
+static FLE_AWS_TEMP_KEY: Lazy<String> = Lazy::new(|| get_env_var("CSFLE_AWS_TEMP_ACCESS_KEY_ID"));
+static FLE_AWS_TEMP_SECRET: Lazy<String> =
+    Lazy::new(|| get_env_var("CSFLE_AWS_TEMP_SECRET_ACCESS_KEY"));
+static FLE_AWS_TEMP_SESSION_TOKEN: Lazy<String> =
+    Lazy::new(|| get_env_var("CSFLE_AWS_TEMP_SESSION_TOKEN"));
 static FLE_AZURE_TENANTID: Lazy<String> = Lazy::new(|| get_env_var("FLE_AZURE_TENANTID"));
 static FLE_AZURE_CLIENTID: Lazy<String> = Lazy::new(|| get_env_var("FLE_AZURE_CLIENTID"));
 static FLE_AZURE_CLIENTSECRET: Lazy<String> = Lazy::new(|| get_env_var("FLE_AZURE_CLIENTSECRET"));
@@ -61,13 +66,16 @@ static CSFLE_TLS_CERT_DIR: Lazy<String> = Lazy::new(|| get_env_var("CSFLE_TLS_CE
 static CRYPT_SHARED_LIB_PATH: Lazy<String> = Lazy::new(|| get_env_var("CRYPT_SHARED_LIB_PATH"));
 
 fn get_env_var(name: &str) -> String {
-    std::env::var(name).unwrap_or_else(|_| {
-        panic!(
-            "Missing environment variable for {}. See src/test/csfle.rs for the list of required \
-             variables and instructions for retrieving them.",
-            name
-        )
-    })
+    match std::env::var(name) {
+        Ok(v) if !v.is_empty() => v,
+        _ => {
+            panic!(
+                "Missing environment variable for {}. See src/test/csfle.rs for the list of \
+                 required variables and instructions for retrieving them.",
+                name
+            )
+        }
+    }
 }
 
 pub(crate) static AWS_KMS: Lazy<KmsInfo> = Lazy::new(|| {
@@ -80,6 +88,17 @@ pub(crate) static AWS_KMS: Lazy<KmsInfo> = Lazy::new(|| {
         None,
     )
 });
+static AWS_TEMP_KMS: Lazy<KmsInfo> = Lazy::new(|| {
+    (
+        KmsProvider::aws(),
+        doc! {
+            "accessKeyId": &*FLE_AWS_TEMP_KEY,
+            "secretAccessKey": &*FLE_AWS_TEMP_SECRET,
+            "sessionToken": &*FLE_AWS_TEMP_SESSION_TOKEN,
+        },
+        None,
+    )
+});
 pub(crate) static AWS_KMS_NAME1: Lazy<KmsInfo> = Lazy::new(|| {
     let aws_info = AWS_KMS.clone();
     (aws_info.0.with_name("name1"), aws_info.1, aws_info.2)
@@ -310,3 +329,39 @@ async fn fle2v2_ok(name: &str) -> bool {
     }
     true
 }
+
+pub(crate) fn fill_kms_placeholders(
+    kms_provider_map: std::collections::HashMap<mongocrypt::ctx::KmsProvider, Document>,
+) -> KmsProviderList {
+    use mongocrypt::ctx::KmsProviderType;
+
+    let placeholder = doc! { "$$placeholder": 1 };
+
+    let mut kms_providers = Vec::new();
+    for (provider, mut config) in kms_provider_map {
+        // AWS uses temp creds if the "sessionToken" key is present in the config
+        let test_kms_provider = if *provider.provider_type() == KmsProviderType::Aws
+            && config.contains_key("sessionToken")
+        {
+            Some(&*AWS_TEMP_KMS)
+        } else {
+            (*ALL_KMS_PROVIDERS).iter().find(|(p, ..)| p == &provider)
+        };
+
+        for (key, value) in config.iter_mut() {
+            if value.as_document() == Some(&placeholder) {
+                let test_kms_provider = test_kms_provider
+                    .unwrap_or_else(|| panic!("missing config for {:?}", provider));
+                let placeholder_value = test_kms_provider.1.get(key).unwrap_or_else(|| {
+                    panic!("provider config {:?} missing key {:?}", provider, key)
+                });
+                *value = placeholder_value.clone();
+            }
+        }
+
+        let tls_options = test_kms_provider.and_then(|(_, _, tls_options)| tls_options.clone());
+        kms_providers.push((provider, config, tls_options));
+    }
+
+    kms_providers
+}

src/test/csfle/spec.rs

@@ -16,8 +16,16 @@ async fn run_unified() {
 
 #[tokio::test(flavor = "multi_thread")]
 async fn run_legacy() {
-    // TODO RUST-528: unskip this file
-    let mut skipped_files = vec!["timeoutMS.json"];
+    let mut skipped_files = vec![
+        // TODO RUST-528: unskip this file
+        "timeoutMS.json",
+        // These files have been migrated to unified tests.
+        // TODO DRIVERS-3178 remove these once the files are gone.
+        "fle2v2-BypassQueryAnalysis.json",
+        "fle2v2-EncryptedFields-vs-EncryptedFieldsMap.json",
+        "localSchema.json",
+        "maxWireVersion.json",
+    ];
     if cfg!(not(feature = "openssl-tls")) {
         skipped_files.push("kmipKMS.json");
     }