Feat: Add admin middleware

4 files changed, 47 insertions, 4 deletions

diff --git a/.env b/.env
new file mode 100644
index 0000000..839a708
--- /dev/null
+++ b/.env
@@ -0,0 +1 @@
+TOKEN=TOKEN
diff --git a/Cargo.lock b/Cargo.lock
index c825f8e..72a89f5 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1166,9 +1166,11 @@ version = "0.1.0"
 dependencies = [
  "anyhow",
  "axum",
+ "dotenvy",
  "maud",
  "serde",
  "sqlx",
+ "subtle",
  "tokio",
  "tower-http",
 ]
diff --git a/Cargo.toml b/Cargo.toml
index c7b1a60..47f2399 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -13,3 +13,5 @@ maud = { version = "0.26.0", features = ["axum"] }
 axum = "0.7.5"
 tower-http = { version = "0.5.2", features = ["full", "fs"] }
 sqlx = { version = "0.8.2", features = ["runtime-tokio", "tls-native-tls", "sqlite", "macros"] }
+dotenvy = "0.15.7"
+subtle = "2.6.1"
diff --git a/src/main.rs b/src/main.rs
index b8e1302..feaa6a7 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,12 +1,15 @@
 use std::env;
 
-use axum::extract::DefaultBodyLimit;
-use axum::http::StatusCode;
-use axum::response::IntoResponse;
+use axum::extract::{DefaultBodyLimit, Request};
+use axum::http::{HeaderMap, StatusCode};
+use axum::middleware::Next;
+use axum::response::{IntoResponse, Response};
 use axum::routing::get;
-use axum::Router;
+use axum::{middleware, Router};
+use dotenvy::dotenv;
 use maud::{html, Markup, DOCTYPE};
 use sqlx::sqlite::SqlitePoolOptions;
+use subtle::ConstantTimeEq;
 use tower_http::services::ServeDir;
 
 mod db;
@@ -34,6 +37,7 @@ where
 
 #[tokio::main]
 async fn main() {
+    dotenv().ok();
     let pool = SqlitePoolOptions::new()
         .connect(
             &env::var("DATABASE_URL").expect("Failed to get environment variable DATABASE_URL"),
@@ -47,6 +51,8 @@ async fn main() {
         .expect("Failed to run migrations");
 
     let app = Router::new()
+        .route("/admin", get(admin_handler))
+        .layer(middleware::from_fn(admin_middleware))
         .route("/", get(home))
         .nest_service("/static", ServeDir::new("src/static"))
         .layer(DefaultBodyLimit::max(10000));
@@ -56,6 +62,38 @@ async fn main() {
         .unwrap();
 }
 
+async fn admin_middleware(
+    headers: HeaderMap,
+    request: Request,
+    next: Next,
+) -> Result<Response, StatusCode> {
+    let expected_token = env::var("TOKEN").map_err(|_| {
+        eprintln!("Error: TOKEN environment variable is not set");
+        StatusCode::INTERNAL_SERVER_ERROR
+    })?;
+
+    let token = headers
+        .get("X-Admin-Token")
+        .and_then(|header| header.to_str().ok())
+        .ok_or(StatusCode::UNAUTHORIZED)?;
+
+    if !expected_token
+        .as_bytes()
+        .ct_eq(token.as_bytes())
+        .unwrap_u8()
+        == 1
+    {
+        eprintln!("Warning: Unauthorized access attempt with invalid token");
+        return Err(StatusCode::UNAUTHORIZED);
+    }
+
+    Ok(next.run(request).await)
+}
+
+async fn admin_handler() -> &'static str {
+    "Hello World"
+}
+
 async fn home() -> Result<impl IntoResponse, AppError> {
     Ok(html!({
         (head("Home"))