feat: ratelimit config from source (#1714)

* feat: extra ratelimit config from source

* feat: refactor to a more general config structure

* chore: address PR comments

* chore: ratelimit -> rate_limit

src/manifest.ts

@@ -6,6 +6,25 @@ import type { InvocationMode } from './function.js'
 import type { FunctionResult } from './utils/format_result.js'
 import type { Route } from './utils/routes.js'
 
+export interface TrafficRules {
+  action: {
+    type: string
+    config: {
+      rateLimitConfig: {
+        algorithm: string
+        windowSize: number
+        windowLimit: number
+      }
+      aggregate: {
+        keys: {
+          type: string
+        }[]
+      }
+      to?: string
+    }
+  }
+}
+
 interface ManifestFunction {
   buildData?: Record<string, unknown>
   invocationMode?: InvocationMode
@@ -20,6 +39,7 @@ interface ManifestFunction {
   bundler?: string
   generator?: string
   priority?: number
+  trafficRules?: TrafficRules
 }
 
 export interface Manifest {
@@ -55,6 +75,7 @@ const formatFunctionForManifest = ({
   name,
   path,
   priority,
+  trafficRules,
   routes,
   runtime,
   runtimeVersion,
@@ -70,6 +91,7 @@ const formatFunctionForManifest = ({
     mainFile,
     name,
     priority,
+    trafficRules,
     runtimeVersion,
     path: resolve(path),
     runtime,

src/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)

src/runtimes/node/in_source_config/index.ts

@@ -1,6 +1,8 @@
 import type { ArgumentPlaceholder, Expression, SpreadElement, JSXNamespacedName } from '@babel/types'
 
 import { InvocationMode, INVOCATION_MODE } from '../../../function.js'
+import { TrafficRules } from '../../../manifest.js'
+import { RateLimitAction, RateLimitAggregator, RateLimitAlgorithm } from '../../../rate_limit.js'
 import { FunctionBundlingUserError } from '../../../utils/error.js'
 import { nonNullable } from '../../../utils/non_nullable.js'
 import { getRoutes, Route } from '../../../utils/routes.js'
@@ -20,6 +22,7 @@ export type ISCValues = {
   routes?: Route[]
   schedule?: string
   methods?: string[]
+  trafficRules?: TrafficRules
 }
 
 export interface StaticAnalysisResult extends ISCValues {
@@ -71,6 +74,60 @@ const normalizeMethods = (input: unknown, name: string): string[] | undefined =>
   })
 }
 
+/**
+ * Extracts the `ratelimit` configuration from the exported config.
+ */
+const getTrafficRulesConfig = (input: unknown, name: string): TrafficRules | undefined => {
+  if (typeof input !== 'object' || input === null) {
+    throw new FunctionBundlingUserError(
+      `Could not parse ratelimit declaration of function '${name}'. Expecting an object, got ${input}`,
+      {
+        functionName: name,
+        runtime: RUNTIME.JAVASCRIPT,
+        bundler: NODE_BUNDLER.ESBUILD,
+      },
+    )
+  }
+
+  const { windowSize, windowLimit, algorithm, aggregateBy, action } = input as Record<string, unknown>
+
+  if (
+    typeof windowSize !== 'number' ||
+    typeof windowLimit !== 'number' ||
+    !Number.isInteger(windowSize) ||
+    !Number.isInteger(windowLimit)
+  ) {
+    throw new FunctionBundlingUserError(
+      `Could not parse ratelimit declaration of function '${name}'. Expecting 'windowSize' and 'limitSize' integer properties, got ${input}`,
+      {
+        functionName: name,
+        runtime: RUNTIME.JAVASCRIPT,
+        bundler: NODE_BUNDLER.ESBUILD,
+      },
+    )
+  }
+
+  const rateLimitAgg = Array.isArray(aggregateBy) ? aggregateBy : [RateLimitAggregator.Domain]
+  const rewriteConfig = 'to' in input && typeof input.to === 'string' ? { to: input.to } : undefined
+
+  return {
+    action: {
+      type: (action as RateLimitAction) || RateLimitAction.Limit,
+      config: {
+        ...rewriteConfig,
+        rateLimitConfig: {
+          windowLimit,
+          windowSize,
+          algorithm: (algorithm as RateLimitAlgorithm) || RateLimitAlgorithm.SlidingWindow,
+        },
+        aggregate: {
+          keys: rateLimitAgg.map((agg) => ({ type: agg })),
+        },
+      },
+    },
+  }
+}
+
 /**
  * Loads a file at a given path, parses it into an AST, and returns a series of
  * data points, such as in-source configuration properties and other metadata.
@@ -131,6 +188,10 @@ export const parseSource = (source: string, { functionName }: FindISCDeclaration
       preferStatic: configExport.preferStatic === true,
     })
 
+    if (configExport.rateLimit !== undefined) {
+      result.trafficRules = getTrafficRulesConfig(configExport.rateLimit, functionName)
+    }
+
     return result
   }
 

src/runtimes/node/index.ts

@@ -141,6 +141,8 @@ const zipFunction: ZipFunction = async function ({
     invocationMode = INVOCATION_MODE.Background
   }
 
+  const { trafficRules } = staticAnalysisResult
+
   const outputModuleFormat =
     extname(finalMainFile) === MODULE_FILE_EXTENSION.MJS ? MODULE_FORMAT.ESM : MODULE_FORMAT.COMMONJS
   const priority = isInternal ? Priority.GeneratedFunction : Priority.UserFunction
@@ -160,6 +162,7 @@ const zipFunction: ZipFunction = async function ({
     nativeNodeModules,
     path: zipPath.path,
     priority,
+    trafficRules,
     runtimeVersion:
       runtimeAPIVersion === 2 ? getNodeRuntimeForV2(config.nodeVersion) : getNodeRuntime(config.nodeVersion),
   }

src/runtimes/runtime.ts

@@ -3,6 +3,7 @@ import type { FunctionConfig } from '../config.js'
 import type { FeatureFlags } from '../feature_flags.js'
 import type { FunctionSource, InvocationMode, SourceFile } from '../function.js'
 import type { ModuleFormat } from '../main.js'
+import { TrafficRules } from '../manifest.js'
 import { ObjectValues } from '../types/utils.js'
 import type { RuntimeCache } from '../utils/cache.js'
 import { Logger } from '../utils/logger.js'
@@ -54,6 +55,7 @@ export interface ZipFunctionResult {
   nativeNodeModules?: object
   path: string
   priority?: number
+  trafficRules?: TrafficRules
   runtimeVersion?: string
   staticAnalysisResult?: StaticAnalysisResult
   entryFilename: string

tests/fixtures/ratelimit/netlify/functions/ratelimit.ts

@@ -0,0 +1,14 @@
+import { Config, Context } from "@netlify/functions";
+
+export default async (req: Request, context: Context) => {
+  return new Response(`Something!`);
+};
+
+export const config: Config = {
+  path: "/ratelimited",
+  rateLimit: {
+    windowLimit: 60,
+    windowSize: 50,
+    aggregateBy: ["ip", "domain"],
+  }
+};

tests/fixtures/ratelimit/netlify/functions/rewrite.ts

@@ -0,0 +1,16 @@
+import { Config, Context } from "@netlify/functions";
+
+export default async (req: Request, context: Context) => {
+  return new Response(`Something!`);
+};
+
+export const config: Config = {
+  path: "/rewrite",
+  rateLimit: {
+    action: "rewrite",
+    to: "/rewritten",
+    windowSize: 20,
+    windowLimit: 200,
+    aggregateBy: ["ip", "domain"],
+  }
+};

tests/main.test.ts

@@ -2898,3 +2898,39 @@ test('Adds a `priority` field to the generated manifest file', async () => {
   const generatedFunction1 = manifest.functions.find((fn) => fn.name === 'function_internal')
   expect(generatedFunction1.priority).toBe(0)
 })
+
+test('Adds a `ratelimit` field to the generated manifest file', async () => {
+  const { path: tmpDir } = await getTmpDir({ prefix: 'zip-it-test' })
+  const fixtureName = 'ratelimit'
+  const manifestPath = join(tmpDir, 'manifest.json')
+  const path = `${fixtureName}/netlify/functions`
+
+  await zipFixture(path, {
+    length: 2,
+    opts: { manifest: manifestPath },
+  })
+
+  const manifest = JSON.parse(await readFile(manifestPath, 'utf-8'))
+
+  expect(manifest.version).toBe(1)
+  expect(manifest.system.arch).toBe(arch)
+  expect(manifest.system.platform).toBe(platform)
+  expect(manifest.timestamp).toBeTypeOf('number')
+
+  const ratelimitFunction = manifest.functions.find((fn) => fn.name === 'ratelimit')
+  const { type: ratelimitType, config: ratelimitConfig } = ratelimitFunction.trafficRules.action
+  expect(ratelimitType).toBe('rate_limit')
+  expect(ratelimitConfig.rateLimitConfig.windowLimit).toBe(60)
+  expect(ratelimitConfig.rateLimitConfig.windowSize).toBe(50)
+  expect(ratelimitConfig.rateLimitConfig.algorithm).toBe('sliding_window')
+  expect(ratelimitConfig.aggregate.keys).toStrictEqual([{ type: 'ip' }, { type: 'domain' }])
+
+  const rewriteFunction = manifest.functions.find((fn) => fn.name === 'rewrite')
+  const { type: rewriteType, config: rewriteConfig } = rewriteFunction.trafficRules.action
+  expect(rewriteType).toBe('rewrite')
+  expect(rewriteConfig.to).toBe('/rewritten')
+  expect(rewriteConfig.rateLimitConfig.windowLimit).toBe(200)
+  expect(rewriteConfig.rateLimitConfig.windowSize).toBe(20)
+  expect(rewriteConfig.rateLimitConfig.algorithm).toBe('sliding_window')
+  expect(rewriteConfig.aggregate.keys).toStrictEqual([{ type: 'ip' }, { type: 'domain' }])
+})