API Keys and Docs

2026-05-21 19:35:02 -04:00
parent 13dff1fb3b
commit 5563c0279c
11 changed files with 1410 additions and 17 deletions
@@ -0,0 +1,236 @@
+use axum::{
+    extract::{Path, Query, State},
+    http::StatusCode,
+    Json,
+};
+use std::collections::HashMap;
+use axum_extra::extract::cookie::SignedCookieJar;
+use sha2::{Sha256, Digest};
+use uuid::Uuid;
+
+use crate::{
+    models::{ApiKeyView, CreateApiKeyRequest, UsagePoint},
+    AppState,
+};
+
+fn get_user_id(jar: &SignedCookieJar) -> Option<String> {
+    jar.get("session_user_id").map(|c| c.value().to_string())
+}
+
+pub fn hash_key(key: &str) -> String {
+    format!("{:x}", Sha256::digest(key.as_bytes()))
+}
+
+fn generate_api_key() -> String {
+    format!(
+        "td_{}{}",
+        Uuid::new_v4().to_string().replace("-", ""),
+        Uuid::new_v4().to_string().replace("-", "")
+    )
+}
+
+pub async fn list_keys(
+    State(state): State<AppState>,
+    jar: SignedCookieJar,
+) -> Result<Json<Vec<ApiKeyView>>, (StatusCode, String)> {
+    let user_id = get_user_id(&jar)
+        .ok_or((StatusCode::UNAUTHORIZED, "Not authenticated".to_string()))?;
+
+    let keys = sqlx::query_as::<_, ApiKeyView>(
+        "SELECT id, name, key_prefix, created_at, last_used_at, rate_limit FROM api_keys WHERE user_id = ? ORDER BY created_at DESC"
+    )
+    .bind(&user_id)
+    .fetch_all(&state.db)
+    .await
+    .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
+
+    Ok(Json(keys))
+}
+
+pub async fn create_key(
+    State(state): State<AppState>,
+    jar: SignedCookieJar,
+    Json(payload): Json<CreateApiKeyRequest>,
+) -> Result<(StatusCode, Json<serde_json::Value>), (StatusCode, String)> {
+    let user_id = get_user_id(&jar)
+        .ok_or((StatusCode::UNAUTHORIZED, "Not authenticated".to_string()))?;
+
+    if payload.name.trim().is_empty() {
+        return Err((StatusCode::BAD_REQUEST, "Key name cannot be empty".to_string()));
+    }
+
+    let existing: Option<(i64,)> = sqlx::query_as("SELECT COUNT(*) FROM api_keys WHERE user_id = ?")
+        .bind(&user_id)
+        .fetch_optional(&state.db)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
+
+    if let Some((count,)) = existing {
+        if count >= 10 {
+            return Err((StatusCode::CONFLICT, "Maximum of 10 API keys per account".to_string()));
+        }
+    }
+
+    let key = generate_api_key();
+    let hash = hash_key(&key);
+    let prefix = key[..11].to_string(); // "td_" + 8 hex chars
+    let id = Uuid::new_v4().to_string();
+    let now = chrono::Utc::now().to_rfc3339();
+
+    sqlx::query(
+        "INSERT INTO api_keys (id, user_id, name, key_hash, key_prefix, created_at) VALUES (?, ?, ?, ?, ?, ?)"
+    )
+    .bind(&id)
+    .bind(&user_id)
+    .bind(&payload.name)
+    .bind(&hash)
+    .bind(&prefix)
+    .bind(&now)
+    .execute(&state.db)
+    .await
+    .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
+
+    Ok((StatusCode::CREATED, Json(serde_json::json!({
+        "id": id,
+        "name": payload.name,
+        "key": key,
+        "prefix": prefix,
+        "created_at": now,
+        "rate_limit": 60,
+    }))))
+}
+
+pub async fn get_aggregate_usage(
+    State(state): State<AppState>,
+    jar: SignedCookieJar,
+    Query(params): Query<HashMap<String, String>>,
+) -> Result<Json<Vec<UsagePoint>>, (StatusCode, String)> {
+    let user_id = get_user_id(&jar)
+        .ok_or((StatusCode::UNAUTHORIZED, "Not authenticated".to_string()))?;
+
+    let period = params.get("period").map(|s| s.as_str()).unwrap_or("1week");
+
+    let usage = match period {
+        "1hr" => {
+            let cutoff = (chrono::Utc::now() - chrono::TimeDelta::hours(1))
+                .format("%Y-%m-%d %H:%M")
+                .to_string();
+            sqlx::query_as::<_, UsagePoint>(
+                "SELECT minute as date, SUM(count) as count
+                 FROM api_key_usage_detail
+                 WHERE key_id IN (SELECT id FROM api_keys WHERE user_id = ?) AND minute >= ?
+                 GROUP BY minute
+                 ORDER BY minute ASC"
+            )
+            .bind(&user_id)
+            .bind(&cutoff)
+            .fetch_all(&state.db)
+            .await
+            .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
+        }
+        "1day" => {
+            let cutoff = (chrono::Utc::now() - chrono::TimeDelta::hours(24))
+                .format("%Y-%m-%d %H:%M")
+                .to_string();
+            sqlx::query_as::<_, UsagePoint>(
+                "SELECT SUBSTR(minute, 1, 13) as date, SUM(count) as count
+                 FROM api_key_usage_detail
+                 WHERE key_id IN (SELECT id FROM api_keys WHERE user_id = ?) AND minute >= ?
+                 GROUP BY SUBSTR(minute, 1, 13)
+                 ORDER BY date ASC"
+            )
+            .bind(&user_id)
+            .bind(&cutoff)
+            .fetch_all(&state.db)
+            .await
+            .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
+        }
+        _ => {
+            let cutoff = (chrono::Utc::now() - chrono::TimeDelta::days(6))
+                .format("%Y-%m-%d")
+                .to_string();
+            sqlx::query_as::<_, UsagePoint>(
+                "SELECT date, SUM(count) as count
+                 FROM api_key_usage
+                 WHERE key_id IN (SELECT id FROM api_keys WHERE user_id = ?) AND date >= ?
+                 GROUP BY date
+                 ORDER BY date ASC"
+            )
+            .bind(&user_id)
+            .bind(&cutoff)
+            .fetch_all(&state.db)
+            .await
+            .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
+        }
+    };
+
+    Ok(Json(usage))
+}
+
+pub async fn regenerate_key(
+    State(state): State<AppState>,
+    jar: SignedCookieJar,
+    Path(key_id): Path<String>,
+) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
+    let user_id = get_user_id(&jar)
+        .ok_or((StatusCode::UNAUTHORIZED, "Not authenticated".to_string()))?;
+
+    let row: Option<(String,)> = sqlx::query_as(
+        "SELECT name FROM api_keys WHERE id = ? AND user_id = ?"
+    )
+    .bind(&key_id)
+    .bind(&user_id)
+    .fetch_optional(&state.db)
+    .await
+    .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
+
+    let (name,) = row.ok_or((StatusCode::NOT_FOUND, "API key not found".to_string()))?;
+
+    let new_key = generate_api_key();
+    let new_hash = hash_key(&new_key);
+    let new_prefix = new_key[..11].to_string();
+    let now = chrono::Utc::now().to_rfc3339();
+
+    sqlx::query(
+        "UPDATE api_keys SET key_hash = ?, key_prefix = ?, created_at = ?, last_used_at = NULL WHERE id = ? AND user_id = ?"
+    )
+    .bind(&new_hash)
+    .bind(&new_prefix)
+    .bind(&now)
+    .bind(&key_id)
+    .bind(&user_id)
+    .execute(&state.db)
+    .await
+    .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
+
+    Ok(Json(serde_json::json!({
+        "id": key_id,
+        "name": name,
+        "key": new_key,
+        "prefix": new_prefix,
+        "created_at": now,
+        "rate_limit": 60,
+    })))
+}
+
+pub async fn delete_key(
+    State(state): State<AppState>,
+    jar: SignedCookieJar,
+    Path(key_id): Path<String>,
+) -> Result<StatusCode, (StatusCode, String)> {
+    let user_id = get_user_id(&jar)
+        .ok_or((StatusCode::UNAUTHORIZED, "Not authenticated".to_string()))?;
+
+    let result = sqlx::query("DELETE FROM api_keys WHERE id = ? AND user_id = ?")
+        .bind(&key_id)
+        .bind(&user_id)
+        .execute(&state.db)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
+
+    if result.rows_affected() == 0 {
+        return Err((StatusCode::NOT_FOUND, "API key not found".to_string()));
+    }
+
+    Ok(StatusCode::NO_CONTENT)
+}
@@ -76,6 +76,35 @@ pub async fn init_schema(pool: &AnyPool) {
            content TEXT NOT NULL,
            created_at TEXT DEFAULT to_char(NOW(), 'YYYY-MM-DD HH24:MI:SS')
        )",
