RUST-1420 Cache AWS credentials received from endpoints

.evergreen/config.yml

@@ -165,7 +165,7 @@ functions:
         working_dir: "src"
         script: |
           ${PREPARE_SHELL}
-          ASYNC_RUNTIME=${ASYNC_RUNTIME} .evergreen/run-aws-tests.sh
+          ASYNC_RUNTIME=${ASYNC_RUNTIME} SKIP_CREDENTIAL_CACHING_TESTS=1 .evergreen/run-aws-tests.sh
 
   "run aws auth test with assume role credentials":
     - command: shell.exec
@@ -203,7 +203,7 @@ functions:
         working_dir: "src"
         script: |
           ${PREPARE_SHELL}
-          ASYNC_RUNTIME=${ASYNC_RUNTIME} .evergreen/run-aws-tests.sh
+          ASYNC_RUNTIME=${ASYNC_RUNTIME} SKIP_CREDENTIAL_CACHING_TESTS=1 .evergreen/run-aws-tests.sh
 
   "run aws auth test with aws EC2 credentials":
     - command: shell.exec
@@ -243,7 +243,7 @@ functions:
         working_dir: "src"
         script: |
           ${PREPARE_SHELL}
-          ASYNC_RUNTIME=${ASYNC_RUNTIME} PROJECT_DIRECTORY=${PROJECT_DIRECTORY} .evergreen/run-aws-tests.sh
+          ASYNC_RUNTIME=${ASYNC_RUNTIME} PROJECT_DIRECTORY=${PROJECT_DIRECTORY} SKIP_CREDENTIAL_CACHING_TESTS=1 .evergreen/run-aws-tests.sh
 
   "run aws auth test with aws credentials and session token as environment variables":
     - command: shell.exec
@@ -265,7 +265,7 @@ functions:
         working_dir: "src"
         script: |
           ${PREPARE_SHELL}
-          ASYNC_RUNTIME=${ASYNC_RUNTIME} .evergreen/run-aws-tests.sh
+          ASYNC_RUNTIME=${ASYNC_RUNTIME} SKIP_CREDENTIAL_CACHING_TESTS=1 .evergreen/run-aws-tests.sh
 
   "run aws ECS auth test":
     - command: shell.exec
@@ -1066,8 +1066,8 @@ tasks:
       - func: "run aws auth test with assume role credentials"
       - func: "run aws auth test with aws credentials as environment variables"
       - func: "run aws auth test with aws credentials and session token as environment variables"
-      - func: "run aws auth test with aws EC2 credentials"
-      - func: "run aws ECS auth test"
+      # - func: "run aws auth test with aws EC2 credentials"
+      # - func: "run aws ECS auth test"
       - func: "run aws assume role with web identity test"
 
   - name: "test-5.0-standalone"
@@ -1127,8 +1127,8 @@ tasks:
       - func: "run aws auth test with assume role credentials"
       - func: "run aws auth test with aws credentials as environment variables"
       - func: "run aws auth test with aws credentials and session token as environment variables"
-      - func: "run aws auth test with aws EC2 credentials"
-      - func: "run aws ECS auth test"
+      # - func: "run aws auth test with aws EC2 credentials"
+      # - func: "run aws ECS auth test"
       - func: "run aws assume role with web identity test"
 
   - name: "test-6.0-standalone"
@@ -1188,8 +1188,8 @@ tasks:
       - func: "run aws auth test with assume role credentials"
       - func: "run aws auth test with aws credentials as environment variables"
       - func: "run aws auth test with aws credentials and session token as environment variables"
-      - func: "run aws auth test with aws EC2 credentials"
-      - func: "run aws ECS auth test"
+      # - func: "run aws auth test with aws EC2 credentials"
+      # - func: "run aws ECS auth test"
       - func: "run aws assume role with web identity test"
 
   - name: "test-7.0-standalone"
@@ -1249,8 +1249,8 @@ tasks:
       - func: "run aws auth test with assume role credentials"
       - func: "run aws auth test with aws credentials as environment variables"
       - func: "run aws auth test with aws credentials and session token as environment variables"
