Merge https://github.com/juspay/tartarus into benches

Cargo.toml

@@ -20,22 +20,26 @@ aws-config = { version = "1.0.1", optional = true }
 aws-sdk-kms = { version = "1.3.0", optional = true }
 base64 = "0.21.2"
 futures = "0.3.28"
-tracing = { version = "0.1.40" }
-tracing-appender = { version = "0.2.2" }
-tracing-attributes = "0.1.27"
-tracing-subscriber = { version = "0.3.17", default-features = true, features = ["env-filter", "json", "registry"] }
 gethostname = "0.4.3"
 rustc-hash = "1.1"
 once_cell = "1.18.0"
 vaultrs = { version = "0.7.0", optional = true }
 
 # Tokio Dependencies
 tokio = { version = "1.33.0", features = ["macros", "rt-multi-thread"] }
-axum = "0.6.20"
-hyper = "0.14.27"
+axum = "0.7.3"
+hyper = "1.0.1"
 tower = { version = "0.4.13", features = ["limit", "buffer", "load-shed"] }
-tower-http = { version = "0.4.4", features = ["trace"] }
-
+tower-http = { version = "0.5.0", features = ["trace"] }
+tracing = { version = "0.1.40" }
+tracing-appender = { version = "0.2.2" }
+tracing-attributes = "0.1.27"
+tracing-subscriber = { version = "0.3.17", default-features = true, features = [
+    "env-filter",
+    "json",
+    "registry",
+] }
+http-body-util = "0.1.0"
 
 diesel = { version = "2.1.3", features = ["postgres", "serde_json", "time"] }
 diesel-async = { version = "0.4.1", features = ["postgres", "deadpool"] }
@@ -61,8 +65,8 @@ argh = "0.1.12"
 
 [dev-dependencies]
 rand = "0.8.5"
-axum-test = "13.0.1"
 criterion = "0.5.1"
+axum-test = "14.2.2"
 
 [build-dependencies]
 cargo_metadata = "0.15.4"

src/app.rs

@@ -1,6 +1,5 @@
 use axum::routing;
 use error_stack::ResultExt;
-use hyper::server::conn;
 use masking::PeekInterface;
 #[cfg(feature = "key_custodian")]
 use tokio::sync::{mpsc::Sender, RwLock};
@@ -52,7 +51,7 @@ pub async fn server1_builder(
     state: Arc<RwLock<AppState>>,
     server_tx: Sender<()>,
 ) -> Result<
-    hyper::Server<conn::AddrIncoming, routing::IntoMakeService<axum::Router>>,
+    axum::serve::Serve<routing::IntoMakeService<axum::Router>, axum::Router>,
     error::ConfigurationError,
 >
 where
@@ -71,7 +70,9 @@ where
         .with_state(shared_state)
         .route("/health", routing::get(routes::health::health));
 
-    let server = axum::Server::try_bind(&socket_addr)?.serve(router.into_make_service());
+    let tcp_listener = tokio::net::TcpListener::bind(&socket_addr).await?;
+    let server = axum::serve(tcp_listener, router.into_make_service());
+
     Ok(server)
 }
 
@@ -82,7 +83,7 @@ where
 pub async fn server2_builder(
     state: &AppState,
 ) -> Result<
-    hyper::Server<conn::AddrIncoming, routing::IntoMakeService<axum::Router>>,
+    axum::serve::Serve<routing::IntoMakeService<axum::Router>, axum::Router>,
     error::ConfigurationError,
 >
 where
@@ -121,8 +122,8 @@ where
                         .level(tracing::Level::ERROR),
                 ),
         );
-
-    let server = axum::Server::try_bind(&socket_addr)?.serve(router.into_make_service());
+    let tcp_listener = tokio::net::TcpListener::bind(&socket_addr).await?;
+    let server = axum::serve(tcp_listener, router.into_make_service());
     Ok(server)
 }
 

src/bin/locker.rs

