117 lines
5.4 KiB
Python
117 lines
5.4 KiB
Python
"""Authentication decisions and the per-client concurrency limit."""
|
|
|
|
import asyncio
|
|
import unittest
|
|
|
|
from mcp_calculator.auth import LOCAL, Failure, Identity, authenticate, extract_token
|
|
from mcp_calculator.clients import Client, Snapshot, hash_token
|
|
from mcp_calculator.executor import ComputationError
|
|
from mcp_calculator.ratelimit import ClientConcurrencyLimiter
|
|
|
|
TOKEN_A = "mcpc_team-a-token-000000000000000000000000"
|
|
TOKEN_SVC = "mcpc_service-token-00000000000000000000000"
|
|
TOKEN_OFF = "mcpc_disabled-token-0000000000000000000000"
|
|
TOKEN_OLD = "mcpc_expired-token-00000000000000000000000"
|
|
|
|
|
|
def snapshot() -> Snapshot:
|
|
clients = {
|
|
"team-a": Client.from_dict("team-a", {"token_sha256": hash_token(TOKEN_A), "max_concurrent": 3}),
|
|
"reporting": Client.from_dict("reporting", {"token_sha256": hash_token(TOKEN_SVC), "allowed_networks": ["10.130.4.0/24"]}),
|
|
"off": Client.from_dict("off", {"token_sha256": hash_token(TOKEN_OFF), "enabled": False}),
|
|
"old": Client.from_dict("old", {"token_sha256": hash_token(TOKEN_OLD), "expires": "2000-01-01"}),
|
|
"batch-net": Client.from_dict("batch-net", {"auth": "network", "allowed_networks": ["10.131.0.0/16"]}),
|
|
"batch-host": Client.from_dict("batch-host", {"auth": "network", "allowed_networks": ["10.131.2.15/32"], "max_concurrent": 1}),
|
|
}
|
|
return Snapshot.build(clients)
|
|
|
|
|
|
class AuthenticateTest(unittest.TestCase):
|
|
def setUp(self):
|
|
self.snapshot = snapshot()
|
|
|
|
def test_valid_token(self):
|
|
self.assertEqual(authenticate(self.snapshot, TOKEN_A, "203.0.113.1"), Identity("team-a", "token", 3))
|
|
|
|
def test_failures(self):
|
|
cases = [
|
|
(None, "203.0.113.1", "missing_credentials"),
|
|
("mcpc_unknown", "203.0.113.1", "invalid_token"),
|
|
(TOKEN_OFF, "203.0.113.1", "disabled"),
|
|
(TOKEN_OLD, "203.0.113.1", "expired"),
|
|
(TOKEN_SVC, "10.130.5.1", "network_denied"),
|
|
(TOKEN_SVC, None, "network_denied"),
|
|
]
|
|
for token, ip, reason in cases:
|
|
with self.subTest(reason=reason):
|
|
result = authenticate(self.snapshot, token, ip)
|
|
self.assertIsInstance(result, Failure)
|
|
self.assertEqual(result.reason, reason)
|
|
|
|
def test_http_status_codes(self):
|
|
self.assertEqual(Failure("missing_credentials").status, 401)
|
|
self.assertEqual(Failure("invalid_token").status, 401)
|
|
self.assertEqual(Failure("disabled").status, 403)
|
|
self.assertEqual(Failure("network_denied").status, 403)
|
|
|
|
def test_token_restricted_to_network(self):
|
|
self.assertEqual(authenticate(self.snapshot, TOKEN_SVC, "10.130.4.77").name, "reporting")
|
|
|
|
def test_network_client_without_token(self):
|
|
self.assertEqual(authenticate(self.snapshot, None, "10.131.9.9"), Identity("batch-net", "network", None))
|
|
|
|
def test_most_specific_network_wins(self):
|
|
self.assertEqual(authenticate(self.snapshot, None, "10.131.2.15"), Identity("batch-host", "network", 1))
|
|
|
|
def test_invalid_token_is_rejected_even_from_trusted_network(self):
|
|
self.assertEqual(authenticate(self.snapshot, "mcpc_wrong", "10.131.2.15").reason, "invalid_token")
|
|
|
|
|
|
class ExtractTokenTest(unittest.TestCase):
|
|
def test_sources(self):
|
|
self.assertEqual(extract_token([(b"authorization", b"Bearer abc")]), "abc")
|
|
self.assertEqual(extract_token([(b"authorization", b"bearer abc ")]), "abc")
|
|
self.assertEqual(extract_token([(b"x-api-key", b"xyz")]), "xyz")
|
|
self.assertIsNone(extract_token([(b"authorization", b"Basic dXNlcjpwYXNz")]))
|
|
self.assertIsNone(extract_token([]))
|
|
|
|
|
|
class ClientConcurrencyTest(unittest.IsolatedAsyncioTestCase):
|
|
async def test_limit_per_client(self):
|
|
limiter = ClientConcurrencyLimiter(default_limit=2)
|
|
client = Identity("team-a", "token", None)
|
|
async with limiter.slot(client), limiter.slot(client):
|
|
self.assertEqual(limiter.active("team-a"), 2)
|
|
with self.assertRaisesRegex(ComputationError, "client busy") as ctx:
|
|
async with limiter.slot(client):
|
|
pass
|
|
self.assertEqual(ctx.exception.kind, "busy")
|
|
# other clients are not affected
|
|
async with limiter.slot(Identity("team-b", "token", None)):
|
|
pass
|
|
self.assertEqual(limiter.active("team-a"), 0)
|
|
|
|
async def test_client_specific_limit_and_unlimited(self):
|
|
limiter = ClientConcurrencyLimiter(default_limit=1)
|
|
self.assertEqual(limiter.limit_for(Identity("x", "token", 5)), 5)
|
|
self.assertEqual(limiter.limit_for(Identity("x", "token", 0)), 0)
|
|
self.assertEqual(limiter.limit_for(LOCAL), 0, "stdio is never limited")
|
|
async with limiter.slot(Identity("x", "token", 0)), limiter.slot(Identity("x", "token", 0)):
|
|
pass
|
|
|
|
async def test_slot_is_released_on_error_and_cancel(self):
|
|
limiter = ClientConcurrencyLimiter(default_limit=1)
|
|
client = Identity("team-a", "token", None)
|
|
with self.assertRaises(ValueError):
|
|
async with limiter.slot(client):
|
|
raise ValueError
|
|
task = asyncio.create_task(self._hold(limiter, client))
|
|
await asyncio.sleep(0.01)
|
|
task.cancel()
|
|
with self.assertRaises(asyncio.CancelledError):
|
|
await task
|
|
self.assertEqual(limiter.active("team-a"), 0)
|
|
|
|
async def _hold(self, limiter, client):
|
|
async with limiter.slot(client):
|
|
await asyncio.sleep(10)
|