From dc7564be01e93bf12fcaf7ecd3668b705dc06636 Mon Sep 17 00:00:00 2001
From: devin ivy <devinivy@gmail.com>
Date: Tue, 5 Mar 2024 20:08:18 -0500
Subject: [PATCH] Appview: support rate limit bypass key for blob resolution
 (#2273)

bsky: support rate limit bypass key for blob resolution
---
 packages/bsky/src/api/blob-resolver.ts | 34 ++++++++++++++++++++++++--
 packages/bsky/src/config.ts            | 21 ++++++++++++++++
 2 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/packages/bsky/src/api/blob-resolver.ts b/packages/bsky/src/api/blob-resolver.ts
index 1a2d1ee560d..facb70c4d6f 100644
--- a/packages/bsky/src/api/blob-resolver.ts
+++ b/packages/bsky/src/api/blob-resolver.ts
@@ -106,7 +106,9 @@ export async function resolveBlob(ctx: AppContext, did: string, cid: CID) {
     throw createError(404, 'Blob not found')
   }
 
-  const blobResult = await retryHttp(() => getBlob({ pds, did, cid: cidStr }))
+  const blobResult = await retryHttp(() =>
+    getBlob(ctx, { pds, did, cid: cidStr }),
+  )
   const imageStream: Readable = blobResult.data
   const verifyCid = new VerifyCidTransform(cid)
 
@@ -119,12 +121,40 @@ export async function resolveBlob(ctx: AppContext, did: string, cid: CID) {
   }
 }
 
-async function getBlob(opts: { pds: string; did: string; cid: string }) {
+async function getBlob(
+  ctx: AppContext,
+  opts: { pds: string; did: string; cid: string },
+) {
   const { pds, did, cid } = opts
   return axios.get(`${pds}/xrpc/com.atproto.sync.getBlob`, {
     params: { did, cid },
     decompress: true,
     responseType: 'stream',
     timeout: 5000, // 5sec of inactivity on the connection
+    headers: getRateLimitBypassHeaders(ctx, pds),
   })
 }
+
+function getRateLimitBypassHeaders(
+  ctx: AppContext,
+  pds: string,
+): { 'x-ratelimit-bypass'?: string } {
+  const {
+    blobRateLimitBypassKey: bypassKey,
+    blobRateLimitBypassHostname: bypassHostname,
+  } = ctx.cfg
+  if (!bypassKey || !bypassHostname) {
+    return {}
+  }
+  const url = new URL(pds)
+  if (bypassHostname.startsWith('.')) {
+    if (url.hostname.endsWith(bypassHostname)) {
+      return { 'x-ratelimit-bypass': bypassKey }
+    }
+  } else {
+    if (url.hostname === bypassHostname) {
+      return { 'x-ratelimit-bypass': bypassKey }
+    }
+  }
+  return {}
+}
diff --git a/packages/bsky/src/config.ts b/packages/bsky/src/config.ts
index 6f9c96776f4..3518f7b42a9 100644
--- a/packages/bsky/src/config.ts
+++ b/packages/bsky/src/config.ts
@@ -21,6 +21,8 @@ export interface ServerConfigValues {
   courierIgnoreBadTls?: boolean
   searchUrl?: string
   cdnUrl?: string
+  blobRateLimitBypassKey?: string
+  blobRateLimitBypassHostname?: string
   // identity
   didPlcUrl: string
   handleResolveNameservers?: string[]
@@ -76,6 +78,15 @@ export class ServerConfig {
     const courierIgnoreBadTls =
       process.env.BSKY_COURIER_IGNORE_BAD_TLS === 'true'
     assert(courierHttpVersion === '1.1' || courierHttpVersion === '2')
+    const blobRateLimitBypassKey =
+      process.env.BSKY_BLOB_RATE_LIMIT_BYPASS_KEY || undefined
+    // single domain would be e.g. "mypds.com", subdomains are supported with a leading dot e.g. ".mypds.com"
+    const blobRateLimitBypassHostname =
+      process.env.BSKY_BLOB_RATE_LIMIT_BYPASS_HOSTNAME || undefined
+    assert(
+      !blobRateLimitBypassKey || blobRateLimitBypassHostname,
+      'must specify a hostname when using a blob rate limit bypass key',
+    )
     const adminPasswords = envList(
       process.env.BSKY_ADMIN_PASSWORDS || process.env.BSKY_ADMIN_PASSWORD,
     )
@@ -106,6 +117,8 @@ export class ServerConfig {
       courierApiKey,
       courierHttpVersion,
       courierIgnoreBadTls,
+      blobRateLimitBypassKey,
+      blobRateLimitBypassHostname,
       adminPasswords,
       modServiceDid,
       ...stripUndefineds(overrides ?? {}),
@@ -197,6 +210,14 @@ export class ServerConfig {
     return this.cfg.cdnUrl
   }
 
+  get blobRateLimitBypassKey() {
+    return this.cfg.blobRateLimitBypassKey
+  }
+
+  get blobRateLimitBypassHostname() {
+    return this.cfg.blobRateLimitBypassHostname
+  }
+
   get didPlcUrl() {
     return this.cfg.didPlcUrl
   }