RUST-1314 Support on-demand AWS credentials for in-use encryption (#831) (#835)

Author: abr-egn

.evergreen/run-csfle-tests.sh

@@ -7,7 +7,7 @@ source ./.evergreen/env.sh
 
 set -o xtrace
 
-FEATURE_FLAGS="in-use-encryption-unstable,${TLS_FEATURE}"
+FEATURE_FLAGS="in-use-encryption-unstable,aws-auth,${TLS_FEATURE}"
 OPTIONS="-- -Z unstable-options --format json --report-time"
 
 if [ "$SINGLE_THREAD" = true ]; then
@@ -38,6 +38,11 @@ set +o errexit
 cargo_test test::csfle > prose.xml
 cargo_test test::spec::client_side_encryption > spec.xml
 
-junit-report-merger results.xml prose.xml spec.xml
+# Unset variables for on-demand credential failure tests.
+unset AWS_ACCESS_KEY_ID
+unset AWS_SECRET_ACCESS_KEY
+cargo_test test::csfle::on_demand_aws_failure > failure.xml
+
+junit-report-merger results.xml prose.xml spec.xml failure.xml
 
 exit ${CARGO_RESULT}

src/client/auth/aws.rs

@@ -115,7 +115,7 @@ pub(super) async fn authenticate_stream(
 
 /// Contains the credentials for MONGODB-AWS authentication.
 #[derive(Debug, Deserialize)]
-struct AwsCredential {
+pub(crate) struct AwsCredential {
     #[serde(rename = "AccessKeyId")]
     access_key: String,
 
@@ -129,7 +129,7 @@ struct AwsCredential {
 impl AwsCredential {
     /// Derives the credentials for an authentication attempt given the set of credentials the user
     /// passed in.
-    async fn get(credential: &Credential, http_client: &HttpClient) -> Result<Self> {
+    pub(crate) async fn get(credential: &Credential, http_client: &HttpClient) -> Result<Self> {
         let access_key = credential
             .username
             .clone()
@@ -348,6 +348,18 @@ impl AwsCredential {
 
         Ok(auth_header)
     }
+
+    pub(crate) fn access_key(&self) -> &str {
+        &self.access_key
+    }
+
+    pub(crate) fn secret_key(&self) -> &str {
+        &self.secret_key
+    }
+
+    pub(crate) fn session_token(&self) -> Option<&str> {
+        self.session_token.as_deref()
+    }
 }
 
 /// The response from the server to the `saslStart` command in a MONGODB-AWS authentication attempt.

src/client/auth/mod.rs

@@ -2,7 +2,7 @@
 //! [`Client`](struct.Client.html).
 
 #[cfg(feature = "aws-auth")]
-mod aws;
+pub(crate) mod aws;
 mod plain;
 mod sasl;
 mod scram;

src/client/csfle.rs

@@ -69,7 +69,7 @@ impl ClientState {
         let exec = CryptExecutor::new_implicit(
             aux_clients.key_vault_client,
             opts.key_vault_namespace.clone(),
-            opts.kms_providers.tls_options().clone(),
+            opts.kms_providers.clone(),
             mongocryptd_opts,
             mongocryptd_client,
             aux_clients.metadata_client,

src/client/csfle/client_encryption.rs

@@ -60,11 +60,12 @@ impl ClientEncryption {
         let kms_providers = KmsProviders::new(kms_providers)?;
         let crypt = Crypt::builder()
             .kms_providers(&kms_providers.credentials_doc()?)?
+            .use_need_kms_credentials_state()
             .build()?;
         let exec = CryptExecutor::new_explicit(
             key_vault_client.weak(),
             key_vault_namespace.clone(),
-            kms_providers.tls_options().clone(),
+            kms_providers,
         )?;
         let key_vault = key_vault_client
             .database(&key_vault_namespace.db)

src/client/csfle/options.rs

@@ -122,11 +122,10 @@ impl KmsProviders {
         Ok(bson::to_document(&self.credentials)?)
     }
 
-    pub(crate) fn tls_options(&self) -> &Option<KmsProvidersTlsOptions> {
-        &self.tls_options
+    pub(crate) fn tls_options(&self) -> Option<&KmsProvidersTlsOptions> {
+        self.tls_options.as_ref()
     }
 
-    #[cfg(test)]
     pub(crate) fn credentials(&self) -> &HashMap<KmsProvider, Document> {
         &self.credentials
     }

src/client/csfle/state_machine.rs

@@ -4,33 +4,33 @@ use std::{
     path::{Path, PathBuf},
 };
 
-use bson::{Document, RawDocument, RawDocumentBuf};
+use bson::{rawdoc, Document, RawDocument, RawDocumentBuf};
 use futures_util::{stream, TryStreamExt};
-use mongocrypt::ctx::{Ctx, State};
+use mongocrypt::ctx::{Ctx, KmsProvider, State};
 use rayon::ThreadPool;
 use tokio::{
     io::{AsyncReadExt, AsyncWriteExt},
     sync::{oneshot, Mutex},
 };
 
 use crate::{
-    client::{options::ServerAddress, WeakClient},
+    client::{auth::Credential, options::ServerAddress, WeakClient},
     coll::options::FindOptions,
     error::{Error, Result},
     operation::{RawOutput, RunCommand},
     options::ReadConcern,
-    runtime::{AsyncStream, Process, TlsConfig},
+    runtime::{AsyncStream, HttpClient, Process, TlsConfig},
     Client,
     Namespace,
 };
 
-use super::options::KmsProvidersTlsOptions;
+use super::options::KmsProviders;
 
 #[derive(Debug)]
 pub(crate) struct CryptExecutor {
     key_vault_client: WeakClient,
     key_vault_namespace: Namespace,
-    tls_options: Option<KmsProvidersTlsOptions>,
+    kms_providers: KmsProviders,
     crypto_threads: ThreadPool,
     mongocryptd: Option<Mongocryptd>,
     mongocryptd_client: Option<Client>,
@@ -41,7 +41,7 @@ impl CryptExecutor {
     pub(crate) fn new_explicit(
         key_vault_client: WeakClient,
         key_vault_namespace: Namespace,
-        tls_options: Option<KmsProvidersTlsOptions>,
+        kms_providers: KmsProviders,
     ) -> Result<Self> {
         // TODO RUST-1492: Replace num_cpus with std::thread::available_parallelism.
         let crypto_threads = rayon::ThreadPoolBuilder::new()
@@ -51,7 +51,7 @@ impl CryptExecutor {
         Ok(Self {
             key_vault_client,
             key_vault_namespace,
-            tls_options,
+            kms_providers,
             crypto_threads,
             mongocryptd: None,
             mongocryptd_client: None,
@@ -62,7 +62,7 @@ impl CryptExecutor {
     pub(crate) async fn new_implicit(
         key_vault_client: WeakClient,
         key_vault_namespace: Namespace,
-        tls_options: Option<KmsProvidersTlsOptions>,
+        kms_providers: KmsProviders,
         mongocryptd_opts: Option<MongocryptdOptions>,
         mongocryptd_client: Option<Client>,
         metadata_client: Option<WeakClient>,
@@ -71,7 +71,7 @@ impl CryptExecutor {
             Some(opts) => Some(Mongocryptd::new(opts).await?),
             None => None,
         };
-        let mut exec = Self::new_explicit(key_vault_client, key_vault_namespace, tls_options)?;
+        let mut exec = Self::new_explicit(key_vault_client, key_vault_namespace, kms_providers)?;
         exec.mongocryptd = mongocryptd;
         exec.mongocryptd_client = mongocryptd_client;
         exec.metadata_client = metadata_client;
@@ -185,8 +185,8 @@ impl CryptExecutor {
                             let addr = ServerAddress::parse(endpoint)?;
                             let provider = kms_ctx.kms_provider()?;
                             let tls_options = self
-                                .tls_options
-                                .as_ref()
+                                .kms_providers
+                                .tls_options()
                                 .and_then(|tls| tls.get(&provider))
                                 .cloned()
                                 .unwrap_or_default();
@@ -208,8 +208,38 @@ impl CryptExecutor {
                         .await?;
                 }
                 State::NeedKmsCredentials => {
-                    // TODO(RUST-1314, RUST-1417): support fetching KMS credentials.
-                    return Err(Error::internal("KMS credentials are not yet supported"));
+                    let ctx = result_mut(&mut ctx)?;
+                    let mut out = rawdoc! {};
+                    if self
+                        .kms_providers
+                        .credentials()
+                        .get(&KmsProvider::Aws)
+                        .map_or(false, |d| d.is_empty())
+                    {
+                        #[cfg(feature = "aws-auth")]
+                        {
+                            let aws_creds = crate::client::auth::aws::AwsCredential::get(
+                                &Credential::default(),
+                                &HttpClient::default(),
+                            )
+                            .await?;
+                            let mut creds = rawdoc! {
+                                "accessKeyId": aws_creds.access_key(),
+                                "secretAccessKey": aws_creds.secret_key(),
+                            };
+                            if let Some(token) = aws_creds.session_token() {
+                                creds.append("sessionToken", token);
+                            }
+                            out.append("aws", creds);
+                        }
+                        #[cfg(not(feature = "aws-auth"))]
+                        {
+                            return Err(Error::invalid_argument(
+                                "On-demand AWS KMS credentials require the `aws-auth` feature.",
+                            ));
+                        }
+                    }
+                    ctx.provide_kms_providers(&out)?;
                 }
                 State::Ready => {
                     let (tx, rx) = oneshot::channel();

src/test/csfle.rs

@@ -2749,7 +2749,70 @@ impl CommandEventHandler for DecryptionEventsHandler {
     }
 }
 
-// TODO RUST-1314: implement prose test 15. On-demand AWS Credentials
+// Prose test 15. On-demand AWS Credentials (failure)
+#[cfg(feature = "aws-auth")]
+#[cfg_attr(feature = "tokio-runtime", tokio::test)]
+#[cfg_attr(feature = "async-std-runtime", async_std::test)]
+async fn on_demand_aws_failure() -> Result<()> {
+    if !check_env("on_demand_aws_failure", false) {
+        return Ok(());
+    }
+    if std::env::var("AWS_ACCESS_KEY_ID").is_ok() && std::env::var("AWS_SECRET_ACCESS_KEY").is_ok()
+    {
+        log_uncaptured("Skipping on_demand_aws_failure: credentials set");
+        return Ok(());
+    }
+    let _guard = LOCK.run_exclusively().await;
+
+    let ce = ClientEncryption::new(
+        Client::test_builder().build().await.into_client(),
+        KV_NAMESPACE.clone(),
+        [(KmsProvider::Aws, doc! {}, None)],
+    )?;
+    let result = ce
+        .create_data_key(MasterKey::Aws {
+            region: "us-east-1".to_string(),
+            key: "arn:aws:kms:us-east-1:579766882180:key/89fcc2c4-08b0-4bd9-9f25-e30687b580d0"
+                .to_string(),
+            endpoint: None,
+        })
+        .run()
+        .await;
+    assert!(
+        result.as_ref().unwrap_err().is_auth_error(),
+        "Expected auth error, got {:?}",
+        result
+    );
+
+    Ok(())
+}
+
+// Prose test 15. On-demand AWS Credentials (success)
+#[cfg(feature = "aws-auth")]
+#[cfg_attr(feature = "tokio-runtime", tokio::test)]
+#[cfg_attr(feature = "async-std-runtime", async_std::test)]
+async fn on_demand_aws_success() -> Result<()> {
+    if !check_env("on_demand_aws_success", false) {
+        return Ok(());
+    }
+    let _guard = LOCK.run_exclusively().await;
+
+    let ce = ClientEncryption::new(
+        Client::test_builder().build().await.into_client(),
+        KV_NAMESPACE.clone(),
+        [(KmsProvider::Aws, doc! {}, None)],
+    )?;
+    ce.create_data_key(MasterKey::Aws {
+        region: "us-east-1".to_string(),
+        key: "arn:aws:kms:us-east-1:579766882180:key/89fcc2c4-08b0-4bd9-9f25-e30687b580d0"
+            .to_string(),
+        endpoint: None,
+    })
+    .run()
+    .await?;
+
+    Ok(())
+}
 
 // TODO RUST-1441: implement prose test 16. Rewrap