fix(): code refinement

adorsys · Sep 25, 2024 · c74773e · c74773e
1 parent 513d476
commit c74773e
Show file tree

Hide file tree

Showing 7 changed files with 107 additions and 56 deletions.
diff --git a/crates/plugins/mediator-coordination/src/forward/error.rs b/crates/plugins/mediator-coordination/src/forward/error.rs
@@ -0,0 +1,25 @@
+use axum::Json;
+use serde_json::{json, Value};
+use thiserror::Error;
+
+#[derive(Debug, Error)]
+pub enum RoutingError {
+    #[error("message body is malformed")]
+    MalformedBody,
+    #[error("Repository not set")]
+    RepostitoryError
+}
+impl RoutingError {
+    /// Converts the error to an axum JSON representation.
+    pub fn json(&self) -> Json<Value> {
+        Json(json!({
+            "error": self.to_string()
+        }))
+    }
+}
+
+impl From<RoutingError> for Json<Value> {
+    fn from(error: RoutingError) -> Self {
+        error.json()
+    }
+}
diff --git a/crates/plugins/mediator-coordination/src/forward/mod.rs b/crates/plugins/mediator-coordination/src/forward/mod.rs
@@ -1 +1,2 @@
+mod error;
 pub mod routing;
diff --git a/crates/plugins/mediator-coordination/src/forward/routing.rs b/crates/plugins/mediator-coordination/src/forward/routing.rs
@@ -1,38 +1,48 @@
+use std::sync::Arc;
+
 use axum::response::{IntoResponse, Response};
 
+use database::Repository;
 use didcomm::{AttachmentData, Message};
 use hyper::StatusCode;
 use mongodb::bson::doc;
 use serde_json::{json, Value};
 
 use crate::{
-
-    model::stateful::entity::RoutedMessage, web::{error::MediationError, AppState, AppStateRepository}
+    model::stateful::entity::{Connection, RoutedMessage},
+    web::{error::MediationError, AppState, AppStateRepository},
 };
 
-/// mediator receives messages of type forward then it unpacks the messages and stores it for pickup
-/// the unpacked message is then repacked for further transmission.
+use super::error::RoutingError;
+
+/// Mediator receives forwarded messages, extract the next field in the message body, and the attachments in the message
+/// then stores the attachment with the next field as key for pickup
 pub async fn mediator_forward_process(
     state: &AppState,
     payload: Message,
 ) -> Result<Message, Response> {
-    let AppStateRepository {
-        message_repository,
-        connection_repository,
-        ..
-    } = state
-        .repository
-        .as_ref()
-        .ok_or_else(|| MediationError::RepostitoryError)
-        .unwrap();
+
+    let result = handler(state, payload).await.unwrap();
+    Ok(result)
+}
 
+async fn checks(
+    message: &Message,
+    connection_repository: &Arc<dyn Repository<Connection>>,
+) -> Result<String, Response> {
+    let next = message.body.get("next").and_then(Value::as_str);
+    match next {
+        Some(next) => next,
+        None => {
+            let response = (StatusCode::BAD_REQUEST, RoutingError::MalformedBody.json());
+            return Err(response.into_response());
+        }
+    };
     // Check if the client's did in mediator's keylist
-
-    let next = payload.body.get("next").and_then(Value::as_str).unwrap();
-
     let _connection = match connection_repository
         .find_one_by(doc! {"keylist": doc!{ "$elemMatch": { "$eq": &next}}})
         .await
+        .map_err(|_| MediationError::RepostitoryError)
         .unwrap()
     {
         Some(connection) => connection,
@@ -44,34 +54,51 @@ pub async fn mediator_forward_process(
             return Err(response.into_response());
         }
     };
+    Ok(next.unwrap().to_string())
+}
 
