Spaces:
Paused
Paused
File size: 2,027 Bytes
86ba74e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | import asyncio
import os
import websockets
import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# The remote HF Space websocket URL (e.g., wss://augment17-redis.hf.space/ws)
REMOTE_WS_URL = os.getenv("REMOTE_REDIS_WS_URL")
LOCAL_PORT = int(os.getenv("LOCAL_REDIS_PORT", 6379))
async def handle_client(reader, writer):
if not REMOTE_WS_URL:
logger.error("REMOTE_REDIS_WS_URL is not set!")
writer.close()
return
logger.info("New local TCP connection. Connecting to remote WebSocket...")
try:
async with websockets.connect(REMOTE_WS_URL) as ws:
logger.info("Connected to remote WebSocket.")
async def pipe_tcp_to_ws():
try:
while True:
data = await reader.read(4096)
if not data:
break
await ws.send(data)
except Exception as e:
logger.info(f"tcp_to_ws closed: {e}")
async def pipe_ws_to_tcp():
try:
while True:
data = await ws.recv()
writer.write(data)
await writer.drain()
except Exception as e:
logger.info(f"ws_to_tcp closed: {e}")
await asyncio.gather(
pipe_tcp_to_ws(),
pipe_ws_to_tcp(),
return_exceptions=True
)
except Exception as e:
logger.error(f"WebSocket connection failed: {e}")
finally:
writer.close()
logger.info("Local TCP connection closed.")
async def main():
server = await asyncio.start_server(handle_client, '127.0.0.1', LOCAL_PORT)
addr = server.sockets[0].getsockname()
logger.info(f"Local WebSocket-to-TCP proxy listening on {addr}")
async with server:
await server.serve_forever()
if __name__ == '__main__':
asyncio.run(main())
|