Spaces:
Paused
Paused
File size: 4,645 Bytes
34fe2fb |
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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 |
# main.py
import asyncio
from fastapi import FastAPI, WebSocket
import uvicorn
# 创建 FastAPI 应用实例
app = FastAPI()
@app.websocket("/")
async def tunnel(websocket: WebSocket):
"""
WebSocket 隧道入口：
1. 接受客户端 WebSocket 连接；
2. 等待客户端发送 CONNECT 请求，解析出目标主机及端口；
3. 尝试与目标主机建立 TCP 连接；
4. 返回 HTTP 200 成功建立隧道；
5. 启动双向数据转发（TCP <--> WebSocket）。
"""
await websocket.accept()
try:
# ---------------------------
# 1. 等待客户端发来的 CONNECT 请求
# ---------------------------
# CONNECT 请求格式示例：
# CONNECT destHost:destPort HTTP/1.1\r\nHost: destHost:destPort\r\n\r\n
request_text = await websocket.receive_text()
lines = request_text.splitlines()
if not lines:
await websocket.send_text("HTTP/1.1 400 Bad Request\r\n\r\n")
await websocket.close()
return
# 解析第一行
first_line = lines[0].strip()
parts = first_line.split()
if len(parts) < 3 or parts[0].upper() != "CONNECT":
await websocket.send_text("HTTP/1.1 400 Bad Request\r\n\r\n")
await websocket.close()
return
# 从第二个字段中获取目标主机及端口，如 destHost:destPort
dest = parts[1]
if ":" not in dest:
await websocket.send_text("HTTP/1.1 400 Bad Request\r\n\r\n")
await websocket.close()
return
dest_parts = dest.split(":", 1)
dest_host = dest_parts[0]
try:
dest_port = int(dest_parts[1])
except Exception:
await websocket.send_text("HTTP/1.1 400 Bad Request\r\n\r\n")
await websocket.close()
return
# ---------------------------
# 2. 建立到目标主机的 TCP 连接
# ---------------------------
try:
reader, writer = await asyncio.open_connection(dest_host, dest_port)
except Exception as e:
err_msg = f"HTTP/1.1 502 Bad Gateway\r\n\r\n无法连接 {dest_host}:{dest_port}，错误：{e}"
await websocket.send_text(err_msg)
await websocket.close()
return
# ---------------------------
# 3. 向客户端返回 200 成功响应
# ---------------------------
await websocket.send_text("HTTP/1.1 200 Connection Established\r\n\r\n")
# ---------------------------
# 4. 双向数据转发
# ---------------------------
async def tcp_to_ws():
"""
从 TCP 连接中读取数据，通过 WebSocket 以二进制方式发送给客户端
"""
try:
while True:
data = await reader.read(1024)
if not data:
break
await websocket.send_bytes(data)
except Exception as e:
# 读取异常或对方关闭连接时退出
print("tcp_to_ws 异常：", e)
async def ws_to_tcp():
"""
从客户端通过 WebSocket 发送的数据写入 TCP 连接
"""
try:
while True:
message = await websocket.receive()
# 接收到的数据可能是文本或二进制，这里尽量以二进制方式处理
if "bytes" in message:
data = message["bytes"]
elif "text" in message:
# 若收到文本数据，则转换为 bytes（可能只在握手阶段出现）
data = message["text"].encode("utf-8")
else:
break
writer.write(data)
await writer.drain()
except Exception as e:
print("ws_to_tcp 异常：", e)
# 并发执行数据转发任务，任一方向关闭则结束隧道
await asyncio.gather(tcp_to_ws(), ws_to_tcp())
except Exception as e:
print("WebSocket 隧道处理异常：", e)
finally:
# 关闭连接
await websocket.close()
# ---------------------------
# 启动服务器：监听 0.0.0.0:7860
# ---------------------------
if __name__ == "__main__":
uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=False)
|