完善服务启动与管理安全

This commit is contained in:
jiesen
2025-12-20 21:04:22 +07:00
parent 0329edaed0
commit 6802dc5a4c
16 changed files with 644 additions and 104 deletions
+1 -1
View File
@@ -3367,7 +3367,7 @@ dependencies = [
[[package]]
name = "proxycast"
version = "0.14.5"
version = "0.14.6"
dependencies = [
"anyhow",
"async-stream",
+7
View File
@@ -1,3 +1,10 @@
fn main() {
// tauri::generate_context! 在编译期会校验 `frontendDist` 路径是否存在。
// 开发/CI 场景下可能只跑 `cargo check/test` 而未先构建前端,从而导致宏 panic。
// 这里提前创建配置中的 `../dist` 目录,避免无关的编译阻塞。
if let Ok(manifest_dir) = std::env::var("CARGO_MANIFEST_DIR") {
let dist_dir = std::path::PathBuf::from(manifest_dir).join("../dist");
let _ = std::fs::create_dir_all(dist_dir);
}
tauri_build::build()
}
+43
View File
@@ -370,6 +370,8 @@ impl HotReloadManager {
/// 验证配置
fn validate_config(&self, config: &Config) -> Result<(), HotReloadError> {
let is_localhost = is_localhost_host(&config.server.host);
// 验证端口范围
if config.server.port == 0 {
return Err(HotReloadError::ValidationError(
@@ -377,6 +379,12 @@ impl HotReloadManager {
));
}
if !is_localhost {
return Err(HotReloadError::ValidationError(
"当前版本仅支持本地监听,请使用 127.0.0.1/localhost/::1".to_string(),
));
}
// 验证重试配置
if config.retry.max_retries > 100 {
return Err(HotReloadError::ValidationError(
@@ -397,6 +405,32 @@ impl HotReloadManager {
));
}
if config.server.api_key.trim().is_empty() {
return Err(HotReloadError::ValidationError(
"API Key 不能为空".to_string(),
));
}
if (!is_localhost || config.remote_management.allow_remote)
&& crate::config::is_default_api_key(&config.server.api_key)
{
return Err(HotReloadError::ValidationError(
"非本地访问场景下禁止使用默认 API Key,请设置强口令".to_string(),
));
}
if config.server.tls.enable {
return Err(HotReloadError::ValidationError(
"当前版本暂不支持 TLS,请关闭 TLS 配置".to_string(),
));
}
if config.remote_management.allow_remote {
return Err(HotReloadError::ValidationError(
"当前版本未启用 TLS,禁止开启远程管理".to_string(),
));
}
Ok(())
}
@@ -437,6 +471,15 @@ impl HotReloadManager {
}
}
fn is_localhost_host(host: &str) -> bool {
if host == "localhost" {
return true;
}
host.parse::<std::net::IpAddr>()
.map(|addr| addr.is_loopback())
.unwrap_or(false)
}
/// 热重载状态
#[derive(Debug, Clone, serde::Serialize)]
pub struct HotReloadStatus {
+20 -11
View File
@@ -313,7 +313,10 @@ impl ImportService {
// 如果是脱敏数据,清理凭证池中的占位符
if bundle.redacted {
Self::clean_redacted_credentials(&mut config);
let server_key_cleared = Self::clean_redacted_credentials(&mut config);
if server_key_cleared {
warnings.push("检测到脱敏的服务器 API Key,已清空,需要手动设置".to_string());
}
}
Ok(ImportResult::success_with_warnings(config, warnings))
@@ -469,7 +472,9 @@ impl ImportService {
/// 清理脱敏的凭证数据
///
/// 移除凭证池中使用占位符的条目
fn clean_redacted_credentials(config: &mut Config) {
fn clean_redacted_credentials(config: &mut Config) -> bool {
let mut server_key_cleared = false;
// 清理 OpenAI 凭证池中的脱敏条目
config
.credential_pool
@@ -490,10 +495,13 @@ impl ImportService {
config.providers.claude.api_key = None;
}
// 清理服务器 API 密钥(如果是脱敏的,恢复默认值)
// 清理服务器 API 密钥(如果是脱敏的,清空并提示手动设置)
if config.server.api_key == REDACTED_PLACEHOLDER {
config.server.api_key = "proxy_cast".to_string();
config.server.api_key = String::new();
server_key_cleared = true;
}
server_key_cleared
}
/// 从文件导入配置
@@ -638,7 +646,7 @@ server:
let current = Config::default();
let yaml = r#"
server:
host: 0.0.0.0
host: 127.0.0.1
port: 9000
api_key: new_key
"#;
@@ -646,7 +654,7 @@ server:
let result = ImportService::import_yaml(yaml, &current, &options).expect("导入应成功");
assert!(result.success);
assert_eq!(result.config.server.host, "0.0.0.0");
assert_eq!(result.config.server.host, "127.0.0.1");
assert_eq!(result.config.server.port, 9000);
assert_eq!(result.config.server.api_key, "new_key");
}
@@ -664,7 +672,7 @@ server:
let yaml = r#"
server:
host: 0.0.0.0
host: 127.0.0.1
port: 9000
api_key: new_key
credential_pool:
@@ -677,7 +685,7 @@ credential_pool:
assert!(result.success);
// 服务器配置应被更新
assert_eq!(result.config.server.host, "0.0.0.0");
assert_eq!(result.config.server.host, "127.0.0.1");
// 凭证池应合并
assert_eq!(result.config.credential_pool.openai.len(), 2);
}
@@ -757,10 +765,11 @@ credential_pool:
proxy_url: None,
});
ImportService::clean_redacted_credentials(&mut config);
let server_key_cleared = ImportService::clean_redacted_credentials(&mut config);
// 服务器 API 密钥应恢复默认值
assert_eq!(config.server.api_key, "proxy_cast");
// 服务器 API 密钥应被清空并提示手动设置
assert!(server_key_cleared);
assert_eq!(config.server.api_key, "");
// Provider API 密钥应被清除
assert!(config.providers.openai.api_key.is_none());
// 凭证池中脱敏的条目应被移除
+5 -5
View File
@@ -21,11 +21,11 @@ pub use hot_reload::{
pub use import::{ImportError, ImportOptions, ImportResult, ImportService, ValidationResult};
pub use path_utils::{collapse_tilde, contains_tilde, expand_tilde};
pub use types::{
AmpConfig, AmpModelMapping, ApiKeyEntry, Config, CredentialEntry, CredentialPoolConfig,
CustomProviderConfig, GeminiApiKeyEntry, IFlowCredentialEntry, InjectionRuleConfig,
InjectionSettings, LoggingConfig, ProviderConfig, ProvidersConfig, QuotaExceededConfig,
RemoteManagementConfig, RetrySettings, RoutingConfig, RoutingRuleConfig, ServerConfig,
TlsConfig, VertexApiKeyEntry, VertexModelAlias,
is_default_api_key, AmpConfig, AmpModelMapping, ApiKeyEntry, Config, CredentialEntry,
CredentialPoolConfig, CustomProviderConfig, GeminiApiKeyEntry, IFlowCredentialEntry,
InjectionRuleConfig, InjectionSettings, LoggingConfig, ProviderConfig, ProvidersConfig,
QuotaExceededConfig, RemoteManagementConfig, RetrySettings, RoutingConfig, RoutingRuleConfig,
ServerConfig, TlsConfig, VertexApiKeyEntry, VertexModelAlias, DEFAULT_API_KEY,
};
pub use yaml::{
load_config, save_config, save_config_yaml, ConfigError, ConfigManager, YamlService,
+20 -24
View File
@@ -15,8 +15,8 @@ use tempfile::NamedTempFile;
fn arb_host() -> impl Strategy<Value = String> {
prop_oneof![
Just("127.0.0.1".to_string()),
Just("0.0.0.0".to_string()),
Just("localhost".to_string()),
Just("::1".to_string()),
"[0-9]{1,3}\\.[0-9]{1,3}\\.[0-9]{1,3}\\.[0-9]{1,3}".prop_map(|s| s),
]
}
@@ -2120,40 +2120,36 @@ proptest! {
/// **Validates: Requirements 5.5**
#[test]
fn prop_export_import_redacted_loses_secrets(config in arb_config_with_secrets()) {
// 导出为 YAML(脱敏)
let yaml = ExportService::export_yaml(&config, true)
// 导出为脱敏 bundle
let options = ExportOptions::redacted();
let bundle = ExportService::export(&config, &options, "1.0.0")
.expect("导出应成功");
// 导入 YAML
// 导入 bundle(脱敏数据会触发清理)
let empty_config = Config::default();
let options = ImportOptions::replace();
let result = ImportService::import_yaml(&yaml, &empty_config, &options)
let import_options = ImportOptions::replace();
let result = ImportService::import(
&bundle,
&empty_config,
&import_options,
&config.auth_dir,
)
.expect("导入应成功");
// 清理脱敏数据
let mut imported = result.config;
ImportService::import(
&ExportBundle::new("1.0.0"),
&imported,
&ImportOptions::merge(),
&config.auth_dir,
).ok(); // 忽略结果,只是为了触发清理
// 验证脱敏后的配置不包含原始敏感信息
// 服务器 API 密钥应为脱敏占位符或默认值
prop_assert!(
imported.server.api_key == REDACTED_PLACEHOLDER ||
imported.server.api_key == "proxy_cast",
"脱敏后服务器 API 密钥应为占位符或默认值: {}",
imported.server.api_key
// 服务器 API 密钥应被清空
prop_assert_eq!(
result.config.server.api_key,
"",
"脱敏后服务器 API 密钥应被清空"
);
// 如果原始配置有 OpenAI API 密钥,导入后应为脱敏占位符
if config.providers.openai.api_key.is_some() {
prop_assert_eq!(
imported.providers.openai.api_key,
Some(REDACTED_PLACEHOLDER.to_string()),
"脱敏后 OpenAI API 密钥应为占位符"
result.config.providers.openai.api_key,
None,
"脱敏后 OpenAI API 密钥应被清空"
);
}
}
+8 -1
View File
@@ -334,8 +334,15 @@ fn default_port() -> u16 {
8999
}
pub const DEFAULT_API_KEY: &str = "proxy_cast";
fn default_api_key() -> String {
"proxy_cast".to_string()
DEFAULT_API_KEY.to_string()
}
/// 是否为默认 API Key
pub fn is_default_api_key(api_key: &str) -> bool {
api_key == DEFAULT_API_KEY
}
impl Default for ServerConfig {
+2 -2
View File
@@ -708,7 +708,7 @@ mod unit_tests {
fn test_parse_yaml_minimal() {
let yaml = r#"
server:
host: "0.0.0.0"
host: "127.0.0.1"
port: 9000
api_key: "test-key"
providers:
@@ -716,7 +716,7 @@ providers:
enabled: true
"#;
let config = ConfigManager::parse_yaml(yaml).unwrap();
assert_eq!(config.server.host, "0.0.0.0");
assert_eq!(config.server.host, "127.0.0.1");
assert_eq!(config.server.port, 9000);
assert_eq!(config.server.api_key, "test-key");
assert!(config.providers.kiro.enabled);
+22 -17
View File
@@ -1,11 +1,8 @@
use rusqlite::Connection;
use serde_json::Value;
/// 从旧的 JSON 配置迁移数据到 SQLite
#[allow(dead_code)]
pub fn migrate_from_json(
conn: &Connection,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
pub fn migrate_from_json(conn: &Connection) -> Result<(), String> {
// 检查是否已经迁移过
let migrated: bool = conn
.query_row(
@@ -20,27 +17,35 @@ pub fn migrate_from_json(
return Ok(());
}
// 读取旧配置文件
let home = dirs::home_dir().ok_or("Cannot find home directory")?;
let config_path = home.join(".proxycast").join("config.json");
// 读取旧配置文件(兼容历史路径)
let mut legacy_paths = Vec::new();
if let Some(config_dir) = dirs::config_dir() {
legacy_paths.push(config_dir.join("proxycast").join("config.json"));
}
if let Some(home) = dirs::home_dir() {
legacy_paths.push(home.join(".proxycast").join("config.json"));
}
if config_path.exists() {
let content = std::fs::read_to_string(&config_path)?;
let _config: Value = serde_json::from_str(&content)?;
if let Some(config_path) = legacy_paths.into_iter().find(|path| path.exists()) {
// 备份旧配置,避免误覆盖
let backup_path = config_path.with_file_name("config.json.backup");
if !backup_path.exists() {
std::fs::copy(&config_path, &backup_path)
.map_err(|e| format!("备份旧配置失败: {}", e))?;
}
// TODO: 解析旧配置并插入到数据库
// 这里需要根据实际的旧配置格式来实现
// 备份旧配置
let backup_path = home.join(".proxycast").join("config.json.backup");
std::fs::copy(&config_path, &backup_path)?;
return Err(format!(
"检测到旧版 config.json({}),当前版本尚未支持自动迁移。请手动导出/重建配置后再启动。",
config_path.display()
));
}
// 标记迁移完成
conn.execute(
"INSERT OR REPLACE INTO settings (key, value) VALUES ('migrated_from_json', 'true')",
[],
)?;
)
.map_err(|e| e.to_string())?;
Ok(())
}
+10 -8
View File
@@ -9,20 +9,22 @@ use std::sync::{Arc, Mutex};
pub type DbConnection = Arc<Mutex<Connection>>;
/// 获取数据库文件路径
pub fn get_db_path() -> PathBuf {
let home = dirs::home_dir().expect("Cannot find home directory");
pub fn get_db_path() -> Result<PathBuf, String> {
let home = dirs::home_dir().ok_or_else(|| "无法获取主目录".to_string())?;
let db_dir = home.join(".proxycast");
std::fs::create_dir_all(&db_dir).expect("Cannot create .proxycast directory");
db_dir.join("proxycast.db")
std::fs::create_dir_all(&db_dir)
.map_err(|e| format!("无法创建数据库目录 {:?}: {}", db_dir, e))?;
Ok(db_dir.join("proxycast.db"))
}
/// 初始化数据库连接
pub fn init_database() -> Result<DbConnection, rusqlite::Error> {
let db_path = get_db_path();
let conn = Connection::open(&db_path)?;
pub fn init_database() -> Result<DbConnection, String> {
let db_path = get_db_path()?;
let conn = Connection::open(&db_path).map_err(|e| e.to_string())?;
// 创建表结构
schema::create_tables(&conn)?;
schema::create_tables(&conn).map_err(|e| e.to_string())?;
migration::migrate_from_json(&conn)?;
Ok(Arc::new(Mutex::new(conn)))
}
+76 -5
View File
@@ -19,6 +19,7 @@ pub mod telemetry;
pub mod tray;
pub mod websocket;
use rand::{distributions::Alphanumeric, Rng};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tauri::{Manager, Runtime};
@@ -175,6 +176,15 @@ mod tests {
pub type AppState = Arc<RwLock<server::ServerState>>;
pub type LogState = Arc<RwLock<logger::LogStore>>;
fn generate_api_key() -> String {
let token: String = rand::thread_rng()
.sample_iter(&Alphanumeric)
.take(32)
.map(char::from)
.collect();
format!("pc_{token}")
}
#[tauri::command]
async fn start_server(
state: tauri::State<'_, AppState>,
@@ -1344,12 +1354,57 @@ async fn test_api(
#[cfg_attr(mobile, tauri::mobile_entry_point)]
pub fn run() {
let config = config::load_config().unwrap_or_default();
let state: AppState = Arc::new(RwLock::new(server::ServerState::new(config)));
let logs: LogState = Arc::new(RwLock::new(logger::LogStore::new()));
let mut config = match config::load_config() {
Ok(cfg) => cfg,
Err(err) => {
tracing::error!("配置加载失败,已中止启动: {}", err);
eprintln!("配置加载失败,已中止启动: {}", err);
return;
}
};
if config.server.api_key == config::DEFAULT_API_KEY {
let new_key = generate_api_key();
config.server.api_key = new_key.clone();
if let Err(err) = config::save_config_yaml(&config) {
tracing::error!("自动生成 API key 失败,无法保存配置,已中止启动: {}", err);
eprintln!("自动生成 API key 失败,无法保存配置,已中止启动: {}", err);
return;
}
tracing::info!("检测到默认 API key,已自动生成并保存新密钥");
eprintln!("检测到默认 API key,已自动生成并保存新密钥");
}
if !is_loopback_host(&config.server.host) {
tracing::error!("当前版本仅支持本地监听,请使用 127.0.0.1/localhost/::1。");
eprintln!("当前版本仅支持本地监听,请使用 127.0.0.1/localhost/::1。");
return;
}
if config.server.api_key == config::DEFAULT_API_KEY {
tracing::error!("检测到使用默认 API key,已中止启动。请配置强密钥。");
eprintln!("检测到使用默认 API key,已中止启动。请配置强密钥。");
return;
}
if config.server.tls.enable {
tracing::error!("检测到 TLS 配置已启用,但当前版本尚未支持 TLS,已中止启动。");
eprintln!("检测到 TLS 配置已启用,但当前版本尚未支持 TLS,已中止启动。");
return;
}
if config.remote_management.allow_remote {
tracing::error!("检测到远程管理已开启,但当前版本未启用 TLS,已中止启动。");
eprintln!("检测到远程管理已开启,但当前版本未启用 TLS,已中止启动。");
return;
}
let state: AppState = Arc::new(RwLock::new(server::ServerState::new(config.clone())));
let logs: LogState = Arc::new(RwLock::new(logger::LogStore::with_config(&config.logging)));
// Initialize database for Switch functionality
let db = database::init_database().expect("Failed to initialize database");
let db = match database::init_database() {
Ok(conn) => conn,
Err(err) => {
tracing::error!("数据库初始化失败,已中止启动: {}", err);
eprintln!("数据库初始化失败,已中止启动: {}", err);
return;
}
};
// Initialize SkillService
let skill_service = SkillService::new().expect("Failed to initialize SkillService");
@@ -1386,8 +1441,14 @@ pub fn run() {
let shared_tokens = Arc::new(parking_lot::RwLock::new(
telemetry::TokenTracker::with_defaults(),
));
let log_rotation = telemetry::LogRotationConfig {
max_memory_logs: 10000,
retention_days: config.logging.retention_days,
max_file_size: 10 * 1024 * 1024,
enable_file_logging: config.logging.enabled,
};
let shared_logger = Arc::new(
telemetry::RequestLogger::with_defaults().expect("Failed to create RequestLogger"),
telemetry::RequestLogger::new(log_rotation).expect("Failed to create RequestLogger"),
);
// Initialize TelemetryState with shared instances
@@ -1749,3 +1810,13 @@ pub fn run() {
.run(tauri::generate_context!())
.expect("error while running tauri application");
}
fn is_loopback_host(host: &str) -> bool {
if host == "localhost" {
return true;
}
match host.parse::<std::net::IpAddr>() {
Ok(addr) => addr.is_loopback(),
Err(_) => false,
}
}
+94 -7
View File
@@ -1,5 +1,5 @@
//! 日志管理模块
use chrono::{Local, Utc};
use chrono::{Duration, Local, Utc};
use serde::{Deserialize, Serialize};
use std::collections::VecDeque;
use std::fs::{self, OpenOptions};
@@ -8,6 +8,25 @@ use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::RwLock;
#[derive(Debug, Clone)]
pub struct LogStoreConfig {
pub max_logs: usize,
pub retention_days: u32,
pub max_file_size: u64,
pub enable_file_logging: bool,
}
impl Default for LogStoreConfig {
fn default() -> Self {
Self {
max_logs: 1000,
retention_days: 7,
max_file_size: 10 * 1024 * 1024,
enable_file_logging: true,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LogEntry {
pub timestamp: String,
@@ -18,6 +37,7 @@ pub struct LogEntry {
pub struct LogStore {
logs: VecDeque<LogEntry>,
max_logs: usize,
config: LogStoreConfig,
log_file_path: Option<PathBuf>,
}
@@ -34,9 +54,12 @@ impl Default for LogStore {
let log_file = log_dir.join("proxycast.log");
let config = LogStoreConfig::default();
Self {
logs: VecDeque::new(),
max_logs: 1000,
max_logs: config.max_logs,
config,
log_file_path: Some(log_file),
}
}
@@ -47,6 +70,14 @@ impl LogStore {
Self::default()
}
pub fn with_config(logging: &crate::config::LoggingConfig) -> Self {
let mut store = Self::default();
store.config.retention_days = logging.retention_days;
store.config.enable_file_logging = logging.enabled;
store.max_logs = store.config.max_logs;
store
}
pub fn add(&mut self, level: &str, message: &str) {
let now = Utc::now();
let entry = LogEntry {
@@ -58,12 +89,16 @@ impl LogStore {
self.logs.push_back(entry.clone());
// 写入日志文件
if let Some(ref path) = self.log_file_path {
let local_time = Local::now().format("%Y-%m-%d %H:%M:%S%.3f");
let log_line = format!("{} [{}] {}\n", local_time, level.to_uppercase(), message);
if self.config.enable_file_logging {
if let Some(ref path) = self.log_file_path {
self.rotate_log_file_if_needed(path);
let local_time = Local::now().format("%Y-%m-%d %H:%M:%S%.3f");
let log_line = format!("{} [{}] {}\n", local_time, level.to_uppercase(), message);
if let Ok(mut file) = OpenOptions::new().create(true).append(true).open(path) {
let _ = file.write_all(log_line.as_bytes());
if let Ok(mut file) = OpenOptions::new().create(true).append(true).open(path) {
let _ = file.write_all(log_line.as_bytes());
}
self.prune_old_logs(path);
}
}
@@ -103,6 +138,58 @@ impl LogStore {
.as_ref()
.map(|p| p.to_string_lossy().to_string())
}
fn rotate_log_file_if_needed(&self, path: &PathBuf) {
let Ok(metadata) = fs::metadata(path) else {
return;
};
if metadata.len() <= self.config.max_file_size {
return;
}
let suffix = Local::now().format("%Y%m%d-%H%M%S");
let rotated = path.with_file_name(format!(
"{}.{}",
path.file_name().unwrap_or_default().to_string_lossy(),
suffix
));
let _ = fs::rename(path, &rotated);
self.prune_old_logs(path);
}
fn prune_old_logs(&self, path: &PathBuf) {
let Some(dir) = path.parent() else {
return;
};
let Ok(entries) = fs::read_dir(dir) else {
return;
};
let cutoff = Utc::now() - Duration::days(self.config.retention_days as i64);
let prefix = format!(
"{}.",
path.file_name().unwrap_or_default().to_string_lossy()
);
for entry in entries.flatten() {
let file_name = entry.file_name();
let file_name = file_name.to_string_lossy();
if !file_name.starts_with(&prefix) {
continue;
}
let Ok(metadata) = entry.metadata() else {
continue;
};
let Ok(modified) = metadata.modified() else {
continue;
};
let modified = chrono::DateTime::<Utc>::from(modified);
if modified < cutoff {
let _ = fs::remove_file(entry.path());
}
}
}
}
#[allow(dead_code)]
@@ -20,10 +20,34 @@ use futures::future::BoxFuture;
use std::{
net::{IpAddr, SocketAddr},
sync::Arc,
sync::Mutex,
task::{Context, Poll},
time::{Duration, Instant},
};
use tower::{Layer, Service};
const MAX_AUTH_FAILURES: u32 = 5;
const FAILURE_WINDOW_SECS: u64 = 60;
const BLOCK_SECS: u64 = 300;
struct FailureState {
count: u32,
window_start: Instant,
blocked_until: Option<Instant>,
}
fn failure_map() -> &'static Mutex<std::collections::HashMap<String, FailureState>> {
static FAILURES: std::sync::OnceLock<Mutex<std::collections::HashMap<String, FailureState>>> =
std::sync::OnceLock::new();
FAILURES.get_or_init(|| Mutex::new(std::collections::HashMap::new()))
}
#[cfg(test)]
pub(crate) fn clear_auth_failure_state() {
let mut map = failure_map().lock().unwrap();
map.clear();
}
/// Management API 认证层
///
/// 用于包装需要认证的管理端点
@@ -98,6 +122,66 @@ impl<S> ManagementAuthService<S> {
.get::<axum::extract::ConnectInfo<SocketAddr>>()
.map(|ci| ci.0)
}
fn get_client_id(req: &Request<Body>) -> String {
if let Some(addr) = Self::get_client_addr(req) {
return addr.ip().to_string();
}
if let Some(forwarded) = req.headers().get("x-forwarded-for") {
if let Ok(value) = forwarded.to_str() {
if let Some(first) = value.split(',').next() {
return first.trim().to_string();
}
}
}
"unknown".to_string()
}
fn check_rate_limit(client_id: &str) -> bool {
let now = Instant::now();
let mut map = failure_map().lock().unwrap();
if let Some(state) = map.get_mut(client_id) {
if let Some(blocked_until) = state.blocked_until {
if blocked_until > now {
return false;
}
state.blocked_until = None;
state.count = 0;
state.window_start = now;
}
if now.duration_since(state.window_start).as_secs() > FAILURE_WINDOW_SECS {
state.count = 0;
state.window_start = now;
}
}
true
}
fn record_failure(client_id: &str) {
let now = Instant::now();
let mut map = failure_map().lock().unwrap();
let entry = map.entry(client_id.to_string()).or_insert(FailureState {
count: 0,
window_start: now,
blocked_until: None,
});
if now.duration_since(entry.window_start).as_secs() > FAILURE_WINDOW_SECS {
entry.count = 0;
entry.window_start = now;
entry.blocked_until = None;
}
entry.count += 1;
if entry.count >= MAX_AUTH_FAILURES {
entry.blocked_until = Some(now + Duration::from_secs(BLOCK_SECS));
}
}
fn record_success(client_id: &str) {
let mut map = failure_map().lock().unwrap();
map.remove(client_id);
}
}
impl<S> Service<Request<Body>> for ManagementAuthService<S>
@@ -118,6 +202,14 @@ where
let mut inner = self.inner.clone();
Box::pin(async move {
let client_id = Self::get_client_id(&req);
if !Self::check_rate_limit(&client_id) {
return Ok(create_error_response(
StatusCode::TOO_MANY_REQUESTS,
"Too many failed authentication attempts",
));
}
// 1. 检查 secret_key 是否为空(禁用管理 API)
let secret_key = match &config.secret_key {
Some(key) if !key.is_empty() => key.clone(),
@@ -151,6 +243,7 @@ where
Some(key) if key == secret_key => {
// 认证成功,继续处理请求
tracing::debug!("[MANAGEMENT_AUTH] Auth successful from {:?}", client_addr);
Self::record_success(&client_id);
inner.call(req).await
}
Some(_) => {
@@ -158,6 +251,7 @@ where
"[MANAGEMENT_AUTH] Invalid secret_key from {:?}",
client_addr
);
Self::record_failure(&client_id);
Ok(create_error_response(
StatusCode::UNAUTHORIZED,
"Invalid secret key",
@@ -168,6 +262,7 @@ where
"[MANAGEMENT_AUTH] Missing secret_key from {:?}",
client_addr
);
Self::record_failure(&client_id);
Ok(create_error_response(
StatusCode::UNAUTHORIZED,
"Missing secret key",
+51 -1
View File
@@ -3,7 +3,9 @@
//! 使用 proptest 进行属性测试
use crate::config::RemoteManagementConfig;
use crate::middleware::management_auth::{ManagementAuthLayer, ManagementAuthService};
use crate::middleware::management_auth::{
clear_auth_failure_state, ManagementAuthLayer, ManagementAuthService,
};
use axum::{
body::Body,
http::{Request, Response, StatusCode},
@@ -102,6 +104,49 @@ fn create_request_with_management_key(key: Option<&str>) -> Request<Body> {
builder.body(Body::empty()).unwrap()
}
/// Helper to create a request with X-Management-Key and X-Forwarded-For headers
fn create_request_with_management_key_and_forwarded(
key: Option<&str>,
forwarded_for: Option<&str>,
) -> Request<Body> {
let mut builder = Request::builder().uri("/v0/management/status");
if let Some(k) = key {
builder = builder.header("x-management-key", k);
}
if let Some(addr) = forwarded_for {
builder = builder.header("x-forwarded-for", addr);
}
builder.body(Body::empty()).unwrap()
}
#[test]
fn test_management_auth_rate_limit_after_failures() {
clear_auth_failure_state();
let config = RemoteManagementConfig {
allow_remote: true,
secret_key: Some("valid_key".to_string()),
disable_control_panel: false,
};
let layer = ManagementAuthLayer::new(config);
let mut service = layer.layer(MockService);
let rt = tokio::runtime::Runtime::new().unwrap();
let client_ip = "203.0.113.10";
for _ in 0..5 {
let req =
create_request_with_management_key_and_forwarded(Some("invalid"), Some(client_ip));
let response = rt.block_on(async { service.call(req).await.unwrap() });
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
}
let req = create_request_with_management_key_and_forwarded(Some("invalid"), Some(client_ip));
let response = rt.block_on(async { service.call(req).await.unwrap() });
assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
}
proptest! {
#![proptest_config(ProptestConfig::with_cases(100))]
@@ -112,6 +157,7 @@ proptest! {
fn prop_management_auth_rejection_missing_key(
secret_key in arb_secret_key()
) {
clear_auth_failure_state();
// Create config with a valid secret_key
let config = RemoteManagementConfig {
allow_remote: true,
@@ -147,6 +193,7 @@ proptest! {
fn prop_management_auth_rejection_invalid_key(
secret_key in arb_secret_key()
) {
clear_auth_failure_state();
// Create config with a valid secret_key
let config = RemoteManagementConfig {
allow_remote: true,
@@ -183,6 +230,7 @@ proptest! {
fn prop_management_auth_acceptance_valid_key(
secret_key in arb_secret_key()
) {
clear_auth_failure_state();
// Create config with a valid secret_key
let config = RemoteManagementConfig {
allow_remote: true,
@@ -218,6 +266,7 @@ proptest! {
fn prop_management_auth_acceptance_x_management_key(
secret_key in arb_secret_key()
) {
clear_auth_failure_state();
// Create config with a valid secret_key
let config = RemoteManagementConfig {
allow_remote: true,
@@ -253,6 +302,7 @@ proptest! {
fn prop_management_auth_rejection_invalid_x_management_key(
secret_key in arb_secret_key()
) {
clear_auth_failure_state();
// Create config with a valid secret_key
let config = RemoteManagementConfig {
allow_remote: true,
+189 -21
View File
@@ -268,6 +268,28 @@ impl ServerState {
let api_key_for_state = api_key.clone(); // 用于保存到 running_api_key
let default_provider_ref = self.default_provider_ref.clone();
if api_key.trim().is_empty() {
return Err("API Key 不能为空".into());
}
if !is_localhost_host(&host) {
return Err("当前版本仅支持本地监听,请使用 127.0.0.1/localhost/::1".into());
}
if (!is_localhost_host(&host) || self.config.remote_management.allow_remote)
&& crate::config::is_default_api_key(&api_key)
{
return Err("非本地访问场景下禁止使用默认 API Key,请设置强口令".into());
}
if self.config.server.tls.enable {
return Err("当前版本暂不支持 TLS,请关闭 TLS 配置".into());
}
if self.config.remote_management.allow_remote {
return Err("当前版本未启用 TLS,禁止开启远程管理".into());
}
// 重新加载凭证
let _ = self.kiro_provider.load_credentials().await;
let kiro = self.kiro_provider.clone();
@@ -330,6 +352,15 @@ impl ServerState {
}
}
fn is_localhost_host(host: &str) -> bool {
if host == "localhost" {
return true;
}
host.parse::<std::net::IpAddr>()
.map(|addr| addr.is_loopback())
.unwrap_or(false)
}
impl Clone for KiroProvider {
fn clone(&self) -> Self {
Self {
@@ -346,6 +377,9 @@ struct AppState {
api_key: String,
base_url: String,
default_provider: Arc<RwLock<String>>,
config: Option<Config>,
config_manager: Option<Arc<std::sync::RwLock<ConfigManager>>>,
start_time: std::time::Instant,
kiro: Arc<RwLock<KiroProvider>>,
logs: Arc<RwLock<LogStore>>,
kiro_refresh_lock: Arc<tokio::sync::Mutex<()>>,
@@ -704,6 +738,9 @@ async fn run_server(
api_key: api_key.to_string(),
base_url,
default_provider,
config: config.clone(),
config_manager: config_manager.clone(),
start_time: std::time::Instant::now(),
kiro: Arc::new(RwLock::new(kiro)),
logs,
kiro_refresh_lock: Arc::new(tokio::sync::Mutex::new(())),
@@ -809,19 +846,54 @@ async fn run_server(
tracing::info!("Server listening on {}", addr);
axum::serve(listener, app)
.with_graceful_shutdown(async move {
let _ = shutdown.await;
})
.await?;
axum::serve(
listener,
app.into_make_service_with_connect_info::<std::net::SocketAddr>(),
)
.with_graceful_shutdown(async move {
let _ = shutdown.await;
})
.await?;
Ok(())
}
async fn health() -> impl IntoResponse {
async fn health(State(state): State<AppState>) -> impl IntoResponse {
let (db_ready, pool_stats) = match &state.db {
Some(db) => {
let db_ok = db
.lock()
.map(|conn| conn.query_row::<i32, _, _>("SELECT 1", [], |row| row.get(0)))
.is_ok();
let stats = if db_ok {
state.pool_service.get_overview(db).ok().map(|items| {
let mut total = 0usize;
let mut healthy = 0usize;
let mut disabled = 0usize;
for item in items {
total += item.stats.total_count;
healthy += item.stats.healthy_count;
disabled += item.stats.disabled_count;
}
serde_json::json!({
"total": total,
"healthy": healthy,
"disabled": disabled
})
})
} else {
None
};
(db_ok, stats)
}
None => (false, None),
};
Json(serde_json::json!({
"status": "healthy",
"version": "0.10.1"
"version": env!("CARGO_PKG_VERSION"),
"db_ready": db_ready,
"pool": pool_stats
}))
}
@@ -4885,21 +4957,32 @@ pub struct UpdateConfigResponse {
pub message: String,
}
fn snapshot_config(state: &AppState) -> Option<Config> {
if let Some(manager) = &state.config_manager {
if let Ok(guard) = manager.read() {
return Some(guard.config().clone());
}
}
state.config.clone()
}
/// GET /v0/management/status - 获取服务器状态
pub async fn management_status(State(state): State<AppState>) -> impl IntoResponse {
let default_provider = state.default_provider.read().await.clone();
// 获取请求数量
let requests = state.processor.stats.read().len() as u64;
let config = snapshot_config(&state).unwrap_or_default();
let uptime_secs = state.start_time.elapsed().as_secs();
let response = ManagementStatusResponse {
running: true,
host: "0.0.0.0".to_string(),
port: 8999,
host: config.server.host,
port: config.server.port,
requests,
uptime_secs: 0, // TODO: Track actual uptime
uptime_secs,
version: env!("CARGO_PKG_VERSION").to_string(),
tls_enabled: false,
tls_enabled: config.server.tls.enable,
default_provider,
};
@@ -5221,26 +5304,32 @@ pub async fn management_get_config(State(state): State<AppState>) -> impl IntoRe
// 获取路由规则数量
let rules_count = state.processor.router.read().await.rules().len();
let config = snapshot_config(&state).unwrap_or_default();
let response = ManagementConfigResponse {
server: ManagementServerConfigInfo {
host: "0.0.0.0".to_string(),
port: 8999,
tls_enabled: false,
host: config.server.host,
port: config.server.port,
tls_enabled: config.server.tls.enable,
},
routing: ManagementRoutingConfigInfo {
default_provider,
rules_count,
},
retry: ManagementRetryConfigInfo {
max_retries: 3,
base_delay_ms: 1000,
max_delay_ms: 30000,
max_retries: config.retry.max_retries,
base_delay_ms: config.retry.base_delay_ms,
max_delay_ms: config.retry.max_delay_ms,
},
remote_management: ManagementRemoteInfo {
allow_remote: false,
has_secret_key: true,
disable_control_panel: false,
allow_remote: config.remote_management.allow_remote,
has_secret_key: config
.remote_management
.secret_key
.as_ref()
.map(|key| !key.is_empty())
.unwrap_or(false),
disable_control_panel: config.remote_management.disable_control_panel,
},
};
@@ -5253,6 +5342,7 @@ pub async fn management_update_config(
Json(request): Json<UpdateConfigRequest>,
) -> impl IntoResponse {
let mut updated = false;
let mut needs_restart = false;
// 更新默认 Provider
if let Some(provider) = request.default_provider {
@@ -5261,6 +5351,29 @@ pub async fn management_update_config(
let mut dp = state.default_provider.write().await;
*dp = provider.clone();
tracing::info!("[MANAGEMENT] Updated default_provider to: {}", provider);
if let Some(manager) = &state.config_manager {
if let Ok(mut guard) = manager.write() {
guard.config_mut().default_provider = provider.clone();
guard.config_mut().routing.default_provider = provider.clone();
if let Err(err) = guard.save() {
return (
StatusCode::INTERNAL_SERVER_ERROR,
Json(UpdateConfigResponse {
success: false,
message: format!("Failed to save config: {}", err),
}),
);
}
} else {
return (
StatusCode::INTERNAL_SERVER_ERROR,
Json(UpdateConfigResponse {
success: false,
message: "Failed to lock config manager".to_string(),
}),
);
}
}
updated = true;
} else {
return (
@@ -5273,12 +5386,67 @@ pub async fn management_update_config(
}
}
// 更新是否允许远程访问(需要重启生效)
if let Some(allow_remote) = request.allow_remote {
if allow_remote {
return (
StatusCode::BAD_REQUEST,
Json(UpdateConfigResponse {
success: false,
message: "当前版本未启用 TLS,禁止开启远程管理".to_string(),
}),
);
}
let manager = match &state.config_manager {
Some(manager) => manager,
None => {
return (
StatusCode::INTERNAL_SERVER_ERROR,
Json(UpdateConfigResponse {
success: false,
message: "Config manager is not available".to_string(),
}),
);
}
};
if let Ok(mut guard) = manager.write() {
guard.config_mut().remote_management.allow_remote = allow_remote;
if let Err(err) = guard.save() {
return (
StatusCode::INTERNAL_SERVER_ERROR,
Json(UpdateConfigResponse {
success: false,
message: format!("Failed to save config: {}", err),
}),
);
}
} else {
return (
StatusCode::INTERNAL_SERVER_ERROR,
Json(UpdateConfigResponse {
success: false,
message: "Failed to lock config manager".to_string(),
}),
);
}
tracing::info!(
"[MANAGEMENT] Updated remote_management.allow_remote to: {}",
allow_remote
);
updated = true;
needs_restart = true;
}
if updated {
(
StatusCode::OK,
Json(UpdateConfigResponse {
success: true,
message: "Configuration updated successfully".to_string(),
message: if needs_restart {
"Configuration updated. Restart required to apply all changes.".to_string()
} else {
"Configuration updated successfully".to_string()
},
}),
)
} else {
+1 -1
View File
@@ -21,7 +21,7 @@
}
],
"security": {
"csp": null
"csp": "default-src 'self'; img-src 'self' data:; style-src 'self' 'unsafe-inline'; script-src 'self' 'unsafe-inline'; connect-src 'self' http://localhost:* ws://localhost:* http://127.0.0.1:* ws://127.0.0.1:* http://[::1]:* ws://[::1]:*; frame-src 'none'; object-src 'none'"
}
},
"bundle": {