opencode-go-proxy v0.1.0: OpenCode Go API 本地转发代理
- 透明转发 /v1/responses、/v1/chat/completions、/v1/messages、/v1/models - 自动注入上游认证 / User-Agent / x-opencode-session 会话头 - 双层密钥:上游 key(-k)+ 客户端访问 key(-c),客户端认证头不透传 - 手写 HTTP/1.1(rustls 仅用于 TLS):keep-alive、chunked、100-continue - SSE 逐块流式透传,上游错误响应详情日志 - mini-socks 同款 systemd 安装(-i/--uninstall),musl 静态单文件 - 37 个单元/集成测试
This commit is contained in:
phptits
@@ -0,0 +1,5 @@
|
||||
[build]
|
||||
target = "x86_64-unknown-linux-musl"
|
||||
|
||||
[target.x86_64-unknown-linux-musl]
|
||||
linker = "musl-gcc"
|
||||
@@ -0,0 +1 @@
|
||||
/target
|
||||
Generated
+344
@@ -0,0 +1,344 @@
|
||||
# This file is automatically @generated by Cargo.
|
||||
# It is not intended for manual editing.
|
||||
version = 4
|
||||
|
||||
[[package]]
|
||||
name = "bytes"
|
||||
version = "1.12.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04"
|
||||
|
||||
[[package]]
|
||||
name = "cc"
|
||||
version = "1.4.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "005ec2760ca554fae18df7a11195552ec576cd665632a881bc011d5bb2fd4d80"
|
||||
dependencies = [
|
||||
"find-msvc-tools",
|
||||
"shlex",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cfg-if"
|
||||
version = "1.0.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801"
|
||||
|
||||
[[package]]
|
||||
name = "find-msvc-tools"
|
||||
version = "0.1.12"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3e0f1c7c3a72c66fd80abe965175f7523475c0489a87d3ff9d6e8c87d87a9d2d"
|
||||
|
||||
[[package]]
|
||||
name = "getrandom"
|
||||
version = "0.2.17"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"libc",
|
||||
"wasi",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "libc"
|
||||
version = "0.2.189"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2"
|
||||
|
||||
[[package]]
|
||||
name = "log"
|
||||
version = "0.4.34"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6"
|
||||
|
||||
[[package]]
|
||||
name = "mio"
|
||||
version = "1.2.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4b18443e9c262bfe8fa82f51666e2642c53393f7e5c27b3e1aeab922cff5b9d8"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"wasi",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "once_cell"
|
||||
version = "1.21.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50"
|
||||
|
||||
[[package]]
|
||||
name = "opencode-go-proxy"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"rustls",
|
||||
"tokio",
|
||||
"tokio-rustls",
|
||||
"webpki-roots",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pin-project-lite"
|
||||
version = "0.2.17"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd"
|
||||
|
||||
[[package]]
|
||||
name = "proc-macro2"
|
||||
version = "1.0.107"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9"
|
||||
dependencies = [
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "quote"
|
||||
version = "1.0.47"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ring"
|
||||
version = "0.17.14"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7"
|
||||
dependencies = [
|
||||
"cc",
|
||||
"cfg-if",
|
||||
"getrandom",
|
||||
"libc",
|
||||
"untrusted",
|
||||
"windows-sys 0.52.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustls"
|
||||
version = "0.23.44"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6725596c3f2c3a0aef021139e145d4eafe314a6623e4680ca83852b2c67ab2ba"
|
||||
dependencies = [
|
||||
"log",
|
||||
"once_cell",
|
||||
"ring",
|
||||
"rustls-pki-types",
|
||||
"rustls-webpki",
|
||||
"subtle",
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustls-pki-types"
|
||||
version = "1.15.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2f4925028c7eb5d1fcdaf196971378ed9d2c1c4efc7dc5d011256f76c99c0a96"
|
||||
dependencies = [
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustls-webpki"
|
||||
version = "0.103.15"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f3c3cf1d8b1e7d4927e2d154c3fcb02979afb9939629c62cd9048d4f07b60ac2"
|
||||
dependencies = [
|
||||
"ring",
|
||||
"rustls-pki-types",
|
||||
"untrusted",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "shlex"
|
||||
version = "2.0.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba"
|
||||
|
||||
[[package]]
|
||||
name = "socket2"
|
||||
version = "0.6.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "subtle"
|
||||
version = "2.6.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292"
|
||||
|
||||
[[package]]
|
||||
name = "syn"
|
||||
version = "3.0.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "12df2e0110f65b775f769bb17ef989067a1d931b2eb822bd4346631eeada89f9"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tokio"
|
||||
version = "1.53.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"libc",
|
||||
"mio",
|
||||
"pin-project-lite",
|
||||
"socket2",
|
||||
"tokio-macros",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tokio-macros"
|
||||
version = "2.7.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "78773a2a397f451582ce068015985c33193cf6dea8b74d2a639fe457b2f07b0e"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tokio-rustls"
|
||||
version = "0.26.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b0c85f2c3ef0b1cd58b36682f4b17aaa995f0e5db534d85692b4903abce21f67"
|
||||
dependencies = [
|
||||
"rustls",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "unicode-ident"
|
||||
version = "1.0.24"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
|
||||
|
||||
[[package]]
|
||||
name = "untrusted"
|
||||
version = "0.9.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1"
|
||||
|
||||
[[package]]
|
||||
name = "wasi"
|
||||
version = "0.11.1+wasi-snapshot-preview1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b"
|
||||
|
||||
[[package]]
|
||||
name = "webpki-roots"
|
||||
version = "1.0.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7dcd9d09a39985f5344844e66b0c530a33843579125f23e21e9f0f220850f22a"
|
||||
dependencies = [
|
||||
"rustls-pki-types",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-link"
|
||||
version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5"
|
||||
|
||||
[[package]]
|
||||
name = "windows-sys"
|
||||
version = "0.52.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d"
|
||||
dependencies = [
|
||||
"windows-targets",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-sys"
|
||||
version = "0.61.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc"
|
||||
dependencies = [
|
||||
"windows-link",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-targets"
|
||||
version = "0.52.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973"
|
||||
dependencies = [
|
||||
"windows_aarch64_gnullvm",
|
||||
"windows_aarch64_msvc",
|
||||
"windows_i686_gnu",
|
||||
"windows_i686_gnullvm",
|
||||
"windows_i686_msvc",
|
||||
"windows_x86_64_gnu",
|
||||
"windows_x86_64_gnullvm",
|
||||
"windows_x86_64_msvc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows_aarch64_gnullvm"
|
||||
version = "0.52.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3"
|
||||
|
||||
[[package]]
|
||||
name = "windows_aarch64_msvc"
|
||||
version = "0.52.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469"
|
||||
|
||||
[[package]]
|
||||
name = "windows_i686_gnu"
|
||||
version = "0.52.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b"
|
||||
|
||||
[[package]]
|
||||
name = "windows_i686_gnullvm"
|
||||
version = "0.52.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66"
|
||||
|
||||
[[package]]
|
||||
name = "windows_i686_msvc"
|
||||
version = "0.52.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66"
|
||||
|
||||
[[package]]
|
||||
name = "windows_x86_64_gnu"
|
||||
version = "0.52.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78"
|
||||
|
||||
[[package]]
|
||||
name = "windows_x86_64_gnullvm"
|
||||
version = "0.52.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d"
|
||||
|
||||
[[package]]
|
||||
name = "windows_x86_64_msvc"
|
||||
version = "0.52.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec"
|
||||
|
||||
[[package]]
|
||||
name = "zeroize"
|
||||
version = "1.9.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e"
|
||||
+19
@@ -0,0 +1,19 @@
|
||||
[package]
|
||||
name = "opencode-go-proxy"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
description = "OpenCode Go API 本地转发代理:自动注入认证 / User-Agent / 会话头,让任意 OpenAI/Anthropic 兼容客户端透明使用 OpenCode Go"
|
||||
license = "MIT"
|
||||
|
||||
[dependencies]
|
||||
tokio = { version = "1", features = ["rt-multi-thread", "macros", "net", "io-util", "time"] }
|
||||
rustls = { version = "0.23", default-features = false, features = ["logging", "ring", "std", "tls12"] }
|
||||
tokio-rustls = { version = "0.26", default-features = false, features = ["ring", "tls12"] }
|
||||
webpki-roots = "1.0"
|
||||
|
||||
[profile.release]
|
||||
opt-level = 3
|
||||
lto = true
|
||||
codegen-units = 1
|
||||
strip = true
|
||||
panic = "abort"
|
||||
@@ -0,0 +1,122 @@
|
||||
# opencode-go-proxy
|
||||
|
||||
OpenCode Go([opencode.ai/zh-cn/go](https://opencode.ai/docs/zh-cn/go/))API 本地转发代理。
|
||||
|
||||
OpenCode Go 要求客户端请求携带专属 `User-Agent` 和稳定的 `x-opencode-session` 会话头,
|
||||
大多数第三方客户端无法自定义这些请求头。本工具在本地起一个透明 HTTP 代理,自动注入
|
||||
OpenCode Go 要求的全部请求头,让任意 OpenAI / Anthropic 兼容客户端零改造接入。
|
||||
|
||||
```
|
||||
客户端(任意 OpenAI/Anthropic SDK)
|
||||
│ http://127.0.0.1:8090/v1/*
|
||||
▼
|
||||
opencode-go-proxy
|
||||
│ 注入 Authorization / User-Agent / x-opencode-session
|
||||
│ 路径映射 /v1/* → https://opencode.ai/zen/go/v1/*
|
||||
▼
|
||||
OpenCode Go
|
||||
```
|
||||
|
||||
## 功能
|
||||
|
||||
- 透明转发四个端点,响应体逐字节流式透传(SSE 首 token 延迟不受影响):
|
||||
- `/v1/responses` — OpenAI Responses API(Grok 4.6、GPT 5.6 Luna、Muse Spark)
|
||||
- `/v1/chat/completions` — OpenAI 兼容(GLM、Kimi、DeepSeek、LongCat、MiMo、Hy3/Hy4、Omen Alpha)
|
||||
- `/v1/messages` — Anthropic 兼容(MiniMax、Qwen 系列)
|
||||
- `/v1/models` — 模型列表
|
||||
- 自动注入:
|
||||
- `Authorization: Bearer <密钥>`(密钥由代理持有,不经过客户端)
|
||||
- `User-Agent: opencode-go-proxy/x.y`(可自定义)
|
||||
- `x-opencode-session`:客户端带了就透传;没带则为每个客户端连接生成稳定 ID
|
||||
(同一 keep-alive 连接内的连续请求共享,近似"每对话一个");也可用 `-s` 固定
|
||||
- `/v1/messages` 额外补 `x-api-key` 与 `anthropic-version: 2023-06-01`(客户端未带时)
|
||||
- 双层密钥:
|
||||
- **上游密钥 `-k`**:OpenCode Go 的 API key,仅存在于代理配置中,转发时注入
|
||||
- **客户端密钥 `-c`**(可选,公网部署强烈建议):配置后所有 API 请求必须携带
|
||||
`Authorization: Bearer <客户端密钥>` 或 `x-api-key`(任一匹配即可),否则 401;
|
||||
`/healthz` 免鉴权供监控使用;`-i` 安装时不指定则自动生成
|
||||
- 支持 keep-alive、`Expect: 100-continue`、chunked 请求体,请求体上限 32 MB
|
||||
- 上游连接/握手超时 15 秒,响应头超时 60 秒,响应体不限时(SSE 长连接)
|
||||
- 客户端自带的认证头不会透传到上游,由代理统一替换
|
||||
|
||||
## 构建
|
||||
|
||||
需要 Rust 1.75+ 与 musl 工具链(或改用默认 gnu target):
|
||||
|
||||
```sh
|
||||
rustup target add x86_64-unknown-linux-musl
|
||||
cargo build --release
|
||||
# 产物:target/x86_64-unknown-linux-musl/release/opencode-go-proxy(静态单文件,约 2 MB)
|
||||
```
|
||||
|
||||
## 使用
|
||||
|
||||
```sh
|
||||
export OPENCODE_GO_API_KEY=sk-xxx # 在 opencode.ai/auth 订阅后获取
|
||||
opencode-go-proxy # 监听 127.0.0.1:8090
|
||||
```
|
||||
|
||||
### systemd 服务安装(mini-socks 方式)
|
||||
|
||||
```sh
|
||||
sudo opencode-go-proxy -i -p 28080 -k sk-xxx # 安装并启动;-c 缺省自动生成客户端密钥
|
||||
sudo opencode-go-proxy -i -p 28080 -k sk-xxx -c ogp-xxx # 或固定客户端密钥
|
||||
sudo opencode-go-proxy --uninstall # 停止并卸载
|
||||
```
|
||||
|
||||
安装会把二进制复制到 `/usr/bin/opencode-go-proxy`,unit 写入
|
||||
`/etc/systemd/system/opencode-go-proxy.service`,`-p/-k/-c/-s/-a/-u` 固化在
|
||||
ExecStart 中。更换密钥:编辑 unit 中的 `-k`(上游)或 `-c`(客户端)后
|
||||
`systemctl restart opencode-go-proxy`。
|
||||
|
||||
```
|
||||
opencode-go-proxy [-l 地址] [-p 端口] [-k 密钥] [-s 会话ID] [-a UA] [-u 上游] [-h]
|
||||
|
||||
-l <地址> 监听地址 (默认 127.0.0.1)
|
||||
-p <端口> 监听端口 (默认 8090)
|
||||
-k <密钥> API 密钥;缺省读环境变量 OPENCODE_GO_API_KEY / OPENCODE_API_KEY
|
||||
-s <会话ID> 固定 x-opencode-session;缺省为每个客户端连接自动生成
|
||||
-a <UA> 上游 User-Agent (默认 opencode-go-proxy/版本)
|
||||
-u <地址> 上游地址 (默认 https://opencode.ai/zen/go;支持 http:// 供本地调试)
|
||||
```
|
||||
|
||||
验证:
|
||||
|
||||
```sh
|
||||
curl http://127.0.0.1:8090/v1/models
|
||||
|
||||
curl http://127.0.0.1:8090/v1/chat/completions \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{"model":"kimi-k3","messages":[{"role":"user","content":"你好"}]}'
|
||||
```
|
||||
|
||||
## 客户端接入
|
||||
|
||||
按各模型所属端点选择协议与模型 ID(完整列表见 `/v1/models`):
|
||||
|
||||
| 客户端 | 配置 |
|
||||
|---|---|
|
||||
| OpenAI SDK / Cherry Studio / LobeChat 等 | Base URL `http://127.0.0.1:8090/v1`,API Key 填客户端密钥(`-c` 的值) |
|
||||
| Claude Code | `ANTHROPIC_BASE_URL=http://127.0.0.1:8090` + `ANTHROPIC_API_KEY=<客户端密钥>` |
|
||||
| curl / 脚本 | `-H "Authorization: Bearer <客户端密钥>"` 直接请求 `/v1/*` |
|
||||
|
||||
模型与端点的对应关系(节选,以 `/v1/models` 为准):
|
||||
|
||||
- **chat/completions**:kimi-k3、glm-5.3-flash、glm-5.3、deepseek-v4-pro、deepseek-v4-flash、mimo-v2.5、longcat-2.0、hy4-preview、omen-alpha
|
||||
- **messages**:minimax-m3、qwen3.8-max、qwen3.8-flash、qwen3.7-max
|
||||
- **responses**:grok-4.6、gpt-5.6-luna、muse-spark-1.3-contributor
|
||||
|
||||
## 日志
|
||||
|
||||
每请求一行(stderr):时间、客户端地址、序号、方法、路径、上游状态码、响应字节数、耗时。
|
||||
API 密钥与请求体不会出现在日志中。
|
||||
|
||||
## 限制
|
||||
|
||||
- 仅支持 HTTP/1.1 客户端
|
||||
- 每个上游请求新建 TLS 连接(AI 请求耗时长,复用收益低,换来实现简单)
|
||||
- 不做协议转换:模型需用其原生协议(如 Claude Code 只能走 `/v1/messages` 端点上的模型)
|
||||
|
||||
## 许可
|
||||
|
||||
MIT
|
||||
+667
@@ -0,0 +1,667 @@
|
||||
// HTTP/1.1 最小实现:仅覆盖转发代理需要的能力。
|
||||
// 请求解析(定长/chunked body、100-continue)、响应头解析、
|
||||
// 头部序列化,以及 body 的流式透传(不整体缓冲,SSE 友好)。
|
||||
|
||||
use std::io;
|
||||
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
|
||||
|
||||
pub const MAX_LINE: usize = 8 * 1024;
|
||||
pub const MAX_HEAD: usize = 64 * 1024;
|
||||
pub const MAX_BODY: u64 = 32 * 1024 * 1024;
|
||||
|
||||
pub fn proto_err(msg: impl Into<String>) -> io::Error {
|
||||
io::Error::new(io::ErrorKind::InvalidData, msg.into())
|
||||
}
|
||||
|
||||
/// 服务器未发 TLS close_notify 就直接关闭:HTTP 语义中这等价于 body 结束
|
||||
/// (rustls 对此返回 UnexpectedEof,见其文档 unexpected-eof 一节)。
|
||||
pub fn is_tls_eof_without_notify(e: &io::Error) -> bool {
|
||||
e.kind() == io::ErrorKind::UnexpectedEof && e.to_string().contains("close_notify")
|
||||
}
|
||||
|
||||
pub fn header_get<'a>(headers: &'a [(String, String)], name: &str) -> Option<&'a str> {
|
||||
headers
|
||||
.iter()
|
||||
.find(|(n, _)| n.eq_ignore_ascii_case(name))
|
||||
.map(|(_, v)| v.as_str())
|
||||
}
|
||||
|
||||
pub const HOP_BY_HOP: &[&str] = &[
|
||||
"connection",
|
||||
"keep-alive",
|
||||
"proxy-authenticate",
|
||||
"proxy-authorization",
|
||||
"te",
|
||||
"trailer",
|
||||
"transfer-encoding",
|
||||
"upgrade",
|
||||
];
|
||||
|
||||
pub fn is_hop_by_hop(name: &str) -> bool {
|
||||
HOP_BY_HOP.iter().any(|h| name.eq_ignore_ascii_case(h))
|
||||
}
|
||||
|
||||
pub struct Request {
|
||||
pub method: String,
|
||||
pub target: String,
|
||||
pub version: String,
|
||||
pub headers: Vec<(String, String)>,
|
||||
pub body: Vec<u8>,
|
||||
}
|
||||
|
||||
pub struct Response {
|
||||
pub status: u16,
|
||||
pub reason: String,
|
||||
pub headers: Vec<(String, String)>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||
pub enum BodyKind {
|
||||
None,
|
||||
Exact(u64),
|
||||
Chunked,
|
||||
Eof,
|
||||
}
|
||||
|
||||
impl Response {
|
||||
pub fn header(&self, name: &str) -> Option<&str> {
|
||||
header_get(&self.headers, name)
|
||||
}
|
||||
|
||||
pub fn body_kind(&self) -> io::Result<BodyKind> {
|
||||
if self.status == 204 || self.status == 304 || (100..200).contains(&self.status) {
|
||||
return Ok(BodyKind::None);
|
||||
}
|
||||
if let Some(te) = self.header("transfer-encoding") {
|
||||
if te.split(',').any(|p| p.trim().eq_ignore_ascii_case("chunked")) {
|
||||
return Ok(BodyKind::Chunked);
|
||||
}
|
||||
return Err(proto_err(format!("上游不支持的 Transfer-Encoding: {te}")));
|
||||
}
|
||||
if let Some(cl) = self.header("content-length") {
|
||||
let n: u64 = cl
|
||||
.parse()
|
||||
.map_err(|_| proto_err(format!("上游非法 Content-Length: {cl}")))?;
|
||||
return Ok(BodyKind::Exact(n));
|
||||
}
|
||||
// 无定长无分块:读到对端关闭为止(HTTP/1.0 语义,部分 SSE 服务也这样)
|
||||
Ok(BodyKind::Eof)
|
||||
}
|
||||
}
|
||||
|
||||
// ---- 带游标缓冲的读取器 -------------------------------------------------------
|
||||
// 与 BufReader 的区别:所有方法消费式返回 owned 数据或直接向写出端转发,
|
||||
// 不存在"读过头又放不回去"的问题,精确控制 body 边界。
|
||||
|
||||
pub struct ConnReader<R> {
|
||||
inner: R,
|
||||
buf: Box<[u8]>,
|
||||
start: usize,
|
||||
end: usize,
|
||||
}
|
||||
|
||||
impl<R: AsyncRead + Unpin> ConnReader<R> {
|
||||
pub fn new(inner: R, capacity: usize) -> Self {
|
||||
Self {
|
||||
inner,
|
||||
buf: vec![0u8; capacity.max(1024)].into_boxed_slice(),
|
||||
start: 0,
|
||||
end: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// 写访问(上游先写后读,顺序使用,缓冲为空时调用)
|
||||
pub fn get_mut(&mut self) -> &mut R {
|
||||
&mut self.inner
|
||||
}
|
||||
|
||||
// 仅在缓冲耗尽时调用:从底层再读一轮
|
||||
async fn load(&mut self) -> io::Result<usize> {
|
||||
debug_assert_eq!(self.start, self.end);
|
||||
let n = self.inner.read(&mut self.buf).await?;
|
||||
self.start = 0;
|
||||
self.end = n;
|
||||
Ok(n)
|
||||
}
|
||||
|
||||
/// 读一行(含 \r\n,原样保留)。EOF 且无任何数据返回 None,行中间断开报错。
|
||||
pub async fn read_line_opt(&mut self, max: usize) -> io::Result<Option<Vec<u8>>> {
|
||||
let mut line = Vec::new();
|
||||
loop {
|
||||
if self.start < self.end {
|
||||
if let Some(pos) = self.buf[self.start..self.end]
|
||||
.iter()
|
||||
.position(|&b| b == b'\n')
|
||||
{
|
||||
let take = pos + 1;
|
||||
line.extend_from_slice(&self.buf[self.start..self.start + take]);
|
||||
self.start += take;
|
||||
if line.len() > max {
|
||||
return Err(proto_err("行长度超过上限"));
|
||||
}
|
||||
return Ok(Some(line));
|
||||
}
|
||||
line.extend_from_slice(&self.buf[self.start..self.end]);
|
||||
self.start = self.end;
|
||||
if line.len() > max {
|
||||
return Err(proto_err("行长度超过上限"));
|
||||
}
|
||||
}
|
||||
let n = self.load().await?;
|
||||
if n == 0 {
|
||||
return if line.is_empty() {
|
||||
Ok(None)
|
||||
} else {
|
||||
Err(proto_err("连接在行中间被截断"))
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn read_line(&mut self, max: usize) -> io::Result<Vec<u8>> {
|
||||
self.read_line_opt(max)
|
||||
.await?
|
||||
.ok_or_else(|| proto_err("连接在读取行时意外关闭"))
|
||||
}
|
||||
|
||||
pub async fn read_exact_n(&mut self, n: usize) -> io::Result<Vec<u8>> {
|
||||
let mut out = Vec::with_capacity(n.min(64 * 1024));
|
||||
while out.len() < n {
|
||||
if self.start < self.end {
|
||||
let take = (self.end - self.start).min(n - out.len());
|
||||
out.extend_from_slice(&self.buf[self.start..self.start + take]);
|
||||
self.start += take;
|
||||
} else {
|
||||
let got = self.load().await?;
|
||||
if got == 0 {
|
||||
return Err(proto_err(format!(
|
||||
"需要 {} 字节,但连接提前关闭",
|
||||
n - out.len()
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// 读取至多 max 字节或到 EOF(错误响应缓冲诊断用)。
|
||||
pub async fn read_upto(&mut self, max: usize) -> io::Result<Vec<u8>> {
|
||||
let mut out = Vec::new();
|
||||
while out.len() < max {
|
||||
if self.start < self.end {
|
||||
let take = (self.end - self.start).min(max - out.len());
|
||||
out.extend_from_slice(&self.buf[self.start..self.start + take]);
|
||||
self.start += take;
|
||||
} else {
|
||||
let got = self.load().await?;
|
||||
if got == 0 {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// 把恰好 n 字节从缓冲/底层流转发到 w,小块直写,不整体缓冲。
|
||||
pub async fn forward_n<W: AsyncWrite + Unpin>(
|
||||
&mut self,
|
||||
w: &mut W,
|
||||
mut n: u64,
|
||||
) -> io::Result<u64> {
|
||||
let mut sent = 0u64;
|
||||
while n > 0 {
|
||||
if self.start < self.end {
|
||||
let avail = (self.end - self.start) as u64;
|
||||
let take = avail.min(n) as usize;
|
||||
w.write_all(&self.buf[self.start..self.start + take]).await?;
|
||||
self.start += take;
|
||||
sent += take as u64;
|
||||
n -= take as u64;
|
||||
} else {
|
||||
let got = self.load().await?;
|
||||
if got == 0 {
|
||||
return Err(proto_err(format!(
|
||||
"还需转发 {} 字节,但上游提前关闭",
|
||||
n
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
w.flush().await?;
|
||||
Ok(sent)
|
||||
}
|
||||
|
||||
/// 转发到对端关闭为止。
|
||||
pub async fn forward_to_eof<W: AsyncWrite + Unpin>(&mut self, w: &mut W) -> io::Result<u64> {
|
||||
let mut sent = 0u64;
|
||||
loop {
|
||||
if self.start < self.end {
|
||||
let take = self.end - self.start;
|
||||
w.write_all(&self.buf[self.start..self.end]).await?;
|
||||
self.start = self.end;
|
||||
sent += take as u64;
|
||||
} else {
|
||||
match self.load().await {
|
||||
Ok(0) => break,
|
||||
Ok(_) => {}
|
||||
Err(e) if is_tls_eof_without_notify(&e) => break,
|
||||
Err(e) => return Err(e),
|
||||
}
|
||||
}
|
||||
}
|
||||
w.flush().await?;
|
||||
Ok(sent)
|
||||
}
|
||||
|
||||
/// chunked 流原样逐块转发(保留分块边界,每块 flush,SSE 延迟不受影响)。
|
||||
pub async fn forward_chunked<W: AsyncWrite + Unpin>(
|
||||
&mut self,
|
||||
w: &mut W,
|
||||
) -> io::Result<u64> {
|
||||
let mut sent = 0u64;
|
||||
loop {
|
||||
let line = self.read_line(MAX_LINE).await?;
|
||||
let size = parse_chunk_size(&line)?;
|
||||
w.write_all(&line).await?;
|
||||
sent += line.len() as u64;
|
||||
if size == 0 {
|
||||
// 末块没有数据与 CRLF,直接是 trailer 区:原样转发直到空行
|
||||
loop {
|
||||
let t = self.read_line(MAX_LINE).await?;
|
||||
w.write_all(&t).await?;
|
||||
sent += t.len() as u64;
|
||||
if t == b"\r\n" || t == b"\n" {
|
||||
break;
|
||||
}
|
||||
}
|
||||
w.flush().await?;
|
||||
return Ok(sent);
|
||||
}
|
||||
self.forward_n(w, size).await?;
|
||||
sent += size;
|
||||
let crlf = self.read_exact_n(2).await?;
|
||||
if crlf != b"\r\n" {
|
||||
return Err(proto_err("chunk 数据后缺少 CRLF"));
|
||||
}
|
||||
w.write_all(&crlf).await?;
|
||||
sent += 2;
|
||||
w.flush().await?;
|
||||
}
|
||||
}
|
||||
|
||||
/// 解码 chunked 请求体(与转发不同,trailers 丢弃)。
|
||||
pub async fn read_chunked_decode(&mut self, limit: u64) -> io::Result<Vec<u8>> {
|
||||
let mut body = Vec::new();
|
||||
loop {
|
||||
let line = self.read_line(MAX_LINE).await?;
|
||||
let size = parse_chunk_size(&line)?;
|
||||
if size == 0 {
|
||||
loop {
|
||||
let t = self.read_line(MAX_LINE).await?;
|
||||
if t == b"\r\n" || t == b"\n" {
|
||||
break;
|
||||
}
|
||||
}
|
||||
return Ok(body);
|
||||
}
|
||||
if body.len() as u64 + size > limit {
|
||||
return Err(proto_err(format!("请求体超过 {} 字节上限", limit)));
|
||||
}
|
||||
let chunk = self.read_exact_n(size as usize).await?;
|
||||
body.extend_from_slice(&chunk);
|
||||
let crlf = self.read_exact_n(2).await?;
|
||||
if crlf != b"\r\n" {
|
||||
return Err(proto_err("chunk 数据后缺少 CRLF"));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_chunk_size(line: &[u8]) -> io::Result<u64> {
|
||||
let s = std::str::from_utf8(line).map_err(|_| proto_err("chunk 行非 UTF-8"))?;
|
||||
let s = s.trim_end_matches(['\r', '\n']);
|
||||
let hex = s.split(';').next().unwrap_or("").trim();
|
||||
u64::from_str_radix(hex, 16).map_err(|_| proto_err(format!("非法 chunk 大小: {hex:?}")))
|
||||
}
|
||||
|
||||
// ---- 请求 / 响应解析与序列化 --------------------------------------------------
|
||||
|
||||
/// 读取一个完整的下游请求。连接干净关闭(EOF)返回 None。
|
||||
/// 客户端带 Expect: 100-continue 且确有 body 时,先代答 100 再收 body。
|
||||
pub async fn read_request<R, W>(r: &mut ConnReader<R>, w: &mut W) -> io::Result<Option<Request>>
|
||||
where
|
||||
R: AsyncRead + Unpin,
|
||||
W: AsyncWrite + Unpin,
|
||||
{
|
||||
let line = match r.read_line_opt(MAX_LINE).await? {
|
||||
None => return Ok(None),
|
||||
Some(l) => l,
|
||||
};
|
||||
let mut head_total = line.len();
|
||||
let line_str = String::from_utf8_lossy(&line);
|
||||
let mut parts = line_str.split_whitespace();
|
||||
let method = parts
|
||||
.next()
|
||||
.ok_or_else(|| proto_err("空请求行"))?
|
||||
.to_string();
|
||||
let target = parts
|
||||
.next()
|
||||
.ok_or_else(|| proto_err("请求行缺少路径"))?
|
||||
.to_string();
|
||||
let version = parts.next().unwrap_or("HTTP/1.0").to_string();
|
||||
if !version.starts_with("HTTP/") {
|
||||
return Err(proto_err(format!("不支持的协议版本: {version}")));
|
||||
}
|
||||
|
||||
let mut headers = Vec::new();
|
||||
loop {
|
||||
let h = r.read_line(MAX_LINE).await?;
|
||||
head_total += h.len();
|
||||
if head_total > MAX_HEAD {
|
||||
return Err(proto_err("请求头超过上限"));
|
||||
}
|
||||
let hs = String::from_utf8_lossy(&h);
|
||||
let hs = hs.trim_end_matches(['\r', '\n']);
|
||||
if hs.is_empty() {
|
||||
break;
|
||||
}
|
||||
let (name, value) = hs
|
||||
.split_once(':')
|
||||
.ok_or_else(|| proto_err(format!("非法头行: {hs:?}")))?;
|
||||
headers.push((name.trim().to_string(), value.trim().to_string()));
|
||||
}
|
||||
|
||||
let expect_continue = header_get(&headers, "expect")
|
||||
.map(|v| v.eq_ignore_ascii_case("100-continue"))
|
||||
.unwrap_or(false);
|
||||
async fn send_continue<W: AsyncWrite + Unpin>(w: &mut W) -> io::Result<()> {
|
||||
w.write_all(b"HTTP/1.1 100 Continue\r\n\r\n").await?;
|
||||
w.flush().await
|
||||
}
|
||||
|
||||
let mut body = Vec::new();
|
||||
if let Some(te) = header_get(&headers, "transfer-encoding") {
|
||||
let is_chunked = te
|
||||
.split(',')
|
||||
.any(|p| p.trim().eq_ignore_ascii_case("chunked"));
|
||||
if !is_chunked {
|
||||
return Err(proto_err(format!("不支持的 Transfer-Encoding: {te}")));
|
||||
}
|
||||
if expect_continue {
|
||||
send_continue(w).await?;
|
||||
}
|
||||
body = r.read_chunked_decode(MAX_BODY).await?;
|
||||
} else if let Some(cl) = header_get(&headers, "content-length") {
|
||||
let n: u64 = cl
|
||||
.parse()
|
||||
.map_err(|_| proto_err(format!("非法 Content-Length: {cl}")))?;
|
||||
if n > MAX_BODY {
|
||||
return Err(proto_err(format!("请求体超过 {} 字节上限", MAX_BODY)));
|
||||
}
|
||||
if expect_continue && n > 0 {
|
||||
send_continue(w).await?;
|
||||
}
|
||||
if n > 0 {
|
||||
body = r.read_exact_n(n as usize).await?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(Some(Request {
|
||||
method,
|
||||
target,
|
||||
version,
|
||||
headers,
|
||||
body,
|
||||
}))
|
||||
}
|
||||
|
||||
pub async fn read_response_head<R: AsyncRead + Unpin>(
|
||||
r: &mut ConnReader<R>,
|
||||
) -> io::Result<Response> {
|
||||
let line = r.read_line(MAX_LINE).await?;
|
||||
let line_str = String::from_utf8_lossy(&line);
|
||||
let mut parts = line_str.split_whitespace();
|
||||
if parts.next().is_none() {
|
||||
return Err(proto_err("上游返回空响应"));
|
||||
}
|
||||
let status = parts
|
||||
.next()
|
||||
.and_then(|s| s.parse::<u16>().ok())
|
||||
.ok_or_else(|| proto_err(format!("非法状态行: {}", line_str.trim())))?;
|
||||
let reason = parts.collect::<Vec<_>>().join(" ");
|
||||
|
||||
let mut headers = Vec::new();
|
||||
let mut head_total = 0usize;
|
||||
loop {
|
||||
let h = r.read_line(MAX_LINE).await?;
|
||||
head_total += h.len();
|
||||
if head_total > MAX_HEAD {
|
||||
return Err(proto_err("上游响应头超过上限"));
|
||||
}
|
||||
let hs = String::from_utf8_lossy(&h);
|
||||
let hs = hs.trim_end_matches(['\r', '\n']);
|
||||
if hs.is_empty() {
|
||||
break;
|
||||
}
|
||||
if let Some((name, value)) = hs.split_once(':') {
|
||||
headers.push((name.trim().to_string(), value.trim().to_string()));
|
||||
}
|
||||
}
|
||||
Ok(Response {
|
||||
status,
|
||||
reason,
|
||||
headers,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn write_request_head<W: AsyncWrite + Unpin>(
|
||||
w: &mut W,
|
||||
method: &str,
|
||||
target: &str,
|
||||
headers: &[(String, String)],
|
||||
) -> io::Result<()> {
|
||||
let mut head = Vec::with_capacity(256 + headers.len() * 48);
|
||||
head.extend_from_slice(format!("{method} {target} HTTP/1.1\r\n").as_bytes());
|
||||
for (n, v) in headers {
|
||||
head.extend_from_slice(format!("{n}: {v}\r\n").as_bytes());
|
||||
}
|
||||
head.extend_from_slice(b"\r\n");
|
||||
w.write_all(&head).await?;
|
||||
w.flush().await
|
||||
}
|
||||
|
||||
pub async fn write_response_head<W: AsyncWrite + Unpin>(
|
||||
w: &mut W,
|
||||
version: &str,
|
||||
status: u16,
|
||||
reason: &str,
|
||||
headers: &[(String, String)],
|
||||
) -> io::Result<()> {
|
||||
let reason = if reason.is_empty() {
|
||||
status_reason(status)
|
||||
} else {
|
||||
reason
|
||||
};
|
||||
let mut head = Vec::with_capacity(256 + headers.len() * 48);
|
||||
head.extend_from_slice(format!("{version} {status} {reason}\r\n").as_bytes());
|
||||
for (n, v) in headers {
|
||||
head.extend_from_slice(format!("{n}: {v}\r\n").as_bytes());
|
||||
}
|
||||
head.extend_from_slice(b"\r\n");
|
||||
w.write_all(&head).await?;
|
||||
w.flush().await
|
||||
}
|
||||
|
||||
pub fn status_reason(code: u16) -> &'static str {
|
||||
match code {
|
||||
200 => "OK",
|
||||
201 => "Created",
|
||||
204 => "No Content",
|
||||
400 => "Bad Request",
|
||||
401 => "Unauthorized",
|
||||
403 => "Forbidden",
|
||||
404 => "Not Found",
|
||||
405 => "Method Not Allowed",
|
||||
408 => "Request Timeout",
|
||||
413 => "Payload Too Large",
|
||||
429 => "Too Many Requests",
|
||||
500 => "Internal Server Error",
|
||||
502 => "Bad Gateway",
|
||||
503 => "Service Unavailable",
|
||||
504 => "Gateway Timeout",
|
||||
_ => "Unknown",
|
||||
}
|
||||
}
|
||||
|
||||
/// 代理自身产生的简短错误响应(带 Connection: close)。
|
||||
pub async fn write_simple_response<W: AsyncWrite + Unpin>(
|
||||
w: &mut W,
|
||||
status: u16,
|
||||
msg: &str,
|
||||
) -> io::Result<()> {
|
||||
let body = format!("{msg}\n");
|
||||
let headers = vec![
|
||||
("Content-Type".to_string(), "text/plain; charset=utf-8".to_string()),
|
||||
("Content-Length".to_string(), body.len().to_string()),
|
||||
("Connection".to_string(), "close".to_string()),
|
||||
];
|
||||
write_response_head(w, "HTTP/1.1", status, status_reason(status), &headers).await?;
|
||||
w.write_all(body.as_bytes()).await?;
|
||||
w.flush().await
|
||||
}
|
||||
|
||||
// ---- 测试 ---------------------------------------------------------------------
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
async fn req_from(raw: &[u8]) -> io::Result<Option<Request>> {
|
||||
let mut r = ConnReader::new(&raw[..], 16);
|
||||
let mut w = Vec::new();
|
||||
read_request(&mut r, &mut w).await
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn parse_get_without_body() {
|
||||
let req = req_from(b"GET /v1/models HTTP/1.1\r\nHost: x\r\nAccept: */*\r\n\r\n")
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(req.method, "GET");
|
||||
assert_eq!(req.target, "/v1/models");
|
||||
assert_eq!(req.version, "HTTP/1.1");
|
||||
assert_eq!(req.body.len(), 0);
|
||||
assert_eq!(header_get(&req.headers, "HOST"), Some("x"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn parse_post_with_content_length() {
|
||||
let raw = b"POST /v1/chat/completions HTTP/1.1\r\nContent-Length: 5\r\n\r\nhello";
|
||||
let req = req_from(raw).await.unwrap().unwrap();
|
||||
assert_eq!(req.body, b"hello");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn parse_chunked_request() {
|
||||
let raw = b"POST /v1/messages HTTP/1.1\r\nTransfer-Encoding: chunked\r\n\r\n\
|
||||
5\r\nhello\r\n3;x=1\r\n wo\r\n0\r\nX-Trailer: t\r\n\r\n";
|
||||
let req = req_from(raw).await.unwrap().unwrap();
|
||||
assert_eq!(req.body, b"hello wo");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn parse_expect_continue_writes_100() {
|
||||
let raw = b"POST /v1/chat/completions HTTP/1.1\r\nExpect: 100-continue\r\nContent-Length: 2\r\n\r\nhi";
|
||||
let mut r = ConnReader::new(&raw[..], 16);
|
||||
let mut w = Vec::new();
|
||||
let req = read_request(&mut r, &mut w).await.unwrap().unwrap();
|
||||
assert_eq!(req.body, b"hi");
|
||||
assert_eq!(w, b"HTTP/1.1 100 Continue\r\n\r\n");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn eof_returns_none() {
|
||||
assert!(req_from(b"").await.unwrap().is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn body_over_limit_rejected() {
|
||||
let raw = b"POST /x HTTP/1.1\r\nContent-Length: 99999999\r\n\r\n";
|
||||
assert!(req_from(raw).await.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn response_head_and_kind() {
|
||||
let raw = b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nTransfer-Encoding: chunked\r\n\r\n";
|
||||
let mut r = ConnReader::new(&raw[..], 16);
|
||||
let resp = read_response_head(&mut r).await.unwrap();
|
||||
assert_eq!(resp.status, 200);
|
||||
assert_eq!(resp.reason, "OK");
|
||||
assert_eq!(resp.header("content-TYPE"), Some("text/event-stream"));
|
||||
assert_eq!(resp.body_kind().unwrap(), BodyKind::Chunked);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn response_content_length_kind() {
|
||||
let raw = b"HTTP/1.1 429 Too Many Requests\r\nContent-Length: 3\r\n\r\n";
|
||||
let mut r = ConnReader::new(&raw[..], 16);
|
||||
let resp = read_response_head(&mut r).await.unwrap();
|
||||
assert_eq!(resp.body_kind().unwrap(), BodyKind::Exact(3));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn forward_chunked_roundtrip() {
|
||||
let payload: &[u8] = b"8\r\nhello, w\r\n4\r\norld\r\n0\r\nX-T: 1\r\n\r\n";
|
||||
let mut r = ConnReader::new(payload, 4);
|
||||
let mut w = Vec::new();
|
||||
let n = r.forward_chunked(&mut w).await.unwrap();
|
||||
assert_eq!(w, payload);
|
||||
assert_eq!(n, payload.len() as u64);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn forward_chunked_without_trailers() {
|
||||
let payload: &[u8] = b"3\r\nabc\r\n0\r\n\r\n";
|
||||
let mut r = ConnReader::new(payload, 64);
|
||||
let mut w = Vec::new();
|
||||
r.forward_chunked(&mut w).await.unwrap();
|
||||
assert_eq!(w, payload);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn forward_exact_and_eof() {
|
||||
let mut r = ConnReader::new(&b"0123456789"[..], 3);
|
||||
let mut w = Vec::new();
|
||||
let n = r.forward_n(&mut w, 4).await.unwrap();
|
||||
assert_eq!(&w[..], b"0123");
|
||||
assert_eq!(n, 4);
|
||||
let mut w2 = Vec::new();
|
||||
let n = r.forward_to_eof(&mut w2).await.unwrap();
|
||||
assert_eq!(&w2[..], b"456789");
|
||||
assert_eq!(n, 6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chunk_size_parsing() {
|
||||
assert_eq!(parse_chunk_size(b"1a\r\n").unwrap(), 26);
|
||||
assert_eq!(parse_chunk_size(b"0;last\r\n").unwrap(), 0);
|
||||
assert!(parse_chunk_size(b"zz\r\n").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hop_by_hop_detection() {
|
||||
assert!(is_hop_by_hop("Transfer-Encoding"));
|
||||
assert!(!is_hop_by_hop("Content-Type"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn simple_response_writes_close() {
|
||||
let mut w = Vec::new();
|
||||
write_simple_response(&mut w, 502, "boom").await.unwrap();
|
||||
let s = String::from_utf8(w).unwrap();
|
||||
assert!(s.starts_with("HTTP/1.1 502 Bad Gateway\r\n"));
|
||||
assert!(s.contains("Connection: close\r\n"));
|
||||
assert!(s.ends_with("boom\n"));
|
||||
}
|
||||
}
|
||||
+334
@@ -0,0 +1,334 @@
|
||||
//! systemd 服务安装(-i)与卸载(-u)。
|
||||
//!
|
||||
//! 安装时复制二进制到 `/usr/bin/opencode-go-proxy`,并把 `-p` / `-k` /
|
||||
//! `-s` / `-a` / `-u` 固化进 unit。全机只允许一个实例:已安装则拒绝再装。
|
||||
|
||||
use std::fs;
|
||||
use std::net::TcpListener;
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
use std::path::Path;
|
||||
use std::process::{Command, exit};
|
||||
|
||||
use crate::VERSION;
|
||||
|
||||
const UNIT_NAME: &str = "opencode-go-proxy";
|
||||
const UNIT_PATH: &str = "/etc/systemd/system/opencode-go-proxy.service";
|
||||
const BIN_PATH: &str = "/usr/bin/opencode-go-proxy";
|
||||
|
||||
/// 安装参数(不含 -i 自身;None 的项不写入 unit,沿用程序默认值)。
|
||||
pub struct InstallSpec {
|
||||
pub port: u16,
|
||||
/// 上游 OpenCode Go API 密钥(必填)
|
||||
pub key: String,
|
||||
/// 客户端访问密钥;None 时安装过程自动生成并打印
|
||||
pub client_key: Option<String>,
|
||||
pub session: Option<String>,
|
||||
pub user_agent: Option<String>,
|
||||
pub upstream: Option<String>,
|
||||
}
|
||||
|
||||
/// 生成 "ogp-<32hex>" 形式的客户端访问密钥,读 /dev/urandom,退化用纳秒时间戳。
|
||||
pub fn generate_client_key() -> String {
|
||||
use std::io::Read;
|
||||
let mut buf = [0u8; 16];
|
||||
match std::fs::File::open("/dev/urandom").as_mut() {
|
||||
Ok(f) => {
|
||||
if f.read_exact(&mut buf).is_err() {
|
||||
let nanos = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.subsec_nanos();
|
||||
buf[..4].copy_from_slice(&nanos.to_le_bytes());
|
||||
}
|
||||
}
|
||||
Err(_) => {
|
||||
let nanos = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.subsec_nanos();
|
||||
buf[..4].copy_from_slice(&nanos.to_le_bytes());
|
||||
}
|
||||
}
|
||||
let hex: String = buf.iter().map(|b| format!("{b:02x}")).collect();
|
||||
format!("ogp-{hex}")
|
||||
}
|
||||
|
||||
/// systemd ExecStart 的参数引用:用双引号包裹并转义,防止空格/引号破坏解析。
|
||||
fn quote(s: &str) -> String {
|
||||
let mut q = String::with_capacity(s.len() + 2);
|
||||
q.push('"');
|
||||
for c in s.chars() {
|
||||
if c == '"' || c == '\\' {
|
||||
q.push('\\');
|
||||
}
|
||||
q.push(c);
|
||||
}
|
||||
q.push('"');
|
||||
q
|
||||
}
|
||||
|
||||
fn require_systemd() {
|
||||
if !Path::new("/run/systemd/system").exists() {
|
||||
eprintln!("opencode-go-proxy: 未检测到 systemd(/run/systemd/system 不存在),无法操作服务。");
|
||||
exit(1);
|
||||
}
|
||||
}
|
||||
|
||||
fn systemctl(args: &[&str]) -> Result<(), String> {
|
||||
match Command::new("systemctl").args(args).output() {
|
||||
Ok(o) if o.status.success() => Ok(()),
|
||||
Ok(o) => {
|
||||
let msg = String::from_utf8_lossy(&o.stderr).trim().to_string();
|
||||
Err(if msg.is_empty() {
|
||||
format!("systemctl {} 失败", args.join(" "))
|
||||
} else {
|
||||
msg
|
||||
})
|
||||
}
|
||||
Err(e) => Err(format!("无法执行 systemctl({e}),请确认系统已安装 systemd")),
|
||||
}
|
||||
}
|
||||
|
||||
fn die(msg: String) -> ! {
|
||||
eprintln!("opencode-go-proxy: {msg}");
|
||||
exit(1)
|
||||
}
|
||||
|
||||
/// 探测端口是否可绑定(TCP 0.0.0.0)。
|
||||
pub fn port_available(port: u16) -> bool {
|
||||
TcpListener::bind(("0.0.0.0", port)).is_ok()
|
||||
}
|
||||
|
||||
fn already_installed() -> bool {
|
||||
Path::new(UNIT_PATH).exists()
|
||||
}
|
||||
|
||||
/// 复制当前可执行文件到 dest,并设置 0755 权限(已存在则覆盖)。
|
||||
pub fn copy_self_to(dest: &Path) -> std::io::Result<()> {
|
||||
let src = std::env::current_exe()?;
|
||||
fs::copy(&src, dest)?;
|
||||
fs::set_permissions(dest, fs::Permissions::from_mode(0o755))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 生成 systemd unit 文件内容。
|
||||
pub fn build_unit(spec: &InstallSpec) -> String {
|
||||
let mut exec = format!("{BIN_PATH} -p {}", spec.port);
|
||||
exec.push_str(&format!(" -k {}", quote(&spec.key)));
|
||||
if let Some(c) = &spec.client_key {
|
||||
exec.push_str(&format!(" -c {}", quote(c)));
|
||||
}
|
||||
if let Some(s) = &spec.session {
|
||||
exec.push_str(&format!(" -s {}", quote(s)));
|
||||
}
|
||||
if let Some(a) = &spec.user_agent {
|
||||
exec.push_str(&format!(" -a {}", quote(a)));
|
||||
}
|
||||
if let Some(u) = &spec.upstream {
|
||||
exec.push_str(&format!(" -u {}", quote(u)));
|
||||
}
|
||||
format!(
|
||||
"# 本文件由 `opencode-go-proxy -i` (v{VERSION}) 生成;需先 -u 再重装
|
||||
[Unit]
|
||||
Description=opencode-go-proxy - OpenCode Go API 本地转发代理 (v{VERSION})
|
||||
After=network-online.target
|
||||
Wants=network-online.target
|
||||
|
||||
[Service]
|
||||
ExecStart={exec}
|
||||
Restart=on-failure
|
||||
RestartSec=3s
|
||||
LimitNOFILE=65535
|
||||
NoNewPrivileges=true
|
||||
|
||||
[Install]
|
||||
WantedBy=multi-user.target
|
||||
"
|
||||
)
|
||||
}
|
||||
|
||||
/// -i:检查未安装 → 检查端口 → 复制二进制 → 写 unit → enable/restart。
|
||||
/// 未提供客户端密钥时自动生成并在安装输出中展示。
|
||||
pub fn install(spec: InstallSpec) {
|
||||
require_systemd();
|
||||
|
||||
if already_installed() {
|
||||
die(format!(
|
||||
"已安装 {UNIT_NAME}.service({UNIT_PATH} 存在),全机仅允许一个实例。\n\
|
||||
请先执行 `sudo opencode-go-proxy -u` 卸载后再安装。"
|
||||
));
|
||||
}
|
||||
if !port_available(spec.port) {
|
||||
die(format!("端口 {} 已被占用,请换一个端口或先释放后再装。", spec.port));
|
||||
}
|
||||
|
||||
let (client_key, client_key_generated) = match spec.client_key {
|
||||
Some(k) => (k, false),
|
||||
None => (generate_client_key(), true),
|
||||
};
|
||||
let unit_spec = InstallSpec {
|
||||
client_key: Some(client_key.clone()),
|
||||
..spec
|
||||
};
|
||||
|
||||
if let Err(e) = copy_self_to(Path::new(BIN_PATH)) {
|
||||
if e.kind() == std::io::ErrorKind::PermissionDenied {
|
||||
die(format!("复制二进制到 {BIN_PATH} 权限不足,请使用 root/sudo 运行。"));
|
||||
}
|
||||
die(format!("复制二进制到 {BIN_PATH} 失败: {e}"));
|
||||
}
|
||||
|
||||
if let Err(e) = fs::write(UNIT_PATH, build_unit(&unit_spec)) {
|
||||
if e.kind() == std::io::ErrorKind::PermissionDenied {
|
||||
die(format!("写入 {UNIT_PATH} 权限不足,请使用 root/sudo 运行。"));
|
||||
}
|
||||
die(format!("写入 {UNIT_PATH} 失败: {e}"));
|
||||
}
|
||||
|
||||
if let Err(e) = systemctl(&["daemon-reload"]) {
|
||||
die(e);
|
||||
}
|
||||
if let Err(e) = systemctl(&["enable", UNIT_NAME]) {
|
||||
die(e);
|
||||
}
|
||||
if let Err(e) = systemctl(&["restart", UNIT_NAME]) {
|
||||
die(format!(
|
||||
"{e}\n(可运行 `journalctl -u {UNIT_NAME} -e` 查看日志,常见原因:端口被占用)"
|
||||
));
|
||||
}
|
||||
|
||||
println!(
|
||||
"服务安装完成并已启动:\n \
|
||||
服务: {UNIT_NAME}.service\n \
|
||||
程序: {BIN_PATH}\n \
|
||||
监听: 127.0.0.1:{}\n \
|
||||
上游: {}\n \
|
||||
健康检查: http://127.0.0.1:{}/healthz\n \
|
||||
客户端密钥: {}{}\n \
|
||||
(客户端以 Authorization: Bearer <客户端密钥> 或 x-api-key 携带)\n\n\
|
||||
常用命令:\n \
|
||||
systemctl status opencode-go-proxy 查看运行状态\n \
|
||||
journalctl -u opencode-go-proxy -f 实时查看日志\n \
|
||||
systemctl restart opencode-go-proxy 重启服务(修改 unit 中的 -k 后需要)\n \
|
||||
sudo opencode-go-proxy -u 停止并卸载服务",
|
||||
unit_spec.port,
|
||||
unit_spec.upstream.as_deref().unwrap_or(crate::DEFAULT_UPSTREAM),
|
||||
unit_spec.port,
|
||||
client_key,
|
||||
if client_key_generated { " (自动生成,已固化进 unit)" } else { " (来自 -c)" },
|
||||
);
|
||||
}
|
||||
|
||||
/// -u:停止 → 禁用 → 删除 unit 文件 → daemon-reload → 删除 /usr/bin 下的二进制。
|
||||
pub fn uninstall() {
|
||||
require_systemd();
|
||||
|
||||
let _ = Command::new("systemctl").args(["stop", UNIT_NAME]).output();
|
||||
let _ = Command::new("systemctl").args(["disable", UNIT_NAME]).output();
|
||||
|
||||
match fs::remove_file(UNIT_PATH) {
|
||||
Ok(()) => {}
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
|
||||
println!("opencode-go-proxy: 服务未安装({UNIT_PATH} 不存在),无需卸载。");
|
||||
return;
|
||||
}
|
||||
Err(e) if e.kind() == std::io::ErrorKind::PermissionDenied => {
|
||||
die(format!("删除 {UNIT_PATH} 权限不足,请使用 root/sudo 运行。"));
|
||||
}
|
||||
Err(e) => die(format!("删除 {UNIT_PATH} 失败: {e}")),
|
||||
}
|
||||
|
||||
let _ = Command::new("systemctl").args(["daemon-reload"]).output();
|
||||
|
||||
match fs::remove_file(BIN_PATH) {
|
||||
Ok(()) => println!("服务已停止并卸载,{BIN_PATH} 已删除。"),
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
|
||||
println!("服务已停止并卸载({BIN_PATH} 不存在,跳过删除)。");
|
||||
}
|
||||
Err(e) if e.kind() == std::io::ErrorKind::PermissionDenied => {
|
||||
die(format!("删除 {BIN_PATH} 权限不足,请手动删除(root/sudo)。"));
|
||||
}
|
||||
Err(e) => die(format!("删除 {BIN_PATH} 失败: {e}")),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn spec() -> InstallSpec {
|
||||
InstallSpec {
|
||||
port: 28080,
|
||||
key: "sk-test".into(),
|
||||
client_key: None,
|
||||
session: None,
|
||||
user_agent: None,
|
||||
upstream: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unit_contains_fixed_args() {
|
||||
let unit = build_unit(&spec());
|
||||
let line = unit.lines().find(|l| l.starts_with("ExecStart=")).expect("应有 ExecStart 行");
|
||||
assert_eq!(line, "ExecStart=/usr/bin/opencode-go-proxy -p 28080 -k \"sk-test\"");
|
||||
assert!(unit.contains("WantedBy=multi-user.target"));
|
||||
assert!(unit.contains("Restart=on-failure"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unit_contains_client_key_when_set() {
|
||||
let s = InstallSpec { client_key: Some("ogp-abc".into()), ..spec() };
|
||||
let line = build_unit(&s)
|
||||
.lines()
|
||||
.find(|l| l.starts_with("ExecStart="))
|
||||
.unwrap()
|
||||
.to_string();
|
||||
assert!(line.contains("-c \"ogp-abc\""), "unit 应固化 -c: {line}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn generated_client_key_shape() {
|
||||
let k = generate_client_key();
|
||||
assert!(k.starts_with("ogp-"));
|
||||
assert_eq!(k.len(), "ogp-".len() + 32);
|
||||
assert_ne!(generate_client_key(), generate_client_key());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unit_quotes_special_chars() {
|
||||
let s = InstallSpec { key: "a\"b\\c".into(), ..spec() };
|
||||
let unit = build_unit(&s);
|
||||
assert!(unit.contains("-k \"a\\\"b\\\\c\""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unit_includes_optional_args() {
|
||||
let s = InstallSpec {
|
||||
session: Some("sess-fixed".into()),
|
||||
user_agent: Some("my-agent/1".into()),
|
||||
upstream: Some("https://example.com/zen/go".into()),
|
||||
..spec()
|
||||
};
|
||||
let unit = build_unit(&s);
|
||||
let line = unit.lines().find(|l| l.starts_with("ExecStart=")).unwrap();
|
||||
assert!(line.contains("-s \"sess-fixed\""));
|
||||
assert!(line.contains("-a \"my-agent/1\""));
|
||||
assert!(line.contains("-u \"https://example.com/zen/go\""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn quote_escapes() {
|
||||
assert_eq!(quote("plain"), "\"plain\"");
|
||||
assert_eq!(quote("a\"b"), "\"a\\\"b\"");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn port_available_roundtrip() {
|
||||
let l = TcpListener::bind(("0.0.0.0", 0)).unwrap();
|
||||
let port = l.local_addr().unwrap().port();
|
||||
assert!(!port_available(port));
|
||||
drop(l);
|
||||
assert!(port_available(port));
|
||||
}
|
||||
}
|
||||
+355
@@ -0,0 +1,355 @@
|
||||
// opencode-go-proxy — OpenCode Go API 本地转发代理
|
||||
//
|
||||
// 透明转发三个协议端点(/v1/responses、/v1/chat/completions、/v1/messages)
|
||||
// 与 /v1/models,自动注入 OpenCode Go 要求的认证、User-Agent 与
|
||||
// x-opencode-session 会话头,让无法自定义请求头的客户端也能直接使用。
|
||||
|
||||
mod http;
|
||||
mod install;
|
||||
mod proxy;
|
||||
|
||||
use std::net::IpAddr;
|
||||
use std::process::exit;
|
||||
use std::sync::Arc;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
pub const VERSION: &str = env!("CARGO_PKG_VERSION");
|
||||
|
||||
pub const DEFAULT_UPSTREAM: &str = "https://opencode.ai/zen/go";
|
||||
const DEFAULT_PORT: u16 = 8090;
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! log_info {
|
||||
($($t:tt)*) => { eprintln!("{} [INFO] {}", $crate::now_str(), format_args!($($t)*)); };
|
||||
}
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! log_warn {
|
||||
($($t:tt)*) => { eprintln!("{} [WARN] {}", $crate::now_str(), format_args!($($t)*)); };
|
||||
}
|
||||
|
||||
/// UTC 时间戳,手写 civil-from-days 算法,不引入 chrono。
|
||||
pub fn now_str() -> String {
|
||||
let d = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default();
|
||||
let secs = d.as_secs() as i64;
|
||||
let (y, m, day) = civil_from_days(secs.div_euclid(86_400));
|
||||
let rem = secs.rem_euclid(86_400);
|
||||
format!(
|
||||
"{:04}-{:02}-{:02} {:02}:{:02}:{:02}",
|
||||
y,
|
||||
m,
|
||||
day,
|
||||
rem / 3600,
|
||||
rem % 3600 / 60,
|
||||
rem % 60
|
||||
)
|
||||
}
|
||||
|
||||
fn civil_from_days(z: i64) -> (i64, u32, u32) {
|
||||
let z = z + 719_468;
|
||||
let era = z.div_euclid(146_097);
|
||||
let doe = z.rem_euclid(146_097);
|
||||
let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
|
||||
let y = yoe + era * 400;
|
||||
let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
|
||||
let mp = (5 * doy + 2) / 153;
|
||||
let d = (doy - (153 * mp + 2) / 5 + 1) as u32;
|
||||
let m = (if mp < 10 { mp + 3 } else { mp - 9 }) as u32;
|
||||
(if m <= 2 { y + 1 } else { y }, m, d)
|
||||
}
|
||||
|
||||
/// 生成 "sess-<16hex>" 形式的会话 ID,读 /dev/urandom,退化用纳秒时间戳。
|
||||
pub fn generate_session_id() -> String {
|
||||
use std::io::Read;
|
||||
let mut buf = [0u8; 8];
|
||||
match std::fs::File::open("/dev/urandom").as_mut() {
|
||||
Ok(f) => {
|
||||
if f.read_exact(&mut buf).is_err() {
|
||||
let nanos = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.subsec_nanos();
|
||||
buf.copy_from_slice(&nanos.to_le_bytes());
|
||||
}
|
||||
}
|
||||
Err(_) => {
|
||||
let nanos = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.subsec_nanos();
|
||||
buf.copy_from_slice(&nanos.to_le_bytes());
|
||||
}
|
||||
}
|
||||
let hex: String = buf.iter().map(|b| format!("{b:02x}")).collect();
|
||||
format!("sess-{hex}")
|
||||
}
|
||||
|
||||
fn usage() -> String {
|
||||
format!(
|
||||
"opencode-go-proxy {VERSION} — OpenCode Go API 本地转发代理
|
||||
|
||||
用法: opencode-go-proxy [选项]
|
||||
|
||||
-l <地址> 监听地址 (默认 127.0.0.1)
|
||||
-p <端口> 监听端口 (默认 {DEFAULT_PORT})
|
||||
-k <密钥> OpenCode Go API 密钥;缺省读环境变量
|
||||
OPENCODE_GO_API_KEY 或 OPENCODE_API_KEY
|
||||
-s <会话ID> 固定 x-opencode-session;缺省为每个客户端连接自动生成
|
||||
-a <UA> 发送给上游的 User-Agent (默认 opencode-go-proxy/{VERSION})
|
||||
-u <地址> 上游地址 (默认 {DEFAULT_UPSTREAM};
|
||||
支持 http:// 用于本地调试,此时不做 TLS)
|
||||
-c <密钥> 客户端访问密钥:开启鉴权,调用方须以 Authorization: Bearer
|
||||
<密钥> 或 x-api-key 携带;-i 安装时缺省则自动生成
|
||||
-i 安装为 systemd 服务(root;固化 -p/-k/-c/-s/-a/-u 进 unit)
|
||||
--uninstall 停止并卸载 systemd 服务 (root)
|
||||
-h 显示本帮助
|
||||
|
||||
转发路径:
|
||||
/v1/responses OpenAI Responses API(Grok 4.6、GPT 5.6 Luna、Muse Spark)
|
||||
/v1/chat/completions OpenAI 兼容(GLM、Kimi、DeepSeek、MiMo 等)
|
||||
/v1/messages Anthropic 兼容(MiniMax、Qwen 系列)
|
||||
/v1/models 模型列表
|
||||
|
||||
示例:
|
||||
opencode-go-proxy -k sk-xxx
|
||||
curl http://127.0.0.1:{DEFAULT_PORT}/v1/chat/completions -d '{{\"model\":\"kimi-k3\",\"messages\":[{{\"role\":\"user\",\"content\":\"hi\"}}]}}'
|
||||
"
|
||||
)
|
||||
}
|
||||
|
||||
struct Args {
|
||||
listen: IpAddr,
|
||||
port: u16,
|
||||
key: Option<String>,
|
||||
client_key: Option<String>,
|
||||
session: Option<String>,
|
||||
user_agent: Option<String>,
|
||||
upstream: Option<String>,
|
||||
install: bool,
|
||||
uninstall: bool,
|
||||
help: bool,
|
||||
}
|
||||
|
||||
fn parse_args<I: Iterator<Item = String>>(it: I) -> Result<Args, String> {
|
||||
let mut args = Args {
|
||||
listen: "127.0.0.1".parse().unwrap(),
|
||||
port: DEFAULT_PORT,
|
||||
key: None,
|
||||
client_key: None,
|
||||
session: None,
|
||||
user_agent: None,
|
||||
upstream: None,
|
||||
install: false,
|
||||
uninstall: false,
|
||||
help: false,
|
||||
};
|
||||
let mut it = it.peekable();
|
||||
while let Some(arg) = it.next() {
|
||||
let need_value = |it: &mut std::iter::Peekable<I>, name: &str| -> Result<String, String> {
|
||||
it.next()
|
||||
.filter(|v| !v.starts_with('-') || v.len() == 1)
|
||||
.ok_or_else(|| format!("{name} 缺少参数值"))
|
||||
};
|
||||
match arg.as_str() {
|
||||
"-l" => {
|
||||
let v = need_value(&mut it, "-l")?;
|
||||
args.listen = v
|
||||
.parse()
|
||||
.map_err(|_| format!("非法监听地址: {v}"))?;
|
||||
}
|
||||
"-p" => {
|
||||
let v = need_value(&mut it, "-p")?;
|
||||
args.port = v.parse().map_err(|_| format!("非法端口: {v}"))?;
|
||||
}
|
||||
"-k" => args.key = Some(need_value(&mut it, "-k")?),
|
||||
"-c" => args.client_key = Some(need_value(&mut it, "-c")?),
|
||||
"-s" => args.session = Some(need_value(&mut it, "-s")?),
|
||||
"-a" => args.user_agent = Some(need_value(&mut it, "-a")?),
|
||||
"-u" => args.upstream = Some(need_value(&mut it, "-u")?),
|
||||
"-i" => args.install = true,
|
||||
"--uninstall" | "-x" => args.uninstall = true,
|
||||
"-h" | "--help" => args.help = true,
|
||||
other => return Err(format!("未知参数: {other}")),
|
||||
}
|
||||
}
|
||||
if args.install && args.uninstall {
|
||||
return Err("-i 与 --uninstall 不能同时使用".into());
|
||||
}
|
||||
Ok(args)
|
||||
}
|
||||
|
||||
/// 解析上游地址为 (是否 TLS, 主机, 端口, 路径前缀)。
|
||||
fn parse_upstream(s: &str) -> Result<proxy::Upstream, String> {
|
||||
let (tls, rest) = if let Some(r) = s.strip_prefix("https://") {
|
||||
(true, r)
|
||||
} else if let Some(r) = s.strip_prefix("http://") {
|
||||
(false, r)
|
||||
} else {
|
||||
(true, s)
|
||||
};
|
||||
let (hostport, path) = match rest.find('/') {
|
||||
Some(i) => (&rest[..i], rest[i..].trim_end_matches('/').to_string()),
|
||||
None => (rest, String::new()),
|
||||
};
|
||||
let (host, port) = match hostport.rsplit_once(':') {
|
||||
Some((h, p)) if !h.is_empty() && !p.is_empty() && p.bytes().all(|c| c.is_ascii_digit()) => {
|
||||
(h.to_string(), p.parse::<u16>().map_err(|_| format!("非法端口: {p}"))?)
|
||||
}
|
||||
_ => (hostport.to_string(), if tls { 443 } else { 80 }),
|
||||
};
|
||||
if host.is_empty() {
|
||||
return Err("上游地址缺少主机名".into());
|
||||
}
|
||||
Ok(proxy::Upstream {
|
||||
tls,
|
||||
host,
|
||||
port,
|
||||
base_path: path,
|
||||
})
|
||||
}
|
||||
|
||||
fn main() {
|
||||
let args = match parse_args(std::env::args().skip(1)) {
|
||||
Ok(a) => a,
|
||||
Err(e) => {
|
||||
eprintln!("opencode-go-proxy: {e}");
|
||||
eprintln!("用 -h 查看用法");
|
||||
exit(2);
|
||||
}
|
||||
};
|
||||
if args.help {
|
||||
println!("{}", usage());
|
||||
return;
|
||||
}
|
||||
|
||||
if args.uninstall {
|
||||
install::uninstall();
|
||||
return;
|
||||
}
|
||||
|
||||
let key = args
|
||||
.key
|
||||
.or_else(|| std::env::var("OPENCODE_GO_API_KEY").ok())
|
||||
.or_else(|| std::env::var("OPENCODE_API_KEY").ok());
|
||||
|
||||
if args.install {
|
||||
let key = key.unwrap_or_else(|| {
|
||||
eprintln!(
|
||||
"opencode-go-proxy: 安装需要 API 密钥;用 -k 指定,或设置环境变量 OPENCODE_GO_API_KEY"
|
||||
);
|
||||
exit(2);
|
||||
});
|
||||
let spec = install::InstallSpec {
|
||||
port: args.port,
|
||||
key,
|
||||
client_key: args.client_key,
|
||||
session: args.session,
|
||||
user_agent: args.user_agent,
|
||||
upstream: args.upstream,
|
||||
};
|
||||
install::install(spec);
|
||||
return;
|
||||
}
|
||||
|
||||
let key = key.unwrap_or_else(|| {
|
||||
eprintln!(
|
||||
"opencode-go-proxy: 未提供 API 密钥;用 -k 指定,或设置环境变量 OPENCODE_GO_API_KEY"
|
||||
);
|
||||
exit(2);
|
||||
});
|
||||
let upstream = match parse_upstream(args.upstream.as_deref().unwrap_or(DEFAULT_UPSTREAM)) {
|
||||
Ok(u) => u,
|
||||
Err(e) => {
|
||||
eprintln!("opencode-go-proxy: {e}");
|
||||
exit(2);
|
||||
}
|
||||
};
|
||||
|
||||
let cfg = Arc::new(proxy::Config {
|
||||
listen: args.listen,
|
||||
port: args.port,
|
||||
api_key: key,
|
||||
client_key: args.client_key,
|
||||
user_agent: args
|
||||
.user_agent
|
||||
.unwrap_or_else(|| format!("opencode-go-proxy/{VERSION}")),
|
||||
fixed_session: args.session,
|
||||
upstream,
|
||||
});
|
||||
|
||||
if let Err(e) = tokio::runtime::Builder::new_multi_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.expect("初始化 tokio 运行时失败")
|
||||
.block_on(proxy::run(cfg))
|
||||
{
|
||||
eprintln!("opencode-go-proxy: {e}");
|
||||
exit(1);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn upstream_parsing() {
|
||||
let u = parse_upstream("https://opencode.ai/zen/go").unwrap();
|
||||
assert!(u.tls);
|
||||
assert_eq!(u.host, "opencode.ai");
|
||||
assert_eq!(u.port, 443);
|
||||
assert_eq!(u.base_path, "/zen/go");
|
||||
|
||||
let u = parse_upstream("http://127.0.0.1:9000").unwrap();
|
||||
assert!(!u.tls);
|
||||
assert_eq!(u.port, 9000);
|
||||
assert_eq!(u.base_path, "");
|
||||
|
||||
let u = parse_upstream("opencode.ai").unwrap();
|
||||
assert!(u.tls);
|
||||
assert_eq!(u.host, "opencode.ai");
|
||||
assert_eq!(u.base_path, "");
|
||||
|
||||
assert!(parse_upstream("https://").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn arg_parsing() {
|
||||
let a = parse_args(
|
||||
["-p", "8080", "-k", "sk-1", "-l", "0.0.0.0", "-h"]
|
||||
.iter()
|
||||
.map(|s| s.to_string()),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(a.port, 8080);
|
||||
assert_eq!(a.key.as_deref(), Some("sk-1"));
|
||||
assert_eq!(a.listen.to_string(), "0.0.0.0");
|
||||
assert!(a.help);
|
||||
|
||||
assert!(parse_args(["-z"].iter().map(|s| s.to_string())).is_err());
|
||||
assert!(parse_args(["-k"].iter().map(|s| s.to_string())).is_err());
|
||||
assert!(parse_args(["-i", "--uninstall"].iter().map(|s| s.to_string())).is_err());
|
||||
let a = parse_args(["-x"].iter().map(|s| s.to_string())).unwrap();
|
||||
assert!(a.uninstall);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_id_shape() {
|
||||
let s = generate_session_id();
|
||||
assert!(s.starts_with("sess-"));
|
||||
assert_eq!(s.len(), "sess-".len() + 16);
|
||||
assert_ne!(generate_session_id(), generate_session_id());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn timestamp_format() {
|
||||
let t = now_str();
|
||||
// "YYYY-MM-DD HH:MM:SS"
|
||||
assert_eq!(t.len(), 19);
|
||||
assert_eq!(t.as_bytes()[4], b'-');
|
||||
assert_eq!(t.as_bytes()[10], b' ');
|
||||
assert_eq!(t.as_bytes()[13], b':');
|
||||
}
|
||||
}
|
||||
+866
@@ -0,0 +1,866 @@
|
||||
// 转发核心:接收本地客户端请求 → 注入认证/UA/会话头 → TLS 连接 opencode.ai →
|
||||
// 响应头重构后回写,body 逐块流式透传。
|
||||
|
||||
use std::io;
|
||||
use std::net::{IpAddr, SocketAddr};
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
|
||||
use tokio::net::{TcpListener, TcpStream};
|
||||
|
||||
use crate::http::{
|
||||
self, header_get, is_hop_by_hop, write_request_head, write_response_head,
|
||||
write_simple_response, BodyKind, ConnReader, Request, Response,
|
||||
};
|
||||
use crate::{generate_session_id, log_info, log_warn};
|
||||
|
||||
const UPSTREAM_CONNECT_TIMEOUT: Duration = Duration::from_secs(15);
|
||||
const RESPONSE_HEAD_TIMEOUT: Duration = Duration::from_secs(60);
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Upstream {
|
||||
pub tls: bool,
|
||||
pub host: String,
|
||||
pub port: u16,
|
||||
/// 上游路径前缀,如 "/zen/go";客户端的 /v1/* 原样接在其后
|
||||
pub base_path: String,
|
||||
}
|
||||
|
||||
pub struct Config {
|
||||
pub listen: IpAddr,
|
||||
pub port: u16,
|
||||
/// 上游 OpenCode Go API 密钥(转发时注入)
|
||||
pub api_key: String,
|
||||
/// 客户端访问密钥:Some 时每个 API 请求必须携带匹配的
|
||||
/// Authorization: Bearer <key> 或 x-api-key(/healthz 除外)
|
||||
pub client_key: Option<String>,
|
||||
pub user_agent: String,
|
||||
/// 强制固定 x-opencode-session;None 时按连接自动生成
|
||||
pub fixed_session: Option<String>,
|
||||
pub upstream: Upstream,
|
||||
}
|
||||
|
||||
type TlsConnector = tokio_rustls::TlsConnector;
|
||||
|
||||
/// 读写统一的擦除流(底层为明文 TcpStream 或 rustls TlsStream)。
|
||||
trait ErasedStream: AsyncRead + AsyncWrite + Unpin + Send {}
|
||||
impl<T: AsyncRead + AsyncWrite + Unpin + Send> ErasedStream for T {}
|
||||
type BoxedStream = Box<dyn ErasedStream>;
|
||||
|
||||
pub fn build_tls_connector(up: &Upstream) -> Result<Option<TlsConnector>, String> {
|
||||
if !up.tls {
|
||||
return Ok(None);
|
||||
}
|
||||
let provider = Arc::new(rustls::crypto::ring::default_provider());
|
||||
let mut roots = rustls::RootCertStore::empty();
|
||||
roots.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
|
||||
let config = rustls::ClientConfig::builder_with_provider(provider)
|
||||
.with_safe_default_protocol_versions()
|
||||
.map_err(|e| format!("TLS 版本配置失败: {e}"))?
|
||||
.with_root_certificates(roots)
|
||||
.with_no_client_auth();
|
||||
Ok(Some(TlsConnector::from(Arc::new(config))))
|
||||
}
|
||||
|
||||
struct UpstreamConn {
|
||||
reader: ConnReader<BoxedStream>,
|
||||
}
|
||||
|
||||
async fn connect_upstream(cfg: &Config, tls: &Option<TlsConnector>) -> io::Result<UpstreamConn> {
|
||||
let host = cfg.upstream.host.as_str();
|
||||
let server_name = rustls::pki_types::ServerName::try_from(host.to_string())
|
||||
.map_err(|e| io::Error::other(format!("无效的上游主机名: {e}")))?;
|
||||
let addr = tokio::net::lookup_host((host, cfg.upstream.port))
|
||||
.await
|
||||
.map_err(|e| io::Error::other(format!("上游 DNS 解析失败: {e}")))?
|
||||
.next()
|
||||
.ok_or_else(|| io::Error::other("上游无可用地址"))?;
|
||||
let tcp = tokio::time::timeout(UPSTREAM_CONNECT_TIMEOUT, TcpStream::connect(addr))
|
||||
.await
|
||||
.map_err(|_| io::Error::new(io::ErrorKind::TimedOut, "上游连接超时"))??;
|
||||
let _ = tcp.set_nodelay(true);
|
||||
let stream: BoxedStream = match tls {
|
||||
None => Box::new(tcp),
|
||||
Some(c) => {
|
||||
let s = tokio::time::timeout(UPSTREAM_CONNECT_TIMEOUT, c.connect(server_name, tcp))
|
||||
.await
|
||||
.map_err(|_| io::Error::new(io::ErrorKind::TimedOut, "TLS 握手超时"))?
|
||||
.map_err(|e| io::Error::other(format!("TLS 握手失败: {e}")))?;
|
||||
Box::new(s)
|
||||
}
|
||||
};
|
||||
Ok(UpstreamConn {
|
||||
reader: ConnReader::new(stream, 64 * 1024),
|
||||
})
|
||||
}
|
||||
|
||||
/// 构造发往上游的请求头:剔除客户端认证与逐跳头,注入代理要求的全部标识。
|
||||
fn build_upstream_headers(req: &Request, cfg: &Config, session: &str) -> Vec<(String, String)> {
|
||||
let default_port = if cfg.upstream.tls { 443 } else { 80 };
|
||||
let host_value = if cfg.upstream.port == default_port {
|
||||
cfg.upstream.host.clone()
|
||||
} else {
|
||||
format!("{}:{}", cfg.upstream.host, cfg.upstream.port)
|
||||
};
|
||||
|
||||
let mut out: Vec<(String, String)> = Vec::with_capacity(req.headers.len() + 8);
|
||||
out.push(("Host".into(), host_value));
|
||||
out.push(("User-Agent".into(), cfg.user_agent.clone()));
|
||||
out.push(("Authorization".into(), format!("Bearer {}", cfg.api_key)));
|
||||
out.push(("X-OpenCode-Session".into(), session.to_string()));
|
||||
|
||||
// Anthropic 兼容端点:x-api-key 为其主要认证头;版本头客户端没带则补默认
|
||||
let is_messages = req
|
||||
.target
|
||||
.split(['?', '#'])
|
||||
.next()
|
||||
.map(|p| p == "/v1/messages")
|
||||
.unwrap_or(false);
|
||||
let mut has_anthropic_version = false;
|
||||
if is_messages {
|
||||
has_anthropic_version = header_get(&req.headers, "anthropic-version").is_some();
|
||||
out.push(("x-api-key".into(), cfg.api_key.clone()));
|
||||
}
|
||||
|
||||
for (n, v) in &req.headers {
|
||||
let ln = n.to_ascii_lowercase();
|
||||
match ln.as_str() {
|
||||
// 由代理统一控制的头,客户端带来的不做透传
|
||||
"host" | "content-length" | "authorization" | "x-api-key" | "user-agent"
|
||||
| "x-opencode-session" => continue,
|
||||
"anthropic-version" => {
|
||||
if has_anthropic_version {
|
||||
out.push((n.clone(), v.clone()));
|
||||
}
|
||||
}
|
||||
_ if is_hop_by_hop(&ln) => continue,
|
||||
_ => out.push((n.clone(), v.clone())),
|
||||
}
|
||||
}
|
||||
if is_messages && !has_anthropic_version {
|
||||
out.push(("anthropic-version".into(), "2023-06-01".into()));
|
||||
}
|
||||
out.push(("Content-Length".into(), req.body.len().to_string()));
|
||||
out
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
enum BodyMode {
|
||||
None,
|
||||
Exact(u64),
|
||||
Chunked,
|
||||
Eof,
|
||||
}
|
||||
|
||||
/// 重构回给客户端的响应头:剥掉逐跳头,按 body 模式重写长度/分块,
|
||||
/// 返回 (响应头, body 转发模式, 是否需要关闭客户端连接)。
|
||||
fn build_client_headers(
|
||||
resp: &Response,
|
||||
method: &str,
|
||||
client_wants_close: bool,
|
||||
) -> (Vec<(String, String)>, BodyMode, bool) {
|
||||
let head_only = method.eq_ignore_ascii_case("HEAD");
|
||||
let mut close = client_wants_close;
|
||||
let mut headers: Vec<(String, String)> = Vec::with_capacity(resp.headers.len() + 2);
|
||||
for (n, v) in &resp.headers {
|
||||
let ln = n.to_ascii_lowercase();
|
||||
match ln.as_str() {
|
||||
"connection" | "keep-alive" | "transfer-encoding" | "content-length" | "trailer" => {
|
||||
continue
|
||||
}
|
||||
_ if is_hop_by_hop(&ln) => continue,
|
||||
_ => headers.push((n.clone(), v.clone())),
|
||||
}
|
||||
}
|
||||
let mode = match resp.body_kind().unwrap_or(BodyKind::Eof) {
|
||||
BodyKind::None => BodyMode::None,
|
||||
BodyKind::Chunked => {
|
||||
headers.push(("Transfer-Encoding".into(), "chunked".into()));
|
||||
BodyMode::Chunked
|
||||
}
|
||||
BodyKind::Exact(n) => {
|
||||
headers.push(("Content-Length".into(), n.to_string()));
|
||||
BodyMode::Exact(n)
|
||||
}
|
||||
BodyKind::Eof => {
|
||||
// 上游以关闭连接界定 body 结束,客户端连接也只能一并关闭
|
||||
close = true;
|
||||
BodyMode::Eof
|
||||
}
|
||||
};
|
||||
let mode = if head_only { BodyMode::None } else { mode };
|
||||
headers.push((
|
||||
"Connection".into(),
|
||||
if close { "close" } else { "keep-alive" }.into(),
|
||||
));
|
||||
(headers, mode, close)
|
||||
}
|
||||
|
||||
enum FwdErr {
|
||||
/// 响应头尚未回给客户端,还可以优雅回 502
|
||||
PreHead(io::Error),
|
||||
/// body 中途失败,响应已不完整,只能断开
|
||||
MidBody(io::Error),
|
||||
}
|
||||
|
||||
/// 常量时间字符串比较(长度差异仍会提前返回,长度本身不构成敏感信息)。
|
||||
fn ct_eq(a: &str, b: &str) -> bool {
|
||||
let (a, b) = (a.as_bytes(), b.as_bytes());
|
||||
if a.len() != b.len() {
|
||||
return false;
|
||||
}
|
||||
a.iter()
|
||||
.zip(b)
|
||||
.fold(0u8, |acc, (x, y)| acc | (x ^ y))
|
||||
== 0
|
||||
}
|
||||
|
||||
/// 提取客户端展示的凭据:Authorization: Bearer <key>(兼容任意大小写 scheme),
|
||||
/// 兜底取 x-api-key。
|
||||
fn client_presented_key(req: &Request) -> Option<&str> {
|
||||
if let Some(v) = header_get(&req.headers, "authorization") {
|
||||
let v = v.trim();
|
||||
if let Some(rest) = v
|
||||
.strip_prefix("Bearer ")
|
||||
.or_else(|| v.strip_prefix("bearer "))
|
||||
{
|
||||
return Some(rest.trim());
|
||||
}
|
||||
return Some(v);
|
||||
}
|
||||
header_get(&req.headers, "x-api-key")
|
||||
}
|
||||
|
||||
/// 401 鉴权失败响应(与上游错误格式一致,OpenAI/Anthropic 客户端均可识别)。
|
||||
async fn write_auth_error(w: &mut (impl AsyncWrite + Unpin)) -> io::Result<()> {
|
||||
let body = r#"{"type":"error","error":{"type":"authentication_error","message":"invalid client API key"}}"#;
|
||||
let headers = [
|
||||
("Content-Type".to_string(), "application/json".to_string()),
|
||||
("Content-Length".to_string(), body.len().to_string()),
|
||||
("Connection".to_string(), "close".to_string()),
|
||||
];
|
||||
write_response_head(w, "HTTP/1.1", 401, "Unauthorized", &headers).await?;
|
||||
w.write_all(body.as_bytes()).await?;
|
||||
w.flush().await
|
||||
}
|
||||
|
||||
async fn forward_request(
|
||||
wr: &mut (impl AsyncWrite + Unpin),
|
||||
req: &Request,
|
||||
cfg: &Config,
|
||||
tls: &Option<TlsConnector>,
|
||||
session: &str,
|
||||
) -> Result<(u16, u64, bool), FwdErr> {
|
||||
let mut up = connect_upstream(cfg, tls).await.map_err(FwdErr::PreHead)?;
|
||||
|
||||
let target = format!("{}{}", cfg.upstream.base_path, req.target);
|
||||
let headers = build_upstream_headers(req, cfg, session);
|
||||
let uw = up.reader.get_mut();
|
||||
write_request_head(uw, &req.method, &target, &headers)
|
||||
.await
|
||||
.map_err(FwdErr::PreHead)?;
|
||||
uw.write_all(&req.body).await.map_err(FwdErr::PreHead)?;
|
||||
uw.flush().await.map_err(FwdErr::PreHead)?;
|
||||
// 注意:这里绝不能 shutdown TLS 写端——close_notify 会被服务器当作
|
||||
// 全双工关闭信号而直接断开;请求边界由 Content-Length 界定
|
||||
|
||||
let mut ur = &mut up.reader;
|
||||
let resp = tokio::time::timeout(RESPONSE_HEAD_TIMEOUT, http::read_response_head(&mut ur))
|
||||
.await
|
||||
.map_err(|_| FwdErr::PreHead(io::Error::new(io::ErrorKind::TimedOut, "上游响应头超时")))?
|
||||
.map_err(FwdErr::PreHead)?;
|
||||
|
||||
let client_wants_close = req.version.eq_ignore_ascii_case("HTTP/1.0")
|
||||
|| header_get(&req.headers, "connection")
|
||||
.map(|v| v.to_ascii_lowercase().contains("close"))
|
||||
.unwrap_or(false);
|
||||
let (resp_headers, mode, close) = build_client_headers(&resp, &req.method, client_wants_close);
|
||||
write_response_head(wr, "HTTP/1.1", resp.status, &resp.reason, &resp_headers)
|
||||
.await
|
||||
.map_err(FwdErr::PreHead)?;
|
||||
|
||||
let bytes = if resp.status >= 400 && matches!(mode, BodyMode::Exact(n) if n <= 1024*1024)
|
||||
|| (resp.status >= 400 && matches!(mode, BodyMode::Chunked | BodyMode::Eof))
|
||||
{
|
||||
// 错误响应一般很小:缓冲读全,记录上游错误详情(诊断 4xx/5xx 用),再整体回写
|
||||
let max: usize = 1024 * 1024;
|
||||
let body = match mode {
|
||||
BodyMode::Exact(n) => up
|
||||
.reader
|
||||
.read_exact_n(n.min(max as u64) as usize)
|
||||
.await
|
||||
.map_err(FwdErr::MidBody)?,
|
||||
BodyMode::Chunked => up
|
||||
.reader
|
||||
.read_chunked_decode(max as u64)
|
||||
.await
|
||||
.map_err(FwdErr::MidBody)?,
|
||||
BodyMode::Eof => up
|
||||
.reader
|
||||
.read_upto(max)
|
||||
.await
|
||||
.map_err(FwdErr::MidBody)?,
|
||||
BodyMode::None => Vec::new(),
|
||||
};
|
||||
if !body.is_empty() {
|
||||
let end = body.len().min(512);
|
||||
log_warn!(
|
||||
"上游 {} 错误详情: {}",
|
||||
resp.status,
|
||||
String::from_utf8_lossy(&body[..end])
|
||||
);
|
||||
}
|
||||
wr.write_all(&body).await.map_err(FwdErr::MidBody)?;
|
||||
wr.flush().await.map_err(FwdErr::MidBody)?;
|
||||
body.len() as u64
|
||||
} else {
|
||||
match mode {
|
||||
BodyMode::None => 0,
|
||||
BodyMode::Exact(n) => up
|
||||
.reader
|
||||
.forward_n(wr, n)
|
||||
.await
|
||||
.map_err(FwdErr::MidBody)?,
|
||||
BodyMode::Chunked => up
|
||||
.reader
|
||||
.forward_chunked(wr)
|
||||
.await
|
||||
.map_err(FwdErr::MidBody)?,
|
||||
BodyMode::Eof => up
|
||||
.reader
|
||||
.forward_to_eof(wr)
|
||||
.await
|
||||
.map_err(FwdErr::MidBody)?,
|
||||
}
|
||||
};
|
||||
Ok((resp.status, bytes, close))
|
||||
}
|
||||
|
||||
async fn handle_conn(stream: TcpStream, peer: SocketAddr, cfg: Arc<Config>, tls: Option<TlsConnector>) {
|
||||
let _ = stream.set_nodelay(true);
|
||||
let (rd, mut wr) = stream.into_split();
|
||||
let mut rd = ConnReader::new(rd, 64 * 1024);
|
||||
|
||||
// 会话 ID:客户端带来则透传;否则该连接内首次生成、后续复用(近似每对话一个)
|
||||
let mut auto_session: Option<String> = cfg.fixed_session.clone();
|
||||
let mut nreq: u64 = 0;
|
||||
|
||||
loop {
|
||||
nreq += 1;
|
||||
let req = match http::read_request(&mut rd, &mut wr).await {
|
||||
Ok(None) => break,
|
||||
Ok(Some(r)) => r,
|
||||
Err(e) => {
|
||||
log_warn!("{peer} 请求解析失败: {e}");
|
||||
let _ = write_simple_response(&mut wr, 400, &format!("请求解析失败: {e}")).await;
|
||||
break;
|
||||
}
|
||||
};
|
||||
let t0 = Instant::now();
|
||||
|
||||
// 本地健康检查端点(不打上游):路径去掉 query 后精确匹配
|
||||
if req.target.split(['?', '#']).next() == Some("/healthz") {
|
||||
let body = b"ok\n";
|
||||
let headers = [
|
||||
("Content-Type".to_string(), "text/plain; charset=utf-8".to_string()),
|
||||
("Content-Length".to_string(), body.len().to_string()),
|
||||
("Connection".to_string(), "keep-alive".to_string()),
|
||||
];
|
||||
if write_response_head(&mut wr, "HTTP/1.1", 200, "OK", &headers).await.is_err() {
|
||||
break;
|
||||
}
|
||||
let _ = wr.write_all(body).await;
|
||||
let _ = wr.flush().await;
|
||||
let client_close = req.version.eq_ignore_ascii_case("HTTP/1.0")
|
||||
|| header_get(&req.headers, "connection")
|
||||
.map(|v| v.to_ascii_lowercase().contains("close"))
|
||||
.unwrap_or(false);
|
||||
if client_close {
|
||||
break;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
// 客户端访问密钥校验:配置后所有 API 请求必须携带匹配凭据
|
||||
if let Some(expected) = &cfg.client_key {
|
||||
let ok = client_presented_key(&req)
|
||||
.map(|k| ct_eq(k, expected))
|
||||
.unwrap_or(false);
|
||||
if !ok {
|
||||
log_warn!(
|
||||
"{peer} #{nreq} {} {} 客户端密钥校验失败",
|
||||
req.method,
|
||||
req.target
|
||||
);
|
||||
let _ = write_auth_error(&mut wr).await;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
let session = match header_get(&req.headers, "x-opencode-session") {
|
||||
Some(s) => s.to_string(),
|
||||
None => match &auto_session {
|
||||
Some(s) => s.clone(),
|
||||
None => {
|
||||
let s = generate_session_id();
|
||||
auto_session = Some(s.clone());
|
||||
s
|
||||
}
|
||||
},
|
||||
};
|
||||
|
||||
match forward_request(&mut wr, &req, &cfg, &tls, &session).await {
|
||||
Ok((status, bytes, close)) => {
|
||||
log_info!(
|
||||
"{peer} #{nreq} {} {} -> {} ({} 字节, {} ms)",
|
||||
req.method,
|
||||
req.target,
|
||||
status,
|
||||
bytes,
|
||||
t0.elapsed().as_millis()
|
||||
);
|
||||
let client_close = req.version.eq_ignore_ascii_case("HTTP/1.0")
|
||||
|| header_get(&req.headers, "connection")
|
||||
.map(|v| v.to_ascii_lowercase().contains("close"))
|
||||
.unwrap_or(false);
|
||||
if close || client_close {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(FwdErr::PreHead(e)) => {
|
||||
log_warn!(
|
||||
"{peer} #{nreq} {} {} 上游失败: {e}",
|
||||
req.method,
|
||||
req.target
|
||||
);
|
||||
let _ = write_simple_response(&mut wr, 502, &format!("上游请求失败: {e}")).await;
|
||||
break;
|
||||
}
|
||||
Err(FwdErr::MidBody(e)) => {
|
||||
log_warn!(
|
||||
"{peer} #{nreq} {} {} 响应体转发中断: {e}",
|
||||
req.method,
|
||||
req.target
|
||||
);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
log_info!("{peer} 连接关闭 (共 {nreq} 个请求)");
|
||||
}
|
||||
|
||||
pub async fn serve(
|
||||
listener: TcpListener,
|
||||
cfg: Arc<Config>,
|
||||
tls: Option<TlsConnector>,
|
||||
) -> io::Result<()> {
|
||||
loop {
|
||||
let (stream, peer) = match listener.accept().await {
|
||||
Ok(x) => x,
|
||||
Err(e) => {
|
||||
log_warn!("accept 失败: {e}");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let cfg = cfg.clone();
|
||||
let tls = tls.clone();
|
||||
tokio::spawn(handle_conn(stream, peer, cfg, tls));
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn run(cfg: Arc<Config>) -> Result<(), String> {
|
||||
let tls = build_tls_connector(&cfg.upstream)?;
|
||||
let listener = TcpListener::bind((cfg.listen, cfg.port))
|
||||
.await
|
||||
.map_err(|e| format!("监听 {}:{} 失败: {e}", cfg.listen, cfg.port))?;
|
||||
let scheme = if cfg.upstream.tls { "https" } else { "http" };
|
||||
log_info!(
|
||||
"opencode-go-proxy http://{}:{} -> {}://{}:{}{}",
|
||||
cfg.listen,
|
||||
cfg.port,
|
||||
scheme,
|
||||
cfg.upstream.host,
|
||||
cfg.upstream.port,
|
||||
cfg.upstream.base_path
|
||||
);
|
||||
serve(listener, cfg, tls)
|
||||
.await
|
||||
.map_err(|e| format!("服务异常退出: {e}"))
|
||||
}
|
||||
|
||||
// ---- 集成测试:mock 明文上游,验证注入/透传/流式/keep-alive ---------------------
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::time::Duration;
|
||||
use tokio::io::AsyncReadExt;
|
||||
|
||||
struct MockReq {
|
||||
method: String,
|
||||
target: String,
|
||||
headers: Vec<(String, String)>,
|
||||
body: Vec<u8>,
|
||||
}
|
||||
|
||||
impl MockReq {
|
||||
fn header(&self, name: &str) -> Option<&str> {
|
||||
header_get(&self.headers, name)
|
||||
}
|
||||
}
|
||||
|
||||
/// 起 mock 上游:每连接循环读请求,handler 返回原始响应字节;连接 EOF 结束。
|
||||
async fn spawn_mock<F, Fut>(handler: F) -> SocketAddr
|
||||
where
|
||||
F: Fn(MockReq) -> Fut + Send + Sync + 'static,
|
||||
Fut: std::future::Future<Output = Vec<u8>> + Send + 'static,
|
||||
{
|
||||
let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let handler = Arc::new(handler);
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
let Ok((sock, _)) = listener.accept().await else { return };
|
||||
let handler = handler.clone();
|
||||
tokio::spawn(async move {
|
||||
let (rd_half, mut wr_half) = sock.into_split();
|
||||
let mut rd = ConnReader::new(rd_half, 4096);
|
||||
loop {
|
||||
let line = match rd.read_line_opt(8192).await {
|
||||
Ok(Some(l)) => l,
|
||||
_ => return,
|
||||
};
|
||||
let line_str = String::from_utf8_lossy(&line).into_owned();
|
||||
let mut it = line_str.split_whitespace();
|
||||
let method = it.next().unwrap_or("").to_string();
|
||||
let target = it.next().unwrap_or("").to_string();
|
||||
let mut headers = Vec::new();
|
||||
loop {
|
||||
let h = rd.read_line(8192).await.unwrap_or_default();
|
||||
let hs = String::from_utf8_lossy(&h).trim_end_matches(['\r', '\n']).to_string();
|
||||
if hs.is_empty() {
|
||||
break;
|
||||
}
|
||||
if let Some((n, v)) = hs.split_once(':') {
|
||||
headers.push((n.trim().to_string(), v.trim().to_string()));
|
||||
}
|
||||
}
|
||||
let body = match header_get(&headers, "content-length") {
|
||||
Some(cl) => rd.read_exact_n(cl.parse().unwrap()).await.unwrap(),
|
||||
None => Vec::new(),
|
||||
};
|
||||
let resp = handler(MockReq { method, target, headers, body }).await;
|
||||
if wr_half.write_all(&resp).await.is_err() {
|
||||
return;
|
||||
}
|
||||
let _ = wr_half.flush().await;
|
||||
}
|
||||
});
|
||||
}
|
||||
});
|
||||
addr
|
||||
}
|
||||
|
||||
fn resp_bytes(status_line: &str, headers: &[&str], body: &[u8]) -> Vec<u8> {
|
||||
let mut out = Vec::new();
|
||||
out.extend_from_slice(format!("{status_line}\r\n").as_bytes());
|
||||
for h in headers {
|
||||
out.extend_from_slice(format!("{h}\r\n").as_bytes());
|
||||
}
|
||||
out.extend_from_slice(format!("Content-Length: {}\r\n", body.len()).as_bytes());
|
||||
out.extend_from_slice(b"\r\n");
|
||||
out.extend_from_slice(body);
|
||||
out
|
||||
}
|
||||
|
||||
async fn http_roundtrip(proxy: SocketAddr, raw_req: &[u8]) -> io::Result<String> {
|
||||
// 统一带 Connection: close:read_to_end 依赖代理在响应后关闭连接
|
||||
let mut s = TcpStream::connect(proxy).await?;
|
||||
let head_end = raw_req
|
||||
.windows(4)
|
||||
.position(|w| w == b"\r\n\r\n")
|
||||
.expect("请求缺少头结束标记");
|
||||
s.write_all(&raw_req[..head_end]).await?;
|
||||
s.write_all(b"\r\nConnection: close\r\n\r\n").await?;
|
||||
s.write_all(&raw_req[head_end + 4..]).await?;
|
||||
let mut buf = Vec::new();
|
||||
s.read_to_end(&mut buf).await?;
|
||||
Ok(String::from_utf8_lossy(&buf).into_owned())
|
||||
}
|
||||
|
||||
async fn start_proxy(
|
||||
upstream: SocketAddr,
|
||||
fixed_session: Option<String>,
|
||||
client_key: Option<String>,
|
||||
) -> SocketAddr {
|
||||
let cfg = Arc::new(Config {
|
||||
listen: "127.0.0.1".parse().unwrap(),
|
||||
port: 0,
|
||||
api_key: "test-key".into(),
|
||||
client_key,
|
||||
user_agent: "opencode-go-proxy/test".into(),
|
||||
fixed_session,
|
||||
upstream: Upstream {
|
||||
tls: false,
|
||||
host: upstream.ip().to_string(),
|
||||
port: upstream.port(),
|
||||
base_path: "/zen/go".into(),
|
||||
},
|
||||
});
|
||||
let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
tokio::spawn(serve(listener, cfg, None));
|
||||
addr
|
||||
}
|
||||
|
||||
const JSON_BODY: &[u8] = br#"{"ok":true}"#;
|
||||
|
||||
#[tokio::test]
|
||||
async fn injects_headers_and_forwards_body() {
|
||||
let upstream = spawn_mock(|req| async move {
|
||||
assert_eq!(req.method, "POST");
|
||||
assert_eq!(req.target, "/zen/go/v1/chat/completions");
|
||||
assert_eq!(req.header("authorization"), Some("Bearer test-key"));
|
||||
assert_eq!(req.header("user-agent"), Some("opencode-go-proxy/test"));
|
||||
assert!(req.header("x-opencode-session").unwrap_or("").starts_with("sess-"));
|
||||
assert!(req.header("content-length").is_some());
|
||||
assert_eq!(req.body, b"{\"p\":1}");
|
||||
resp_bytes("HTTP/1.1 200 OK", &["Content-Type: application/json"], JSON_BODY)
|
||||
})
|
||||
.await;
|
||||
let proxy = start_proxy(upstream, None, None).await;
|
||||
let raw = b"POST /v1/chat/completions HTTP/1.1\r\nHost: anything\r\n\
|
||||
Authorization: Bearer client-fake\r\nUser-Agent: some-client/9\r\n\
|
||||
Content-Type: application/json\r\nContent-Length: 7\r\n\r\n{\"p\":1}";
|
||||
let resp = http_roundtrip(proxy, raw).await.unwrap();
|
||||
assert!(resp.starts_with("HTTP/1.1 200 OK\r\n"));
|
||||
assert!(resp.contains("Content-Type: application/json\r\n"));
|
||||
assert!(resp.ends_with("{\"ok\":true}"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn session_header_passthrough() {
|
||||
let upstream = spawn_mock(|req| async move {
|
||||
assert_eq!(req.header("x-opencode-session"), Some("my-session-123"));
|
||||
resp_bytes("HTTP/1.1 200 OK", &[], b"")
|
||||
})
|
||||
.await;
|
||||
let proxy = start_proxy(upstream, Some("fixed-but-ignored".into()), None).await;
|
||||
let raw = b"POST /v1/chat/completions HTTP/1.1\r\n\
|
||||
x-opencode-session: my-session-123\r\nContent-Length: 0\r\n\r\n";
|
||||
let resp = http_roundtrip(proxy, raw).await.unwrap();
|
||||
assert!(resp.starts_with("HTTP/1.1 200 OK"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fixed_session_used_when_missing() {
|
||||
let upstream = spawn_mock(|req| async move {
|
||||
assert_eq!(req.header("x-opencode-session"), Some("fixed-sess"));
|
||||
resp_bytes("HTTP/1.1 200 OK", &[], b"")
|
||||
})
|
||||
.await;
|
||||
let proxy = start_proxy(upstream, Some("fixed-sess".into()), None).await;
|
||||
let resp = http_roundtrip(
|
||||
proxy,
|
||||
b"POST /v1/chat/completions HTTP/1.1\r\nContent-Length: 0\r\n\r\n",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(resp.starts_with("HTTP/1.1 200 OK"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn messages_endpoint_gets_anthropic_headers() {
|
||||
let upstream = spawn_mock(|req| async move {
|
||||
assert_eq!(req.target, "/zen/go/v1/messages");
|
||||
assert_eq!(req.header("x-api-key"), Some("test-key"));
|
||||
assert_eq!(req.header("anthropic-version"), Some("2023-06-01"));
|
||||
resp_bytes("HTTP/1.1 200 OK", &[], b"")
|
||||
})
|
||||
.await;
|
||||
let proxy = start_proxy(upstream, None, None).await;
|
||||
let raw = b"POST /v1/messages HTTP/1.1\r\nContent-Length: 0\r\n\r\n";
|
||||
let resp = http_roundtrip(proxy, raw).await.unwrap();
|
||||
assert!(resp.starts_with("HTTP/1.1 200 OK"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn chunked_response_streamed_intact() {
|
||||
let payload: &[u8] = b"data: {\"a\":1}\r\n\r\ndata: {\"b\":2}\r\n\r\n";
|
||||
let upstream = spawn_mock(|req| async move {
|
||||
let _ = req;
|
||||
let mut out = Vec::new();
|
||||
out.extend_from_slice(
|
||||
b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nTransfer-Encoding: chunked\r\n\r\n",
|
||||
);
|
||||
out.extend_from_slice(b"11\r\ndata: {\"a\":1}\r\n\r\n\r\n");
|
||||
tokio::time::sleep(Duration::from_millis(30)).await;
|
||||
out.extend_from_slice(b"11\r\ndata: {\"b\":2}\r\n\r\n\r\n");
|
||||
out.extend_from_slice(b"0\r\n\r\n");
|
||||
out
|
||||
})
|
||||
.await;
|
||||
let proxy = start_proxy(upstream, None, None).await;
|
||||
let raw = b"POST /v1/chat/completions HTTP/1.1\r\nContent-Length: 0\r\n\r\n";
|
||||
let resp = http_roundtrip(proxy, raw).await.unwrap();
|
||||
assert!(resp.contains("Transfer-Encoding: chunked"));
|
||||
// body 原样透传的是 chunked 字节流,解码后应还原出完整 payload
|
||||
let head_end = resp.find("\r\n\r\n").unwrap() + 4;
|
||||
assert_eq!(
|
||||
&resp[head_end..],
|
||||
"11\r\ndata: {\"a\":1}\r\n\r\n\r\n11\r\ndata: {\"b\":2}\r\n\r\n\r\n0\r\n\r\n"
|
||||
);
|
||||
let mut decoded = Vec::new();
|
||||
let mut r = ConnReader::new(&resp.as_bytes()[head_end..], 64);
|
||||
decoded.extend_from_slice(&r.read_chunked_decode(u64::MAX).await.unwrap());
|
||||
assert_eq!(decoded, payload);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn keep_alive_serves_two_requests() {
|
||||
let upstream = spawn_mock(|req| async move {
|
||||
assert_eq!(req.target, "/zen/go/v1/models");
|
||||
resp_bytes("HTTP/1.1 200 OK", &[], b"m")
|
||||
})
|
||||
.await;
|
||||
let proxy = start_proxy(upstream, None, None).await;
|
||||
let mut s = TcpStream::connect(proxy).await.unwrap();
|
||||
s.write_all(b"GET /v1/models HTTP/1.1\r\nHost: x\r\n\r\n")
|
||||
.await
|
||||
.unwrap();
|
||||
let mut buf = [0u8; 512];
|
||||
let n = s.read(&mut buf).await.unwrap();
|
||||
assert!(String::from_utf8_lossy(&buf[..n]).starts_with("HTTP/1.1 200 OK"));
|
||||
s.write_all(b"GET /v1/models HTTP/1.1\r\nHost: x\r\nConnection: close\r\n\r\n")
|
||||
.await
|
||||
.unwrap();
|
||||
let mut rest = Vec::new();
|
||||
s.read_to_end(&mut rest).await.unwrap();
|
||||
assert!(String::from_utf8_lossy(&rest).starts_with("HTTP/1.1 200 OK"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn healthz_answered_locally() {
|
||||
// 不起 mock 上游:/healthz 必须由代理本地应答
|
||||
let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let cfg = Arc::new(Config {
|
||||
listen: "127.0.0.1".parse().unwrap(),
|
||||
port: 0,
|
||||
api_key: "k".into(),
|
||||
client_key: None,
|
||||
user_agent: "ua".into(),
|
||||
fixed_session: None,
|
||||
// 上游指向必然拒绝连接的地址,若转发则失败
|
||||
upstream: Upstream { tls: false, host: "127.0.0.1".into(), port: 1, base_path: "/zen/go".into() },
|
||||
});
|
||||
tokio::spawn(serve(listener, cfg, None));
|
||||
|
||||
let resp = http_roundtrip(addr, b"GET /healthz HTTP/1.1\r\nHost: x\r\n\r\n").await.unwrap();
|
||||
assert!(resp.starts_with("HTTP/1.1 200 OK\r\n"));
|
||||
assert!(resp.ends_with("ok\n"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn client_key_enforced() {
|
||||
let upstream = spawn_mock(|req| async move {
|
||||
// 通过鉴权的请求应带着代理注入的上游 key 到达
|
||||
assert_eq!(req.header("authorization"), Some("Bearer test-key"));
|
||||
resp_bytes("HTTP/1.1 200 OK", &[], b"ok")
|
||||
})
|
||||
.await;
|
||||
let proxy = start_proxy(upstream, None, Some("ogp-secret".into())).await;
|
||||
|
||||
// 无凭据 → 401
|
||||
let resp = http_roundtrip(
|
||||
proxy,
|
||||
b"POST /v1/chat/completions HTTP/1.1\r\nContent-Length: 0\r\n\r\n",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(resp.starts_with("HTTP/1.1 401 "));
|
||||
assert!(resp.contains("authentication_error"));
|
||||
|
||||
// 错误凭据 → 401
|
||||
let resp = http_roundtrip(
|
||||
proxy,
|
||||
b"POST /v1/chat/completions HTTP/1.1\r\nAuthorization: Bearer wrong\r\nContent-Length: 0\r\n\r\n",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(resp.starts_with("HTTP/1.1 401 "));
|
||||
|
||||
// 正确 Bearer → 放行;正确 x-api-key → 放行
|
||||
for header in [
|
||||
"Authorization: Bearer ogp-secret".to_string(),
|
||||
"x-api-key: ogp-secret".to_string(),
|
||||
] {
|
||||
let raw = format!(
|
||||
"POST /v1/chat/completions HTTP/1.1\r\n{header}\r\nContent-Length: 0\r\n\r\n"
|
||||
);
|
||||
let resp = http_roundtrip(proxy, raw.as_bytes()).await.unwrap();
|
||||
assert!(resp.starts_with("HTTP/1.1 200 OK"), "头 {header} 应放行: {resp}");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn client_key_untouched_when_absent() {
|
||||
// 未配置客户端密钥:任意/缺失凭据都放行
|
||||
let upstream = spawn_mock(|req| async move {
|
||||
let _ = &req;
|
||||
resp_bytes("HTTP/1.1 200 OK", &[], b"ok")
|
||||
})
|
||||
.await;
|
||||
let proxy = start_proxy(upstream, None, None).await;
|
||||
let resp = http_roundtrip(
|
||||
proxy,
|
||||
b"GET /v1/models HTTP/1.1\r\nAuthorization: Bearer whatever\r\n\r\n",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(resp.starts_with("HTTP/1.1 200 OK"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ct_eq_constant_compare() {
|
||||
assert!(ct_eq("abc", "abc"));
|
||||
assert!(!ct_eq("abc", "abd"));
|
||||
assert!(!ct_eq("abc", "ab"));
|
||||
assert!(ct_eq("", ""));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn upstream_error_maps_to_502() {
|
||||
// 指向无人监听的端口
|
||||
let dead = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
|
||||
let dead_addr = dead.local_addr().unwrap();
|
||||
drop(dead);
|
||||
let proxy = start_proxy(dead_addr, None, None).await;
|
||||
let resp = http_roundtrip(
|
||||
proxy,
|
||||
b"GET /v1/models HTTP/1.1\r\nHost: x\r\n\r\n",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(resp.starts_with("HTTP/1.1 502"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn upstream_status_forwarded() {
|
||||
let upstream = spawn_mock(|req| async move {
|
||||
let _ = &req;
|
||||
resp_bytes("HTTP/1.1 401 Unauthorized", &["WWW-Authenticate: Bearer"], b"no key")
|
||||
})
|
||||
.await;
|
||||
let proxy = start_proxy(upstream, None, None).await;
|
||||
let resp = http_roundtrip(
|
||||
proxy,
|
||||
b"GET /v1/models HTTP/1.1\r\nHost: x\r\n\r\n",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(resp.starts_with("HTTP/1.1 401 Unauthorized\r\n"));
|
||||
assert!(resp.contains("WWW-Authenticate: Bearer\r\n"));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user