@upstash/ratelimit 0.2.0 → 0.3.0-canary.0

package/dist/index.mjs

@@ -0,0 +1,707 @@
+// src/duration.ts
+function ms(d) {
+  const [timeString, duration] = d.split(" ");
+  const time = parseFloat(timeString);
+  switch (duration) {
+    case "ms":
+      return time;
+    case "s":
+      return time * 1e3;
+    case "m":
+      return time * 1e3 * 60;
+    case "h":
+      return time * 1e3 * 60 * 60;
+    case "d":
+      return time * 1e3 * 60 * 60 * 24;
+    default:
+      throw new Error(`Unable to parse window size: ${d}`);
+  }
+}
+
+// src/analytics.ts
+var Analytics = class {
+  redis;
+  prefix;
+  constructor(config) {
+    this.redis = config.redis;
+    this.prefix = config.prefix ?? "@upstash/ratelimit";
+  }
+  extractGeo(req) {
+    if (typeof req.geo !== "undefined") {
+      return req.geo;
+    }
+    if (typeof req.cf !== "undefined") {
+      return req.cf;
+    }
+    return {};
+  }
+  async record(event) {
+    const bucket = new Date().setUTCHours(0, 0, 0, 0).toFixed(0);
+    const key = [this.prefix, "events", bucket].join(":");
+    await this.redis.hincrby(
+      key,
+      JSON.stringify({
+        ...event,
+        time: void 0
+      }),
+      1
+    );
+  }
+  async aggregate(aggregateBy, cutoff = 0) {
+    const keys = [];
+    let cursor = 0;
+    do {
+      const [nextCursor, found] = await this.redis.scan(cursor, {
+        match: [this.prefix, "events", "*"].join(":"),
+        count: 1e3
+      });
+      cursor = nextCursor;
+      for (const key of found) {
+        const timestamp = parseInt(key.split(":").pop());
+        if (timestamp >= cutoff) {
+          keys.push(key);
+        }
+      }
+    } while (cursor !== 0);
+    const days = {};
+    await Promise.all(
+      keys.sort().map(async (key) => {
+        const fields = await this.redis.hgetall(key);
+        if (!fields) {
+          return;
+        }
+        const day = {};
+        for (const [field, count] of Object.entries(fields)) {
+          const r = JSON.parse(field);
+          for (const [k, v] of Object.entries(r)) {
+            if (k !== aggregateBy) {
+              continue;
+            }
+            if (!day[v]) {
+              day[v] = {
+                success: 0,
+                blocked: 0
+              };
+            }
+            if (r.success) {
+              day[v].success += count;
+            } else {
+              day[v].blocked += count;
+            }
+          }
+        }
+        days[key.split(":")[2]] = day;
+      })
+    );
+    return days;
+  }
+  async series(aggregateBy, cutoff = 0) {
+    const keys = [];
+    let cursor = 0;
+    do {
+      const [nextCursor, found] = await this.redis.scan(cursor, {
+        match: [this.prefix, "events", "*"].join(":"),
+        count: 1e3
+      });
+      cursor = nextCursor;
+      for (const key of found) {
+        const timestamp = parseInt(key.split(":").pop());
+        if (timestamp >= cutoff) {
+          keys.push(key);
+        }
+      }
+    } while (cursor !== 0);
+    const days = await Promise.all(
+      keys.sort().map(async (key) => {
+        const fields = await this.redis.hgetall(key);
+        const day = { time: parseInt(key.split(":")[2]) };
+        if (!fields) {
+          return day;
+        }
+        for (const [field, count] of Object.entries(fields)) {
+          const r = JSON.parse(field);
+          for (const [k, v] of Object.entries(r)) {
+            console.log({ k, v });
+            if (k !== aggregateBy) {
+              continue;
+            }
+            if (!day[v]) {
+              day[v] = 0;
+            }
+            day[v] += count;
+          }
+        }
+        return day;
+      })
+    );
+    return days;
+  }
+  async getUsage(cutoff = 0) {
+    const records = await this.aggregate("identifier", cutoff);
+    const usage = {};
+    for (const day of Object.values(records)) {
+      for (const [k, v] of Object.entries(day)) {
+        if (!usage[k]) {
+          usage[k] = { success: 0, blocked: 0 };
+        }
+        usage[k].success += v.success;
+        usage[k].blocked += v.blocked;
+      }
+    }
+    return usage;
+  }
+};
+
+// src/cache.ts
+var Cache = class {
+  cache;
+  constructor(cache) {
+    this.cache = cache;
+  }
+  isBlocked(identifier) {
+    if (!this.cache.has(identifier)) {
+      return { blocked: false, reset: 0 };
+    }
+    const reset = this.cache.get(identifier);
+    if (reset < Date.now()) {
+      this.cache.delete(identifier);
+      return { blocked: false, reset: 0 };
+    }
+    return { blocked: true, reset };
+  }
+  blockUntil(identifier, reset) {
+    this.cache.set(identifier, reset);
+  }
+  set(key, value) {
+    this.cache.set(key, value);
+  }
+  get(key) {
+    return this.cache.get(key) || null;
+  }
+  incr(key) {
+    let value = this.cache.get(key) ?? 0;
+    value += 1;
+    this.cache.set(key, value);
+    return value;
+  }
+};
+
+// src/ratelimit.ts
+var Ratelimit = class {
+  limiter;
+  ctx;
+  prefix;
+  timeout;
+  analytics;
+  constructor(config) {
+    this.ctx = config.ctx;
+    this.limiter = config.limiter;
+    this.timeout = config.timeout;
+    this.prefix = config.prefix ?? "@upstash/ratelimit";
+    this.analytics = config.analytics !== false ? new Analytics({
+      redis: Array.isArray(this.ctx.redis) ? this.ctx.redis[0] : this.ctx.redis,
+      prefix: this.prefix
+    }) : void 0;
+    if (config.ephemeralCache instanceof Map) {
+      this.ctx.cache = new Cache(config.ephemeralCache);
+    } else if (typeof config.ephemeralCache === "undefined") {
+      this.ctx.cache = new Cache(/* @__PURE__ */ new Map());
+    }
+  }
+  limit = async (identifier, req) => {
+    const key = [this.prefix, identifier].join(":");
+    let timeoutId = null;
+    try {
+      const arr = [this.limiter(this.ctx, key)];
+      if (this.timeout) {
+        arr.push(
+          new Promise((resolve) => {
+            timeoutId = setTimeout(() => {
+              resolve({
+                success: true,
+                limit: 0,
+                remaining: 0,
+                reset: 0,
+                pending: Promise.resolve()
+              });
+            }, this.timeout);
+          })
+        );
+      }
+      const res = await Promise.race(arr);
+      if (this.analytics) {
+        const geo = req ? this.analytics.extractGeo(req) : void 0;
+        const analyticsP = this.analytics.record({
+          identifier,
+          time: Date.now(),
+          success: res.success,
+          ...geo
+        });
+        res.pending = Promise.all([res.pending, analyticsP]);
+      }
+      return res;
+    } finally {
+      if (timeoutId) {
+        clearTimeout(timeoutId);
+      }
+    }
+  };
+  blockUntilReady = async (identifier, timeout) => {
+    if (timeout <= 0) {
+      throw new Error("timeout must be positive");
+    }
+    let res;
+    const deadline = Date.now() + timeout;
+    while (true) {
+      res = await this.limit(identifier);
+      if (res.success) {
+        break;
+      }
+      if (res.reset === 0) {
+        throw new Error("This should not happen");
+      }
+      const wait = Math.min(res.reset, deadline) - Date.now();
+      await new Promise((r) => setTimeout(r, wait));
+      if (Date.now() > deadline) {
+        break;
+      }
+    }
+    return res;
+  };
+};
+
+// src/single.ts
+var RegionRatelimit = class extends Ratelimit {
+  constructor(config) {
+    super({
+      prefix: config.prefix,
+      limiter: config.limiter,
+      timeout: config.timeout,
+      analytics: config.analytics,
+      ctx: {
+        redis: config.redis
+      },
+      ephemeralCache: config.ephemeralCache
+    });
+  }
+  static fixedWindow(tokens, window) {
+    const windowDuration = ms(window);
+    const script = `
+    local key     = KEYS[1]
+    local window  = ARGV[1]
+    
+    local r = redis.call("INCR", key)
+    if r == 1 then 
+    -- The first time this key is set, the value will be 1.
+    -- So we only need the expire command once
+    redis.call("PEXPIRE", key, window)
+    end
+    
+    return r
+    `;
+    return async function(ctx, identifier) {
+      const bucket = Math.floor(Date.now() / windowDuration);
+      const key = [identifier, bucket].join(":");
+      if (ctx.cache) {
+        const { blocked, reset: reset2 } = ctx.cache.isBlocked(identifier);
+        if (blocked) {
+          return {
+            success: false,
+            limit: tokens,
+            remaining: 0,
+            reset: reset2,
+            pending: Promise.resolve()
+          };
+        }
+      }
+      const usedTokensAfterUpdate = await ctx.redis.eval(script, [key], [windowDuration]);
+      const success = usedTokensAfterUpdate <= tokens;
+      const reset = (bucket + 1) * windowDuration;
+      if (ctx.cache && !success) {
+        ctx.cache.blockUntil(identifier, reset);
+      }
+      return {
+        success,
+        limit: tokens,
+        remaining: tokens - usedTokensAfterUpdate,
+        reset,
+        pending: Promise.resolve()
+      };
+    };
+  }
+  static slidingWindow(tokens, window) {
+    const script = `
+      local currentKey  = KEYS[1]           -- identifier including prefixes
+      local previousKey = KEYS[2]           -- key of the previous bucket
+      local tokens      = tonumber(ARGV[1]) -- tokens per window
+      local now         = ARGV[2]           -- current timestamp in milliseconds
+      local window      = ARGV[3]           -- interval in milliseconds
+
+      local requestsInCurrentWindow = redis.call("GET", currentKey)
+      if requestsInCurrentWindow == false then
+        requestsInCurrentWindow = 0
+      end
+
+
+      local requestsInPreviousWindow = redis.call("GET", previousKey)
+      if requestsInPreviousWindow == false then
+        requestsInPreviousWindow = 0
+      end
+      local percentageInCurrent = ( now % window) / window
+      if requestsInPreviousWindow * ( 1 - percentageInCurrent ) + requestsInCurrentWindow >= tokens then
+        return 0
+      end
+
+      local newValue = redis.call("INCR", currentKey)
+      if newValue == 1 then 
+        -- The first time this key is set, the value will be 1.
+        -- So we only need the expire command once
+        redis.call("PEXPIRE", currentKey, window * 2 + 1000) -- Enough time to overlap with a new window + 1 second
+      end
+      return tokens - newValue
+      `;
+    const windowSize = ms(window);
+    return async function(ctx, identifier) {
+      const now = Date.now();
+      const currentWindow = Math.floor(now / windowSize);
+      const currentKey = [identifier, currentWindow].join(":");
+      const previousWindow = currentWindow - windowSize;
+      const previousKey = [identifier, previousWindow].join(":");
+      if (ctx.cache) {
+        const { blocked, reset: reset2 } = ctx.cache.isBlocked(identifier);
+        if (blocked) {
+          return {
+            success: false,
+            limit: tokens,
+            remaining: 0,
+            reset: reset2,
+            pending: Promise.resolve()
+          };
+        }
+      }
+      const remaining = await ctx.redis.eval(script, [currentKey, previousKey], [tokens, now, windowSize]);
+      const success = remaining > 0;
+      const reset = (currentWindow + 1) * windowSize;
+      if (ctx.cache && !success) {
+        ctx.cache.blockUntil(identifier, reset);
+      }
+      return {
+        success,
+        limit: tokens,
+        remaining,
+        reset,
+        pending: Promise.resolve()
+      };
+    };
+  }
+  static tokenBucket(refillRate, interval, maxTokens) {
+    const script = `
+        local key         = KEYS[1]           -- identifier including prefixes
+        local maxTokens   = tonumber(ARGV[1]) -- maximum number of tokens
+        local interval    = tonumber(ARGV[2]) -- size of the window in milliseconds
+        local refillRate  = tonumber(ARGV[3]) -- how many tokens are refilled after each interval
+        local now         = tonumber(ARGV[4]) -- current timestamp in milliseconds
+        local remaining   = 0
+        
+        local bucket = redis.call("HMGET", key, "updatedAt", "tokens")
+        
+        if bucket[1] == false then
+          -- The bucket does not exist yet, so we create it and add a ttl.
+          remaining = maxTokens - 1
+          
+          redis.call("HMSET", key, "updatedAt", now, "tokens", remaining)
+          redis.call("PEXPIRE", key, interval)
+  
+          return {remaining, now + interval}
+        end
+
+        -- The bucket does exist
+  
+        local updatedAt = tonumber(bucket[1])
+        local tokens = tonumber(bucket[2])
+  
+        if now >= updatedAt + interval then
+          remaining = math.min(maxTokens, tokens + refillRate) - 1
+          
+          redis.call("HMSET", key, "updatedAt", now, "tokens", remaining)
+          return {remaining, now + interval}
+        end
+  
+        if tokens > 0 then
+          remaining = tokens - 1
+          redis.call("HMSET", key, "updatedAt", now, "tokens", remaining)
+        end
+  
+        return {remaining, updatedAt + interval}
+       `;
+    const intervalDuration = ms(interval);
+    return async function(ctx, identifier) {
+      if (ctx.cache) {
+        const { blocked, reset: reset2 } = ctx.cache.isBlocked(identifier);
+        if (blocked) {
+          return {
+            success: false,
+            limit: maxTokens,
+            remaining: 0,
+            reset: reset2,
+            pending: Promise.resolve()
+          };
+        }
+      }
+      const now = Date.now();
+      const key = [identifier, Math.floor(now / intervalDuration)].join(":");
+      const [remaining, reset] = await ctx.redis.eval(
+        script,
+        [key],
+        [maxTokens, intervalDuration, refillRate, now]
+      );
+      const success = remaining > 0;
+      if (ctx.cache && !success) {
+        ctx.cache.blockUntil(identifier, reset);
+      }
+      return {
+        success,
+        limit: maxTokens,
+        remaining,
+        reset,
+        pending: Promise.resolve()
+      };
+    };
+  }
+  static cachedFixedWindow(tokens, window) {
+    const windowDuration = ms(window);
+    const script = `
+      local key     = KEYS[1]
+      local window  = ARGV[1]
+      
+      local r = redis.call("INCR", key)
+      if r == 1 then 
+      -- The first time this key is set, the value will be 1.
+      -- So we only need the expire command once
+      redis.call("PEXPIRE", key, window)
+      end
+      
+      return r
+      `;
+    return async function(ctx, identifier) {
+      if (!ctx.cache) {
+        throw new Error("This algorithm requires a cache");
+      }
+      const bucket = Math.floor(Date.now() / windowDuration);
+      const key = [identifier, bucket].join(":");
+      const reset = (bucket + 1) * windowDuration;
+      const hit = typeof ctx.cache.get(key) === "number";
+      if (hit) {
+        const cachedTokensAfterUpdate = ctx.cache.incr(key);
+        const success = cachedTokensAfterUpdate < tokens;
+        const pending = success ? ctx.redis.eval(script, [key], [windowDuration]).then((t) => {
+          ctx.cache.set(key, t);
+        }) : Promise.resolve();
+        return {
+          success,
+          limit: tokens,
+          remaining: tokens - cachedTokensAfterUpdate,
+          reset,
+          pending
+        };
+      }
+      const usedTokensAfterUpdate = await ctx.redis.eval(script, [key], [windowDuration]);
+      ctx.cache.set(key, usedTokensAfterUpdate);
+      const remaining = tokens - usedTokensAfterUpdate;
+      return {
+        success: remaining >= 0,
+        limit: tokens,
+        remaining,
+        reset,
+        pending: Promise.resolve()
+      };
+    };
+  }
+};
+
+// src/multi.ts
+function randomId() {
+  let result = "";
+  const characters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
+  const charactersLength = characters.length;
+  for (let i = 0; i < 16; i++) {
+    result += characters.charAt(Math.floor(Math.random() * charactersLength));
+  }
+  return result;
+}
+var MultiRegionRatelimit = class extends Ratelimit {
+  constructor(config) {
+    super({
+      prefix: config.prefix,
+      limiter: config.limiter,
+      timeout: config.timeout,
+      analytics: config.analytics,
+      ctx: {
+        redis: config.redis,
+        cache: config.ephemeralCache ? new Cache(config.ephemeralCache) : void 0
+      }
+    });
+  }
+  static fixedWindow(tokens, window) {
+    const windowDuration = ms(window);
+    const script = `
+    local key     = KEYS[1]
+    local id      = ARGV[1]
+    local window  = ARGV[2]
+    
+    redis.call("SADD", key, id)
+    local members = redis.call("SMEMBERS", key)
+    if #members == 1 then
+    -- The first time this key is set, the value will be 1.
+    -- So we only need the expire command once
+      redis.call("PEXPIRE", key, window)
+    end
+    
+    return members
+`;
+    return async function(ctx, identifier) {
+      if (ctx.cache) {
+        const { blocked, reset: reset2 } = ctx.cache.isBlocked(identifier);
+        if (blocked) {
+          return {
+            success: false,
+            limit: tokens,
+            remaining: 0,
+            reset: reset2,
+            pending: Promise.resolve()
+          };
+        }
+      }
+      const requestID = randomId();
+      const bucket = Math.floor(Date.now() / windowDuration);
+      const key = [identifier, bucket].join(":");
+      const dbs = ctx.redis.map((redis) => ({
+        redis,
+        request: redis.eval(script, [key], [requestID, windowDuration])
+      }));
+      const firstResponse = await Promise.any(dbs.map((s) => s.request));
+      const usedTokens = firstResponse.length;
+      const remaining = tokens - usedTokens - 1;
+      async function sync() {
+        const individualIDs = await Promise.all(dbs.map((s) => s.request));
+        const allIDs = Array.from(new Set(individualIDs.flatMap((_) => _)).values());
+        for (const db of dbs) {
+          const ids = await db.request;
+          if (ids.length >= tokens) {
+            continue;
+          }
+          const diff = allIDs.filter((id) => !ids.includes(id));
+          if (diff.length === 0) {
+            continue;
+          }
+          await db.redis.sadd(key, ...allIDs);
+        }
+      }
+      const success = remaining > 0;
+      const reset = (bucket + 1) * windowDuration;
+      if (ctx.cache && !success) {
+        ctx.cache.blockUntil(identifier, reset);
+      }
+      return {
+        success,
+        limit: tokens,
+        remaining,
+        reset,
+        pending: sync()
+      };
+    };
+  }
+  static slidingWindow(tokens, window) {
+    const windowSize = ms(window);
+    const script = `
+      local currentKey  = KEYS[1]           -- identifier including prefixes
+      local previousKey = KEYS[2]           -- key of the previous bucket
+      local tokens      = tonumber(ARGV[1]) -- tokens per window
+      local now         = ARGV[2]           -- current timestamp in milliseconds
+      local window      = ARGV[3]           -- interval in milliseconds
+      local requestID   = ARGV[4]           -- uuid for this request
+
+
+      local currentMembers = redis.call("SMEMBERS", currentKey)
+      local requestsInCurrentWindow = #currentMembers
+      local previousMembers = redis.call("SMEMBERS", previousKey)
+      local requestsInPreviousWindow = #previousMembers
+
+      local percentageInCurrent = ( now % window) / window
+      if requestsInPreviousWindow * ( 1 - percentageInCurrent ) + requestsInCurrentWindow >= tokens then
+        return {currentMembers, previousMembers}
+      end
+
+      redis.call("SADD", currentKey, requestID)
+      table.insert(currentMembers, requestID)
+      if requestsInCurrentWindow == 0 then 
+        -- The first time this key is set, the value will be 1.
+        -- So we only need the expire command once
+        redis.call("PEXPIRE", currentKey, window * 2 + 1000) -- Enough time to overlap with a new window + 1 second
+      end
+      return {currentMembers, previousMembers}
+      `;
+    const windowDuration = ms(window);
+    return async function(ctx, identifier) {
+      if (ctx.cache) {
+        const { blocked, reset: reset2 } = ctx.cache.isBlocked(identifier);
+        if (blocked) {
+          return {
+            success: false,
+            limit: tokens,
+            remaining: 0,
+            reset: reset2,
+            pending: Promise.resolve()
+          };
+        }
+      }
+      const requestID = randomId();
+      const now = Date.now();
+      const currentWindow = Math.floor(now / windowSize);
+      const currentKey = [identifier, currentWindow].join(":");
+      const previousWindow = currentWindow - windowSize;
+      const previousKey = [identifier, previousWindow].join(":");
+      const dbs = ctx.redis.map((redis) => ({
+        redis,
+        request: redis.eval(script, [currentKey, previousKey], [tokens, now, windowDuration, requestID])
+      }));
+      const percentageInCurrent = now % windowDuration / windowDuration;
+      const [current, previous] = await Promise.any(dbs.map((s) => s.request));
+      const usedTokens = previous.length * (1 - percentageInCurrent) + current.length;
+      const remaining = tokens - usedTokens;
+      async function sync() {
+        const [individualIDs] = await Promise.all(dbs.map((s) => s.request));
+        const allIDs = Array.from(new Set(individualIDs.flatMap((_) => _)).values());
+        for (const db of dbs) {
+          const [ids] = await db.request;
+          if (ids.length >= tokens) {
+            continue;
+          }
+          const diff = allIDs.filter((id) => !ids.includes(id));
+          if (diff.length === 0) {
+            continue;
+          }
+          await db.redis.sadd(currentKey, ...allIDs);
+        }
+      }
+      const success = remaining > 0;
+      const reset = (currentWindow + 1) * windowDuration;
+      if (ctx.cache && !success) {
+        ctx.cache.blockUntil(identifier, reset);
+      }
+      return {
+        success,
+        limit: tokens,
+        remaining,
+        reset,
+        pending: sync()
+      };
+    };
+  }
+};
+export {
+  Analytics,
+  MultiRegionRatelimit,
+  RegionRatelimit as Ratelimit
+};
+//# sourceMappingURL=index.mjs.map