-      - func: "run aws auth test with aws EC2 credentials"
-      - func: "run aws ECS auth test"
+      # - func: "run aws auth test with aws EC2 credentials"
+      # - func: "run aws ECS auth test"
       - func: "run aws assume role with web identity test"
 
   - name: "test-rapid-standalone"
@@ -1310,8 +1310,8 @@ tasks:
       - func: "run aws auth test with assume role credentials"
       - func: "run aws auth test with aws credentials as environment variables"
       - func: "run aws auth test with aws credentials and session token as environment variables"
-      - func: "run aws auth test with aws EC2 credentials"
-      - func: "run aws ECS auth test"
+      # - func: "run aws auth test with aws EC2 credentials"
+      # - func: "run aws ECS auth test"
       - func: "run aws assume role with web identity test"
 
   - name: "test-latest-standalone"
@@ -1372,8 +1372,8 @@ tasks:
       - func: "run aws auth test with assume role credentials"
       - func: "run aws auth test with aws credentials as environment variables"
       - func: "run aws auth test with aws credentials and session token as environment variables"
-      - func: "run aws auth test with aws EC2 credentials"
-      - func: "run aws ECS auth test"
+      # - func: "run aws auth test with aws EC2 credentials"
+      # - func: "run aws ECS auth test"
       - func: "run aws assume role with web identity test"
 
   - name: "test-connection-string"

.evergreen/run-aws-tests.sh

@@ -44,5 +44,5 @@ set -o errexit
 
 source ./.evergreen/configure-rust.sh
 
-RUST_BACKTRACE=1 cargo test --features aws-auth auth_aws::auth_aws
+RUST_BACKTRACE=1 cargo test --features aws-auth auth_aws
 RUST_BACKTRACE=1 cargo test --features aws-auth lambda_examples::auth::test_handler

src/bson_util/mod.rs

@@ -235,6 +235,17 @@ pub(crate) fn serialize_result_error_as_string<S: Serializer, T: Serialize>(
         .serialize(serializer)
 }
 
