Spaces:
Sleeping
Sleeping
| // Basic Authentication Middleware for Admin UI and API | |
| use axum::{ | |
| extract::Request, | |
| http::{header, StatusCode}, | |
| middleware::Next, | |
| response::{IntoResponse, Response, Html}, | |
| }; | |
| use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64}; | |
| use once_cell::sync::Lazy; | |
| /// Admin credentials from environment variables | |
| struct AdminCredentials { | |
| username: String, | |
| password: String, | |
| enabled: bool, | |
| } | |
| /// Load admin credentials from environment | |
| fn load_admin_credentials() -> AdminCredentials { | |
| let username = std::env::var("ADMIN_USER") | |
| .or_else(|_| std::env::var("ADMIN_USERNAME")) | |
| .unwrap_or_default(); | |
| let password = std::env::var("ADMIN_PASS") | |
| .or_else(|_| std::env::var("ADMIN_PASSWORD")) | |
| .unwrap_or_default(); | |
| let enabled = !username.is_empty() && !password.is_empty(); | |
| if enabled { | |
| tracing::info!("Admin authentication enabled for user: {}", username); | |
| } else { | |
| tracing::warn!("Admin authentication DISABLED - set ADMIN_USER and ADMIN_PASS environment variables"); | |
| } | |
| AdminCredentials { | |
| username, | |
| password, | |
| enabled, | |
| } | |
| } | |
| /// Static admin credentials (loaded once at startup) | |
| static ADMIN_CREDS: Lazy<AdminCredentials> = Lazy::new(load_admin_credentials); | |
| /// Check if admin auth is enabled | |
| pub fn is_admin_auth_enabled() -> bool { | |
| ADMIN_CREDS.enabled | |
| } | |
| /// Validate Basic Auth credentials | |
| fn validate_basic_auth(auth_header: &str) -> bool { | |
| if !ADMIN_CREDS.enabled { | |
| return true; | |
| } | |
| // Parse "Basic <base64>" | |
| let encoded = match auth_header.strip_prefix("Basic ") { | |
| Some(e) => e, | |
| None => return false, | |
| }; | |
| // Decode base64 | |
| let decoded = match BASE64.decode(encoded) { | |
| Ok(d) => d, | |
| Err(_) => return false, | |
| }; | |
| // Parse "username:password" | |
| let decoded_str = match String::from_utf8(decoded) { | |
| Ok(s) => s, | |
| Err(_) => return false, | |
| }; | |
| let parts: Vec<&str> = decoded_str.splitn(2, ':').collect(); | |
| if parts.len() != 2 { | |
| return false; | |
| } | |
| parts[0] == ADMIN_CREDS.username && parts[1] == ADMIN_CREDS.password | |
| } | |
| /// Basic Auth middleware for admin pages and API | |
| pub async fn basic_auth_middleware(request: Request, next: Next) -> Result<Response, Response> { | |
| let path = request.uri().path(); | |
| // Skip auth for proxy endpoints (/v1/*) - they use API key auth | |
| if path.starts_with("/v1") { | |
| return Ok(next.run(request).await); | |
| } | |
| // Skip auth for health check | |
| if path == "/healthz" { | |
| return Ok(next.run(request).await); | |
| } | |
| // If admin auth is disabled, allow all | |
| if !ADMIN_CREDS.enabled { | |
| return Ok(next.run(request).await); | |
| } | |
| // Check Authorization header | |
| let auth_header = request | |
| .headers() | |
| .get(header::AUTHORIZATION) | |
| .and_then(|h| h.to_str().ok()); | |
| match auth_header { | |
| Some(header) if validate_basic_auth(header) => { | |
| Ok(next.run(request).await) | |
| } | |
| _ => { | |
| // Return 401 with WWW-Authenticate header | |
| Err(unauthorized_response(path)) | |
| } | |
| } | |
| } | |
| /// Generate 401 response with login prompt | |
| fn unauthorized_response(path: &str) -> Response { | |
| // For API endpoints, return JSON | |
| if path.starts_with("/api/") { | |
| let body = serde_json::json!({ | |
| "success": false, | |
| "error": "Authentication required. Set ADMIN_USER and ADMIN_PASS in HuggingFace Secrets." | |
| }); | |
| return ( | |
| StatusCode::UNAUTHORIZED, | |
| [ | |
| (header::CONTENT_TYPE, "application/json"), | |
| (header::WWW_AUTHENTICATE, "Basic realm=\"Antigravity Admin\""), | |
| ], | |
| serde_json::to_string(&body).unwrap(), | |
| ) | |
| .into_response(); | |
| } | |
| // For web pages, return HTML login prompt | |
| ( | |
| StatusCode::UNAUTHORIZED, | |
| [ | |
| (header::CONTENT_TYPE, "text/html; charset=utf-8"), | |
| (header::WWW_AUTHENTICATE, "Basic realm=\"Antigravity Admin\""), | |
| ], | |
| Html(r#" | |
| <!DOCTYPE html> | |
| <html> | |
| <head> | |
| <title>Login Required</title> | |
| <style> | |
| body { font-family: system-ui; display: flex; justify-content: center; align-items: center; height: 100vh; margin: 0; background: #f5f5f5; } | |
| .box { background: white; padding: 2rem; border-radius: 8px; box-shadow: 0 2px 10px rgba(0,0,0,0.1); text-align: center; } | |
| h1 { color: #333; margin-bottom: 1rem; } | |
| p { color: #666; } | |
| </style> | |
| </head> | |
| <body> | |
| <div class="box"> | |
| <h1>Login Required</h1> | |
| <p>Please enter your credentials to access Antigravity Manager.</p> | |
| <p style="font-size: 0.8rem; color: #999;">Configure ADMIN_USER and ADMIN_PASS in HuggingFace Secrets</p> | |
| </div> | |
| </body> | |
| </html> | |
| "#.to_string()), | |
| ) | |
| .into_response() | |
| } | |