refactor(torii): add accept header extractor

Signed-off-by: Shanin Roman <[email protected]>
hyperledger-iroha · Jul 2, 2024 · ea2fc50 · ea2fc50
1 parent 6910e37
commit ea2fc50
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 9 deletions.
diff --git a/torii/src/lib.rs b/torii/src/lib.rs
@@ -40,7 +40,7 @@ use tower_http::{
     trace::{DefaultMakeSpan, TraceLayer},
 };
 use utils::{
-    body::{ClientQueryRequestExtractor, ScaleVersioned},
+    extractors::{ClientQueryRequestExtractor, ExtractAccept, ScaleVersioned},
     Scale,
 };
 
@@ -118,11 +118,10 @@ impl Torii {
                 &format!("{}/*tail", uri::STATUS),
                 get({
                     let metrics_reporter = self.metrics_reporter.clone();
-                    move |headers: axum::http::header::HeaderMap, axum::extract::Path(tail): axum::extract::Path<String>| {
-                        let accept = headers.get(axum::http::header::ACCEPT);
+                    move |accept: Option<ExtractAccept>, axum::extract::Path(tail): axum::extract::Path<String>| {
                         core::future::ready(routing::handle_status(
                             &metrics_reporter,
-                            accept,
+                            accept.map(|extract| extract.0),
                             Some(&tail),
                         ))
                     }
@@ -132,9 +131,8 @@ impl Torii {
                 uri::STATUS,
                 get({
                     let metrics_reporter = self.metrics_reporter.clone();
-                    move |headers: axum::http::header::HeaderMap| {
-                        let accept = headers.get(axum::http::header::ACCEPT);
-                        core::future::ready(routing::handle_status(&metrics_reporter, accept, None))
+                    move |accept: Option<ExtractAccept>| {
+                        core::future::ready(routing::handle_status(&metrics_reporter, accept.map(|extract| extract.0), None))
                     }
                 }),
             )

diff --git a/torii/src/utils.rs b/torii/src/utils.rs
@@ -24,11 +24,12 @@ impl<T: Encode + Send> IntoResponse for Scale<T> {
     }
 }
 
-pub mod body {
+pub mod extractors {
     use axum::{
         async_trait,
         body::Bytes,
         extract::{FromRequest, FromRequestParts, Query, Request},
+        http::StatusCode,
     };
     use iroha_data_model::query::cursor::ForwardCursor;
 
@@ -57,7 +58,7 @@ pub mod body {
                 .map_err(|err| {
                     (
                         axum::http::StatusCode::BAD_REQUEST,
-                        format!("Transaction Rejected (Malformed), Reason : '{err}'"),
+                        format!("Could not decode request: {err}"),
                     )
                         .into_response()
                 })
@@ -93,4 +94,27 @@ pub mod body {
                 .map(ClientQueryRequestExtractor)
         }
     }
+
+    /// Extractor of Accept header
+    pub struct ExtractAccept(pub HeaderValue);
+
+    #[async_trait]
+    impl<S> FromRequestParts<S> for ExtractAccept
+    where
+        S: Send + Sync,
+    {
+        type Rejection = (StatusCode, &'static str);
+
+        async fn from_request_parts(
+            parts: &mut axum::http::request::Parts,
+            _state: &S,
+        ) -> Result<Self, Self::Rejection> {
+            parts
+                .headers
+                .get(axum::http::header::ACCEPT)
+                .cloned()
+                .map(ExtractAccept)
+                .ok_or((StatusCode::BAD_REQUEST, "`Accept` header is missing"))
+        }
+    }
 }