+        "CREATE TABLE IF NOT EXISTS api_keys (
+            id TEXT PRIMARY KEY,
+            user_id TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
+            name TEXT NOT NULL,
+            key_hash TEXT NOT NULL UNIQUE,
+            key_prefix TEXT NOT NULL,
+            created_at TEXT NOT NULL,
+            last_used_at TEXT,
+            rate_limit INTEGER NOT NULL DEFAULT 60
+        )",
+        "CREATE TABLE IF NOT EXISTS api_render_cache (
+            id TEXT PRIMARY KEY,
+            content_hash TEXT NOT NULL UNIQUE,
+            format TEXT NOT NULL,
+            data BYTEA NOT NULL,
+            created_at TEXT NOT NULL
+        )",
+        "CREATE TABLE IF NOT EXISTS api_key_usage (
+            key_id TEXT NOT NULL REFERENCES api_keys(id) ON DELETE CASCADE,
+            date TEXT NOT NULL,
+            count INTEGER NOT NULL DEFAULT 1,
+            PRIMARY KEY(key_id, date)
+        )",
+        "CREATE TABLE IF NOT EXISTS api_key_usage_detail (
+            key_id TEXT NOT NULL REFERENCES api_keys(id) ON DELETE CASCADE,
+            minute TEXT NOT NULL,
+            count INTEGER NOT NULL DEFAULT 1,
+            PRIMARY KEY(key_id, minute)
+        )",
    ];

    for stmt in &statements {
@@ -81,6 +81,35 @@ pub async fn init_schema(pool: &AnyPool) {
            content TEXT NOT NULL,
            created_at TEXT DEFAULT (strftime('%Y-%m-%d %H:%M:%S', 'now'))
        )",
+        "CREATE TABLE IF NOT EXISTS api_keys (
+            id TEXT PRIMARY KEY,
+            user_id TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
+            name TEXT NOT NULL,
+            key_hash TEXT NOT NULL UNIQUE,
+            key_prefix TEXT NOT NULL,
+            created_at TEXT NOT NULL,
+            last_used_at TEXT,
+            rate_limit INTEGER NOT NULL DEFAULT 60
+        )",
+        "CREATE TABLE IF NOT EXISTS api_render_cache (
+            id TEXT PRIMARY KEY,
+            content_hash TEXT NOT NULL UNIQUE,
+            format TEXT NOT NULL,
+            data BLOB NOT NULL,
+            created_at TEXT NOT NULL
+        )",
+        "CREATE TABLE IF NOT EXISTS api_key_usage (
+            key_id TEXT NOT NULL REFERENCES api_keys(id) ON DELETE CASCADE,
+            date TEXT NOT NULL,
+            count INTEGER NOT NULL DEFAULT 1,
+            PRIMARY KEY(key_id, date)
+        )",
+        "CREATE TABLE IF NOT EXISTS api_key_usage_detail (
+            key_id TEXT NOT NULL REFERENCES api_keys(id) ON DELETE CASCADE,
+            minute TEXT NOT NULL,
+            count INTEGER NOT NULL DEFAULT 1,
+            PRIMARY KEY(key_id, minute)
+        )",
    ];

    for stmt in &statements {
@@ -13,6 +13,7 @@ use tower_http::trace::TraceLayer;
 use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};

 mod admin;
