hobokenchicken
9b8483e797
This commit introduces: - AES-256-GCM encryption for LLM provider API keys in the database. - HMAC-SHA256 signed session tokens with activity-based refresh logic. - Standardized frontend XSS protection using a global escapeHtml utility. - Hardened security headers and request body size limits. - Improved database integrity with foreign key enforcement and atomic transactions. - Integration tests for the full encrypted key storage and proxy usage lifecycle.
519 lines
18 KiB
Rust
519 lines
18 KiB
Rust
use axum::{
|
|
extract::{Path, State},
|
|
response::Json,
|
|
};
|
|
use chrono;
|
|
use rand::Rng;
|
|
use serde::Deserialize;
|
|
use serde_json;
|
|
use sqlx::Row;
|
|
use tracing::warn;
|
|
use uuid;
|
|
|
|
use super::{ApiResponse, DashboardState};
|
|
|
|
/// Generate a random API token: sk-{48 hex chars}
|
|
fn generate_token() -> String {
|
|
let mut rng = rand::rng();
|
|
let bytes: Vec<u8> = (0..24).map(|_| rng.random::<u8>()).collect();
|
|
format!("sk-{}", hex::encode(bytes))
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
pub(super) struct CreateClientRequest {
|
|
pub(super) name: String,
|
|
pub(super) client_id: Option<String>,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
pub(super) struct UpdateClientPayload {
|
|
pub(super) name: Option<String>,
|
|
pub(super) description: Option<String>,
|
|
pub(super) is_active: Option<bool>,
|
|
pub(super) rate_limit_per_minute: Option<i64>,
|
|
}
|
|
|
|
pub(super) async fn handle_get_clients(State(state): State<DashboardState>) -> Json<ApiResponse<serde_json::Value>> {
|
|
let pool = &state.app_state.db_pool;
|
|
|
|
let result = sqlx::query(
|
|
r#"
|
|
SELECT
|
|
client_id as id,
|
|
name,
|
|
description,
|
|
created_at,
|
|
total_requests,
|
|
total_tokens,
|
|
total_cost,
|
|
is_active,
|
|
rate_limit_per_minute
|
|
FROM clients
|
|
ORDER BY created_at DESC
|
|
"#,
|
|
)
|
|
.fetch_all(pool)
|
|
.await;
|
|
|
|
match result {
|
|
Ok(rows) => {
|
|
let clients: Vec<serde_json::Value> = rows
|
|
.into_iter()
|
|
.map(|row| {
|
|
serde_json::json!({
|
|
"id": row.get::<String, _>("id"),
|
|
"name": row.get::<Option<String>, _>("name").unwrap_or_else(|| "Unnamed".to_string()),
|
|
"description": row.get::<Option<String>, _>("description"),
|
|
"created_at": row.get::<chrono::DateTime<chrono::Utc>, _>("created_at"),
|
|
"requests_count": row.get::<i64, _>("total_requests"),
|
|
"total_tokens": row.get::<i64, _>("total_tokens"),
|
|
"total_cost": row.get::<f64, _>("total_cost"),
|
|
"status": if row.get::<bool, _>("is_active") { "active" } else { "inactive" },
|
|
"rate_limit_per_minute": row.get::<Option<i64>, _>("rate_limit_per_minute"),
|
|
})
|
|
})
|
|
.collect();
|
|
|
|
Json(ApiResponse::success(serde_json::json!(clients)))
|
|
}
|
|
Err(e) => {
|
|
warn!("Failed to fetch clients: {}", e);
|
|
Json(ApiResponse::error("Failed to fetch clients".to_string()))
|
|
}
|
|
}
|
|
}
|
|
|
|
pub(super) async fn handle_create_client(
|
|
State(state): State<DashboardState>,
|
|
headers: axum::http::HeaderMap,
|
|
Json(payload): Json<CreateClientRequest>,
|
|
) -> Json<ApiResponse<serde_json::Value>> {
|
|
let (session, _) = match super::auth::require_admin(&state, &headers).await {
|
|
Ok((session, new_token)) => (session, new_token),
|
|
Err(e) => return e,
|
|
};
|
|
|
|
let pool = &state.app_state.db_pool;
|
|
|
|
let client_id = payload
|
|
.client_id
|
|
.unwrap_or_else(|| format!("client-{}", &uuid::Uuid::new_v4().to_string()[..8]));
|
|
|
|
let result = sqlx::query(
|
|
r#"
|
|
INSERT INTO clients (client_id, name, is_active)
|
|
VALUES (?, ?, TRUE)
|
|
RETURNING *
|
|
"#,
|
|
)
|
|
.bind(&client_id)
|
|
.bind(&payload.name)
|
|
.fetch_one(pool)
|
|
.await;
|
|
|
|
match result {
|
|
Ok(row) => {
|
|
// Auto-generate a token for the new client
|
|
let token = generate_token();
|
|
let token_result = sqlx::query(
|
|
"INSERT INTO client_tokens (client_id, token, name) VALUES (?, ?, 'default')",
|
|
)
|
|
.bind(&client_id)
|
|
.bind(&token)
|
|
.execute(pool)
|
|
.await;
|
|
|
|
if let Err(e) = token_result {
|
|
warn!("Client created but failed to generate token: {}", e);
|
|
}
|
|
|
|
Json(ApiResponse::success(serde_json::json!({
|
|
"id": row.get::<String, _>("client_id"),
|
|
"name": row.get::<Option<String>, _>("name"),
|
|
"created_at": row.get::<chrono::DateTime<chrono::Utc>, _>("created_at"),
|
|
"status": "active",
|
|
"token": token,
|
|
})))
|
|
}
|
|
Err(e) => {
|
|
warn!("Failed to create client: {}", e);
|
|
Json(ApiResponse::error(format!("Failed to create client: {}", e)))
|
|
}
|
|
}
|
|
}
|
|
|
|
pub(super) async fn handle_get_client(
|
|
State(state): State<DashboardState>,
|
|
Path(id): Path<String>,
|
|
) -> Json<ApiResponse<serde_json::Value>> {
|
|
let pool = &state.app_state.db_pool;
|
|
|
|
let result = sqlx::query(
|
|
r#"
|
|
SELECT
|
|
c.client_id as id,
|
|
c.name,
|
|
c.description,
|
|
c.is_active,
|
|
c.rate_limit_per_minute,
|
|
c.created_at,
|
|
COALESCE(c.total_tokens, 0) as total_tokens,
|
|
COALESCE(c.total_cost, 0.0) as total_cost,
|
|
COUNT(r.id) as total_requests,
|
|
MAX(r.timestamp) as last_request
|
|
FROM clients c
|
|
LEFT JOIN llm_requests r ON c.client_id = r.client_id
|
|
WHERE c.client_id = ?
|
|
GROUP BY c.client_id
|
|
"#,
|
|
)
|
|
.bind(&id)
|
|
.fetch_optional(pool)
|
|
.await;
|
|
|
|
match result {
|
|
Ok(Some(row)) => Json(ApiResponse::success(serde_json::json!({
|
|
"id": row.get::<String, _>("id"),
|
|
"name": row.get::<Option<String>, _>("name").unwrap_or_else(|| "Unnamed".to_string()),
|
|
"description": row.get::<Option<String>, _>("description"),
|
|
"is_active": row.get::<bool, _>("is_active"),
|
|
"rate_limit_per_minute": row.get::<Option<i64>, _>("rate_limit_per_minute"),
|
|
"created_at": row.get::<chrono::DateTime<chrono::Utc>, _>("created_at"),
|
|
"total_tokens": row.get::<i64, _>("total_tokens"),
|
|
"total_cost": row.get::<f64, _>("total_cost"),
|
|
"total_requests": row.get::<i64, _>("total_requests"),
|
|
"last_request": row.get::<Option<chrono::DateTime<chrono::Utc>>, _>("last_request"),
|
|
"status": if row.get::<bool, _>("is_active") { "active" } else { "inactive" },
|
|
}))),
|
|
Ok(None) => Json(ApiResponse::error(format!("Client '{}' not found", id))),
|
|
Err(e) => {
|
|
warn!("Failed to fetch client: {}", e);
|
|
Json(ApiResponse::error(format!("Failed to fetch client: {}", e)))
|
|
}
|
|
}
|
|
}
|
|
|
|
pub(super) async fn handle_update_client(
|
|
State(state): State<DashboardState>,
|
|
headers: axum::http::HeaderMap,
|
|
Path(id): Path<String>,
|
|
Json(payload): Json<UpdateClientPayload>,
|
|
) -> Json<ApiResponse<serde_json::Value>> {
|
|
let (session, _) = match super::auth::require_admin(&state, &headers).await {
|
|
Ok((session, new_token)) => (session, new_token),
|
|
Err(e) => return e,
|
|
};
|
|
|
|
let pool = &state.app_state.db_pool;
|
|
|
|
// Build dynamic UPDATE query from provided fields
|
|
let mut sets = Vec::new();
|
|
let mut binds: Vec<String> = Vec::new();
|
|
|
|
if let Some(ref name) = payload.name {
|
|
sets.push("name = ?");
|
|
binds.push(name.clone());
|
|
}
|
|
if let Some(ref desc) = payload.description {
|
|
sets.push("description = ?");
|
|
binds.push(desc.clone());
|
|
}
|
|
if payload.is_active.is_some() {
|
|
sets.push("is_active = ?");
|
|
}
|
|
if payload.rate_limit_per_minute.is_some() {
|
|
sets.push("rate_limit_per_minute = ?");
|
|
}
|
|
|
|
if sets.is_empty() {
|
|
return Json(ApiResponse::error("No fields to update".to_string()));
|
|
}
|
|
|
|
// Always update updated_at
|
|
sets.push("updated_at = CURRENT_TIMESTAMP");
|
|
|
|
let sql = format!("UPDATE clients SET {} WHERE client_id = ?", sets.join(", "));
|
|
let mut query = sqlx::query(&sql);
|
|
|
|
// Bind in the same order as sets
|
|
for b in &binds {
|
|
query = query.bind(b);
|
|
}
|
|
if let Some(active) = payload.is_active {
|
|
query = query.bind(active);
|
|
}
|
|
if let Some(rate) = payload.rate_limit_per_minute {
|
|
query = query.bind(rate);
|
|
}
|
|
query = query.bind(&id);
|
|
|
|
match query.execute(pool).await {
|
|
Ok(result) => {
|
|
if result.rows_affected() == 0 {
|
|
return Json(ApiResponse::error(format!("Client '{}' not found", id)));
|
|
}
|
|
// Return the updated client
|
|
let row = sqlx::query(
|
|
r#"
|
|
SELECT client_id as id, name, description, is_active, rate_limit_per_minute,
|
|
created_at, total_requests, total_tokens, total_cost
|
|
FROM clients WHERE client_id = ?
|
|
"#,
|
|
)
|
|
.bind(&id)
|
|
.fetch_one(pool)
|
|
.await;
|
|
|
|
match row {
|
|
Ok(row) => Json(ApiResponse::success(serde_json::json!({
|
|
"id": row.get::<String, _>("id"),
|
|
"name": row.get::<Option<String>, _>("name").unwrap_or_else(|| "Unnamed".to_string()),
|
|
"description": row.get::<Option<String>, _>("description"),
|
|
"is_active": row.get::<bool, _>("is_active"),
|
|
"rate_limit_per_minute": row.get::<Option<i64>, _>("rate_limit_per_minute"),
|
|
"created_at": row.get::<chrono::DateTime<chrono::Utc>, _>("created_at"),
|
|
"total_requests": row.get::<i64, _>("total_requests"),
|
|
"total_tokens": row.get::<i64, _>("total_tokens"),
|
|
"total_cost": row.get::<f64, _>("total_cost"),
|
|
"status": if row.get::<bool, _>("is_active") { "active" } else { "inactive" },
|
|
}))),
|
|
Err(e) => {
|
|
warn!("Failed to fetch updated client: {}", e);
|
|
// Update succeeded but fetch failed — still report success
|
|
Json(ApiResponse::success(serde_json::json!({ "message": "Client updated" })))
|
|
}
|
|
}
|
|
}
|
|
Err(e) => {
|
|
warn!("Failed to update client: {}", e);
|
|
Json(ApiResponse::error(format!("Failed to update client: {}", e)))
|
|
}
|
|
}
|
|
}
|
|
|
|
pub(super) async fn handle_delete_client(
|
|
State(state): State<DashboardState>,
|
|
headers: axum::http::HeaderMap,
|
|
Path(id): Path<String>,
|
|
) -> Json<ApiResponse<serde_json::Value>> {
|
|
let (session, _) = match super::auth::require_admin(&state, &headers).await {
|
|
Ok((session, new_token)) => (session, new_token),
|
|
Err(e) => return e,
|
|
};
|
|
|
|
let pool = &state.app_state.db_pool;
|
|
|
|
// Don't allow deleting the default client
|
|
if id == "default" {
|
|
return Json(ApiResponse::error("Cannot delete default client".to_string()));
|
|
}
|
|
|
|
let result = sqlx::query("DELETE FROM clients WHERE client_id = ?")
|
|
.bind(id)
|
|
.execute(pool)
|
|
.await;
|
|
|
|
match result {
|
|
Ok(_) => Json(ApiResponse::success(serde_json::json!({ "message": "Client deleted" }))),
|
|
Err(e) => Json(ApiResponse::error(format!("Failed to delete client: {}", e))),
|
|
}
|
|
}
|
|
|
|
pub(super) async fn handle_client_usage(
|
|
State(state): State<DashboardState>,
|
|
Path(id): Path<String>,
|
|
) -> Json<ApiResponse<serde_json::Value>> {
|
|
let pool = &state.app_state.db_pool;
|
|
|
|
// Get per-model breakdown for this client
|
|
let result = sqlx::query(
|
|
r#"
|
|
SELECT
|
|
model,
|
|
provider,
|
|
COUNT(*) as request_count,
|
|
SUM(prompt_tokens) as prompt_tokens,
|
|
SUM(completion_tokens) as completion_tokens,
|
|
SUM(total_tokens) as total_tokens,
|
|
SUM(cost) as total_cost,
|
|
AVG(duration_ms) as avg_duration_ms
|
|
FROM llm_requests
|
|
WHERE client_id = ?
|
|
GROUP BY model, provider
|
|
ORDER BY total_cost DESC
|
|
"#,
|
|
)
|
|
.bind(&id)
|
|
.fetch_all(pool)
|
|
.await;
|
|
|
|
match result {
|
|
Ok(rows) => {
|
|
let breakdown: Vec<serde_json::Value> = rows
|
|
.into_iter()
|
|
.map(|row| {
|
|
serde_json::json!({
|
|
"model": row.get::<String, _>("model"),
|
|
"provider": row.get::<String, _>("provider"),
|
|
"request_count": row.get::<i64, _>("request_count"),
|
|
"prompt_tokens": row.get::<i64, _>("prompt_tokens"),
|
|
"completion_tokens": row.get::<i64, _>("completion_tokens"),
|
|
"total_tokens": row.get::<i64, _>("total_tokens"),
|
|
"total_cost": row.get::<f64, _>("total_cost"),
|
|
"avg_duration_ms": row.get::<f64, _>("avg_duration_ms"),
|
|
})
|
|
})
|
|
.collect();
|
|
|
|
Json(ApiResponse::success(serde_json::json!({
|
|
"client_id": id,
|
|
"breakdown": breakdown,
|
|
})))
|
|
}
|
|
Err(e) => {
|
|
warn!("Failed to fetch client usage: {}", e);
|
|
Json(ApiResponse::error(format!("Failed to fetch client usage: {}", e)))
|
|
}
|
|
}
|
|
}
|
|
|
|
// ── Token management endpoints ──────────────────────────────────────
|
|
|
|
pub(super) async fn handle_get_client_tokens(
|
|
State(state): State<DashboardState>,
|
|
Path(id): Path<String>,
|
|
) -> Json<ApiResponse<serde_json::Value>> {
|
|
let pool = &state.app_state.db_pool;
|
|
|
|
let result = sqlx::query(
|
|
r#"
|
|
SELECT id, token, name, is_active, created_at, last_used_at
|
|
FROM client_tokens
|
|
WHERE client_id = ?
|
|
ORDER BY created_at DESC
|
|
"#,
|
|
)
|
|
.bind(&id)
|
|
.fetch_all(pool)
|
|
.await;
|
|
|
|
match result {
|
|
Ok(rows) => {
|
|
let tokens: Vec<serde_json::Value> = rows
|
|
.into_iter()
|
|
.map(|row| {
|
|
let token: String = row.get("token");
|
|
// Mask all but last 8 chars: sk-••••abcd1234
|
|
let masked = if token.len() > 8 {
|
|
format!("{}••••{}", &token[..3], &token[token.len() - 8..])
|
|
} else {
|
|
"••••".to_string()
|
|
};
|
|
serde_json::json!({
|
|
"id": row.get::<i64, _>("id"),
|
|
"token_masked": masked,
|
|
"name": row.get::<Option<String>, _>("name").unwrap_or_else(|| "default".to_string()),
|
|
"is_active": row.get::<bool, _>("is_active"),
|
|
"created_at": row.get::<chrono::DateTime<chrono::Utc>, _>("created_at"),
|
|
"last_used_at": row.get::<Option<chrono::DateTime<chrono::Utc>>, _>("last_used_at"),
|
|
})
|
|
})
|
|
.collect();
|
|
|
|
Json(ApiResponse::success(serde_json::json!(tokens)))
|
|
}
|
|
Err(e) => {
|
|
warn!("Failed to fetch client tokens: {}", e);
|
|
Json(ApiResponse::error(format!("Failed to fetch client tokens: {}", e)))
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
pub(super) struct CreateTokenRequest {
|
|
pub(super) name: Option<String>,
|
|
}
|
|
|
|
pub(super) async fn handle_create_client_token(
|
|
State(state): State<DashboardState>,
|
|
headers: axum::http::HeaderMap,
|
|
Path(id): Path<String>,
|
|
Json(payload): Json<CreateTokenRequest>,
|
|
) -> Json<ApiResponse<serde_json::Value>> {
|
|
let (session, _) = match super::auth::require_admin(&state, &headers).await {
|
|
Ok((session, new_token)) => (session, new_token),
|
|
Err(e) => return e,
|
|
};
|
|
|
|
let pool = &state.app_state.db_pool;
|
|
|
|
// Verify client exists
|
|
let exists: Option<(i64,)> = sqlx::query_as("SELECT 1 as x FROM clients WHERE client_id = ?")
|
|
.bind(&id)
|
|
.fetch_optional(pool)
|
|
.await
|
|
.unwrap_or(None);
|
|
|
|
if exists.is_none() {
|
|
return Json(ApiResponse::error(format!("Client '{}' not found", id)));
|
|
}
|
|
|
|
let token = generate_token();
|
|
let token_name = payload.name.unwrap_or_else(|| "default".to_string());
|
|
|
|
let result = sqlx::query(
|
|
"INSERT INTO client_tokens (client_id, token, name) VALUES (?, ?, ?) RETURNING id, created_at",
|
|
)
|
|
.bind(&id)
|
|
.bind(&token)
|
|
.bind(&token_name)
|
|
.fetch_one(pool)
|
|
.await;
|
|
|
|
match result {
|
|
Ok(row) => Json(ApiResponse::success(serde_json::json!({
|
|
"id": row.get::<i64, _>("id"),
|
|
"token": token,
|
|
"name": token_name,
|
|
"created_at": row.get::<chrono::DateTime<chrono::Utc>, _>("created_at"),
|
|
}))),
|
|
Err(e) => {
|
|
warn!("Failed to create client token: {}", e);
|
|
Json(ApiResponse::error(format!("Failed to create token: {}", e)))
|
|
}
|
|
}
|
|
}
|
|
|
|
pub(super) async fn handle_delete_client_token(
|
|
State(state): State<DashboardState>,
|
|
headers: axum::http::HeaderMap,
|
|
Path((client_id, token_id)): Path<(String, i64)>,
|
|
) -> Json<ApiResponse<serde_json::Value>> {
|
|
let (session, _) = match super::auth::require_admin(&state, &headers).await {
|
|
Ok((session, new_token)) => (session, new_token),
|
|
Err(e) => return e,
|
|
};
|
|
|
|
let pool = &state.app_state.db_pool;
|
|
|
|
let result = sqlx::query("DELETE FROM client_tokens WHERE id = ? AND client_id = ?")
|
|
.bind(token_id)
|
|
.bind(&client_id)
|
|
.execute(pool)
|
|
.await;
|
|
|
|
match result {
|
|
Ok(r) => {
|
|
if r.rows_affected() == 0 {
|
|
Json(ApiResponse::error("Token not found".to_string()))
|
|
} else {
|
|
Json(ApiResponse::success(serde_json::json!({ "message": "Token revoked" })))
|
|
}
|
|
}
|
|
Err(e) => {
|
|
warn!("Failed to delete client token: {}", e);
|
|
Json(ApiResponse::error(format!("Failed to revoke token: {}", e)))
|
|
}
|
|
}
|
|
}
|