diff --git a/package.json b/package.json index d1075b233..7b2502bf9 100644 --- a/package.json +++ b/package.json @@ -19,7 +19,7 @@ "test": "vitest run", "test:watch": "vitest", "test:edge": "deno test --allow-read --allow-env=DEMO_CODE_PEPPER,EXA_API_KEY,SUPABASE_URL,SUPABASE_SERVICE_ROLE_KEY,SPARRING_INTEGRITY_SECRET,OPENROUTER_API_KEY --allow-net=127.0.0.1 supabase/functions/", - "test:db": "supabase test db", + "test:db": "supabase test db && bash scripts/test-daily-quota.sh", "test:e2e": "playwright test --project=desktop-chromium --project=mobile-chromium", "test:e2e:webkit": "playwright test --project=webkit", "test:e2e:staging": "playwright test --project=staging", diff --git a/scripts/test-daily-quota.sh b/scripts/test-daily-quota.sh new file mode 100644 index 000000000..224dc13cc --- /dev/null +++ b/scripts/test-daily-quota.sh @@ -0,0 +1,153 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +database_url="${SUPABASE_DB_URL:-postgresql://postgres:postgres@127.0.0.1:54322/postgres}" +database_password="${SUPABASE_DB_PASSWORD:-postgres}" +test_user_id="66666666-6666-4666-8666-666666666666" +result_dir="$(mktemp -d)" +first_pid="" +second_pid="" + +cleanup() { + touch "$result_dir/release-first" + if [ -n "$first_pid" ]; then + wait "$first_pid" 2>/dev/null || true + fi + if [ -n "$second_pid" ]; then + wait "$second_pid" 2>/dev/null || true + fi + PGPASSWORD="$database_password" psql "$database_url" -qX --set ON_ERROR_STOP=1 \ + --command "DELETE FROM auth.users WHERE id = '$test_user_id'" >/dev/null + rm -rf "$result_dir" +} +trap cleanup EXIT + +( + cd "$repo_root" + PGPASSWORD="$database_password" psql "$database_url" -qX --set ON_ERROR_STOP=1 <<'SQL' +BEGIN; +INSERT INTO auth.users (id) +VALUES ('77777777-7777-4777-8777-777777777777'); +DROP FUNCTION public.consume_daily_quota(uuid, integer, text); +CREATE FUNCTION public.consume_daily_quota(uuid, integer) +RETURNS TABLE (allowed boolean, current_count integer) +LANGUAGE sql +AS $$ SELECT false, 0 $$; +ALTER TABLE public.daily_usage_counters DROP CONSTRAINT daily_usage_counters_pkey; +ALTER TABLE public.daily_usage_counters DROP COLUMN quota_key; +ALTER TABLE public.daily_usage_counters + ADD CONSTRAINT daily_usage_counters_pkey PRIMARY KEY (user_id, usage_date); +INSERT INTO public.daily_usage_counters (user_id, usage_date, request_count) +VALUES ('77777777-7777-4777-8777-777777777777', CURRENT_DATE - 1, 7); +\i supabase/migrations/20260920125555_scope_daily_usage_counters.sql +DO $$ +BEGIN + IF NOT EXISTS ( + SELECT 1 + FROM public.daily_usage_counters + WHERE user_id = '77777777-7777-4777-8777-777777777777' + AND usage_date = CURRENT_DATE - 1 + AND quota_key = 'ai' + AND request_count = 7 + ) THEN + RAISE EXCEPTION 'legacy daily usage row was not preserved under the ai quota'; + END IF; +END; +$$; +ROLLBACK; +SQL +) +echo "Legacy quota row migration: count 7 preserved under ai" + +PGPASSWORD="$database_password" psql "$database_url" -qX --set ON_ERROR_STOP=1 \ + --command "DELETE FROM auth.users WHERE id = '$test_user_id'; INSERT INTO auth.users (id) VALUES ('$test_user_id')" \ + >/dev/null + +reserve_first() { + PGPASSWORD="$database_password" psql "$database_url" -qAtX --set ON_ERROR_STOP=1 <"$result_dir/first" & +first_pid=$! + +for _ in {1..100}; do + if [ -e "$result_dir/first-holds-lock" ]; then + break + fi + if ! kill -0 "$first_pid" 2>/dev/null; then + wait "$first_pid" + fi + sleep 0.05 +done +if [ ! -e "$result_dir/first-holds-lock" ]; then + echo "first quota reservation did not acquire its row lock" + exit 1 +fi + +reserve_second >"$result_dir/second" & +second_pid=$! + +second_is_waiting=false +for _ in {1..100}; do + if [ "$(PGPASSWORD="$database_password" psql "$database_url" -qAtX --set ON_ERROR_STOP=1 --command \ + "SELECT EXISTS (SELECT 1 FROM pg_stat_activity WHERE pid <> pg_backend_pid() AND wait_event_type = 'Lock' AND query LIKE '%daily_quota_concurrency_second%')")" = "t" ]; then + second_is_waiting=true + break + fi + if ! kill -0 "$second_pid" 2>/dev/null; then + break + fi + sleep 0.05 +done +if [ "$second_is_waiting" != true ]; then + echo "second quota reservation did not wait on the first row lock" + exit 1 +fi + +touch "$result_dir/release-first" + +first_result=0 +second_result=0 +wait "$first_pid" || first_result=$? +wait "$second_pid" || second_result=$? +if [ "$first_result" -ne 0 ] || [ "$second_result" -ne 0 ]; then + exit 1 +fi + +actual_results="$(grep -hE '^[tf]\|[0-9]+$' "$result_dir/first" "$result_dir/second" | sort)" +expected_results=$'f|1\nt|1' +if [ "$actual_results" != "$expected_results" ]; then + echo "unexpected concurrent quota results:" + echo "$actual_results" + exit 1 +fi + +stored_count="$(PGPASSWORD="$database_password" psql "$database_url" -qAtX --set ON_ERROR_STOP=1 --command \ + "SELECT request_count FROM public.daily_usage_counters WHERE user_id = '$test_user_id' AND usage_date = CURRENT_DATE AND quota_key = 'concurrent'")" +if [ "$stored_count" != "1" ]; then + echo "expected one stored concurrent reservation, got: $stored_count" + exit 1 +fi + +echo "Concurrent quota reservations: 1 allowed, 1 rejected, stored count 1" diff --git a/supabase/functions/import-source/index.ts b/supabase/functions/import-source/index.ts index e6f5977b0..801a2e38f 100644 --- a/supabase/functions/import-source/index.ts +++ b/supabase/functions/import-source/index.ts @@ -49,6 +49,7 @@ async function consumeDailyImportQuota(userId: string): Promise<{ allowed: boole const { data, error } = await supabase.rpc("consume_daily_quota", { p_user_id: userId, p_limit: DAILY_IMPORT_LIMIT, + p_quota_key: "import_source", }); if (error) throw error; const row = Array.isArray(data) ? data[0] : data; @@ -238,7 +239,7 @@ function normalizeText(text: string) { .slice(0, MAX_CONTENT_CHARS); } -Deno.serve(async (req) => { +export async function handleRequest(req: Request): Promise { const corsHeaders = getCorsHeaders(req.headers.get("origin")); if (req.method === "OPTIONS") { @@ -331,4 +332,8 @@ Deno.serve(async (req) => { await logImport(user.id, "error", "exception", Date.now() - startedAt); return jsonResponse({ error: message }, 400, corsHeaders); } -}); +} + +if (import.meta.main) { + Deno.serve(handleRequest); +} diff --git a/supabase/functions/import-source/index_test.ts b/supabase/functions/import-source/index_test.ts new file mode 100644 index 000000000..099fe1a56 --- /dev/null +++ b/supabase/functions/import-source/index_test.ts @@ -0,0 +1,70 @@ +const environment = { + SUPABASE_URL: "https://supabase.test", + SUPABASE_SERVICE_ROLE_KEY: "test-service-role-key", +}; + +for (const [key, value] of Object.entries(environment)) { + Deno.env.set(key, value); +} + +const { handleRequest } = await import("./index.ts"); + +function jsonResponse(body: unknown, status = 200): Response { + return Response.json(body, { status }); +} + +Deno.test("import-source reserves only import_source quota", async () => { + const originalFetch = globalThis.fetch; + let quotaRequest: Record | undefined; + + globalThis.fetch = + (async (input: Request | URL | string, init?: RequestInit) => { + const url = new URL(input instanceof Request ? input.url : String(input)); + + if (url.origin !== "https://supabase.test") { + throw new Error(`Unexpected request: ${url}`); + } + if (url.pathname === "/auth/v1/user") { + return jsonResponse({ id: "user-1", aud: "authenticated" }); + } + if (url.pathname.includes("/rest/v1/user_roles")) { + return jsonResponse({ role: "user" }); + } + if (url.pathname.includes("/rest/v1/rpc/consume_daily_quota")) { + quotaRequest = JSON.parse(String(init?.body)) as Record< + string, + unknown + >; + return jsonResponse([{ allowed: false, current_count: 50 }]); + } + throw new Error(`Unexpected Supabase request: ${url}`); + }) as typeof fetch; + + try { + const response = await handleRequest( + new Request("http://127.0.0.1/functions/v1/import-source", { + method: "POST", + headers: { + Authorization: "Bearer test-user-token", + "Content-Type": "application/json", + }, + body: JSON.stringify({ url: "https://example.com/article" }), + }), + ); + + if (response.status !== 429) { + throw new Error( + `import-source returned ${response.status}: ${await response.text()}`, + ); + } + if (quotaRequest?.p_quota_key !== "import_source") { + throw new Error( + `import-source reserved the wrong quota: ${ + JSON.stringify(quotaRequest) + }`, + ); + } + } finally { + globalThis.fetch = originalFetch; + } +}); diff --git a/supabase/migrations/20260920125555_scope_daily_usage_counters.sql b/supabase/migrations/20260920125555_scope_daily_usage_counters.sql new file mode 100644 index 000000000..602a4c460 --- /dev/null +++ b/supabase/migrations/20260920125555_scope_daily_usage_counters.sql @@ -0,0 +1,59 @@ +-- Keep independently configured daily quotas in separate atomic counters. +ALTER TABLE public.daily_usage_counters + ADD COLUMN quota_key text NOT NULL DEFAULT 'ai'; + +ALTER TABLE public.daily_usage_counters + DROP CONSTRAINT daily_usage_counters_pkey; + +ALTER TABLE public.daily_usage_counters + ADD CONSTRAINT daily_usage_counters_pkey + PRIMARY KEY (user_id, usage_date, quota_key); + +-- The new defaulted parameter changes the PostgreSQL function signature, so +-- remove the old overload before creating its replacement. +DROP FUNCTION public.consume_daily_quota(uuid, integer); + +CREATE FUNCTION public.consume_daily_quota( + p_user_id uuid, + p_limit integer, + p_quota_key text DEFAULT 'ai' +) +RETURNS TABLE (allowed boolean, current_count integer) +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = public +AS $$ +DECLARE + v_quota_key text := nullif(trim(p_quota_key), ''); + v_current_count integer; +BEGIN + IF auth.role() <> 'service_role' OR p_user_id IS NULL OR p_limit < 1 OR v_quota_key IS NULL THEN + RAISE EXCEPTION 'quota reservation is restricted to trusted server code'; + END IF; + + INSERT INTO public.daily_usage_counters (user_id, usage_date, quota_key, request_count) + VALUES (p_user_id, CURRENT_DATE, v_quota_key, 1) + ON CONFLICT (user_id, usage_date, quota_key) DO UPDATE + SET request_count = public.daily_usage_counters.request_count + 1 + WHERE public.daily_usage_counters.request_count < p_limit + RETURNING request_count INTO v_current_count; + + IF FOUND THEN + RETURN QUERY SELECT true, v_current_count; + RETURN; + END IF; + + SELECT request_count INTO v_current_count + FROM public.daily_usage_counters + WHERE user_id = p_user_id + AND usage_date = CURRENT_DATE + AND quota_key = v_quota_key; + + RETURN QUERY SELECT false, COALESCE(v_current_count, 0); +END; +$$; + +REVOKE ALL ON FUNCTION public.consume_daily_quota(uuid, integer, text) + FROM PUBLIC, anon, authenticated; +GRANT EXECUTE ON FUNCTION public.consume_daily_quota(uuid, integer, text) + TO service_role; diff --git a/supabase/tests/daily_usage_quota_test.sql b/supabase/tests/daily_usage_quota_test.sql new file mode 100644 index 000000000..7219d4070 --- /dev/null +++ b/supabase/tests/daily_usage_quota_test.sql @@ -0,0 +1,116 @@ +BEGIN; +SELECT plan(15); + +INSERT INTO auth.users (id) +VALUES + ('44444444-4444-4444-8444-444444444444'), + ('55555555-5555-4555-8555-555555555555'); + +SELECT has_column( + 'public', + 'daily_usage_counters', + 'quota_key', + 'daily usage counters identify the quota being consumed' +); +SELECT col_not_null( + 'public', + 'daily_usage_counters', + 'quota_key', + 'quota keys cannot be null' +); +SELECT col_default_is( + 'public', + 'daily_usage_counters', + 'quota_key', + 'ai', + 'existing and omitted quota keys default to ai' +); +SELECT results_eq( + $$SELECT pg_get_constraintdef(oid) + FROM pg_constraint + WHERE conrelid = 'public.daily_usage_counters'::regclass + AND contype = 'p'$$, + ARRAY['PRIMARY KEY (user_id, usage_date, quota_key)'], + 'the primary key scopes counters by quota key' +); +SELECT results_eq( + $$SELECT count(*)::bigint + FROM pg_proc + WHERE pronamespace = 'public'::regnamespace + AND proname = 'consume_daily_quota'$$, + ARRAY[1::bigint], + 'the migration leaves one consume_daily_quota overload' +); +SELECT ok( + NOT has_function_privilege('anon', 'public.consume_daily_quota(uuid, integer, text)', 'EXECUTE'), + 'anonymous callers cannot reserve quota' +); +SELECT ok( + NOT has_function_privilege('authenticated', 'public.consume_daily_quota(uuid, integer, text)', 'EXECUTE'), + 'authenticated callers cannot reserve quota' +); +SELECT ok( + has_function_privilege('service_role', 'public.consume_daily_quota(uuid, integer, text)', 'EXECUTE'), + 'service role can reserve quota' +); + +SET LOCAL ROLE authenticated; +SELECT set_config('request.jwt.claim.role', 'authenticated', true); +SELECT throws_ok( + $$SELECT * FROM public.consume_daily_quota('44444444-4444-4444-8444-444444444444', 3)$$, + '42501', + NULL, + 'unauthorized callers still cannot execute the quota RPC' +); + +RESET ROLE; +SET LOCAL ROLE service_role; +SELECT set_config('request.jwt.claim.role', 'service_role', true); + +SELECT results_eq( + $$SELECT allowed, current_count + FROM public.consume_daily_quota('44444444-4444-4444-8444-444444444444', 3)$$, + $$VALUES (true, 1)$$, + 'calls that omit the quota key reserve ai quota' +); +SELECT results_eq( + $$SELECT allowed, current_count + FROM public.consume_daily_quota('44444444-4444-4444-8444-444444444444', 1, 'import_source')$$, + $$VALUES (true, 1)$$, + 'import reservations increment the import_source counter' +); +SELECT results_eq( + $$SELECT allowed, current_count + FROM public.consume_daily_quota('44444444-4444-4444-8444-444444444444', 1, 'import_source')$$, + $$VALUES (false, 1)$$, + 'the import_source counter enforces its own limit' +); +SELECT results_eq( + $$SELECT allowed, current_count + FROM public.consume_daily_quota('44444444-4444-4444-8444-444444444444', 3)$$, + $$VALUES (true, 2)$$, + 'exhausting import quota does not exhaust ai quota' +); +SELECT results_eq( + $$SELECT quota_key, request_count + FROM public.daily_usage_counters + WHERE user_id = '44444444-4444-4444-8444-444444444444' + ORDER BY quota_key$$, + $$VALUES ('ai'::text, 2), ('import_source'::text, 1)$$, + 'ai and import reservations update only their own counters' +); + +RESET ROLE; +INSERT INTO public.daily_usage_counters (user_id, usage_date, request_count) +VALUES ('55555555-5555-4555-8555-555555555555', CURRENT_DATE - 1, 7); +SELECT results_eq( + $$SELECT quota_key, request_count + FROM public.daily_usage_counters + WHERE user_id = '55555555-5555-4555-8555-555555555555' + AND usage_date = CURRENT_DATE - 1$$, + $$VALUES ('ai'::text, 7)$$, + 'rows without an explicit quota key retain their count under ai' +); + +SELECT * FROM finish(); +ROLLBACK;