import hashlib import json import logging import threading import time import uuid from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from constants import DEVICE_INVALID_MESSAGE from mcp_protocol import McpProtocolHandler from services.request_context import RequestContextParser from utils.rate_limiter import SimpleRateLimiter logger = logging.getLogger(__name__) def extract_client_ip(headers, client_address): if client_address and len(client_address) > 0: return str(client_address[0]) return '' class PublicMcpHttpHandler: def __init__(self, gateway_app, context_parser=None, rate_limiter=None): self.gateway_app = gateway_app self.context_parser = context_parser or RequestContextParser() self.rate_limiter = rate_limiter # Cache registered tool names at startup for rate-key validation; # unknown names fall back to the bare IP bucket, preventing bucket explosion self._known_tools = frozenset(gateway_app.registered_tool_names()) @staticmethod def _build_trace_request_id(headers): return 'rq_http_{0}'.format(uuid.uuid4().hex[:16]) @staticmethod def _client_request_id(headers): incoming = '' for key, value in (headers or {}).items(): if str(key).lower() == 'x-request-id': incoming = str(value or '').strip() break if not incoming: return '' return hashlib.sha256(incoming.encode('utf-8')).hexdigest()[:16] def _check_rate_limit( self, rate_key, method, trace_request_id, jsonrpc_id=None, tool_name='', *, log_identity=None ): """Returns an error response if rate limit exceeded, else None. rate_key — the session-based bucket key used for tools/call log_identity — optional string shown in warning logs (e.g. client_ip) """ if self.rate_limiter and rate_key and not self.rate_limiter.is_allowed(rate_key): logger.warning( "MCP public rate limit exceeded", extra={ 'request_id': trace_request_id, 'jsonrpc_id': jsonrpc_id, 'protocol_method': method, 'tool_code': tool_name, 'client_identity': log_identity or rate_key, 'protocol_code': -32029, 'diagnostic_reason': 'RATE_LIMIT_EXCEEDED', }, ) return McpProtocolHandler._error_response( jsonrpc_id, -32029, 'Rate limit exceeded. Please try again later.', trace_request_id, ) return None def handle_json_rpc( self, headers, message, client_ip='', trace_request_id='', ): request_id = message.get('id') if isinstance(message, dict) else None trace_request_id = str(trace_request_id or '').strip() \ or self._build_trace_request_id(headers) method = str((message or {}).get('method') or '').strip() message_params = (message or {}).get('params') or {} tool_name = str(message_params.get('name') or '').strip() \ if isinstance(message_params, dict) else '' try: if method == 'initialize': # initialize is a stateless handshake that only returns server metadata; # rate-limiting it would block clients from connecting at all, so we skip it. return McpProtocolHandler._success_response(request_id, { 'protocolVersion': McpProtocolHandler.protocol_version, 'capabilities': {'tools': {'listChanged': False}}, 'serverInfo': { 'name': McpProtocolHandler.server_name, 'version': McpProtocolHandler.server_version, }, }) if method == 'tools/list': context = self.context_parser.parse(headers or {}) if not context.has_session(): logger.warning( 'MCP public device session unavailable', extra={ 'request_id': trace_request_id, 'jsonrpc_id': request_id, 'protocol_method': method, 'tool_code': '', 'protocol_code': -32001, 'diagnostic_reason': 'GATEWAY_SESSION_NOT_FOUND', }, ) return McpProtocolHandler._error_response( request_id, -32001, DEVICE_INVALID_MESSAGE, trace_request_id, ) tools = [] for tool in self.gateway_app.list_tools( context.gateway_session_id, request_id=trace_request_id, ): normalized = dict(tool) if 'input_schema' in normalized: normalized['inputSchema'] = normalized.pop('input_schema') tools.append(normalized) return McpProtocolHandler._success_response(request_id, {'tools': tools}) if method == 'tools/call': if not isinstance(message_params, dict): raise ValueError('tool parameters must be an object') # Parse context first — tools/call always requires a valid session context = self.context_parser.parse(headers or {}) if not context.has_session(): logger.warning( 'MCP public device session unavailable', extra={ 'request_id': trace_request_id, 'jsonrpc_id': request_id, 'protocol_method': method, 'tool_code': tool_name, 'protocol_code': -32001, 'diagnostic_reason': 'GATEWAY_SESSION_NOT_FOUND', }, ) return McpProtocolHandler._error_response( request_id, -32001, DEVICE_INVALID_MESSAGE, trace_request_id, ) # Session-based rate limiting: each employee gets an independent quota per tool. # Unknown tool names fall back to the bare session bucket to prevent key explosion. session_id = context.gateway_session_id rate_key = '{0}:{1}'.format(session_id, tool_name) if tool_name in self._known_tools else session_id blocked = self._check_rate_limit( rate_key, method, trace_request_id, request_id, tool_name, log_identity=client_ip, ) if blocked: return blocked params = message_params result = self.gateway_app.call_tool( session_id, params.get('name'), params.get('arguments') or {}, request_id=trace_request_id, client_ip=client_ip, ) return McpProtocolHandler._tool_call_response( request_id, tool_name, result, ) return McpProtocolHandler._error_response( request_id, -32601, 'Method not found: {0}'.format(method), trace_request_id, ) except Exception as exc: if method == 'tools/call': diagnostic_reason = ( 'PARAM_VALIDATION_FAILED' if isinstance(exc, ValueError) else 'UNEXPECTED_EXCEPTION' ) logger.error( 'MCP public tool request failed', extra={ 'request_id': trace_request_id, 'jsonrpc_id': request_id, 'protocol_method': method, 'tool_code': tool_name, 'response_code': 'MCP_9001', 'diagnostic_reason': diagnostic_reason, 'exception_class': exc.__class__.__name__, }, ) return McpProtocolHandler._tool_exception_response( request_id, tool_name, exc, trace_request_id, ) logger.error( "MCP public request failed", extra={ 'request_id': trace_request_id, 'jsonrpc_id': request_id, 'protocol_method': method, 'tool_code': tool_name, 'protocol_code': -32000, 'diagnostic_reason': 'UNEXPECTED_EXCEPTION', 'exception_class': exc.__class__.__name__, }, ) return McpProtocolHandler._error_response( request_id, -32000, 'Gateway request failed. Please try again later.', trace_request_id, ) def create_http_handler(gateway_app, rate_limiter=None): rpc_handler = PublicMcpHttpHandler(gateway_app, rate_limiter=rate_limiter) class Handler(BaseHTTPRequestHandler): def log_message(self, format, *args): # Override to use Python logging instead of stderr logger.info(f"{self.address_string()} - {format % args}") def do_GET(self): if self.path == '/health': self._write_json({'ok': True}) return self.send_response(404) self.end_headers() def do_POST(self): request_headers = dict(self.headers.items()) client_ip = extract_client_ip(request_headers, self.client_address) trace_request_id = rpc_handler._build_trace_request_id(request_headers) client_request_id_hash = rpc_handler._client_request_id( request_headers ) length = int(self.headers.get('Content-Length') or '0') if self.path != '/mcp': logger.warning(f"[HTTP] 404: ip={client_ip}, path={self.path}") if length > 0: self.rfile.read(length) self.send_response(404) self.end_headers() return body = self.rfile.read(length).decode('utf-8-sig') try: message = json.loads(body) except json.JSONDecodeError as exc: logger.warning( 'MCP public invalid JSON', extra={ 'request_id': trace_request_id, 'jsonrpc_id': None, 'protocol_method': '', 'tool_code': '', 'protocol_code': -32700, 'diagnostic_reason': 'PARAM_VALIDATION_FAILED', 'exception_class': exc.__class__.__name__, 'client_identity': client_ip, }, ) self._write_json(McpProtocolHandler._error_response( None, -32700, 'Parse error', trace_request_id, )) return method = message.get('method', '') request_id = message.get('id') if message.get('id') is not None else '' logger.info( 'MCP public request', extra={ 'request_id': trace_request_id, 'jsonrpc_id': request_id, 'protocol_method': method, 'client_identity': client_ip, 'client_request_id_hash': client_request_id_hash, }, ) response = rpc_handler.handle_json_rpc( request_headers, message, client_ip, trace_request_id=trace_request_id, ) self._write_json(response) def _write_json(self, payload): raw = json.dumps(payload, ensure_ascii=False).encode('utf-8') self.send_response(200) self.send_header('Content-Type', 'application/json; charset=utf-8') self.send_header('Content-Length', str(len(raw))) self.end_headers() self.wfile.write(raw) return Handler def serve_public(gateway_app, host='0.0.0.0', port=8765, enable_rate_limit=True, rate_limit_max_requests=60, rate_limit_window_seconds=60): # Configure logging logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', datefmt='%Y-%m-%d %H:%M:%S' ) # Configure rate limiting rate_limiter = None if enable_rate_limit: rate_limiter = SimpleRateLimiter( max_requests=rate_limit_max_requests, window_seconds=rate_limit_window_seconds, ) logger.info(f"Rate limiting enabled: {rate_limit_max_requests} requests/{rate_limit_window_seconds}s per session and tool (tools/call only)") logger.info(f"Starting public MCP Gateway on {host}:{port}") server = ThreadingHTTPServer((host, int(port)), create_http_handler(gateway_app, rate_limiter)) # Schedule periodic cleanup to prevent unbounded memory growth in the rate limiter if rate_limiter is not None: def _cleanup_loop(): while True: time.sleep(300) # Use the actual window as max_age to avoid deleting entries still within the window rate_limiter.cleanup(max_age_seconds=rate_limit_window_seconds) t = threading.Thread(target=_cleanup_loop, daemon=True) t.start() try: server.serve_forever() except KeyboardInterrupt: logger.info("Shutting down public MCP Gateway") server.shutdown()