feat(api): per-key scope/format lockdown (keyRestriction flag) (firecrawl#3960)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Author: mogery

apps/api/src/__tests__/snips/v2/key-restriction.test.ts

@@ -0,0 +1,176 @@
+import { eq } from "drizzle-orm";
+import { idmux, map, scrape, scrapeRaw, scrapeTimeout } from "./lib";
+import { createTestIdUrl, describeIf, Identity, TEST_PRODUCTION } from "../lib";
+import { db } from "../../../db/connection";
+import * as schema from "../../../db/schema";
+import { parseApi } from "../../../lib/parseApi";
+
+// Needs idmux (to create teams with the keyRestriction flag) and direct DB
+// access (to seed key_restriction_config), so production-suite only.
+describeIf(TEST_PRODUCTION)(
+  "Key restriction (keyRestriction team flag)",
+  () => {
+    const seededKeyIds: number[] = [];
+
+    async function seedRestriction(
+      identity: Identity,
+      restriction: { allowedFormats?: string[]; allowedEndpoints?: string[] },
+    ) {
+      const [keyRow] = await db
+        .select({ id: schema.api_keys.id })
+        .from(schema.api_keys)
+        .where(eq(schema.api_keys.key, parseApi(identity.apiKey)))
+        .limit(1);
+      expect(keyRow).toBeDefined();
+
+      await db.insert(schema.key_restriction_config).values({
+        api_key_id: keyRow!.id,
+        team_id: identity.teamId,
+        allowed_formats: restriction.allowedFormats ?? [],
+        allowed_endpoints: restriction.allowedEndpoints ?? [],
+      });
+      seededKeyIds.push(keyRow!.id);
+    }
+
+    afterAll(async () => {
+      for (const keyId of seededKeyIds) {
+        await db
+          .delete(schema.key_restriction_config)
+          .where(eq(schema.key_restriction_config.api_key_id, keyId));
+      }
+    });
+
+    it.concurrent(
+      "allows markdown scrapes on a markdown-only key",
+      async () => {
+        const identity = await idmux({
+          name: "key-restriction/markdown-allowed",
+          credits: 10000,
+          flags: { keyRestriction: true },
+        });
+        await seedRestriction(identity, { allowedFormats: ["markdown"] });
+
+        // Explicit markdown and the implicit default must both pass.
+        const doc = await scrape(
+          { url: createTestIdUrl(), formats: ["markdown"] },
+          identity,
+        );
+        expect(doc.markdown).toBeDefined();
+
+        const docDefault = await scrape({ url: createTestIdUrl() }, identity);
+        expect(docDefault.markdown).toBeDefined();
+      },
+      scrapeTimeout * 2,
+    );
+
+    it.concurrent(
+      "rejects non-allowed formats on a markdown-only key",
+      async () => {
+        const identity = await idmux({
+          name: "key-restriction/format-blocked",
+          credits: 10000,
+          flags: { keyRestriction: true },
+        });
+        await seedRestriction(identity, { allowedFormats: ["markdown"] });
+
+        const response = await scrapeRaw(
+          { url: createTestIdUrl(), formats: ["markdown", "rawHtml"] },
+          identity,
+        );
+
+        expect(response.statusCode).toBe(403);
+        expect(response.body.success).toBe(false);
+        expect(response.body.error).toContain(
+          "restricted to the following formats",
+        );
+      },
+      scrapeTimeout,
+    );
+
+    it.concurrent(
+      "rejects content-returning actions on a format-restricted key",
+      async () => {
+        const identity = await idmux({
+          name: "key-restriction/action-blocked",
+          credits: 10000,
+          flags: { keyRestriction: true },
+        });
+        await seedRestriction(identity, { allowedFormats: ["markdown"] });
+
+        const response = await scrapeRaw(
+          {
+            url: createTestIdUrl(),
+            formats: ["markdown"],
+            actions: [{ type: "screenshot" }],
+          },
+          identity,
+        );
+
+        expect(response.statusCode).toBe(403);
+        expect(response.body.success).toBe(false);
+        expect(response.body.error).toContain("screenshot");
+      },
+      scrapeTimeout,
+    );
+
+    it.concurrent(
+      "enforces the endpoint allowlist",
+      async () => {
+        const identity = await idmux({
+          name: "key-restriction/endpoint",
+          credits: 10000,
+          flags: { keyRestriction: true },
+        });
+        await seedRestriction(identity, { allowedEndpoints: ["scrape"] });
+
+        // Allowed endpoint group works...
+        await scrape({ url: createTestIdUrl() }, identity);
+
+        // ...anything else is rejected at auth.
+        const response = await map({ url: "https://firecrawl.dev" }, identity);
+        expect(response.statusCode).toBe(403);
+        expect(response.body.success).toBe(false);
+        expect(response.body.error).toContain(
+          "restricted to the following endpoints",
+        );
+      },
+      scrapeTimeout * 2,
+    );
+
+    it.concurrent(
+      "does not restrict a flagged team before it configures a restriction",
+      async () => {
+        const identity = await idmux({
+          name: "key-restriction/no-config",
+          credits: 10000,
+          flags: { keyRestriction: true },
+        });
+
+        const doc = await scrape(
+          { url: createTestIdUrl(), formats: ["rawHtml"] },
+          identity,
+        );
+        expect(doc.rawHtml).toBeDefined();
+      },
+      scrapeTimeout,
+    );
+
+    it.concurrent(
+      "does not restrict when the flag is off even if a config exists",
+      async () => {
+        const identity = await idmux({
+          name: "key-restriction/no-flag",
+          credits: 10000,
+        });
+        await seedRestriction(identity, { allowedFormats: ["markdown"] });
+
+        const doc = await scrape(
+          { url: createTestIdUrl(), formats: ["rawHtml"] },
+          identity,
+        );
+        expect(doc.rawHtml).toBeDefined();
+      },
+      scrapeTimeout,
+    );
+  },
+);

apps/api/src/controllers/auth.ts

@@ -17,6 +17,7 @@ import {
 } from "../lib/keyless";
 import { isKeylessIpSuspicious } from "../lib/spur";
 import { checkIpRestriction } from "../lib/ip-restriction";
+import { checkKeyEndpointRestriction } from "../lib/key-restriction";
 import { deleteKey, getValue, setValue } from "../services/redis";
 import { redlock } from "../services/redlock";
 import { eq } from "drizzle-orm";
@@ -817,6 +818,23 @@ async function supaAuthenticateUser(
     }
   }
 
+  if (chunk?.flags?.keyRestriction) {
+    // Enforced here rather than in route middleware so every authenticated
+    // surface (v0/v1/v2, websocket status) goes through the same gate.
+    const endpointCheck = await checkKeyEndpointRestriction(
+      req.originalUrl ?? req.url ?? "",
+      chunk.api_key_id,
+      chunk.flags,
+    );
+    if (!endpointCheck.allowed) {
+      return {
+        success: false,
+        error: endpointCheck.error,
+        status: endpointCheck.status,
+      };
+    }
+  }
+
   const team_endpoint_token = token === config.PREVIEW_TOKEN ? iptoken : teamId;
 
   try {

apps/api/src/controllers/v0/admin/key-restriction-cache-clear.ts

@@ -0,0 +1,29 @@
+import { Request, Response } from "express";
+import { clearKeyRestrictionCache } from "../../../lib/key-restriction";
+import { logger } from "../../../lib/logger";
+
+// Called by the dashboard after it writes key_restriction_config so
+// restriction edits take effect immediately instead of after the 60s
+// cache TTL.
+export async function keyRestrictionCacheClearController(
+  req: Request,
+  res: Response,
+) {
+  try {
+    const api_key_id = req.body?.api_key_id;
+
+    if (typeof api_key_id !== "number" || !Number.isInteger(api_key_id)) {
+      return res.status(400).json({ error: "api_key_id is required" });
+    }
+
+    await clearKeyRestrictionCache(api_key_id);
+
+    logger.info("Key restriction config cache cleared", { api_key_id });
+    res.json({ ok: true });
+  } catch (error) {
+    logger.error("Error clearing key restriction cache via API route", {
+      error,
+    });
+    res.status(500).json({ error: "Internal server error" });
+  }
+}

apps/api/src/controllers/v1/batch-scrape.ts

@@ -26,6 +26,11 @@ import { UNSUPPORTED_SITE_MESSAGE } from "../../lib/strings";
 import { isUrlBlocked } from "../../scraper/WebScraper/utils/blocklist";
 import { fromV1ScrapeOptions } from "../v2/types";
 import { checkPermissions } from "../../lib/permissions";
+import {
+  actionTypesOf,
+  checkKeyFormatRestriction,
+  formatTypesOf,
+} from "../../lib/key-restriction";
 import {
   crawlGroup,
   resolveNewGroupBackend,
@@ -52,6 +57,19 @@ export async function batchScrapeController(
     });
   }
 
+  const keyRestriction = await checkKeyFormatRestriction(
+    formatTypesOf(req.body.formats),
+    actionTypesOf(req.body.actions),
+    req.acuc?.api_key_id,
+    req.acuc?.flags ?? null,
+  );
+  if (!keyRestriction.allowed) {
+    return res.status(keyRestriction.status).json({
+      success: false,
+      error: keyRestriction.error,
+    });
+  }
+
   const zeroDataRetention =
     getScrapeZDR(req.acuc?.flags) === "forced" || req.body.zeroDataRetention;
 

