feat: ratelimit config from source (netlify/edge-bundler#583)

* feat: ratelimit config from source

* feat: config consistent with lambda

* chore: ratelimit -> rate_limit

packages/edge-bundler/node/config.test.ts

@@ -13,6 +13,7 @@ import { bundle } from './bundler.js'
 import { FunctionConfig, getFunctionConfig } from './config.js'
 import type { Declaration } from './declaration.js'
 import { ImportMap } from './import_map.js'
+import { RateLimitAction, RateLimitAggregator } from './rate_limit.js'
 
 const importMapFile = {
   baseURL: new URL('file:///some/path/import-map.json'),
@@ -83,7 +84,7 @@ const functions: TestFunctions[] = [
   },
   {
     testName: 'config with wrong onError',
-    name: 'func7',
+    name: 'func6',
     source: `
       export default async () => new Response("Hello from function two")
       export const config = { onError: "foo" }
@@ -93,7 +94,7 @@ const functions: TestFunctions[] = [
   {
     testName: 'config with `path`',
     expectedConfig: { path: '/home' },
-    name: 'func6',
+    name: 'func7',
     source: `
         export default async () => new Response("Hello from function three")
 
@@ -108,17 +109,74 @@ const functions: TestFunctions[] = [
       name: 'a displayName',
       onError: 'bypass',
     },
-    name: 'func6',
+    name: 'func8',
     source: `
         export default async () => new Response("Hello from function three")
 
-        export const config = { path: "/home",
+        export const config = {
+          path: "/home",
           generator: '@netlify/fake-plugin@1.0.0',
           name: 'a displayName',
           onError: 'bypass',
         }
       `,
   },
+  {
+    testName: 'config with ratelimit',
+    expectedConfig: {
+      path: '/ratelimit',
+      name: 'a limit rate',
+      rateLimit: {
+        windowSize: 10,
+        windowLimit: 100,
+        aggregateBy: [RateLimitAggregator.IP, RateLimitAggregator.Domain],
+      },
+    },
+    name: 'func9',
+    source: `
+        export default async () => new Response("Rate my limits")
+
+        export const config = {
+          path: "/ratelimit",
+          rateLimit: {
+            windowSize: 10,
+            windowLimit: 100,
+            aggregateBy: ["ip", "domain"],
+          },
+          name: 'a limit rate',
+        }
+      `,
+  },
+  {
+    testName: 'config with rewrite',
+    expectedConfig: {
+      path: '/rewrite',
+      name: 'a limit rewrite',
+      rateLimit: {
+        action: RateLimitAction.Rewrite,
+        to: '/rewritten',
+        windowSize: 20,
+        windowLimit: 200,
+        aggregateBy: [RateLimitAggregator.IP, RateLimitAggregator.Domain],
+      },
+    },
+    name: 'func9',
+    source: `
+        export default async () => new Response("Rate my limits")
+
+        export const config = {
+          path: "/rewrite",
+          rateLimit: {
+            action: "rewrite",
+            to: "/rewritten",
+            windowSize: 20,
+            windowLimit: 200,
+            aggregateBy: ["ip", "domain"],
+          },
+          name: 'a limit rewrite',
+        }
+      `,
+  },
 ]
 describe('`getFunctionConfig` extracts configuration properties from function file', () => {
   test.each(functions)('$testName', async (func) => {

packages/edge-bundler/node/config.ts

@@ -10,6 +10,7 @@ import { EdgeFunction } from './edge_function.js'
 import { ImportMap } from './import_map.js'
 import { Logger } from './logger.js'
 import { getPackagePath } from './package_json.js'
+import { RateLimit } from './rate_limit.js'
 
 enum ConfigExitCode {
   Success = 0,
@@ -46,6 +47,7 @@ export interface FunctionConfig {
   name?: string
   generator?: string
   method?: HTTPMethod | HTTPMethod[]
+  rateLimit?: RateLimit
 }
 
 const getConfigExtractor = () => {

packages/edge-bundler/node/manifest.test.ts

@@ -9,6 +9,7 @@ import { BundleError } from './bundle_error.js'
 import { Cache, FunctionConfig } from './config.js'
 import { Declaration } from './declaration.js'
 import { generateManifest } from './manifest.js'
+import { RateLimitAction, RateLimitAggregator } from './rate_limit.js'
 
 test('Generates a manifest with different bundles', () => {
   const bundle1 = {
@@ -486,3 +487,88 @@ test('Returns functions without a declaration and unrouted functions', () => {
   expect(declarationsWithoutFunction).toEqual(['func-3'])
   expect(unroutedFunctions).toEqual(['func-2', 'func-4'])
 })
+
+test('Generates a manifest with rate limit config', () => {
+  const functions = [{ name: 'func-1', path: '/path/to/func-1.ts' }]
+  const declarations: Declaration[] = [{ function: 'func-1', path: '/f1/*' }]
+
+  const userFunctionConfig: Record<string, FunctionConfig> = {
+    'func-1': { rateLimit: { windowLimit: 100, windowSize: 60 } },
+  }
+  const { manifest } = generateManifest({
+    bundles: [],
+    declarations,
+    functions,
+    userFunctionConfig,
+  })
+
+  const expectedRoutes = [{ function: 'func-1', pattern: '^/f1(?:/(.*))/?$', excluded_patterns: [], path: '/f1/*' }]
+  const expectedFunctionConfig = {
+    'func-1': {
+      traffic_rules: {
+        action: {
+          type: 'rate_limit',
+          config: {
+            rate_limit_config: {
+              window_limit: 100,
+              window_size: 60,
+              algorithm: 'sliding_window',
+            },
+            aggregate: {
+              keys: [{ type: 'domain' }],
+            },
+          },
+        },
+      },
+    },
+  }
+  expect(manifest.routes).toEqual(expectedRoutes)
+  expect(manifest.function_config).toEqual(expectedFunctionConfig)
+})
+
+test('Generates a manifest with rewrite config', () => {
+  const functions = [{ name: 'func-1', path: '/path/to/func-1.ts' }]
+  const declarations: Declaration[] = [{ function: 'func-1', path: '/f1/*' }]
+
+  const userFunctionConfig: Record<string, FunctionConfig> = {
+    'func-1': {
+      rateLimit: {
+        action: RateLimitAction.Rewrite,
+        to: '/new_path',
+        windowLimit: 100,
+        windowSize: 60,
+        aggregateBy: [RateLimitAggregator.Domain, RateLimitAggregator.IP],
+      },
+    },
+  }
+  const { manifest } = generateManifest({
+    bundles: [],
+    declarations,
+    functions,
+    userFunctionConfig,
+  })
+
+  const expectedRoutes = [{ function: 'func-1', pattern: '^/f1(?:/(.*))/?$', excluded_patterns: [], path: '/f1/*' }]
+  const expectedFunctionConfig = {
+    'func-1': {
+      traffic_rules: {
+        action: {
+          type: 'rewrite',
+          config: {
+            to: '/new_path',
+            rate_limit_config: {
+              window_limit: 100,
+              window_size: 60,
+              algorithm: 'sliding_window',
+            },
+            aggregate: {
+              keys: [{ type: 'domain' }, { type: 'ip' }],
+            },
+          },
+        },
+      },
+    },
+  }
+  expect(manifest.routes).toEqual(expectedRoutes)
+  expect(manifest.function_config).toEqual(expectedFunctionConfig)
+})

packages/edge-bundler/node/manifest.ts

@@ -9,6 +9,7 @@ import { EdgeFunction } from './edge_function.js'
 import { FeatureFlags } from './feature_flags.js'
 import { Layer } from './layer.js'
 import { getPackageVersion } from './package_json.js'
+import { RateLimit, RateLimitAction, RateLimitAlgorithm, RateLimitAggregator } from './rate_limit.js'
 import { nonNullable } from './utils/non_nullable.js'
 import { ExtendedURLPattern } from './utils/urlpattern.js'
 
@@ -20,12 +21,33 @@ interface Route {
   methods?: string[]
 }
 
+interface TrafficRules {
+  action: {
+    type: string
+    config: {
+      rate_limit_config: {
+        algorithm: string
+        window_size: number
+        window_limit: number
+      }
+      aggregate: {
+        keys: {
+          type: string
+        }[]
+      }
+      to?: string
+    }
+  }
+}
+
 export interface EdgeFunctionConfig {
   excluded_patterns: string[]
   on_error?: string
   generator?: string
   name?: string
+  traffic_rules?: TrafficRules
 }
+
 interface Manifest {
   bundler_version: string
   bundles: { asset: string; format: string }[]
@@ -122,26 +144,35 @@ const generateManifest = ({
   const routedFunctions = new Set<string>()
   const declarationsWithoutFunction = new Set<string>()
 
-  for (const [name, { excludedPath, onError }] of Object.entries(userFunctionConfig)) {
+  for (const [name, { excludedPath, onError, rateLimit }] of Object.entries(userFunctionConfig)) {
     // If the config block is for a function that is not defined, discard it.
     if (manifestFunctionConfig[name] === undefined) {
       continue
     }
 
     addExcludedPatterns(name, manifestFunctionConfig, excludedPath)
 
-    manifestFunctionConfig[name] = { ...manifestFunctionConfig[name], on_error: onError }
+    manifestFunctionConfig[name] = {
+      ...manifestFunctionConfig[name],
+      on_error: onError,
+      traffic_rules: getTrafficRulesConfig(rateLimit),
+    }
   }
 
-  for (const [name, { excludedPath, path, onError, ...rest }] of Object.entries(internalFunctionConfig)) {
+  for (const [name, { excludedPath, path, onError, rateLimit, ...rest }] of Object.entries(internalFunctionConfig)) {
     // If the config block is for a function that is not defined, discard it.
     if (manifestFunctionConfig[name] === undefined) {
       continue
     }
 
     addExcludedPatterns(name, manifestFunctionConfig, excludedPath)
 
-    manifestFunctionConfig[name] = { ...manifestFunctionConfig[name], on_error: onError, ...rest }
+    manifestFunctionConfig[name] = {
+      ...manifestFunctionConfig[name],
+      on_error: onError,
+      traffic_rules: getTrafficRulesConfig(rateLimit),
+      ...rest,
+    }
   }
 
   declarations.forEach((declaration) => {
@@ -202,6 +233,32 @@ const generateManifest = ({
   return { declarationsWithoutFunction: [...declarationsWithoutFunction], manifest, unroutedFunctions }
 }
 
+const getTrafficRulesConfig = (rl: RateLimit | undefined) => {
+  if (rl === undefined) {
+    return
+  }
+
+  const rateLimitAgg = Array.isArray(rl.aggregateBy) ? rl.aggregateBy : [RateLimitAggregator.Domain]
+  const rewriteConfig = 'to' in rl && typeof rl.to === 'string' ? { to: rl.to } : undefined
+
+  return {
+    action: {
+      type: rl.action || RateLimitAction.Limit,
+      config: {
+        ...rewriteConfig,
+        rate_limit_config: {
+          window_limit: rl.windowLimit,
+          window_size: rl.windowSize,
+          algorithm: RateLimitAlgorithm.SlidingWindow,
+        },
+        aggregate: {
+          keys: rateLimitAgg.map((agg) => ({ type: agg })),
+        },
+      },
+    },
+  }
+}
+
 const pathToRegularExpression = (path: string) => {
   if (!path) {
     return null

packages/edge-bundler/node/rate_limit.ts

@@ -0,0 +1,30 @@
+export enum RateLimitAlgorithm {
+  SlidingWindow = 'sliding_window',
+}
+
+export enum RateLimitAggregator {
+  Domain = 'domain',
+  IP = 'ip',
+}
+
+export enum RateLimitAction {
+  Limit = 'rate_limit',
+  Rewrite = 'rewrite',
+}
+
+interface SlidingWindow {
+  windowLimit: number
+  windowSize: number
+}
+
+export type RewriteActionConfig = SlidingWindow & {
+  to: string
+}
+
+interface RateLimitConfig {
+  action?: RateLimitAction
+  aggregateBy?: RateLimitAggregator | RateLimitAggregator[]
+  algorithm?: RateLimitAlgorithm
+}
+
+export type RateLimit = RateLimitConfig & (SlidingWindow | RewriteActionConfig)