@@ -19,11 +19,11 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     {
         let state_lock = state.clone();
 
-        let (server1_tx, mut server1_rx) = tokio::sync::mpsc::channel::<()>(1);
+        let (server1_tx, server1_rx) = tokio::sync::mpsc::channel::<()>(1);
 
         let server1 = tartarus::app::server1_builder(state_lock, server1_tx.clone())
             .await?
-            .with_graceful_shutdown(graceful_shutdown_server1(&mut server1_rx));
+            .with_graceful_shutdown(graceful_shutdown_server1(server1_rx));
 
         logger::info!(
             "Key Custodian started [{:?}] [{:?}]",
@@ -42,7 +42,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 }
 
 #[cfg(feature = "key_custodian")]
-async fn graceful_shutdown_server1(recv: &mut tokio::sync::mpsc::Receiver<()>) {
+async fn graceful_shutdown_server1(mut recv: tokio::sync::mpsc::Receiver<()>) {
     recv.recv().await;
     logger::info!("Shutting down the server1 gracefully.");
 }

src/bin/utils.rs

@@ -80,8 +80,13 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
             let pub_key = read_file_to_string(
                 &public_key.ok_or(error::CryptoError::InvalidData("public key not found"))?,
             )?;
-            jwe_operation(|x| {
-                JWEncryption::new(priv_key, pub_key, jwe::RSA_OAEP_256, jwe::RSA_OAEP).encrypt(x)
+            jwe_operation(|payload| {
+                JWEncryption::new(priv_key, pub_key, jwe::RSA_OAEP_256, jwe::RSA_OAEP)
+                    .encrypt(payload)
+                    .and_then(|payload| {
+                        Ok(serde_json::to_vec(&payload)
+                            .map_err(error::CryptoError::SerdeJsonError)?)
+                    })
             })?;
         }
         SubCommand::JweDecrypt(JweD {
@@ -94,8 +99,15 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
             let pub_key = read_file_to_string(
                 &public_key.ok_or(error::CryptoError::InvalidData("private key not found"))?,
             )?;
-            jwe_operation(|x| {
-                JWEncryption::new(priv_key, pub_key, jwe::RSA_OAEP_256, jwe::RSA_OAEP).decrypt(x)
+            jwe_operation(|payload| {
+                serde_json::from_slice(&payload)
+                    .map_err(error::CryptoError::SerdeJsonError)
+                    .map_err(Into::into)
+                    .and_then(|payload| {
+                        JWEncryption::new(priv_key, pub_key, jwe::RSA_OAEP_256, jwe::RSA_OAEP)
+                            .decrypt(payload)
+                    })
+                // (x)
             })?;
         }
     }

src/crypto/jw.rs

@@ -53,7 +53,7 @@ impl JwsBody {
     }
 }
 
-#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
 #[serde(rename_all = "camelCase")]
 pub struct JweBody {
     pub header: String,
@@ -90,10 +90,10 @@ impl JweBody {
     }
 }
 
-impl super::Encryption<Vec<u8>, Vec<u8>> for JWEncryption {
+impl super::Encryption<Vec<u8>, JweBody> for JWEncryption {
     type ReturnType<'a, T> = Result<T, ContainerError<error::CryptoError>>;
 
-    fn encrypt(&self, input: Vec<u8>) -> Self::ReturnType<'_, Vec<u8>> {
+    fn encrypt(&self, input: Vec<u8>) -> Self::ReturnType<'_, JweBody> {
         let payload = input;
         let jws_encoded = jws_sign_payload(&payload, self.private_key.peek().as_bytes())?;
         let jws_body = JwsBody::from_dotted_str(&jws_encoded).ok_or(
@@ -107,13 +107,11 @@ impl super::Encryption<Vec<u8>, Vec<u8>> for JWEncryption {
         )?;
         let jwe_body = JweBody::from_str(&jwe_encrypted)
             .ok_or(error::CryptoError::InvalidData("JWE data incomplete"))?;
-        Ok(serde_json::to_vec(&jwe_body).map_err(error::CryptoError::from)?)
+        Ok(jwe_body)
     }
 
-    fn decrypt(&self, input: Vec<u8>) -> Self::ReturnType<'_, Vec<u8>> {
-        let jwe_body: JweBody = serde_json::from_slice(&input).map_err(error::CryptoError::from)?;
-        let jwe_encoded = jwe_body.get_dotted_jwe();
-        // let algo = jwe::RSA_OAEP_256;
+    fn decrypt(&self, input: JweBody) -> Self::ReturnType<'_, Vec<u8>> {
+        let jwe_encoded = input.get_dotted_jwe();
         let jwe_decrypted =
             decrypt_jwe(&jwe_encoded, self.private_key.peek(), self.decryption_algo)?;
 

src/error.rs

@@ -15,6 +15,8 @@ pub enum ConfigurationError {
     ServerError(#[from] hyper::Error),
     #[error("invalid host for socket")]
     AddressError(#[from] std::net::AddrParseError),
+    #[error("invalid host for socket")]
+    IOError(#[from] std::io::Error),
     #[error("Error while connecting/creating database pool")]
     DatabaseError,
     #[error("Failed to KMS decrypt: {0}")]

src/middleware.rs

@@ -1,63 +1,50 @@
 use crate::app::AppState;
-use crate::crypto::jw::JWEncryption;
+use crate::crypto::jw::{self, JWEncryption};
 use crate::crypto::Encryption;
 use crate::error::{self, ContainerError, ResultContainerExt};
-use axum::{
-    body::BoxBody,
-    extract,
-    http::{Request, Response},
-    middleware::Next,
-};
-use error_stack::ResultExt;
-use hyper::body::HttpBody;
-use hyper::Body;
+use axum::body::Body;
+use axum::http::{request, response};
+use axum::{extract, http::Request, middleware::Next};
+
+use http_body_util::BodyExt;
 use josekit::jwe;
 
 /// Middleware providing implementation to perform JWE + JWS encryption and decryption around the
 /// card APIs
 pub async fn middleware(
     extract::State(state): extract::State<AppState>,
-    request: Request<Body>,
-    next: Next<Body>,
-) -> Result<Response<BoxBody>, ContainerError<error::ApiError>> {
-    let (parts, body) = request.into_parts();
-
-    let request_body =
-        hyper::body::to_bytes(body)
-            .await
-            .change_error(error::ApiError::RequestMiddlewareError(
-                "Failed to read request body for jwe decryption",
-            ))?;
-
+    parts: request::Parts,
+    axum::Json(jwe_body): axum::Json<jw::JweBody>,
+    next: Next,
+) -> Result<(response::Parts, axum::Json<jw::JweBody>), ContainerError<error::ApiError>> {
     let keys = JWEncryption {
         private_key: state.config.secrets.locker_private_key,
         public_key: state.config.secrets.tenant_public_key,
         encryption_algo: jwe::RSA_OAEP,
         decryption_algo: jwe::RSA_OAEP_256,
     };
 
-    let jwe_decrypted = keys.decrypt(request_body.to_vec())?;
+    let jwe_decrypted = keys.decrypt(jwe_body)?;
 
     let next_layer_payload = Request::from_parts(parts, Body::from(jwe_decrypted));
 
-    let response = next.run(next_layer_payload).await;
-
-    let (parts, body) = response.into_parts();
+    let (mut parts, body) = next.run(next_layer_payload).await.into_parts();
 
-    let response_body = hyper::body::to_bytes(body).await.change_error(
-        error::ApiError::ResponseMiddlewareError("Failed to read response body for jws signing"),
-    )?;
+    let response_body = body
+        .collect()
+        .await
+        .change_error(error::ApiError::ResponseMiddlewareError(
+            "Failed to read response body for jws signing",
+        ))?
+        .to_bytes();
 
-    let jws_signed = keys.encrypt(response_body.to_vec())?;
+    let jwe_payload = keys.encrypt(response_body.to_vec())?;
 
-    let jwt = String::from_utf8(jws_signed).change_error(
-        error::ApiError::ResponseMiddlewareError("Could not convert to UTF-8"),
-    )?;
+    parts.headers = hyper::HeaderMap::new();
+    parts.headers.append(
+        hyper::header::CONTENT_TYPE,
+        axum::http::HeaderValue::from_static("application/json"),
+    );
 
-    Ok(axum::http::response::Builder::new()
-        .status(parts.status)
-        .body(jwt.map_err(axum::Error::new).boxed_unsync())
-        .change_context(error::ApiError::ResponseMiddlewareError(
-            "failed while generating the response",
-        ))?)
+    Ok((parts, axum::Json(jwe_payload)))
 }

src/routes/data.rs

@@ -8,6 +8,8 @@ use axum::middleware;
 
 use masking::ExposeInterface;
 
+use types::StoreCardResponse;
+
 use crate::{
     app::AppState,
     crypto::{aes::GcmAes256, sha::Sha512},
@@ -94,7 +96,7 @@ pub async fn add_card(
 
     let optional_hash_table = state.db.find_by_data_hash(&hash_data).await?;
 
-    let output = match optional_hash_table {
+    let (duplication_check, output) = match optional_hash_table {
         Some(hash_table) => {
             let stored_data = state
                 .db
@@ -107,7 +109,10 @@ pub async fn add_card(
                 )
                 .await?;
 
-            match stored_data {
+            let duplication_check =
+                transformers::validate_card_metadata(stored_data.as_ref(), &request.data)?;
+
+            let output = match stored_data {
                 Some(data) => data,
                 None => {
                     state
@@ -123,12 +128,14 @@ pub async fn add_card(
                         )
                         .await?
                 }
-            }
+            };
+
+            (duplication_check, output)
         }
         None => {
             let hash_table = state.db.insert_hash(hash_data).await?;
 
-            state
+            let output = state
                 .db
                 .insert_or_get_from_locker(
                     (
@@ -139,11 +146,13 @@ pub async fn add_card(
                         .try_into()?,
                     &merchant_dek,
                 )
-                .await?
+                .await?;
+
+            (None, output)
         }
     };
 
-    Ok(Json(output.into()))
+    Ok(Json(StoreCardResponse::from((duplication_check, output))))
 }
 
 /// `/data/delete` handling the requirement of deleting cards