async def lavalink_mitm():
await websocket.accept()
headers = websocket.headers
authorization = headers.get("authorization", None)
num_shards = headers.get("num-shards", None)
user_id = headers.get("user-id", None)
print(headers)
if authorization and \
num_shards and \
user_id:
async with aiohttp.ClientSession() as _session:
async with _session.ws_connect(lavalink_address,
headers={
"Authorization": authorization,
"Num-Shards": num_shards,
"User-Id": user_id
}) as lavalink_ws:
loop = asyncio.get_event_loop()
try:
await asyncio.gather(
copy_current_websocket_context(mitm_sender)(
lavalink_ws, authorization, num_shards, user_id),
copy_current_websocket_context(mitm_receiver)(
lavalink_ws, authorization, num_shards, user_id))
except Exception as e:
rockutils.prefix_print(str(e),
prefix="IPC Slave",
prefix_colour="light red",
text_colour="red")
else:
return abort(403)
async def ipc_slave(cluster, auth):
if auth != config['ipc']['auth_key']:
return "Invalid authentication", 403
else:
await websocket.accept()
cluster = str(cluster).lower()
ipc_jobs[cluster] = []
clusters_initialized.add(cluster)
if not cluster in cluster_jobs:
cluster_jobs[cluster] = []
rockutils.prefix_print(f"Connected to cluster {cluster}")
last_ping = time.time()
loop = asyncio.get_event_loop()
try:
await asyncio.gather(
copy_current_websocket_context(sending)(cluster),
copy_current_websocket_context(receiving)(cluster))
except Exception as e:
rockutils.prefix_print(str(e),
prefix="IPC Slave",
prefix_colour="light red",
text_colour="red")
def test_copy_current_websocket_context_error() -> None:
with pytest.raises(RuntimeError):
copy_current_websocket_context(lambda: None)()