Feat: Use random token for admin authentication

2 files changed, 40 insertions, 9 deletions

diff --git a/src/handlers/admin.rs b/src/handlers/admin.rs
index ca18ebc..af7aa42 100644
--- a/src/handlers/admin.rs
+++ b/src/handlers/admin.rs
@@ -1,30 +1,43 @@
 use std::collections::HashMap;
-use std::env;
 
 use anyhow::Context;
 use axum::body::Bytes;
-use axum::extract::{Multipart, Request};
+use axum::extract::{Multipart, Request, State};
 use axum::http::{HeaderMap, StatusCode};
 use axum::middleware::Next;
 use axum::response::{IntoResponse, Response};
 use axum::{Extension, Json};
+use rand::Rng;
 use subtle::ConstantTimeEq;
 
 use crate::db::story::Story;
 use crate::db::ProductDatabaseHandler;
 use crate::error::AppError;
 
-use super::ApiContext;
+use super::{ApiContext, AppState};
+
+pub fn generate_token() -> String {
+    const CHARSET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZ\
+                            abcdefghijklmnopqrstuvwxyz\
+                            0123456789)(*&^%$#@!~";
+    let mut rng = rand::thread_rng();
+    let token: String = (0..32)
+        .map(|_| {
+            let idx = rng.gen_range(0..CHARSET.len());
+            CHARSET[idx] as char
+        })
+        .collect();
+    println!("New Token generated: {}", token);
+    token
+}
 
 pub async fn admin_middleware(
+    State(state): State<AppState>,
     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 mut expected_token = state.admin_token.lock().await;
 
     let token = headers
         .get("X-Admin-Token")
@@ -38,8 +51,10 @@ pub async fn admin_middleware(
         == 1
     {
         eprintln!("Warning: Unauthorized access attempt with invalid token");
+        *expected_token = generate_token();
         return Err(StatusCode::UNAUTHORIZED);
     }
+    *expected_token = generate_token();
 
     Ok(next.run(request).await)
 }
diff --git a/src/handlers/mod.rs b/src/handlers/mod.rs
index 4a31909..34481e7 100644
--- a/src/handlers/mod.rs
+++ b/src/handlers/mod.rs
@@ -1,3 +1,6 @@
+use std::sync::Arc;
+use tokio::sync::Mutex;
+
 use axum::extract::DefaultBodyLimit;
 use axum::middleware::{self};
 use axum::routing::{get, post};
@@ -6,7 +9,7 @@ use sqlx::SqlitePool;
 use tower_http::add_extension::AddExtensionLayer;
 use tower_http::services::ServeDir;
 
-use admin::{admin_healthcheck, admin_middleware, upload_story};
+use admin::{admin_healthcheck, admin_middleware, generate_token, upload_story};
 use frontend_builder::home;
 
 use self::frontend_builder::{build_page, DummyProduct};
@@ -19,11 +22,24 @@ struct ApiContext {
     pool: SqlitePool,
 }
 
+#[derive(Clone)]
+struct AppState {
+    admin_token: Arc<Mutex<String>>,
+}
+
 pub async fn router(pool: SqlitePool) -> Router {
+    let state = AppState {
+        admin_token: Arc::new(Mutex::new(generate_token())),
+    };
+
     Router::new()
         .route("/admin/story", post(upload_story))
         .route("/admin", get(admin_healthcheck))
-        .layer(middleware::from_fn(admin_middleware))
+        .layer(middleware::from_fn_with_state(
+            state.clone(),
+            admin_middleware,
+        ))
+        .with_state(state)
         .route("/", get(home))
         .route("/stories", get(build_page::<DummyProduct>))
         .nest_service("/static", ServeDir::new("static"))