forked from boos/BoosAPI
feat: Rust refactor — full backend rewrite (Phases 1-6)
Phase 1-3: Scaffold + types + auth (JWT, Argon2, API keys) Phase 4a: Admin API (sites/accounts/tokens/downstream-keys/settings) Phase 5: Proxy core (/v1/chat, /v1/images, /v1/models) Phase 6: ComfyUI proxy routes migrated - 80+ source files, all cargo check + 18 tests green - Drop-in compatible with Node.js API surface
This commit is contained in:
@@ -0,0 +1,195 @@
|
||||
use crate::{
|
||||
db::models::account::Account,
|
||||
error::{AppError, AppResult},
|
||||
AppState,
|
||||
};
|
||||
use axum::{
|
||||
extract::{Path, Query, State},
|
||||
http::StatusCode,
|
||||
Json,
|
||||
};
|
||||
use serde::Deserialize;
|
||||
use serde_json::{json, Value};
|
||||
use tracing::info;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct ListAccountsQuery {
|
||||
pub site_id: Option<i64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct CreateAccountRequest {
|
||||
pub site_id: i64,
|
||||
pub username: Option<String>,
|
||||
pub access_token: String,
|
||||
pub api_token: Option<String>,
|
||||
pub status: Option<String>,
|
||||
pub checkin_enabled: Option<bool>,
|
||||
pub extra_config: Option<String>,
|
||||
pub oauth_provider: Option<String>,
|
||||
pub oauth_account_key: Option<String>,
|
||||
pub oauth_project_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct UpdateAccountRequest {
|
||||
pub username: Option<String>,
|
||||
pub access_token: Option<String>,
|
||||
pub api_token: Option<String>,
|
||||
pub status: Option<String>,
|
||||
pub checkin_enabled: Option<bool>,
|
||||
pub extra_config: Option<String>,
|
||||
pub oauth_provider: Option<String>,
|
||||
pub oauth_account_key: Option<String>,
|
||||
pub oauth_project_id: Option<String>,
|
||||
}
|
||||
|
||||
/// GET /api/accounts?site_id=
|
||||
pub async fn list_accounts(
|
||||
State(state): State<AppState>,
|
||||
Query(query): Query<ListAccountsQuery>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
info!("Listing accounts, site_id filter: {:?}", query.site_id);
|
||||
|
||||
let accounts: Vec<Account> = if let Some(site_id) = query.site_id {
|
||||
sqlx::query_as::<_, Account>(
|
||||
"SELECT * FROM accounts WHERE site_id = ? ORDER BY sort_order, id",
|
||||
)
|
||||
.bind(site_id)
|
||||
.fetch_all(&*state.db)
|
||||
.await?
|
||||
} else {
|
||||
sqlx::query_as::<_, Account>("SELECT * FROM accounts ORDER BY site_id, sort_order, id")
|
||||
.fetch_all(&*state.db)
|
||||
.await?
|
||||
};
|
||||
|
||||
Ok(Json(json!({ "data": accounts })))
|
||||
}
|
||||
|
||||
/// POST /api/accounts
|
||||
pub async fn create_account(
|
||||
State(state): State<AppState>,
|
||||
Json(payload): Json<CreateAccountRequest>,
|
||||
) -> AppResult<(StatusCode, Json<Value>)> {
|
||||
info!("Creating account for site_id: {}", payload.site_id);
|
||||
|
||||
if payload.access_token.trim().is_empty() {
|
||||
return Err(AppError::BadRequest("access_token is required".into()));
|
||||
}
|
||||
|
||||
let site_exists: Option<(i64,)> = sqlx::query_as("SELECT id FROM sites WHERE id = ?")
|
||||
.bind(payload.site_id)
|
||||
.fetch_optional(&*state.db)
|
||||
.await?;
|
||||
|
||||
if site_exists.is_none() {
|
||||
return Err(AppError::BadRequest(format!(
|
||||
"Site with id {} not found",
|
||||
payload.site_id
|
||||
)));
|
||||
}
|
||||
|
||||
let status = payload.status.unwrap_or_else(|| "active".to_string());
|
||||
|
||||
let result = sqlx::query(
|
||||
"INSERT INTO accounts (site_id, username, access_token, api_token, status, checkin_enabled, extra_config, oauth_provider, oauth_account_key, oauth_project_id) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"
|
||||
)
|
||||
.bind(payload.site_id)
|
||||
.bind(&payload.username)
|
||||
.bind(&payload.access_token)
|
||||
.bind(&payload.api_token)
|
||||
.bind(&status)
|
||||
.bind(payload.checkin_enabled)
|
||||
.bind(&payload.extra_config)
|
||||
.bind(&payload.oauth_provider)
|
||||
.bind(&payload.oauth_account_key)
|
||||
.bind(&payload.oauth_project_id)
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
|
||||
let account_id = result.last_insert_rowid();
|
||||
let account: Account = sqlx::query_as::<_, Account>("SELECT * FROM accounts WHERE id = ?")
|
||||
.bind(account_id)
|
||||
.fetch_one(&*state.db)
|
||||
.await?;
|
||||
|
||||
Ok((StatusCode::CREATED, Json(json!({ "data": account }))))
|
||||
}
|
||||
|
||||
/// GET /api/accounts/{id}
|
||||
pub async fn get_account(
|
||||
State(state): State<AppState>,
|
||||
Path(id): Path<i64>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
info!("Getting account: {}", id);
|
||||
|
||||
let account: Account = sqlx::query_as::<_, Account>("SELECT * FROM accounts WHERE id = ?")
|
||||
.bind(id)
|
||||
.fetch_optional(&*state.db)
|
||||
.await?
|
||||
.ok_or_else(|| AppError::NotFound(format!("Account with id {} not found", id)))?;
|
||||
|
||||
Ok(Json(json!({ "data": account })))
|
||||
}
|
||||
|
||||
/// PUT /api/accounts/{id}
|
||||
pub async fn update_account(
|
||||
State(state): State<AppState>,
|
||||
Path(id): Path<i64>,
|
||||
Json(payload): Json<UpdateAccountRequest>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
info!("Updating account: {}", id);
|
||||
|
||||
let existing: Account = sqlx::query_as::<_, Account>("SELECT * FROM accounts WHERE id = ?")
|
||||
.bind(id)
|
||||
.fetch_optional(&*state.db)
|
||||
.await?
|
||||
.ok_or_else(|| AppError::NotFound(format!("Account with id {} not found", id)))?;
|
||||
|
||||
sqlx::query(
|
||||
"UPDATE accounts SET username = ?, access_token = ?, api_token = ?, status = ?, checkin_enabled = ?, extra_config = ?, oauth_provider = ?, oauth_account_key = ?, oauth_project_id = ?, updated_at = datetime('now') WHERE id = ?"
|
||||
)
|
||||
.bind(payload.username.or(existing.username))
|
||||
.bind(payload.access_token.as_ref().unwrap_or(&existing.access_token))
|
||||
.bind(payload.api_token.or(existing.api_token))
|
||||
.bind(payload.status.unwrap_or(existing.status))
|
||||
.bind(payload.checkin_enabled.or(existing.checkin_enabled))
|
||||
.bind(payload.extra_config.or(existing.extra_config))
|
||||
.bind(payload.oauth_provider.or(existing.oauth_provider))
|
||||
.bind(payload.oauth_account_key.or(existing.oauth_account_key))
|
||||
.bind(payload.oauth_project_id.or(existing.oauth_project_id))
|
||||
.bind(id)
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
|
||||
let account: Account = sqlx::query_as::<_, Account>("SELECT * FROM accounts WHERE id = ?")
|
||||
.bind(id)
|
||||
.fetch_one(&*state.db)
|
||||
.await?;
|
||||
|
||||
Ok(Json(json!({ "data": account })))
|
||||
}
|
||||
|
||||
/// DELETE /api/accounts/{id}
|
||||
pub async fn delete_account(
|
||||
State(state): State<AppState>,
|
||||
Path(id): Path<i64>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
info!("Deleting account: {}", id);
|
||||
|
||||
let affected = sqlx::query("DELETE FROM accounts WHERE id = ?")
|
||||
.bind(id)
|
||||
.execute(&*state.db)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
if affected == 0 {
|
||||
return Err(AppError::NotFound(format!(
|
||||
"Account with id {} not found",
|
||||
id
|
||||
)));
|
||||
}
|
||||
|
||||
Ok(Json(json!({ "data": { "deleted": true, "id": id } })))
|
||||
}
|
||||
@@ -0,0 +1,196 @@
|
||||
use crate::{
|
||||
db::models::downstream_api_key::DownstreamApiKey,
|
||||
error::{AppError, AppResult},
|
||||
AppState,
|
||||
};
|
||||
use axum::{
|
||||
extract::{Path, State},
|
||||
http::StatusCode,
|
||||
Json,
|
||||
};
|
||||
use serde::Deserialize;
|
||||
use serde_json::{json, Value};
|
||||
use tracing::info;
|
||||
use uuid::Uuid;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct CreateDownstreamKeyRequest {
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
pub group_name: Option<String>,
|
||||
pub tags: Option<String>,
|
||||
pub enabled: Option<bool>,
|
||||
pub expires_at: Option<String>,
|
||||
pub max_cost: Option<f64>,
|
||||
pub max_requests: Option<i64>,
|
||||
pub supported_models: Option<Vec<String>>,
|
||||
pub allowed_route_ids: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct UpdateDownstreamKeyRequest {
|
||||
pub name: Option<String>,
|
||||
pub description: Option<String>,
|
||||
pub group_name: Option<String>,
|
||||
pub tags: Option<String>,
|
||||
pub enabled: Option<bool>,
|
||||
pub expires_at: Option<String>,
|
||||
pub max_cost: Option<f64>,
|
||||
pub max_requests: Option<i64>,
|
||||
pub supported_models: Option<Vec<String>>,
|
||||
pub allowed_route_ids: Option<String>,
|
||||
}
|
||||
|
||||
/// Helper to validate models from a JSON string against a Vec<String>
|
||||
fn validate_supported_models(models: &[String]) -> AppResult<()> {
|
||||
for model in models {
|
||||
if model.trim().is_empty() {
|
||||
return Err(AppError::BadRequest(
|
||||
"supported_models contains empty model name".into(),
|
||||
));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// GET /api/downstream-api-keys
|
||||
pub async fn list_downstream_keys(State(state): State<AppState>) -> AppResult<Json<Value>> {
|
||||
info!("Listing downstream API keys");
|
||||
|
||||
let keys: Vec<DownstreamApiKey> =
|
||||
sqlx::query_as::<_, DownstreamApiKey>("SELECT * FROM downstream_api_keys ORDER BY name")
|
||||
.fetch_all(&*state.db)
|
||||
.await?;
|
||||
|
||||
Ok(Json(json!({ "data": keys })))
|
||||
}
|
||||
|
||||
/// POST /api/downstream-api-keys
|
||||
pub async fn create_downstream_key(
|
||||
State(state): State<AppState>,
|
||||
Json(payload): Json<CreateDownstreamKeyRequest>,
|
||||
) -> AppResult<(StatusCode, Json<Value>)> {
|
||||
info!("Creating downstream API key: {}", payload.name);
|
||||
|
||||
if payload.name.trim().is_empty() {
|
||||
return Err(AppError::BadRequest("name is required".into()));
|
||||
}
|
||||
|
||||
// Auto-generate key
|
||||
let key = format!("sk-{}", Uuid::new_v4().to_string().replace('-', ""));
|
||||
|
||||
// Serialize supported_models to JSON string if provided
|
||||
let supported_models: Option<String> = if let Some(ref models) = payload.supported_models {
|
||||
validate_supported_models(models)?;
|
||||
Some(
|
||||
serde_json::to_string(models)
|
||||
.map_err(|e| AppError::BadRequest(format!("Invalid supported_models: {}", e)))?,
|
||||
)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let enabled = payload.enabled.unwrap_or(true);
|
||||
|
||||
let result = sqlx::query(
|
||||
"INSERT INTO downstream_api_keys (name, key, description, group_name, tags, enabled, expires_at, max_cost, max_requests, supported_models, allowed_route_ids) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"
|
||||
)
|
||||
.bind(&payload.name)
|
||||
.bind(&key)
|
||||
.bind(&payload.description)
|
||||
.bind(&payload.group_name)
|
||||
.bind(&payload.tags)
|
||||
.bind(enabled)
|
||||
.bind(&payload.expires_at)
|
||||
.bind(payload.max_cost)
|
||||
.bind(payload.max_requests)
|
||||
.bind(&supported_models)
|
||||
.bind(&payload.allowed_route_ids)
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
|
||||
let key_id = result.last_insert_rowid();
|
||||
let api_key: DownstreamApiKey =
|
||||
sqlx::query_as::<_, DownstreamApiKey>("SELECT * FROM downstream_api_keys WHERE id = ?")
|
||||
.bind(key_id)
|
||||
.fetch_one(&*state.db)
|
||||
.await?;
|
||||
|
||||
Ok((StatusCode::CREATED, Json(json!({ "data": api_key }))))
|
||||
}
|
||||
|
||||
/// PUT /api/downstream-api-keys/{id}
|
||||
pub async fn update_downstream_key(
|
||||
State(state): State<AppState>,
|
||||
Path(id): Path<i64>,
|
||||
Json(payload): Json<UpdateDownstreamKeyRequest>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
info!("Updating downstream API key: {}", id);
|
||||
|
||||
let existing: DownstreamApiKey =
|
||||
sqlx::query_as::<_, DownstreamApiKey>("SELECT * FROM downstream_api_keys WHERE id = ?")
|
||||
.bind(id)
|
||||
.fetch_optional(&*state.db)
|
||||
.await?
|
||||
.ok_or_else(|| {
|
||||
AppError::NotFound(format!("Downstream key with id {} not found", id))
|
||||
})?;
|
||||
|
||||
let supported_models: Option<String> = if let Some(ref models) = payload.supported_models {
|
||||
validate_supported_models(models)?;
|
||||
Some(
|
||||
serde_json::to_string(models)
|
||||
.map_err(|e| AppError::BadRequest(format!("Invalid supported_models: {}", e)))?,
|
||||
)
|
||||
} else {
|
||||
existing.supported_models.clone()
|
||||
};
|
||||
|
||||
sqlx::query(
|
||||
"UPDATE downstream_api_keys SET name = ?, description = ?, group_name = ?, tags = ?, enabled = ?, expires_at = ?, max_cost = ?, max_requests = ?, supported_models = ?, allowed_route_ids = ?, updated_at = datetime('now') WHERE id = ?"
|
||||
)
|
||||
.bind(payload.name.unwrap_or(existing.name))
|
||||
.bind(payload.description.or(existing.description))
|
||||
.bind(payload.group_name.or(existing.group_name))
|
||||
.bind(payload.tags.or(existing.tags))
|
||||
.bind(payload.enabled.unwrap_or(existing.enabled.unwrap_or(true)))
|
||||
.bind(payload.expires_at.or(existing.expires_at))
|
||||
.bind(payload.max_cost.or(existing.max_cost))
|
||||
.bind(payload.max_requests.or(existing.max_requests))
|
||||
.bind(&supported_models)
|
||||
.bind(payload.allowed_route_ids.or(existing.allowed_route_ids))
|
||||
.bind(id)
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
|
||||
let api_key: DownstreamApiKey =
|
||||
sqlx::query_as::<_, DownstreamApiKey>("SELECT * FROM downstream_api_keys WHERE id = ?")
|
||||
.bind(id)
|
||||
.fetch_one(&*state.db)
|
||||
.await?;
|
||||
|
||||
Ok(Json(json!({ "data": api_key })))
|
||||
}
|
||||
|
||||
/// DELETE /api/downstream-api-keys/{id}
|
||||
pub async fn delete_downstream_key(
|
||||
State(state): State<AppState>,
|
||||
Path(id): Path<i64>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
info!("Deleting downstream API key: {}", id);
|
||||
|
||||
let affected = sqlx::query("DELETE FROM downstream_api_keys WHERE id = ?")
|
||||
.bind(id)
|
||||
.execute(&*state.db)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
if affected == 0 {
|
||||
return Err(AppError::NotFound(format!(
|
||||
"Downstream key with id {} not found",
|
||||
id
|
||||
)));
|
||||
}
|
||||
|
||||
Ok(Json(json!({ "data": { "deleted": true, "id": id } })))
|
||||
}
|
||||
@@ -0,0 +1,158 @@
|
||||
use crate::{
|
||||
db::models::setting::Setting,
|
||||
error::{AppError, AppResult},
|
||||
AppState,
|
||||
};
|
||||
use axum::{
|
||||
extract::State,
|
||||
http::HeaderMap,
|
||||
middleware,
|
||||
response::Response,
|
||||
routing::{get, put},
|
||||
Json, Router,
|
||||
};
|
||||
use serde_json::{json, Value};
|
||||
use tracing::info;
|
||||
|
||||
mod accounts;
|
||||
mod downstream_keys;
|
||||
mod sites;
|
||||
mod tokens;
|
||||
|
||||
/// Middleware to require admin authentication via Bearer token.
|
||||
async fn require_admin_auth(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
request: axum::http::Request<axum::body::Body>,
|
||||
next: axum::middleware::Next,
|
||||
) -> Result<Response, AppError> {
|
||||
let auth_header = headers
|
||||
.get("Authorization")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|v| v.strip_prefix("Bearer "))
|
||||
.ok_or_else(|| AppError::Unauthorized("Missing or invalid Authorization header".into()))?;
|
||||
|
||||
if auth_header != state.config.auth.auth_token {
|
||||
return Err(AppError::Unauthorized("Invalid admin token".into()));
|
||||
}
|
||||
|
||||
Ok(next.run(request).await)
|
||||
}
|
||||
|
||||
/// Build all admin routes. Returns Router<AppState> (without calling .with_state).
|
||||
pub fn routes(state: AppState) -> Router<AppState> {
|
||||
let sites_routes = Router::new()
|
||||
.route(
|
||||
"/api/sites",
|
||||
get(sites::list_sites).post(sites::create_site),
|
||||
)
|
||||
.route(
|
||||
"/api/sites/{id}",
|
||||
get(sites::get_site)
|
||||
.put(sites::update_site)
|
||||
.delete(sites::delete_site),
|
||||
);
|
||||
|
||||
let accounts_routes = Router::new()
|
||||
.route(
|
||||
"/api/accounts",
|
||||
get(accounts::list_accounts).post(accounts::create_account),
|
||||
)
|
||||
.route(
|
||||
"/api/accounts/{id}",
|
||||
get(accounts::get_account)
|
||||
.put(accounts::update_account)
|
||||
.delete(accounts::delete_account),
|
||||
);
|
||||
|
||||
let tokens_routes = Router::new()
|
||||
.route(
|
||||
"/api/tokens",
|
||||
get(tokens::list_tokens).post(tokens::create_token),
|
||||
)
|
||||
.route(
|
||||
"/api/tokens/refresh-route",
|
||||
get(tokens::refresh_route).post(tokens::refresh_route),
|
||||
)
|
||||
.route("/api/tokens/route-decision", get(tokens::route_decision))
|
||||
.route(
|
||||
"/api/tokens/{id}",
|
||||
get(tokens::get_token)
|
||||
.put(tokens::update_token)
|
||||
.delete(tokens::delete_token),
|
||||
);
|
||||
|
||||
let downstream_routes = Router::new()
|
||||
.route(
|
||||
"/api/downstream-api-keys",
|
||||
get(downstream_keys::list_downstream_keys).post(downstream_keys::create_downstream_key),
|
||||
)
|
||||
.route(
|
||||
"/api/downstream-api-keys/{id}",
|
||||
put(downstream_keys::update_downstream_key)
|
||||
.delete(downstream_keys::delete_downstream_key),
|
||||
);
|
||||
|
||||
let settings_routes =
|
||||
Router::new().route("/api/settings", get(get_settings).put(update_setting));
|
||||
|
||||
Router::new()
|
||||
.merge(sites_routes)
|
||||
.merge(accounts_routes)
|
||||
.merge(tokens_routes)
|
||||
.merge(downstream_routes)
|
||||
.merge(settings_routes)
|
||||
.route_layer(middleware::from_fn_with_state(state, require_admin_auth))
|
||||
}
|
||||
|
||||
// ── Settings handlers ────────────────────────────────────────────────
|
||||
|
||||
/// GET /api/settings — return all settings as a key-value map.
|
||||
async fn get_settings(State(state): State<AppState>) -> AppResult<Json<Value>> {
|
||||
info!("Getting all settings");
|
||||
|
||||
let rows: Vec<Setting> =
|
||||
sqlx::query_as::<_, Setting>("SELECT key, value FROM settings ORDER BY key")
|
||||
.fetch_all(&*state.db)
|
||||
.await?;
|
||||
|
||||
let map: std::collections::BTreeMap<String, Option<String>> =
|
||||
rows.into_iter().map(|s| (s.key, s.value)).collect();
|
||||
|
||||
Ok(Json(json!({ "data": map })))
|
||||
}
|
||||
|
||||
/// PUT /api/settings — upsert a setting by key.
|
||||
/// Expects JSON body: { "key": "...", "value": "..." }
|
||||
async fn update_setting(
|
||||
State(state): State<AppState>,
|
||||
Json(payload): Json<serde_json::Value>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let key = payload
|
||||
.get("key")
|
||||
.and_then(|k| k.as_str())
|
||||
.ok_or_else(|| AppError::BadRequest("key is required".into()))?;
|
||||
|
||||
let value = payload.get("value").and_then(|v| {
|
||||
if v.is_null() {
|
||||
None
|
||||
} else {
|
||||
v.as_str().map(|s| s.to_string())
|
||||
}
|
||||
});
|
||||
|
||||
info!("Upserting setting: {}", key);
|
||||
|
||||
sqlx::query("INSERT OR REPLACE INTO settings (key, value) VALUES (?, ?)")
|
||||
.bind(key)
|
||||
.bind(&value)
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
|
||||
Ok(Json(json!({
|
||||
"data": {
|
||||
"key": key,
|
||||
"value": value
|
||||
}
|
||||
})))
|
||||
}
|
||||
@@ -0,0 +1,168 @@
|
||||
use crate::{
|
||||
db::models::site::Site,
|
||||
error::{AppError, AppResult},
|
||||
AppState,
|
||||
};
|
||||
use axum::{
|
||||
extract::{Path, State},
|
||||
http::StatusCode,
|
||||
Json,
|
||||
};
|
||||
use serde::Deserialize;
|
||||
use serde_json::{json, Value};
|
||||
use tracing::info;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct CreateSiteRequest {
|
||||
pub name: String,
|
||||
pub url: String,
|
||||
pub platform: String,
|
||||
pub status: Option<String>,
|
||||
pub api_key: Option<String>,
|
||||
pub proxy_url: Option<String>,
|
||||
pub use_system_proxy: Option<bool>,
|
||||
pub custom_headers: Option<String>,
|
||||
pub global_weight: Option<f64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct UpdateSiteRequest {
|
||||
pub name: Option<String>,
|
||||
pub url: Option<String>,
|
||||
pub platform: Option<String>,
|
||||
pub status: Option<String>,
|
||||
pub api_key: Option<String>,
|
||||
pub proxy_url: Option<String>,
|
||||
pub use_system_proxy: Option<bool>,
|
||||
pub custom_headers: Option<String>,
|
||||
pub global_weight: Option<f64>,
|
||||
}
|
||||
|
||||
/// GET /api/sites
|
||||
pub async fn list_sites(State(state): State<AppState>) -> AppResult<Json<Value>> {
|
||||
info!("Listing all sites");
|
||||
|
||||
let sites: Vec<Site> =
|
||||
sqlx::query_as::<_, Site>("SELECT * FROM sites ORDER BY sort_order, name")
|
||||
.fetch_all(&*state.db)
|
||||
.await?;
|
||||
|
||||
Ok(Json(json!({ "data": sites })))
|
||||
}
|
||||
|
||||
/// POST /api/sites
|
||||
pub async fn create_site(
|
||||
State(state): State<AppState>,
|
||||
Json(payload): Json<CreateSiteRequest>,
|
||||
) -> AppResult<(StatusCode, Json<Value>)> {
|
||||
info!("Creating site: {}", payload.name);
|
||||
|
||||
if payload.name.trim().is_empty() {
|
||||
return Err(AppError::BadRequest("name is required".into()));
|
||||
}
|
||||
if payload.url.trim().is_empty() {
|
||||
return Err(AppError::BadRequest("url is required".into()));
|
||||
}
|
||||
if payload.platform.trim().is_empty() {
|
||||
return Err(AppError::BadRequest("platform is required".into()));
|
||||
}
|
||||
|
||||
let status = payload.status.unwrap_or_else(|| "active".to_string());
|
||||
|
||||
let result = sqlx::query(
|
||||
"INSERT INTO sites (name, url, platform, status, api_key, proxy_url, use_system_proxy, custom_headers, global_weight) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)"
|
||||
)
|
||||
.bind(&payload.name)
|
||||
.bind(&payload.url)
|
||||
.bind(&payload.platform)
|
||||
.bind(&status)
|
||||
.bind(&payload.api_key)
|
||||
.bind(&payload.proxy_url)
|
||||
.bind(payload.use_system_proxy)
|
||||
.bind(&payload.custom_headers)
|
||||
.bind(payload.global_weight)
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
|
||||
let site_id = result.last_insert_rowid();
|
||||
let site: Site = sqlx::query_as::<_, Site>("SELECT * FROM sites WHERE id = ?")
|
||||
.bind(site_id)
|
||||
.fetch_one(&*state.db)
|
||||
.await?;
|
||||
|
||||
Ok((StatusCode::CREATED, Json(json!({ "data": site }))))
|
||||
}
|
||||
|
||||
/// GET /api/sites/{id}
|
||||
pub async fn get_site(
|
||||
State(state): State<AppState>,
|
||||
Path(id): Path<i64>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
info!("Getting site: {}", id);
|
||||
|
||||
let site: Site = sqlx::query_as::<_, Site>("SELECT * FROM sites WHERE id = ?")
|
||||
.bind(id)
|
||||
.fetch_optional(&*state.db)
|
||||
.await?
|
||||
.ok_or_else(|| AppError::NotFound(format!("Site with id {} not found", id)))?;
|
||||
|
||||
Ok(Json(json!({ "data": site })))
|
||||
}
|
||||
|
||||
/// PUT /api/sites/{id}
|
||||
pub async fn update_site(
|
||||
State(state): State<AppState>,
|
||||
Path(id): Path<i64>,
|
||||
Json(payload): Json<UpdateSiteRequest>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
info!("Updating site: {}", id);
|
||||
|
||||
let existing: Site = sqlx::query_as::<_, Site>("SELECT * FROM sites WHERE id = ?")
|
||||
.bind(id)
|
||||
.fetch_optional(&*state.db)
|
||||
.await?
|
||||
.ok_or_else(|| AppError::NotFound(format!("Site with id {} not found", id)))?;
|
||||
|
||||
sqlx::query(
|
||||
"UPDATE sites SET name = ?, url = ?, platform = ?, status = ?, api_key = ?, proxy_url = ?, use_system_proxy = ?, custom_headers = ?, global_weight = ?, updated_at = datetime('now') WHERE id = ?"
|
||||
)
|
||||
.bind(payload.name.unwrap_or(existing.name))
|
||||
.bind(payload.url.unwrap_or(existing.url))
|
||||
.bind(payload.platform.unwrap_or(existing.platform))
|
||||
.bind(payload.status.unwrap_or(existing.status))
|
||||
.bind(payload.api_key.or(existing.api_key))
|
||||
.bind(payload.proxy_url.or(existing.proxy_url))
|
||||
.bind(payload.use_system_proxy.or(existing.use_system_proxy))
|
||||
.bind(payload.custom_headers.or(existing.custom_headers))
|
||||
.bind(payload.global_weight.or(existing.global_weight))
|
||||
.bind(id)
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
|
||||
let site: Site = sqlx::query_as::<_, Site>("SELECT * FROM sites WHERE id = ?")
|
||||
.bind(id)
|
||||
.fetch_one(&*state.db)
|
||||
.await?;
|
||||
|
||||
Ok(Json(json!({ "data": site })))
|
||||
}
|
||||
|
||||
/// DELETE /api/sites/{id}
|
||||
pub async fn delete_site(
|
||||
State(state): State<AppState>,
|
||||
Path(id): Path<i64>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
info!("Deleting site: {}", id);
|
||||
|
||||
let affected = sqlx::query("DELETE FROM sites WHERE id = ?")
|
||||
.bind(id)
|
||||
.execute(&*state.db)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
if affected == 0 {
|
||||
return Err(AppError::NotFound(format!("Site with id {} not found", id)));
|
||||
}
|
||||
|
||||
Ok(Json(json!({ "data": { "deleted": true, "id": id } })))
|
||||
}
|
||||
@@ -0,0 +1,286 @@
|
||||
use crate::{
|
||||
db::models::account::AccountToken,
|
||||
error::{AppError, AppResult},
|
||||
AppState,
|
||||
};
|
||||
use axum::{
|
||||
extract::{Path, Query, State},
|
||||
http::StatusCode,
|
||||
Json,
|
||||
};
|
||||
use serde::Deserialize;
|
||||
use serde_json::{json, Value};
|
||||
use tracing::info;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct ListTokensQuery {
|
||||
pub account_id: Option<i64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct CreateTokenRequest {
|
||||
pub account_id: i64,
|
||||
pub name: String,
|
||||
pub token: String,
|
||||
pub source: Option<String>,
|
||||
pub token_group: Option<String>,
|
||||
pub enabled: Option<bool>,
|
||||
pub is_default: Option<bool>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct UpdateTokenRequest {
|
||||
pub name: Option<String>,
|
||||
pub token: Option<String>,
|
||||
pub source: Option<String>,
|
||||
pub token_group: Option<String>,
|
||||
pub enabled: Option<bool>,
|
||||
pub is_default: Option<bool>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct RefreshRouteRequest {
|
||||
pub model_pattern: Option<String>,
|
||||
pub route_id: Option<i64>,
|
||||
}
|
||||
|
||||
/// GET /api/tokens?account_id=
|
||||
pub async fn list_tokens(
|
||||
State(state): State<AppState>,
|
||||
Query(query): Query<ListTokensQuery>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
info!("Listing tokens, account_id filter: {:?}", query.account_id);
|
||||
|
||||
let tokens: Vec<AccountToken> = if let Some(account_id) = query.account_id {
|
||||
sqlx::query_as::<_, AccountToken>(
|
||||
"SELECT * FROM account_tokens WHERE account_id = ? ORDER BY id",
|
||||
)
|
||||
.bind(account_id)
|
||||
.fetch_all(&*state.db)
|
||||
.await?
|
||||
} else {
|
||||
sqlx::query_as::<_, AccountToken>("SELECT * FROM account_tokens ORDER BY account_id, id")
|
||||
.fetch_all(&*state.db)
|
||||
.await?
|
||||
};
|
||||
|
||||
Ok(Json(json!({ "data": tokens })))
|
||||
}
|
||||
|
||||
/// POST /api/tokens
|
||||
pub async fn create_token(
|
||||
State(state): State<AppState>,
|
||||
Json(payload): Json<CreateTokenRequest>,
|
||||
) -> AppResult<(StatusCode, Json<Value>)> {
|
||||
info!("Creating token for account_id: {}", payload.account_id);
|
||||
|
||||
// Validate
|
||||
if payload.name.trim().is_empty() {
|
||||
return Err(AppError::BadRequest("name is required".into()));
|
||||
}
|
||||
if payload.token.trim().is_empty() {
|
||||
return Err(AppError::BadRequest("token is required".into()));
|
||||
}
|
||||
|
||||
// Verify account exists
|
||||
let account_exists: Option<(i64,)> = sqlx::query_as("SELECT id FROM accounts WHERE id = ?")
|
||||
.bind(payload.account_id)
|
||||
.fetch_optional(&*state.db)
|
||||
.await?;
|
||||
|
||||
if account_exists.is_none() {
|
||||
return Err(AppError::BadRequest(format!(
|
||||
"Account with id {} not found",
|
||||
payload.account_id
|
||||
)));
|
||||
}
|
||||
|
||||
let source = payload.source.unwrap_or_else(|| "manual".to_string());
|
||||
let enabled = payload.enabled.unwrap_or(true);
|
||||
let is_default = payload.is_default.unwrap_or(false);
|
||||
|
||||
let result = sqlx::query(
|
||||
"INSERT INTO account_tokens (account_id, name, token, source, token_group, enabled, is_default) VALUES (?, ?, ?, ?, ?, ?, ?)"
|
||||
)
|
||||
.bind(payload.account_id)
|
||||
.bind(&payload.name)
|
||||
.bind(&payload.token)
|
||||
.bind(&source)
|
||||
.bind(&payload.token_group)
|
||||
.bind(enabled)
|
||||
.bind(is_default)
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
|
||||
let token_id = result.last_insert_rowid();
|
||||
let token: AccountToken =
|
||||
sqlx::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
|
||||
.bind(token_id)
|
||||
.fetch_one(&*state.db)
|
||||
.await?;
|
||||
|
||||
Ok((StatusCode::CREATED, Json(json!({ "data": token }))))
|
||||
}
|
||||
|
||||
/// GET /api/tokens/{id}
|
||||
pub async fn get_token(
|
||||
State(state): State<AppState>,
|
||||
Path(id): Path<i64>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
info!("Getting token: {}", id);
|
||||
|
||||
let token: AccountToken =
|
||||
sqlx::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
|
||||
.bind(id)
|
||||
.fetch_optional(&*state.db)
|
||||
.await?
|
||||
.ok_or_else(|| AppError::NotFound(format!("Token with id {} not found", id)))?;
|
||||
|
||||
Ok(Json(json!({ "data": token })))
|
||||
}
|
||||
|
||||
/// PUT /api/tokens/{id}
|
||||
pub async fn update_token(
|
||||
State(state): State<AppState>,
|
||||
Path(id): Path<i64>,
|
||||
Json(payload): Json<UpdateTokenRequest>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
info!("Updating token: {}", id);
|
||||
|
||||
let existing: AccountToken =
|
||||
sqlx::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
|
||||
.bind(id)
|
||||
.fetch_optional(&*state.db)
|
||||
.await?
|
||||
.ok_or_else(|| AppError::NotFound(format!("Token with id {} not found", id)))?;
|
||||
|
||||
sqlx::query(
|
||||
"UPDATE account_tokens SET name = ?, token = ?, source = ?, token_group = ?, enabled = ?, is_default = ?, updated_at = datetime('now') WHERE id = ?"
|
||||
)
|
||||
.bind(payload.name.unwrap_or(existing.name))
|
||||
.bind(payload.token.as_ref().unwrap_or(&existing.token))
|
||||
.bind(payload.source.unwrap_or(existing.source.unwrap_or_else(|| "manual".to_string())))
|
||||
.bind(payload.token_group.or(existing.token_group))
|
||||
.bind(payload.enabled.unwrap_or(existing.enabled.unwrap_or(true)))
|
||||
.bind(payload.is_default.unwrap_or(existing.is_default.unwrap_or(false)))
|
||||
.bind(id)
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
|
||||
let token: AccountToken =
|
||||
sqlx::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
|
||||
.bind(id)
|
||||
.fetch_one(&*state.db)
|
||||
.await?;
|
||||
|
||||
Ok(Json(json!({ "data": token })))
|
||||
}
|
||||
|
||||
/// DELETE /api/tokens/{id}
|
||||
pub async fn delete_token(
|
||||
State(state): State<AppState>,
|
||||
Path(id): Path<i64>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
info!("Deleting token: {}", id);
|
||||
|
||||
let affected = sqlx::query("DELETE FROM account_tokens WHERE id = ?")
|
||||
.bind(id)
|
||||
.execute(&*state.db)
|
||||
.await?
|
||||
.rows_affected();
|
||||
|
||||
if affected == 0 {
|
||||
return Err(AppError::NotFound(format!(
|
||||
"Token with id {} not found",
|
||||
id
|
||||
)));
|
||||
}
|
||||
|
||||
Ok(Json(json!({ "data": { "deleted": true, "id": id } })))
|
||||
}
|
||||
|
||||
/// POST /api/tokens/refresh-route
|
||||
pub async fn refresh_route(
|
||||
State(state): State<AppState>,
|
||||
Json(payload): Json<RefreshRouteRequest>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
info!("Refreshing route decision");
|
||||
|
||||
if let Some(route_id) = payload.route_id {
|
||||
sqlx::query(
|
||||
"UPDATE token_routes SET decision_refreshed_at = datetime('now'), updated_at = datetime('now') WHERE id = ?"
|
||||
)
|
||||
.bind(route_id)
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
|
||||
let route: Option<crate::db::models::token_route::TokenRoute> =
|
||||
sqlx::query_as("SELECT * FROM token_routes WHERE id = ?")
|
||||
.bind(route_id)
|
||||
.fetch_optional(&*state.db)
|
||||
.await?;
|
||||
|
||||
return Ok(Json(
|
||||
json!({ "data": { "refreshed": true, "route": route } }),
|
||||
));
|
||||
}
|
||||
|
||||
if let Some(ref model_pattern) = payload.model_pattern {
|
||||
let result = sqlx::query(
|
||||
"UPDATE token_routes SET decision_refreshed_at = datetime('now'), updated_at = datetime('now') WHERE model_pattern LIKE ?"
|
||||
)
|
||||
.bind(format!("%{}%", model_pattern))
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
|
||||
return Ok(Json(json!({
|
||||
"data": {
|
||||
"refreshed": true,
|
||||
"affected_routes": result.rows_affected()
|
||||
}
|
||||
})));
|
||||
}
|
||||
|
||||
// Refresh all routes
|
||||
let result = sqlx::query(
|
||||
"UPDATE token_routes SET decision_refreshed_at = datetime('now'), updated_at = datetime('now')"
|
||||
)
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
|
||||
Ok(Json(json!({
|
||||
"data": {
|
||||
"refreshed": true,
|
||||
"affected_routes": result.rows_affected()
|
||||
}
|
||||
})))
|
||||
}
|
||||
|
||||
/// GET /api/tokens/route-decision
|
||||
pub async fn route_decision(
|
||||
State(state): State<AppState>,
|
||||
Query(query): Query<RefreshRouteRequest>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
info!("Getting route decisions");
|
||||
|
||||
let routes: Vec<crate::db::models::token_route::TokenRoute> =
|
||||
if let Some(route_id) = query.route_id {
|
||||
sqlx::query_as("SELECT * FROM token_routes WHERE id = ? AND enabled = 1")
|
||||
.bind(route_id)
|
||||
.fetch_optional(&*state.db)
|
||||
.await?
|
||||
.map(|r| vec![r])
|
||||
.unwrap_or_default()
|
||||
} else if let Some(ref model_pattern) = query.model_pattern {
|
||||
sqlx::query_as("SELECT * FROM token_routes WHERE model_pattern LIKE ? AND enabled = 1")
|
||||
.bind(format!("%{}%", model_pattern))
|
||||
.fetch_all(&*state.db)
|
||||
.await?
|
||||
} else {
|
||||
sqlx::query_as("SELECT * FROM token_routes WHERE enabled = 1 ORDER BY model_pattern")
|
||||
.fetch_all(&*state.db)
|
||||
.await?
|
||||
};
|
||||
|
||||
Ok(Json(json!({ "data": routes })))
|
||||
}
|
||||
@@ -0,0 +1,315 @@
|
||||
// ComfyUI proxy routes — forward to local ComfyUI instance
|
||||
//
|
||||
// These routes proxy requests to a local ComfyUI server for AI video/image workflow editing.
|
||||
// All business routes require downstream API key or admin token authentication.
|
||||
// The /api/comfyui/proxy/* route serves static assets (public).
|
||||
|
||||
use crate::{
|
||||
auth::api_key::{ authorize_downstream_token, DownstreamAuthResult },
|
||||
error::{ AppError, AppResult },
|
||||
AppState,
|
||||
};
|
||||
use axum::{
|
||||
extract::{ Path, Query, State },
|
||||
http::{ HeaderMap, StatusCode },
|
||||
response::{ Html, IntoResponse, Response },
|
||||
routing::{ get, post },
|
||||
Json, Router,
|
||||
};
|
||||
use reqwest::Client;
|
||||
use serde_json::Value;
|
||||
use std::collections::HashMap;
|
||||
use tracing::info;
|
||||
|
||||
/// Default ComfyUI backend address.
|
||||
const DEFAULT_COMFYUI_BASE: &str = "http://127.0.0.1:8188";
|
||||
|
||||
/// Shared HTTP client
|
||||
fn http_client() -> Client {
|
||||
Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(120))
|
||||
.build()
|
||||
.expect("Failed to build HTTP client")
|
||||
}
|
||||
|
||||
// ── Auth helper (same pattern as proxy module) ────────────────────────
|
||||
|
||||
async fn check_auth(state: &AppState, headers: &HeaderMap) -> Result<(), AppError> {
|
||||
let auth = headers
|
||||
.get("Authorization")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|v| v.strip_prefix("Bearer "))
|
||||
.and_then(|v| v.trim().split_whitespace().next())
|
||||
.ok_or_else(|| AppError::Unauthorized("Missing Authorization header".into()))?;
|
||||
|
||||
// Admin token passes
|
||||
if auth == state.config.auth.auth_token {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Downstream API key
|
||||
if auth.starts_with("sk-") {
|
||||
let result = authorize_downstream_token(auth, &state.db).await
|
||||
.map_err(|e| AppError::Unauthorized(e))?;
|
||||
|
||||
match result {
|
||||
DownstreamAuthResult::Managed { key, .. } => {
|
||||
if !key.enabled.unwrap_or(true) {
|
||||
return Err(AppError::Forbidden("API key is disabled".into()));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
DownstreamAuthResult::Global { .. } => Ok(()),
|
||||
}
|
||||
} else {
|
||||
Err(AppError::Unauthorized("Invalid token".into()))
|
||||
}
|
||||
}
|
||||
|
||||
/// Get ComfyUI base URL from env or default
|
||||
fn comfyui_base() -> String {
|
||||
std::env::var("COMFYUI_BASE").unwrap_or_else(|_| DEFAULT_COMFYUI_BASE.to_string())
|
||||
}
|
||||
|
||||
// ── Route builder ────────────────────────────────────────────────────
|
||||
|
||||
pub fn routes() -> Router<AppState> {
|
||||
Router::new()
|
||||
// Business routes (require auth)
|
||||
.route("/api/comfyui/prompt", post(proxy_post_prompt))
|
||||
.route("/api/comfyui/history", get(proxy_get_history))
|
||||
.route("/api/comfyui/history/{id}", get(proxy_get_history_by_id))
|
||||
.route("/api/comfyui/queue", get(proxy_get_queue))
|
||||
.route("/api/comfyui/view", get(proxy_get_view))
|
||||
// Static proxy (public — assets/JS/CSS)
|
||||
.route("/api/comfyui/proxy/*path", get(proxy_static))
|
||||
// HTML page (public)
|
||||
.route("/comfyui", get(comfyui_page))
|
||||
}
|
||||
|
||||
// ── Business route handlers (require auth) ───────────────────────────
|
||||
|
||||
/// POST /api/comfyui/prompt
|
||||
async fn proxy_post_prompt(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Json(body): Json<Value>,
|
||||
) -> AppResult<Response> {
|
||||
info!("POST /api/comfyui/prompt");
|
||||
check_auth(&state, &headers).await?;
|
||||
proxy_to_comfyui("POST", "/prompt", None, Some(body)).await
|
||||
}
|
||||
|
||||
/// GET /api/comfyui/history
|
||||
async fn proxy_get_history(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
) -> AppResult<Response> {
|
||||
info!("GET /api/comfyui/history");
|
||||
check_auth(&state, &headers).await?;
|
||||
proxy_to_comfyui("GET", "/history", None, None).await
|
||||
}
|
||||
|
||||
/// GET /api/comfyui/history/{id}
|
||||
async fn proxy_get_history_by_id(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Path(id): Path<String>,
|
||||
) -> AppResult<Response> {
|
||||
info!("GET /api/comfyui/history/{}", id);
|
||||
check_auth(&state, &headers).await?;
|
||||
let path = format!("/history/{}", urlencoding(&id));
|
||||
proxy_to_comfyui("GET", &path, None, None).await
|
||||
}
|
||||
|
||||
/// GET /api/comfyui/queue
|
||||
async fn proxy_get_queue(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
) -> AppResult<Response> {
|
||||
info!("GET /api/comfyui/queue");
|
||||
check_auth(&state, &headers).await?;
|
||||
proxy_to_comfyui("GET", "/queue", None, None).await
|
||||
}
|
||||
|
||||
/// GET /api/comfyui/view?filename=...&type=...&subfolder=...
|
||||
async fn proxy_get_view(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Query(params): Query<HashMap<String, String>>,
|
||||
) -> AppResult<Response> {
|
||||
info!("GET /api/comfyui/view");
|
||||
check_auth(&state, &headers).await?;
|
||||
|
||||
let qs = build_query_string(¶ms);
|
||||
let path = format!("/view{}", qs);
|
||||
proxy_to_comfyui_binary("GET", &path).await
|
||||
}
|
||||
|
||||
// ── Static proxy (public — assets/JS/CSS for iframe) ─────────────────
|
||||
|
||||
/// GET /api/comfyui/proxy/{path} — proxy everything (HTML, assets)
|
||||
async fn proxy_static(
|
||||
Path(path): Path<String>,
|
||||
) -> AppResult<Response> {
|
||||
info!("GET /api/comfyui/proxy/{}", path);
|
||||
let target_path = if path.is_empty() || path == "/" {
|
||||
"/".to_string()
|
||||
} else {
|
||||
format!("/{}", path)
|
||||
};
|
||||
|
||||
let client = http_client();
|
||||
let url = format!("{}{}", comfyui_base(), target_path);
|
||||
let resp = client.get(&url).send().await
|
||||
.map_err(|e| AppError::UpstreamError(format!("ComfyUI unreachable: {e}")))?;
|
||||
|
||||
let status = resp.status();
|
||||
let content_type = resp.headers()
|
||||
.get("content-type")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.unwrap_or("application/octet-stream")
|
||||
.to_string();
|
||||
|
||||
if content_type.contains("html") {
|
||||
let text = resp.text().await
|
||||
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
|
||||
// Rebrand: ComfyUI → BoosAPI
|
||||
let rebranded = text
|
||||
.replace("ComfyUI", "BoosAPI")
|
||||
.replace("comfyui", "BoosAPI");
|
||||
// Rewrite URLs
|
||||
let rewritten = rewrite_proxy_urls(&rebranded);
|
||||
Ok((StatusCode::from_u16(status.as_u16()).unwrap(), [(axum::http::header::CONTENT_TYPE, "text/html; charset=utf-8")], rewritten).into_response())
|
||||
} else if content_type.contains("text") || content_type.contains("json") {
|
||||
let text = resp.text().await
|
||||
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
|
||||
Ok((StatusCode::from_u16(status.as_u16()).unwrap(), [(axum::http::header::CONTENT_TYPE, content_type)], text).into_response())
|
||||
} else {
|
||||
let bytes = resp.bytes().await
|
||||
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
|
||||
Ok((StatusCode::from_u16(status.as_u16()).unwrap(), [(axum::http::header::CONTENT_TYPE, content_type)], bytes).into_response())
|
||||
}
|
||||
}
|
||||
|
||||
/// GET /comfyui — HTML page with iframe
|
||||
async fn comfyui_page() -> Html<&'static str> {
|
||||
Html(r##"<!DOCTYPE html>
|
||||
<html lang="zh-CN">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>BoosAPI - ComfyUI</title>
|
||||
<style>
|
||||
* { margin: 0; padding: 0; box-sizing: border-box; }
|
||||
html, body { width: 100%; height: 100%; overflow: hidden; background: #1a1a2e; }
|
||||
iframe { width: 100%; height: 100%; border: none; }
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<iframe src="/api/comfyui/proxy/" title="BoosAPI ComfyUI"></iframe>
|
||||
</body>
|
||||
</html>"##)
|
||||
}
|
||||
|
||||
// ── Helpers ──────────────────────────────────────────────────────────
|
||||
|
||||
/// Rewrite URLs in HTML so asset paths go through the proxy
|
||||
fn rewrite_proxy_urls(html: &str) -> String {
|
||||
// Replace src="/...", href="/...", action="/..." with /api/comfyui/proxy/...
|
||||
let mut result = String::with_capacity(html.len() * 12 / 10);
|
||||
let mut chars = html.char_indices().peekable();
|
||||
|
||||
while let Some((_, ch)) = chars.next() {
|
||||
if ch == 's' || ch == 'h' || ch == 'a' {
|
||||
// Look ahead for rc="/", ref="/", ction="/"
|
||||
let snippet: String = chars.clone().take(6).map(|(_, c)| c).collect();
|
||||
let full = format!("{}{}", ch, snippet);
|
||||
if full.starts_with("src=\"/") || full.starts_with("href=\"/") || full.starts_with("action=\"/") {
|
||||
let attr_end = if full.starts_with("src=\"") { 4 } else if full.starts_with("href=\"") { 5 } else { 7 };
|
||||
result.push_str(&full[..attr_end]);
|
||||
result.push_str("/api/comfyui/proxy/");
|
||||
// Skip the rest of the attribute
|
||||
for _ in 0..snippet.len() { chars.next(); }
|
||||
continue;
|
||||
}
|
||||
}
|
||||
result.push(ch);
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
/// Build query string from params
|
||||
fn build_query_string(params: &HashMap<String, String>) -> String {
|
||||
if params.is_empty() {
|
||||
return String::new();
|
||||
}
|
||||
let mut parts: Vec<String> = params
|
||||
.iter()
|
||||
.map(|(k, v)| format!("{}={}", urlencoding(k), urlencoding(v)))
|
||||
.collect();
|
||||
parts.sort(); // deterministic ordering
|
||||
format!("?{}", parts.join("&"))
|
||||
}
|
||||
|
||||
/// Simple URL encoding (handles the chars we need)
|
||||
fn urlencoding(s: &str) -> String {
|
||||
s.replace('%', "%25")
|
||||
.replace(' ', "%20")
|
||||
.replace('?', "%3F")
|
||||
.replace('&', "%26")
|
||||
.replace('#', "%23")
|
||||
.replace('/', "%2F")
|
||||
}
|
||||
|
||||
/// Proxy to ComfyUI — returns JSON/text response
|
||||
async fn proxy_to_comfyui(
|
||||
method: &str,
|
||||
path: &str,
|
||||
_query: Option<&HashMap<String, String>>,
|
||||
body: Option<Value>,
|
||||
) -> AppResult<Response> {
|
||||
let client = http_client();
|
||||
let url = format!("{}{}", comfyui_base(), path);
|
||||
|
||||
let req = match method {
|
||||
"POST" => {
|
||||
let body_str = body.map(|b| b.to_string()).unwrap_or_else(|| "{}".to_string());
|
||||
client.post(&url)
|
||||
.header("Content-Type", "application/json")
|
||||
.body(body_str)
|
||||
}
|
||||
_ => client.get(&url),
|
||||
};
|
||||
|
||||
let resp = req.send().await
|
||||
.map_err(|e| AppError::UpstreamError(format!("ComfyUI unreachable: {e}")))?;
|
||||
|
||||
let status = resp.status();
|
||||
let text = resp.text().await
|
||||
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
|
||||
|
||||
Ok((StatusCode::from_u16(status.as_u16()).unwrap(), [(axum::http::header::CONTENT_TYPE, "application/json")], text).into_response())
|
||||
}
|
||||
|
||||
/// Proxy to ComfyUI — returns binary response (for images)
|
||||
async fn proxy_to_comfyui_binary(_method: &str, path: &str) -> AppResult<Response> {
|
||||
let client = http_client();
|
||||
let url = format!("{}{}", comfyui_base(), path);
|
||||
|
||||
let req = client.get(&url);
|
||||
let resp = req.send().await
|
||||
.map_err(|e| AppError::UpstreamError(format!("ComfyUI unreachable: {e}")))?;
|
||||
|
||||
let status = resp.status();
|
||||
let content_type = resp.headers()
|
||||
.get("content-type")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.unwrap_or("application/octet-stream")
|
||||
.to_string();
|
||||
let bytes = resp.bytes().await
|
||||
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
|
||||
|
||||
Ok((StatusCode::from_u16(status.as_u16()).unwrap(), [(axum::http::header::CONTENT_TYPE, content_type)], bytes).into_response())
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
use crate::{error::AppResult, AppState};
|
||||
use axum::{extract::State, Json};
|
||||
use serde_json::{json, Value};
|
||||
|
||||
pub async fn health_check(State(state): State<AppState>) -> AppResult<Json<Value>> {
|
||||
// Simple DB ping
|
||||
sqlx::query("SELECT 1").execute(&*state.db).await?;
|
||||
|
||||
Ok(Json(json!({
|
||||
"status": "ok",
|
||||
"database": "ok",
|
||||
"version": env!("CARGO_PKG_VERSION"),
|
||||
})))
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
use crate::AppState;
|
||||
use axum::{routing::get, Router};
|
||||
|
||||
mod admin;
|
||||
mod comfyui;
|
||||
mod health;
|
||||
mod proxy;
|
||||
|
||||
pub fn build_router(state: AppState) -> Router {
|
||||
Router::new()
|
||||
.route("/healthz", get(health::health_check))
|
||||
.merge(admin::routes(state.clone()))
|
||||
.merge(proxy::routes())
|
||||
.merge(comfyui::routes())
|
||||
.with_state(state)
|
||||
}
|
||||
@@ -0,0 +1,190 @@
|
||||
// Proxy routes — OpenAI-compatible /v1/ endpoints
|
||||
//
|
||||
// These routes handle:
|
||||
// POST /v1/chat/completions — streaming and non-streaming
|
||||
// POST /v1/images/generations — image generation
|
||||
// GET /v1/models — list available models
|
||||
|
||||
use crate::{
|
||||
auth::api_key::{ authorize_downstream_token, DownstreamAuthResult },
|
||||
error::{ AppError, AppResult },
|
||||
proxy::executors::{ self, UpstreamRoute },
|
||||
AppState,
|
||||
};
|
||||
use axum::{
|
||||
extract::State,
|
||||
http::HeaderMap,
|
||||
response::{ IntoResponse, Response, Sse },
|
||||
routing::{ get, post },
|
||||
Json, Router,
|
||||
};
|
||||
use futures::stream;
|
||||
use futures::StreamExt;
|
||||
use serde_json::{ json, Value };
|
||||
use std::convert::Infallible;
|
||||
use tracing::info;
|
||||
|
||||
// ── Auth types ────────────────────────────────────────────────────────
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ProxyAuth {
|
||||
pub token: String,
|
||||
pub key_id: Option<i64>,
|
||||
pub source: DownstreamSource,
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub enum DownstreamSource {
|
||||
Managed { key_id: i64 },
|
||||
Global,
|
||||
}
|
||||
|
||||
// ── Route definitions ─────────────────────────────────────────────────
|
||||
|
||||
/// Build the /v1/ proxy router
|
||||
pub fn routes() -> Router<AppState> {
|
||||
Router::new()
|
||||
.route("/v1/chat/completions", post(chat_completions))
|
||||
.route("/v1/images/generations", post(images_generations))
|
||||
.route("/v1/models", get(list_models))
|
||||
}
|
||||
|
||||
// ── Handlers ───────────────────────────────────────────────────────────
|
||||
|
||||
/// POST /v1/chat/completions
|
||||
async fn chat_completions(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Json(body): Json<Value>,
|
||||
) -> Result<Response, AppError> {
|
||||
info!("POST /v1/chat/completions");
|
||||
|
||||
check_auth(&state, &headers).await?;
|
||||
|
||||
let model = body
|
||||
.get("model")
|
||||
.and_then(|m| m.as_str())
|
||||
.ok_or_else(|| AppError::BadRequest("model is required".into()))?
|
||||
.to_string();
|
||||
|
||||
let route = resolve_route(&state, &model).await?;
|
||||
let stream = body.get("stream").and_then(|s| s.as_bool()).unwrap_or(false);
|
||||
|
||||
if stream {
|
||||
let chunks = executors::execute_chat_completions_stream(state, route, body).await?;
|
||||
let sse = Sse::new(stream::iter(chunks).map(|chunk| {
|
||||
Ok::<_, Infallible>(axum::response::sse::Event::default().data(chunk))
|
||||
}));
|
||||
Ok(sse.into_response())
|
||||
} else {
|
||||
let result = executors::execute_chat_completions(state, route, body).await?;
|
||||
Ok(Json(result).into_response())
|
||||
}
|
||||
}
|
||||
|
||||
/// POST /v1/images/generations
|
||||
async fn images_generations(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Json(body): Json<Value>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
info!("POST /v1/images/generations");
|
||||
|
||||
check_auth(&state, &headers).await?;
|
||||
|
||||
let model = body
|
||||
.get("model")
|
||||
.and_then(|m| m.as_str())
|
||||
.ok_or_else(|| AppError::BadRequest("model is required".into()))?
|
||||
.to_string();
|
||||
|
||||
let route = resolve_route(&state, &model).await?;
|
||||
let result = executors::execute_image_generation(state, route, body).await?;
|
||||
Ok(Json(result))
|
||||
}
|
||||
|
||||
/// GET /v1/models
|
||||
async fn list_models(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
) -> AppResult<Json<Value>> {
|
||||
info!("GET /v1/models");
|
||||
|
||||
check_auth(&state, &headers).await?;
|
||||
|
||||
let rows: Vec<(String, String)> = sqlx::query_as(
|
||||
"SELECT tr.name, tr.model FROM token_routes tr WHERE tr.enabled = 1 ORDER BY tr.name"
|
||||
)
|
||||
.fetch_all(&*state.db)
|
||||
.await?;
|
||||
|
||||
let models: Vec<Value> = rows
|
||||
.into_iter()
|
||||
.map(|(name, model)| {
|
||||
json!({
|
||||
"id": model,
|
||||
"object": "model",
|
||||
"created": 0,
|
||||
"owned_by": name,
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(Json(json!({
|
||||
"object": "list",
|
||||
"data": models,
|
||||
})))
|
||||
}
|
||||
|
||||
// ── Auth helper ───────────────────────────────────────────────────────
|
||||
|
||||
async fn check_auth(state: &AppState, headers: &HeaderMap) -> Result<(), AppError> {
|
||||
let auth = headers
|
||||
.get("Authorization")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|v| v.strip_prefix("Bearer "))
|
||||
.and_then(|v| v.trim().split_whitespace().next())
|
||||
.ok_or_else(|| AppError::Unauthorized("Missing Authorization header".into()))?;
|
||||
|
||||
let result = authorize_downstream_token(auth, &state.db).await
|
||||
.map_err(|e| AppError::Unauthorized(e))?;
|
||||
|
||||
match result {
|
||||
DownstreamAuthResult::Managed { key, .. } => {
|
||||
if !key.enabled.unwrap_or(true) {
|
||||
return Err(AppError::Forbidden("API key is disabled".into()));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
DownstreamAuthResult::Global { .. } => Ok(()),
|
||||
}
|
||||
}
|
||||
|
||||
// ── Route Resolution ──────────────────────────────────────────────────
|
||||
|
||||
/// Resolve upstream endpoint + API key for a given model.
|
||||
async fn resolve_route(state: &AppState, model: &str) -> AppResult<UpstreamRoute> {
|
||||
let row: Option<(String, String)> = sqlx::query_as(
|
||||
"SELECT s.base_url, COALESCE(a.api_key, s.default_key, '') as api_key
|
||||
FROM token_routes tr
|
||||
JOIN sites s ON s.id = tr.site_id
|
||||
LEFT JOIN accounts a ON a.id = tr.account_id AND a.site_id = s.id
|
||||
WHERE tr.enabled = 1 AND tr.model = ?
|
||||
LIMIT 1"
|
||||
)
|
||||
.bind(model)
|
||||
.fetch_optional(&*state.db)
|
||||
.await?;
|
||||
|
||||
match row {
|
||||
Some((base_url, api_key)) => Ok(UpstreamRoute {
|
||||
base_url,
|
||||
api_key,
|
||||
model: model.to_string(),
|
||||
}),
|
||||
None => Err(AppError::BadRequest(format!(
|
||||
"No route configured for model '{}'",
|
||||
model
|
||||
))),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
// Routes - user
|
||||
Reference in New Issue
Block a user