Merge pull request #44 from vaisu-bhut/rate_limit

Rate-limitting with Redis

Author: vaisu-bhut

docker-compose.yml

@@ -25,6 +25,7 @@ services:
       - db
       - sentinel
       - datadog-agent
+      - redis
     volumes:
       - ./services/gateway/app:/app/app
       - dogstatsd-socket:/var/run/datadog
@@ -196,6 +197,26 @@ services:
       timeout: 5s
       retries: 5
 
+  # Redis - Rate Limiting & Caching
+  redis:
+    image: redis:7-alpine
+    restart: always
+    ports:
+      - "6379:6379"
+    volumes:
+      - redis_data:/data
+    networks:
+      - clestiq-network
+    healthcheck:
+      test: ["CMD", "redis-cli", "ping"]
+      interval: 10s
+      timeout: 5s
+      retries: 5
+    labels:
+      com.datadoghq.tags.service: "clestiq-shield-redis"
+      com.datadoghq.tags.env: "development"
+      com.datadoghq.tags.version: "7.0.0"
+
   # Datadog Agent
   datadog-agent:
     image: gcr.io/datadoghq/agent:latest
@@ -257,3 +278,4 @@ networks:
 volumes:
   postgres_data:
   dogstatsd-socket:
+  redis_data:

scripts/test_rate_limits.ps1

