feat: ratelimit config from source

Author: paulo

Diff for: 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 './ratelimit.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/[email protected]',
           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) => {

Diff for: 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 './ratelimit.js'
 
 enum ConfigExitCode {
   Success = 0,
@@ -46,6 +47,7 @@ export interface FunctionConfig {
   name?: string
   generator?: string
   method?: HTTPMethod | HTTPMethod[]
+  ratelimit?: Ratelimit
 }
 
 const getConfigExtractor = () => {

Diff for: 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 './ratelimit.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_config: {
+        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_config: {
+        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)
+})

Diff for: node/manifest.ts

@@ -9,6 +9,13 @@ 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,
+  RewriteActionConfig,
+  RatelimitAction,
+  RatelimitAlgorithm,
+  RatelimitAggregator,
+} from './ratelimit.js'
 import { nonNullable } from './utils/non_nullable.js'
 import { ExtendedURLPattern } from './utils/urlpattern.js'
 
@@ -20,12 +27,33 @@ interface Route {
   methods?: string[]
 }
 
+interface TrafficRulesConfig {
+  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_config?: TrafficRulesConfig
 }
+
 interface Manifest {
   bundler_version: string
   bundles: { asset: string; format: string }[]
@@ -122,26 +150,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_config: 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_config: getTrafficRulesConfig(ratelimit),
+      ...rest,
+    }
   }
 
   declarations.forEach((declaration) => {
@@ -202,6 +239,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 = (rl as RewriteActionConfig).to ? { to: (rl as RewriteActionConfig).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

Diff for: node/ratelimit.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)