How to Implement Token Bucket Rate Limiting In NestJS

Rate limiting is a crucial aspect of building scalable and secure web applications. It helps prevent abuse and ensures the fair usage of resources. In this blog post, we will explore how to implement token bucket rate limiting in a NestJS application with a reset interval of 24 hours.

What is Token Bucket Rate Limiting?

Token bucket rate limiting is an algorithm that controls the rate at which a system processes requests. The main idea is that a fixed number of tokens are added to a bucket regularly. When a request arrives, a token is removed from the bucket. If there are no tokens left, the request is rejected.

Setting up a NestJS Application

To get started, let’s create a new NestJS application. First, you need to install the Nest CLI globally:

npm i -g @nestjs/cli

Next, create a new project using the following command:

nest new nestjs-rate-limiter

Navigate to the project directory:

cd nestjs-rate-limiter

Now, install the required dependencies:

npm install

Implementing Token Bucket Rate Limiting

Before we continue: This implementation does not use a database or persist the values in any way. It also relies on the user’s IP address, not a user ID. In a real-world application, you would store token counts in a database and compare them per user.

To implement token bucket rate limiting, we will create a custom RateLimiterInterceptor that will be responsible for managing tokens and limiting requests.

1. Create a RateLimiterInterceptor

First, create a new file called rate-limiter.interceptor.ts in the src directory and then open the file up and add in the following:

import {
  Injectable,
  NestInterceptor,
  ExecutionContext,
  CallHandler,
  HttpException,
  HttpStatus,
} from '@nestjs/common';
import { Observable } from 'rxjs';

Next, define the RateLimiterInterceptor class:

@Injectable()
export class RateLimiterInterceptor implements NestInterceptor {
  // ...
}

2. Implement the Token Bucket Algorithm

Inside the RateLimiterInterceptor class, add the following properties:

private readonly bucketSize: number;
private readonly tokens: Map<string, number>;
private readonly lastRefill: Map<string, number>;

Then, initialize these properties in the constructor:

constructor(bucketSize: number) {
  this.bucketSize = bucketSize;
  this.tokens = new Map();
  this.lastRefill = new Map();
}

Now, implement a refillTokens method to add tokens to the bucket:

private refillTokens(ip: string, now: number): void {
  const lastRefillTime = this.lastRefill.get(ip) || now;
  const elapsedTime = now - lastRefillTime;

  if (elapsedTime >= 24 * 60 * 60 * 1000) { // 24 hours in milliseconds
    this.tokens.set(ip, this.bucketSize);
    this.lastRefill.set(ip, now);
  }
}

Finally, implement the intercept method to handle incoming requests:

intercept(context: ExecutionContext, next: CallHandler): Observable<any> {
  const httpContext = context.switchToHttp();
  const request = httpContext.getRequest();
  const ip = request.ip;
  const now = Date.now();

  if (!this.tokens.has(ip)) {
    this.tokens.set(ip, this.bucketSize);
  }

  this.refillTokens(ip, now);

  const tokens = this.tokens.get(ip);

  if (tokens > 0) {
    this.tokens.set(ip, tokens - 1);
    return next.handle();
  } else {
    throw new HttpException(
      'Too many requests',
      HttpStatus.TOO_MANY_REQUESTS,
    );
  }
}

3. Register the RateLimiterInterceptor

Now, let’s register the RateLimiterInterceptor in the AppModule. Open src/app.module.ts and import the RateLimiterInterceptor:

import { RateLimiterInterceptor } from './rate-limiter.interceptor';

Then, add the interceptor to the providers array. Also, make sure you import APP_INTERCEPTOR from @nestjs/core as well:

providers: [
  {
    provide: APP_INTERCEPTOR,
    useValue: new RateLimiterInterceptor(5),
  },
],

In this example, we set the bucket size to 5 tokens and the refill rate to one token per 24 hours.

Testing the Application

To test the rate-limiting functionality, start the application:

npm run start

Now, in your browser, your NestJS application should be running on port 3000 so visit http://localhost:3000 and refresh the page more than 5 times. You should eventually get a rate-limiting 429 response.

Further Improvements

As mentioned in the implementation section, we’ve implemented an idea, but not something you would deploy into production. The flaw is everything is stored in memory, and IP addresses can be spoofed. In a real application, usage would be restricted on a per-user token basis.

• Implement a UserService, which contains a method for getting the user’s available tokens. It would accept a user ID.
• In the UserService, there would be an update method to update the number of tokens.
• You would get the token count from the request > user object inside the interceptor. You would also update the count every time a request is made.

Conclusion

In this blog post, we have successfully implemented token bucket rate limiting in a NestJS application. We created a custom RateLimiterInterceptor to manage the token buckets and limit requests based on the client’s IP address. We also configured the interceptor to reset the token bucket every 24 hours.

By implementing rate limiting, you can protect your application against abuse and ensure a fair distribution of resources among your users. The token bucket algorithm provides a flexible and efficient way to control the rate at which requests are processed. You can easily adjust the bucket size and refill rate to suit the needs of your application.