diff --git a/src/lua/record.rs b/src/lua/record.rs index 7409841..9ee52aa 100644 --- a/src/lua/record.rs +++ b/src/lua/record.rs @@ -17,6 +17,7 @@ const INTERNAL_FIELDS: &[&str] = &[ "_schema", "_key_type", "_rkey", + "_repo_override", ]; /// Register the `Record` global constructor and static methods. @@ -42,6 +43,8 @@ pub fn register_record_api( async move { let collection: String = this.raw_get("_collection")?; let schema: mlua::Value = this.raw_get("_schema")?; + let repo_override: Option = this.raw_get("_repo_override")?; + let repo = repo_override.as_deref().unwrap_or_else(|| claims.did()); // Validate required fields against schema if let mlua::Value::Table(ref schema_table) = schema { @@ -62,7 +65,7 @@ pub fn register_record_api( .to_string(); let pds_body = json!({ - "repo": claims.did(), + "repo": repo, "collection": collection, "rkey": rkey, "record": data, @@ -105,7 +108,7 @@ pub fn register_record_api( indexed_at = NOW()"#, ) .bind(uri) - .bind(claims.did()) + .bind(repo) .bind(&collection) .bind(&rkey) .bind(&data) @@ -120,7 +123,7 @@ pub fn register_record_api( // CREATE let rkey: Option = this.raw_get("_rkey")?; let mut pds_body = json!({ - "repo": claims.did(), + "repo": repo, "collection": collection, "record": data, }); @@ -165,7 +168,7 @@ pub fn register_record_api( cid = EXCLUDED.cid"#, ) .bind(uri) - .bind(claims.did()) + .bind(repo) .bind(&collection) .bind(rkey) .bind(&data) @@ -207,6 +210,8 @@ pub fn register_record_api( mlua::Error::runtime("cannot delete a Record that has no _uri") })?; let collection: String = this.raw_get("_collection")?; + let repo_override: Option = this.raw_get("_repo_override")?; + let repo = repo_override.as_deref().unwrap_or_else(|| claims.did()); let rkey = uri .split('/') @@ -215,7 +220,7 @@ pub fn register_record_api( .to_string(); let pds_body = json!({ - "repo": claims.did(), + "repo": repo, "collection": collection, "rkey": rkey, }); @@ -284,6 +289,18 @@ pub fn register_record_api( methods.set("set_rkey", set_rkey_fn)?; } + // :set_repo(did) + { + let set_repo_fn = lua.create_function(|_lua, (this, did): (mlua::Table, String)| { + if did.is_empty() { + return Err(mlua::Error::runtime("did must be a non-empty string")); + } + this.raw_set("_repo_override", did)?; + Ok(this) + })?; + methods.set("set_repo", set_repo_fn)?; + } + // :generate_rkey() { let generate_rkey_fn = lua.create_function(|_lua, this: mlua::Table| { @@ -384,6 +401,7 @@ pub fn register_record_api( table.raw_set("_uri", mlua::Value::Nil)?; table.raw_set("_cid", mlua::Value::Nil)?; table.raw_set("_schema", schema_value.clone())?; + table.raw_set("_repo_override", mlua::Value::Nil)?; // Auto-set _key_type from the lexicon's record_key match lexicon.as_ref().and_then(|l| l.record_key.as_deref()) { @@ -427,7 +445,7 @@ pub fn register_record_api( let session = session.clone(); async move { // Extract save data from each record (sync) - type SaveItem = (mlua::Table, String, Option, Option, Value); + type SaveItem = (mlua::Table, String, Option, Option, Option, Value); let mut save_items: Vec = Vec::new(); for pair in records_table.sequence_values::() { @@ -435,6 +453,7 @@ pub fn register_record_api( let collection: String = record_table.raw_get("_collection")?; let existing_uri: Option = record_table.raw_get("_uri")?; let rkey: Option = record_table.raw_get("_rkey")?; + let repo_override: Option = record_table.raw_get("_repo_override")?; // Validate let schema: mlua::Value = record_table.raw_get("_schema")?; @@ -443,19 +462,21 @@ pub fn register_record_api( } let data = extract_record_data(&lua, &record_table, &collection)?; - save_items.push((record_table, collection, existing_uri, rkey, data)); + save_items.push((record_table, collection, existing_uri, rkey, repo_override, data)); } // Parallel PDS calls - let futs = save_items.iter().map(|(_, collection, existing_uri, rkey, data)| { + let futs = save_items.iter().map(|(_, collection, existing_uri, rkey, repo_override, data)| { let state = state.clone(); let claims = claims.clone(); let session = session.clone(); let collection = collection.clone(); let existing_uri = existing_uri.clone(); let rkey = rkey.clone(); + let repo_override = repo_override.clone(); let data = data.clone(); async move { + let repo = repo_override.as_deref().unwrap_or_else(|| claims.did()); if let Some(ref uri) = existing_uri { let rkey = uri .split('/') @@ -464,7 +485,7 @@ pub fn register_record_api( .to_string(); let pds_body = json!({ - "repo": claims.did(), + "repo": repo, "collection": collection, "rkey": rkey, "record": data, @@ -511,7 +532,7 @@ pub fn register_record_api( indexed_at = NOW()"#, ) .bind(uri.as_str()) - .bind(claims.did()) + .bind(repo) .bind(&collection) .bind(&rkey) .bind(&data) @@ -524,7 +545,7 @@ pub fn register_record_api( Ok(result) } else { let mut pds_body = json!({ - "repo": claims.did(), + "repo": repo, "collection": collection, "record": data, }); @@ -576,7 +597,7 @@ pub fn register_record_api( cid = EXCLUDED.cid"#, ) .bind(uri) - .bind(claims.did()) + .bind(repo) .bind(&collection) .bind(rkey) .bind(&data) @@ -595,7 +616,7 @@ pub fn register_record_api( let results = try_join_all(futs).await?; // Write back _uri and _cid (sync) - for (i, (record_table, _, _, _, _)) in save_items.iter().enumerate() { + for (i, (record_table, _, _, _, _, _)) in save_items.iter().enumerate() { if let Some(result) = results.get(i) { if let Some(uri) = result.get("uri").and_then(|v| v.as_str()) { record_table.raw_set("_uri", uri.to_string())?; @@ -645,6 +666,7 @@ pub fn register_record_api( table.raw_set("_schema", schema_value)?; table.raw_set("_key_type", mlua::Value::Nil)?; table.raw_set("_rkey", mlua::Value::Nil)?; + table.raw_set("_repo_override", mlua::Value::Nil)?; // Copy record fields if let Some(obj) = record.as_object() { @@ -717,6 +739,7 @@ pub fn register_record_api( table.raw_set("_schema", schema_value)?; table.raw_set("_key_type", mlua::Value::Nil)?; table.raw_set("_rkey", mlua::Value::Nil)?; + table.raw_set("_repo_override", mlua::Value::Nil)?; if let Some(obj) = record.as_object() { for (k, v) in obj { diff --git a/web/src/app/dashboard/settings/api-keys/page.tsx b/web/src/app/dashboard/settings/api-keys/page.tsx index d3dd182..c65bf9d 100644 --- a/web/src/app/dashboard/settings/api-keys/page.tsx +++ b/web/src/app/dashboard/settings/api-keys/page.tsx @@ -47,6 +47,7 @@ const PERMISSION_CATEGORIES: Record = { Users: ["users:create", "users:read", "users:update", "users:delete"], "API Keys": ["api-keys:create", "api-keys:read", "api-keys:delete"], Backfill: ["backfill:create", "backfill:read"], + "Rate Limits": ["rate-limits:read", "rate-limits:create", "rate-limits:delete"], System: ["stats:read", "events:read"], }; diff --git a/web/src/app/dashboard/settings/page.tsx b/web/src/app/dashboard/settings/page.tsx index a2c245f..c86cc49 100644 --- a/web/src/app/dashboard/settings/page.tsx +++ b/web/src/app/dashboard/settings/page.tsx @@ -16,6 +16,8 @@ export default function SettingsPage() { router.replace("/dashboard/settings/env-variables"); } else if (hasPermission("api-keys:read")) { router.replace("/dashboard/settings/api-keys"); + } else if (hasPermission("rate-limits:read")) { + router.replace("/dashboard/settings/rate-limits"); } }, [router, hasPermission]); diff --git a/web/src/app/dashboard/settings/rate-limits/page.tsx b/web/src/app/dashboard/settings/rate-limits/page.tsx new file mode 100644 index 0000000..3f8a983 --- /dev/null +++ b/web/src/app/dashboard/settings/rate-limits/page.tsx @@ -0,0 +1,530 @@ +"use client"; + +import { useCallback, useEffect, useState } from "react"; +import { Trash2 } from "lucide-react"; + +import { useAuth } from "@/lib/auth-context"; +import { useCurrentUser } from "@/hooks/use-current-user"; +import { + getRateLimits, + upsertRateLimit, + deleteRateLimit, + setRateLimitEnabled, + addAllowlistEntry, + removeAllowlistEntry, +} from "@/lib/api"; +import type { RateLimitSummary, AllowlistEntry } from "@/types/rate-limits"; +import { SiteHeader } from "@/components/site-header"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { Switch } from "@/components/ui/switch"; +import { + ResponsiveDialog, + ResponsiveDialogClose, + ResponsiveDialogContent, + ResponsiveDialogDescription, + ResponsiveDialogFooter, + ResponsiveDialogHeader, + ResponsiveDialogTitle, + ResponsiveDialogTrigger, +} from "@/components/ui/responsive-dialog"; +import { + Table, + TableBody, + TableCell, + TableHead, + TableHeader, + TableRow, +} from "@/components/ui/table"; + +export default function RateLimitsPage() { + const { getToken } = useAuth(); + const { hasPermission } = useCurrentUser(); + const [enabled, setEnabled] = useState(false); + const [limits, setLimits] = useState([]); + const [allowlist, setAllowlist] = useState([]); + const [error, setError] = useState(null); + const [toggling, setToggling] = useState(false); + + const load = useCallback(() => { + getRateLimits(getToken) + .then((data) => { + setEnabled(data.enabled); + setLimits(data.limits); + setAllowlist(data.allowlist); + }) + .catch((e) => setError(e.message)); + }, [getToken]); + + useEffect(() => { + load(); + }, [load]); + + async function handleToggleEnabled(checked: boolean) { + setToggling(true); + try { + await setRateLimitEnabled(getToken, { enabled: checked }); + setEnabled(checked); + } catch (e: unknown) { + setError(e instanceof Error ? e.message : String(e)); + } finally { + setToggling(false); + } + } + + async function handleDeleteLimit(id: number) { + try { + await deleteRateLimit(getToken, id); + load(); + } catch (e: unknown) { + setError(e instanceof Error ? e.message : String(e)); + } + } + + async function handleRemoveAllowlistEntry(id: number) { + try { + await removeAllowlistEntry(getToken, id); + load(); + } catch (e: unknown) { + setError(e instanceof Error ? e.message : String(e)); + } + } + + return ( + <> + +
+ {error &&

{error}

} + + {/* Global toggle */} +
+ + +
+ + {/* Rate limit rules */} +
+
+
+

