feat(server,dashboard): 引入多工作流数据隔离、安全中间件与前端 ESM 模块化重构
- server: 实现按 workflow_name 的多工作流数据隔离与旧数据库平滑迁移机制 - server: 新增 API Key 认证(auth)、限流中间件(rate_limit)与运维备份接口(admin) - server: 统一 AppError 错误处理体系，重构调度器 scheduler 支持工作流级重置与抢占 - node: 节点 ID 缺失时自动生成随机 UUID，原生支持 `docker compose --scale node=N` 动态扩容 - dashboard: 前端模块化重构(state/api/components)，升级 CSS 变量设计系统与 Toast 通知 - docker/docs: 更新 /healthz 健康检查、部署脚本 IP 配置及数据库设计文档
This commit is contained in:
@@ -25,5 +25,8 @@ chrono.workspace = true
|
||||
uuid.workspace = true
|
||||
tempfile.workspace = true
|
||||
dotenvy.workspace = true
|
||||
sha2.workspace = true
|
||||
hex.workspace = true
|
||||
subtle = "2"
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,154 @@
|
||||
//! 管理 API（Admin 角色）。
|
||||
//!
|
||||
//! 提供 node 凭据的可视化与运维操作，供 Dashboard 管理界面调用：
|
||||
//! - 列出所有节点及其凭据状态（在线/token 是否有效/吊销/颁发时间）
|
||||
//! - 吊销指定节点的专属 token（立即失效，不影响其他节点）
|
||||
//! - 重新颁发指定节点的专属 token（返回新明文，旧 token 失效）
|
||||
//!
|
||||
//! 这些端点均要求 Admin 角色（见 mod.rs 授权矩阵），node 自身无权操作他人或自身凭据，
|
||||
//! 从而保证「吊销/重发」是管理员主动行为，避免被攻陷节点篡改凭据体系。
|
||||
|
||||
use super::{is_valid_node_id, AppState};
|
||||
use axum::{
|
||||
extract::{Path as AxumPath, State},
|
||||
http::StatusCode,
|
||||
response::IntoResponse,
|
||||
Json,
|
||||
};
|
||||
use serde_json::json;
|
||||
use tracing::{info, warn};
|
||||
|
||||
/// GET /api/admin/nodes — 列出全部节点及凭据状态。
|
||||
pub async fn list_nodes(
|
||||
State(state): State<AppState>,
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
match state.db.list_nodes_with_credentials().await {
|
||||
Ok(list) => Ok((
|
||||
StatusCode::OK,
|
||||
Json(json!({ "success": true, "message": "成功获取节点列表", "data": list })),
|
||||
)),
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
|
||||
/// POST /api/admin/nodes/:node_id/revoke — 吊销指定节点的专属 token。
|
||||
///
|
||||
/// 吊销后该 node 的现有 token 立即失效，须重新走注册流程领取新 token。
|
||||
/// 操作幂等：对无凭据记录或已吊销的节点调用不会报错。
|
||||
pub async fn revoke_node(
|
||||
State(state): State<AppState>,
|
||||
AxumPath(node_id): AxumPath<String>,
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
// node_id 白名单校验，防止注入或异常输入（与 register_node 的 node_id 来源口径一致）
|
||||
if !is_valid_node_id(&node_id) {
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
"非法的节点 ID 参数".to_string(),
|
||||
));
|
||||
}
|
||||
match state.db.revoke_node_token(&node_id).await {
|
||||
Ok(_) => {
|
||||
info!("管理员已吊销节点 {} 的专属 token", node_id);
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
Json(
|
||||
json!({ "success": true, "message": format!("节点 '{}' 的 token 已吊销", node_id) }),
|
||||
),
|
||||
))
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("吊销节点 {} token 失败: {}", node_id, e);
|
||||
Err(e.into())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// POST /api/admin/nodes/:node_id/reissue — 重新颁发指定节点的专属 token。
|
||||
///
|
||||
/// 旧 token 立即失效，返回新 token 明文（仅此一次，DB 只存 hash）。
|
||||
/// 节点需用新 token 重新注册或由管理员手动同步到节点本地 `.node_token`。
|
||||
pub async fn reissue_node(
|
||||
State(state): State<AppState>,
|
||||
AxumPath(node_id): AxumPath<String>,
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
if !is_valid_node_id(&node_id) {
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
"非法的节点 ID 参数".to_string(),
|
||||
));
|
||||
}
|
||||
// 仅允许对已注册的节点重发 token（防止凭据表被写入幽灵 node_id）
|
||||
match state.db.get_node_exists(&node_id).await {
|
||||
Ok(false) => {
|
||||
return Err(crate::api::AppError::NotFound(format!(
|
||||
"节点 '{}' 不存在，请先注册",
|
||||
node_id
|
||||
)));
|
||||
}
|
||||
Ok(true) => {}
|
||||
Err(e) => return Err(e.into()),
|
||||
}
|
||||
match state.db.issue_node_token(&node_id).await {
|
||||
Ok(new_token) => {
|
||||
info!("管理员已为节点 {} 重新颁发专属 token", node_id);
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
Json(json!({
|
||||
"success": true,
|
||||
"message": format!("节点 '{}' 的 token 已重新颁发，请将新 token 同步到该节点", node_id),
|
||||
"node_token": new_token,
|
||||
})),
|
||||
))
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("为节点 {} 重新颁发 token 失败: {}", node_id, e);
|
||||
Err(e.into())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// POST /api/admin/nodes/:node_id/approve — 管理员同意节点接入申请。
|
||||
pub async fn approve_node(
|
||||
State(state): State<AppState>,
|
||||
AxumPath(node_id): AxumPath<String>,
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
if !is_valid_node_id(&node_id) {
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
"非法的节点 ID 参数".to_string(),
|
||||
));
|
||||
}
|
||||
match state.db.approve_node(&node_id).await {
|
||||
Ok(_token) => {
|
||||
info!("管理员已同意节点 {} 的接入申请并生成专属 Token", node_id);
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
Json(
|
||||
json!({ "success": true, "message": format!("节点 '{}' 已授权加入集群", node_id) }),
|
||||
),
|
||||
))
|
||||
}
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
|
||||
/// POST /api/admin/nodes/:node_id/reject — 管理员拒绝节点接入申请。
|
||||
pub async fn reject_node(
|
||||
State(state): State<AppState>,
|
||||
AxumPath(node_id): AxumPath<String>,
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
if !is_valid_node_id(&node_id) {
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
"非法的节点 ID 参数".to_string(),
|
||||
));
|
||||
}
|
||||
match state.db.reject_node(&node_id).await {
|
||||
Ok(_) => {
|
||||
info!("管理员已拒绝节点 {} 的接入申请并移除", node_id);
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
Json(
|
||||
json!({ "success": true, "message": format!("已拒绝节点 '{}' 的接入申请", node_id) }),
|
||||
),
|
||||
))
|
||||
}
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,121 @@
|
||||
//! 管理员表单登录与凭据校验 API。
|
||||
//!
|
||||
//! 提供基于短密码的身份认证服务：
|
||||
//! - POST /api/login：校验管理员密码，成功后返回 Admin Token，并记录 IP 错误次数防止暴力破解。
|
||||
//! - GET /api/auth/check：由 auth_middleware 保护，供前端初始化时检测当前保存的 Token 是否有效。
|
||||
|
||||
use super::{ct_eq_str, AppState};
|
||||
use axum::{
|
||||
extract::{ConnectInfo, State},
|
||||
http::StatusCode,
|
||||
response::IntoResponse,
|
||||
Json,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::net::SocketAddr;
|
||||
use tracing::{info, warn};
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct LoginRequest {
|
||||
pub password: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct LoginResponse {
|
||||
pub success: bool,
|
||||
pub message: String,
|
||||
pub token: Option<String>,
|
||||
}
|
||||
|
||||
/// POST /api/login — 管理员密码登录端点。
|
||||
pub async fn login(
|
||||
State(state): State<AppState>,
|
||||
ConnectInfo(addr): ConnectInfo<SocketAddr>,
|
||||
Json(req): Json<LoginRequest>,
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
let client_ip = addr.ip();
|
||||
|
||||
// 限流检查：5 分钟内最多允许 5 次失败尝试（基于 RateLimiter 防暴力破解）
|
||||
if state.rate_limiter.is_rate_limited(client_ip) {
|
||||
warn!("客户端 IP {} 登录失败次数过多，已临时封禁锁定", client_ip);
|
||||
return Err(crate::api::AppError::TooManyRequests(
|
||||
"登录失败次数过多，已被临时锁定，请 5 分钟后再试".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let admin_token = match state.admin_token.as_deref() {
|
||||
Some(t) if !t.is_empty() => t,
|
||||
_ => {
|
||||
warn!("系统未配置 admin_token 且鉴权未禁用，拒绝登录");
|
||||
return Err(crate::api::AppError::Forbidden(
|
||||
"服务端未配置管理员凭据，请检查配置文件".to_string(),
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
// 恒定时间密码比对（防时序旁路攻击）
|
||||
if ct_eq_str(&req.password, admin_token) {
|
||||
// 生成随机 64 位 Session Token
|
||||
let session_token = format!(
|
||||
"{}{}",
|
||||
uuid::Uuid::new_v4().simple(),
|
||||
uuid::Uuid::new_v4().simple()
|
||||
);
|
||||
let expiry = std::time::Instant::now() + std::time::Duration::from_secs(24 * 3600);
|
||||
|
||||
// 存储 Token 到内存中（带容量上限清理）
|
||||
{
|
||||
let mut sessions = state.admin_sessions.write().await;
|
||||
let now = std::time::Instant::now();
|
||||
// 1. 清理已过期的 session
|
||||
sessions.retain(|_, exp| *exp > now);
|
||||
// 2. 若超出容量限制，淘汰最老/最快过期的 session
|
||||
while sessions.len() >= crate::api::MAX_ADMIN_SESSIONS {
|
||||
if let Some(oldest_key) = sessions
|
||||
.iter()
|
||||
.min_by_key(|(_, exp)| **exp)
|
||||
.map(|(k, _)| k.clone())
|
||||
{
|
||||
sessions.remove(&oldest_key);
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
sessions.insert(session_token.clone(), expiry);
|
||||
}
|
||||
|
||||
info!(
|
||||
"客户端 IP {} 密码验证成功，已颁发 Admin Session Token",
|
||||
client_ip
|
||||
);
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
Json(LoginResponse {
|
||||
success: true,
|
||||
message: "登录成功".to_string(),
|
||||
token: Some(session_token),
|
||||
}),
|
||||
))
|
||||
} else {
|
||||
warn!("客户端 IP {} 登录密码校验失败", client_ip);
|
||||
// 记录一次失败
|
||||
state.rate_limiter.record_failure(client_ip);
|
||||
Err(crate::api::AppError::Unauthorized(
|
||||
"管理员密码错误，请重新输入".to_string(),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
/// GET /api/auth/check — 校验当前 Admin Token 是否有效。
|
||||
///
|
||||
/// 放在 auth_middleware（Role::Admin）之后，只要到达此 handler 说明 Token 校验必定成功。
|
||||
pub async fn check_auth() -> impl IntoResponse {
|
||||
(
|
||||
StatusCode::OK,
|
||||
Json(serde_json::json!({
|
||||
"success": true,
|
||||
"message": "Token 验证有效",
|
||||
"authenticated": true
|
||||
})),
|
||||
)
|
||||
}
|
||||
@@ -1,37 +1,81 @@
|
||||
use axum::{
|
||||
body::Body,
|
||||
extract::Path as AxumPath,
|
||||
http::{header, StatusCode},
|
||||
response::IntoResponse,
|
||||
};
|
||||
use axum::{body::Body, extract::Path as AxumPath, http::header, response::IntoResponse};
|
||||
use std::path::{Path, PathBuf};
|
||||
use tokio::fs::File;
|
||||
use tokio_util::io::ReaderStream;
|
||||
|
||||
pub async fn download_single_data_file(AxumPath(filename): AxumPath<String>) -> axum::response::Response {
|
||||
pub async fn download_single_data_file(
|
||||
AxumPath(filename): AxumPath<String>,
|
||||
) -> Result<axum::response::Response, crate::api::AppError> {
|
||||
let safe_name = Path::new(&filename)
|
||||
.file_name()
|
||||
.map(|s| s.to_string_lossy().to_string())
|
||||
.unwrap_or_default();
|
||||
|
||||
if safe_name.is_empty() || safe_name.starts_with('.') {
|
||||
return (StatusCode::BAD_REQUEST, "无效的数据文件名").into_response();
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
"无效的数据文件名".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
// 严苛白名单过滤：严防 `..`、特殊符号注入及路径穿透攻击
|
||||
if !safe_name.chars().all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-' || c == '+' || c == '@') {
|
||||
tracing::warn!("拦截到疑似非法字符构造的敏感及越界资源抓取行为: {}", safe_name);
|
||||
return (StatusCode::BAD_REQUEST, "参数非法，请求的文件包含系统不许可的危险专属占位或路径重定向字符").into_response();
|
||||
if !safe_name.chars().all(|c| {
|
||||
c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-' || c == '+' || c == '@'
|
||||
}) {
|
||||
tracing::warn!(
|
||||
"拦截到疑似非法字符构造的敏感及越界资源抓取行为: {}",
|
||||
safe_name
|
||||
);
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
"参数非法，请求的文件包含系统不许可的危险专属占位或路径重定向字符".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let rel_path = format!("assets/data/{}", safe_name);
|
||||
tracing::debug!("服务端处理数据文件下载请求: {}", safe_name);
|
||||
stream_asset_file(&rel_path, "application/octet-stream").await.into_response()
|
||||
stream_asset_file(&rel_path, "application/octet-stream").await
|
||||
}
|
||||
|
||||
pub async fn download_linelist() -> axum::response::Response {
|
||||
let linelist_path = std::env::var("DCTS_LINELIST_PATH").unwrap_or_else(|_| "assets/gfVIS99.dat".to_string());
|
||||
stream_asset_file(&linelist_path, "application/octet-stream").await.into_response()
|
||||
pub async fn download_linelist() -> Result<axum::response::Response, crate::api::AppError> {
|
||||
let linelist_path =
|
||||
std::env::var("DCTS_LINELIST_PATH").unwrap_or_else(|_| "assets/gfVIS99.dat".to_string());
|
||||
|
||||
// 路径规约校验：DCTS_LINELIST_PATH 解析后的绝对路径必须落在 assets 根目录内，
|
||||
// 防止环境变量被设为 ../../etc/passwd 之类导致任意文件流出。
|
||||
// assets 根目录优先取 DCTS_ASSETS_DIR，回退到相对路径 assets。
|
||||
let assets_root = std::env::var("DCTS_ASSETS_DIR").unwrap_or_else(|_| "assets".to_string());
|
||||
if !is_path_within_assets(&linelist_path, &assets_root) {
|
||||
tracing::warn!(
|
||||
"DCTS_LINELIST_PATH '{}' 不在 assets 根目录 '{}' 内，拒绝下载",
|
||||
linelist_path,
|
||||
assets_root
|
||||
);
|
||||
return Err(crate::api::AppError::Forbidden(
|
||||
"请求的谱线文件路径越界，已被拒绝".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
stream_asset_file(&linelist_path, "application/octet-stream").await
|
||||
}
|
||||
|
||||
/// 校验 target 路径（经 canonicalize 后）是否落在 assets 根目录之内。
|
||||
/// 对不存在的路径（canonicalize 失败）回退到 starts_with 的词法比较，宁可偏严。
|
||||
fn is_path_within_assets(target: &str, assets_root: &str) -> bool {
|
||||
// 严防 `..` 词法穿透
|
||||
if target.contains("..") {
|
||||
return false;
|
||||
}
|
||||
|
||||
let target_path = Path::new(target);
|
||||
let root_path = Path::new(assets_root);
|
||||
|
||||
let target_abs = std::fs::canonicalize(target_path).ok();
|
||||
let root_abs = std::fs::canonicalize(root_path).ok();
|
||||
|
||||
match (target_abs, root_abs) {
|
||||
(Some(t), Some(r)) => t.starts_with(&r),
|
||||
// 路径尚未存在时用词法前缀比较（canonicalize 需要文件存在）
|
||||
_ => target_path.starts_with(root_path),
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_asset(rel_path: &str) -> Option<PathBuf> {
|
||||
@@ -65,10 +109,17 @@ fn resolve_asset(rel_path: &str) -> Option<PathBuf> {
|
||||
None
|
||||
}
|
||||
|
||||
async fn stream_asset_file(rel_path: &str, content_type: &'static str) -> impl IntoResponse {
|
||||
async fn stream_asset_file(
|
||||
rel_path: &str,
|
||||
content_type: &'static str,
|
||||
) -> Result<axum::response::Response, crate::api::AppError> {
|
||||
let resolved_path = match resolve_asset(rel_path) {
|
||||
Some(p) => p,
|
||||
None => return (StatusCode::NOT_FOUND, "资源数据文件不存在").into_response(),
|
||||
None => {
|
||||
return Err(crate::api::AppError::NotFound(
|
||||
"资源数据文件不存在".to_string(),
|
||||
))
|
||||
}
|
||||
};
|
||||
|
||||
match File::open(&resolved_path).await {
|
||||
@@ -89,9 +140,11 @@ async fn stream_asset_file(rel_path: &str, content_type: &'static str) -> impl I
|
||||
(header::CONTENT_DISPOSITION, disposition),
|
||||
];
|
||||
|
||||
(headers, body).into_response()
|
||||
Ok((headers, body).into_response())
|
||||
}
|
||||
Err(e) => {
|
||||
let boxed_err: anyhow::Error = e.into();
|
||||
Err(crate::api::AppError::Internal(boxed_err))
|
||||
}
|
||||
Err(_) => (StatusCode::INTERNAL_SERVER_ERROR, "无法读取资源数据文件").into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,51 @@
|
||||
use axum::{
|
||||
http::StatusCode,
|
||||
response::{IntoResponse, Response},
|
||||
Json,
|
||||
};
|
||||
use serde_json::json;
|
||||
use tracing::error;
|
||||
|
||||
pub enum AppError {
|
||||
BadRequest(String),
|
||||
Unauthorized(String),
|
||||
Forbidden(String),
|
||||
NotFound(String),
|
||||
Conflict(String),
|
||||
TooManyRequests(String),
|
||||
Internal(anyhow::Error),
|
||||
}
|
||||
|
||||
impl IntoResponse for AppError {
|
||||
fn into_response(self) -> Response {
|
||||
let (status, error_message) = match self {
|
||||
AppError::BadRequest(msg) => (StatusCode::BAD_REQUEST, msg),
|
||||
AppError::Unauthorized(msg) => (StatusCode::UNAUTHORIZED, msg),
|
||||
AppError::Forbidden(msg) => (StatusCode::FORBIDDEN, msg),
|
||||
AppError::NotFound(msg) => (StatusCode::NOT_FOUND, msg),
|
||||
AppError::Conflict(msg) => (StatusCode::CONFLICT, msg),
|
||||
AppError::TooManyRequests(msg) => (StatusCode::TOO_MANY_REQUESTS, msg),
|
||||
AppError::Internal(err) => {
|
||||
error!("Internal server error: {:?}", err);
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
"Internal server error".to_string(),
|
||||
)
|
||||
}
|
||||
};
|
||||
|
||||
let body = Json(json!({
|
||||
"success": false,
|
||||
"message": error_message,
|
||||
"data": serde_json::Value::Null
|
||||
}));
|
||||
|
||||
(status, body).into_response()
|
||||
}
|
||||
}
|
||||
|
||||
impl From<anyhow::Error> for AppError {
|
||||
fn from(inner: anyhow::Error) -> Self {
|
||||
AppError::Internal(inner)
|
||||
}
|
||||
}
|
||||
+249
-41
@@ -1,20 +1,28 @@
|
||||
pub mod admin;
|
||||
pub mod auth;
|
||||
pub mod data;
|
||||
pub mod error;
|
||||
pub mod node;
|
||||
pub mod rate_limit;
|
||||
pub mod seed;
|
||||
pub mod status;
|
||||
pub mod task;
|
||||
pub mod workflow;
|
||||
|
||||
pub use error::AppError;
|
||||
|
||||
use crate::db::Database;
|
||||
use crate::scheduler::GridScheduler;
|
||||
use axum::{
|
||||
extract::State,
|
||||
http::{header, Request, StatusCode},
|
||||
middleware::Next,
|
||||
response::IntoResponse,
|
||||
};
|
||||
use crate::db::Database;
|
||||
use crate::scheduler::GridScheduler;
|
||||
use mq::sqlite_queue::SqliteTaskQueue;
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::sync::Arc;
|
||||
use subtle::ConstantTimeEq;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct AppState {
|
||||
@@ -22,53 +30,253 @@ pub struct AppState {
|
||||
pub queue: Arc<SqliteTaskQueue>,
|
||||
pub scheduler: Arc<GridScheduler>,
|
||||
pub results_dir: String,
|
||||
/// 限流与密码防暴破限速器
|
||||
pub rate_limiter: rate_limit::RateLimiter,
|
||||
/// 兼容字段：Some 表示「已启用某种鉴权」，用于 main.rs 决定是否挂载鉴权中间件。
|
||||
pub auth_token: Option<String>,
|
||||
/// Admin 凭据（管理 Dashboard / workflow 写操作）。
|
||||
pub admin_token: Option<String>,
|
||||
/// 应急开关：跳过全部鉴权（仅本地调试）。
|
||||
pub auth_disabled: bool,
|
||||
/// 动态 Session Tokens（登录后发放），设置 24 小时过期
|
||||
pub admin_sessions:
|
||||
Arc<tokio::sync::RwLock<std::collections::HashMap<String, std::time::Instant>>>,
|
||||
}
|
||||
|
||||
/// 固定时间敏感字符串一致性核验函数，彻底消解时序测信道猜测危险
|
||||
fn constant_time_eq(a: &str, b: &str) -> bool {
|
||||
let a_bytes = a.as_bytes();
|
||||
let b_bytes = b.as_bytes();
|
||||
let mut diff = (a_bytes.len() ^ b_bytes.len()) as u64;
|
||||
// 遍历目标 secret (b_bytes) 的完整长度，使耗时仅受 server 预期 token 长度决定
|
||||
for (i, &y) in b_bytes.iter().enumerate() {
|
||||
let x = if i < a_bytes.len() { a_bytes[i] } else { 0 };
|
||||
diff |= (x ^ y) as u64;
|
||||
/// Admin Session 最大保存上限
|
||||
pub const MAX_ADMIN_SESSIONS: usize = 100;
|
||||
|
||||
/// 已认证的 Node 身份（中间件校验 node token 通过后注入 request extension）。
|
||||
///
|
||||
/// 下游 handler（heartbeat / claim / report）通过 `Extension<AuthenticatedNode>` 取出，
|
||||
/// 用于校验请求体里声称的 node_id 与 token 绑定的 node_id 一致，杜绝跨节点冒充。
|
||||
#[derive(Clone)]
|
||||
pub struct AuthenticatedNode {
|
||||
pub node_id: String,
|
||||
}
|
||||
|
||||
/// 授权角色：决定某条路径需要哪类主体才能访问。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum Role {
|
||||
/// 公开免鉴权端点：登录 /login，节点注册申请 /node/register，审批状态查询 /node/check_status。
|
||||
Public,
|
||||
/// Node 运行态：心跳/领任务/上报/下载种子与数据。需 node 专属 token。
|
||||
Node,
|
||||
/// 管理操作：workflow CRUD、起停、查看 status、审批节点。需 admin token。
|
||||
Admin,
|
||||
}
|
||||
|
||||
/// 路径 → 角色授权矩阵。
|
||||
///
|
||||
/// 设计依据（最小权限）：
|
||||
/// - Admin 写操作（workflow CRUD / start / stop / status / approve / reject）只对 admin token 开放。
|
||||
/// - Node 运行态接口只认 node 专属 token（管理员在 Dashboard 审批后颁发，绑定 node_id，可吊销）。
|
||||
/// - 注册端点 /node/register 和状态轮询 /node/check_status 为 Public 免凭据（提交申请 ➔ 待管理员审批）。
|
||||
///
|
||||
/// 注意：路径已去掉 `/api` 前缀（nest 挂载后中间件看到的 path 不含 nest 前缀）。
|
||||
fn required_role(path: &str, method: &axum::http::Method) -> Option<Role> {
|
||||
use axum::http::Method;
|
||||
// 公开免鉴权端点
|
||||
if (path == "/login" || path == "/node/register" || path == "/node/check_status")
|
||||
&& method == Method::POST
|
||||
{
|
||||
return Some(Role::Public);
|
||||
}
|
||||
diff == 0
|
||||
// 校验身份与状态 -> Admin
|
||||
if path == "/auth/check" && method == Method::GET {
|
||||
return Some(Role::Admin);
|
||||
}
|
||||
// 写操作 → Admin
|
||||
if path == "/workflows" && (method == Method::POST || method == Method::GET) {
|
||||
return Some(Role::Admin);
|
||||
}
|
||||
if path.starts_with("/workflows/") {
|
||||
// GET/PUT/DELETE /workflows/:name, POST /start|stop → Admin
|
||||
return Some(Role::Admin);
|
||||
}
|
||||
if path == "/status" && method == Method::GET {
|
||||
return Some(Role::Admin);
|
||||
}
|
||||
// 管理 API（节点凭据查看/审批/吊销/重发）→ Admin
|
||||
if path.starts_with("/admin/") {
|
||||
return Some(Role::Admin);
|
||||
}
|
||||
// Node 运行态 → Node
|
||||
if path == "/node/heartbeat" && method == Method::POST {
|
||||
return Some(Role::Node);
|
||||
}
|
||||
if path == "/task/claim" && method == Method::POST {
|
||||
return Some(Role::Node);
|
||||
}
|
||||
if path == "/task/report" && method == Method::POST {
|
||||
return Some(Role::Node);
|
||||
}
|
||||
if path.starts_with("/seed/") && method == Method::GET {
|
||||
return Some(Role::Node);
|
||||
}
|
||||
if (path.starts_with("/data/file/") || path == "/data/linelist") && method == Method::GET {
|
||||
return Some(Role::Node);
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Axum 鉴权中间件：若 AppState 中配置了 auth_token 则强制校验 Bearer Token 或 X-API-Key
|
||||
pub async fn auth_middleware(
|
||||
State(state): State<AppState>,
|
||||
req: Request<axum::body::Body>,
|
||||
next: Next,
|
||||
) -> impl IntoResponse {
|
||||
if let Some(ref expected_token) = state.auth_token {
|
||||
let auth_header = req
|
||||
.headers()
|
||||
.get(header::AUTHORIZATION)
|
||||
.and_then(|v| v.to_str().ok());
|
||||
let api_key_header = req
|
||||
.headers()
|
||||
.get("x-api-key")
|
||||
.and_then(|v| v.to_str().ok());
|
||||
/// 恒定时间字符串比较。
|
||||
///
|
||||
/// 先对两个串各自做 SHA-256，再比较等长摘要（32 字节），彻底消除长度时序旁路——
|
||||
/// 任意长度的输入都产生相同长度的摘要，比较耗时固定，攻击者无法通过响应时间探得 token 长度。
|
||||
fn ct_eq_str(a: &str, b: &str) -> bool {
|
||||
let ha = {
|
||||
let mut h = Sha256::new();
|
||||
h.update(a.as_bytes());
|
||||
h.finalize()
|
||||
};
|
||||
let hb = {
|
||||
let mut h = Sha256::new();
|
||||
h.update(b.as_bytes());
|
||||
h.finalize()
|
||||
};
|
||||
ha.ct_eq(&hb).into()
|
||||
}
|
||||
|
||||
let token_valid = match (auth_header, api_key_header) {
|
||||
(Some(auth), _) if auth.starts_with("Bearer ") => constant_time_eq(&auth[7..], expected_token),
|
||||
(Some(auth), _) => constant_time_eq(auth, expected_token),
|
||||
(_, Some(key)) => constant_time_eq(key, expected_token),
|
||||
_ => false,
|
||||
};
|
||||
/// node_id 白名单：字母、数字、点、下划线、连字符，长度 1-128。
|
||||
/// 用于 register_node / admin revoke / reissue 统一入口校验，与 Dashboard XSS 防护口径一致。
|
||||
pub(crate) fn is_valid_node_id(id: &str) -> bool {
|
||||
!id.is_empty()
|
||||
&& id.len() <= 128
|
||||
&& id
|
||||
.chars()
|
||||
.all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-')
|
||||
}
|
||||
|
||||
if !token_valid {
|
||||
return (
|
||||
StatusCode::UNAUTHORIZED,
|
||||
"Unauthorized: Invalid or missing authentication token",
|
||||
)
|
||||
.into_response();
|
||||
/// host_name 白名单：可打印 ASCII（排除控制字符），长度 1-128。
|
||||
/// 防止 host_name 携带 HTML/控制字符进入管理 Dashboard 触发存储型 XSS 或污染显示。
|
||||
pub(crate) fn is_valid_host_name(name: &str) -> bool {
|
||||
!name.is_empty()
|
||||
&& name.len() <= 128
|
||||
&& name.chars().all(|c| c.is_ascii() && !c.is_ascii_control())
|
||||
}
|
||||
|
||||
/// 从请求头提取凭据原文（支持 `Authorization: Bearer <t>` 与 `X-API-Key: <t>`）。
|
||||
///
|
||||
/// 安全：非 `Bearer ` 前缀的 Authorization 一律视为无 token（不再回退为裸头值比较），
|
||||
/// 避免 `Authorization: Basic ...` 之类的上游代理头被误送入 token 比对。
|
||||
fn extract_token(req: &Request<axum::body::Body>) -> Option<String> {
|
||||
if let Some(auth) = req
|
||||
.headers()
|
||||
.get(header::AUTHORIZATION)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
{
|
||||
if let Some(rest) = auth.strip_prefix("Bearer ") {
|
||||
if !rest.is_empty() {
|
||||
return Some(rest.to_string());
|
||||
}
|
||||
}
|
||||
// 非 Bearer 前缀或空值：不作为 token
|
||||
}
|
||||
if let Some(key) = req.headers().get("x-api-key").and_then(|v| v.to_str().ok()) {
|
||||
if !key.is_empty() {
|
||||
return Some(key.to_string());
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Axum 鉴权中间件（L2）。
|
||||
///
|
||||
/// 流程：
|
||||
/// 1. 应急关闭（auth_disabled）→ 直接放行。
|
||||
/// 2. 路径不在授权矩阵 → 视为未公开接口，拒绝（401）。
|
||||
/// 3. 按角色校验对应凭据：
|
||||
/// - Admin: admin token 恒定时间比对。
|
||||
/// - Node: node 专属 token 经 DB 反查 node_id（token 只存 hash）。
|
||||
/// 4. Node 角色额外校验：请求声称的 node_id 须与 token 绑定的 node_id 一致
|
||||
/// （防 A 节点用 B 节点的 token 越权操作）。claim_task / data 下载无 node_id
|
||||
/// 输入，则仅校验 token 有效即可。
|
||||
pub async fn auth_middleware(
|
||||
State(state): State<AppState>,
|
||||
mut req: Request<axum::body::Body>,
|
||||
next: Next,
|
||||
) -> impl IntoResponse {
|
||||
// 应急关闭：本地调试专用，绕过全部校验
|
||||
if state.auth_disabled {
|
||||
return next.run(req).await.into_response();
|
||||
}
|
||||
|
||||
let path = req.uri().path().to_string();
|
||||
let method = req.method().clone();
|
||||
let role = match required_role(&path, &method) {
|
||||
Some(Role::Public) => {
|
||||
return next.run(req).await.into_response();
|
||||
}
|
||||
Some(r) => r,
|
||||
None => {
|
||||
// 未在矩阵中的路径一律拒绝（默认拒绝原则）
|
||||
return (StatusCode::UNAUTHORIZED, "Unauthorized").into_response();
|
||||
}
|
||||
};
|
||||
|
||||
let token = match extract_token(&req) {
|
||||
Some(t) => t,
|
||||
None => {
|
||||
return (StatusCode::UNAUTHORIZED, "Unauthorized: missing token").into_response();
|
||||
}
|
||||
};
|
||||
|
||||
// 审计日志：仅记录写操作（POST/PUT/DELETE）的「谁、做了什么」，不记请求体（防泄露）。
|
||||
// 在校验通过后记录 subject；校验失败由 401 分支体现，不单独审计。
|
||||
use axum::http::Method;
|
||||
let is_write = matches!(method, Method::POST | Method::PUT | Method::DELETE);
|
||||
|
||||
match role {
|
||||
Role::Public => unreachable!(),
|
||||
Role::Admin => {
|
||||
let mut valid = false;
|
||||
if let Some(ref admin) = state.admin_token {
|
||||
if ct_eq_str(&token, admin) {
|
||||
valid = true;
|
||||
}
|
||||
}
|
||||
if !valid {
|
||||
let mut sessions = state.admin_sessions.write().await;
|
||||
let now = std::time::Instant::now();
|
||||
sessions.retain(|_, expiry| *expiry > now);
|
||||
if sessions.contains_key(&token) {
|
||||
valid = true;
|
||||
}
|
||||
}
|
||||
|
||||
if valid {
|
||||
if is_write {
|
||||
tracing::info!(target: "dcts_audit", "AUDIT subject=admin method={} path={}", method, path);
|
||||
}
|
||||
return next.run(req).await.into_response();
|
||||
}
|
||||
(
|
||||
StatusCode::UNAUTHORIZED,
|
||||
"Unauthorized: invalid admin token",
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
Role::Node => {
|
||||
// 用 token 反查所属 node_id（DB 只存 hash，明文不落库）
|
||||
match state.db.find_node_by_token(&token).await {
|
||||
Some(token_node_id) => {
|
||||
// 仅对非例行高频请求（如任务结果上报 /task/report）记录 AUDIT 审计日志，
|
||||
// 成功的例行心跳 (/node/heartbeat) 与空闲领任务 (/task/claim) 静默跳过。
|
||||
if is_write && path != "/node/heartbeat" && path != "/task/claim" {
|
||||
tracing::info!(target: "dcts_audit", "AUDIT subject=node:{} method={} path={}", token_node_id, method, path);
|
||||
}
|
||||
req.extensions_mut().insert(AuthenticatedNode {
|
||||
node_id: token_node_id,
|
||||
});
|
||||
next.run(req).await.into_response()
|
||||
}
|
||||
None => (
|
||||
StatusCode::UNAUTHORIZED,
|
||||
"Unauthorized: invalid or revoked node token",
|
||||
)
|
||||
.into_response(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
next.run(req).await.into_response()
|
||||
}
|
||||
|
||||
@@ -1,25 +1,152 @@
|
||||
use super::AppState;
|
||||
use axum::{extract::State, response::IntoResponse, Json};
|
||||
use super::{is_valid_host_name, is_valid_node_id, AppState, AuthenticatedNode};
|
||||
use axum::{
|
||||
extract::{Extension, State},
|
||||
http::StatusCode,
|
||||
response::IntoResponse,
|
||||
Json,
|
||||
};
|
||||
use common::models::{NodeHeartbeatRequest, NodeRegisterRequest};
|
||||
use serde_json::json;
|
||||
use tracing::{info, warn};
|
||||
|
||||
pub async fn register_node(
|
||||
State(state): State<AppState>,
|
||||
auth_node: Option<Extension<AuthenticatedNode>>,
|
||||
Json(req): Json<NodeRegisterRequest>,
|
||||
) -> impl IntoResponse {
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
// 入口白名单校验
|
||||
if !is_valid_node_id(&req.node_id) {
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
"非法的节点 ID（仅允许字母、数字、点、下划线、连字符，长度 1-128）".to_string(),
|
||||
));
|
||||
}
|
||||
if !is_valid_host_name(&req.host_name) {
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
"非法的主机名（仅允许可打印 ASCII，长度 1-128）".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
// 已认证已拿到 Token 的节点刷新元数据配置
|
||||
if let Some(Extension(ref auth)) = auth_node {
|
||||
if auth.node_id == req.node_id {
|
||||
let _ = state.db.register_node(&req).await;
|
||||
info!("已授权节点 {} 刷新配置成功", req.node_id);
|
||||
return Ok((
|
||||
StatusCode::OK,
|
||||
Json(json!({
|
||||
"status": "approved",
|
||||
"message": "节点配置更新成功",
|
||||
"node_token": null,
|
||||
})),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
// 申请注册新节点（免凭据提交申请，进入 pending_approval 状态）
|
||||
match state.db.register_node(&req).await {
|
||||
Ok(_) => Json(json!({"status": "ok", "message": "节点注册成功"})),
|
||||
Err(e) => Json(json!({"status": "error", "message": e.to_string()})),
|
||||
Ok(true) => {
|
||||
info!(
|
||||
"接收到新节点 {} 的注册申请，已加入待审批 (pending_approval) 队列",
|
||||
req.node_id
|
||||
);
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
Json(json!({
|
||||
"status": "pending_approval",
|
||||
"message": "节点注册申请已成功提交！请在管理 Dashboard 控制台上点击【同意接入】授权该节点",
|
||||
"node_token": null,
|
||||
})),
|
||||
))
|
||||
}
|
||||
Ok(false) => {
|
||||
// 节点已处于待审批或已存在列表
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
Json(json!({
|
||||
"status": "pending_approval",
|
||||
"message": "节点注册申请等待管理员审批中",
|
||||
"node_token": null,
|
||||
})),
|
||||
))
|
||||
}
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
pub struct CheckNodeStatusRequest {
|
||||
pub node_id: String,
|
||||
}
|
||||
|
||||
/// POST /api/node/check_status — Node 端轮询检查审批结果。
|
||||
pub async fn check_node_status(
|
||||
State(state): State<AppState>,
|
||||
Json(req): Json<CheckNodeStatusRequest>,
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
if !is_valid_node_id(&req.node_id) {
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
"非法的节点 ID 参数".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
// 尝试拉取取走即焚的暂存明文 Token
|
||||
match state.db.take_pending_node_token(&req.node_id).await {
|
||||
Ok(Some(raw_token)) => {
|
||||
info!(
|
||||
"节点 {} 的注册申请已被管理员审批同意，下发专属 Token",
|
||||
req.node_id
|
||||
);
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
Json(json!({
|
||||
"status": "approved",
|
||||
"message": "节点已通过审批授权",
|
||||
"node_token": raw_token,
|
||||
})),
|
||||
))
|
||||
}
|
||||
Ok(None) | Err(_) => {
|
||||
// 查节点表状态
|
||||
match state.db.get_node_exists(&req.node_id).await {
|
||||
Ok(true) => Ok((
|
||||
StatusCode::OK,
|
||||
Json(json!({
|
||||
"status": "pending_approval",
|
||||
"message": "等待管理员在控制台点击同意",
|
||||
"node_token": null,
|
||||
})),
|
||||
)),
|
||||
_ => Ok((
|
||||
StatusCode::OK,
|
||||
Json(json!({
|
||||
"status": "rejected",
|
||||
"message": "节点注册申请未通过或已被移除",
|
||||
"node_token": null,
|
||||
})),
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn heartbeat_node(
|
||||
State(state): State<AppState>,
|
||||
Extension(auth_node): Extension<AuthenticatedNode>,
|
||||
Json(req): Json<NodeHeartbeatRequest>,
|
||||
) -> impl IntoResponse {
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
// 身份绑定校验：请求体声称的 node_id 必须与 token 绑定的 node_id 一致，
|
||||
// 杜绝「持有 A 节点 token 却冒充 B 节点发心跳」的跨节点越权。
|
||||
if req.node_id != auth_node.node_id {
|
||||
warn!(
|
||||
"节点心跳身份校验失败：token 绑定 node={}，但请求体声称 node_id={}",
|
||||
auth_node.node_id, req.node_id
|
||||
);
|
||||
return Err(crate::api::AppError::Forbidden(
|
||||
"node_id 与凭据不匹配".to_string(),
|
||||
));
|
||||
}
|
||||
match state.db.heartbeat_node(&req).await {
|
||||
Ok(_) => Json(json!({"status": "ok"})),
|
||||
Err(e) => Json(json!({"status": "error", "message": e.to_string()})),
|
||||
Ok(_) => Ok(Json(json!({"status": "ok"}))),
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,196 @@
|
||||
//! 鉴权失败速率限制中间件（防 token 在线暴力）。
|
||||
//!
|
||||
//! 设计：对返回 401 的请求按客户端 IP 维护滑动窗口失败计数。当某 IP 在窗口内
|
||||
//! 累计失败超过阈值，后续请求直接返回 429（持续到窗口内计数回落）。
|
||||
//!
|
||||
//! 仅作用于鉴权路径（与 auth_middleware 叠加），不影响已认证的正常业务流。
|
||||
//! 已认证请求返回 2xx，不计入失败窗口，因此合法节点/管理员的高频调用不受影响。
|
||||
//!
|
||||
//! IP 来源：优先取 `X-Forwarded-For` 首段（反代场景），回退到连接的 `ConnectInfo<SocketAddr>`
|
||||
//! （需 main.rs 用 `into_make_service_with_connect_info` 启动）。两者都拿不到时按"未知 IP"聚合。
|
||||
|
||||
use axum::{
|
||||
extract::{ConnectInfo, State},
|
||||
http::Request,
|
||||
middleware::Next,
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
use std::collections::{HashMap, VecDeque};
|
||||
use std::net::{IpAddr, SocketAddr};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{Duration, Instant};
|
||||
use tracing::warn;
|
||||
|
||||
/// 限流状态：按 IP 维护近窗口内的失败时间戳队列。
|
||||
#[derive(Clone)]
|
||||
pub struct RateLimiter {
|
||||
inner: Arc<Mutex<HashMap<IpAddr, VecDeque<Instant>>>>,
|
||||
window: Duration,
|
||||
max_failures: usize,
|
||||
/// 计数策略：
|
||||
/// - `false`（默认，通用 API 限流器）：仅对鉴权失败（400/401/403）的响应计数。
|
||||
/// - `true`（注册端点专用限流器）：对匹配路径（如 `/node/register`）的**所有**响应计数，
|
||||
/// 无论成败——这是对注册接口的独立节流设计，防止恶意频繁注册。
|
||||
///
|
||||
/// 历史问题：此前中间件对所有 `/node/register` 请求无条件计数，导致该 limiter 若复用为
|
||||
/// 通用 API 限流器时，20 次成功注册会把整个 IP 锁出所有 `/api/*` 端点（跨端点连锁）。
|
||||
/// 引入此标志把两种语义显式分离。
|
||||
count_all: bool,
|
||||
}
|
||||
|
||||
impl RateLimiter {
|
||||
/// 构造通用限流器：仅在鉴权失败（400/401/403）时计数。
|
||||
pub fn new(max_failures: usize, window: Duration) -> Self {
|
||||
Self {
|
||||
inner: Arc::new(Mutex::new(HashMap::new())),
|
||||
window,
|
||||
max_failures,
|
||||
count_all: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// 构造「全量计数」限流器：对匹配路径的所有响应（无论成败）计数。
|
||||
/// 用于注册端点专用节流。
|
||||
pub fn new_count_all(max_failures: usize, window: Duration) -> Self {
|
||||
Self {
|
||||
inner: Arc::new(Mutex::new(HashMap::new())),
|
||||
window,
|
||||
max_failures,
|
||||
count_all: true,
|
||||
}
|
||||
}
|
||||
|
||||
/// 检查该 IP 是否已被限流（窗口内失败次数超阈值）。不修改计数。
|
||||
pub(crate) fn is_rate_limited(&self, ip: IpAddr) -> bool {
|
||||
let now = Instant::now();
|
||||
let mut map = match self.inner.lock() {
|
||||
Ok(g) => g,
|
||||
Err(e) => e.into_inner(), // poisoned：仍尽力返回判断，避免鉴权因锁中毒全部放行
|
||||
};
|
||||
if let Some(queue) = map.get_mut(&ip) {
|
||||
// 清理过期时间戳
|
||||
while let Some(front) = queue.front() {
|
||||
if now.duration_since(*front) > self.window {
|
||||
queue.pop_front();
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
if queue.is_empty() {
|
||||
map.remove(&ip);
|
||||
return false;
|
||||
}
|
||||
return queue.len() >= self.max_failures;
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
/// 记录一次失败（追加时间戳）。
|
||||
pub(crate) fn record_failure(&self, ip: IpAddr) {
|
||||
let now = Instant::now();
|
||||
let mut map = match self.inner.lock() {
|
||||
Ok(g) => g,
|
||||
Err(e) => e.into_inner(),
|
||||
};
|
||||
let queue = map.entry(ip).or_default();
|
||||
queue.push_back(now);
|
||||
// 顺带清理，防止队列无限增长
|
||||
while let Some(front) = queue.front() {
|
||||
if now.duration_since(*front) > self.window {
|
||||
queue.pop_front();
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
if queue.is_empty() {
|
||||
map.remove(&ip);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 判断 IP 是否为本地环回或私有网段 IP。
|
||||
fn is_private_or_loopback_ip(ip: IpAddr) -> bool {
|
||||
match ip {
|
||||
IpAddr::V4(v4) => v4.is_loopback() || v4.is_private(),
|
||||
IpAddr::V6(v6) => v6.is_loopback(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 从请求中提取客户端 IP。
|
||||
/// 仅当底层连接 (ConnectInfo) 为本地环回或私有网段时才信任反向代理传递的 X-Forwarded-For / X-Real-IP。
|
||||
fn extract_client_ip(req: &Request<axum::body::Body>) -> Option<IpAddr> {
|
||||
let direct_ip = req
|
||||
.extensions()
|
||||
.get::<ConnectInfo<SocketAddr>>()
|
||||
.map(|ci| ci.0.ip());
|
||||
|
||||
// 如果有直连 IP 且不是私有/环回地址，说明未经过可信反代，直接返回直连 IP 拒绝盲信 X-Forwarded-For
|
||||
if let Some(ip) = direct_ip {
|
||||
if !is_private_or_loopback_ip(ip) {
|
||||
return Some(ip);
|
||||
}
|
||||
}
|
||||
|
||||
// 只有处于本地/私有网络反代之后时，才尝试提取 X-Forwarded-For
|
||||
if let Some(xff) = req
|
||||
.headers()
|
||||
.get("x-forwarded-for")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
{
|
||||
if let Some(first) = xff.split(',').map(|s| s.trim()).next() {
|
||||
if !first.is_empty() {
|
||||
if let Ok(ip) = first.parse::<IpAddr>() {
|
||||
return Some(ip);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// 回退：X-Real-IP
|
||||
if let Some(xri) = req.headers().get("x-real-ip").and_then(|v| v.to_str().ok()) {
|
||||
if let Ok(ip) = xri.parse::<IpAddr>() {
|
||||
return Some(ip);
|
||||
}
|
||||
}
|
||||
// 回退：直连 IP
|
||||
direct_ip
|
||||
}
|
||||
|
||||
/// 限流中间件：在鉴权之前检查该 IP 是否已被限流。
|
||||
///
|
||||
/// 放在 auth_middleware **之前**（外层）：被限流的 IP 直接 429，不进鉴权逻辑。
|
||||
/// 是否记入失败窗口，由 auth_middleware 的结果决定——为此 auth 中间件会把 401 的 IP
|
||||
/// 通过本 limiter 记录。但为避免跨中间件传参的复杂性，这里采用「先放行让 auth 判定，
|
||||
/// 若返回 401 再记录」的方式：见下方包装函数 `rate_limit_with_auth`。
|
||||
pub async fn rate_limit_middleware(
|
||||
State(limiter): State<RateLimiter>,
|
||||
req: Request<axum::body::Body>,
|
||||
next: Next,
|
||||
) -> Response {
|
||||
let ip = extract_client_ip(&req).unwrap_or(IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED));
|
||||
let is_register = req.uri().path().ends_with("/node/register");
|
||||
|
||||
if limiter.is_rate_limited(ip) {
|
||||
warn!("客户端 IP {} 鉴权失败次数过多，已限流（429）", ip);
|
||||
return (
|
||||
axum::http::StatusCode::TOO_MANY_REQUESTS,
|
||||
"鉴权失败次数过多，请稍后重试",
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
|
||||
let resp = next.run(req).await;
|
||||
|
||||
// 计入速率窗口的条件：
|
||||
// - 鉴权失败（401/403/400）：通用与专用限流器都计；
|
||||
// - 或 limiter 配置为 count_all 且请求落在专用节流路径（如 /node/register）：
|
||||
// 这种情况下成功响应也计，作为对注册接口本身的独立节流（防恶意频繁注册）。
|
||||
// 通用 API 限流器（count_all=false）不会因 is_register 把成功请求计入，
|
||||
// 避免了「成功注册连锁锁出整个 /api/*」的历史缺陷。
|
||||
let status = resp.status().as_u16();
|
||||
let auth_failed = status == 401 || status == 403 || status == 400;
|
||||
if auth_failed || (limiter.count_all && is_register) {
|
||||
limiter.record_failure(ip);
|
||||
}
|
||||
|
||||
resp
|
||||
}
|
||||
@@ -2,7 +2,7 @@ use super::AppState;
|
||||
use axum::{
|
||||
body::Body,
|
||||
extract::{Path as AxumPath, State},
|
||||
http::{header, StatusCode},
|
||||
http::header,
|
||||
response::Response,
|
||||
};
|
||||
use tokio::fs::File;
|
||||
@@ -12,13 +12,20 @@ use tracing::warn;
|
||||
pub async fn download_seed(
|
||||
State(state): State<AppState>,
|
||||
AxumPath(name): AxumPath<String>,
|
||||
) -> Response {
|
||||
if name.is_empty() || name.starts_with('.') || !name.chars().all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-' || c == '+' || c == '@') {
|
||||
warn!("拒绝可能包含路径穿越或特别注入序列的非法种子下载请求: {}", name);
|
||||
return Response::builder()
|
||||
.status(StatusCode::BAD_REQUEST)
|
||||
.body(Body::from("非法的种子名称参数"))
|
||||
.unwrap();
|
||||
) -> Result<Response, crate::api::AppError> {
|
||||
if name.is_empty()
|
||||
|| name.starts_with('.')
|
||||
|| !name.chars().all(|c| {
|
||||
c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-' || c == '+' || c == '@'
|
||||
})
|
||||
{
|
||||
warn!(
|
||||
"拒绝可能包含路径穿越或特别注入序列的非法种子下载请求: {}",
|
||||
name
|
||||
);
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
"非法的种子名称参数".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let seed_file_path = std::path::Path::new(&state.results_dir)
|
||||
@@ -27,32 +34,27 @@ pub async fn download_seed(
|
||||
|
||||
if !seed_file_path.is_file() {
|
||||
warn!("客户端请求的种子文件不存在: {}", seed_file_path.display());
|
||||
return Response::builder()
|
||||
.status(StatusCode::NOT_FOUND)
|
||||
.body(Body::from("请求的种子文件不存在"))
|
||||
.unwrap();
|
||||
return Err(crate::api::AppError::NotFound(
|
||||
"请求的种子文件不存在".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let file = match File::open(&seed_file_path).await {
|
||||
Ok(file) => file,
|
||||
Err(_) => {
|
||||
return Response::builder()
|
||||
.status(StatusCode::INTERNAL_SERVER_ERROR)
|
||||
.body(Body::from("无法打开种子文件"))
|
||||
.unwrap();
|
||||
Err(e) => {
|
||||
return Err(crate::api::AppError::Internal(e.into()));
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
let stream = ReaderStream::new(file);
|
||||
let body = Body::from_stream(stream);
|
||||
|
||||
Response::builder()
|
||||
Ok(Response::builder()
|
||||
.header(header::CONTENT_TYPE, "application/octet-stream")
|
||||
.header(
|
||||
header::CONTENT_DISPOSITION,
|
||||
format!("attachment; filename=\"{}.7\"", name),
|
||||
)
|
||||
.body(body)
|
||||
.unwrap()
|
||||
.unwrap())
|
||||
}
|
||||
|
||||
@@ -2,21 +2,37 @@ use super::AppState;
|
||||
use axum::{extract::State, response::IntoResponse, Json};
|
||||
use serde_json::json;
|
||||
|
||||
pub async fn get_status(State(state): State<AppState>) -> impl IntoResponse {
|
||||
/// 轻量健康检查端点（不走鉴权）。
|
||||
///
|
||||
/// 供 docker healthcheck、负载均衡、外部监控探测。刻意只返回固定 ok，
|
||||
/// 不触碰数据库或调度器，避免健康检查本身拖累系统或因 DB 瞬时锁导致误判不健康。
|
||||
pub async fn healthz() -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
Ok(Json(json!({ "status": "ok" })))
|
||||
}
|
||||
|
||||
pub async fn get_status(
|
||||
State(state): State<AppState>,
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
let nodes = state.db.get_active_nodes().await.unwrap_or_default();
|
||||
let total_active_slots: i32 = nodes.iter().map(|n| n.active_slots).sum();
|
||||
let total_max_slots: i32 = nodes.iter().map(|n| n.max_slots).sum();
|
||||
|
||||
let grid_stats = state.db.get_grid_summary_stats().await.unwrap_or(serde_json::json!({
|
||||
"total": 0, "pending": 0, "running": 0, "converged": 0, "failed": 0
|
||||
}));
|
||||
// dashboard 全局概览：聚合全部工作流的 grid_points（多工作流分区后仍提供全局合计）。
|
||||
// 若需单工作流进度，可扩展为按 workflow 查询参数分别聚合。
|
||||
let grid_stats = state
|
||||
.db
|
||||
.get_grid_summary_stats(None)
|
||||
.await
|
||||
.unwrap_or(serde_json::json!({
|
||||
"total": 0, "pending": 0, "running": 0, "converged": 0, "failed": 0
|
||||
}));
|
||||
|
||||
Json(json!({
|
||||
Ok(Json(json!({
|
||||
"status": "online",
|
||||
"nodes_online": nodes.len(),
|
||||
"total_active_slots": total_active_slots,
|
||||
"total_max_slots": total_max_slots,
|
||||
"nodes": nodes,
|
||||
"grid_stats": grid_stats,
|
||||
}))
|
||||
})))
|
||||
}
|
||||
|
||||
+125
-54
@@ -1,6 +1,6 @@
|
||||
use super::AppState;
|
||||
use super::{AppState, AuthenticatedNode};
|
||||
use axum::{
|
||||
extract::{Multipart, State},
|
||||
extract::{Extension, Multipart, State},
|
||||
response::IntoResponse,
|
||||
Json,
|
||||
};
|
||||
@@ -12,27 +12,38 @@ use tracing::{info, warn};
|
||||
|
||||
use axum::http::StatusCode;
|
||||
|
||||
pub async fn claim_task(State(state): State<AppState>) -> impl IntoResponse {
|
||||
match state.queue.pop_task().await {
|
||||
pub async fn claim_task(
|
||||
State(state): State<AppState>,
|
||||
Extension(auth_node): Extension<AuthenticatedNode>,
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
// 领用时记录任务归属：pop_task 写入 claimed_by_node_id，
|
||||
// report 阶段据此校验「上报者确为领用者」，杜绝跨节点伪造结果。
|
||||
match state.queue.pop_task(&auth_node.node_id).await {
|
||||
Ok(Some(task)) => {
|
||||
if let Err(e) = state.db.mark_grid_point_running(&task.point_name).await {
|
||||
// 多工作流分区：mark_grid_point_running 须带 workflow_name，避免按 name 全局更新
|
||||
// 误改其他工作流的同名点。TaskSpec.workflow_name 在调度时已绑定。
|
||||
let wf = task.workflow_name.as_deref().unwrap_or("");
|
||||
if let Err(e) = state.db.mark_grid_point_running(&task.point_name, wf).await {
|
||||
warn!("领用任务 {} 后同步变更为 running 状态遇到异常: {}. 后置 stale 定时自取检索引索将介入修复维护", task.task_id, e);
|
||||
}
|
||||
(StatusCode::OK, Json(json!({"status": "ok", "task": task}))).into_response()
|
||||
Ok((StatusCode::OK, Json(json!({"status": "ok", "task": task}))))
|
||||
}
|
||||
Ok(None) => Ok((
|
||||
StatusCode::OK,
|
||||
Json(json!({"status": "empty", "task": null})),
|
||||
)),
|
||||
Err(e) => {
|
||||
tracing::error!("领用任务数据库异常: {}", e);
|
||||
Err(crate::api::AppError::Internal(e))
|
||||
}
|
||||
Ok(None) => (StatusCode::OK, Json(json!({"status": "empty", "task": null}))).into_response(),
|
||||
Err(e) => (
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(json!({"status": "error", "message": format!("领用任务失败: {}", e)})),
|
||||
)
|
||||
.into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn report_task(
|
||||
State(state): State<AppState>,
|
||||
Extension(auth_node): Extension<AuthenticatedNode>,
|
||||
mut multipart: Multipart,
|
||||
) -> impl IntoResponse {
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
let mut report_json: Option<TaskReport> = None;
|
||||
let mut seed_file_data: Option<Vec<u8>> = None;
|
||||
|
||||
@@ -51,53 +62,98 @@ pub async fn report_task(
|
||||
}
|
||||
}
|
||||
|
||||
let report = match report_json {
|
||||
let mut report = match report_json {
|
||||
Some(r) => r,
|
||||
None => {
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(json!({"status": "error", "message": "请求中缺少 report 字段"})),
|
||||
)
|
||||
.into_response();
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
"请求中缺少 report 字段".to_string(),
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
// ── 任务归属校验（S1 核心，防跨节点伪造结果投毒）──
|
||||
// 1. 该 task_id 必须由当前鉴权 node 领用（claim 时记录的 claimed_by_node_id 匹配）。
|
||||
// 2. 上报的 point_name 必须与该 task 绑定的 point_name 一致（防跨点上报）。
|
||||
// 3. 忽略 body 里声称的 node_id，统一以鉴权 node_id 写库（修复审计归因断裂）。
|
||||
// 4. 取 task 绑定的 workflow_name，用于定向更新该工作流的 grid_points（多工作流分区）。
|
||||
let (claimed_point, claimed_workflow) = match state
|
||||
.queue
|
||||
.verify_task_claim(&report.task_id.to_string(), &auth_node.node_id)
|
||||
.await
|
||||
{
|
||||
Ok(Some((p, w))) => (p, w),
|
||||
Ok(None) => {
|
||||
warn!(
|
||||
"任务归属校验失败：node={} 上报 task_id={} 但未领用或已被清理",
|
||||
auth_node.node_id, report.task_id
|
||||
);
|
||||
return Err(crate::api::AppError::Forbidden(
|
||||
"任务未由本节点领用或已上报过".to_string(),
|
||||
));
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!("校验任务归属数据库异常: {}", e);
|
||||
return Err(crate::api::AppError::Internal(e));
|
||||
}
|
||||
};
|
||||
if claimed_point != report.point_name {
|
||||
warn!(
|
||||
"任务点名校验失败：task_id={} 领用 point={} 但上报 point={}",
|
||||
report.task_id, claimed_point, report.point_name
|
||||
);
|
||||
return Err(crate::api::AppError::Forbidden(
|
||||
"上报的网格点与领用任务不匹配".to_string(),
|
||||
));
|
||||
}
|
||||
// 统一以鉴权 node_id 覆盖 body 里的 node_id，保证归因可信
|
||||
report.node_id = auth_node.node_id.clone();
|
||||
// workflow_name 以领用记录为准（claim 时从 TaskSpec 落库），body 无权声称。
|
||||
let workflow_name = claimed_workflow.unwrap_or_default();
|
||||
|
||||
let name = report.point_name.clone();
|
||||
|
||||
if name.is_empty() || name.starts_with('.') || !name.chars().all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-' || c == '+' || c == '@') {
|
||||
warn!("拒绝可能包含路径穿越或特殊非常规编码号攻击的网格点名称请求: {}", name);
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(json!({"status": "error", "message": "非法的网格点名称参数"})),
|
||||
)
|
||||
.into_response();
|
||||
if name.is_empty()
|
||||
|| name.starts_with('.')
|
||||
|| !name.chars().all(|c| {
|
||||
c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-' || c == '+' || c == '@'
|
||||
})
|
||||
{
|
||||
warn!(
|
||||
"拒绝可能包含路径穿越或特殊非常规编码号攻击的网格点名称请求: {}",
|
||||
name
|
||||
);
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
"非法的网格点名称参数".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let params = match extract_params(&report) {
|
||||
Some(p) => p,
|
||||
None => {
|
||||
warn!("网格点 {} 汇报数据解析失败: 无法解析 params 或 summary_json", name);
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(json!({"status": "error", "message": "无法解析 params 或 summary_json"})),
|
||||
)
|
||||
.into_response();
|
||||
warn!(
|
||||
"网格点 {} 汇报数据解析失败: 无法解析 params 或 summary_json",
|
||||
name
|
||||
);
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
"无法解析 params 或 summary_json".to_string(),
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
// Record in DB
|
||||
if let Err(e) = state.db.record_task_report(&report).await {
|
||||
warn!("记录网格点 {} 任务结果到数据库失败: {}", name, e);
|
||||
return (
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(json!({"status": "error", "message": format!("记录数据库失败: {}", e)})),
|
||||
)
|
||||
.into_response();
|
||||
// Record in DB（带 workflow_name 定向更新该工作流的 grid_points）
|
||||
if let Err(e) = state.db.record_task_report(&report, &workflow_name).await {
|
||||
// DB 错误细节进日志，对客户端只返回通用消息（避免泄露表结构/内部错误给未授权方）
|
||||
tracing::error!("记录网格点 {} 任务结果到数据库失败: {}", name, e);
|
||||
return Err(crate::api::AppError::Internal(e));
|
||||
}
|
||||
|
||||
// Clean up task from task_queue table to prevent queue DB bloat
|
||||
if let Err(e) = state.queue.remove_task(&report.task_id.to_string()).await {
|
||||
tracing::warn!("从任务队列中清理已上报任务记录 {} 失败: {}", report.task_id, e);
|
||||
tracing::warn!(
|
||||
"从任务队列中清理已上报任务记录 {} 失败: {}",
|
||||
report.task_id,
|
||||
e
|
||||
);
|
||||
}
|
||||
|
||||
// 采用原子写入模式保持 conv.json 与核心二进制数据完整落地后才揭晓真实文件名
|
||||
@@ -112,30 +168,46 @@ pub async fn report_task(
|
||||
// Save seed file .7 using atomic temporary writing strategy
|
||||
if report.converged && !report.atmosphere_has_nan {
|
||||
if let Some(bytes) = seed_file_data {
|
||||
let seed_tmp = model_dir.join(format!("{}.7.{}.tmp", name, uuid::Uuid::new_v4().simple()));
|
||||
let seed_tmp =
|
||||
model_dir.join(format!("{}.7.{}.tmp", name, uuid::Uuid::new_v4().simple()));
|
||||
let seed_path = model_dir.join(format!("{}.7", name));
|
||||
if fs::write(&seed_tmp, bytes).await.is_ok() {
|
||||
if fs::rename(&seed_tmp, &seed_path).await.is_ok() {
|
||||
info!("成功保持原子写入落地并保存网格点 {} 的收敛种子文件: {}", name, seed_path.display());
|
||||
let _ = state
|
||||
.db
|
||||
.insert_seed(¶ms, &seed_path.to_string_lossy())
|
||||
.await;
|
||||
}
|
||||
if fs::write(&seed_tmp, bytes).await.is_ok()
|
||||
&& fs::rename(&seed_tmp, &seed_path).await.is_ok()
|
||||
{
|
||||
info!(
|
||||
"成功保持原子写入落地并保存网格点 {} 的收敛种子文件: {}",
|
||||
name,
|
||||
seed_path.display()
|
||||
);
|
||||
let _ = state
|
||||
.db
|
||||
.insert_seed(¶ms, &seed_path.to_string_lossy())
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if report.status == TaskStatus::Failed || report.status == TaskStatus::Timeout || report.atmosphere_has_nan {
|
||||
if !report.converged
|
||||
|| report.atmosphere_has_nan
|
||||
|| report.status == TaskStatus::Failed
|
||||
|| report.status == TaskStatus::Timeout
|
||||
{
|
||||
// Task did not succeed -> check if seed_step fallback should be triggered
|
||||
info!("网格点 {} 计算未成功完成，检查种子回退机制...", name);
|
||||
if let Err(e) = state.scheduler.trigger_seed_step_fallback(¶ms).await {
|
||||
if let Err(e) = state
|
||||
.scheduler
|
||||
.trigger_seed_step_fallback(¶ms, &workflow_name)
|
||||
.await
|
||||
{
|
||||
warn!("网格点 {} 触发种子回退机制失败: {}", name, e);
|
||||
}
|
||||
}
|
||||
|
||||
(StatusCode::OK, Json(json!({"status": "ok", "message": "上报成功"}))).into_response()
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
Json(json!({"status": "ok", "message": "上报成功"})),
|
||||
))
|
||||
}
|
||||
|
||||
fn extract_params(report: &TaskReport) -> Option<GridPointParams> {
|
||||
@@ -146,4 +218,3 @@ fn extract_params(report: &TaskReport) -> Option<GridPointParams> {
|
||||
.ok()
|
||||
.map(|summary| summary.params)
|
||||
}
|
||||
|
||||
|
||||
+141
-138
@@ -23,162 +23,166 @@ pub struct ApiResponse<T> {
|
||||
pub data: Option<T>,
|
||||
}
|
||||
|
||||
pub async fn list_workflows(State(state): State<AppState>) -> impl IntoResponse {
|
||||
pub async fn list_workflows(
|
||||
State(state): State<AppState>,
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
match state.db.list_workflows().await {
|
||||
Ok(list) => (StatusCode::OK, Json(ApiResponse { success: true, message: "成功获取工作流列表".to_string(), data: Some(list) })),
|
||||
Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, Json(ApiResponse { success: false, message: format!("获取工作流列表失败: {}", e), data: None })),
|
||||
Ok(list) => Ok((
|
||||
StatusCode::OK,
|
||||
Json(ApiResponse {
|
||||
success: true,
|
||||
message: "成功获取工作流列表".to_string(),
|
||||
data: Some(list),
|
||||
}),
|
||||
)),
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn get_workflow(
|
||||
State(state): State<AppState>,
|
||||
AxumPath(name): AxumPath<String>,
|
||||
) -> impl IntoResponse {
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
match state.db.get_workflow(&name).await {
|
||||
Ok(Some(item)) => (StatusCode::OK, Json(ApiResponse { success: true, message: "成功获取工作流详情".to_string(), data: Some(item) })),
|
||||
Ok(None) => (StatusCode::NOT_FOUND, Json(ApiResponse { success: false, message: format!("工作流 '{}' 未找到", name), data: None })),
|
||||
Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, Json(ApiResponse { success: false, message: format!("获取工作流详情失败: {}", e), data: None })),
|
||||
Ok(Some(item)) => Ok((
|
||||
StatusCode::OK,
|
||||
Json(ApiResponse {
|
||||
success: true,
|
||||
message: "成功获取工作流详情".to_string(),
|
||||
data: Some(item),
|
||||
}),
|
||||
)),
|
||||
Ok(None) => Err(crate::api::AppError::NotFound(format!(
|
||||
"工作流 '{}' 未找到",
|
||||
name
|
||||
))),
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
|
||||
/// 工作流名称白名单：仅允许字母、数字、点、下划线、连字符，长度 1-64。
|
||||
/// 与 report_task/download_seed 的网格点名校验口径保持一致，从源头阻止
|
||||
/// 名称携带 HTML/JS 特殊字符进入 Dashboard 渲染（存储型 XSS 根因之一）。
|
||||
fn is_valid_workflow_name(name: &str) -> bool {
|
||||
!name.is_empty()
|
||||
&& name.len() <= 64
|
||||
&& name
|
||||
.chars()
|
||||
.all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-')
|
||||
}
|
||||
|
||||
pub async fn save_workflow(
|
||||
State(state): State<AppState>,
|
||||
Json(req): Json<CreateWorkflowRequest>,
|
||||
) -> impl IntoResponse {
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
// 名称白名单校验（优先于 YAML 校验，拒绝携带特殊字符的名称）
|
||||
if !is_valid_workflow_name(&req.name) {
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
"工作流名称仅允许字母、数字、点(.)、下划线(_)、连字符(-)，长度 1-64".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
// Validate YAML config string
|
||||
if let Err(e) = serde_yaml::from_str::<GridConfig>(&req.config_yaml) {
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("无效的 YAML 配置: {}", e),
|
||||
data: None,
|
||||
}),
|
||||
);
|
||||
return Err(crate::api::AppError::BadRequest(format!(
|
||||
"无效的 YAML 配置: {}",
|
||||
e
|
||||
)));
|
||||
}
|
||||
|
||||
// 检查被编辑的工作流是否正处于激活运行中
|
||||
if let Ok(Some(existing)) = state.db.get_workflow(&req.name).await {
|
||||
if existing.status == "running" || existing.status == "initializing" {
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("工作流 '{}' 正处在运行或初始加载流程中，严禁原地覆写参数重设至 IDLE；如待变更参数请先调 API 显式触发停止后再保存", req.name),
|
||||
data: None,
|
||||
}),
|
||||
);
|
||||
return Err(crate::api::AppError::BadRequest(
|
||||
format!("工作流 '{}' 正处在运行或初始加载流程中，严禁原地覆写参数重设至 IDLE；如待变更参数请先调 API 显式触发停止后再保存", req.name)
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
match state.db.upsert_workflow(&req.name, req.description.as_deref(), &req.config_yaml, "idle").await {
|
||||
match state
|
||||
.db
|
||||
.upsert_workflow(
|
||||
&req.name,
|
||||
req.description.as_deref(),
|
||||
&req.config_yaml,
|
||||
"idle",
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(_) => {
|
||||
info!("成功注册/更新工作流配置: {}", req.name);
|
||||
(
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
Json(ApiResponse::<()> {
|
||||
success: true,
|
||||
message: format!("工作流 '{}' 保存成功", req.name),
|
||||
data: None,
|
||||
}),
|
||||
)
|
||||
))
|
||||
}
|
||||
Err(e) => (
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("保存工作流失败: {}", e),
|
||||
data: None,
|
||||
}),
|
||||
),
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn delete_workflow(
|
||||
State(state): State<AppState>,
|
||||
AxumPath(name): AxumPath<String>,
|
||||
) -> impl IntoResponse {
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
// 拦截正在运行或初始加载中的工作流删除请求
|
||||
if let Ok(Some(existing)) = state.db.get_workflow(&name).await {
|
||||
if existing.status == "running" || existing.status == "initializing" {
|
||||
return Err(crate::api::AppError::BadRequest(format!(
|
||||
"工作流 '{}' 当前处于 '{}' 状态，无法直接删除。请先显式暂停/停止该工作流。",
|
||||
name, existing.status
|
||||
)));
|
||||
}
|
||||
}
|
||||
|
||||
let _ = state.queue.clear_queue_by_workflow(&name).await;
|
||||
match state.db.delete_workflow(&name).await {
|
||||
Ok(_) => (
|
||||
Ok(_) => Ok((
|
||||
StatusCode::OK,
|
||||
Json(ApiResponse::<()> {
|
||||
success: true,
|
||||
message: format!("工作流 '{}' 已删除", name),
|
||||
data: None,
|
||||
}),
|
||||
),
|
||||
Err(e) => (
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("删除工作流失败: {}", e),
|
||||
data: None,
|
||||
}),
|
||||
),
|
||||
)),
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn start_workflow(
|
||||
State(state): State<AppState>,
|
||||
AxumPath(name): AxumPath<String>,
|
||||
) -> impl IntoResponse {
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
let item = match state.db.get_workflow(&name).await {
|
||||
Ok(Some(item)) => item,
|
||||
Ok(None) => {
|
||||
return (
|
||||
StatusCode::NOT_FOUND,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("工作流 '{}' 未找到", name),
|
||||
data: None,
|
||||
}),
|
||||
)
|
||||
}
|
||||
Err(e) => {
|
||||
return (
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("获取工作流失败: {}", e),
|
||||
data: None,
|
||||
}),
|
||||
)
|
||||
return Err(crate::api::AppError::NotFound(format!(
|
||||
"工作流 '{}' 未找到",
|
||||
name
|
||||
)))
|
||||
}
|
||||
Err(e) => return Err(e.into()),
|
||||
};
|
||||
|
||||
if item.status == "running" || item.status == "initializing" {
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("工作流 '{}' 已处在初始建立状态中或者已处于运行状态，无需且不允许进行并行重置启动", name),
|
||||
data: None,
|
||||
}),
|
||||
);
|
||||
return Err(crate::api::AppError::BadRequest(format!(
|
||||
"工作流 '{}' 已处在初始建立状态中或者已处于运行状态，无需且不允许进行并行重置启动",
|
||||
name
|
||||
)));
|
||||
}
|
||||
|
||||
// 通过原子性抢占更新将状态切换为 initializing，拦截同名流上的多并发调用导致的双重加载破坏性竞态
|
||||
match state.db.transition_workflow_to_initializing(&name).await {
|
||||
Ok(false) => {
|
||||
return (
|
||||
StatusCode::CONFLICT,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("工作流 '{}' 初始化抢占挂起异常，表明已在另一会话上下文中顺利推入启动通道", name),
|
||||
data: None,
|
||||
}),
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
return (
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("原子化抢占和迁移工作流状态发生异常: {}", e),
|
||||
data: None,
|
||||
}),
|
||||
);
|
||||
return Err(crate::api::AppError::Conflict(format!(
|
||||
"工作流 '{}' 初始化抢占挂起异常，表明已在另一会话上下文中顺利推入启动通道",
|
||||
name
|
||||
)));
|
||||
}
|
||||
Err(e) => return Err(e.into()),
|
||||
Ok(true) => {}
|
||||
}
|
||||
|
||||
@@ -186,70 +190,69 @@ pub async fn start_workflow(
|
||||
Ok(cfg) => cfg,
|
||||
Err(e) => {
|
||||
let _ = state.db.update_workflow_status(&name, "idle").await;
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("解析工作流 YAML 发生语法或参数解析异常: {}", e),
|
||||
data: None,
|
||||
}),
|
||||
);
|
||||
return Err(crate::api::AppError::BadRequest(format!(
|
||||
"解析工作流 YAML 发生语法或参数解析异常: {}",
|
||||
e
|
||||
)));
|
||||
}
|
||||
};
|
||||
|
||||
info!("成功占据独享启动权，开始启动工作流 '{}'，系统进行 64/32 维深度平展开网格结构计算化推列并推送队列...", name);
|
||||
match state.scheduler.initialize_grid(&grid_cfg).await {
|
||||
Ok(_) => {
|
||||
let _ = state.db.update_workflow_status(&name, "running").await;
|
||||
let _ = state.scheduler.schedule_pending_tasks().await;
|
||||
(
|
||||
StatusCode::OK,
|
||||
Json(ApiResponse::<()> {
|
||||
success: true,
|
||||
message: format!("工作流 '{}' 建立与挂载成功并已接续排班", name),
|
||||
data: None,
|
||||
}),
|
||||
)
|
||||
info!("成功占据独享启动权，开始异步启动工作流 '{}'，系统将在后台进行 64/32 维深度平展开网格结构计算化推列并推送队列...", name);
|
||||
|
||||
let bg_state = state.clone();
|
||||
let bg_name = name.clone();
|
||||
let bg_grid_cfg = grid_cfg;
|
||||
|
||||
tokio::spawn(async move {
|
||||
match bg_state
|
||||
.scheduler
|
||||
.initialize_grid(&bg_grid_cfg, &bg_name)
|
||||
.await
|
||||
{
|
||||
Ok(_) => {
|
||||
let _ = bg_state
|
||||
.db
|
||||
.update_workflow_status(&bg_name, "running")
|
||||
.await;
|
||||
let _ = bg_state.scheduler.schedule_pending_tasks().await;
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("工作流 {} 网格初始化中途失败，已回退为 idle；已写入的点保留，重新启动会幂等补齐: {}", bg_name, e);
|
||||
let _ = bg_state.db.update_workflow_status(&bg_name, "idle").await;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = state.db.update_workflow_status(&name, "idle").await;
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("展开与挂载初始化任务点到系统队列失败: {}", e),
|
||||
data: None,
|
||||
}),
|
||||
)
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
Json(ApiResponse::<()> {
|
||||
success: true,
|
||||
message: format!("工作流 '{}' 已进入后台异步建立与挂载流程", name),
|
||||
data: None,
|
||||
}),
|
||||
))
|
||||
}
|
||||
|
||||
pub async fn stop_workflow(
|
||||
State(state): State<AppState>,
|
||||
AxumPath(name): AxumPath<String>,
|
||||
) -> impl IntoResponse {
|
||||
) -> Result<impl IntoResponse, crate::api::AppError> {
|
||||
match state.db.update_workflow_status(&name, "paused").await {
|
||||
Ok(_) => {
|
||||
let _ = state.queue.clear_queue().await;
|
||||
let _ = state.db.reset_queued_grid_points_to_pending().await;
|
||||
(
|
||||
// 多工作流分区：清理与重置都限定在本工作流内，避免误伤其他并发运行的工作流。
|
||||
// - clear_queue_by_workflow：只删本工作流的排队任务。
|
||||
// - reset_queued_grid_points_to_pending(&name)：只把本工作流的 queued 点打回 pending。
|
||||
let _ = state.queue.clear_queue_by_workflow(&name).await;
|
||||
let _ = state.db.reset_queued_grid_points_to_pending(&name).await;
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
Json(ApiResponse::<()> {
|
||||
success: true,
|
||||
message: format!("工作流 '{}' 已暂停，排队任务已暂停调度", name),
|
||||
data: None,
|
||||
}),
|
||||
)
|
||||
))
|
||||
}
|
||||
Err(e) => (
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(ApiResponse::<()> {
|
||||
success: false,
|
||||
message: format!("暂停工作流失败: {}", e),
|
||||
data: None,
|
||||
}),
|
||||
),
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,58 @@
|
||||
use axum::http::HeaderValue;
|
||||
use tower_http::cors::{AllowOrigin, CorsLayer};
|
||||
use tracing::info;
|
||||
|
||||
/// 构建 CORS 中间件层。
|
||||
///
|
||||
/// 严格安全策略：仅允许**同源**（Origin 匹配请求头的 Host）或**本地 Origin**（localhost / 127.0.0.1 / [::1]）。
|
||||
pub fn build_cors_layer() -> CorsLayer {
|
||||
info!("CORS 策略：仅允许同源或本地 Origin（localhost / 127.0.0.1 / [::1]）");
|
||||
|
||||
CorsLayer::new()
|
||||
.allow_origin(AllowOrigin::predicate(
|
||||
|origin: &HeaderValue, head: &axum::http::request::Parts| {
|
||||
let Ok(origin_str) = origin.to_str() else {
|
||||
return false;
|
||||
};
|
||||
let Ok(uri) = origin_str.parse::<axum::http::Uri>() else {
|
||||
return false;
|
||||
};
|
||||
let Some(host) = uri.host() else {
|
||||
return false;
|
||||
};
|
||||
let clean_host = host.trim_start_matches('[').trim_end_matches(']');
|
||||
|
||||
// 1. 本地来源 (localhost / 127.0.0.1 / [::1])
|
||||
if clean_host == "localhost"
|
||||
|| clean_host == "127.0.0.1"
|
||||
|| clean_host == "::1"
|
||||
|| clean_host.starts_with("127.")
|
||||
{
|
||||
return true;
|
||||
}
|
||||
|
||||
// 2. 同源来源 (Origin 匹配请求头的 Host)
|
||||
if let Some(host_header) = head.headers.get(axum::http::header::HOST) {
|
||||
if let Ok(host_str) = host_header.to_str() {
|
||||
if let Some(authority) = uri.authority() {
|
||||
if authority.as_str().eq_ignore_ascii_case(host_str) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
false
|
||||
},
|
||||
))
|
||||
.allow_methods([
|
||||
axum::http::Method::GET,
|
||||
axum::http::Method::POST,
|
||||
axum::http::Method::PUT,
|
||||
axum::http::Method::DELETE,
|
||||
])
|
||||
.allow_headers([
|
||||
axum::http::header::AUTHORIZATION,
|
||||
axum::http::header::CONTENT_TYPE,
|
||||
])
|
||||
}
|
||||
+1354
-106
File diff suppressed because it is too large
Load Diff
@@ -1,4 +1,4 @@
|
||||
pub mod api;
|
||||
pub mod cors;
|
||||
pub mod db;
|
||||
pub mod scheduler;
|
||||
|
||||
|
||||
+280
-43
@@ -3,8 +3,9 @@ use server::api::{self, AppState};
|
||||
use server::db::Database;
|
||||
use server::scheduler::GridScheduler;
|
||||
|
||||
|
||||
use axum::{
|
||||
extract::DefaultBodyLimit,
|
||||
http::HeaderValue,
|
||||
routing::{get, post},
|
||||
Router,
|
||||
};
|
||||
@@ -16,12 +17,15 @@ use std::net::SocketAddr;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
use tokio::time::{sleep, Duration};
|
||||
use tower_http::cors::CorsLayer;
|
||||
use tower_http::services::{ServeDir, ServeFile};
|
||||
use tracing::info;
|
||||
|
||||
#[derive(Parser, Debug)]
|
||||
#[command(name = "server", version = "0.1.0", about = "Distributed Computing TLUSTY/SYNSPEC (DCTS) Server")]
|
||||
#[command(
|
||||
name = "server",
|
||||
version = "0.1.0",
|
||||
about = "Distributed Computing TLUSTY/SYNSPEC (DCTS) Server"
|
||||
)]
|
||||
struct CliArgs {
|
||||
/// Optional path to workflow configuration YAML file to auto-register on startup
|
||||
#[arg(short = 'w', long = "workflow")]
|
||||
@@ -61,12 +65,15 @@ async fn main() -> Result<()> {
|
||||
let default_wf_path = Path::new(&server_cfg.grid_config);
|
||||
if default_wf_path.is_file() {
|
||||
if let Ok(yaml_content) = std::fs::read_to_string(default_wf_path) {
|
||||
if let Err(e) = db.upsert_workflow(
|
||||
"sdB_cno",
|
||||
Some("sdB CNO 6D Stellar Atmosphere Grid"),
|
||||
&yaml_content,
|
||||
"idle",
|
||||
).await {
|
||||
if let Err(e) = db
|
||||
.upsert_workflow(
|
||||
"sdB_cno",
|
||||
Some("sdB CNO 6D Stellar Atmosphere Grid"),
|
||||
&yaml_content,
|
||||
"idle",
|
||||
)
|
||||
.await
|
||||
{
|
||||
tracing::warn!("预注册默认工作流失败: {}", e);
|
||||
} else {
|
||||
info!("已在数据库中成功预注册默认工作流 'sdB_cno'");
|
||||
@@ -74,15 +81,38 @@ async fn main() -> Result<()> {
|
||||
}
|
||||
}
|
||||
|
||||
// 弱口令凭据安全警告检测
|
||||
let is_weak_token = |t: Option<&str>| -> bool {
|
||||
match t {
|
||||
Some(s) => {
|
||||
s.len() < 12 || s == "fmqi123" || s == "admin" || s == "123456" || s == "secret"
|
||||
}
|
||||
None => false,
|
||||
}
|
||||
};
|
||||
if is_weak_token(server_cfg.auth_token.as_deref())
|
||||
|| is_weak_token(server_cfg.admin_token.as_deref())
|
||||
{
|
||||
tracing::warn!("⚠️ 检测到系统当前正在使用弱口令凭据或默认 Token！建议生产环境在 .env 中配置使用 openssl rand -hex 32 生成的高强度 Token！");
|
||||
}
|
||||
|
||||
let rate_limiter = api::rate_limit::RateLimiter::new(5, std::time::Duration::from_secs(300));
|
||||
|
||||
let state = AppState {
|
||||
db,
|
||||
queue: queue.clone(),
|
||||
scheduler: scheduler.clone(),
|
||||
results_dir: server_cfg.results_dir,
|
||||
results_dir: server_cfg.results_dir.clone(),
|
||||
rate_limiter,
|
||||
auth_token: server_cfg.auth_token.clone(),
|
||||
admin_token: server_cfg.admin_token.clone(),
|
||||
auth_disabled: server_cfg.auth_disabled,
|
||||
admin_sessions: std::sync::Arc::new(tokio::sync::RwLock::new(
|
||||
std::collections::HashMap::new(),
|
||||
)),
|
||||
};
|
||||
|
||||
// Background loop for stale task requeueing, offline node detection, and scheduler checking
|
||||
// Background maintenance & scheduling with Exponential Backoff
|
||||
let bg_db = state.db.clone();
|
||||
let bg_queue = queue.clone();
|
||||
let bg_scheduler = scheduler.clone();
|
||||
@@ -90,58 +120,226 @@ async fn main() -> Result<()> {
|
||||
let node_stale_sec = server_cfg.node_stale_sec;
|
||||
|
||||
tokio::spawn(async move {
|
||||
let mut fail_count: u32 = 0;
|
||||
let mut first_run = true;
|
||||
loop {
|
||||
sleep(Duration::from_secs(30)).await;
|
||||
if let Ok(requeued_points) = bg_queue.requeue_stale_tasks(stale_sec).await {
|
||||
if !requeued_points.is_empty() {
|
||||
info!("重新将 {} 个超时/掉线任务放回待计算队列", requeued_points.len());
|
||||
let _ = bg_db.reset_specific_grid_points_to_pending(&requeued_points).await;
|
||||
}
|
||||
if first_run {
|
||||
first_run = false;
|
||||
} else {
|
||||
let base_delay = 30u64;
|
||||
let current_delay = if fail_count == 0 {
|
||||
base_delay
|
||||
} else {
|
||||
(base_delay * (1u64 << fail_count.min(4))).min(300)
|
||||
};
|
||||
sleep(Duration::from_secs(current_delay)).await;
|
||||
}
|
||||
if let Ok(offline) = bg_db.mark_stale_nodes_offline(node_stale_sec).await {
|
||||
if offline > 0 {
|
||||
info!("已标记 {} 个心跳超时的计算节点为离线状态", offline);
|
||||
|
||||
let bg_db_clone = bg_db.clone();
|
||||
let bg_queue_clone = bg_queue.clone();
|
||||
let bg_scheduler_clone = bg_scheduler.clone();
|
||||
|
||||
let join_handle = tokio::spawn(async move {
|
||||
let mut has_error = false;
|
||||
match bg_queue_clone.requeue_stale_tasks(stale_sec).await {
|
||||
Ok(requeued) => {
|
||||
if !requeued.is_empty() {
|
||||
info!("重新将 {} 个超时/掉线任务放回待计算队列", requeued.len());
|
||||
let mut by_wf: std::collections::HashMap<String, Vec<String>> =
|
||||
std::collections::HashMap::new();
|
||||
for (point, wf) in &requeued {
|
||||
by_wf
|
||||
.entry(wf.clone().unwrap_or_default())
|
||||
.or_default()
|
||||
.push(point.clone());
|
||||
}
|
||||
for (wf, points) in by_wf {
|
||||
let _ = bg_db_clone
|
||||
.reset_specific_grid_points_to_pending(&points, &wf)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("重投超时任务失败: {}", e);
|
||||
has_error = true;
|
||||
}
|
||||
}
|
||||
|
||||
match bg_db_clone.mark_stale_nodes_offline(node_stale_sec).await {
|
||||
Ok(offline) => {
|
||||
if offline > 0 {
|
||||
info!("已标记 {} 个心跳超时的计算节点为离线状态", offline);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("标记超时节点离线失败: {}", e);
|
||||
has_error = true;
|
||||
}
|
||||
}
|
||||
|
||||
if let Err(e) = bg_scheduler_clone.schedule_pending_tasks().await {
|
||||
tracing::warn!("后台定时性任务调度检测失败: {}", e);
|
||||
has_error = true;
|
||||
}
|
||||
|
||||
if let Err(e) = bg_db_clone.sync_all_running_workflows_completion().await {
|
||||
tracing::warn!("后台同步已完成工作流状态失败: {}", e);
|
||||
has_error = true;
|
||||
}
|
||||
|
||||
has_error
|
||||
});
|
||||
|
||||
match join_handle.await {
|
||||
Ok(has_error) => {
|
||||
if has_error {
|
||||
fail_count = fail_count.saturating_add(1);
|
||||
} else {
|
||||
fail_count = 0;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!("后台维护任务内部发生 Panic: {:?}", e);
|
||||
fail_count = fail_count.saturating_add(1);
|
||||
}
|
||||
}
|
||||
if let Err(e) = bg_scheduler.schedule_pending_tasks().await {
|
||||
tracing::warn!("后台定时性任务调度检测失败: {}", e);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// 每天自动触发一次数据库备份。
|
||||
// 备份目录跟随 server_cfg.backup_dir（DCTS_BACKUP_DIR，默认 data/backups），
|
||||
// 与 DB_PATH 解耦，避免 DB 卷与备份卷不一致时备份落到未持久化层。
|
||||
// 首次延迟 1 小时，避免频繁重启（如调试阶段）短时间堆积备份文件；backup_database
|
||||
// 自身还带有 7 天保留期清理兜底。
|
||||
let backup_db = state.db.clone();
|
||||
let backup_dir = server_cfg.backup_dir.clone();
|
||||
tokio::spawn(async move {
|
||||
sleep(Duration::from_secs(3600)).await;
|
||||
loop {
|
||||
if let Err(e) = backup_db.backup_database(&backup_dir).await {
|
||||
tracing::warn!("自动备份数据库失败: {}", e);
|
||||
}
|
||||
sleep(Duration::from_secs(24 * 3600)).await;
|
||||
}
|
||||
});
|
||||
|
||||
// 大体积上传端点单独拎出，套用更宽松的 body limit（256MB，覆盖收敛种子 .7 文件量级）
|
||||
// 并限制并发数：每个 report 请求最多 256MB 驻留内存，无并发上限时 N 个请求可耗尽内存。
|
||||
// 限流后超出并发数的请求排队等待（而非直接拒绝），保证正常业务不被误伤。
|
||||
// 其余 API 用 10MB 默认上限，防止大文件内存耗尽 DoS。
|
||||
const REPORT_BODY_LIMIT: usize = 256 * 1024 * 1024;
|
||||
const DEFAULT_BODY_LIMIT: usize = 10 * 1024 * 1024;
|
||||
const REPORT_MAX_CONCURRENCY: usize = 4;
|
||||
|
||||
let report_router = Router::new()
|
||||
.route("/task/report", post(api::task::report_task))
|
||||
.layer(DefaultBodyLimit::max(REPORT_BODY_LIMIT))
|
||||
.layer(tower::ServiceBuilder::new().concurrency_limit(REPORT_MAX_CONCURRENCY));
|
||||
|
||||
// 节点注册接口独立 IP 限流保护（每分钟最多 10 次申请，无论成败都计数，防恶意频繁注册）
|
||||
// 使用 new_count_all：此 limiter 专挂 /node/register，对注册路径的所有响应计入窗口。
|
||||
// 通用 API 限流器（见下方 auth_enabled 分支）用 new 构造（count_all=false），不会因
|
||||
// 成功注册把 IP 锁出整个 /api/*，避免跨端点连锁限流。
|
||||
let register_limiter =
|
||||
api::rate_limit::RateLimiter::new_count_all(10, std::time::Duration::from_secs(60));
|
||||
let register_rate_limit_layer = axum::middleware::from_fn_with_state(
|
||||
register_limiter,
|
||||
api::rate_limit::rate_limit_middleware,
|
||||
);
|
||||
|
||||
let api_router = Router::new()
|
||||
// Auth API
|
||||
.route("/login", post(api::auth::login))
|
||||
.route("/auth/check", get(api::auth::check_auth))
|
||||
// Core Node & Task API
|
||||
.route("/node/register", post(api::node::register_node))
|
||||
.route(
|
||||
"/node/register",
|
||||
post(api::node::register_node).layer(register_rate_limit_layer),
|
||||
)
|
||||
.route("/node/check_status", post(api::node::check_node_status))
|
||||
.route("/node/heartbeat", post(api::node::heartbeat_node))
|
||||
.route("/task/claim", post(api::task::claim_task))
|
||||
.route("/task/report", post(api::task::report_task))
|
||||
.route("/seed/:name", get(api::seed::download_seed))
|
||||
.route("/status", get(api::status::get_status))
|
||||
// Static Data API
|
||||
.route("/data/file/*filename", get(api::data::download_single_data_file))
|
||||
.route(
|
||||
"/data/file/*filename",
|
||||
get(api::data::download_single_data_file),
|
||||
)
|
||||
.route("/data/linelist", get(api::data::download_linelist))
|
||||
// Workflow Management CRUD API
|
||||
.route("/workflows", get(api::workflow::list_workflows).post(api::workflow::save_workflow))
|
||||
.route("/workflows/:name", get(api::workflow::get_workflow).put(api::workflow::save_workflow).delete(api::workflow::delete_workflow))
|
||||
.route("/workflows/:name/start", post(api::workflow::start_workflow))
|
||||
.route("/workflows/:name/stop", post(api::workflow::stop_workflow));
|
||||
.route(
|
||||
"/workflows",
|
||||
get(api::workflow::list_workflows).post(api::workflow::save_workflow),
|
||||
)
|
||||
.route(
|
||||
"/workflows/:name",
|
||||
get(api::workflow::get_workflow)
|
||||
.put(api::workflow::save_workflow)
|
||||
.delete(api::workflow::delete_workflow),
|
||||
)
|
||||
.route(
|
||||
"/workflows/:name/start",
|
||||
post(api::workflow::start_workflow),
|
||||
)
|
||||
.route("/workflows/:name/stop", post(api::workflow::stop_workflow))
|
||||
// Admin Management API（节点凭据查看/审批/吊销/重发，均要求 Admin 角色）
|
||||
.route("/admin/nodes", get(api::admin::list_nodes))
|
||||
.route(
|
||||
"/admin/nodes/:node_id/approve",
|
||||
post(api::admin::approve_node),
|
||||
)
|
||||
.route(
|
||||
"/admin/nodes/:node_id/reject",
|
||||
post(api::admin::reject_node),
|
||||
)
|
||||
.route(
|
||||
"/admin/nodes/:node_id/revoke",
|
||||
post(api::admin::revoke_node),
|
||||
)
|
||||
.route(
|
||||
"/admin/nodes/:node_id/reissue",
|
||||
post(api::admin::reissue_node),
|
||||
)
|
||||
// 合并大体积上报路由（继承各自的 body limit）
|
||||
.merge(report_router)
|
||||
.layer(DefaultBodyLimit::max(DEFAULT_BODY_LIMIT));
|
||||
|
||||
let api_router = if state.auth_token.is_some() {
|
||||
info!("已为 DCTS 服务端 API 路由启用 Bearer Token / X-API-Key 访问控制鉴权");
|
||||
// 鉴权启用条件：未应急关闭，且配置了 admin 凭据。
|
||||
let auth_enabled = !state.auth_disabled && state.admin_token.is_some();
|
||||
|
||||
let api_router = if auth_enabled {
|
||||
info!("已启用 API 身份鉴权保护（Admin 端点需 admin token 验证；Node 节点免 Token 提交申请，经 Dashboard 管理员审批授权下发）");
|
||||
// 鉴权失败限流（防 token 在线暴力）：外层先判 IP 限流，内层再做鉴权。
|
||||
// 限流状态为 20 次/分钟（按 IP），超阈值返回 429。
|
||||
let limiter = api::rate_limit::RateLimiter::new(20, std::time::Duration::from_secs(60));
|
||||
let rate_limit_layer =
|
||||
axum::middleware::from_fn_with_state(limiter, api::rate_limit::rate_limit_middleware);
|
||||
let auth_layer = axum::middleware::from_fn_with_state(state.clone(), api::auth_middleware);
|
||||
api_router.layer(auth_layer)
|
||||
api_router.layer(auth_layer).layer(rate_limit_layer)
|
||||
} else {
|
||||
tracing::warn!("⚠️ 警告：未检测到 DCTS_AUTH_TOKEN 环境变量，服务端目前运行在【内网无鉴权模式】！所有 REST API 接口均为公开可访问状态。");
|
||||
tracing::warn!(
|
||||
"⚠️ 警告：未配置 DCTS_ADMIN_TOKEN / DCTS_ENROLLMENT_TOKEN（且未启用 DCTS_AUTH_DISABLE），\
|
||||
服务端运行在【无鉴权模式】！公网部署务必配置凭据。"
|
||||
);
|
||||
api_router
|
||||
};
|
||||
|
||||
// Host Dashboard SPA static files from dashboard/dist if directory exists or fallback to index.html
|
||||
let serve_dir = ServeDir::new("dashboard/dist")
|
||||
.fallback(ServeFile::new("dashboard/dist/index.html"));
|
||||
let serve_dir =
|
||||
ServeDir::new("dashboard/dist").fallback(ServeFile::new("dashboard/dist/index.html"));
|
||||
|
||||
// 安全响应头（CSP / nosniff / DENY / Referrer-Policy）。
|
||||
let security_headers = axum::middleware::from_fn(security_headers_middleware);
|
||||
|
||||
let app = Router::new()
|
||||
// 独立健康检查端点：不走鉴权、不走 CORS/body 限制，专供 docker healthcheck 与外部监控探测。
|
||||
// 开启鉴权后 /api/status 会返回 401，导致容器被判定不健康而反复重启，故单独提供 /healthz。
|
||||
.route("/healthz", get(api::status::healthz))
|
||||
.nest("/api", api_router)
|
||||
.layer(CorsLayer::permissive())
|
||||
.layer(server::cors::build_cors_layer())
|
||||
.layer(security_headers)
|
||||
.fallback_service(serve_dir)
|
||||
.with_state(state);
|
||||
|
||||
@@ -149,14 +347,53 @@ async fn main() -> Result<()> {
|
||||
info!("DCTS 服务端已在 http://{} 启动监听", addr);
|
||||
|
||||
let listener = tokio::net::TcpListener::bind(addr).await?;
|
||||
axum::serve(listener, app)
|
||||
.with_graceful_shutdown(async {
|
||||
let _ = tokio::signal::ctrl_c().await;
|
||||
info!("收到 Ctrl+C 终止信号，DCTS 服务端准备优雅关闭...");
|
||||
})
|
||||
.await?;
|
||||
// into_make_service_with_connect_info：让限流中间件能从连接拿到客户端 IP（反代场景则用 X-Forwarded-For）
|
||||
axum::serve(
|
||||
listener,
|
||||
app.into_make_service_with_connect_info::<SocketAddr>(),
|
||||
)
|
||||
.with_graceful_shutdown(async {
|
||||
let _ = tokio::signal::ctrl_c().await;
|
||||
info!("收到 Ctrl+C 终止信号，DCTS 服务端准备优雅关闭...");
|
||||
})
|
||||
.await?;
|
||||
|
||||
info!("DCTS 服务端已安全关闭。");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 注入安全响应头的中间件函数。
|
||||
async fn security_headers_middleware(
|
||||
req: axum::http::Request<axum::body::Body>,
|
||||
next: axum::middleware::Next,
|
||||
) -> axum::response::Response {
|
||||
let mut resp = next.run(req).await;
|
||||
|
||||
let headers = resp.headers_mut();
|
||||
// CSP：default-src 'self'；放行 Google Fonts（index.html 引用）；允许 data: 图片。
|
||||
// 已移除 'unsafe-eval'：dashboard 构建产物不使用 eval/new Function（已核实），保留它会
|
||||
// 显著削弱 CSP 的脚本注入防护。'unsafe-inline' 暂留（静态 SPA 内联脚本/handler 需要），
|
||||
// 彻底方案需前端改造为外链 + per-request nonce 注入，见 docs TODO。
|
||||
headers
|
||||
.entry(axum::http::header::CONTENT_SECURITY_POLICY)
|
||||
.or_insert_with(|| {
|
||||
HeaderValue::from_static(
|
||||
"default-src 'self'; script-src 'self' 'unsafe-inline'; \
|
||||
style-src 'self' 'unsafe-inline' https://fonts.googleapis.com; \
|
||||
font-src 'self' data: https://fonts.gstatic.com; \
|
||||
connect-src 'self'; img-src 'self' data: blob:; \
|
||||
frame-ancestors 'none'",
|
||||
)
|
||||
});
|
||||
headers
|
||||
.entry(axum::http::header::X_CONTENT_TYPE_OPTIONS)
|
||||
.or_insert_with(|| HeaderValue::from_static("nosniff"));
|
||||
headers
|
||||
.entry(axum::http::header::X_FRAME_OPTIONS)
|
||||
.or_insert_with(|| HeaderValue::from_static("DENY"));
|
||||
headers
|
||||
.entry(axum::http::HeaderName::from_static("referrer-policy"))
|
||||
.or_insert_with(|| HeaderValue::from_static("strict-origin-when-cross-origin"));
|
||||
|
||||
resp
|
||||
}
|
||||
|
||||
+333
-52
@@ -23,13 +23,32 @@ impl GridScheduler {
|
||||
}
|
||||
}
|
||||
|
||||
/// Expands grid points from config and registers them into the database
|
||||
pub async fn initialize_grid(&self, cfg: &GridConfig) -> Result<()> {
|
||||
if let Err(e) = self.queue.clear_queue().await {
|
||||
tracing::warn!("初始化网格时清理闲置排队记录发生警告: {}", e);
|
||||
/// Expands grid points from config and registers them into the database.
|
||||
///
|
||||
/// 多工作流分区（#3 修复）：
|
||||
/// - 仅清理**本工作流**的排队任务（clear_queue_by_workflow），不再 clear_queue() 全局清空，
|
||||
/// 避免启动工作流 B 时误删工作流 A 的在队任务。
|
||||
/// - 仅重置**本工作流**的 queued 点为 pending（reset_queued_grid_points_to_pending 带 wf），
|
||||
/// 避免误伤其他工作流。
|
||||
/// - upsert 带 workflow_name，使同一物理点可属于多个工作流。
|
||||
pub async fn initialize_grid(&self, cfg: &GridConfig, workflow_name: &str) -> Result<()> {
|
||||
if let Err(e) = self.queue.clear_queue_by_workflow(workflow_name).await {
|
||||
tracing::warn!(
|
||||
"初始化工作流 {} 网格时清理该流闲置排队记录发生警告: {}",
|
||||
workflow_name,
|
||||
e
|
||||
);
|
||||
}
|
||||
if let Err(e) = self.db.reset_queued_grid_points_to_pending().await {
|
||||
tracing::warn!("重置网格状态到 pending 处理过程遇到异常: {}", e);
|
||||
if let Err(e) = self
|
||||
.db
|
||||
.reset_queued_grid_points_to_pending(workflow_name)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
"重置工作流 {} 网格状态到 pending 处理过程遇到异常: {}",
|
||||
workflow_name,
|
||||
e
|
||||
);
|
||||
}
|
||||
let mut points = Vec::new();
|
||||
|
||||
@@ -59,9 +78,21 @@ impl GridScheduler {
|
||||
a.cno_sum()
|
||||
.partial_cmp(&b.cno_sum())
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
.then_with(|| a.teff.partial_cmp(&b.teff).unwrap_or(std::cmp::Ordering::Equal))
|
||||
.then_with(|| b.logg.partial_cmp(&a.logg).unwrap_or(std::cmp::Ordering::Equal))
|
||||
.then_with(|| a.loghe.partial_cmp(&b.loghe).unwrap_or(std::cmp::Ordering::Equal))
|
||||
.then_with(|| {
|
||||
a.teff
|
||||
.partial_cmp(&b.teff)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
})
|
||||
.then_with(|| {
|
||||
b.logg
|
||||
.partial_cmp(&a.logg)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
})
|
||||
.then_with(|| {
|
||||
a.loghe
|
||||
.partial_cmp(&b.loghe)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
})
|
||||
});
|
||||
|
||||
// Group into Waves by cno_sum
|
||||
@@ -79,46 +110,100 @@ impl GridScheduler {
|
||||
current_cno = Some(cno);
|
||||
}
|
||||
|
||||
self.db.upsert_grid_point(pt, wave_idx).await?;
|
||||
// upsert 是幂等的 ON CONFLICT DO NOTHING：若 initialize_grid 中途失败，
|
||||
// 重新 start 该工作流会自然补齐（#4 半初始化回退由幂等性消解）。
|
||||
self.db
|
||||
.upsert_grid_point(pt, wave_idx, workflow_name)
|
||||
.await?;
|
||||
}
|
||||
|
||||
info!("已在数据库中成功初始化并记录 {} 个恒星大气网格点", points.len());
|
||||
info!(
|
||||
"已在数据库中成功初始化并记录工作流 {} 的 {} 个恒星大气网格点",
|
||||
workflow_name,
|
||||
points.len()
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn get_active_timeout_sec(&self) -> u64 {
|
||||
if let Ok(yamls) = self.db.get_running_workflow_config_yamls().await {
|
||||
for yaml in yamls {
|
||||
if let Ok(cfg) = serde_yaml::from_str::<GridConfig>(&yaml) {
|
||||
return cfg.timeout_sec;
|
||||
}
|
||||
/// 读取指定工作流的 timeout_sec（按工作流分区：多工作流各有自己的超时配置）。
|
||||
async fn get_workflow_timeout_sec(&self, workflow_name: &str) -> u64 {
|
||||
if let Ok(Some(wf)) = self.db.get_workflow(workflow_name).await {
|
||||
if let Ok(cfg) = serde_yaml::from_str::<GridConfig>(&wf.config_yaml) {
|
||||
return cfg.timeout_sec;
|
||||
}
|
||||
}
|
||||
7200
|
||||
}
|
||||
|
||||
/// Enqueues pending grid points into MQ with active seed detection and batching
|
||||
/// 读取指定工作流的 seed_step_fallback 配置。
|
||||
async fn get_workflow_seed_step_fallback(&self, workflow_name: &str) -> bool {
|
||||
if let Ok(Some(wf)) = self.db.get_workflow(workflow_name).await {
|
||||
if let Ok(cfg) = serde_yaml::from_str::<GridConfig>(&wf.config_yaml) {
|
||||
return cfg.seed_step_fallback;
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
/// Enqueues pending grid points into MQ with active seed detection and batching.
|
||||
///
|
||||
/// 多工作流分区（#3 修复）：对**每个** running/initializing 工作流分别派发任务，
|
||||
/// 替代原来「全局只一个 running workflow」的 LIMIT 1 假设。各工作流独立 batch、
|
||||
/// 独立 seed 匹配（seeds 仍是全局共享的物理资源池）。
|
||||
pub async fn schedule_pending_tasks(&self) -> Result<usize> {
|
||||
if !self.db.has_running_workflow().await? {
|
||||
let workflows = self.db.get_running_workflow_names().await?;
|
||||
if workflows.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
let timeout_sec = self.get_active_timeout_sec().await;
|
||||
let batch_limit: usize = std::env::var("DCTS_BATCH_LIMIT")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(100);
|
||||
|
||||
// SQL 层直接附加 LIMIT = batch_limit 筛选，完全免除数万点位无谓内存反序列化和空耗对象释放开销
|
||||
let pending = self.db.get_pending_grid_points_limit(batch_limit).await?;
|
||||
let mut total_dispatched = 0;
|
||||
|
||||
for wf in &workflows {
|
||||
let dispatched = self
|
||||
.schedule_pending_tasks_for_workflow(wf, batch_limit)
|
||||
.await?;
|
||||
total_dispatched += dispatched;
|
||||
}
|
||||
|
||||
if total_dispatched > 0 {
|
||||
info!(
|
||||
"已成功将 {} 个待计算网格点推进任务队列（跨 {} 个工作流）",
|
||||
total_dispatched,
|
||||
workflows.len()
|
||||
);
|
||||
}
|
||||
|
||||
Ok(total_dispatched)
|
||||
}
|
||||
|
||||
/// 为单个工作流派发 pending 点。
|
||||
async fn schedule_pending_tasks_for_workflow(
|
||||
&self,
|
||||
workflow_name: &str,
|
||||
batch_limit: usize,
|
||||
) -> Result<usize> {
|
||||
let timeout_sec = self.get_workflow_timeout_sec(workflow_name).await;
|
||||
|
||||
// SQL 层直接附加 LIMIT = batch_limit + workflow_name 筛选，完全免除数万点位无谓内存反序列化
|
||||
let pending = self
|
||||
.db
|
||||
.get_pending_grid_points_limit(batch_limit, workflow_name)
|
||||
.await?;
|
||||
let mut dispatched = 0;
|
||||
|
||||
for (name, params, _wave) in pending {
|
||||
|
||||
// Check if any seed is available in DB for active SeedStep scheduling
|
||||
// Check if any seed is available in DB for active SeedStep scheduling（seeds 全局共享）
|
||||
let (task_type, seed_name) = match self.db.find_best_seed_from_db(¶ms).await {
|
||||
Ok(Some(seed_match)) => {
|
||||
info!("网格点 {} 匹配到数据库近邻种子 {} (距离: {:.2})，安排 SeedStep 热启动调度", name, seed_match.name, seed_match.distance);
|
||||
info!(
|
||||
"工作流 {} 网格点 {} 匹配到数据库近邻种子 {} (距离: {:.2})，安排 SeedStep 热启动调度",
|
||||
workflow_name, name, seed_match.name, seed_match.distance
|
||||
);
|
||||
(TaskType::SeedStep, Some(seed_match.name))
|
||||
}
|
||||
_ => (TaskType::ColdRun, None),
|
||||
@@ -131,48 +216,96 @@ impl GridScheduler {
|
||||
task_type,
|
||||
seed_point_name: seed_name,
|
||||
timeout_sec,
|
||||
workflow_name: Some(workflow_name.to_string()),
|
||||
};
|
||||
|
||||
self.db.insert_task(&task_spec).await?;
|
||||
// 采用先标记 DB 状态为 Queued 后发 MQ 的时序，防止推入 MQ 后数据库修改异常导向下一轮误重投
|
||||
self.db.update_grid_status(&name, common::models::GridPointStatus::Queued).await?;
|
||||
self.db
|
||||
.update_grid_status(
|
||||
&name,
|
||||
common::models::GridPointStatus::Queued,
|
||||
workflow_name,
|
||||
)
|
||||
.await?;
|
||||
match self.queue.push_task(&task_spec).await {
|
||||
Ok(_) => {
|
||||
dispatched += 1;
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("将任务 {} 推入 MQ 队列失败，回滚网格点状态: {}", name, e);
|
||||
let _ = self.db.update_grid_status(&name, common::models::GridPointStatus::Pending).await;
|
||||
tracing::warn!(
|
||||
"将任务 {} 推入 MQ 队列失败，执行严格状态回滚以避免脏数据: {}",
|
||||
name,
|
||||
e
|
||||
);
|
||||
if let Err(db_e) = self
|
||||
.db
|
||||
.update_grid_status(
|
||||
&name,
|
||||
common::models::GridPointStatus::Pending,
|
||||
workflow_name,
|
||||
)
|
||||
.await
|
||||
{
|
||||
tracing::error!(
|
||||
"关键性回滚异常：任务 {} 无法重置回 Pending: {}",
|
||||
name,
|
||||
db_e
|
||||
);
|
||||
}
|
||||
let _ = self.queue.remove_task(&task_spec.task_id.to_string()).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if dispatched > 0 {
|
||||
info!("已成功将 {} 个待计算网格点推进任务队列", dispatched);
|
||||
}
|
||||
|
||||
Ok(dispatched)
|
||||
}
|
||||
|
||||
/// Triggers seed_step fallback for a failed point if a seed is available
|
||||
pub async fn trigger_seed_step_fallback(&self, params: &GridPointParams) -> Result<bool> {
|
||||
if !self.db.has_running_workflow().await? {
|
||||
/// Triggers seed_step fallback for a failed point if a seed is available.
|
||||
///
|
||||
/// 多工作流分区（#3 修复）：传入 `workflow_name` 明确该失败点所属工作流，
|
||||
/// 用该工作流自身的 timeout / seed_step_fallback 配置，并把 TaskSpec.workflow_name
|
||||
/// 绑定到该工作流。
|
||||
///
|
||||
/// 语义（种子回退仅一次）：
|
||||
/// - 仅当该工作流配置 `seed_step_fallback: true` 时才考虑回退；
|
||||
/// - 仅当该点**尚未**派发过任何 seed_step 任务时才回退一次；
|
||||
/// - 找不到合适近邻种子则不回退，由调用方保持 failed 终态。
|
||||
pub async fn trigger_seed_step_fallback(
|
||||
&self,
|
||||
params: &GridPointParams,
|
||||
workflow_name: &str,
|
||||
) -> Result<bool> {
|
||||
// 该工作流须仍处于 running 态才回退（避免 stop 后继续派发）
|
||||
let still_running = self
|
||||
.db
|
||||
.get_running_workflow_names()
|
||||
.await?
|
||||
.iter()
|
||||
.any(|w| w == workflow_name);
|
||||
if !still_running {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
if !self.get_workflow_seed_step_fallback(workflow_name).await {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
let name = params.model_name();
|
||||
if let Ok(Some((status, attempt_count))) = self.db.get_grid_point_status(&name).await {
|
||||
if status == "failed" || attempt_count >= 3 {
|
||||
info!("网格点 {} 已达到最大重试次数 ({}) 或处于 failed 状态，跳过种子热启动回退", name, attempt_count);
|
||||
return Ok(false);
|
||||
}
|
||||
// 种子回退仅一次：该点在该工作流中已经派发过 seed_step 任务就不再触发新的回退
|
||||
if self.db.has_seed_step_attempt(&name, workflow_name).await? {
|
||||
info!(
|
||||
"网格点 {} 已使用过一次种子热启动回退，不再重复回退，保持 failed 终态",
|
||||
name
|
||||
);
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
// seeds 全局共享：跨工作流复用已收敛的邻近种子
|
||||
let seed_match_opt = self.db.find_best_seed_from_db(params).await.ok().flatten();
|
||||
|
||||
if let Some(seed_match) = seed_match_opt {
|
||||
let timeout_sec = self.get_active_timeout_sec().await;
|
||||
let timeout_sec = self.get_workflow_timeout_sec(workflow_name).await;
|
||||
let name = params.model_name();
|
||||
let task_spec = TaskSpec {
|
||||
task_id: Uuid::new_v4(),
|
||||
@@ -181,22 +314,38 @@ impl GridScheduler {
|
||||
task_type: TaskType::SeedStep,
|
||||
seed_point_name: Some(seed_match.name.clone()),
|
||||
timeout_sec,
|
||||
workflow_name: Some(workflow_name.to_string()),
|
||||
};
|
||||
|
||||
self.db.insert_task(&task_spec).await?;
|
||||
self.db.update_grid_status(&name, common::models::GridPointStatus::Queued).await?;
|
||||
self.db
|
||||
.update_grid_status(
|
||||
&name,
|
||||
common::models::GridPointStatus::Queued,
|
||||
workflow_name,
|
||||
)
|
||||
.await?;
|
||||
if let Err(e) = self.queue.push_task(&task_spec).await {
|
||||
let _ = self.db.update_grid_status(&name, common::models::GridPointStatus::Pending).await;
|
||||
let _ = self
|
||||
.db
|
||||
.update_grid_status(
|
||||
&name,
|
||||
common::models::GridPointStatus::Pending,
|
||||
workflow_name,
|
||||
)
|
||||
.await;
|
||||
let _ = self.queue.remove_task(&task_spec.task_id.to_string()).await;
|
||||
return Err(e);
|
||||
}
|
||||
info!("触发种子步进 (seed_step)：网格点 {} 将使用 6 维近邻种子 {} 热启动重试", name, seed_match.name);
|
||||
info!(
|
||||
"触发种子步进 (seed_step)：工作流 {} 网格点 {} 将使用 6 维近邻种子 {} 热启动重试",
|
||||
workflow_name, name, seed_match.name
|
||||
);
|
||||
Ok(true)
|
||||
} else {
|
||||
Ok(false)
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -212,8 +361,16 @@ mod tests {
|
||||
let results_dir = temp_dir.path().join("results");
|
||||
|
||||
let db = Database::new(&db_path.to_string_lossy()).await.unwrap();
|
||||
let queue = Arc::new(SqliteTaskQueue::new(&queue_db_path.to_string_lossy()).await.unwrap());
|
||||
let scheduler = GridScheduler::new(db.clone(), queue.clone(), results_dir.to_string_lossy().to_string());
|
||||
let queue = Arc::new(
|
||||
SqliteTaskQueue::new(&queue_db_path.to_string_lossy())
|
||||
.await
|
||||
.unwrap(),
|
||||
);
|
||||
let scheduler = GridScheduler::new(
|
||||
db.clone(),
|
||||
queue.clone(),
|
||||
results_dir.to_string_lossy().to_string(),
|
||||
);
|
||||
|
||||
let cfg = GridConfig {
|
||||
grid: GridAxesConfig {
|
||||
@@ -238,17 +395,141 @@ mod tests {
|
||||
linelist: None,
|
||||
};
|
||||
|
||||
scheduler.initialize_grid(&cfg).await.unwrap();
|
||||
db.upsert_workflow("test_wf", None, "", "running").await.unwrap();
|
||||
scheduler.initialize_grid(&cfg, "test_wf").await.unwrap();
|
||||
db.upsert_workflow("test_wf", None, "", "running")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let pending = db.get_pending_grid_points().await.unwrap();
|
||||
let pending = db.get_pending_grid_points("test_wf").await.unwrap();
|
||||
assert_eq!(pending.len(), 1);
|
||||
|
||||
let dispatched = scheduler.schedule_pending_tasks().await.unwrap();
|
||||
assert_eq!(dispatched, 1);
|
||||
|
||||
let popped = queue.pop_task().await.unwrap();
|
||||
let popped = queue.pop_task("test-node").await.unwrap();
|
||||
assert!(popped.is_some());
|
||||
}
|
||||
}
|
||||
|
||||
/// 多工作流分区调度测试（#3 修复验证）：
|
||||
/// 1. wf_a 调度推入队列的任务，在初始化 wf_b 后依然存在（initialize_grid 改用
|
||||
/// clear_queue_by_workflow，不再全局 clear_queue）。
|
||||
/// 2. 两个 running 工作流的 pending 点都能被 schedule_pending_tasks 派发。
|
||||
#[tokio::test]
|
||||
async fn test_multi_workflow_dispatch_isolation() {
|
||||
let temp_dir = tempfile::tempdir().unwrap();
|
||||
let db = Database::new(&temp_dir.path().join("mw_db.db").to_string_lossy())
|
||||
.await
|
||||
.unwrap();
|
||||
let queue = Arc::new(
|
||||
SqliteTaskQueue::new(&temp_dir.path().join("mw_queue.db").to_string_lossy())
|
||||
.await
|
||||
.unwrap(),
|
||||
);
|
||||
let scheduler = GridScheduler::new(db.clone(), queue.clone(), "results".to_string());
|
||||
|
||||
let mk_cfg = |teff: f64| GridConfig {
|
||||
grid: GridAxesConfig {
|
||||
teff: vec![teff],
|
||||
logg: vec![5.5],
|
||||
loghe: vec![-1.0],
|
||||
logc: vec![-2.0],
|
||||
logn: vec![-2.0],
|
||||
logo: vec![-2.0],
|
||||
},
|
||||
chain: vec![],
|
||||
synspec: None,
|
||||
nworkers: 4,
|
||||
timeout_sec: 3600,
|
||||
resume: true,
|
||||
seed_step_fallback: true,
|
||||
results: None,
|
||||
itek_fallback: vec![],
|
||||
niter: Some(100),
|
||||
template: None,
|
||||
fort55: None,
|
||||
linelist: None,
|
||||
};
|
||||
|
||||
// wf_a 初始化并推入队列
|
||||
scheduler
|
||||
.initialize_grid(&mk_cfg(35000.0), "wf_a")
|
||||
.await
|
||||
.unwrap();
|
||||
db.upsert_workflow("wf_a", None, "", "running")
|
||||
.await
|
||||
.unwrap();
|
||||
let d_a = scheduler.schedule_pending_tasks().await.unwrap();
|
||||
assert_eq!(d_a, 1);
|
||||
// 任务已在队
|
||||
assert!(queue.pop_task("node-a").await.unwrap().is_some());
|
||||
|
||||
// 重新推一个 wf_a 任务（上一行 pop 掉了），再初始化 wf_b
|
||||
db.update_grid_status(
|
||||
&GridPointParams {
|
||||
teff: 35000.0,
|
||||
logg: 5.5,
|
||||
loghe: -1.0,
|
||||
logc: -2.0,
|
||||
logn: -2.0,
|
||||
logo: -2.0,
|
||||
}
|
||||
.model_name(),
|
||||
common::models::GridPointStatus::Pending,
|
||||
"wf_a",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let _ = scheduler
|
||||
.schedule_pending_tasks_for_workflow("wf_a", 100)
|
||||
.await
|
||||
.unwrap();
|
||||
// 此时 wf_a 队列里应有一个任务
|
||||
assert_eq!(
|
||||
queue
|
||||
.pop_task("node-a")
|
||||
.await
|
||||
.unwrap()
|
||||
.and_then(|t| t.workflow_name),
|
||||
Some("wf_a".to_string())
|
||||
);
|
||||
|
||||
// 关键断言：把 wf_a 任务重新推回队列后，初始化 wf_b 不应清空它。
|
||||
db.update_grid_status(
|
||||
&GridPointParams {
|
||||
teff: 35000.0,
|
||||
logg: 5.5,
|
||||
loghe: -1.0,
|
||||
logc: -2.0,
|
||||
logn: -2.0,
|
||||
logo: -2.0,
|
||||
}
|
||||
.model_name(),
|
||||
common::models::GridPointStatus::Pending,
|
||||
"wf_a",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let _ = scheduler
|
||||
.schedule_pending_tasks_for_workflow("wf_a", 100)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// 初始化 wf_b（内部 clear_queue_by_workflow("wf_b")，不该动 wf_a 的任务）
|
||||
scheduler
|
||||
.initialize_grid(&mk_cfg(40000.0), "wf_b")
|
||||
.await
|
||||
.unwrap();
|
||||
db.upsert_workflow("wf_b", None, "", "running")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// wf_a 的任务仍在队：可被 node 弹出，且 workflow_name == wf_a
|
||||
let popped_a = queue.pop_task("node-a").await.unwrap();
|
||||
assert!(popped_a.is_some(), "初始化 wf_b 不应清空 wf_a 的队列任务");
|
||||
assert_eq!(popped_a.unwrap().workflow_name, Some("wf_a".to_string()));
|
||||
|
||||
// wf_b 的点也能被调度（两个 running 工作流并存）
|
||||
let d_b = scheduler.schedule_pending_tasks().await.unwrap();
|
||||
assert!(d_b >= 1, "wf_b 的 pending 点应被派发");
|
||||
}
|
||||
}
|
||||
|
||||
+1175
-17
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user