From cdff534549e466deb1aabbc7100b30b007c22cb7 Mon Sep 17 00:00:00 2001
From: Thanh Nguyen <thanh@timweri.me>
Date: Thu, 19 Dec 2024 12:36:55 -0500
Subject: [PATCH] Allow KMS IAM auth via default config (#57)

* Allow KMS IAM auth via default config

* Fix comment

* Update error
---
 rustica/src/signing/amazon_kms.rs | 52 +++++++++++++++----------------
 rustica/src/signing/mod.rs        |  4 ++-
 2 files changed, 29 insertions(+), 27 deletions(-)

diff --git a/rustica/src/signing/amazon_kms.rs b/rustica/src/signing/amazon_kms.rs
index 6197ea8..d8983c6 100644
--- a/rustica/src/signing/amazon_kms.rs
+++ b/rustica/src/signing/amazon_kms.rs
@@ -1,11 +1,10 @@
 use std::time::Duration;
 
 use aws_config::timeout::TimeoutConfig;
-use aws_credential_types::provider::future;
-use aws_credential_types::provider::ProvideCredentials;
+use aws_credential_types::Credentials;
 use aws_sdk_kms::types::SigningAlgorithmSpec;
 use aws_sdk_kms::{
-    config::{Credentials, Region},
+    config::Region,
     primitives::Blob,
     Client,
 };
@@ -39,9 +38,9 @@ pub struct KmsKeyDefinition {
 #[derive(Clone, Debug, Deserialize)]
 pub struct Config {
     /// The AWS access key that can access the KMS keys
-    aws_access_key_id: String,
+    aws_access_key_id: Option<String>,
     /// The secret corresponding to the AWS access key
-    aws_secret_access_key: String,
+    aws_secret_access_key: Option<String>,
     /// The region to be used
     aws_region: String,
 
@@ -91,21 +90,6 @@ pub struct AmazonKMSSigner {
     client: Client,
 }
 
-impl ProvideCredentials for Config {
-    fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a>
-    where
-        Self: 'a,
-    {
-        future::ProvideCredentials::ready(Ok(Credentials::new(
-            self.aws_access_key_id.clone(),
-            self.aws_secret_access_key.clone(),
-            None,
-            None,
-            "AmazonKMSSigner",
-        )))
-    }
-}
-
 pub struct KmsRcgenRemoteSigner {
     x509_public_key: Vec<u8>,
     /// The id to lookup the private portion of the X509 key in Amazon KMS
@@ -275,12 +259,28 @@ impl SignerConfig for Config {
             .operation_attempt_timeout(Duration::from_secs(10))
             .operation_timeout(Duration::from_secs(10))
             .build();
-        let aws_config = aws_config::from_env()
-            .timeout_config(timeout_config)
-            .region(Region::new(self.aws_region.clone()))
-            .credentials_provider(self.clone())
-            .load()
-            .await;
+        let aws_config = match (self.aws_access_key_id, self.aws_secret_access_key) {
+            (Some(_), None) => return Err(
+                SigningError::InvalidAwsConfig("aws_access_key_id is defined but aws_secret_access_key is not defined".to_string())
+            ),
+            (None, Some(_)) => return Err(
+                SigningError::InvalidAwsConfig("aws_secret_access_key is defined but aws_access_key_id is not defined".to_string())
+            ),
+            (Some(access_key_id), Some(secret_access_key)) => aws_config::from_env()
+                .region(Region::new(self.aws_region.clone()))
+                .credentials_provider(
+                    Credentials::new(access_key_id, secret_access_key, None, None, "AmazonKMSSigner")
+                )
+                .load()
+                .await,
+            // If access key is not defined, use the default config
+            (None, None) => aws_config::from_env()
+                .timeout_config(timeout_config)
+                .region(Region::new(self.aws_region.clone()))
+                .load()
+                .await,
+        };
+
         let client = Client::new(&aws_config);
 
         let ssh_keys = match (self.user_key, self.host_key) {
diff --git a/rustica/src/signing/mod.rs b/rustica/src/signing/mod.rs
index 4213e8c..3e0c4c4 100644
--- a/rustica/src/signing/mod.rs
+++ b/rustica/src/signing/mod.rs
@@ -125,6 +125,7 @@ pub enum SigningError {
     IdenticalUserAndHostKey(String),
     SignerDoesNotHaveSSHKeys,
     SignerDoesNotAllRequiredSSHKeys,
+    InvalidAwsConfig(String),
 }
 
 impl std::fmt::Display for SigningError {
@@ -137,7 +138,8 @@ impl std::fmt::Display for SigningError {
             Self::DuplicatedKey(a1, a2) => write!(f, "Authorities {a1} and {a2} share at least one key. This is not allowed as it almost always a misconfiguration leading to access that is not correctly restricted"),
             Self::IdenticalUserAndHostKey(authority) => write!(f, "Authority {authority} has an identical key for both user and host certificates. This is not allowed as it's much safer to use separate keys for both."),
             Self::SignerDoesNotHaveSSHKeys => write!(f, "Signer was not configured with SSH keys so it cannot create an SSH certificate"),
-            Self::SignerDoesNotAllRequiredSSHKeys => write!(f, "Signer did not have both user and host keys defined")
+            Self::SignerDoesNotAllRequiredSSHKeys => write!(f, "Signer did not have both user and host keys defined"),
+            Self::InvalidAwsConfig(e) => write!(f, "Invalid AWS config: {e}"),
         }
     }
 }