feat: refine BoosAPI console and user workspace

src/routes/comfyui/mod.rs

@@ -6,6 +6,7 @@
 
 use crate::{
     auth::api_key::{authorize_downstream_token, DownstreamAuthResult},
+    auth::jwt::validate_jwt,
     error::{AppError, AppResult},
     AppState,
 };
@@ -47,6 +48,15 @@ async fn check_auth(state: &AppState, headers: &HeaderMap) -> Result<(), AppErro
         return Ok(());
     }
 
+    if let Ok(claims) = validate_jwt(&state.config.auth.jwt_secret, auth) {
+        if claims.role == "admin" {
+            return Ok(());
+        }
+        return Err(AppError::Forbidden(
+            "Admin account is required for this route".into(),
+        ));
+    }
+
     // Downstream API key
     if auth.starts_with("sk-") {
         let proxy_token = state.runtime_proxy_token();
@@ -69,6 +79,105 @@ async fn check_auth(state: &AppState, headers: &HeaderMap) -> Result<(), AppErro
     }
 }
 
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{
+        auth::jwt::generate_jwt,
+        config::{AppConfig, AuthConfig, ComfyUIConfig, DbConfig, ProxyConfig, ServerConfig},
+        services::{builtin_keys::BuiltinKeys, runtime_auth::RuntimeAuthSettings},
+    };
+    use axum::http::HeaderValue;
+    use sqlx::sqlite::SqlitePoolOptions;
+
+    async fn test_state() -> AppState {
+        let pool = SqlitePoolOptions::new()
+            .max_connections(1)
+            .connect("sqlite::memory:")
+            .await
+            .expect("sqlite pool");
+
+        AppState::new(
+            AppConfig {
+                server: ServerConfig {
+                    host: "127.0.0.1".to_string(),
+                    port: 4000,
+                    log_level: "info".to_string(),
+                    cors_origins: Vec::new(),
+                },
+                db: DbConfig {
+                    url: "sqlite::memory:".to_string(),
+                    max_connections: 1,
+                    ssl: false,
+                    runtime_db_url: None,
+                    runtime_db_dialect: None,
+                    runtime_db_ssl: None,
+                },
+                auth: AuthConfig {
+                    jwt_secret: "test-jwt-secret".to_string(),
+                    jwt_expiry_hours: 24,
+                    auth_token: "test-admin-token".to_string(),
+                    proxy_token: "sk-test-proxy".to_string(),
+                },
+                proxy: ProxyConfig {
+                    upstream_timeout_secs: 30,
+                    first_byte_timeout_secs: 30,
+                    max_retries: 1,
+                    proxy_log_retention_days: 7,
+                    template_proxy_url: None,
+                },
+                comfyui: ComfyUIConfig::default(),
+            },
+            pool,
+            BuiltinKeys {
+                comfyui_agent: (0, "sk-comfyui-test".to_string()),
+            },
+            RuntimeAuthSettings::from_token_for_test("sk-test-proxy"),
+        )
+    }
+
+    fn bearer_headers(token: &str) -> HeaderMap {
+        let mut headers = HeaderMap::new();
+        headers.insert(
+            "Authorization",
+            HeaderValue::from_str(&format!("Bearer {token}")).expect("valid bearer header"),
+        );
+        headers
+    }
+
+    #[tokio::test]
+    async fn comfyui_auth_accepts_admin_jwt() {
+        let state = test_state().await;
+        let token = generate_jwt(
+            &state.config.auth.jwt_secret,
+            "42",
+            "admin",
+            state.config.auth.jwt_expiry_hours,
+        )
+        .expect("admin jwt");
+
+        check_auth(&state, &bearer_headers(&token))
+            .await
+            .expect("admin jwt can access comfyui proxy routes");
+    }
+
+    #[tokio::test]
+    async fn comfyui_auth_rejects_user_jwt() {
+        let state = test_state().await;
+        let token = generate_jwt(
+            &state.config.auth.jwt_secret,
+            "43",
+            "user",
+            state.config.auth.jwt_expiry_hours,
+        )
+        .expect("user jwt");
+
+        let result = check_auth(&state, &bearer_headers(&token)).await;
+
+        assert!(matches!(result, Err(AppError::Forbidden(_))));
+    }
+}
+
 /// Get ComfyUI base URL from embedded/external runtime configuration.
 fn comfyui_base(state: &AppState) -> String {
     state.comfyui.internal_base_url()