Add rate limiting functionality using @upstash/ratelimit and integrate with TRPC procedures

app/api/trpc/[trpc]/route.ts

@@ -9,7 +9,29 @@ const handler = (req: Request) =>
     endpoint: "/api/trpc",
     req,
     router: appRouter,
-    createContext,
+    createContext: () => createContext({ req }),
+    responseMeta: ({ errors }) => {
+      const err = errors.find((e) => e.code === "TOO_MANY_REQUESTS");
+      if (!err) return {};
+
+      // Prefer formatted data from errorFormatter for consistent headers
+      const data = (
+        err as {
+          data?: { retryAfter?: number; limit?: number; remaining?: number };
+        }
+      ).data;
+      const retryAfter = Math.max(1, Math.round(Number(data?.retryAfter ?? 1)));
+      const headers: Record<string, string> = {
+        "Retry-After": String(retryAfter),
+        "Cache-Control": "no-cache, no-store",
+      };
+      if (typeof data?.limit === "number")
+        headers["X-RateLimit-Limit"] = String(data.limit);
+      if (typeof data?.remaining === "number")
+        headers["X-RateLimit-Remaining"] = String(data.remaining);
+
+      return { headers, status: 429 };
+    },
     onError: ({ path, error }) => {
       // Development logging
       if (process.env.NODE_ENV === "development") {

lib/analytics/server.ts

@@ -65,7 +65,11 @@ export const captureServer = async (
   });
 
   // flush events to posthog in background
-  waitUntil(client.shutdown());
+  try {
+    waitUntil?.(client.shutdown());
+  } catch {
+    // no-op
+  }
 };
 
 export const captureServerException = async (
@@ -85,5 +89,9 @@ export const captureServerException = async (
   );
 
   // flush events to posthog in background
-  waitUntil(client.shutdown());
+  try {
+    waitUntil?.(client.shutdown());
+  } catch {
+    // no-op
+  }
 };

lib/schemas/internal/storage.ts

@@ -1,5 +1,7 @@
 import { z } from "zod";
 
 export const StorageKindSchema = z.enum(["favicon", "screenshot", "social"]);
+export const StorageUrlSchema = z.object({ url: z.string().url().nullable() });
 
 export type StorageKind = z.infer<typeof StorageKindSchema>;
+export type StorageUrl = z.infer<typeof StorageUrlSchema>;

package.json

@@ -40,6 +40,7 @@
     "@trpc/client": "^11.6.0",
     "@trpc/server": "^11.6.0",
     "@trpc/tanstack-react-query": "^11.6.0",
+    "@upstash/ratelimit": "^2.0.6",
     "@upstash/redis": "^1.35.6",
     "@vercel/analytics": "^1.5.0",
     "@vercel/functions": "^3.1.4",

server/ratelimit.ts

@@ -0,0 +1,75 @@
+import "server-only";
+
+import { TRPCError } from "@trpc/server";
+import { Ratelimit } from "@upstash/ratelimit";
+import { waitUntil } from "@vercel/functions";
+import { redis } from "@/lib/redis";
+import { t } from "@/trpc/init";
+
+export const SERVICE_LIMITS = {
+  dns: { points: 60, window: "1 m" },
+  headers: { points: 60, window: "1 m" },
+  certs: { points: 30, window: "1 m" },
+  registration: { points: 4, window: "1 m" },
+  screenshot: { points: 3, window: "1 m" },
+  favicon: { points: 120, window: "1 h" },
+  seo: { points: 30, window: "1 m" },
+  hosting: { points: 30, window: "1 m" },
+  pricing: { points: 30, window: "1 m" },
+} as const;
+
+export type ServiceName = keyof typeof SERVICE_LIMITS;
+
+const limiters = Object.fromEntries(
+  Object.entries(SERVICE_LIMITS).map(([service, cfg]) => [
+    service,
+    new Ratelimit({
+      redis,
+      limiter: Ratelimit.slidingWindow(
+        cfg.points,
+        cfg.window as `${number} ${"s" | "m" | "h"}`,
+      ),
+      prefix: `@upstash/ratelimit:${service}`,
+      analytics: true,
+    }),
+  ]),
+) as Record<ServiceName, Ratelimit>;
+
+export async function assertRateLimit(service: ServiceName, ip: string) {
+  const res = await limiters[service].limit(ip);
+
+  if (!res.success) {
+    const retryAfterSec = Math.max(
+      1,
+      Math.ceil((res.reset - Date.now()) / 1000),
+    );
+
+    throw new TRPCError({
+      code: "TOO_MANY_REQUESTS",
+      message: `Rate limit exceeded for ${service}. Try again in ${retryAfterSec}s.`,
+      cause: {
+        retryAfter: retryAfterSec,
+        service,
+        limit: res.limit,
+        remaining: res.remaining,
+        reset: res.reset,
+      },
+    });
+  }
+
+  // allow ratelimit analytics to be sent in background
+  try {
+    waitUntil?.(res.pending);
+  } catch {
+    // no-op
+  }
+
+  return { limit: res.limit, remaining: res.remaining, reset: res.reset };
+}
+
+export const rateLimitMiddleware = t.middleware(async ({ ctx, next, meta }) => {
+  const service = (meta?.service ?? "") as ServiceName;
+  if (!service || !(service in SERVICE_LIMITS) || !ctx.ip) return next();
+  await assertRateLimit(service, ctx.ip);
+  return next();
+});

