feat: add ratelimiting, add leaky bucket ratelimiter plugin

.vscode/settings.json

@@ -7,11 +7,11 @@
     "files.encoding": "utf8",
     "typescript.tsdk": "node_modules/typescript/lib",
     "editor.codeActionsOnSave": {
-        "source.fixAll.eslint": true
+        "source.fixAll.eslint": "explicit"
     },
     "[typescript]": {
         "editor.defaultFormatter": "esbenp.prettier-vscode"
     },
     "typescript.preferences.importModuleSpecifier": "relative",
     "prettier.prettierPath": "./node_modules/prettier"
-}
+}

playground/express/ip-based-mitigation.playground.ts → playground/expressjs.playground.ts

@@ -1,9 +1,10 @@
 import express from 'express';
 
-import { AbuseIPDBPlugin } from '../../src/ip-based-mitigation/plugins/abuse-ipdb/abuse-ipdb.plugin';
-import { VirusTotalIpAddressResponseDataAttributesResultEntryResult } from '../../src/ip-based-mitigation/plugins/virus-total/client/enums/virus-total-ip-address.enum';
-import { VirusTotalPlugin } from '../../src/ip-based-mitigation/plugins/virus-total/virus-total.plugin';
-import umbress from '../../src/main';
+import { AbuseIPDBPlugin } from '../src/ip-based-mitigation/plugins/abuse-ipdb/abuse-ipdb.plugin';
+import { VirusTotalIpAddressResponseDataAttributesResultEntryResult } from '../src/ip-based-mitigation/plugins/virus-total/client/enums/virus-total-ip-address.enum';
+import { VirusTotalPlugin } from '../src/ip-based-mitigation/plugins/virus-total/virus-total.plugin';
+import umbress from '../src/main';
+import { LeakyBucketRatelimiterPlugin } from '../src/ratelimiter/plugins/leaky-bucket/leaky-bucket-ratelimiter.plugin';
 
 const app = express();
 
@@ -36,6 +37,11 @@ app.use(
         },
       }),
     ],
+    ratelimiter: new LeakyBucketRatelimiterPlugin({
+      capacity: 60,
+      rate: 1,
+      action: (request, response) => response.status(429).end(),
+    }),
   }),
 );
 

src/ip-based-mitigation/ip-based-mitigation.service.ts

@@ -35,7 +35,7 @@ export class IpBasedMitigationService<R, S> implements BasePluginService<R, S> {
       }
 
       if (executionStyle === IpBasedMitigationPluginExecutionStyleEnum.ASYNC) {
-        plugins.forEach((plugin) => plugin.shouldBan(ipAddress));
+        this.#asyncronousExecutionFlow(request, response, plugins, ipAddress, store);
       }
     }
   }

src/main.ts

@@ -4,12 +4,15 @@ import { UmbressOptions } from './options/interfaces/options.interface';
 import { optionsServiceInstance } from './options/options.instance';
 import { OptionsService } from './options/options.service';
 import { ProcessorService } from './processor/processor.service';
