add redis
This commit is contained in:
saidsurucu
+127
-39
@@ -1,5 +1,6 @@
|
||||
"""
|
||||
Simplified MCP OAuth HTTP adapter - only Clerk JWT based authentication
|
||||
Uses Redis for authorization code storage to support multi-machine deployment
|
||||
"""
|
||||
|
||||
import os
|
||||
@@ -10,6 +11,9 @@ from urllib.parse import urlencode, quote
|
||||
from fastapi import APIRouter, Request, Query, HTTPException
|
||||
from fastapi.responses import RedirectResponse, JSONResponse
|
||||
|
||||
# Import Redis session store
|
||||
from redis_session_store import get_redis_store
|
||||
|
||||
# Try to import Clerk SDK
|
||||
try:
|
||||
from clerk_backend_api import Clerk
|
||||
@@ -26,14 +30,50 @@ router = APIRouter()
|
||||
BASE_URL = os.getenv("BASE_URL", "https://api.yargimcp.com")
|
||||
CLERK_DOMAIN = os.getenv("CLERK_DOMAIN", "accounts.yargimcp.com")
|
||||
|
||||
# Initialize Redis store
|
||||
redis_store = None
|
||||
|
||||
def get_redis_session_store():
|
||||
"""Get Redis store instance with lazy initialization."""
|
||||
global redis_store
|
||||
if redis_store is None:
|
||||
try:
|
||||
import concurrent.futures
|
||||
import functools
|
||||
|
||||
# Use thread pool with timeout to prevent hanging
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
|
||||
future = executor.submit(get_redis_store)
|
||||
try:
|
||||
# 5 second timeout for Redis initialization
|
||||
redis_store = future.result(timeout=5.0)
|
||||
if redis_store:
|
||||
logger.info("Redis session store initialized for OAuth handler")
|
||||
else:
|
||||
logger.warning("Redis store initialization returned None")
|
||||
except concurrent.futures.TimeoutError:
|
||||
logger.error("Redis initialization timed out after 5 seconds")
|
||||
redis_store = None
|
||||
future.cancel() # Try to cancel the hanging operation
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to initialize Redis store: {e}")
|
||||
redis_store = None
|
||||
|
||||
if redis_store is None:
|
||||
# Fall back to in-memory storage with warning
|
||||
logger.warning("Falling back to in-memory storage - multi-machine deployment will not work")
|
||||
|
||||
return redis_store
|
||||
|
||||
@router.get("/.well-known/oauth-authorization-server")
|
||||
async def get_oauth_metadata():
|
||||
"""OAuth 2.0 Authorization Server Metadata (RFC 8414)"""
|
||||
return JSONResponse({
|
||||
"issuer": BASE_URL,
|
||||
"authorization_endpoint": f"{BASE_URL}/auth/login",
|
||||
"token_endpoint": f"{BASE_URL}/auth/callback",
|
||||
"registration_endpoint": f"{BASE_URL}/auth/register",
|
||||
"authorization_endpoint": "https://yargimcp.com/mcp-callback",
|
||||
"token_endpoint": f"{BASE_URL}/token",
|
||||
"registration_endpoint": f"{BASE_URL}/register",
|
||||
"response_types_supported": ["code"],
|
||||
"grant_types_supported": ["authorization_code"],
|
||||
"code_challenge_methods_supported": ["S256"],
|
||||
@@ -165,24 +205,38 @@ async def oauth_callback(
|
||||
# Generate authorization code
|
||||
auth_code = f"clerk_auth_{os.urandom(16).hex()}"
|
||||
|
||||
# Store code with JWT token mapping (in production, use proper storage)
|
||||
# For now, we'll use a simple in-memory storage
|
||||
# Prepare code data
|
||||
import time
|
||||
code_data = {
|
||||
"user_id": user_id,
|
||||
"session_id": session_id,
|
||||
"real_jwt_token": real_jwt_token,
|
||||
"user_authenticated": user_authenticated,
|
||||
"created_at": time.time(),
|
||||
"expires_at": time.time() + 300 # 5 minutes expiry
|
||||
"client_id": client_id,
|
||||
"redirect_uri": redirect_uri,
|
||||
"scope": scope or "read search"
|
||||
}
|
||||
|
||||
# Store in module-level dict (in production, use Redis or database)
|
||||
if not hasattr(oauth_callback, '_code_storage'):
|
||||
oauth_callback._code_storage = {}
|
||||
oauth_callback._code_storage[auth_code] = code_data
|
||||
|
||||
logger.info(f"Stored authorization code with real JWT token")
|
||||
# Try to store in Redis, fall back to in-memory if Redis unavailable
|
||||
store = get_redis_session_store()
|
||||
if store:
|
||||
# Store in Redis with automatic expiration
|
||||
success = store.set_oauth_code(auth_code, code_data)
|
||||
if success:
|
||||
logger.info(f"Stored authorization code {auth_code[:10]}... in Redis with real JWT token")
|
||||
else:
|
||||
logger.error(f"Failed to store authorization code in Redis, falling back to in-memory")
|
||||
# Fall back to in-memory storage
|
||||
if not hasattr(oauth_callback, '_code_storage'):
|
||||
oauth_callback._code_storage = {}
|
||||
oauth_callback._code_storage[auth_code] = code_data
|
||||
else:
|
||||
# Fall back to in-memory storage
|
||||
logger.warning("Redis not available, using in-memory storage")
|
||||
if not hasattr(oauth_callback, '_code_storage'):
|
||||
oauth_callback._code_storage = {}
|
||||
oauth_callback._code_storage[auth_code] = code_data
|
||||
logger.info(f"Stored authorization code in memory (fallback)")
|
||||
|
||||
# Redirect back to client with authorization code
|
||||
redirect_params = {
|
||||
@@ -272,15 +326,25 @@ async def oauth_callback_post(request: Request):
|
||||
content={"error": "invalid_grant", "error_description": "Invalid authorization code"}
|
||||
)
|
||||
|
||||
# TODO: In production, validate code against stored session
|
||||
# For now, we'll return a placeholder response
|
||||
|
||||
# Retrieve stored JWT token using authorization code
|
||||
# Retrieve stored JWT token using authorization code from Redis or in-memory fallback
|
||||
stored_code_data = None
|
||||
|
||||
# Get stored code data from authorization flow
|
||||
if hasattr(oauth_callback, '_code_storage'):
|
||||
# Try to get from Redis first, then fall back to in-memory
|
||||
store = get_redis_session_store()
|
||||
if store:
|
||||
stored_code_data = store.get_oauth_code(code, delete_after_use=True)
|
||||
if stored_code_data:
|
||||
logger.info(f"Retrieved authorization code {code[:10]}... from Redis")
|
||||
else:
|
||||
logger.warning(f"Authorization code {code[:10]}... not found in Redis")
|
||||
|
||||
# Fall back to in-memory storage if Redis unavailable or code not found
|
||||
if not stored_code_data and hasattr(oauth_callback, '_code_storage'):
|
||||
stored_code_data = oauth_callback._code_storage.get(code)
|
||||
if stored_code_data:
|
||||
# Clean up in-memory storage
|
||||
oauth_callback._code_storage.pop(code, None)
|
||||
logger.info(f"Retrieved authorization code {code[:10]}... from in-memory storage")
|
||||
|
||||
if not stored_code_data:
|
||||
logger.error(f"No stored data found for authorization code: {code}")
|
||||
@@ -289,13 +353,11 @@ async def oauth_callback_post(request: Request):
|
||||
content={"error": "invalid_grant", "error_description": "Authorization code not found or expired"}
|
||||
)
|
||||
|
||||
# Check if code is expired
|
||||
# Note: Redis TTL handles expiration automatically, but check for manual expiration for in-memory fallback
|
||||
import time
|
||||
if time.time() > stored_code_data.get("expires_at", 0):
|
||||
expires_at = stored_code_data.get("expires_at", 0)
|
||||
if expires_at and time.time() > expires_at:
|
||||
logger.error(f"Authorization code expired: {code}")
|
||||
# Clean up expired code
|
||||
if hasattr(oauth_callback, '_code_storage'):
|
||||
oauth_callback._code_storage.pop(code, None)
|
||||
return JSONResponse(
|
||||
status_code=400,
|
||||
content={"error": "invalid_grant", "error_description": "Authorization code expired"}
|
||||
@@ -306,7 +368,7 @@ async def oauth_callback_post(request: Request):
|
||||
|
||||
if real_jwt_token:
|
||||
logger.info("Returning real Clerk JWT token")
|
||||
# Clean up used code
|
||||
# Note: Code already deleted from Redis, clean up in-memory fallback if used
|
||||
if hasattr(oauth_callback, '_code_storage'):
|
||||
oauth_callback._code_storage.pop(code, None)
|
||||
|
||||
@@ -334,6 +396,26 @@ async def oauth_callback_post(request: Request):
|
||||
content={"error": "server_error", "error_description": str(e)}
|
||||
)
|
||||
|
||||
@router.post("/register")
|
||||
async def register_client(request: Request):
|
||||
"""Dynamic Client Registration (RFC 7591)"""
|
||||
|
||||
data = await request.json()
|
||||
logger.info(f"Client registration request: {data}")
|
||||
|
||||
# Simple dynamic registration - accept any client
|
||||
client_id = f"mcp-client-{os.urandom(8).hex()}"
|
||||
|
||||
return JSONResponse({
|
||||
"client_id": client_id,
|
||||
"client_secret": None, # Public client
|
||||
"redirect_uris": data.get("redirect_uris", []),
|
||||
"grant_types": ["authorization_code"],
|
||||
"response_types": ["code"],
|
||||
"client_name": data.get("client_name", "MCP Client"),
|
||||
"token_endpoint_auth_method": "none"
|
||||
})
|
||||
|
||||
@router.post("/token")
|
||||
async def token_endpoint(request: Request):
|
||||
"""OAuth 2.1 Token Endpoint - exchanges code for Clerk JWT"""
|
||||
@@ -369,17 +451,25 @@ async def token_endpoint(request: Request):
|
||||
content={"error": "invalid_grant", "error_description": "Invalid authorization code"}
|
||||
)
|
||||
|
||||
# In a real implementation, you would:
|
||||
# 1. Validate the code against stored session
|
||||
# 2. Extract user info from the session
|
||||
# 3. Return the actual Clerk JWT token
|
||||
|
||||
# Retrieve stored JWT token using authorization code
|
||||
# Retrieve stored JWT token using authorization code from Redis or in-memory fallback
|
||||
stored_code_data = None
|
||||
|
||||
# Get stored code data from authorization flow
|
||||
if hasattr(oauth_callback, '_code_storage'):
|
||||
# Try to get from Redis first, then fall back to in-memory
|
||||
store = get_redis_session_store()
|
||||
if store:
|
||||
stored_code_data = store.get_oauth_code(code, delete_after_use=True)
|
||||
if stored_code_data:
|
||||
logger.info(f"Retrieved authorization code {code[:10]}... from Redis (/token endpoint)")
|
||||
else:
|
||||
logger.warning(f"Authorization code {code[:10]}... not found in Redis (/token endpoint)")
|
||||
|
||||
# Fall back to in-memory storage if Redis unavailable or code not found
|
||||
if not stored_code_data and hasattr(oauth_callback, '_code_storage'):
|
||||
stored_code_data = oauth_callback._code_storage.get(code)
|
||||
if stored_code_data:
|
||||
# Clean up in-memory storage
|
||||
oauth_callback._code_storage.pop(code, None)
|
||||
logger.info(f"Retrieved authorization code {code[:10]}... from in-memory storage (/token endpoint)")
|
||||
|
||||
if not stored_code_data:
|
||||
logger.error(f"No stored data found for authorization code: {code}")
|
||||
@@ -388,13 +478,11 @@ async def token_endpoint(request: Request):
|
||||
content={"error": "invalid_grant", "error_description": "Authorization code not found or expired"}
|
||||
)
|
||||
|
||||
# Check if code is expired
|
||||
# Note: Redis TTL handles expiration automatically, but check for manual expiration for in-memory fallback
|
||||
import time
|
||||
if time.time() > stored_code_data.get("expires_at", 0):
|
||||
expires_at = stored_code_data.get("expires_at", 0)
|
||||
if expires_at and time.time() > expires_at:
|
||||
logger.error(f"Authorization code expired: {code}")
|
||||
# Clean up expired code
|
||||
if hasattr(oauth_callback, '_code_storage'):
|
||||
oauth_callback._code_storage.pop(code, None)
|
||||
return JSONResponse(
|
||||
status_code=400,
|
||||
content={"error": "invalid_grant", "error_description": "Authorization code expired"}
|
||||
@@ -405,7 +493,7 @@ async def token_endpoint(request: Request):
|
||||
|
||||
if real_jwt_token:
|
||||
logger.info("Returning real Clerk JWT token from /token endpoint")
|
||||
# Clean up used code
|
||||
# Note: Code already deleted from Redis, clean up in-memory fallback if used
|
||||
if hasattr(oauth_callback, '_code_storage'):
|
||||
oauth_callback._code_storage.pop(code, None)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user