+#[cfg(feature = "aws-auth")]
+pub(crate) fn deserialize_datetime_option_from_double<'de, D>(
+    deserializer: D,
+) -> std::result::Result<Option<bson::DateTime>, D::Error>
+where
+    D: Deserializer<'de>,
+{
+    let millis = f64::deserialize(deserializer)? * 1000.0;
+    Ok(Some(bson::DateTime::from_millis(millis as i64)))
+}
+
 #[cfg(test)]
 mod test {
     use crate::bson_util::num_decimal_digits;

src/client/auth/aws.rs

@@ -1,13 +1,16 @@
-use std::{fs::File, io::Read};
+use std::{fs::File, io::Read, time::Duration};
 
 use chrono::{offset::Utc, DateTime};
 use hmac::Hmac;
+use lazy_static::lazy_static;
 use rand::distributions::{Alphanumeric, DistString};
 use serde::Deserialize;
 use sha2::{Digest, Sha256};
+use tokio::sync::Mutex;
 
 use crate::{
     bson::{doc, rawdoc, spec::BinarySubtype, Binary, Bson, Document},
+    bson_util::deserialize_datetime_option_from_double,
     client::{
         auth::{
             self,
@@ -27,12 +30,31 @@ const AWS_EC2_IP: &str = "169.254.169.254";
 const AWS_LONG_DATE_FMT: &str = "%Y%m%dT%H%M%SZ";
 const MECH_NAME: &str = "MONGODB-AWS";
 
+lazy_static! {
+    static ref CACHED_CREDENTIAL: Mutex<Option<AwsCredential>> = Mutex::new(None);
+}
+
 /// Performs MONGODB-AWS authentication for a given stream.
 pub(super) async fn authenticate_stream(
     conn: &mut Connection,
     credential: &Credential,
     server_api: Option<&ServerApi>,
     http_client: &HttpClient,
+) -> Result<()> {
+    match authenticate_stream_inner(conn, credential, server_api, http_client).await {
+        Ok(()) => Ok(()),
+        Err(error) => {
+            *CACHED_CREDENTIAL.lock().await = None;
+            Err(error)
+        }
+    }
+}
+
+async fn authenticate_stream_inner(
+    conn: &mut Connection,
+    credential: &Credential,
+    server_api: Option<&ServerApi>,
+    http_client: &HttpClient,
 ) -> Result<()> {
     let source = match credential.source.as_deref() {
         Some("$external") | None => "$external",
@@ -68,7 +90,23 @@ pub(super) async fn authenticate_stream(
     let server_first = ServerFirst::parse(server_first_response.auth_response_body(MECH_NAME)?)?;
     server_first.validate(&nonce)?;
 
-    let aws_credential = AwsCredential::get(credential, http_client).await?;
+    let aws_credential = {
+        // Limit scope of this variable to avoid holding onto the lock for the duration of
+        // authenticate_stream.
+        let cached_credential = CACHED_CREDENTIAL.lock().await;
+        match *cached_credential {
+            Some(ref aws_credential) if !aws_credential.is_expired() => aws_credential.clone(),
+            _ => {
+                // From the spec: the driver MUST not place a lock on making a request.
+                drop(cached_credential);
+                let aws_credential = AwsCredential::get(credential, http_client).await?;
+                if aws_credential.expiration.is_some() {
+                    *CACHED_CREDENTIAL.lock().await = Some(aws_credential.clone());
+                }
+                aws_credential
+            }
+        }
+    };
 
     let date = Utc::now();
 
@@ -117,7 +155,7 @@ pub(super) async fn authenticate_stream(
 }
 
 /// Contains the credentials for MONGODB-AWS authentication.
-#[derive(Debug, Deserialize)]
+#[derive(Clone, Debug, Deserialize)]
 #[serde(rename_all = "PascalCase")]
 pub(crate) struct AwsCredential {
     access_key_id: String,
@@ -126,6 +164,9 @@ pub(crate) struct AwsCredential {
 
     #[serde(alias = "Token")]
     session_token: Option<String>,
+
+    #[serde(default, deserialize_with = "deserialize_datetime_option_from_double")]
+    expiration: Option<bson::DateTime>,
 }
 
 impl AwsCredential {
@@ -157,6 +198,7 @@ impl AwsCredential {
                 access_key_id: access_key,
                 secret_access_key: secret_key,
                 session_token,
+                expiration: None,
             });
         }
 
@@ -419,6 +461,16 @@ impl AwsCredential {
     pub(crate) fn session_token(&self) -> Option<&str> {
         self.session_token.as_deref()
     }
+
+    fn is_expired(&self) -> bool {
+        match self.expiration {
+            Some(expiration) => {
+                expiration.saturating_duration_since(bson::DateTime::now())
+                    < Duration::from_secs(5 * 60)
+            }
+            None => true,
+        }
+    }
 }
 
 /// The response from the server to the `saslStart` command in a MONGODB-AWS authentication attempt.
@@ -496,3 +548,45 @@ impl ServerFirst {
         }
     }
 }
+
+#[cfg(test)]
+pub(crate) mod test_utils {
+    use super::{AwsCredential, CACHED_CREDENTIAL};
+
+    pub(crate) async fn cached_credential() -> Option<AwsCredential> {
+        CACHED_CREDENTIAL.lock().await.clone()
+    }
+
+    pub(crate) async fn clear_cached_credential() {
+        *CACHED_CREDENTIAL.lock().await = None;
+    }
+
+    pub(crate) async fn poison_cached_credential() {
+        CACHED_CREDENTIAL
+            .lock()
+            .await
+            .as_mut()
+            .unwrap()
+            .access_key_id = "bad".into();
+    }
+
+    pub(crate) async fn cached_access_key_id() -> String {
+        cached_credential().await.unwrap().access_key_id
+    }
+
+    pub(crate) async fn cached_secret_access_key() -> String {
+        cached_credential().await.unwrap().secret_access_key
+    }
+
+    pub(crate) async fn cached_session_token() -> Option<String> {
+        cached_credential().await.unwrap().session_token
+    }
+
+    pub(crate) async fn cached_expiration() -> bson::DateTime {
+        cached_credential().await.unwrap().expiration.unwrap()
+    }
+
+    pub(crate) async fn set_cached_expiration(expiration: bson::DateTime) {
+        CACHED_CREDENTIAL.lock().await.as_mut().unwrap().expiration = Some(expiration);
+    }
+}