-    let attachments = payload.attachments.unwrap_or_default();
-    for att in attachments {
-        let attached = match att.data {
-            AttachmentData::Json { value: val } => val.json,
-            _ => json!(0),
+async fn handler(state: &AppState, message: Message) -> Result<Message, MediationError> {
+    let AppStateRepository {
+        message_repository,
+        connection_repository,
+        ..
+    } = state
+        .repository
+        .as_ref()
+        .ok_or_else(|| MediationError::RepostitoryError)?;
+    let next = checks(&message, connection_repository).await.unwrap();
+
+    let attachments = message.attachments.unwrap_or_default();
+    for attachment in attachments {
+        let attached = match attachment.data {
+            AttachmentData::Json { value: data } => data.json,
+            AttachmentData::Base64 { value: data } => json!(data.base64),
+            AttachmentData::Links { value: data } => json!(data.links),
         };
         message_repository
             .store(RoutedMessage {
                 id: None,
                 message: json!(attached),
-                recipient_did: next.to_string(),
+                recipient_did: next.clone(),
             })
             .await
-            .map_err(|_| MediationError::PersisenceError)
-            .unwrap();
+            .map_err(|_| MediationError::PersisenceError)?;
     }
-
     Ok(Message::build("".to_string(), "".to_string(), json!("")).finalize())
 }
-
 #[cfg(test)]
 mod test {
 
     use std::sync::Arc;
 
     use crate::{
-        didcomm::bridge::LocalSecretsResolver, model::stateful::entity::Connection, repository::stateful::tests::{MockConnectionRepository, MockMessagesRepository, MockSecretsRepository}, util::{self, MockFileSystem}, web::AppStateRepository
+        didcomm::bridge::LocalSecretsResolver,
+        model::stateful::entity::Connection,
+        repository::stateful::tests::{
+            MockConnectionRepository, MockMessagesRepository, MockSecretsRepository,
+        },
+        util::{self, MockFileSystem},
+        web::AppStateRepository,
     };
 
     use super::*;

diff --git a/crates/plugins/mediator-coordination/src/web/coord/handler/stateful.rs b/crates/plugins/mediator-coordination/src/web/coord/handler/stateful.rs
@@ -345,8 +345,8 @@ pub async fn process_plain_keylist_update_message(
                 updated: confirmations
             }),
         )
-        .to(sender.clone())
-        .from(mediator_did.clone())
+        .to(sender)
+        .from(mediator_did.to_owned())
         .finalize(),
     )
 }

