Building a rate limiter in NodeJS using the token bucket algorithm

The Problem...

Recently, I picked up a task that reads like this "set up a rate limiter that allows x number of requests within x minutes. Block requests from an IP address that exceeds the limit for x minutes". If this is your first time hearing about rate limiting, then you know exactly how I felt when I got this task. Confused. Luckily, I have a new best friend that I could consult. Google. After doing some research, I was able to implement a solution using a third-party library. However, while researching, I came across an article listing various algorithms for building rate limiters. A particular algorithm seemed very interesting and I had to try it out.

The algorithm is known as the token bucket algorithm. According to the aforementioned article and Wikipedia.), the algorithm is used in packet-switched and telecommunication networks to either shape or police traffic. It is used to check if data transmissions conform to the defined limits on bandwidth and burstiness. While the mechanism and usage of the algorithm are quite robust, in our case, we will be borrowing the algorithm in its simplest form. Making use of a single record (bucket) to determine the limit of requests that can be made to our API from an IP address.

For this implementation, we will be making use of a remote Redis connection to store request logs. You can create one here. Also, we will be using Moment to keep track of request timestamps.

Installation of required packages

npm install redis moment --save

Next is the creation of rate limiter middleware rateLimiter.js. the code below is the middleware that provides the rate-limiting functionality.

import moment from 'moment';
import { createClient } from 'redis';

const redisClient = createClient({
  url: `redis://${process.env.REDIS_USERNAME}:${process.env.REDIS_PASSWORD}@${process.env.REDIS_HOST}:${process.env.REDIS_PORT}`,
});

redisClient.connect()

const WINDOW_SIZE_IN_MINUTES = process.env.WINDOW_SIZE_IN_MINUTES || 1;
const MAX_WINDOW_REQUEST_COUNT = process.env.MAX_WINDOW_REQUEST_COUNT || 100;
const REQUEST_BLOCK_DURATION_IN_MINUTES = process.env.REQUEST_BLOCK_DURATION_IN_MINUTES || 15;

const rateLimiter = async (req, res, next) => {
  const currentRequestTime = moment();
  const requestIpAddress = req.ip;
  const existingRequestLog = await redisClient.get(requestIpAddress);

  const newRequestLog = {
    firstRequestTimeStamp: currentRequestTime.unix(),
    requestCount: 1,
    ipAddressBlocked: false,
  };

  if (existingRequestLog === null) {
    await redisClient.set(requestIpAddress, JSON.stringify(newRequestLog));
    await next();
  } else if (existingRequestLog) {
    let requestLog = JSON.parse(existingRequestLog);

    if (requestLog.ipAddressBlocked) {
      // calculate time stamp for elapsed block
      let elapsedBlockTimeStamp = currentRequestTime
        .subtract(REQUEST_BLOCK_DURATION_IN_MINUTES, 'minutes')
        .unix();

      // if block duration has not passed, block request. if passed, reset request log
      if (requestLog.ipBlockTimeStamp > elapsedBlockTimeStamp) {
        return res.status(429).json({
          status: false,
          message: 'Request Blocked. Please try again later,
        });
      } else {
        await redisClient.set(requestIpAddress, JSON.stringify(newRequestLog));
        await next();
      }
    } else if (!requestLog.ipAddressBlocked) {
      // check if the request window has elapsed
      let windowStartTimestamp = moment().subtract(WINDOW_SIZE_IN_MINUTES, 'minutes').unix();

      if (requestLog.firstRequestTimeStamp < windowStartTimestamp) {
        await redisClient.set(requestIpAddress, JSON.stringify(newRequestLog));
        await next();
      } else if (
        requestLog.firstRequestTimeStamp >= windowStartTimestamp &&
        requestLog.requestCount >= MAX_WINDOW_REQUEST_COUNT
      ) {
        let blockedRequestLog = {
          ...requestLog,
          ipAddressBlocked: true,
          ipBlockTimeStamp: currentRequestTime.unix(),
        };

        await redisClient.set(requestIpAddress, JSON.stringify(blockedRequestLog));

        return res.status(429).json({
          status: false,
          message: 'Maximum attempts exceeded. Please try later.',
        });
      } else if (
        requestLog.firstRequestTimeStamp >= windowStartTimestamp &&
        requestLog.requestCount < MAX_WINDOW_REQUEST_COUNT
      ) {
        requestLog.requestCount++;

        await redisClient.set(requestIpAddress, JSON.stringify(requestLog));
        await next();
      }
    }
  }
};

module.exports = rateLimiter;

Let's do a step-by-step walkthrough of what we have.

First, we imported the previously installed packages and initiated a connection to the redis database.

Next, we defined the constraints for our rate-limiting middleware. WINDOW_SIZE_IN_MINUTES is the size of the window for our request limits. MAX_WINDOW_REQUEST_COUNT is the number of requests users can make within a window before they are blocked. REQUEST_BLOCK_DURATION_IN_MINUTES is the duration of the block that will be enforced when users exceed the request count within a window.

Next, we implement the rate limiting logic. The first thing we do is to save the current request's timestamp currentRequestTime. Next, we fetch the user's record existingRequestLog from Redis using their IP address req.ip. Then we initialize a variable newRequestLog. As the name suggests, this is a new request log that could be saved at different points in our rate limiter.

If null is returned by the query to fetch the user request log, this means that the user has never made a request to our API. We then store a new log for this user on Redis and move the user on.

If a record is found by the query, we parse the returned value so its content can be accessed. Next, we check this parsed value requestLog to see if the user is blocked requestLog.ipAddressBlocked. If blocked, we calculate the timestamp for an elapsed block elapsedBlockTimeStamp i.e. the time a user's block must have started for it to have elapsed. If the block is still active requestLog.ipBlockTimeStamp > elapsedBlockTimeStamp, we send the appropriate response to the user. Otherwise, we reset the user's request log to newRequestLog and move the user on.

If the user is not blocked !requestLog.ipAddressBlocked, we check if the user's request window has elapsed. If the user's first request was sent before the current window start time, we reset the user's log and start a new request window and move the user on.

If the user's first request falls within the current window and the user's request count is at the maximum, we block the user, update the user's log to reflect this block and send the appropriate response. Otherwise, if the current window is still valid and the user's request count is less than the maximum, we increment the request count, update the user's IP log and move the user on.

Then we export the middleware for use in other parts of our application