+mod api_keys;
 mod auth;
 mod compiler;
 mod db;
@@ -21,6 +22,7 @@ mod folders;
 mod files;
 mod handlers;
 mod models;
+mod public_api;
 mod setup;
 mod world;
 mod collab;
@@ -28,6 +30,8 @@ mod collab;
 use compiler::TypstCompiler;
 use handlers::{compile_handler, export_handler, yjs_handler};

+pub type RateLimiterMap = Arc<Mutex<HashMap<String, (u32, std::time::Instant)>>>;
+
 #[derive(Clone)]
 pub struct AppState {
    pub compiler: Arc<Mutex<TypstCompiler>>,
@@ -35,6 +39,7 @@ pub struct AppState {
    pub db: AnyPool,
    pub key: Key,
    pub registration_enabled: bool,
+    pub rate_limiter: RateLimiterMap,
 }

 impl axum::extract::FromRef<AppState> for Key {
@@ -84,6 +89,7 @@ async fn main() {
        db,
        key,
        registration_enabled,
+        rate_limiter: Arc::new(Mutex::new(HashMap::new())),
    };

    let api_routes = Router::new()
@@ -114,7 +120,14 @@ async fn main() {
        .route("/docs/{id}/invite", post(collab::invite_collaborator))
        .route("/docs/{id}/comments", get(collab::get_comments).post(collab::add_comment))
        .route("/docs/{id}/versions", get(collab::get_versions).post(collab::create_version))
-        .route("/comments/{id}", patch(collab::update_comment).delete(collab::delete_comment));
+        .route("/comments/{id}", patch(collab::update_comment).delete(collab::delete_comment))
+        .route("/keys", get(api_keys::list_keys).post(api_keys::create_key))
+        .route("/keys/usage", get(api_keys::get_aggregate_usage))
+        .route("/keys/{id}", delete(api_keys::delete_key))
+        .route("/keys/{id}/regenerate", post(api_keys::regenerate_key));
+
+    let v1_routes = Router::new()
+        .route("/render", post(public_api::render_handler));

    let yjs_routes = Router::new()
        .route("/{id}", get(yjs_handler));
@@ -123,6 +136,7 @@ async fn main() {

    let app = Router::new()
        .nest("/api", api_routes.layer(TraceLayer::new_for_http()))
+        .nest("/v1", v1_routes.layer(TraceLayer::new_for_http()))
        .nest("/yjs", yjs_routes.layer(TraceLayer::new_for_http()))
        .fallback_service(ServeDir::new(&static_dir).fallback(ServeFile::new(format!("{}/index.html", static_dir))))
        .with_state(state);
@@ -195,6 +195,27 @@ pub struct AdminCreateUserRequest {
    pub is_admin: Option<bool>,
 }

+#[derive(Debug, Serialize, Deserialize, FromRow)]
+pub struct UsagePoint {
+    pub date: String,
+    pub count: i64,
+}
+
+#[derive(Debug, Serialize, Deserialize, FromRow)]
+pub struct ApiKeyView {
+    pub id: String,
+    pub name: String,
+    pub key_prefix: String,
+    pub created_at: String,
+    pub last_used_at: Option<String>,
+    pub rate_limit: i64,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct CreateApiKeyRequest {
+    pub name: String,
+}
+
 #[derive(Debug, Serialize, Deserialize)]
 pub struct SetupRequest {
    pub username: String,
@@ -0,0 +1,211 @@
+use axum::{
+    extract::State,
+    http::{header, StatusCode},
+    response::IntoResponse,
+    Json,
+};
+use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64};
+use serde::Deserialize;
+use sha2::{Sha256, Digest};
+use std::collections::HashMap;
+use uuid::Uuid;
+
+use crate::{api_keys::hash_key, AppState};
+
+#[derive(Deserialize)]
+pub struct RenderRequest {
+    pub code: String,
+    pub format: String,
+    pub files: Option<Vec<InlineFile>>,
+}
+
+#[derive(Deserialize)]
+pub struct InlineFile {
+    pub name: String,
+    pub data: String, // base64-encoded
+}
+
+fn compute_cache_key(format: &str, code: &str, files: &Option<Vec<InlineFile>>) -> String {
+    let mut hasher = Sha256::new();
+    hasher.update(format.as_bytes());
+    hasher.update(b"\x00");
+    hasher.update(code.as_bytes());
+    if let Some(files) = files {
+        let mut pairs: Vec<_> = files.iter().map(|f| (f.name.as_str(), f.data.as_str())).collect();
+        pairs.sort_by_key(|(n, _)| *n);
+        for (name, data) in pairs {
+            hasher.update(b"\x01");
+            hasher.update(name.as_bytes());
+            hasher.update(data.as_bytes());
+        }
+    }
+    format!("{:x}", hasher.finalize())
+}
+
+pub async fn render_handler(
+    State(state): State<AppState>,
+    headers: axum::http::HeaderMap,
+    Json(payload): Json<RenderRequest>,
+) -> impl IntoResponse {
+    // Extract Bearer token
+    let api_key = match headers
+        .get("Authorization")
+        .and_then(|v| v.to_str().ok())
+        .filter(|v| v.starts_with("Bearer "))
+        .map(|v| v[7..].to_string())
+    {
+        Some(k) => k,
+        None => return (StatusCode::UNAUTHORIZED, "Missing or invalid Authorization header. Use: Authorization: Bearer <api-key>").into_response(),
+    };
+
+    if payload.format != "png" && payload.format != "pdf" {
+        return (StatusCode::BAD_REQUEST, "Invalid format. Must be 'png' or 'pdf'").into_response();
+    }
+
+    if payload.code.trim().is_empty() {
+        return (StatusCode::BAD_REQUEST, "code cannot be empty").into_response();
+    }
+
+    let key_hash = hash_key(&api_key);
+
+    let key_row = sqlx::query_as::<_, (String, String, i64)>(
+        "SELECT id, user_id, rate_limit FROM api_keys WHERE key_hash = ?"
+    )
+    .bind(&key_hash)
+    .fetch_optional(&state.db)
+    .await;
+
+    let (key_id, user_id, rate_limit) = match key_row {
+        Ok(Some(row)) => row,
+        Ok(None) => return (StatusCode::UNAUTHORIZED, "Invalid API key").into_response(),
+        Err(e) => {
+            let msg = format!("Database error: {}", e);
+            return (StatusCode::INTERNAL_SERVER_ERROR, msg).into_response();
+        }
+    };
+
+    // Rate limiting: fixed window of 60 seconds
+    {
+        let mut limiter = state.rate_limiter.lock().await;
+        let now = std::time::Instant::now();
+        let window = std::time::Duration::from_secs(60);
+        let entry = limiter.entry(key_id.clone()).or_insert((0u32, now));
+        if now.duration_since(entry.1) > window {
+            entry.0 = 1;
+            entry.1 = now;
+        } else if entry.0 < rate_limit as u32 {
+            entry.0 += 1;
+        } else {
+            return (StatusCode::TOO_MANY_REQUESTS, "Rate limit exceeded. Max requests per minute reached.").into_response();
+        }
+    }
+
+    let now = chrono::Utc::now();
+    let now_str = now.to_rfc3339();
+    let today = now.format("%Y-%m-%d").to_string();
+
+    let _ = sqlx::query("UPDATE api_keys SET last_used_at = ? WHERE id = ?")
+        .bind(&now_str)
+        .bind(&key_id)
+        .execute(&state.db)
+        .await;
+
+    let _ = sqlx::query(
+        "INSERT INTO api_key_usage (key_id, date, count) VALUES (?, ?, 1) \
+         ON CONFLICT (key_id, date) DO UPDATE SET count = api_key_usage.count + 1"
+    )
+    .bind(&key_id)
+    .bind(&today)
+    .execute(&state.db)
+    .await;
+
+    let minute_str = now.format("%Y-%m-%d %H:%M").to_string();
+    let _ = sqlx::query(
+        "INSERT INTO api_key_usage_detail (key_id, minute, count) VALUES (?, ?, 1) \
+         ON CONFLICT (key_id, minute) DO UPDATE SET count = api_key_usage_detail.count + 1"
+    )
+    .bind(&key_id)
+    .bind(&minute_str)
+    .execute(&state.db)
+    .await;
+
+    let cutoff_minute = (chrono::Utc::now() - chrono::TimeDelta::hours(25))
+        .format("%Y-%m-%d %H:%M")
+        .to_string();
+    let _ = sqlx::query("DELETE FROM api_key_usage_detail WHERE minute < ?")
+        .bind(&cutoff_minute)
+        .execute(&state.db)
+        .await;
+
+    // Check cache
+    let cache_key = compute_cache_key(&payload.format, &payload.code, &payload.files);
+    let content_type: &'static str = if payload.format == "pdf" { "application/pdf" } else { "image/png" };
+
+    if let Ok(Some((data, created_at))) = sqlx::query_as::<_, (Vec<u8>, String)>(
+        "SELECT data, created_at FROM api_render_cache WHERE content_hash = ? AND format = ?"
+    )
+    .bind(&cache_key)
+    .bind(&payload.format)
+    .fetch_optional(&state.db)
+    .await
+    {
+        if let Ok(parsed) = chrono::DateTime::parse_from_rfc3339(&created_at) {
+            let age = chrono::Utc::now().signed_duration_since(parsed.with_timezone(&chrono::Utc));
+            if age.num_seconds() < 3600 {
+                return (StatusCode::OK, [(header::CONTENT_TYPE, content_type)], data).into_response();
+            }
+        }
+    }
+
+    // Load user's account files
+    let mut files_map: HashMap<String, Vec<u8>> = HashMap::new();
+    if let Ok(files) = sqlx::query_as::<_, (String, Vec<u8>)>(
+        "SELECT name, data FROM files WHERE owner_id = ?"
+    )
+    .bind(&user_id)
+    .fetch_all(&state.db)
+    .await
+    {
+        for (name, data) in files {
+            files_map.insert(name, data);
+        }
+    }
+
+    // Inline files override account files
+    if let Some(inline_files) = &payload.files {
+        for f in inline_files {
+            if let Ok(decoded) = BASE64.decode(&f.data) {
+                files_map.insert(f.name.clone(), decoded);
+            }
+        }
+    }
+
+    // Compile
+    let compiler = state.compiler.lock().await;
+    let result = match payload.format.as_str() {
+        "pdf" => compiler.export_pdf(payload.code.clone(), files_map),
+        "png" => compiler.export_png(payload.code.clone(), files_map),
+        _ => unreachable!(),
+    };
+    drop(compiler);
+
+    match result {
+        Ok(data) => {
+            // Store in cache (ignore errors — concurrent inserts are fine)
+            let cache_id = Uuid::new_v4().to_string();
+            let _ = sqlx::query(
+                "INSERT INTO api_render_cache (id, content_hash, format, data, created_at) VALUES (?, ?, ?, ?, ?)"
+            )
+            .bind(&cache_id)
+            .bind(&cache_key)
+            .bind(&payload.format)
+            .bind(&data)
+            .bind(&now_str)
+            .execute(&state.db)
+            .await;
+
+            (StatusCode::OK, [(header::CONTENT_TYPE, content_type)], data).into_response()
+        }
+        Err(_) => (StatusCode::UNPROCESSABLE_ENTITY, "Typst compilation failed. Check your code for errors.").into_response(),
+    }
+}