@@ -0,0 +1,215 @@
+# Test-RateLimits.ps1
+# Verifies Rate Limiting for Tokens, App Creation, Key Creation, and Penalties.
+
+$GatewayUrl = "http://localhost:8000"
+$Username = "ratetestuser_$(Get-Random)"
+$Password = "TestPass123!"
+
+function Invoke-RestMethodWithMetrics {
+    param(
+        [string]$Uri,
+        [string]$Method,
+        [hashtable]$Headers,
+        [object]$Body,
+        [bool]$SkipError = $true
+    )
+    try {
+        if ($Body) {
+            $jsonBody = $Body | ConvertTo-Json -Depth 10
+            return Invoke-RestMethod -Uri $Uri -Method $Method -Headers $Headers -Body $jsonBody -ContentType "application/json" -ErrorAction Stop
+        } else {
+            return Invoke-RestMethod -Uri $Uri -Method $Method -Headers $Headers -ContentType "application/json" -ErrorAction Stop
+        }
+    } catch {
+        if ($SkipError) {
+             if ($_.Exception.Response) {
+                # Attempt to read the error stream
+                $stream = $_.Exception.Response.GetResponseStream()
+                if ($stream) {
+                    $reader = New-Object System.IO.StreamReader($stream)
+                    $errorBody = $reader.ReadToEnd()
+                    try {
+                        # Add a fake property to mimic Invoke-RestMethod object so downstream checks pass
+                        $errObj = $errorBody | ConvertFrom-Json
+                        # We need to attach StatusCode somehow or just return the raw PSCustomObject
+                        # But our checks look for .StatusCode or .status_code property validation wrapper
+                        # Let's return a custom object
+                        return [PSCustomObject]@{
+                            StatusCode = [int]$_.Exception.Response.StatusCode
+                            status_code = [int]$_.Exception.Response.StatusCode # For compatibility
+                            Body = $errObj
+                            Raw = $errorBody
+                            IsError = $true
+                        }
+                    } catch {
+                         return [PSCustomObject]@{
+                            StatusCode = [int]$_.Exception.Response.StatusCode
+                            status_code = [int]$_.Exception.Response.StatusCode
+                            Body = $errorBody
+                            IsError = $true
+                        }
+                    }
+                }
+                return $_.Exception.Response
+             } else {
+                Write-Host "Error request failed with no response object: $($_.Exception.Message)" -ForegroundColor Red
+                return $null
+             }
+        } else {
+            throw $_
+        }
+    }
+}
+
+Write-Host "--- Rate Limit Verification Script ---" -ForegroundColor Cyan
+
+# 1. Setup User & Auth
+Write-Host "`n[Setup] Creating User and Logging in..."
+$userBody = @{ email = "$Username@example.com"; password = $Password; full_name = "Rate Test User" }
+$signup = Invoke-RestMethodWithMetrics -Uri "$GatewayUrl/api/v1/auth/register" -Method Post -Body $userBody -SkipError $false
+Write-Host "User created: $($signup.id)"
+
+$loginBody = @{ username = "$Username@example.com"; password = $Password }
+# Note: Login expects form-data usually, but let's try JSON or adjust if needed. EagleEye auth often uses OAuth2 form.
+# If Gateway proxies /api/v1/auth/login, it might expect form data.
+$formBody = "username=$Username@example.com&password=$Password"
+try {
+    $tokenResponse = Invoke-RestMethod -Uri "$GatewayUrl/api/v1/auth/login" -Method Post -Body $formBody -ContentType "application/x-www-form-urlencoded"
+} catch {
+    Write-Error "Login failed. Ensure Gateway is running and proxies to EagleEye."
+    exit 1
+}
+$token = $tokenResponse.access_token
+$authHeader = @{ "Authorization" = "Bearer $token" }
+Write-Host "Got Token."
+
+# 2. App Creation Limit Test (Limit: 2)
+Write-Host "`n[Test 1] App Creation Limit (Target: 2)"
+for ($i = 1; $i -le 3; $i++) {
+    $rand = Get-Random
+    $appBody = @{ name = "App_$($i)_$rand"; description = "Test App" }
+    $response = Invoke-RestMethodWithMetrics -Uri "$GatewayUrl/api/v1/apps/" -Method Post -Headers $authHeader -Body $appBody
+    
+    if ($null -eq $response) { continue }
+    if ($response.IsError -eq $true -or $response.GetType().Name -eq "HttpResponseMessageWrapper") {
+         # Error response
+         $code = if ($response.StatusCode) { $response.StatusCode } else { $response.status_code } 
+         if ([int]$code -eq 429) {
+             Write-Host "[$i] Request blocked as expected (429)." -ForegroundColor Green
+         } else {
+             Write-Host "[$i] Request failed with unexpected code: $code" -ForegroundColor Red
+             Write-Host "DEBUG Info: Type=$($response.GetType().Name)" -ForegroundColor DarkGray
+             # Attempt to print body if exists
+             try { Write-Host "DEBUG Body: $($response.Body | ConvertTo-Json -Depth 2)" -ForegroundColor DarkGray } catch {}
+         }
+    } else {
+        Write-Host "[$i] App created: $($response.id)" -ForegroundColor Yellow
+        if ($i -eq 1) { $global:appId = $response.id; $global:appName = $response.name }
+    }
+}
+
+# 3. Key Creation Limit Test (Limit: 4)
+Write-Host "`n[Test 2] Key Creation Limit (Target: 4)"
+if (-not $global:appId) {
+    Write-Warning "Skipping Test 2: No App ID available from Test 1."
+} else {
+    # Use the first app created
+    for ($i = 1; $i -le 5; $i++) {
+    $keyBody = @{ name = "Key_$i" }
+    $response = Invoke-RestMethodWithMetrics -Uri "$GatewayUrl/api/v1/apps/$global:appId/keys" -Method Post -Headers $authHeader -Body $keyBody
+    
+    if ($null -eq $response) { continue }
+    if ($response.IsError -eq $true -or $response.GetType().Name -eq "HttpResponseMessageWrapper") {
+         $code = if ($response.StatusCode) { $response.StatusCode } else { $response.status_code }
+         if ([int]$code -eq 429) {
+             Write-Host "[$i] Request blocked as expected (429)." -ForegroundColor Green
+         } else {
+             Write-Host "[$i] Request failed with unexpected code: $code" -ForegroundColor Red
+         }
+    } else {
+         Write-Host "[$i] Key created: $($response.key_prefix)..." -ForegroundColor Yellow
+         if ($i -eq 1) { $global:apiKey = $response.api_key; $global:keyId = $response.id }
+    }
+}
+}
+
+# 4. Token Limit & Penalty Test
+Write-Host "`n[Test 3] Token Usage & Penalty (Limit: 5k/5min, 2 Strikes)"
+if (-not $global:apiKey) {
+    Write-Warning "Skipping Test 3: No API Key available from Test 2."
+} else {
+$headers = @{ "X-API-Key" = $global:apiKey; "Content-Type" = "application/json" }
+$chatBody = @{
+    query = "Write a 500 word story about a space adventure to Mars." 
+    model = "gemini-3-flash-preview"
+    moderation = "moderate"
+    max_output_tokens = 1000
+}
+
+# We need to loop until we hit 10k tokens.
+# Assuming each request uses ~100 tokens. 100 requests.
+$simulatedTokens = 0
+$limit = 5000
+$count = 0
+
+while ($simulatedTokens -lt $limit + 2000) { # Go a bit over
+    $count++
+    $response = Invoke-RestMethodWithMetrics -Uri "$GatewayUrl/chat/" -Method Post -Headers $headers -Body $chatBody
+    
+    if ($null -eq $response) { continue }
+    if ($response.IsError -eq $true -or $response.GetType().Name -eq "HttpResponseMessageWrapper") {
+         $code = if ($response.StatusCode) { $response.StatusCode } else { $response.status_code }
+         
+         if ([int]$code -eq 429) {
+             Write-Host "[$count] Rate Limit Hit (429)!" -ForegroundColor Green
+             # Checking penalty
+             # To trigger penalty (disable key), we need to hit 429 TWICE.
+             # We just hit it once. We should wait a second and hit it again.
+             Start-Sleep -Seconds 1
+             Write-Host "Attempting to trigger 2nd strike..."
+             $response2 = Invoke-RestMethodWithMetrics -Uri "$GatewayUrl/chat/" -Method Post -Headers $headers -Body $chatBody
+             $code2 = if ($response2.StatusCode) { $response2.StatusCode } else { $response2.status_code }
+             
+             if ([int]$code2 -eq 429) {
+                 Write-Host "2nd Strike recorded." -ForegroundColor Green
+             }
+             
+             # Now, 3rd attempt should be 403 Forbidden (Key Disabled)
+             Write-Host "Verifying Key Disabling..."
+             $response3 = Invoke-RestMethodWithMetrics -Uri "$GatewayUrl/chat/" -Method Post -Headers $headers -Body $chatBody
+             $code3 = if ($response3.StatusCode) { $response3.StatusCode } else { $response3.status_code }
+             
+             if ([int]$code3 -eq 403) {
+                 Write-Host "SUCCESS: Key has been disabled (403: Blocked by app)." -ForegroundColor Green
+                 
+                 # Try to print the message
+                 if ($response3.detail) {
+                    Write-Host "Message: $($response3.detail)" -ForegroundColor Cyan
+                 } elseif ($response3.message) {
+                    Write-Host "Message: $($response3.message)" -ForegroundColor Cyan
+                 } else {
+                    # Fallback for raw stream reading if needed, but Invoke-RestMethod parsed it
+                    Write-Host "Response Body: $($response3 | ConvertTo-Json -Depth 1 -Compress)" -ForegroundColor Cyan
+                 }
+                 break
+             } else {
+                 Write-Host "FAILURE: Key was not disabled. Code: $code3" -ForegroundColor Red
+                 break
+             }
+         } else {
+             Write-Host "[$count] Failed: $code" -ForegroundColor Red
+             Write-Host "DEBUG Info: Type=$($response.GetType().Name)" -ForegroundColor DarkGray
+             try { Write-Host "DEBUG Body: $($response.Body | ConvertTo-Json -Depth 2)" -ForegroundColor DarkGray } catch {}
+             break
+         }
+    } else {
+        # Valid response
+        $usage = $response.metrics.token_usage.total_tokens
+        $simulatedTokens += $usage
+        Write-Host "[$count] OK. Used $usage tokens. Total: $simulatedTokens / $limit" -NoNewline
+        if ($count % 5 -eq 0) { Write-Host "" } else { Write-Host " | " -NoNewline }
+    }
+}
+}
+
+Write-Host "`nTest Complete."

