Asfmq
b91f1e4fa5
- 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 配置及数据库设计文档
1315 lines
43 KiB
Rust
1315 lines
43 KiB
Rust
use axum::{
|
||
body::Body,
|
||
http::{Request, StatusCode},
|
||
};
|
||
use mq::sqlite_queue::SqliteTaskQueue;
|
||
use server::{api::AppState, db::Database, scheduler::GridScheduler};
|
||
use std::sync::Arc;
|
||
use tower::ServiceExt; // for oneshot
|
||
|
||
#[tokio::test]
|
||
async fn test_server_api_flow() {
|
||
let temp_dir = tempfile::tempdir().unwrap();
|
||
let db_path = temp_dir.path().join("api_db.db");
|
||
let queue_db_path = temp_dir.path().join("api_queue.db");
|
||
let results_dir = temp_dir.path().join("results");
|
||
std::fs::create_dir_all(&results_dir).unwrap();
|
||
|
||
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 = Arc::new(GridScheduler::new(
|
||
db.clone(),
|
||
queue.clone(),
|
||
results_dir.to_string_lossy().to_string(),
|
||
));
|
||
|
||
let state = AppState {
|
||
db,
|
||
queue,
|
||
scheduler,
|
||
results_dir: results_dir.to_string_lossy().to_string(),
|
||
rate_limiter: server::api::rate_limit::RateLimiter::new(
|
||
5,
|
||
std::time::Duration::from_secs(300),
|
||
),
|
||
auth_token: None,
|
||
admin_token: None,
|
||
auth_disabled: false,
|
||
admin_sessions: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
|
||
};
|
||
|
||
let app = axum::Router::new()
|
||
.route(
|
||
"/api/node/register",
|
||
axum::routing::post(server::api::node::register_node),
|
||
)
|
||
.route(
|
||
"/api/node/heartbeat",
|
||
axum::routing::post(server::api::node::heartbeat_node),
|
||
)
|
||
.route(
|
||
"/api/task/claim",
|
||
axum::routing::post(server::api::task::claim_task),
|
||
)
|
||
.route(
|
||
"/api/status",
|
||
axum::routing::get(server::api::status::get_status),
|
||
)
|
||
.route(
|
||
"/api/workflows",
|
||
axum::routing::get(server::api::workflow::list_workflows)
|
||
.post(server::api::workflow::save_workflow),
|
||
)
|
||
.with_state(state);
|
||
|
||
// 1. Check status API
|
||
let response = app
|
||
.clone()
|
||
.oneshot(
|
||
Request::builder()
|
||
.uri("/api/status")
|
||
.body(Body::empty())
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
|
||
assert_eq!(response.status(), StatusCode::OK);
|
||
|
||
// 2. Register node API
|
||
let reg_body = serde_json::json!({
|
||
"node_id": "test-node-api",
|
||
"host_name": "api-host",
|
||
"max_slots": 8
|
||
});
|
||
let response = app
|
||
.clone()
|
||
.oneshot(
|
||
Request::builder()
|
||
.method("POST")
|
||
.uri("/api/node/register")
|
||
.header("content-type", "application/json")
|
||
.body(Body::from(serde_json::to_vec(®_body).unwrap()))
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
|
||
assert_eq!(response.status(), StatusCode::OK);
|
||
|
||
// 3. Save Workflow API
|
||
let wf_body = serde_json::json!({
|
||
"name": "test_api_wf",
|
||
"description": "Test Workflow Description",
|
||
"config_yaml": "grid:\n teff: [35000]\n logg: [5.5]\n loghe: [-1]\n logc: [-2]\n logn: [-2]\n logo: [-2]"
|
||
});
|
||
let response = app
|
||
.clone()
|
||
.oneshot(
|
||
Request::builder()
|
||
.method("POST")
|
||
.uri("/api/workflows")
|
||
.header("content-type", "application/json")
|
||
.body(Body::from(serde_json::to_vec(&wf_body).unwrap()))
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
|
||
assert_eq!(response.status(), StatusCode::OK);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_auth_middleware_scope_and_running_status() {
|
||
let temp_dir = tempfile::tempdir().unwrap();
|
||
let db_path = temp_dir.path().join("auth_db.db");
|
||
let queue_db_path = temp_dir.path().join("auth_queue.db");
|
||
let results_dir = temp_dir.path().join("results");
|
||
std::fs::create_dir_all(&results_dir).unwrap();
|
||
|
||
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 = Arc::new(GridScheduler::new(
|
||
db.clone(),
|
||
queue.clone(),
|
||
results_dir.to_string_lossy().to_string(),
|
||
));
|
||
|
||
let state = AppState {
|
||
db: db.clone(),
|
||
queue: queue.clone(),
|
||
scheduler,
|
||
results_dir: results_dir.to_string_lossy().to_string(),
|
||
rate_limiter: server::api::rate_limit::RateLimiter::new(
|
||
5,
|
||
std::time::Duration::from_secs(300),
|
||
),
|
||
auth_token: Some("secret_token_123".to_string()),
|
||
admin_token: Some("secret_token_123".to_string()),
|
||
auth_disabled: false,
|
||
admin_sessions: std::sync::Arc::new(tokio::sync::RwLock::new(
|
||
std::collections::HashMap::new(),
|
||
)),
|
||
};
|
||
|
||
let api_router = axum::Router::new().route(
|
||
"/status",
|
||
axum::routing::get(server::api::status::get_status),
|
||
);
|
||
|
||
let auth_layer =
|
||
axum::middleware::from_fn_with_state(state.clone(), server::api::auth_middleware);
|
||
let api_router = api_router.layer(auth_layer);
|
||
|
||
let app = axum::Router::new()
|
||
.nest("/api", api_router)
|
||
.with_state(state);
|
||
|
||
// Unauthenticated API request -> 401 Unauthorized
|
||
let res = app
|
||
.clone()
|
||
.oneshot(
|
||
Request::builder()
|
||
.uri("/api/status")
|
||
.body(Body::empty())
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(res.status(), StatusCode::UNAUTHORIZED);
|
||
|
||
// Authenticated API request -> 200 OK
|
||
let res = app
|
||
.clone()
|
||
.oneshot(
|
||
Request::builder()
|
||
.uri("/api/status")
|
||
.header("authorization", "Bearer secret_token_123")
|
||
.body(Body::empty())
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(res.status(), StatusCode::OK);
|
||
|
||
// Test mark_grid_point_running
|
||
let params = common::models::GridPointParams {
|
||
teff: 35000.0,
|
||
logg: 5.5,
|
||
loghe: -1.0,
|
||
logc: -2.0,
|
||
logn: -2.0,
|
||
logo: -2.0,
|
||
};
|
||
db.upsert_grid_point(¶ms, 0, "test_wf").await.unwrap();
|
||
db.mark_grid_point_running(¶ms.model_name(), "test_wf")
|
||
.await
|
||
.unwrap();
|
||
|
||
let stats = db.get_grid_summary_stats(None).await.unwrap();
|
||
assert_eq!(stats["running"], 1);
|
||
}
|
||
|
||
/// L2 鉴权核心流程测试：
|
||
/// 注册节点 → 颁发专属 token → 用 token 调 heartbeat（200）→ 吊销 → 再调（401）。
|
||
#[tokio::test]
|
||
async fn test_l2_node_token_issue_revoke_flow() {
|
||
let temp_dir = tempfile::tempdir().unwrap();
|
||
let db_path = temp_dir.path().join("l2_db.db");
|
||
|
||
let db = Database::new(&db_path.to_string_lossy()).await.unwrap();
|
||
|
||
// 1. 注册节点
|
||
let reg = common::models::NodeRegisterRequest {
|
||
node_id: "node-l2-test".to_string(),
|
||
host_name: "l2-host".to_string(),
|
||
max_slots: 4,
|
||
};
|
||
db.register_node(®).await.unwrap();
|
||
|
||
// 2. 颁发专属 token，返回明文
|
||
let token = db.issue_node_token("node-l2-test").await.unwrap();
|
||
assert!(!token.is_empty());
|
||
|
||
// 2b. 一次性取走暂存明文（take_pending_node_token）：首次取到与颁发一致的明文，
|
||
// 再次取为 None（取走即焚）。验证 #4 简化为单一 UPDATE...RETURNING 后行为一致。
|
||
let pending = db.take_pending_node_token("node-l2-test").await.unwrap();
|
||
assert_eq!(pending.as_deref(), Some(token.as_str()));
|
||
let pending2 = db.take_pending_node_token("node-l2-test").await.unwrap();
|
||
assert!(
|
||
pending2.is_none(),
|
||
"取走即焚：第二次 take_pending 必须返回 None"
|
||
);
|
||
// 不存在的 node take 也应返回 None（不报错）
|
||
assert!(db
|
||
.take_pending_node_token("node-not-exist")
|
||
.await
|
||
.unwrap()
|
||
.is_none());
|
||
|
||
// 3. token 可反查到 node_id（此调用会把 token 写入内存缓存）
|
||
let found = db.find_node_by_token(&token).await;
|
||
assert_eq!(found.as_deref(), Some("node-l2-test"));
|
||
|
||
// 4. 错误 token 查不到
|
||
assert!(db.find_node_by_token("wrong-token").await.is_none());
|
||
|
||
// 5. 吊销后 token 立即失效
|
||
db.revoke_node_token("node-l2-test").await.unwrap();
|
||
assert!(db.find_node_by_token(&token).await.is_none());
|
||
|
||
// 6. 重新颁发后恢复（覆盖式）
|
||
let token2 = db.issue_node_token("node-l2-test").await.unwrap();
|
||
assert_eq!(
|
||
db.find_node_by_token(&token2).await.as_deref(),
|
||
Some("node-l2-test")
|
||
);
|
||
// 旧 token 仍失效（已被覆盖）
|
||
assert!(db.find_node_by_token(&token).await.is_none());
|
||
}
|
||
|
||
/// 验证中间件对 node token 的端到端鉴权：
|
||
/// 用真实 node token 调 /node/heartbeat 应 200；吊销后再调应 401。
|
||
#[tokio::test]
|
||
async fn test_middleware_node_role_with_token() {
|
||
let temp_dir = tempfile::tempdir().unwrap();
|
||
let db_path = temp_dir.path().join("l2_mw_db.db");
|
||
let queue_db_path = temp_dir.path().join("l2_mw_queue.db");
|
||
let results_dir = temp_dir.path().join("results");
|
||
std::fs::create_dir_all(&results_dir).unwrap();
|
||
|
||
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 = Arc::new(GridScheduler::new(
|
||
db.clone(),
|
||
queue.clone(),
|
||
results_dir.to_string_lossy().to_string(),
|
||
));
|
||
|
||
// 注册并颁发 token
|
||
let reg = common::models::NodeRegisterRequest {
|
||
node_id: "node-mw-test".to_string(),
|
||
host_name: "mw-host".to_string(),
|
||
max_slots: 2,
|
||
};
|
||
db.register_node(®).await.unwrap();
|
||
let token = db.issue_node_token("node-mw-test").await.unwrap();
|
||
|
||
let state = AppState {
|
||
db: db.clone(),
|
||
queue,
|
||
scheduler,
|
||
results_dir: results_dir.to_string_lossy().to_string(),
|
||
rate_limiter: server::api::rate_limit::RateLimiter::new(
|
||
5,
|
||
std::time::Duration::from_secs(300),
|
||
),
|
||
auth_token: Some("placeholder".to_string()),
|
||
admin_token: None,
|
||
auth_disabled: false,
|
||
admin_sessions: std::sync::Arc::new(tokio::sync::RwLock::new(
|
||
std::collections::HashMap::new(),
|
||
)),
|
||
};
|
||
|
||
let api_router = axum::Router::new().route(
|
||
"/node/heartbeat",
|
||
axum::routing::post(server::api::node::heartbeat_node),
|
||
);
|
||
let auth_layer =
|
||
axum::middleware::from_fn_with_state(state.clone(), server::api::auth_middleware);
|
||
let api_router = api_router.layer(auth_layer);
|
||
let app = axum::Router::new()
|
||
.nest("/api", api_router)
|
||
.with_state(state);
|
||
|
||
// 用有效 node token 调 heartbeat → 200
|
||
let hb_body = serde_json::json!({
|
||
"node_id": "node-mw-test", "active_slots": 1, "cpu_usage": 10.0, "memory_usage": 20.0
|
||
});
|
||
let res = app
|
||
.clone()
|
||
.oneshot(
|
||
Request::builder()
|
||
.method("POST")
|
||
.uri("/api/node/heartbeat")
|
||
.header("authorization", format!("Bearer {}", token))
|
||
.header("content-type", "application/json")
|
||
.body(Body::from(serde_json::to_vec(&hb_body).unwrap()))
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(res.status(), StatusCode::OK);
|
||
|
||
// 吊销后同样请求 → 401
|
||
db.revoke_node_token("node-mw-test").await.unwrap();
|
||
let res = app
|
||
.oneshot(
|
||
Request::builder()
|
||
.method("POST")
|
||
.uri("/api/node/heartbeat")
|
||
.header("authorization", format!("Bearer {}", token))
|
||
.header("content-type", "application/json")
|
||
.body(Body::from(serde_json::to_vec(&hb_body).unwrap()))
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(res.status(), StatusCode::UNAUTHORIZED);
|
||
}
|
||
|
||
/// 管理 API（Admin）端到端测试：
|
||
/// 列出节点 → 吊销 → 重发 token → 鉴权（admin 放行，node/匿名 401）。
|
||
#[tokio::test]
|
||
async fn test_admin_node_management_api() {
|
||
let temp_dir = tempfile::tempdir().unwrap();
|
||
let db_path = temp_dir.path().join("admin_db.db");
|
||
let queue_db_path = temp_dir.path().join("admin_queue.db");
|
||
let results_dir = temp_dir.path().join("results");
|
||
std::fs::create_dir_all(&results_dir).unwrap();
|
||
|
||
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 = Arc::new(GridScheduler::new(
|
||
db.clone(),
|
||
queue.clone(),
|
||
results_dir.to_string_lossy().to_string(),
|
||
));
|
||
|
||
// 预置一个已注册并颁发 token 的节点
|
||
let reg = common::models::NodeRegisterRequest {
|
||
node_id: "node-admin-test".to_string(),
|
||
host_name: "admin-host".to_string(),
|
||
max_slots: 2,
|
||
};
|
||
db.register_node(®).await.unwrap();
|
||
let original_token = db.issue_node_token("node-admin-test").await.unwrap();
|
||
|
||
let state = AppState {
|
||
db: db.clone(),
|
||
queue,
|
||
scheduler,
|
||
results_dir: results_dir.to_string_lossy().to_string(),
|
||
rate_limiter: server::api::rate_limit::RateLimiter::new(
|
||
5,
|
||
std::time::Duration::from_secs(300),
|
||
),
|
||
auth_token: Some("admin-secret".to_string()),
|
||
admin_token: Some("admin-secret".to_string()),
|
||
auth_disabled: false,
|
||
admin_sessions: std::sync::Arc::new(tokio::sync::RwLock::new(
|
||
std::collections::HashMap::new(),
|
||
)),
|
||
};
|
||
|
||
let api_router = axum::Router::new()
|
||
.route(
|
||
"/admin/nodes",
|
||
axum::routing::get(server::api::admin::list_nodes),
|
||
)
|
||
.route(
|
||
"/admin/nodes/:node_id/revoke",
|
||
axum::routing::post(server::api::admin::revoke_node),
|
||
)
|
||
.route(
|
||
"/admin/nodes/:node_id/reissue",
|
||
axum::routing::post(server::api::admin::reissue_node),
|
||
);
|
||
let auth_layer =
|
||
axum::middleware::from_fn_with_state(state.clone(), server::api::auth_middleware);
|
||
let app = axum::Router::new()
|
||
.nest("/api", api_router.layer(auth_layer))
|
||
.with_state(state);
|
||
|
||
// 1. 匿名访问 /admin/nodes → 401
|
||
let res = app
|
||
.clone()
|
||
.oneshot(
|
||
Request::builder()
|
||
.uri("/api/admin/nodes")
|
||
.body(Body::empty())
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(res.status(), StatusCode::UNAUTHORIZED);
|
||
|
||
// 2. admin 访问 /admin/nodes → 200，且能看到预置节点 token_status=active
|
||
let res = app
|
||
.clone()
|
||
.oneshot(
|
||
Request::builder()
|
||
.uri("/api/admin/nodes")
|
||
.header("authorization", "Bearer admin-secret")
|
||
.body(Body::empty())
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(res.status(), StatusCode::OK);
|
||
let body: serde_json::Value = serde_json::from_slice(
|
||
&axum::body::to_bytes(res.into_body(), usize::MAX)
|
||
.await
|
||
.unwrap(),
|
||
)
|
||
.unwrap();
|
||
let nodes = body["data"].as_array().expect("data 应为数组");
|
||
let target = nodes
|
||
.iter()
|
||
.find(|n| n["node_id"] == "node-admin-test")
|
||
.expect("应包含预置节点");
|
||
assert_eq!(target["token_status"], "active");
|
||
|
||
// 3. admin 吊销节点 token → 200
|
||
let res = app
|
||
.clone()
|
||
.oneshot(
|
||
Request::builder()
|
||
.method("POST")
|
||
.uri("/api/admin/nodes/node-admin-test/revoke")
|
||
.header("authorization", "Bearer admin-secret")
|
||
.body(Body::empty())
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(res.status(), StatusCode::OK);
|
||
// 旧 token 立即失效
|
||
assert!(db.find_node_by_token(&original_token).await.is_none());
|
||
|
||
// 4. admin 重发 token → 200，返回新明文 token
|
||
let res = app
|
||
.clone()
|
||
.oneshot(
|
||
Request::builder()
|
||
.method("POST")
|
||
.uri("/api/admin/nodes/node-admin-test/reissue")
|
||
.header("authorization", "Bearer admin-secret")
|
||
.body(Body::empty())
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(res.status(), StatusCode::OK);
|
||
let body: serde_json::Value = serde_json::from_slice(
|
||
&axum::body::to_bytes(res.into_body(), usize::MAX)
|
||
.await
|
||
.unwrap(),
|
||
)
|
||
.unwrap();
|
||
let new_token = body["node_token"].as_str().expect("应返回 node_token");
|
||
assert!(!new_token.is_empty());
|
||
assert_ne!(new_token, original_token);
|
||
// 新 token 可用
|
||
assert_eq!(
|
||
db.find_node_by_token(new_token).await.as_deref(),
|
||
Some("node-admin-test")
|
||
);
|
||
|
||
// 5. node token 不能访问 admin API → 401（即便持有有效 node token）
|
||
let res = app
|
||
.oneshot(
|
||
Request::builder()
|
||
.uri("/api/admin/nodes")
|
||
.header("authorization", format!("Bearer {}", new_token))
|
||
.body(Body::empty())
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(res.status(), StatusCode::UNAUTHORIZED);
|
||
}
|
||
|
||
/// S1 身份绑定测试：节点 A 用自己 token 冒充节点 B 发心跳 → 403；一致时 → 200。
|
||
#[tokio::test]
|
||
async fn test_node_heartbeat_node_id_binding() {
|
||
let temp_dir = tempfile::tempdir().unwrap();
|
||
let db_path = temp_dir.path().join("bind_db.db");
|
||
let queue_db_path = temp_dir.path().join("bind_queue.db");
|
||
let results_dir = temp_dir.path().join("results");
|
||
std::fs::create_dir_all(&results_dir).unwrap();
|
||
|
||
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 = Arc::new(GridScheduler::new(
|
||
db.clone(),
|
||
queue,
|
||
results_dir.to_string_lossy().to_string(),
|
||
));
|
||
|
||
// 预置节点 A，颁发 token
|
||
let reg = common::models::NodeRegisterRequest {
|
||
node_id: "node-A".to_string(),
|
||
host_name: "h".to_string(),
|
||
max_slots: 2,
|
||
};
|
||
db.register_node(®).await.unwrap();
|
||
let token_a = db.issue_node_token("node-A").await.unwrap();
|
||
|
||
let state = AppState {
|
||
db: db.clone(),
|
||
queue: Arc::new(SqliteTaskQueue::new(":memory:").await.unwrap()),
|
||
scheduler,
|
||
results_dir: results_dir.to_string_lossy().to_string(),
|
||
rate_limiter: server::api::rate_limit::RateLimiter::new(
|
||
5,
|
||
std::time::Duration::from_secs(300),
|
||
),
|
||
auth_token: Some("x".into()),
|
||
admin_token: None,
|
||
auth_disabled: false,
|
||
admin_sessions: std::sync::Arc::new(tokio::sync::RwLock::new(
|
||
std::collections::HashMap::new(),
|
||
)),
|
||
};
|
||
let api_router = axum::Router::new().route(
|
||
"/node/heartbeat",
|
||
axum::routing::post(server::api::node::heartbeat_node),
|
||
);
|
||
let auth_layer =
|
||
axum::middleware::from_fn_with_state(state.clone(), server::api::auth_middleware);
|
||
let app = axum::Router::new()
|
||
.nest("/api", api_router.layer(auth_layer))
|
||
.with_state(state);
|
||
|
||
// A 用自己 token，但 body 声称 node_id=node-B（冒充）→ 403
|
||
let hb =
|
||
serde_json::json!({"node_id":"node-B","active_slots":1,"cpu_usage":0.0,"memory_usage":0.0});
|
||
let res = app
|
||
.clone()
|
||
.oneshot(
|
||
Request::builder()
|
||
.method("POST")
|
||
.uri("/api/node/heartbeat")
|
||
.header("authorization", format!("Bearer {}", token_a))
|
||
.header("content-type", "application/json")
|
||
.body(Body::from(serde_json::to_vec(&hb).unwrap()))
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(res.status(), StatusCode::FORBIDDEN);
|
||
|
||
// A 用自己 token，body node_id=node-A（一致）→ 200
|
||
let hb_ok =
|
||
serde_json::json!({"node_id":"node-A","active_slots":1,"cpu_usage":0.0,"memory_usage":0.0});
|
||
let res = app
|
||
.oneshot(
|
||
Request::builder()
|
||
.method("POST")
|
||
.uri("/api/node/heartbeat")
|
||
.header("authorization", format!("Bearer {}", token_a))
|
||
.header("content-type", "application/json")
|
||
.body(Body::from(serde_json::to_vec(&hb_ok).unwrap()))
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(res.status(), StatusCode::OK);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_admin_login_flow() {
|
||
let temp_dir = tempfile::tempdir().unwrap();
|
||
let db_path = temp_dir.path().join("login_db.db");
|
||
let queue_db_path = temp_dir.path().join("login_queue.db");
|
||
let results_dir = temp_dir.path().join("results");
|
||
std::fs::create_dir_all(&results_dir).unwrap();
|
||
|
||
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 = Arc::new(GridScheduler::new(
|
||
db.clone(),
|
||
queue.clone(),
|
||
results_dir.to_string_lossy().to_string(),
|
||
));
|
||
|
||
let admin_pass = "my_short_admin_password_123";
|
||
let state = AppState {
|
||
db,
|
||
queue,
|
||
scheduler,
|
||
results_dir: results_dir.to_string_lossy().to_string(),
|
||
rate_limiter: server::api::rate_limit::RateLimiter::new(
|
||
5,
|
||
std::time::Duration::from_secs(300),
|
||
),
|
||
auth_token: Some(admin_pass.to_string()),
|
||
admin_token: Some(admin_pass.to_string()),
|
||
auth_disabled: false,
|
||
admin_sessions: std::sync::Arc::new(tokio::sync::RwLock::new(
|
||
std::collections::HashMap::new(),
|
||
)),
|
||
};
|
||
|
||
let api_router = axum::Router::new()
|
||
.route("/login", axum::routing::post(server::api::auth::login))
|
||
.route(
|
||
"/auth/check",
|
||
axum::routing::get(server::api::auth::check_auth),
|
||
)
|
||
.layer(axum::middleware::from_fn_with_state(
|
||
state.clone(),
|
||
server::api::auth_middleware,
|
||
));
|
||
|
||
let app = axum::Router::new()
|
||
.nest("/api", api_router)
|
||
.with_state(state);
|
||
|
||
// 1. 密码错误 ➔ 401
|
||
let wrong_body = serde_json::json!({ "password": "wrong_password" });
|
||
let res = app
|
||
.clone()
|
||
.oneshot(
|
||
Request::builder()
|
||
.method("POST")
|
||
.uri("/api/login")
|
||
.header("content-type", "application/json")
|
||
.extension(axum::extract::ConnectInfo(std::net::SocketAddr::from((
|
||
[127, 0, 0, 1],
|
||
12345,
|
||
))))
|
||
.body(Body::from(serde_json::to_vec(&wrong_body).unwrap()))
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(res.status(), StatusCode::UNAUTHORIZED);
|
||
|
||
// 2. 正确密码 ➔ 200 + token
|
||
let right_body = serde_json::json!({ "password": admin_pass });
|
||
let res = app
|
||
.clone()
|
||
.oneshot(
|
||
Request::builder()
|
||
.method("POST")
|
||
.uri("/api/login")
|
||
.header("content-type", "application/json")
|
||
.extension(axum::extract::ConnectInfo(std::net::SocketAddr::from((
|
||
[127, 0, 0, 1],
|
||
12345,
|
||
))))
|
||
.body(Body::from(serde_json::to_vec(&right_body).unwrap()))
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(res.status(), StatusCode::OK);
|
||
let bytes = axum::body::to_bytes(res.into_body(), 1024 * 1024)
|
||
.await
|
||
.unwrap();
|
||
let json: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
|
||
assert_eq!(json["success"], true);
|
||
let token = json["token"].as_str().unwrap();
|
||
assert!(token.len() == 64);
|
||
|
||
// 3. 携带拿到的 Token 访问受保护的 /api/auth/check ➔ 200
|
||
let res = app
|
||
.oneshot(
|
||
Request::builder()
|
||
.method("GET")
|
||
.uri("/api/auth/check")
|
||
.header("authorization", format!("Bearer {}", token))
|
||
.body(Body::empty())
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(res.status(), StatusCode::OK);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_node_approval_workflow() {
|
||
let temp_dir = tempfile::tempdir().unwrap();
|
||
let db_path = temp_dir.path().join("appr_db.db");
|
||
let queue_db_path = temp_dir.path().join("appr_queue.db");
|
||
let results_dir = temp_dir.path().join("results");
|
||
std::fs::create_dir_all(&results_dir).unwrap();
|
||
|
||
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 = Arc::new(GridScheduler::new(
|
||
db.clone(),
|
||
queue.clone(),
|
||
results_dir.to_string_lossy().to_string(),
|
||
));
|
||
|
||
let admin_pass = "admin_approval_secret";
|
||
let state = AppState {
|
||
db,
|
||
queue,
|
||
scheduler,
|
||
results_dir: results_dir.to_string_lossy().to_string(),
|
||
rate_limiter: server::api::rate_limit::RateLimiter::new(
|
||
5,
|
||
std::time::Duration::from_secs(300),
|
||
),
|
||
auth_token: Some(admin_pass.to_string()),
|
||
admin_token: Some(admin_pass.to_string()),
|
||
auth_disabled: false,
|
||
admin_sessions: std::sync::Arc::new(tokio::sync::RwLock::new(
|
||
std::collections::HashMap::new(),
|
||
)),
|
||
};
|
||
|
||
let api_router = axum::Router::new()
|
||
.route(
|
||
"/node/register",
|
||
axum::routing::post(server::api::node::register_node),
|
||
)
|
||
.route(
|
||
"/node/check_status",
|
||
axum::routing::post(server::api::node::check_node_status),
|
||
)
|
||
.route(
|
||
"/node/heartbeat",
|
||
axum::routing::post(server::api::node::heartbeat_node),
|
||
)
|
||
.route(
|
||
"/admin/nodes/:node_id/approve",
|
||
axum::routing::post(server::api::admin::approve_node),
|
||
)
|
||
.layer(axum::middleware::from_fn_with_state(
|
||
state.clone(),
|
||
server::api::auth_middleware,
|
||
));
|
||
|
||
let app = axum::Router::new()
|
||
.nest("/api", api_router)
|
||
.with_state(state);
|
||
|
||
// 1. 新节点免凭据申请注册 ➔ 200 + status: pending_approval
|
||
let reg_body = serde_json::json!({
|
||
"node_id": "node-pending-01",
|
||
"host_name": "worker-host",
|
||
"max_slots": 4
|
||
});
|
||
let res = app
|
||
.clone()
|
||
.oneshot(
|
||
Request::builder()
|
||
.method("POST")
|
||
.uri("/api/node/register")
|
||
.header("content-type", "application/json")
|
||
.body(Body::from(serde_json::to_vec(®_body).unwrap()))
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(res.status(), StatusCode::OK);
|
||
let bytes = axum::body::to_bytes(res.into_body(), 1024 * 1024)
|
||
.await
|
||
.unwrap();
|
||
let json: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
|
||
assert_eq!(json["status"], "pending_approval");
|
||
|
||
// 2. Node 端轮询查状态 ➔ status: pending_approval
|
||
let check_body = serde_json::json!({ "node_id": "node-pending-01" });
|
||
let res = app
|
||
.clone()
|
||
.oneshot(
|
||
Request::builder()
|
||
.method("POST")
|
||
.uri("/api/node/check_status")
|
||
.header("content-type", "application/json")
|
||
.body(Body::from(serde_json::to_vec(&check_body).unwrap()))
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(res.status(), StatusCode::OK);
|
||
let bytes = axum::body::to_bytes(res.into_body(), 1024 * 1024)
|
||
.await
|
||
.unwrap();
|
||
let json: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
|
||
assert_eq!(json["status"], "pending_approval");
|
||
|
||
// 3. 管理员在 Dashboard 点击同意 ➔ 200
|
||
let res = app
|
||
.clone()
|
||
.oneshot(
|
||
Request::builder()
|
||
.method("POST")
|
||
.uri("/api/admin/nodes/node-pending-01/approve")
|
||
.header("authorization", format!("Bearer {}", admin_pass))
|
||
.body(Body::empty())
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(res.status(), StatusCode::OK);
|
||
let bytes = axum::body::to_bytes(res.into_body(), 1024 * 1024)
|
||
.await
|
||
.unwrap();
|
||
println!("Step 3 body: {}", String::from_utf8_lossy(&bytes));
|
||
|
||
// 4. Node 端再次轮询查状态 ➔ status: approved + 获取 node_token
|
||
let res = app
|
||
.clone()
|
||
.oneshot(
|
||
Request::builder()
|
||
.method("POST")
|
||
.uri("/api/node/check_status")
|
||
.header("content-type", "application/json")
|
||
.body(Body::from(serde_json::to_vec(&check_body).unwrap()))
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(res.status(), StatusCode::OK);
|
||
let bytes = axum::body::to_bytes(res.into_body(), 1024 * 1024)
|
||
.await
|
||
.unwrap();
|
||
let json: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
|
||
println!("Step 4 status returned: {:?}", json);
|
||
|
||
assert_eq!(json["status"], "approved");
|
||
let node_token = json["node_token"].as_str().unwrap();
|
||
|
||
// 5. Node 携带拿到到的专属 Token 发送心跳 ➔ 200
|
||
let hb_body = serde_json::json!({
|
||
"node_id": "node-pending-01",
|
||
"active_slots": 1,
|
||
"cpu_usage": 10.0,
|
||
"memory_usage": 20.0
|
||
});
|
||
let res = app
|
||
.oneshot(
|
||
Request::builder()
|
||
.method("POST")
|
||
.uri("/api/node/heartbeat")
|
||
.header("authorization", format!("Bearer {}", node_token))
|
||
.header("content-type", "application/json")
|
||
.body(Body::from(serde_json::to_vec(&hb_body).unwrap()))
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(res.status(), StatusCode::OK);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_admin_sessions_capacity_limit() {
|
||
let temp_dir = tempfile::tempdir().unwrap();
|
||
let db_path = temp_dir.path().join("sess_db.db");
|
||
let queue_db_path = temp_dir.path().join("sess_queue.db");
|
||
let results_dir = temp_dir.path().join("results");
|
||
std::fs::create_dir_all(&results_dir).unwrap();
|
||
|
||
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 = Arc::new(GridScheduler::new(
|
||
db.clone(),
|
||
queue.clone(),
|
||
results_dir.to_string_lossy().to_string(),
|
||
));
|
||
|
||
let admin_pass = "admin_capacity_secret";
|
||
let state = AppState {
|
||
db,
|
||
queue,
|
||
scheduler,
|
||
results_dir: results_dir.to_string_lossy().to_string(),
|
||
rate_limiter: server::api::rate_limit::RateLimiter::new(
|
||
1000,
|
||
std::time::Duration::from_secs(300),
|
||
),
|
||
auth_token: Some(admin_pass.to_string()),
|
||
admin_token: Some(admin_pass.to_string()),
|
||
auth_disabled: false,
|
||
admin_sessions: std::sync::Arc::new(tokio::sync::RwLock::new(
|
||
std::collections::HashMap::new(),
|
||
)),
|
||
};
|
||
|
||
let api_router =
|
||
axum::Router::new().route("/login", axum::routing::post(server::api::auth::login));
|
||
let app = axum::Router::new()
|
||
.nest("/api", api_router)
|
||
.with_state(state.clone());
|
||
|
||
// 连续登录 105 次，超过 100 容量上限
|
||
for _ in 0..105 {
|
||
let right_body = serde_json::json!({ "password": admin_pass });
|
||
let res = app
|
||
.clone()
|
||
.oneshot(
|
||
Request::builder()
|
||
.method("POST")
|
||
.uri("/api/login")
|
||
.header("content-type", "application/json")
|
||
.extension(axum::extract::ConnectInfo(std::net::SocketAddr::from((
|
||
[127, 0, 0, 1],
|
||
12345,
|
||
))))
|
||
.body(Body::from(serde_json::to_vec(&right_body).unwrap()))
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(res.status(), StatusCode::OK);
|
||
}
|
||
|
||
// 验证 sessions 集合保存的数量不超过 MAX_ADMIN_SESSIONS (100)
|
||
let sessions = state.admin_sessions.read().await;
|
||
assert!(sessions.len() <= server::api::MAX_ADMIN_SESSIONS);
|
||
assert_eq!(sessions.len(), 100);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_node_register_rate_limit() {
|
||
let temp_dir = tempfile::tempdir().unwrap();
|
||
let db_path = temp_dir.path().join("reg_limit_db.db");
|
||
let queue_db_path = temp_dir.path().join("reg_limit_queue.db");
|
||
let results_dir = temp_dir.path().join("results");
|
||
std::fs::create_dir_all(&results_dir).unwrap();
|
||
|
||
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 = Arc::new(GridScheduler::new(
|
||
db.clone(),
|
||
queue.clone(),
|
||
results_dir.to_string_lossy().to_string(),
|
||
));
|
||
|
||
let state = AppState {
|
||
db,
|
||
queue,
|
||
scheduler,
|
||
results_dir: results_dir.to_string_lossy().to_string(),
|
||
rate_limiter: server::api::rate_limit::RateLimiter::new(
|
||
5,
|
||
std::time::Duration::from_secs(300),
|
||
),
|
||
auth_token: None,
|
||
admin_token: None,
|
||
auth_disabled: false,
|
||
admin_sessions: std::sync::Arc::new(tokio::sync::RwLock::new(
|
||
std::collections::HashMap::new(),
|
||
)),
|
||
};
|
||
|
||
// 注册端点专用限流器（count_all=true）：对成功请求也计数，模拟生产 main.rs 的注册节流配置。
|
||
let register_limiter =
|
||
server::api::rate_limit::RateLimiter::new_count_all(5, std::time::Duration::from_secs(60));
|
||
let register_rate_limit_layer = axum::middleware::from_fn_with_state(
|
||
register_limiter,
|
||
server::api::rate_limit::rate_limit_middleware,
|
||
);
|
||
|
||
let app = axum::Router::new()
|
||
.route(
|
||
"/api/node/register",
|
||
axum::routing::post(server::api::node::register_node).layer(register_rate_limit_layer),
|
||
)
|
||
.with_state(state);
|
||
|
||
let reg_body = serde_json::json!({
|
||
"node_id": "test-limit-node",
|
||
"host_name": "limit-host",
|
||
"max_slots": 4
|
||
});
|
||
|
||
// 5 次以内的注册尝试 ➔ 200
|
||
for _ in 0..5 {
|
||
let res = app
|
||
.clone()
|
||
.oneshot(
|
||
Request::builder()
|
||
.method("POST")
|
||
.uri("/api/node/register")
|
||
.header("content-type", "application/json")
|
||
.extension(axum::extract::ConnectInfo(std::net::SocketAddr::from((
|
||
[127, 0, 0, 1],
|
||
12345,
|
||
))))
|
||
.body(Body::from(serde_json::to_vec(®_body).unwrap()))
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(res.status(), StatusCode::OK);
|
||
}
|
||
|
||
// 第 6 次触发限流 ➔ 429 TOO_MANY_REQUESTS
|
||
let res = app
|
||
.oneshot(
|
||
Request::builder()
|
||
.method("POST")
|
||
.uri("/api/node/register")
|
||
.header("content-type", "application/json")
|
||
.extension(axum::extract::ConnectInfo(std::net::SocketAddr::from((
|
||
[127, 0, 0, 1],
|
||
12345,
|
||
))))
|
||
.body(Body::from(serde_json::to_vec(®_body).unwrap()))
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(res.status(), StatusCode::TOO_MANY_REQUESTS);
|
||
}
|
||
|
||
/// 回归保护：通用 API 限流器（count_all=false）不应因 /node/register 的成功响应计数。
|
||
/// 历史缺陷：此前中间件对 is_register 无条件计数，导致通用限流器复用时成功注册会把 IP
|
||
/// 锁出整个 /api/*（跨端点连锁）。此测试断言多次成功注册后通用限流器仍放行。
|
||
#[tokio::test]
|
||
async fn test_general_limiter_does_not_count_successful_register() {
|
||
let temp_dir = tempfile::tempdir().unwrap();
|
||
let db_path = temp_dir.path().join("gen_limit_db.db");
|
||
let queue_db_path = temp_dir.path().join("gen_limit_queue.db");
|
||
let results_dir = temp_dir.path().join("results");
|
||
std::fs::create_dir_all(&results_dir).unwrap();
|
||
|
||
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 = Arc::new(GridScheduler::new(
|
||
db.clone(),
|
||
queue.clone(),
|
||
results_dir.to_string_lossy().to_string(),
|
||
));
|
||
|
||
let state = AppState {
|
||
db,
|
||
queue,
|
||
scheduler,
|
||
results_dir: results_dir.to_string_lossy().to_string(),
|
||
rate_limiter: server::api::rate_limit::RateLimiter::new(
|
||
5,
|
||
std::time::Duration::from_secs(300),
|
||
),
|
||
auth_token: None,
|
||
admin_token: None,
|
||
auth_disabled: false,
|
||
admin_sessions: std::sync::Arc::new(tokio::sync::RwLock::new(
|
||
std::collections::HashMap::new(),
|
||
)),
|
||
};
|
||
|
||
// 通用限流器（new = count_all=false），阈值仅 3，远低于下面的请求数。
|
||
// 若仍对成功注册计数，第 4 次就会 429。
|
||
let general_limiter =
|
||
server::api::rate_limit::RateLimiter::new(3, std::time::Duration::from_secs(60));
|
||
let rate_limit_layer = axum::middleware::from_fn_with_state(
|
||
general_limiter,
|
||
server::api::rate_limit::rate_limit_middleware,
|
||
);
|
||
|
||
let app = axum::Router::new()
|
||
.route(
|
||
"/api/node/register",
|
||
axum::routing::post(server::api::node::register_node),
|
||
)
|
||
.layer(rate_limit_layer)
|
||
.with_state(state);
|
||
|
||
let reg_body = serde_json::json!({
|
||
"node_id": "gen-limit-node",
|
||
"host_name": "gen-host",
|
||
"max_slots": 4
|
||
});
|
||
|
||
// 连续 6 次成功注册（远超阈值 3）：通用限流器不应计数成功响应，全部应为 200。
|
||
for i in 0..6 {
|
||
// 每次用不同 node_id 避免重复注册逻辑干扰
|
||
let body = serde_json::json!({
|
||
"node_id": format!("gen-limit-node-{}", i),
|
||
"host_name": "gen-host",
|
||
"max_slots": 4
|
||
});
|
||
let res = app
|
||
.clone()
|
||
.oneshot(
|
||
Request::builder()
|
||
.method("POST")
|
||
.uri("/api/node/register")
|
||
.header("content-type", "application/json")
|
||
.extension(axum::extract::ConnectInfo(std::net::SocketAddr::from((
|
||
[127, 0, 0, 1],
|
||
12345,
|
||
))))
|
||
.body(Body::from(serde_json::to_vec(&body).unwrap()))
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(
|
||
res.status(),
|
||
StatusCode::OK,
|
||
"第 {} 次注册应成功，通用限流器不应计数成功响应",
|
||
i + 1
|
||
);
|
||
}
|
||
|
||
// 静默 reg_body 未使用的警告（保留以对齐其它测试结构）
|
||
let _ = ®_body;
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_cors_same_origin_and_local_policy() {
|
||
let temp_dir = tempfile::tempdir().unwrap();
|
||
let db_path = temp_dir.path().join("cors_db.db");
|
||
let queue_db_path = temp_dir.path().join("cors_queue.db");
|
||
let results_dir = temp_dir.path().join("results");
|
||
std::fs::create_dir_all(&results_dir).unwrap();
|
||
|
||
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 = Arc::new(GridScheduler::new(
|
||
db.clone(),
|
||
queue.clone(),
|
||
results_dir.to_string_lossy().to_string(),
|
||
));
|
||
|
||
let state = AppState {
|
||
db,
|
||
queue,
|
||
scheduler,
|
||
results_dir: results_dir.to_string_lossy().to_string(),
|
||
rate_limiter: server::api::rate_limit::RateLimiter::new(
|
||
100,
|
||
std::time::Duration::from_secs(300),
|
||
),
|
||
auth_token: None,
|
||
admin_token: None,
|
||
auth_disabled: false,
|
||
admin_sessions: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
|
||
};
|
||
|
||
let app = axum::Router::new()
|
||
.route(
|
||
"/api/status",
|
||
axum::routing::get(server::api::status::get_status),
|
||
)
|
||
.layer(server::cors::build_cors_layer())
|
||
.with_state(state);
|
||
|
||
// 1. 本地 localhost 请求 -> 允许 CORS
|
||
let res = app
|
||
.clone()
|
||
.oneshot(
|
||
Request::builder()
|
||
.uri("/api/status")
|
||
.header("origin", "http://localhost:5173")
|
||
.body(Body::empty())
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(res.status(), StatusCode::OK);
|
||
assert_eq!(
|
||
res.headers()
|
||
.get("access-control-allow-origin")
|
||
.unwrap()
|
||
.to_str()
|
||
.unwrap(),
|
||
"http://localhost:5173"
|
||
);
|
||
|
||
// 2. 本地 127.0.0.1 请求 -> 允许 CORS
|
||
let res = app
|
||
.clone()
|
||
.oneshot(
|
||
Request::builder()
|
||
.uri("/api/status")
|
||
.header("origin", "http://127.0.0.1:3000")
|
||
.body(Body::empty())
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(res.status(), StatusCode::OK);
|
||
assert_eq!(
|
||
res.headers()
|
||
.get("access-control-allow-origin")
|
||
.unwrap()
|
||
.to_str()
|
||
.unwrap(),
|
||
"http://127.0.0.1:3000"
|
||
);
|
||
|
||
// 3. 同源请求 (Origin 匹配 Host 标头) -> 允许 CORS
|
||
let res = app
|
||
.clone()
|
||
.oneshot(
|
||
Request::builder()
|
||
.uri("/api/status")
|
||
.header("host", "192.168.1.100:8090")
|
||
.header("origin", "http://192.168.1.100:8090")
|
||
.body(Body::empty())
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(res.status(), StatusCode::OK);
|
||
assert_eq!(
|
||
res.headers()
|
||
.get("access-control-allow-origin")
|
||
.unwrap()
|
||
.to_str()
|
||
.unwrap(),
|
||
"http://192.168.1.100:8090"
|
||
);
|
||
|
||
// 4. 外部非法跨源请求 -> 拒绝 CORS
|
||
let res = app
|
||
.oneshot(
|
||
Request::builder()
|
||
.uri("/api/status")
|
||
.header("host", "192.168.1.100:8090")
|
||
.header("origin", "https://attacker.example.com")
|
||
.body(Body::empty())
|
||
.unwrap(),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(res.status(), StatusCode::OK);
|
||
assert!(res.headers().get("access-control-allow-origin").is_none());
|
||
}
|