客户端

.gitignore

@@ -0,0 +1,3 @@
+.env
+venv/
+node_modules/

README.md

@@ -0,0 +1,36 @@
+# Client agent demo
+
+This is a minimal client agent that connects to the control server WebSocket and forwards TCP connections from a local port to the server using a simple message framing protocol. This demo is for prototyping only and not production-ready.
+
+## Features
+- Connects to control server WebSocket (`/ws`) and performs handshake
+- Supports a simple multiplexing prototype: multiple local connections are assigned short connection IDs and multiplexed over a single WebSocket
+- Can forward local TCP connections to the server using a framed protocol
+- Offers an optional simple HTTP page response for browser tests (not full HTTP proxy)
+
+## Running locally
+
+1. Create virtual environment and install dependencies:
+
+```bash
+cd client
+python -m venv .venv
+source .venv/bin/activate
+pip install -r requirements.txt
+```
+
+2. Run the client pointing to local server:
+
+```bash
+python client.py --server ws://localhost:8000/ws --local-port 9000
+```
+
+3. Test with a simple TCP client (keeps connection open):
+
+```bash
+python test_local_tcp.py
+```
+
+4. For browser test: open `http://localhost:9000` — the client will return a simple demo page (prototype only).
+
+Note: For remote HF Space server, use `wss://<your-space-host>/ws` as the `--server` URL.

client.py

