fix(proxy): forward whitelisted headers in WebSocket proxy upstream handshake (#865)

docs/_specs/proxy-enhancements/01_requirement.md

@@ -7,10 +7,16 @@ ROCK Admin 目前提供两类代理能力：
 1. **WebSocket Proxy** (`/sandboxes/{id}/proxy/ws[/{path}]`)：将 WebSocket 连接转发到沙箱内的固定 SERVER 端口（`Port.SERVER = 8080`），不支持用户指定目标端口。
 2. **HTTP Proxy** (`/sandboxes/{sandbox_id}/proxy[/{path}]`)：将 HTTP 请求转发到沙箱内的固定 SERVER 端口（`Port.SERVER = 8080`），且硬编码为 `method="POST"`，不支持 GET / PUT / DELETE / PATCH 等其他 HTTP 方法，也不支持用户指定目标端口。
 
-这三个限制阻碍了以下场景：
+除了端口和 method 能力不足之外，WebSocket proxy 还有一个上下文丢失问题：
+- Admin 在转发 WebSocket 握手到下游服务时，目前只处理 `Sec-WebSocket-Protocol` 子协议，不透传通用请求头
+- 当下游服务依赖 `Origin`、`Authorization`、`Cookie`、`X-Forwarded-*` 等头做来源校验、认证、会话恢复或审计时，二跳握手会丢失这些上下文
+- 典型失败现象是下游日志出现 `origin not allowed`，或者因为缺失 token / cookie 导致握手或后续鉴权失败
+
+这些限制阻碍了以下场景：
 - 沙箱内运行了多个 WebSocket 服务（如 Jupyter Kernel、VS Code Server、自定义推理服务），需要连接到不同端口
 - 沙箱内服务使用 RESTful 风格 API，需要 GET 查询、PUT 更新、DELETE 删除
 - 沙箱内运行了多个 HTTP 服务，需要访问非 8080 端口
+- 沙箱内 WebSocket 服务依赖浏览器来源校验、认证头、cookie 或链路追踪头，要求代理保留客户端请求上下文
 
 ---
 
@@ -33,12 +39,21 @@ ROCK Admin 目前提供两类代理能力：
    - 在现有 `/sandboxes/{sandbox_id}/proxy[/{path}]` 路由上，允许通过 query param `port` 指定目标 HTTP 端口
    - 当 `port` 未指定时，保持现有行为（使用 `Port.SERVER = 8080`）
 
+4. **WebSocket Proxy 支持黑名单过滤透传通用请求头**
+   - 通过 `/sandboxes/{id}/proxy/{path:path}` 建立 WebSocket 代理时，默认将客户端请求中的通用 header 转发到下游服务，通过黑名单排除不应转发的头
+   - 黑名单排除的头包括：WebSocket 握手专用头（`Host`、`Connection`、`Upgrade`、`Sec-WebSocket-Key`、`Sec-WebSocket-Version`、`Sec-WebSocket-Extensions`、`Sec-WebSocket-Protocol`）和 hop-by-hop 头（`Transfer-Encoding`、`TE`、`Trailer`、`Keep-Alive`、`Proxy-Authorization`、`Proxy-Connection`、`Content-Length`）
+   - `Origin` 是必须支持的关键头，因为下游服务可能使用来源白名单（例如 `gateway.controlUi.allowedOrigins`）校验 WebSocket 握手
+   - `Sec-WebSocket-Protocol` 继续通过 WebSocket 子协议协商传递，不作为普通 header 透传
+   - 采用黑名单策略，确保用户自定义 header 能被透传到下游服务
+
 ### Out（本次不做的）
 
 - WebSocket Proxy 的认证/鉴权增强
 - HTTP Proxy 的 multipart/form-data 支持（upload 接口已单独处理）
 - `host_proxy` 的 method 扩展（范围外）
 - SDK 客户端侧的封装更新
+- 自动伪造、补默认值或重写任意 `Origin`
+- 透传 WebSocket 握手专用头（如 `Host`、`Connection`、`Upgrade`、`Sec-WebSocket-Key`）
 
 ---
 
@@ -53,6 +68,11 @@ ROCK Admin 目前提供两类代理能力：
 - **AC7**：SSE streaming（`text/event-stream`）在所有 method 下仍然正常工作
 - **AC8**：`GET /sandboxes/{sandbox_id}/proxy/api/health?port=9000` 能成功代理到沙箱内 9000 端口的 HTTP 服务
 - **AC9**：HTTP proxy 不带 `port` 参数时行为不变（向后兼容）
+- **AC10**：当客户端 WebSocket 握手包含 `Origin` 时，代理发起到下游的二跳握手必须携带相同 `Origin`，以满足下游来源校验
+- **AC11**：当客户端握手包含 `Authorization`、`Cookie`、`X-Forwarded-*`、`X-Request-Id`、`Traceparent` 或任意自定义头时，代理发起到下游的二跳握手必须一并转发
+- **AC12**：`Sec-WebSocket-Protocol` 必须继续通过现有 `subprotocols` 机制转发和协商，不能降级为普通 header 透传
+- **AC13**：黑名单头（`Host`、`Connection`、`Upgrade`、`Sec-WebSocket-Key`、`Sec-WebSocket-Version`、`Sec-WebSocket-Extensions`、`Transfer-Encoding`、`TE`、`Trailer`、`Keep-Alive`、`Proxy-Authorization`、`Proxy-Connection`、`Content-Length`）不得被转发到下游
+- **AC14**：当客户端未携带任何白名单头时，WebSocket proxy 的默认行为与现状保持一致（向后兼容）
 
 ---
 
@@ -62,11 +82,16 @@ ROCK Admin 目前提供两类代理能力：
 - 不修改 `Port.SERVER` / `Port.PROXY` 枚举定义
 - `port_validation` 逻辑复用现有 `validate_port_forward_port`（但 WebSocket proxy 端口校验需新增对 `Port.SERVER = 8080` 的允许，或直接用同一校验函数）
 - 保持现有 portforward 端点（`/sandboxes/{id}/portforward`）不变
+- WebSocket header 转发采用**黑名单**策略，排除握手专用头和 hop-by-hop 头，允许用户自定义 header 透传
+- `Origin` 通过 WebSocket 客户端库的显式参数透传，不与普通 `additional_headers` 混用
+- `Sec-WebSocket-Protocol` 继续通过 `subprotocols` 参数处理，不走通用 header 透传逻辑
 
 ---
 
 ## Risks & Rollout
 
 - **风险**：WebSocket proxy 中用户可以指定任意端口访问沙箱内服务，存在横向访问风险 → 通过 `sandbox_id` 鉴权已覆盖，端口范围校验作为防护层
-- **回滚**：修改仅在 `sandbox_proxy_api.py` 和 `sandbox_proxy_service.py` 内，回滚只需还原这两个文件
-- **上线策略**：无数据库变更，直接部署
+- **风险**：黑名单策略下，未列入黑名单的头会默认转发 → 通过单元测试确保握手专用头和 hop-by-hop 头被正确过滤
+- **风险**：不同下游服务对 `Origin`、`Cookie`、`Authorization` 的要求不同，透传后暴露出原本被代理层掩盖的问题 → 以“尽可能保留客户端上下文、但不伪造默认值”为原则
+- **回滚**：修改集中在 `sandbox_proxy_service.py` 和对应单元测试；如需回滚，可仅还原 WebSocket header 透传逻辑
+- **上线策略**：无数据库变更，直接部署；建议先在依赖 `Origin` 校验的控制台类服务上验证

docs/_specs/proxy-enhancements/03_implementation.md

@@ -7,15 +7,21 @@ admin 与 sandbox 不在同一 K8s 集群，`host_ip` 为宿主机 IP，容器
 - **WebSocket proxy 自定义端口**：复用 rocklet 现有的 `/portforward` WebSocket 端点中转（与 `/sandboxes/{id}/portforward` 相同机制）
 - **HTTP proxy 自定义端口**：需在 rocklet 新增 `/http_proxy` HTTP 端点，admin 转发请求给 rocklet，rocklet 在容器内访问目标服务
 
+除了端口和 method 能力之外，当前 WebSocket proxy 还存在握手上下文丢失问题：
+- `sandbox_proxy_service.websocket_proxy()` 在调用 `websockets.connect(...)` 时，目前只传了 `subprotocols`
+- 下游服务收到的是 Admin 重新发起的二跳握手，请求来源会表现为 `Python websockets/...`，而不是客户端原始握手上下文
+- 结果是 `Origin`、`Authorization`、`Cookie`、`X-Forwarded-*` 等头全部丢失，依赖这些头的服务会报 `origin not allowed` 或鉴权失败
+
 ---
 
 ## File Changes
 
 | 文件 | 修改类型 | 说明 |
 |------|------|------|
 | `rock/rocklet/local_api.py` | **新增** | 新增 `ANY /http_proxy/{path:path}?port={port}` 端点 |
-| `rock/sandbox/service/sandbox_proxy_service.py` | 修改 | `http_proxy` 有 `port` 时改走 rocklet `/http_proxy` 中转；WebSocket proxy 有 `port` 时改走 rocklet `/portforward` 中转 |
-| `rock/admin/entrypoints/sandbox_proxy_api.py` | 无变更 | 已支持，无需修改 |
+| `rock/sandbox/service/sandbox_proxy_service.py` | 修改 | `http_proxy` 有 `port` 时改走 rocklet `/http_proxy` 中转；WebSocket proxy 有 `port` 时改走 rocklet `/portforward` 中转；补充 WebSocket 通用 headers 白名单透传 |
+| `rock/admin/entrypoints/sandbox_proxy_api.py` | 无变更 | 路由签名保持不变，无需调整 |
+| `tests/unit/sandbox/test_websocket_proxy_subprotocol.py` | 修改 | 增加 `Origin` / `additional_headers` 透传与禁转头测试 |
 
 ---
 
@@ -110,6 +116,88 @@ async def http_proxy(self, sandbox_id, target_path, body, headers, method="POST"
     ...
 ```
 
+### 变更 4：WebSocket proxy 通用 headers 黑名单过滤透传
+
+需要在 `sandbox_proxy_service.websocket_proxy()` 内新增一层 header 提取和过滤逻辑，将客户端握手里的”通用请求头”转为上游二跳握手参数。
+
+**设计要点**：
+
+1. **将 `Origin` 单独处理**
+   - `websockets.connect()` 在 15.0.1 版本里提供 `origin=` 参数
+   - `Origin` 不作为普通 `additional_headers` 重复透传，避免语义混乱和重复 header
+
+2. **其余头走黑名单过滤**
+   - 黑名单：WebSocket 握手专用头（`Host`、`Connection`、`Upgrade`、`Sec-WebSocket-Key`、`Sec-WebSocket-Version`、`Sec-WebSocket-Extensions`、`Sec-WebSocket-Protocol`）和 hop-by-hop 头（`Transfer-Encoding`、`TE`、`Trailer`、`Keep-Alive`、`Proxy-Authorization`、`Proxy-Connection`、`Content-Length`）
+   - 不在黑名单中的头默认转发，确保用户自定义 header 能到达下游
+
+3. **无可转发头时保持兼容**
+   - 若客户端未携带任何非黑名单头，则 `origin=None`、`additional_headers=None`
+   - 这样代理行为与当前实现保持一致，不引入额外副作用
+
+**实现代码**（`rock/sandbox/utils/proxy.py`）：
+
+```python
+BLOCKED_WS_HEADER_NAMES = {
+    “host”,
+    “connection”,
+    “upgrade”,
+    “sec-websocket-key”,
+    “sec-websocket-version”,
+    “sec-websocket-extensions”,
+    “sec-websocket-protocol”,
+    “transfer-encoding”,
+    “te”,
+    “trailer”,
+    “keep-alive”,
+    “proxy-authorization”,
+    “proxy-connection”,
+    “content-length”,
+}
+
+
+def build_upstream_ws_headers(client_websocket):
+    origin = client_websocket.headers.get(“origin”) or client_websocket.headers.get(“Origin”)
+    additional_headers = []
+
+    for key, value in client_websocket.headers.items():
+        lower_key = key.lower()
+        if lower_key == “origin”:
+            continue
+        if lower_key in BLOCKED_WS_HEADER_NAMES:
+            continue
+        additional_headers.append((key, value))
+
+    return origin, additional_headers or None
+```
+
+接入方式：
+
+```python
+origin, additional_headers = build_upstream_ws_headers(client_websocket)
+
+async with websockets.connect(
+    target_url,
+    ping_interval=None,
+    ping_timeout=None,
+    origin=origin,
+    additional_headers=additional_headers,
+    subprotocols=upstream_subprotocols,
+) as target_websocket:
+    ...
+```
+
+### 变更 5：测试覆盖扩展
+
+现有测试主要覆盖子协议转发，需要补充 header 透传相关单测。
+
+**新增测试点**：
+- `Origin` 存在时，`websockets.connect()` 收到相同 `origin=`
+- 已知头（`Authorization`、`Cookie`、`X-Forwarded-*`、`X-Request-Id`、`Traceparent`、`EagleEye-*`）存在时，`websockets.connect()` 收到 `additional_headers`
+- 用户自定义头（如 `x-my-custom`）能被正常转发到下游
+- 黑名单头（`Host`、`Connection`、`Upgrade`、`Sec-WebSocket-Key`、`Sec-WebSocket-Version`、`Sec-WebSocket-Extensions` 等）不得出现在 `additional_headers`
+- `Sec-WebSocket-Protocol` 继续通过 `subprotocols=` 转发，不能出现在 `additional_headers`
+- 无可转发头时，`origin` / `additional_headers` 为 `None`，保持向后兼容
+
 ---
 
 ## Execution Plan
@@ -131,14 +219,25 @@ async def http_proxy(self, sandbox_id, target_path, body, headers, method="POST"
 - 有 `port` 时，改用 `http://{host_ip}:{rocklet_mapped_port}/http_proxy/{path}?port={port}`
 - 无 `port` 时保持原逻辑不变
 
+### Step 4：新增 WebSocket 通用 header 黑名单过滤与透传逻辑
+- 文件：`rock/sandbox/utils/proxy.py`（独立模块）、`rock/sandbox/service/sandbox_proxy_service.py`（调用方）
+- 新增 `build_upstream_ws_headers()` helper，负责从 `client_websocket.headers` 中提取 `Origin` 并通过黑名单过滤 `additional_headers`
+- 在 `websocket_proxy()` 调用 `websockets.connect()` 时传入 `origin=` 和 `additional_headers=`
+- 保持现有 `subprotocols=` 协商逻辑不变
+
+### Step 5：补充 WebSocket header 透传测试
+- 文件：`tests/unit/sandbox/test_websocket_proxy_headers.py`
+- 新增 `Origin` 透传、已知 header 透传、自定义 header 透传、黑名单 header 过滤、兼容性测试
+
 ---
 
 ## Rollback & Compatibility
 
 - **向后兼容**：`rock_target_port` 未指定时，所有逻辑路径与原实现完全一致
+- **向后兼容**：客户端未携带任何白名单 header 时，WebSocket 二跳握手行为与现状一致
 - **回滚**：
   - rocklet：还原 `local_api.py`，重新发布镜像
-  - admin：还原 `sandbox_proxy_service.py`
+  - admin：还原 `sandbox_proxy_service.py` 和对应单元测试
 
 ---
 
@@ -147,3 +246,6 @@ async def http_proxy(self, sandbox_id, target_path, body, headers, method="POST"
 - WebSocket proxy 自定义端口时，`path` 参数不生效（rocklet portforward 是纯 TCP 隧道，不感知 HTTP path）
 - rocklet `/http_proxy` 端点的 `port` 参数需要校验（复用 `validate_port_forward_port`）
 - rocklet 镜像需要重新发布才能生效
+- WebSocket header 透传采用黑名单策略，排除握手专用头和 hop-by-hop 头，允许用户自定义 header 透传
+- `Origin` 应通过 `websockets.connect(origin=...)` 传入；不要与 `additional_headers` 重复
+- `Sec-WebSocket-Protocol` 必须继续通过 `subprotocols=` 传递，避免和普通 header 透传逻辑冲突

rock/sandbox/service/sandbox_proxy_service.py

@@ -40,6 +40,7 @@
 from rock.common.port_validation import validate_port_forward_port
 from rock.logger import init_logger
 from rock.sandbox.sandbox_meta_store import SandboxMetaStore
+from rock.sandbox.utils.proxy import build_upstream_ws_headers
 from rock.sandbox.utils.timeout import SandboxTimeoutHelper
 from rock.sdk.common.exceptions import BadRequestRockError
 from rock.utils import EAGLE_EYE_TRACE_ID, trace_id_ctx_var
@@ -208,12 +209,15 @@ async def websocket_proxy(
 
         client_subprotocols = getattr(client_websocket, "subprotocols", []) or []
         upstream_subprotocols = client_subprotocols if client_subprotocols else ["binary", "base64"]
+        origin, additional_headers = build_upstream_ws_headers(client_websocket)
 
         try:
             async with websockets.connect(
                 target_url,
                 ping_interval=None,
                 ping_timeout=None,
+                origin=origin,
+                additional_headers=additional_headers,
                 subprotocols=upstream_subprotocols,
             ) as target_websocket:
                 negotiated = getattr(target_websocket, "subprotocol", None)

rock/sandbox/utils/proxy.py

@@ -0,0 +1,31 @@
+BLOCKED_WS_HEADER_NAMES = {
+    "host",
+    "connection",
+    "upgrade",
+    "sec-websocket-key",
+    "sec-websocket-version",
+    "sec-websocket-extensions",
+    "sec-websocket-protocol",
+    "transfer-encoding",
+    "te",
+    "trailer",
+    "keep-alive",
+    "proxy-authorization",
+    "proxy-connection",
+    "content-length",
+}
+
+
+def build_upstream_ws_headers(client_websocket):
+    origin = client_websocket.headers.get("origin") or client_websocket.headers.get("Origin")
+    additional_headers = []
+
+    for key, value in client_websocket.headers.items():
+        lower_key = key.lower()
+        if lower_key == "origin":
+            continue
+        if lower_key in BLOCKED_WS_HEADER_NAMES:
+            continue
+        additional_headers.append((key, value))
+
+    return origin, additional_headers or None