feat: exclude bot traffic from egress quotas

README.md

@@ -15,6 +15,12 @@ SECRET_HEADER_VALUE=SecretToken
 CHAINALYSIS_API_KEY=YourChainalysisApiKey
 ```
 
+Create `piece-retriever/.dev.vars` file with the following content:
+
+```
+BOT_TOKENS="{\"secret\":\"dev\"}"
+```
+
 Create `terminator/.dev.vars` file with the following content:
 
 ```

db/migrations/0018_retrieval_logs_from_bot.sql

@@ -0,0 +1 @@
+ALTER TABLE retrieval_logs ADD COLUMN bot_name TEXT;

piece-retriever/bin/piece-retriever.js

@@ -58,7 +58,7 @@ export default {
     const workerStartedAt = performance.now()
     const requestCountryCode = request.headers.get('CF-IPCountry')
 
-    const { payerWalletAddress, pieceCid } = parseRequest(request, env)
+    const { payerWalletAddress, pieceCid, botName } = parseRequest(request, env)
 
     httpAssert(payerWalletAddress && pieceCid, 400, 'Missing required fields')
     httpAssert(
@@ -118,6 +118,7 @@ export default {
             requestCountryCode,
             timestamp: requestTimestamp,
             dataSetId,
+            botName,
           }),
         )
         const response = new Response(
@@ -145,6 +146,7 @@ export default {
             requestCountryCode,
             timestamp: requestTimestamp,
             dataSetId,
+            botName,
           }),
         )
         const response = new Response(
@@ -184,6 +186,7 @@ export default {
               workerTtfb: firstByteAt - workerStartedAt,
             },
             dataSetId,
+            botName,
           })
 
           await updateDataSetStats(env, {
@@ -219,6 +222,7 @@ export default {
           requestCountryCode,
           timestamp: requestTimestamp,
           dataSetId: null,
+          botName,
         }),
       )
 

piece-retriever/lib/request.js

@@ -6,12 +6,14 @@ import { httpAssert } from './http-assert.js'
  * @param {Request} request
  * @param {object} options
  * @param {string} options.DNS_ROOT
+ * @param {string} options.BOT_TOKENS
  * @returns {{
  *   payerWalletAddress?: string
  *   pieceCid?: string
+ *   botName?: string
  * }}
  */
