diff --git a/.github/workflows/deno.yml b/.github/workflows/deno.yml deleted file mode 100644 index d53a2b9..0000000 --- a/.github/workflows/deno.yml +++ /dev/null @@ -1,51 +0,0 @@ -# This workflow uses actions that are not certified by GitHub. -# They are provided by a third-party and are governed by -# separate terms of service, privacy policy, and support -# documentation. - -# This workflow will install Deno and run tests across stable and canary builds on Windows, Ubuntu and macOS. -# For more information see: https://github.com/denoland/setup-deno - -name: Deno - -on: - push: - branches: [main] - pull_request: - branches: [main] - -jobs: - test: - runs-on: ${{ matrix.os }} # runs a test on Ubuntu, Windows and macOS - - strategy: - matrix: - os: [macOS-latest, windows-latest, ubuntu-latest] - - steps: - - name: Setup repo - uses: actions/checkout@v2 - - - name: Setup Deno - uses: denoland/setup-deno@v1 - with: - deno-version: v1.x - - - name: Format Files - run: deno fmt - working-directory: src - - - name: Run linter - run: deno lint - working-directory: src - - - # Currently we have no dependencies - # - name: Cache dependencies - # run: deno cache deps.ts - # working-directory: src - - # Currently we have no need for testing - # - name: Run tests - # run: deno test -A --unstable - # working-directory: src diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml new file mode 100644 index 0000000..719511a --- /dev/null +++ b/.github/workflows/test.yml @@ -0,0 +1,29 @@ +name: Test and Build + +on: + push: + branches: [main] + pull_request: + branches: [main] + +jobs: + test: + runs-on: ubuntu-latest + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Setup Deno + uses: denoland/setup-deno@v1 + with: + deno-version: v1.x + + - name: Check formatting + run: deno fmt --check + + - name: Run linter + run: deno lint + + - name: Run Deno tests + run: deno test --allow-net \ No newline at end of file diff --git a/.gitignore b/.gitignore index b9cd518..afb4dd3 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,21 @@ -/.idea/ +# Build output +dist/ +storages.js +storages.d.ts + +# Node.js node_modules/ package-lock.json -*.tsbuildinfo +npm-debug.log* +yarn-debug.log* +yarn-error.log* + +# Deno +deno.lock + +# Editor .vscode/ -node-compatible/out/ + +# OS +.DS_Store +Thumbs.db \ No newline at end of file diff --git a/LICENSE b/LICENSE index dd9284f..95911c2 100644 --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,6 @@ MIT License -Copyright (c) 2022-2024 Amir Zouerami +Copyright (c) 2022-2025 Amir Zouerami Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/RATELIMITER.svg b/RATELIMITER.svg deleted file mode 100644 index b4b2317..0000000 --- a/RATELIMITER.svg +++ /dev/null @@ -1,14 +0,0 @@ - - RATELIMITER-svg - - - - - - - - \ No newline at end of file diff --git a/README.md b/README.md index dd5ecc9..8b13d74 100644 --- a/README.md +++ b/README.md @@ -1,127 +1,27 @@ -[![Test Deno](https://github.com/Amir-Zouerami/ratelimiter/actions/workflows/deno.yml/badge.svg)](https://github.com/Amir-Zouerami/ratelimiter/actions/workflows/deno.yml) +[![Test and Build](https://github.com/Amir-Zouerami/ratelimiter/actions/workflows/test.yml/badge.svg)](https://github.com/Amir-Zouerami/ratelimiter/actions/workflows/test.yml) -# ratelimiter +# Rate Limit Users (`ratelimiter`) + +`ratelimiter` is an advanced and flexible middleware for the grammY framework, designed to protect +Telegram bots from spam and resource abuse.

- - ratelimiter Logo + + grammY rate limiter cover