services/gateway/app/api/deps.py

@@ -30,15 +30,32 @@ async def get_api_key(
     result = await db.execute(
         select(ApiKey)
         .options(selectinload(ApiKey.application))
-        .filter(ApiKey.key_hash == hashed_key, ApiKey.is_active)
+        .filter(ApiKey.key_hash == hashed_key)
     )
     api_key_obj = result.scalars().first()
 
-    if not api_key_obj or not api_key_obj.application:
-        logger.warning("Authentication failed", api_key_prefix=api_key[:4] + "...")
+    if not api_key_obj:
+        logger.warning(
+            "Authentication failed: Key not found", api_key_prefix=api_key[:4] + "..."
+        )
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,
             detail="Invalid API Key",
         )
 
+    if not api_key_obj.is_active:
+        logger.warning(
+            "Authentication failed: Key disabled", api_key_prefix=api_key[:4] + "..."
+        )
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail="API Key blocked by application",
+        )
+
+    if not api_key_obj.application:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid API Key (No App)",
+        )
+
     return api_key_obj

services/gateway/app/api/v1/endpoints/chat.py

@@ -16,6 +16,7 @@
     TokenUsage,
 )
 from app.core.telemetry import telemetry
+from app.main import rate_limiter
 
 router = APIRouter()
 logger = structlog.get_logger()
