Javascript API rate-limit utility using lru-cache

ratelimit.utils.ts

import LRU from "lru-cache"

type RateLimitOptions = {
  uniqueTokenPerInterval: number
  interval: number
}

// This rateLimit utility utilizes an lru cache to keep track of
// request limits for the API.
//
// RateLimitOptions
// uniqueTokenPerInterval: number of unique tokens (in most cases this will be IP addresses)
// that make requests in a given time interval
// interval: time window after which X-RateLimit-Remaining will be refreshed to the provided limit
//
// The following example allows a maximum of 10,000 unique users to each make 100 requests per hour
// const limiter = rateLimit({
//   interval: 3600 * 1000, // 1 hour
//   uniqueTokenPerInterval: 10000, // Max 10000 tokens per interval
// })
//
// try {
//   await limiter.check(res, 100, "CACHE_TOKEN")
// } catch (error) {
//   return res.status(429).json({ message: "Rate limit exceeded" })
// }
//

const rateLimit = (options: RateLimitOptions) => {
  const tokenCache = new LRU({
    max: options.uniqueTokenPerInterval || 500, //default 500 unique tokens per interval
    ttl: options.interval || 60000, //default 1 minute interval
  })

  return {
    check: (res, limit, token) =>
      new Promise<void>((resolve, reject) => {
        //purges stale tokens
        tokenCache.purgeStale()

        // if the amount of tokens in the cache reaches max, reject and set the
        // most recent time a token will be purged
        if (tokenCache.size >= tokenCache.max) {
          const remainingTTL = tokenCache.getRemainingTTL(tokenCache.rkeys().next().value)
          const retryAfterSeconds = new Date(remainingTTL).getSeconds()
          res.setHeader("RetryAfter", retryAfterSeconds)
          return reject()
        }

        const tokenCount = tokenCache.get(token) || [0]
        if (tokenCount[0] === 0) {
          tokenCache.set(token, tokenCount)
        }
        tokenCount[0] += 1

        const currentUsage = tokenCount[0]
        const isRateLimited = currentUsage > parseInt(limit, 10)

        res.setHeader("X-RateLimit-Limit", limit)
        res.setHeader("X-RateLimit-Remaining", isRateLimited ? 0 : limit - currentUsage)

        //set Retry-After for user if they've been rate limited
        if (isRateLimited) {
          const rlRemainingTTL = tokenCache.getRemainingTTL(token)
          const rlRetryAfterSeconds = new Date(rlRemainingTTL).getSeconds()
          res.setHeader("RetryAfter", rlRetryAfterSeconds)
        }

        return isRateLimited ? reject() : resolve()
      }),
  }
}

export default rateLimit