+import { ratelimiterServiceInstance } from './ratelimiter/ratelimiter.instance';
+import { RatelimiterService } from './ratelimiter/ratelimiter.service';
 
 export default function umbress<R, S>(userOptions: UmbressOptions<R, S>): (request: R, response: S) => void {
   const processor = new ProcessorService<R, S>(
     userOptions,
     optionsServiceInstance as OptionsService<R, S>,
     ipBasedMitigationServiceInstance as IpBasedMitigationService<R, S>,
+    ratelimiterServiceInstance as RatelimiterService<R, S>,
   );
 
   return processor.process;

src/options/interfaces/options.interface.ts

@@ -2,6 +2,7 @@ import type { Cache } from 'cache-manager';
 
 import { IpBasedMitigationPluginExecutionStyleEnum } from '../../ip-based-mitigation/enums/execution-style.enum';
 import { BaseIpBasedMitigationPlugin } from '../../ip-based-mitigation/plugins/base-adapter';
+import { BaseRatelimiterPlugin } from '../../ratelimiter/plugins/base-adapter';
 
 type CachingOptions = {
   caching?: Promise<Cache>;
@@ -11,4 +12,5 @@ export type UmbressOptions<R, S> = CachingOptions & {
   ipAddressExtractor(request: R): string;
   ipBasedMitigation?: BaseIpBasedMitigationPlugin<R, S>[];
   ipBasedMitigationExecutionStyle?: IpBasedMitigationPluginExecutionStyleEnum;
+  ratelimiter?: BaseRatelimiterPlugin<R, S>;
 };

src/options/options.service.ts

@@ -13,8 +13,8 @@ export class OptionsService<R, S> {
     caching: caching('memory'),
   };
 
-  mergeDefaultsAndUserProvidedOptions<T extends Record<string, unknown> = UmbressOptions<R, S>>(userOptions: T): T {
+  mergeDefaultsAndUserProvidedOptions<T extends UmbressOptions<R, S>>(userOptions: T): T {
     // TODO: fix generic types
-    return combine<T, T>(this.#defaultOptions as unknown as T, userOptions);
+    return combine<T, T>(this.#defaultOptions as T, userOptions) as T;
   }
 }

src/processor/processor.service.ts

@@ -1,6 +1,7 @@
 import { IpBasedMitigationService } from '../ip-based-mitigation/ip-based-mitigation.service';
 import { UmbressOptions } from '../options/interfaces/options.interface';
 import { OptionsService } from '../options/options.service';
+import { RatelimiterService } from '../ratelimiter/ratelimiter.service';
 
 export class ProcessorService<R, S> {
   #options: UmbressOptions<R, S>;
@@ -9,14 +10,18 @@ export class ProcessorService<R, S> {
 
   #ipBasedMitigationService: IpBasedMitigationService<R, S>;
 
+  #ratelimiterService: RatelimiterService<R, S>;
+
   constructor(
     userOptions: UmbressOptions<R, S>,
     optionsService: OptionsService<R, S>,
     ipBasedMitigationService: IpBasedMitigationService<R, S>,
+    ratelimiterService: RatelimiterService<R, S>,
   ) {
     this.#optionsService = optionsService;
     this.#options = this.#optionsService.mergeDefaultsAndUserProvidedOptions(userOptions);
     this.#ipBasedMitigationService = ipBasedMitigationService;
+    this.#ratelimiterService = ratelimiterService;
   }
 
   async process(request: R, response: S): Promise<S | void> {
@@ -33,5 +38,9 @@ export class ProcessorService<R, S> {
         store,
       );
     }
+
+    if (this.#options.ratelimiter) {
+      this.#ratelimiterService.execute(request, response, this.#options.ratelimiter, ipAddress, store);
+    }
   }
 }

src/ratelimiter/plugins/base-adapter.ts

@@ -0,0 +1,9 @@
+export abstract class BaseRatelimiterPlugin<R, S> {
+  abstract get name(): string;
+
+  abstract action(request: R, response: S, secondsOverdue: number): S | void;
+
+  abstract get ttl(): number;
+
+  abstract get capacity(): number;
+}

src/ratelimiter/plugins/leaky-bucket/interfaces/leaky-bucket-ratelimiter-options.interface.ts

@@ -0,0 +1,5 @@
+export interface LeakyBucketRatelimiterPluginOptions<R, S> {
+  capacity: number;
+  rate: number;
+  action: (request: R, response: S, secondsOverdue: number) => void;
+}

src/ratelimiter/plugins/leaky-bucket/leaky-bucket-ratelimiter.plugin.ts

@@ -0,0 +1,32 @@
+import { LeakyBucketRatelimiterPluginOptions } from './interfaces/leaky-bucket-ratelimiter-options.interface';
+import { BaseRatelimiterPlugin } from '../base-adapter';
+
+export class LeakyBucketRatelimiterPlugin<R, S> implements BaseRatelimiterPlugin<R, S> {
+  #capacity: number;
+
+  #rate: number;
+
+  #action: typeof this.action;
+
+  constructor(options: LeakyBucketRatelimiterPluginOptions<R, S>) {
+    this.#capacity = options.capacity;
+    this.#rate = options.rate;
+    this.#action = options.action;
+  }
+
+  get name(): string {
+    return this.constructor.name;
+  }
+
+  get ttl(): number {
+    return this.#capacity * 1000;
+  }
+
+  get capacity(): number {
+    return this.#capacity;
+  }
+
+  action(request: R, response: S, secondsOverdue: number): void | S {
+    return this.#action(request, response, secondsOverdue);
+  }
+}

src/ratelimiter/ratelimiter.instance.ts

@@ -0,0 +1,3 @@
+import { RatelimiterService } from './ratelimiter.service';
+
+export const ratelimiterServiceInstance = new RatelimiterService();

src/ratelimiter/ratelimiter.service.ts

@@ -0,0 +1,41 @@
+import { Cache } from 'cache-manager';
+
+import { BaseRatelimiterPlugin } from './plugins/base-adapter';
+import { BasePluginService } from '../base-plugin/base-plugin.service';
+
+export class RatelimiterService<R, S> implements BasePluginService<R, S> {
+  static cacheKeyPrefix = 'leaky-bucket-ratelimiter';
+
+  #buildCacheKey(ipAddress: string): string {
+    return `${RatelimiterService.cacheKeyPrefix}-${ipAddress}`;
+  }
+
+  async execute(
+    request: R,
+    response: S,
+    plugin: BaseRatelimiterPlugin<R, S>,
+    ipAddress: string,
+    store: Cache,
+  ): Promise<void | S> {
+    const now = Date.now();
+    const cacheKey = this.#buildCacheKey(ipAddress);
+    const cached = await store.get<string>(cacheKey);
+
+    if (cached) {
+      const parsed = JSON.parse(cached) as number[];
+      const leaked = parsed.filter((timestamp) => timestamp < now + plugin.ttl);
+
+      if (leaked.length > plugin.capacity) {
+        plugin.action(request, response, leaked.length - plugin.capacity);
+      }
+
+      await this.#save(store, cacheKey, [...leaked, now], plugin.ttl);
+    } else {
+      await this.#save(store, cacheKey, [now], plugin.ttl);
+    }
+  }
+
+  async #save(store: Cache, cacheKey: string, timestamps: number[], ttl: number): Promise<void> {
+    await store.set(cacheKey, JSON.stringify(timestamps), ttl);
+  }
+}