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)()