From ba011050585f0561da02baa1514fdefbe558487b Mon Sep 17 00:00:00 2001
From: Mo8it <mo8it@proton.me>
Date: Wed, 2 Nov 2022 00:24:17 +0100
Subject: [PATCH] Add address from config, rate limiter and tracing

---
 Cargo.toml     |  5 +++++
 src/config.rs  | 14 ++++++++++++++
 src/errors.rs  |  2 +-
 src/logging.rs | 19 +++++++++++++++++++
 src/main.rs    | 45 +++++++++++++++++++++++++++++++++++++++------
 5 files changed, 78 insertions(+), 7 deletions(-)
 create mode 100644 src/logging.rs

diff --git a/Cargo.toml b/Cargo.toml
index 74cd1d2..8dfa202 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -20,3 +20,8 @@ lettre = "0.10"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 tokio = { version = "1.21", features = ["full"] }
+tower = { version = "0.4", features = ["limit", "buffer"] }
+tower-http = { version = "0.3", features = ["trace"] }
+tracing = "0.1"
+tracing-appender = "0.2"
+tracing-subscriber = "0.3"
diff --git a/src/config.rs b/src/config.rs
index 600f120..04248f0 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -4,6 +4,12 @@ use std::env;
 use std::fs::File;
 use std::io::BufReader;
 
+#[derive(Deserialize)]
+pub struct SocketAddress {
+    pub address: [u8; 4],
+    pub port: u16,
+}
+
 #[derive(Deserialize)]
 pub struct EmailServer {
     pub server_name: String,
@@ -18,12 +24,20 @@ pub struct Address {
     pub domain: String,
 }
 
+#[derive(Deserialize)]
+pub struct Logging {
+    pub directory: String,
+    pub filename_prefix: String,
+}
+
 #[derive(Deserialize)]
 pub struct Config {
     pub path_prefix: String,
+    pub socket_address: SocketAddress,
     pub email_server: EmailServer,
     pub email_from: Address,
     pub email_to: Address,
+    pub logging: Logging,
 }
 
 impl Config {
diff --git a/src/errors.rs b/src/errors.rs
index 74bc4c9..e55ebae 100644
--- a/src/errors.rs
+++ b/src/errors.rs
@@ -11,6 +11,6 @@ impl IntoResponse for AppError {
 
 impl From<anyhow::Error> for AppError {
     fn from(err: anyhow::Error) -> Self {
-        Self(err.into())
+        Self(err)
     }
 }
diff --git a/src/logging.rs b/src/logging.rs
new file mode 100644
index 0000000..01b899a
--- /dev/null
+++ b/src/logging.rs
@@ -0,0 +1,19 @@
+use tracing_appender::non_blocking::WorkerGuard;
+use tracing_subscriber::filter::LevelFilter;
+
+use crate::config;
+
+pub fn init_logger(logging_config: &config::Logging) -> WorkerGuard {
+    let file_appender = tracing_appender::rolling::hourly(
+        &logging_config.directory,
+        &logging_config.filename_prefix,
+    );
+    let (non_blocking, guard) = tracing_appender::non_blocking(file_appender);
+
+    tracing_subscriber::fmt()
+        .with_max_level(LevelFilter::INFO)
+        .with_writer(non_blocking)
+        .init();
+
+    guard
+}
diff --git a/src/main.rs b/src/main.rs
index 006ed55..903368c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,6 +2,7 @@ mod captcha_solutions;
 mod config;
 mod errors;
 mod forms;
+mod logging;
 mod mailer;
 mod routes;
 mod templates;
@@ -9,14 +10,34 @@ mod templates;
 use anyhow::Result;
 use axum::extract::Extension;
 use axum::routing::{get, post};
+use axum::{error_handling::HandleErrorLayer, http::StatusCode, BoxError};
 use axum::{Router, Server};
+use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 use std::process;
 use std::sync::Arc;
+use std::time::Duration;
+use tower::buffer::BufferLayer;
+use tower::limit::RateLimitLayer;
+use tower::ServiceBuilder;
+use tower_http::trace::{DefaultOnResponse, TraceLayer};
+use tracing::Level;
+use tracing_appender::non_blocking::WorkerGuard;
 
-async fn init() -> Result<()> {
+async fn init() -> Result<WorkerGuard> {
     let mut config = config::Config::new()?;
     let path_prefix = config.path_prefix.clone();
     let mailer = Arc::new(mailer::Mailer::new(&mut config)?);
+
+    let address = config.socket_address.address;
+    let socket_address = SocketAddr::new(
+        IpAddr::V4(Ipv4Addr::new(
+            address[0], address[1], address[2], address[3],
+        )),
+        config.socket_address.port,
+    );
+
+    let tracing_worker_gurad = logging::init_logger(&config.logging);
+
     let config = Arc::new(config);
     let captcha_solutions = Arc::new(captcha_solutions::SharedCaptchaSolutions::new());
 
@@ -24,24 +45,36 @@ async fn init() -> Result<()> {
         .route("/", get(routes::index))
         .route("/submit", post(routes::submit))
         .route("/success", get(routes::success))
+        .layer(TraceLayer::new_for_http().on_response(DefaultOnResponse::new().level(Level::INFO)))
+        .layer(
+            ServiceBuilder::new()
+                .layer(HandleErrorLayer::new(|err: BoxError| async move {
+                    (
+                        StatusCode::INTERNAL_SERVER_ERROR,
+                        format!("Unhandled error: {}", err),
+                    )
+                }))
+                .layer(BufferLayer::new(1024))
+                .layer(RateLimitLayer::new(2, Duration::from_secs(3))),
+        )
         .layer(Extension(config))
         .layer(Extension(mailer))
         .layer(Extension(captcha_solutions));
 
     let app = Router::new().nest(&path_prefix, routes);
 
-    Server::bind(&([127, 0, 0, 1], 8080).into())
+    Server::bind(&socket_address)
         .serve(app.into_make_service())
         .await
         .unwrap();
 
-    Ok(())
+    Ok(tracing_worker_gurad)
 }
 
 #[tokio::main]
 async fn main() {
-    init().await.unwrap_or_else(|e| {
-        eprintln!("{e}");
+    let _tracing_worker_gurad = init().await.unwrap_or_else(|e| {
+        eprintln!("{e:?}");
         process::exit(1);
-    })
+    });
 }