Source code
Revision control
Copy as Markdown
Other Tools
#!/usr/bin/env python3
"""
Example of using token refresh middleware with aiohttp client.
This example shows how to implement a middleware that handles JWT token
refresh automatically. The middleware:
- Adds bearer tokens to requests
- Detects when tokens are expired
- Automatically refreshes tokens when needed
- Handles concurrent requests during token refresh
This example includes a test server that simulates a JWT auth system.
Note: This is a simplified example for demonstration purposes.
In production, use proper JWT libraries and secure token storage.
"""
import asyncio
import hashlib
import json
import logging
import secrets
import time
from http import HTTPStatus
from typing import TYPE_CHECKING, Any, Coroutine, Dict, List, Union
from aiohttp import (
ClientHandlerType,
ClientRequest,
ClientResponse,
ClientSession,
hdrs,
web,
)
logging.basicConfig(level=logging.INFO)
_LOGGER = logging.getLogger(__name__)
class TokenRefreshMiddleware:
"""Middleware that handles JWT token refresh automatically."""
def __init__(self, token_endpoint: str, refresh_token: str) -> None:
self.token_endpoint = token_endpoint
self.refresh_token = refresh_token
self.access_token: Union[str, None] = None
self.token_expires_at: Union[float, None] = None
self._refresh_lock = asyncio.Lock()
async def _refresh_access_token(self, session: ClientSession) -> str:
"""Refresh the access token using the refresh token."""
async with self._refresh_lock:
# Check if another coroutine already refreshed the token
if (
self.token_expires_at
and time.time() < self.token_expires_at
and self.access_token
):
_LOGGER.debug("Token already refreshed by another request")
return self.access_token
_LOGGER.info("Refreshing access token...")
# Make refresh request without middleware to avoid recursion
async with session.post(
self.token_endpoint,
json={"refresh_token": self.refresh_token},
middlewares=(), # Disable middleware for this request
) as resp:
resp.raise_for_status()
data = await resp.json()
if "access_token" not in data:
raise ValueError("No access_token in refresh response")
self.access_token = data["access_token"]
# Token expires in 5 minutes for demo, refresh 30 seconds early
expires_in = data.get("expires_in", 300)
self.token_expires_at = time.time() + expires_in - 30
_LOGGER.info(
"Token refreshed successfully, expires in %s seconds", expires_in
)
if TYPE_CHECKING:
assert self.access_token is not None # Just assigned above
return self.access_token
async def __call__(
self,
request: ClientRequest,
handler: ClientHandlerType,
) -> ClientResponse:
"""Add auth token to request, refreshing if needed."""
# Skip token for refresh endpoint to avoid recursion
if str(request.url).endswith("/token/refresh"):
return await handler(request)
# Refresh token if needed
if not self.access_token or (
self.token_expires_at and time.time() >= self.token_expires_at
):
await self._refresh_access_token(request.session)
# Add token to request
request.headers[hdrs.AUTHORIZATION] = f"Bearer {self.access_token}"
_LOGGER.debug("Added Bearer token to request")
# Execute request
response = await handler(request)
# If we get 401, try refreshing token once
if response.status == HTTPStatus.UNAUTHORIZED:
_LOGGER.info("Got 401, attempting token refresh...")
await self._refresh_access_token(request.session)
request.headers[hdrs.AUTHORIZATION] = f"Bearer {self.access_token}"
response = await handler(request)
return response
class TestServer:
"""Test server with JWT-like token authentication."""
def __init__(self) -> None:
self.tokens_db: Dict[str, Dict[str, Union[str, float]]] = {}
self.refresh_tokens_db: Dict[str, Dict[str, Union[str, float]]] = {
# Hash of refresh token -> user data
hashlib.sha256(b"demo_refresh_token_12345").hexdigest(): {
"user_id": "user123",
"username": "testuser",
"issued_at": time.time(),
}
}
def generate_access_token(self) -> str:
"""Generate a secure random access token."""
return secrets.token_urlsafe(32)
async def _process_token_refresh(self, data: Dict[str, str]) -> web.Response:
"""Process the token refresh request."""
refresh_token = data.get("refresh_token")
if not refresh_token:
return web.json_response({"error": "refresh_token required"}, status=400)
# Hash the refresh token to look it up
refresh_token_hash = hashlib.sha256(refresh_token.encode()).hexdigest()
if refresh_token_hash not in self.refresh_tokens_db:
return web.json_response({"error": "Invalid refresh token"}, status=401)
user_data = self.refresh_tokens_db[refresh_token_hash]
# Generate new access token
access_token = self.generate_access_token()
expires_in = 300 # 5 minutes for demo
# Store the access token with expiry
token_hash = hashlib.sha256(access_token.encode()).hexdigest()
self.tokens_db[token_hash] = {
"user_id": user_data["user_id"],
"username": user_data["username"],
"expires_at": time.time() + expires_in,
"issued_at": time.time(),
}
# Clean up expired tokens periodically
current_time = time.time()
self.tokens_db = {
k: v
for k, v in self.tokens_db.items()
if isinstance(v["expires_at"], float) and v["expires_at"] > current_time
}
return web.json_response(
{
"access_token": access_token,
"token_type": "Bearer",
"expires_in": expires_in,
}
)
async def handle_token_refresh(self, request: web.Request) -> web.Response:
"""Handle token refresh requests."""
try:
data = await request.json()
return await self._process_token_refresh(data)
except json.JSONDecodeError:
return web.json_response({"error": "Invalid request"}, status=400)
async def verify_bearer_token(
self, request: web.Request
) -> Union[Dict[str, Union[str, float]], None]:
"""Verify bearer token and return user data if valid."""
auth_header = request.headers.get(hdrs.AUTHORIZATION, "")
if not auth_header.startswith("Bearer "):
return None
token = auth_header[7:] # Remove "Bearer "
token_hash = hashlib.sha256(token.encode()).hexdigest()
# Check if token exists and is not expired
if token_hash in self.tokens_db:
token_data = self.tokens_db[token_hash]
if (
isinstance(token_data["expires_at"], float)
and token_data["expires_at"] > time.time()
):
return token_data
return None
async def handle_protected_resource(self, request: web.Request) -> web.Response:
"""Protected endpoint that requires valid bearer token."""
user_data = await self.verify_bearer_token(request)
if not user_data:
return web.json_response({"error": "Invalid or expired token"}, status=401)
return web.json_response(
{
"message": "Access granted to protected resource",
"user": user_data["username"],
"data": "Secret information",
}
)
async def handle_user_info(self, request: web.Request) -> web.Response:
"""Another protected endpoint."""
user_data = await self.verify_bearer_token(request)
if not user_data:
return web.json_response({"error": "Invalid or expired token"}, status=401)
return web.json_response(
{
"user_id": user_data["user_id"],
"username": user_data["username"],
"email": f"{user_data['username']}@example.com",
"roles": ["user", "admin"],
}
)
async def run_test_server() -> web.AppRunner:
"""Run a test server with JWT auth endpoints."""
test_server = TestServer()
app = web.Application()
app.router.add_post("/token/refresh", test_server.handle_token_refresh)
app.router.add_get("/api/protected", test_server.handle_protected_resource)
app.router.add_get("/api/user", test_server.handle_user_info)
runner = web.AppRunner(app)
await runner.setup()
site = web.TCPSite(runner, "localhost", 8080)
await site.start()
return runner
async def run_tests() -> None:
"""Run all token refresh middleware tests."""
# Create token refresh middleware
# In a real app, this refresh token would be securely stored
token_middleware = TokenRefreshMiddleware(
refresh_token="demo_refresh_token_12345",
)
async with ClientSession(middlewares=(token_middleware,)) as session:
print("=== Test 1: First request (will trigger token refresh) ===")
if resp.status == 200:
data = await resp.json()
print(f"Success! Response: {data}")
else:
print(f"Failed with status: {resp.status}")
print("\n=== Test 2: Second request (uses cached token) ===")
if resp.status == 200:
data = await resp.json()
print(f"User info: {data}")
else:
print(f"Failed with status: {resp.status}")
print("\n=== Test 3: Multiple concurrent requests ===")
print("(Should only refresh token once)")
coros: List[Coroutine[Any, Any, ClientResponse]] = []
for i in range(3):
coros.append(coro)
responses = await asyncio.gather(*coros)
for i, resp in enumerate(responses):
async with resp:
if resp.status == 200:
print(f"Request {i + 1}: Success")
else:
print(f"Request {i + 1}: Failed with {resp.status}")
print("\n=== Test 4: Simulate token expiry ===")
# For demo purposes, force token expiry
token_middleware.token_expires_at = time.time() - 1
print("Token expired, next request should trigger refresh...")
if resp.status == 200:
data = await resp.json()
print(f"Success after token refresh! Response: {data}")
else:
print(f"Failed with status: {resp.status}")
print("\n=== Test 5: Request without middleware (no auth) ===")
# Make a request without any middleware to show the difference
async with session.get(
middlewares=(), # Bypass all middleware for this request
) as resp:
print(f"Status: {resp.status}")
if resp.status == 401:
error = await resp.json()
print(f"Failed as expected without auth: {error}")
async def main() -> None:
# Start test server
server = await run_test_server()
try:
await run_tests()
finally:
await server.cleanup()
if __name__ == "__main__":
asyncio.run(main())