A rate limiter made with Dynamodb

rateLimiter.ts

// 1. Interface Definition
interface RateLimitOptions {
  identifier: string;  // e.g., IP address or user ID
  points: number;      // max number of requests
  duration: number;    // time window in seconds
}

// 2. Core Rate Limit Check Function
export const checkRateLimit = async ({
  identifier,
  points,
  duration,
}: RateLimitOptions): Promise<{ allowed: boolean; remaining: number; reset: number }> => {
  const now = Math.floor(Date.now() / 1000);
  const expiresAt = now + duration;

  const params = {
    TableName: 'rate-limiter',
    Key: { identifier },
    UpdateExpression:
      'SET #count = if_not_exists(#count, :start) + :inc, #firstRequest = if_not_exists(#firstRequest, :now)',
    ConditionExpression: 'attribute_not_exists(#expiresAt) OR #expiresAt < :now',
    ExpressionAttributeNames: {
      '#count': 'count',
      '#firstRequest': 'firstRequest',
      '#expiresAt': 'expiresAt',
    },
    ExpressionAttributeValues: {
      ':inc': 1,
      ':now': now,
      ':start': 0,
    },
    ReturnValues: 'ALL_NEW',
  };

  try {
    const result = await dynamoDb.update(params).promise();
    const count = result.Attributes?.count || 1;
    const firstRequest = result.Attributes?.firstRequest || now;

    // Set TTL for new records
    if (count === 1) {
      await setTTL(identifier, expiresAt);
    }

    return {
      allowed: count <= points,
      remaining: Math.max(0, points - count),
      reset: firstRequest + duration - now
    };
  } catch (error) {
    if (error.code === 'ConditionalCheckFailedException') {
      return await handleExistingWindow(identifier, points);
    }
    // Fail open to prevent service disruption
    return { allowed: true, remaining: points, reset: duration };
  }
};

// 3. TTL Management
const setTTL = async (identifier: string, expiresAt: number) => {
  const ttlParams = {
    TableName: 'rate-limiter',
    Key: { identifier },
    UpdateExpression: 'SET #expiresAt = :expiresAt',
    ExpressionAttributeNames: {
      '#expiresAt': 'expiresAt',
    },
    ExpressionAttributeValues: {
      ':expiresAt': expiresAt,
    },
  };
  return dynamoDb.update(ttlParams).promise();
};

// 4. Handling Existing Time Windows
const handleExistingWindow = async (identifier: string, points: number) => {
  const getParams = {
    TableName: 'rate-limiter',
    Key: { identifier },
  };
  
  const getResult = await dynamoDb.get(getParams).promise();
  const now = Math.floor(Date.now() / 1000);
  const count = getResult.Item?.count || 0;
  const firstRequest = getResult.Item?.firstRequest || now;
  const duration = getResult.Item?.expiresAt - firstRequest;

  return {
    allowed: count < points,
    remaining: Math.max(0, points - count),
    reset: firstRequest + duration - now
  };
};

// 5. DynamoDB Table Definition (CloudFormation/SAM)
const RateLimiterTable = {
  Type: 'AWS::DynamoDB::Table',
  Properties: {
    TableName: 'rate-limiter',
    BillingMode: 'PAY_PER_REQUEST',
    AttributeDefinitions: [
      {
        AttributeName: 'identifier',
        AttributeType: 'S',
      },
    ],
    KeySchema: [
      {
        AttributeName: 'identifier',
        KeyType: 'HASH',
      },
    ],
    TimeToLiveSpecification: {
      AttributeName: 'expiresAt',
      Enabled: true,
    },
  },
};

// 6. Example Usage in API Route
export async function POST(req: Request) {
  const ip = req.headers.get('x-forwarded-for')?.split(',')[0] || 'unknown';
  
  const rateLimit = await checkRateLimit({
    identifier: ip,
    points: 100,    // 100 requests
    duration: 3600, // per hour
  });

  if (!rateLimit.allowed) {
    return new Response(
      JSON.stringify({
        error: 'Too Many Requests',
        message: `Rate limit exceeded. Try again in ${rateLimit.reset} seconds.`
      }),
      {
        status: 429,
        headers: {
          'X-RateLimit-Limit': '100',
          'X-RateLimit-Remaining': rateLimit.remaining.toString(),
          'X-RateLimit-Reset': rateLimit.reset.toString(),
        },
      }
    );
  }

  // Continue with API logic...
}