Security & Rate Limiting

pip install fastapi uvicorn slowapi pydantic

Step 1: The Vulnerable Server

First, let's create the basic server without protection to establish our baseline.

from fastapi import FastAPI
import uvicorn

app = FastAPI()

# A mock function simulating an LLM call
def mock_llm(prompt: str):
    return f"AI Response to: {prompt}"

@app.post("/chat")
def chat_endpoint(prompt: str):
    # This is currently vulnerable. 
    # Any length, any content, any speed is allowed.
    response = mock_llm(prompt)
    return {"response": response}

if __name__ == "__main__":
    # We run on port 8000
    uvicorn.run(app, host="127.0.0.1", port=8000)

You can run this, but don't get too attached. It's unsafe. Anyone can send a prompt with 1 million characters, or hit it 500 times a second.

Step 2: Adding Rate Limiting

from fastapi import FastAPI, Request
from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.util import get_remote_address
from slowapi.errors import RateLimitExceeded
import uvicorn

# 1. Initialize the Limiter
# get_remote_address identifies users by their IP address
limiter = Limiter(key_func=get_remote_address)

app = FastAPI()

# 2. Connect the limiter to the app so it can catch errors
app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)

def mock_llm(prompt: str):
    return f"AI Response to: {prompt}"

@app.post("/chat")
# 3. Apply the limit decorator
# This specific endpoint allows 5 requests per minute
@limiter.limit("5/minute") 
def chat_endpoint(request: Request, prompt: str):
    # Note: We MUST add 'request: Request' as a parameter 
    # so slowapi knows who is asking.
    
    response = mock_llm(prompt)
    return {"response": response}

if __name__ == "__main__":
    print("Server starting... try sending more than 5 requests!")
    uvicorn.run(app, host="127.0.0.1", port=8000)

Run this command 6 times quickly in your terminal:

On the 6th try, you won't get an AI response. You will get:

You have just saved your wallet.

Step 3: Input Validation (The Bouncer)

Rate limiting stops speed, but it doesn't stop bad data. What if someone sends a prompt that is empty, or 50,000 characters long?

from fastapi import FastAPI, Request, HTTPException
from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.util import get_remote_address
from slowapi.errors import RateLimitExceeded
from pydantic import BaseModel, Field, validator
import uvicorn

limiter = Limiter(key_func=get_remote_address)
app = FastAPI()
app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)

# 1. Define the Rules for Input
class ChatRequest(BaseModel):
    prompt: str = Field(..., min_length=5, max_length=100)
    
    # 2. Custom validator for specific abuse checks
    @validator('prompt')
    def check_abuse(cls, v):
        abuse_keywords = ["hack", "ignore instructions", "system prompt"]
        for word in abuse_keywords:
            if word in v.lower():
                raise ValueError("Unsafe content detected.")
        return v

@app.post("/chat")
@limiter.limit("10/minute") 
def chat_endpoint(request: Request, chat_data: ChatRequest):
    # chat_data is now guaranteed to be valid.
    # If it was too long, too short, or contained 'hack', 
    # FastAPI rejected it automatically before this line.
    
    return {"response": f"Processed safe prompt: {chat_data.prompt}"}

if __name__ == "__main__":
    uvicorn.run(app, host="127.0.0.1", port=8000)

Step 4: PII Redaction (Privacy Shield)

Users often accidentally paste emails or phone numbers into chatbots. We should scrub this data before sending it to an external LLM provider to comply with privacy laws and general safety.

import re

# ... (imports and setup stay the same)

def redact_pii(text: str) -> str:
    # Regex to find phone numbers (simple version: 123-456-7890)
    phone_pattern = r"\d{3}-\d{3}-\d{4}"
    # Regex to find emails
    email_pattern = r"[\w\.-]+@[\w\.-]+\.\w+"
    
    clean_text = re.sub(phone_pattern, "[PHONE REDACTED]", text)
    clean_text = re.sub(email_pattern, "[EMAIL REDACTED]", clean_text)
    return clean_text

@app.post("/chat")
@limiter.limit("10/minute") 
def chat_endpoint(request: Request, chat_data: ChatRequest):
    # 1. Validate (Happens automatically via Pydantic)
    
    # 2. Redact
    safe_prompt = redact_pii(chat_data.prompt)
    
    # 3. Simulate sending to LLM
    print(f"Sending to LLM: {safe_prompt}")
    
    return {
        "original_length": len(chat_data.prompt),
        "sanitized_prompt": safe_prompt,
        "response": "AI processed your request."
    }

Send a request with: "Call me at 555-019-2834 regarding help@example.com immediately."

The API returns: "Call me at [PHONE REDACTED] regarding [EMAIL REDACTED] immediately."

Step 5: The Complete Secure Server

Here is the final, runnable code combining everything.

from fastapi import FastAPI, Request, HTTPException
from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.util import get_remote_address
from slowapi.errors import RateLimitExceeded
from pydantic import BaseModel, Field, validator
import uvicorn
import re

# SETUP
limiter = Limiter(key_func=get_remote_address)
app = FastAPI()
app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)

# VALIDATION MODELS
class ChatRequest(BaseModel):
    # Enforce length limits to prevent token overflow
    prompt: str = Field(..., min_length=2, max_length=200)
    
    @validator('prompt')
    def no_injection(cls, v):
        # Basic check for prompt injection attempts
        forbidden = ["ignore previous", "system prompt", "sudo mode"]
        if any(bad in v.lower() for bad in forbidden):
            raise ValueError("Potential prompt injection detected.")
        return v

# UTILITIES
def redact_pii(text: str) -> str:
    # Redact Phone Numbers (XXX-XXX-XXXX)
    text = re.sub(r"\d{3}-\d{3}-\d{4}", "", text)
    # Redact Emails
    text = re.sub(r"[\w\.-]+@[\w\.-]+\.\w+", "", text)
    return text

# ENDPOINTS
@app.post("/secure_chat")
@limiter.limit("5/minute") # STRICT LIMIT
def secure_chat(request: Request, body: ChatRequest):
    
    # 1. Redact PII
    safe_prompt = redact_pii(body.prompt)
    
    # 2. Logic (Simulated LLM call)
    # In a real app, you would pass safe_prompt to OpenAI here
    
    return {
        "status": "success",
        "processed_prompt": safe_prompt,
        "message": "Sent to LLM securely."
    }

if __name__ == "__main__":
    uvicorn.run(app, host="127.0.0.1", port=8000)

Now You Try

You have a secure server. Now, extend its capabilities.

Challenge Project: Security Audit

You are now the Chief Security Officer for your own code.

Take the "Customer Support Bot" you built in Phase 2 or 3 (or create a simple one that takes a customer complaint and categorizes it).

* Validation: The system detects "Ignore rules" and rejects the request OR the System Prompt is robust enough to categorize it as "Urgent/Sentiment Negative" without turning into a pirate.

What You Learned

Today you moved from "making it work" to "making it survivable."