server/routers/domain.ts

@@ -9,8 +9,10 @@ import {
   PricingSchema,
   RegistrationSchema,
   SeoResponseSchema,
+  StorageUrlSchema,
 } from "@/lib/schemas";
 import { inngest } from "@/server/inngest/client";
+import { rateLimitMiddleware } from "@/server/ratelimit";
 import { getCertificates } from "@/server/services/certificates";
 import { resolveAll } from "@/server/services/dns";
 import { getOrCreateFaviconBlobUrl } from "@/server/services/favicon";
@@ -20,9 +22,9 @@ import { getPricingForTld } from "@/server/services/pricing";
 import { getRegistration } from "@/server/services/registration";
 import { getOrCreateScreenshotBlobUrl } from "@/server/services/screenshot";
 import { getSeo } from "@/server/services/seo";
-import { createTRPCRouter, loggedProcedure } from "@/trpc/init";
+import { createTRPCRouter, publicProcedure } from "@/trpc/init";
 
-export const domainInput = z
+const DomainInputSchema = z
   .object({ domain: z.string().min(1) })
   .transform(({ domain }) => ({ domain: normalizeDomainInput(domain) }))
   .refine(({ domain }) => toRegistrableDomain(domain) !== null, {
@@ -31,16 +33,22 @@ export const domainInput = z
   });
 
 export const domainRouter = createTRPCRouter({
-  registration: loggedProcedure
-    .input(domainInput)
+  registration: publicProcedure
+    .meta({ service: "registration" })
+    .use(rateLimitMiddleware)
+    .input(DomainInputSchema)
     .output(RegistrationSchema)
     .query(({ input }) => getRegistration(input.domain)),
-  pricing: loggedProcedure
-    .input(domainInput)
+  pricing: publicProcedure
+    .meta({ service: "pricing" })
+    .use(rateLimitMiddleware)
+    .input(DomainInputSchema)
     .output(PricingSchema)
     .query(({ input }) => getPricingForTld(input.domain)),
-  dns: loggedProcedure
-    .input(domainInput)
+  dns: publicProcedure
+    .meta({ service: "dns" })
+    .use(rateLimitMiddleware)
+    .input(DomainInputSchema)
     .output(DnsResolveResultSchema)
     .query(async ({ input }) => {
       const result = await resolveAll(input.domain);
@@ -51,26 +59,40 @@ export const domainRouter = createTRPCRouter({
       });
       return result;
     }),
-  hosting: loggedProcedure
-    .input(domainInput)
+  hosting: publicProcedure
+    .meta({ service: "hosting" })
+    .use(rateLimitMiddleware)
+    .input(DomainInputSchema)
     .output(HostingSchema)
     .query(({ input }) => detectHosting(input.domain)),
-  certificates: loggedProcedure
-    .input(domainInput)
+  certificates: publicProcedure
+    .meta({ service: "certs" })
+    .use(rateLimitMiddleware)
+    .input(DomainInputSchema)
     .output(CertificatesSchema)
     .query(({ input }) => getCertificates(input.domain)),
-  headers: loggedProcedure
-    .input(domainInput)
+  headers: publicProcedure
+    .meta({ service: "headers" })
+    .use(rateLimitMiddleware)
+    .input(DomainInputSchema)
     .output(HttpHeadersSchema)
     .query(({ input }) => probeHeaders(input.domain)),
-  seo: loggedProcedure
-    .input(domainInput)
+  seo: publicProcedure
+    .meta({ service: "seo" })
+    .use(rateLimitMiddleware)
+    .input(DomainInputSchema)
     .output(SeoResponseSchema)
     .query(({ input }) => getSeo(input.domain)),
-  favicon: loggedProcedure
-    .input(domainInput)
+  favicon: publicProcedure
+    .meta({ service: "favicon" })
+    .use(rateLimitMiddleware)
+    .input(DomainInputSchema)
+    .output(StorageUrlSchema)
     .query(({ input }) => getOrCreateFaviconBlobUrl(input.domain)),
-  screenshot: loggedProcedure
-    .input(domainInput)
+  screenshot: publicProcedure
+    .meta({ service: "screenshot" })
+    .use(rateLimitMiddleware)
+    .input(DomainInputSchema)
+    .output(StorageUrlSchema)
     .query(({ input }) => getOrCreateScreenshotBlobUrl(input.domain)),
 });

