fix: move domain access tracking from middleware to query handlers

Author: jakejarvis

server/routers/domain.ts

@@ -20,7 +20,11 @@ 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, domainProcedure } from "@/trpc/init";
+import {
+  createTRPCRouter,
+  domainProcedure,
+  trackDomainAccess,
+} from "@/trpc/init";
 
 const DomainInputSchema = z
   .object({ domain: z.string().min(1) })
@@ -35,7 +39,10 @@ export const domainRouter = createTRPCRouter({
     .meta({ service: "registration" })
     .input(DomainInputSchema)
     .output(RegistrationSchema)
-    .query(({ input }) => getRegistration(input.domain)),
+    .query(({ input }) => {
+      trackDomainAccess(input.domain);
+      return getRegistration(input.domain);
+    }),
   pricing: domainProcedure
     .meta({ service: "pricing" })
     .input(DomainInputSchema)
@@ -45,27 +52,42 @@ export const domainRouter = createTRPCRouter({
     .meta({ service: "dns" })
     .input(DomainInputSchema)
     .output(DnsResolveResultSchema)
-    .query(({ input }) => resolveAll(input.domain)),
+    .query(({ input }) => {
+      trackDomainAccess(input.domain);
+      return resolveAll(input.domain);
+    }),
   hosting: domainProcedure
     .meta({ service: "hosting" })
     .input(DomainInputSchema)
     .output(HostingSchema)
-    .query(({ input }) => detectHosting(input.domain)),
+    .query(({ input }) => {
+      trackDomainAccess(input.domain);
+      return detectHosting(input.domain);
+    }),
   certificates: domainProcedure
     .meta({ service: "certs" })
     .input(DomainInputSchema)
     .output(CertificatesSchema)
-    .query(({ input }) => getCertificates(input.domain)),
+    .query(({ input }) => {
+      trackDomainAccess(input.domain);
+      return getCertificates(input.domain);
+    }),
   headers: domainProcedure
     .meta({ service: "headers" })
     .input(DomainInputSchema)
     .output(HttpHeadersResponseSchema)
-    .query(({ input }) => probeHeaders(input.domain)),
+    .query(({ input }) => {
+      trackDomainAccess(input.domain);
+      return probeHeaders(input.domain);
+    }),
   seo: domainProcedure
     .meta({ service: "seo" })
     .input(DomainInputSchema)
     .output(SeoResponseSchema)
-    .query(({ input }) => getSeo(input.domain)),
+    .query(({ input }) => {
+      trackDomainAccess(input.domain);
+      return getSeo(input.domain);
+    }),
   favicon: domainProcedure
     .meta({ service: "favicon" })
     .input(DomainInputSchema)

trpc/init.ts

@@ -95,40 +95,38 @@ const withRatelimit = t.middleware(async ({ ctx, next, meta }) => {
   return next();
 });
 
-/**
- * Middleware to record domain access for decay calculation.
- * Expects input to have a `domain` field containing the domain string.
- */
-const withAccess = t.middleware(async ({ input, next }) => {
-  // Extract domain from input and record access if it's a valid registrable domain
-  const domain =
-    typeof input === "object" &&
-    input !== null &&
-    "domain" in input &&
-    typeof input.domain === "string"
-      ? input.domain
-      : null;
-
-  if (domain) {
-    const registrable = toRegistrableDomain(domain);
-    if (registrable) {
-      recordDomainAccess(registrable);
-    }
-  }
-
-  return next();
-});
-
 /**
  * Public procedure with logging.
  * Use this for all public endpoints (e.g. health check, etc).
  */
 export const publicProcedure = t.procedure.use(withLogging);
 
 /**
- * Domain-specific procedure with rate limiting and access tracking.
+ * Domain-specific procedure with rate limiting.
  * Use this for all domain data endpoints (dns, hosting, seo, etc).
+ *
+ * Note: Access tracking is done in the query handler, not middleware,
+ * because the input hasn't been parsed yet at the middleware stage.
+ */
+export const domainProcedure = publicProcedure.use(withRatelimit);
+
+/**
+ * Helper to track domain access in query handlers.
+ * Call this at the start of domain-related query handlers.
+ *
+ * @param domain - The domain from the input
+ * @example
+ * ```ts
+ * .query(({ input }) => {
+ *   trackDomainAccess(input.domain);
+ *   return getSeo(input.domain);
+ * })
+ * ```
  */
-export const domainProcedure = publicProcedure
-  .use(withRatelimit)
-  .use(withAccess);
+export function trackDomainAccess(domain: string): void {
+  const registrable = toRegistrableDomain(domain);
+  if (registrable) {
+    console.debug(`[trpc] recording access for domain: ${registrable}`);
+    recordDomainAccess(registrable);
+  }
+}