Move get credentials logic to a public method

Author: JamieTsai1024

src/client/auth/aws.rs

@@ -99,45 +99,7 @@ async fn authenticate_stream_inner(
     let server_first = ServerFirst::parse(server_first_response.auth_response_body(MECH_NAME)?)?;
     server_first.validate(&nonce)?;
 
-    // Find credentials using MongoDB URI or AWS SDK
-    let aws_credential = if let (Some(access_key), Some(secret_key)) =
-        (&credential.username, &credential.password)
-    {
-        // Look for credentials in the MongoDB URI
-        AwsCredential::from_sdk_creds(
-            access_key.clone(),
-            secret_key.clone(),
-            credential
-                .mechanism_properties
-                .as_ref()
-                .and_then(|mp| mp.get_str("AWS_SESSION_TOKEN").ok())
-                .map(str::to_owned),
-            None,
-        )
-    } else {
-        // If credentials are not provided in the URI, use the AWS SDK to load
-        let creds = aws_config::load_defaults(BehaviorVersion::latest())
-            .await
-            .credentials_provider()
-            .ok_or_else(|| {
-                Error::authentication_error(
-                    MECH_NAME,
-                    &format!("no credential provider configured"),
-                )
-            })?
-            .provide_credentials()
-            .await
-            .map_err(|e| {
-                Error::authentication_error(MECH_NAME, &format!("failed to get creds: {e}"))
-            })?;
-
-        AwsCredential::from_sdk_creds(
-            creds.access_key_id().to_string(),
-            creds.secret_access_key().to_string(),
-            creds.session_token().map(|s| s.to_string()),
-            None,
-        )
-    };
+    let aws_credential = get_aws_credentials(credential).await?;
 
     // Find credentials using original implementation without AWS SDK
     // let aws_credential = {
@@ -203,6 +165,46 @@ async fn authenticate_stream_inner(
     Ok(())
 }
 
+// Find credentials using MongoDB URI or AWS SDK
+pub async fn get_aws_credentials(credential: &Credential) -> Result<AwsCredential> {
+    if let (Some(access_key), Some(secret_key)) = (&credential.username, &credential.password) {
+        // Look for credentials in the MongoDB URI
+        Ok(AwsCredential::from_sdk_creds(
+            access_key.clone(),
+            secret_key.clone(),
+            credential
+                .mechanism_properties
+                .as_ref()
+                .and_then(|mp| mp.get_str("AWS_SESSION_TOKEN").ok())
+                .map(str::to_owned),
+            None,
+        ))
+    } else {
+        // If credentials are not provided in the URI, use the AWS SDK to load
+        let creds = aws_config::load_defaults(BehaviorVersion::latest())
+            .await
+            .credentials_provider()
+            .ok_or_else(|| {
+                Error::authentication_error(
+                    MECH_NAME,
+                    &format!("no credential provider configured"),
+                )
+            })?
+            .provide_credentials()
+            .await
+            .map_err(|e| {
+                Error::authentication_error(MECH_NAME, &format!("failed to get creds: {e}"))
+            })?;
+
+        Ok(AwsCredential::from_sdk_creds(
+            creds.access_key_id().to_string(),
+            creds.secret_access_key().to_string(),
+            creds.session_token().map(|s| s.to_string()),
+            None,
+        ))
+    }
+}
+
 /// Contains the credentials for MONGODB-AWS authentication.
 // RUST-1529 note: dead_code tag added to avoid unused warnings on expiration field
 #[allow(dead_code)]

src/client/csfle/state_machine.rs

@@ -255,11 +255,10 @@ impl CryptExecutor {
                                         runtime::HttpClient,
                                     };
 
-                                    let aws_creds = AwsCredential::get(
-                                        &Credential::default(),
-                                        &HttpClient::default(),
-                                    )
-                                    .await?;
+                                    // Original implementation using the AWS SDK
+                                    let aws_creds =
+                                        get_aws_credentials(&Credential::default()).await?;
+
                                     let mut creds = rawdoc! {
                                         "accessKeyId": aws_creds.access_key(),
                                         "secretAccessKey": aws_creds.secret_key(),