From 1c42157c593f8915dfb4bebfe37e3dddb956de31 Mon Sep 17 00:00:00 2001
From: Daniel Holmgren <dtholmgren@gmail.com>
Date: Thu, 14 Sep 2023 13:54:47 -0500
Subject: [PATCH] Improve xrpc server error handling (#1597)

improve xrpc server error handling
---
 packages/xrpc-server/src/rate-limiter.ts | 21 ++++++++++++++-------
 packages/xrpc-server/src/server.ts       | 23 ++++++++++++++++++-----
 packages/xrpc-server/src/types.ts        |  4 ++--
 3 files changed, 34 insertions(+), 14 deletions(-)

diff --git a/packages/xrpc-server/src/rate-limiter.ts b/packages/xrpc-server/src/rate-limiter.ts
index 82719101674..ab6d79d07af 100644
--- a/packages/xrpc-server/src/rate-limiter.ts
+++ b/packages/xrpc-server/src/rate-limiter.ts
@@ -61,7 +61,7 @@ export class RateLimiter implements RateLimiterI {
   async consume(
     ctx: XRPCReqContext,
     opts?: { calcKey?: CalcKeyFn; calcPoints?: CalcPointsFn },
-  ): Promise<RateLimiterStatus | null> {
+  ): Promise<RateLimiterStatus | RateLimitExceededError | null> {
     if (
       this.byPassSecret &&
       ctx.req.header('x-ratelimit-bypass') === this.byPassSecret
@@ -82,7 +82,7 @@ export class RateLimiter implements RateLimiterI {
       // yes this library rejects with a res not an error
       if (err instanceof RateLimiterRes) {
         const status = formatLimiterStatus(this.limiter, err)
-        throw new RateLimitExceededError(status)
+        return new RateLimitExceededError(status)
       } else {
         if (this.failClosed) {
           throw err
@@ -119,12 +119,18 @@ export const formatLimiterStatus = (
 export const consumeMany = async (
   ctx: XRPCReqContext,
   fns: RateLimiterConsume[],
-): Promise<void> => {
-  if (fns.length === 0) return
+): Promise<RateLimiterStatus | RateLimitExceededError | null> => {
+  if (fns.length === 0) return null
   const results = await Promise.all(fns.map((fn) => fn(ctx)))
   const tightestLimit = getTightestLimit(results)
-  if (tightestLimit !== null) {
+  if (tightestLimit === null) {
+    return null
+  } else if (tightestLimit instanceof RateLimitExceededError) {
+    setResHeaders(ctx, tightestLimit.status)
+    return tightestLimit
+  } else {
     setResHeaders(ctx, tightestLimit)
+    return tightestLimit
   }
 }
 
@@ -142,11 +148,12 @@ export const setResHeaders = (
 }
 
 export const getTightestLimit = (
-  resps: (RateLimiterStatus | null)[],
-): RateLimiterStatus | null => {
+  resps: (RateLimiterStatus | RateLimitExceededError | null)[],
+): RateLimiterStatus | RateLimitExceededError | null => {
   let lowest: RateLimiterStatus | null = null
   for (const resp of resps) {
     if (resp === null) continue
+    if (resp instanceof RateLimitExceededError) return resp
     if (lowest === null || resp.remainingPoints < lowest.remainingPoints) {
       lowest = resp
     }
diff --git a/packages/xrpc-server/src/server.ts b/packages/xrpc-server/src/server.ts
index 40f4baac771..5d27e3e45a8 100644
--- a/packages/xrpc-server/src/server.ts
+++ b/packages/xrpc-server/src/server.ts
@@ -34,6 +34,7 @@ import {
   RateLimiterI,
   RateLimiterConsume,
   isShared,
+  RateLimitExceededError,
 } from './types'
 import {
   decodeQueryParams,
@@ -247,7 +248,10 @@ export class Server {
 
         // handle rate limits
         if (consumeRateLimit) {
-          await consumeRateLimit(reqCtx)
+          const result = await consumeRateLimit(reqCtx)
+          if (result instanceof RateLimitExceededError) {
+            return next(result)
+          }
         }
 
         // run the handler
@@ -475,14 +479,23 @@ function createAuthMiddleware(verifier: AuthVerifier): RequestHandler {
 const errorMiddleware: ErrorRequestHandler = function (err, req, res, next) {
   const locals: RequestLocals | undefined = req[kRequestLocals]
   const methodSuffix = locals ? ` method ${locals.nsid}` : ''
-  if (err instanceof XRPCError) {
-    log.error(err, `error in xrpc${methodSuffix}`)
-  } else {
+  const xrpcError = XRPCError.fromError(err)
+  if (xrpcError instanceof InternalServerError) {
+    // log trace for unhandled exceptions
     log.error(err, `unhandled exception in xrpc${methodSuffix}`)
+  } else {
+    // do not log trace for known xrpc errors
+    log.error(
+      {
+        status: xrpcError.type,
+        message: xrpcError.message,
+        name: xrpcError.customErrorName,
+      },
+      `error in xrpc${methodSuffix}`,
+    )
   }
   if (res.headersSent) {
     return next(err)
   }
-  const xrpcError = XRPCError.fromError(err)
   return res.status(xrpcError.type).json(xrpcError.payload)
 }
diff --git a/packages/xrpc-server/src/types.ts b/packages/xrpc-server/src/types.ts
index 801c8baa6f2..829f8a270d2 100644
--- a/packages/xrpc-server/src/types.ts
+++ b/packages/xrpc-server/src/types.ts
@@ -95,7 +95,7 @@ export interface RateLimiterI {
 export type RateLimiterConsume = (
   ctx: XRPCReqContext,
   opts?: { calcKey?: CalcKeyFn; calcPoints?: CalcPointsFn },
-) => Promise<RateLimiterStatus | null>
+) => Promise<RateLimiterStatus | RateLimitExceededError | null>
 
 export type RateLimiterCreator = (opts: {
   keyPrefix: string
@@ -224,7 +224,7 @@ export class ForbiddenError extends XRPCError {
 
 export class RateLimitExceededError extends XRPCError {
   constructor(
-    status: RateLimiterStatus,
+    public status: RateLimiterStatus,
     errorMessage?: string,
     customErrorName?: string,
   ) {