fix(webapp): query & API endpoint hardening (query authz, prompt-override, ai-title rate limit) (#44)

Author: d-cs

.server-changes/query-and-api-endpoint-hardening.md

@@ -0,0 +1,6 @@
+---
+area: webapp
+type: fix
+---
+
+Harden the query and prompt-override APIs and rate-limit the query ai-title endpoint

apps/webapp/app/routes/api.v1.prompts.$slug.override.ts

@@ -13,7 +13,11 @@ const CreateBody = z.object({
   textContent: z.string(),
   model: z.string().optional(),
   commitMessage: z.string().optional(),
-  source: z.string().optional(),
+  // `code` is reserved for deploy-created versions.
+  source: z
+    .string()
+    .refine((source) => source !== "code")
+    .optional(),
 });
 
 const UpdateBody = z.object({

apps/webapp/app/routes/api.v1.query.ts

@@ -5,6 +5,7 @@ import { createActionApiRoute, everyResource } from "~/services/routeBuilders/ap
 import { executeQuery, type QueryScope } from "~/services/queryService.server";
 import { logger } from "~/services/logger.server";
 import { rowsToCSV } from "~/utils/dataExport";
+import { detectQueryTables } from "~/v3/detectQueryTables";
 import { querySchemas } from "~/v3/querySchemas";
 
 const BodySchema = z.object({
@@ -16,14 +17,12 @@ const BodySchema = z.object({
   format: z.enum(["json", "csv"]).default("json"),
 });
 
-/** Extract table names from a TRQL query for authorization */
-function detectTables(query: string): string[] {
-  return querySchemas
-    .filter((s) => {
-      const escaped = s.name.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
-      return new RegExp(`\\bFROM\\s+${escaped}\\b`, "i").test(query);
-    })
-    .map((s) => s.name);
+const allowedQueryTables = new Set(querySchemas.map((s) => s.name));
+
+/** Every table the query reads, for per-table JWT-scope authorization.
+ * `null` means unparseable — callers deny by default. */
+function detectTables(query: string): string[] | null {
+  return detectQueryTables(query, allowedQueryTables);
 }
 
 const { action, loader } = createActionApiRoute(
@@ -34,12 +33,14 @@ const { action, loader } = createActionApiRoute(
     findResource: async () => 1,
     authorization: {
       action: "read",
-      // A multi-table query reads from every detected table. Wrap with
-      // everyResource so a JWT scoped to one table can't pass auth for
-      // a query that also reads tables it isn't scoped to (would be the
-      // same OR-loophole the batch trigger route had pre-fix).
+      // A multi-table query reads from every detected table, so wrap with
+      // everyResource: a JWT scoped to one table must not pass auth for a
+      // query that also reads tables it isn't scoped to.
       resource: (_, __, ___, body) => {
         const tables = detectTables(body.query);
+        // Unparseable query → deny. It must not fall through to the
+        // permissive {type:"query",id:"all"} branch.
+        if (tables === null) return { type: "query", id: "__unparseable__" };
         return tables.length > 0
           ? everyResource(tables.map((id) => ({ type: "query", id })))
           : { type: "query", id: "all" };

apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.$projectParam.env.$envParam.query.ai-title.tsx

@@ -8,17 +8,34 @@ import { findProjectBySlug } from "~/models/project.server";
 import { findEnvironmentBySlug } from "~/models/runtimeEnvironment.server";
 import { requireUserId } from "~/services/session.server";
 import { EnvironmentParamSchema } from "~/utils/pathBuilder";
+import { aiTitleRateLimiter } from "~/v3/services/aiTitleRateLimiter.server";
 import { AIQueryTitleService } from "~/v3/services/aiQueryTitleService.server";
 
+// `/resources/*` isn't covered by the global apiRateLimiter (`/api/*` only),
+// so this endpoint needs its own per-route limiter and length cap.
+const MAX_QUERY_LENGTH = 10_000;
+
 const RequestSchema = z.object({
-  query: z.string().min(1, "Query is required"),
+  query: z.string().min(1, "Query is required").max(MAX_QUERY_LENGTH),
   queryId: z.string().optional(),
 });
 
 export async function action({ request, params }: ActionFunctionArgs) {
   const userId = await requireUserId(request);
   const { organizationSlug, projectParam, envParam } = EnvironmentParamSchema.parse(params);
 
+  const limit = await aiTitleRateLimiter.limit(`user:${userId}`);
+  if (!limit.success) {
+    return json(
+      {
+        success: false as const,
+        error: "Too many requests — please wait a moment and try again.",
+        title: null,
+      },
+      { status: 429 }
+    );
+  }
+
   // Parse the request body
   const [error, data] = await tryCatch(request.json());
   if (error) {

apps/webapp/app/v3/detectQueryTables.ts

@@ -0,0 +1,113 @@
+import {
+  parseTSQLSelect,
+  SyntaxError as TSQLSyntaxError,
+  type Field,
+  type JoinExpr,
+  type SelectQuery,
+  type SelectSetQuery,
+} from "@internal/tsql";
+
+/**
+ * Extract every known table a TRQL query reads — the FROM table, every JOIN in
+ * the chain, and any subqueries — for per-table JWT-scope authorization.
+ *
+ * `allowedTableNames` is the set of recognised table names (matched
+ * case-insensitively); anything not in it is ignored. Injected so this stays
+ * dependency-free (the caller derives it from the query schemas).
+ *
+ * Returns `null` when the query can't be parsed; callers MUST treat `null` as
+ * deny-by-default.
+ */
+export function detectQueryTables(query: string, allowedTableNames: Set<string>): string[] | null {
+  let ast: SelectQuery | SelectSetQuery;
+  try {
+    ast = parseTSQLSelect(query);
+  } catch (err) {
+    if (err instanceof TSQLSyntaxError) return null;
+    throw err;
+  }
+
+  const allowed = new Map(Array.from(allowedTableNames, (n) => [n.toLowerCase(), n]));
+  const seen = new Set<string>();
+  const scanned = new WeakSet<object>();
+
+  function visitSelect(q: SelectQuery): void {
+    // CTE bodies: `WITH r AS (SELECT ... FROM <table>) ...` — the table is
+    // read by the CTE even when the outer query only references the CTE alias.
+    if (q.ctes) {
+      for (const cte of Object.values(q.ctes)) {
+        scanForSubqueries(cte.expr);
+      }
+    }
+    // FROM / JOIN chain (tables + FROM-position subqueries).
+    if (q.select_from) visitJoin(q.select_from);
+    // Subqueries anywhere else (WHERE, SELECT list, GROUP BY, ORDER BY, etc.)
+    // can each embed a SELECT that reads a real table, e.g.
+    // `WHERE id IN (SELECT … FROM runs)`.
+    scanForSubqueries(q.select);
+    scanForSubqueries(q.where);
+    scanForSubqueries(q.prewhere);
+    scanForSubqueries(q.having);
+    scanForSubqueries(q.group_by);
+    scanForSubqueries(q.array_join_list);
+    scanForSubqueries(q.order_by);
+    scanForSubqueries(q.limit);
+    scanForSubqueries(q.offset);
+    scanForSubqueries(q.limit_by);
+    scanForSubqueries(q.window_exprs);
+  }
+  // Shape-agnostic walk of an expression subtree: descends every nested
+  // object/array and hands any embedded SELECT to the query visitors, so a new
+  // node shape can't silently reintroduce a detection gap. The WeakSet guards
+  // against back-reference cycles the AST might carry.
+  function scanForSubqueries(node: unknown): void {
+    if (node === null || typeof node !== "object") return;
+    if (scanned.has(node)) return;
+    scanned.add(node);
+    if (Array.isArray(node)) {
+      for (const item of node) scanForSubqueries(item);
+      return;
+    }
+    const expressionType = (node as { expression_type?: string }).expression_type;
+    if (expressionType === "select_query") {
+      visitSelect(node as SelectQuery);
+      return;
+    }
+    if (expressionType === "select_set_query") {
+      visitSelectSet(node as SelectSetQuery);
+      return;
+    }
+    for (const value of Object.values(node)) scanForSubqueries(value);
+  }
+  function visitSelectSet(q: SelectSetQuery): void {
+    visitAny(q.initial_select_query);
+    for (const node of q.subsequent_select_queries ?? []) {
+      visitAny(node.select_query);
+    }
+  }
+  function visitAny(q: SelectQuery | SelectSetQuery): void {
+    if (q.expression_type === "select_query") visitSelect(q);
+    else visitSelectSet(q);
+  }
+  function visitJoin(node: JoinExpr): void {
+    const tableExpr = node.table;
+    if (tableExpr) {
+      if ((tableExpr as Field).expression_type === "field") {
+        const name = (tableExpr as Field).chain[0];
+        const canonicalName =
+          typeof name === "string" ? allowed.get(name.toLowerCase()) : undefined;
+        if (canonicalName) seen.add(canonicalName);
+      } else if ((tableExpr as SelectQuery).expression_type === "select_query") {
+        visitSelect(tableExpr as SelectQuery);
+      } else if ((tableExpr as SelectSetQuery).expression_type === "select_set_query") {
+        visitSelectSet(tableExpr as SelectSetQuery);
+      }
+    }
+    if (node.next_join) visitJoin(node.next_join);
+  }
+
+  if (ast.expression_type === "select_set_query") visitSelectSet(ast);
+  else visitSelect(ast);
+
+  return Array.from(seen);
+}

apps/webapp/app/v3/services/aiTitleRateLimiter.server.ts

@@ -0,0 +1,25 @@
+import { Ratelimit } from "@upstash/ratelimit";
+import { type RedisWithClusterOptions } from "~/redis.server";
+import { createRedisRateLimitClient, RateLimiter } from "~/services/rateLimiter.server";
+import { singleton } from "~/utils/singleton";
+
+// The query ai-title endpoint lives under `/resources/*`, which the global
+// apiRateLimiter (only `/api/*`) does not cover, so it needs its own per-user
+// cap. Exported so the policy is asserted in tests rather than re-encoded.
+export const AI_TITLE_RATE_LIMIT_ATTEMPTS = 30;
+export const AI_TITLE_RATE_LIMIT_WINDOW = "10 m" as const;
+
+/**
+ * Build the ai-title per-user rate limiter. Production uses the env-derived
+ * rate-limit Redis; tests inject a container Redis.
+ */
+export function createAITitleRateLimiter(redisOptions?: RedisWithClusterOptions): RateLimiter {
+  return new RateLimiter({
+    ...(redisOptions ? { redisClient: createRedisRateLimitClient(redisOptions) } : {}),
+    keyPrefix: "query.ai-title",
+    limiter: Ratelimit.slidingWindow(AI_TITLE_RATE_LIMIT_ATTEMPTS, AI_TITLE_RATE_LIMIT_WINDOW),
+    logFailure: true,
+  });
+}
+
+export const aiTitleRateLimiter = singleton("aiTitleRateLimiter", () => createAITitleRateLimiter());

apps/webapp/app/v3/services/promptOverrideSource.ts

@@ -0,0 +1,7 @@
+// `source: "code"` is reserved for the deploy path; no other caller may write
+// it. Normalize anything that isn't a legitimate caller-supplied value to
+// "dashboard". Dependency-free so the rule can be unit-tested directly; it
+// backs the service-layer check (the route layer also constrains `source`).
+export function normalizePromptOverrideSource(source: string | undefined | null): string {
+  return source && source !== "code" ? source : "dashboard";
+}

apps/webapp/app/v3/services/promptService.server.ts

@@ -1,6 +1,7 @@
 import { createHash } from "crypto";
 import { prisma } from "~/db.server";
 import { BaseService, ServiceValidationError } from "./baseService.server";
+import { normalizePromptOverrideSource } from "./promptOverrideSource";
 
 export class PromptService extends BaseService {
   async promoteVersion(promptId: string, versionId: string, options?: { sourceGuard?: boolean }) {
@@ -62,13 +63,17 @@ export class PromptService extends BaseService {
         WHERE "promptId" = ${promptId} AND 'override' = ANY("labels")
       `;
 
+      // Defence in depth: reject the reserved `code` source regardless of
+      // caller, in case a future caller skips the route-layer check.
+      const safeSource = normalizePromptOverrideSource(data.source);
+
       await tx.promptVersion.create({
         data: {
           promptId,
           version: nextVersion,
           textContent: data.textContent,
           model: data.model || null,
-          source: data.source || "dashboard",
+          source: safeSource,
           commitMessage: data.commitMessage || null,
           contentHash,
           labels: ["override"],

apps/webapp/test/aiTitleRateLimiter.test.ts

@@ -0,0 +1,54 @@
+import { redisTest } from "@internal/testcontainers";
+import { type RedisOptions } from "ioredis";
+import { describe, expect, vi } from "vitest";
+import { type RedisWithClusterOptions } from "../app/redis.server.js";
+import {
+  AI_TITLE_RATE_LIMIT_ATTEMPTS,
+  createAITitleRateLimiter,
+} from "../app/v3/services/aiTitleRateLimiter.server.js";
+
+vi.setConfig({ testTimeout: 60_000 });
+
+// Plaintext container: without tlsDisabled the client attempts TLS, the
+// connection fails, and @upstash/ratelimit fails open (allowing everything).
+const toRedisOptions = (o: RedisOptions): RedisWithClusterOptions => ({
+  host: o.host,
+  port: o.port,
+  username: o.username,
+  password: o.password,
+  tlsDisabled: true,
+});
+
+let seq = 0;
+const userKey = (label: string) => `user:${label}-${seq++}`;
+
+// The query ai-title endpoint isn't covered by the global apiRateLimiter, so
+// this per-user limiter is the only thing bounding it.
+describe("aiTitleRateLimiter", () => {
+  redisTest("allows up to the limit then blocks further attempts", async ({ redisOptions }) => {
+    const limiter = createAITitleRateLimiter(toRedisOptions(redisOptions));
+    const key = userKey("loop");
+
+    for (let i = 0; i < AI_TITLE_RATE_LIMIT_ATTEMPTS; i++) {
+      const r = await limiter.limit(key);
+      expect(r.success).toBe(true);
+    }
+
+    const blocked = await limiter.limit(key);
+    expect(blocked.success).toBe(false);
+  });
+
+  redisTest("scopes the limit per user", async ({ redisOptions }) => {
+    const limiter = createAITitleRateLimiter(toRedisOptions(redisOptions));
+    const victim = userKey("victim");
+    const bystander = userKey("bystander");
+
+    for (let i = 0; i < AI_TITLE_RATE_LIMIT_ATTEMPTS; i++) {
+      await limiter.limit(victim);
+    }
+    expect((await limiter.limit(victim)).success).toBe(false);
+
+    // A different user is unaffected.
+    expect((await limiter.limit(bystander)).success).toBe(true);
+  });
+});

apps/webapp/test/detectQueryTables.test.ts

@@ -0,0 +1,59 @@
+import { describe, expect, it } from "vitest";
+import { detectQueryTables } from "../app/v3/detectQueryTables.js";
+
+const allowed = new Set(["runs", "tasks"]);
+const sorted = (xs: string[] | null) => (xs === null ? null : [...xs].sort());
+
+// detectQueryTables backs per-table JWT-scope authorization for /api/v1/query.
+// Key behaviours over the old FROM-only regex: it sees JOINed and subquery
+// tables, and returns null for unparseable queries so the caller denies by
+// default.
+describe("detectQueryTables", () => {
+  it("detects the FROM table", () => {
+    expect(sorted(detectQueryTables("SELECT * FROM runs", allowed))).toEqual(["runs"]);
+  });
+
+  it("returns the canonical table name when query casing differs", () => {
+    expect(sorted(detectQueryTables("SELECT * FROM RUNS", allowed))).toEqual(["runs"]);
+  });
+
+  it("detects every JOINed table, not just FROM", () => {
+    expect(
+      sorted(detectQueryTables("SELECT * FROM runs JOIN tasks ON runs.id = tasks.run_id", allowed))
+    ).toEqual(["runs", "tasks"]);
+  });
+
+  it("detects tables inside a FROM subquery", () => {
+    expect(sorted(detectQueryTables("SELECT * FROM (SELECT * FROM runs) AS r", allowed))).toEqual([
+      "runs",
+    ]);
+  });
+
+  it("detects a table read only inside a CTE body", () => {
+    expect(
+      sorted(detectQueryTables("WITH r AS (SELECT * FROM runs) SELECT * FROM r", allowed))
+    ).toEqual(["runs"]);
+  });
+
+  it("detects a table read only inside a WHERE subquery", () => {
+    expect(
+      sorted(
+        detectQueryTables("SELECT * FROM tasks WHERE id IN (SELECT run_id FROM runs)", allowed)
+      )
+    ).toEqual(["runs", "tasks"]);
+  });
+
+  it("detects a table read only inside a SELECT-list subquery", () => {
+    expect(
+      sorted(detectQueryTables("SELECT (SELECT count() FROM runs) AS c FROM tasks", allowed))
+    ).toEqual(["runs", "tasks"]);
+  });
+
+  it("ignores tables that aren't in the allowed schema set", () => {
+    expect(detectQueryTables("SELECT * FROM runs", new Set(["tasks"]))).toEqual([]);
+  });
+
+  it("returns null for an unparseable query (caller denies by default)", () => {
+    expect(detectQueryTables("definitely not a valid query !!!", allowed)).toBeNull();
+  });
+});