From 78d61e7e0ccdfa0166e6190bebcea6b3c516d66c Mon Sep 17 00:00:00 2001
From: Neur0toxine <pashok9825@gmail.com>
Date: Wed, 29 May 2024 19:59:02 +0300
Subject: [PATCH] wip: json errors

---
 Cargo.toml        |  2 ++
 src/json_error.rs | 47 +++++++++++++++++++++++++++++++++++++++++++++++
 src/main.rs       |  1 +
 src/server.rs     | 28 +++++++++++++++++-----------
 4 files changed, 67 insertions(+), 11 deletions(-)
 create mode 100644 src/json_error.rs

diff --git a/Cargo.toml b/Cargo.toml
index 62b8dbd..2e79983 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -20,3 +20,5 @@ http = "1.1.0"
 tokio-util = { version = "0.7.11", features = ["io"] }
 futures-util = "0.3.30"
 infer = "0.15.0"
+tower = { version = "0.4.13", features = ["timeout"] }
+tower-service = "0.3.2"
diff --git a/src/json_error.rs b/src/json_error.rs
new file mode 100644
index 0000000..6addf03
--- /dev/null
+++ b/src/json_error.rs
@@ -0,0 +1,47 @@
+use tower::Layer;
+
+use crate::json_error::private::JsonErrorService;
+
+#[derive(Debug, Clone)]
+#[must_use]
+pub struct JsonError;
+
+impl<S> Layer<S> for JsonError {
+    type Service = JsonErrorService<S>;
+
+    fn layer(&self, inner: S) -> Self::Service {
+        JsonErrorService { inner }
+    }
+}
+
+mod private {
+    use std::task::Context;
+
+    use http::Request;
+    use tower_service::Service;
+
+    #[derive(Debug, Clone, Copy)]
+    pub struct JsonErrorService<S> {
+        pub(super) inner: S,
+    }
+
+    impl<B, S> Service<Request<B>> for JsonErrorService<S>
+        where
+            S: Service<Request<B>>,
+    {
+        type Response = S::Response;
+        type Error = S::Error;
+        type Future = S::Future;
+
+        #[inline]
+        fn poll_ready(&mut self, cx: &mut Context<'_>) -> std::task::Poll<Result<(), Self::Error>> {
+            self.inner.poll_ready(cx)
+        }
+
+        #[inline]
+        fn call(&mut self, mut req: Request<B>) -> Self::Future {
+            let response = self.inner.call(req);
+            response
+        }
+    }
+}
\ No newline at end of file
diff --git a/src/main.rs b/src/main.rs
index 764217f..c91dfbc 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -11,6 +11,7 @@ mod server;
 mod task;
 mod thread_pool;
 mod transcoder;
+mod json_error;
 
 const WORK_DIR_IN_OUT_LIFETIME: u64 = 60 * 60;
 
diff --git a/src/server.rs b/src/server.rs
index 8e1670d..28ec529 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -1,33 +1,35 @@
 use std::env;
 use std::ffi::OsStr;
+use std::sync::Arc;
 use std::time::{Duration, SystemTime};
 
+use axum::{Json, Router};
+use axum::body::Body;
+use axum::body::Bytes;
 use axum::extract::{DefaultBodyLimit, Path, State};
 use axum::http::StatusCode;
 use axum::response::IntoResponse;
 use axum::routing::{get, post};
-use axum::{Json, Router};
 use axum_typed_multipart::TypedMultipart;
+use futures_util::StreamExt;
 use tokio::fs;
+use tokio::fs::File;
+use tokio::io::AsyncReadExt;
 use tokio::net::TcpListener;
 use tokio::time::interval;
+use tokio_util::io::ReaderStream;
+use tower::ServiceBuilder;
 use tower_http::trace::TraceLayer;
 use tracing::{debug, error};
 use uuid::Uuid;
 
 use crate::dto::{ConvertRequest, ConvertResponse, ConvertURLRequest, ErrorResponse};
-use crate::task::{Task, TaskParams};
-use crate::thread_pool::ThreadPool;
 
 use crate::filepath;
 use crate::filepath::{in_file_path, out_file_path};
-use axum::body::Body;
-use axum::body::Bytes;
-use futures_util::StreamExt;
-use std::sync::Arc;
-use tokio::fs::File;
-use tokio::io::AsyncReadExt;
-use tokio_util::io::ReaderStream;
+use crate::json_error::JsonError;
+use crate::task::{Task, TaskParams};
+use crate::thread_pool::ThreadPool;
 
 const CONTENT_LENGTH_LIMIT: usize = 100 * 1024 * 1024;
 
@@ -73,7 +75,11 @@ impl Server {
             .route("/enqueue_url", post(enqueue_url))
             .route("/get/:identifier", get(download_file))
             .with_state(this)
-            .layer(TraceLayer::new_for_http())
+            .layer(
+                ServiceBuilder::new()
+                    .layer(TraceLayer::new_for_http())
+                    .layer(JsonError)
+            )
             .fallback(handler_not_found);
 
         tracing::info!("listening on {addr}");