trpc/client.tsx

@@ -11,6 +11,7 @@ import { useState } from "react";
 import superjson from "superjson";
 import { TRPCProvider as Provider } from "@/lib/trpc/client";
 import type { AppRouter } from "@/server/routers/_app";
+import { errorToastLink } from "@/trpc/error-toast-link";
 import { getQueryClient } from "@/trpc/query-client";
 
 let browserQueryClient: ReturnType<typeof getQueryClient> | undefined;
@@ -37,6 +38,7 @@ export function TRPCProvider({ children }: { children: React.ReactNode }) {
             process.env.NODE_ENV === "development" ||
             (opts.direction === "down" && opts.result instanceof Error),
         }),
+        errorToastLink(),
         httpBatchStreamLink({
           url: `${getBaseUrl()}/api/trpc`,
           transformer: superjson,

trpc/error-toast-link.tsx

@@ -0,0 +1,60 @@
+"use client";
+
+import { TRPCClientError, type TRPCLink } from "@trpc/client";
+import type { AnyRouter } from "@trpc/server";
+import { observable } from "@trpc/server/observable";
+import { Siren } from "lucide-react";
+import { toast } from "sonner";
+
+function formatWait(seconds: number): string {
+  if (!Number.isFinite(seconds) || seconds <= 1) return "a moment";
+  const s = Math.round(seconds);
+  const m = Math.floor(s / 60);
+  const sec = s % 60;
+  if (m <= 0) return `${sec}s`;
+  if (m < 60) return sec ? `${m}m ${sec}s` : `${m}m`;
+  const h = Math.floor(m / 60);
+  const rm = m % 60;
+  return rm ? `${h}h ${rm}m` : `${h}h`;
+}
+
+export function errorToastLink<
+  TRouter extends AnyRouter = AnyRouter,
+>(): TRPCLink<TRouter> {
+  return () =>
+    ({ next, op }) =>
+      observable((observer) => {
+        const sub = next(op).subscribe({
+          next(value) {
+            observer.next(value);
+          },
+          error(err) {
+            if (err instanceof TRPCClientError) {
+              const code = err.data?.code;
+              if (code === "TOO_MANY_REQUESTS") {
+                const retryAfterSec = Math.max(
+                  1,
+                  Math.round(Number(err.data?.retryAfter ?? 1)),
+                );
+                const service = err.data?.service as string | undefined;
+                const friendly = formatWait(retryAfterSec);
+                const title = service
+                  ? `Too many ${service} requests`
+                  : "You're doing that too much";
+                toast.error(title, {
+                  id: "rate-limit",
+                  description: `Try again in ${friendly}.`,
+                  icon: <Siren className="h-4 w-4" />,
+                  position: "top-center",
+                });
+              }
+            }
+            observer.error(err);
+          },
+          complete() {
+            observer.complete();
+          },
+        });
+        return () => sub.unsubscribe();
+      });
+}