mirror of
https://github.com/Av0idD4rk/packmate.git
synced 2026-10-23 14:31:01 +00:00
302 lines
18 KiB
Python
302 lines
18 KiB
Python
"""Real TCP + real MCP SDK + full payload ledger. Disposable Compose DB only.
|
|
|
|
The host runner restarts MCP when restart-request.json appears. The test does not
|
|
mount the Docker socket or give MCP container-control privileges.
|
|
"""
|
|
import asyncio
|
|
import base64
|
|
from contextlib import asynccontextmanager
|
|
import json
|
|
import math
|
|
import os
|
|
from pathlib import Path
|
|
import struct
|
|
import sys
|
|
import time
|
|
from types import SimpleNamespace
|
|
|
|
import httpx
|
|
import psycopg
|
|
from mcp import ClientSession
|
|
from mcp.client.streamable_http import streamable_http_client
|
|
|
|
sys.path.insert(0, "/stress")
|
|
from network_load import run as generate
|
|
from verify_ledger import expected_digest
|
|
|
|
API = "http://172.31.250.10:65000"
|
|
REPORTS = Path("/reports")
|
|
TOKEN = "mcp-integration-test-token-000000000000"
|
|
|
|
|
|
@asynccontextmanager
|
|
async def session():
|
|
async with httpx.AsyncClient(headers={"Authorization": "Bearer " + TOKEN}, timeout=15) as transport:
|
|
async with streamable_http_client("http://mcp:65002/mcp", http_client=transport) as (read, write, _):
|
|
async with ClientSession(read, write) as client:
|
|
await client.initialize()
|
|
yield client
|
|
|
|
|
|
async def call(client, name, args=None):
|
|
response = await client.call_tool(name, args or {})
|
|
assert not response.isError, (name, response.content)
|
|
assert response.structuredContent is not None, name
|
|
return response.structuredContent
|
|
|
|
|
|
async def finish(client, job):
|
|
async with asyncio.timeout(15):
|
|
while True:
|
|
status = await call(client, "get_job", {"job_id": job["job_id"]})
|
|
if status["status"] != "running":
|
|
assert status["status"] in ("completed", "partial"), status
|
|
return status
|
|
await asyncio.sleep(0.1)
|
|
|
|
|
|
async def main():
|
|
REPORTS.mkdir(exist_ok=True)
|
|
requested = int(os.getenv("MCP_STRESS_STREAMS", "50000"))
|
|
duration = float(os.getenv("MCP_STRESS_SECONDS", "120"))
|
|
assert requested >= 5000 and duration >= 40, "profile too small to test restart and overlap"
|
|
report = {"checks": [], "baseline_streams": 2000, "live_streams": requested}
|
|
db = await psycopg.AsyncConnection.connect("host=db dbname=packmate user=packmate password=network-stress", autocommit=True)
|
|
|
|
async def rows(sql, args=()):
|
|
return await (await db.execute(sql, args)).fetchall()
|
|
|
|
async def scalar(sql, args=()):
|
|
return (await rows(sql, args))[0][0]
|
|
|
|
async def wait_count(expected):
|
|
async with asyncio.timeout(60):
|
|
while await scalar("SELECT count(*) FROM stream") != expected:
|
|
await asyncio.sleep(0.5)
|
|
|
|
async def load(count, run_id, name, seconds=0):
|
|
result = await generate(SimpleNamespace(host="172.31.250.10", port=18080, streams=count,
|
|
frames=8, payload_bytes=1024, concurrency=128, run_id=run_id, shard=0,
|
|
stream_offset=0, report=str(REPORTS / (name + "-load.json")), abort_file=None,
|
|
minimum_duration_seconds=seconds))
|
|
assert result["successfulStreams"] == count and not result["sampleErrors"], result
|
|
return result
|
|
|
|
async def ledger(run_id, count, lower=0):
|
|
# Aggregate in PostgreSQL, transfer hashes, independently reconstruct expected bytes.
|
|
entries = await rows("""
|
|
WITH payloads AS (
|
|
SELECT s.id, string_agg(p.content, ''::bytea ORDER BY p.id) FILTER (WHERE p.incoming) AS inc,
|
|
string_agg(p.content, ''::bytea ORDER BY p.id) FILTER (WHERE NOT p.incoming) AS out,
|
|
bool_and(CASE WHEN p.incoming THEN p.destination_ip='172.31.250.10' AND p.source_ip<>'172.31.250.10'
|
|
ELSE p.source_ip='172.31.250.10' AND p.destination_ip<>'172.31.250.10' END) AS ips,
|
|
count(p.id)=s.packets_count AND sum(octet_length(p.content))=s.size_bytes
|
|
AND s.service_id=18080 AND s.protocol=0 AS metadata
|
|
FROM stream s JOIN packet p ON p.stream_id=s.id WHERE s.id > %s
|
|
GROUP BY s.id)
|
|
SELECT id, substring(inc FROM 5 FOR 24), octet_length(inc), md5(inc),
|
|
octet_length(out), md5(out), ips, metadata FROM payloads ORDER BY id
|
|
""", (lower,))
|
|
assert len(entries) == count, (len(entries), count)
|
|
seen = set()
|
|
for sid, marker, insize, indigest, outsize, outdigest, ips, metadata in entries:
|
|
run, connection, shard, sequence = struct.unpack("!QQII", marker)
|
|
assert run == run_id and shard == sequence == 0 and 0 <= connection < count, (sid, marker)
|
|
assert connection not in seen, ("duplicate", connection)
|
|
seen.add(connection)
|
|
assert insize == outsize == 8 * 1028 and ips and metadata, sid
|
|
assert indigest == outdigest == expected_digest(run_id, connection, 0, 8, 1024), sid
|
|
return {"verified_streams": len(seen), "bytes_per_direction": count * 8 * 1028}
|
|
|
|
async with httpx.AsyncClient(auth=("stress", "stress"), timeout=10) as rest:
|
|
for _ in range(90):
|
|
try:
|
|
if (await rest.get(API + "/api/analysis/overview")).status_code == 200:
|
|
break
|
|
except httpx.HTTPError:
|
|
pass
|
|
await asyncio.sleep(1)
|
|
else:
|
|
raise AssertionError("backend not ready")
|
|
assert await scalar("SELECT count(*) FROM stream") == 0, "requires empty disposable database"
|
|
async with session() as client:
|
|
settings = {"name": "MCP capacity echo", "http": False, "urldecode_http_requests": False, "merge_adjacent_packets": False}
|
|
await call(client, "create_service", {"port": 18080, "settings": settings})
|
|
assert not (await call(client, "create_service", {"port": 18080, "settings": settings}))["created"]
|
|
pattern = await call(client, "create_pattern", {"settings": {"name": "capacity marker",
|
|
"match": {"type": "substring", "value": "PMSTRESS"}, "service_port": 18080}})
|
|
pattern_id = pattern["pattern"]["id"]
|
|
report["baseline_load"] = await load(2000, 8123001, "baseline")
|
|
await wait_count(2000)
|
|
report["baseline_ledger"] = await ledger(8123001, 2000)
|
|
baseline_end = await scalar("SELECT max(id) FROM stream")
|
|
preserved = [r[0] for r in await rows("SELECT id FROM stream ORDER BY id LIMIT 10")]
|
|
assert len((await call(client, "set_favorites", {"stream_ids": preserved}))["updated"]) == 10
|
|
|
|
# Real scheduler: both slots full, reject a third, cancel active work,
|
|
# keep partial results, and independently hit both budget guards.
|
|
query = {"last_minutes": 60, "service_ports": [18080]}
|
|
first = await call(client, "aggregate_traffic", {"query": query})
|
|
second = await call(client, "aggregate_traffic", {"query": query})
|
|
busy = await client.call_tool("aggregate_traffic", {"query": query})
|
|
assert busy.isError and "BUSY" in str(busy.content), busy
|
|
await asyncio.sleep(0.2)
|
|
cancelled = await call(client, "cancel_job", {"job_id": first["job_id"]})
|
|
assert cancelled["status"] == "cancelled" and cancelled["scanned_packets"] > 0, cancelled
|
|
assert (await call(client, "get_job_results", {"job_id": first["job_id"]}))["items"]
|
|
timed = await finish(client, second)
|
|
assert timed["reason"] == "time_budget" and not timed["coverage_complete"], timed
|
|
byte_job = await finish(client, await call(client, "search_traffic", {"query": query | {
|
|
"match": {"type": "substring", "value": "NEVER_PRESENT_31337"}}}))
|
|
assert byte_job["reason"] == "byte_budget" and not byte_job["coverage_complete"], byte_job
|
|
report["budget_checks"] = {"cancelled_packets": cancelled["scanned_packets"],
|
|
"time_reason": timed["reason"], "byte_reason": byte_job["reason"]}
|
|
page1 = await finish(client, await call(client, "search_traffic", {"query": query, "limit": 20}))
|
|
resumed = page1["query"] | {"cursor": page1["next_scan_cursor"]}
|
|
page2 = await finish(client, await call(client, "search_traffic", {"query": resumed, "limit": 20}))
|
|
one = (await call(client, "get_job_results", {"job_id": page1["job_id"]}))["items"]
|
|
two = (await call(client, "get_job_results", {"job_id": page2["job_id"]}))["items"]
|
|
assert len(one) == len(two) == 20 and not ({r["packet_id"] for r in one} & {r["packet_id"] for r in two})
|
|
lost_job_id = page1["job_id"]
|
|
report["checks"].append("real MCP: concurrent job cap, cancellation, retained partial results, byte/time budgets and cursor resume")
|
|
for job in (first, second, byte_job, page1, page2):
|
|
await call(client, "delete_job", {"job_id": job["job_id"]})
|
|
lost_job_id = (await finish(client, await call(client, "search_traffic", {"query": query, "limit": 1})))["job_id"]
|
|
|
|
await db.execute("UPDATE stream SET start_timestamp=start_timestamp-7200000, end_timestamp=end_timestamp-7200000, persisted_at=persisted_at-7200000 WHERE NOT favorite")
|
|
done = asyncio.Event()
|
|
restarted = asyncio.Event()
|
|
latencies, drops, remaining, job_reasons = [], [], [], []
|
|
started = time.monotonic()
|
|
reconnects = 0
|
|
|
|
async def traffic():
|
|
try:
|
|
report["live_load"] = await load(requested, 8123002, "live", duration)
|
|
finally:
|
|
done.set()
|
|
|
|
async def monitor():
|
|
while not done.is_set():
|
|
before = time.monotonic()
|
|
response = await rest.get(API + "/api/pcap/metrics")
|
|
response.raise_for_status()
|
|
latencies.append((time.monotonic() - before) * 1000)
|
|
metric = response.json()
|
|
assert all(metric[k] == 0 for k in ("droppedPackets", "interfaceDroppedPackets", "applicationDroppedPackets")), metric
|
|
drops.append(metric)
|
|
remaining.append(await scalar("SELECT count(*) FROM stream WHERE id<=%s AND NOT favorite", (baseline_end,)))
|
|
await asyncio.sleep(0.4)
|
|
|
|
async def investigations():
|
|
nonlocal reconnects
|
|
while not done.is_set():
|
|
try:
|
|
async with session() as client:
|
|
while not done.is_set():
|
|
job = await finish(client, await call(client, "search_traffic", {"query": {
|
|
"last_minutes": 60, "match": {"value": "(a+)+$"}}, "limit": 1000}))
|
|
job_reasons.append(job["reason"])
|
|
await call(client, "delete_job", {"job_id": job["job_id"]})
|
|
except Exception:
|
|
# Only the deliberate restart window permits transport failures.
|
|
assert (REPORTS / "restart-request.json").exists() and not restarted.is_set()
|
|
reconnects += 1
|
|
await asyncio.sleep(0.5)
|
|
|
|
async def restart_check():
|
|
await asyncio.sleep(min(30, duration / 3))
|
|
(REPORTS / "restart-request.json").write_text(json.dumps({"elapsed_seconds": time.monotonic() - started}))
|
|
async with asyncio.timeout(30):
|
|
while not (REPORTS / "restart-done.json").exists():
|
|
await asyncio.sleep(0.2)
|
|
async with session() as client:
|
|
missing = await client.call_tool("get_job", {"job_id": lost_job_id})
|
|
assert missing.isError and "NOT_FOUND" in str(missing.content)
|
|
assert any(s["port"] == 18080 for s in (await call(client, "list_services"))["items"])
|
|
assert any(p["id"] == pattern_id for p in (await call(client, "list_patterns"))["items"])
|
|
assert (await call(client, "get_stream", {"stream_id": preserved[0]}))["favorite"]
|
|
# Allow investigation task to finish reconnecting before closing the window.
|
|
await asyncio.sleep(2)
|
|
restarted.set()
|
|
|
|
await asyncio.gather(traffic(), monitor(), investigations(), restart_check())
|
|
assert restarted.is_set() and any(0 < n < 1990 for n in remaining), remaining[:10]
|
|
await wait_count(requested + 10)
|
|
report["live_ledger"] = await ledger(8123002, requested, baseline_end)
|
|
assert await scalar("SELECT count(*) FROM stream WHERE id<=%s", (baseline_end,)) == 10
|
|
assert await scalar("SELECT count(*) FROM stream_found_patterns WHERE pattern_id=%s", (pattern_id,)) == requested + 10
|
|
assert await scalar("SELECT count(*) FROM packet p LEFT JOIN stream s ON s.id=p.stream_id WHERE s.id IS NULL") == 0
|
|
p99 = sorted(latencies)[math.ceil(len(latencies) * .99) - 1]
|
|
assert p99 < 500, p99
|
|
report.update(api_p99_ms=round(p99, 3), api_samples=len(latencies), metrics=drops[-1],
|
|
mcp_jobs=len(job_reasons), job_reasons=sorted({str(x) for x in job_reasons}), reconnect_attempts=reconnects,
|
|
cleanup_overlap=True, restart_during_traffic=True)
|
|
report["checks"].append("full TCP ledger, live markers, favorites and orphan checks pass across cleaner + MCP restart")
|
|
|
|
# Actual large binary traffic, merged as a single saved message per direction.
|
|
async with session() as client:
|
|
settings["merge_adjacent_packets"] = True
|
|
await call(client, "update_service", {"port": 18080, "settings": settings})
|
|
async def edge_payload(payload):
|
|
prior = await scalar("SELECT max(id) FROM stream")
|
|
reader, writer = await asyncio.open_connection("172.31.250.10", 18080)
|
|
writer.write(struct.pack("!I", len(payload)) + payload)
|
|
await writer.drain()
|
|
size = struct.unpack("!I", await reader.readexactly(4))[0]
|
|
assert await reader.readexactly(size) == payload
|
|
writer.close()
|
|
await writer.wait_closed()
|
|
async with asyncio.timeout(15):
|
|
while not (saved := await rows("SELECT id,start_timestamp FROM stream WHERE id>%s", (prior,))):
|
|
await asyncio.sleep(.2)
|
|
assert len(saved) == 1
|
|
sid, timestamp = saved[0]
|
|
return sid, {"from_ms": timestamp, "to_ms": timestamp + 1, "direction": "incoming"}
|
|
|
|
binary = bytes(range(256)) * 384
|
|
sid, query = await edge_payload(binary)
|
|
stream = await call(client, "get_stream", {"stream_id": sid})
|
|
packet = next(p for p in stream["packets"] if p["incoming"])
|
|
assert packet["size_bytes"] == len(binary) + 4
|
|
chunks, offset = [], 0
|
|
while offset is not None:
|
|
part = await call(client, "read_packet", {"packet_id": packet["packet_id"], "offset": offset, "length": 16384, "format": "base64"})
|
|
chunks.append(base64.b64decode(part["content"]))
|
|
offset = part["next_offset"]
|
|
assert b"".join(chunks) == struct.pack("!I", len(binary)) + binary
|
|
skipped = await finish(client, await call(client, "search_traffic", {"query": query | {"match": {"type": "hex", "value": "00010203"}}}))
|
|
assert skipped["skipped_oversize_packets"] == 1 and skipped["matched_packets"] == 0 and not skipped["coverage_complete"], skipped
|
|
report["checks"].append("96 KiB binary message: exact paged read; search reports oversized skip without false match")
|
|
sid, query = await edge_payload(b"x" * 40000)
|
|
# Legacy UI-created literals must receive the same bounded lookback as MCP patterns.
|
|
response = await rest.post(API + "/api/pattern/", json={"name": "legacy literal", "value": "x", "color": "#123456",
|
|
"searchType": "SUBSTRING", "searchTarget": "CONTENT", "directionType": "INPUT", "actionType": "FIND", "serviceId": 18080})
|
|
response.raise_for_status()
|
|
legacy = response.json()["id"]
|
|
await call(client, "set_pattern_enabled", {"pattern_id": legacy, "enabled": False})
|
|
for _ in range(2):
|
|
applied = await finish(client, await call(client, "start_lookback", {"pattern_id": legacy, "query": query}))
|
|
assert applied["applied_packets"] == 1, applied
|
|
assert await scalar("SELECT count(*) FROM found_pattern WHERE pattern_id=%s", (legacy,)) == 256
|
|
ip_pattern = await call(client, "create_pattern", {"settings": {"name": "visible IP regex", "match": {
|
|
"target": "destination_ip", "value": r"^172\.31\.250\.10$"}, "direction": "INPUT", "service_port": 18080}})
|
|
ipsid, _ = await edge_payload(b"ip-visibility")
|
|
visible = (await rest.post(API + f"/api/packet/{ipsid}", json={"pageSize": 20})).json()
|
|
assert any(any(m["patternId"] == ip_pattern["pattern"]["id"] for m in p["matches"]) for p in visible), visible
|
|
await call(client, "update_pattern", {"pattern_id": legacy, "name": "renamed legacy", "color": "#654321"})
|
|
await call(client, "delete_pattern", {"pattern_id": legacy})
|
|
assert not any(p["id"] == legacy for p in (await call(client, "list_patterns"))["items"])
|
|
assert any(p["id"] == legacy and p["deleted"] for p in (await call(client, "list_patterns", {"include_deleted": True}))["items"])
|
|
await call(client, "create_service", {"port": 18081, "settings": settings | {"name": "temporary"}})
|
|
await call(client, "delete_service", {"port": 18081})
|
|
assert not any(s["port"] == 18081 for s in (await call(client, "list_services"))["items"])
|
|
report["checks"].append("legacy literal lookback capped/idempotent; IP regex visible in UI; service/pattern mutations persist")
|
|
await db.close()
|
|
(REPORTS / "capacity.json").write_text(json.dumps(report, indent=2))
|
|
print(json.dumps(report, indent=2), flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|