-## ❓ What does it do? -**** -🔌 [ratelimiter](https://github.com/grammyjs/ratelimiter) is a rate-limiting middleware for Telegram bots made with [grammY](https://grammy.dev/) or [Telegraf](https://github.com/telegraf/telegraf) bot frameworks. It rate limits users and stop them from spamming requests to your bot. You should note that this package **does not** rate limit the incoming requests from telegram servers, instead, it tracks the incoming requests by `from.id` and dismisses them on arrival so no further processing load is added to your servers. - -Under normal circumstances, every request will be processed & answered by your bot which means spamming it will not be that difficult. Adding this middleware to your bot limits the number of requests a specific Telegram user can send during a certain time frame. - -## 🔧 Customizability -This middleware exposes 5 customizable options: -- `timeFrame`: The time frame during which the requests will be monitored (defaults to `1000` ms). -- `limit`: The number of requests allowed within each `timeFrame` (defaults to `1`). -- `storageClient`: The type of storage to use for keeping track of users and their requests. It supports Redis as well. The default value is `MEMORY_STORE` which uses an in-memory Map, but you can also pass in a Redis client from [ioredis](https://github.com/luin/ioredis) or [redis](https://deno.land/x/redis) packages. Other redis drivers might work as well, but I have not tested them. -- `onLimitExceeded`: A function that describes what to do if the user exceeds the limit (ignores the extra requests by default). -- `keyGenerator`: A function that returns a unique key generated for each user (it uses `from.id` by default). This key is used to identify the user, therefore it should be unique and user specific. -- `keyPrefix`: The prefix to be added to your key (returned from `keyGenerator`). Will be "RATE_LIMITER" if no value passed. - -> Note: You must have redis-server **2.6.0** and above on your server to use Redis storage client with ratelimiter. Older versions of Redis are not supported. - -## 💻 Runtime Support -This plugin supports both [grammY](https://grammy.dev/) and [Telegraf](https://telegraf.js.org/) bot frameworks, therefore Deno and Node are both supported. The following examples use [express](https://github.com/expressjs/express) but you can use ratelimiter with any grammy/telegraf supported framework or with no frameworks at all. - -## 💻 How to Use -There are two ways of using ratelimiter: -- Accepting the defaults (Default Configuration). -- Passing a custom object containing your settings (Manual Configuration). - -### ✅ Default Configuration - -The following example uses [express](https://github.com/expressjs/express) as the webserver and [webhooks](https://grammy.dev/guide/deployment-types.html) to rate-limit users. This snippet demonstrates the easiest way of using ratelimiter which is accepting the default behavior: - -``` typescript -import express from "express"; -import { Bot } from "grammy"; -import { limit } from "@grammyjs/ratelimiter" - -const app = express(); -const bot = new Bot("YOUR BOT TOKEN HERE"); - -app.use(express.json()); -bot.use(limit()); - -app.listen(3000, () => { - bot.api.setWebhook("YOUR DOMAIN HERRE", { drop_pending_updates: true }); - console.log('The application is listening on port 3000!'); -}) -``` - -### ✅ Manual Configuration - -As mentioned before, you can pass an `Options` object to the `limit()` function to alter ratelimiter's behaviors. In the following snippet, I have decided to use Redis as my storage option: - -``` typescript -import express from "express"; -import { Bot } from "grammy"; -import { limit } from "@grammyjs/ratelimiter" -import Redis from "ioredis"; +At its core, `ratelimiter` acts as a configurable gatekeeper for incoming updates. It allows +developers to define precise rules for how many messages a user or chat (or any arbitrary entity) +can send in a given period, ensuring the bot remains responsive and server resources are protected +from overload. +The plugin inspects each incoming message, identifies its source, and decides if it should be +processed or dismissed based on the rules you set. -const app = express(); -const bot = new Bot("YOUR BOT TOKEN HERE"); -const redis = new Redis(); - - -app.use(express.json()); -bot.use(limit({ - timeFrame: 2000, - - limit: 3, - - // "MEMORY_STORAGE" is the default mode. Therefore if you want to use Redis, do not pass storageClient at all. - storageClient: redis, - - onLimitExceeded: ctx => { ctx?.reply("Please refrain from sending too many requests!") }, - - // Note that the key should be a number in string format such as "123456789" - keyGenerator: ctx => { return ctx.from?.id.toString() } -})); - -app.listen(3000, () => { - bot.api.setWebhook("YOUR DOMAIN HERRE", { drop_pending_updates: true }); - console.log('The application is listening on port 3000!'); -}) -``` -As you can see in the above example, each user is allowed to send 3 requests every 2 seconds. If said user sends more requests, the bot replies with _Please refrain from sending too many requests_. That request will not travel further and dies immediately as we do not call `next()`. - -> Note: To avoid flooding Telegram servers, `onLimitExceeded` is only executed once in every `timeFrame`. - -Another use case would be limiting the incoming requests from a chat instead of a specific user: -``` typescript -import express from "express"; -import { Bot } from "grammy"; -import { limit } from "@grammyjs/ratelimiter" - -const app = express(); -const bot = new Bot("YOUR BOT TOKEN HERE"); - -app.use(express.json()); -bot.use(limit({ - keyGenerator: (ctx) => { - if (ctx.chat?.type === "group" || ctx.chat?.type === "supergroup") { - // Note that the key should be a number in string format such as "123456789" - return ctx.chat.id.toString(); - } - } -})); - -app.listen(3000, () => { - bot.api.setWebhook("YOUR DOMAIN HERRE", { drop_pending_updates: true }); - console.log('The application is listening on port 3000!'); -}) -``` -In this example, I have used `chat.id` as the unique key for rate-limiting. - -## Acknowledgements -This package was heavily inspired by [telegraf-ratelimit](https://github.com/telegraf/telegraf-ratelimit). +> **For more information and how-to instructions, please visit** +> [**the official grammY ratelimiter documentation.**](https://grammy.dev/plugins/ratelimiter) ## License + Distributed under the MIT License. See `LICENSE` for more information. diff --git a/deno.jsonc b/deno.jsonc new file mode 100644 index 0000000..c02902b --- /dev/null +++ b/deno.jsonc @@ -0,0 +1,27 @@ +{ + "name": "@grammyjs/ratelimiter", + "version": "2.0.0", + "exports": "./mod.ts", + "tasks": { + "check": "deno check **/*.ts", + "test": "deno test --allow-all", + "dev": "deno run --watch --allow-all examples/basic.ts" + }, + "fmt": { + "options": { + "useTabs": true, + "lineWidth": 100, + "indentWidth": 4, + "singleQuote": true + } + }, + "compilerOptions": { + "strict": true, + "lib": ["deno.ns", "deno.window"] + }, + "imports": { + "grammy": "npm:grammy@^1.24.0", + "grammy/types": "npm:grammy@^1.24.0/types", + "@std/assert": "https://deno.land/std@0.224.0/assert/mod.ts" + } +} diff --git a/grammy-ratelimiter-cover.png b/grammy-ratelimiter-cover.png new file mode 100644 index 0000000..6b3d96f Binary files /dev/null and b/grammy-ratelimiter-cover.png differ diff --git a/mod.ts b/mod.ts new file mode 100644 index 0000000..4603451 --- /dev/null +++ b/mod.ts @@ -0,0 +1,32 @@ +/** + * # Rate Limiter for grammY + * + * This is the main entry point for the grammY rate-limiter middleware. + * + * @module + */ + +export { limit } from './src/core/middleware.ts'; +export { Limiter } from './src/core/builder.ts'; + +// available strategies. +export { + FixedWindowStrategy, + type FixedWindowStrategyOptions, +} from './src/strategies/fixed_window.ts'; +export { + TokenBucketStrategy, + type TokenBucketStrategyOptions, +} from './src/strategies/token_bucket.ts'; + +// all core types and interfaces that developers might need for type annotations. +export type { + GrammyContext, + ILimiterStrategy, + IStorageEngine, + KeyGenerator, + LimitResult, + NextFunction, + OnLimitExceeded, + PenaltyDurationGenerator, +} from './src/types.ts'; diff --git a/node-compatible/.gitignore b/node-compatible/.gitignore deleted file mode 100644 index 1f72880..0000000 --- a/node-compatible/.gitignore +++ /dev/null @@ -1,6 +0,0 @@ -node_modules/ -package-lock.json -*.tsbuildinfo -.vscode/ -./src -*.tgz \ No newline at end of file diff --git a/node-compatible/README.md b/node-compatible/README.md deleted file mode 100644 index de23b02..0000000 --- a/node-compatible/README.md +++ /dev/null @@ -1,3 +0,0 @@ -# Rate limiting for grammY - -Please visit [the main README file](https://github.com/grammyjs/rateLimiter) of the repository to learn more about this [grammY](https://grammy.dev) plugin. diff --git a/node-compatible/package.json b/node-compatible/package.json deleted file mode 100644 index 489030a..0000000 --- a/node-compatible/package.json +++ /dev/null @@ -1,38 +0,0 @@ -{ - "name": "@grammyjs/ratelimiter", - "version": "1.2.0", - "description": "This is a plugin for grammY and Telegraf Telegram bot frameworks to rate limit users and deflect heavy spamming in your bots.", - "scripts": { - "prepare": "deno2node tsconfig.json" - }, - "repository": { - "type": "git", - "url": "git+https://github.com/grammyjs/rateLimiter.git" - }, - "main": "./out/mod.js", - "files": [ - "out/" - ], - "keywords": [ - "grammyY", - "grammyY plugin", - "Telegraf", - "Telegraf plugin", - "bot", - "rateLimiter", - "Telegram", - "Telegram bot", - "bot framework", - "framework" - ], - "author": "Amir Zouerami", - "license": "MIT", - "bugs": { - "url": "https://github.com/grammyjs/rateLimiter/issues" - }, - "homepage": "https://github.com/grammyjs/rateLimiter#readme", - "devDependencies": { - "@types/node": "^12.20.55", - "deno2node": "~1.3.0" - } -} diff --git a/node-compatible/tsconfig.json b/node-compatible/tsconfig.json deleted file mode 100644 index 614e815..0000000 --- a/node-compatible/tsconfig.json +++ /dev/null @@ -1,18 +0,0 @@ -{ - "compilerOptions": { - "forceConsistentCasingInFileNames": true, - "newLine": "lf", - "noFallthroughCasesInSwitch": true, - "noImplicitReturns": true, - "noUncheckedIndexedAccess": true, - "noUnusedParameters": false, - "rootDir": "../src", - "strict": true, - "declaration": true, - "module": "CommonJS", - "outDir": "./out", - "skipLibCheck": true, - "target": "es2019" - }, - "include": ["../src"] -} diff --git a/package.json b/package.json new file mode 100644 index 0000000..62bd48a --- /dev/null +++ b/package.json @@ -0,0 +1,56 @@ +{ + "name": "@grammyjs/ratelimiter", + "version": "2.0.0", + "description": "The most advanced and flexible rate-limiter for grammY.", + "type": "module", + "scripts": { + "build": "deno2node --project tsconfig.json", + "test:deno": "deno test --allow-net", + "test:node": "node --test dist/tests/", + "test": "npm run test:deno && npm run build && npm run test:node" + }, + "repository": { + "type": "git", + "url": "git+https://github.com/grammyjs/rateLimiter.git" + }, + "files": [ + "dist/" + ], + "main": "./dist/mod.js", + "module": "./dist/mod.js", + "types": "./dist/mod.d.ts", + "exports": { + ".": { + "types": "./dist/mod.d.ts", + "import": "./dist/mod.js", + "require": "./dist/mod.js" + }, + "./storages": { + "types": "./dist/storages.d.ts", + "import": "./dist/storages.js", + "require": "./dist/storages.js" + } + }, + "keywords": [ + "grammy", + "grammy-plugin", + "telegraf", + "bot", + "ratelimiter", + "rate-limit", + "telegram", + "telegram-bot", + "middleware", + "spam-protection" + ], + "author": "Amir Zouerami", + "license": "MIT", + "bugs": { + "url": "https://github.com/grammyjs/rateLimiter/issues" + }, + "homepage": "https://github.com/grammyjs/rateLimiter#readme", + "devDependencies": { + "@types/node": "^20.11.20", + "deno2node": "^1.5.0" + } +} diff --git a/src/README.md b/src/README.md deleted file mode 100644 index de23b02..0000000 --- a/src/README.md +++ /dev/null @@ -1,3 +0,0 @@ -# Rate limiting for grammY - -Please visit [the main README file](https://github.com/grammyjs/rateLimiter) of the repository to learn more about this [grammY](https://grammy.dev) plugin. diff --git a/src/core/builder.ts b/src/core/builder.ts new file mode 100644 index 0000000..e441e4f --- /dev/null +++ b/src/core/builder.ts @@ -0,0 +1,204 @@ +import { Rule } from './rule.ts'; +import { EventEmitter } from '../lib/event_emitter.ts'; +import { FixedWindowStrategy } from '../strategies/fixed_window.ts'; +import { TokenBucketStrategy } from '../strategies/token_bucket.ts'; +import type { + DynamicLimitGenerator, + GrammyContext, + ILimiterStrategy, + IStorageEngine, + KeyGenerator, + LimiterEvents, + LimitResult, + Mutable, + OnLimitExceeded, + PenaltyDurationGenerator, +} from '../types.ts'; + +/** + * A fluent API builder for constructing rate-limiter rules. + */ +export class Limiter { + private config: Partial>> = {}; + private readonly events = new EventEmitter>(); + + constructor() { + this.config.events = this.events; + } + + /** + * Sets the limiting strategy to Fixed Window. + * @param options The configuration for the fixed window. + */ + public fixedWindow( + options: { limit: number | DynamicLimitGenerator; timeFrame: number }, + ): this { + this.config.strategy = new FixedWindowStrategy({ + limit: typeof options.limit === 'number' ? options.limit : 1, + timeFrame: options.timeFrame, + }); + + if (typeof options.limit !== 'number') { + this.config.dynamicLimitGenerator = options.limit; + } + + return this; + } + + /** + * Sets the limiting strategy to Token Bucket. + * @param options The configuration for the token bucket. + */ + public tokenBucket( + options: { bucketSize: number; tokensPerInterval: number; interval: number }, + ): this { + this.config.strategy = new TokenBucketStrategy(options); + return this; + } + + /** + * Sets a custom limiting strategy. + * @param strategy An instance of a class that implements `ILimiterStrategy`. + */ + public customStrategy(strategy: ILimiterStrategy): this { + this.config.strategy = strategy; + return this; + } + + /** + * Sets the storage engine for the rule. + * + * @param storage An instance of a storage engine. + */ + public useStorage(storage: IStorageEngine): this { + this.config.storage = storage; + return this; + } + + /** + * Defines what entity to rate-limit. + * + * @param scope A predefined scope ("user", "chat", "global") or a custom key generator function. + */ + public limitFor(scope: 'user' | 'chat' | 'global' | KeyGenerator): this { + if (typeof scope === 'function') { + this.config.keyGenerator = scope; + } else { + switch (scope) { + case 'user': + this.config.keyGenerator = (ctx: C) => ctx.from?.id.toString(); + break; + case 'chat': + this.config.keyGenerator = (ctx: C) => ctx.chat?.id.toString(); + break; + case 'global': + this.config.keyGenerator = () => '___GLOBAL___'; + break; + } + } + + return this; + } + + /** + * Registers an event listener for this limiter instance. + * This allows for observing the limiter's behavior for logging or analytics. + * + * @param eventName The event to listen for. + * @param listener The callback function. + */ + public on>( + eventName: E, + listener: (...args: LimiterEvents[E]) => void, + ): this { + this.events.on(eventName, listener); + return this; + } + + /** + * Unregisters an event listener. + * @param eventName The event to stop listening for. + * @param listener The callback function to remove. + */ + public off>( + eventName: E, + listener: (...args: LimiterEvents[E]) => void, + ): this { + this.events.off(eventName, listener); + return this; + } + + /** + * Sets a prefix for all keys generated for this rule. + * Useful for preventing key collisions between different rules in the same storage. + * + * @param prefix The string to prepend to the key. + */ + public withKeyPrefix(prefix: string): this { + this.config.keyPrefix = prefix; + return this; + } + + /** + * Rate limiter will only run if this function returns `true`. + * + * @param predicate A function that takes the context and returns a boolean. + */ + public onlyIf(predicate: (ctx: C) => boolean | Promise): this { + this.config.filter = predicate; + return this; + } + + /** + * Defines the callback to execute when a request is throttled (rate-limited). + * + * @param handler The function to call when the limit is exceeded. + */ + public onThrottled(handler: OnLimitExceeded): this { + this.config.onLimitExceeded = handler; + return this; + } + + /** + * Enables and configures the "Penalty Box" feature. + * + * @param options The penalty configuration. + */ + public withPenalty( + options: { penaltyTime: number | PenaltyDurationGenerator; penaltyKeyPrefix?: string }, + ): this { + let generator: PenaltyDurationGenerator; + + if (typeof options.penaltyTime === 'number') { + const penaltyValue = options.penaltyTime; + generator = (_ctx: C, _info: LimitResult) => penaltyValue; + } else { + generator = options.penaltyTime; + } + + this.config.penalty = { + generator: generator, + keyPrefix: options.penaltyKeyPrefix ?? 'GRAMMY:RATELIMITER:PENALTY', + }; + + return this; + } + + /** + * Finalizes the configuration and returns a validated `Rule` instance. + */ + public build(): Rule { + if (!this.config.keyPrefix) { + console.warn( + ` +[grammy-ratelimiter] WARNING: No .withKeyPrefix() was set for this limiter. +Using the default prefix is not recommended when using multiple limiters, as it can lead to data collisions. +Please assign a unique prefix for each rule, e.g., .withKeyPrefix('my-rule'). + +`, + ); + } + + return new Rule(this.config); + } +} diff --git a/src/core/middleware.ts b/src/core/middleware.ts new file mode 100644 index 0000000..2fda831 --- /dev/null +++ b/src/core/middleware.ts @@ -0,0 +1,84 @@ +import { FixedWindowStrategy } from '../strategies/fixed_window.ts'; +import type { GrammyContext, NextFunction } from '../types.ts'; +import { Limiter } from './builder.ts'; +import type { Rule } from './rule.ts'; + +/** + * The main middleware function generator. + * + * It accepts either a pre-built `Rule` object or a `Limiter` builder instance. + * + * @param ruleOrBuilder An instance of a `Rule` or a `Limiter` builder. + * @returns A grammY-compatible middleware function. + */ +export function limit( + ruleOrBuilder: Rule | Limiter, +): (ctx: C, next: NextFunction) => Promise { + const rule = ruleOrBuilder instanceof Limiter ? ruleOrBuilder.build() : ruleOrBuilder; + + return async (ctx: C, next: NextFunction): Promise => { + if (rule.penalty) { + const baseKey = rule.keyGenerator(ctx); + + if (baseKey) { + const penaltyKey = `${rule.penalty.keyPrefix}:${baseKey}`; + const isPenalized = await rule.storage.checkPenalty(penaltyKey); + + if (isPenalized) { + return; + } + } + } + + const applies = await rule.filter(ctx); + if (!applies) { + return await next(); + } + + const entityKey = rule.keyGenerator(ctx); + if (entityKey === undefined) { + return await next(); + } + + const storageKey = `${rule.keyPrefix}:${entityKey}`; + + let strategy = rule.strategy; + if (rule.dynamicLimitGenerator && rule.strategy instanceof FixedWindowStrategy) { + const newLimit = rule.dynamicLimitGenerator(ctx); + + strategy = new FixedWindowStrategy({ + ...rule.strategy.options, + limit: newLimit, + }); + } + + const result = await strategy.check(storageKey, rule.storage); + + if (result.isAllowed) { + if (rule.events.hasListeners('allowed')) { + rule.events.emit('allowed', ctx, result); + } + + await next(); + } else { + if (rule.events.hasListeners('throttled')) { + rule.events.emit('throttled', ctx, result); + } + + await rule.onLimitExceeded(ctx, result, rule.storage); + + if (rule.penalty) { + const penaltyDuration = rule.penalty.generator(ctx, result); + + if (penaltyDuration > 0) { + const penaltyKey = `${rule.penalty.keyPrefix}:${entityKey}`; + await rule.storage.setPenalty(penaltyKey, penaltyDuration); + + if (rule.events.hasListeners('penaltyApplied')) { + rule.events.emit('penaltyApplied', ctx, entityKey, penaltyDuration); + } + } + } + } + }; +} diff --git a/src/core/rule.ts b/src/core/rule.ts new file mode 100644 index 0000000..ba7a1c7 --- /dev/null +++ b/src/core/rule.ts @@ -0,0 +1,70 @@ +import type { EventEmitter } from '../lib/event_emitter.ts'; +import type { + DynamicLimitGenerator, + GrammyContext, + ILimiterStrategy, + IStorageEngine, + KeyGenerator, + LimiterEvents, + OnLimitExceeded, + PenaltyDurationGenerator, +} from '../types.ts'; + +/** + * Represents a ready-to-use rate-limiting rule. + * + * **NOTE**: The properties are read-only so you cannot modify them after construction. + */ +export class Rule { + public readonly strategy: ILimiterStrategy; + public readonly storage: IStorageEngine; + public readonly keyGenerator: KeyGenerator; + public readonly events: EventEmitter>; + public readonly keyPrefix: string; + public readonly onLimitExceeded: OnLimitExceeded; + + // determines if the limiter should run for a given request. + public readonly filter: (ctx: C) => boolean | Promise; + + // dynamically determines the limit for Fixed Window strategies. + public readonly dynamicLimitGenerator?: DynamicLimitGenerator; + + public readonly penalty?: { + generator: PenaltyDurationGenerator; + keyPrefix: string; + }; + + constructor(config: Partial>) { + if (!config.strategy) { + throw new Error( + 'Cannot build rule: A limiting strategy must be defined. Use .fixedWindow(), .tokenBucket(), or .customStrategy() on the builder.', + ); + } + + if (!config.storage) { + throw new Error( + 'Cannot build rule: A storage engine must be provided. Use .useStorage() on the builder. It is recommended to create one store instance and share it across all rules.', + ); + } + + if (!config.keyGenerator) { + throw new Error( + 'Cannot build rule: A key generation strategy must be defined. Use .limitFor() on the builder.', + ); + } + + if (!config.events) { + throw new Error('[INTERNAL] Cannot build rule: An event emitter instance is missing.'); + } + + this.strategy = config.strategy; + this.storage = config.storage; + this.keyGenerator = config.keyGenerator; + this.events = config.events; + this.keyPrefix = config.keyPrefix ?? 'GRAMMY:RATELIMITER'; + this.filter = config.filter ?? (() => true); + this.onLimitExceeded = config.onLimitExceeded ?? (() => {}); + this.dynamicLimitGenerator = config.dynamicLimitGenerator; + this.penalty = config.penalty; + } +} diff --git a/src/lib/event_emitter.ts b/src/lib/event_emitter.ts new file mode 100644 index 0000000..2bc2691 --- /dev/null +++ b/src/lib/event_emitter.ts @@ -0,0 +1,72 @@ +import type { EventMap } from '../types.ts'; + +/** + * A type alias for a generic listener function. + */ +// deno-lint-ignore no-explicit-any +type Listener = (...args: any[]) => void; + +/** + * A simple event emitter. + */ +export class EventEmitter { + private readonly listeners = new Map>(); + + /** + * Registers a listener for a given event. + * + * @param eventName The name of the event to listen to. + * @param listener The function to call when the event is emitted. + */ + public on(eventName: E, listener: (...args: T[E]) => void): this { + const eventListeners = this.listeners.get(eventName) ?? new Set(); + + eventListeners.add(listener); + this.listeners.set(eventName, eventListeners); + + return this; + } + + /** + * Unregisters a listener for a given event. + * + * @param eventName The name of the event. + * @param listener The listener function to remove. + */ + public off(eventName: E, listener: (...args: T[E]) => void): this { + const eventListeners = this.listeners.get(eventName); + + if (eventListeners) { + eventListeners.delete(listener); + } + + return this; + } + + /** + * Emits an event. + * + * @param eventName The name of the event to emit. + * @param args The arguments to pass to the listeners. + */ + public emit(eventName: E, ...args: T[E]): void { + const eventListeners = this.listeners.get(eventName); + + if (eventListeners) { + for (const listener of eventListeners) { + listener(...args); + } + } + } + + /** + * Checks if there are any listeners for a specific event. + * + * @param eventName The name of the event. + * @returns True if at least one listener is registered. + */ + public hasListeners(eventName: E): boolean { + const eventListeners = this.listeners.get(eventName); + return eventListeners ? eventListeners.size > 0 : false; + } +} diff --git a/src/memoryStore.ts b/src/memoryStore.ts deleted file mode 100644 index f3dcad7..0000000 --- a/src/memoryStore.ts +++ /dev/null @@ -1,18 +0,0 @@ -import { unref } from "./platform.deno.ts"; - -export class MemoryStore { - hits = new Map(); - - constructor(timeFrame: number) { - unref(setInterval(() => { - this.hits.clear(); - }, timeFrame)); - } - - increment(key: string): number { - let counter = this.hits.get(key) ?? 0; - counter++; - this.hits.set(key, counter); - return counter; - } -} diff --git a/src/mod.ts b/src/mod.ts deleted file mode 100644 index 33cdc3d..0000000 --- a/src/mod.ts +++ /dev/null @@ -1 +0,0 @@ -export * from "./rateLimiter.ts"; diff --git a/src/platform.deno.ts b/src/platform.deno.ts index 70d7102..a668a42 100644 --- a/src/platform.deno.ts +++ b/src/platform.deno.ts @@ -1 +1 @@ -export const unref = Deno.unrefTimer \ No newline at end of file +export const unref = Deno.unrefTimer; diff --git a/src/platform.node.ts b/src/platform.node.ts index 6e3db6e..8892ffe 100644 --- a/src/platform.node.ts +++ b/src/platform.node.ts @@ -1,3 +1,7 @@ -export const unref = (interval: ReturnType) => { - interval.unref(); +interface UnrefableTimer { + unref(): void; +} + +export const unref = (timer: number | object): void => { + (timer as UnrefableTimer).unref(); }; diff --git a/src/rateLimiter.ts b/src/rateLimiter.ts deleted file mode 100644 index afa8069..0000000 --- a/src/rateLimiter.ts +++ /dev/null @@ -1,57 +0,0 @@ -import { MemoryStore } from "./memoryStore.ts"; -import { - Context, - defaultOptions, - NextFunction, - OptionsInterface, - RedisType, -} from "./typesAndDefaults.ts"; -import { RedisStore } from "./redisStore.ts"; - -/** - * - * @param UserOptions an object of rateLimiter options: - * ```ts - * { - * timeFrame: 1000, - * limit: 1, - * onLimitExceeded: (ctx, next) => {}, - * storageClient: "MEMORY_STORE", - * keyGenerator: (ctx) => ctx.from && ctx.from.id.toString(), - * } - * ``` - * - * as explained in [customizability](https://github.com/Amir-Zouerami/rateLimiter#-customizability) - * @description A middleware function generator - * @returns a middleware function to be passed to `bot.use()` - */ - -export const limit = ( - userOptions?: OptionsInterface, -) => { - const options = { ...defaultOptions, ...userOptions }; - const store = options.storageClient === "MEMORY_STORE" - ? new MemoryStore(options.timeFrame) - : new RedisStore(options.storageClient as RT, options.timeFrame); - - const keyPrefix = userOptions?.keyPrefix ?? defaultOptions.keyPrefix; - - const middlewareFunc = async (ctx: C, next: NextFunction) => { - const key = options.keyGenerator(ctx); - if (!key) { - return await next(); - } - - const hits = await store.increment(keyPrefix + key); - - if (hits === options.limit + 1 || (options.alwaysReply && hits > options.limit)) { - return options.onLimitExceeded(ctx, next); - } - - if (hits <= options.limit) { - return await next(); - } - }; - - return middlewareFunc; -}; diff --git a/src/redisStore.ts b/src/redisStore.ts deleted file mode 100644 index da4daeb..0000000 --- a/src/redisStore.ts +++ /dev/null @@ -1,20 +0,0 @@ -import { RedisType } from "./typesAndDefaults.ts"; - -export class RedisStore { - private client: RedisType; - timeFrame: number; - - constructor(client: RedisType, timeFrame: number) { - this.client = client; - this.timeFrame = timeFrame; - } - - async increment(key: string): Promise { - const counter = await this.client.incr(key); - - if (counter === 1) { - await this.client.pexpire(key, this.timeFrame); - } - return counter; - } -} diff --git a/src/stores/memory.ts b/src/stores/memory.ts new file mode 100644 index 0000000..fa57dc3 --- /dev/null +++ b/src/stores/memory.ts @@ -0,0 +1,158 @@ +import type { IStorageEngine, TokenBucketState } from '../types.ts'; +import { unref } from '../platform.deno.ts'; + +/** + * An in-memory storage record that bundles the state with its expiry timestamp. + * @template T The type of the state being stored. + */ +interface MemoryRecord { + state: T; + expiresAt: number; +} + +/** + * An in-memory storage engine that uses a JavaScript `Map`. + */ +export class MemoryStore implements IStorageEngine { + private readonly store = new Map>(); + // @ts-ignore NodeJS runtime is not available in deno. However deno2node will use this for node.js builds. + private readonly cleanupIntervalId?: number | NodeJS.Timeout; + + /** + * Constructs a new `MemoryStore`. + * **WARNING**: DO NOT PASS `null` to the constructor. It is used for testing purposes only. + * + * @param cleanupIntervalMs The interval (in milliseconds) at which to sweep + * for and remove expired keys. Pass `null` to disable. Defaults to 30 seconds. + */ + constructor(cleanupIntervalMs: number | null = 30 * 1000) { + if (cleanupIntervalMs && cleanupIntervalMs > 0) { + this.cleanupIntervalId = setInterval(() => this.sweep(), cleanupIntervalMs); + unref(this.cleanupIntervalId); + } + } + + // ... rest of the file is unchanged ... + /** + * The internal garbage collector. It iterates through all stored records + * and removes any that have passed their expiration time. + */ + private sweep(): void { + const now = Date.now(); + + for (const [key, record] of this.store.entries()) { + if (record.expiresAt <= now) { + this.store.delete(key); + } + } + } + + /** + * A private helper to retrieve a record from the store. + * + * It includes a proactive check to delete the record if it has expired + * at the moment of retrieval. + * + * @template T The expected type of the state within the record. + * @param key The key of the record to retrieve. + * @returns The `MemoryRecord` if it exists and has not expired, otherwise `undefined`. + */ + private getRecord(key: string): MemoryRecord | undefined { + const record = this.store.get(key) as MemoryRecord | undefined; + if (!record) { + return undefined; + } + + if (record.expiresAt <= Date.now()) { + this.store.delete(key); + return undefined; + } + return record; + } + + /** + * Retrieves the state for a key, used by the Token Bucket strategy. + * + * @param key The unique key for the record. + * @returns A promise that resolves with the `TokenBucketState` or `undefined`. + */ + public get(key: string): Promise { + const record = this.getRecord(key); + return Promise.resolve(record?.state); + } + + /** + * Sets the state for a key, used by the Token Bucket strategy. + * + * @param key The unique key for the record. + * @param state The `TokenBucketState` to store. + * @param ttl The time-to-live for the record in milliseconds. + */ + public set(key: string, state: TokenBucketState, ttl: number): Promise { + const record: MemoryRecord = { + state, + expiresAt: Date.now() + ttl, + }; + + this.store.set(key, record); + return Promise.resolve(); + } + + /** + * Deletes a record from the store. + * + * @param key The unique key for the record to delete. + */ + public delete(key: string): Promise { + this.store.delete(key); + return Promise.resolve(); + } + + /** + * Atomically increments the hit count for a key, used by the Fixed Window strategy. + * + * @param key The unique key for the record. + * @param ttl The time-to-live for the record in milliseconds. + * @returns A promise that resolves with the new hit count. + */ + public increment(key: string, ttl: number): Promise { + let record = this.getRecord(key); + + if (!record) { + record = { state: 0, expiresAt: 0 }; + } + + record.state += 1; + record.expiresAt = Date.now() + ttl; + + this.store.set(key, record); + return Promise.resolve(record.state); + } + + /** + * Sets a penalty for a key, used by the Penalty Box feature. + * + * @param key The unique key for the penalty. + * @param ttl The duration of the penalty in milliseconds. + */ + public setPenalty(key: string, ttl: number): Promise { + const record: MemoryRecord = { + state: true, + expiresAt: Date.now() + ttl, + }; + + this.store.set(key, record); + return Promise.resolve(); + } + + /** + * Checks if a penalty exists for a key. + * + * @param key The unique key for the penalty. + * @returns A promise that resolves to `true` if a penalty is active, otherwise `false`. + */ + public checkPenalty(key: string): Promise { + const record = this.getRecord(key); + return Promise.resolve(record !== undefined); + } +} diff --git a/src/stores/redis.ts b/src/stores/redis.ts new file mode 100644 index 0000000..751e0e6 --- /dev/null +++ b/src/stores/redis.ts @@ -0,0 +1,132 @@ +import type { IStorageEngine, TokenBucketState } from '../types.ts'; + +/** + * Defines the contract for a consumer-provided object that has the methods + * needed by `RedisStore` to communicate with Redis. + */ +export interface IRedisClient { + /** + * Loads a Lua script into the Redis script cache and returns its SHA1 hash. + */ + scriptLoad(script: string): Promise; + + /** + * Executes a pre-loaded Lua script by its SHA1 hash. + * + * @param sha The SHA1 hash of the loaded script + * @param keys Array of Redis keys (becomes KEYS[] in Lua) + * @param args Array of arguments (becomes ARGV[] in Lua) - should be flat array like [ttl] + */ + evalsha(sha: string, keys: string[], args: (string | number)[]): Promise; + + /** + * Retrieves a value for a key. + */ + get(key: string): Promise; + + /** + * Sets a key with a value and millisecond expiry. + * + * Implementation examples: + * - ioredis: `set(key, value, 'PX', ttlMilliseconds)` + * - node-redis: `pSetEx(key, ttlMilliseconds, value)` + * - deno-redis: `set(key, value, { px: ttlMilliseconds })` + */ + setWithExpiry(key: string, value: string, ttlMilliseconds: number): Promise; + + /** + * Checks for the existence of a key. + * Should return 1 if key exists, 0 if it doesn't. + */ + exists(key: string): Promise; + + /** + * Deletes a key. + */ + del(key: string): Promise; +} + +const LUA_SCRIPT_ATOMIC_INCREMENT = ` + local current = redis.call('INCR', KEYS[1]) + if current == 1 then + redis.call('PEXPIRE', KEYS[1], ARGV[1]) + end + return current +`; + +/** + * Tries to determine if a Redis client error is a "NOSCRIPT" error. + * + * @param error The error thrown by the Redis client. + */ +function isNoscriptError(error: unknown): boolean { + if (error instanceof Error && typeof error.message === 'string') { + return error.message.includes('NOSCRIPT'); + } + return false; +} + +/** + * A storage engine that uses Redis as the backend. + */ +export class RedisStore implements IStorageEngine { + private readonly client: IRedisClient; + private scriptSha: string | null = null; + + /** + * Constructs a new `RedisStore`. + * + * @param client An object that conforms to the `IRedisClient` interface. + */ + constructor(client: IRedisClient) { + this.client = client; + } + + public async get(key: string): Promise { + const state = await this.client.get(key); + return state ? JSON.parse(state) : undefined; + } + + public async set(key: string, state: TokenBucketState, ttl: number): Promise { + await this.client.setWithExpiry(key, JSON.stringify(state), ttl); + } + + public async delete(key: string): Promise { + await this.client.del(key); + } + + public async checkPenalty(key: string): Promise { + const exists = await this.client.exists(key); + return exists === 1; + } + + public async setPenalty(key: string, ttl: number): Promise { + await this.client.setWithExpiry(key, '1', ttl); + } + + /** + * Atomically increments a key using the Lua script above. + * + * If the Redis script cache is ever flushed (e.g., on a server restart), + * this method will try re-caching the script and retrying the `EVALSHA` command. + */ + public async increment(key: string, ttl: number): Promise { + if (!this.scriptSha) { + this.scriptSha = await this.client.scriptLoad(LUA_SCRIPT_ATOMIC_INCREMENT); + } + + try { + const result = await this.client.evalsha(this.scriptSha, [key], [ttl]); + return result as number; + } catch (error) { + if (isNoscriptError(error)) { + this.scriptSha = await this.client.scriptLoad(LUA_SCRIPT_ATOMIC_INCREMENT); + const result = await this.client.evalsha(this.scriptSha, [key], [ttl]); + + return result as number; + } + + throw error; + } + } +} diff --git a/src/strategies/fixed_window.ts b/src/strategies/fixed_window.ts new file mode 100644 index 0000000..981d4ab --- /dev/null +++ b/src/strategies/fixed_window.ts @@ -0,0 +1,39 @@ +import type { FixedWindowState, ILimiterStrategy, IStorageEngine, LimitResult } from '../types.ts'; + +export interface FixedWindowStrategyOptions extends Record { + // The maximum number of requests allowed in a single window. + limit: number; + // The duration of the window in milliseconds. + timeFrame: number; +} + +/** + * Implements the Fixed Window rate-limiting algorithm. + */ +export class FixedWindowStrategy implements ILimiterStrategy { + public readonly options: FixedWindowStrategyOptions; + + constructor(options: FixedWindowStrategyOptions) { + if (options.limit <= 0 || options.timeFrame <= 0) { + throw new Error('FixedWindowStrategy: limit and timeFrame must be positive numbers.'); + } + + this.options = options; + } + + /** + * Checks if a request is allowed to pass. + * + * @param key The unique key for the entity being limited. + * @param storage The storage engine to use for tracking hits. + * @returns A promise that resolves to a `LimitResult`. + */ + public async check(key: string, storage: IStorageEngine): Promise { + const hits = await storage.increment(key, this.options.timeFrame); + const isAllowed = hits <= this.options.limit; + const remaining = Math.max(0, this.options.limit - hits); + const reset = this.options.timeFrame; + + return { isAllowed, remaining, reset }; + } +} diff --git a/src/strategies/token_bucket.ts b/src/strategies/token_bucket.ts new file mode 100644 index 0000000..968ba47 --- /dev/null +++ b/src/strategies/token_bucket.ts @@ -0,0 +1,120 @@ +import type { ILimiterStrategy, IStorageEngine, LimitResult, TokenBucketState } from '../types.ts'; + +/** + * Configuration options for the TokenBucketStrategy. + */ +export interface TokenBucketStrategyOptions extends Record { + // The maximum number of tokens the bucket can hold. Defines the burst limit. + bucketSize: number; + // The time interval (in milliseconds) for token refills. + interval: number; + // The number of tokens added to the bucket per interval. Defines the sustained rate. + tokensPerInterval: number; +} + +/** + * Implements the Token Bucket rate-limiting algorithm. + */ +export class TokenBucketStrategy implements ILimiterStrategy { + public readonly options: TokenBucketStrategyOptions; + + /** + * The Time-To-Live for records in storage. This should be long enough to + * ensure a user's state doesn't vanish while they are active. + */ + private readonly storageTtl: number; + + /** + * Constructs a new `TokenBucketStrategy`. + * + * @param options The configuration for the token bucket algorithm. + */ + constructor(options: TokenBucketStrategyOptions) { + if (options.bucketSize <= 0 || options.interval <= 0 || options.tokensPerInterval <= 0) { + throw new Error( + 'TokenBucketStrategy: bucketSize, interval, and tokensPerInterval must be positive numbers.', + ); + } + + this.options = options; + + const timeToFill = Math.ceil(this.options.bucketSize / this.options.tokensPerInterval) * + this.options.interval; + this.storageTtl = timeToFill; + } + + /** + * Executes the token bucket algorithm for a given request. + * + * This method gets the user's current state, refills their token bucket + * based on the elapsed time, consumes a token if available, and saves + * the new state back to storage. + * + * @param key The unique key for the entity being limited. + * @param storage The storage engine used to persist the token bucket state. + * @returns A promise that resolves to a `LimitResult`. + */ + public async check(key: string, storage: IStorageEngine): Promise { + const now = Date.now(); + const currentState = await storage.get(key); + + const state = currentState ?? { + tokens: this.options.bucketSize, + lastRefill: now, + }; + + this.refill(state, now); + + let isAllowed: boolean; + + if (state.tokens >= 1) { + isAllowed = true; + state.tokens -= 1; + } else { + isAllowed = false; + } + + await storage.set(key, state, this.storageTtl); + + const remaining = Math.floor(state.tokens); + const reset = this.calculateReset(state.tokens); + + return { isAllowed, remaining, reset }; + } + + /** + * Calculates the number of tokens to add to the bucket based on elapsed time + * + * @param state The current `TokenBucketState` for the user. + * @param now The current timestamp from `Date.now()`. + */ + private refill(state: TokenBucketState, now: number): void { + const elapsed = now - state.lastRefill; + + if (elapsed <= 0) { + return; + } + + const tokensToAdd = (elapsed / this.options.interval) * this.options.tokensPerInterval; + state.tokens = Math.min(this.options.bucketSize, state.tokens + tokensToAdd); + state.lastRefill = now; + } + + /** + * Calculates the time in milliseconds until the user will have at least one token. + * If the user already has one or more tokens, it returns 0. + * + * @param currentTokens The user's current token count (can be fractional). + * @returns The time in milliseconds until the next token is available. + */ + private calculateReset(currentTokens: number): number { + if (currentTokens >= 1) { + return 0; + } + + const needed = 1 - currentTokens; + + const timePerToken = this.options.interval / this.options.tokensPerInterval; + return Math.ceil(needed * timePerToken); + } +} diff --git a/src/types.ts b/src/types.ts new file mode 100644 index 0000000..6ae5be5 --- /dev/null +++ b/src/types.ts @@ -0,0 +1,170 @@ +/** + * A minimal, structurally-compatible representation of the grammY Context object. + */ +export interface GrammyContext { + from?: { + id: number; + is_bot: boolean; + first_name: string; + last_name?: string; + username?: string; + language_code?: string; + }; + chat?: { + id: number; + type: 'private' | 'group' | 'supergroup' | 'channel'; + title?: string; + username?: string; + first_name?: string; + }; +} + +export type NextFunction = () => Promise; + +// ==================== Core Interfaces ==================== + +/** + * Defines the contract for a storage engine. + */ +export interface IStorageEngine { + /** + * Retrieves the state for a given key. + * + * @param key The unique identifier for the entity being limited. + * @returns A promise that resolves to the state, or undefined if not found. + */ + get(key: string): Promise; + + /** + * Sets the state for a given key. + * + * @param key The unique identifier. + * @param state The new state to store. + * @param ttl The Time To Live for the record, in milliseconds. + */ + set(key: string, state: TokenBucketState, ttl: number): Promise; + + /** + * Deletes the state for a given key. + * + * @param key The unique identifier. + */ + delete(key: string): Promise; + + /** + * Atomically increments a key and sets its expiry on the first increment. + * + * @param key The unique identifier. + * @param ttl The Time To Live for the record, in milliseconds. + * @returns The new value of the counter. + */ + increment(key: string, ttl: number): Promise; + + /** + * Sets a key with a simple "true" value. + * + * @param key The unique identifier for the penalty. + * @param ttl The duration of the penalty in milliseconds. + */ + setPenalty(key: string, ttl: number): Promise; + + /** + * Checks for the existence of a penalty key. + * + * @param key The unique identifier for the penalty. + * @returns A promise that resolves to true if the penalty exists. + */ + checkPenalty(key: string): Promise; +} + +/** + * Defines the contract for a rate-limiting algorithm (Strategy). + */ +export interface ILimiterStrategy { + /** + * The configuration options used to create this strategy instance. + */ + readonly options?: Record; + + /** + * Called on every incoming request that matches the rule. + * + * @param key The unique identifier for the entity being limited. + * @param storage The storage engine used to persist state. + * @returns A promise that resolves to a `LimitResult`. + */ + check(key: string, storage: IStorageEngine): Promise; +} + +// ==================== Result & State Types ==================== + +/** + * The rich result object returned by a limiter strategy. + */ +export interface LimitResult { + isAllowed: boolean; + remaining: number; + reset: number; +} + +/** The data structure used by the Token Bucket strategy. */ +export interface TokenBucketState { + tokens: number; + lastRefill: number; +} + +/** The data structure used by the Fixed Window strategy. */ +export interface FixedWindowState { + hits: number; +} + +// ==================== Configuration & Helper Types ==================== + +/** A map of event names to their corresponding listener argument types. */ +// deno-lint-ignore no-explicit-any +export type EventMap = Record; + +/** + * Describes the events that can be emitted by the limiter, mapping + * event names to the arguments their listeners will receive. + */ +export interface LimiterEvents extends EventMap { + /** + * Fired when a request is allowed to pass. + */ + allowed: [ctx: C, info: LimitResult]; + + /** + * Fired when a request is throttled (rate-limited). + */ + throttled: [ctx: C, info: LimitResult]; + + /** + * Fired when a penalty is applied to a key. + */ + penaltyApplied: [ctx: C, key: string, duration: number]; +} + +/** + * A function executed when a user is being rate-limited. + * + * @param ctx The grammY context object. + * @param info Information about the limit that was hit. + * @param storage The storage engine instance, for advanced use cases like notification locks. + */ +export type OnLimitExceeded = ( + ctx: C, + info: LimitResult, + storage: IStorageEngine, +) => unknown; + +export type KeyGenerator = (ctx: C) => string | undefined; +export type DynamicLimitGenerator = (ctx: C) => number; +export type PenaltyDurationGenerator = ( + ctx: C, + info: LimitResult, +) => number; + +export type Mutable = { + -readonly [P in keyof T]: T[P]; +}; diff --git a/src/typesAndDefaults.ts b/src/typesAndDefaults.ts deleted file mode 100644 index 9f878e1..0000000 --- a/src/typesAndDefaults.ts +++ /dev/null @@ -1,88 +0,0 @@ -interface User { - id: number; - "is_bot": boolean; - "first_name": string; - "last_name"?: string; - "username"?: string; - "language_code"?: string; -} - -export interface Context { - from: User | undefined; -} - -export interface RedisType { - incr(key: string): Promise; - pexpire(key: string, milliseconds: number): Promise; -} - -export type NextFunction = () => Promise; - -export interface OptionsInterface { - /** - * @default 1000 (miliseconds) - * @description The time frame during which the requests will be monitored - */ - timeFrame?: number; - - /** - * @default 1 - * @description the number of requests a user is allowed to send in a specific timeFrame - */ - limit?: number; - - /** - * @default MEMORY_STORE - * @param MEMORY_STORE which uses the a Map() in memory - * @param REDIS which is your redis client you get from [ioredis](https://github.com/luin/ioredis) or [x\redis](https://deno.land/x/redis) library. - * @description The type of storage to use for keeping track of users and their requests. - */ - storageClient?: "MEMORY_STORE" | RT; - - /** - * @param ctx Is the context object you get from grammy/telegraf. - * @param next Is the next function you get from grammy/telegraf. - * @description Executed only once for each limit excess unless alwaysReply is explicitly set to true. By default, it does nothing, meaning that the user is not notified when they exceed the limit. The middleware simply ignores excessive requests and the user is required to wait. - */ - onLimitExceeded?: (ctx: C, next: NextFunction) => void; - - /** - * @default false - * @description Whether to always call onLimitExceeded or not. - */ - alwaysReply?: boolean; - - /** - * @param ctx Is the context object you get from grammy/telegraf. - * @returns A unique **string** key (identifier). - * @description A function to generate a unique key for every user. You could set it as any key you want (e.g. group id) - * @see [Getting Started](https://github.com/Amir-Zouerami/rateLimiter#-how-to-use) - */ - keyGenerator?: (ctx: C) => string | undefined; - - /** - * @default "RATE_LIMITER" - * @description A string prefix that is getting added to the storage key after calling the `keyGenerator()`. - */ - keyPrefix?: string | undefined; -} - -type DefaultOptions = Required, ( - | 'timeFrame' - | 'limit' - | 'onLimitExceeded' - | 'storageClient' - | 'keyGenerator' - | 'keyPrefix' - )> ->; - -export const defaultOptions: DefaultOptions = { - timeFrame: 1000, - limit: 1, - onLimitExceeded: (_ctx: Context, _next: NextFunction) => {}, - storageClient: "MEMORY_STORE", - keyGenerator: (ctx: Context) => ctx.from?.id.toString(), - keyPrefix: "RATE_LIMITER", -}; diff --git a/storages.ts b/storages.ts new file mode 100644 index 0000000..a3e5052 --- /dev/null +++ b/storages.ts @@ -0,0 +1,10 @@ +/** + * # Storage Engines + * + * This module exports all the available storage engines for the rate-limiter. + * + * @module + */ + +export { MemoryStore } from './src/stores/memory.ts'; +export { type IRedisClient, RedisStore } from './src/stores/redis.ts'; diff --git a/tests/core.test.ts b/tests/core.test.ts new file mode 100644 index 0000000..3d08d29 --- /dev/null +++ b/tests/core.test.ts @@ -0,0 +1,337 @@ +import { limit, Limiter } from '../mod.ts'; +import { MemoryStore } from '../src/stores/memory.ts'; +import { assert, assertEquals, assertRejects } from '@std/assert'; +import type { GrammyContext, NextFunction } from '../src/types.ts'; + +const createMockCtx = (fromId: number): GrammyContext => ({ + from: { id: fromId, is_bot: false, first_name: 'test' }, + chat: { id: fromId, type: 'private', first_name: 'test' }, +}); + +const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms)); + +const createMemoryStore = () => new MemoryStore(null); + +Deno.test('Core Rate Limiter Tests', async (t) => { + await t.step('Builder should throw if essential components are missing', async () => { + await assertRejects( + // deno-lint-ignore require-await + async () => new Limiter().build(), + Error, + 'A limiting strategy must be defined', + ); + + await assertRejects( + // deno-lint-ignore require-await + async () => new Limiter().fixedWindow({ limit: 1, timeFrame: 1000 }).build(), + Error, + 'A storage engine must be provided', + ); + + await assertRejects( + // deno-lint-ignore require-await + async () => + new Limiter().fixedWindow({ limit: 1, timeFrame: 1000 }).useStorage( + createMemoryStore(), + ).build(), + Error, + 'A key generation strategy must be defined', + ); + }); + + await t.step('FixedWindowStrategy should limit requests correctly', async () => { + const storage = createMemoryStore(); + const limiter = new Limiter().useStorage(storage).fixedWindow({ limit: 2, timeFrame: 1000 }) + .limitFor('user'); + + const middleware = limit(limiter); + let nextCalled = 0; + const next: NextFunction = () => { + nextCalled++; + return Promise.resolve(); + }; + + await middleware(createMockCtx(100), next); + await middleware(createMockCtx(100), next); + assertEquals(nextCalled, 2); + + await middleware(createMockCtx(100), next); + assertEquals(nextCalled, 2); + + await middleware(createMockCtx(200), next); + assertEquals(nextCalled, 3); + + await sleep(1100); + + await middleware(createMockCtx(100), next); + assertEquals(nextCalled, 4); + }); + + await t.step('TokenBucketStrategy should handle bursts and refills', async () => { + const storage = createMemoryStore(); + const limiter = new Limiter() + .useStorage(storage) + .tokenBucket({ bucketSize: 3, tokensPerInterval: 1, interval: 1000 }) + .limitFor('user'); + + const middleware = limit(limiter); + let nextCalled = 0; + const next: NextFunction = () => { + nextCalled++; + return Promise.resolve(); + }; + + await middleware(createMockCtx(100), next); + await middleware(createMockCtx(100), next); + await middleware(createMockCtx(100), next); + assertEquals(nextCalled, 3); + + await middleware(createMockCtx(100), next); + assertEquals(nextCalled, 3); + + await sleep(1100); + + await middleware(createMockCtx(100), next); + assertEquals(nextCalled, 4); + + await middleware(createMockCtx(100), next); + assertEquals(nextCalled, 4); + }); + + await t.step('onThrottled handler should be called', async () => { + let throttled = false; + const limiter = new Limiter() + .useStorage(createMemoryStore()) + .fixedWindow({ limit: 1, timeFrame: 1000 }) + .limitFor('user') + .onThrottled(() => { + throttled = true; + }); + + const middleware = limit(limiter); + const next: NextFunction = () => Promise.resolve(); + + await middleware(createMockCtx(100), next); + assert(!throttled, 'onThrottled should not be called for allowed request'); + + await middleware(createMockCtx(100), next); + assert(throttled, 'onThrottled should have been called'); + }); + + await t.step('PenaltyBox should mute a user', async () => { + const storage = createMemoryStore(); + const limiter = new Limiter() + .useStorage(storage) + .fixedWindow({ limit: 1, timeFrame: 1000 }) + .limitFor('user') + .withPenalty({ penaltyTime: 1000 }); + + const middleware = limit(limiter); + let nextCalled = 0; + const next: NextFunction = () => { + nextCalled++; + return Promise.resolve(); + }; + + await middleware(createMockCtx(100), next); + assertEquals(nextCalled, 1); + + await middleware(createMockCtx(100), next); + assertEquals(nextCalled, 1); + + await middleware(createMockCtx(100), next); + assertEquals(nextCalled, 1); + + await sleep(1100); + + await middleware(createMockCtx(100), next); + assertEquals(nextCalled, 2); + }); + + await t.step('onlyIf() should conditionally apply the limiter', async () => { + const storage = createMemoryStore(); + // Only limit users with ID 100 + const limiter = new Limiter() + .useStorage(storage) + .fixedWindow({ limit: 1, timeFrame: 1000 }) + .limitFor('user') + .onlyIf((ctx) => ctx.from?.id === 100); + + const middleware = limit(limiter); + let nextCalled = 0; + + const next: NextFunction = () => { + nextCalled++; + return Promise.resolve(); + }; + + // User 200 should not be limited at all + await middleware(createMockCtx(200), next); + await middleware(createMockCtx(200), next); + assertEquals(nextCalled, 2); + + // User 100 should be limited + await middleware(createMockCtx(100), next); + assertEquals(nextCalled, 3); + // Second call for user 100 should be throttled + await middleware(createMockCtx(100), next); + assertEquals(nextCalled, 3); + }); + + await t.step('FixedWindow should support dynamic limits', async () => { + const storage = createMemoryStore(); + const isAdmin = (ctx: GrammyContext) => ctx.from?.id === 200; // User 200 is an admin + + const limiter = new Limiter() + .useStorage(storage) + .fixedWindow({ + // Admins get 5 requests, others get 1 + limit: (ctx) => (isAdmin(ctx) ? 5 : 1), + timeFrame: 1000, + }) + .limitFor('user'); + + const middleware = limit(limiter); + let nextCalled = 0; + const next: NextFunction = () => { + nextCalled++; + return Promise.resolve(); + }; + + // Regular user (ID 100) + await middleware(createMockCtx(100), next); + assertEquals(nextCalled, 1); + await middleware(createMockCtx(100), next); + assertEquals(nextCalled, 1); // Throttled + + // Admin user (ID 200) + await middleware(createMockCtx(200), next); + await middleware(createMockCtx(200), next); + assertEquals(nextCalled, 3); + }); + + await t.step("limitFor('global') should apply the same limit to all users", async () => { + const storage = createMemoryStore(); + const limiter = new Limiter().useStorage(storage).fixedWindow({ limit: 1, timeFrame: 1000 }) + .limitFor('global'); + + const middleware = limit(limiter); + let nextCalled = 0; + const next: NextFunction = () => { + nextCalled++; + return Promise.resolve(); + }; + + // First call from user 100 + await middleware(createMockCtx(100), next); + assertEquals(nextCalled, 1); + + // Second call from a DIFFERENT user (200) should be throttled + await middleware(createMockCtx(200), next); + assertEquals(nextCalled, 1); + }); + + await t.step('Event emitters for "allowed" and "penaltyApplied" should fire', async () => { + let allowedFired = false; + let penaltyFired = false; + let throttledFired = false; + + const limiter = new Limiter() + .useStorage(createMemoryStore()) + .fixedWindow({ limit: 1, timeFrame: 1000 }) + .limitFor('user') + .withPenalty({ penaltyTime: 5000 }); + + limiter.on('allowed', () => (allowedFired = true)); + limiter.on('throttled', () => (throttledFired = true)); + limiter.on('penaltyApplied', () => (penaltyFired = true)); + + const middleware = limit(limiter); + const next: NextFunction = () => Promise.resolve(); + + // First call should trigger 'allowed' + await middleware(createMockCtx(100), next); + assert(allowedFired, '"allowed" event did not fire'); + assert(!throttledFired, '"throttled" should not have fired yet'); + assert(!penaltyFired, '"penaltyApplied" should not have fired yet'); + + // Second call should trigger 'throttled' and 'penaltyApplied' + await middleware(createMockCtx(100), next); + assert(throttledFired, '"throttled" event did not fire'); + assert(penaltyFired, '"penaltyApplied" event did not fire'); + }); + + await t.step('Custom key generator should create separate limits', async () => { + interface ContextWithCommand extends GrammyContext { + message?: { + text?: string; + }; + } + + const storage = createMemoryStore(); + const limiter = new Limiter() + .useStorage(storage) + .fixedWindow({ limit: 1, timeFrame: 1000 }) + .limitFor((ctx) => { + const userId = ctx.from?.id; + // This is now fully type-safe, no `any` cast needed. + const command = ctx.message?.text?.split(' ')[0]; + return userId && command ? `${userId}:${command}` : undefined; + }); + + const middleware = limit(limiter); + let nextCalled = 0; + const next: NextFunction = () => { + nextCalled++; + return Promise.resolve(); + }; + + const baseCtx = createMockCtx(100); + + // First call for user 100 on /start + await middleware({ ...baseCtx, message: { text: '/start' } }, next); + assertEquals(nextCalled, 1); + + // Second call for user 100 on /start should be throttled + await middleware({ ...baseCtx, message: { text: '/start' } }, next); + assertEquals(nextCalled, 1); + + // But a call for the same user on /help should pass + await middleware({ ...baseCtx, message: { text: '/help' } }, next); + assertEquals(nextCalled, 2); + }); + + await t.step('PenaltyBox should support dynamic penalty time', async () => { + const storage = createMemoryStore(); + const limiter = new Limiter() + .useStorage(storage) + .fixedWindow({ limit: 1, timeFrame: 1000 }) + .limitFor('user') + .withPenalty({ + // Penalty is 1000ms times the number of remaining "hits" below zero + penaltyTime: (_ctx, info) => 1000 * Math.abs(info.remaining), + }); + + const middleware = limit(limiter); + const next: NextFunction = () => Promise.resolve(); + + await middleware(createMockCtx(100), next); // remaining: 0, allowed + await middleware(createMockCtx(100), next); // remaining: -1, throttled, penalty applied for 1000ms + + // This third call should be ignored due to the penalty + let penaltyCheck = false; + await middleware(createMockCtx(100), () => { + penaltyCheck = true; // This should not be called + return Promise.resolve(); + }); + assert(!penaltyCheck); + + // Wait for the penalty to expire + await sleep(1100); + + // The user should be able to make a request again + await middleware(createMockCtx(100), next); // This works + const wasPenalized = await storage.checkPenalty('GRAMMY:RATELIMITER:PENALTY:100'); + assert(!wasPenalized, 'Penalty key should have expired'); + }); +}); diff --git a/tsconfig.json b/tsconfig.json new file mode 100644 index 0000000..59c2e7e --- /dev/null +++ b/tsconfig.json @@ -0,0 +1,13 @@ +{ + "compilerOptions": { + "declaration": true, + "forceConsistentCasingInFileNames": true, + "module": "NodeNext", + "moduleResolution": "NodeNext", + "outDir": "dist", + "skipLibCheck": true, + "strict": true, + "target": "ES2022" + }, + "include": ["src/**/*.ts", "mod.ts", "storages.ts"] +}