Introduction

Welcome to the fourth lesson of the "Implementing Rate Limiting" course! In this lesson, we will explore the concept of role-based rate limiting, a crucial aspect of API security. Role-based access control (RBAC) allows us to assign different permissions and access levels to users based on their roles, such as anonymous, standard, premium, and admin users. This lesson will guide you through implementing and testing role-based rate limiting using custom decorators in a Python-based REST API with FastAPI. By the end of this lesson, you'll be equipped to tailor API access based on user roles, ensuring both security and optimal resource usage. Let's dive in! 🚀

Understanding Role-Based Rate Limiting

Role-based rate limiting is a method of controlling the number of requests a user can make to an API based on their assigned role. Unlike global or endpoint-specific rate limiting, which applies the same limits to all users or specific endpoints, role-based rate limiting allows for more granular control. This approach is essential for ensuring that different user roles have access to the resources they need without compromising the security or performance of the API.

For instance, an admin user might require a higher request limit than a standard user due to their need to perform more administrative tasks. By creating custom decorators in Python with FastAPI, we can implement these tailored limits, enhancing both security and user experience.

Understanding JWT for Authentication

Before we implement role-based rate limiting, it's important to understand how we'll identify user roles. We'll use JSON Web Tokens (JWT) for this purpose. JWT is an open standard that defines a compact and self-contained way for securely transmitting information between parties as a JSON object.

In our context, when a user logs in, they receive a JWT that contains:

  • Their user ID
  • Their role (admin, premium, standard, etc.)
  • An expiration time

This token is signed with a secret key known only to the server, making it secure and tamper-proof. When the user makes requests to the API, they include this token in the Authorization header, allowing the server to:

  1. Verify the token's authenticity
  2. Extract information about the user (including their role)
  3. Apply the appropriate rate limits based on that role
Setting Up the Rate Limiter

Now, let's implement role-based rate limiting using custom Python decorators. We'll start by setting up the rate limiter to assign different request limits based on user roles.

from functools import wraps
from fastapi import Request, HTTPException
from collections import defaultdict
import time
from jose import jwt, ExpiredSignatureError, JWTError
from ..config import JWT_SECRET_KEY

# Store request counts per user/IP with timestamps
request_store = defaultdict(list)

def role_based_limiter(func):
    """
    Role-based rate limiter decorator
    - Admin users: 100 requests per 30 seconds
    - Premium users: 30 requests per 30 seconds
    - Standard users: 10 requests per 30 seconds
    - Anonymous users: 5 requests per 30 seconds
    """
    @wraps(func)
    async def wrapper(request: Request, *args, **kwargs):
        current_time = time.time()
        window_seconds = 30
        
        # Check if user is authenticated and get role
        auth_header = request.headers.get('authorization')
        if auth_header and auth_header.startswith('Bearer '):
            try:
                token = auth_header.split(' ')[1]
                decoded = jwt.decode(token, JWT_SECRET_KEY, algorithms=['HS256'])
                
                if decoded:
                    # Apply different limits based on user role
                    role = decoded.get('role')
                    if role == 'admin':
                        max_requests = 100  # Admin users: 100 requests/30 seconds
                    elif role == 'premium':
                        max_requests = 30   # Premium users: 30 requests/30 seconds
                    else:
                        max_requests = 10   # Standard users: 10 requests/30 seconds
                    
                    key = f"user:{decoded['userId']}"
            except (ExpiredSignatureError, JWTError):
                # Token verification failed, treat as anonymous
                key = f"ip:{request.client.host}"
                max_requests = 5  # Anonymous users: 5 requests/30 seconds
        else:
            key = f"ip:{request.client.host}"
            max_requests = 5  # Anonymous users: 5 requests/30 seconds
        
        # Clean old requests outside the window
        request_store[key] = [req_time for req_time in request_store[key] 
                            if current_time - req_time < window_seconds]
        
        # Check if limit exceeded
        if len(request_store[key]) >= max_requests:
            retry_after = int(window_seconds - (current_time - min(request_store[key])))
            
            raise HTTPException(
                status_code=429,
                detail={
                    "error": "Too Many Requests",
                    "message": "You have reached your request limit. Please try again later.",
                    "retryAfter": retry_after
                }
            )
        
        # Add current request
        request_store[key].append(current_time)
        
        # Call the original function
        return await func(request, *args, **kwargs)
    
    return wrapper

In this code, we configure the rate limiter decorator to apply different request limits based on the user's role. We use a simple in-memory store to track requests per user/IP address within a sliding window. The decorator extracts the user's role from the JWT token and assigns the appropriate limit: 100 requests for admin, 30 for premium, 10 for standard, and 5 for unauthenticated users.

Applying the Rate Limiter

Next, we apply the configured rate limiter decorator to the /premium-content endpoint.

from fastapi import APIRouter, Request
from ..middleware.rate_limiters import role_based_limiter

router = APIRouter()

@router.get("/premium-content")
@role_based_limiter
async def premium_content(request: Request):
    return {"message": "This is premium content with role-based rate limiting"}

