mcp/tests/test_public_server.py

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 = []

    def list_tools(self):
        return [{'name': 'query_order', 'description': 'query order', '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 {'code': 'MCP_0000', 'data': {'ok': True}}


class PublicMcpHttpHandlerTest(unittest.TestCase):
    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_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)
        handler.handle_json_rpc({}, {'jsonrpc': '2.0', 'id': 1, 'method': 'tools/list'}, client_ip='1.2.3.4')
        response = handler.handle_json_rpc({}, {'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({}, {'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()