mirror of
https://github.com/Av0idD4rk/packmate.git
synced 2026-10-23 14:31:01 +00:00
41 lines
1.2 KiB
Python
41 lines
1.2 KiB
Python
import asyncio
|
|
import os
|
|
import struct
|
|
|
|
|
|
HOST = os.getenv("STRESS_TARGET_HOST", "0.0.0.0")
|
|
PORT = int(os.getenv("STRESS_TARGET_PORT", "18080"))
|
|
MAX_FRAME_BYTES = 1_048_576
|
|
|
|
|
|
async def handle(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
|
|
try:
|
|
while True:
|
|
header = await reader.readexactly(4)
|
|
size = struct.unpack("!I", header)[0]
|
|
if size <= 0 or size > MAX_FRAME_BYTES:
|
|
raise ValueError(f"invalid frame size: {size}")
|
|
payload = await reader.readexactly(size)
|
|
writer.write(struct.pack("!I", len(payload)) + payload)
|
|
await writer.drain()
|
|
except (asyncio.IncompleteReadError, ConnectionError, ValueError):
|
|
pass
|
|
finally:
|
|
writer.close()
|
|
try:
|
|
await writer.wait_closed()
|
|
except ConnectionError:
|
|
pass
|
|
|
|
|
|
async def main() -> None:
|
|
server = await asyncio.start_server(handle, HOST, PORT, backlog=4096, limit=2 * MAX_FRAME_BYTES)
|
|
addresses = ", ".join(str(sock.getsockname()) for sock in server.sockets or [])
|
|
print(f"TCP stress target listening on {addresses}", flush=True)
|
|
async with server:
|
|
await server.serve_forever()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|