commit a8198f47e427d5fda4851213512488c1424cb45e Author: phptits Date: Thu Sep 10 07:29:00 2026 +0000 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 个单元/集成测试 diff --git a/.cargo/config.toml b/.cargo/config.toml new file mode 100644 index 0000000..84af28c --- /dev/null +++ b/.cargo/config.toml @@ -0,0 +1,5 @@ +[build] +target = "x86_64-unknown-linux-musl" + +[target.x86_64-unknown-linux-musl] +linker = "musl-gcc" diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ea8c4bf --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +/target diff --git a/Cargo.lock b/Cargo.lock new file mode 100644 index 0000000..c2ada13 --- /dev/null +++ b/Cargo.lock @@ -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" diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..92e4ad3 --- /dev/null +++ b/Cargo.toml @@ -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" diff --git a/README.md b/README.md new file mode 100644 index 0000000..876d64c --- /dev/null +++ b/README.md @@ -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 上游 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 diff --git a/src/http.rs b/src/http.rs new file mode 100644 index 0000000..3b6acec --- /dev/null +++ b/src/http.rs @@ -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) -> 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, +} + +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 { + 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 { + inner: R, + buf: Box<[u8]>, + start: usize, + end: usize, +} + +impl ConnReader { + 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 { + 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>> { + 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> { + self.read_line_opt(max) + .await? + .ok_or_else(|| proto_err("连接在读取行时意外关闭")) + } + + pub async fn read_exact_n(&mut self, n: usize) -> io::Result> { + 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> { + 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( + &mut self, + w: &mut W, + mut n: u64, + ) -> io::Result { + 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(&mut self, w: &mut W) -> io::Result { + 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( + &mut self, + w: &mut W, + ) -> io::Result { + 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> { + 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 { + 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: &mut ConnReader, w: &mut W) -> io::Result> +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: &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: &mut ConnReader, +) -> io::Result { + 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::().ok()) + .ok_or_else(|| proto_err(format!("非法状态行: {}", line_str.trim())))?; + let reason = parts.collect::>().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: &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: &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: &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> { + 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")); + } +} diff --git a/src/install.rs b/src/install.rs new file mode 100644 index 0000000..c1dcb7f --- /dev/null +++ b/src/install.rs @@ -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, + pub session: Option, + pub user_agent: Option, + pub upstream: Option, +} + +/// 生成 "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)); + } +} diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..cbf8f2e --- /dev/null +++ b/src/main.rs @@ -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 发送给上游的 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, + client_key: Option, + session: Option, + user_agent: Option, + upstream: Option, + install: bool, + uninstall: bool, + help: bool, +} + +fn parse_args>(it: I) -> Result { + 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, name: &str| -> Result { + 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 { + 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::().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':'); + } +} diff --git a/src/proxy.rs b/src/proxy.rs new file mode 100644 index 0000000..b77d8ad --- /dev/null +++ b/src/proxy.rs @@ -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 或 x-api-key(/healthz 除外) + pub client_key: Option, + pub user_agent: String, + /// 强制固定 x-opencode-session;None 时按连接自动生成 + pub fixed_session: Option, + pub upstream: Upstream, +} + +type TlsConnector = tokio_rustls::TlsConnector; + +/// 读写统一的擦除流(底层为明文 TcpStream 或 rustls TlsStream)。 +trait ErasedStream: AsyncRead + AsyncWrite + Unpin + Send {} +impl ErasedStream for T {} +type BoxedStream = Box; + +pub fn build_tls_connector(up: &Upstream) -> Result, 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, +} + +async fn connect_upstream(cfg: &Config, tls: &Option) -> io::Result { + 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 (兼容任意大小写 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, + 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, tls: Option) { + 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 = 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, + tls: Option, +) -> 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) -> 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, + } + + impl MockReq { + fn header(&self, name: &str) -> Option<&str> { + header_get(&self.headers, name) + } + } + + /// 起 mock 上游:每连接循环读请求,handler 返回原始响应字节;连接 EOF 结束。 + async fn spawn_mock(handler: F) -> SocketAddr + where + F: Fn(MockReq) -> Fut + Send + Sync + 'static, + Fut: std::future::Future> + 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 { + 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 { + // 统一带 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, + client_key: Option, + ) -> 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")); + } +}