refactor: purge Node.js, Rust is the one true backend
- Moved boosapi-rs/* to repo root - Deleted all Node.js source (src/server/, src/web/, scripts/) - Deleted node_modules, package.json, tsconfig, vite.config - Updated .gitignore and AGENTS.md for Rust-only stack - Fixed static file path for frontend dist/web/ - Dockerfile stays the same (already Rust-only) - Kept dist/web/ for Rust static serving
This commit is contained in:
@@ -0,0 +1,42 @@
|
||||
use crate::{db::models::site_announcement::SiteAnnouncement, error::AppResult};
|
||||
use sqlx::SqlitePool;
|
||||
use std::sync::Arc;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct AnnouncementService {
|
||||
db: Arc<SqlitePool>,
|
||||
}
|
||||
|
||||
impl AnnouncementService {
|
||||
pub fn new(db: Arc<SqlitePool>) -> Self {
|
||||
Self { db }
|
||||
}
|
||||
|
||||
pub async fn get_active_announcements(&self) -> AppResult<Vec<SiteAnnouncement>> {
|
||||
self.expire_old_announcements().await?;
|
||||
let rows = sqlx::query_as::<_, SiteAnnouncement>(
|
||||
"SELECT *
|
||||
FROM site_announcements
|
||||
WHERE dismissed_at IS NULL
|
||||
AND (starts_at IS NULL OR starts_at <= datetime('now'))
|
||||
AND (ends_at IS NULL OR ends_at > datetime('now'))
|
||||
ORDER BY first_seen_at DESC, id DESC",
|
||||
)
|
||||
.fetch_all(&*self.db)
|
||||
.await?;
|
||||
Ok(rows)
|
||||
}
|
||||
|
||||
pub async fn expire_old_announcements(&self) -> AppResult<u64> {
|
||||
let result = sqlx::query(
|
||||
"UPDATE site_announcements
|
||||
SET dismissed_at = datetime('now')
|
||||
WHERE dismissed_at IS NULL
|
||||
AND ends_at IS NOT NULL
|
||||
AND ends_at <= datetime('now')",
|
||||
)
|
||||
.execute(&*self.db)
|
||||
.await?;
|
||||
Ok(result.rows_affected())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,282 @@
|
||||
use crate::{
|
||||
error::{AppError, AppResult},
|
||||
services::shared::{enforce_key_request_limit, increment_key_usage},
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sqlx::{FromRow, SqlitePool};
|
||||
use std::sync::Arc;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct UsageTracker {
|
||||
db: Arc<SqlitePool>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, FromRow, PartialEq)]
|
||||
pub struct UsageSummary {
|
||||
pub downstream_key_id: i64,
|
||||
pub request_count: i64,
|
||||
pub prompt_tokens: i64,
|
||||
pub completion_tokens: i64,
|
||||
pub total_tokens: i64,
|
||||
pub total_cost: f64,
|
||||
}
|
||||
|
||||
impl UsageTracker {
|
||||
pub fn new(db: Arc<SqlitePool>) -> Self {
|
||||
Self { db }
|
||||
}
|
||||
|
||||
pub async fn record_usage(
|
||||
&self,
|
||||
key_id: i64,
|
||||
model: impl Into<String>,
|
||||
prompt_tokens: i64,
|
||||
completion_tokens: i64,
|
||||
cost: f64,
|
||||
) -> AppResult<()> {
|
||||
Self::ensure_schema(&self.db).await?;
|
||||
let prompt_tokens = prompt_tokens.max(0);
|
||||
let completion_tokens = completion_tokens.max(0);
|
||||
let total_tokens = prompt_tokens + completion_tokens;
|
||||
|
||||
sqlx::query(
|
||||
"INSERT INTO usage_logs
|
||||
(downstream_key_id, model, prompt_tokens, completion_tokens, total_tokens, cost)
|
||||
VALUES (?, ?, ?, ?, ?, ?)",
|
||||
)
|
||||
.bind(key_id)
|
||||
.bind(model.into())
|
||||
.bind(prompt_tokens)
|
||||
.bind(completion_tokens)
|
||||
.bind(total_tokens)
|
||||
.bind(cost.max(0.0))
|
||||
.execute(&*self.db)
|
||||
.await?;
|
||||
|
||||
increment_key_usage(&self.db, key_id, cost.max(0.0)).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn get_key_usage(
|
||||
&self,
|
||||
key_id: i64,
|
||||
since: Option<&str>,
|
||||
until: Option<&str>,
|
||||
) -> AppResult<UsageSummary> {
|
||||
Self::ensure_schema(&self.db).await?;
|
||||
|
||||
let summary = sqlx::query_as::<_, UsageSummary>(
|
||||
"SELECT
|
||||
? AS downstream_key_id,
|
||||
COUNT(*) AS request_count,
|
||||
COALESCE(SUM(prompt_tokens), 0) AS prompt_tokens,
|
||||
COALESCE(SUM(completion_tokens), 0) AS completion_tokens,
|
||||
COALESCE(SUM(total_tokens), 0) AS total_tokens,
|
||||
COALESCE(SUM(cost), 0.0) AS total_cost
|
||||
FROM usage_logs
|
||||
WHERE downstream_key_id = ?
|
||||
AND (? IS NULL OR created_at >= ?)
|
||||
AND (? IS NULL OR created_at <= ?)",
|
||||
)
|
||||
.bind(key_id)
|
||||
.bind(key_id)
|
||||
.bind(since)
|
||||
.bind(since)
|
||||
.bind(until)
|
||||
.bind(until)
|
||||
.fetch_one(&*self.db)
|
||||
.await?;
|
||||
|
||||
Ok(summary)
|
||||
}
|
||||
|
||||
pub async fn check_rate_limit(&self, key_id: i64) -> AppResult<()> {
|
||||
enforce_key_request_limit(&self.db, key_id, 1).await?;
|
||||
|
||||
let row: Option<(Option<f64>, Option<f64>, Option<bool>)> = sqlx::query_as(
|
||||
"SELECT max_cost, used_cost, enabled
|
||||
FROM downstream_api_keys
|
||||
WHERE id = ?",
|
||||
)
|
||||
.bind(key_id)
|
||||
.fetch_optional(&*self.db)
|
||||
.await?;
|
||||
|
||||
let (max_cost, used_cost, enabled) =
|
||||
row.ok_or_else(|| AppError::Unauthorized("API key not found".into()))?;
|
||||
if !enabled.unwrap_or(true) {
|
||||
return Err(AppError::Forbidden("API key is disabled".into()));
|
||||
}
|
||||
if let Some(max_cost) = max_cost {
|
||||
if used_cost.unwrap_or(0.0) >= max_cost {
|
||||
return Err(AppError::RateLimited);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn create_daily_usage_snapshot(&self) -> AppResult<()> {
|
||||
Self::create_daily_usage_snapshot_for_db(&self.db).await
|
||||
}
|
||||
|
||||
pub async fn create_daily_usage_snapshot_for_db(db: &SqlitePool) -> AppResult<()> {
|
||||
Self::ensure_schema(db).await?;
|
||||
sqlx::query(
|
||||
"INSERT INTO usage_daily_snapshots
|
||||
(local_day, downstream_key_id, request_count, prompt_tokens,
|
||||
completion_tokens, total_tokens, total_cost, updated_at)
|
||||
SELECT
|
||||
date(created_at, 'localtime') AS local_day,
|
||||
downstream_key_id,
|
||||
COUNT(*) AS request_count,
|
||||
COALESCE(SUM(prompt_tokens), 0) AS prompt_tokens,
|
||||
COALESCE(SUM(completion_tokens), 0) AS completion_tokens,
|
||||
COALESCE(SUM(total_tokens), 0) AS total_tokens,
|
||||
COALESCE(SUM(cost), 0.0) AS total_cost,
|
||||
datetime('now') AS updated_at
|
||||
FROM usage_logs
|
||||
GROUP BY date(created_at, 'localtime'), downstream_key_id
|
||||
ON CONFLICT(local_day, downstream_key_id) DO UPDATE SET
|
||||
request_count = excluded.request_count,
|
||||
prompt_tokens = excluded.prompt_tokens,
|
||||
completion_tokens = excluded.completion_tokens,
|
||||
total_tokens = excluded.total_tokens,
|
||||
total_cost = excluded.total_cost,
|
||||
updated_at = datetime('now')",
|
||||
)
|
||||
.execute(db)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn ensure_schema(db: &SqlitePool) -> AppResult<()> {
|
||||
sqlx::query(
|
||||
"CREATE TABLE IF NOT EXISTS usage_logs (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
downstream_key_id INTEGER,
|
||||
model TEXT,
|
||||
prompt_tokens INTEGER DEFAULT 0,
|
||||
completion_tokens INTEGER DEFAULT 0,
|
||||
total_tokens INTEGER DEFAULT 0,
|
||||
cost REAL DEFAULT 0.0,
|
||||
created_at TEXT DEFAULT (datetime('now'))
|
||||
)",
|
||||
)
|
||||
.execute(db)
|
||||
.await?;
|
||||
sqlx::query(
|
||||
"CREATE INDEX IF NOT EXISTS usage_logs_downstream_key_created_idx
|
||||
ON usage_logs (downstream_key_id, created_at)",
|
||||
)
|
||||
.execute(db)
|
||||
.await?;
|
||||
sqlx::query(
|
||||
"CREATE TABLE IF NOT EXISTS usage_daily_snapshots (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
local_day TEXT NOT NULL,
|
||||
downstream_key_id INTEGER,
|
||||
request_count INTEGER NOT NULL DEFAULT 0,
|
||||
prompt_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
completion_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
total_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
total_cost REAL NOT NULL DEFAULT 0.0,
|
||||
created_at TEXT DEFAULT (datetime('now')),
|
||||
updated_at TEXT DEFAULT (datetime('now')),
|
||||
UNIQUE(local_day, downstream_key_id)
|
||||
)",
|
||||
)
|
||||
.execute(db)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use sqlx::sqlite::SqlitePoolOptions;
|
||||
|
||||
async fn pool() -> SqlitePool {
|
||||
let pool = SqlitePoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect("sqlite::memory:")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
sqlx::query(
|
||||
"CREATE TABLE downstream_api_keys (
|
||||
id INTEGER PRIMARY KEY,
|
||||
name TEXT NOT NULL,
|
||||
key TEXT NOT NULL,
|
||||
enabled INTEGER DEFAULT 1,
|
||||
max_cost REAL,
|
||||
used_cost REAL DEFAULT 0,
|
||||
max_requests INTEGER,
|
||||
used_requests INTEGER DEFAULT 0,
|
||||
last_used_at TEXT,
|
||||
updated_at TEXT DEFAULT (datetime('now'))
|
||||
)",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"INSERT INTO downstream_api_keys
|
||||
(id, name, key, enabled, max_cost, used_cost, max_requests, used_requests)
|
||||
VALUES (1, 'test', 'sk-test', 1, 1.0, 0.0, 5, 0)",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
pool
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn records_and_aggregates_usage() {
|
||||
let tracker = UsageTracker::new(Arc::new(pool().await));
|
||||
|
||||
tracker
|
||||
.record_usage(1, "gpt-test", 10, 20, 0.25)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let summary = tracker.get_key_usage(1, None, None).await.unwrap();
|
||||
assert_eq!(summary.request_count, 1);
|
||||
assert_eq!(summary.prompt_tokens, 10);
|
||||
assert_eq!(summary.completion_tokens, 20);
|
||||
assert_eq!(summary.total_tokens, 30);
|
||||
assert_eq!(summary.total_cost, 0.25);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rate_limit_checks_request_and_cost_limits() {
|
||||
let tracker = UsageTracker::new(Arc::new(pool().await));
|
||||
tracker.check_rate_limit(1).await.unwrap();
|
||||
|
||||
sqlx::query("UPDATE downstream_api_keys SET used_requests = 5 WHERE id = 1")
|
||||
.execute(&*tracker.db)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
tracker.check_rate_limit(1).await,
|
||||
Err(AppError::RateLimited)
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn creates_daily_snapshot() {
|
||||
let tracker = UsageTracker::new(Arc::new(pool().await));
|
||||
tracker
|
||||
.record_usage(1, "gpt-test", 2, 3, 0.1)
|
||||
.await
|
||||
.unwrap();
|
||||
tracker.create_daily_usage_snapshot().await.unwrap();
|
||||
|
||||
let count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM usage_daily_snapshots")
|
||||
.fetch_one(&*tracker.db)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(count, 1);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,202 @@
|
||||
use crate::{error::AppResult, proxy::token_router::TokenRouter, services::billing::UsageTracker};
|
||||
use sqlx::SqlitePool;
|
||||
use std::sync::{
|
||||
atomic::{AtomicBool, Ordering},
|
||||
Arc,
|
||||
};
|
||||
use tokio::{
|
||||
task::JoinHandle,
|
||||
time::{interval, Duration},
|
||||
};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct CronScheduler {
|
||||
db: Arc<SqlitePool>,
|
||||
running: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
pub struct CronSchedulerHandle {
|
||||
running: Arc<AtomicBool>,
|
||||
tasks: Vec<JoinHandle<()>>,
|
||||
}
|
||||
|
||||
impl CronScheduler {
|
||||
pub fn new(db: Arc<SqlitePool>) -> Self {
|
||||
Self {
|
||||
db,
|
||||
running: Arc::new(AtomicBool::new(false)),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn start(&self) -> CronSchedulerHandle {
|
||||
self.running.store(true, Ordering::SeqCst);
|
||||
|
||||
let tasks = vec![
|
||||
self.spawn_periodic(
|
||||
"usage aggregation",
|
||||
Duration::from_secs(300),
|
||||
Self::aggregate_usage,
|
||||
),
|
||||
self.spawn_periodic(
|
||||
"route refresh",
|
||||
Duration::from_secs(600),
|
||||
Self::refresh_routes,
|
||||
),
|
||||
self.spawn_periodic("model probe", Duration::from_secs(900), Self::probe_models),
|
||||
self.spawn_periodic(
|
||||
"channel recovery",
|
||||
Duration::from_secs(120),
|
||||
Self::recover_channels,
|
||||
),
|
||||
];
|
||||
|
||||
CronSchedulerHandle {
|
||||
running: self.running.clone(),
|
||||
tasks,
|
||||
}
|
||||
}
|
||||
|
||||
fn spawn_periodic<F, Fut>(&self, name: &'static str, every: Duration, task: F) -> JoinHandle<()>
|
||||
where
|
||||
F: Fn(Arc<SqlitePool>) -> Fut + Send + Sync + Copy + 'static,
|
||||
Fut: std::future::Future<Output = AppResult<()>> + Send + 'static,
|
||||
{
|
||||
let db = self.db.clone();
|
||||
let running = self.running.clone();
|
||||
tokio::spawn(async move {
|
||||
let mut ticker = interval(every);
|
||||
while running.load(Ordering::SeqCst) {
|
||||
ticker.tick().await;
|
||||
if !running.load(Ordering::SeqCst) {
|
||||
break;
|
||||
}
|
||||
if let Err(err) = task(db.clone()).await {
|
||||
tracing::warn!(%err, task = name, "cron task failed");
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
async fn aggregate_usage(db: Arc<SqlitePool>) -> AppResult<()> {
|
||||
UsageTracker::create_daily_usage_snapshot_for_db(&db).await?;
|
||||
sqlx::query(
|
||||
"UPDATE downstream_api_keys
|
||||
SET updated_at = datetime('now')
|
||||
WHERE used_requests > 0 OR used_cost > 0",
|
||||
)
|
||||
.execute(&*db)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn refresh_routes(db: Arc<SqlitePool>) -> AppResult<()> {
|
||||
sqlx::query(
|
||||
"UPDATE token_routes
|
||||
SET decision_refreshed_at = datetime('now'), updated_at = datetime('now')
|
||||
WHERE enabled = 1",
|
||||
)
|
||||
.execute(&*db)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn probe_models(db: Arc<SqlitePool>) -> AppResult<()> {
|
||||
let router = TokenRouter::new(db.clone());
|
||||
let models: Vec<String> = sqlx::query_scalar(
|
||||
"SELECT model_pattern FROM token_routes WHERE enabled = 1 ORDER BY id LIMIT 25",
|
||||
)
|
||||
.fetch_all(&*db)
|
||||
.await?;
|
||||
|
||||
for model in models {
|
||||
let probe_model = model.replace('*', "probe");
|
||||
let _ = router.resolve_selected(&probe_model).await;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn recover_channels(db: Arc<SqlitePool>) -> AppResult<()> {
|
||||
sqlx::query(
|
||||
"UPDATE route_channels
|
||||
SET cooldown_until = NULL, cooldown_level = 0, consecutive_fail_count = 0
|
||||
WHERE cooldown_until IS NOT NULL AND cooldown_until <= datetime('now')",
|
||||
)
|
||||
.execute(&*db)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl CronSchedulerHandle {
|
||||
pub fn stop(&self) {
|
||||
self.running.store(false, Ordering::SeqCst);
|
||||
for task in &self.tasks {
|
||||
task.abort();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for CronSchedulerHandle {
|
||||
fn drop(&mut self) {
|
||||
self.running.store(false, Ordering::SeqCst);
|
||||
for task in &self.tasks {
|
||||
task.abort();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use sqlx::sqlite::SqlitePoolOptions;
|
||||
|
||||
#[tokio::test]
|
||||
async fn scheduler_starts_and_stops_background_tasks() {
|
||||
let pool = SqlitePoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect("sqlite::memory:")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
sqlx::query(
|
||||
"CREATE TABLE downstream_api_keys (
|
||||
id INTEGER PRIMARY KEY,
|
||||
used_requests INTEGER DEFAULT 0,
|
||||
used_cost REAL DEFAULT 0,
|
||||
updated_at TEXT
|
||||
)",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"CREATE TABLE token_routes (
|
||||
id INTEGER PRIMARY KEY,
|
||||
model_pattern TEXT NOT NULL,
|
||||
enabled INTEGER DEFAULT 1,
|
||||
decision_refreshed_at TEXT,
|
||||
updated_at TEXT
|
||||
)",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"CREATE TABLE route_channels (
|
||||
id INTEGER PRIMARY KEY,
|
||||
cooldown_until TEXT,
|
||||
cooldown_level INTEGER DEFAULT 0,
|
||||
consecutive_fail_count INTEGER DEFAULT 0
|
||||
)",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let scheduler = CronScheduler::new(Arc::new(pool));
|
||||
let handle = scheduler.start();
|
||||
assert!(handle.running.load(Ordering::SeqCst));
|
||||
handle.stop();
|
||||
assert!(!handle.running.load(Ordering::SeqCst));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
pub mod announcement;
|
||||
pub mod billing;
|
||||
pub mod cron;
|
||||
pub mod notification;
|
||||
pub mod oauth;
|
||||
pub mod shared;
|
||||
pub mod update_center;
|
||||
@@ -0,0 +1,544 @@
|
||||
use crate::error::AppResult;
|
||||
use chrono::{FixedOffset, SecondsFormat, TimeZone, Utc};
|
||||
use lettre::{
|
||||
message::Mailbox, transport::smtp::authentication::Credentials, Message, SmtpTransport,
|
||||
Transport,
|
||||
};
|
||||
use reqwest::Url;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::json;
|
||||
use sqlx::SqlitePool;
|
||||
use std::{collections::HashMap, sync::Arc};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct NotificationService {
|
||||
pub db: Arc<SqlitePool>,
|
||||
pub client: reqwest::Client,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum NotificationLevel {
|
||||
Info,
|
||||
Warning,
|
||||
Error,
|
||||
Critical,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Notification {
|
||||
pub title: String,
|
||||
pub message: String,
|
||||
pub level: NotificationLevel,
|
||||
pub channel: Option<String>,
|
||||
pub signature: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
struct NotificationConfig {
|
||||
telegram_bot_token: Option<String>,
|
||||
telegram_chat_id: Option<String>,
|
||||
telegram_thread_id: Option<i64>,
|
||||
wecom_webhook: Option<String>,
|
||||
feishu_webhook: Option<String>,
|
||||
webhook_url: Option<String>,
|
||||
bark_url: Option<String>,
|
||||
serverchan_key: Option<String>,
|
||||
smtp: SmtpConfig,
|
||||
cooldown_sec: i64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
struct SmtpConfig {
|
||||
host: Option<String>,
|
||||
port: Option<u16>,
|
||||
user: Option<String>,
|
||||
pass: Option<String>,
|
||||
from: Option<String>,
|
||||
to: Option<String>,
|
||||
}
|
||||
|
||||
impl NotificationService {
|
||||
pub fn new(db: Arc<SqlitePool>) -> Self {
|
||||
Self {
|
||||
db,
|
||||
client: reqwest::Client::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_client(db: Arc<SqlitePool>, client: reqwest::Client) -> Self {
|
||||
Self { db, client }
|
||||
}
|
||||
|
||||
pub async fn send(&self, notif: Notification) -> AppResult<()> {
|
||||
Self::ensure_schema(&self.db).await?;
|
||||
let config = self.load_config().await?;
|
||||
|
||||
if self.should_throttle(¬if, config.cooldown_sec).await? {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let mut channels = self.enabled_channels(&config);
|
||||
if let Some(channel) = notif.channel.as_deref() {
|
||||
channels.retain(|candidate| candidate == &normalize_channel(channel));
|
||||
}
|
||||
|
||||
if channels.is_empty() {
|
||||
tracing::debug!(
|
||||
channel = notif.channel.as_deref().unwrap_or("all"),
|
||||
"notification skipped because no matching channel is configured"
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
for channel in channels {
|
||||
if let Err(err) = self.send_channel(&channel, ¬if, &config).await {
|
||||
tracing::warn!(%err, channel, "notification channel failed");
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn load_config(&self) -> AppResult<NotificationConfig> {
|
||||
let rows: Vec<(String, Option<String>)> =
|
||||
sqlx::query_as("SELECT key, value FROM settings WHERE key LIKE 'notify_%'")
|
||||
.fetch_all(&*self.db)
|
||||
.await?;
|
||||
let map: HashMap<String, String> = rows
|
||||
.into_iter()
|
||||
.filter_map(|(key, value)| value.map(|value| (key, value)))
|
||||
.filter(|(_, value)| !value.trim().is_empty())
|
||||
.collect();
|
||||
|
||||
Ok(NotificationConfig {
|
||||
telegram_bot_token: setting(&map, "notify_telegram_bot_token"),
|
||||
telegram_chat_id: setting(&map, "notify_telegram_chat_id"),
|
||||
telegram_thread_id: setting(&map, "notify_telegram_thread_id")
|
||||
.and_then(|value| value.parse().ok()),
|
||||
wecom_webhook: setting(&map, "notify_wecom_webhook"),
|
||||
feishu_webhook: setting(&map, "notify_feishu_webhook"),
|
||||
webhook_url: setting(&map, "notify_webhook_url"),
|
||||
bark_url: setting(&map, "notify_bark_url"),
|
||||
serverchan_key: setting(&map, "notify_serverchan_key"),
|
||||
smtp: SmtpConfig {
|
||||
host: setting(&map, "notify_smtp_host"),
|
||||
port: setting(&map, "notify_smtp_port").and_then(|value| value.parse().ok()),
|
||||
user: setting(&map, "notify_smtp_user"),
|
||||
pass: setting(&map, "notify_smtp_pass"),
|
||||
from: setting(&map, "notify_smtp_from"),
|
||||
to: setting(&map, "notify_smtp_to"),
|
||||
},
|
||||
cooldown_sec: setting(&map, "notify_cooldown_sec")
|
||||
.and_then(|value| value.parse().ok())
|
||||
.unwrap_or(60),
|
||||
})
|
||||
}
|
||||
|
||||
async fn should_throttle(&self, notif: &Notification, cooldown_sec: i64) -> AppResult<bool> {
|
||||
let Some(signature) = notif.signature.as_deref().map(str::trim) else {
|
||||
return Ok(false);
|
||||
};
|
||||
if signature.is_empty() || cooldown_sec <= 0 {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
let updated = sqlx::query(
|
||||
"UPDATE notification_throttle
|
||||
SET count = count + 1,
|
||||
title = ?,
|
||||
message = ?,
|
||||
level = ?,
|
||||
updated_at = datetime('now')
|
||||
WHERE signature = ?
|
||||
AND updated_at >= datetime('now', ?)",
|
||||
)
|
||||
.bind(¬if.title)
|
||||
.bind(¬if.message)
|
||||
.bind(level_name(¬if.level))
|
||||
.bind(signature)
|
||||
.bind(format!("-{} seconds", cooldown_sec))
|
||||
.execute(&*self.db)
|
||||
.await?;
|
||||
|
||||
if updated.rows_affected() > 0 {
|
||||
tracing::info!(signature, "notification merged into cooldown throttle");
|
||||
return Ok(true);
|
||||
}
|
||||
|
||||
sqlx::query(
|
||||
"INSERT INTO notification_throttle
|
||||
(signature, title, message, level, count, updated_at)
|
||||
VALUES (?, ?, ?, ?, 1, datetime('now'))
|
||||
ON CONFLICT(signature) DO UPDATE SET
|
||||
title = excluded.title,
|
||||
message = excluded.message,
|
||||
level = excluded.level,
|
||||
count = 1,
|
||||
updated_at = datetime('now')",
|
||||
)
|
||||
.bind(signature)
|
||||
.bind(¬if.title)
|
||||
.bind(¬if.message)
|
||||
.bind(level_name(¬if.level))
|
||||
.execute(&*self.db)
|
||||
.await?;
|
||||
Ok(false)
|
||||
}
|
||||
|
||||
fn enabled_channels(&self, config: &NotificationConfig) -> Vec<String> {
|
||||
let mut channels = Vec::new();
|
||||
if config.telegram_bot_token.is_some() && config.telegram_chat_id.is_some() {
|
||||
channels.push("telegram".to_string());
|
||||
}
|
||||
if config.wecom_webhook.is_some() {
|
||||
channels.push("wecom".to_string());
|
||||
}
|
||||
if config.feishu_webhook.is_some() {
|
||||
channels.push("feishu".to_string());
|
||||
}
|
||||
if config.webhook_url.is_some() {
|
||||
channels.push("webhook".to_string());
|
||||
}
|
||||
if config.bark_url.is_some() {
|
||||
channels.push("bark".to_string());
|
||||
}
|
||||
if config.serverchan_key.is_some() {
|
||||
channels.push("serverchan".to_string());
|
||||
}
|
||||
if config.smtp.is_configured() {
|
||||
channels.push("smtp".to_string());
|
||||
}
|
||||
channels
|
||||
}
|
||||
|
||||
async fn send_channel(
|
||||
&self,
|
||||
channel: &str,
|
||||
notif: &Notification,
|
||||
config: &NotificationConfig,
|
||||
) -> anyhow::Result<()> {
|
||||
match channel {
|
||||
"telegram" => self.send_telegram(notif, config).await,
|
||||
"wecom" => self.send_wecom(notif, config).await,
|
||||
"feishu" => self.send_feishu(notif, config).await,
|
||||
"webhook" => self.send_webhook(notif, config).await,
|
||||
"bark" => self.send_bark(notif, config).await,
|
||||
"serverchan" => self.send_serverchan(notif, config).await,
|
||||
"smtp" => self.send_smtp(notif, config).await,
|
||||
_ => Ok(()),
|
||||
}
|
||||
}
|
||||
|
||||
async fn send_telegram(
|
||||
&self,
|
||||
notif: &Notification,
|
||||
config: &NotificationConfig,
|
||||
) -> anyhow::Result<()> {
|
||||
let token = required(config.telegram_bot_token.as_deref(), "telegram token")?;
|
||||
let chat_id = required(config.telegram_chat_id.as_deref(), "telegram chat id")?;
|
||||
let mut body = json!({
|
||||
"chat_id": chat_id,
|
||||
"text": truncate_chars(&format_notification_text(notif), 3900),
|
||||
});
|
||||
if let Some(thread_id) = config.telegram_thread_id {
|
||||
body["message_thread_id"] = json!(thread_id);
|
||||
}
|
||||
self.client
|
||||
.post(format!("https://api.telegram.org/bot{token}/sendMessage"))
|
||||
.json(&body)
|
||||
.send()
|
||||
.await?
|
||||
.error_for_status()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn send_wecom(
|
||||
&self,
|
||||
notif: &Notification,
|
||||
config: &NotificationConfig,
|
||||
) -> anyhow::Result<()> {
|
||||
let url = required(config.wecom_webhook.as_deref(), "wecom webhook")?;
|
||||
self.client
|
||||
.post(url)
|
||||
.json(&json!({
|
||||
"msgtype": "text",
|
||||
"text": { "content": format_notification_text(notif) },
|
||||
}))
|
||||
.send()
|
||||
.await?
|
||||
.error_for_status()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn send_feishu(
|
||||
&self,
|
||||
notif: &Notification,
|
||||
config: &NotificationConfig,
|
||||
) -> anyhow::Result<()> {
|
||||
let url = required(config.feishu_webhook.as_deref(), "feishu webhook")?;
|
||||
self.client
|
||||
.post(url)
|
||||
.json(&json!({
|
||||
"msg_type": "text",
|
||||
"content": { "text": format_notification_text(notif) },
|
||||
}))
|
||||
.send()
|
||||
.await?
|
||||
.error_for_status()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn send_webhook(
|
||||
&self,
|
||||
notif: &Notification,
|
||||
config: &NotificationConfig,
|
||||
) -> anyhow::Result<()> {
|
||||
let url = required(config.webhook_url.as_deref(), "webhook url")?;
|
||||
self.client
|
||||
.post(url)
|
||||
.json(&json!({
|
||||
"title": ¬if.title,
|
||||
"message": ¬if.message,
|
||||
"level": level_name(¬if.level),
|
||||
}))
|
||||
.send()
|
||||
.await?
|
||||
.error_for_status()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn send_bark(
|
||||
&self,
|
||||
notif: &Notification,
|
||||
config: &NotificationConfig,
|
||||
) -> anyhow::Result<()> {
|
||||
let base = required(config.bark_url.as_deref(), "bark url")?;
|
||||
let mut url = Url::parse(base)?;
|
||||
{
|
||||
let mut segments = url
|
||||
.path_segments_mut()
|
||||
.map_err(|_| anyhow::anyhow!("bark url cannot be a base url"))?;
|
||||
segments.pop_if_empty();
|
||||
segments.push(¬if.title);
|
||||
segments.push(¬if.message);
|
||||
}
|
||||
self.client.get(url).send().await?.error_for_status()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn send_serverchan(
|
||||
&self,
|
||||
notif: &Notification,
|
||||
config: &NotificationConfig,
|
||||
) -> anyhow::Result<()> {
|
||||
let key = required(config.serverchan_key.as_deref(), "serverchan key")?;
|
||||
self.client
|
||||
.post(format!("https://sctapi.ftqq.com/{key}.send"))
|
||||
.form(&[
|
||||
("title", notif.title.as_str()),
|
||||
("desp", notif.message.as_str()),
|
||||
])
|
||||
.send()
|
||||
.await?
|
||||
.error_for_status()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn send_smtp(
|
||||
&self,
|
||||
notif: &Notification,
|
||||
config: &NotificationConfig,
|
||||
) -> anyhow::Result<()> {
|
||||
let smtp = config.smtp.clone();
|
||||
let notif = notif.clone();
|
||||
tokio::task::spawn_blocking(move || smtp.send(¬if)).await??;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn ensure_schema(db: &SqlitePool) -> AppResult<()> {
|
||||
sqlx::query(
|
||||
"CREATE TABLE IF NOT EXISTS notification_throttle (
|
||||
signature TEXT PRIMARY KEY NOT NULL,
|
||||
title TEXT NOT NULL,
|
||||
message TEXT NOT NULL,
|
||||
level TEXT NOT NULL,
|
||||
count INTEGER NOT NULL DEFAULT 1,
|
||||
updated_at TEXT DEFAULT (datetime('now'))
|
||||
)",
|
||||
)
|
||||
.execute(db)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl SmtpConfig {
|
||||
fn is_configured(&self) -> bool {
|
||||
self.host.is_some()
|
||||
&& self.port.is_some()
|
||||
&& self.user.is_some()
|
||||
&& self.pass.is_some()
|
||||
&& self.from.is_some()
|
||||
}
|
||||
|
||||
fn send(self, notif: &Notification) -> anyhow::Result<()> {
|
||||
let host = required(self.host.as_deref(), "smtp host")?;
|
||||
let port = self.port.unwrap_or(587);
|
||||
let user = required(self.user.as_deref(), "smtp user")?;
|
||||
let pass = required(self.pass.as_deref(), "smtp pass")?;
|
||||
let from = required(self.from.as_deref(), "smtp from")?;
|
||||
let to = self.to.as_deref().unwrap_or(from);
|
||||
|
||||
let email = Message::builder()
|
||||
.from(from.parse::<Mailbox>()?)
|
||||
.to(to.parse::<Mailbox>()?)
|
||||
.subject(format!(
|
||||
"[BoosAPI][{}] {}",
|
||||
level_name(¬if.level),
|
||||
notif.title
|
||||
))
|
||||
.body(format_notification_text(notif))?;
|
||||
|
||||
let mailer = SmtpTransport::relay(host)?
|
||||
.port(port)
|
||||
.credentials(Credentials::new(user.to_string(), pass.to_string()))
|
||||
.build();
|
||||
mailer.send(&email)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn setting(map: &HashMap<String, String>, key: &str) -> Option<String> {
|
||||
map.get(key).map(|value| value.trim().to_string())
|
||||
}
|
||||
|
||||
fn required<'a>(value: Option<&'a str>, name: &str) -> anyhow::Result<&'a str> {
|
||||
value.ok_or_else(|| anyhow::anyhow!("{name} is not configured"))
|
||||
}
|
||||
|
||||
fn normalize_channel(channel: &str) -> String {
|
||||
channel.trim().to_ascii_lowercase()
|
||||
}
|
||||
|
||||
fn level_name(level: &NotificationLevel) -> &'static str {
|
||||
match level {
|
||||
NotificationLevel::Info => "info",
|
||||
NotificationLevel::Warning => "warning",
|
||||
NotificationLevel::Error => "error",
|
||||
NotificationLevel::Critical => "critical",
|
||||
}
|
||||
}
|
||||
|
||||
fn format_notification_text(notif: &Notification) -> String {
|
||||
let utc = Utc::now();
|
||||
let shanghai = FixedOffset::east_opt(8 * 3600)
|
||||
.expect("valid Shanghai offset")
|
||||
.from_utc_datetime(&utc.naive_utc());
|
||||
format!(
|
||||
"[BoosAPI][{}] {}\n\n{}\n\nLevel: {}\nLocal Time: {} (Asia/Shanghai)\nUTC Time: {}",
|
||||
level_name(¬if.level).to_ascii_uppercase(),
|
||||
notif.title,
|
||||
notif.message,
|
||||
level_name(¬if.level),
|
||||
shanghai.format("%Y-%m-%d %H:%M:%S"),
|
||||
utc.to_rfc3339_opts(SecondsFormat::Millis, true)
|
||||
)
|
||||
}
|
||||
|
||||
fn truncate_chars(value: &str, max_chars: usize) -> String {
|
||||
if value.chars().count() <= max_chars {
|
||||
return value.to_string();
|
||||
}
|
||||
value.chars().take(max_chars).collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use sqlx::sqlite::SqlitePoolOptions;
|
||||
|
||||
async fn pool() -> SqlitePool {
|
||||
let pool = SqlitePoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect("sqlite::memory:")
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query("CREATE TABLE settings (key TEXT PRIMARY KEY NOT NULL, value TEXT)")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
pool
|
||||
}
|
||||
|
||||
async fn set(pool: &SqlitePool, key: &str, value: &str) {
|
||||
sqlx::query("INSERT OR REPLACE INTO settings (key, value) VALUES (?, ?)")
|
||||
.bind(key)
|
||||
.bind(value)
|
||||
.execute(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn skips_when_no_channel_is_configured() {
|
||||
let pool = pool().await;
|
||||
let service = NotificationService::new(Arc::new(pool));
|
||||
service
|
||||
.send(Notification {
|
||||
title: "Disk".to_string(),
|
||||
message: "Almost full".to_string(),
|
||||
level: NotificationLevel::Warning,
|
||||
channel: Some("webhook".to_string()),
|
||||
signature: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn throttles_duplicate_signature() {
|
||||
let pool = pool().await;
|
||||
set(&pool, "notify_cooldown_sec", "60").await;
|
||||
|
||||
let service = NotificationService::new(Arc::new(pool));
|
||||
let notification = Notification {
|
||||
title: "Route down".to_string(),
|
||||
message: "primary failed".to_string(),
|
||||
level: NotificationLevel::Error,
|
||||
channel: Some("webhook".to_string()),
|
||||
signature: Some("route:1".to_string()),
|
||||
};
|
||||
service.send(notification.clone()).await.unwrap();
|
||||
service.send(notification).await.unwrap();
|
||||
|
||||
let count: i64 = sqlx::query_scalar(
|
||||
"SELECT count FROM notification_throttle WHERE signature = 'route:1'",
|
||||
)
|
||||
.fetch_one(&*service.db)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(count, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn formats_telegram_compatible_text() {
|
||||
let text = format_notification_text(&Notification {
|
||||
title: "Title".to_string(),
|
||||
message: "Message".to_string(),
|
||||
level: NotificationLevel::Critical,
|
||||
channel: None,
|
||||
signature: None,
|
||||
});
|
||||
|
||||
assert!(text.starts_with("[BoosAPI][CRITICAL] Title\n\nMessage"));
|
||||
assert!(text.contains("Local Time:"));
|
||||
assert!(text.contains("UTC Time:"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncates_by_char_count() {
|
||||
assert_eq!(truncate_chars("abcdef", 3), "abc");
|
||||
assert_eq!(truncate_chars("短消息", 2), "短消");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,207 @@
|
||||
use crate::error::{AppError, AppResult};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use sqlx::{FromRow, SqlitePool};
|
||||
use std::sync::Arc;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct OAuthService {
|
||||
db: Arc<SqlitePool>,
|
||||
client: reqwest::Client,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, FromRow)]
|
||||
pub struct OAuthRouteUnit {
|
||||
pub id: i64,
|
||||
pub site_id: i64,
|
||||
pub provider: String,
|
||||
pub name: String,
|
||||
pub strategy: String,
|
||||
pub enabled: Option<bool>,
|
||||
pub created_at: Option<String>,
|
||||
pub updated_at: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct RefreshedToken {
|
||||
pub account_id: i64,
|
||||
pub access_token: String,
|
||||
pub expires_at: Option<String>,
|
||||
}
|
||||
|
||||
impl OAuthService {
|
||||
pub fn new(db: Arc<SqlitePool>) -> Self {
|
||||
Self {
|
||||
db,
|
||||
client: reqwest::Client::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn refresh_token(&self, account_id: i64, site_id: i64) -> AppResult<RefreshedToken> {
|
||||
let row: Option<(String, Option<String>, Option<String>)> = sqlx::query_as(
|
||||
"SELECT access_token, api_token, extra_config
|
||||
FROM accounts
|
||||
WHERE id = ? AND site_id = ? AND oauth_provider IS NOT NULL",
|
||||
)
|
||||
.bind(account_id)
|
||||
.bind(site_id)
|
||||
.fetch_optional(&*self.db)
|
||||
.await?;
|
||||
|
||||
let (current_access_token, refresh_token, extra_config) =
|
||||
row.ok_or_else(|| AppError::NotFound("OAuth account not found".into()))?;
|
||||
|
||||
if !self.token_expired(extra_config.as_deref()) {
|
||||
return Ok(RefreshedToken {
|
||||
account_id,
|
||||
access_token: current_access_token,
|
||||
expires_at: token_expiry(extra_config.as_deref()),
|
||||
});
|
||||
}
|
||||
|
||||
let refresh_token = refresh_token
|
||||
.filter(|token| !token.trim().is_empty())
|
||||
.ok_or_else(|| AppError::BadRequest("OAuth account has no refresh token".into()))?;
|
||||
|
||||
let mut refreshed = self.perform_refresh(site_id, &refresh_token).await?;
|
||||
refreshed.account_id = account_id;
|
||||
sqlx::query(
|
||||
"UPDATE accounts
|
||||
SET access_token = ?, extra_config = ?, updated_at = datetime('now')
|
||||
WHERE id = ? AND site_id = ?",
|
||||
)
|
||||
.bind(&refreshed.access_token)
|
||||
.bind(
|
||||
serde_json::to_string(&serde_json::json!({
|
||||
"expires_at": refreshed.expires_at,
|
||||
"refreshed_at": chrono::Utc::now().to_rfc3339(),
|
||||
}))
|
||||
.map_err(|err| AppError::Internal(err.into()))?,
|
||||
)
|
||||
.bind(account_id)
|
||||
.bind(site_id)
|
||||
.execute(&*self.db)
|
||||
.await?;
|
||||
|
||||
Ok(refreshed)
|
||||
}
|
||||
|
||||
pub fn token_expired(&self, extra_config: Option<&str>) -> bool {
|
||||
token_expired(extra_config)
|
||||
}
|
||||
|
||||
pub async fn list_route_units(&self, site_id: Option<i64>) -> AppResult<Vec<OAuthRouteUnit>> {
|
||||
if let Some(site_id) = site_id {
|
||||
sqlx::query_as("SELECT * FROM oauth_route_units WHERE site_id = ? ORDER BY id")
|
||||
.bind(site_id)
|
||||
.fetch_all(&*self.db)
|
||||
.await
|
||||
.map_err(Into::into)
|
||||
} else {
|
||||
sqlx::query_as("SELECT * FROM oauth_route_units ORDER BY id")
|
||||
.fetch_all(&*self.db)
|
||||
.await
|
||||
.map_err(Into::into)
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn create_route_unit(
|
||||
&self,
|
||||
site_id: i64,
|
||||
provider: &str,
|
||||
name: &str,
|
||||
strategy: Option<&str>,
|
||||
) -> AppResult<OAuthRouteUnit> {
|
||||
let result = sqlx::query(
|
||||
"INSERT INTO oauth_route_units (site_id, provider, name, strategy)
|
||||
VALUES (?, ?, ?, ?)",
|
||||
)
|
||||
.bind(site_id)
|
||||
.bind(provider)
|
||||
.bind(name)
|
||||
.bind(strategy.unwrap_or("round_robin"))
|
||||
.execute(&*self.db)
|
||||
.await?;
|
||||
|
||||
sqlx::query_as("SELECT * FROM oauth_route_units WHERE id = ?")
|
||||
.bind(result.last_insert_rowid())
|
||||
.fetch_one(&*self.db)
|
||||
.await
|
||||
.map_err(Into::into)
|
||||
}
|
||||
|
||||
async fn perform_refresh(
|
||||
&self,
|
||||
site_id: i64,
|
||||
refresh_token: &str,
|
||||
) -> AppResult<RefreshedToken> {
|
||||
let token_endpoint: Option<String> = sqlx::query_scalar(
|
||||
"SELECT json_extract(custom_headers, '$.oauth_token_url') FROM sites WHERE id = ?",
|
||||
)
|
||||
.bind(site_id)
|
||||
.fetch_optional(&*self.db)
|
||||
.await?;
|
||||
|
||||
let Some(token_endpoint) = token_endpoint.filter(|url| !url.trim().is_empty()) else {
|
||||
return Err(AppError::BadRequest(
|
||||
"OAuth token endpoint is not configured for this site".into(),
|
||||
));
|
||||
};
|
||||
|
||||
let response: Value = self
|
||||
.client
|
||||
.post(token_endpoint)
|
||||
.form(&[
|
||||
("grant_type", "refresh_token"),
|
||||
("refresh_token", refresh_token),
|
||||
])
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| AppError::UpstreamError(format!("OAuth refresh failed: {err}")))?
|
||||
.json()
|
||||
.await
|
||||
.map_err(|err| AppError::UpstreamError(format!("Invalid OAuth response: {err}")))?;
|
||||
|
||||
let access_token = response
|
||||
.get("access_token")
|
||||
.and_then(Value::as_str)
|
||||
.ok_or_else(|| AppError::UpstreamError("OAuth response missing access_token".into()))?
|
||||
.to_string();
|
||||
let expires_at = response
|
||||
.get("expires_at")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_string);
|
||||
|
||||
Ok(RefreshedToken {
|
||||
account_id: 0,
|
||||
access_token,
|
||||
expires_at,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn token_expired(extra_config: Option<&str>) -> bool {
|
||||
let Some(extra_config) = extra_config else {
|
||||
return true;
|
||||
};
|
||||
let Ok(value) = serde_json::from_str::<Value>(extra_config) else {
|
||||
return true;
|
||||
};
|
||||
let Some(expires_at) = value.get("expires_at").and_then(Value::as_str) else {
|
||||
return true;
|
||||
};
|
||||
chrono::DateTime::parse_from_rfc3339(expires_at)
|
||||
.map(|expires_at| expires_at <= chrono::Utc::now() + chrono::Duration::minutes(5))
|
||||
.unwrap_or(true)
|
||||
}
|
||||
|
||||
fn token_expiry(extra_config: Option<&str>) -> Option<String> {
|
||||
extra_config
|
||||
.and_then(|config| serde_json::from_str::<Value>(config).ok())
|
||||
.and_then(|value| {
|
||||
value
|
||||
.get("expires_at")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_string)
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
use crate::{error::AppError, error::AppResult};
|
||||
use sqlx::SqlitePool;
|
||||
|
||||
pub async fn enforce_key_request_limit(
|
||||
db: &SqlitePool,
|
||||
key_id: i64,
|
||||
increment: i64,
|
||||
) -> AppResult<()> {
|
||||
let row: Option<(Option<i64>, Option<i64>, Option<bool>)> = sqlx::query_as(
|
||||
"SELECT max_requests, used_requests, enabled
|
||||
FROM downstream_api_keys
|
||||
WHERE id = ?",
|
||||
)
|
||||
.bind(key_id)
|
||||
.fetch_optional(db)
|
||||
.await?;
|
||||
|
||||
let (max_requests, used_requests, enabled) =
|
||||
row.ok_or_else(|| AppError::Unauthorized("API key not found".into()))?;
|
||||
if !enabled.unwrap_or(true) {
|
||||
return Err(AppError::Forbidden("API key is disabled".into()));
|
||||
}
|
||||
if let Some(max_requests) = max_requests {
|
||||
if used_requests.unwrap_or(0) + increment > max_requests {
|
||||
return Err(AppError::RateLimited);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn increment_key_usage(db: &SqlitePool, key_id: i64, cost: f64) -> AppResult<()> {
|
||||
sqlx::query(
|
||||
"UPDATE downstream_api_keys
|
||||
SET used_requests = COALESCE(used_requests, 0) + 1,
|
||||
used_cost = COALESCE(used_cost, 0) + ?,
|
||||
last_used_at = datetime('now'),
|
||||
updated_at = datetime('now')
|
||||
WHERE id = ?",
|
||||
)
|
||||
.bind(cost)
|
||||
.bind(key_id)
|
||||
.execute(db)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,115 @@
|
||||
use crate::error::{AppError, AppResult};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use sqlx::SqlitePool;
|
||||
use std::sync::Arc;
|
||||
|
||||
const UPDATE_CONFIG_KEY: &str = "update_center_config";
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct UpdateCenter {
|
||||
db: Arc<SqlitePool>,
|
||||
client: reqwest::Client,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct UpdateInfo {
|
||||
pub current_version: String,
|
||||
pub latest_version: Option<String>,
|
||||
pub update_available: bool,
|
||||
pub release_notes: Option<String>,
|
||||
pub download_url: Option<String>,
|
||||
}
|
||||
|
||||
impl UpdateCenter {
|
||||
pub fn new(db: Arc<SqlitePool>) -> Self {
|
||||
Self {
|
||||
db,
|
||||
client: reqwest::Client::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn check_for_updates(&self) -> AppResult<UpdateInfo> {
|
||||
let config = self.load_config().await?;
|
||||
let current_version = env!("CARGO_PKG_VERSION").to_string();
|
||||
let Some(endpoint) = config
|
||||
.get("check_url")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|url| !url.trim().is_empty())
|
||||
else {
|
||||
return Ok(UpdateInfo {
|
||||
current_version,
|
||||
latest_version: None,
|
||||
update_available: false,
|
||||
release_notes: None,
|
||||
download_url: None,
|
||||
});
|
||||
};
|
||||
|
||||
let payload: Value = self
|
||||
.client
|
||||
.get(endpoint)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| AppError::UpstreamError(format!("Update check failed: {err}")))?
|
||||
.json()
|
||||
.await
|
||||
.map_err(|err| AppError::UpstreamError(format!("Invalid update response: {err}")))?;
|
||||
|
||||
let latest_version = payload
|
||||
.get("version")
|
||||
.or_else(|| payload.get("latestVersion"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_string);
|
||||
let update_available = latest_version
|
||||
.as_deref()
|
||||
.map(|latest| latest != current_version)
|
||||
.unwrap_or(false);
|
||||
|
||||
Ok(UpdateInfo {
|
||||
current_version,
|
||||
latest_version,
|
||||
update_available,
|
||||
release_notes: payload
|
||||
.get("release_notes")
|
||||
.or_else(|| payload.get("releaseNotes"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_string),
|
||||
download_url: payload
|
||||
.get("download_url")
|
||||
.or_else(|| payload.get("downloadUrl"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_string),
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn apply_update(&self, version: &str) -> AppResult<String> {
|
||||
let task_id = format!("update-{}", uuid::Uuid::new_v4());
|
||||
sqlx::query(
|
||||
"INSERT INTO background_tasks
|
||||
(id, task_type, title, status, message, dedupe_key, created_at, updated_at)
|
||||
VALUES (?, 'update-center-deploy', ?, 'pending', ?, ?, datetime('now'), datetime('now'))",
|
||||
)
|
||||
.bind(&task_id)
|
||||
.bind(format!("Apply BoosAPI {version}"))
|
||||
.bind("Update request accepted; an external supervisor should perform download and restart")
|
||||
.bind(format!("update-center:deploy:{version}"))
|
||||
.execute(&*self.db)
|
||||
.await?;
|
||||
Ok(task_id)
|
||||
}
|
||||
|
||||
async fn load_config(&self) -> AppResult<Value> {
|
||||
let raw: Option<String> = sqlx::query_scalar("SELECT value FROM settings WHERE key = ?")
|
||||
.bind(UPDATE_CONFIG_KEY)
|
||||
.fetch_optional(&*self.db)
|
||||
.await?;
|
||||
|
||||
raw.as_deref()
|
||||
.map(serde_json::from_str)
|
||||
.transpose()
|
||||
.map_err(|err| AppError::BadRequest(format!("Invalid update center config: {err}")))?
|
||||
.map(Ok)
|
||||
.unwrap_or_else(|| Ok(Value::Object(Default::default())))
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user