diff --git a/crates/plugins/mediator-coordination/src/web/coord/handler/stateless.rs b/crates/plugins/mediator-coordination/src/web/coord/handler/stateless.rs
@@ -157,7 +157,7 @@ mod tests {
         let response = app
             .oneshot(
                 Request::builder()
-                    .uri(String::from("/mediate"))
+                    .uri(String::from("/"))
                     .method(Method::POST)
                     .header(CONTENT_TYPE, DIDCOMM_ENCRYPTED_MIME_TYPE)
                     .body(Body::from(packed_msg))
@@ -245,7 +245,7 @@ mod tests {
                 .clone()
                 .oneshot(
                     Request::builder()
-                        .uri(String::from("/mediate"))
+                        .uri(String::from("/"))
                         .method(Method::POST)
                         .header(CONTENT_TYPE, DIDCOMM_ENCRYPTED_MIME_TYPE)
                         .body(Body::from(packed_msg))
@@ -312,7 +312,7 @@ mod tests {
         let response = app
             .oneshot(
                 Request::builder()
-                    .uri(String::from("/mediate"))
+                    .uri(String::from("/"))
                     .method(Method::POST)
                     .header(CONTENT_TYPE, DIDCOMM_ENCRYPTED_MIME_TYPE)
                     .body(Body::from(packed_msg))
@@ -375,7 +375,7 @@ mod tests {
         let response = app
             .oneshot(
                 Request::builder()
-                    .uri(String::from("/mediate"))
+                    .uri(String::from("/"))
                     .method(Method::POST)
                     .body(Body::from(packed_msg))
                     .unwrap(),
@@ -434,7 +434,7 @@ mod tests {
         let response = app
             .oneshot(
                 Request::builder()
-                    .uri(String::from("/mediate"))
+                    .uri(String::from("/"))
                     .method(Method::POST)
                     .header(CONTENT_TYPE, DIDCOMM_ENCRYPTED_MIME_TYPE)
                     .body(Body::from(packed_msg))
@@ -493,7 +493,7 @@ mod tests {
         let response = app
             .oneshot(
                 Request::builder()
-                    .uri(String::from("/mediate"))
+                    .uri(String::from("/"))
                     .method(Method::POST)
                     .header(CONTENT_TYPE, DIDCOMM_ENCRYPTED_MIME_TYPE)
                     .body(Body::from(packed_msg))
@@ -545,7 +545,7 @@ mod tests {
         let response = app
             .oneshot(
                 Request::builder()
-                    .uri(String::from("/mediate"))
+                    .uri(String::from("/"))
                     .method(Method::POST)
                     .header(CONTENT_TYPE, DIDCOMM_ENCRYPTED_MIME_TYPE)
                     .body(Body::from(packed_msg))
@@ -596,7 +596,7 @@ mod tests {
         let response = app
             .oneshot(
                 Request::builder()
-                    .uri(String::from("/mediate"))
+                    .uri(String::from("/"))
                     .method(Method::POST)
                     .header(CONTENT_TYPE, DIDCOMM_ENCRYPTED_MIME_TYPE)
                     .body(Body::from(packed_msg))

diff --git a/crates/plugins/mediator-coordination/src/web/handler.rs b/crates/plugins/mediator-coordination/src/web/handler.rs
@@ -11,8 +11,9 @@ use crate::{
     constant::{
         DIDCOMM_ENCRYPTED_MIME_TYPE, KEYLIST_QUERY_2_0, KEYLIST_UPDATE_2_0, MEDIATE_FORWARD_2_0,
         MEDIATE_REQUEST_2_0,
-    }, forward::routing::mediator_forward_process, web::{self, error::MediationError, AppState}
-
+    },
+    forward::routing::mediator_forward_process,
+    web::{self, error::MediationError, AppState},
 };
 
 #[axum::debug_handler]
@@ -21,12 +22,7 @@ pub(crate) async fn process_didcomm_message(
     Extension(message): Extension<Message>,
 ) -> Response {
     if message.type_ == MEDIATE_FORWARD_2_0 {
-        let response = mediator_forward_process(&state, message)
-            .await
-            .map(|_| StatusCode::ACCEPTED.into_response())
-            .map_err(|err| err);
-
-        return match response {
+        return match mediator_forward_process(&state, message).await {
             Ok(_message) => StatusCode::ACCEPTED.into_response(),
             Err(response) => response,
         };
@@ -38,20 +34,19 @@ pub(crate) async fn process_didcomm_message(
                 message,
             )
             .await
-        },
+        }
         KEYLIST_QUERY_2_0 => {
             web::coord::handler::stateful::process_plain_keylist_query_message(
                 Arc::clone(&state),
                 message,
             )
             .await
         }
-     
+
         MEDIATE_REQUEST_2_0 => {
             web::coord::handler::stateful::process_mediate_request(&state, &message).await
         }
-
-
+
         _ => {
             let response = (
                 StatusCode::BAD_REQUEST,
@@ -98,8 +93,11 @@ pub mod tests {
     use web::AppStateRepository;
 
     use crate::{
-        didcomm::bridge::LocalSecretsResolver, repository::stateful::tests::{MockConnectionRepository, MockMessagesRepository, MockSecretsRepository}, util::{self, MockFileSystem}
-
+        didcomm::bridge::LocalSecretsResolver,
+        repository::stateful::tests::{
+            MockConnectionRepository, MockMessagesRepository, MockSecretsRepository,
+        },
+        util::{self, MockFileSystem},
     };
 
     pub fn setup() -> (Router, Arc<AppState>) {
@@ -205,8 +203,8 @@ pub mod tests {
 mod tests2 {
     use super::{tests as global, *};
     use crate::{
-        constant::KEYLIST_UPDATE_RESPONSE_2_0, repository::stateful::tests::MockConnectionRepository,
-
+        constant::KEYLIST_UPDATE_RESPONSE_2_0,
+        repository::stateful::tests::MockConnectionRepository,
     };
     use axum::Router;
     use hyper::{Body, Method, Request};
@@ -291,7 +289,7 @@ mod tests2 {
         let response = app
             .oneshot(
                 Request::builder()
-                    .uri(String::from("/mediate"))
+                    .uri(String::from("/"))
                     .method(Method::POST)
                     .header(CONTENT_TYPE, DIDCOMM_ENCRYPTED_MIME_TYPE)
                     .body(Body::from(packed_msg))
@@ -341,4 +339,4 @@ mod tests2 {
             })
         );
     }
-}
+}
diff --git a/crates/plugins/mediator-coordination/src/web/mod.rs b/crates/plugins/mediator-coordination/src/web/mod.rs
@@ -16,7 +16,7 @@ use crate::{
 pub fn routes(state: Arc<AppState>) -> Router {
     Router::new()
         // Unified route for all DIDComm messages
-        .route("/mediate", post(handler::process_didcomm_message))
+        .route("/", post(handler::process_didcomm_message))
         .route_layer(middleware::from_fn_with_state(
             state.clone(),
             midlw::unpack_didcomm_message,