refactor(workflow_ui): 抽出 auth 子模块

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
boos
2026-06-18 14:05:15 +08:00
parent fd4c1f0ac1
commit 8a55efbfb6
2 changed files with 67 additions and 57 deletions
+62
View File
@@ -0,0 +1,62 @@
use super::{parse_workflow_cookie, WorkflowCaller};
use crate::auth::jwt::validate_jwt;
use crate::error::{AppError, AppResult};
use crate::AppState;
use axum::http::HeaderMap;
pub(super) fn caller_from_headers(
state: &AppState,
headers: &HeaderMap,
) -> AppResult<WorkflowCaller> {
let token = headers
.get(axum::http::header::AUTHORIZATION)
.and_then(|value| value.to_str().ok())
.and_then(|value| {
value
.strip_prefix("Bearer ")
.or_else(|| value.strip_prefix("bearer "))
})
.or_else(|| {
headers
.get(axum::http::header::COOKIE)
.and_then(|value| value.to_str().ok())
.and_then(parse_workflow_cookie)
})
.map(str::trim)
.filter(|value| !value.is_empty())
.ok_or_else(|| AppError::Unauthorized("Missing authorization token".into()))?;
let proxy_token = state.runtime_proxy_token();
if token == state.config.auth.auth_token || token == proxy_token {
return Ok(WorkflowCaller::Admin);
}
let claims = validate_jwt(&state.config.auth.jwt_secret, token)
.map_err(|_| AppError::Unauthorized("Invalid or expired token".into()))?;
if claims.role == "admin" {
return Ok(WorkflowCaller::Admin);
}
if claims.role == "user" {
let user_id = claims
.sub
.parse::<i64>()
.map_err(|_| AppError::Unauthorized("Invalid user ID in token".into()))?;
return Ok(WorkflowCaller::User { user_id });
}
Err(AppError::Forbidden("Workflow access denied".into()))
}
pub(super) fn caller_user_id(caller: &WorkflowCaller) -> Option<i64> {
match caller {
WorkflowCaller::Admin => None,
WorkflowCaller::User { user_id } => Some(*user_id),
}
}
pub(super) fn caller_is_admin(caller: &WorkflowCaller) -> bool {
matches!(caller, WorkflowCaller::Admin)
}
pub(super) async fn auth_caller(state: &AppState, headers: &HeaderMap) -> AppResult<WorkflowCaller> {
caller_from_headers(state, headers)
}
+5 -57
View File
@@ -1,6 +1,5 @@
use crate::workflow::executor::{WorkflowProgressCallback, WorkflowProgressEvent}; use crate::workflow::executor::{WorkflowProgressCallback, WorkflowProgressEvent};
use crate::{ use crate::{
auth::jwt::validate_jwt,
error::{AppError, AppResult}, error::{AppError, AppResult},
services::workspace::{self, WorkspaceAssetInput}, services::workspace::{self, WorkspaceAssetInput},
workflow::{ workflow::{
@@ -53,7 +52,7 @@ static EMBEDDED_WORKFLOW_UI_CLOUD: Dir<'_> =
include_dir!("$CARGO_MANIFEST_DIR/dist/workflow-ui-cloud"); include_dir!("$CARGO_MANIFEST_DIR/dist/workflow-ui-cloud");
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
enum WorkflowCaller { pub(super) enum WorkflowCaller {
Admin, Admin,
User { user_id: i64 }, User { user_id: i64 },
} }
@@ -121,6 +120,9 @@ struct WorkflowValidateQuery {
check_files: bool, check_files: bool,
} }
mod auth;
use auth::{auth_caller, caller_from_headers, caller_is_admin, caller_user_id};
pub fn routes() -> Router<AppState> { pub fn routes() -> Router<AppState> {
Router::new() Router::new()
.route("/workflow-ui", get(workflow_ui_index)) .route("/workflow-ui", get(workflow_ui_index))
@@ -310,56 +312,6 @@ fn workflow_cloud_api_key(config: &crate::config::ComfyUIConfig) -> Option<&str>
.filter(|value| !value.is_empty()) .filter(|value| !value.is_empty())
} }
fn caller_from_headers(state: &AppState, headers: &HeaderMap) -> AppResult<WorkflowCaller> {
let token = headers
.get(axum::http::header::AUTHORIZATION)
.and_then(|value| value.to_str().ok())
.and_then(|value| {
value
.strip_prefix("Bearer ")
.or_else(|| value.strip_prefix("bearer "))
})
.or_else(|| {
headers
.get(axum::http::header::COOKIE)
.and_then(|value| value.to_str().ok())
.and_then(parse_workflow_cookie)
})
.map(str::trim)
.filter(|value| !value.is_empty())
.ok_or_else(|| AppError::Unauthorized("Missing authorization token".into()))?;
let proxy_token = state.runtime_proxy_token();
if token == state.config.auth.auth_token || token == proxy_token {
return Ok(WorkflowCaller::Admin);
}
let claims = validate_jwt(&state.config.auth.jwt_secret, token)
.map_err(|_| AppError::Unauthorized("Invalid or expired token".into()))?;
if claims.role == "admin" {
return Ok(WorkflowCaller::Admin);
}
if claims.role == "user" {
let user_id = claims
.sub
.parse::<i64>()
.map_err(|_| AppError::Unauthorized("Invalid user ID in token".into()))?;
return Ok(WorkflowCaller::User { user_id });
}
Err(AppError::Forbidden("Workflow access denied".into()))
}
fn caller_user_id(caller: &WorkflowCaller) -> Option<i64> {
match caller {
WorkflowCaller::Admin => None,
WorkflowCaller::User { user_id } => Some(*user_id),
}
}
fn caller_is_admin(caller: &WorkflowCaller) -> bool {
matches!(caller, WorkflowCaller::Admin)
}
fn default_history_limit() -> i64 { fn default_history_limit() -> i64 {
50 50
} }
@@ -1872,10 +1824,6 @@ fn cache_control_for_path(path: &StdPath) -> &'static str {
} }
} }
async fn auth_caller(state: &AppState, headers: &HeaderMap) -> AppResult<WorkflowCaller> {
caller_from_headers(state, headers)
}
fn auth_caller_optional(state: &AppState, headers: &HeaderMap) -> Option<WorkflowCaller> { fn auth_caller_optional(state: &AppState, headers: &HeaderMap) -> Option<WorkflowCaller> {
caller_from_headers(state, headers).ok() caller_from_headers(state, headers).ok()
} }
@@ -6083,7 +6031,7 @@ fn inject_enhanced_bootstrap(html: String) -> String {
html.replace("</head>", &format!("{bootstrap}\n</head>")) html.replace("</head>", &format!("{bootstrap}\n</head>"))
} }
fn parse_workflow_cookie(raw: &str) -> Option<&str> { pub(super) fn parse_workflow_cookie(raw: &str) -> Option<&str> {
raw.split(';') raw.split(';')
.map(str::trim) .map(str::trim)
.find_map(|pair| pair.strip_prefix("boosapi_workflow_token=")) .find_map(|pair| pair.strip_prefix("boosapi_workflow_token="))