@@ -64,6 +65,51 @@ async def chat_request(
         moderation=body.moderation,
     )
 
+    # --- RATE LIMIT CHECK (Token Usage) ---
+    # Convert UUID to str for Redis key
+    key_id = str(api_key.id)
+    # 1. Check if disabled (handled by deps.get_api_key/DB, but we double check or just trust DB)
+
+    # 2. Check Token Limit: 10k per 5 mins (300s)
+    token_limit_key = f"rate:tokens:{key_id}"
+    TOKEN_LIMIT = 5000
+    TOKEN_WINDOW = 300
+
+    is_allowed = await rate_limiter.check_current_usage(token_limit_key, TOKEN_LIMIT)
+    if not is_allowed:
+        # Check penalties
+        violation_key = f"rate:violations:{key_id}"
+        VIOLATION_WINDOW = 1200  # 20 mins
+
+        violations = await rate_limiter.record_violation(
+            violation_key, VIOLATION_WINDOW
+        )
+        logger.warning("Rate limit exceeded", key_id=key_id, violations=violations)
+
+        if violations >= 2:
+            # DISABLE KEY
+            logger.critical(
+                "Disabling API Key due to repeated violations",
+                key_id=key_id,
+                app_id=str(current_app.id),
+            )
+            api_key.is_active = False
+            await db.commit()
+
+            telemetry.increment(
+                "clestiq.gateway.keys_disabled", tags=[f"app:{current_app.name}"]
+            )
+
+            raise HTTPException(
+                status_code=status.HTTP_403_FORBIDDEN,
+                detail="API Key disabled due to repeated rate limit violations",
+            )
+
+        raise HTTPException(
+            status_code=status.HTTP_429_TOO_MANY_REQUESTS,
+            detail="Token rate limit exceeded (10k tokens / 5 mins)",
+        )
+
     # Get client info
     client_ip = request.client.host if request.client else None
     user_agent = request.headers.get("user-agent")
@@ -213,6 +259,9 @@ async def chat_request(
     from sqlalchemy import func
 
     api_key.last_used_at = func.now()
+    api_key.last_used_at = func.now()
+    if api_key.request_count is None:
+        api_key.request_count = 0
     api_key.request_count += 1
 
     # Update usage_data JSON
@@ -272,6 +321,15 @@ async def chat_request(
             tags=[f"app:{current_app.name}", f"model:{model_used}", "type:total"],
         )
 
+        # --- UPDATE RATE LIMITER ---
+        # Increment token usage
+        # We count total tokens (input + output)
+        total_tokens_used = response_metrics.token_usage.total_tokens
+        if total_tokens_used > 0:
+            await rate_limiter.increment_and_check(
+                token_limit_key, total_tokens_used, TOKEN_LIMIT, TOKEN_WINDOW
+            )
+
     # 4. Tokens Saved (Efficiency)
     if response_metrics.tokens_saved > 0:
         telemetry.increment(