forked from Hyperion/HyperionOS
117 lines
3.7 KiB
Python
117 lines
3.7 KiB
Python
import asyncio
|
|
import struct
|
|
import ipaddress
|
|
import websockets
|
|
import random
|
|
import time
|
|
|
|
SERVER = "ws://localhost:8765"
|
|
CLIENT_COUNT = 5000
|
|
SEND_INTERVAL = (0.2, 0.5) # safer backpressure
|
|
|
|
IP_PACKET_HEADER = "!4s4sI"
|
|
HEADER_SIZE = struct.calcsize(IP_PACKET_HEADER)
|
|
|
|
# Opcodes
|
|
OP_IP_REQUEST = 0x01
|
|
OP_IP_ASSIGN = 0x02
|
|
OP_DATA = 0x10
|
|
ASSIGN_OK = 0x00
|
|
|
|
STATIC_IP_BASE = ipaddress.IPv4Address("10.0.0.10")
|
|
assigned_ips: set[bytes] = set()
|
|
assigned_ips_lock = asyncio.Lock()
|
|
|
|
|
|
def random_payload(my_ip: bytes) -> bytes:
|
|
return f"hello from {ipaddress.IPv4Address(my_ip)} @ {time.time():.2f}".encode()
|
|
|
|
|
|
async def virtual_node(node_id: int):
|
|
my_ip = None
|
|
try:
|
|
async with websockets.connect(SERVER, max_size=2**20) as ws:
|
|
# -------------------------
|
|
# CONTROL PLANE: IP REQUEST
|
|
# -------------------------
|
|
role = random.random()
|
|
if role < 0.70:
|
|
await ws.send(bytes([OP_IP_REQUEST, 0x00]))
|
|
elif role < 0.90:
|
|
static_ip = (STATIC_IP_BASE + node_id).packed
|
|
await ws.send(bytes([OP_IP_REQUEST, 0x01]) + static_ip)
|
|
elif role < 0.95:
|
|
bad_ip = ipaddress.IPv4Address(random.getrandbits(32)).packed
|
|
await ws.send(bytes([OP_IP_REQUEST, 0x01]) + bad_ip)
|
|
else:
|
|
fake_packet = bytes([OP_DATA]) + b"\x00" * (HEADER_SIZE + 10)
|
|
await ws.send(fake_packet)
|
|
return
|
|
|
|
reply = await ws.recv()
|
|
if not isinstance(reply, bytes) or len(reply) != 6:
|
|
return
|
|
if reply[0] != OP_IP_ASSIGN or reply[1] != ASSIGN_OK:
|
|
return
|
|
|
|
my_ip = reply[2:6]
|
|
|
|
async with assigned_ips_lock:
|
|
assigned_ips.add(my_ip)
|
|
|
|
# -------------------------
|
|
# RECEIVER
|
|
# -------------------------
|
|
async def receiver():
|
|
try:
|
|
async for msg in ws:
|
|
if not isinstance(msg, bytes):
|
|
continue
|
|
if msg[0] != OP_DATA:
|
|
continue
|
|
if len(msg) < 1 + HEADER_SIZE:
|
|
continue
|
|
src, dst, length = struct.unpack(
|
|
IP_PACKET_HEADER, msg[1:1 + HEADER_SIZE]
|
|
)
|
|
payload = msg[1 + HEADER_SIZE:]
|
|
if len(payload) != length:
|
|
continue
|
|
except websockets.ConnectionClosed:
|
|
pass
|
|
|
|
# -------------------------
|
|
# SENDER
|
|
# -------------------------
|
|
async def sender():
|
|
while True:
|
|
await asyncio.sleep(random.uniform(*SEND_INTERVAL))
|
|
async with assigned_ips_lock:
|
|
if not assigned_ips:
|
|
continue
|
|
dst = random.choice(list(assigned_ips))
|
|
if dst == my_ip:
|
|
continue
|
|
payload = random_payload(my_ip)
|
|
packet = (
|
|
bytes([OP_DATA])
|
|
+ struct.pack(IP_PACKET_HEADER, my_ip, dst, len(payload))
|
|
+ payload
|
|
)
|
|
await ws.send(packet)
|
|
|
|
await asyncio.gather(receiver(), sender())
|
|
|
|
finally:
|
|
async with assigned_ips_lock:
|
|
assigned_ips.discard(my_ip)
|
|
|
|
|
|
async def main():
|
|
print(f"🚀 Starting {CLIENT_COUNT} virtual nodes")
|
|
tasks = [asyncio.create_task(virtual_node(i)) for i in range(CLIENT_COUNT)]
|
|
await asyncio.gather(*tasks)
|
|
|
|
|
|
asyncio.run(main())
|