Implementing Rate Limiting per Socket

In this lesson, you will learn how to protect your WebSocket server from message flooding by applying a token-bucket rate limiter to individual client connections.

The Token Bucket Concept

In a real-time system, a single malicious or buggy client could overwhelm your server by emitting thousands of events per second. To prevent this, we use the Token Bucket algorithm. Imagine each connected socket has a "bucket" that holds a maximum number of tokens. Every time a client sends a message, they must consume a token. These tokens replenish at a fixed rate over time. If the bucket is empty, the server simply ignores or rejects the incoming message.

Implementation per Socket

Because we want to enforce this per connection, we can instantiate a limiter object for every new socket and store it within the socket's metadata.

// Simple Token Bucket state per socket
function createLimiter(tokens, refillRate) {
  let bucket = {
    tokens,
    lastRefill: Date.now(),
    refillRate // ms per token
  };

  return {
    consume: () => {
      const now = Date.now();
      const elapsed = now - bucket.lastRefill;
      // Replenish tokens based on time passed
      bucket.tokens = Math.min(10, bucket.tokens + (elapsed / bucket.refillRate));
      bucket.lastRefill = now;

      if (bucket.tokens >= 1) {
        bucket.tokens -= 1;
        return true;
      }
      return false;
    }
  };
}

// Applying to the socket
io.on('connection', (socket) => {
  socket.limiter = createLimiter(10, 1000); // Max 10, refill 1 per second

  socket.on('message', (data) => {
    if (!socket.limiter.consume()) {
      return socket.emit('error', 'Rate limit exceeded');
    }
    // Process legitimate message...
  });
});

By attaching the limiter to the socket object, we ensure that each user is gated by their own independent bucket, preventing one noisy client from impacting others. This approach ensures your event-driven architecture remains responsive even under heavy load, as the server only spends resources processing events that fit within the defined throughput policy.

Key Takeaway: Implement rate limiting per socket by attaching a token-bucket state to each connection, ensuring you only process messages when the bucket has available tokens.