feat(kms): reduce redundant kms calls (#1264)

crates/router/src/configs.rs

@@ -1,3 +1,5 @@
 mod defaults;
+#[cfg(feature = "kms")]
+pub(super) mod kms;
 pub mod settings;
 mod validations;

crates/router/src/configs/kms.rs

@@ -0,0 +1,60 @@
+use common_utils::errors::CustomResult;
+use external_services::kms;
+use masking::ExposeInterface;
+
+use crate::configs::settings;
+
+#[async_trait::async_trait]
+// This trait performs inplace decryption of the structure on which this is implemented
+pub(crate) trait KmsDecrypt {
+    async fn decrypt_inner(self, kms_config: &kms::KmsConfig) -> CustomResult<Self, kms::KmsError>
+    where
+        Self: Sized;
+}
+
+#[async_trait::async_trait]
+impl KmsDecrypt for settings::Jwekey {
+    async fn decrypt_inner(self, kms_config: &kms::KmsConfig) -> CustomResult<Self, kms::KmsError> {
+        let client = kms::get_kms_client(kms_config).await;
+
+        // If this pattern required repetition, a macro approach needs to be deviced
+        let (
+            locker_encryption_key1,
+            locker_encryption_key2,
+            locker_decryption_key1,
+            locker_decryption_key2,
+            vault_encryption_key,
+            vault_private_key,
+            tunnel_private_key,
+        ) = tokio::try_join!(
+            client.decrypt(self.locker_encryption_key1),
+            client.decrypt(self.locker_encryption_key2),
+            client.decrypt(self.locker_decryption_key1),
+            client.decrypt(self.locker_decryption_key2),
+            client.decrypt(self.vault_encryption_key),
+            client.decrypt(self.vault_private_key),
+            client.decrypt(self.tunnel_private_key),
+        )?;
+
+        Ok(Self {
+            locker_key_identifier1: self.locker_key_identifier1,
+            locker_key_identifier2: self.locker_key_identifier2,
+            locker_encryption_key1,
+            locker_encryption_key2,
+            locker_decryption_key1,
+            locker_decryption_key2,
+            vault_encryption_key,
+            vault_private_key,
+            tunnel_private_key,
+        })
+    }
+}
+
+#[async_trait::async_trait]
+impl KmsDecrypt for settings::ActiveKmsSecrets {
+    async fn decrypt_inner(self, kms_config: &kms::KmsConfig) -> CustomResult<Self, kms::KmsError> {
+        Ok(Self {
+            jwekey: self.jwekey.expose().decrypt_inner(kms_config).await?.into(),
+        })
+    }
+}

crates/router/src/configs/settings.rs

@@ -39,6 +39,16 @@ pub enum Subcommand {
     GenerateOpenapiSpec,
 }
 
+#[cfg(feature = "kms")]
+/// Store the decrypted kms secret values for active use in the application
+/// Currently using `StrongSecret` won't have any effect as this struct have smart pointers to heap
+/// allocations.
+/// note: we can consider adding such behaviour in the future with custom implementation
+#[derive(Clone)]
+pub struct ActiveKmsSecrets {
+    pub jwekey: masking::Secret<Jwekey>,
+}
+
 #[derive(Debug, Deserialize, Clone, Default)]
 #[serde(default)]
 pub struct Settings {

crates/router/src/core/payment_methods/cards.rs

@@ -250,10 +250,10 @@ pub async fn add_card_hs(
     merchant_account: &storage::MerchantAccount,
 ) -> errors::CustomResult<(api::PaymentMethodResponse, bool), errors::VaultError> {
     let locker = &state.conf.locker;
+    #[cfg(not(feature = "kms"))]
     let jwekey = &state.conf.jwekey;
-
     #[cfg(feature = "kms")]
-    let kms_config = &state.conf.kms;
+    let jwekey = &state.kms_secrets;
 
     let db = &*state.store;
     let merchant_id = &merchant_account.merchant_id;
@@ -264,15 +264,8 @@ pub async fn add_card_hs(
         .get_required_value("locker_id")
         .change_context(errors::VaultError::SaveCardFailed)?;
 
-    let request = payment_methods::mk_add_card_request_hs(
+    let request =
-        jwekey,
+        payment_methods::mk_add_card_request_hs(jwekey, locker, &card, &customer_id, merchant_id)
-        locker,
-        &card,
-        &customer_id,
-        merchant_id,
-        #[cfg(feature = "kms")]
-        kms_config,
-    )
             .await?;
 
     let stored_card_response = if !locker.mock_locker {
@@ -284,12 +277,7 @@ pub async fn add_card_hs(
             .get_response_inner("JweBody")
             .change_context(errors::VaultError::FetchCardFailed)?;
 
-        let decrypted_payload = payment_methods::get_decrypted_response_payload(
+        let decrypted_payload = payment_methods::get_decrypted_response_payload(jwekey, jwe_body)
-            jwekey,
-            jwe_body,
-            #[cfg(feature = "kms")]
-            kms_config,
-        )
             .await
             .change_context(errors::VaultError::SaveCardFailed)
             .attach_printable("Error getting decrypted response payload")?;
@@ -394,10 +382,10 @@ pub async fn get_card_from_hs_locker<'a>(
     card_reference: &'a str,
 ) -> errors::CustomResult<payment_methods::Card, errors::VaultError> {
     let locker = &state.conf.locker;
+    #[cfg(not(feature = "kms"))]
     let jwekey = &state.conf.jwekey;
-
     #[cfg(feature = "kms")]
-    let kms_config = &state.conf.kms;
+    let jwekey = &state.kms_secrets;
 
     let request = payment_methods::mk_get_card_request_hs(
         jwekey,
@@ -405,8 +393,6 @@ pub async fn get_card_from_hs_locker<'a>(
         customer_id,
         merchant_id,
         card_reference,
-        #[cfg(feature = "kms")]
-        kms_config,
     )
     .await
     .change_context(errors::VaultError::FetchCardFailed)
@@ -419,12 +405,7 @@ pub async fn get_card_from_hs_locker<'a>(
         let jwe_body: services::JweBody = response
             .get_response_inner("JweBody")
             .change_context(errors::VaultError::FetchCardFailed)?;
-        let decrypted_payload = payment_methods::get_decrypted_response_payload(
+        let decrypted_payload = payment_methods::get_decrypted_response_payload(jwekey, jwe_body)
-            jwekey,
-            jwe_body,
-            #[cfg(feature = "kms")]
-            kms_config,
-        )
             .await
             .change_context(errors::VaultError::FetchCardFailed)
             .attach_printable("Error getting decrypted response payload for get card")?;
@@ -483,10 +464,10 @@ pub async fn delete_card_from_hs_locker<'a>(
     card_reference: &'a str,
 ) -> errors::RouterResult<payment_methods::DeleteCardResp> {
     let locker = &state.conf.locker;
+    #[cfg(not(feature = "kms"))]
     let jwekey = &state.conf.jwekey;
-
     #[cfg(feature = "kms")]
-    let kms_config = &state.conf.kms;
+    let jwekey = &state.kms_secrets;
 
     let request = payment_methods::mk_delete_card_request_hs(
         jwekey,
@@ -494,8 +475,6 @@ pub async fn delete_card_from_hs_locker<'a>(
         customer_id,
         merchant_id,
         card_reference,
-        #[cfg(feature = "kms")]
-        kms_config,
     )
     .await
     .change_context(errors::ApiErrorResponse::InternalServerError)
@@ -507,12 +486,7 @@ pub async fn delete_card_from_hs_locker<'a>(
             .change_context(errors::ApiErrorResponse::InternalServerError)
             .attach_printable("Failed while executing call_connector_api for delete card");
         let jwe_body: services::JweBody = response.get_response_inner("JweBody")?;
-        let decrypted_payload = payment_methods::get_decrypted_response_payload(
+        let decrypted_payload = payment_methods::get_decrypted_response_payload(jwekey, jwe_body)
-            jwekey,
-            jwe_body,
-            #[cfg(feature = "kms")]
-            kms_config,
-        )
             .await
             .change_context(errors::ApiErrorResponse::InternalServerError)
             .attach_printable("Error getting decrypted response payload for delete card")?;

crates/router/src/core/payment_methods/transformers.rs

@@ -2,8 +2,6 @@ use std::str::FromStr;
 
 use common_utils::{ext_traits::StringExt, pii::Email};
 use error_stack::ResultExt;
-#[cfg(feature = "kms")]
-use external_services::kms;
 use josekit::jwe;
 use serde::{Deserialize, Serialize};
 
@@ -151,24 +149,14 @@ pub fn get_dotted_jws(jws: encryption::JwsBody) -> String {
 }
 
 pub async fn get_decrypted_response_payload(
-    jwekey: &settings::Jwekey,
+    #[cfg(not(feature = "kms"))] jwekey: &settings::Jwekey,
+    #[cfg(feature = "kms")] jwekey: &settings::ActiveKmsSecrets,
     jwe_body: encryption::JweBody,
-    #[cfg(feature = "kms")] kms_config: &kms::KmsConfig,
 ) -> CustomResult<String, errors::VaultError> {
     #[cfg(feature = "kms")]
-    let public_key = kms::get_kms_client(kms_config)
+    let public_key = jwekey.jwekey.peek().vault_encryption_key.clone();
-        .await
-        .decrypt(&jwekey.vault_encryption_key)
-        .await
-        .change_context(errors::VaultError::SaveCardFailed)
-        .attach_printable("Fails to get public key of vault")?;
     #[cfg(feature = "kms")]
-    let private_key = kms::get_kms_client(kms_config)
+    let private_key = jwekey.jwekey.peek().vault_private_key.clone();
-        .await
-        .decrypt(&jwekey.vault_private_key)
-        .await
-        .change_context(errors::VaultError::SaveCardFailed)
-        .attach_printable("Error getting private key for signing jws")?;
 
     #[cfg(not(feature = "kms"))]
     let public_key = jwekey.vault_encryption_key.to_owned();
@@ -199,9 +187,9 @@ pub async fn get_decrypted_response_payload(
 }
 
 pub async fn mk_basilisk_req(
-    jwekey: &settings::Jwekey,
+    #[cfg(feature = "kms")] jwekey: &settings::ActiveKmsSecrets,
+    #[cfg(not(feature = "kms"))] jwekey: &settings::Jwekey,
     jws: &str,
-    #[cfg(feature = "kms")] kms_config: &kms::KmsConfig,
 ) -> CustomResult<encryption::JweBody, errors::VaultError> {
     let jws_payload: Vec<&str> = jws.split('.').collect();
 
@@ -219,12 +207,7 @@ pub async fn mk_basilisk_req(
         .change_context(errors::VaultError::SaveCardFailed)?;
 
     #[cfg(feature = "kms")]
-    let public_key = kms::get_kms_client(kms_config)
+    let public_key = jwekey.jwekey.peek().vault_encryption_key.clone();
-        .await
-        .decrypt(&jwekey.vault_encryption_key)
-        .await
-        .change_context(errors::VaultError::SaveCardFailed)
-        .attach_printable("Fails to get encryption key of vault")?;
 
     #[cfg(not(feature = "kms"))]
     let public_key = jwekey.vault_encryption_key.to_owned();
@@ -251,12 +234,12 @@ pub async fn mk_basilisk_req(
 }
 
 pub async fn mk_add_card_request_hs(
-    jwekey: &settings::Jwekey,
+    #[cfg(not(feature = "kms"))] jwekey: &settings::Jwekey,
+    #[cfg(feature = "kms")] jwekey: &settings::ActiveKmsSecrets,
     locker: &settings::Locker,
     card: &api::CardDetail,
     customer_id: &str,
     merchant_id: &str,
-    #[cfg(feature = "kms")] kms_config: &kms::KmsConfig,
 ) -> CustomResult<services::Request, errors::VaultError> {
     let merchant_customer_id = if cfg!(feature = "sandbox") {
         format!("{customer_id}::{merchant_id}")
@@ -281,12 +264,7 @@ pub async fn mk_add_card_request_hs(
         .change_context(errors::VaultError::RequestEncodingFailed)?;
 
     #[cfg(feature = "kms")]
-    let private_key = kms::get_kms_client(kms_config)
+    let private_key = jwekey.jwekey.peek().vault_private_key.clone();
-        .await
-        .decrypt(&jwekey.vault_private_key)
-        .await
-        .change_context(errors::VaultError::SaveCardFailed)
-        .attach_printable("Error getting private key for signing jws")?;
 
     #[cfg(not(feature = "kms"))]
     let private_key = jwekey.vault_private_key.to_owned();
@@ -295,13 +273,7 @@ pub async fn mk_add_card_request_hs(
         .await
         .change_context(errors::VaultError::RequestEncodingFailed)?;
 
-    let jwe_payload = mk_basilisk_req(
+    let jwe_payload = mk_basilisk_req(jwekey, &jws).await?;
-        jwekey,
-        &jws,
-        #[cfg(feature = "kms")]
-        kms_config,
-    )
-    .await?;
 
     let body = utils::Encode::<encryption::JweBody>::encode_to_value(&jwe_payload)
         .change_context(errors::VaultError::RequestEncodingFailed)?;
@@ -416,12 +388,12 @@ pub fn mk_add_card_request(
 }
 
 pub async fn mk_get_card_request_hs(
-    jwekey: &settings::Jwekey,
+    #[cfg(not(feature = "kms"))] jwekey: &settings::Jwekey,
+    #[cfg(feature = "kms")] jwekey: &settings::ActiveKmsSecrets,
     locker: &settings::Locker,
     customer_id: &str,
     merchant_id: &str,
     card_reference: &str,
-    #[cfg(feature = "kms")] kms_config: &kms::KmsConfig,
 ) -> CustomResult<services::Request, errors::VaultError> {
     let merchant_customer_id = if cfg!(feature = "sandbox") {
         format!("{customer_id}::{merchant_id}")
@@ -437,12 +409,7 @@ pub async fn mk_get_card_request_hs(
         .change_context(errors::VaultError::RequestEncodingFailed)?;
 
     #[cfg(feature = "kms")]
-    let private_key = kms::get_kms_client(kms_config)
+    let private_key = jwekey.jwekey.peek().vault_private_key.clone();
-        .await
-        .decrypt(&jwekey.vault_private_key)
-        .await
-        .change_context(errors::VaultError::SaveCardFailed)
-        .attach_printable("Error getting private key for signing jws")?;
 
     #[cfg(not(feature = "kms"))]
     let private_key = jwekey.vault_private_key.to_owned();
@@ -451,13 +418,7 @@ pub async fn mk_get_card_request_hs(
         .await
         .change_context(errors::VaultError::RequestEncodingFailed)?;
 
-    let jwe_payload = mk_basilisk_req(
+    let jwe_payload = mk_basilisk_req(jwekey, &jws).await?;
-        jwekey,
-        &jws,
-        #[cfg(feature = "kms")]
-        kms_config,
-    )
-    .await?;
 
     let body = utils::Encode::<encryption::JweBody>::encode_to_value(&jwe_payload)
         .change_context(errors::VaultError::RequestEncodingFailed)?;
@@ -508,12 +469,12 @@ pub fn mk_get_card_response(card: GetCardResponse) -> errors::RouterResult<Card>
 }
 
 pub async fn mk_delete_card_request_hs(
-    jwekey: &settings::Jwekey,
+    #[cfg(feature = "kms")] jwekey: &settings::ActiveKmsSecrets,
+    #[cfg(not(feature = "kms"))] jwekey: &settings::Jwekey,
     locker: &settings::Locker,
     customer_id: &str,
     merchant_id: &str,
     card_reference: &str,
-    #[cfg(feature = "kms")] kms_config: &kms::KmsConfig,
 ) -> CustomResult<services::Request, errors::VaultError> {
     let merchant_customer_id = if cfg!(feature = "sandbox") {
         format!("{customer_id}::{merchant_id}")
@@ -529,12 +490,7 @@ pub async fn mk_delete_card_request_hs(
         .change_context(errors::VaultError::RequestEncodingFailed)?;
 
     #[cfg(feature = "kms")]
-    let private_key = kms::get_kms_client(kms_config)
+    let private_key = jwekey.jwekey.peek().vault_private_key.clone();
-        .await
-        .decrypt(&jwekey.vault_private_key)
-        .await
-        .change_context(errors::VaultError::SaveCardFailed)
-        .attach_printable("Error getting private key for signing jws")?;
 
     #[cfg(not(feature = "kms"))]
     let private_key = jwekey.vault_private_key.to_owned();
@@ -543,13 +499,7 @@ pub async fn mk_delete_card_request_hs(
         .await
         .change_context(errors::VaultError::RequestEncodingFailed)?;
 
-    let jwe_payload = mk_basilisk_req(
+    let jwe_payload = mk_basilisk_req(jwekey, &jws).await?;
-        jwekey,
-        &jws,
-        #[cfg(feature = "kms")]
-        kms_config,
-    )
-    .await?;
 
     let body = utils::Encode::<encryption::JweBody>::encode_to_value(&jwe_payload)
         .change_context(errors::VaultError::RequestEncodingFailed)?;

crates/router/src/core/payment_methods/vault.rs

@@ -3,8 +3,6 @@ use common_utils::generate_id_with_default_len;
 use error_stack::report;
 use error_stack::{IntoReport, ResultExt};
 #[cfg(feature = "basilisk")]
-use external_services::kms;
-#[cfg(feature = "basilisk")]
 use josekit::jwe;
 use masking::PeekInterface;
 use router_env::{instrument, tracing};
@@ -470,11 +468,11 @@ pub fn get_key_id(keys: &settings::Jwekey) -> &str {
 
 #[cfg(feature = "basilisk")]
 async fn get_locker_jwe_keys(
-    keys: &settings::Jwekey,
+    keys: &settings::ActiveKmsSecrets,
-    kms_config: &kms::KmsConfig,
 ) -> CustomResult<(String, String), errors::EncryptionError> {
+    let keys = keys.jwekey.peek();
     let key_id = get_key_id(keys);
-    let (encryption_key, decryption_key) = if key_id == keys.locker_key_identifier1 {
+    let (public_key, private_key) = if key_id == keys.locker_key_identifier1 {
         (&keys.locker_encryption_key1, &keys.locker_decryption_key1)
     } else if key_id == keys.locker_key_identifier2 {
         (&keys.locker_encryption_key2, &keys.locker_decryption_key2)
@@ -482,18 +480,7 @@ async fn get_locker_jwe_keys(
         return Err(errors::EncryptionError.into());
     };
 
-    let public_key = kms::get_kms_client(kms_config)
+    Ok((public_key.to_string(), private_key.to_string()))
-        .await
-        .decrypt(encryption_key)
-        .await
-        .change_context(errors::EncryptionError)?;
-    let private_key = kms::get_kms_client(kms_config)
-        .await
-        .decrypt(decryption_key)
-        .await
-        .change_context(errors::EncryptionError)?;
-
-    Ok((public_key, private_key))
 }
 
 #[cfg(feature = "basilisk")]
@@ -515,7 +502,7 @@ pub async fn create_tokenize(
     )
     .change_context(errors::ApiErrorResponse::InternalServerError)?;
 
-    let (public_key, private_key) = get_locker_jwe_keys(&state.conf.jwekey, &state.conf.kms)
+    let (public_key, private_key) = get_locker_jwe_keys(&state.kms_secrets)
         .await
         .change_context(errors::ApiErrorResponse::InternalServerError)
         .attach_printable("Error getting Encryption key")?;
@@ -592,7 +579,7 @@ pub async fn get_tokenized_data(
     let payload = serde_json::to_string(&payload_to_be_encrypted)
         .map_err(|_x| errors::ApiErrorResponse::InternalServerError)?;
 
-    let (public_key, private_key) = get_locker_jwe_keys(&state.conf.jwekey, &state.conf.kms)
+    let (public_key, private_key) = get_locker_jwe_keys(&state.kms_secrets)
         .await
         .change_context(errors::ApiErrorResponse::InternalServerError)
         .attach_printable("Error getting Encryption key")?;
@@ -663,7 +650,7 @@ pub async fn delete_tokenized_data(
     let payload = serde_json::to_string(&payload_to_be_encrypted)
         .map_err(|_x| errors::ApiErrorResponse::InternalServerError)?;
 
-    let (public_key, _private_key) = get_locker_jwe_keys(&state.conf.jwekey, &state.conf.kms)
+    let (public_key, _private_key) = get_locker_jwe_keys(&state.kms_secrets)
         .await
         .change_context(errors::ApiErrorResponse::InternalServerError)
         .attach_printable("Error getting Encryption key")?;

crates/router/src/core/payments/helpers.rs

@@ -6,8 +6,6 @@ use common_utils::{
 };
 // TODO : Evaluate all the helper functions ()
 use error_stack::{report, IntoReport, ResultExt};
-#[cfg(feature = "kms")]
-use external_services::kms;
 use josekit::jwe;
 use masking::{ExposeOptionInterface, PeekInterface};
 use router_env::{instrument, tracing};
@@ -1607,15 +1605,7 @@ pub async fn get_merchant_connector_account(
                 )?;
 
             #[cfg(feature = "kms")]
-            let kms_config = &state.conf.kms;
+            let private_key = state.kms_secrets.jwekey.peek().tunnel_private_key.clone();
-
-            #[cfg(feature = "kms")]
-            let private_key = kms::get_kms_client(kms_config)
-                .await
-                .decrypt(state.conf.jwekey.tunnel_private_key.to_owned())
-                .await
-                .change_context(errors::ApiErrorResponse::InternalServerError)
-                .attach_printable("Error getting tunnel private key")?;
 
             #[cfg(not(feature = "kms"))]
             let private_key = state.conf.jwekey.tunnel_private_key.to_owned();

crates/router/src/routes/app.rs

@@ -12,8 +12,10 @@ use super::{admin::*, api_keys::*, disputes::*, files::*};
 use super::{configs::*, customers::*, mandates::*, payments::*, payouts::*, refunds::*};
 #[cfg(feature = "oltp")]
 use super::{ephemeral_key::*, payment_methods::*, webhooks::*};
+#[cfg(feature = "kms")]
+use crate::configs::kms;
 use crate::{
-    configs::settings::Settings,
+    configs::settings,
     db::{MockDb, StorageImpl, StorageInterface},
     routes::cards_info::card_iin_info,
     services::Store,
@@ -23,13 +25,15 @@ use crate::{
 pub struct AppState {
     pub flow_name: String,
     pub store: Box<dyn StorageInterface>,
-    pub conf: Settings,
+    pub conf: settings::Settings,
     #[cfg(feature = "email")]
     pub email_client: Box<dyn EmailClient>,
+    #[cfg(feature = "kms")]
+    pub kms_secrets: settings::ActiveKmsSecrets,
 }
 
 pub trait AppStateInfo {
-    fn conf(&self) -> Settings;
+    fn conf(&self) -> settings::Settings;
     fn flow_name(&self) -> String;
     fn store(&self) -> Box<dyn StorageInterface>;
     #[cfg(feature = "email")]
@@ -37,7 +41,7 @@ pub trait AppStateInfo {
 }
 
 impl AppStateInfo for AppState {
-    fn conf(&self) -> Settings {
+    fn conf(&self) -> settings::Settings {
         self.conf.to_owned()
     }
     fn flow_name(&self) -> String {
@@ -54,7 +58,7 @@ impl AppStateInfo for AppState {
 
 impl AppState {
     pub async fn with_storage(
-        conf: Settings,
+        conf: settings::Settings,
         storage_impl: StorageImpl,
         shut_down_signal: oneshot::Sender<()>,
     ) -> Self {
@@ -66,6 +70,17 @@ impl AppState {
             StorageImpl::Mock => Box::new(MockDb::new(&conf).await),
         };
 
+        #[cfg(feature = "kms")]
+        #[allow(clippy::expect_used)]
+        let kms_secrets = kms::KmsDecrypt::decrypt_inner(
+            settings::ActiveKmsSecrets {
+                jwekey: conf.jwekey.clone().into(),
+            },
+            &conf.kms,
+        )
+        .await
+        .expect("Failed while performing KMS decryption");
+
         #[cfg(feature = "email")]
         #[allow(clippy::expect_used)]
         let email_client = Box::new(AwsSes::new(&conf.email).await);
@@ -75,10 +90,12 @@ impl AppState {
             conf,
             #[cfg(feature = "email")]
             email_client,
+            #[cfg(feature = "kms")]
+            kms_secrets,
         }
     }
 
-    pub async fn new(conf: Settings, shut_down_signal: oneshot::Sender<()>) -> Self {
+    pub async fn new(conf: settings::Settings, shut_down_signal: oneshot::Sender<()>) -> Self {
         Self::with_storage(conf, StorageImpl::Postgresql, shut_down_signal).await
     }
 }