Here, we attach the @role_based_limiter decorator to the /premium-content route. This ensures that requests to this endpoint are subject to the role-based rate limits we defined earlier. In FastAPI, decorators are applied above the route definition.

Testing and Analyzing Role-Based Rate Limiting

With the role-based rate limiting in place, it's time to test its effectiveness. We'll simulate requests from different user roles and analyze the outcomes.

import urllib.request
import urllib.parse
import urllib.error
import json
import jwt
import time

JWT_SECRET_KEY = 'jwt-secret-key'

def create_token(user_id, role):
    return jwt.encode({'userId': user_id, 'role': role}, JWT_SECRET_KEY, algorithm='HS256')

def main():
    standard_token = create_token('user123', 'user')
    premium_token = create_token('user456', 'premium')
    admin_token = create_token('admin789', 'admin')

    roles = [
        {'name': 'anonymous', 'token': None, 'requests': 7},
        {'name': 'standard', 'token': standard_token, 'requests': 12},
        {'name': 'premium', 'token': premium_token, 'requests': 32},
        {'name': 'admin', 'token': admin_token, 'requests': 20}
    ]

    for role in roles:
        print(f"\nTesting {role['name']} role:")
        for i in range(1, role['requests'] + 1):
            try:
                req = urllib.request.Request('http://localhost:3000/api/snippets/premium-content')
                
                if role['token']:
                    req.add_header('Authorization', f'Bearer {role["token"]}')
                
                with urllib.request.urlopen(req) as response:
                    response_data = json.loads(response.read().decode())
                    print(f"Request {i}: {response.status} Success")
            except urllib.error.HTTPError as e:
                try:
                    error_response = e.read().decode()
                    error_data = json.loads(error_response)
                    print(f"Request {i}: {e.code} {json.dumps(error_data)}")
                except:
                    print(f"Request {i}: {e.code} {e.reason}")
            
            time.sleep(0.05)

if __name__ == "__main__":
    main()

When you run this test script, you should see output similar to the following:

Testing anonymous role:
Request 1: 200 Success
Request 2: 200 Success
Request 3: 200 Success
Request 4: 200 Success
Request 5: 200 Success
Request 6: 429 {"error": "Too Many Requests", "message": "You have reached your request limit. Please try again later.", "retryAfter": 30}
Request 7: 429 {"error": "Too Many Requests", "message": "You have reached your request limit. Please try again later.", "retryAfter": 30}

Testing standard role:
Request 1: 200 Success
Request 2: 200 Success
Request 3: 200 Success
Request 4: 200 Success
Request 5: 200 Success
Request 6: 200 Success
Request 7: 200 Success
Request 8: 200 Success
Request 9: 200 Success
Request 10: 200 Success
Request 11: 429 {"error": "Too Many Requests", "message": "You have reached your request limit. Please try again later.", "retryAfter": 29}
Request 12: 429 {"error": "Too Many Requests", "message": "You have reached your request limit. Please try again later.", "retryAfter": 29}

Testing premium role:
Request 1: 200 Success
Request 2: 200 Success
...
Request 29: 200 Success
Request 30: 200 Success
Request 31: 429 {"error": "Too Many Requests", "message": "You have reached your request limit. Please try again later.", "retryAfter": 28}
Request 32: 429 {"error": "Too Many Requests", "message": "You have reached your request limit. Please try again later.", "retryAfter": 28}

Testing admin role:
Request 1: 200 Success
Request 2: 200 Success
...
Request 19: 200 Success
Request 20: 200 Success

The output demonstrates that:

  • Anonymous users are limited to 5 requests per 30 seconds
  • Standard users can make 10 requests per 30 seconds
  • Premium users can make 30 requests per 30 seconds
  • Admin users can make up to 100 requests per 30 seconds (only 20 tested here)

Notice that once a user exceeds their limit, subsequent requests receive a 429 Too Many Requests error with a retryAfter value indicating when they can try again.

In this test script, we create JWT tokens for different roles and simulate requests to the /premium-content endpoint. By analyzing the response logs, we can verify that the rate limits are enforced correctly, with each role receiving the appropriate number of requests before hitting the limit.

In this test script, we create JWT tokens for different roles and simulate requests to the /premium-content endpoint. By analyzing the response logs, we can verify that the rate limits are enforced correctly, with each role receiving the appropriate number of requests before hitting the limit.

Conclusion and Next Steps

In this lesson, we explored the concept of role-based rate limiting and its implementation in a Python REST API using FastAPI. We demonstrated how to assign different request limits based on user roles using custom decorators and tested the effectiveness of this approach. By tailoring access according to user roles, we can enhance both security and user experience.

As you move on to the practice exercises, you'll have the opportunity to apply what you've learned and further solidify your understanding. In the upcoming lessons, we'll continue to build on these concepts, exploring additional security measures to protect your API. Keep up the great work, and let's continue to secure our APIs! 🎉

Previous
Sign up
Join the 1M+ learners on CodeSignal
Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal