fix: harden v2 custom widgets and Workshop (#6480)

packages/custom-widgets/src/server/errors.ts

@@ -10,17 +10,19 @@ export type CustomWidgetDomainErrorCode =
 
 export class CustomWidgetDomainError extends Error {
   public readonly code: CustomWidgetDomainErrorCode;
+  public readonly reason?: "timeout";
   public readonly retryAfterMs?: number;
 
   public constructor(input: {
     code: CustomWidgetDomainErrorCode;
     message: string;
-    cause?: unknown;
+    reason?: "timeout";
     retryAfterMs?: number;
   }) {
-    super(input.message, { cause: input.cause });
+    super(input.message);
     this.name = "CustomWidgetDomainError";
     this.code = input.code;
+    this.reason = input.reason;
     this.retryAfterMs = input.retryAfterMs;
   }
 }

packages/custom-widgets/src/server/index.ts

@@ -16,7 +16,12 @@ export {
   validateCustomWidgetUrl,
 } from "./request-executor";
 export type { CustomWidgetAuthConfig, CustomWidgetHttpRequest, CustomWidgetHttpResponse } from "./request-executor";
-export { CUSTOM_WIDGET_REQUEST_CONCURRENCY_TTL_MS, CustomWidgetRequestLimiter } from "./request-limits";
+export {
+  CUSTOM_WIDGET_DEFINITION_CONCURRENCY_LIMIT,
+  CUSTOM_WIDGET_REQUEST_CONCURRENCY_TTL_MS,
+  CUSTOM_WIDGET_USER_ITEM_CONCURRENCY_LIMIT,
+  CustomWidgetRequestLimiter,
+} from "./request-limits";
 export type { RequestLimitInput, RequestLimiterOptions, RequestLimitStore } from "./request-limits";
 export { CustomWidgetPreviewSessionService } from "./preview-sessions";
 export type {

packages/custom-widgets/src/server/network-policy.ts

@@ -66,6 +66,11 @@ loopbackAddresses.addAddress("::1", "ipv6");
 
 type AddressFamily = 4 | 6;
 export type ResolvedAddress = { address: string; family: AddressFamily };
+export type HostResolver = (hostname: string) => Promise<ResolvedAddress[]>;
+export interface ResolveHostOptions {
+  signal?: AbortSignal;
+  resolver?: HostResolver;
+}
 type AddressClass = "public" | "private" | "loopback" | "blocked";
 const normalizeHostname = (value: string) =>
   value.startsWith("[") && value.endsWith("]") ? value.slice(1, -1) : value;
@@ -83,12 +88,19 @@ export function classifyAddress(address: string): AddressClass {
 export async function resolveAndValidateHost(
   hostname: string,
   scope: CustomJsxNetworkScope,
+  options: ResolveHostOptions = {},
 ): Promise<ResolvedAddress[]> {
   const normalized = normalizeHostname(hostname);
   const family = isIP(normalized) as AddressFamily | 0;
+  throwIfAborted(options.signal);
   const addresses = family
     ? [{ address: normalized, family }]
-    : ((await lookup(normalized, { all: true, verbatim: true })) as ResolvedAddress[]);
+    : await abortable(
+        options.resolver?.(normalized) ??
+          (lookup(normalized, { all: true, verbatim: true }) as Promise<ResolvedAddress[]>),
+        options.signal,
+      );
+  throwIfAborted(options.signal);
   if (!addresses.length)
     throw new CustomWidgetDomainError({ code: "BAD_REQUEST", message: "Target host did not resolve" });
   for (const address of addresses) {
@@ -106,6 +118,39 @@ export async function resolveAndValidateHost(
   return addresses;
 }
 
+function throwIfAborted(signal?: AbortSignal): void {
+  if (signal?.aborted) throw signal.reason ?? createAbortError();
+}
+
+function abortable<T>(operation: Promise<T>, signal?: AbortSignal): Promise<T> {
+  if (!signal) return operation;
+  if (signal.aborted) return Promise.reject(signal.reason ?? createAbortError());
+
+  return new Promise<T>((resolve, reject) => {
+    const onAbort = () => {
+      signal.removeEventListener("abort", onAbort);
+      reject(signal.reason ?? createAbortError());
+    };
+    signal.addEventListener("abort", onAbort, { once: true });
+    operation.then(
+      (value) => {
+        signal.removeEventListener("abort", onAbort);
+        resolve(value);
+      },
+      (error: unknown) => {
+        signal.removeEventListener("abort", onAbort);
+        reject(error);
+      },
+    );
+  });
+}
+
+function createAbortError(): Error {
+  const error = new Error("Operation aborted");
+  error.name = "AbortError";
+  return error;
+}
+
 export function validateCustomWidgetUrl(value: string | URL): URL {
   let url: URL;
   try {

packages/custom-widgets/src/server/request-dispatcher-lifecycle.ts

@@ -0,0 +1,62 @@
+import type { createPinnedAgent } from "./network-policy";
+
+export const DISPATCHER_CLOSE_GRACE_MS = 1_000;
+
+type CustomWidgetDispatcher = ReturnType<typeof createPinnedAgent>;
+
+export async function closeDispatcher(
+  dispatcher: CustomWidgetDispatcher,
+  deadlineSignal: AbortSignal,
+): Promise<boolean> {
+  const cleanupController = new AbortController();
+  const timeout = setTimeout(() => cleanupController.abort(), DISPATCHER_CLOSE_GRACE_MS);
+  try {
+    await abortable(dispatcher.close(), AbortSignal.any([deadlineSignal, cleanupController.signal]));
+    return true;
+  } catch (error) {
+    detachDispatcherDestroy(dispatcher);
+    if (cleanupController.signal.aborted && !deadlineSignal.aborted) return false;
+    throw error;
+  } finally {
+    clearTimeout(timeout);
+  }
+}
+
+function detachDispatcherDestroy(dispatcher: CustomWidgetDispatcher): void {
+  try {
+    void dispatcher.destroy().catch(() => undefined);
+  } catch {
+    // Cleanup is best effort after its bounded grace period.
+  }
+}
+
+function abortable<T>(operation: Promise<T>, signal: AbortSignal): Promise<T> {
+  if (signal.aborted) {
+    void operation.catch(() => undefined);
+    return Promise.reject(signal.reason ?? createAbortError());
+  }
+
+  return new Promise<T>((resolve, reject) => {
+    const onAbort = () => {
+      signal.removeEventListener("abort", onAbort);
+      reject(signal.reason ?? createAbortError());
+    };
+    signal.addEventListener("abort", onAbort, { once: true });
+    operation.then(
+      (value) => {
+        signal.removeEventListener("abort", onAbort);
+        resolve(value);
+      },
+      (error: unknown) => {
+        signal.removeEventListener("abort", onAbort);
+        reject(error);
+      },
+    );
+  });
+}
+
+function createAbortError(): Error {
+  const error = new Error("Operation aborted");
+  error.name = "AbortError";
+  return error;
+}

packages/custom-widgets/src/server/request-executor.ts

@@ -12,6 +12,7 @@ import {
   resolveSameOriginTarget,
   validateCustomWidgetUrl,
 } from "./network-policy";
+import { closeDispatcher } from "./request-dispatcher-lifecycle";
 import { parseResponseBody } from "./response";
 
 export {
@@ -33,6 +34,8 @@ const REQUEST_TIMEOUT_MS = 10_000;
 const MAX_QUERY_REDIRECTS = 3;
 const MAX_RESPONSE_CACHE_ENTRIES = 1_000;
 export const MAX_REQUEST_DURATION_MS = 45_000;
+const TIMEOUT_ERROR_CODES = new Set(["UND_ERR_CONNECT_TIMEOUT", "UND_ERR_HEADERS_TIMEOUT", "UND_ERR_BODY_TIMEOUT"]);
+const RESPONSE_TOO_LARGE_ERROR_CODE = "UND_ERR_RES_EXCEEDED_MAX_SIZE";
 
 export interface CustomWidgetAuthConfig {
   type: string;
@@ -51,7 +54,7 @@ export interface CustomWidgetHttpRequest {
   kind: "query" | "action";
   cacheKey?: string;
   cacheTtlSeconds?: number;
-  logError?: (event: { origin: string; method: CustomWidgetMethod; errorName: string }) => void;
+  logError?: (event: { origin: string; method: CustomWidgetMethod; errorName: string; reason?: "timeout" }) => void;
 }
 
 export interface CustomWidgetHttpResponse {
@@ -64,6 +67,10 @@ const cache = new Map<string, { expiresAt: number; response: CustomWidgetHttpRes
 const inFlight = new Map<string, Promise<CustomWidgetHttpResponse>>();
 let cacheEpoch = 0;
 
+type RequestHopResult =
+  | { kind: "response"; response: CustomWidgetHttpResponse }
+  | { kind: "redirect"; location: string | null; statusCode: number };
+
 async function performRequest(input: CustomWidgetHttpRequest): Promise<CustomWidgetHttpResponse> {
   const controller = new AbortController();
   const timeout = setTimeout(() => controller.abort(), MAX_REQUEST_DURATION_MS);
@@ -74,11 +81,19 @@ async function performRequest(input: CustomWidgetHttpRequest): Promise<CustomWid
       throw new CustomWidgetDomainError({
         code: "BAD_GATEWAY",
         message: "External request exceeded the total time limit",
-        cause: error,
+        reason: "timeout",
       });
     }
     if (error instanceof CustomWidgetDomainError) throw error;
-    throw error;
+    input.logError?.({
+      origin: URL.canParse(input.baseUrl) ? new URL(input.baseUrl).origin : "invalid",
+      method: input.method,
+      errorName: "TransportError",
+    });
+    throw new CustomWidgetDomainError({
+      code: "BAD_GATEWAY",
+      message: "External request failed",
+    });
   } finally {
     clearTimeout(timeout);
   }
@@ -96,15 +111,16 @@ async function performRequestWithinDeadline(
   const maxRedirects = input.kind === "query" ? MAX_QUERY_REDIRECTS : 0;
   for (let redirects = 0; ; redirects += 1) {
     const dispatcher = createPinnedAgent(
-      await resolveAndValidateHost(currentUrl.hostname, input.networkScope),
+      await resolveAndValidateHost(currentUrl.hostname, input.networkScope, { signal: deadlineSignal }),
       REQUEST_TIMEOUT_MS,
     );
     const headers = buildHeaders(input, currentUrl, currentBody);
     const controller = new AbortController();
     const timeout = setTimeout(() => controller.abort(), REQUEST_TIMEOUT_MS);
-    let responseData: Awaited<ReturnType<typeof dispatcher.request>>;
+    let result: RequestHopResult | undefined;
+    let lifecycleFailure: { error: unknown } | undefined;
     try {
-      responseData = await dispatcher.request({
+      const responseData = await dispatcher.request({
         origin: currentUrl.origin,
         path: `${currentUrl.pathname}${currentUrl.search}`,
         method: currentMethod,
@@ -112,54 +128,88 @@ async function performRequestWithinDeadline(
         body: currentBody,
         signal: AbortSignal.any([deadlineSignal, controller.signal]),
       });
-    } catch (error) {
-      await dispatcher.close();
-      if (error instanceof CustomWidgetDomainError) throw error;
-      input.logError?.({
-        origin: currentUrl.origin,
-        method: currentMethod,
-        errorName: error instanceof Error ? error.name : "UnknownError",
-      });
-      throw new CustomWidgetDomainError({
-        code: "BAD_GATEWAY",
-        message: error instanceof Error ? error.message : "External request failed",
-        cause: error,
-      });
-    } finally {
-      clearTimeout(timeout);
-    }
-    if (![301, 302, 303, 307, 308].includes(responseData.statusCode)) {
-      try {
+      if (![301, 302, 303, 307, 308].includes(responseData.statusCode)) {
         const body = await responseData.body.arrayBuffer();
         const response = new Response(body.byteLength > 0 ? body : null, {
           status: responseData.statusCode,
           statusText: STATUS_CODES[responseData.statusCode] ?? "",
           headers: normalizeResponseHeaders(responseData.headers),
         });
-        return {
-          ok: response.ok,
-          status: response.status,
-          statusText: response.statusText,
-          data: await parseResponseBody(response),
+        result = {
+          kind: "response",
+          response: {
+            ok: response.ok,
+            status: response.status,
+            statusText: response.statusText,
+            data: await parseResponseBody(response),
+          },
         };
-      } finally {
-        await dispatcher.close();
+      } else {
+        await responseData.body.dump();
+        result = {
+          kind: "redirect",
+          location: normalizeResponseHeaders(responseData.headers).get("location"),
+          statusCode: responseData.statusCode,
+        };
+      }
+    } catch (error) {
+      lifecycleFailure = { error };
+    } finally {
+      clearTimeout(timeout);
+      try {
+        const closed = await closeDispatcher(dispatcher, deadlineSignal);
+        if (!closed)
+          input.logError?.({
+            origin: currentUrl.origin,
+            method: currentMethod,
+            errorName: "DispatcherCloseTimeout",
+          });
+      } catch {
+        input.logError?.({
+          origin: currentUrl.origin,
+          method: currentMethod,
+          errorName: "DispatcherCloseError",
+        });
       }
     }
-    try {
-      await responseData.body.dump();
-    } finally {
-      await dispatcher.close();
+
+    if (lifecycleFailure) {
+      const { error } = lifecycleFailure;
+      if (error instanceof CustomWidgetDomainError) throw error;
+      if (error instanceof Error && "code" in error && error.code === RESPONSE_TOO_LARGE_ERROR_CODE) {
+        throw new CustomWidgetDomainError({
+          code: "PAYLOAD_TOO_LARGE",
+          message: "Response exceeds the 1 MiB limit",
+        });
+      }
+      const timedOut = isCustomWidgetRequestTimeoutError(error, deadlineSignal, controller.signal);
+      input.logError?.({
+        origin: currentUrl.origin,
+        method: currentMethod,
+        errorName: timedOut ? "RequestTimeout" : "TransportError",
+        ...(timedOut ? { reason: "timeout" as const } : {}),
+      });
+      throw new CustomWidgetDomainError({
+        code: "BAD_GATEWAY",
+        message: timedOut ? "External request timed out" : "External request failed",
+        ...(timedOut ? { reason: "timeout" as const } : {}),
+      });
     }
+
+    if (!result)
+      throw new CustomWidgetDomainError({
+        code: "INTERNAL_SERVER_ERROR",
+        message: "External request did not complete",
+      });
+    if (result.kind === "response") return result.response;
     if (redirects >= maxRedirects)
       throw new CustomWidgetDomainError({ code: "BAD_GATEWAY", message: "Upstream redirect limit exceeded" });
-    const location = normalizeResponseHeaders(responseData.headers).get("location");
-    if (!location)
+    if (!result.location)
       throw new CustomWidgetDomainError({ code: "BAD_GATEWAY", message: "Upstream redirect is missing a location" });
-    const redirected = validateCustomWidgetUrl(new URL(location, currentUrl));
+    const redirected = validateCustomWidgetUrl(new URL(result.location, currentUrl));
     if (redirected.origin !== baseUrl.origin)
       throw new CustomWidgetDomainError({ code: "FORBIDDEN", message: "Cross-origin redirects are not allowed" });
-    if (responseData.statusCode === 303 || ([301, 302].includes(responseData.statusCode) && currentMethod === "POST")) {
+    if (result.statusCode === 303 || ([301, 302].includes(result.statusCode) && currentMethod === "POST")) {
       currentMethod = "GET";
       currentBody = undefined;
     }
@@ -167,6 +217,14 @@ async function performRequestWithinDeadline(
   }
 }
 
+export function isCustomWidgetRequestTimeoutError(error: unknown, ...signals: readonly AbortSignal[]): boolean {
+  if (signals.some((signal) => signal.aborted)) return true;
+  if (!(error instanceof Error)) return false;
+  if (error.name.includes("Timeout")) return true;
+  const code = "code" in error && typeof error.code === "string" ? error.code : undefined;
+  return code !== undefined && TIMEOUT_ERROR_CODES.has(code);
+}
+
 function normalizeResponseHeaders(values: Record<string, string | string[] | undefined>) {
   const headers = new Headers();
   for (const [name, value] of Object.entries(values)) {

packages/custom-widgets/src/server/request-limits.ts

@@ -7,12 +7,16 @@ const WINDOW_MS = 60_000;
 // Keep the lease above that total and identify every acquisition so a late
 // release can never remove capacity owned by a newer request.
 export const CUSTOM_WIDGET_REQUEST_CONCURRENCY_TTL_MS = 60_000;
+export const CUSTOM_WIDGET_USER_ITEM_CONCURRENCY_LIMIT = 4;
+export const CUSTOM_WIDGET_DEFINITION_CONCURRENCY_LIMIT = 8;
 const CATEGORY_LIMITS = { query: 60, action: 10, delete: 3 } as const;
+const DEFINITION_CATEGORY_LIMITS = { query: 240, action: 40, delete: 12 } as const;
 type RequestCategory = keyof typeof CATEGORY_LIMITS;
 
 export interface RequestLimitInput {
   category: RequestCategory;
   userId?: string;
+  anonymousId?: string;
   itemId: string;
   definitionId: string;
 }
@@ -35,33 +39,45 @@ export class CustomWidgetRequestLimiter {
   public constructor(private readonly options: RequestLimiterOptions = {}) {}
 
   public async acquire(input: RequestLimitInput): Promise<() => Promise<void>> {
-    const identity = input.userId ?? "anonymous";
+    const identity = input.userId ? `user:${input.userId}` : `anonymous:${input.anonymousId ?? "unknown"}`;
     const rateKey = `custom-widget:rate:${input.category}:${hashKey(identity, input.itemId)}`;
+    const definitionRateKey = `custom-widget:rate:v2:definition:${input.category}:${hashKey(input.definitionId)}`;
     const userKey = `custom-widget:concurrency:v2:user-item:${hashKey(identity, input.itemId)}`;
     const definitionKey = `custom-widget:concurrency:v2:definition:${hashKey(input.definitionId)}`;
-    if (!this.options.store) return this.acquireLocal(rateKey, userKey, definitionKey, input.category);
+    if (!this.options.store)
+      return this.acquireLocal(rateKey, definitionRateKey, userKey, definitionKey, input.category);
     try {
-      return await this.acquireStored(rateKey, userKey, definitionKey, input.category, this.options.store);
+      return await this.acquireStored(
+        rateKey,
+        definitionRateKey,
+        userKey,
+        definitionKey,
+        input.category,
+        this.options.store,
+      );
     } catch (error) {
       if (error instanceof CustomWidgetDomainError) throw error;
       this.options.onStoreError?.(error);
       throw new CustomWidgetDomainError({
         code: "INTERNAL_SERVER_ERROR",
         message: "Request limiter is unavailable",
-        cause: error,
       });
     }
   }
 
-  private acquireLocal(rateKey: string, userKey: string, definitionKey: string, category: RequestCategory) {
+  private acquireLocal(
+    rateKey: string,
+    definitionRateKey: string,
+    userKey: string,
+    definitionKey: string,
+    category: RequestCategory,
+  ) {
     const now = Date.now();
-    const recent = (this.rateBuckets.get(rateKey) ?? []).filter((timestamp) => now - timestamp < WINDOW_MS);
-    if (recent.length >= CATEGORY_LIMITS[category]) throwLimit(WINDOW_MS - (now - (recent[0] ?? now)));
-    recent.push(now);
-    this.rateBuckets.set(rateKey, recent);
-    this.acquireLocalConcurrency(userKey, 4);
+    this.acquireLocalRate(rateKey, CATEGORY_LIMITS[category], now);
+    this.acquireLocalRate(definitionRateKey, DEFINITION_CATEGORY_LIMITS[category], now);
+    this.acquireLocalConcurrency(userKey, CUSTOM_WIDGET_USER_ITEM_CONCURRENCY_LIMIT);
     try {
-      this.acquireLocalConcurrency(definitionKey, 8);
+      this.acquireLocalConcurrency(definitionKey, CUSTOM_WIDGET_DEFINITION_CONCURRENCY_LIMIT);
     } catch (error) {
       this.releaseLocalConcurrency(userKey);
       throw error;
@@ -74,6 +90,7 @@ export class CustomWidgetRequestLimiter {
 
   private async acquireStored(
     rateKey: string,
+    definitionRateKey: string,
     userKey: string,
     definitionKey: string,
     category: RequestCategory,
@@ -82,11 +99,27 @@ export class CustomWidgetRequestLimiter {
     const ownerId = randomUUID();
     const rate = await store.incrementRate(rateKey, WINDOW_MS);
     if (rate.count > CATEGORY_LIMITS[category]) throwLimit(rate.retryAfterMs);
-    if (!(await store.acquireConcurrency(userKey, ownerId, 4, CUSTOM_WIDGET_REQUEST_CONCURRENCY_TTL_MS))) {
+    const definitionRate = await store.incrementRate(definitionRateKey, WINDOW_MS);
+    if (definitionRate.count > DEFINITION_CATEGORY_LIMITS[category]) throwLimit(definitionRate.retryAfterMs);
+    if (
+      !(await store.acquireConcurrency(
+        userKey,
+        ownerId,
+        CUSTOM_WIDGET_USER_ITEM_CONCURRENCY_LIMIT,
+        CUSTOM_WIDGET_REQUEST_CONCURRENCY_TTL_MS,
+      ))
+    ) {
       throwLimit(1_000);
     }
     try {
-      if (!(await store.acquireConcurrency(definitionKey, ownerId, 8, CUSTOM_WIDGET_REQUEST_CONCURRENCY_TTL_MS))) {
+      if (
+        !(await store.acquireConcurrency(
+          definitionKey,
+          ownerId,
+          CUSTOM_WIDGET_DEFINITION_CONCURRENCY_LIMIT,
+          CUSTOM_WIDGET_REQUEST_CONCURRENCY_TTL_MS,
+        ))
+      ) {
         throwLimit(1_000);
       }
     } catch (error) {
@@ -98,6 +131,13 @@ export class CustomWidgetRequestLimiter {
     };
   }
 
+  private acquireLocalRate(key: string, limit: number, now: number): void {
+    const recent = (this.rateBuckets.get(key) ?? []).filter((timestamp) => now - timestamp < WINDOW_MS);
+    if (recent.length >= limit) throwLimit(WINDOW_MS - (now - (recent[0] ?? now)));
+    recent.push(now);
+    this.rateBuckets.set(key, recent);
+  }
+
   private acquireLocalConcurrency(key: string, limit: number): void {
     const count = this.concurrency.get(key) ?? 0;
     if (count >= limit) throwLimit(1_000);