-export function parseRequest(request, { DNS_ROOT }) {
+export function parseRequest(request, { DNS_ROOT, BOT_TOKENS }) {
   const url = new URL(request.url)
   console.log('retrieval request', { DNS_ROOT, url })
 
@@ -31,5 +33,32 @@ export function parseRequest(request, { DNS_ROOT }) {
     `Invalid CID: ${pieceCid}. It is not a valid CommP (v1 or v2).`,
   )
 
-  return { payerWalletAddress, pieceCid }
+  const botName = checkBotAuthorization(request, { BOT_TOKENS })
+
+  return { payerWalletAddress, pieceCid, botName }
+}
+
+/**
+ * @param {Request} request
+ * @param {object} args
+ * @param {string} args.BOT_TOKENS
+ * @returns {string | undefined} Bot name or the access token
+ */
+export function checkBotAuthorization(request, { BOT_TOKENS }) {
+  const botTokens = JSON.parse(BOT_TOKENS)
+
+  const auth = request.headers.get('authorization')
+  if (!auth) return undefined
+
+  const [prefix, token, ...rest] = auth.split(' ')
+
+  httpAssert(
+    prefix === 'Bearer' && token && rest.length === 0,
+    401,
+    'Unauthorized: Authorization header must use Bearer scheme',
+  )
+
+  httpAssert(token in botTokens, 401, 'Unauthorized: Invalid Access Token')
+
+  return botTokens[token]
 }

piece-retriever/lib/store.js

@@ -21,6 +21,8 @@ import { httpAssert } from './http-assert.js'
  *   request originated from
  * @param {string | null} params.dataSetId - The data set ID associated with the
  *   retrieval
+ * @param {string | undefined} params.botName - The name of the bot making the
+ *   request, or null for anonymous requests
  * @returns {Promise<void>} - A promise that resolves when the log is inserted.
  */
 export async function logRetrievalResult(env, params) {
@@ -33,6 +35,7 @@ export async function logRetrievalResult(env, params) {
     performanceStats,
     requestCountryCode,
     dataSetId,
+    botName,
   } = params
 
   try {
@@ -47,9 +50,10 @@ export async function logRetrievalResult(env, params) {
         fetch_ttlb,
         worker_ttfb,
         request_country_code,
-        data_set_id
+        data_set_id,
+        bot_name
       )
-      VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
+      VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
       `,
     )
       .bind(
@@ -62,6 +66,7 @@ export async function logRetrievalResult(env, params) {
         performanceStats?.workerTtfb ?? null,
         requestCountryCode,
         dataSetId,
+        botName ?? null,
       )
       .run()
   } catch (error) {

piece-retriever/test/request.test.js

@@ -1,51 +1,104 @@
 import { describe, it, expect } from 'vitest'
-import { parseRequest } from '../lib/request.js'
+import { parseRequest, checkBotAuthorization } from '../lib/request.js'
 
 const DNS_ROOT = '.filbeam.io'
 const TEST_WALLET = 'abc123'
 const TEST_CID = 'baga123'
+const BOT_TOKENS = JSON.stringify({ secret: 'bot1' })
 
 describe('parseRequest', () => {
   it('should parse payerWalletAddress and pieceCid from a URL with both params', () => {
-    const request = { url: `https://${TEST_WALLET}${DNS_ROOT}/${TEST_CID}` }
-    const result = parseRequest(request, { DNS_ROOT })
+    const request = new Request(`https://${TEST_WALLET}${DNS_ROOT}/${TEST_CID}`)
+    const result = parseRequest(request, { DNS_ROOT, BOT_TOKENS })
     expect(result).toEqual({
       payerWalletAddress: TEST_WALLET,
       pieceCid: TEST_CID,
     })
   })
 
   it('should parse payerWalletAddress and pieceCid from a URL with leading slash', () => {
-    const request = { url: `https://${TEST_WALLET}${DNS_ROOT}//${TEST_CID}` }
-    const result = parseRequest(request, { DNS_ROOT })
+    const request = new Request(
+      `https://${TEST_WALLET}${DNS_ROOT}//${TEST_CID}`,
+    )
+    const result = parseRequest(request, { DNS_ROOT, BOT_TOKENS })
     expect(result).toEqual({
       payerWalletAddress: TEST_WALLET,
       pieceCid: TEST_CID,
     })
   })
 
   it('should return descriptive error for missing pieceCid', () => {
-    const request = { url: `https://${TEST_WALLET}${DNS_ROOT}/` }
-    expect(() => parseRequest(request, { DNS_ROOT })).toThrowError(
+    const request = new Request(`https://${TEST_WALLET}${DNS_ROOT}/`)
+    expect(() => parseRequest(request, { DNS_ROOT, BOT_TOKENS })).toThrowError(
       'Missing required path element: `/{CID}`',
     )
   })
 
   it('should return undefined for both if no params in path', () => {
-    const request = { url: 'https://filbeam.io' }
-    expect(() => parseRequest(request, { DNS_ROOT })).toThrowError(
+    const request = new Request('https://filbeam.io')
+    expect(() => parseRequest(request, { DNS_ROOT, BOT_TOKENS })).toThrowError(
       'Invalid hostname: filbeam.io. It must end with .filbeam.io.',
     )
   })
 
   it('should ignore query parameters', () => {
-    const request = {
-      url: `https://${TEST_WALLET}${DNS_ROOT}/${TEST_CID}?foo=bar`,
-    }
-    const result = parseRequest(request, { DNS_ROOT })
+    const request = new Request(
+      `https://${TEST_WALLET}${DNS_ROOT}/${TEST_CID}?foo=bar`,
+    )
+    const result = parseRequest(request, { DNS_ROOT, BOT_TOKENS })
     expect(result).toEqual({
       payerWalletAddress: TEST_WALLET,
       pieceCid: TEST_CID,
     })
   })
 })
+
+describe('checkBotAuthorization', () => {
+  it('should return undefined when no authorization header is present', () => {
+    const request = new Request('https://example.com', {
+      headers: {},
+    })
+    const result = checkBotAuthorization(request, { BOT_TOKENS })
+    expect(result).toBeUndefined()
+  })
+
+  it('should throw 401 error when authorization header is not Bearer format', () => {
+    const request = new Request('https://example.com', {
+      headers: { authorization: 'Basic sometoken' },
+    })
+    expect(() => checkBotAuthorization(request, { BOT_TOKENS })).toThrowError(
+      'Unauthorized: Authorization header must use Bearer scheme',
+    )
+  })
+
+  it('should throw 401 error when authorization header has no token after Bearer', () => {
+    const request = new Request('https://example.com', {
+      headers: { authorization: 'Bearer' },
+    })
+    expect(() => checkBotAuthorization(request, { BOT_TOKENS })).toThrowError(
+      'Unauthorized: Authorization header must use Bearer scheme',
+    )
+  })
+
+  it('should throw 401 error when token is not in BOT_TOKENS list', () => {
+    const request = new Request('https://example.com', {
+      headers: { authorization: 'Bearer invalid_token' },
+    })
+    expect(() => checkBotAuthorization(request, { BOT_TOKENS })).toThrowError(
+      'Unauthorized: Invalid Access Token',
+    )
+  })
+
+  it('should return token prefix when valid token is provided', () => {
+    const request = new Request('https://example.com', {
+      headers: { authorization: 'Bearer secret' },
+    })
+    const result = checkBotAuthorization(request, {
+      BOT_TOKENS: JSON.stringify({
+        secret: 'bot1',
+        secret_2: 'bot2',
+      }),
+    })
+    expect(result).toBe('bot1')
+  })
+})