feat(custodian): add support for authorization for accessing keys (#24)

Author: NishantJoshi00

Cargo.lock

@@ -53,6 +53,7 @@ rustc-hash = "1.1.0"
 rayon = "1.10.0"
 once_cell = "1.19.0"
 hyper = "1.3.1"
+blake3 = "1.5.4"
 
 [build-dependencies]
 cargo_metadata = "0.18.1"

Cargo.toml

@@ -24,4 +24,5 @@ log_format = "console"
 
 [secrets]
 master_key = "6d761d32f1b14ef34cf016d726b29b02b5cfce92a8959f1bfb65995c8100925e"
-
+access_token = "secret123"
+hash_context = "keymanager:hyperswitch"

config/development.toml

@@ -0,0 +1,2 @@
+-- This file should undo anything in `up.sql`
+ALTER TABLE data_key_store DROP COLUMN IF EXISTS token;

migrations/2024-09-30-125631_add-token-column/down.sql

@@ -0,0 +1,2 @@
+-- Your SQL goes here
+ALTER TABLE data_key_store ADD COLUMN IF NOT EXISTS token VARCHAR(255);

migrations/2024-09-30-125631_add-token-column/up.sql

@@ -128,6 +128,8 @@ pub struct Secrets {
     pub master_key: GcmAes256,
     #[cfg(feature = "aws")]
     pub kms_config: AwsKmsConfig,
+    pub access_token: masking::Secret<String>,
+    pub hash_context: masking::Secret<String>,
 }
 
 #[derive(Deserialize, Debug)]

src/config.rs

@@ -1,4 +1,5 @@
 mod crux;
+pub(crate) mod custodian;
 mod decryption;
 mod encryption;
 
@@ -17,14 +18,17 @@ use axum::extract::{Json, State};
 use opentelemetry::KeyValue;
 use std::sync::Arc;
 
+use self::custodian::Custodian;
+
 pub async fn encrypt_data(
     State(state): State<Arc<AppState>>,
+    custodian: Custodian,
     Json(req): Json<EncryptDataRequest>,
 ) -> errors::ApiResponseResult<Json<EncryptionResponse>> {
     let (data_identifier, key_identifier) = req.identifier.get_identifier();
 
     utils::record_api_operation(
-        encryption::encryption(state, req),
+        encryption::encryption(state, custodian, req),
         &metrics::ENCRYPTION_API_LATENCY,
         &[
             KeyValue::new("data_identifier", data_identifier),
@@ -36,12 +40,13 @@ pub async fn encrypt_data(
 
 pub async fn decrypt_data(
     State(state): State<Arc<AppState>>,
+    custodian: Custodian,
     Json(req): Json<DecryptionRequest>,
 ) -> errors::ApiResponseResult<Json<DecryptionResponse>> {
     let (data_identifier, key_identifier) = req.identifier.get_identifier();
 
     utils::record_api_operation(
-        decryption::decryption(state, req),
+        decryption::decryption(state, custodian, req),
         &metrics::DECRYPTION_API_LATENCY,
         &[
             KeyValue::new("data_identifier", data_identifier),

src/core/crypto.rs

@@ -1,3 +1,4 @@
+use error_stack::ensure;
 use masking::PeekInterface;
 use rayon::prelude::*;
 
@@ -15,6 +16,8 @@ use crate::{
     },
 };
 
+use super::custodian::Custodian;
+
 #[async_trait::async_trait]
 pub trait KeyEncrypter<ToType> {
     async fn encrypt(self, state: &AppState) -> errors::CustomResult<ToType, errors::CryptoError>;
@@ -47,6 +50,7 @@ impl KeyEncrypter<DataKeyNew> for Key {
                 time::OffsetDateTime::now_utc().date(),
                 time::OffsetDateTime::now_utc().time(),
             ),
+            token: self.token,
         })
     }
 }
@@ -71,6 +75,7 @@ impl KeyDecrypter<Key> for DataKey {
             version: self.version,
             key: decrypted_key.into(),
             source,
+            token: self.token,
         })
     }
 }
@@ -81,6 +86,7 @@ pub trait DataEncrypter<ToType> {
         self,
         state: &AppState,
         identifier: &Identifier,
+        custodian: Custodian,
     ) -> errors::CustomResult<ToType, errors::CryptoError>;
 }
 