apps/api/src/controllers/v1/crawl.ts

@@ -18,6 +18,11 @@ import { _addScrapeJobToBullMQ } from "../../services/queue-jobs";
 import { logger as _logger } from "../../lib/logger";
 import { fromV1ScrapeOptions } from "../v2/types";
 import { checkPermissions } from "../../lib/permissions";
+import {
+  actionTypesOf,
+  checkKeyFormatRestriction,
+  formatTypesOf,
+} from "../../lib/key-restriction";
 import {
   crawlGroup,
   resolveNewGroupBackend,
@@ -43,6 +48,19 @@ export async function crawlController(
     });
   }
 
+  const keyRestriction = await checkKeyFormatRestriction(
+    formatTypesOf(req.body.scrapeOptions?.formats),
+    actionTypesOf(req.body.scrapeOptions?.actions),
+    req.acuc?.api_key_id,
+    req.acuc?.flags ?? null,
+  );
+  if (!keyRestriction.allowed) {
+    return res.status(keyRestriction.status).json({
+      success: false,
+      error: keyRestriction.error,
+    });
+  }
+
   const zeroDataRetention =
     getScrapeZDR(req.acuc?.flags) === "forced" || req.body.zeroDataRetention;
 

apps/api/src/controllers/v1/scrape.ts

@@ -14,6 +14,11 @@ import { fromV1ScrapeOptions } from "../v2/types";
 import { TransportableError } from "../../lib/error";
 import { NuQJob } from "../../services/worker/nuq";
 import { checkPermissions } from "../../lib/permissions";
+import {
+  actionTypesOf,
+  checkKeyFormatRestriction,
+  formatTypesOf,
+} from "../../lib/key-restriction";
 import { includesFormat } from "../../lib/format-utils";
 import { teamConcurrencySemaphore } from "../../services/worker/team-semaphore";
 import { processJobInternal } from "../../services/worker/scrape-worker";
@@ -53,6 +58,19 @@ export async function scrapeController(
     });
   }
 
+  const keyRestriction = await checkKeyFormatRestriction(
+    formatTypesOf(req.body.formats),
+    actionTypesOf(req.body.actions),
+    req.acuc?.api_key_id,
+    req.acuc?.flags ?? null,
+  );
+  if (!keyRestriction.allowed) {
+    return res.status(keyRestriction.status).json({
+      success: false,
+      error: keyRestriction.error,
+    });
+  }
+
   const zeroDataRetention =
     getScrapeZDR(req.acuc?.flags) === "forced" || req.body.zeroDataRetention;
 

apps/api/src/controllers/v1/search.ts

@@ -12,6 +12,11 @@ import { v7 as uuidv7 } from "uuid";
 import { logSearch, logRequest } from "../../services/logging/log_job";
 import { search } from "../../search";
 import { logger as _logger } from "../../lib/logger";
+import {
+  actionTypesOf,
+  checkKeyFormatRestriction,
+  formatTypesOf,
+} from "../../lib/key-restriction";
 import type { Logger } from "winston";
 import { ScrapeJobTimeoutError } from "../../lib/error";
 import { captureExceptionWithZdrCheck } from "../../services/sentry";
@@ -122,6 +127,24 @@ export async function searchController(
   try {
     req.body = searchRequestSchema.parse(req.body);
 
+    const requestedFormats = formatTypesOf(req.body.scrapeOptions?.formats);
+    const keyRestriction = await checkKeyFormatRestriction(
+      requestedFormats,
+      // Search only scrapes (and only runs actions) when formats are
+      // requested; without them scrapeOptions is ignored entirely.
+      requestedFormats.length > 0
+        ? actionTypesOf(req.body.scrapeOptions?.actions)
+        : [],
+      req.acuc?.api_key_id,
+      req.acuc?.flags ?? null,
+    );
+    if (!keyRestriction.allowed) {
+      return res.status(keyRestriction.status).json({
+        success: false,
+        error: keyRestriction.error,
+      });
+    }
+
     logger = logger.child({
       version: "v1",
       query: req.body.query,

apps/api/src/controllers/v1/types.ts

@@ -1313,6 +1313,8 @@ export type TeamFlags = {
   ipWhitelist?: boolean;
   // gates the per-team API key IP allowlist (ip_restriction_config table)
   ipRestriction?: boolean;
+  // gates the per-key scope/format lockdown (key_restriction_config table)
+  keyRestriction?: boolean;
   skipCountryCheck?: boolean;
   browserBeta?: boolean;
   bypassCreditChecks?: boolean;