Rate Limit Rules

+

+ Configure global defaults and per-method overrides. +

+
+ {hasPermission("rate-limits:create") && ( + + )} +
+ +
+ + + + Method + Capacity + Refill Rate + Updated + + + + + {limits.length === 0 && ( + + + No rate limit rules yet. + + + )} + {limits.map((limit) => ( + + + {limit.method ?? ( + + Global default + + )} + + {limit.capacity} + {limit.refill_rate} tokens/sec + + {new Date(limit.updated_at).toLocaleString()} + + +
+ {hasPermission("rate-limits:create") && ( + + )} + {hasPermission("rate-limits:delete") && limit.method !== null && ( + handleDeleteLimit(limit.id)} + /> + )} +
+
+
+ ))} +
+
+
+
+ + {/* IP allowlist */} +
+
+
+

IP Allowlist

+

+ IPs or CIDRs that bypass rate limiting. +

+
+ {hasPermission("rate-limits:create") && ( + + )} +
+ +
+ + + + CIDR + Note + Created + + + + + {allowlist.length === 0 && ( + + + No allowlist entries yet. + + + )} + {allowlist.map((entry) => ( + + + {entry.cidr} + + + {entry.note ?? "—"} + + + {new Date(entry.created_at).toLocaleString()} + + + {hasPermission("rate-limits:delete") && ( + handleRemoveAllowlistEntry(entry.id)} + /> + )} + + + ))} + +
+
+
+
+ + ); +} + +function UpsertRuleDiag({ + getToken, + onSuccess, + existing, +}: { + getToken: () => Promise; + onSuccess: () => void; + existing?: RateLimitSummary; +}) { + const [method, setMethod] = useState(existing?.method ?? ""); + const [capacity, setCapacity] = useState(String(existing?.capacity ?? "")); + const [refillRate, setRefillRate] = useState( + String(existing?.refill_rate ?? "") + ); + const [error, setError] = useState(null); + const [open, setOpen] = useState(false); + + const isEdit = !!existing; + + async function handleSubmit() { + setError(null); + const cap = Number(capacity); + const rate = Number(refillRate); + if (!cap || cap <= 0 || !rate || rate <= 0) { + setError("Capacity and refill rate must be positive numbers."); + return; + } + try { + const body: { method?: string; capacity: number; refill_rate: number } = { + capacity: cap, + refill_rate: rate, + }; + if (isEdit && existing.method !== null) { + body.method = existing.method; + } else if (!isEdit && method.trim()) { + body.method = method.trim(); + } + await upsertRateLimit(getToken, body); + setOpen(false); + if (!isEdit) { + setMethod(""); + setCapacity(""); + setRefillRate(""); + } + onSuccess(); + } catch (e: unknown) { + setError(e instanceof Error ? e.message : String(e)); + } + } + + return ( + { + setOpen(o); + if (o) { + setMethod(existing?.method ?? ""); + setCapacity(String(existing?.capacity ?? "")); + setRefillRate(String(existing?.refill_rate ?? "")); + setError(null); + } + }} + > + + {isEdit ? ( + + ) : ( + + )} + + + + + {isEdit ? "Edit Rule" : "Add Rule"} + + + {isEdit + ? `Update rate limit for ${existing.method ?? "global default"}.` + : "Add a rate limit rule. Leave method empty for global default."} + + +
+ {error &&

{error}

} + {!isEdit && ( +
+ + setMethod(e.target.value)} + placeholder="com.atproto.sync.getRepo" + className="font-mono" + /> +
+ )} +
+ + setCapacity(e.target.value)} + placeholder="100" + /> +
+
+ + setRefillRate(e.target.value)} + placeholder="10" + /> +
+
+ + + + + + +
+
+ ); +} + +function AddAllowlistDialog({ + getToken, + onSuccess, +}: { + getToken: () => Promise; + onSuccess: () => void; +}) { + const [cidr, setCidr] = useState(""); + const [note, setNote] = useState(""); + const [error, setError] = useState(null); + const [open, setOpen] = useState(false); + + async function handleAdd() { + setError(null); + try { + const body: { cidr: string; note?: string } = { cidr: cidr.trim() }; + if (note.trim()) body.note = note.trim(); + await addAllowlistEntry(getToken, body); + setCidr(""); + setNote(""); + setOpen(false); + onSuccess(); + } catch (e: unknown) { + setError(e instanceof Error ? e.message : String(e)); + } + } + + return ( + { + setOpen(o); + if (o) { + setCidr(""); + setNote(""); + setError(null); + } + }} + > + + + + + + Add Allowlist Entry + + Add an IP or CIDR range that bypasses rate limiting. + + +
+ {error &&

{error}

} +
+ + setCidr(e.target.value)} + placeholder="10.0.0.0/8" + className="font-mono" + /> +
+
+ + setNote(e.target.value)} + placeholder="Internal network" + /> +
+
+ + + + + + +
+
+ ); +} + +function DeleteConfirmDialog({ + title, + description, + onConfirm, +}: { + title: string; + description: string; + onConfirm: () => void; +}) { + const [open, setOpen] = useState(false); + const [deleting, setDeleting] = useState(false); + + async function handleConfirm() { + setDeleting(true); + try { + await onConfirm(); + setOpen(false); + } finally { + setDeleting(false); + } + } + + return ( + + + + + + + {title} + + {description} + + + + + + + + + + + ); +} diff --git a/web/src/app/dashboard/settings/users/page.tsx b/web/src/app/dashboard/settings/users/page.tsx index 93cf207..6103cd1 100644 --- a/web/src/app/dashboard/settings/users/page.tsx +++ b/web/src/app/dashboard/settings/users/page.tsx @@ -55,6 +55,7 @@ const PERMISSION_CATEGORIES: Record = { Users: ["users:create", "users:read", "users:update", "users:delete"], "API Keys": ["api-keys:create", "api-keys:read", "api-keys:delete"], Backfill: ["backfill:create", "backfill:read"], + "Rate Limits": ["rate-limits:read", "rate-limits:create", "rate-limits:delete"], System: ["stats:read", "events:read"], }; diff --git a/web/src/components/app-sidebar.tsx b/web/src/components/app-sidebar.tsx index 22cb215..58768e7 100644 --- a/web/src/components/app-sidebar.tsx +++ b/web/src/components/app-sidebar.tsx @@ -13,6 +13,7 @@ import { IconVariable, IconTag, IconChevronRight, + IconShield, } from "@tabler/icons-react" import Image from "next/image" import Link from "next/link" @@ -53,6 +54,7 @@ const settingsSubItems = [ { title: "ENV Variables", url: "/dashboard/settings/env-variables", icon: IconVariable, requiredPermissions: ["script-variables:read"] }, { title: "API Keys", url: "/dashboard/settings/api-keys", icon: IconKey, requiredPermissions: ["api-keys:read"] }, { title: "Labelers", url: "/dashboard/settings/labelers", icon: IconTag, requiredPermissions: ["labelers:read"] }, + { title: "Rate Limits", url: "/dashboard/settings/rate-limits", icon: IconShield, requiredPermissions: ["rate-limits:read"] }, ] as const export function AppSidebar({ diff --git a/web/src/lib/api.ts b/web/src/lib/api.ts index 5da82c4..2143296 100644 --- a/web/src/lib/api.ts +++ b/web/src/lib/api.ts @@ -11,6 +11,7 @@ import type { AdminListRecordsResponse } from "@/types/records" import type { EventsListResponse } from "@/types/events" import type { ScriptVariableSummary } from "@/types/script-variables" import type { LabelerSummary } from "@/types/labelers" +import type { RateLimitsResponse } from "@/types/rate-limits" export type { ApiKeySummary, CreateApiKeyResponse } from "@/types/api-keys" export type { CollectionStat, StatsResponse } from "@/types/stats" @@ -24,6 +25,7 @@ export type { EventLogEntry, EventsListResponse } from "@/types/events" export type { ScriptVariableSummary } from "@/types/script-variables" export type { LabelerSummary } from "@/types/labelers" export type { RecordLabel } from "@/types/records" +export type { RateLimitSummary, AllowlistEntry, RateLimitsResponse } from "@/types/rate-limits" // The DPoP proof for admin API calls must target AIP's userinfo URL, // because the backend forwards the proof to AIP for token validation. @@ -361,6 +363,59 @@ export function deleteLabeler( }) } +// Rate Limits +export function getRateLimits(getToken: () => Promise) { + return apiFetch("/admin/rate-limits", getToken) +} + +export function upsertRateLimit( + getToken: () => Promise, + body: { method?: string; capacity: number; refill_rate: number } +) { + return apiFetch("/admin/rate-limits", getToken, { + method: "POST", + body: JSON.stringify(body), + }) +} + +export function deleteRateLimit( + getToken: () => Promise, + id: number +) { + return apiFetch(`/admin/rate-limits/${encodeURIComponent(id)}`, getToken, { + method: "DELETE", + }) +} + +export function setRateLimitEnabled( + getToken: () => Promise, + body: { enabled: boolean } +) { + return apiFetch("/admin/rate-limits/enabled", getToken, { + method: "PUT", + body: JSON.stringify(body), + }) +} + +export function addAllowlistEntry( + getToken: () => Promise, + body: { cidr: string; note?: string } +) { + return apiFetch("/admin/rate-limits/allowlist", getToken, { + method: "POST", + body: JSON.stringify(body), + }) +} + +export function removeAllowlistEntry( + getToken: () => Promise, + id: number +) { + return apiFetch(`/admin/rate-limits/allowlist/${encodeURIComponent(id)}`, getToken, { + method: "DELETE", + }) +} + // Event Logs export function getEvents( getToken: () => Promise, diff --git a/web/src/types/rate-limits.ts b/web/src/types/rate-limits.ts new file mode 100644 index 0000000..8babcdf --- /dev/null +++ b/web/src/types/rate-limits.ts @@ -0,0 +1,21 @@ +export interface RateLimitSummary { + id: number + method: string | null + capacity: number + refill_rate: number + created_at: string + updated_at: string +} + +export interface AllowlistEntry { + id: number + cidr: string + note: string | null + created_at: string +} + +export interface RateLimitsResponse { + enabled: boolean + limits: RateLimitSummary[] + allowlist: AllowlistEntry[] +}