@@ -90,6 +96,7 @@ pub trait DataDecrypter<ToType> {
         self,
         state: &AppState,
         identifier: &Identifier,
+        custodian: Custodian,
     ) -> errors::CustomResult<ToType, errors::CryptoError>;
 }
 
@@ -99,11 +106,20 @@ impl DataEncrypter<EncryptedDataGroup> for DecryptedDataGroup {
         self,
         state: &AppState,
         identifier: &Identifier,
+        custodian: Custodian,
     ) -> errors::CustomResult<EncryptedDataGroup, errors::CryptoError> {
         let version = Version::get_latest(identifier, state).await;
         let decrypted_key = Key::get_key(state, identifier, version).await.switch()?;
         let key = GcmAes256::new(decrypted_key.key)?;
 
+        let stored_token = decrypted_key.token;
+        let provided_token = custodian.into_access_token(state);
+
+        ensure!(
+            !identifier.is_entity() || (stored_token.eq(&provided_token)),
+            errors::CryptoError::AuthenticationFailed
+        );
+
         state.thread_pool.install(|| {
             self.0
                 .into_par_iter()
@@ -125,12 +141,21 @@ impl DataDecrypter<DecryptedDataGroup> for EncryptedDataGroup {
         self,
         state: &AppState,
         identifier: &Identifier,
+        custodian: Custodian,
     ) -> errors::CustomResult<DecryptedDataGroup, errors::CryptoError> {
         let version = FxHashSet::from_iter(self.0.values().map(|d| d.version));
         let decrypted_keys = Key::get_multiple_keys(state, identifier, version)
             .await
             .switch()?;
 
+        let mut stored_tokens = decrypted_keys.values().map(|k| &k.token);
+        let provided_token = custodian.into_access_token(state);
+
+        ensure!(
+            !identifier.is_entity() || stored_tokens.all(|t| t.eq(&provided_token)),
+            errors::CryptoError::AuthenticationFailed
+        );
+
         state.thread_pool.install(|| {
             self
             .0
@@ -160,9 +185,19 @@ impl DataEncrypter<EncryptedData> for DecryptedData {
         self,
         state: &AppState,
         identifier: &Identifier,
+        custodian: Custodian,
     ) -> errors::CustomResult<EncryptedData, errors::CryptoError> {
         let version = Version::get_latest(identifier, state).await;
         let decrypted_key = Key::get_key(state, identifier, version).await.switch()?;
+
+        let stored_token = decrypted_key.token;
+        let provided_token = custodian.into_access_token(state);
+
+        ensure!(
+            !identifier.is_entity() || (stored_token.eq(&provided_token)),
+            errors::CryptoError::AuthenticationFailed
+        );
+
         let key = GcmAes256::new(decrypted_key.key)?;
 
         let encrypted_data = key.encrypt(self.inner())?;
@@ -180,9 +215,19 @@ impl DataDecrypter<DecryptedData> for EncryptedData {
         self,
         state: &AppState,
         identifier: &Identifier,
+        custodian: Custodian,
     ) -> errors::CustomResult<DecryptedData, errors::CryptoError> {
         let version = self.version;
         let decrypted_key = Key::get_key(state, identifier, version).await.switch()?;
+
+        let stored_token = decrypted_key.token;
+        let provided_token = custodian.into_access_token(state);
+
+        ensure!(
+            !identifier.is_entity() || (stored_token.eq(&provided_token)),
+            errors::CryptoError::AuthenticationFailed
+        );
+
         let key = GcmAes256::new(decrypted_key.key)?;
 
         let decrypted_data = key.decrypt(self.inner())?;

src/core/crypto/crux.rs

@@ -0,0 +1,87 @@
+use std::sync::Arc;
+
+use axum::extract::FromRequestParts;
+use axum::http::request;
+use base64::Engine;
+use error_stack::{ensure, ResultExt};
+use hyper::header;
+use masking::{PeekInterface, Secret, StrongSecret};
+
+use crate::app::AppState;
+use crate::consts::base64::BASE64_ENGINE;
+use crate::errors::{ApiErrorContainer, CustomResult, ParsingError, SwitchError, ToContainerError};
+
+pub struct Custodian {
+    pub keys: Option<(StrongSecret<String>, StrongSecret<String>)>,
+}
+
+impl Custodian {
+    fn new(keys: Option<(String, String)>) -> Self {
+        let keys = keys.map(|(key1, key2)| (StrongSecret::new(key1), StrongSecret::new(key2)));
+        Self { keys }
+    }
+
+    pub fn into_access_token(self, state: &AppState) -> Option<StrongSecret<String>> {
+        self.keys
+            .map(|(x, y)| format!("{}:{}", x.peek(), y.peek()))
+            .map(|key| crate::crypto::blake3::Blake3::hash(state, Secret::new(key)))
+            .map(hex::encode)
+            .map(StrongSecret::new)
+    }
+}
+
+#[axum::async_trait]
+impl FromRequestParts<Arc<AppState>> for Custodian {
+    type Rejection = ApiErrorContainer;
+    async fn from_request_parts(
+        parts: &mut request::Parts,
+        _state: &Arc<AppState>,
+    ) -> Result<Self, Self::Rejection> {
+        parts
+            .headers
+            .get(header::AUTHORIZATION)
+            .map(extract_credential)
+            .transpose()
+            .switch()
+            .to_container_error()
+            .map(Self::new)
+    }
+}
+
+fn extract_credential(
+    header: &header::HeaderValue,
+) -> CustomResult<(String, String), ParsingError> {
+    let header = header.to_str().change_context(ParsingError::ParsingFailed(
+        "Failed while converting header to string".to_string(),
+    ))?;
+
+    let credential = header
+        .strip_prefix("Basic ")
+        .ok_or(ParsingError::ParsingFailed(
+            "Authorization scheme is not basic".to_string(),
+        ))?;
+    let credential = credential.trim();
+    let credential =
+        BASE64_ENGINE
+            .decode(credential)
+            .change_context(ParsingError::DecodingFailed(
+                "Failed while base64 decoding the authorization header".to_string(),
+            ))?;
+    let credential = String::from_utf8(credential).change_context(ParsingError::DecodingFailed(
+        "Failed while converting base64 to utf8".to_string(),
+    ))?;
+    let mut parts = credential.split(':');
+    let key1 = parts.next().ok_or(ParsingError::ParsingFailed(
+        "Failed while extracting key1 from credential".to_string(),
+    ))?;
+    let key2 = parts.next().ok_or(ParsingError::ParsingFailed(
+        "Failed while extracting key2 from credential".to_string(),
+    ))?;
+
+    ensure!(
+        parts.next().is_none(),
+        ParsingError::ParsingFailed("Credential has more than 2 parts".to_string())
+    );
+
+    Ok((key1.to_string(), key2.to_string()))
+}

src/core/crypto/custodian.rs

@@ -9,14 +9,17 @@ use crate::{
 };
 use opentelemetry::KeyValue;
 
+use super::custodian::Custodian;
+
 pub(super) async fn decryption(
     state: Arc<AppState>,
+    custodian: Custodian,
     req: DecryptionRequest,
 ) -> errors::CustomResult<DecryptionResponse, errors::ApplicationErrorResponse> {
     let identifier = req.identifier.clone();
     let decrypted_data = req
         .data
-        .decrypt(&state, &identifier)
+        .decrypt(&state, &identifier, custodian)
         .await
         .map_err(|err| {
             logger::error!(encryption_error=?err);