fix: persist runtime proxy token

This commit is contained in:
boos
2026-06-06 17:11:51 +08:00
parent b3eb86d5ea
commit f026582008
17 changed files with 276 additions and 27 deletions
+8 -1
View File
@@ -19,6 +19,7 @@ pub struct AppState {
pub db: Arc<SqlitePool>,
pub comfyui: Arc<services::comfyui_process::ComfyUIProcess>,
pub builtin_keys: Arc<services::builtin_keys::BuiltinKeys>,
pub runtime_auth: Arc<services::runtime_auth::RuntimeAuthSettings>,
pub register_captcha: Arc<services::register_captcha::RegisterCaptchaStore>,
pub registration_limiter: Arc<services::registration_limiter::RegistrationLimiter>,
}
@@ -28,6 +29,7 @@ impl AppState {
config: config::AppConfig,
db: SqlitePool,
builtin_keys: services::builtin_keys::BuiltinKeys,
runtime_auth: services::runtime_auth::RuntimeAuthSettings,
) -> Self {
let boosapi_host = if config.server.host == "0.0.0.0" || config.server.host == "::" {
"127.0.0.1".to_string()
@@ -39,7 +41,7 @@ impl AppState {
let comfyui = services::comfyui_process::ComfyUIProcess::new(
config.comfyui.clone(),
boosapi_url.clone(),
config.auth.proxy_token.clone(),
runtime_auth.proxy_token(),
comfyui_agent_secret.clone(),
);
@@ -48,10 +50,15 @@ impl AppState {
db: Arc::new(db),
comfyui,
builtin_keys: Arc::new(builtin_keys),
runtime_auth: Arc::new(runtime_auth),
register_captcha: Arc::new(services::register_captcha::RegisterCaptchaStore::new()),
registration_limiter: Arc::new(
services::registration_limiter::RegistrationLimiter::new(),
),
}
}
pub fn runtime_proxy_token(&self) -> String {
self.runtime_auth.proxy_token()
}
}
+13 -2
View File
@@ -1,6 +1,11 @@
use anyhow::Context;
use axum::serve;
use boosapi::{config::AppConfig, db, services::builtin_keys, AppState};
use boosapi::{
config::AppConfig,
db,
services::{builtin_keys, runtime_auth},
AppState,
};
use clap::{Parser, Subcommand};
use std::ffi::OsString;
use std::net::SocketAddr;
@@ -139,9 +144,15 @@ async fn main() -> anyhow::Result<()> {
let builtin = builtin_keys::BuiltinKeys {
comfyui_agent: comfyui_agent_key,
};
let runtime_auth = runtime_auth::RuntimeAuthSettings::load(
&pool,
config.auth.proxy_token.clone(),
)
.await
.context("failed to load runtime auth settings")?;
// Build app state
let state = AppState::new(config.clone(), pool, builtin);
let state = AppState::new(config.clone(), pool, builtin, runtime_auth);
if let Err(err) = state.comfyui.start().await {
tracing::warn!(%err, "ComfyUI embedded process did not start");
}
+115 -2
View File
@@ -966,6 +966,8 @@ async fn change_auth_token() -> Json<Value> {
async fn runtime_settings(State(state): State<AppState>) -> Json<Value> {
Json(json!({
"proxyTokenMasked": state.runtime_auth.proxy_token_masked(),
"proxyTokenSource": if state.runtime_auth.using_default_proxy_token() { "config" } else { "runtime" },
"server": {
"host": state.config.server.host,
"port": state.config.server.port,
@@ -989,8 +991,34 @@ async fn runtime_settings(State(state): State<AppState>) -> Json<Value> {
}))
}
async fn update_runtime_settings() -> Json<Value> {
Json(json!({ "success": true }))
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
struct RuntimeSettingsUpdateRequest {
proxy_token: Option<String>,
}
async fn update_runtime_settings(
State(state): State<AppState>,
Json(req): Json<RuntimeSettingsUpdateRequest>,
) -> AppResult<Json<Value>> {
if let Some(proxy_token) = req.proxy_token {
let updated = state
.runtime_auth
.set_proxy_token(&state.db, proxy_token)
.await
.map_err(|err| AppError::BadRequest(err.to_string()))?;
return Ok(Json(json!({
"success": true,
"proxyTokenMasked": crate::services::runtime_auth::mask_token(&updated),
"proxyTokenSource": if state.runtime_auth.using_default_proxy_token() { "config" } else { "runtime" },
})));
}
Ok(Json(json!({
"success": true,
"proxyTokenMasked": state.runtime_auth.proxy_token_masked(),
"proxyTokenSource": if state.runtime_auth.using_default_proxy_token() { "config" } else { "runtime" },
})))
}
async fn brand_list() -> Json<Value> {
@@ -1166,3 +1194,88 @@ async fn clear_usage_data(State(state): State<AppState>) -> AppResult<Json<Value
async fn test_notification() -> Json<Value> {
Json(json!({ "success": true, "message": "notification endpoint available" }))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{
config::{AppConfig, AuthConfig, DbConfig, ProxyConfig, ServerConfig},
services::{builtin_keys::BuiltinKeys, runtime_auth::RuntimeAuthSettings},
};
use sqlx::sqlite::SqlitePoolOptions;
async fn test_state() -> AppState {
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.unwrap();
sqlx::query("CREATE TABLE settings (key TEXT PRIMARY KEY NOT NULL, value TEXT)")
.execute(&pool)
.await
.unwrap();
let runtime_auth = RuntimeAuthSettings::load(&pool, "sk-default-token".to_string())
.await
.unwrap();
AppState::new(
AppConfig {
server: ServerConfig {
host: "127.0.0.1".into(),
port: 0,
log_level: "info".into(),
cors_origins: vec![],
},
db: DbConfig {
url: "sqlite::memory:".into(),
max_connections: 1,
ssl: false,
runtime_db_url: None,
runtime_db_dialect: None,
runtime_db_ssl: None,
},
auth: AuthConfig {
jwt_secret: "secret".into(),
jwt_expiry_hours: 1,
auth_token: "auth".into(),
proxy_token: "sk-default-token".into(),
},
proxy: ProxyConfig {
upstream_timeout_secs: 30,
first_byte_timeout_secs: 30,
max_retries: 3,
proxy_log_retention_days: 7,
template_proxy_url: None,
},
comfyui: Default::default(),
},
pool,
BuiltinKeys {
comfyui_agent: (1, "sk-comfyui-test".into()),
},
runtime_auth,
)
}
#[tokio::test]
async fn runtime_settings_reports_and_updates_proxy_token() {
let state = test_state().await;
let Json(initial) = runtime_settings(State(state.clone())).await;
assert_eq!(initial["proxyTokenMasked"], "••••oken");
assert_eq!(initial["proxyTokenSource"], "config");
let Json(updated) = update_runtime_settings(
State(state.clone()),
Json(RuntimeSettingsUpdateRequest {
proxy_token: Some("sk-new-runtime-token".into()),
}),
)
.await
.unwrap();
assert_eq!(updated["success"], true);
assert_eq!(updated["proxyTokenMasked"], "••••oken");
assert_eq!(updated["proxyTokenSource"], "runtime");
assert_eq!(state.runtime_proxy_token(), "sk-new-runtime-token");
}
}
+1
View File
@@ -417,6 +417,7 @@ mod tests {
BuiltinKeys {
comfyui_agent: (1, "sk-comfyui-abcdef0123456789".into()),
},
crate::services::runtime_auth::RuntimeAuthSettings::from_token_for_test("proxy"),
)
}
+14 -2
View File
@@ -217,6 +217,7 @@ const HIDDEN_SETTINGS_PREFIX: &str = "builtin_key.";
fn is_hidden_setting_key(key: &str) -> bool {
key.starts_with(HIDDEN_SETTINGS_PREFIX)
|| key == crate::services::runtime_auth::PROXY_TOKEN_SETTING_KEY
}
/// GET /api/settings — return all settings as a key-value map.
@@ -313,6 +314,12 @@ mod tests {
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO settings (key, value) VALUES ('runtime.auth.proxy_token', 'sk-runtime-secret')",
)
.execute(&pool)
.await
.unwrap();
AppState::new(
AppConfig {
@@ -349,11 +356,12 @@ mod tests {
BuiltinKeys {
comfyui_agent: (1, "sk-comfyui-secret".into()),
},
crate::services::runtime_auth::RuntimeAuthSettings::from_token_for_test("proxy"),
)
}
/// 5.2 — `GET /api/settings` returns regular settings but never any
/// `builtin_key.*` rows.
/// 5.2 — `GET /api/settings` returns regular settings but never protected
/// secret rows.
#[tokio::test]
async fn list_settings_filters_builtin_key_prefix() {
let state = settings_state().await;
@@ -364,6 +372,10 @@ mod tests {
!data.contains_key("builtin_key.comfyui_agent.secret"),
"builtin_key.* rows must not appear in /api/settings response"
);
assert!(
!data.contains_key("runtime.auth.proxy_token"),
"runtime proxy token must not appear in /api/settings response"
);
}
/// 5.3 — `PUT /api/settings` rejects writes targeting `builtin_key.*`
+2 -1
View File
@@ -49,8 +49,9 @@ async fn check_auth(state: &AppState, headers: &HeaderMap) -> Result<(), AppErro
// Downstream API key
if auth.starts_with("sk-") {
let proxy_token = state.runtime_proxy_token();
let result =
authorize_downstream_token(auth, &state.db, Some(&state.config.auth.proxy_token))
authorize_downstream_token(auth, &state.db, Some(&proxy_token))
.await
.map_err(AppError::Unauthorized)?;
+5 -3
View File
@@ -243,11 +243,12 @@ async fn resolve_agent_auth(
)
})?;
if token == state.config.auth.auth_token || token == state.config.auth.proxy_token {
let proxy_token = state.runtime_proxy_token();
if token == state.config.auth.auth_token || token == proxy_token {
return Ok(AgentAuthContext {
user_id: None,
role: "admin".to_string(),
outbound_key: state.config.auth.proxy_token.clone(),
outbound_key: proxy_token,
key_name: "global-proxy".to_string(),
});
}
@@ -282,7 +283,7 @@ async fn resolve_agent_auth(
Err(_) if claims.role == "admin" => Ok(AgentAuthContext {
user_id: Some(user_id),
role: claims.role,
outbound_key: state.config.auth.proxy_token.clone(),
outbound_key: state.runtime_proxy_token(),
key_name: "global-proxy".to_string(),
}),
Err(message) => Err((
@@ -1245,6 +1246,7 @@ mod tests {
BuiltinKeys {
comfyui_agent: (1, "sk-comfyui-test".into()),
},
crate::services::runtime_auth::RuntimeAuthSettings::from_token_for_test("proxy"),
)
}
+4 -3
View File
@@ -104,11 +104,12 @@ async fn resolve_auth(headers: &HeaderMap, state: &AppState) -> Result<DigitalHu
let token = extract_bearer_token(headers)
.ok_or_else(|| AppError::Unauthorized("Missing Authorization header".into()))?;
if token == state.config.auth.auth_token || token == state.config.auth.proxy_token {
let proxy_token = state.runtime_proxy_token();
if token == state.config.auth.auth_token || token == proxy_token {
return Ok(DigitalHumanAuth {
user_id: None,
role: "admin".to_string(),
outbound_key: state.config.auth.proxy_token.clone(),
outbound_key: proxy_token,
});
}
@@ -131,7 +132,7 @@ async fn resolve_auth(headers: &HeaderMap, state: &AppState) -> Result<DigitalHu
Err(_) if claims.role == "admin" => Ok(DigitalHumanAuth {
user_id: Some(user_id),
role: claims.role,
outbound_key: state.config.auth.proxy_token.clone(),
outbound_key: state.runtime_proxy_token(),
}),
Err(message) => Err(AppError::Forbidden(message)),
}
+4 -3
View File
@@ -88,11 +88,12 @@ async fn resolve_auth(headers: &HeaderMap, state: &AppState) -> Result<MoneyPrin
let token = extract_bearer_token(headers)
.ok_or_else(|| AppError::Unauthorized("Missing Authorization header".into()))?;
if token == state.config.auth.auth_token || token == state.config.auth.proxy_token {
let proxy_token = state.runtime_proxy_token();
if token == state.config.auth.auth_token || token == proxy_token {
return Ok(MoneyPrinterAuth {
user_id: None,
role: "admin".to_string(),
outbound_key: state.config.auth.proxy_token.clone(),
outbound_key: proxy_token,
key_name: "global-proxy".to_string(),
});
}
@@ -117,7 +118,7 @@ async fn resolve_auth(headers: &HeaderMap, state: &AppState) -> Result<MoneyPrin
Err(_) if claims.role == "admin" => Ok(MoneyPrinterAuth {
user_id: Some(user_id),
role: claims.role,
outbound_key: state.config.auth.proxy_token.clone(),
outbound_key: state.runtime_proxy_token(),
key_name: "global-proxy".to_string(),
}),
Err(message) => Err(AppError::Forbidden(message)),
+4 -3
View File
@@ -123,11 +123,12 @@ async fn resolve_picture_agent_auth(
let token = extract_bearer_token(headers)
.ok_or_else(|| AppError::Unauthorized("Missing Authorization header".into()))?;
if token == state.config.auth.auth_token || token == state.config.auth.proxy_token {
let proxy_token = state.runtime_proxy_token();
if token == state.config.auth.auth_token || token == proxy_token {
return Ok(PictureAgentAuth {
user_id: None,
role: "admin".to_string(),
outbound_key: state.config.auth.proxy_token.clone(),
outbound_key: proxy_token,
key_name: "global-proxy".to_string(),
});
}
@@ -153,7 +154,7 @@ async fn resolve_picture_agent_auth(
Err(_) if claims.role == "admin" => Ok(PictureAgentAuth {
user_id: Some(user_id),
role: claims.role,
outbound_key: state.config.auth.proxy_token.clone(),
outbound_key: state.runtime_proxy_token(),
key_name: "global-proxy".to_string(),
}),
Err(message) => Err(AppError::Forbidden(message)),
+3 -1
View File
@@ -610,7 +610,8 @@ async fn check_auth(state: &AppState, headers: &HeaderMap) -> Result<Option<i64>
.and_then(|v| v.split_whitespace().next())
.ok_or_else(|| AppError::Unauthorized("Missing Authorization header".into()))?;
let result = authorize_downstream_token(auth, &state.db, Some(&state.config.auth.proxy_token))
let proxy_token = state.runtime_proxy_token();
let result = authorize_downstream_token(auth, &state.db, Some(&proxy_token))
.await
.map_err(AppError::Unauthorized)?;
@@ -1693,6 +1694,7 @@ mod tests {
crate::services::builtin_keys::BuiltinKeys {
comfyui_agent: (0, "sk-comfyui-test".to_string()),
},
crate::services::runtime_auth::RuntimeAuthSettings::from_token_for_test("proxy"),
)
}
+3 -3
View File
@@ -273,7 +273,7 @@ async fn call_llm(state: &AppState, prompt: &str, system: &str) -> AppResult<Str
let resp = client
.post(url)
.bearer_auth(&state.config.auth.proxy_token)
.bearer_auth(state.runtime_proxy_token())
.json(&json!({
"model": &state.config.comfyui.agent_model,
"messages": [
@@ -319,7 +319,7 @@ async fn call_image_gen(state: &AppState, prompt: &str) -> AppResult<String> {
let resp = client
.post(url)
.bearer_auth(&state.config.auth.proxy_token)
.bearer_auth(state.runtime_proxy_token())
.json(&json!({
"model": "gpt-image-2",
"prompt": prompt,
@@ -362,7 +362,7 @@ async fn call_tts(state: &AppState, text: &str, voice: &str) -> AppResult<Vec<u8
let resp = client
.post(url)
.bearer_auth(&state.config.auth.proxy_token)
.bearer_auth(state.runtime_proxy_token())
.json(&json!({
"model": "tts-1",
"input": text,
+2 -1
View File
@@ -307,7 +307,8 @@ fn caller_from_headers(state: &AppState, headers: &HeaderMap) -> AppResult<Workf
.filter(|value| !value.is_empty())
.ok_or_else(|| AppError::Unauthorized("Missing authorization token".into()))?;
if token == state.config.auth.auth_token || token == state.config.auth.proxy_token {
let proxy_token = state.runtime_proxy_token();
if token == state.config.auth.auth_token || token == proxy_token {
return Ok(WorkflowCaller::Admin);
}
+2 -1
View File
@@ -91,7 +91,8 @@ fn resolve_caller(
let token = bearer(headers)
.ok_or_else(|| AppError::Unauthorized("Missing Authorization header".into()))?;
if token == state.config.auth.auth_token || token == state.config.auth.proxy_token {
let proxy_token = state.runtime_proxy_token();
if token == state.config.auth.auth_token || token == proxy_token {
return Ok(WorkspaceCaller::Admin);
}
+1
View File
@@ -8,6 +8,7 @@ pub mod oauth;
pub mod register_captcha;
pub mod registration_limiter;
pub mod route_rebuild;
pub mod runtime_auth;
pub mod runtime_logs;
pub mod shared;
pub mod workspace;
+93
View File
@@ -0,0 +1,93 @@
use sqlx::SqlitePool;
use parking_lot::RwLock;
pub const PROXY_TOKEN_SETTING_KEY: &str = "runtime.auth.proxy_token";
#[derive(Debug)]
pub struct RuntimeAuthSettings {
default_proxy_token: String,
proxy_token: RwLock<String>,
}
impl RuntimeAuthSettings {
#[cfg(test)]
pub fn from_token_for_test(proxy_token: impl Into<String>) -> Self {
let proxy_token = proxy_token.into();
Self {
default_proxy_token: proxy_token.clone(),
proxy_token: RwLock::new(proxy_token),
}
}
pub async fn load(pool: &SqlitePool, default_proxy_token: String) -> anyhow::Result<Self> {
let stored: Option<String> =
sqlx::query_scalar("SELECT value FROM settings WHERE key = ?")
.bind(PROXY_TOKEN_SETTING_KEY)
.fetch_optional(pool)
.await?;
let proxy_token = stored
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(default_proxy_token.trim())
.to_string();
Ok(Self {
default_proxy_token,
proxy_token: RwLock::new(proxy_token),
})
}
pub fn proxy_token(&self) -> String {
self.proxy_token.read().clone()
}
pub async fn set_proxy_token(
&self,
pool: &SqlitePool,
proxy_token: String,
) -> anyhow::Result<String> {
let normalized = normalize_proxy_token(&proxy_token)?;
sqlx::query("INSERT OR REPLACE INTO settings (key, value) VALUES (?, ?)")
.bind(PROXY_TOKEN_SETTING_KEY)
.bind(&normalized)
.execute(pool)
.await?;
*self.proxy_token.write() = normalized.clone();
Ok(normalized)
}
pub fn proxy_token_masked(&self) -> String {
mask_token(&self.proxy_token())
}
pub fn using_default_proxy_token(&self) -> bool {
self.proxy_token() == self.default_proxy_token
}
}
pub fn normalize_proxy_token(value: &str) -> anyhow::Result<String> {
let token = value.trim();
if !token.starts_with("sk-") {
anyhow::bail!("下游访问令牌必须以 sk- 开头");
}
if token.chars().count() < 6 {
anyhow::bail!("下游访问令牌至少 6 位(含 sk-)");
}
Ok(token.to_string())
}
pub fn mask_token(value: &str) -> String {
let token = value.trim();
if token.is_empty() {
return String::new();
}
let suffix: String = token
.chars()
.rev()
.take(4)
.collect::<Vec<_>>()
.into_iter()
.rev()
.collect();
format!("••••{}", suffix)
}
+2 -1
View File
@@ -97,7 +97,8 @@ fn bearer(headers: &HeaderMap) -> Option<&str> {
fn resolve_caller(headers: &HeaderMap, state: &AppState) -> AppResult<WorkflowCaller> {
let token = bearer(headers)
.ok_or_else(|| AppError::Unauthorized("Missing Authorization header".into()))?;
if token == state.config.auth.auth_token || token == state.config.auth.proxy_token {
let proxy_token = state.runtime_proxy_token();
if token == state.config.auth.auth_token || token == proxy_token {
return Ok(WorkflowCaller::Admin);
}
let claims = validate_jwt(&state.config.auth.jwt_secret, token)