@@ -0,0 +1,250 @@
+import argparse
+import asyncio
+import struct
+import websockets
+import socket
+import uuid
+
+# Simple length-prefixed framing
+# frame: 4-byte big-endian length + payload
+
+async def websocket_send_frame(ws, data: bytes, frame_type: bytes = b'D'):  # frame_type: b'D' for data, b'O' for open, etc.
+    # frame format: 1-byte type + 4-byte length + payload
+    header = frame_type + struct.pack('>I', len(data))
+    await ws.send(header + data)
+
+async def websocket_recv_frame(ws, recv_lock: asyncio.Lock):
+    # Serialize calls to ws.recv using recv_lock to avoid concurrent recv errors
+    async with recv_lock:
+        data = await ws.recv()
+    if isinstance(data, str):
+        data = data.encode()
+    if len(data) < 5:
+        raise ValueError('frame too short')
+    frame_type = data[0:1]
+    length = struct.unpack('>I', data[1:5])[0]
+    payload = data[5:5+length]
+    return frame_type, payload
+
+async def handle_local_connection(local_reader, local_writer, ws, recv_lock: asyncio.Lock):
+    peername = None
+    try:
+        sock = local_writer.get_extra_info('socket')
+        if sock:
+            peername = sock.getpeername()
+    except Exception:
+        peername = None
+
+    print(f'New local connection from {peername}')
+
+    async def read_local_then_send():
+        try:
+            while True:
+                chunk = await local_reader.read(4096)
+                if not chunk:
+                    # EOF
+                    print('local EOF, sending close frame')
+                    await websocket_send_frame(ws, b'__CLOSE__')
+                    break
+                await websocket_send_frame(ws, chunk)
+        except asyncio.CancelledError:
+            pass
+        except Exception as e:
+            print('local read/send error', e)
+
+    async def recv_then_write_local():
+        try:
+            while True:
+                try:
+                    frame_type, payload = await websocket_recv_frame(ws, recv_lock)
+                except Exception as e:
+                    print('ws recv/local write error', e)
+                    break
+                if payload == b'__CLOSE__':
+                    print('received close frame from server')
+                    break
+                try:
+                    local_writer.write(payload)
+                    await local_writer.drain()
+                except Exception as e:
+                    print('local writer error while writing payload', e)
+                    break
+        finally:
+            try:
+                local_writer.close()
+                await local_writer.wait_closed()
+            except Exception:
+                pass
+
+    # run both tasks and ensure exceptions are handled
+    task_send = asyncio.create_task(read_local_then_send())
+    task_recv = asyncio.create_task(recv_then_write_local())
+    done, pending = await asyncio.wait({task_send, task_recv}, return_when=asyncio.FIRST_EXCEPTION)
+    for t in pending:
+        t.cancel()
+    for t in done:
+        if t.exception():
+            print('task exception in handle_local_connection:', t.exception())
+    print(f'Connection from {peername} closed')
+
+
+async def handle_local_connection_multiplex(local_reader, local_writer, ws, recv_lock: asyncio.Lock, conn_id: bytes, conn_queue: asyncio.Queue, conn_queues: dict):
+    # Similar to handle_local_connection but prefixes payload with conn_id
+    peername = None
+    try:
+        sock = local_writer.get_extra_info('socket')
+        if sock:
+            peername = sock.getpeername()
+    except Exception:
+        peername = None
+
+    print(f'New multiplexed connection {conn_id.decode()} from {peername}')
+
+    async def read_local_then_send():
+        try:
+            while True:
+                chunk = await local_reader.read(4096)
+                if not chunk:
+                    # EOF
+                    print(f'local EOF for {conn_id.decode()}, sending close frame')
+                    await websocket_send_frame(ws, conn_id + b'__CLOSE__', frame_type=b'D')
+                    break
+                await websocket_send_frame(ws, conn_id + chunk, frame_type=b'D')
+        except asyncio.CancelledError:
+            pass
+        except Exception as e:
+            print('local read/send error', e)
+
+    async def recv_then_write_local():
+        try:
+            while True:
+                try:
+                    payload = await conn_queue.get()
+                except Exception as e:
+                    print('conn_queue get error', e)
+                    break
+                if payload == b'__CLOSE__':
+                    print(f'received close for {conn_id.decode()}')
+                    break
+                try:
+                    local_writer.write(payload)
+                    await local_writer.drain()
+                except Exception as e:
+                    print('local writer error while writing payload', e)
+                    break
+        finally:
+            try:
+                local_writer.close()
+                await local_writer.wait_closed()
+            except Exception:
+                pass
+
+    task_send = asyncio.create_task(read_local_then_send())
+    task_recv = asyncio.create_task(recv_then_write_local())
+    done, pending = await asyncio.wait({task_send, task_recv}, return_when=asyncio.FIRST_EXCEPTION)
+    for t in pending:
+        t.cancel()
+    for t in done:
+        if t.exception():
+            print('task exception in handle_local_connection_multiplex:', t.exception())
+    # cleanup queue registration
+    try:
+        conn_queues.pop(conn_id, None)
+    except Exception:
+        pass
+    print(f'Multiplexed connection {conn_id.decode()} from {peername} closed')
+
+
+async def run_client(server_url: str, local_port: int):
+    async with websockets.connect(server_url) as ws:
+        print('Connected to server')
+        # create a lock to serialize ws.recv calls
+        recv_lock = asyncio.Lock()
+
+        # handshake using the same recv_lock
+        await ws.send('ping')
+        async with recv_lock:
+            resp = await ws.recv()
+        print('handshake response:', resp)
+
+        # multiplexing: assign a unique client session id
+        session_id = uuid.uuid4().hex[:8].encode()
+        print('session id:', session_id.decode())
+
+        # mapping of conn_id bytes -> asyncio.Queue for demuxing incoming frames
+        conn_queues: dict[bytes, asyncio.Queue] = {}
+
+        async def demuxer_loop():
+            try:
+                while True:
+                    frame_type, payload = await websocket_recv_frame(ws, recv_lock)
+                    # only handle data frames (type 'D') for demuxing
+                    if frame_type != b'D':
+                        # ignore other frame types for now
+                        continue
+                    if len(payload) < 8:
+                        print('demux: payload too short')
+                        continue
+                    cid = payload[:8]
+                    body = payload[8:]
+                    q = conn_queues.get(cid)
+                    if q:
+                        await q.put(body)
+                    else:
+                        print('demux: no queue for', cid)
+            except Exception as e:
+                print('demuxer exiting:', e)
+
+        demux_task = asyncio.create_task(demuxer_loop())
+
+        # start local server to accept connections and forward (each connection gets an ephemeral id)
+        async def accept_callback(r, w):
+            try:
+                data = await asyncio.wait_for(r.read(1024), timeout=1.5)
+            except Exception:
+                data = b''
+            # Browser quick response
+            if data.startswith((b'GET ', b'POST ', b'HEAD ')):
+                body = b"<html><body><h1>cpolar demo</h1><p>Connected</p></body></html>"
+                html = b"HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=utf-8\r\nContent-Length: " + str(len(body)).encode() + b"\r\n\r\n" + body
+                try:
+                    w.write(html)
+                    await w.drain()
+                except Exception:
+                    pass
+                try:
+                    w.close()
+                    await w.wait_closed()
+                except Exception:
+                    pass
+                return
+
+            # when a new local connection arrives, send an OPEN frame with connection id
+            conn_id = uuid.uuid4().hex[:8].encode()
+            q: asyncio.Queue = asyncio.Queue()
+            conn_queues[conn_id] = q
+            await websocket_send_frame(ws, conn_id, frame_type=b'O')
+            # First chunk (if any) is data, send it immediately as conn_id+data (if non-empty)
+            if data:
+                await websocket_send_frame(ws, conn_id + data, frame_type=b'D')
+            # start handler with conn-specific id by passing queue and full conn_queues
+            asyncio.create_task(handle_local_connection_multiplex(r, w, ws, recv_lock, conn_id, q, conn_queues))
+
+        server = await asyncio.start_server(accept_callback, '127.0.0.1', local_port)
+        addrs = ', '.join(str(sock.getsockname()) for sock in server.sockets)
+        print(f'Listening on {addrs}, forward to server (multiplexed)')
+
+        async with server:
+            await server.serve_forever()
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--server', required=True, help='WebSocket server URL, e.g. ws://localhost:8000/ws')
+    parser.add_argument('--local-port', type=int, default=9000, help='Local port to listen on')
+    args = parser.parse_args()
+
+    try:
+        asyncio.run(run_client(args.server, args.local_port))
+    except KeyboardInterrupt:
+        print('client exiting')
+

requirements.txt

@@ -0,0 +1,3 @@
+websockets==11.0.3
+anyio==3.7.0
+python-socketio==5.9.0