A simple rate limit function for Node.js using MongoDB

Rate limiting can help secure your server or API from attacks, misuse or overuse. By limiting repeat requests you have a little more control. I recently had a bot attack on one of my sites, and rate limiting is one extra feature I implemented - on top of some other security measures.

For this blog post, I'm going to build a rate limit utility that:

• Sets rate limits for specific actions and time windows
• Checks rate limits
• Consumes tokens when an action happens
• Stores records in MongoDB
• Resets limits when the time window expires
• Automatically deletes old stale limits from the database

That might sound like a lot to get through, but the code for this ends up at only 59 lines (including comments!).

Here we go!

Set limits

The first thing we need to do is decide what our limits are.

I always try to set limits above the level I expect a reasonable user to reach. Ideally, I never want a usual user to hit a rate limit.

I also try to set rate limits per user where possible, so that one user's actions can't affect other users. In some cases, this isn't possible, like public forms (contact forms, registration forms etc), so I would add other security measures like CAPTCHA, honey pots etc as an additional barrier to try to stop bad requests using up the rate limit.

Let's look at a couple of rate limits that would be common across online services.

For example, maybe we want to prevent an attack on our registration form, so we think a rate limit of 60 new accounts per hour (averaging one a minute) would be reasonable for our small API.

And we want to prevent a cart testing attack, so we will set a limit of 5 card declines per day, per user, on our checkout.

Let's start with these two limits.

// dates
const milliseconds = require('date-fns/milliseconds');

// limit tokens
const limits = {
  "card-decline": (userId) => ({ window: milliseconds({ days: 1 }), tokens: userId ? 5 : 100 }),
  "register": () => ({ window: milliseconds({ hours: 1 }), tokens: 60 }),
};

A couple of things:

1. I'm using milliseconds from date-fns, so I don't have to work out the milliseconds myself.
2. I'm using a function to create the limits so that I can have both user-specific and site-wide tokens for a limit (e.g. a user can have 5 card declines in a day, and across the site, there's a limit of 100 in a day).

How the limits will work is there is a window property, which will tell us how many milliseconds a limit is valid for.

