File size: 4,236 Bytes
6ead73b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | use axum::{
Router,
routing::{get, post},
extract::State,
Json,
http::{StatusCode, HeaderMap},
};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use crate::AppState;
use crate::rate_limit::RateLimitResult;
use crate::hf_router::HfError;
#[derive(Deserialize)]
pub struct PlayRequest {
pub repo_id: String,
pub path: String,
#[serde(rename = "type")]
pub req_type: Option<String>,
}
#[derive(Serialize)]
pub struct PlayResponse {
pub url: String,
pub expires_in: u64,
}
pub fn routes() -> Router<Arc<AppState>> {
Router::new()
.route("/play", post(play_handler))
.route("/play/bot", post(play_handler))
.route("/play/tma", post(play_handler))
.route("/play/channel", post(play_handler))
.route("/deliver", post(deliver_handler))
.route("/keys", get(keys_handler))
.route("/stats", get(stats_handler))
}
async fn play_handler(
State(state): State<Arc<AppState>>,
headers: HeaderMap,
Json(req): Json<PlayRequest>,
) -> Result<Json<PlayResponse>, (StatusCode, Json<serde_json::Value>)> {
let api_key = headers
.get("authorization")
.and_then(|v| v.to_str().ok())
.and_then(|v| v.strip_prefix("Bearer "))
.ok_or_else(|| {
(StatusCode::UNAUTHORIZED, Json(serde_json::json!({"error": "Missing Authorization header"})))
})?;
let client = state.auth.validate(api_key).await
.ok_or_else(|| {
(StatusCode::UNAUTHORIZED, Json(serde_json::json!({"error": "Invalid API key"})))
})?;
if !client.active {
return Err((StatusCode::FORBIDDEN, Json(serde_json::json!({"error": "Client account is suspended"}))));
}
let rate_result = state.rate_limiter.check_and_increment(
&client.id,
client.rate_limit_per_minute,
client.monthly_limit,
);
match rate_result {
RateLimitResult::MinuteExceeded => {
return Err((StatusCode::TOO_MANY_REQUESTS, Json(serde_json::json!({
"error": "Rate limit exceeded", "limit": "per_minute"
}))));
}
RateLimitResult::MonthlyExceeded => {
return Err((StatusCode::TOO_MANY_REQUESTS, Json(serde_json::json!({
"error": "Monthly quota exceeded", "limit": "monthly"
}))));
}
RateLimitResult::Ok => {}
}
let resolve_result = state.hf_router.resolve(&req.repo_id, &req.path, req.req_type.as_deref()).await
.map_err(|e| {
let (code, msg) = match e {
HfError::NotFound => (StatusCode::NOT_FOUND, "File not found on Hugging Face".to_string()),
HfError::RateLimited => (StatusCode::TOO_MANY_REQUESTS, "HF API rate limited".to_string()),
HfError::Network(s) => (StatusCode::BAD_GATEWAY, format!("HF network error: {}", s)),
HfError::NoRedirect(s) => (StatusCode::BAD_GATEWAY, format!("HF redirect error: {}", s)),
};
(code, Json(serde_json::json!({"error": msg})))
})?;
Ok(Json(PlayResponse {
url: resolve_result.url,
expires_in: resolve_result.expires_in,
}))
}
#[derive(Deserialize)]
pub struct DeliverRequest {
pub user_id: i64,
pub chat_id: i64,
pub repo_id: String,
pub path: String,
}
async fn deliver_handler(
State(_state): State<Arc<AppState>>,
Json(_req): Json<DeliverRequest>,
) -> Result<Json<serde_json::Value>, (StatusCode, Json<serde_json::Value>)> {
Err((StatusCode::NOT_IMPLEMENTED, Json(serde_json::json!({"error": "Not implemented"}))))
}
#[derive(Serialize)]
pub struct KeyInfo {
pub client_id: String,
pub plan: String,
pub keys: Vec<crate::auth::ClientSubKey>,
}
async fn keys_handler(
State(_state): State<Arc<AppState>>,
) -> Result<Json<KeyInfo>, (StatusCode, Json<serde_json::Value>)> {
Err((StatusCode::NOT_IMPLEMENTED, Json(serde_json::json!({"error": "Not implemented"}))))
}
async fn stats_handler(
State(_state): State<Arc<AppState>>,
) -> Result<Json<serde_json::Value>, (StatusCode, Json<serde_json::Value>)> {
Err((StatusCode::NOT_IMPLEMENTED, Json(serde_json::json!({"error": "Not implemented"}))))
}
|