From ef24cfedbff489bff73d6e1e6114fdd527235acb Mon Sep 17 00:00:00 2001
From: Tyler Murphy <tylermurphy534@gmail.com>
Date: Tue, 31 Jan 2023 00:00:13 -0500
Subject: rate limiting

---
 src/api/mod.rs | 43 +++++++++++++++++++++++++++++++++++++++++++
 src/main.rs    | 24 +++---------------------
 2 files changed, 46 insertions(+), 21 deletions(-)

(limited to 'src')

diff --git a/src/api/mod.rs b/src/api/mod.rs
index c347207..d36b127 100644
--- a/src/api/mod.rs
+++ b/src/api/mod.rs
@@ -1,5 +1,48 @@
+use crate::types::extract::RouterURI;
+use axum::{Extension, Router, BoxError, error_handling::HandleErrorLayer};
+use tower::ServiceBuilder;
+use tower_governor::{governor::GovernorConfigBuilder, GovernorLayer, errors::display_error, key_extractor::SmartIpKeyExtractor};
+
 pub mod admin;
 pub mod auth;
 pub mod pages;
 pub mod posts;
 pub mod users;
+
+pub fn router() -> Router {
+    let governor_conf = Box::new(
+        GovernorConfigBuilder::default()
+            .burst_size(10)
+            .per_second(1)
+            .key_extractor(SmartIpKeyExtractor)
+            .finish()
+            .expect("Failed to create rate limiter"),
+    );
+
+    Router::new()
+        .nest(
+            "/admin",
+            admin::router().layer(Extension(RouterURI("/api/admin"))),
+        )
+        .nest(
+            "/auth",
+            auth::router().layer(Extension(RouterURI("/api/auth"))),
+        )
+        .nest(
+            "/users",
+            users::router().layer(Extension(RouterURI("/api/users"))),
+        )
+        .nest(
+            "/posts",
+            posts::router().layer(Extension(RouterURI("/api/posts"))),
+        )
+        .layer(
+            ServiceBuilder::new()
+                .layer(HandleErrorLayer::new(|e: BoxError| async move {
+                    display_error(e)
+                }))
+                .layer(GovernorLayer {
+                    config: Box::leak(governor_conf),
+                }),
+        )
+}
diff --git a/src/main.rs b/src/main.rs
index 03cee25..a72ec5f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -3,7 +3,7 @@ use axum::{
     http::{Request, StatusCode},
     middleware::{self, Next},
     response::Response,
-    Extension, RequestExt, Router,
+    RequestExt, Router,
 };
 use axum_client_ip::ClientIp;
 use std::{net::SocketAddr, process::exit};
@@ -14,10 +14,7 @@ use tracing_subscriber::{
 };
 use types::http::ResponseCode;
 
-use crate::{
-    api::{auth, pages, posts, users},
-    types::extract::RouterURI,
-};
+use crate::api::pages;
 
 mod admin;
 mod api;
@@ -77,22 +74,7 @@ async fn main() {
         .layer(middleware::from_fn(log))
         .layer(middleware::from_fn(serve))
         .nest("/", pages::router())
-        .nest(
-            "/api/admin",
-            api::admin::router().layer(Extension(RouterURI("/api/admin"))),
-        )
-        .nest(
-            "/api/auth",
-            auth::router().layer(Extension(RouterURI("/api/auth"))),
-        )
-        .nest(
-            "/api/users",
-            users::router().layer(Extension(RouterURI("/api/users"))),
-        )
-        .nest(
-            "/api/posts",
-            posts::router().layer(Extension(RouterURI("/api/posts"))),
-        )
+        .nest("/api", api::router())
         .layer(CookieManagerLayer::new());
 
     let Ok(addr) = "[::]:8080".parse::<std::net::SocketAddr>() else {
-- 
cgit v1.2.3-freya