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,108 @@
|
||||
use argon2::password_hash::{PasswordHash, PasswordHasher, PasswordVerifier, SaltString};
|
||||
use argon2::Argon2;
|
||||
use sqlx::SqlitePool;
|
||||
|
||||
use crate::db::models::downstream_api_key::DownstreamApiKey;
|
||||
|
||||
/// Result of downstream token authorization.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum DownstreamAuthResult {
|
||||
/// Found in downstream_api_keys table.
|
||||
Managed {
|
||||
key: DownstreamApiKey,
|
||||
},
|
||||
/// Matched the global proxy token from config.
|
||||
Global {
|
||||
token: String,
|
||||
},
|
||||
}
|
||||
|
||||
/// Authorize a downstream token (sk-xxx or global proxy token).
|
||||
///
|
||||
/// Looks up the token in:
|
||||
/// 1. downstream_api_keys table (for managed keys)
|
||||
/// 2. Falls back to checking against the global proxy token
|
||||
pub async fn authorize_downstream_token(
|
||||
token: &str,
|
||||
pool: &SqlitePool,
|
||||
) -> Result<DownstreamAuthResult, String> {
|
||||
// Try managed keys first
|
||||
if token.starts_with("sk-") {
|
||||
let key: Option<DownstreamApiKey> = sqlx::query_as::<_, DownstreamApiKey>(
|
||||
"SELECT * FROM downstream_api_keys WHERE key = ? AND enabled = 1"
|
||||
)
|
||||
.bind(token)
|
||||
.fetch_optional(pool)
|
||||
.await
|
||||
.map_err(|e| format!("Database error: {e}"))?;
|
||||
|
||||
if let Some(k) = key {
|
||||
return Ok(DownstreamAuthResult::Managed { key: k });
|
||||
}
|
||||
|
||||
// Key not found in table
|
||||
return Err("Invalid API key".to_string());
|
||||
}
|
||||
|
||||
// Not an sk- key, can't be a downstream key
|
||||
Err("Invalid API key format".to_string())
|
||||
}
|
||||
|
||||
/// Generate 16 random bytes for a salt using `uuid::Uuid::new_v4()`.
|
||||
fn random_salt_bytes() -> [u8; 16] {
|
||||
*uuid::Uuid::new_v4().as_bytes()
|
||||
}
|
||||
|
||||
/// Hash an API key using Argon2id with a random salt.
|
||||
pub fn hash_api_key(key: &str) -> anyhow::Result<String> {
|
||||
let salt = SaltString::encode_b64(&random_salt_bytes())
|
||||
.map_err(|e| anyhow::anyhow!("Failed to generate salt: {e}"))?;
|
||||
let hash = Argon2::default()
|
||||
.hash_password(key.as_bytes(), &salt)
|
||||
.map_err(|e| anyhow::anyhow!("Failed to hash API key: {e}"))?;
|
||||
Ok(hash.to_string())
|
||||
}
|
||||
|
||||
/// Verify an API key against an Argon2id hash.
|
||||
pub fn verify_api_key(key: &str, hash: &str) -> anyhow::Result<bool> {
|
||||
let parsed_hash =
|
||||
PasswordHash::new(hash).map_err(|e| anyhow::anyhow!("Invalid hash format: {e}"))?;
|
||||
Ok(Argon2::default()
|
||||
.verify_password(key.as_bytes(), &parsed_hash)
|
||||
.is_ok())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_hash_and_verify_api_key() {
|
||||
let key = "sk-test-api-key-12345";
|
||||
let hash = hash_api_key(key).unwrap();
|
||||
assert!(verify_api_key(key, &hash).unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_wrong_key_fails() {
|
||||
let key = "sk-test-api-key-12345";
|
||||
let hash = hash_api_key(key).unwrap();
|
||||
assert!(!verify_api_key("wrong-key", &hash).unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_invalid_hash_fails() {
|
||||
let result = verify_api_key("some-key", "not-a-valid-hash");
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_different_hashes_for_same_key() {
|
||||
let key = "sk-test-key";
|
||||
let hash1 = hash_api_key(key).unwrap();
|
||||
let hash2 = hash_api_key(key).unwrap();
|
||||
assert_ne!(hash1, hash2);
|
||||
assert!(verify_api_key(key, &hash1).unwrap());
|
||||
assert!(verify_api_key(key, &hash2).unwrap());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
use chrono::{Duration, Utc};
|
||||
use jsonwebtoken::{decode, encode, DecodingKey, EncodingKey, Header, Validation};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct Claims {
|
||||
pub sub: String, // user_id
|
||||
pub role: String, // "admin" | "user"
|
||||
pub exp: usize, // expiry timestamp
|
||||
pub iat: usize, // issued at
|
||||
}
|
||||
|
||||
pub fn generate_jwt(
|
||||
secret: &str,
|
||||
user_id: &str,
|
||||
role: &str,
|
||||
expiry_hours: i64,
|
||||
) -> anyhow::Result<String> {
|
||||
let now = Utc::now();
|
||||
let claims = Claims {
|
||||
sub: user_id.to_string(),
|
||||
role: role.to_string(),
|
||||
iat: now.timestamp() as usize,
|
||||
exp: (now + Duration::hours(expiry_hours)).timestamp() as usize,
|
||||
};
|
||||
Ok(encode(
|
||||
&Header::default(),
|
||||
&claims,
|
||||
&EncodingKey::from_secret(secret.as_bytes()),
|
||||
)?)
|
||||
}
|
||||
|
||||
pub fn validate_jwt(secret: &str, token: &str) -> anyhow::Result<Claims> {
|
||||
let data = decode::<Claims>(
|
||||
token,
|
||||
&DecodingKey::from_secret(secret.as_bytes()),
|
||||
&Validation::default(),
|
||||
)?;
|
||||
Ok(data.claims)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_generate_and_validate_jwt() {
|
||||
let secret = "test-secret-key";
|
||||
let token = generate_jwt(secret, "user-123", "admin", 24).unwrap();
|
||||
let claims = validate_jwt(secret, &token).unwrap();
|
||||
assert_eq!(claims.sub, "user-123");
|
||||
assert_eq!(claims.role, "admin");
|
||||
assert!(claims.iat > 0);
|
||||
assert!(claims.exp > claims.iat);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_invalid_token_fails() {
|
||||
let secret = "test-secret-key";
|
||||
let result = validate_jwt(secret, "invalid-token");
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_wrong_secret_fails() {
|
||||
let token = generate_jwt("secret1", "user-123", "user", 24).unwrap();
|
||||
let result = validate_jwt("secret2", &token);
|
||||
assert!(result.is_err());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
// Auth module — JWT, API key, and password authentication
|
||||
|
||||
pub mod api_key;
|
||||
pub mod jwt;
|
||||
pub mod password;
|
||||
|
||||
/// Represents the authentication context for a request.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum AuthContext {
|
||||
/// Admin token-based authentication.
|
||||
Admin { token: String },
|
||||
/// Authenticated user session with JWT.
|
||||
User { user_id: String, role: String },
|
||||
/// Downstream API key authentication.
|
||||
DownstreamKey { key_id: i64, name: String },
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
use argon2::password_hash::{PasswordHash, PasswordHasher, PasswordVerifier, SaltString};
|
||||
use argon2::Argon2;
|
||||
|
||||
/// Generate 16 random bytes for a salt using `uuid::Uuid::new_v4()`.
|
||||
fn random_salt_bytes() -> [u8; 16] {
|
||||
*uuid::Uuid::new_v4().as_bytes()
|
||||
}
|
||||
|
||||
/// Hash a user password using Argon2id with a random salt.
|
||||
pub fn hash_password(password: &str) -> anyhow::Result<String> {
|
||||
let salt = SaltString::encode_b64(&random_salt_bytes())
|
||||
.map_err(|e| anyhow::anyhow!("Failed to generate salt: {e}"))?;
|
||||
let hash = Argon2::default()
|
||||
.hash_password(password.as_bytes(), &salt)
|
||||
.map_err(|e| anyhow::anyhow!("Failed to hash password: {e}"))?;
|
||||
Ok(hash.to_string())
|
||||
}
|
||||
|
||||
/// Verify a password against an Argon2id hash.
|
||||
pub fn verify_password(password: &str, hash: &str) -> anyhow::Result<bool> {
|
||||
let parsed_hash =
|
||||
PasswordHash::new(hash).map_err(|e| anyhow::anyhow!("Invalid hash format: {e}"))?;
|
||||
Ok(Argon2::default()
|
||||
.verify_password(password.as_bytes(), &parsed_hash)
|
||||
.is_ok())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_hash_and_verify_password() {
|
||||
let password = "secure-password-123!";
|
||||
let hash = hash_password(password).unwrap();
|
||||
assert!(verify_password(password, &hash).unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_wrong_password_fails() {
|
||||
let password = "correct-password";
|
||||
let hash = hash_password(password).unwrap();
|
||||
assert!(!verify_password("wrong-password", &hash).unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_different_hashes_for_same_password() {
|
||||
let password = "my-password";
|
||||
let hash1 = hash_password(password).unwrap();
|
||||
let hash2 = hash_password(password).unwrap();
|
||||
assert_ne!(hash1, hash2);
|
||||
assert!(verify_password(password, &hash1).unwrap());
|
||||
assert!(verify_password(password, &hash2).unwrap());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_password() {
|
||||
let password = "";
|
||||
let hash = hash_password(password).unwrap();
|
||||
assert!(verify_password(password, &hash).unwrap());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
use figment::{
|
||||
providers::{Env, Format, Toml},
|
||||
Figment,
|
||||
};
|
||||
use serde::Deserialize;
|
||||
|
||||
#[derive(Debug, Deserialize, Clone)]
|
||||
pub struct AppConfig {
|
||||
pub server: ServerConfig,
|
||||
pub db: DbConfig,
|
||||
pub auth: AuthConfig,
|
||||
pub proxy: ProxyConfig,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Clone)]
|
||||
pub struct ServerConfig {
|
||||
pub host: String,
|
||||
pub port: u16,
|
||||
pub log_level: String,
|
||||
pub cors_origins: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Clone)]
|
||||
pub struct DbConfig {
|
||||
pub url: String,
|
||||
pub max_connections: u32,
|
||||
pub ssl: bool,
|
||||
pub runtime_db_url: Option<String>,
|
||||
pub runtime_db_dialect: Option<String>,
|
||||
pub runtime_db_ssl: Option<bool>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Clone)]
|
||||
pub struct AuthConfig {
|
||||
pub jwt_secret: String,
|
||||
pub jwt_expiry_hours: i64,
|
||||
pub auth_token: String,
|
||||
pub proxy_token: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Clone)]
|
||||
pub struct ProxyConfig {
|
||||
pub upstream_timeout_secs: u64,
|
||||
pub first_byte_timeout_secs: u64,
|
||||
pub max_retries: u32,
|
||||
pub proxy_log_retention_days: u32,
|
||||
}
|
||||
|
||||
impl AppConfig {
|
||||
pub fn load() -> anyhow::Result<Self> {
|
||||
dotenvy::dotenv().ok();
|
||||
Ok(Figment::new()
|
||||
.merge(Toml::file("config.toml").nested())
|
||||
.merge(Env::prefixed("BOOSAPI_").split("__"))
|
||||
.extract()?)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
use sqlx::sqlite::{SqlitePool, SqlitePoolOptions};
|
||||
|
||||
pub mod models;
|
||||
|
||||
pub async fn create_pool(database_url: &str, max_connections: u32) -> anyhow::Result<SqlitePool> {
|
||||
let pool = SqlitePoolOptions::new()
|
||||
.max_connections(max_connections)
|
||||
.connect(database_url)
|
||||
.await?;
|
||||
Ok(pool)
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sqlx::FromRow;
|
||||
|
||||
/// Represents an account bound to a site.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, FromRow)]
|
||||
pub struct Account {
|
||||
pub id: i64,
|
||||
pub site_id: i64,
|
||||
pub username: Option<String>,
|
||||
pub access_token: String,
|
||||
pub api_token: Option<String>,
|
||||
pub balance: Option<f64>,
|
||||
pub balance_used: Option<f64>,
|
||||
pub quota: Option<f64>,
|
||||
pub unit_cost: Option<f64>,
|
||||
pub value_score: Option<f64>,
|
||||
pub status: String,
|
||||
pub checkin_enabled: Option<bool>,
|
||||
pub last_checkin_at: Option<String>,
|
||||
pub last_balance_refresh: Option<String>,
|
||||
pub extra_config: Option<String>,
|
||||
pub is_pinned: Option<bool>,
|
||||
pub sort_order: Option<i64>,
|
||||
pub oauth_provider: Option<String>,
|
||||
pub oauth_account_key: Option<String>,
|
||||
pub oauth_project_id: Option<String>,
|
||||
pub created_at: Option<String>,
|
||||
pub updated_at: Option<String>,
|
||||
}
|
||||
|
||||
/// Represents an account token (from account_tokens table).
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, FromRow)]
|
||||
pub struct AccountToken {
|
||||
pub id: i64,
|
||||
pub account_id: i64,
|
||||
pub name: String,
|
||||
pub token: String,
|
||||
pub source: Option<String>,
|
||||
pub token_group: Option<String>,
|
||||
pub value_status: String,
|
||||
pub enabled: Option<bool>,
|
||||
pub is_default: Option<bool>,
|
||||
pub created_at: Option<String>,
|
||||
pub updated_at: Option<String>,
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sqlx::FromRow;
|
||||
|
||||
/// Represents a downstream API key for external consumers.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, FromRow)]
|
||||
pub struct DownstreamApiKey {
|
||||
pub id: i64,
|
||||
pub name: String,
|
||||
pub key: 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 used_cost: Option<f64>,
|
||||
pub max_requests: Option<i64>,
|
||||
pub used_requests: Option<i64>,
|
||||
pub supported_models: Option<String>,
|
||||
pub allowed_route_ids: Option<String>,
|
||||
pub site_weight_multipliers: Option<String>,
|
||||
pub excluded_site_ids: Option<String>,
|
||||
pub excluded_credential_refs: Option<String>,
|
||||
pub last_used_at: Option<String>,
|
||||
pub user_id: Option<i64>,
|
||||
pub created_at: Option<String>,
|
||||
pub updated_at: Option<String>,
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
// Database models
|
||||
|
||||
pub mod account;
|
||||
pub mod downstream_api_key;
|
||||
pub mod route_channel;
|
||||
pub mod setting;
|
||||
pub mod site;
|
||||
pub mod token_route;
|
||||
@@ -0,0 +1,27 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sqlx::FromRow;
|
||||
|
||||
/// Represents a route channel linking routes, accounts, and tokens.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, FromRow)]
|
||||
pub struct RouteChannel {
|
||||
pub id: i64,
|
||||
pub route_id: i64,
|
||||
pub account_id: i64,
|
||||
pub token_id: Option<i64>,
|
||||
pub oauth_route_unit_id: Option<i64>,
|
||||
pub source_model: Option<String>,
|
||||
pub priority: Option<i64>,
|
||||
pub weight: Option<i64>,
|
||||
pub enabled: Option<bool>,
|
||||
pub manual_override: Option<bool>,
|
||||
pub success_count: Option<i64>,
|
||||
pub fail_count: Option<i64>,
|
||||
pub total_latency_ms: Option<i64>,
|
||||
pub total_cost: Option<f64>,
|
||||
pub last_used_at: Option<String>,
|
||||
pub last_selected_at: Option<String>,
|
||||
pub last_fail_at: Option<String>,
|
||||
pub consecutive_fail_count: i64,
|
||||
pub cooldown_level: i64,
|
||||
pub cooldown_until: Option<String>,
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sqlx::FromRow;
|
||||
|
||||
/// Represents a key-value setting.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, FromRow)]
|
||||
pub struct Setting {
|
||||
pub key: String,
|
||||
pub value: Option<String>,
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sqlx::FromRow;
|
||||
|
||||
/// Represents a site (AI provider endpoint).
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, FromRow)]
|
||||
pub struct Site {
|
||||
pub id: i64,
|
||||
pub name: String,
|
||||
pub url: String,
|
||||
pub platform: String,
|
||||
pub status: String,
|
||||
pub api_key: Option<String>,
|
||||
pub proxy_url: Option<String>,
|
||||
pub use_system_proxy: Option<bool>,
|
||||
pub custom_headers: Option<String>,
|
||||
pub external_checkin_url: Option<String>,
|
||||
pub global_weight: Option<f64>,
|
||||
pub post_refresh_probe_enabled: Option<bool>,
|
||||
pub post_refresh_probe_model: Option<String>,
|
||||
pub post_refresh_probe_scope: Option<String>,
|
||||
pub post_refresh_probe_latency_threshold_ms: Option<i64>,
|
||||
pub is_pinned: Option<bool>,
|
||||
pub sort_order: Option<i64>,
|
||||
pub created_at: Option<String>,
|
||||
pub updated_at: Option<String>,
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sqlx::FromRow;
|
||||
|
||||
/// Represents a token route (model routing rule).
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, FromRow)]
|
||||
pub struct TokenRoute {
|
||||
pub id: i64,
|
||||
pub model_pattern: String,
|
||||
pub model_mapping: Option<String>,
|
||||
pub display_name: Option<String>,
|
||||
pub display_icon: Option<String>,
|
||||
pub decision_snapshot: Option<String>,
|
||||
pub decision_refreshed_at: Option<String>,
|
||||
pub routing_strategy: Option<String>,
|
||||
pub route_mode: Option<String>,
|
||||
pub enabled: Option<bool>,
|
||||
pub created_at: Option<String>,
|
||||
pub updated_at: Option<String>,
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
// Database queries
|
||||
@@ -0,0 +1,59 @@
|
||||
use axum::{
|
||||
http::StatusCode,
|
||||
response::{IntoResponse, Json, Response},
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum AppError {
|
||||
#[error("Not found: {0}")]
|
||||
NotFound(String),
|
||||
|
||||
#[error("Unauthorized: {0}")]
|
||||
Unauthorized(String),
|
||||
|
||||
#[error("Forbidden: {0}")]
|
||||
Forbidden(String),
|
||||
|
||||
#[error("Bad request: {0}")]
|
||||
BadRequest(String),
|
||||
|
||||
#[error("Upstream error: {0}")]
|
||||
UpstreamError(String),
|
||||
|
||||
#[error("Rate limited")]
|
||||
RateLimited,
|
||||
|
||||
#[error("Internal error: {0}")]
|
||||
Internal(#[from] anyhow::Error),
|
||||
|
||||
#[error("Database error: {0}")]
|
||||
Database(#[from] sqlx::Error),
|
||||
}
|
||||
|
||||
impl IntoResponse for AppError {
|
||||
fn into_response(self) -> Response {
|
||||
let (status, message) = match &self {
|
||||
AppError::NotFound(msg) => (StatusCode::NOT_FOUND, msg.clone()),
|
||||
AppError::Unauthorized(msg) => (StatusCode::UNAUTHORIZED, msg.clone()),
|
||||
AppError::Forbidden(msg) => (StatusCode::FORBIDDEN, msg.clone()),
|
||||
AppError::BadRequest(msg) => (StatusCode::BAD_REQUEST, msg.clone()),
|
||||
AppError::UpstreamError(msg) => (StatusCode::BAD_GATEWAY, msg.clone()),
|
||||
AppError::RateLimited => (StatusCode::TOO_MANY_REQUESTS, "rate_limited".into()),
|
||||
AppError::Internal(e) => (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()),
|
||||
AppError::Database(e) => (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()),
|
||||
};
|
||||
(
|
||||
status,
|
||||
Json(json!({
|
||||
"error": {
|
||||
"message": message,
|
||||
"type": status.canonical_reason(),
|
||||
}
|
||||
})),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
}
|
||||
|
||||
pub type AppResult<T> = Result<T, AppError>;
|
||||
@@ -0,0 +1,27 @@
|
||||
pub mod auth;
|
||||
pub mod config;
|
||||
pub mod db;
|
||||
pub mod error;
|
||||
pub mod proxy;
|
||||
pub mod routes;
|
||||
pub mod services;
|
||||
pub mod shared;
|
||||
pub mod transformers;
|
||||
|
||||
use sqlx::SqlitePool;
|
||||
use std::sync::Arc;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct AppState {
|
||||
pub config: config::AppConfig,
|
||||
pub db: Arc<SqlitePool>,
|
||||
}
|
||||
|
||||
impl AppState {
|
||||
pub fn new(config: config::AppConfig, db: SqlitePool) -> Self {
|
||||
Self {
|
||||
config,
|
||||
db: Arc::new(db),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
use axum::serve;
|
||||
use boosapi::{config::AppConfig, db, AppState};
|
||||
use std::net::SocketAddr;
|
||||
use tokio::signal;
|
||||
use tracing_subscriber::EnvFilter;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
// Initialize tracing
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(EnvFilter::from_default_env())
|
||||
.init();
|
||||
|
||||
// Load config
|
||||
let config = AppConfig::load()?;
|
||||
tracing::info!("Loaded configuration");
|
||||
|
||||
// Create database pool
|
||||
let pool = db::create_pool(&config.db.url, config.db.max_connections).await?;
|
||||
tracing::info!("Database connected");
|
||||
|
||||
// Run migrations
|
||||
sqlx::migrate!("./migrations").run(&pool).await?;
|
||||
tracing::info!("Database migrations applied");
|
||||
|
||||
// Build app state
|
||||
let state = AppState::new(config.clone(), pool);
|
||||
|
||||
// Build router
|
||||
let app = boosapi::routes::build_router(state);
|
||||
|
||||
// Start server
|
||||
let addr = SocketAddr::new(config.server.host.parse()?, config.server.port);
|
||||
tracing::info!("Starting server on {addr}");
|
||||
|
||||
serve(
|
||||
tokio::net::TcpListener::bind(addr).await?,
|
||||
app.into_make_service(),
|
||||
)
|
||||
.with_graceful_shutdown(shutdown_signal())
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn shutdown_signal() {
|
||||
let ctrl_c = async {
|
||||
signal::ctrl_c()
|
||||
.await
|
||||
.expect("failed to install Ctrl+C handler");
|
||||
};
|
||||
|
||||
#[cfg(unix)]
|
||||
let terminate = async {
|
||||
signal::unix::signal(signal::unix::SignalKind::terminate())
|
||||
.expect("failed to install SIGTERM handler")
|
||||
.recv()
|
||||
.await;
|
||||
};
|
||||
|
||||
#[cfg(not(unix))]
|
||||
let terminate = std::future::pending::<()>();
|
||||
|
||||
tokio::select! {
|
||||
_ = ctrl_c => {},
|
||||
_ = terminate => {},
|
||||
}
|
||||
|
||||
tracing::info!("Shutdown signal received, starting graceful shutdown");
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
// Proxy capabilities
|
||||
@@ -0,0 +1 @@
|
||||
// Proxy CLI profiles
|
||||
@@ -0,0 +1,157 @@
|
||||
// Proxy executors — forward requests to upstream APIs
|
||||
//
|
||||
// Each executor:
|
||||
// 1. Transforms the OpenAI-compatible request to the upstream format
|
||||
// 2. Sends the HTTP request
|
||||
// 3. Transforms the response back to OpenAI-compatible format
|
||||
|
||||
use crate::error::{ AppError, AppResult };
|
||||
use crate::AppState;
|
||||
use reqwest::Client;
|
||||
use serde_json::Value;
|
||||
use std::time::Duration;
|
||||
|
||||
/// Resolved upstream info for a given model request.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct UpstreamRoute {
|
||||
pub base_url: String,
|
||||
pub api_key: String,
|
||||
pub model: String,
|
||||
}
|
||||
|
||||
/// Create a shared reqwest client with sensible defaults.
|
||||
fn http_client() -> Client {
|
||||
Client::builder()
|
||||
.timeout(Duration::from_secs(300))
|
||||
.pool_max_idle_per_host(32)
|
||||
.build()
|
||||
.expect("Failed to build HTTP client")
|
||||
}
|
||||
|
||||
/// Execute a non-streaming chat completions request.
|
||||
pub async fn execute_chat_completions(
|
||||
_state: AppState,
|
||||
route: UpstreamRoute,
|
||||
body: Value,
|
||||
) -> AppResult<Value> {
|
||||
let client = http_client();
|
||||
let upstream_url = format!("{}/chat/completions", route.base_url.trim_end_matches('/'));
|
||||
|
||||
let resp = client
|
||||
.post(&upstream_url)
|
||||
.header("Authorization", format!("Bearer {}", route.api_key))
|
||||
.header("Content-Type", "application/json")
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| AppError::UpstreamError(format!("Upstream request failed: {e}")))?;
|
||||
|
||||
let status = resp.status();
|
||||
let text = resp
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| AppError::UpstreamError(format!("Failed to read upstream response: {e}")))?;
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(AppError::UpstreamError(format!(
|
||||
"Upstream returned {status}: {text}"
|
||||
)));
|
||||
}
|
||||
|
||||
let json: Value = serde_json::from_str(&text)
|
||||
.map_err(|e| AppError::UpstreamError(format!("Invalid JSON from upstream: {e}")))?;
|
||||
|
||||
Ok(json)
|
||||
}
|
||||
|
||||
/// Execute a streaming chat completions request.
|
||||
/// Returns a vector of SSE data chunks as strings.
|
||||
pub async fn execute_chat_completions_stream(
|
||||
_state: AppState,
|
||||
route: UpstreamRoute,
|
||||
body: Value,
|
||||
) -> AppResult<Vec<String>> {
|
||||
let client = http_client();
|
||||
let upstream_url = format!("{}/chat/completions", route.base_url.trim_end_matches('/'));
|
||||
|
||||
let resp = client
|
||||
.post(&upstream_url)
|
||||
.header("Authorization", format!("Bearer {}", route.api_key))
|
||||
.header("Content-Type", "application/json")
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| AppError::UpstreamError(format!("Upstream request failed: {e}")))?;
|
||||
|
||||
let status = resp.status();
|
||||
if !status.is_success() {
|
||||
let text = resp
|
||||
.text()
|
||||
.await
|
||||
.unwrap_or_else(|_| "unknown error".to_string());
|
||||
return Err(AppError::UpstreamError(format!(
|
||||
"Upstream returned {status}: {text}"
|
||||
)));
|
||||
}
|
||||
|
||||
// Read the streaming response as chunks
|
||||
let mut chunks = Vec::new();
|
||||
let mut stream = resp.bytes_stream();
|
||||
use futures::StreamExt;
|
||||
while let Some(chunk_result) = stream.next().await {
|
||||
match chunk_result {
|
||||
Ok(bytes) => {
|
||||
let text = String::from_utf8_lossy(&bytes).to_string();
|
||||
// Parse SSE lines from the chunk
|
||||
for line in text.lines() {
|
||||
if line.starts_with("data: ") {
|
||||
chunks.push(line.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
return Err(AppError::UpstreamError(format!(
|
||||
"Stream error: {e}"
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(chunks)
|
||||
}
|
||||
|
||||
/// Execute an image generation request.
|
||||
pub async fn execute_image_generation(
|
||||
_state: AppState,
|
||||
route: UpstreamRoute,
|
||||
body: Value,
|
||||
) -> AppResult<Value> {
|
||||
let client = http_client();
|
||||
let upstream_url = format!("{}/images/generations", route.base_url.trim_end_matches('/'));
|
||||
|
||||
let resp = client
|
||||
.post(&upstream_url)
|
||||
.header("Authorization", format!("Bearer {}", route.api_key))
|
||||
.header("Content-Type", "application/json")
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| AppError::UpstreamError(format!("Upstream request failed: {e}")))?;
|
||||
|
||||
let status = resp.status();
|
||||
let text = resp
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| AppError::UpstreamError(format!("Failed to read upstream response: {e}")))?;
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(AppError::UpstreamError(format!(
|
||||
"Upstream returned {status}: {text}"
|
||||
)));
|
||||
}
|
||||
|
||||
let json: Value = serde_json::from_str(&text)
|
||||
.map_err(|e| AppError::UpstreamError(format!("Invalid JSON from upstream: {e}")))?;
|
||||
|
||||
Ok(json)
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
// Proxy core — OpenAI-compatible request routing
|
||||
//
|
||||
// Architecture:
|
||||
// /v1/chat/completions → auth (downstream key) → route resolution → upstream call → response
|
||||
// /v1/images/generations → ...same flow...
|
||||
// /v1/audio/speech → ...same flow...
|
||||
|
||||
pub mod executors;
|
||||
pub mod providers;
|
||||
pub mod runtime;
|
||||
pub mod surfaces;
|
||||
@@ -0,0 +1 @@
|
||||
// Proxy providers
|
||||
@@ -0,0 +1 @@
|
||||
// Proxy runtime
|
||||
@@ -0,0 +1 @@
|
||||
// Proxy surfaces
|
||||
@@ -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
|
||||
@@ -0,0 +1 @@
|
||||
// Services - announcement
|
||||
@@ -0,0 +1 @@
|
||||
// Services - cron
|
||||
@@ -0,0 +1 @@
|
||||
// Services module
|
||||
@@ -0,0 +1 @@
|
||||
// Services - OAuth
|
||||
@@ -0,0 +1 @@
|
||||
// Services - shared
|
||||
@@ -0,0 +1 @@
|
||||
// Services - update center
|
||||
@@ -0,0 +1,90 @@
|
||||
use serde::de::DeserializeOwned;
|
||||
use serde::Serialize;
|
||||
use serde_json::Value;
|
||||
|
||||
/// Shortcut for `serde_json::from_value` with a descriptive error type.
|
||||
pub fn from_value<T: DeserializeOwned>(value: Value) -> Result<T, serde_json::Error> {
|
||||
serde_json::from_value(value)
|
||||
}
|
||||
|
||||
/// Shortcut for `serde_json::to_value` with a descriptive error type.
|
||||
pub fn to_value<T: Serialize>(value: &T) -> Result<Value, serde_json::Error> {
|
||||
serde_json::to_value(value)
|
||||
}
|
||||
|
||||
/// Pretty-print a JSON value as a string.
|
||||
pub fn pretty(value: &Value) -> String {
|
||||
serde_json::to_string_pretty(value).unwrap_or_default()
|
||||
}
|
||||
|
||||
/// Compact JSON string from a serializable value.
|
||||
pub fn to_string<T: Serialize>(value: &T) -> Result<String, serde_json::Error> {
|
||||
serde_json::to_string(value)
|
||||
}
|
||||
|
||||
/// Parse a JSON string into a value.
|
||||
pub fn from_str<T: DeserializeOwned>(s: &str) -> Result<T, serde_json::Error> {
|
||||
serde_json::from_str(s)
|
||||
}
|
||||
|
||||
/// Check whether a JSON value is null or absent.
|
||||
pub fn is_null(value: &Value) -> bool {
|
||||
value.is_null()
|
||||
}
|
||||
|
||||
/// Get a string field from a JSON object, returning None if missing or wrong type.
|
||||
pub fn get_string(value: &Value, key: &str) -> Option<String> {
|
||||
value.get(key).and_then(|v| v.as_str()).map(String::from)
|
||||
}
|
||||
|
||||
/// Get an i64 field from a JSON object.
|
||||
pub fn get_i64(value: &Value, key: &str) -> Option<i64> {
|
||||
value.get(key).and_then(|v| v.as_i64())
|
||||
}
|
||||
|
||||
/// Get a f64 field from a JSON object.
|
||||
pub fn get_f64(value: &Value, key: &str) -> Option<f64> {
|
||||
value.get(key).and_then(|v| v.as_f64())
|
||||
}
|
||||
|
||||
/// Get a bool field from a JSON object.
|
||||
pub fn get_bool(value: &Value, key: &str) -> Option<bool> {
|
||||
value.get(key).and_then(|v| v.as_bool())
|
||||
}
|
||||
|
||||
/// Get a nested value by a dotted path (e.g. "a.b.c").
|
||||
pub fn get_nested<'a>(value: &'a Value, path: &str) -> Option<&'a Value> {
|
||||
let mut current = value;
|
||||
for segment in path.split('.') {
|
||||
current = current.get(segment)?;
|
||||
}
|
||||
Some(current)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn test_get_string() {
|
||||
let v = json!({"name": "hello"});
|
||||
assert_eq!(get_string(&v, "name"), Some("hello".into()));
|
||||
assert_eq!(get_string(&v, "missing"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_nested() {
|
||||
let v = json!({"a": {"b": {"c": 42}}});
|
||||
assert_eq!(get_nested(&v, "a.b.c"), Some(&json!(42)));
|
||||
assert_eq!(get_nested(&v, "a.b.missing"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_roundtrip() {
|
||||
let val = json!({"x": 1, "y": "z"});
|
||||
let s = to_string(&val).unwrap();
|
||||
let back: Value = from_str(&s).unwrap();
|
||||
assert_eq!(val, back);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
// Shared utilities
|
||||
|
||||
pub mod json;
|
||||
pub mod pagination;
|
||||
pub mod sse;
|
||||
pub mod time;
|
||||
|
||||
pub use json::*;
|
||||
pub use pagination::*;
|
||||
pub use sse::*;
|
||||
pub use time::*;
|
||||
@@ -0,0 +1,71 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Query parameters for paginated list endpoints.
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct PaginationParams {
|
||||
pub page: Option<i64>,
|
||||
pub page_size: Option<i64>,
|
||||
pub sort_by: Option<String>,
|
||||
pub sort_order: Option<String>,
|
||||
}
|
||||
|
||||
impl PaginationParams {
|
||||
/// Returns the page number (1-indexed), defaulting to 1.
|
||||
pub fn page(&self) -> i64 {
|
||||
self.page.filter(|&p| p > 0).unwrap_or(1)
|
||||
}
|
||||
|
||||
/// Returns the page size, defaulting to 20, clamped to [1, 100].
|
||||
pub fn page_size(&self) -> i64 {
|
||||
self.page_size
|
||||
.filter(|&s| s > 0)
|
||||
.map(|s| s.min(100))
|
||||
.unwrap_or(20)
|
||||
}
|
||||
|
||||
/// Returns the offset for SQL LIMIT/OFFSET queries.
|
||||
pub fn offset(&self) -> i64 {
|
||||
(self.page() - 1) * self.page_size()
|
||||
}
|
||||
|
||||
/// Returns the sort column, defaulting to "id".
|
||||
pub fn sort_by(&self) -> &str {
|
||||
self.sort_by.as_deref().unwrap_or("id")
|
||||
}
|
||||
|
||||
/// Returns the sort direction, defaulting to "asc".
|
||||
pub fn sort_order(&self) -> &str {
|
||||
self.sort_order.as_deref().unwrap_or("asc")
|
||||
}
|
||||
}
|
||||
|
||||
/// Standard paginated response wrapper.
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct PaginatedResponse<T: Serialize> {
|
||||
pub data: Vec<T>,
|
||||
pub total: i64,
|
||||
pub page: i64,
|
||||
pub page_size: i64,
|
||||
pub total_pages: i64,
|
||||
}
|
||||
|
||||
impl<T: Serialize> PaginatedResponse<T> {
|
||||
/// Build a paginated response from data, total count, and query params.
|
||||
pub fn new(data: Vec<T>, total: i64, params: &PaginationParams) -> Self {
|
||||
let page = params.page();
|
||||
let page_size = params.page_size();
|
||||
let total_pages = if total == 0 {
|
||||
0
|
||||
} else {
|
||||
(total as f64 / page_size as f64).ceil() as i64
|
||||
};
|
||||
|
||||
Self {
|
||||
data,
|
||||
total,
|
||||
page,
|
||||
page_size,
|
||||
total_pages,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,51 @@
|
||||
use axum::response::sse::Event;
|
||||
use serde::Serialize;
|
||||
use std::convert::Infallible;
|
||||
|
||||
/// Create an SSE event from a serializable data type.
|
||||
pub fn sse_event<T: Serialize>(event_type: &str, data: &T) -> Result<Event, Infallible> {
|
||||
Ok(Event::default()
|
||||
.event(event_type)
|
||||
.data(serde_json::to_string(data).unwrap_or_default()))
|
||||
}
|
||||
|
||||
/// Angles for character prompts (front, back, left, right).
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct Angles {
|
||||
pub front: String,
|
||||
pub back: String,
|
||||
pub left: String,
|
||||
pub right: String,
|
||||
}
|
||||
|
||||
/// SSE event data discriminated by the `type` field.
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(tag = "type")]
|
||||
pub enum SseData {
|
||||
Text {
|
||||
content: String,
|
||||
},
|
||||
CharacterPrompts {
|
||||
character: String,
|
||||
angles: Angles,
|
||||
},
|
||||
Image {
|
||||
url: String,
|
||||
character: String,
|
||||
angle: String,
|
||||
},
|
||||
Scene {
|
||||
id: u32,
|
||||
description: String,
|
||||
prompt: String,
|
||||
},
|
||||
Tts {
|
||||
character: String,
|
||||
text: String,
|
||||
url: String,
|
||||
},
|
||||
ComfyuiWorkflow {
|
||||
json: serde_json::Value,
|
||||
},
|
||||
Done,
|
||||
}
|
||||
@@ -0,0 +1,92 @@
|
||||
use chrono::{DateTime, NaiveDateTime, Utc};
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
/// Returns the current UTC time as a `DateTime<Utc>`.
|
||||
pub fn now() -> DateTime<Utc> {
|
||||
Utc::now()
|
||||
}
|
||||
|
||||
/// Returns the current Unix timestamp in seconds.
|
||||
pub fn timestamp() -> i64 {
|
||||
Utc::now().timestamp()
|
||||
}
|
||||
|
||||
/// Returns the current Unix timestamp in milliseconds.
|
||||
pub fn timestamp_ms() -> i64 {
|
||||
Utc::now().timestamp_millis()
|
||||
}
|
||||
|
||||
/// Format a `DateTime<Utc>` as an RFC 3339 / ISO 8601 string.
|
||||
pub fn format_rfc3339(dt: &DateTime<Utc>) -> String {
|
||||
dt.to_rfc3339()
|
||||
}
|
||||
|
||||
/// Format a `DateTime<Utc>` as a human-friendly string.
|
||||
pub fn format_human(dt: &DateTime<Utc>) -> String {
|
||||
dt.format("%Y-%m-%d %H:%M:%S UTC").to_string()
|
||||
}
|
||||
|
||||
/// Parse an RFC 3339 / ISO 8601 string into a `DateTime<Utc>`.
|
||||
pub fn parse_rfc3339(s: &str) -> Result<DateTime<Utc>, chrono::ParseError> {
|
||||
s.parse::<DateTime<Utc>>()
|
||||
}
|
||||
|
||||
/// Convert a Unix timestamp (seconds) to `DateTime<Utc>`.
|
||||
pub fn from_timestamp(secs: i64) -> Option<DateTime<Utc>> {
|
||||
DateTime::from_timestamp(secs, 0)
|
||||
}
|
||||
|
||||
/// Convert a `SystemTime` to `DateTime<Utc>`.
|
||||
pub fn from_system_time(t: SystemTime) -> Option<DateTime<Utc>> {
|
||||
match t.duration_since(UNIX_EPOCH) {
|
||||
Ok(d) => DateTime::from_timestamp(d.as_secs() as i64, d.subsec_nanos()),
|
||||
Err(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert a `NaiveDateTime` (assumed UTC) to `DateTime<Utc>`.
|
||||
pub fn from_naive_utc(dt: NaiveDateTime) -> DateTime<Utc> {
|
||||
DateTime::from_naive_utc_and_offset(dt, Utc)
|
||||
}
|
||||
|
||||
/// Compute the duration between two timestamps.
|
||||
pub fn duration_secs(start: DateTime<Utc>, end: DateTime<Utc>) -> f64 {
|
||||
(end - start).num_milliseconds() as f64 / 1000.0
|
||||
}
|
||||
|
||||
/// Check if a timestamp is older than the given duration.
|
||||
pub fn is_older_than(dt: DateTime<Utc>, duration: Duration) -> bool {
|
||||
let cutoff = Utc::now() - chrono::Duration::from_std(duration).unwrap_or_default();
|
||||
dt < cutoff
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_now_and_timestamp() {
|
||||
let ts = timestamp();
|
||||
assert!(ts > 1_700_000_000); // sanity check for 2024+
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_format_roundtrip() {
|
||||
let dt = now();
|
||||
let s = format_rfc3339(&dt);
|
||||
let parsed = parse_rfc3339(&s).unwrap();
|
||||
assert_eq!(dt.timestamp(), parsed.timestamp());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_from_timestamp() {
|
||||
let dt = from_timestamp(0).unwrap();
|
||||
assert_eq!(dt.timestamp(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_older_than() {
|
||||
let old = Utc::now() - chrono::Duration::hours(2);
|
||||
assert!(is_older_than(old, Duration::from_secs(3600)));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
// Transformers - Anthropic Messages
|
||||
@@ -0,0 +1 @@
|
||||
// Transformers - Anthropic
|
||||
@@ -0,0 +1 @@
|
||||
// Transformers - canonical
|
||||
@@ -0,0 +1 @@
|
||||
// Transformers - Gemini Generate Content
|
||||
@@ -0,0 +1 @@
|
||||
// Transformers - Gemini
|
||||
@@ -0,0 +1 @@
|
||||
// Transformers module
|
||||
@@ -0,0 +1 @@
|
||||
// Transformers - OpenAI Chat
|
||||
@@ -0,0 +1 @@
|
||||
// Transformers - OpenAI
|
||||
@@ -0,0 +1 @@
|
||||
// Transformers - OpenAI Responses
|
||||
@@ -0,0 +1 @@
|
||||
// Transformers - shared
|
||||
Reference in New Issue
Block a user