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
2026-09-10 07:29:00 +00:00
commit a8198f47e4
9 changed files with 2713 additions and 0 deletions
+5
View File
@@ -0,0 +1,5 @@
[build]
target = "x86_64-unknown-linux-musl"
[target.x86_64-unknown-linux-musl]
linker = "musl-gcc"
+1
View File
@@ -0,0 +1 @@
/target
Generated
+344
View File
@@ -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
View File
@@ -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"
+122
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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"));
}
}