For a rate-limited action, like creating an account or getting a card payment declined, a rate limit will be created with a timestamp. We can use the window for the limit to check if the rate limit window has expired when we consume tokens in the future (I'll cover this when we create a function to check rate limits shortly).

The tokens are the number of times that action can happen until the limit window is reached.

In the case of the "register" limit, 60 tokens can be used within that hour window. If they are all used within 30 minutes, new users will have to wait until the window expires before they can get new tokens to register an account.

That's the rate limit - 60 per hour max.

Create a MongoDB collection

We're going to need to store our rate limits somewhere, so we know how many credits have been consumed. Since I was already using MongoDB for my database, I'm going with that!

Let's create a new collection called rate_limits.

You can create a new collection in the MongoDB Atlas Web UI by selecting your database and clicking Collections > Create Collection, or with the following command:

db.createCollection("rate_limits");

Now we can create the functions to save our rate limits to the rate_limits collection.

Check and create limits

Ok, our limits document stored in the database is going to look like this:

{
  key: "card-decline", // the rate limit key
  id: ObjectId("542c2b97bac0595474108b48"), // OPTIONAL id e.g. a user id
  timestamp: 1689090116748, // the time the rate limit was started (in milliseconds)
  tokens: 59 // how many tokens are left
}

The id key won't always be used, for example, if it's the "register" rate limit, we won't give it an id. But if it's a "card-decline" limit, then if user123 uses their 5 token allowance, I still want user234 to be able to make a payment. I'll pass the user id as the id to give them separate rate limits in this case.

Here's the function we can use the check a rate limit.

// database connection
const {db} = require('../server');

const checkLimit = async function({ key, id }) {
  // a query to check in the db for an existing rate limit
  let findQuery = { key, ...(id ? { id } : {}) };
  // look for the rate limit and just return the timestamp and token
  let limit = await db.collection('rate_limits').findOne(findQuery, { projection: { timestamp: 1, tokens: 1, _id: 0 } });
  // if there is no existing limit, or if the rate limit window has expired we need to create a new one
  if (!limit || limit.timestamp < (Date.now() - limits[key](id).window)) {
    limit = { timestamp: Date.now(), tokens: limits[key](id).tokens };
  }
  // if there are no tokens left then the rate limit has been reached
  if (limit.tokens < 1) {
    throw new Error('Too Many Requests'); // throw an error
  }
  // return the limit so we can consume it
  return limit;
};

Currently, we can use our checkLimit function to check if a rate limit has been reached checkLimit({ key: "card-decline", id: userId }) or checkLimit({ key: "register" }) but it only tells us if a rate limit is reached. And since we are not consuming any tokens yet, it never will be!

Let's fix that.

Consuming tokens

Each time a rate-limited action happens we want to consume a token. Once the tokens go down to zero, our checkLimit function above will throw an error, and we can stop the action.

Had a card payment declined 5 times? Hold your horses buddy - you've got a problem!

Let's create a new consumeLimit function to use up tokens. Like the checkToken function, it will take a key and optional id, but it will also take the limit object we returned from checkLimit - because you can't consume a limit without checking it first*!

*We will put these two functions together shortly, but I've created them separately for some particular use cases we will look at in the examples.

// consume a limit token
const consumeLimit = async function({ key, id, limit }) {
  // the query for the rate limit document
  let findQuery = { key, ...(id ? { id } : {}) };
  // consume a token
  limit.tokens -= 1;
  // update the db (or insert if the limit doesn't exist)
  await db.collection('rate_limits').updateOne(findQuery, { $set: limit }, { upsert: true });
  // return the limit
  return limit;
};

I've commented out the function to show what each line is doing, but essentially it reduces the tokens by one and updates (with updateOne) the database. Since the checkToken function throws an error if we have no tokens left to consume, this function should only be reached if we still have a rate limit allowance for the action.

Putting checkLimit and consumeLimit together

In most cases, we can run these two functions together. Here's a rateLimit function that checks the rate limit, and then consumes a token:

// use a rate limit
const rateLimit = async function({ id, key }) {
  const limit = await checkLimit({ id, key });
  return await consumeLimit({ id, key, limit });
};

Let's say, we have the account register API route. Before registering an account, we can consume a token (if one is available).

// rate limit
try {
    await rateLimit({ key: "register" });
} catch (error) {
    return res.status(429).send({ error: "Temporarily unable to register new accounts, please try again later."});
}
// carry on registering the account

If an error is thrown, we can respond to the user with the error, and not register the new account.

You might be thinking that you could use findOneAndUpdate and skip the extra database call. And you could. But I've kept checkLimit and consumeLimit available as individual functions because in some use cases, you need them to be.

For example, our "card decline" rate limit. We need to check if any tokens are available before attempting a card payment, but we only know if we need to consume a token after the card payment - because we only use a token if it declines!

That might look something like this:

let limit;
// check user has any tokens left
try {
  limit = await checkLimit({ id: userId, key: "card-decline" });
} catch (error) {
  throw new Error("Too many failed card attempts, please check your bank.");
}
// try to take a payment
chargeCard(paymentInfo, async function(err, success) {
  if (err) {
    // if card is declined use a token
    if (err.type === 'CardError') {
      await consumeLimit({ id: userId, key: "card-decline", limit });
    }
    // throw the error
    throw new Error(err.message);
  } else {
    // payment success
    return success.payment;
  }
});

Adding helper keys

A pattern I like (but is optional) is adding a helper keys object so that when I am using the rate limits around my code, in other files, I don't need to remember the key for each limit.

// helper keys
const limitKeys = {
  cardDecline: "card-decline",
  register: "register",
};

I can use this object to get the string for each key, for example, limitKeys.cardDecline is "card-decline".

Now when I need to use a rate limit in my code, instead of await checkLimit({ id: userId, key: "card-decline" }), I can do await checkLimit({ id: userId, key: limitKeys.cardDecline }).

It has a few benefits:

• I don't need to keep checking back what string I'm using for each key
• If I make a typo I'll get an error instead of creating a new key with the typo in my database
• It automatically pops up with the keys making it less likely I'll make a typo in the first place
• If I want to change a key in the future, I can just do it in one place (the limitKeys object) and all my code will start using that new key

Create a TTL index

We don't need these rate limits to stick around forever, since each window is only a short period of time. If a user becomes inactive, I don't need to be storing a rate limit from 3 months ago.

Let's save space in our database by only storing the data we need!

Most of my rate limits are by the hour or by the day, but I'd like to have some flexibility in case I add weekly or monthly limits in the future, so I'm thinking of keeping the data for three months.

I was originally considering deleting a rate limit document if I access it and the time had expired, but that would still lead to stale data since inactive users won't trigger any checks to the rate limit collection.

Then I thought about some kind of cron job to run on the database each night to look for stale documents. But that was going to be some extra set-up, and I wanted a quick and easy solution.

Luckily, MongoDB has the perfect index for this use case - TTL (time to live). We can tell the index how long the documents in our collection should live, and it will delete them once they expire. No extra code is required!

Sounds perfect!

To create a TTL index, we give the index our timestamp field and a expireAfterSeconds option with the number of seconds. In this case, I'm using 7776000 for 90 days worth of seconds - you can adjust this to whatever value you require.

db.collection("rate_limits").createIndex({ timestamp: 1 }, { expireAfterSeconds: 7776000 });

Any old rate limits will now be removed from the database when no longer needed.

The rateLimit.util.js file

Here's the full file for my new rate limit util. I keep this on my server in /utils (I also have a blog post about my Node.js file structure).

// database connection
const {db} = require('../server');

// dates
const milliseconds = require('date-fns/milliseconds');

// limit tokens
const limits = {
  "card-decline": (userId) => ({ window: milliseconds({ days: 1 }), tokens: userId ? 5 : 100 }),
  "register": () => ({ window: milliseconds({ hours: 1 }), tokens: 60 }),
};

// helper keys
const limitKeys = {
  cardDecline: "card-decline",
  register: "register",
};

// check a limit
const checkLimit = async function({ key, id }) {
  // a query to check in the db for an existing rate limit
  let findQuery = { key, ...(id ? { id } : {}) };
  // look for the rate limit and just return the timestamp and token
  let limit = await db.collection('rate_limits').findOne(findQuery, { projection: { timestamp: 1, tokens: 1, _id: 0 } });
  // if there is no existing limit, or if the rate limit window has expired we need to create a new one
  if (!limit || limit.timestamp < (Date.now() - limits[key](id).window)) {
    limit = { timestamp: Date.now(), tokens: limits[key](id).tokens };
  }
  // if there are no tokens left then the rate limit has been reached
  if (limit.tokens < 1) {
    throw new Error('Too Many Requests'); // throw an error
  }
  // return the limit so we can consume it
  return limit;
};

// consume a limit token
const consumeLimit = async function({ key, id, limit }) {
  // the query for the rate limit document
  let findQuery = { key, ...(id ? { id } : {}) };
  // consume a token
  limit.tokens -= 1;
  // update the db (or insert if the limit doesn't exist)
  await db.collection('rate_limits').updateOne(findQuery, { $set: limit }, { upsert: true });
  // return the limit
  return limit;
};

// use a rate limit
const rateLimit = async function({ id, key }) {
  const limit = await checkLimit({ id, key });
  return await consumeLimit({ id, key, limit });
};

// exports
exports.limitKeys = limitKeys;
exports.checkLimit = checkLimit;
exports.consumeLimit = consumeLimit;
exports.rateLimit = rateLimit;