import unittest from public_server import PublicMcpHttpHandler, extract_client_ip from utils.rate_limiter import SimpleRateLimiter class FakeContext: def __init__(self, gateway_session_id): self.gateway_session_id = gateway_session_id def has_session(self): return bool(self.gateway_session_id) class FakeParser: def parse(self, headers): return FakeContext(headers.get('X-Gateway-Session', '')) class FakeGateway: def __init__(self): self.calls = [] self.list_calls = [] self.tool_result = {'code': 'MCP_0000', 'data': {'ok': True}} def registered_tool_names(self): return ('query_order', 'query_track') def list_tools(self, gateway_session_id): self.list_calls.append(gateway_session_id) return [{'name': 'query_track', 'description': 'query track', 'input_schema': {'type': 'object'}}] def call_tool(self, gateway_session_id, name, arguments=None, request_id='', client_ip=''): self.calls.append((gateway_session_id, name, arguments, request_id, client_ip)) return self.tool_result class PublicMcpHttpHandlerTest(unittest.TestCase): def test_constructor_uses_registered_names_without_loading_dynamic_list(self): gateway = FakeGateway() PublicMcpHttpHandler(gateway, context_parser=FakeParser()) self.assertEqual([], gateway.list_calls) def test_handle_tools_list_passes_session_and_returns_dynamic_tools(self): gateway = FakeGateway() handler = PublicMcpHttpHandler(gateway, context_parser=FakeParser()) response = handler.handle_json_rpc( headers={'X-Gateway-Session': 'GWS_A'}, message={ 'jsonrpc': '2.0', 'id': 2, 'method': 'tools/list', 'params': {}, }, client_ip='10.0.0.5', ) self.assertEqual(['GWS_A'], gateway.list_calls) self.assertEqual( ['query_track'], [tool['name'] for tool in response['result']['tools']], ) self.assertIn('inputSchema', response['result']['tools'][0]) def test_missing_session_on_tools_list_returns_protocol_error(self): gateway = FakeGateway() handler = PublicMcpHttpHandler(gateway, context_parser=FakeParser()) response = handler.handle_json_rpc( headers={}, message={ 'jsonrpc': '2.0', 'id': 3, 'method': 'tools/list', 'params': {}, }, client_ip='10.0.0.5', ) self.assertEqual(-32000, response['error']['code']) self.assertIn('Workbuddy', response['error']['message']) self.assertEqual([], gateway.list_calls) def test_handle_tools_call_passes_gateway_session_to_public_gateway(self): gateway = FakeGateway() handler = PublicMcpHttpHandler(gateway, context_parser=FakeParser()) response = handler.handle_json_rpc( headers={'X-Gateway-Session': 'GWS_A'}, message={ 'jsonrpc': '2.0', 'id': 1, 'method': 'tools/call', 'params': { 'name': 'query_order', 'arguments': {'keyword': 'USC'}, }, }, client_ip='10.0.0.5' ) self.assertEqual(False, response['result']['isError']) self.assertEqual('GWS_A', gateway.calls[0][0]) self.assertEqual('query_order', gateway.calls[0][1]) self.assertEqual('10.0.0.5', gateway.calls[0][4]) def test_backend_business_error_returns_tool_error_result(self): gateway = FakeGateway() gateway.tool_result = { 'code': 'MCP_1401', 'msg': 'invalid exact order query conditions', 'data': [], 'meta': {'request_id': 'rq_backend'}, } handler = PublicMcpHttpHandler(gateway, context_parser=FakeParser()) response = handler.handle_json_rpc( headers={'X-Gateway-Session': 'GWS_A'}, message={ 'jsonrpc': '2.0', 'id': 7, 'method': 'tools/call', 'params': { 'name': 'query_order', 'arguments': {'keyword': 'USC'}, }, }, client_ip='10.0.0.5', ) self.assertNotIn('error', response) self.assertTrue(response['result']['isError']) self.assertIn('MCP_1401', response['result']['content'][0]['text']) self.assertEqual( { 'code': 'MCP_1401', 'msg': 'invalid exact order query conditions', 'meta': {'request_id': 'rq_backend'}, }, response['result']['structuredContent'], ) def test_missing_session_on_tool_call_returns_error_content(self): gateway = FakeGateway() handler = PublicMcpHttpHandler(gateway, context_parser=FakeParser()) response = handler.handle_json_rpc( headers={}, message={ 'jsonrpc': '2.0', 'id': 1, 'method': 'tools/call', 'params': {'name': 'query_order', 'arguments': {'keyword': 'USC'}}, }, ) self.assertTrue(response['result']['isError']) self.assertIn('这台设备的 Workbuddy 配置已失效，请重新生成配置', response['result']['content'][0]['text']) def test_extract_client_ip_ignores_spoofable_forwarded_for_header(self): client_ip = extract_client_ip( {'X-Forwarded-For': '203.0.113.9'}, ('10.0.0.5', 54321), ) self.assertEqual('10.0.0.5', client_ip) class RateLimitTest(unittest.TestCase): def _make_handler(self, max_requests=2): gateway = FakeGateway() limiter = SimpleRateLimiter(max_requests=max_requests, window_seconds=60) return PublicMcpHttpHandler(gateway, context_parser=FakeParser(), rate_limiter=limiter) def _tools_call_msg(self, tool_name='query_order'): return { 'jsonrpc': '2.0', 'id': 1, 'method': 'tools/call', 'params': {'name': tool_name, 'arguments': {'keyword': 'test'}}, } # --- tools/call per-tool独立限流 --- def test_known_tool_rate_limited_after_quota_exhausted(self): handler = self._make_handler(max_requests=1) handler.handle_json_rpc({'X-Gateway-Session': 'GWS_A'}, self._tools_call_msg(), client_ip='1.2.3.4') response = handler.handle_json_rpc({'X-Gateway-Session': 'GWS_A'}, self._tools_call_msg(), client_ip='1.2.3.4') self.assertIn('error', response) self.assertIn('Rate limit', response['error']['message']) def test_two_known_tools_have_independent_quotas(self): # query_order 限流不影响 query_track handler = self._make_handler(max_requests=1) handler.handle_json_rpc({'X-Gateway-Session': 'GWS_A'}, self._tools_call_msg('query_order'), client_ip='1.2.3.4') # query_order quota exhausted, query_track should still work response = handler.handle_json_rpc({'X-Gateway-Session': 'GWS_A'}, self._tools_call_msg('query_track'), client_ip='1.2.3.4') self.assertNotIn('error', response) def test_unknown_tool_name_uses_shared_session_bucket_not_new_bucket(self): # 未知工具名应归入 session_id 桶，同一会话的不同未知工具共享同一配额，不能无限新建桶 handler = self._make_handler(max_requests=1) # 先用 session 桶打一次（用未知工具名触发） handler.handle_json_rpc({'X-Gateway-Session': 'GWS_A'}, self._tools_call_msg('fake_tool_1'), client_ip='1.2.3.4') # 再用同会话的另一个未知工具名，应该命中同一个 session 桶，被限流 response = handler.handle_json_rpc({'X-Gateway-Session': 'GWS_A'}, self._tools_call_msg('fake_tool_2'), client_ip='1.2.3.4') self.assertIn('error', response) self.assertIn('Rate limit', response['error']['message']) def test_different_sessions_have_independent_quotas(self): # 不同会话（不同员工）即使来自同一IP，也拥有独立的配额 handler = self._make_handler(max_requests=1) handler.handle_json_rpc({'X-Gateway-Session': 'GWS_A'}, self._tools_call_msg(), client_ip='1.1.1.1') # 不同 session，即使同一 IP，quota 独立 → 应该允许 response = handler.handle_json_rpc({'X-Gateway-Session': 'GWS_B'}, self._tools_call_msg(), client_ip='1.1.1.1') self.assertNotIn('error', response) def test_same_session_from_different_ips_shares_quota(self): # 同一会话从不同IP发来（如移动网络切换），仍共享同一会话配额 handler = self._make_handler(max_requests=1) handler.handle_json_rpc({'X-Gateway-Session': 'GWS_A'}, self._tools_call_msg(), client_ip='1.1.1.1') # 同 session，不同 IP → 命中同一 session 桶，被限流 response = handler.handle_json_rpc({'X-Gateway-Session': 'GWS_A'}, self._tools_call_msg(), client_ip='2.2.2.2') self.assertIn('error', response) self.assertIn('Rate limit', response['error']['message']) # --- initialize 不受限流，tools/list 有独立 bucket --- def test_initialize_is_never_rate_limited(self): # initialize 是握手协议，无论请求多少次都不应被限流 handler = self._make_handler(max_requests=1) for i in range(5): response = handler.handle_json_rpc({}, {'jsonrpc': '2.0', 'id': i, 'method': 'initialize'}, client_ip='1.2.3.4') self.assertNotIn('error', response, f'initialize should never be rate limited (attempt {i})') self.assertIn('protocolVersion', response['result']) def test_tools_list_is_rate_limited(self): handler = self._make_handler(max_requests=1) headers = {'X-Gateway-Session': 'GWS_A'} handler.handle_json_rpc(headers, {'jsonrpc': '2.0', 'id': 1, 'method': 'tools/list'}, client_ip='1.2.3.4') response = handler.handle_json_rpc(headers, {'jsonrpc': '2.0', 'id': 2, 'method': 'tools/list'}, client_ip='1.2.3.4') self.assertIn('error', response) self.assertIn('Rate limit', response['error']['message']) def test_tools_list_quota_independent_from_tools_call_quota(self): # tools/list 和 tools/call 使用不同 bucket，互不影响 handler = self._make_handler(max_requests=1) # 用掉 tools/list 的1个 slot handler.handle_json_rpc({'X-Gateway-Session': 'GWS_A'}, {'jsonrpc': '2.0', 'id': 1, 'method': 'tools/list'}, client_ip='1.2.3.4') # tools/call 应仍可正常执行（使用独立的 client_ip:tool_name bucket） response = handler.handle_json_rpc( {'X-Gateway-Session': 'GWS_A'}, {'jsonrpc': '2.0', 'id': 2, 'method': 'tools/call', 'params': {'name': 'query_order', 'arguments': {'keyword': 'test'}}}, client_ip='1.2.3.4', ) self.assertNotIn('error', response) if __name__ == '__main__': unittest.main()