集成令牌桶(Token Bucket)算法来实现请求频率限制

在 Next.js 后端应用中集成令牌桶(Token Bucket)算法来实现请求频率限制,可以帮助保护 API 免受恶意攻击,避免服务过载.Next.js 支持 API 路由(API Routes),这些路由可以处理后端请求,因此可以在其中集成令牌桶算法.

1. 创建 TokenBucket 类

// utils/tokenBucket.ts

class TokenBucket {
  constructor(tokenGenerationRate, bucketCapacity) {
    this.tokenGenerationRate = tokenGenerationRate; // 每秒生成的令牌数量
    this.bucketCapacity = bucketCapacity; // 桶的最大容量
    this.tokens = 0; // 当前令牌数
    this.lastCheck = Date.now(); // 上次生成令牌的时间
  }

  // 更新桶中的令牌数
  updateTokens() {
    const now = Date.now();
    const timeElapsed = (now - this.lastCheck) / 1000; // 计算时间间隔,单位为秒

    // 根据时间间隔生成新的令牌
    this.tokens = Math.min(this.bucketCapacity, this.tokens + timeElapsed * this.tokenGenerationRate);
    this.lastCheck = now;
  }

  // 尝试获取一个令牌,返回是否可以请求通过
  tryConsume() {
    this.updateTokens();
    
    if (this.tokens >= 1) {
      this.tokens -= 1; // 消耗一个令牌
      return true; // 允许通过
    } else {
      return false; // 拒绝请求
    }
  }
}

module.exports = TokenBucket;

2. 在 API 路由中使用令牌桶

接下来,在 Next.js 中的 API 路由中集成这个令牌桶算法.假设你有一个限制请求频率的 API 路由,你可以在该路由中检查请求是否符合令牌桶的规则.

// pages/api/limitedRequest.js
import TokenBucket from '../../utils/tokenBucket';

// 创建一个 TokenBucket 实例,令牌生成速率为每秒 5 个令牌,桶容量为 10 个令牌
const tokenBucket = new TokenBucket(5, 10);

export default async function handler(req, res) {
  // 每次请求前,尝试从令牌桶获取一个令牌
  if (tokenBucket.tryConsume()) {
    // 如果令牌足够,处理请求
    res.status(200).json({ message: '请求成功' });
  } else {
    // 如果令牌不足,拒绝请求
    res.status(429).json({ error: '请求频率过高,请稍后再试' });
  }
}

3. 前端请求

一旦后端 API 实现了请求限制,前端可以像通常那样调用该 API.使用 fetch 或任何其他 HTTP 请求工具,前端发送请求到上述 API 路由.

async function makeRequest() {
  const response = await fetch('/api/limitedRequest');
  
  if (response.ok) {
    const data = await response.json();
    console.log(data.message);
  } else if (response.status === 429) {
    const error = await response.json();
    console.error(error.error);  // 输出 "请求频率过高,请稍后再试"
  }
}

4. 配置缓存或会话管理(可选)

为了在多个请求之间保留令牌桶的状态,你可能需要考虑将令牌桶的状态存储在缓存中,或者与用户会话相关联.对于高并发或分布式应用,使用 Redis 或其他缓存技术是常见的做法,这样令牌桶状态能够跨多个请求和服务器实例共享.

import Redis from 'ioredis';

const redis = new Redis(); // 连接到本地的 Redis 实例

class TokenBucket {
  constructor(tokenGenerationRate, bucketCapacity, userId) {
    this.tokenGenerationRate = tokenGenerationRate;
    this.bucketCapacity = bucketCapacity;
    this.userId = userId; // 假设每个用户都有唯一的 userId
    this.redisKey = `token_bucket:${userId}`;
  }

  async updateTokens() {
    const lastCheck = await redis.get(`${this.redisKey}:lastCheck`);
    const now = Date.now();
    const timeElapsed = (now - lastCheck) / 1000;

    const currentTokens = await redis.get(`${this.redisKey}:tokens`);
    const tokens = Math.min(this.bucketCapacity, parseInt(currentTokens || '0') + timeElapsed * this.tokenGenerationRate);

    await redis.set(`${this.redisKey}:tokens`, tokens);
    await redis.set(`${this.redisKey}:lastCheck`, now);
  }

  async tryConsume() {
    await this.updateTokens();

    const tokens = await redis.get(`${this.redisKey}:tokens`);

    if (parseInt(tokens || '0') >= 1) {
      await redis.decr(`${this.redisKey}:tokens`);
      return true;
    } else {
      return false;
    }
  }
}

export default TokenBucket;