feat: propagate a 429 status code if email-provider rate limiting occurs

src/extensions/email_provider_rate_limit_extension.ts

@@ -0,0 +1,22 @@
+import { flattenErrors, statusCodeForError } from "lib/graphqlErrorHandler"
+
+// Log in `extensions` if this request was by an email provider and resulted
+// in a 429 rate limit error from a downstream request made using `fetch`.
+export const emailProviderRateLimitExtension = (_documentAST, result) => {
+  const extensions = {}
+
+  if (result.errors && result.errors.length) {
+    result.errors.some((err) => {
+      flattenErrors(err).some((e) => {
+        const httpStatusCode = statusCodeForError(e)
+        if (httpStatusCode === 429) {
+          extensions["emailProviderLimited"] = true
+          return true
+        }
+        return false
+      })
+    })
+  }
+
+  return extensions
+}

src/index.ts

@@ -38,6 +38,7 @@ import { principalFieldDirectiveValidation } from "directives/principleField/pri
 import * as Sentry from "@sentry/node"
 import { bodyParserMiddleware } from "lib/bodyParserMiddleware"
 import { initilizeFeatureFlags } from "lib/featureFlags"
+import { emailProviderRateLimitExtension } from "extensions/email_provider_rate_limit_extension"
 
 // Initialize Unleash feature flags as early as possible
 initilizeFeatureFlags()
@@ -84,10 +85,17 @@ function createExtensions(document, result, requestID, userAgent) {
     ? fetchLoggerRequestDone(requestID, userAgent)
     : {}
 
+  // Should this introspect xapp token/roles for `email_provider` instead?
+  const isEmailProviderRequest = userAgent.includes("Braze")
+  const emailProviderExtensions = isEmailProviderRequest
+    ? emailProviderRateLimitExtension(document, result)
+    : {}
+
   const extensions = {
     ...optionalFieldsExtensions,
     ...principalFieldExtensions,
     ...requestLoggerExtensions,
+    ...emailProviderExtensions,
   }
 
   // Instead of an empty hash, which will include `extensions: {}`
@@ -248,6 +256,42 @@ export const supportedV2RouteHandler = (req, res, next, server) => {
   return server(req, res, next)
 }
 
+// Middleware that must be mounted on our GraphQL route _before_
+// our regular GraphQL handler. This is so we can register a new `res.end`
+// method that will intercept the response and check if any email provider
+// rate limting occurred.
+//
+// If so, we return a 429 status code and a JSON error message.
+app.use("/v2", function (req, res, next) {
+  const userAgent = req.headers["user-agent"] || ""
+  if (!userAgent.includes("Braze")) {
+    return next()
+  }
+
+  const originalEnd = res.end
+
+  res.end = function (...args) {
+    res.end = originalEnd
+    const [chunk, encoding, cb] = args
+
+    const contentType = res.get("content-type") || "text/html"
+
+    if (contentType.indexOf("application/json") === 0) {
+      const response = JSON.parse(chunk)
+
+      if (response.extensions && response.extensions.emailProviderLimited) {
+        res.status(429).json({ error: "rate limit reached, try again later" })
+      } else {
+        res.end(chunk, encoding, cb)
+      }
+    } else {
+      res.end(chunk, encoding, cb)
+    }
+  }
+
+  return next()
+})
+
 app.use("/v2", (req, res, next) => {
   supportedV2RouteHandler(req, res, next, graphqlServer)
 })

src/integration/__tests__/email_provider_rate_limit.test.ts

@@ -0,0 +1,38 @@
+jest.mock("lib/apis/fetch", () => jest.fn())
+import { HTTPError } from "lib/HTTPError"
+import fetch from "lib/apis/fetch"
+const mockFetch = fetch as jest.Mock
+
+describe("rate limiting custom status code and response", () => {
+  const request = require("supertest")
+  const app = require("../../index").default
+  const gql = require("lib/gql").default
+
+  beforeEach(() => {
+    mockFetch.mockClear()
+    mockFetch.mockReset()
+  })
+
+  it("propagates a 429 for email provider rate limiting", async () => {
+    mockFetch.mockRejectedValueOnce(
+      new HTTPError("Braze rate limit reached", 429)
+    )
+
+    const response = await request(app)
+      .post("/v2")
+      .set("Accept", "application/json")
+      .set("User-Agent", "Braze")
+      .send({
+        query: gql`
+          {
+            artist(id: "banksy") {
+              name
+            }
+          }
+        `,
+      })
+
+    expect(response.statusCode).toBe(429)
+    expect(response.body.error).toBe("rate limit reached, try again later")
+  })
+})