starred/codex-proxy-rs
1
0
Fork 0
mirror of https://github.com/zyycn/codex-proxy-rs.git synced 2026-11-06 08:06:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: address audit findings

Browse Source 
Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: Zyy <No Email>
This commit is contained in:
Zyy
2026-07-03 18:20:05 +08:00
co-authored by Codex
parent 71279877df
commit 112bd76f66
115 changed files with 7754 additions and 4718 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/ci.yml
+21 -6
View File
@@ -18,15 +18,15 @@ jobs:
runs-on: ubuntu-latest
steps:
- name: Checkout
uses: actions/checkout@v7
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7
- name: Setup Rust
uses: dtolnay/rust-toolchain@stable
uses: dtolnay/rust-toolchain@4be7066ada62dd38de10e7b70166bc74ed198c30 # stable
with:
components: rustfmt,clippy
- name: Cache Rust
uses: Swatinem/rust-cache@v2
uses: Swatinem/rust-cache@e18b497796c12c097a38f9edb9d0641fb99eee32 # v2
with:
workspaces: backend -> target
@@ -43,15 +43,15 @@ jobs:
runs-on: ubuntu-latest
steps:
- name: Checkout
uses: actions/checkout@v7
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7
- name: Setup pnpm
uses: pnpm/action-setup@v6
uses: pnpm/action-setup@0ebf47130e4866e96fce0953f49152a61190b271 # v6
with:
version: 11.7.0
- name: Setup Node
uses: actions/setup-node@v6
uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6
with:
node-version: 24
cache: pnpm
@@ -60,5 +60,20 @@ jobs:
- name: Install frontend dependencies
run: pnpm --dir frontend install --frozen-lockfile
- name: Frontend format
run: pnpm --dir frontend format:check
- name: Frontend build
run: pnpm --dir frontend build
deploy:
runs-on: ubuntu-latest
steps:
- name: Checkout
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7
- name: Validate Docker Compose
run: docker compose -f deploy/docker-compose.yml config
- name: Build runtime image
run: docker build -f deploy/Dockerfile --target runtime -t codex-proxy-rs:ci .
.github/workflows/release.yml
+75 -24
View File
@@ -12,8 +12,7 @@ on:
type: string
permissions:
contents: write
packages: write
contents: read
concurrency:
group: release-${{ github.event_name == 'workflow_dispatch' && inputs.tag || github.ref_name }}
@@ -37,7 +36,7 @@ jobs:
steps:
- name: Checkout
uses: actions/checkout@v7
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7
with:
ref: ${{ github.event_name == 'workflow_dispatch' && inputs.tag || github.ref }}
fetch-depth: 0
@@ -147,17 +146,17 @@ jobs:
steps:
- name: Checkout
uses: actions/checkout@v7
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7
with:
ref: ${{ needs.release-info.outputs.tag }}
- name: Setup pnpm
uses: pnpm/action-setup@v6
uses: pnpm/action-setup@0ebf47130e4866e96fce0953f49152a61190b271 # v6
with:
version: 10
- name: Setup Node.js
uses: actions/setup-node@v6
uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6
with:
node-version: 24
cache: pnpm
@@ -172,7 +171,7 @@ jobs:
working-directory: frontend
- name: Upload frontend artifact
uses: actions/upload-artifact@v7
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7
with:
name: frontend-dist
path: frontend/dist
@@ -189,12 +188,12 @@ jobs:
steps:
- name: Checkout
uses: actions/checkout@v7
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7
with:
ref: ${{ needs.release-info.outputs.tag }}
- name: Cache Cargo
uses: actions/cache@v6
uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6
with:
path: |
~/.cargo/registry
@@ -205,7 +204,7 @@ jobs:
cargo-${{ runner.os }}-${{ matrix.target }}-
- name: Download frontend artifact
uses: actions/download-artifact@v8
uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8
with:
name: frontend-dist
path: backend/web/dist
@@ -228,7 +227,7 @@ jobs:
CPR_BUILD_TYPE: release
- name: Upload binary artifact
uses: actions/upload-artifact@v7
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7
with:
name: codex-proxy-rs-${{ matrix.os }}-${{ matrix.arch }}
path: dist/binaries/${{ matrix.os }}-${{ matrix.arch }}/
@@ -238,21 +237,25 @@ jobs:
package-assets:
needs: [release-info, build-frontend, build-binary]
runs-on: ubuntu-latest
permissions:
contents: read
id-token: write
attestations: write
steps:
- name: Checkout
uses: actions/checkout@v7
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7
with:
ref: ${{ needs.release-info.outputs.tag }}
- name: Download frontend artifact
uses: actions/download-artifact@v8
uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8
with:
name: frontend-dist
path: dist/frontend
- name: Download binary artifacts
uses: actions/download-artifact@v8
uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8
with:
pattern: codex-proxy-rs-*
path: dist/binaries
@@ -336,8 +339,13 @@ jobs:
file.write(f"{sha256(archive)} {archive.name}\n")
PY
- name: Generate release asset provenance
uses: actions/attest@a1948c3f048ba23858d222213b7c278aabede763 # v4
with:
subject-path: dist/release/*
- name: Upload release assets
uses: actions/upload-artifact@v7
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7
with:
name: release-assets
path: dist/release
@@ -347,27 +355,32 @@ jobs:
publish:
needs: [release-info, build-frontend, build-binary, package-assets]
runs-on: ubuntu-latest
permissions:
contents: write
packages: write
id-token: write
attestations: write
steps:
- name: Checkout
uses: actions/checkout@v7
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7
with:
ref: ${{ needs.release-info.outputs.tag }}
- name: Download frontend artifact
uses: actions/download-artifact@v8
uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8
with:
name: frontend-dist
path: dist/docker/web/dist
- name: Download Docker binary artifacts
uses: actions/download-artifact@v8
uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8
with:
pattern: codex-proxy-rs-linux-*
path: dist/docker-artifacts
- name: Download release assets
uses: actions/download-artifact@v8
uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8
with:
name: release-assets
path: dist/release
@@ -409,20 +422,24 @@ jobs:
PY
- name: Setup QEMU
uses: docker/setup-qemu-action@v4
uses: docker/setup-qemu-action@96fe6ef7f33517b61c61be40b68a1882f3264fb8 # v4
- name: Setup Docker Buildx
uses: docker/setup-buildx-action@v4
uses: docker/setup-buildx-action@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c # v4
- name: Install Cosign
uses: sigstore/cosign-installer@6f9f17788090df1f26f669e9d70d6ae9567deba6 # v4.1.2
- name: Login to GHCR
uses: docker/login-action@v4
uses: docker/login-action@c99871dec2022cc055c062a10cc1a1310835ceb4 # v4
with:
registry: ${{ env.REGISTRY }}
username: ${{ github.actor }}
password: ${{ secrets.GITHUB_TOKEN }}
- name: Build and push app image
uses: docker/build-push-action@v7
id: build-and-push
uses: docker/build-push-action@53b7df96c91f9c12dcc8a07bcb9ccacbed38856a # v7
with:
context: .
file: deploy/Dockerfile
@@ -433,6 +450,8 @@ jobs:
CPR_PREBUILT_DIR=dist/docker
cache-from: type=gha,scope=runtime-prebuilt
cache-to: type=gha,mode=max,scope=runtime-prebuilt
provenance: mode=max
sbom: true
labels: |
org.opencontainers.image.title=Codex Proxy RS
org.opencontainers.image.version=${{ needs.release-info.outputs.version }}
@@ -443,8 +462,38 @@ jobs:
${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:latest
${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:sha-${{ needs.release-info.outputs.git_sha }}
- name: Generate container provenance attestation
uses: actions/attest@a1948c3f048ba23858d222213b7c278aabede763 # v4
with:
subject-name: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
subject-digest: ${{ steps.build-and-push.outputs.digest }}
push-to-registry: true
- name: Sign app image
shell: bash
env:
DIGEST: ${{ steps.build-and-push.outputs.digest }}
IMAGE_TAGS: |
${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:${{ needs.release-info.outputs.version }}
${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:latest
${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:sha-${{ needs.release-info.outputs.git_sha }}
run: |
set -euo pipefail
if [ -z "${DIGEST}" ]; then
echo "missing image digest"
exit 1
fi
images=()
while IFS= read -r tag; do
[ -n "${tag}" ] || continue
images+=("${tag}@${DIGEST}")
done <<< "${IMAGE_TAGS}"
cosign sign --yes "${images[@]}"
- name: Publish GitHub Release
uses: softprops/action-gh-release@v3
uses: softprops/action-gh-release@718ea10b132b3b2eba29c1007bb80653f286566b # v3
with:
tag_name: ${{ needs.release-info.outputs.tag }}
generate_release_notes: true
@@ -461,6 +510,8 @@ jobs:
**Manual download:**
Download the appropriate archive for your platform from the assets below.
Release archives are covered by GitHub artifact attestations. Docker images are published with SBOM/provenance attestations and keyless cosign signatures.
## 📚 Documentation
- [GitHub Repository](https://github.com/${{ github.repository }})
.github/workflows/security-scan.yml
+25 -6
View File
@@ -21,13 +21,13 @@ jobs:
timeout-minutes: 20
steps:
- name: Checkout
uses: actions/checkout@v7
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7
- name: Setup Rust
uses: dtolnay/rust-toolchain@stable
uses: dtolnay/rust-toolchain@4be7066ada62dd38de10e7b70166bc74ed198c30 # stable
- name: Cache Rust
uses: Swatinem/rust-cache@v2
uses: Swatinem/rust-cache@e18b497796c12c097a38f9edb9d0641fb99eee32 # v2
with:
workspaces: backend -> target
@@ -43,15 +43,15 @@ jobs:
timeout-minutes: 15
steps:
- name: Checkout
uses: actions/checkout@v7
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7
- name: Setup pnpm
uses: pnpm/action-setup@v6
uses: pnpm/action-setup@0ebf47130e4866e96fce0953f49152a61190b271 # v6
with:
version: 11.7.0
- name: Setup Node
uses: actions/setup-node@v6
uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6
with:
node-version: 24
cache: pnpm
@@ -62,3 +62,22 @@ jobs:
- name: Run pnpm audit
run: pnpm --dir frontend audit --prod --audit-level=high
docker-security:
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- name: Checkout
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7
- name: Build runtime image
run: docker build -f deploy/Dockerfile --target runtime -t codex-proxy-rs:security-scan .
- name: Scan runtime image
uses: aquasecurity/trivy-action@915b19bbe73b92a6cf82a1bc12b087c9a19a5fe2 # v0.28.0
with:
image-ref: codex-proxy-rs:security-scan
severity: HIGH,CRITICAL
vuln-type: os,library
ignore-unfixed: true
exit-code: '1'
README.md
+6 -3
View File
@@ -23,18 +23,18 @@
cargo run --manifest-path backend/Cargo.toml
```
启动后监听 `http://0.0.0.0:8080`，配置从运行目录下的 `config.yaml` 读取。首次启动会自动创建管理端账号（默认用户名密码见配置文件）。
启动后监听 `http://0.0.0.0:8080`，默认配置从运行目录下的 `config.yaml` 读取，也可以用 `CPR_CONFIG_FILE=/path/to/config.yaml` 显式指定。首次启动会自动创建管理端账号；`admin.default_password` 必须显式设置为至少 12 位的非默认密码。
Docker 部署：
```bash
cp deploy/config.example.yaml deploy/config.yaml
# 编辑 deploy/config.yaml，至少改掉 admin.default_password
# 编辑 deploy/config.yaml，设置 admin.default_password
docker compose -f deploy/docker-compose.yml build
docker compose -f deploy/docker-compose.yml up -d
```
挂载关系：`deploy/config.yaml` → `/app/config.yaml`，`deploy/data` → `/app/data`，`deploy/logs` → `/app/logs`。
默认只监听宿主机 `127.0.0.1:8080`。挂载关系：`deploy/config.yaml` → `/app/config.yaml`，数据和日志写入 Compose 命名卷 `cpr-data`、`cpr-logs`。
## API
@@ -90,6 +90,9 @@ database:
logging:
directory: .runtime/logs
retention_days: 14
admin:
default_password: '<set-a-long-random-password>'
```
完整配置项见 `deploy/config.example.yaml`，包括 TLS、WebSocket 连接池、请求指纹、账号额度策略等。
backend/Cargo.lock
Generated
+2 -132
View File
@@ -329,6 +329,7 @@ dependencies = [
"futures",
"hex",
"indexmap",
"ipnet",
"rand 0.10.1",
"rand_core 0.6.4",
"reqwest",
@@ -339,8 +340,7 @@ dependencies = [
"semver",
"serde",
"serde_json",
"serde_yml",
"sha2 0.11.0",
"sha2 0.10.9",
"sqlx",
"subtle",
"tar",
@@ -355,7 +355,6 @@ dependencies = [
"tracing-subscriber",
"tungstenite",
"uuid",
"validator",
"wiremock",
]
@@ -408,12 +407,6 @@ dependencies = [
"yaml-rust2",
]
[[package]]
name = "const-oid"
version = "0.10.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c"
[[package]]
name = "const-random"
version = "0.1.18"
@@ -588,41 +581,6 @@ dependencies = [
"cmov",
]
[[package]]
name = "darling"
version = "0.20.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee"
dependencies = [
"darling_core",
"darling_macro",
]
[[package]]
name = "darling_core"
version = "0.20.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0d00b9596d185e565c2207a0b01f8bd1a135483d02d9b7b0a54b11da8d53412e"
dependencies = [
"fnv",
"ident_case",
"proc-macro2",
"quote",
"strsim",
"syn",
]
[[package]]
name = "darling_macro"
version = "0.20.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead"
dependencies = [
"darling_core",
"quote",
"syn",
]
[[package]]
name = "data-encoding"
version = "2.11.0"
@@ -671,7 +629,6 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2"
dependencies = [
"block-buffer 0.12.1",
"const-oid",
"crypto-common 0.2.2",
"ctutils",
]
@@ -1343,12 +1300,6 @@ dependencies = [
"zerovec",
]
[[package]]
name = "ident_case"
version = "1.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39"
[[package]]
name = "idna"
version = "1.1.0"
@@ -1566,19 +1517,6 @@ dependencies = [
"windows-sys 0.61.2",
]
[[package]]
name = "noyalib"
version = "0.0.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e493c05128df7a83b9676b709d590e0ebc285c7ed3152bc679668e8c1e506af5"
dependencies = [
"indexmap",
"memchr",
"rustc-hash",
"serde",
"smallvec",
]
[[package]]
name = "nu-ansi-term"
version = "0.50.3"
@@ -1772,28 +1710,6 @@ dependencies = [
"zerocopy",
]
[[package]]
name = "proc-macro-error-attr2"
version = "2.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "96de42df36bb9bba5542fe9f1a054b8cc87e172759a1868aa05c1f3acc89dfc5"
dependencies = [
"proc-macro2",
"quote",
]
[[package]]
name = "proc-macro-error2"
version = "2.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "11ec05c52be0a07b08061f7dd003e7d7092e0472bc731b4af7bb1ef876109802"
dependencies = [
"proc-macro-error-attr2",
"proc-macro2",
"quote",
"syn",
]
[[package]]
name = "proc-macro2"
version = "1.0.106"
@@ -2299,16 +2215,6 @@ dependencies = [
"serde",
]
[[package]]
name = "serde_yml"
version = "0.0.13"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "909764a65f86829ccdb5eea9ab355843aa02c019a7bfd47465092953565caa05"
dependencies = [
"noyalib",
"serde",
]
[[package]]
name = "sha1"
version = "0.10.6"
@@ -2616,12 +2522,6 @@ dependencies = [
"unicode-properties",
]
[[package]]
name = "strsim"
version = "0.11.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f"
[[package]]
name = "subtle"
version = "2.6.1"
@@ -3169,36 +3069,6 @@ dependencies = [
"wasm-bindgen",
]
[[package]]
name = "validator"
version = "0.20.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "43fb22e1a008ece370ce08a3e9e4447a910e92621bb49b85d6e48a45397e7cfa"
dependencies = [
"idna",
"once_cell",
"regex",
"serde",
"serde_derive",
"serde_json",
"url",
"validator_derive",
]
[[package]]
name = "validator_derive"
version = "0.20.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b7df16e474ef958526d1205f6dda359fdfab79d9aa6d54bafcb92dcd07673dca"
dependencies = [
"darling",
"once_cell",
"proc-macro-error2",
"proc-macro2",
"quote",
"syn",
]
[[package]]
name = "valuable"
version = "0.1.1"
backend/Cargo.toml
+2 -3
View File
@@ -19,6 +19,7 @@ flate2 = "1.1.9"
futures = "0.3.32"
hex = "0.4.3"
indexmap = { version = "2.0", features = ["serde"] }
ipnet = "2.12.0"
rand = "0.10.1"
rand_core = { version = "0.6.4", features = ["getrandom"] }
reqwest = { version = "=0.12.28", default-features = false, features = ["json", "stream", "cookies", "rustls-tls-native-roots", "gzip", "brotli", "zstd", "deflate", "http2"] }
@@ -29,7 +30,7 @@ secrecy = "0.10.3"
semver = "1.0.27"
serde = { version = "1.0.228", features = ["derive"] }
serde_json = "1.0.150"
sha2 = "0.11.0"
sha2 = "0.10.9"
sqlx = { version = "0.9.0", features = ["runtime-tokio", "tls-rustls", "sqlite", "chrono", "uuid"] }
subtle = "2.6.1"
tar = "0.4.44"
@@ -42,10 +43,8 @@ tracing-appender = "0.2.5"
tracing-subscriber = { version = "0.3.23", features = ["env-filter", "json"] }
tungstenite = { version = "0.27.0", features = ["deflate", "proxy"] }
uuid = { version = "1.23.3", features = ["v4", "v7", "serde"] }
validator = { version = "0.20.0", features = ["derive"] }
[dev-dependencies]
serde_yml = "0.0.13"
tempfile = "3.27.0"
tower = { version = "0.5.3", features = ["util"] }
wiremock = "0.6.5"
backend/src/admin/accounts/mod.rs
+1
View File
@@ -1,4 +1,5 @@
//! 管理端账号 HTTP 模块。
pub(crate) mod quota_view;
pub mod routes;
pub mod service;
backend/src/admin/accounts/quota_view.rs
+375
View File
@@ -0,0 +1,375 @@
//! 管理端账号 quota 展示视图。
use chrono::{DateTime, Utc};
use serde::Serialize;
use serde_json::Value;
use crate::infra::{
format::format_percent,
time::{china_datetime, china_relative_time},
};
const FIVE_HOUR_WINDOW_SECONDS: u64 = 18_000;
const WEEK_WINDOW_SECONDS: u64 = 604_800;
const MONTH_WINDOW_SECONDS: u64 = 2_592_000;
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct AdminAccountQuotaData {
refreshed_at_display: String,
windows: Vec<AdminAccountQuotaWindowData>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AdminAccountQuotaWindowData {
key: String,
group: String,
window_seconds: Option<u64>,
label_display: String,
used_percent: Option<f64>,
used_percent_display: String,
reset_at_display: String,
window_used_display: String,
}
impl Default for AdminAccountQuotaData {
fn default() -> Self {
Self {
refreshed_at_display: "-".to_string(),
windows: Vec::new(),
}
}
}
impl AdminAccountQuotaData {
pub(crate) fn has_high_usage(&self) -> bool {
self.windows
.iter()
.any(|window| window.used_percent.is_some_and(|percent| percent >= 80.0))
}
}
pub(crate) fn quota_data(
quota_json: &str,
fetched_at: Option<DateTime<Utc>>,
) -> AdminAccountQuotaData {
let quota = serde_json::from_str::<Value>(quota_json).unwrap_or(Value::Null);
let windows = quota_windows(&quota);
AdminAccountQuotaData {
refreshed_at_display: china_relative_time(fetched_at, Utc::now()),
windows,
}
}
fn quota_windows(quota: &Value) -> Vec<AdminAccountQuotaWindowData> {
let mut windows = Vec::new();
let has_monthly_limit = push_monthly_quota_window(&mut windows, quota.get("monthly_limit"));
if let Some(snapshots) = quota.get("snapshots").and_then(Value::as_array) {
for snapshot in snapshots {
push_snapshot_quota_windows(&mut windows, snapshot, has_monthly_limit);
}
}
windows.sort_by_key(quota_window_sort_key);
windows
}
fn push_monthly_quota_window(
windows: &mut Vec<AdminAccountQuotaWindowData>,
monthly_limit: Option<&Value>,
) -> bool {
let Some(monthly_limit) = monthly_limit.filter(|value| !value.is_null()) else {
return false;
};
let used_percent = monthly_limit
.get("used_percent")
.and_then(number_value)
.map(|value| value.clamp(0.0, 100.0));
let reset_at = monthly_limit
.get("reset_at")
.and_then(Value::as_i64)
.and_then(|value| DateTime::<Utc>::from_timestamp(value, 0));
let window_seconds = monthly_limit
.get("window_minutes")
.and_then(Value::as_u64)
.and_then(|minutes| minutes.checked_mul(60))
.or(Some(MONTH_WINDOW_SECONDS));
if used_percent.is_none() && reset_at.is_none() {
return false;
}
let key = monthly_limit
.get("key")
.and_then(Value::as_str)
.unwrap_or("monthly");
windows.push(AdminAccountQuotaWindowData {
key: quota_key_segment(key),
group: "monthly".to_string(),
window_seconds,
label_display: "月限额".to_string(),
used_percent,
used_percent_display: used_percent.map_or_else(|| "-".to_string(), format_percent),
reset_at_display: reset_at
.as_ref()
.map_or_else(|| "-".to_string(), china_datetime),
window_used_display: quota_window_used_display(reset_at, window_seconds),
});
true
}
fn push_snapshot_quota_windows(
windows: &mut Vec<AdminAccountQuotaWindowData>,
snapshot: &Value,
skip_core_monthly: bool,
) {
let source = snapshot
.get("source")
.and_then(Value::as_str)
.unwrap_or("quota");
let source_key = snapshot_source_key(source, snapshot);
let label_prefix = snapshot_label(snapshot);
for role in ["primary", "secondary"] {
let Some(window) = snapshot.get(role).filter(|value| !value.is_null()) else {
continue;
};
let window_seconds = window
.get("window_minutes")
.and_then(Value::as_u64)
.and_then(|minutes| minutes.checked_mul(60));
if skip_core_monthly
&& source == "core"
&& window_seconds
.is_some_and(|seconds| quota_window_matches(seconds, MONTH_WINDOW_SECONDS))
{
continue;
}
push_quota_window(
windows,
&source_key,
role,
label_prefix.as_deref(),
Some(window),
);
}
}
fn push_quota_window(
windows: &mut Vec<AdminAccountQuotaWindowData>,
source_key: &str,
role: &str,
label_prefix: Option<&str>,
window: Option<&Value>,
) {
let Some(window) = window.filter(|value| !value.is_null()) else {
return;
};
let used_percent = window
.get("used_percent")
.and_then(number_value)
.map(|value| value.clamp(0.0, 100.0));
let reset_at = window
.get("reset_at")
.and_then(Value::as_i64)
.and_then(|value| DateTime::<Utc>::from_timestamp(value, 0));
let window_seconds = window
.get("window_minutes")
.and_then(Value::as_u64)
.and_then(|minutes| minutes.checked_mul(60));
let label_prefix = label_prefix
.map(str::trim)
.filter(|value| !value.is_empty());
if reset_at.is_none()
&& window_seconds.is_none()
&& (used_percent.is_none() || label_prefix.is_none())
{
return;
}
let base_label = quota_window_label_display(window_seconds);
let label_display = label_prefix
.map(|value| format!("{value} · {base_label}"))
.unwrap_or(base_label);
windows.push(AdminAccountQuotaWindowData {
key: unique_quota_window_key(windows, source_key, role, window_seconds),
group: quota_window_group(window_seconds).to_string(),
window_seconds,
label_display,
used_percent,
used_percent_display: used_percent.map_or_else(|| "-".to_string(), format_percent),
reset_at_display: reset_at
.as_ref()
.map_or_else(|| "-".to_string(), china_datetime),
window_used_display: quota_window_used_display(reset_at, window_seconds),
});
}
fn snapshot_source_key(source: &str, snapshot: &Value) -> String {
let label = snapshot
.get("limit_name")
.and_then(Value::as_str)
.or_else(|| snapshot.get("metered_feature").and_then(Value::as_str))
.unwrap_or(source);
format!("{}-{}", quota_key_segment(source), quota_key_segment(label))
}
fn snapshot_label(snapshot: &Value) -> Option<String> {
let source = snapshot.get("source").and_then(Value::as_str);
if source == Some("core") {
return None;
}
let label = snapshot
.get("limit_name")
.and_then(Value::as_str)
.or_else(|| snapshot.get("metered_feature").and_then(Value::as_str))
.map(str::trim)
.filter(|value| !value.is_empty())?;
if label.eq_ignore_ascii_case("codex") {
return None;
}
if is_review_limit_label(Some(label)) {
return Some("代码审查".to_string());
}
Some(label.to_string())
}
fn unique_quota_window_key(
windows: &[AdminAccountQuotaWindowData],
source_key: &str,
role: &str,
window_seconds: Option<u64>,
) -> String {
let bucket = quota_window_key_part(window_seconds).unwrap_or(role);
let key = format!("{source_key}-{bucket}");
if windows.iter().any(|window| window.key == key) {
format!("{key}-{role}")
} else {
key
}
}
fn quota_window_key_part(window_seconds: Option<u64>) -> Option<&'static str> {
match window_seconds {
Some(seconds) if quota_window_matches(seconds, FIVE_HOUR_WINDOW_SECONDS) => {
Some("five-hour")
}
Some(seconds) if quota_window_matches(seconds, WEEK_WINDOW_SECONDS) => Some("weekly"),
Some(seconds) if quota_window_matches(seconds, MONTH_WINDOW_SECONDS) => Some("monthly"),
_ => None,
}
}
fn quota_key_segment(value: &str) -> String {
let mut segment = String::with_capacity(value.len());
for ch in value.chars() {
if ch.is_ascii_alphanumeric() {
segment.push(ch.to_ascii_lowercase());
} else if !segment.ends_with('-') {
segment.push('-');
}
}
let segment = segment.trim_matches('-');
if segment.is_empty() {
"quota".to_string()
} else {
segment.to_string()
}
}
fn quota_window_sort_key(window: &AdminAccountQuotaWindowData) -> (u8, u64, String) {
let group_order = match window.group.as_str() {
"monthly" => 0,
"shortTerm" => 1,
_ => 2,
};
(
group_order,
window.window_seconds.unwrap_or(0),
window.key.clone(),
)
}
fn quota_window_group(window_seconds: Option<u64>) -> &'static str {
match window_seconds {
Some(seconds) if quota_window_matches(seconds, MONTH_WINDOW_SECONDS) => "monthly",
Some(seconds)
if quota_window_matches(seconds, FIVE_HOUR_WINDOW_SECONDS)
|| quota_window_matches(seconds, WEEK_WINDOW_SECONDS) =>
{
"shortTerm"
}
_ => "other",
}
}
fn quota_window_matches(actual: u64, expected: u64) -> bool {
actual > 0 && actual.abs_diff(expected) <= expected / 20
}
fn quota_window_label_display(window_seconds: Option<u64>) -> String {
let Some(window_seconds) = window_seconds.filter(|seconds| *seconds > 0) else {
return "额度".to_string();
};
match window_seconds {
seconds if quota_window_matches(seconds, FIVE_HOUR_WINDOW_SECONDS) => {
"5小时限额".to_string()
}
seconds if quota_window_matches(seconds, WEEK_WINDOW_SECONDS) => "周限额".to_string(),
seconds if quota_window_matches(seconds, MONTH_WINDOW_SECONDS) => "月限额".to_string(),
seconds if seconds % 86_400 == 0 => format!("{}天限额", seconds / 86_400),
seconds if seconds % 3_600 == 0 => format!("{}小时限额", seconds / 3_600),
seconds => format!("{}分钟限额", seconds.div_ceil(60)),
}
}
fn quota_window_used_display(
reset_at: Option<DateTime<Utc>>,
window_seconds: Option<u64>,
) -> String {
let (Some(reset_at), Some(window_seconds)) = (reset_at, window_seconds) else {
return "-".to_string();
};
let remaining = reset_at
.signed_duration_since(Utc::now())
.num_seconds()
.max(0)
.cast_unsigned();
let used = window_seconds.saturating_sub(remaining);
format!(
"{} / {}",
format_duration_days(used),
format_duration_days(window_seconds)
)
}
fn format_duration_days(seconds: u64) -> String {
let days = seconds as f64 / 86_400.0;
if days >= 1.0 {
format!("{days:.1}d")
} else {
format!("{:.1}h", seconds as f64 / 3_600.0)
}
}
fn number_value(value: &Value) -> Option<f64> {
value
.as_f64()
.or_else(|| value.as_str().and_then(|value| value.parse::<f64>().ok()))
.filter(|value| value.is_finite())
}
fn is_review_limit_label(value: Option<&str>) -> bool {
let normalized = value
.unwrap_or_default()
.trim()
.to_ascii_lowercase()
.replace(['-', ' '], "_");
normalized == "review"
|| normalized == "code_review"
|| normalized == "codex_review"
|| normalized == "codex_code_review"
|| normalized.contains("code_review")
|| normalized.contains("codex_review")
}
backend/src/admin/accounts/routes.rs
+171 -429
View File
@@ -2,7 +2,7 @@
use axum::{
body::Body,
extract::{Query, State},
extract::{Extension, Query, State},
http::{
header::{CACHE_CONTROL, CONNECTION, CONTENT_TYPE},
HeaderMap, StatusCode,
@@ -14,24 +14,30 @@ use std::collections::HashMap;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use serde_json::{json, Value};
use sqlx::{QueryBuilder, Row, Sqlite};
use validator::{Validate, ValidationError};
use crate::{
admin::auth::session::require_admin_auth,
admin::auth::session::AdminAuth,
admin::response::{
AdminEnvelope, AdminError, AdminResponse, BatchDeleteData, CursorPageMeta,
NumberedPageMeta, PageMeta, ADMIN_OK_CODE, ADMIN_OK_MESSAGE,
},
admin::update_payload::{parse_editable_update, EditableUpdateMessages},
admin::{
accounts::quota_view::{quota_data, AdminAccountQuotaData},
accounts::service::{
AdminAccountError, AdminAccountHealthCheck, AdminAccountMetadata,
AdminAccountRefreshOutcome, AdminAccountRefreshResult, AdminAccountUpdate,
OAuthExchangeInput,
},
monitoring::{billing, service::AdminUsageRecord},
monitoring::{
account_usage_service::AdminUsageRecord,
billing,
usage_record_model::{UsageRecord, UsageRecordLevel},
},
},
http::middleware::request_id::RequestId,
infra::{
format::{
format_cost, format_percent, format_plain_number, format_tokens, nonnegative_i64_to_u64,
@@ -44,9 +50,7 @@ use crate::{
};
const ACCOUNT_STATS_PAGE_LIMIT: u32 = 200;
const FIVE_HOUR_WINDOW_SECONDS: u64 = 18_000;
const WEEK_WINDOW_SECONDS: u64 = 604_800;
const MONTH_WINDOW_SECONDS: u64 = 2_592_000;
const ACCOUNT_EXPORT_CONFIRMATION: &str = "export_sensitive_accounts";
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
@@ -87,6 +91,7 @@ pub(crate) struct AccountIdQuery {
#[serde(rename_all = "camelCase")]
pub(crate) struct AccountExportQuery {
ids: Option<String>,
confirm: Option<String>,
}
#[derive(Debug, Deserialize)]
@@ -174,35 +179,6 @@ impl AdminAccountData {
}
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AdminAccountQuotaData {
refreshed_at_display: String,
windows: Vec<AdminAccountQuotaWindowData>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AdminAccountQuotaWindowData {
key: String,
group: String,
window_seconds: Option<u64>,
label_display: String,
used_percent: Option<f64>,
used_percent_display: String,
reset_at_display: String,
window_used_display: String,
}
impl Default for AdminAccountQuotaData {
fn default() -> Self {
Self {
refreshed_at_display: "-".to_string(),
windows: Vec::new(),
}
}
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AdminAccountUsageData {
@@ -494,10 +470,9 @@ struct AccountOAuthAuthorizeData {
/// `GET /api/admin/accounts`
pub(crate) async fn accounts(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
Query(params): Query<AccountsQuery>,
) -> Result<impl IntoResponse, AdminError> {
require_admin_auth(&state, &headers).await?;
let limit = clamp_limit(params.page_size.or(params.limit).unwrap_or(50));
let use_numbered_page = params.page.is_some() || params.page_size.is_some();
let stats = account_list_stats(&state).await;
@@ -557,10 +532,9 @@ pub(crate) async fn accounts(
/// `POST /api/admin/accounts`
pub(crate) async fn create_account(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
Json(payload): Json<CreateAccountRequest>,
) -> Result<impl IntoResponse, AdminError> {
require_admin_auth(&state, &headers).await?;
match state
.services
.admin_accounts
@@ -578,32 +552,65 @@ pub(crate) async fn create_account(
/// `GET /api/admin/accounts/export`
pub(crate) async fn export_accounts(
State(state): State<AppState>,
request_id: Option<Extension<RequestId>>,
_auth: AdminAuth,
headers: HeaderMap,
Query(query): Query<AccountExportQuery>,
) -> Result<impl IntoResponse, AdminError> {
require_admin_auth(&state, &headers).await?;
match state
.services
.admin_accounts
.export(account_export_ids(query.ids.as_deref()))
.await
{
Ok(result) => Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(result),
)),
Err(error) => Err(account_error(&error)),
let request_id = request_id.map(|Extension(request_id)| request_id.as_str().to_string());
let source = account_action_source(&headers);
let ids = account_export_ids(query.ids.as_deref());
if query.confirm.as_deref() != Some(ACCOUNT_EXPORT_CONFIRMATION) {
record_account_export_audit(
&state,
request_id.as_deref(),
&source,
&ids,
AccountExportAuditOutcome::MissingConfirmation,
)
.await;
return Err(AdminError::bad_request(
"account export requires confirm=export_sensitive_accounts",
));
}
match state.services.admin_accounts.export(ids.clone()).await {
Ok(result) => {
record_account_export_audit(
&state,
request_id.as_deref(),
&source,
&ids,
AccountExportAuditOutcome::Success,
)
.await;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(result),
))
}
Err(error) => {
record_account_export_audit(
&state,
request_id.as_deref(),
&source,
&ids,
AccountExportAuditOutcome::from_error(&error),
)
.await;
Err(account_error(&error))
}
}
}
/// `POST /api/admin/accounts/refresh`
pub(crate) async fn refresh_account(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
Json(payload): Json<AccountActionRequest>,
) -> Result<impl IntoResponse, AdminError> {
let account_id = payload.id;
require_admin_auth(&state, &headers).await?;
match state
.services
.admin_accounts
@@ -621,10 +628,9 @@ pub(crate) async fn refresh_account(
/// `POST /api/admin/accounts/health-check`
pub(crate) async fn health_check_accounts(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
Json(payload): Json<AccountHealthCheckRequest>,
) -> Result<impl IntoResponse, AdminError> {
require_admin_auth(&state, &headers).await?;
let stagger_ms = payload.stagger_ms.unwrap_or(3_000);
if !(500..=30_000).contains(&stagger_ms) && stagger_ms != 0 {
return Err(AdminError::bad_request(
@@ -655,11 +661,10 @@ pub(crate) async fn health_check_accounts(
/// `GET /api/admin/accounts/quota`
pub(crate) async fn account_quota(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
Query(query): Query<AccountIdQuery>,
) -> Result<impl IntoResponse, AdminError> {
let account_id = query.id;
require_admin_auth(&state, &headers).await?;
match state
.services
.admin_accounts
@@ -700,10 +705,9 @@ pub(crate) async fn account_quota(
/// `POST /api/admin/accounts/import`
pub(crate) async fn import_accounts(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
Json(payload): Json<Value>,
) -> Result<impl IntoResponse, AdminError> {
require_admin_auth(&state, &headers).await?;
match state.services.admin_accounts.import(payload).await {
Ok(result) => Ok(AdminResponse::new(
StatusCode::OK,
@@ -720,9 +724,8 @@ pub(crate) async fn import_accounts(
/// `POST /api/admin/accounts/oauth/authorize`
pub(crate) async fn oauth_authorize_account(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
) -> Result<impl IntoResponse, AdminError> {
require_admin_auth(&state, &headers).await?;
match state.services.admin_accounts.oauth_authorize().await {
Ok(result) => Ok(AdminResponse::new(
StatusCode::OK,
@@ -740,10 +743,9 @@ pub(crate) async fn oauth_authorize_account(
/// `POST /api/admin/accounts/oauth/exchange`
pub(crate) async fn oauth_exchange_account(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
Json(payload): Json<AccountOAuthExchangeRequest>,
) -> Result<impl IntoResponse, AdminError> {
require_admin_auth(&state, &headers).await?;
match state
.services
.admin_accounts
@@ -771,10 +773,9 @@ pub(crate) async fn oauth_exchange_account(
pub(crate) async fn test_account_connection(
State(state): State<AppState>,
Query(query): Query<AccountTestQuery>,
headers: HeaderMap,
_auth: AdminAuth,
) -> Result<Response, AdminError> {
let account_id = query.id;
require_admin_auth(&state, &headers).await?;
let model = query
.model_id
@@ -804,10 +805,9 @@ pub(crate) async fn test_account_connection(
pub(crate) async fn account_models(
State(state): State<AppState>,
Query(query): Query<AccountIdQuery>,
headers: HeaderMap,
_auth: AdminAuth,
) -> Result<impl IntoResponse, AdminError> {
let account_id = query.id;
require_admin_auth(&state, &headers).await?;
match state
.services
.admin_accounts
@@ -833,10 +833,9 @@ pub(crate) async fn account_models(
/// `POST /api/admin/accounts/delete`
pub(crate) async fn batch_delete_accounts(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
Json(payload): Json<BatchDeleteAccountsRequest>,
) -> Result<impl IntoResponse, AdminError> {
require_admin_auth(&state, &headers).await?;
match state
.services
.admin_accounts
@@ -857,11 +856,9 @@ pub(crate) async fn batch_delete_accounts(
/// `POST /api/admin/accounts/update`
pub(crate) async fn update_account(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
Json(payload): Json<Value>,
) -> Result<impl IntoResponse, AdminError> {
require_admin_auth(&state, &headers).await?;
let ParsedAccountUpdate { id, update } = parse_account_update(&payload)?;
match state
@@ -962,10 +959,7 @@ async fn list_all_account_metadata(state: &AppState) -> Vec<AdminAccountMetadata
}
fn account_quota_has_high_usage(quota: &AdminAccountQuotaData) -> bool {
quota
.windows
.iter()
.any(|window| window.used_percent.is_some_and(|percent| percent >= 80.0))
quota.has_high_usage()
}
fn account_summary_needs_attention(status: AccountStatus) -> bool {
@@ -1165,356 +1159,21 @@ fn model_usage_data(usage: AccountModelUsageRecord) -> AdminAccountModelUsageDat
}
}
fn quota_data(quota_json: &str, fetched_at: Option<DateTime<Utc>>) -> AdminAccountQuotaData {
let quota = serde_json::from_str::<Value>(quota_json).unwrap_or(Value::Null);
let windows = quota_windows(&quota);
AdminAccountQuotaData {
refreshed_at_display: china_relative_time(fetched_at, Utc::now()),
windows,
}
}
fn quota_windows(quota: &Value) -> Vec<AdminAccountQuotaWindowData> {
let mut windows = Vec::new();
let has_monthly_limit = push_monthly_quota_window(&mut windows, quota.get("monthly_limit"));
if let Some(snapshots) = quota.get("snapshots").and_then(Value::as_array) {
for snapshot in snapshots {
push_snapshot_quota_windows(&mut windows, snapshot, has_monthly_limit);
}
}
windows.sort_by_key(quota_window_sort_key);
windows
}
fn push_monthly_quota_window(
windows: &mut Vec<AdminAccountQuotaWindowData>,
monthly_limit: Option<&Value>,
) -> bool {
let Some(monthly_limit) = monthly_limit.filter(|value| !value.is_null()) else {
return false;
};
let used_percent = monthly_limit
.get("used_percent")
.and_then(number_value)
.map(|value| value.clamp(0.0, 100.0));
let reset_at = monthly_limit
.get("reset_at")
.and_then(Value::as_i64)
.and_then(|value| DateTime::<Utc>::from_timestamp(value, 0));
let window_seconds = monthly_limit
.get("window_minutes")
.and_then(Value::as_u64)
.and_then(|minutes| minutes.checked_mul(60))
.or(Some(MONTH_WINDOW_SECONDS));
if used_percent.is_none() && reset_at.is_none() {
return false;
}
let key = monthly_limit
.get("key")
.and_then(Value::as_str)
.unwrap_or("monthly");
windows.push(AdminAccountQuotaWindowData {
key: quota_key_segment(key),
group: "monthly".to_string(),
window_seconds,
label_display: "月限额".to_string(),
used_percent,
used_percent_display: used_percent.map_or_else(|| "-".to_string(), format_percent),
reset_at_display: reset_at
.as_ref()
.map_or_else(|| "-".to_string(), china_datetime),
window_used_display: quota_window_used_display(reset_at, window_seconds),
});
true
}
fn push_snapshot_quota_windows(
windows: &mut Vec<AdminAccountQuotaWindowData>,
snapshot: &Value,
skip_core_monthly: bool,
) {
let source = snapshot
.get("source")
.and_then(Value::as_str)
.unwrap_or("quota");
let source_key = snapshot_source_key(source, snapshot);
let label_prefix = snapshot_label(snapshot);
for role in ["primary", "secondary"] {
let Some(window) = snapshot.get(role).filter(|value| !value.is_null()) else {
continue;
};
let window_seconds = window
.get("window_minutes")
.and_then(Value::as_u64)
.and_then(|minutes| minutes.checked_mul(60));
if skip_core_monthly
&& source == "core"
&& window_seconds
.is_some_and(|seconds| quota_window_matches(seconds, MONTH_WINDOW_SECONDS))
{
continue;
}
push_quota_window(
windows,
&source_key,
role,
label_prefix.as_deref(),
Some(window),
);
}
}
fn push_quota_window(
windows: &mut Vec<AdminAccountQuotaWindowData>,
source_key: &str,
role: &str,
label_prefix: Option<&str>,
window: Option<&Value>,
) {
let Some(window) = window.filter(|value| !value.is_null()) else {
return;
};
let used_percent = window
.get("used_percent")
.and_then(number_value)
.map(|value| value.clamp(0.0, 100.0));
let reset_at = window
.get("reset_at")
.and_then(Value::as_i64)
.and_then(|value| DateTime::<Utc>::from_timestamp(value, 0));
let window_seconds = window
.get("window_minutes")
.and_then(Value::as_u64)
.and_then(|minutes| minutes.checked_mul(60));
let label_prefix = label_prefix
.map(str::trim)
.filter(|value| !value.is_empty());
if reset_at.is_none()
&& window_seconds.is_none()
&& (used_percent.is_none() || label_prefix.is_none())
{
return;
}
let base_label = quota_window_label_display(window_seconds);
let label_display = label_prefix
.map(|value| format!("{value} · {base_label}"))
.unwrap_or(base_label);
windows.push(AdminAccountQuotaWindowData {
key: unique_quota_window_key(windows, source_key, role, window_seconds),
group: quota_window_group(window_seconds).to_string(),
window_seconds,
label_display,
used_percent,
used_percent_display: used_percent.map_or_else(|| "-".to_string(), format_percent),
reset_at_display: reset_at
.as_ref()
.map_or_else(|| "-".to_string(), china_datetime),
window_used_display: quota_window_used_display(reset_at, window_seconds),
});
}
fn snapshot_source_key(source: &str, snapshot: &Value) -> String {
let label = snapshot
.get("limit_name")
.and_then(Value::as_str)
.or_else(|| snapshot.get("metered_feature").and_then(Value::as_str))
.unwrap_or(source);
format!("{}-{}", quota_key_segment(source), quota_key_segment(label))
}
fn snapshot_label(snapshot: &Value) -> Option<String> {
let source = snapshot.get("source").and_then(Value::as_str);
if source == Some("core") {
return None;
}
let label = snapshot
.get("limit_name")
.and_then(Value::as_str)
.or_else(|| snapshot.get("metered_feature").and_then(Value::as_str))
.map(str::trim)
.filter(|value| !value.is_empty())?;
if label.eq_ignore_ascii_case("codex") {
return None;
}
if is_review_limit_label(Some(label)) {
return Some("代码审查".to_string());
}
Some(label.to_string())
}
fn unique_quota_window_key(
windows: &[AdminAccountQuotaWindowData],
source_key: &str,
role: &str,
window_seconds: Option<u64>,
) -> String {
let bucket = quota_window_key_part(window_seconds).unwrap_or(role);
let key = format!("{source_key}-{bucket}");
if windows.iter().any(|window| window.key == key) {
format!("{key}-{role}")
} else {
key
}
}
fn quota_window_key_part(window_seconds: Option<u64>) -> Option<&'static str> {
match window_seconds {
Some(seconds) if quota_window_matches(seconds, FIVE_HOUR_WINDOW_SECONDS) => {
Some("five-hour")
}
Some(seconds) if quota_window_matches(seconds, WEEK_WINDOW_SECONDS) => Some("weekly"),
Some(seconds) if quota_window_matches(seconds, MONTH_WINDOW_SECONDS) => Some("monthly"),
_ => None,
}
}
fn quota_key_segment(value: &str) -> String {
let mut segment = String::with_capacity(value.len());
for ch in value.chars() {
if ch.is_ascii_alphanumeric() {
segment.push(ch.to_ascii_lowercase());
} else if !segment.ends_with('-') {
segment.push('-');
}
}
let segment = segment.trim_matches('-');
if segment.is_empty() {
"quota".to_string()
} else {
segment.to_string()
}
}
fn quota_window_sort_key(window: &AdminAccountQuotaWindowData) -> (u8, u64, String) {
let group_order = match window.group.as_str() {
"monthly" => 0,
"shortTerm" => 1,
_ => 2,
};
(
group_order,
window.window_seconds.unwrap_or(0),
window.key.clone(),
)
}
fn quota_window_group(window_seconds: Option<u64>) -> &'static str {
match window_seconds {
Some(seconds) if quota_window_matches(seconds, MONTH_WINDOW_SECONDS) => "monthly",
Some(seconds)
if quota_window_matches(seconds, FIVE_HOUR_WINDOW_SECONDS)
|| quota_window_matches(seconds, WEEK_WINDOW_SECONDS) =>
{
"shortTerm"
}
_ => "other",
}
}
fn quota_window_matches(actual: u64, expected: u64) -> bool {
actual > 0 && actual.abs_diff(expected) <= expected / 20
}
fn quota_window_label_display(window_seconds: Option<u64>) -> String {
let Some(window_seconds) = window_seconds.filter(|seconds| *seconds > 0) else {
return "额度".to_string();
};
match window_seconds {
seconds if quota_window_matches(seconds, FIVE_HOUR_WINDOW_SECONDS) => {
"5小时限额".to_string()
}
seconds if quota_window_matches(seconds, WEEK_WINDOW_SECONDS) => "周限额".to_string(),
seconds if quota_window_matches(seconds, MONTH_WINDOW_SECONDS) => "月限额".to_string(),
seconds if seconds % 86_400 == 0 => format!("{}天限额", seconds / 86_400),
seconds if seconds % 3_600 == 0 => format!("{}小时限额", seconds / 3_600),
seconds => format!("{}分钟限额", seconds.div_ceil(60)),
}
}
fn quota_window_used_display(
reset_at: Option<DateTime<Utc>>,
window_seconds: Option<u64>,
) -> String {
let (Some(reset_at), Some(window_seconds)) = (reset_at, window_seconds) else {
return "-".to_string();
};
let remaining = reset_at
.signed_duration_since(Utc::now())
.num_seconds()
.max(0)
.cast_unsigned();
let used = window_seconds.saturating_sub(remaining);
format!(
"{} / {}",
format_duration_days(used),
format_duration_days(window_seconds)
)
}
fn format_duration_days(seconds: u64) -> String {
let days = seconds as f64 / 86_400.0;
if days >= 1.0 {
format!("{days:.1}d")
} else {
format!("{:.1}h", seconds as f64 / 3_600.0)
}
}
fn number_value(value: &Value) -> Option<f64> {
value
.as_f64()
.or_else(|| value.as_str().and_then(|value| value.parse::<f64>().ok()))
.filter(|value| value.is_finite())
}
fn is_review_limit_label(value: Option<&str>) -> bool {
let normalized = value
.unwrap_or_default()
.trim()
.to_ascii_lowercase()
.replace(['-', ' '], "_");
normalized == "review"
|| normalized == "code_review"
|| normalized == "codex_review"
|| normalized == "codex_code_review"
|| normalized.contains("code_review")
|| normalized.contains("codex_review")
}
struct ParsedAccountUpdate {
id: String,
update: AdminAccountUpdate,
}
#[derive(Debug, Deserialize, Validate)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct AccountUpdateRequest {
#[validate(custom(function = "validate_non_empty_trimmed"))]
id: String,
label: Option<Option<String>>,
#[validate(custom(function = "validate_non_empty_trimmed"))]
status: Option<String>,
}
fn parse_account_update(payload: &Value) -> Result<ParsedAccountUpdate, AdminError> {
let object = payload
.as_object()
.ok_or_else(|| AdminError::bad_request("Account update request must be an object"))?;
for field in object.keys() {
if !matches!(field.as_str(), "id" | "label" | "status") {
return Err(AdminError::bad_request(format!("{field} is not editable")));
}
}
let payload = serde_json::from_value::<AccountUpdateRequest>(payload.clone())
.map_err(|_| AdminError::bad_request("Invalid account update request"))?;
payload
.validate()
.map_err(|_| AdminError::bad_request("Invalid account update request"))?;
let payload = parse_editable_update(
payload,
EditableUpdateMessages {
object_required: "Account update request must be an object",
invalid: "Invalid account update request",
empty_update: "Account update request must include editable fields",
unknown_field_editable: true,
},
)?;
let update = AdminAccountUpdate {
label: payload.label.map(|label| {
label.and_then(|value| {
@@ -1535,13 +1194,6 @@ fn parse_account_update(payload: &Value) -> Result<ParsedAccountUpdate, AdminErr
})
}
fn validate_non_empty_trimmed(value: &str) -> Result<(), ValidationError> {
if value.trim().is_empty() {
return Err(ValidationError::new("required"));
}
Ok(())
}
fn account_error(error: &AdminAccountError) -> AdminError {
match error {
AdminAccountError::InvalidStatus(_)
@@ -1576,6 +1228,96 @@ fn account_not_found() -> AdminError {
AdminError::not_found("Account not found")
}
enum AccountExportAuditOutcome {
Success,
MissingConfirmation,
EmptyIds,
NotFound,
Failed,
}
impl AccountExportAuditOutcome {
fn from_error(error: &AdminAccountError) -> Self {
match error {
AdminAccountError::EmptyIds => Self::EmptyIds,
AdminAccountError::NotFound => Self::NotFound,
_ => Self::Failed,
}
}
fn level(&self) -> UsageRecordLevel {
match self {
Self::Success => UsageRecordLevel::Warn,
Self::MissingConfirmation | Self::EmptyIds | Self::NotFound => UsageRecordLevel::Warn,
Self::Failed => UsageRecordLevel::Error,
}
}
fn message(&self) -> &'static str {
match self {
Self::Success => "Admin account export succeeded",
Self::MissingConfirmation => "Admin account export missing confirmation",
Self::EmptyIds => "Admin account export missing account ids",
Self::NotFound => "Admin account export account not found",
Self::Failed => "Admin account export failed",
}
}
fn status_code(&self) -> i64 {
match self {
Self::Success => i64::from(StatusCode::OK.as_u16()),
Self::MissingConfirmation | Self::EmptyIds => {
i64::from(StatusCode::BAD_REQUEST.as_u16())
}
Self::NotFound => i64::from(StatusCode::NOT_FOUND.as_u16()),
Self::Failed => i64::from(StatusCode::INTERNAL_SERVER_ERROR.as_u16()),
}
}
fn failure_class(&self) -> Option<&'static str> {
match self {
Self::Success => None,
Self::MissingConfirmation => Some("missing_confirmation"),
Self::EmptyIds => Some("empty_ids"),
Self::NotFound => Some("account_not_found"),
Self::Failed => Some("export_failed"),
}
}
}
async fn record_account_export_audit(
state: &AppState,
request_id: Option<&str>,
source: &str,
ids: &[String],
outcome: AccountExportAuditOutcome,
) {
let mut record = UsageRecord::new("admin_account_export", outcome.level(), outcome.message());
record.request_id = request_id.map(ToString::to_string);
record.route = Some("/api/admin/accounts/export".to_string());
record.status_code = Some(outcome.status_code());
record.failure_class = outcome.failure_class().map(ToString::to_string);
record.metadata = json!({
"source": source,
"accountIds": ids,
"accountCount": ids.len(),
});
if let Err(error) = state.services.usage_records.record_audit(record).await {
tracing::warn!(%error, "failed to record account export audit event");
}
}
fn account_action_source(headers: &HeaderMap) -> String {
headers
.get("x-real-ip")
.and_then(|value| value.to_str().ok())
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("unknown")
.to_string()
}
fn account_export_ids(value: Option<&str>) -> Vec<String> {
value
.into_iter()
backend/src/admin/accounts/service/exporting.rs
+4 -24
View File
@@ -9,15 +9,12 @@ use crate::{
upstream::accounts::{model::AccountStatus, store::StoredAccount},
};
const EXPORT_PAGE_LIMIT: u32 = 200;
impl AdminAccountService {
pub async fn export(&self, ids: Vec<String>) -> Result<ExportedAccounts, AdminAccountError> {
let accounts = if ids.is_empty() {
self.export_all_accounts().await?
} else {
self.export_accounts_by_id(ids).await?
};
if ids.is_empty() {
return Err(AdminAccountError::EmptyIds);
}
let accounts = self.export_accounts_by_id(ids).await?;
Ok(ExportedAccounts {
source_format: "cpr",
@@ -25,23 +22,6 @@ impl AdminAccountService {
})
}
async fn export_all_accounts(&self) -> Result<Vec<StoredAccount>, AdminAccountError> {
let mut accounts = Vec::new();
let mut cursor = None;
loop {
let page = self
.store
.list(cursor, EXPORT_PAGE_LIMIT)
.await
.map_err(|_| AdminAccountError::Export)?;
accounts.extend(page.items);
let Some(next_cursor) = page.next_cursor else {
return Ok(accounts);
};
cursor = Some(next_cursor);
}
}
async fn export_accounts_by_id(
&self,
ids: Vec<String>,
backend/src/admin/accounts/service/lifecycle.rs
+4 -5
View File
@@ -9,7 +9,10 @@ use secrecy::{ExposeSecret, SecretString};
use tokio::time::sleep;
use crate::{
infra::json::{NumberedPage, Page},
infra::{
json::{NumberedPage, Page},
time::elapsed_millis_i64,
},
upstream::accounts::store::AccountStore,
};
@@ -483,7 +486,3 @@ fn manual_refresh_failure_is_permanent(
| crate::upstream::accounts::token_refresh::RefreshFailure::Transport => false,
}
}
fn elapsed_millis_i64(started_at: std::time::Instant) -> i64 {
i64::try_from(started_at.elapsed().as_millis()).unwrap_or(i64::MAX)
}
backend/src/admin/auth/service.rs
+88
View File
@@ -1,12 +1,19 @@
//! 管理员会话存储与认证领域逻辑。
use std::{collections::HashMap, sync::Arc};
use chrono::{DateTime, Duration, Utc};
use sqlx::SqlitePool;
use thiserror::Error;
use tokio::sync::Mutex;
use uuid::Uuid;
use crate::infra::identity::{hash_admin_password, verify_admin_password};
const LOGIN_FAILURE_WINDOW: Duration = Duration::minutes(15);
const LOGIN_LOCK_DURATION: Duration = Duration::minutes(15);
const LOGIN_MAX_FAILURES: u32 = 5;
/// 管理员登录成功后的会话。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AdminLoginSession {
@@ -40,6 +47,9 @@ pub enum AdminSessionError {
/// 删除管理员会话失败。
#[error("failed to delete admin session")]
DeleteSession,
/// 管理员登录尝试过于频繁。
#[error("too many admin login attempts")]
LoginThrottled,
}
/// 管理员会话服务。
@@ -48,6 +58,14 @@ pub struct AdminSessionService {
store: SqliteAdminSessionStore,
default_username: String,
session_ttl_minutes: u64,
login_failures: Arc<Mutex<HashMap<String, LoginFailureState>>>,
}
#[derive(Debug, Clone)]
struct LoginFailureState {
failures: u32,
first_failed_at: DateTime<Utc>,
locked_until: Option<DateTime<Utc>>,
}
impl AdminSessionService {
@@ -61,6 +79,7 @@ impl AdminSessionService {
store,
default_username,
session_ttl_minutes,
login_failures: Arc::new(Mutex::new(HashMap::new())),
}
}
@@ -88,11 +107,16 @@ impl AdminSessionService {
/// 登录并创建管理员会话。
pub async fn login(
&self,
source: &str,
username: Option<&str>,
password: &str,
) -> Result<Option<AdminLoginSession>, AdminSessionError> {
let source = login_source_key(source);
self.ensure_login_allowed(&source).await?;
let username = username.unwrap_or(&self.default_username);
if username != self.default_username.as_str() {
self.record_failed_login(&source).await;
return Ok(None);
}
let Some(admin) = self
@@ -101,13 +125,16 @@ impl AdminSessionService {
.await
.map_err(|_| AdminSessionError::LoadAdmin)?
else {
self.record_failed_login(&source).await;
return Ok(None);
};
let password_matches = verify_admin_password(password, &admin.password_hash)
.map_err(|_| AdminSessionError::VerifyPassword)?;
if !password_matches {
self.record_failed_login(&source).await;
return Ok(None);
}
self.clear_failed_login(&source).await;
let session_id = format!("sess_{}", Uuid::new_v4().simple());
let ttl_minutes = self.session_ttl_minutes.min(i64::MAX as u64) as i64;
let expires_at = Utc::now() + Duration::minutes(ttl_minutes);
@@ -128,6 +155,67 @@ impl AdminSessionService {
.await
.map_err(|_| AdminSessionError::DeleteSession)
}
async fn ensure_login_allowed(&self, source: &str) -> Result<(), AdminSessionError> {
let now = Utc::now();
let mut failures = self.login_failures.lock().await;
retain_active_login_failures(&mut failures, now);
if failures
.get(source)
.and_then(|state| state.locked_until)
.is_some_and(|locked_until| locked_until > now)
{
return Err(AdminSessionError::LoginThrottled);
}
Ok(())
}
async fn record_failed_login(&self, source: &str) {
let now = Utc::now();
let mut failures = self.login_failures.lock().await;
retain_active_login_failures(&mut failures, now);
let state = failures
.entry(source.to_string())
.or_insert_with(|| LoginFailureState {
failures: 0,
first_failed_at: now,
locked_until: None,
});
if now - state.first_failed_at > LOGIN_FAILURE_WINDOW {
state.failures = 0;
state.first_failed_at = now;
state.locked_until = None;
}
state.failures += 1;
if state.failures >= LOGIN_MAX_FAILURES {
state.locked_until = Some(now + LOGIN_LOCK_DURATION);
}
}
async fn clear_failed_login(&self, source: &str) {
self.login_failures.lock().await.remove(source);
}
}
fn login_source_key(source: &str) -> String {
let source = source.trim();
if source.is_empty() {
"unknown".to_string()
} else {
source.to_string()
}
}
fn retain_active_login_failures(
failures: &mut HashMap<String, LoginFailureState>,
now: DateTime<Utc>,
) {
failures.retain(|_, state| {
state
.locked_until
.is_some_and(|locked_until| locked_until > now)
|| now - state.first_failed_at <= LOGIN_FAILURE_WINDOW
});
}
// ---------------------------------------------------------------------------
backend/src/admin/auth/session.rs
+164 -10
View File
@@ -1,20 +1,42 @@
//! 管理端会话处理器。
use axum::{
extract::State,
http::{header::SET_COOKIE, HeaderMap, HeaderValue, StatusCode},
extract::{Extension, FromRequestParts, State},
http::{header::SET_COOKIE, request::Parts, HeaderMap, HeaderValue, StatusCode},
response::{IntoResponse, Response},
Json,
};
use serde::{Deserialize, Serialize};
use serde_json::json;
use crate::{
admin::response::{AdminEnvelope, AdminError, AdminResponse},
admin::{
auth::service::AdminSessionError,
monitoring::usage_record_model::{UsageRecord, UsageRecordLevel},
response::{AdminEnvelope, AdminError, AdminResponse},
},
http::middleware::request_id::RequestId,
infra::time::china_rfc3339,
runtime::state::AppState,
};
const ADMIN_SESSION_COOKIE: &str = "cpr_admin_session";
const ADMIN_SESSION_COOKIE_ATTRS: &str = "Path=/; Secure; HttpOnly; SameSite=Lax";
/// 已通过管理员会话或管理员 API Key 鉴权的请求。
pub(crate) struct AdminAuth;
impl FromRequestParts<AppState> for AdminAuth {
type Rejection = AdminError;
async fn from_request_parts(
parts: &mut Parts,
state: &AppState,
) -> Result<Self, Self::Rejection> {
require_admin_auth(state, &parts.headers).await?;
Ok(Self)
}
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
@@ -38,15 +60,71 @@ struct AdminSessionStatusData {
/// `POST /api/admin/login`
pub(crate) async fn login(
State(state): State<AppState>,
request_id: Option<Extension<RequestId>>,
headers: HeaderMap,
Json(payload): Json<AdminLoginRequest>,
) -> Result<Response, AdminError> {
let session = state
let source = admin_login_source(&headers);
let request_id = request_id.map(|Extension(request_id)| request_id.as_str().to_string());
let username_provided = payload
.username
.as_deref()
.is_some_and(|username| !username.trim().is_empty());
let login_result = state
.services
.admin_sessions
.login(payload.username.as_deref(), payload.password.as_str())
.await
.map_err(|_| AdminError::internal("Failed to create admin session"))?
.ok_or_else(AdminError::invalid_admin_credentials)?;
.login(
&source,
payload.username.as_deref(),
payload.password.as_str(),
)
.await;
let session = match login_result {
Ok(Some(session)) => {
record_admin_login_audit(
&state,
request_id.as_deref(),
&source,
username_provided,
AdminLoginAuditOutcome::Success,
)
.await;
session
}
Ok(None) => {
record_admin_login_audit(
&state,
request_id.as_deref(),
&source,
username_provided,
AdminLoginAuditOutcome::InvalidCredentials,
)
.await;
return Err(AdminError::invalid_admin_credentials());
}
Err(AdminSessionError::LoginThrottled) => {
record_admin_login_audit(
&state,
request_id.as_deref(),
&source,
username_provided,
AdminLoginAuditOutcome::Throttled,
)
.await;
return Err(AdminError::too_many_login_attempts());
}
Err(_) => {
record_admin_login_audit(
&state,
request_id.as_deref(),
&source,
username_provided,
AdminLoginAuditOutcome::InternalError,
)
.await;
return Err(AdminError::internal("Failed to create admin session"));
}
};
let mut response = AdminResponse::new(
StatusCode::OK,
@@ -56,7 +134,7 @@ pub(crate) async fn login(
)
.into_response();
let cookie = format!(
"{ADMIN_SESSION_COOKIE}={}; Path=/; HttpOnly; SameSite=Lax",
"{ADMIN_SESSION_COOKIE}={}; {ADMIN_SESSION_COOKIE_ATTRS}",
session.session_id
);
response.headers_mut().insert(
@@ -67,6 +145,72 @@ pub(crate) async fn login(
Ok(response)
}
enum AdminLoginAuditOutcome {
Success,
InvalidCredentials,
Throttled,
InternalError,
}
impl AdminLoginAuditOutcome {
fn level(&self) -> UsageRecordLevel {
match self {
Self::Success => UsageRecordLevel::Info,
Self::InvalidCredentials | Self::Throttled => UsageRecordLevel::Warn,
Self::InternalError => UsageRecordLevel::Error,
}
}
fn message(&self) -> &'static str {
match self {
Self::Success => "Admin login succeeded",
Self::InvalidCredentials => "Admin login failed",
Self::Throttled => "Admin login throttled",
Self::InternalError => "Admin login failed internally",
}
}
fn status_code(&self) -> i64 {
match self {
Self::Success => i64::from(StatusCode::OK.as_u16()),
Self::InvalidCredentials => i64::from(StatusCode::UNAUTHORIZED.as_u16()),
Self::Throttled => i64::from(StatusCode::TOO_MANY_REQUESTS.as_u16()),
Self::InternalError => i64::from(StatusCode::INTERNAL_SERVER_ERROR.as_u16()),
}
}
fn failure_class(&self) -> Option<&'static str> {
match self {
Self::Success => None,
Self::InvalidCredentials => Some("invalid_credentials"),
Self::Throttled => Some("login_throttled"),
Self::InternalError => Some("session_error"),
}
}
}
async fn record_admin_login_audit(
state: &AppState,
request_id: Option<&str>,
source: &str,
username_provided: bool,
outcome: AdminLoginAuditOutcome,
) {
let mut record = UsageRecord::new("admin_auth", outcome.level(), outcome.message());
record.request_id = request_id.map(ToString::to_string);
record.route = Some("/api/admin/login".to_string());
record.status_code = Some(outcome.status_code());
record.failure_class = outcome.failure_class().map(ToString::to_string);
record.metadata = json!({
"source": source,
"usernameProvided": username_provided,
});
if let Err(error) = state.services.usage_records.record_audit(record).await {
tracing::warn!(%error, "failed to record admin login audit event");
}
}
/// `GET /api/admin/auth/status`
pub(crate) async fn session_status(
State(state): State<AppState>,
@@ -106,7 +250,7 @@ pub(crate) async fn logout(
)
.into_response();
let cookie = format!("{ADMIN_SESSION_COOKIE}=; Path=/; HttpOnly; SameSite=Lax; Max-Age=0");
let cookie = format!("{ADMIN_SESSION_COOKIE}=; {ADMIN_SESSION_COOKIE_ATTRS}; Max-Age=0");
response.headers_mut().insert(
SET_COOKIE,
HeaderValue::from_str(&cookie)
@@ -161,3 +305,13 @@ fn admin_api_key_header(headers: &HeaderMap) -> Option<String> {
let value = headers.get("x-api-key")?.to_str().ok()?.trim();
(!value.is_empty()).then(|| value.to_string())
}
fn admin_login_source(headers: &HeaderMap) -> String {
headers
.get("x-real-ip")
.and_then(|value| value.to_str().ok())
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("unknown")
.to_string()
}
backend/src/admin/keys/routes.rs
+16 -44
View File
@@ -2,20 +2,20 @@
use axum::{
extract::{Query, State},
http::{HeaderMap, StatusCode},
http::StatusCode,
response::IntoResponse,
Json,
};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use validator::{Validate, ValidationError};
use crate::{
admin::auth::session::require_admin_auth,
admin::auth::session::AdminAuth,
admin::keys::service::{AdminClientApiKey, AdminClientKeyError},
admin::response::{
AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse, BatchDeleteData,
},
admin::update_payload::{parse_editable_update, EditableUpdateMessages},
infra::{
json::{clamp_limit, Page},
time::{china_relative_time_str, china_rfc3339_str},
@@ -88,16 +88,6 @@ fn client_key_not_found() -> AdminError {
AdminError::not_found("Client API key not found")
}
#[derive(Debug, Deserialize, Validate)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct ClientApiKeyUpdateRequest {
#[validate(custom(function = "validate_non_empty_trimmed"))]
id: String,
label: Option<Option<String>>,
#[validate(custom(function = "validate_non_empty_trimmed"))]
status: Option<String>,
}
struct ParsedClientApiKeyUpdate {
id: String,
label: Option<Option<String>>,
@@ -105,21 +95,15 @@ struct ParsedClientApiKeyUpdate {
}
fn parse_client_api_key_update(payload: &Value) -> Result<ParsedClientApiKeyUpdate, AdminError> {
if !payload.is_object() {
return Err(AdminError::bad_request(
"Client API key update request must be an object",
));
}
let update = serde_json::from_value::<ClientApiKeyUpdateRequest>(payload.clone())
.map_err(|_| AdminError::bad_request("Invalid client API key update request"))?;
update
.validate()
.map_err(|_| AdminError::bad_request("Invalid client API key update request"))?;
if update.label.is_none() && update.status.is_none() {
return Err(AdminError::bad_request(
"Client API key update request must include label or status",
));
}
let update = parse_editable_update(
payload,
EditableUpdateMessages {
object_required: "Client API key update request must be an object",
invalid: "Invalid client API key update request",
empty_update: "Client API key update request must include label or status",
unknown_field_editable: false,
},
)?;
Ok(ParsedClientApiKeyUpdate {
id: update.id,
label: update.label,
@@ -127,20 +111,12 @@ fn parse_client_api_key_update(payload: &Value) -> Result<ParsedClientApiKeyUpda
})
}
fn validate_non_empty_trimmed(value: &str) -> Result<(), ValidationError> {
if value.trim().is_empty() {
return Err(ValidationError::new("required"));
}
Ok(())
}
/// `GET /api/admin/keys`
pub(crate) async fn api_keys(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
Query(query): Query<ApiKeysQuery>,
) -> Result<impl IntoResponse, AdminError> {
require_admin_auth(&state, &headers).await?;
let limit = clamp_limit(query.limit.unwrap_or(50));
match state
.services
@@ -165,10 +141,9 @@ pub(crate) async fn api_keys(
/// `POST /api/admin/keys`
pub(crate) async fn create_api_key(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
Json(payload): Json<CreateApiKeyRequest>,
) -> Result<impl IntoResponse, AdminError> {
require_admin_auth(&state, &headers).await?;
match state.services.admin_client_keys.create(&payload.name).await {
Ok(created) => Ok(AdminResponse::new(
StatusCode::OK,
@@ -181,11 +156,9 @@ pub(crate) async fn create_api_key(
/// `POST /api/admin/keys/update`
pub(crate) async fn update_api_key(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
Json(payload): Json<Value>,
) -> Result<impl IntoResponse, AdminError> {
require_admin_auth(&state, &headers).await?;
let update = parse_client_api_key_update(&payload)?;
if let Some(label) = update.label {
match state
@@ -225,10 +198,9 @@ pub(crate) async fn update_api_key(
/// `POST /api/admin/keys/delete`
pub(crate) async fn batch_delete_api_keys(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
Json(payload): Json<BatchDeleteClientApiKeysRequest>,
) -> Result<impl IntoResponse, AdminError> {
require_admin_auth(&state, &headers).await?;
match state
.services
.admin_client_keys
backend/src/admin/mod.rs
+1
View File
@@ -16,3 +16,4 @@ pub mod router;
pub mod settings;
/// 系统版本与在线更新处理器。
pub mod system;
pub(crate) mod update_payload;
backend/src/admin/monitoring/service.rs → backend/src/admin/monitoring/account_usage_service.rs
+30 -4
View File
@@ -2,10 +2,11 @@
use chrono::{DateTime, Utc};
use thiserror::Error;
use crate::admin::monitoring::usage_store::{
SqliteUsageStore, UsageListRecord, UsageSummary, UsageTimeBucketRecord,
use crate::admin::monitoring::{
account_usage_store::{SqliteUsageStore, UsageListRecord, UsageSummary, UsageTimeBucketRecord},
billing,
};
use crate::infra::json::Page;
use crate::infra::{format::nonnegative_i64_to_u64, json::Page};
/// 管理端用量统计服务。
#[derive(Clone)]
@@ -121,7 +122,7 @@ pub struct AdminUsageSummary {
}
/// 管理端时间桶用量记录。
#[derive(Debug, Clone, PartialEq, Eq)]
#[derive(Debug, Clone, PartialEq)]
pub struct AdminUsageTimeBucketRecord {
pub bucket_start: DateTime<Utc>,
pub model: String,
@@ -137,6 +138,7 @@ pub struct AdminUsageTimeBucketRecord {
pub latency_count: i64,
pub max_latency_ms: i64,
pub min_latency_ms: i64,
pub cost_usd: Option<f64>,
}
impl From<UsageSummary> for AdminUsageSummary {
@@ -160,6 +162,7 @@ impl From<UsageSummary> for AdminUsageSummary {
impl From<UsageTimeBucketRecord> for AdminUsageTimeBucketRecord {
fn from(record: UsageTimeBucketRecord) -> Self {
let cost_usd = time_bucket_cost_usd(&record);
Self {
bucket_start: record.bucket_start,
model: record.model,
@@ -175,10 +178,33 @@ impl From<UsageTimeBucketRecord> for AdminUsageTimeBucketRecord {
latency_count: record.latency_count,
max_latency_ms: record.max_latency_ms,
min_latency_ms: record.min_latency_ms,
cost_usd,
}
}
}
fn time_bucket_cost_usd(record: &UsageTimeBucketRecord) -> Option<f64> {
let input_tokens = nonnegative_i64_to_u64(record.input_tokens);
let output_tokens = nonnegative_i64_to_u64(record.output_tokens);
let cached_tokens = nonnegative_i64_to_u64(record.cached_tokens);
if input_tokens == 0 && output_tokens == 0 && cached_tokens == 0 {
return None;
}
let model = record.model.trim();
if model.is_empty() {
return None;
}
Some(billing::calculate_cost(
input_tokens,
output_tokens,
cached_tokens,
model,
record.service_tier.as_deref(),
))
}
impl From<UsageListRecord> for AdminUsageRecord {
fn from(usage: UsageListRecord) -> Self {
Self {
backend/src/admin/monitoring/usage_store.rs → backend/src/admin/monitoring/account_usage_store.rs
+24 -60
View File
@@ -1,12 +1,17 @@
//! SQLite 用量存储。
use chrono::{DateTime, Utc};
use sqlx::{sqlite::SqliteRow, Row, SqlitePool};
use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite, SqlitePool};
use thiserror::Error;
use crate::infra::json::{decode_cursor, encode_cursor, Page};
use crate::infra::{
json::{decode_cursor, encode_cursor, Page},
time::{
parse_optional_rfc3339_utc as parse_optional_rfc3339, parse_rfc3339_utc as parse_rfc3339,
},
};
const LIST_USAGE_AFTER_CURSOR_SQL: &str = r"
const LIST_USAGE_SELECT_SQL: &str = r"
select
au.account_id,
a.email,
@@ -32,41 +37,7 @@ select
au.limit_window_seconds,
au.last_used_at
from account_usage au
left join accounts a on a.id = au.account_id
where au.last_used_at < ?
or (au.last_used_at = ? and au.account_id < ?)
order by au.last_used_at desc, au.account_id desc
limit ?";
const LIST_USAGE_SQL: &str = r"
select
au.account_id,
a.email,
a.label,
a.plan_type,
au.request_count,
au.empty_response_count,
au.input_tokens,
au.output_tokens,
au.cached_tokens,
au.reasoning_tokens,
au.total_tokens,
au.image_input_tokens,
au.image_output_tokens,
au.image_request_count,
au.image_request_failed_count,
au.window_request_count,
au.window_input_tokens,
au.window_output_tokens,
au.window_cached_tokens,
au.window_started_at,
au.window_reset_at,
au.limit_window_seconds,
au.last_used_at
from account_usage au
left join accounts a on a.id = au.account_id
order by au.last_used_at desc, au.account_id desc
limit ?";
left join accounts a on a.id = au.account_id";
const USAGE_SUMMARY_SQL: &str = r"
select
@@ -254,24 +225,24 @@ impl SqliteUsageStore {
limit: u32,
) -> SqliteUsageStoreResult<Page<UsageListRecord>> {
let limit = limit.clamp(1, 200);
let fetch_limit = i64::from(limit) + 1;
let mut builder = QueryBuilder::<Sqlite>::new(LIST_USAGE_SELECT_SQL);
if let Some(cursor) = cursor {
let (last_used_at, account_id) =
decode_cursor(&cursor).ok_or(SqliteUsageStoreError::InvalidCursor)?;
let rows = sqlx::query(LIST_USAGE_AFTER_CURSOR_SQL)
.bind(&last_used_at)
.bind(&last_used_at)
.bind(&account_id)
.bind(limit + 1)
.fetch_all(&self.pool)
.await?;
Ok(to_page(&rows, limit))
} else {
let rows = sqlx::query(LIST_USAGE_SQL)
.bind(limit + 1)
.fetch_all(&self.pool)
.await?;
Ok(to_page(&rows, limit))
builder.push("\nwhere au.last_used_at < ");
builder.push_bind(last_used_at.clone());
builder.push("\n or (au.last_used_at = ");
builder.push_bind(last_used_at);
builder.push(" and au.account_id < ");
builder.push_bind(account_id);
builder.push(")");
}
builder.push("\norder by au.last_used_at desc, au.account_id desc\nlimit ");
builder.push_bind(fetch_limit);
let rows = builder.build().fetch_all(&self.pool).await?;
Ok(to_page(&rows, limit))
}
/// 汇总账号用量。
@@ -347,8 +318,7 @@ fn usage_list_from_row(row: &SqliteRow) -> SqliteUsageStoreResult<UsageListRecor
fn usage_time_bucket_from_row(row: &SqliteRow) -> SqliteUsageStoreResult<UsageTimeBucketRecord> {
let service_tier = row.get::<String, _>("service_tier").trim().to_string();
Ok(UsageTimeBucketRecord {
bucket_start: DateTime::parse_from_rfc3339(row.get::<&str, _>("bucket_start"))?
.with_timezone(&Utc),
bucket_start: parse_rfc3339(row.get::<&str, _>("bucket_start"))?,
model: row.get("model"),
service_tier: (!service_tier.is_empty()).then_some(service_tier),
request_count: row.get("request_count"),
@@ -365,12 +335,6 @@ fn usage_time_bucket_from_row(row: &SqliteRow) -> SqliteUsageStoreResult<UsageTi
})
}
fn parse_optional_rfc3339(value: Option<&str>) -> SqliteUsageStoreResult<Option<DateTime<Utc>>> {
value
.map(|value| Ok(DateTime::parse_from_rfc3339(value)?.with_timezone(&Utc)))
.transpose()
}
fn to_page(rows: &[SqliteRow], limit: u32) -> Page<UsageListRecord> {
let has_more = rows.len() > limit as usize;
let mut items = Vec::with_capacity(limit as usize);
backend/src/admin/monitoring/dashboard.rs
+17 -37
View File
@@ -4,7 +4,7 @@ use std::{cmp::Reverse, collections::HashMap};
use axum::{
extract::{Query, State},
http::{HeaderMap, StatusCode},
http::StatusCode,
response::IntoResponse,
};
use chrono::{DateTime, Duration, Utc};
@@ -13,15 +13,16 @@ use serde_json::Value;
use crate::{
admin::{
auth::session::require_admin_auth,
auth::session::AdminAuth,
monitoring::{
billing,
account_usage_service::{
AdminUsageRecord, AdminUsageSummary, AdminUsageTimeBucketRecord,
},
diagnostics::{
fingerprint_diagnostics, AccountCapacityDiagnostics, AccountPoolDiagnostics,
},
service::{AdminUsageRecord, AdminUsageSummary, AdminUsageTimeBucketRecord},
usage_record_store::AdminUsageRecordFilter,
usage_records::{account_email_map, usage_record_items, UsageRecordData},
usage_record_routes::{usage_record_items, UsageRecordData},
usage_record_service::AdminUsageRecordFilter,
},
response::{AdminEnvelope, AdminError, AdminResponse},
},
@@ -250,10 +251,8 @@ impl UsageWindow {
/// `GET /api/admin/dashboard/summary`
pub(crate) async fn dashboard_summary(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
) -> Result<impl IntoResponse, AdminError> {
require_admin_auth(&state, &headers).await?;
let accounts = state
.services
.accounts
@@ -293,7 +292,12 @@ pub(crate) async fn dashboard_summary(
let quota_used_by_account = account_quota_used_percent_by_id(&state, &usage_records).await;
let settings = state.services.settings.current();
let recent_usage_records = recent_events.into_iter().take(10).collect::<Vec<_>>();
let account_emails = account_email_map(&state, &recent_usage_records).await?;
let account_emails = state
.services
.usage_records
.account_email_map(&recent_usage_records)
.await
.map_err(|error| AdminError::usage_record_accounts_failed(error.to_string()))?;
let dashboard_usage_records = usage_record_items(recent_usage_records, &account_emails);
Ok(AdminResponse::new(
@@ -315,10 +319,9 @@ pub(crate) async fn dashboard_summary(
/// `GET /api/admin/dashboard/trend?kind=usage|latency|errors`
pub(crate) async fn dashboard_trend(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
Query(query): Query<DashboardTrendQuery>,
) -> Result<impl IntoResponse, AdminError> {
require_admin_auth(&state, &headers).await?;
let time_buckets = dashboard_time_buckets(&state, Utc::now()).await;
Ok(AdminResponse::new(
@@ -736,7 +739,7 @@ fn cost_window(
let mut has_usage = false;
for record in records {
if record.bucket_start >= start && record.bucket_start < end {
if let Some(cost) = bucket_cost(record) {
if let Some(cost) = record.cost_usd {
total += cost;
has_usage = true;
}
@@ -749,7 +752,7 @@ fn total_cost(records: &[AdminUsageTimeBucketRecord]) -> f64 {
let mut total = 0.0;
let mut has_usage = false;
for record in records {
if let Some(cost) = bucket_cost(record) {
if let Some(cost) = record.cost_usd {
total += cost;
has_usage = true;
}
@@ -761,29 +764,6 @@ fn total_cost(records: &[AdminUsageTimeBucketRecord]) -> f64 {
0.0
}
fn bucket_cost(record: &AdminUsageTimeBucketRecord) -> Option<f64> {
let input_tokens = nonnegative_i64_to_u64(record.input_tokens);
let output_tokens = nonnegative_i64_to_u64(record.output_tokens);
let cached_tokens = nonnegative_i64_to_u64(record.cached_tokens);
if input_tokens == 0 && output_tokens == 0 && cached_tokens == 0 {
return None;
}
let model = record.model.trim();
if model.is_empty() {
return None;
}
let service_tier = record.service_tier.as_deref();
Some(billing::calculate_cost(
input_tokens,
output_tokens,
cached_tokens,
model,
service_tier,
))
}
fn apply_bucket(window: &mut UsageWindow, record: &AdminUsageTimeBucketRecord) {
window.requests += nonnegative_i64_to_u64(record.request_count);
window.input_tokens += nonnegative_i64_to_u64(record.input_tokens);
backend/src/admin/monitoring/mod.rs
+5 -4
View File
@@ -1,10 +1,11 @@
//! 管理面监控模块：日志、用量统计和 Dashboard 聚合。
pub mod account_usage_service;
pub mod account_usage_store;
pub(crate) mod billing;
pub mod dashboard;
pub mod diagnostics;
pub mod service;
pub mod usage_record;
pub mod usage_record_model;
pub mod usage_record_routes;
pub mod usage_record_service;
pub mod usage_record_store;
pub mod usage_records;
pub mod usage_store;
backend/src/admin/monitoring/usage_record.rs → backend/src/admin/monitoring/usage_record_model.rs
+21 -1
View File
@@ -100,9 +100,29 @@ pub(in crate::admin) fn metadata_service_tier(metadata: &Value) -> Option<&str>
.filter(|value| !value.is_empty())
}
pub(in crate::admin) fn metadata_string(metadata: &Value, keys: &[&str]) -> Option<String> {
keys.iter().find_map(|key| {
metadata
.get(*key)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToString::to_string)
})
}
pub(in crate::admin) fn metadata_i64(metadata: &Value, keys: &[&str]) -> Option<i64> {
keys.iter()
.find_map(|key| metadata.get(*key).and_then(Value::as_i64))
}
pub(in crate::admin) fn metadata_nonnegative_i64(metadata: &Value, keys: &[&str]) -> Option<i64> {
metadata_i64(metadata, keys).filter(|value| *value >= 0)
}
/// 响应事件记录（供 dispatch 模块使用）。
pub struct ResponseUsageRecord<'a> {
pub usage_records: &'a crate::admin::monitoring::usage_record_store::AdminUsageRecordService,
pub usage_records: &'a crate::admin::monitoring::usage_record_service::AdminUsageRecordService,
pub request_id: &'a str,
pub account_id: &'a str,
pub route: &'a str,
backend/src/admin/monitoring/usage_records.rs → backend/src/admin/monitoring/usage_record_routes.rs
+66 -138
View File
@@ -4,25 +4,27 @@ use std::collections::HashMap;
use axum::{
extract::{Query, State},
http::{HeaderMap, StatusCode},
http::StatusCode,
response::IntoResponse,
};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use sqlx::{QueryBuilder, Row, Sqlite};
use crate::{
admin::monitoring::{
billing,
usage_record::{metadata_service_tier, UsageRecord, UsageRecordLevel},
usage_record_model::{
metadata_nonnegative_i64, metadata_string, UsageRecord, UsageRecordLevel,
},
usage_record_service::{
usage_record_cost_details, AdminUsageRecordError, AdminUsageRecordFilter,
},
usage_record_store::{
AdminUsageRecordError, AdminUsageRecordFilter, UsageRecordBreakdown,
UsageRecordEndpointSource, UsageRecordModelSource, UsageRecordSummary,
UsageRecordTrendPoint,
UsageRecordBreakdown, UsageRecordEndpointSource, UsageRecordModelSource,
UsageRecordSummary, UsageRecordTrendPoint,
},
},
admin::{
auth::session::require_admin_auth,
auth::session::AdminAuth,
response::{AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse},
},
infra::{
@@ -200,10 +202,9 @@ struct UsageRecordTrendPointData {
/// `GET /api/admin/usage/records`
pub(crate) async fn usage_records(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
Query(query): Query<UsageRecordsQuery>,
) -> Result<impl IntoResponse, AdminError> {
require_admin_auth(&state, &headers).await?;
let limit = clamp_limit(query.page_size.or(query.limit).unwrap_or(50));
let page = query.page;
let use_numbered_page = page.is_some() || query.page_size.is_some();
@@ -218,7 +219,12 @@ pub(crate) async fn usage_records(
.await
{
Ok(page) => {
let account_emails = account_email_map(&state, &page.items).await?;
let account_emails = state
.services
.usage_records
.account_email_map(&page.items)
.await
.map_err(|error| log_error(&error))?;
let page = NumberedPage {
items: usage_record_items(page.items, &account_emails),
total: page.total,
@@ -241,7 +247,12 @@ pub(crate) async fn usage_records(
.await
{
Ok(page) => {
let account_emails = account_email_map(&state, &page.items).await?;
let account_emails = state
.services
.usage_records
.account_email_map(&page.items)
.await
.map_err(|error| log_error(&error))?;
let page = Page {
items: usage_record_items(page.items, &account_emails),
next_cursor: page.next_cursor,
@@ -258,10 +269,9 @@ pub(crate) async fn usage_records(
/// `GET /api/admin/usage/records/summary`
pub(crate) async fn usage_records_summary(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
Query(query): Query<UsageRecordsQuery>,
) -> Result<impl IntoResponse, AdminError> {
require_admin_auth(&state, &headers).await?;
let filter = filter_from_query(query)?;
match state.services.usage_records.summary(filter).await {
Ok(summary) => Ok(AdminResponse::new(
@@ -275,10 +285,9 @@ pub(crate) async fn usage_records_summary(
/// `GET /api/admin/usage/records/insights/models`
pub(crate) async fn usage_records_model_distribution(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
Query(query): Query<UsageRecordDistributionQuery>,
) -> Result<impl IntoResponse, AdminError> {
require_admin_auth(&state, &headers).await?;
let source = model_source_from_query(query.source)?;
let filter = filter_from_query(query.records)?;
match state
@@ -298,10 +307,9 @@ pub(crate) async fn usage_records_model_distribution(
/// `GET /api/admin/usage/records/insights/endpoints`
pub(crate) async fn usage_records_endpoint_distribution(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
Query(query): Query<UsageRecordDistributionQuery>,
) -> Result<impl IntoResponse, AdminError> {
require_admin_auth(&state, &headers).await?;
let source = endpoint_source_from_query(query.source)?;
let filter = filter_from_query(query.records)?;
match state
@@ -321,10 +329,9 @@ pub(crate) async fn usage_records_endpoint_distribution(
/// `GET /api/admin/usage/records/insights/token-trend`
pub(crate) async fn usage_records_token_trend(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
Query(query): Query<UsageRecordsQuery>,
) -> Result<impl IntoResponse, AdminError> {
require_admin_auth(&state, &headers).await?;
let filter = filter_from_query(query)?;
match state.services.usage_records.token_trend(filter).await {
Ok(points) => Ok(AdminResponse::new(
@@ -343,10 +350,9 @@ pub(crate) async fn usage_records_token_trend(
/// `GET /api/admin/usage/records/insights/latency-trend`
pub(crate) async fn usage_records_latency_trend(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
Query(query): Query<UsageRecordsQuery>,
) -> Result<impl IntoResponse, AdminError> {
require_admin_auth(&state, &headers).await?;
let filter = filter_from_query(query)?;
match state.services.usage_records.latency_trend(filter).await {
Ok(points) => Ok(AdminResponse::new(
@@ -365,13 +371,17 @@ pub(crate) async fn usage_records_latency_trend(
/// `GET /api/admin/usage/records/detail`
pub(crate) async fn usage_record_detail(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
Query(query): Query<UsageRecordDetailQuery>,
) -> Result<impl IntoResponse, AdminError> {
require_admin_auth(&state, &headers).await?;
match state.services.usage_records.get(&query.id).await {
Ok(Some(log)) => {
let account_emails = account_email_map(&state, std::slice::from_ref(&log)).await?;
let account_emails = state
.services
.usage_records
.account_email_map(std::slice::from_ref(&log))
.await
.map_err(|error| log_error(&error))?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(usage_record_data(log, &account_emails)),
@@ -385,9 +395,8 @@ pub(crate) async fn usage_record_detail(
/// `POST /api/admin/usage/records/delete`
pub(crate) async fn clear_usage_records(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
) -> Result<impl IntoResponse, AdminError> {
require_admin_auth(&state, &headers).await?;
match state.services.usage_records.clear().await {
Ok(cleared) => Ok(AdminResponse::new(
StatusCode::OK,
@@ -462,6 +471,9 @@ fn log_error(error: &AdminUsageRecordError) -> AdminError {
| AdminUsageRecordError::Clear
| AdminUsageRecordError::Append
| AdminUsageRecordError::Trim => AdminError::internal(error.to_string()),
AdminUsageRecordError::Accounts => {
AdminError::usage_record_accounts_failed(error.to_string())
}
}
}
@@ -484,50 +496,6 @@ fn non_empty(value: Option<String>) -> Option<String> {
.filter(|value| !value.is_empty())
}
pub(crate) async fn account_email_map(
state: &AppState,
items: &[UsageRecord],
) -> Result<HashMap<String, String>, AdminError> {
let mut account_ids = items
.iter()
.filter_map(|item| item.account_id.as_deref())
.map(str::to_string)
.collect::<Vec<_>>();
account_ids.sort_unstable();
account_ids.dedup();
if account_ids.is_empty() {
return Ok(HashMap::new());
}
let mut builder = QueryBuilder::<Sqlite>::new("select id, email from accounts where id in (");
let mut separated = builder.separated(", ");
for account_id in &account_ids {
separated.push_bind(account_id);
}
separated.push_unseparated(")");
let rows = builder
.build()
.fetch_all(state.services.background_tasks.accounts.pool())
.await
.map_err(|error| {
AdminError::usage_record_accounts_failed(format!(
"Failed to load usage record accounts: {error}"
))
})?;
Ok(rows
.into_iter()
.filter_map(|row| {
row.get::<Option<String>, _>("email")
.map(|email| email.trim().to_string())
.filter(|email| !email.is_empty())
.map(|email| (row.get::<String, _>("id"), email))
})
.collect())
}
pub(crate) fn usage_record_items(
items: Vec<UsageRecord>,
account_emails: &HashMap<String, String>,
@@ -555,7 +523,7 @@ fn usage_record_data(
metadata_string(&record.metadata, &["reasoningEffort", "reasoning_effort"]);
let token_details = usage_token_details(&record.metadata);
let cost_details = usage_cost_details(&record, upstream_model.as_deref(), &token_details);
let first_token_latency_ms = metadata_i64(
let first_token_latency_ms = metadata_nonnegative_i64(
&record.metadata,
&[
"firstTokenMs",
@@ -626,51 +594,37 @@ fn usage_cost_details(
upstream_model: Option<&str>,
tokens: &UsageRecordTokenDetailsData,
) -> Option<UsageRecordCostDetailsData> {
if tokens.input_tokens == 0 && tokens.output_tokens == 0 && tokens.cached_tokens == 0 {
return None;
}
let model = upstream_model
.or(record.model.as_deref())
.map(str::trim)
.filter(|value| !value.is_empty())?;
let service_tier = metadata_service_tier(&record.metadata);
let breakdown = billing::calculate_cost_breakdown(
let cost = usage_record_cost_details(
record,
upstream_model,
tokens.input_tokens,
tokens.output_tokens,
tokens.cached_tokens,
model,
service_tier,
);
let original_cost = breakdown.input_cost + breakdown.output_cost + breakdown.cache_read_cost;
)?;
Some(UsageRecordCostDetailsData {
input_cost: breakdown.input_cost,
output_cost: breakdown.output_cost,
cache_read_cost: breakdown.cache_read_cost,
total_cost: breakdown.total_cost,
billed_cost: breakdown.total_cost,
original_cost,
input_price_per_mtoken: breakdown.input_price_per_mtoken,
output_price_per_mtoken: breakdown.output_price_per_mtoken,
cache_read_price_per_mtoken: breakdown.cache_read_price_per_mtoken,
service_tier: breakdown.service_tier.clone(),
service_tier_display: breakdown
.service_tier
.as_deref()
.map(format_service_tier)
.unwrap_or_else(|| "Default".to_string()),
multiplier: breakdown.tier_multiplier,
input_cost_display: format_precise_cost(breakdown.input_cost),
output_cost_display: format_precise_cost(breakdown.output_cost),
cache_read_cost_display: format_precise_cost(breakdown.cache_read_cost),
total_cost_display: format_precise_cost(breakdown.total_cost),
billed_cost_display: format_precise_cost(breakdown.total_cost),
original_cost_display: format_precise_cost(original_cost),
input_price_display: format_token_price(breakdown.input_price_per_mtoken),
output_price_display: format_token_price(breakdown.output_price_per_mtoken),
cache_read_price_display: format_token_price(breakdown.cache_read_price_per_mtoken),
multiplier_display: format_multiplier(breakdown.tier_multiplier),
input_cost: cost.input_cost,
output_cost: cost.output_cost,
cache_read_cost: cost.cache_read_cost,
total_cost: cost.total_cost,
billed_cost: cost.billed_cost,
original_cost: cost.original_cost,
input_price_per_mtoken: cost.input_price_per_mtoken,
output_price_per_mtoken: cost.output_price_per_mtoken,
cache_read_price_per_mtoken: cost.cache_read_price_per_mtoken,
service_tier: cost.service_tier,
service_tier_display: cost.service_tier_display,
multiplier: cost.multiplier,
input_cost_display: format_precise_cost(cost.input_cost),
output_cost_display: format_precise_cost(cost.output_cost),
cache_read_cost_display: format_precise_cost(cost.cache_read_cost),
total_cost_display: format_precise_cost(cost.total_cost),
billed_cost_display: format_precise_cost(cost.billed_cost),
original_cost_display: format_precise_cost(cost.original_cost),
input_price_display: format_token_price(cost.input_price_per_mtoken),
output_price_display: format_token_price(cost.output_price_per_mtoken),
cache_read_price_display: format_token_price(cost.cache_read_price_per_mtoken),
multiplier_display: format_multiplier(cost.multiplier),
})
}
@@ -688,32 +642,6 @@ fn metadata_u64_at(value: &Value, key: &str) -> Option<u64> {
.and_then(|value| value.as_u64().or_else(|| value.as_i64()?.try_into().ok()))
}
fn metadata_string(metadata: &Value, keys: &[&str]) -> Option<String> {
keys.iter().find_map(|key| {
metadata
.get(*key)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToString::to_string)
})
}
fn metadata_i64(metadata: &Value, keys: &[&str]) -> Option<i64> {
keys.iter()
.find_map(|key| metadata.get(*key).and_then(Value::as_i64))
.filter(|value| *value >= 0)
}
fn format_service_tier(value: &str) -> String {
match value {
"priority" | "fast" => "Fast".to_string(),
"flex" => "Flex".to_string(),
"default" => "Default".to_string(),
other => other.to_string(),
}
}
impl From<UsageRecordSummary> for UsageRecordSummaryData {
fn from(summary: UsageRecordSummary) -> Self {
Self {
backend/src/admin/monitoring/usage_record_service.rs
+374
View File
@@ -0,0 +1,374 @@
//! 管理端使用记录服务。
use std::{collections::HashMap, sync::Arc};
use chrono::{DateTime, Utc};
use thiserror::Error;
use crate::{
admin::monitoring::{
billing,
usage_record_model::metadata_service_tier,
usage_record_model::{UsageRecord, UsageRecordLevel},
usage_record_store::{
SqliteUsageRecordStore, UsageRecordBreakdown, UsageRecordEndpointSource,
UsageRecordFilter, UsageRecordModelSource, UsageRecordSummary, UsageRecordTrendPoint,
},
},
infra::json::{NumberedPage, Page},
};
/// 管理端使用记录查询过滤器。
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct AdminUsageRecordFilter {
/// 事件类别。
pub kind: Option<String>,
/// 事件等级。
pub level: Option<UsageRecordLevel>,
/// 请求 ID。
pub request_id: Option<String>,
/// 账号 ID。
pub account_id: Option<String>,
/// 路由。
pub route: Option<String>,
/// 模型。
pub model: Option<String>,
/// HTTP 状态码。
pub status_code: Option<i64>,
/// 上游传输方式。
pub transport: Option<String>,
/// 同一请求内的上游尝试序号。
pub attempt_index: Option<i64>,
/// 上游 HTTP 状态码。
pub upstream_status_code: Option<i64>,
/// 失败分类。
pub failure_class: Option<String>,
/// 上游响应 ID。
pub response_id: Option<String>,
/// 上游请求 ID。
pub upstream_request_id: Option<String>,
/// 搜索关键词。
pub search: Option<String>,
/// 起始时间。
pub start_time: Option<DateTime<Utc>>,
/// 结束时间。
pub end_time: Option<DateTime<Utc>>,
}
impl From<AdminUsageRecordFilter> for UsageRecordFilter {
fn from(filter: AdminUsageRecordFilter) -> Self {
Self {
kind: filter.kind,
level: filter.level,
request_id: filter.request_id,
account_id: filter.account_id,
route: filter.route,
model: filter.model,
status_code: filter.status_code,
transport: filter.transport,
attempt_index: filter.attempt_index,
upstream_status_code: filter.upstream_status_code,
failure_class: filter.failure_class,
response_id: filter.response_id,
upstream_request_id: filter.upstream_request_id,
search: filter.search,
start_time: filter.start_time,
end_time: filter.end_time,
}
}
}
/// 清空使用记录结果。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AdminClearUsageRecords {
/// 清理数量。
pub cleared: u64,
}
/// 管理端使用记录费用明细。
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct AdminUsageRecordCostDetails {
pub input_cost: f64,
pub output_cost: f64,
pub cache_read_cost: f64,
pub total_cost: f64,
pub billed_cost: f64,
pub original_cost: f64,
pub input_price_per_mtoken: f64,
pub output_price_per_mtoken: f64,
pub cache_read_price_per_mtoken: f64,
pub service_tier: Option<String>,
pub service_tier_display: String,
pub multiplier: f64,
}
/// 管理端使用记录错误。
#[derive(Debug, Error)]
pub enum AdminUsageRecordError {
/// 列表失败。
#[error("failed to list usage records")]
List,
/// 读取失败。
#[error("failed to get usage record")]
Get,
/// 清空失败。
#[error("failed to clear usage records")]
Clear,
/// 写入失败。
#[error("failed to append usage record")]
Append,
/// 裁剪失败。
#[error("failed to trim usage records")]
Trim,
/// 账号关联信息读取失败。
#[error("failed to load usage record accounts")]
Accounts,
}
/// 管理端使用记录服务。
#[derive(Clone)]
pub struct AdminUsageRecordService {
store: SqliteUsageRecordStore,
settings: Arc<tokio::sync::RwLock<AdminUsageRecordSettings>>,
}
#[derive(Debug, Clone, Copy)]
struct AdminUsageRecordSettings {
enabled: bool,
capacity: u32,
capture_body: bool,
}
impl AdminUsageRecordService {
/// 构造管理端使用记录服务。
pub fn new(
store: SqliteUsageRecordStore,
enabled: bool,
capacity: u32,
capture_body: bool,
) -> Self {
Self {
store,
settings: Arc::new(tokio::sync::RwLock::new(AdminUsageRecordSettings {
enabled,
capacity,
capture_body,
})),
}
}
/// 分页查询日志。
pub async fn list(
&self,
cursor: Option<String>,
limit: u32,
filter: AdminUsageRecordFilter,
) -> Result<Page<UsageRecord>, AdminUsageRecordError> {
self.store
.list(filter.into(), cursor, limit)
.await
.map_err(|_| AdminUsageRecordError::List)
}
/// 按页码查询日志。
pub async fn list_page(
&self,
page: u32,
page_size: u32,
filter: AdminUsageRecordFilter,
) -> Result<NumberedPage<UsageRecord>, AdminUsageRecordError> {
self.store
.list_page(filter.into(), page, page_size)
.await
.map_err(|_| AdminUsageRecordError::List)
}
/// 按 ID 读取日志。
pub async fn get(&self, id: &str) -> Result<Option<UsageRecord>, AdminUsageRecordError> {
self.store
.get(id)
.await
.map_err(|_| AdminUsageRecordError::Get)
}
/// 读取使用记录关联账号的邮箱映射。
pub async fn account_email_map(
&self,
items: &[UsageRecord],
) -> Result<HashMap<String, String>, AdminUsageRecordError> {
self.store
.account_email_map(items)
.await
.map_err(|_| AdminUsageRecordError::Accounts)
}
/// 汇总使用记录。
pub async fn summary(
&self,
filter: AdminUsageRecordFilter,
) -> Result<UsageRecordSummary, AdminUsageRecordError> {
self.store
.summary(filter.into())
.await
.map_err(|_| AdminUsageRecordError::List)
}
/// 按模型来源聚合使用记录分布。
pub async fn model_distribution(
&self,
filter: AdminUsageRecordFilter,
source: UsageRecordModelSource,
) -> Result<Vec<UsageRecordBreakdown>, AdminUsageRecordError> {
self.store
.model_distribution(filter.into(), source)
.await
.map_err(|_| AdminUsageRecordError::List)
}
/// 按端点来源聚合使用记录分布。
pub async fn endpoint_distribution(
&self,
filter: AdminUsageRecordFilter,
source: UsageRecordEndpointSource,
) -> Result<Vec<UsageRecordBreakdown>, AdminUsageRecordError> {
self.store
.endpoint_distribution(filter.into(), source)
.await
.map_err(|_| AdminUsageRecordError::List)
}
/// 聚合 Token 趋势。
pub async fn token_trend(
&self,
filter: AdminUsageRecordFilter,
) -> Result<Vec<UsageRecordTrendPoint>, AdminUsageRecordError> {
self.store
.trend(filter.into())
.await
.map_err(|_| AdminUsageRecordError::List)
}
/// 聚合延迟趋势。
pub async fn latency_trend(
&self,
filter: AdminUsageRecordFilter,
) -> Result<Vec<UsageRecordTrendPoint>, AdminUsageRecordError> {
self.store
.trend(filter.into())
.await
.map_err(|_| AdminUsageRecordError::List)
}
/// 清空日志。
pub async fn clear(&self) -> Result<AdminClearUsageRecords, AdminUsageRecordError> {
self.store
.clear()
.await
.map(|cleared| AdminClearUsageRecords { cleared })
.map_err(|_| AdminUsageRecordError::Clear)
}
/// 记录使用记录。
pub async fn record(&self, event: UsageRecord) -> Result<(), AdminUsageRecordError> {
let settings = *self.settings.read().await;
if !settings.enabled && event.level != UsageRecordLevel::Error {
return Ok(());
}
self.append_with_settings(event, settings).await
}
/// 记录安全审计事件，不受普通使用记录开关影响。
pub async fn record_audit(&self, event: UsageRecord) -> Result<(), AdminUsageRecordError> {
let settings = *self.settings.read().await;
self.append_with_settings(event, settings).await
}
async fn append_with_settings(
&self,
mut event: UsageRecord,
settings: AdminUsageRecordSettings,
) -> Result<(), AdminUsageRecordError> {
apply_capture_body_policy(&mut event, settings.capture_body);
self.store
.append(&event)
.await
.map_err(|_| AdminUsageRecordError::Append)?;
self.store
.trim_to_capacity(settings.capacity)
.await
.map_err(|_| AdminUsageRecordError::Trim)?;
Ok(())
}
}
fn apply_capture_body_policy(event: &mut UsageRecord, capture_body: bool) {
if capture_body {
return;
}
let Some(metadata) = event.metadata.as_object_mut() else {
return;
};
for key in [
"body",
"rawBody",
"requestBody",
"responseBody",
"upstreamBody",
] {
metadata.remove(key);
}
}
pub(crate) fn usage_record_cost_details(
record: &UsageRecord,
upstream_model: Option<&str>,
input_tokens: u64,
output_tokens: u64,
cached_tokens: u64,
) -> Option<AdminUsageRecordCostDetails> {
if input_tokens == 0 && output_tokens == 0 && cached_tokens == 0 {
return None;
}
let model = upstream_model
.or(record.model.as_deref())
.map(str::trim)
.filter(|value| !value.is_empty())?;
let service_tier = metadata_service_tier(&record.metadata);
let breakdown = billing::calculate_cost_breakdown(
input_tokens,
output_tokens,
cached_tokens,
model,
service_tier,
);
let original_cost = breakdown.input_cost + breakdown.output_cost + breakdown.cache_read_cost;
Some(AdminUsageRecordCostDetails {
input_cost: breakdown.input_cost,
output_cost: breakdown.output_cost,
cache_read_cost: breakdown.cache_read_cost,
total_cost: breakdown.total_cost,
billed_cost: breakdown.total_cost,
original_cost,
input_price_per_mtoken: breakdown.input_price_per_mtoken,
output_price_per_mtoken: breakdown.output_price_per_mtoken,
cache_read_price_per_mtoken: breakdown.cache_read_price_per_mtoken,
service_tier: breakdown.service_tier.clone(),
service_tier_display: breakdown
.service_tier
.as_deref()
.map(format_service_tier)
.unwrap_or_else(|| "Default".to_string()),
multiplier: breakdown.tier_multiplier,
})
}
fn format_service_tier(value: &str) -> String {
match value {
"priority" | "fast" => "Fast".to_string(),
"flex" => "Flex".to_string(),
"default" => "Default".to_string(),
other => other.to_string(),
}
}
backend/src/admin/monitoring/usage_record_store.rs
+83 -317
View File
@@ -1,6 +1,6 @@
//! 使用记录存储实现（SQLite）。
use std::collections::BTreeMap;
use std::collections::{BTreeMap, HashMap};
use chrono::{DateTime, Utc};
use serde_json::Value;
@@ -9,10 +9,15 @@ use thiserror::Error;
use crate::admin::monitoring::{
billing,
usage_record::{metadata_service_tier, UsageRecord, UsageRecordLevel},
usage_record_model::{
metadata_i64, metadata_service_tier, metadata_string, UsageRecord, UsageRecordLevel,
},
};
use crate::infra::{
format::optional_nonnegative_i64_to_u64,
json::{decode_cursor, encode_cursor, page_offset, NumberedPage, Page},
time::{china_quarter_hour_start, parse_rfc3339_utc as parse_rfc3339},
};
use crate::infra::json::{decode_cursor, encode_cursor, page_offset, NumberedPage, Page};
use crate::infra::time::china_quarter_hour_start;
/// 默认保留最新使用记录数量。
pub const DEFAULT_USAGE_RECORD_CAPACITY: u32 = 2_000;
@@ -44,6 +49,7 @@ pub type SqliteUsageRecordStoreResult<T> = Result<T, SqliteUsageRecordStoreError
const USAGE_RECORD_SERVICE_TIER_SQL: &str =
"nullif(trim(json_extract(metadata_json, '$.serviceTier')), '')";
const USAGE_RECORD_SELECT_SQL: &str = "select id, request_id, kind, level, account_id, route, model, status_code, transport, attempt_index, upstream_status_code, failure_class, response_id, upstream_request_id, latency_ms, message, metadata_json, created_at from usage_records";
/// 使用记录查询过滤器。
#[derive(Debug, Clone, Default, PartialEq, Eq)]
@@ -112,9 +118,7 @@ impl SqliteUsageRecordStore {
limit: u32,
) -> SqliteUsageRecordStoreResult<Page<UsageRecord>> {
let fetch_limit = i64::from(limit) + 1;
let mut builder = QueryBuilder::<Sqlite>::new(
"select id, request_id, kind, level, account_id, route, model, status_code, transport, attempt_index, upstream_status_code, failure_class, response_id, upstream_request_id, latency_ms, message, metadata_json, created_at from usage_records",
);
let mut builder = QueryBuilder::<Sqlite>::new(USAGE_RECORD_SELECT_SQL);
push_filter(&mut builder, &filter, cursor.as_deref())?;
builder.push(" order by created_at desc, id desc limit ");
builder.push_bind(fetch_limit);
@@ -148,9 +152,7 @@ impl SqliteUsageRecordStore {
let page_size = page_size.clamp(1, 200);
let total = count_usage_records(&self.pool, &filter).await?;
let offset = page_offset(page, page_size);
let mut builder = QueryBuilder::<Sqlite>::new(
"select id, request_id, kind, level, account_id, route, model, status_code, transport, attempt_index, upstream_status_code, failure_class, response_id, upstream_request_id, latency_ms, message, metadata_json, created_at from usage_records",
);
let mut builder = QueryBuilder::<Sqlite>::new(USAGE_RECORD_SELECT_SQL);
push_filter(&mut builder, &filter, None)?;
builder.push(" order by created_at desc, id desc limit ");
builder.push_bind(i64::from(page_size));
@@ -173,15 +175,21 @@ impl SqliteUsageRecordStore {
/// 按 ID 读取使用记录。
pub async fn get(&self, id: &str) -> SqliteUsageRecordStoreResult<Option<UsageRecord>> {
let row = sqlx::query(
"select id, request_id, kind, level, account_id, route, model, status_code, transport, attempt_index, upstream_status_code, failure_class, response_id, upstream_request_id, latency_ms, message, metadata_json, created_at from usage_records where id = ?",
)
.bind(id)
.fetch_optional(&self.pool)
.await?;
let mut builder = QueryBuilder::<Sqlite>::new(USAGE_RECORD_SELECT_SQL);
builder.push(" where id = ");
builder.push_bind(id);
let row = builder.build().fetch_optional(&self.pool).await?;
row.map(|row| usage_record_from_row(&row)).transpose()
}
/// 读取使用记录关联账号的邮箱映射。
pub async fn account_email_map(
&self,
items: &[UsageRecord],
) -> SqliteUsageRecordStoreResult<HashMap<String, String>> {
usage_record_account_email_map(&self.pool, items).await
}
/// 汇总事件使用记录。
pub async fn summary(
&self,
@@ -208,16 +216,8 @@ impl SqliteUsageRecordStore {
usage_endpoint_distribution(&self.pool, &filter, source, 8).await
}
/// 聚合 Token 使用趋势。
pub async fn token_trend(
&self,
filter: UsageRecordFilter,
) -> SqliteUsageRecordStoreResult<Vec<UsageRecordTrendPoint>> {
usage_trend(&self.pool, &filter).await
}
/// 聚合延迟趋势。
pub async fn latency_trend(
/// 聚合使用记录趋势。
pub async fn trend(
&self,
filter: UsageRecordFilter,
) -> SqliteUsageRecordStoreResult<Vec<UsageRecordTrendPoint>> {
@@ -284,21 +284,6 @@ fn usage_record_summary_fields(event: &UsageRecord) -> UsageRecordSummaryFields
}
}
fn metadata_string(metadata: &Value, keys: &[&str]) -> Option<String> {
keys.iter().find_map(|key| {
metadata
.get(*key)
.and_then(Value::as_str)
.filter(|value| !value.trim().is_empty())
.map(ToString::to_string)
})
}
fn metadata_i64(metadata: &Value, keys: &[&str]) -> Option<i64> {
keys.iter()
.find_map(|key| metadata.get(*key).and_then(Value::as_i64))
}
async fn append_event(pool: &SqlitePool, event: &UsageRecord) -> SqliteUsageRecordStoreResult<()> {
let summary = usage_record_summary_fields(event);
sqlx::query(
@@ -473,13 +458,17 @@ async fn usage_summary(
push_filter(&mut builder, filter, None)?;
let row = builder.build().fetch_one(pool).await?;
let input_tokens = nonnegative_i64(row.get::<Option<i64>, _>("input_tokens"));
let output_tokens = nonnegative_i64(row.get::<Option<i64>, _>("output_tokens"));
let cached_tokens = nonnegative_i64(row.get::<Option<i64>, _>("cached_tokens"));
let input_tokens = optional_nonnegative_i64_to_u64(row.get::<Option<i64>, _>("input_tokens"));
let output_tokens = optional_nonnegative_i64_to_u64(row.get::<Option<i64>, _>("output_tokens"));
let cached_tokens = optional_nonnegative_i64_to_u64(row.get::<Option<i64>, _>("cached_tokens"));
Ok(UsageRecordSummary {
total_requests: nonnegative_i64(row.get::<Option<i64>, _>("total_requests")),
error_requests: nonnegative_i64(row.get::<Option<i64>, _>("error_requests")),
total_requests: optional_nonnegative_i64_to_u64(
row.get::<Option<i64>, _>("total_requests"),
),
error_requests: optional_nonnegative_i64_to_u64(
row.get::<Option<i64>, _>("error_requests"),
),
input_tokens,
output_tokens,
cached_tokens,
@@ -488,8 +477,39 @@ async fn usage_summary(
})
}
fn nonnegative_i64(value: Option<i64>) -> u64 {
value.unwrap_or_default().max(0).cast_unsigned()
async fn usage_record_account_email_map(
pool: &SqlitePool,
items: &[UsageRecord],
) -> SqliteUsageRecordStoreResult<HashMap<String, String>> {
let mut account_ids = items
.iter()
.filter_map(|item| item.account_id.as_deref())
.map(str::to_string)
.collect::<Vec<_>>();
account_ids.sort_unstable();
account_ids.dedup();
if account_ids.is_empty() {
return Ok(HashMap::new());
}
let mut builder = QueryBuilder::<Sqlite>::new("select id, email from accounts where id in (");
let mut separated = builder.separated(", ");
for account_id in &account_ids {
separated.push_bind(account_id);
}
separated.push_unseparated(")");
let rows = builder.build().fetch_all(pool).await?;
Ok(rows
.into_iter()
.filter_map(|row| {
row.get::<Option<String>, _>("email")
.map(|email| email.trim().to_string())
.filter(|email| !email.is_empty())
.map(|email| (row.get::<String, _>("id"), email))
})
.collect())
}
async fn usage_model_distribution(
@@ -590,9 +610,12 @@ async fn usage_breakdown(
let mut items = Vec::<UsageRecordBreakdown>::new();
for row in rows {
let name: String = row.get(group_alias);
let input_tokens = nonnegative_i64(row.get::<Option<i64>, _>("input_tokens"));
let output_tokens = nonnegative_i64(row.get::<Option<i64>, _>("output_tokens"));
let cached_tokens = nonnegative_i64(row.get::<Option<i64>, _>("cached_tokens"));
let input_tokens =
optional_nonnegative_i64_to_u64(row.get::<Option<i64>, _>("input_tokens"));
let output_tokens =
optional_nonnegative_i64_to_u64(row.get::<Option<i64>, _>("output_tokens"));
let cached_tokens =
optional_nonnegative_i64_to_u64(row.get::<Option<i64>, _>("cached_tokens"));
let cost = usage_breakdown_cost(
input_tokens,
output_tokens,
@@ -600,7 +623,8 @@ async fn usage_breakdown(
&row.get::<String, _>("billing_model"),
row.get::<Option<String>, _>("service_tier").as_deref(),
);
let request_count = nonnegative_i64(row.get::<Option<i64>, _>("request_count"));
let request_count =
optional_nonnegative_i64_to_u64(row.get::<Option<i64>, _>("request_count"));
let average_latency_ms = row.get::<Option<f64>, _>("average_latency_ms");
if let Some(item) = items.iter_mut().find(|item| item.name == name) {
@@ -678,11 +702,14 @@ async fn usage_trend(
for row in rows {
let date: String = row.get("date");
let input_tokens = nonnegative_i64(row.get::<Option<i64>, _>("input_tokens"));
let output_tokens = nonnegative_i64(row.get::<Option<i64>, _>("output_tokens"));
let input_tokens =
optional_nonnegative_i64_to_u64(row.get::<Option<i64>, _>("input_tokens"));
let output_tokens =
optional_nonnegative_i64_to_u64(row.get::<Option<i64>, _>("output_tokens"));
let cache_creation_tokens =
nonnegative_i64(row.get::<Option<i64>, _>("cache_creation_tokens"));
let cached_tokens = nonnegative_i64(row.get::<Option<i64>, _>("cached_tokens"));
optional_nonnegative_i64_to_u64(row.get::<Option<i64>, _>("cache_creation_tokens"));
let cached_tokens =
optional_nonnegative_i64_to_u64(row.get::<Option<i64>, _>("cached_tokens"));
let cost = usage_breakdown_cost(
input_tokens,
output_tokens,
@@ -932,48 +959,6 @@ fn level_from_db(value: &str) -> SqliteUsageRecordStoreResult<UsageRecordLevel>
}
}
fn parse_rfc3339(value: &str) -> SqliteUsageRecordStoreResult<DateTime<Utc>> {
Ok(DateTime::parse_from_rfc3339(value)?.with_timezone(&Utc))
}
/// 管理端使用记录服务相关类型。
/// 使用记录查询过滤器。
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct AdminUsageRecordFilter {
/// 事件类别。
pub kind: Option<String>,
/// 事件等级。
pub level: Option<UsageRecordLevel>,
/// 请求 ID。
pub request_id: Option<String>,
/// 账号 ID。
pub account_id: Option<String>,
/// 路由。
pub route: Option<String>,
/// 模型。
pub model: Option<String>,
/// HTTP 状态码。
pub status_code: Option<i64>,
/// 上游传输方式。
pub transport: Option<String>,
/// 同一请求内的上游尝试序号。
pub attempt_index: Option<i64>,
/// 上游 HTTP 状态码。
pub upstream_status_code: Option<i64>,
/// 失败分类。
pub failure_class: Option<String>,
/// 上游响应 ID。
pub response_id: Option<String>,
/// 上游请求 ID。
pub upstream_request_id: Option<String>,
/// 搜索关键词。
pub search: Option<String>,
/// 起始时间。
pub start_time: Option<DateTime<Utc>>,
/// 结束时间。
pub end_time: Option<DateTime<Utc>>,
}
/// 使用记录汇总。
#[derive(Debug, Clone, Copy, Default, PartialEq)]
pub struct UsageRecordSummary {
@@ -1062,222 +1047,3 @@ pub struct UsageRecordTrendPoint {
/// 平均耗时。
pub average_latency_ms: Option<f64>,
}
impl From<AdminUsageRecordFilter> for UsageRecordFilter {
fn from(filter: AdminUsageRecordFilter) -> Self {
Self {
kind: filter.kind,
level: filter.level,
request_id: filter.request_id,
account_id: filter.account_id,
route: filter.route,
model: filter.model,
status_code: filter.status_code,
transport: filter.transport,
attempt_index: filter.attempt_index,
upstream_status_code: filter.upstream_status_code,
failure_class: filter.failure_class,
response_id: filter.response_id,
upstream_request_id: filter.upstream_request_id,
search: filter.search,
start_time: filter.start_time,
end_time: filter.end_time,
}
}
}
/// 清空使用记录结果。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AdminClearUsageRecords {
/// 清理数量。
pub cleared: u64,
}
/// 管理端使用记录错误。
#[derive(Debug, Error)]
pub enum AdminUsageRecordError {
/// 列表失败。
#[error("failed to list usage records")]
List,
/// 读取失败。
#[error("failed to get usage record")]
Get,
/// 清空失败。
#[error("failed to clear usage records")]
Clear,
/// 写入失败。
#[error("failed to append usage record")]
Append,
/// 裁剪失败。
#[error("failed to trim usage records")]
Trim,
}
/// 管理端使用记录服务。
#[derive(Clone)]
pub struct AdminUsageRecordService {
store: SqliteUsageRecordStore,
settings: std::sync::Arc<tokio::sync::RwLock<AdminUsageRecordSettings>>,
}
#[derive(Debug, Clone, Copy)]
struct AdminUsageRecordSettings {
enabled: bool,
capacity: u32,
capture_body: bool,
}
impl AdminUsageRecordService {
/// 构造管理端使用记录服务。
pub fn new(
store: SqliteUsageRecordStore,
enabled: bool,
capacity: u32,
capture_body: bool,
) -> Self {
Self {
store,
settings: std::sync::Arc::new(tokio::sync::RwLock::new(AdminUsageRecordSettings {
enabled,
capacity,
capture_body,
})),
}
}
/// 分页查询日志。
pub async fn list(
&self,
cursor: Option<String>,
limit: u32,
filter: AdminUsageRecordFilter,
) -> Result<Page<UsageRecord>, AdminUsageRecordError> {
self.store
.list(filter.into(), cursor, limit)
.await
.map_err(|_| AdminUsageRecordError::List)
}
/// 按页码查询日志。
pub async fn list_page(
&self,
page: u32,
page_size: u32,
filter: AdminUsageRecordFilter,
) -> Result<NumberedPage<UsageRecord>, AdminUsageRecordError> {
self.store
.list_page(filter.into(), page, page_size)
.await
.map_err(|_| AdminUsageRecordError::List)
}
/// 按 ID 读取日志。
pub async fn get(&self, id: &str) -> Result<Option<UsageRecord>, AdminUsageRecordError> {
self.store
.get(id)
.await
.map_err(|_| AdminUsageRecordError::Get)
}
/// 汇总使用记录。
pub async fn summary(
&self,
filter: AdminUsageRecordFilter,
) -> Result<UsageRecordSummary, AdminUsageRecordError> {
self.store
.summary(filter.into())
.await
.map_err(|_| AdminUsageRecordError::List)
}
/// 按模型来源聚合使用记录分布。
pub async fn model_distribution(
&self,
filter: AdminUsageRecordFilter,
source: UsageRecordModelSource,
) -> Result<Vec<UsageRecordBreakdown>, AdminUsageRecordError> {
self.store
.model_distribution(filter.into(), source)
.await
.map_err(|_| AdminUsageRecordError::List)
}
/// 按端点来源聚合使用记录分布。
pub async fn endpoint_distribution(
&self,
filter: AdminUsageRecordFilter,
source: UsageRecordEndpointSource,
) -> Result<Vec<UsageRecordBreakdown>, AdminUsageRecordError> {
self.store
.endpoint_distribution(filter.into(), source)
.await
.map_err(|_| AdminUsageRecordError::List)
}
/// 聚合 Token 趋势。
pub async fn token_trend(
&self,
filter: AdminUsageRecordFilter,
) -> Result<Vec<UsageRecordTrendPoint>, AdminUsageRecordError> {
self.store
.token_trend(filter.into())
.await
.map_err(|_| AdminUsageRecordError::List)
}
/// 聚合延迟趋势。
pub async fn latency_trend(
&self,
filter: AdminUsageRecordFilter,
) -> Result<Vec<UsageRecordTrendPoint>, AdminUsageRecordError> {
self.store
.latency_trend(filter.into())
.await
.map_err(|_| AdminUsageRecordError::List)
}
/// 清空日志。
pub async fn clear(&self) -> Result<AdminClearUsageRecords, AdminUsageRecordError> {
self.store
.clear()
.await
.map(|cleared| AdminClearUsageRecords { cleared })
.map_err(|_| AdminUsageRecordError::Clear)
}
/// 记录使用记录。
pub async fn record(&self, mut event: UsageRecord) -> Result<(), AdminUsageRecordError> {
let settings = *self.settings.read().await;
if !settings.enabled && event.level != UsageRecordLevel::Error {
return Ok(());
}
apply_capture_body_policy(&mut event, settings.capture_body);
self.store
.append(&event)
.await
.map_err(|_| AdminUsageRecordError::Append)?;
self.store
.trim_to_capacity(settings.capacity)
.await
.map_err(|_| AdminUsageRecordError::Trim)?;
Ok(())
}
}
fn apply_capture_body_policy(event: &mut UsageRecord, capture_body: bool) {
if capture_body {
return;
}
let Some(metadata) = event.metadata.as_object_mut() else {
return;
};
for key in [
"body",
"rawBody",
"requestBody",
"responseBody",
"upstreamBody",
] {
metadata.remove(key);
}
}
backend/src/admin/response.rs
+9
View File
@@ -19,6 +19,7 @@ const ADMIN_INVALID_ENDPOINT_SOURCE_CODE: u32 = 40004;
const ADMIN_SESSION_REQUIRED_CODE: u32 = 40101;
const ADMIN_INVALID_CREDENTIALS_CODE: u32 = 40102;
const ADMIN_INVALID_API_KEY_CODE: u32 = 40103;
const ADMIN_TOO_MANY_LOGIN_ATTEMPTS_CODE: u32 = 42901;
const ADMIN_NOT_FOUND_CODE: u32 = 40401;
const ADMIN_CONFLICT_CODE: u32 = 40901;
const ADMIN_SETTINGS_PERSIST_CODE: u32 = 50000;
@@ -98,6 +99,14 @@ impl AdminError {
)
}
pub fn too_many_login_attempts() -> Self {
Self::new(
StatusCode::TOO_MANY_REQUESTS,
ADMIN_TOO_MANY_LOGIN_ATTEMPTS_CODE,
"Too many admin login attempts",
)
}
pub fn not_found(message: impl Into<String>) -> Self {
Self::new(StatusCode::NOT_FOUND, ADMIN_NOT_FOUND_CODE, message)
}
backend/src/admin/router.rs
+1 -1
View File
@@ -17,7 +17,7 @@ use super::{
keys::routes::{api_keys, batch_delete_api_keys, create_api_key, update_api_key},
monitoring::{
dashboard::{dashboard_summary, dashboard_trend},
usage_records::{
usage_record_routes::{
clear_usage_records, usage_record_detail, usage_records,
usage_records_endpoint_distribution, usage_records_latency_trend,
usage_records_model_distribution, usage_records_summary, usage_records_token_trend,
backend/src/admin/settings/routes.rs
+7 -18
View File
@@ -1,17 +1,12 @@
//! 管理端设置处理器。
use axum::{
body::Bytes,
extract::State,
http::{HeaderMap, StatusCode},
response::IntoResponse,
};
use axum::{body::Bytes, extract::State, http::StatusCode, response::IntoResponse};
use serde::Serialize;
use serde_json::Value;
use std::collections::BTreeMap;
use crate::{
admin::auth::session::require_admin_auth,
admin::auth::session::AdminAuth,
admin::response::{AdminEnvelope, AdminError, AdminResponse},
config::{
settings::{AdminSettingsPatch, RuntimeSettingsError, SettingsServiceError},
@@ -59,9 +54,8 @@ impl AdminSettingsData {
/// `GET /api/admin/settings`
pub(crate) async fn settings(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
) -> Result<impl IntoResponse, AdminError> {
require_admin_auth(&state, &headers).await?;
let config = state.services.settings.current();
Ok(AdminResponse::new(
StatusCode::OK,
@@ -72,11 +66,9 @@ pub(crate) async fn settings(
/// `POST /api/admin/settings`
pub(crate) async fn update_settings(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
body: Bytes,
) -> Result<impl IntoResponse, AdminError> {
require_admin_auth(&state, &headers).await?;
let patch = parse_settings_patch(&body)?;
let config = state
.services
@@ -100,9 +92,8 @@ pub(crate) async fn update_settings(
/// `GET /api/admin/settings/admin-api-key`
pub(crate) async fn admin_api_key_status(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
) -> Result<impl IntoResponse, AdminError> {
require_admin_auth(&state, &headers).await?;
let status = state
.services
.settings
@@ -119,9 +110,8 @@ pub(crate) async fn admin_api_key_status(
/// `POST /api/admin/settings/admin-api-key/regenerate`
pub(crate) async fn regenerate_admin_api_key(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
) -> Result<impl IntoResponse, AdminError> {
require_admin_auth(&state, &headers).await?;
let key = state
.services
.settings
@@ -138,9 +128,8 @@ pub(crate) async fn regenerate_admin_api_key(
/// `DELETE /api/admin/settings/admin-api-key`
pub(crate) async fn delete_admin_api_key(
State(state): State<AppState>,
headers: HeaderMap,
_auth: AdminAuth,
) -> Result<impl IntoResponse, AdminError> {
require_admin_auth(&state, &headers).await?;
state
.services
.settings
backend/src/admin/system/mod.rs
+2
View File
@@ -1,3 +1,5 @@
//! 系统版本与在线更新处理器。
pub mod routes;
mod state;
mod updater;
backend/src/admin/system/routes.rs
+3 -1530
View File
File diff suppressed because it is too large Load Diff
backend/src/admin/system/state.rs
+234
View File
@@ -0,0 +1,234 @@
//! 系统自更新状态与文件锁。
use std::{
fs, io,
io::Write,
path::{Path, PathBuf},
time::Duration,
};
use chrono::Utc;
use serde::{Deserialize, Serialize};
use crate::admin::response::AdminError;
#[derive(Debug, Deserialize, Serialize, Default)]
#[serde(rename_all = "camelCase")]
pub(crate) struct SystemUpdateStatusData {
pub previous_version: Option<String>,
pub current_version: Option<String>,
#[serde(default)]
pub operation: SystemOperationState,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(crate) enum SystemOperationKind {
Update,
Rollback,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "camelCase")]
enum SystemOperationStatus {
#[default]
Idle,
Running,
Succeeded,
Failed,
}
#[derive(Debug, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct SystemOperationState {
operation_id: Option<String>,
kind: Option<SystemOperationKind>,
status: SystemOperationStatus,
target_version: Option<String>,
message: Option<String>,
error: Option<String>,
started_at: Option<String>,
finished_at: Option<String>,
}
impl Default for SystemOperationState {
fn default() -> Self {
Self {
operation_id: None,
kind: None,
status: SystemOperationStatus::Idle,
target_version: None,
message: None,
error: None,
started_at: None,
finished_at: None,
}
}
}
#[derive(Debug)]
pub(crate) struct OperationLock {
path: PathBuf,
}
pub(crate) fn read_state(path: &Path) -> Result<SystemUpdateStatusData, AdminError> {
if !path.exists() {
return Ok(SystemUpdateStatusData::default());
}
let data =
fs::read_to_string(path).map_err(internal_error_with("Failed to read update state"))?;
serde_json::from_str(&data).map_err(internal_error_with("Invalid update state"))
}
fn write_state(path: &Path, state: &SystemUpdateStatusData) -> Result<(), AdminError> {
if let Some(parent) = path
.parent()
.filter(|parent| !parent.as_os_str().is_empty())
{
fs::create_dir_all(parent).map_err(internal_error_with(
"Failed to create update state directory",
))?;
}
let data = serde_json::to_string_pretty(state)
.map_err(internal_error_with("Failed to encode update state"))?;
fs::write(path, data).map_err(internal_error_with("Failed to write update state"))
}
pub(crate) fn set_operation_running(
path: &Path,
operation_id: &str,
kind: SystemOperationKind,
version: Option<&str>,
) -> Result<(), AdminError> {
let mut state = read_state(path)?;
state.operation = SystemOperationState {
operation_id: Some(operation_id.to_string()),
kind: Some(kind),
status: SystemOperationStatus::Running,
target_version: version.map(ToString::to_string),
message: Some("operation running".to_string()),
error: None,
started_at: Some(Utc::now().to_rfc3339()),
finished_at: None,
};
write_state(path, &state)
}
pub(crate) fn finish_operation(
path: &Path,
operation_id: &str,
kind: SystemOperationKind,
version: Option<String>,
error: Option<String>,
) {
let mut state = match read_state(path) {
Ok(state) => state,
Err(error) => {
eprintln!("failed to read update state: {error}");
return;
}
};
if state.operation.operation_id.as_deref() != Some(operation_id) {
return;
}
if let Some(error) = error {
state.operation.status = SystemOperationStatus::Failed;
state.operation.message = Some("operation failed".to_string());
state.operation.error = Some(error);
} else {
state.operation.status = SystemOperationStatus::Succeeded;
state.operation.message = Some("operation succeeded".to_string());
state.operation.error = None;
if kind == SystemOperationKind::Update {
state.previous_version = state.current_version.take();
state.current_version = version.clone();
}
state.operation.target_version = version;
}
state.operation.finished_at = Some(Utc::now().to_rfc3339());
if let Err(error) = write_state(path, &state) {
eprintln!("failed to write update operation state: {error}");
}
}
pub(crate) fn operation_id(kind: &str) -> String {
format!("sysop-{kind}-{}", Utc::now().timestamp_millis())
}
impl OperationLock {
pub(crate) fn acquire(path: &Path) -> Result<Self, AdminError> {
if let Some(parent) = path
.parent()
.filter(|parent| !parent.as_os_str().is_empty())
{
fs::create_dir_all(parent).map_err(internal_error_with(
"Failed to create update lock directory",
))?;
}
match Self::try_create(path) {
Ok(()) => Ok(Self {
path: path.to_path_buf(),
}),
Err(error) if error.kind() == io::ErrorKind::AlreadyExists => {
if stale_lock(path)? {
fs::remove_file(path)
.map_err(internal_error_with("Failed to remove stale update lock"))?;
Self::try_create(path)
.map_err(internal_error_with("Failed to create update lock"))?;
return Ok(Self {
path: path.to_path_buf(),
});
}
Err(AdminError::conflict("System update already running"))
}
Err(error) => Err(internal_error("Failed to create update lock", error)),
}
}
fn try_create(path: &Path) -> io::Result<()> {
let mut file = fs::OpenOptions::new()
.write(true)
.create_new(true)
.open(path)?;
writeln!(
file,
"pid={}\ncreated_at={}",
std::process::id(),
Utc::now().to_rfc3339()
)
}
}
impl Drop for OperationLock {
fn drop(&mut self) {
if let Err(error) = fs::remove_file(&self.path) {
if error.kind() != io::ErrorKind::NotFound {
eprintln!("failed to remove update lock: {error}");
}
}
}
}
fn stale_lock(path: &Path) -> Result<bool, AdminError> {
let metadata = fs::metadata(path).map_err(internal_error_with("Failed to read update lock"))?;
let modified = metadata
.modified()
.map_err(internal_error_with("Failed to read update lock timestamp"))?;
modified
.elapsed()
.map(|age| age > Duration::from_secs(30 * 60))
.map_err(internal_error_with("Failed to calculate update lock age"))
}
fn internal_error(context: &str, error: impl std::fmt::Display) -> AdminError {
AdminError::internal(format!("{context}: {error}"))
}
fn internal_error_with<E: std::fmt::Display>(
context: &'static str,
) -> impl FnOnce(E) -> AdminError {
move |error| internal_error(context, error)
}
backend/src/admin/system/updater.rs
+665
View File
@@ -0,0 +1,665 @@
//! 管理端系统版本与在线更新路由。
use std::{
convert::Infallible,
env, fmt, fs, io,
path::{Path, PathBuf},
sync::OnceLock,
time::Duration,
};
use axum::{
extract::Query,
http::StatusCode,
response::{
sse::{Event, KeepAlive, Sse},
IntoResponse,
},
};
use chrono::Utc;
use futures::Stream;
use serde::{Deserialize, Serialize};
use tokio::sync::{broadcast, Mutex};
use crate::{
admin::auth::session::AdminAuth,
admin::response::{AdminEnvelope, AdminError, AdminResponse},
runtime::shutdown::request_shutdown,
};
use super::state::{
finish_operation, operation_id, read_state, set_operation_running, OperationLock,
SystemOperationKind,
};
mod download;
use download::{
download_file, validate_download_url, validate_github_api_base, verify_checksum,
DownloadProgress,
};
mod archive;
use archive::{extract_release_archive, replace_release_files, rollback_release_update};
mod release;
use release::{
check_latest_release, fetch_latest_release, select_release_archive, update_info_from_release,
GitHubRelease,
};
const APP_BINARY_NAME: &str = "codex-proxy-rs";
const DEFAULT_WEB_DIST_DIR: &str = "/app/web/dist";
const GITHUB_API_BASE: &str = "https://api.github.com/repos";
const MAX_DOWNLOAD_SIZE: u64 = 500 * 1024 * 1024;
static SYSTEM_OPERATION_LOCK: Mutex<()> = Mutex::const_new(());
static UPDATE_EVENT_SENDER: OnceLock<broadcast::Sender<SystemUpdateEvent>> = OnceLock::new();
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct CheckUpdateQuery {
force: Option<bool>,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct VersionData {
version: String,
git_sha: String,
build_time: String,
deployment_mode: String,
deployment_mode_label: String,
update_channel: String,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct UpdateStartedData {
operation_id: String,
deployment_mode: String,
message: String,
need_restart: bool,
target_version: String,
}
#[derive(Debug, Clone, Copy, Serialize)]
#[serde(rename_all = "lowercase")]
enum UpdateLogLevel {
Info,
Success,
Warning,
Error,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct SystemUpdateEvent {
id: String,
operation_id: Option<String>,
level: UpdateLogLevel,
step: Option<String>,
message: String,
at: String,
}
#[derive(Debug, Clone)]
struct SystemUpdateConfig {
version: String,
git_sha: String,
build_time: String,
deployment_mode: String,
build_type: String,
update_channel: String,
update_repository: Option<String>,
github_api_base: String,
executable_path: Option<PathBuf>,
web_dist_dir: PathBuf,
update_state_file: PathBuf,
update_lock_file: PathBuf,
}
/// `GET /api/admin/system/version`
pub(crate) async fn version(_auth: AdminAuth) -> Result<impl IntoResponse, AdminError> {
let config = SystemUpdateConfig::from_env();
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(config.version_data()),
))
}
/// `GET /api/admin/system/check-updates`
pub(crate) async fn check_updates(
_auth: AdminAuth,
Query(query): Query<CheckUpdateQuery>,
) -> Result<impl IntoResponse, AdminError> {
let config = SystemUpdateConfig::from_env();
let info = check_latest_release(&config, query.force.unwrap_or(false)).await;
Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(info)))
}
/// `GET /api/admin/system/update-events`
pub(crate) async fn update_event_stream(
_auth: AdminAuth,
) -> Result<Sse<impl Stream<Item = Result<Event, Infallible>>>, AdminError> {
let receiver = update_event_sender().subscribe();
let stream = futures::stream::unfold(receiver, |mut receiver| async move {
loop {
match receiver.recv().await {
Ok(message) => {
let id = message.id.clone();
let data = serde_json::to_string(&message).unwrap_or_else(|_| "{}".to_string());
let event = Event::default().event("update").id(id).data(data);
return Some((Ok(event), receiver));
}
Err(broadcast::error::RecvError::Lagged(_)) => continue,
Err(broadcast::error::RecvError::Closed) => return None,
}
}
});
Ok(Sse::new(stream).keep_alive(KeepAlive::default()))
}
/// `POST /api/admin/system/update`
pub(crate) async fn perform_update(_auth: AdminAuth) -> Result<impl IntoResponse, AdminError> {
let _guard = SYSTEM_OPERATION_LOCK
.try_lock()
.map_err(|_| AdminError::conflict("System update already running"))?;
let config = SystemUpdateConfig::from_env();
if let Some(reason) = config.update_support_error() {
emit_update_event(
UpdateLogLevel::Error,
None,
Some("preflight"),
reason.clone(),
);
return Err(AdminError::conflict(reason));
}
emit_update_event(
UpdateLogLevel::Info,
None,
Some("release"),
"正在获取最新 Release 信息",
);
let release = fetch_latest_release(
&config.github_api_base,
config
.update_repository
.as_deref()
.ok_or_else(|| AdminError::conflict("Update checks require CPR_UPDATE_REPOSITORY"))?,
)
.await
.map_err(|error| {
emit_update_event(UpdateLogLevel::Error, None, Some("release"), error.clone());
AdminError::bad_gateway(error)
})?;
let info = update_info_from_release(&config, release.clone());
if !info.has_update {
emit_update_event(
UpdateLogLevel::Warning,
None,
Some("release"),
"当前版本已是最新",
);
return Err(AdminError::conflict("Already up to date"));
}
let target_version = info.latest_version.clone();
let operation_id = operation_id("update");
let lock = OperationLock::acquire(&config.update_lock_file)?;
set_operation_running(
&config.update_state_file,
&operation_id,
SystemOperationKind::Update,
Some(&target_version),
)?;
emit_update_event(
UpdateLogLevel::Info,
Some(&operation_id),
Some("prepare"),
format!("准备更新到 v{target_version}"),
);
let result = perform_release_update(&config, &release, &target_version, &operation_id).await;
if let Err(error) = &result {
emit_update_event(
UpdateLogLevel::Error,
Some(&operation_id),
Some("failed"),
error.to_string(),
);
} else {
emit_update_event(
UpdateLogLevel::Success,
Some(&operation_id),
Some("done"),
"更新文件已替换，等待服务重启生效",
);
}
finish_operation(
&config.update_state_file,
&operation_id,
SystemOperationKind::Update,
result.as_ref().ok().map(|_| target_version.clone()),
result.as_ref().err().map(ToString::to_string),
);
drop(lock);
result?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(UpdateStartedData {
operation_id,
deployment_mode: config.deployment_mode,
message: "更新完成，请重启服务。".to_string(),
need_restart: true,
target_version,
}),
))
}
/// `GET /api/admin/system/update-status`
pub(crate) async fn update_status(_auth: AdminAuth) -> Result<impl IntoResponse, AdminError> {
let config = SystemUpdateConfig::from_env();
let status = read_state(&config.update_state_file)?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(status),
))
}
/// `POST /api/admin/system/rollback`
pub(crate) async fn rollback(_auth: AdminAuth) -> Result<impl IntoResponse, AdminError> {
let _guard = SYSTEM_OPERATION_LOCK
.try_lock()
.map_err(|_| AdminError::conflict("系统操作正在执行中"))?;
let config = SystemUpdateConfig::from_env();
if let Some(reason) = config.update_support_error() {
return Err(AdminError::conflict(reason));
}
let operation_id = operation_id("rollback");
let lock = OperationLock::acquire(&config.update_lock_file)?;
set_operation_running(
&config.update_state_file,
&operation_id,
SystemOperationKind::Rollback,
None,
)?;
let result = rollback_release_update(&config).await;
finish_operation(
&config.update_state_file,
&operation_id,
SystemOperationKind::Rollback,
None,
result.as_ref().err().map(ToString::to_string),
);
drop(lock);
result?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(serde_json::json!({
"message": "回滚完成，请重启服务。",
"needRestart": true,
"operationId": operation_id
})),
))
}
/// `POST /api/admin/system/restart`
pub(crate) async fn restart(_auth: AdminAuth) -> Result<impl IntoResponse, AdminError> {
if env_string("CPR_ENABLE_SELF_RESTART").as_deref() != Some("true") {
return Err(AdminError::conflict(
"自重启未启用，请设置 CPR_ENABLE_SELF_RESTART=true",
));
}
tokio::spawn(async {
tokio::time::sleep(Duration::from_millis(500)).await;
request_shutdown();
});
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(serde_json::json!({
"message": "已安排重启",
"operationId": operation_id("restart")
})),
))
}
impl SystemUpdateConfig {
fn from_env() -> Self {
let update_repository = env_string("CPR_UPDATE_REPOSITORY");
let state_file = env_string("CPR_UPDATE_STATE_FILE")
.map(PathBuf::from)
.unwrap_or_else(|| PathBuf::from("/app/data/update-state.json"));
let update_lock_file = env_string("CPR_UPDATE_LOCK_FILE")
.map(PathBuf::from)
.unwrap_or_else(|| state_file.with_extension("lock"));
Self {
version: build_version(),
git_sha: build_git_sha(),
build_time: build_time(),
deployment_mode: env_string("CPR_DEPLOYMENT_MODE")
.unwrap_or_else(|| "source".to_string()),
build_type: build_type(),
update_channel: env_string("CPR_UPDATE_CHANNEL")
.unwrap_or_else(|| "stable".to_string()),
update_repository,
github_api_base: env_string("CPR_GITHUB_API_BASE")
.unwrap_or_else(|| GITHUB_API_BASE.to_string()),
executable_path: env_string("CPR_UPDATE_EXE_PATH").map(PathBuf::from),
web_dist_dir: env_string("CPR_WEB_DIST_DIR")
.map(PathBuf::from)
.unwrap_or_else(|| PathBuf::from(DEFAULT_WEB_DIST_DIR)),
update_state_file: state_file,
update_lock_file,
}
}
fn version_data(&self) -> VersionData {
VersionData {
version: self.version.clone(),
git_sha: self.git_sha.clone(),
build_time: self.build_time.clone(),
deployment_mode: self.deployment_mode.clone(),
deployment_mode_label: deployment_mode_label(&self.deployment_mode).to_string(),
update_channel: self.update_channel.clone(),
}
}
fn update_support_error(&self) -> Option<String> {
if self.build_type != "release" {
return Some("一键更新需要正式构建包".to_string());
}
if self.update_repository.is_none() {
return Some("检查更新需要配置 CPR_UPDATE_REPOSITORY".to_string());
}
if let Err(error) = validate_github_api_base(&self.github_api_base) {
return Some(error);
}
None
}
fn release_cache_key(&self) -> String {
format!(
"{}|{}|{}|{}|{}",
self.update_repository.as_deref().unwrap_or_default(),
self.version,
self.deployment_mode,
self.build_type,
self.update_channel,
)
}
fn executable_path(&self) -> Result<PathBuf, AdminError> {
if let Some(path) = &self.executable_path {
return Ok(path.clone());
}
env::current_exe()
.and_then(fs::canonicalize)
.map_err(internal_error_with("Failed to resolve executable"))
}
}
fn internal_error(context: &str, error: impl fmt::Display) -> AdminError {
AdminError::internal(format!("{context}: {error}"))
}
fn internal_error_with<E: fmt::Display>(context: &'static str) -> impl FnOnce(E) -> AdminError {
move |error| internal_error(context, error)
}
fn bad_request_with<E: fmt::Display>(context: &'static str) -> impl FnOnce(E) -> AdminError {
move |error| AdminError::bad_request(format!("{context}: {error}"))
}
fn bad_gateway_with<E: fmt::Display>(context: &'static str) -> impl FnOnce(E) -> AdminError {
move |error| AdminError::bad_gateway(format!("{context}: {error}"))
}
async fn perform_release_update(
config: &SystemUpdateConfig,
release: &GitHubRelease,
version: &str,
operation_id: &str,
) -> Result<(), AdminError> {
emit_update_event(
UpdateLogLevel::Info,
Some(operation_id),
Some("asset"),
"正在选择匹配当前平台的更新包",
);
let archive = select_release_archive(release, version)?;
emit_update_event(
UpdateLogLevel::Info,
Some(operation_id),
Some("asset"),
format!(
"已选择更新包 {} ({})",
archive.name,
format_bytes(archive.size)
),
);
emit_update_event(
UpdateLogLevel::Info,
Some(operation_id),
Some("verify"),
"正在校验下载地址",
);
validate_download_url(&archive.browser_download_url, &config.github_api_base)?;
if archive.size > MAX_DOWNLOAD_SIZE {
return Err(AdminError::bad_request("Release archive is too large"));
}
let checksum = release
.assets
.iter()
.find(|asset| asset.name == "checksums.txt")
.ok_or_else(|| AdminError::bad_gateway("Release checksums.txt is required"))?;
validate_download_url(&checksum.browser_download_url, &config.github_api_base)?;
emit_update_event(
UpdateLogLevel::Info,
Some(operation_id),
Some("prepare"),
"正在创建临时更新目录",
);
let exe_path = config.executable_path()?;
let exe_dir = exe_path
.parent()
.ok_or_else(|| AdminError::internal("Executable path has no parent"))?;
let temp_dir = fs::canonicalize(exe_dir)
.and_then(|dir| fs::create_dir_all(&dir).map(|()| dir))
.and_then(|dir| tempfile_dir_in(&dir, ".codex-proxy-rs-update-"))
.map_err(internal_error_with("Failed to create update temp dir"))?;
let archive_path = temp_dir.join(&archive.name);
emit_update_event(
UpdateLogLevel::Info,
Some(operation_id),
Some("download"),
"开始下载更新包",
);
download_file(
&archive.browser_download_url,
&archive_path,
MAX_DOWNLOAD_SIZE,
Some(DownloadProgress {
operation_id,
total_size: archive.size,
}),
)
.await?;
emit_update_event(
UpdateLogLevel::Success,
Some(operation_id),
Some("download"),
"更新包下载完成",
);
emit_update_event(
UpdateLogLevel::Info,
Some(operation_id),
Some("checksum"),
"正在校验 checksum",
);
verify_checksum(&archive_path, &archive.name, &checksum.browser_download_url).await?;
emit_update_event(
UpdateLogLevel::Success,
Some(operation_id),
Some("checksum"),
"checksum 校验通过",
);
emit_update_event(
UpdateLogLevel::Info,
Some(operation_id),
Some("extract"),
"正在解压更新包",
);
let extracted = extract_release_archive(&archive_path, &temp_dir)?;
emit_update_event(
UpdateLogLevel::Success,
Some(operation_id),
Some("extract"),
"更新包解压完成",
);
emit_update_event(
UpdateLogLevel::Info,
Some(operation_id),
Some("replace"),
"正在替换应用文件",
);
replace_release_files(&exe_path, &config.web_dist_dir, extracted)?;
emit_update_event(
UpdateLogLevel::Success,
Some(operation_id),
Some("replace"),
"应用文件替换完成",
);
let _ = fs::remove_dir_all(temp_dir);
Ok(())
}
fn tempfile_dir_in(parent: &Path, prefix: &str) -> io::Result<PathBuf> {
for attempt in 0..100 {
let path = parent.join(format!(
"{prefix}{}-{attempt}",
Utc::now().timestamp_nanos_opt().unwrap_or_default()
));
match fs::create_dir(&path) {
Ok(()) => return Ok(path),
Err(error) if error.kind() == io::ErrorKind::AlreadyExists => continue,
Err(error) => return Err(error),
}
}
Err(io::Error::new(
io::ErrorKind::AlreadyExists,
"failed to create unique temp dir",
))
}
fn env_string(key: &str) -> Option<String> {
env::var(key)
.ok()
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
}
fn runtime_build_info_override(key: &str) -> Option<String> {
if cfg!(debug_assertions) {
env_string(key)
} else {
None
}
}
fn build_version() -> String {
runtime_build_info_override("CPR_VERSION").unwrap_or_else(|| {
option_env!("CPR_VERSION")
.unwrap_or(env!("CARGO_PKG_VERSION"))
.to_string()
})
}
fn build_git_sha() -> String {
runtime_build_info_override("CPR_GIT_SHA")
.unwrap_or_else(|| option_env!("CPR_GIT_SHA").unwrap_or("unknown").to_string())
}
fn build_time() -> String {
runtime_build_info_override("CPR_BUILD_TIME").unwrap_or_else(|| {
option_env!("CPR_BUILD_TIME")
.unwrap_or("unknown")
.to_string()
})
}
fn build_type() -> String {
runtime_build_info_override("CPR_BUILD_TYPE").unwrap_or_else(|| {
option_env!("CPR_BUILD_TYPE")
.unwrap_or("source")
.to_string()
})
}
fn update_event_sender() -> &'static broadcast::Sender<SystemUpdateEvent> {
UPDATE_EVENT_SENDER.get_or_init(|| {
let (sender, _receiver) = broadcast::channel(256);
sender
})
}
fn emit_update_event(
level: UpdateLogLevel,
operation_id: Option<&str>,
step: Option<&str>,
message: impl Into<String>,
) {
let now = Utc::now();
let event = SystemUpdateEvent {
id: format!(
"update-log-{}",
now.timestamp_nanos_opt()
.unwrap_or_else(|| now.timestamp_millis())
),
operation_id: operation_id.map(ToString::to_string),
level,
step: step.map(ToString::to_string),
message: message.into(),
at: now.to_rfc3339(),
};
let _ = update_event_sender().send(event);
}
fn format_bytes(bytes: u64) -> String {
const KIB: u64 = 1024;
const MIB: u64 = KIB * 1024;
if bytes >= MIB {
return format!("{:.1} MiB", bytes as f64 / MIB as f64);
}
if bytes >= KIB {
return format!("{:.1} KiB", bytes as f64 / KIB as f64);
}
format!("{bytes} B")
}
fn deployment_mode_label(value: &str) -> &str {
match value {
"docker" => "Docker",
"source" => "源码部署",
"binary" => "二进制部署",
_ => value,
}
}
fn build_type_label(value: &str) -> &str {
match value {
"release" => "正式构建",
"source" => "源码构建",
"dev" => "开发构建",
_ => value,
}
}
backend/src/admin/system/updater/archive.rs
+262
View File
@@ -0,0 +1,262 @@
//! 自更新归档解压、文件替换与回滚。
use std::{
fs, io,
path::{Component, Path, PathBuf},
};
use flate2::read::GzDecoder;
use tar::Archive;
use crate::admin::response::AdminError;
use super::{internal_error, internal_error_with, SystemUpdateConfig, APP_BINARY_NAME};
#[derive(Debug)]
pub(super) struct ExtractedRelease {
binary_path: PathBuf,
web_dist_dir: Option<PathBuf>,
}
pub(super) async fn rollback_release_update(config: &SystemUpdateConfig) -> Result<(), AdminError> {
let exe_path = config.executable_path()?;
let backup_path = backup_path_for(&exe_path);
if !backup_path.exists() {
return Err(AdminError::conflict("No binary backup found for rollback"));
}
fs::rename(&backup_path, &exe_path).map_err(internal_error_with("Binary rollback failed"))?;
let web_backup = backup_path_for(&config.web_dist_dir);
if web_backup.exists() {
if config.web_dist_dir.exists() {
fs::remove_dir_all(&config.web_dist_dir)
.map_err(internal_error_with("Web rollback cleanup failed"))?;
}
fs::rename(&web_backup, &config.web_dist_dir)
.map_err(internal_error_with("Web rollback failed"))?;
}
Ok(())
}
pub(super) fn extract_release_archive(
archive_path: &Path,
temp_dir: &Path,
) -> Result<ExtractedRelease, AdminError> {
let file =
fs::File::open(archive_path).map_err(internal_error_with("Failed to open archive"))?;
let decoder = GzDecoder::new(file);
let mut archive = Archive::new(decoder);
let extract_dir = temp_dir.join("extracted");
fs::create_dir_all(&extract_dir)
.map_err(internal_error_with("Failed to create extract dir"))?;
let binary_path = temp_dir.join(APP_BINARY_NAME);
let web_dist_dir = temp_dir.join("web-dist");
let mut found_binary = false;
let mut found_web = false;
for entry in archive
.entries()
.map_err(internal_error_with("Failed to read archive"))?
{
let mut entry = entry.map_err(internal_error_with("Invalid archive entry"))?;
if !entry.header().entry_type().is_file() {
continue;
}
let path = entry
.path()
.map_err(internal_error_with("Invalid archive path"))?
.to_path_buf();
if unsafe_archive_path(&path) {
return Err(AdminError::bad_request("Unsafe archive path"));
}
if path.file_name().is_some_and(|name| name == APP_BINARY_NAME) {
entry
.unpack(&binary_path)
.map_err(internal_error_with("Failed to extract binary"))?;
found_binary = true;
continue;
}
if let Some(relative) = web_dist_relative_path(&path) {
let target = web_dist_dir.join(relative);
if let Some(parent) = target.parent() {
fs::create_dir_all(parent)
.map_err(internal_error_with("Failed to create web asset dir"))?;
}
entry
.unpack(&target)
.map_err(internal_error_with("Failed to extract web asset"))?;
found_web = true;
}
}
if !found_binary {
return Err(AdminError::bad_request(
"Release archive does not contain codex-proxy-rs",
));
}
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
fs::set_permissions(&binary_path, fs::Permissions::from_mode(0o755))
.map_err(internal_error_with("Failed to chmod binary"))?;
}
Ok(ExtractedRelease {
binary_path,
web_dist_dir: found_web.then_some(web_dist_dir),
})
}
pub(super) fn replace_release_files(
exe_path: &Path,
web_dist_dir: &Path,
extracted: ExtractedRelease,
) -> Result<(), AdminError> {
let web_backup = backup_path_for(web_dist_dir);
let web_replaced = if let Some(new_web) = extracted.web_dist_dir {
replace_dir(web_dist_dir, &web_backup, &new_web)?;
true
} else {
false
};
let binary_backup = backup_path_for(exe_path);
if binary_backup.exists() {
if let Err(error) = fs::remove_file(&binary_backup) {
if web_replaced {
let _ = restore_dir(web_dist_dir, &web_backup);
}
return Err(internal_error("Failed to remove old binary backup", error));
}
}
if let Err(error) = move_file(exe_path, &binary_backup) {
if web_replaced {
let _ = restore_dir(web_dist_dir, &web_backup);
}
return Err(internal_error("Binary backup failed", error));
}
if let Err(error) = move_file(&extracted.binary_path, exe_path) {
let _ = fs::remove_file(exe_path);
let _ = move_file(&binary_backup, exe_path);
if web_replaced {
let _ = restore_dir(web_dist_dir, &web_backup);
}
return Err(internal_error("Binary replace failed", error));
}
Ok(())
}
fn replace_dir(current: &Path, backup: &Path, replacement: &Path) -> Result<(), AdminError> {
if backup.exists() {
fs::remove_dir_all(backup)
.map_err(internal_error_with("Failed to remove old web backup"))?;
}
if current.exists() {
move_dir(current, backup).map_err(internal_error_with("Failed to backup web assets"))?;
}
if let Err(error) = move_dir(replacement, current) {
if backup.exists() {
let _ = restore_dir(current, backup);
}
return Err(internal_error("Failed to replace web assets", error));
}
Ok(())
}
fn restore_dir(current: &Path, backup: &Path) -> io::Result<()> {
if current.exists() {
fs::remove_dir_all(current)?;
}
if backup.exists() {
move_dir(backup, current)?;
}
Ok(())
}
fn move_file(from: &Path, to: &Path) -> io::Result<()> {
match fs::rename(from, to) {
Ok(()) => Ok(()),
Err(error) if is_cross_device_link(&error) => {
fs::copy(from, to)?;
fs::remove_file(from)
}
Err(error) => Err(error),
}
}
fn move_dir(from: &Path, to: &Path) -> io::Result<()> {
match fs::rename(from, to) {
Ok(()) => Ok(()),
Err(error) if is_cross_device_link(&error) => {
copy_dir_all(from, to)?;
fs::remove_dir_all(from)
}
Err(error) => Err(error),
}
}
fn copy_dir_all(from: &Path, to: &Path) -> io::Result<()> {
fs::create_dir_all(to)?;
for entry in fs::read_dir(from)? {
let entry = entry?;
let file_type = entry.file_type()?;
let target = to.join(entry.file_name());
if file_type.is_dir() {
copy_dir_all(&entry.path(), &target)?;
} else if file_type.is_file() {
fs::copy(entry.path(), &target)?;
fs::set_permissions(&target, entry.metadata()?.permissions())?;
} else {
return Err(io::Error::new(
io::ErrorKind::Unsupported,
format!("unsupported file type in {}", entry.path().display()),
));
}
}
Ok(())
}
fn is_cross_device_link(error: &io::Error) -> bool {
error.kind() == io::ErrorKind::CrossesDevices
}
fn backup_path_for(path: &Path) -> PathBuf {
let mut backup = path.as_os_str().to_os_string();
backup.push(".backup");
PathBuf::from(backup)
}
fn web_dist_relative_path(path: &Path) -> Option<PathBuf> {
let components = path
.components()
.filter_map(|component| match component {
Component::Normal(value) => Some(value.to_string_lossy().to_string()),
_ => None,
})
.collect::<Vec<_>>();
for index in 0..components.len().saturating_sub(1) {
if components[index] == "web"
&& components
.get(index + 1)
.is_some_and(|value| value == "dist")
{
return Some(components[index + 2..].iter().collect());
}
if components[index] == "dist" {
return Some(components[index + 1..].iter().collect());
}
}
None
}
fn unsafe_archive_path(path: &Path) -> bool {
path.is_absolute()
|| path.components().any(|component| {
matches!(
component,
Component::ParentDir | Component::Prefix(_) | Component::RootDir
)
})
}
backend/src/admin/system/updater/download.rs
+226
View File
@@ -0,0 +1,226 @@
//! 自更新下载、checksum 与 URL 校验。
use std::{
fs,
io::{Read, Write},
path::Path,
time::Duration,
};
use sha2::{Digest, Sha256};
use crate::admin::response::AdminError;
use super::{
bad_gateway_with, bad_request_with, emit_update_event, format_bytes, internal_error_with,
UpdateLogLevel,
};
#[derive(Debug, Clone, Copy)]
pub(super) struct DownloadProgress<'a> {
pub operation_id: &'a str,
pub total_size: u64,
}
pub(super) async fn download_file(
url: &str,
dest: &Path,
max_size: u64,
progress: Option<DownloadProgress<'_>>,
) -> Result<(), AdminError> {
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(120))
.build()
.map_err(internal_error_with("Failed to create HTTP client"))?;
let mut response = client
.get(url)
.send()
.await
.map_err(bad_gateway_with("Download failed"))?;
if !response.status().is_success() {
return Err(AdminError::bad_gateway(format!(
"Download failed with {}",
response.status()
)));
}
let mut file =
fs::File::create(dest).map_err(internal_error_with("Failed to create download file"))?;
let mut downloaded = 0_u64;
let mut next_progress = 10_u64;
while let Some(chunk) = response
.chunk()
.await
.map_err(bad_gateway_with("Download stream failed"))?
{
downloaded += chunk.len() as u64;
if downloaded > max_size {
return Err(AdminError::bad_request("Download exceeds max allowed size"));
}
file.write_all(&chunk)
.map_err(internal_error_with("Failed to write download"))?;
if let Some(progress) = progress {
let percent = download_percent(downloaded, progress.total_size);
if progress.total_size > 0
&& (percent >= next_progress || downloaded >= progress.total_size)
{
emit_update_event(
UpdateLogLevel::Info,
Some(progress.operation_id),
Some("download"),
format!(
"已下载 {} / {} ({percent}%)",
format_bytes(downloaded),
format_bytes(progress.total_size)
),
);
next_progress = percent.saturating_add(10);
}
}
}
Ok(())
}
pub(super) async fn verify_checksum(
file_path: &Path,
file_name: &str,
checksum_url: &str,
) -> Result<(), AdminError> {
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(30))
.build()
.map_err(internal_error_with("Failed to create HTTP client"))?;
let body = client
.get(checksum_url)
.send()
.await
.map_err(bad_gateway_with("Checksum download failed"))?;
let status = body.status();
if !status.is_success() {
return Err(AdminError::bad_gateway(format!(
"Checksum download failed with {status}"
)));
}
let body = body
.text()
.await
.map_err(bad_gateway_with("Checksum read failed"))?;
let expected = body.lines().find_map(|line| {
let mut parts = line.split_whitespace();
let hash = parts.next()?;
let name = parts.next()?;
(Path::new(name).file_name()?.to_string_lossy() == file_name).then(|| hash.to_string())
});
let expected = expected.ok_or_else(|| AdminError::bad_gateway("Checksum not found"))?;
let actual = sha256_file(file_path)?;
if expected.eq_ignore_ascii_case(&actual) {
Ok(())
} else {
Err(AdminError::bad_gateway("Checksum mismatch"))
}
}
pub(super) fn validate_download_url(
raw_url: &str,
github_api_base: &str,
) -> Result<(), AdminError> {
let url = reqwest::Url::parse(raw_url).map_err(bad_request_with("Invalid download URL"))?;
if local_update_test_download_allowed(&url, github_api_base) {
return Ok(());
}
if url.scheme() != "https" {
return Err(AdminError::bad_request(
"Only HTTPS release downloads are allowed",
));
}
let host = url
.host_str()
.ok_or_else(|| AdminError::bad_request("Download URL is missing host"))?;
let allowed = host == "github.com"
|| host.ends_with(".github.com")
|| host == "objects.githubusercontent.com"
|| host.ends_with(".objects.githubusercontent.com");
if allowed {
Ok(())
} else {
Err(AdminError::bad_request(format!(
"Download host is not allowed: {host}"
)))
}
}
pub(super) fn validate_github_api_base(raw_url: &str) -> Result<(), String> {
let url = reqwest::Url::parse(raw_url)
.map_err(|error| format!("Invalid GitHub API base: {error}"))?;
if !url.username().is_empty()
|| url.password().is_some()
|| url.query().is_some()
|| url.fragment().is_some()
{
return Err("GitHub API base must not include credentials, query, or fragment".to_string());
}
if url.path().trim_end_matches('/') != "/repos" {
return Err("GitHub API base path must be /repos".to_string());
}
if url_host_is_loopback(&url) {
return Ok(());
}
if url.scheme() != "https" {
return Err("GitHub API base must use HTTPS".to_string());
}
if url.host_str() != Some("api.github.com") {
return Err("GitHub API base must be https://api.github.com/repos".to_string());
}
Ok(())
}
fn sha256_file(path: &Path) -> Result<String, AdminError> {
let mut file =
fs::File::open(path).map_err(internal_error_with("Failed to open checksum file"))?;
let mut hasher = Sha256::new();
let mut buffer = [0_u8; 8192];
loop {
let read = file
.read(&mut buffer)
.map_err(internal_error_with("Failed to read checksum file"))?;
if read == 0 {
break;
}
hasher.update(&buffer[..read]);
}
Ok(hex::encode(hasher.finalize()))
}
fn local_update_test_download_allowed(url: &reqwest::Url, github_api_base: &str) -> bool {
if url.scheme() != "http" || !url_host_is_loopback(url) {
return false;
}
let Ok(api_base) = reqwest::Url::parse(github_api_base) else {
return false;
};
api_base.scheme() == "http"
&& url_host_is_loopback(&api_base)
&& url.host_str() == api_base.host_str()
&& url.port_or_known_default() == api_base.port_or_known_default()
}
fn url_host_is_loopback(url: &reqwest::Url) -> bool {
url.host_str().is_some_and(|host| {
host == "localhost"
|| host == "127.0.0.1"
|| host == "::1"
|| host
.parse::<std::net::IpAddr>()
.is_ok_and(|ip| ip.is_loopback())
})
}
fn download_percent(downloaded: u64, total_size: u64) -> u64 {
if total_size == 0 {
return 0;
}
downloaded
.saturating_mul(100)
.saturating_div(total_size)
.min(100)
}
backend/src/admin/system/updater/release.rs
+281
View File
@@ -0,0 +1,281 @@
//! 自更新 Release 查询、缓存和 asset 选择。
use std::{
env,
time::{Duration, Instant},
};
use semver::Version;
use serde::{Deserialize, Serialize};
use tokio::sync::Mutex;
use crate::admin::response::AdminError;
use super::{
build_type_label, deployment_mode_label, download::validate_github_api_base,
SystemUpdateConfig, APP_BINARY_NAME,
};
const UPDATE_CACHE_TTL: Duration = Duration::from_secs(20 * 60);
static RELEASE_CACHE: Mutex<Option<CachedUpdateInfo>> = Mutex::const_new(None);
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub(super) struct UpdateInfoData {
pub current_version: String,
pub latest_version: String,
pub has_update: bool,
pub deployment_mode: String,
pub deployment_mode_label: String,
pub build_type: String,
pub build_type_label: String,
pub release_url: Option<String>,
pub notes: Option<String>,
pub cached: bool,
pub update_supported: bool,
pub unsupported_reason: Option<String>,
pub warning: Option<String>,
}
#[derive(Debug, Clone)]
struct CachedUpdateInfo {
key: String,
info: UpdateInfoData,
cached_at: Instant,
}
#[derive(Debug, Clone, Deserialize)]
pub(super) struct GitHubRelease {
pub tag_name: String,
pub name: Option<String>,
pub body: Option<String>,
pub html_url: Option<String>,
pub prerelease: bool,
#[serde(default)]
pub assets: Vec<GitHubAsset>,
}
#[derive(Debug, Clone, Deserialize)]
pub(super) struct GitHubAsset {
pub name: String,
pub browser_download_url: String,
pub size: u64,
}
pub(super) async fn check_latest_release(
config: &SystemUpdateConfig,
force: bool,
) -> UpdateInfoData {
let Some(repository) = config.update_repository.as_deref() else {
return UpdateInfoData {
current_version: config.version.clone(),
latest_version: config.version.clone(),
has_update: false,
deployment_mode: config.deployment_mode.clone(),
deployment_mode_label: deployment_mode_label(&config.deployment_mode).to_string(),
build_type: config.build_type.clone(),
build_type_label: build_type_label(&config.build_type).to_string(),
release_url: None,
notes: None,
cached: false,
update_supported: false,
unsupported_reason: Some("检查更新需要配置 CPR_UPDATE_REPOSITORY".to_string()),
warning: None,
};
};
if let Err(reason) = validate_github_api_base(&config.github_api_base) {
return UpdateInfoData {
current_version: config.version.clone(),
latest_version: config.version.clone(),
has_update: false,
deployment_mode: config.deployment_mode.clone(),
deployment_mode_label: deployment_mode_label(&config.deployment_mode).to_string(),
build_type: config.build_type.clone(),
build_type_label: build_type_label(&config.build_type).to_string(),
release_url: None,
notes: None,
cached: false,
update_supported: false,
unsupported_reason: Some(reason),
warning: None,
};
}
let cache_key = config.release_cache_key();
if !force {
if let Some(info) = cached_release_info(&cache_key).await {
return info;
}
}
match fetch_latest_release(&config.github_api_base, repository).await {
Ok(release) => {
let info = update_info_from_release(config, release);
cache_release_info(cache_key, &info).await;
info
}
Err(error) => cached_release_info(&cache_key)
.await
.unwrap_or_else(|| UpdateInfoData {
current_version: config.version.clone(),
latest_version: config.version.clone(),
has_update: false,
deployment_mode: config.deployment_mode.clone(),
deployment_mode_label: deployment_mode_label(&config.deployment_mode).to_string(),
build_type: config.build_type.clone(),
build_type_label: build_type_label(&config.build_type).to_string(),
release_url: None,
notes: None,
cached: false,
update_supported: false,
unsupported_reason: config.update_support_error(),
warning: Some(error),
}),
}
}
pub(super) async fn fetch_latest_release(
api_base: &str,
repository: &str,
) -> Result<GitHubRelease, String> {
let url = format!(
"{}/{repository}/releases/latest",
api_base.trim_end_matches('/')
);
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(30))
.build()
.map_err(|error| error.to_string())?;
let response = client
.get(url)
.header(reqwest::header::ACCEPT, "application/vnd.github+json")
.header(reqwest::header::USER_AGENT, APP_BINARY_NAME)
.send()
.await
.map_err(|error| error.to_string())?;
let status = response.status();
if !status.is_success() {
return Err(format!("GitHub release check failed with {status}"));
}
response
.json::<GitHubRelease>()
.await
.map_err(|error| error.to_string())
}
pub(super) fn update_info_from_release(
config: &SystemUpdateConfig,
release: GitHubRelease,
) -> UpdateInfoData {
let latest_version = normalize_version_tag(&release.tag_name);
let has_update = release_allowed(config, &release)
&& version_is_newer(&config.version, &latest_version).unwrap_or(false);
let unsupported_reason = config.update_support_error();
let update_supported = unsupported_reason.is_none();
let notes = release.body.or(release.name);
UpdateInfoData {
current_version: config.version.clone(),
latest_version,
has_update,
deployment_mode: config.deployment_mode.clone(),
deployment_mode_label: deployment_mode_label(&config.deployment_mode).to_string(),
build_type: config.build_type.clone(),
build_type_label: build_type_label(&config.build_type).to_string(),
release_url: release.html_url,
notes,
cached: false,
update_supported,
unsupported_reason,
warning: None,
}
}
pub(super) fn select_release_archive<'a>(
release: &'a GitHubRelease,
version: &str,
) -> Result<&'a GitHubAsset, AdminError> {
let os_aliases = platform_os_aliases();
let arch_aliases = platform_arch_aliases();
let normalized = normalize_version_tag(version);
release
.assets
.iter()
.find(|asset| {
let name = asset.name.as_str();
name.contains(APP_BINARY_NAME)
&& name.contains(&normalized)
&& asset_matches_current_platform(name, os_aliases, arch_aliases)
&& !name.ends_with(".txt")
})
.or_else(|| {
release.assets.iter().find(|asset| {
let name = asset.name.as_str();
asset_matches_current_platform(name, os_aliases, arch_aliases)
&& !name.ends_with(".txt")
})
})
.ok_or_else(|| {
AdminError::conflict(format!(
"No compatible release archive found for {}/{}",
env::consts::OS,
env::consts::ARCH
))
})
}
async fn cached_release_info(cache_key: &str) -> Option<UpdateInfoData> {
let cache = RELEASE_CACHE.lock().await;
let cached = cache.as_ref()?;
if cached.key != cache_key || cached.cached_at.elapsed() > UPDATE_CACHE_TTL {
return None;
}
let mut info = cached.info.clone();
info.cached = true;
Some(info)
}
async fn cache_release_info(cache_key: String, info: &UpdateInfoData) {
let mut cache = RELEASE_CACHE.lock().await;
*cache = Some(CachedUpdateInfo {
key: cache_key,
info: info.clone(),
cached_at: Instant::now(),
});
}
fn release_allowed(config: &SystemUpdateConfig, release: &GitHubRelease) -> bool {
config.update_channel != "stable" || !release.prerelease
}
fn version_is_newer(current: &str, latest: &str) -> Option<bool> {
let current = Version::parse(&normalize_version_tag(current)).ok()?;
let latest = Version::parse(&normalize_version_tag(latest)).ok()?;
Some(latest > current)
}
fn platform_os_aliases() -> &'static [&'static str] {
match env::consts::OS {
"macos" => &["macos", "darwin"],
"windows" => &["windows", "win32"],
"linux" => &["linux"],
_ => &[env::consts::OS],
}
}
fn platform_arch_aliases() -> &'static [&'static str] {
match env::consts::ARCH {
"x86_64" => &["x86_64", "amd64"],
"aarch64" => &["aarch64", "arm64"],
_ => &[env::consts::ARCH],
}
}
fn asset_matches_current_platform(name: &str, os_aliases: &[&str], arch_aliases: &[&str]) -> bool {
os_aliases.iter().any(|os| name.contains(os))
&& arch_aliases.iter().any(|arch| name.contains(arch))
}
fn normalize_version_tag(version: &str) -> String {
version.trim().trim_start_matches('v').to_string()
}
backend/src/admin/update_payload.rs
+62
View File
@@ -0,0 +1,62 @@
//! 管理端通用更新请求解析。
use serde::Deserialize;
use serde_json::Value;
use crate::admin::response::AdminError;
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct EditableUpdateRequest {
id: String,
label: Option<Option<String>>,
status: Option<String>,
}
pub(crate) struct ParsedEditableUpdate {
pub id: String,
pub label: Option<Option<String>>,
pub status: Option<String>,
}
pub(crate) struct EditableUpdateMessages<'a> {
pub object_required: &'a str,
pub invalid: &'a str,
pub empty_update: &'a str,
pub unknown_field_editable: bool,
}
pub(crate) fn parse_editable_update(
payload: &Value,
messages: EditableUpdateMessages<'_>,
) -> Result<ParsedEditableUpdate, AdminError> {
let object = payload
.as_object()
.ok_or_else(|| AdminError::bad_request(messages.object_required))?;
if messages.unknown_field_editable {
for field in object.keys() {
if !matches!(field.as_str(), "id" | "label" | "status") {
return Err(AdminError::bad_request(format!("{field} is not editable")));
}
}
}
let update = serde_json::from_value::<EditableUpdateRequest>(payload.clone())
.map_err(|_| AdminError::bad_request(messages.invalid))?;
if is_blank(&update.id) || update.status.as_deref().is_some_and(is_blank) {
return Err(AdminError::bad_request(messages.invalid));
}
if update.label.is_none() && update.status.is_none() {
return Err(AdminError::bad_request(messages.empty_update));
}
Ok(ParsedEditableUpdate {
id: update.id,
label: update.label,
status: update.status,
})
}
fn is_blank(value: &str) -> bool {
value.trim().is_empty()
}
backend/src/config/loader.rs
+19 -7
View File
@@ -1,19 +1,31 @@
use std::path::Path;
use std::{env, path::Path};
use super::types::AppConfig;
const CONFIG_FILE_ENV: &str = "CPR_CONFIG_FILE";
impl AppConfig {
/// 从当前目录加载配置文件。
/// 从 `CPR_CONFIG_FILE` 或当前目录 `config.yaml` 加载配置文件。
pub fn load() -> Result<Self, ::config::ConfigError> {
Self::load_from_dir(".")
let config_file = env::var(CONFIG_FILE_ENV)
.ok()
.filter(|value| !value.trim().is_empty());
match config_file {
Some(path) => load_file(path),
None => Self::load_from_dir("."),
}
}
/// 从指定目录加载配置文件。
pub fn load_from_dir(config_dir: impl AsRef<Path>) -> Result<Self, ::config::ConfigError> {
let config_dir = config_dir.as_ref();
::config::Config::builder()
.add_source(::config::File::from(config_dir.join("config.yaml")).required(true))
.build()?
.try_deserialize()
load_file(config_dir.join("config.yaml"))
}
}
fn load_file(config_file: impl AsRef<Path>) -> Result<AppConfig, ::config::ConfigError> {
::config::Config::builder()
.add_source(::config::File::from(config_file.as_ref()).required(true))
.build()?
.try_deserialize()
}
backend/src/config/types.rs
+40 -1
View File
@@ -2,6 +2,17 @@ use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
const WEAK_ADMIN_PASSWORDS: &[&str] = &[
"",
"admin",
"123456",
"password",
"changeme",
"change-me",
"replace-me",
"codex-proxy-rs",
];
/// 应用运行总配置。
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
@@ -43,11 +54,15 @@ pub struct AppConfig {
/// HTTP 监听配置。
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct ServerConfig {
/// 监听主机。
pub host: String,
/// 监听端口。
pub port: u16,
/// 可信反向代理 IP 或 CIDR。只有这些 peer 的转发 IP 头会被采信。
#[serde(default)]
pub trusted_proxies: Vec<String>,
}
/// 上游 API 地址配置。
@@ -276,6 +291,30 @@ impl Default for AdminConfig {
}
}
impl AdminConfig {
/// 校验首次启动创建管理员时使用的默认密码。
pub fn validate_default_password(&self) -> Result<(), AdminConfigError> {
let password = self.default_password.trim();
if password.len() < 12 {
return Err(AdminConfigError::WeakDefaultPassword);
}
if WEAK_ADMIN_PASSWORDS.contains(&password.to_ascii_lowercase().as_str()) {
return Err(AdminConfigError::WeakDefaultPassword);
}
Ok(())
}
}
/// 管理员配置校验错误。
#[derive(Debug, Clone, Copy, thiserror::Error, PartialEq, Eq)]
pub enum AdminConfigError {
/// 默认管理员密码为空、过短或属于常见弱口令。
#[error(
"admin.default_password must be set to a non-default password with at least 12 characters"
)]
WeakDefaultPassword,
}
fn default_session_cleanup_interval() -> u64 {
3600 // 默认每小时清理一次过期会话
}
@@ -285,7 +324,7 @@ fn default_admin_username() -> String {
}
fn default_admin_password() -> String {
"admin".to_string()
String::new()
}
/// 日志持久化配置。
backend/src/http/middleware/request_id.rs
+142 -16
View File
@@ -1,17 +1,56 @@
//! 请求上下文中间件。
use std::net::SocketAddr;
use std::net::{IpAddr, SocketAddr};
use axum::{
extract::{connect_info::ConnectInfo, Request},
http::HeaderValue,
extract::{connect_info::ConnectInfo, Request, State},
http::{HeaderMap, HeaderValue},
middleware::Next,
response::Response,
};
use ipnet::IpNet;
use uuid::Uuid;
/// 请求 ID 头。
pub const REQUEST_ID_HEADER: &str = "x-request-id";
const MAX_REQUEST_ID_LEN: usize = 128;
const CF_CONNECTING_IP_HEADER: &str = "cf-connecting-ip";
const REAL_IP_HEADER: &str = "x-real-ip";
const X_FORWARDED_FOR_HEADER: &str = "x-forwarded-for";
const UNTRUSTED_CLIENT_IP_HEADERS: &[&str] = &[
CF_CONNECTING_IP_HEADER,
REAL_IP_HEADER,
X_FORWARDED_FOR_HEADER,
];
/// 可信反向代理配置。
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct TrustedProxyConfig {
networks: Vec<IpNet>,
}
impl TrustedProxyConfig {
/// 从 IP 或 CIDR 字符串构造可信代理配置。
pub fn from_entries(entries: &[String]) -> Result<Self, TrustedProxyConfigError> {
let networks = entries
.iter()
.map(|entry| parse_trusted_proxy_entry(entry))
.collect::<Result<Vec<_>, _>>()?;
Ok(Self { networks })
}
fn is_trusted(&self, ip: IpAddr) -> bool {
self.networks.iter().any(|network| network.contains(&ip))
}
}
/// 可信代理配置错误。
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum TrustedProxyConfigError {
/// 配置项不是 IP 地址或 CIDR。
#[error("invalid trusted proxy entry `{entry}`; expected an IP address or CIDR network")]
InvalidEntry { entry: String },
}
/// 请求 ID。
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -26,7 +65,7 @@ impl RequestId {
/// 从 HTTP 头解析请求 ID。
pub fn from_header(value: &HeaderValue) -> Option<Self> {
let value = value.to_str().ok()?.trim();
if value.is_empty() {
if !valid_request_id(value) {
return None;
}
Some(Self(value.to_string()))
@@ -38,15 +77,37 @@ impl RequestId {
}
}
fn valid_request_id(value: &str) -> bool {
!value.is_empty()
&& value.len() <= MAX_REQUEST_ID_LEN
&& value
.bytes()
.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-' | b'.'))
}
/// 为请求附加请求 ID，并在响应头中回写。
pub async fn attach_request_id(mut request: Request, next: Next) -> Response {
attach_request_context(&TrustedProxyConfig::default(), &mut request);
run_with_request_id(request, next).await
}
/// 为请求附加请求上下文，并按可信代理配置解析真实客户端 IP。
pub async fn attach_request_id_with_proxy_config(
State(trusted_proxies): State<TrustedProxyConfig>,
mut request: Request,
next: Next,
) -> Response {
attach_request_context(&trusted_proxies, &mut request);
run_with_request_id(request, next).await
}
async fn run_with_request_id(mut request: Request, next: Next) -> Response {
let request_id = request
.headers()
.get(REQUEST_ID_HEADER)
.and_then(RequestId::from_header)
.unwrap_or_else(RequestId::generate);
attach_real_ip_from_connection(&mut request);
request.extensions_mut().insert(request_id.clone());
let mut response = next.run(request).await;
if let Ok(value) = HeaderValue::from_str(request_id.as_str()) {
@@ -55,19 +116,84 @@ pub async fn attach_request_id(mut request: Request, next: Next) -> Response {
response
}
fn attach_real_ip_from_connection(request: &mut Request) {
if request.headers().contains_key("cf-connecting-ip")
|| request.headers().contains_key("x-real-ip")
|| request.headers().contains_key("x-forwarded-for")
{
return;
fn attach_request_context(trusted_proxies: &TrustedProxyConfig, request: &mut Request) {
attach_real_ip_from_connection(trusted_proxies, request);
}
fn attach_real_ip_from_connection(trusted_proxies: &TrustedProxyConfig, request: &mut Request) {
let peer_ip = request
.extensions()
.get::<ConnectInfo<SocketAddr>>()
.map(|ConnectInfo(addr)| addr.ip());
let client_ip =
peer_ip.map(|peer_ip| client_ip_from_peer(peer_ip, request.headers(), trusted_proxies));
for header in UNTRUSTED_CLIENT_IP_HEADERS {
request.headers_mut().remove(*header);
}
let Some(ConnectInfo(addr)) = request.extensions().get::<ConnectInfo<SocketAddr>>() else {
let Some(client_ip) = client_ip else {
return;
};
let Ok(value) = HeaderValue::from_str(&addr.ip().to_string()) else {
return;
};
request.headers_mut().insert("x-real-ip", value);
if let Ok(value) = HeaderValue::from_str(&client_ip.to_string()) {
request.headers_mut().insert(REAL_IP_HEADER, value);
}
}
fn client_ip_from_peer(
peer_ip: IpAddr,
headers: &HeaderMap,
trusted_proxies: &TrustedProxyConfig,
) -> IpAddr {
if trusted_proxies.is_trusted(peer_ip) {
trusted_forwarded_client_ip(headers).unwrap_or(peer_ip)
} else {
peer_ip
}
}
fn trusted_forwarded_client_ip(headers: &HeaderMap) -> Option<IpAddr> {
header_ip(headers, CF_CONNECTING_IP_HEADER)
.or_else(|| first_forwarded_for_ip(headers))
.or_else(|| header_ip(headers, REAL_IP_HEADER))
}
fn header_ip(headers: &HeaderMap, name: &str) -> Option<IpAddr> {
parse_ip(headers.get(name)?.to_str().ok()?)
}
fn first_forwarded_for_ip(headers: &HeaderMap) -> Option<IpAddr> {
let value = headers.get(X_FORWARDED_FOR_HEADER)?.to_str().ok()?;
parse_ip(value.split(',').next()?)
}
fn parse_ip(value: &str) -> Option<IpAddr> {
value.trim().parse().ok()
}
fn parse_trusted_proxy_entry(entry: &str) -> Result<IpNet, TrustedProxyConfigError> {
let entry = entry.trim();
if entry.is_empty() {
return Err(invalid_trusted_proxy_entry(entry));
}
if let Ok(network) = entry.parse::<IpNet>() {
return Ok(network);
}
let ip = entry
.parse::<IpAddr>()
.map_err(|_| invalid_trusted_proxy_entry(entry))?;
IpNet::new(
ip,
match ip {
IpAddr::V4(_) => 32,
IpAddr::V6(_) => 128,
},
)
.map_err(|_| invalid_trusted_proxy_entry(entry))
}
fn invalid_trusted_proxy_entry(entry: &str) -> TrustedProxyConfigError {
TrustedProxyConfigError::InvalidEntry {
entry: entry.to_string(),
}
}
backend/src/http/middleware/trace.rs
+35 -2
View File
@@ -36,20 +36,22 @@ pub fn http_trace_layer() -> HttpTraceLayer {
fn make_http_span(request: &Request<Body>) -> Span {
let request_id = request_id(request);
let uri = sanitized_uri(request);
tracing::info_span!(
"http",
request_id = request_id.as_deref(),
method = %request.method(),
uri = %request.uri()
uri = %uri
)
}
fn on_http_request(request: &Request<Body>, _span: &Span) {
let request_id = request_id(request);
let uri = sanitized_uri(request);
info!(
request_id = request_id.as_deref(),
method = %request.method(),
uri = %request.uri(),
uri = %uri,
"received HTTP request"
);
}
@@ -80,3 +82,34 @@ fn request_id(request: &Request<Body>) -> Option<String> {
.get::<RequestId>()
.map(|rid| rid.as_str().to_string())
}
fn sanitized_uri(request: &Request<Body>) -> String {
let uri = request.uri();
let path = uri.path();
uri.query().map_or_else(
|| path.to_string(),
|query| format!("{path}?{}", sanitized_query(query)),
)
}
fn sanitized_query(query: &str) -> String {
let mut redacted = String::new();
for (index, part) in query.split('&').enumerate() {
if index > 0 {
redacted.push('&');
}
redacted.push_str(&sanitized_query_part(part));
}
redacted
}
fn sanitized_query_part(part: &str) -> String {
let Some((key, _value)) = part.split_once('=') else {
return "<redacted>".to_string();
};
if key.is_empty() {
"<redacted>=<redacted>".to_string()
} else {
format!("{key}=<redacted>")
}
}
backend/src/http/router.rs
+37 -4
View File
@@ -2,10 +2,13 @@
use std::path::Path;
use axum::{middleware, Router};
use axum::{extract::State, http::StatusCode, middleware, routing::get, Router};
use crate::admin;
use crate::http::middleware::{request_id::attach_request_id, trace::http_trace_layer};
use crate::http::middleware::{
request_id::{attach_request_id_with_proxy_config, TrustedProxyConfig},
trace::http_trace_layer,
};
use crate::proxy;
use crate::runtime::state::AppState;
use crate::web;
@@ -15,15 +18,45 @@ pub const DEFAULT_ASSET_DIST_DIR: &str = "web/dist";
/// 构造整个 HTTP 服务路由。
pub fn router() -> Router<AppState> {
router_with_assets(DEFAULT_ASSET_DIST_DIR)
router_with_trusted_proxies(TrustedProxyConfig::default())
}
/// 使用指定可信代理配置构造整个 HTTP 服务路由。
pub fn router_with_trusted_proxies(trusted_proxies: TrustedProxyConfig) -> Router<AppState> {
router_with_assets_and_trusted_proxies(DEFAULT_ASSET_DIST_DIR, trusted_proxies)
}
/// 使用指定前端构建产物目录构造整个 HTTP 服务路由。
pub fn router_with_assets(dist_dir: impl AsRef<Path>) -> Router<AppState> {
router_with_assets_and_trusted_proxies(dist_dir, TrustedProxyConfig::default())
}
/// 使用指定前端构建产物目录和可信代理配置构造整个 HTTP 服务路由。
pub fn router_with_assets_and_trusted_proxies(
dist_dir: impl AsRef<Path>,
trusted_proxies: TrustedProxyConfig,
) -> Router<AppState> {
Router::new()
.route("/healthz", get(healthz))
.merge(proxy::router::router())
.merge(admin::router::router())
.fallback_service(web::assets::spa_router(dist_dir))
.layer(http_trace_layer())
.layer(middleware::from_fn(attach_request_id))
.layer(middleware::from_fn_with_state(
trusted_proxies,
attach_request_id_with_proxy_config,
))
}
async fn healthz(State(state): State<AppState>) -> StatusCode {
match sqlx::query("select 1")
.execute(state.services.background_tasks.client_keys.pool())
.await
{
Ok(_) => StatusCode::NO_CONTENT,
Err(error) => {
tracing::warn!(error = %error, "health check failed");
StatusCode::SERVICE_UNAVAILABLE
}
}
}
backend/src/infra/database.rs
+141 -4
View File
@@ -1,10 +1,23 @@
//! SQLite 数据库连接与 Schema 初始化。
//! SQLite 数据库连接与迁移初始化。
use std::{fs, path::Path, str::FromStr, time::Duration};
use sqlx::{sqlite::SqliteConnectOptions, SqlitePool};
/// 连接 SQLite 并初始化 schema。
const CURRENT_SCHEMA_VERSION: i64 = 1;
const MIGRATIONS: &[Migration] = &[Migration {
version: 1,
name: "initial",
sql: include_str!("migrations/0001_initial.sql"),
}];
struct Migration {
version: i64,
name: &'static str,
sql: &'static str,
}
/// 连接 SQLite 并应用未执行的迁移。
pub async fn connect_sqlite(database_url: &str) -> Result<SqlitePool, sqlx::Error> {
ensure_sqlite_parent_dir(database_url)?;
let options = SqliteConnectOptions::from_str(database_url)?
@@ -21,11 +34,135 @@ pub async fn connect_sqlite(database_url: &str) -> Result<SqlitePool, sqlx::Erro
}
async fn initialize_schema(pool: &SqlitePool) -> Result<(), sqlx::Error> {
let schema = include_str!("schema.sql");
sqlx::raw_sql(schema).execute(pool).await?;
let mut tx = pool.begin().await?;
ensure_migrations_table(&mut tx).await?;
let applied_versions = applied_migration_versions(&mut tx).await?;
if applied_versions.is_empty() {
reject_unversioned_existing_schema(&mut tx).await?;
}
validate_applied_migrations(&applied_versions)?;
for migration in MIGRATIONS {
if applied_versions.contains(&migration.version) {
continue;
}
sqlx::raw_sql(migration.sql).execute(&mut *tx).await?;
record_migration(&mut tx, migration).await?;
}
tx.commit().await?;
Ok(())
}
async fn ensure_migrations_table(
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
) -> Result<(), sqlx::Error> {
if table_exists(tx, "schema_migrations").await? {
return Ok(());
}
sqlx::query(
"create table schema_migrations (
version integer not null check (version > 0),
name text not null,
applied_at text not null,
primary key (version)
)",
)
.execute(&mut **tx)
.await?;
Ok(())
}
async fn table_exists(
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
table_name: &str,
) -> Result<bool, sqlx::Error> {
let row: Option<(i64,)> = sqlx::query_as(
"select 1
from sqlite_master
where type = 'table' and name = ?
limit 1",
)
.bind(table_name)
.fetch_optional(&mut **tx)
.await?;
Ok(row.is_some())
}
async fn applied_migration_versions(
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
) -> Result<Vec<i64>, sqlx::Error> {
let rows: Vec<(i64,)> =
sqlx::query_as("select version from schema_migrations order by version")
.fetch_all(&mut **tx)
.await?;
Ok(rows.into_iter().map(|row| row.0).collect())
}
fn validate_applied_migrations(applied_versions: &[i64]) -> Result<(), sqlx::Error> {
for version in applied_versions {
if MIGRATIONS
.iter()
.any(|migration| migration.version == *version)
{
continue;
}
return Err(unsupported_schema_version(*version));
}
if let Some(version) = applied_versions
.last()
.copied()
.filter(|version| *version > CURRENT_SCHEMA_VERSION)
{
return Err(unsupported_schema_version(version));
}
Ok(())
}
async fn reject_unversioned_existing_schema(
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
) -> Result<(), sqlx::Error> {
let (table_count,): (i64,) = sqlx::query_as(
"select count(*)
from sqlite_master
where type = 'table'
and name <> 'schema_migrations'
and name not like 'sqlite_%'",
)
.fetch_one(&mut **tx)
.await?;
if table_count > 0 {
return Err(sqlx::Error::Protocol(
"unversioned sqlite schema is unsupported; recreate the database or add a migration"
.to_string(),
));
}
Ok(())
}
async fn record_migration(
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
migration: &Migration,
) -> Result<(), sqlx::Error> {
sqlx::query(
"insert into schema_migrations (version, name, applied_at)
values (?, ?, strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))",
)
.bind(migration.version)
.bind(migration.name)
.execute(&mut **tx)
.await?;
Ok(())
}
fn unsupported_schema_version(version: i64) -> sqlx::Error {
sqlx::Error::Protocol(format!(
"unsupported sqlite schema version {version}; expected {CURRENT_SCHEMA_VERSION}"
))
}
fn ensure_sqlite_parent_dir(database_url: &str) -> Result<(), sqlx::Error> {
let Some(path) = sqlite_file_path(database_url) else {
return Ok(());
backend/src/infra/format.rs
+5
View File
@@ -5,6 +5,11 @@ pub fn nonnegative_i64_to_u64(value: i64) -> u64 {
u64::try_from(value).unwrap_or(0)
}
/// 将可空负数归零后转换为无符号展示数值。
pub fn optional_nonnegative_i64_to_u64(value: Option<i64>) -> u64 {
value.map(nonnegative_i64_to_u64).unwrap_or_default()
}
/// 使用中文界面常用千分位格式展示整数。
pub fn format_number(value: u64) -> String {
let text = value.to_string();
backend/src/infra/schema.sql → backend/src/infra/migrations/0001_initial.sql
+49 -49
View File
@@ -3,7 +3,7 @@ pragma foreign_keys = on;
-- ============================================
-- Admin Users
-- ============================================
create table if not exists admin_users (
create table admin_users (
id text primary key,
password_hash text not null,
created_at text not null,
@@ -13,19 +13,19 @@ create table if not exists admin_users (
-- ============================================
-- Admin Sessions
-- ============================================
create table if not exists admin_sessions (
create table admin_sessions (
id text primary key,
user_id text not null references admin_users(id) on delete cascade,
expires_at text not null,
created_at text not null
);
create index if not exists idx_admin_sessions_expires on admin_sessions(expires_at);
create index idx_admin_sessions_expires on admin_sessions(expires_at);
-- ============================================
-- Client API Keys
-- ============================================
create table if not exists client_api_keys (
create table client_api_keys (
id text primary key,
name text not null,
prefix text not null,
@@ -36,13 +36,13 @@ create table if not exists client_api_keys (
last_used_at text
);
create index if not exists idx_client_api_keys_key_enabled on client_api_keys(key) where enabled = 1;
create index if not exists idx_client_api_keys_created_id on client_api_keys(created_at desc, id desc);
create index idx_client_api_keys_key_enabled on client_api_keys(key) where enabled = 1;
create index idx_client_api_keys_created_id on client_api_keys(created_at desc, id desc);
-- ============================================
-- Runtime Settings
-- ============================================
create table if not exists runtime_settings (
create table runtime_settings (
id integer primary key check (id = 1),
model_aliases_json text not null default '{}',
refresh_margin_seconds integer not null check (refresh_margin_seconds > 0),
@@ -57,7 +57,7 @@ create table if not exists runtime_settings (
-- ============================================
-- Accounts
-- ============================================
create table if not exists accounts (
create table accounts (
id text primary key,
email text,
chatgpt_account_id text,
@@ -79,28 +79,28 @@ create table if not exists accounts (
updated_at text not null
);
create index if not exists idx_accounts_status on accounts(status);
create index if not exists idx_accounts_added_id on accounts(added_at desc, id desc);
create unique index if not exists ux_accounts_chatgpt_identity
create index idx_accounts_status on accounts(status);
create index idx_accounts_added_id on accounts(added_at desc, id desc);
create unique index ux_accounts_chatgpt_identity
on accounts(chatgpt_account_id, coalesce(chatgpt_user_id, ''))
where chatgpt_account_id is not null;
-- ============================================
-- Account Refresh Leases
-- ============================================
create table if not exists account_refresh_leases (
create table account_refresh_leases (
account_id text primary key references accounts(id) on delete cascade,
owner text not null,
expires_at text not null,
updated_at text not null
);
create index if not exists idx_account_refresh_leases_expires on account_refresh_leases(expires_at);
create index idx_account_refresh_leases_expires on account_refresh_leases(expires_at);
-- ============================================
-- Account Usage Statistics
-- ============================================
create table if not exists account_usage (
create table account_usage (
account_id text primary key references accounts(id) on delete cascade,
request_count integer not null default 0 check (request_count >= 0),
empty_response_count integer not null default 0 check (empty_response_count >= 0),
@@ -127,13 +127,13 @@ create table if not exists account_usage (
last_used_at text
);
create index if not exists idx_account_usage_last_used_account
create index idx_account_usage_last_used_account
on account_usage(last_used_at desc, account_id desc);
-- ============================================
-- Account Model Usage Statistics
-- ============================================
create table if not exists account_model_usage (
create table account_model_usage (
account_id text not null references accounts(id) on delete cascade,
model text not null,
request_count integer not null default 0 check (request_count >= 0),
@@ -145,13 +145,13 @@ create table if not exists account_model_usage (
primary key (account_id, model)
);
create index if not exists idx_account_model_usage_last_used
create index idx_account_model_usage_last_used
on account_model_usage(last_used_at desc, account_id, model);
-- ============================================
-- Usage Time Buckets
-- ============================================
create table if not exists usage_time_buckets (
create table usage_time_buckets (
bucket_start text not null,
account_id text not null default '',
model text not null default '',
@@ -171,15 +171,15 @@ create table if not exists usage_time_buckets (
primary key (bucket_start, account_id, model, service_tier)
);
create index if not exists idx_usage_time_buckets_bucket
create index idx_usage_time_buckets_bucket
on usage_time_buckets(bucket_start);
create index if not exists idx_usage_time_buckets_model_bucket
create index idx_usage_time_buckets_model_bucket
on usage_time_buckets(model, bucket_start);
-- ============================================
-- Model Account Routes
-- ============================================
create table if not exists model_account_routes (
create table model_account_routes (
model text not null,
account_id text not null references accounts(id) on delete cascade,
priority integer not null default 0 check (priority >= 0),
@@ -189,15 +189,15 @@ create table if not exists model_account_routes (
primary key (model, account_id)
);
create index if not exists idx_model_account_routes_account
create index idx_model_account_routes_account
on model_account_routes(account_id);
create index if not exists idx_model_account_routes_enabled_model
create index idx_model_account_routes_enabled_model
on model_account_routes(enabled, model, priority asc, account_id);
-- ============================================
-- Account Cookies
-- ============================================
create table if not exists account_cookies (
create table account_cookies (
id text primary key,
account_id text not null references accounts(id) on delete cascade,
domain text not null,
@@ -209,16 +209,16 @@ create table if not exists account_cookies (
unique(account_id, domain, name, path)
);
create index if not exists idx_account_cookies_account on account_cookies(account_id);
create index if not exists idx_account_cookies_account_domain on account_cookies(account_id, domain);
create index if not exists idx_account_cookies_expires
create index idx_account_cookies_account on account_cookies(account_id);
create index idx_account_cookies_account_domain on account_cookies(account_id, domain);
create index idx_account_cookies_expires
on account_cookies(expires_at)
where expires_at is not null;
-- ============================================
-- Device Fingerprints
-- ============================================
create table if not exists fingerprints (
create table fingerprints (
id text primary key,
originator text not null,
app_version text not null,
@@ -234,7 +234,7 @@ create table if not exists fingerprints (
updated_at text not null
);
create table if not exists fingerprint_update_history (
create table fingerprint_update_history (
id text primary key,
current_fingerprint_id text not null references fingerprints(id) on delete cascade,
app_version text not null,
@@ -245,13 +245,13 @@ create table if not exists fingerprint_update_history (
created_at text not null
);
create index if not exists idx_fingerprint_update_history_created_id
create index idx_fingerprint_update_history_created_id
on fingerprint_update_history(created_at desc, id desc);
-- ============================================
-- Usage Records
-- ============================================
create table if not exists usage_records (
create table usage_records (
id text primary key,
request_id text,
kind text not null,
@@ -272,24 +272,24 @@ create table if not exists usage_records (
created_at text not null
);
create index if not exists idx_usage_records_created_id on usage_records(created_at desc, id desc);
create index if not exists idx_usage_records_kind_created on usage_records(kind, created_at desc);
create index if not exists idx_usage_records_request_id on usage_records(request_id);
create index if not exists idx_usage_records_account on usage_records(account_id, created_at desc) where account_id is not null;
create index if not exists idx_usage_records_transport on usage_records(transport, created_at desc) where transport is not null;
create index if not exists idx_usage_records_failure_class on usage_records(failure_class, created_at desc) where failure_class is not null;
create index if not exists idx_usage_records_response_id on usage_records(response_id) where response_id is not null;
create index if not exists idx_usage_records_upstream_request_id on usage_records(upstream_request_id) where upstream_request_id is not null;
create index if not exists idx_usage_records_level_created on usage_records(level, created_at desc);
create index if not exists idx_usage_records_route_created on usage_records(route, created_at desc) where route is not null;
create index if not exists idx_usage_records_model_created on usage_records(model, created_at desc) where model is not null;
create index if not exists idx_usage_records_status_created on usage_records(status_code, created_at desc) where status_code is not null;
create index if not exists idx_usage_records_upstream_status_created on usage_records(upstream_status_code, created_at desc) where upstream_status_code is not null;
create index idx_usage_records_created_id on usage_records(created_at desc, id desc);
create index idx_usage_records_kind_created on usage_records(kind, created_at desc);
create index idx_usage_records_request_id on usage_records(request_id);
create index idx_usage_records_account on usage_records(account_id, created_at desc) where account_id is not null;
create index idx_usage_records_transport on usage_records(transport, created_at desc) where transport is not null;
create index idx_usage_records_failure_class on usage_records(failure_class, created_at desc) where failure_class is not null;
create index idx_usage_records_response_id on usage_records(response_id) where response_id is not null;
create index idx_usage_records_upstream_request_id on usage_records(upstream_request_id) where upstream_request_id is not null;
create index idx_usage_records_level_created on usage_records(level, created_at desc);
create index idx_usage_records_route_created on usage_records(route, created_at desc) where route is not null;
create index idx_usage_records_model_created on usage_records(model, created_at desc) where model is not null;
create index idx_usage_records_status_created on usage_records(status_code, created_at desc) where status_code is not null;
create index idx_usage_records_upstream_status_created on usage_records(upstream_status_code, created_at desc) where upstream_status_code is not null;
-- ============================================
-- Model Plan Snapshots
-- ============================================
create table if not exists model_plan_snapshots (
create table model_plan_snapshots (
plan_type text primary key,
models_json text not null,
fetched_at text not null
@@ -298,7 +298,7 @@ create table if not exists model_plan_snapshots (
-- ============================================
-- Session Affinities
-- ============================================
create table if not exists session_affinities (
create table session_affinities (
response_id text primary key,
account_id text not null references accounts(id) on delete cascade,
conversation_id text not null,
@@ -311,6 +311,6 @@ create table if not exists session_affinities (
created_at text not null
);
create index if not exists idx_session_affinities_conversation on session_affinities(conversation_id, created_at desc);
create index if not exists idx_session_affinities_expires on session_affinities(expires_at);
create index if not exists idx_session_affinities_active_order on session_affinities(expires_at, created_at, response_id);
create index idx_session_affinities_conversation on session_affinities(conversation_id, created_at desc);
create index idx_session_affinities_expires on session_affinities(expires_at);
create index idx_session_affinities_active_order on session_affinities(expires_at, created_at, response_id);
backend/src/infra/time.rs
+19 -8
View File
@@ -1,5 +1,7 @@
//! 时间格式化辅助。
use std::time::Instant;
use chrono::{DateTime, FixedOffset, SecondsFormat, TimeZone, Timelike, Utc};
use serde::Serializer;
@@ -10,6 +12,23 @@ pub fn china_rfc3339(value: &DateTime<Utc>) -> String {
value.with_timezone(&china_offset()).to_rfc3339()
}
/// 将 RFC3339 字符串解析为 UTC 时间。
pub fn parse_rfc3339_utc(value: &str) -> Result<DateTime<Utc>, chrono::ParseError> {
Ok(DateTime::parse_from_rfc3339(value)?.with_timezone(&Utc))
}
/// 将可选 RFC3339 字符串解析为 UTC 时间。
pub fn parse_optional_rfc3339_utc(
value: Option<&str>,
) -> Result<Option<DateTime<Utc>>, chrono::ParseError> {
value.map(parse_rfc3339_utc).transpose()
}
/// 返回从起点到当前的毫秒数，并在超过 i64 上限时饱和。
pub fn elapsed_millis_i64(started_at: Instant) -> i64 {
started_at.elapsed().as_millis().min(i64::MAX as u128) as i64
}
/// 将 UTC 时间输出为中国时区毫秒精度 RFC3339 字符串。
pub fn china_rfc3339_millis(value: &DateTime<Utc>) -> String {
value
@@ -54,14 +73,6 @@ pub fn china_date(value: &DateTime<Utc>) -> String {
.to_string()
}
/// 将 UTC 时间输出为中国时区时间。
pub fn china_time(value: &DateTime<Utc>) -> String {
value
.with_timezone(&china_offset())
.format("%H:%M:%S")
.to_string()
}
/// 将 UTC 时间输出为适合文件名的中国时区时间戳。
pub fn china_filename_timestamp_millis(value: &DateTime<Utc>) -> String {
value
backend/src/proxy/auth.rs
+84 -12
View File
@@ -2,30 +2,102 @@
//!
//! 提供 Bearer token 提取与客户端 API key 校验。
use axum::http::HeaderMap;
use axum::http::{header::AUTHORIZATION, HeaderMap};
use crate::runtime::state::AppState;
/// client API key 鉴权失败原因。
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ClientApiKeyAuthError {
/// 缺失 Authorization 头。
MissingAuthorization,
/// Authorization 头不是合法的 Bearer token。
MalformedAuthorization,
/// Bearer token 不是 client API key 格式。
InvalidKeyFormat,
/// client API key 不存在或已禁用。
InvalidKey,
/// 鉴权存储失败。
StoreUnavailable { message: String },
}
impl ClientApiKeyAuthError {
/// 返回可用于日志和指标的稳定失败原因。
pub fn reason(&self) -> &'static str {
match self {
Self::MissingAuthorization => "missing_authorization",
Self::MalformedAuthorization => "malformed_authorization",
Self::InvalidKeyFormat => "invalid_key_format",
Self::InvalidKey => "invalid_key",
Self::StoreUnavailable { .. } => "store_unavailable",
}
}
}
/// 从请求头提取 Bearer 客户端 API key。
pub fn bearer_client_api_key(headers: &HeaderMap) -> Option<String> {
pub fn bearer_client_api_key(headers: &HeaderMap) -> Result<&str, ClientApiKeyAuthError> {
let raw = headers
.get(axum::http::header::AUTHORIZATION)?
.get(AUTHORIZATION)
.ok_or(ClientApiKeyAuthError::MissingAuthorization)?
.to_str()
.ok()?;
let token = raw.strip_prefix("Bearer ")?.trim();
(!token.is_empty()).then(|| token.to_string())
.map_err(|_| ClientApiKeyAuthError::MalformedAuthorization)?;
let token = raw
.strip_prefix("Bearer ")
.ok_or(ClientApiKeyAuthError::MalformedAuthorization)?
.trim();
if token.is_empty() {
return Err(ClientApiKeyAuthError::MalformedAuthorization);
}
if !token.starts_with("sk_") {
return Err(ClientApiKeyAuthError::InvalidKeyFormat);
}
Ok(token)
}
/// 校验客户端 API key。
pub async fn authorize_client_api_key(state: &AppState, headers: &HeaderMap) -> bool {
let Some(api_key) = bearer_client_api_key(headers) else {
return false;
};
match authorize_client_api_key_result(state, headers).await {
Ok(()) => true,
Err(error) => {
log_client_api_key_auth_failure(&error);
false
}
}
}
state
/// 校验客户端 API key，返回可观测的失败原因。
pub async fn authorize_client_api_key_result(
state: &AppState,
headers: &HeaderMap,
) -> Result<(), ClientApiKeyAuthError> {
let api_key = bearer_client_api_key(headers)?;
let verified = state
.services
.client_keys
.verify(&api_key)
.verify(api_key)
.await
.unwrap_or(false)
.map_err(|error| ClientApiKeyAuthError::StoreUnavailable {
message: error.to_string(),
})?;
verified
.then_some(())
.ok_or(ClientApiKeyAuthError::InvalidKey)
}
fn log_client_api_key_auth_failure(error: &ClientApiKeyAuthError) {
match error {
ClientApiKeyAuthError::StoreUnavailable { message } => {
tracing::warn!(
auth_failure = error.reason(),
error = %message,
"client api key authorization failed"
);
}
_ => {
tracing::warn!(
auth_failure = error.reason(),
"client api key authorization failed"
);
}
}
}
backend/src/proxy/dispatch/auth_recovery.rs
+28
View File
@@ -0,0 +1,28 @@
use std::sync::Arc;
use crate::upstream::{
accounts::{model::AccountStatus, token_refresh::RuntimeTokenRefreshService},
token_client::OpenAiTokenClient,
};
pub(super) fn trigger_refresh_after_auth_failure(
token_refresh: &Arc<RuntimeTokenRefreshService<OpenAiTokenClient>>,
account_id: &str,
status: AccountStatus,
) {
if status != AccountStatus::Expired {
return;
}
let token_refresh = Arc::clone(token_refresh);
let account_id = account_id.to_string();
tokio::spawn(async move {
if let Err(error) = token_refresh.trigger_account_refresh_now(&account_id).await {
tracing::warn!(
account_id = %account_id,
error = %error,
"reactive token refresh after upstream auth failure failed"
);
}
});
}
backend/src/proxy/dispatch/chat.rs
+24 -148
View File
@@ -9,10 +9,11 @@ use thiserror::Error;
use crate::{
admin::monitoring::{
usage_record::{ResponseUsageRecord, UsageRecordLevel},
usage_record_store::AdminUsageRecordService,
usage_record_model::{ResponseUsageRecord, UsageRecordLevel},
usage_record_service::AdminUsageRecordService,
},
proxy::dispatch::{
auth_recovery::trigger_refresh_after_auth_failure,
cloudflare::{
cloudflare_challenge_error_message, cloudflare_path_block_error_message,
is_cloudflare_challenge_upstream_error, is_cloudflare_path_block_upstream_error,
@@ -29,6 +30,10 @@ use crate::{
QuotaVerificationContext, QuotaVerificationDecision,
QUOTA_VERIFY_LIMIT_REACHED_MESSAGE,
},
usage_events::{
reasoning_effort_from_request, record_dispatch_error_event, record_response_event,
DispatchErrorUsageRecord,
},
},
proxy::openai::chat::ChatStreamTranslationError,
upstream::accounts::{
@@ -80,24 +85,6 @@ impl ChatDispatchService {
}
}
fn trigger_refresh_after_auth_failure(&self, account_id: &str, status: AccountStatus) {
if status != AccountStatus::Expired {
return;
}
let token_refresh = self.token_refresh.clone();
let account_id = account_id.to_string();
tokio::spawn(async move {
if let Err(error) = token_refresh.trigger_account_refresh_now(&account_id).await {
tracing::warn!(
account_id = %account_id,
error = %error,
"reactive token refresh after upstream auth failure failed"
);
}
});
}
/// 调度非流式 Chat Completions 请求到 Codex Responses 上游。
pub async fn complete(
&self,
@@ -311,7 +298,11 @@ impl ChatDispatchService {
self.account_pool
.set_status(&release_account_id, account_status)
.await;
self.trigger_refresh_after_auth_failure(&release_account_id, account_status);
trigger_refresh_after_auth_failure(
&self.token_refresh,
&release_account_id,
account_status,
);
excluded_account_ids.push(release_account_id);
}
Err(error) if is_cloudflare_challenge_upstream_error(&error) => {
@@ -554,8 +545,6 @@ struct ChatDispatchErrorEventRecord<'a> {
error: &'a ChatDispatchError,
}
use crate::admin::monitoring::usage_record::UsageRecord;
async fn record_chat_dispatch_error_event(record: ChatDispatchErrorEventRecord<'_>) {
let mut metadata = serde_json::json!({
"route": record.route,
@@ -571,127 +560,18 @@ async fn record_chat_dispatch_error_event(record: ChatDispatchErrorEventRecord<'
object.insert("transport".to_string(), serde_json::json!(transport));
}
}
let mut event = UsageRecord::new("v1.chat", UsageRecordLevel::Error, "chat dispatch failed");
event.request_id = Some(record.request_id.to_string());
event.account_id = record.account_id.map(ToString::to_string);
event.route = Some(record.route.to_string());
event.model = Some(record.model.to_string());
event.status_code = Some(record.error.http_status_code() as i64);
event.latency_ms = Some(elapsed_millis_i64(record.started_at));
event.metadata = metadata;
if let Err(error) = record.usage_records.record(event).await {
tracing::warn!(error = %error, "failed to record chat dispatch error event");
}
}
async fn record_response_event(record: ResponseUsageRecord<'_>) {
let mut event = UsageRecord::new("v1.chat", record.level, record.message);
event.request_id = Some(record.request_id.to_string());
event.account_id = Some(record.account_id.to_string());
event.route = Some(record.route.to_string());
event.model = Some(record.model.to_string());
event.status_code = Some(record.status_code);
event.latency_ms = Some(elapsed_millis_i64(record.started_at));
event.metadata = record.metadata;
if let Some(object) = event.metadata.as_object_mut() {
object
.entry("route".to_string())
.or_insert_with(|| serde_json::json!(record.route));
object
.entry("apiKind".to_string())
.or_insert_with(|| serde_json::json!("chat"));
enrich_usage_record_identity(
object,
record.requested_model,
record.model,
record.client_ip,
record.client_user_agent,
record.reasoning_effort,
record.service_tier,
);
}
let rate_limit_headers = record.rate_limit_headers;
if !rate_limit_headers.is_empty() {
if let Some(object) = event.metadata.as_object_mut() {
object.insert(
"rateLimitHeaders".to_string(),
serde_json::json!(rate_limit_headers),
);
}
}
if let Err(error) = record.usage_records.record(event).await {
tracing::warn!(account_id = %record.account_id, error = %error, "failed to record response event");
}
}
fn enrich_usage_record_identity(
object: &mut serde_json::Map<String, Value>,
requested_model: Option<&str>,
upstream_model: &str,
client_ip: Option<&str>,
client_user_agent: Option<&str>,
reasoning_effort: Option<&str>,
service_tier: Option<&str>,
) {
let upstream_model = upstream_model.trim();
let requested_model = requested_model
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(upstream_model);
object.insert(
"requestedModel".to_string(),
Value::String(requested_model.to_string()),
);
object.insert(
"upstreamModel".to_string(),
Value::String(upstream_model.to_string()),
);
if let Some(client_ip) = client_ip.map(str::trim).filter(|value| !value.is_empty()) {
object.insert("clientIp".to_string(), Value::String(client_ip.to_string()));
}
if let Some(user_agent) = client_user_agent
.map(str::trim)
.filter(|value| !value.is_empty())
{
object.insert(
"userAgent".to_string(),
Value::String(user_agent.to_string()),
);
}
if let Some(reasoning_effort) = reasoning_effort
.map(str::trim)
.filter(|value| !value.is_empty())
{
object.insert(
"reasoningEffort".to_string(),
Value::String(reasoning_effort.to_string()),
);
}
if let Some(service_tier) = service_tier
.map(str::trim)
.filter(|value| !value.is_empty())
{
object.insert(
"serviceTier".to_string(),
Value::String(service_tier.to_string()),
);
}
}
fn reasoning_effort_from_request(
request: &crate::upstream::protocol::responses::CodexResponsesRequest,
) -> Option<&str> {
request
.reasoning
.as_ref()?
.get("effort")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
record_dispatch_error_event(DispatchErrorUsageRecord {
usage_records: record.usage_records,
request_id: record.request_id,
account_id: record.account_id,
route: record.route,
model: record.model,
started_at: record.started_at,
status_code: i64::from(record.error.http_status_code()),
message: "chat dispatch failed",
metadata,
})
.await;
}
fn enrich_chat_error_metadata(
@@ -789,7 +669,3 @@ fn enrich_chat_error_metadata(
object.insert("upstreamStatus".to_string(), serde_json::json!(status));
}
}
fn elapsed_millis_i64(started_at: Instant) -> i64 {
started_at.elapsed().as_millis().min(i64::MAX as u128) as i64
}
backend/src/proxy/dispatch/mod.rs
+2
View File
@@ -1,5 +1,6 @@
//! 请求调度模块（账号选择、回退、恢复、会话亲和性）。
mod auth_recovery;
pub mod chat;
pub mod cloudflare;
pub mod errors;
@@ -8,3 +9,4 @@ pub mod reasoning_replay;
pub mod responses;
pub mod session_affinity;
pub mod upstream;
mod usage_events;
backend/src/proxy/dispatch/responses.rs
+66 -730
View File
@@ -7,26 +7,24 @@
//! - 账号回退与错误恢复
//! - 配额验证
use std::{
pin::Pin,
sync::Arc,
task::{Context, Poll},
time::Instant,
};
use std::{pin::Pin, sync::Arc, time::Instant};
use axum::body::Bytes;
use chrono::{DateTime, Duration, Utc};
use futures::{stream::Stream, StreamExt};
use serde_json::{json, Value};
use thiserror::Error;
use tokio::sync::{mpsc, oneshot, Mutex};
use tokio::sync::Mutex;
use tracing;
use crate::{
admin::monitoring::{
usage_record::{ResponseUsageRecord, UsageRecord, UsageRecordLevel},
usage_record_store::AdminUsageRecordService,
usage_record_model::{ResponseUsageRecord, UsageRecord, UsageRecordLevel},
usage_record_service::AdminUsageRecordService,
},
infra::time::elapsed_millis_i64,
proxy::dispatch::{
auth_recovery::trigger_refresh_after_auth_failure,
cloudflare::{
cloudflare_challenge_error_message, cloudflare_path_block_error_message,
is_cloudflare_challenge_upstream_error, is_cloudflare_path_block_upstream_error,
@@ -34,11 +32,10 @@ use crate::{
},
errors::{
auth_failure_account_status, backend_transport_name, is_auth_upstream_error,
is_history_recovery_signal, is_history_recovery_upstream_error,
is_invalid_encrypted_content_signal, is_model_unsupported_signal,
is_model_unsupported_upstream_error, is_quota_exhausted_upstream_error,
is_rate_limit_upstream_error, rate_limit_cooldown_until, upstream_error_body,
upstream_error_http_status, upstream_error_set_cookie_headers,
is_history_recovery_upstream_error, is_model_unsupported_upstream_error,
is_quota_exhausted_upstream_error, is_rate_limit_upstream_error,
rate_limit_cooldown_until, upstream_error_body, upstream_error_http_status,
upstream_error_set_cookie_headers,
},
reasoning_replay::ReasoningReplayCache,
session_affinity::{
@@ -52,6 +49,12 @@ use crate::{
QuotaVerificationContext, QuotaVerificationDecision,
QUOTA_VERIFY_LIMIT_REACHED_MESSAGE,
},
usage_events::{
enrich_event_route_metadata, enrich_usage_record_identity,
event_kind as response_event_kind, reasoning_effort_from_compact_request,
reasoning_effort_from_request, record_dispatch_error_event, record_response_event,
DispatchErrorUsageRecord,
},
},
upstream::accounts::{
model::{Account, AccountStatus},
@@ -63,20 +66,15 @@ use crate::{
protocol::{
events::{extract_sse_usage, extract_usage, TokenUsage},
responses::{
apply_response_model_options, completed_response_metadata,
reconvert_responses_sse_event_tuple_values, response_body_has_first_event,
response_from_codex_sse, CodexCompactRequest, CodexResponsesRequest,
CollectedResponse, ResponsesSseFailure,
},
sse::{
encode_sse_event, parse_sse_events, sse_body_has_done, SseError, DONE_SSE_FRAME,
apply_response_model_options, completed_response_metadata, response_from_codex_sse,
CodexCompactRequest, CodexResponsesRequest, CollectedResponse, ResponsesSseFailure,
},
sse::{parse_sse_events, sse_frame_end, SseError},
},
token_client::OpenAiTokenClient,
transport::{
backend_transport_for_response_request, is_banned_auth_signal,
is_banned_upstream_error, CodexBackendClient, CodexBackendResponse,
CodexBackendSseStream, CodexBackendTransport, CodexClientError,
backend_transport_for_response_request, is_banned_upstream_error, CodexBackendClient,
CodexBackendResponse, CodexBackendSseStream, CodexBackendTransport, CodexClientError,
CodexRateLimitHeaderUpdates, CodexTurnStateUpdate, WebSocketPoolDecision,
},
},
@@ -88,6 +86,18 @@ use super::implicit_resume::{
use crate::proxy::openai::responses::response_failed_sse_event_with_id;
mod live_stream;
use live_stream::spawn_live_response_stream;
mod sse_failure;
use sse_failure::{
auth_sse_failure_account_status, client_error_invalid_reasoning_replay, first_sse_failure,
is_auth_sse_failure, is_history_recovery_sse_failure, is_model_unsupported_sse_failure,
is_quota_exhausted_sse_failure, sse_failure_error_body, sse_failure_invalid_reasoning_replay,
status_code_for_stream_failure, stream_failure_http_status, stream_failure_metadata,
stream_failure_source, synthetic_stream_disconnected_detail, STREAM_DISCONNECTED_CODE,
STREAM_DISCONNECTED_MESSAGE,
};
#[derive(Clone, Copy)]
enum ExhaustedAccountClass {
QuotaExhausted,
@@ -138,29 +148,6 @@ pub struct ResponseDispatchStream {
pub body: ResponseBodyStream,
}
struct MpscResponseBodyStream {
receiver: mpsc::Receiver<Result<Bytes, ResponseDispatchStreamError>>,
cancel: Option<oneshot::Sender<()>>,
}
impl Drop for MpscResponseBodyStream {
fn drop(&mut self) {
if let Some(cancel) = self.cancel.take() {
let _ = cancel.send(());
}
}
}
impl Stream for MpscResponseBodyStream {
type Item = Result<Bytes, ResponseDispatchStreamError>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.receiver.poll_recv(cx)
}
}
use axum::body::Bytes;
impl ResponseDispatchService {
pub(crate) fn new(parts: ResponseDispatchServiceParts) -> Self {
Self {
@@ -178,24 +165,6 @@ impl ResponseDispatchService {
}
}
fn trigger_refresh_after_auth_failure(&self, account_id: &str, status: AccountStatus) {
if status != AccountStatus::Expired {
return;
}
let token_refresh = self.token_refresh.clone();
let account_id = account_id.to_string();
tokio::spawn(async move {
if let Err(error) = token_refresh.trigger_account_refresh_now(&account_id).await {
tracing::warn!(
account_id = %account_id,
error = %error,
"reactive token refresh after upstream auth failure failed"
);
}
});
}
async fn prepare_response_session(
&self,
request: &mut CodexResponsesRequest,
@@ -736,7 +705,8 @@ impl ResponseDispatchService {
self.account_pool
.set_status(&release_account_id, account_status)
.await;
self.trigger_refresh_after_auth_failure(
trigger_refresh_after_auth_failure(
&self.token_refresh,
&release_account_id,
account_status,
);
@@ -800,7 +770,11 @@ impl ResponseDispatchService {
self.account_pool
.set_status(&release_account_id, account_status)
.await;
self.trigger_refresh_after_auth_failure(&release_account_id, account_status);
trigger_refresh_after_auth_failure(
&self.token_refresh,
&release_account_id,
account_status,
);
excluded_account_ids.push(release_account_id);
}
Err(error) if is_cloudflare_challenge_upstream_error(&error) => {
@@ -1368,7 +1342,8 @@ impl ResponseDispatchService {
self.account_pool
.set_status(&release_account_id, account_status)
.await;
self.trigger_refresh_after_auth_failure(
trigger_refresh_after_auth_failure(
&self.token_refresh,
&release_account_id,
account_status,
);
@@ -1481,7 +1456,11 @@ impl ResponseDispatchService {
self.account_pool
.set_status(&release_account_id, account_status)
.await;
self.trigger_refresh_after_auth_failure(&release_account_id, account_status);
trigger_refresh_after_auth_failure(
&self.token_refresh,
&release_account_id,
account_status,
);
excluded_account_ids.push(release_account_id);
}
Err(error) if is_cloudflare_challenge_upstream_error(&error) => {
@@ -1872,7 +1851,11 @@ impl ResponseDispatchService {
self.account_pool
.set_status(&release_account_id, account_status)
.await;
self.trigger_refresh_after_auth_failure(&release_account_id, account_status);
trigger_refresh_after_auth_failure(
&self.token_refresh,
&release_account_id,
account_status,
);
excluded_account_ids.push(release_account_id);
}
Err(error) if is_cloudflare_challenge_upstream_error(&error) => {
@@ -1964,91 +1947,6 @@ impl ResponseDispatchService {
}
}
fn sse_failure_error_body(failure: &ResponsesSseFailure) -> String {
match failure.upstream_code.as_deref() {
Some(code) => serde_json::json!({
"error": {
"code": code,
"message": failure.message.as_str(),
}
})
.to_string(),
None => failure.message.clone(),
}
}
fn is_quota_exhausted_sse_failure(failure: &ResponsesSseFailure) -> bool {
failure
.upstream_code
.as_deref()
.is_some_and(|code| matches!(code, "quota_exceeded" | "insufficient_quota"))
|| failure.message.to_ascii_lowercase().contains("quota")
}
fn is_auth_sse_failure(failure: &ResponsesSseFailure) -> bool {
failure.upstream_code.as_deref().is_some_and(|code| {
let code = code.to_ascii_lowercase();
matches!(
code.as_str(),
"token_invalid"
| "token_expired"
| "token_revoked"
| "account_deactivated"
| "unauthorized"
| "invalid_api_key"
)
}) || {
let message = failure.message.to_ascii_lowercase();
message.contains("token revoked")
|| message.contains("token invalid")
|| message.contains("token expired")
}
}
fn is_model_unsupported_sse_failure(failure: &ResponsesSseFailure) -> bool {
failure
.upstream_code
.as_deref()
.is_some_and(is_model_unsupported_signal)
|| is_model_unsupported_signal(&failure.message)
}
fn is_history_recovery_sse_failure(failure: &ResponsesSseFailure) -> bool {
failure
.upstream_code
.as_deref()
.is_some_and(is_history_recovery_signal)
|| is_history_recovery_signal(&failure.message)
}
fn sse_failure_invalid_reasoning_replay(failure: &ResponsesSseFailure) -> bool {
failure
.upstream_code
.as_deref()
.is_some_and(is_invalid_encrypted_content_signal)
|| is_invalid_encrypted_content_signal(&failure.message)
}
fn client_error_invalid_reasoning_replay(error: &CodexClientError) -> bool {
matches!(
error,
CodexClientError::Upstream { body, .. } if is_invalid_encrypted_content_signal(body)
)
}
fn auth_sse_failure_account_status(failure: &ResponsesSseFailure) -> AccountStatus {
if failure
.upstream_code
.as_deref()
.is_some_and(is_banned_auth_signal)
|| is_banned_auth_signal(&failure.message)
{
AccountStatus::Banned
} else {
AccountStatus::Expired
}
}
fn strip_request_history(request: &mut CodexResponsesRequest) {
request.previous_response_id = None;
request.turn_state = None;
@@ -2060,7 +1958,7 @@ async fn prefetch_first_sse_chunk(
mut body: CodexBackendSseStream,
) -> Result<(Bytes, CodexBackendSseStream), ResponseDispatchError> {
let mut prefetched = Vec::new();
while !contains_sse_event_separator(&prefetched) {
while sse_frame_end(&prefetched).is_none() {
let Some(next) = body.next().await else {
if prefetched.is_empty() {
return Err(ResponseDispatchError::EmptyUpstreamResponse);
@@ -2081,21 +1979,6 @@ async fn prefetch_first_sse_chunk(
Ok((Bytes::from(prefetched), body))
}
fn contains_sse_event_separator(bytes: &[u8]) -> bool {
bytes.windows(2).any(|window| window == b"\n\n")
|| bytes.windows(4).any(|window| window == b"\r\n\r\n")
}
fn first_sse_failure(prefetched: &[u8]) -> Result<Option<ResponsesSseFailure>, SseError> {
let body = String::from_utf8_lossy(prefetched);
match response_from_codex_sse(&body, None)? {
CollectedResponse::Failed(failure) => Ok(Some(failure)),
CollectedResponse::Completed(_)
| CollectedResponse::MissingCompleted
| CollectedResponse::Empty => Ok(None),
}
}
struct LiveResponseStreamContext {
account_pool: Arc<RuntimeAccountPoolService>,
session_affinity: Arc<RuntimeSessionAffinityService>,
@@ -2121,296 +2004,6 @@ struct LiveResponseStreamContext {
started_at: Instant,
}
fn spawn_live_response_stream(
context: LiveResponseStreamContext,
prefetched: Bytes,
mut body: CodexBackendSseStream,
) -> ResponseDispatchStream {
let (sender, receiver) = mpsc::channel(8);
let (cancel_sender, mut cancel_receiver) = oneshot::channel();
tokio::spawn(async move {
let mut tuple_transformer = context
.tuple_schema
.clone()
.map(TupleSseEventTransformer::new);
let mut body_bytes = Vec::new();
let mut first_token_ms = None;
if !send_live_response_stream_chunk(
&sender,
&mut body_bytes,
tuple_transformer.as_mut(),
prefetched,
)
.await
{
context.account_pool.release(&context.account_id).await;
return;
}
update_first_token_ms(&context, &body_bytes, &mut first_token_ms);
loop {
let next = tokio::select! {
_ = &mut cancel_receiver => {
context.account_pool.release(&context.account_id).await;
return;
}
next = body.next() => next,
};
let Some(next) = next else {
break;
};
match next {
Ok(chunk) => {
if !send_live_response_stream_chunk(
&sender,
&mut body_bytes,
tuple_transformer.as_mut(),
chunk,
)
.await
{
context.account_pool.release(&context.account_id).await;
return;
}
update_first_token_ms(&context, &body_bytes, &mut first_token_ms);
}
Err(error) => {
if !flush_live_response_stream_transformer(
&sender,
&mut body_bytes,
tuple_transformer.as_mut(),
)
.await
{
context.account_pool.release(&context.account_id).await;
return;
}
let detail = error.to_string();
let Some(body_text) =
send_live_response_stream_tail(&sender, &mut body_bytes, Some(&detail))
.await
else {
context.account_pool.release(&context.account_id).await;
return;
};
finalize_live_response_stream(context, body_text, first_token_ms).await;
return;
}
}
}
if !flush_live_response_stream_transformer(
&sender,
&mut body_bytes,
tuple_transformer.as_mut(),
)
.await
{
context.account_pool.release(&context.account_id).await;
return;
}
let Some(body_text) = send_live_response_stream_tail(&sender, &mut body_bytes, None).await
else {
context.account_pool.release(&context.account_id).await;
return;
};
finalize_live_response_stream(context, body_text, first_token_ms).await;
});
ResponseDispatchStream {
body: Box::pin(MpscResponseBodyStream {
receiver,
cancel: Some(cancel_sender),
}),
}
}
fn update_first_token_ms(
context: &LiveResponseStreamContext,
body_bytes: &[u8],
first_token_ms: &mut Option<i64>,
) {
if first_token_ms.is_none() && response_body_has_first_event(body_bytes) {
*first_token_ms = Some(elapsed_millis_i64(context.started_at).max(1));
}
}
async fn send_live_response_stream_chunk(
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
body_bytes: &mut Vec<u8>,
transformer: Option<&mut TupleSseEventTransformer>,
chunk: Bytes,
) -> bool {
let chunks = match transformer {
Some(transformer) => transformer.push(&chunk),
None => vec![chunk],
};
send_live_response_stream_chunks(sender, body_bytes, chunks).await
}
async fn flush_live_response_stream_transformer(
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
body_bytes: &mut Vec<u8>,
transformer: Option<&mut TupleSseEventTransformer>,
) -> bool {
let Some(transformer) = transformer else {
return true;
};
send_live_response_stream_chunks(sender, body_bytes, transformer.finish()).await
}
async fn send_live_response_stream_chunks(
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
body_bytes: &mut Vec<u8>,
chunks: Vec<Bytes>,
) -> bool {
for chunk in chunks {
body_bytes.extend_from_slice(&chunk);
if sender.send(Ok(chunk)).await.is_err() {
return false;
}
}
true
}
struct TupleSseEventTransformer {
tuple_schema: Value,
pending: Vec<u8>,
}
impl TupleSseEventTransformer {
fn new(tuple_schema: Value) -> Self {
Self {
tuple_schema,
pending: Vec::new(),
}
}
fn push(&mut self, chunk: &[u8]) -> Vec<Bytes> {
self.pending.extend_from_slice(chunk);
let mut chunks = Vec::new();
while let Some(frame_end) = next_sse_frame_end(&self.pending) {
let frame = self.pending.drain(..frame_end).collect::<Vec<_>>();
chunks.push(self.transform_frame(&frame));
}
chunks
}
fn finish(&mut self) -> Vec<Bytes> {
if self.pending.is_empty() {
return Vec::new();
}
let frame = std::mem::take(&mut self.pending);
vec![self.transform_frame(&frame)]
}
fn transform_frame(&self, frame: &[u8]) -> Bytes {
let frame_text = String::from_utf8_lossy(frame);
let Ok(events) = parse_sse_events(&frame_text) else {
return Bytes::copy_from_slice(frame);
};
let [event] = events.as_slice() else {
return Bytes::copy_from_slice(frame);
};
let Ok(data) = serde_json::from_str::<Value>(&event.data) else {
return Bytes::copy_from_slice(frame);
};
let transformed = reconvert_responses_sse_event_tuple_values(
event.event.as_deref(),
data,
&self.tuple_schema,
);
Bytes::from(encode_sse_event(
event.event.as_deref().unwrap_or_default(),
&transformed.to_string(),
))
}
}
fn next_sse_frame_end(bytes: &[u8]) -> Option<usize> {
let lf_lf = bytes
.windows(2)
.position(|window| window == b"\n\n")
.map(|position| position + 2);
let crlf_crlf = bytes
.windows(4)
.position(|window| window == b"\r\n\r\n")
.map(|position| position + 4);
match (lf_lf, crlf_crlf) {
(Some(left), Some(right)) => Some(left.min(right)),
(Some(end), None) | (None, Some(end)) => Some(end),
(None, None) => None,
}
}
async fn send_live_response_stream_tail(
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
body_bytes: &mut Vec<u8>,
failure_detail: Option<&str>,
) -> Option<String> {
let mut body_text = String::from_utf8_lossy(body_bytes).to_string();
if !sse_body_has_terminal_event(&body_text) {
if let Some(separator) = missing_sse_event_separator(&body_text) {
body_text.push_str(separator);
body_bytes.extend_from_slice(separator.as_bytes());
if sender
.send(Ok(Bytes::copy_from_slice(separator.as_bytes())))
.await
.is_err()
{
return None;
}
}
let failure =
premature_close_failed_event(latest_response_id(&body_text).as_deref(), failure_detail);
body_text.push_str(&failure);
body_bytes.extend_from_slice(failure.as_bytes());
if sender.send(Ok(Bytes::from(failure))).await.is_err() {
return None;
}
}
if !sse_body_has_done(&body_text) {
body_text.push_str(DONE_SSE_FRAME);
body_bytes.extend_from_slice(DONE_SSE_FRAME.as_bytes());
if sender
.send(Ok(Bytes::from_static(DONE_SSE_FRAME.as_bytes())))
.await
.is_err()
{
return None;
}
}
Some(body_text)
}
fn sse_body_has_terminal_event(body: &str) -> bool {
parse_sse_events(body).is_ok_and(|events| {
events.iter().any(|event| {
matches!(
event.event.as_deref(),
Some("response.completed" | "response.failed" | "error")
)
})
})
}
fn missing_sse_event_separator(body: &str) -> Option<&'static str> {
if body.is_empty()
|| body.ends_with("\n\n")
|| body.ends_with("\r\n\r\n")
|| body.ends_with("\r\r")
{
None
} else if body.ends_with('\n') || body.ends_with('\r') {
Some("\n")
} else {
Some("\n\n")
}
}
fn latest_response_id(body: &str) -> Option<String> {
parse_sse_events(body).ok().and_then(|events| {
events.iter().rev().find_map(|event| {
@@ -2426,9 +2019,6 @@ fn latest_response_id(body: &str) -> Option<String> {
})
}
const STREAM_DISCONNECTED_CODE: &str = "stream_disconnected";
const STREAM_DISCONNECTED_MESSAGE: &str = "Upstream stream closed before response.completed";
fn premature_close_failed_event(response_id: Option<&str>, detail: Option<&str>) -> String {
let message = match detail.filter(|value| !value.trim().is_empty()) {
Some(detail) => format!("{STREAM_DISCONNECTED_MESSAGE}: {detail}"),
@@ -2660,26 +2250,18 @@ async fn record_response_dispatch_error_event(record: ResponseDispatchErrorEvent
record.transport,
);
enrich_response_dispatch_error_metadata(&mut metadata, record.error);
enrich_event_route_metadata(&mut metadata, record.route);
let mut event = UsageRecord::new(
response_event_kind(record.route),
UsageRecordLevel::Error,
"v1 responses dispatch failed",
);
event.request_id = Some(record.request_id.to_string());
event.account_id = record.account_id.map(ToString::to_string);
event.route = Some(record.route.to_string());
event.model = Some(record.model.to_string());
event.status_code = Some(i64::from(record.error.http_status_code()));
event.latency_ms = Some(elapsed_millis_i64(record.started_at));
event.metadata = metadata;
if let Err(error) = record.usage_records.record(event).await {
tracing::warn!(
account_id = record.account_id.unwrap_or(""),
error = %error,
"failed to record response dispatch error event"
);
}
record_dispatch_error_event(DispatchErrorUsageRecord {
usage_records: record.usage_records,
request_id: record.request_id,
account_id: record.account_id,
route: record.route,
model: record.model,
started_at: record.started_at,
status_code: i64::from(record.error.http_status_code()),
message: "v1 responses dispatch failed",
metadata,
})
.await;
}
async fn record_response_upstream_error_event(record: ResponseUpstreamErrorEventRecord<'_>) {
@@ -2757,150 +2339,6 @@ async fn record_prefetched_response_stream_failure_event(
.await;
}
async fn record_response_event(record: ResponseUsageRecord<'_>) {
let mut metadata = record.metadata;
enrich_event_route_metadata(&mut metadata, record.route);
let mut event = UsageRecord::new(
response_event_kind(record.route),
record.level,
record.message,
);
event.request_id = Some(record.request_id.to_string());
event.account_id = Some(record.account_id.to_string());
event.route = Some(record.route.to_string());
event.model = Some(record.model.to_string());
event.status_code = Some(record.status_code);
event.latency_ms = Some(elapsed_millis_i64(record.started_at));
enrich_usage_record_identity(
&mut metadata,
record.requested_model,
record.model,
record.client_ip,
record.client_user_agent,
record.reasoning_effort,
record.service_tier,
);
event.metadata = metadata;
let rate_limit_headers = record.rate_limit_headers;
if !rate_limit_headers.is_empty() {
if let Some(object) = event.metadata.as_object_mut() {
object.insert(
"rateLimitHeaders".to_string(),
serde_json::json!(rate_limit_headers),
);
}
}
if let Err(error) = record.usage_records.record(event).await {
tracing::warn!(account_id = %record.account_id, error = %error, "failed to record response event");
}
}
fn response_event_kind(route: &str) -> &'static str {
if route == "/v1/chat/completions" {
"v1.chat"
} else {
"v1.response"
}
}
fn response_api_kind(route: &str) -> &'static str {
if route == "/v1/chat/completions" {
"chat"
} else {
"responses"
}
}
fn enrich_event_route_metadata(metadata: &mut Value, route: &str) {
let Some(object) = metadata.as_object_mut() else {
return;
};
object
.entry("route".to_string())
.or_insert_with(|| Value::String(route.to_string()));
object
.entry("apiKind".to_string())
.or_insert_with(|| Value::String(response_api_kind(route).to_string()));
}
fn enrich_usage_record_identity(
metadata: &mut Value,
requested_model: Option<&str>,
upstream_model: &str,
client_ip: Option<&str>,
client_user_agent: Option<&str>,
reasoning_effort: Option<&str>,
service_tier: Option<&str>,
) {
let Some(object) = metadata.as_object_mut() else {
return;
};
let upstream_model = upstream_model.trim();
let requested_model = requested_model
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(upstream_model);
object.insert(
"requestedModel".to_string(),
Value::String(requested_model.to_string()),
);
object.insert(
"upstreamModel".to_string(),
Value::String(upstream_model.to_string()),
);
if let Some(client_ip) = client_ip.map(str::trim).filter(|value| !value.is_empty()) {
object.insert("clientIp".to_string(), Value::String(client_ip.to_string()));
}
if let Some(user_agent) = client_user_agent
.map(str::trim)
.filter(|value| !value.is_empty())
{
object.insert(
"userAgent".to_string(),
Value::String(user_agent.to_string()),
);
}
if let Some(reasoning_effort) = reasoning_effort
.map(str::trim)
.filter(|value| !value.is_empty())
{
object.insert(
"reasoningEffort".to_string(),
Value::String(reasoning_effort.to_string()),
);
}
if let Some(service_tier) = service_tier
.map(str::trim)
.filter(|value| !value.is_empty())
{
object.insert(
"serviceTier".to_string(),
Value::String(service_tier.to_string()),
);
}
}
fn reasoning_effort_from_request(request: &CodexResponsesRequest) -> Option<&str> {
reasoning_effort_from_value(request.reasoning.as_ref())
}
fn reasoning_effort_from_compact_request(request: &CodexCompactRequest) -> Option<&str> {
reasoning_effort_from_value(request.reasoning.as_ref())
}
fn reasoning_effort_from_value(reasoning: Option<&Value>) -> Option<&str> {
reasoning?
.get("effort")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn ensure_stream_metadata_flag(metadata: &mut Value) {
let Some(object) = metadata.as_object_mut() else {
*metadata = serde_json::json!({ "stream": true });
@@ -3066,57 +2504,6 @@ async fn live_response_turn_state(context: &LiveResponseStreamContext) -> Option
context.turn_state.clone()
}
fn stream_failure_metadata(failure: &ResponsesSseFailure, usage: Option<TokenUsage>) -> Value {
let mut metadata = serde_json::json!({
"stream": true,
"failed": true,
"failureEvent": failure.event,
"failureMessage": failure.message,
"upstreamCode": failure.upstream_code,
"usage": usage,
});
enrich_stream_failure_source_metadata(&mut metadata, failure);
metadata
}
fn enrich_stream_failure_source_metadata(metadata: &mut Value, failure: &ResponsesSseFailure) {
let Some(object) = metadata.as_object_mut() else {
return;
};
object.insert(
"failureSource".to_string(),
Value::String(stream_failure_source(failure).to_string()),
);
if let Some(detail) = synthetic_stream_disconnected_detail(failure) {
object.insert("synthetic".to_string(), Value::Bool(true));
if !detail.is_empty() {
object.insert("failureDetail".to_string(), Value::String(detail));
}
}
}
fn stream_failure_source(failure: &ResponsesSseFailure) -> &'static str {
if synthetic_stream_disconnected_detail(failure).is_some() {
"proxy"
} else {
"upstream"
}
}
fn synthetic_stream_disconnected_detail(failure: &ResponsesSseFailure) -> Option<String> {
if failure.upstream_code.as_deref() != Some(STREAM_DISCONNECTED_CODE) {
return None;
}
let detail = failure
.message
.strip_prefix(STREAM_DISCONNECTED_MESSAGE)?
.strip_prefix(": ")
.unwrap_or_default()
.trim()
.to_string();
Some(detail)
}
fn insert_first_token_ms(metadata: &mut Value, first_token_ms: Option<i64>) {
let Some(first_token_ms) = first_token_ms else {
return;
@@ -3138,53 +2525,6 @@ fn insert_websocket_pool_decision(metadata: &mut Value, decision: Option<WebSock
}
}
fn status_code_for_stream_failure(failure: &ResponsesSseFailure) -> i64 {
let code = failure
.upstream_code
.as_deref()
.unwrap_or("error")
.to_ascii_lowercase();
if code.contains("model") && (code.contains("not_supported") || code.contains("not_available"))
{
return 400;
}
if code.contains("invalid_request") || code.contains("not_found") {
return 400;
}
if code.contains("context_window")
|| code.contains("invalid_prompt")
|| code.contains("cyber_policy")
|| code.contains("bad_request")
{
return 400;
}
if code.contains("rate_limit") || code.contains("usage_limit") {
return 429;
}
if code.contains("unauthorized")
|| code.contains("invalid_api_key")
|| code == "token_invalid"
|| code == "token_expired"
|| code == "account_deactivated"
{
return 401;
}
if code.contains("forbidden") || code.contains("banned") {
return 403;
}
if code.contains("payment") || code.contains("quota") {
return 402;
}
if code.contains("server_overloaded") {
return 503;
}
502
}
fn stream_failure_http_status(failure: &ResponsesSseFailure) -> u16 {
u16::try_from(status_code_for_stream_failure(failure)).unwrap_or(502)
}
// ====================================================================
// Affinity + replay recording
// ====================================================================
@@ -3494,7 +2834,3 @@ fn enrich_response_dispatch_error_metadata(metadata: &mut Value, error: &Respons
object.insert("upstreamStatus".to_string(), json!(status));
}
}
fn elapsed_millis_i64(started_at: Instant) -> i64 {
started_at.elapsed().as_millis().min(i64::MAX as u128) as i64
}
backend/src/proxy/dispatch/responses/live_stream.rs
+317
View File
@@ -0,0 +1,317 @@
use std::{
pin::Pin,
task::{Context, Poll},
};
use axum::body::Bytes;
use futures::{stream::Stream, StreamExt};
use serde_json::Value;
use tokio::sync::{mpsc, oneshot};
use crate::{
proxy::dispatch::responses::{
finalize_live_response_stream, latest_response_id, premature_close_failed_event,
LiveResponseStreamContext, ResponseDispatchStream, ResponseDispatchStreamError,
},
upstream::{
protocol::{
responses::{
reconvert_responses_sse_event_tuple_values, update_first_response_event_ms,
},
sse::{
encode_sse_event, parse_sse_events, sse_body_has_done, sse_frame_end,
DONE_SSE_FRAME,
},
},
transport::CodexBackendSseStream,
},
};
struct MpscResponseBodyStream {
receiver: mpsc::Receiver<Result<Bytes, ResponseDispatchStreamError>>,
cancel: Option<oneshot::Sender<()>>,
}
impl Drop for MpscResponseBodyStream {
fn drop(&mut self) {
if let Some(cancel) = self.cancel.take() {
let _ = cancel.send(());
}
}
}
impl Stream for MpscResponseBodyStream {
type Item = Result<Bytes, ResponseDispatchStreamError>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.receiver.poll_recv(cx)
}
}
pub(super) fn spawn_live_response_stream(
context: LiveResponseStreamContext,
prefetched: Bytes,
mut body: CodexBackendSseStream,
) -> ResponseDispatchStream {
let (sender, receiver) = mpsc::channel(8);
let (cancel_sender, mut cancel_receiver) = oneshot::channel();
tokio::spawn(async move {
let mut tuple_transformer = context
.tuple_schema
.clone()
.map(TupleSseEventTransformer::new);
let mut body_bytes = Vec::new();
let mut first_token_ms = None;
if !send_live_response_stream_chunk(
&sender,
&mut body_bytes,
tuple_transformer.as_mut(),
prefetched,
)
.await
{
context.account_pool.release(&context.account_id).await;
return;
}
update_first_response_event_ms(context.started_at, &body_bytes, &mut first_token_ms);
loop {
let next = tokio::select! {
_ = &mut cancel_receiver => {
context.account_pool.release(&context.account_id).await;
return;
}
next = body.next() => next,
};
let Some(next) = next else {
break;
};
match next {
Ok(chunk) => {
if !send_live_response_stream_chunk(
&sender,
&mut body_bytes,
tuple_transformer.as_mut(),
chunk,
)
.await
{
context.account_pool.release(&context.account_id).await;
return;
}
update_first_response_event_ms(
context.started_at,
&body_bytes,
&mut first_token_ms,
);
}
Err(error) => {
if !flush_live_response_stream_transformer(
&sender,
&mut body_bytes,
tuple_transformer.as_mut(),
)
.await
{
context.account_pool.release(&context.account_id).await;
return;
}
let detail = error.to_string();
let Some(body_text) =
send_live_response_stream_tail(&sender, &mut body_bytes, Some(&detail))
.await
else {
context.account_pool.release(&context.account_id).await;
return;
};
finalize_live_response_stream(context, body_text, first_token_ms).await;
return;
}
}
}
if !flush_live_response_stream_transformer(
&sender,
&mut body_bytes,
tuple_transformer.as_mut(),
)
.await
{
context.account_pool.release(&context.account_id).await;
return;
}
let Some(body_text) = send_live_response_stream_tail(&sender, &mut body_bytes, None).await
else {
context.account_pool.release(&context.account_id).await;
return;
};
finalize_live_response_stream(context, body_text, first_token_ms).await;
});
ResponseDispatchStream {
body: Box::pin(MpscResponseBodyStream {
receiver,
cancel: Some(cancel_sender),
}),
}
}
async fn send_live_response_stream_chunk(
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
body_bytes: &mut Vec<u8>,
transformer: Option<&mut TupleSseEventTransformer>,
chunk: Bytes,
) -> bool {
let chunks = match transformer {
Some(transformer) => transformer.push(&chunk),
None => vec![chunk],
};
send_live_response_stream_chunks(sender, body_bytes, chunks).await
}
async fn flush_live_response_stream_transformer(
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
body_bytes: &mut Vec<u8>,
transformer: Option<&mut TupleSseEventTransformer>,
) -> bool {
let Some(transformer) = transformer else {
return true;
};
send_live_response_stream_chunks(sender, body_bytes, transformer.finish()).await
}
async fn send_live_response_stream_chunks(
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
body_bytes: &mut Vec<u8>,
chunks: Vec<Bytes>,
) -> bool {
for chunk in chunks {
body_bytes.extend_from_slice(&chunk);
if sender.send(Ok(chunk)).await.is_err() {
return false;
}
}
true
}
struct TupleSseEventTransformer {
tuple_schema: Value,
pending: Vec<u8>,
}
impl TupleSseEventTransformer {
fn new(tuple_schema: Value) -> Self {
Self {
tuple_schema,
pending: Vec::new(),
}
}
fn push(&mut self, chunk: &[u8]) -> Vec<Bytes> {
self.pending.extend_from_slice(chunk);
let mut chunks = Vec::new();
while let Some(frame_end) = sse_frame_end(&self.pending) {
let frame = self.pending.drain(..frame_end).collect::<Vec<_>>();
chunks.push(self.transform_frame(&frame));
}
chunks
}
fn finish(&mut self) -> Vec<Bytes> {
if self.pending.is_empty() {
return Vec::new();
}
let frame = std::mem::take(&mut self.pending);
vec![self.transform_frame(&frame)]
}
fn transform_frame(&self, frame: &[u8]) -> Bytes {
let frame_text = String::from_utf8_lossy(frame);
let Ok(events) = parse_sse_events(&frame_text) else {
return Bytes::copy_from_slice(frame);
};
let [event] = events.as_slice() else {
return Bytes::copy_from_slice(frame);
};
let Ok(data) = serde_json::from_str::<Value>(&event.data) else {
return Bytes::copy_from_slice(frame);
};
let transformed = reconvert_responses_sse_event_tuple_values(
event.event.as_deref(),
data,
&self.tuple_schema,
);
Bytes::from(encode_sse_event(
event.event.as_deref().unwrap_or_default(),
&transformed.to_string(),
))
}
}
async fn send_live_response_stream_tail(
sender: &mpsc::Sender<Result<Bytes, ResponseDispatchStreamError>>,
body_bytes: &mut Vec<u8>,
failure_detail: Option<&str>,
) -> Option<String> {
let mut body_text = String::from_utf8_lossy(body_bytes).to_string();
if !sse_body_has_terminal_event(&body_text) {
if let Some(separator) = missing_sse_event_separator(&body_text) {
body_text.push_str(separator);
body_bytes.extend_from_slice(separator.as_bytes());
if sender
.send(Ok(Bytes::copy_from_slice(separator.as_bytes())))
.await
.is_err()
{
return None;
}
}
let failure =
premature_close_failed_event(latest_response_id(&body_text).as_deref(), failure_detail);
body_text.push_str(&failure);
body_bytes.extend_from_slice(failure.as_bytes());
if sender.send(Ok(Bytes::from(failure))).await.is_err() {
return None;
}
}
if !sse_body_has_done(&body_text) {
body_text.push_str(DONE_SSE_FRAME);
body_bytes.extend_from_slice(DONE_SSE_FRAME.as_bytes());
if sender
.send(Ok(Bytes::from_static(DONE_SSE_FRAME.as_bytes())))
.await
.is_err()
{
return None;
}
}
Some(body_text)
}
fn sse_body_has_terminal_event(body: &str) -> bool {
parse_sse_events(body).is_ok_and(|events| {
events.iter().any(|event| {
matches!(
event.event.as_deref(),
Some("response.completed" | "response.failed" | "error")
)
})
})
}
fn missing_sse_event_separator(body: &str) -> Option<&'static str> {
if body.is_empty()
|| body.ends_with("\n\n")
|| body.ends_with("\r\n\r\n")
|| body.ends_with("\r\r")
{
None
} else if body.ends_with('\n') || body.ends_with('\r') {
Some("\n")
} else {
Some("\n\n")
}
}
backend/src/proxy/dispatch/responses/sse_failure.rs
+231
View File
@@ -0,0 +1,231 @@
use serde_json::Value;
use crate::{
proxy::dispatch::errors::{
is_history_recovery_signal, is_invalid_encrypted_content_signal,
is_model_unsupported_signal,
},
upstream::{
accounts::model::AccountStatus,
protocol::{
events::TokenUsage,
responses::{response_from_codex_sse, CollectedResponse, ResponsesSseFailure},
sse::SseError,
},
transport::{is_banned_auth_signal, CodexClientError},
},
};
pub(super) const STREAM_DISCONNECTED_CODE: &str = "stream_disconnected";
pub(super) const STREAM_DISCONNECTED_MESSAGE: &str =
"Upstream stream closed before response.completed";
pub(super) fn sse_failure_error_body(failure: &ResponsesSseFailure) -> String {
match failure.upstream_code.as_deref() {
Some(code) => serde_json::json!({
"error": {
"code": code,
"message": failure.message.as_str(),
}
})
.to_string(),
None => failure.message.clone(),
}
}
fn sse_failure_matches<F>(failure: &ResponsesSseFailure, signal: F) -> bool
where
F: Fn(&str) -> bool,
{
failure.upstream_code.as_deref().is_some_and(&signal) || signal(&failure.message)
}
fn sse_failure_matches_parts<C, M>(
failure: &ResponsesSseFailure,
code_signal: C,
message_signal: M,
) -> bool
where
C: Fn(&str) -> bool,
M: Fn(&str) -> bool,
{
failure.upstream_code.as_deref().is_some_and(code_signal) || message_signal(&failure.message)
}
pub(super) fn is_quota_exhausted_sse_failure(failure: &ResponsesSseFailure) -> bool {
sse_failure_matches_parts(
failure,
|code| matches!(code, "quota_exceeded" | "insufficient_quota"),
|message| message.to_ascii_lowercase().contains("quota"),
)
}
fn is_auth_sse_failure_code(code: &str) -> bool {
let code = code.to_ascii_lowercase();
matches!(
code.as_str(),
"token_invalid"
| "token_expired"
| "token_revoked"
| "account_deactivated"
| "unauthorized"
| "invalid_api_key"
)
}
fn is_auth_sse_failure_message(message: &str) -> bool {
let message = message.to_ascii_lowercase();
message.contains("token revoked")
|| message.contains("token invalid")
|| message.contains("token expired")
}
pub(super) fn is_auth_sse_failure(failure: &ResponsesSseFailure) -> bool {
sse_failure_matches_parts(
failure,
is_auth_sse_failure_code,
is_auth_sse_failure_message,
)
}
pub(super) fn is_model_unsupported_sse_failure(failure: &ResponsesSseFailure) -> bool {
sse_failure_matches(failure, is_model_unsupported_signal)
}
pub(super) fn is_history_recovery_sse_failure(failure: &ResponsesSseFailure) -> bool {
sse_failure_matches(failure, is_history_recovery_signal)
}
pub(super) fn sse_failure_invalid_reasoning_replay(failure: &ResponsesSseFailure) -> bool {
sse_failure_matches(failure, is_invalid_encrypted_content_signal)
}
pub(super) fn client_error_invalid_reasoning_replay(error: &CodexClientError) -> bool {
matches!(
error,
CodexClientError::Upstream { body, .. } if is_invalid_encrypted_content_signal(body)
)
}
pub(super) fn auth_sse_failure_account_status(failure: &ResponsesSseFailure) -> AccountStatus {
if sse_failure_matches(failure, is_banned_auth_signal) {
AccountStatus::Banned
} else {
AccountStatus::Expired
}
}
pub(super) fn first_sse_failure(
prefetched: &[u8],
) -> Result<Option<ResponsesSseFailure>, SseError> {
let body = String::from_utf8_lossy(prefetched);
match response_from_codex_sse(&body, None)? {
CollectedResponse::Failed(failure) => Ok(Some(failure)),
CollectedResponse::Completed(_)
| CollectedResponse::MissingCompleted
| CollectedResponse::Empty => Ok(None),
}
}
pub(super) fn stream_failure_metadata(
failure: &ResponsesSseFailure,
usage: Option<TokenUsage>,
) -> Value {
let mut metadata = serde_json::json!({
"stream": true,
"failed": true,
"failureEvent": failure.event,
"failureMessage": failure.message,
"upstreamCode": failure.upstream_code,
"usage": usage,
});
enrich_stream_failure_source_metadata(&mut metadata, failure);
metadata
}
fn enrich_stream_failure_source_metadata(metadata: &mut Value, failure: &ResponsesSseFailure) {
let Some(object) = metadata.as_object_mut() else {
return;
};
object.insert(
"failureSource".to_string(),
Value::String(stream_failure_source(failure).to_string()),
);
if let Some(detail) = synthetic_stream_disconnected_detail(failure) {
object.insert("synthetic".to_string(), Value::Bool(true));
if !detail.is_empty() {
object.insert("failureDetail".to_string(), Value::String(detail));
}
}
}
pub(super) fn stream_failure_source(failure: &ResponsesSseFailure) -> &'static str {
if synthetic_stream_disconnected_detail(failure).is_some() {
"proxy"
} else {
"upstream"
}
}
pub(super) fn synthetic_stream_disconnected_detail(
failure: &ResponsesSseFailure,
) -> Option<String> {
if failure.upstream_code.as_deref() != Some(STREAM_DISCONNECTED_CODE) {
return None;
}
let detail = failure
.message
.strip_prefix(STREAM_DISCONNECTED_MESSAGE)?
.strip_prefix(": ")
.unwrap_or_default()
.trim()
.to_string();
Some(detail)
}
pub(super) fn status_code_for_stream_failure(failure: &ResponsesSseFailure) -> i64 {
let code = failure
.upstream_code
.as_deref()
.unwrap_or("error")
.to_ascii_lowercase();
if code.contains("model") && (code.contains("not_supported") || code.contains("not_available"))
{
return 400;
}
if code.contains("invalid_request") || code.contains("not_found") {
return 400;
}
if code.contains("context_window")
|| code.contains("invalid_prompt")
|| code.contains("cyber_policy")
|| code.contains("bad_request")
{
return 400;
}
if code.contains("rate_limit") || code.contains("usage_limit") {
return 429;
}
if code.contains("unauthorized")
|| code.contains("invalid_api_key")
|| code == "token_invalid"
|| code == "token_expired"
|| code == "account_deactivated"
{
return 401;
}
if code.contains("forbidden") || code.contains("banned") {
return 403;
}
if code.contains("payment") || code.contains("quota") {
return 402;
}
if code.contains("server_overloaded") {
return 503;
}
502
}
pub(super) fn stream_failure_http_status(failure: &ResponsesSseFailure) -> u16 {
u16::try_from(status_code_for_stream_failure(failure)).unwrap_or(502)
}
backend/src/proxy/dispatch/session_affinity.rs
+4 -5
View File
@@ -9,7 +9,10 @@ use sqlx::SqlitePool;
use thiserror::Error;
use uuid::Uuid;
use crate::upstream::protocol::responses::CodexResponsesRequest;
use crate::{
infra::time::parse_rfc3339_utc as parse_rfc3339,
upstream::protocol::responses::CodexResponsesRequest,
};
// ====================================================================
// 亲和性核心类型
@@ -518,10 +521,6 @@ fn stored_session_affinity_from_row(
})
}
fn parse_rfc3339(value: &str) -> SqliteSessionAffinityStoreResult<DateTime<Utc>> {
Ok(DateTime::parse_from_rfc3339(value)?.with_timezone(&Utc))
}
fn optional_nonnegative_i64_to_u64(value: Option<i64>) -> Option<u64> {
value.and_then(|value| u64::try_from(value).ok())
}
backend/src/proxy/dispatch/upstream.rs
+158 -118
View File
@@ -1,6 +1,6 @@
//! 调度层共享的账号级 Codex 上游调用。
use std::time::Instant;
use std::{future::Future, time::Instant};
use serde_json::Value;
@@ -254,6 +254,108 @@ async fn create_compact_response_with_account(
const MAX_UPSTREAM_5XX_RETRIES_PER_ACCOUNT: usize = 2;
async fn retry_upstream_5xx<T, F, Fut>(
request_id: &str,
account_id: &str,
endpoint: Option<&str>,
retry_message: &'static str,
failure_message: &'static str,
mut operation: F,
) -> Result<T, CodexClientError>
where
F: FnMut() -> Fut,
Fut: Future<Output = Result<T, CodexClientError>>,
{
let mut retries = 0;
loop {
let result = operation().await;
match result {
Err(error)
if is_retryable_upstream_5xx_error(&error)
&& retries < MAX_UPSTREAM_5XX_RETRIES_PER_ACCOUNT =>
{
log_upstream_retry(
request_id,
account_id,
endpoint,
retries + 1,
&error,
retry_message,
);
retries += 1;
}
Ok(response) => return Ok(response),
Err(error) => {
log_upstream_failure(
request_id,
account_id,
endpoint,
retries,
&error,
failure_message,
);
return Err(error);
}
}
}
}
fn log_upstream_retry(
request_id: &str,
account_id: &str,
endpoint: Option<&str>,
retry: usize,
error: &CodexClientError,
message: &'static str,
) {
if let Some(endpoint) = endpoint {
tracing::warn!(
request_id,
account_id = %account_id,
endpoint,
retry,
error = %error,
message
);
} else {
tracing::warn!(
request_id,
account_id = %account_id,
retry,
error = %error,
message
);
}
}
fn log_upstream_failure(
request_id: &str,
account_id: &str,
endpoint: Option<&str>,
retries: usize,
error: &CodexClientError,
message: &'static str,
) {
if let Some(endpoint) = endpoint {
tracing::warn!(
request_id,
account_id = %account_id,
endpoint,
retries,
error = %error,
message
);
} else {
tracing::warn!(
request_id,
account_id = %account_id,
retries,
error = %error,
message
);
}
}
pub(crate) async fn create_response_with_account_retrying_5xx(
codex: &CodexBackendClient,
installation_id: Option<&str>,
@@ -263,45 +365,25 @@ pub(crate) async fn create_response_with_account_retrying_5xx(
account: &Account,
started_at: Instant,
) -> Result<CodexBackendResponse, CodexClientError> {
let mut retries = 0;
loop {
let result = create_response_with_account(
codex,
installation_id,
cloudflare,
request,
request_id,
account,
started_at,
)
.await;
match result {
Err(error)
if is_retryable_upstream_5xx_error(&error)
&& retries < MAX_UPSTREAM_5XX_RETRIES_PER_ACCOUNT =>
{
tracing::warn!(
request_id,
account_id = %account.id,
retry = retries + 1,
error = %error,
"upstream response request failed with retryable 5xx"
);
retries += 1;
}
Ok(response) => return Ok(response),
Err(error) => {
tracing::warn!(
request_id,
account_id = %account.id,
retries,
error = %error,
"upstream response request failed"
);
return Err(error);
}
}
}
retry_upstream_5xx(
request_id,
&account.id,
None,
"upstream response request failed with retryable 5xx",
"upstream response request failed",
|| {
create_response_with_account(
codex,
installation_id,
cloudflare,
request,
request_id,
account,
started_at,
)
},
)
.await
}
pub(crate) async fn create_response_stream_with_account_retrying_5xx(
@@ -312,44 +394,24 @@ pub(crate) async fn create_response_stream_with_account_retrying_5xx(
request_id: &str,
account: &Account,
) -> Result<CodexBackendStreamingResponse, CodexClientError> {
let mut retries = 0;
loop {
let result = create_response_stream_with_account(
codex,
installation_id,
cloudflare,
request,
request_id,
account,
)
.await;
match result {
Err(error)
if is_retryable_upstream_5xx_error(&error)
&& retries < MAX_UPSTREAM_5XX_RETRIES_PER_ACCOUNT =>
{
tracing::warn!(
request_id,
account_id = %account.id,
retry = retries + 1,
error = %error,
"upstream response stream request failed with retryable 5xx"
);
retries += 1;
}
Ok(response) => return Ok(response),
Err(error) => {
tracing::warn!(
request_id,
account_id = %account.id,
retries,
error = %error,
"upstream response stream request failed"
);
return Err(error);
}
}
}
retry_upstream_5xx(
request_id,
&account.id,
None,
"upstream response stream request failed with retryable 5xx",
"upstream response stream request failed",
|| {
create_response_stream_with_account(
codex,
installation_id,
cloudflare,
request,
request_id,
account,
)
},
)
.await
}
pub(crate) async fn create_compact_response_with_account_retrying_5xx(
@@ -360,44 +422,22 @@ pub(crate) async fn create_compact_response_with_account_retrying_5xx(
request_id: &str,
account: &Account,
) -> Result<CodexCompactResponse, CodexClientError> {
let mut retries = 0;
loop {
let result = create_compact_response_with_account(
codex,
installation_id,
cloudflare,
request,
request_id,
account,
)
.await;
match result {
Err(error)
if is_retryable_upstream_5xx_error(&error)
&& retries < MAX_UPSTREAM_5XX_RETRIES_PER_ACCOUNT =>
{
tracing::warn!(
request_id,
account_id = %account.id,
endpoint = "compact",
retry = retries + 1,
error = %error,
"upstream compact request failed with retryable 5xx"
);
retries += 1;
}
Ok(response) => return Ok(response),
Err(error) => {
tracing::warn!(
request_id,
account_id = %account.id,
endpoint = "compact",
retries,
error = %error,
"upstream compact request failed"
);
return Err(error);
}
}
}
retry_upstream_5xx(
request_id,
&account.id,
Some("compact"),
"upstream compact request failed with retryable 5xx",
"upstream compact request failed",
|| {
create_compact_response_with_account(
codex,
installation_id,
cloudflare,
request,
request_id,
account,
)
},
)
.await
}
backend/src/proxy/dispatch/usage_events.rs
+175
View File
@@ -0,0 +1,175 @@
//! Dispatch 使用记录事件辅助。
use std::time::Instant;
use serde_json::Value;
use crate::{
admin::monitoring::{
usage_record_model::{ResponseUsageRecord, UsageRecord, UsageRecordLevel},
usage_record_service::AdminUsageRecordService,
},
infra::time::elapsed_millis_i64,
upstream::protocol::responses::{CodexCompactRequest, CodexResponsesRequest},
};
pub(super) struct DispatchErrorUsageRecord<'a> {
pub usage_records: &'a AdminUsageRecordService,
pub request_id: &'a str,
pub account_id: Option<&'a str>,
pub route: &'a str,
pub model: &'a str,
pub started_at: Instant,
pub status_code: i64,
pub message: &'a str,
pub metadata: Value,
}
pub(super) async fn record_dispatch_error_event(record: DispatchErrorUsageRecord<'_>) {
let mut metadata = record.metadata;
enrich_event_route_metadata(&mut metadata, record.route);
let mut event = UsageRecord::new(
event_kind(record.route),
UsageRecordLevel::Error,
record.message,
);
event.request_id = Some(record.request_id.to_string());
event.account_id = record.account_id.map(ToString::to_string);
event.route = Some(record.route.to_string());
event.model = Some(record.model.to_string());
event.status_code = Some(record.status_code);
event.latency_ms = Some(elapsed_millis_i64(record.started_at));
event.metadata = metadata;
if let Err(error) = record.usage_records.record(event).await {
tracing::warn!(
account_id = record.account_id.unwrap_or(""),
error = %error,
"failed to record dispatch error event"
);
}
}
pub(super) async fn record_response_event(record: ResponseUsageRecord<'_>) {
let mut metadata = record.metadata;
enrich_event_route_metadata(&mut metadata, record.route);
let mut event = UsageRecord::new(event_kind(record.route), record.level, record.message);
event.request_id = Some(record.request_id.to_string());
event.account_id = Some(record.account_id.to_string());
event.route = Some(record.route.to_string());
event.model = Some(record.model.to_string());
event.status_code = Some(record.status_code);
event.latency_ms = Some(elapsed_millis_i64(record.started_at));
enrich_usage_record_identity(
&mut metadata,
record.requested_model,
record.model,
record.client_ip,
record.client_user_agent,
record.reasoning_effort,
record.service_tier,
);
event.metadata = metadata;
if !record.rate_limit_headers.is_empty() {
if let Some(object) = event.metadata.as_object_mut() {
object.insert(
"rateLimitHeaders".to_string(),
serde_json::json!(record.rate_limit_headers),
);
}
}
if let Err(error) = record.usage_records.record(event).await {
tracing::warn!(account_id = %record.account_id, error = %error, "failed to record response event");
}
}
pub(super) fn event_kind(route: &str) -> &'static str {
if route == "/v1/chat/completions" {
"v1.chat"
} else {
"v1.response"
}
}
pub(super) fn api_kind(route: &str) -> &'static str {
if route == "/v1/chat/completions" {
"chat"
} else {
"responses"
}
}
pub(super) fn enrich_event_route_metadata(metadata: &mut Value, route: &str) {
let Some(object) = metadata.as_object_mut() else {
return;
};
object
.entry("route".to_string())
.or_insert_with(|| Value::String(route.to_string()));
object
.entry("apiKind".to_string())
.or_insert_with(|| Value::String(api_kind(route).to_string()));
}
pub(super) fn enrich_usage_record_identity(
metadata: &mut Value,
requested_model: Option<&str>,
upstream_model: &str,
client_ip: Option<&str>,
client_user_agent: Option<&str>,
reasoning_effort: Option<&str>,
service_tier: Option<&str>,
) {
let Some(object) = metadata.as_object_mut() else {
return;
};
let upstream_model = upstream_model.trim();
let requested_model = requested_model
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(upstream_model);
object.insert(
"requestedModel".to_string(),
Value::String(requested_model.to_string()),
);
object.insert(
"upstreamModel".to_string(),
Value::String(upstream_model.to_string()),
);
insert_trimmed_string(object, "clientIp", client_ip);
insert_trimmed_string(object, "userAgent", client_user_agent);
insert_trimmed_string(object, "reasoningEffort", reasoning_effort);
insert_trimmed_string(object, "serviceTier", service_tier);
}
pub(super) fn reasoning_effort_from_request(request: &CodexResponsesRequest) -> Option<&str> {
reasoning_effort_from_value(request.reasoning.as_ref())
}
pub(super) fn reasoning_effort_from_compact_request(request: &CodexCompactRequest) -> Option<&str> {
reasoning_effort_from_value(request.reasoning.as_ref())
}
fn insert_trimmed_string(
object: &mut serde_json::Map<String, Value>,
key: &str,
value: Option<&str>,
) {
if let Some(value) = value.map(str::trim).filter(|value| !value.is_empty()) {
object.insert(key.to_string(), Value::String(value.to_string()));
}
}
fn reasoning_effort_from_value(reasoning: Option<&Value>) -> Option<&str> {
reasoning?
.get("effort")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
backend/src/proxy/openai/chat.rs
+4 -12
View File
@@ -29,7 +29,9 @@ use crate::{
protocol::{
responses::{apply_response_model_options, CodexResponsesRequest},
schema::{prepare_schema, reconvert_tuple_values},
sse::{encode_sse_event, parse_sse_events, SseError, DONE_SSE_FRAME},
sse::{
encode_sse_event, parse_sse_events, sse_frame_separator, SseError, DONE_SSE_FRAME,
},
},
},
};
@@ -171,7 +173,7 @@ impl ChatCompletionStreamTranslator {
self.pending.push_str(chunk);
let mut output = String::new();
while let Some((event_end, separator_len)) = sse_event_separator(&self.pending) {
while let Some((event_end, separator_len)) = sse_frame_separator(&self.pending) {
let event_frame = self.pending[..event_end].to_string();
self.pending.drain(..event_end + separator_len);
if event_frame.trim().is_empty() {
@@ -1207,16 +1209,6 @@ fn openai_stream_usage(usage: Option<&Value>) -> Value {
openai_usage
}
fn sse_event_separator(input: &str) -> Option<(usize, usize)> {
let lf = input.find("\n\n").map(|index| (index, 2));
let crlf = input.find("\r\n\r\n").map(|index| (index, 4));
match (lf, crlf) {
(Some(left), Some(right)) => Some(if left.0 <= right.0 { left } else { right }),
(Some(found), None) | (None, Some(found)) => Some(found),
(None, None) => None,
}
}
fn chat_completion_stream_id() -> String {
let suffix = Uuid::new_v4()
.simple()
backend/src/runtime/bootstrap.rs
+4 -1
View File
@@ -5,6 +5,7 @@ use chrono::Utc;
use crate::config::settings::RuntimeSettingsService;
use crate::config::types::AppConfig;
use crate::http::middleware::request_id::TrustedProxyConfig;
use crate::infra::database::connect_sqlite;
use crate::infra::logging::{init_tracing, LogGuard, RotationConfig};
use crate::infra::paths::{ensure_data_dir, load_or_create_installation_id};
@@ -59,6 +60,7 @@ async fn build_application(
),
};
config.admin.validate_default_password()?;
let default_admin_password = config.admin.default_password.clone();
let data_dir = ensure_data_dir()?;
let installation_id = load_or_create_installation_id(Some(&data_dir))?;
@@ -97,10 +99,11 @@ async fn build_application(
);
let task_coordinator = TaskCoordinator::start(&runtime_config, &services);
let trusted_proxies = TrustedProxyConfig::from_entries(&config.server.trusted_proxies)?;
let app_state = AppState { config, services };
Ok((
crate::http::router::router().with_state(app_state),
crate::http::router::router_with_trusted_proxies(trusted_proxies).with_state(app_state),
task_coordinator,
))
}
backend/src/runtime/services.rs
+17 -14
View File
@@ -4,12 +4,13 @@ use std::sync::Arc as StdArc;
use crate::{
admin::monitoring::{
service::AdminUsageService,
account_usage_service::AdminUsageService,
account_usage_store::SqliteUsageStore,
usage_record_service::AdminUsageRecordService,
usage_record_store::{
AdminUsageRecordService, SqliteUsageRecordStore, DEFAULT_USAGE_RECORD_CAPACITY,
SqliteUsageRecordStore, DEFAULT_USAGE_RECORD_CAPACITY,
DEFAULT_USAGE_RECORD_CAPTURE_BODY,
},
usage_store::SqliteUsageStore,
},
admin::{
accounts::service::{AdminAccountService, AdminAccountServiceParts},
@@ -43,6 +44,7 @@ use crate::{
token_client::{default_openai_token_client, OpenAiTokenClient, TokenClientConfig},
transport::{
build_reqwest_client, tls::CustomCaError, CodexBackendClient, CodexModelCatalogClient,
CodexWebSocketPool, CodexWebSocketPoolConfig,
},
},
};
@@ -96,6 +98,7 @@ pub struct Services {
pub responses: StdArc<ResponseDispatchService>,
pub session_affinity: StdArc<RuntimeSessionAffinityService>,
pub codex: StdArc<CodexBackendClient>,
pub websocket_pool: Option<StdArc<CodexWebSocketPool>>,
pub fingerprint: Fingerprint,
pub installation_id: Option<String>,
pub background_tasks: BackgroundTaskStores,
@@ -158,23 +161,22 @@ impl Services {
let installation_id = installation_id.filter(|id| !id.trim().is_empty());
let account_store_trait =
StdArc::new(stores.accounts.clone()) as StdArc<dyn AccountStoreTrait>;
let websocket_pool = config.ws_pool.enabled.then(|| {
StdArc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig {
enabled: true,
max_age: std::time::Duration::from_millis(config.ws_pool.max_age_ms),
max_per_account: config.ws_pool.max_per_account,
..CodexWebSocketPoolConfig::default()
}))
});
let codex = {
let client = CodexBackendClient::new(
build_reqwest_client(config.tls.force_http11)?,
config.api.base_url.clone(),
fingerprint.clone(),
);
if config.ws_pool.enabled {
let pool =
StdArc::new(crate::upstream::transport::CodexWebSocketPool::with_config(
crate::upstream::transport::CodexWebSocketPoolConfig {
enabled: config.ws_pool.enabled,
max_age: std::time::Duration::from_millis(config.ws_pool.max_age_ms),
max_per_account: config.ws_pool.max_per_account,
..crate::upstream::transport::CodexWebSocketPoolConfig::default()
},
));
StdArc::new(client.with_websocket_pool(pool))
if let Some(pool) = &websocket_pool {
StdArc::new(client.with_websocket_pool(pool.clone()))
} else {
StdArc::new(client)
}
@@ -292,6 +294,7 @@ impl Services {
responses,
session_affinity,
codex,
websocket_pool,
fingerprint,
installation_id,
background_tasks: stores,
backend/src/runtime/shutdown.rs
+19
View File
@@ -1,9 +1,15 @@
//! 优雅关闭信号处理。
use std::sync::OnceLock;
use tokio::signal;
use tokio::sync::broadcast;
static SHUTDOWN_REQUESTS: OnceLock<broadcast::Sender<()>> = OnceLock::new();
/// 等待进程关闭信号（Ctrl+C 或 SIGTERM）。
pub async fn shutdown_signal() {
let mut requested = shutdown_sender().subscribe();
let ctrl_c = async {
signal::ctrl_c()
.await
@@ -22,7 +28,20 @@ pub async fn shutdown_signal() {
let terminate = std::future::pending::<()>();
tokio::select! {
_ = requested.recv() => {},
() = ctrl_c => {},
() = terminate => {},
}
}
/// 请求进程按优雅关闭路径退出。
pub fn request_shutdown() {
let _ = shutdown_sender().send(());
}
fn shutdown_sender() -> &'static broadcast::Sender<()> {
SHUTDOWN_REQUESTS.get_or_init(|| {
let (sender, _receiver) = broadcast::channel(16);
sender
})
}
backend/src/runtime/tasks/cleanup.rs
+92
View File
@@ -0,0 +1,92 @@
//! 通用过期数据清理任务骨架。
use std::{fmt::Display, future::Future, time::Duration};
use chrono::{DateTime, Utc};
use tokio::time::interval;
use tracing::{debug, info, warn};
use super::coordinator::SchedulerHandle;
/// 支持按时间清理过期记录的存储。
pub(crate) trait ExpiredCleanupStore: Send + Sync + 'static {
type Error: Display + Send + Sync + 'static;
fn delete_expired_at(
&self,
now: DateTime<Utc>,
) -> impl Future<Output = Result<u64, Self::Error>> + Send + '_;
}
#[derive(Clone, Copy)]
pub(crate) struct CleanupTaskMessages {
pub started: &'static str,
pub deleted: &'static str,
pub empty: &'static str,
pub failed: &'static str,
pub stopped: &'static str,
}
/// 周期性过期数据清理任务。
pub(crate) struct CleanupTask<S> {
store: S,
interval_secs: u64,
messages: CleanupTaskMessages,
}
impl<S> CleanupTask<S>
where
S: ExpiredCleanupStore,
{
pub(crate) fn new(store: S, interval_secs: u64, messages: CleanupTaskMessages) -> Self {
Self {
store,
interval_secs,
messages,
}
}
pub(crate) fn start(self) -> SchedulerHandle {
let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1);
tokio::spawn(async move {
let mut ticker = interval(Duration::from_secs(self.interval_secs));
info!(
interval_secs = self.interval_secs,
"{}", self.messages.started
);
loop {
tokio::select! {
_ = ticker.tick() => {
match self.cleanup_once().await {
Ok(deleted) if deleted > 0 => {
info!(deleted, "{}", self.messages.deleted);
}
Ok(_) => {
debug!("{}", self.messages.empty);
}
Err(error) => {
warn!(error = %error, "{}", self.messages.failed);
}
}
}
_ = shutdown_rx.recv() => {
info!("{}", self.messages.stopped);
break;
}
}
}
});
SchedulerHandle::new(shutdown_tx)
}
pub(crate) async fn cleanup_once(&self) -> Result<u64, S::Error> {
self.cleanup_once_at(Utc::now()).await
}
pub(crate) async fn cleanup_once_at(&self, now: DateTime<Utc>) -> Result<u64, S::Error> {
self.store.delete_expired_at(now).await
}
}
backend/src/runtime/tasks/cookie_cleanup.rs
+28 -40
View File
@@ -1,73 +1,61 @@
//! Cookie 清理任务。
use std::time::Duration;
use std::future::Future;
use chrono::{DateTime, Utc};
use tokio::time::interval;
use tracing::{debug, info, warn};
use crate::upstream::accounts::cookies::{SqliteCookieStore, SqliteCookieStoreError};
use super::coordinator::SchedulerHandle;
use super::{
cleanup::{CleanupTask, CleanupTaskMessages, ExpiredCleanupStore},
coordinator::SchedulerHandle,
};
const DEFAULT_INTERVAL_SECS: u64 = 5 * 60;
const MESSAGES: CleanupTaskMessages = CleanupTaskMessages {
started: "Cookie 清理任务已启动",
deleted: "已清理过期 Cookie",
empty: "没有需要清理的过期 Cookie",
failed: "清理过期 Cookie 失败",
stopped: "Cookie 清理任务已关闭",
};
/// Cookie 清理任务。
pub struct CookieCleanupTask {
store: SqliteCookieStore,
interval_secs: u64,
task: CleanupTask<SqliteCookieStore>,
}
impl CookieCleanupTask {
/// 构造默认 Cookie 清理任务。
pub fn new(store: SqliteCookieStore) -> Self {
Self {
store,
interval_secs: DEFAULT_INTERVAL_SECS,
task: CleanupTask::new(store, DEFAULT_INTERVAL_SECS, MESSAGES),
}
}
/// 启动后台清理任务。
pub fn start(self) -> SchedulerHandle {
let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1);
tokio::spawn(async move {
let mut ticker = interval(Duration::from_secs(self.interval_secs));
info!(interval_secs = self.interval_secs, "Cookie 清理任务已启动");
loop {
tokio::select! {
_ = ticker.tick() => {
match self.cleanup_once().await {
Ok(deleted) if deleted > 0 => {
info!(deleted, "已清理过期 Cookie");
}
Ok(_) => {
debug!("没有需要清理的过期 Cookie");
}
Err(error) => {
warn!(error = %error, "清理过期 Cookie 失败");
}
}
}
_ = shutdown_rx.recv() => {
info!("Cookie 清理任务已关闭");
break;
}
}
}
});
SchedulerHandle::new(shutdown_tx)
self.task.start()
}
/// 执行一次清理。
pub async fn cleanup_once(&self) -> Result<u64, SqliteCookieStoreError> {
self.cleanup_once_at(Utc::now()).await
self.task.cleanup_once().await
}
/// 在指定时间点执行一次清理。
pub async fn cleanup_once_at(&self, now: DateTime<Utc>) -> Result<u64, SqliteCookieStoreError> {
self.store.cleanup_expired(now).await
self.task.cleanup_once_at(now).await
}
}
impl ExpiredCleanupStore for SqliteCookieStore {
type Error = SqliteCookieStoreError;
fn delete_expired_at(
&self,
now: DateTime<Utc>,
) -> impl Future<Output = Result<u64, Self::Error>> + Send + '_ {
self.cleanup_expired(now)
}
}
backend/src/runtime/tasks/coordinator.rs
+17 -1
View File
@@ -1,6 +1,6 @@
//! 后台任务协调器。
use std::path::PathBuf;
use std::{path::PathBuf, sync::Arc};
use tokio::task::JoinHandle;
@@ -96,6 +96,12 @@ impl TaskCoordinator {
services.fingerprint.build_number.clone(),
),
);
if let Some(pool) = &services.websocket_pool {
coordinator.push(
"websocket_pool",
SchedulerHandle::from_websocket_pool(pool.clone()),
);
}
coordinator
}
@@ -117,6 +123,8 @@ pub enum SchedulerHandle {
Channel(tokio::sync::mpsc::Sender<()>),
/// 直接持有 tokio 任务句柄。
JoinHandle(JoinHandle<()>),
/// 关闭上游 WebSocket 连接池。
WebSocketPool(Arc<crate::upstream::transport::CodexWebSocketPool>),
}
impl SchedulerHandle {
@@ -130,6 +138,11 @@ impl SchedulerHandle {
Self::JoinHandle(handle)
}
/// 使用 WebSocket 连接池构造句柄。
pub fn from_websocket_pool(pool: Arc<crate::upstream::transport::CodexWebSocketPool>) -> Self {
Self::WebSocketPool(pool)
}
/// 关闭任务。
pub async fn shutdown(self) {
match self {
@@ -139,6 +152,9 @@ impl SchedulerHandle {
Self::JoinHandle(handle) => {
handle.abort();
}
Self::WebSocketPool(pool) => {
pool.shutdown().await;
}
}
}
}
backend/src/runtime/tasks/mod.rs
+1
View File
@@ -1,5 +1,6 @@
//! 后台任务子模块。
mod cleanup;
pub mod cookie_cleanup;
pub mod coordinator;
pub mod fingerprint_update;
backend/src/runtime/tasks/session_affinity_cleanup.rs
+29 -43
View File
@@ -1,72 +1,47 @@
//! 会话亲和性清理任务。
use std::time::Duration;
use std::future::Future;
use chrono::{DateTime, Utc};
use tokio::time::interval;
use tracing::{debug, info, warn};
use crate::proxy::dispatch::session_affinity::{
SqliteSessionAffinityStore, SqliteSessionAffinityStoreError,
};
use super::coordinator::SchedulerHandle;
use super::{
cleanup::{CleanupTask, CleanupTaskMessages, ExpiredCleanupStore},
coordinator::SchedulerHandle,
};
const MESSAGES: CleanupTaskMessages = CleanupTaskMessages {
started: "会话亲和性清理任务已启动",
deleted: "已清理过期会话亲和性",
empty: "没有需要清理的过期会话亲和性",
failed: "清理过期会话亲和性失败",
stopped: "会话亲和性清理任务已关闭",
};
/// 会话亲和性清理任务。
pub struct SessionAffinityCleanupTask {
store: SqliteSessionAffinityStore,
interval_secs: u64,
task: CleanupTask<SqliteSessionAffinityStore>,
}
impl SessionAffinityCleanupTask {
/// 构造会话亲和性清理任务。
pub fn new(store: SqliteSessionAffinityStore, interval_secs: u64) -> Self {
Self {
store,
interval_secs,
task: CleanupTask::new(store, interval_secs, MESSAGES),
}
}
/// 启动后台清理任务。
pub fn start(self) -> SchedulerHandle {
let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1);
tokio::spawn(async move {
let mut ticker = interval(Duration::from_secs(self.interval_secs));
info!(
interval_secs = self.interval_secs,
"会话亲和性清理任务已启动"
);
loop {
tokio::select! {
_ = ticker.tick() => {
match self.cleanup_once().await {
Ok(deleted) if deleted > 0 => {
info!(deleted, "已清理过期会话亲和性");
}
Ok(_) => {
debug!("没有需要清理的过期会话亲和性");
}
Err(error) => {
warn!(error = %error, "清理过期会话亲和性失败");
}
}
}
_ = shutdown_rx.recv() => {
info!("会话亲和性清理任务已关闭");
break;
}
}
}
});
SchedulerHandle::new(shutdown_tx)
self.task.start()
}
/// 执行一次清理。
pub async fn cleanup_once(&self) -> Result<u64, SqliteSessionAffinityStoreError> {
self.cleanup_once_at(Utc::now()).await
self.task.cleanup_once().await
}
/// 在指定时间点执行一次清理。
@@ -74,6 +49,17 @@ impl SessionAffinityCleanupTask {
&self,
now: DateTime<Utc>,
) -> Result<u64, SqliteSessionAffinityStoreError> {
self.store.delete_expired(now).await
self.task.cleanup_once_at(now).await
}
}
impl ExpiredCleanupStore for SqliteSessionAffinityStore {
type Error = SqliteSessionAffinityStoreError;
fn delete_expired_at(
&self,
now: DateTime<Utc>,
) -> impl Future<Output = Result<u64, Self::Error>> + Send + '_ {
self.delete_expired(now)
}
}
backend/src/runtime/tasks/session_cleanup.rs
+29 -43
View File
@@ -1,74 +1,60 @@
//! 会话清理任务。
use std::time::Duration;
use std::future::Future;
use chrono::{DateTime, Utc};
use tokio::time::interval;
use tracing::{debug, info, warn};
use crate::admin::auth::service::SqliteAdminSessionStore;
use super::coordinator::SchedulerHandle;
use super::{
cleanup::{CleanupTask, CleanupTaskMessages, ExpiredCleanupStore},
coordinator::SchedulerHandle,
};
const MESSAGES: CleanupTaskMessages = CleanupTaskMessages {
started: "管理员会话清理任务已启动",
deleted: "已清理过期管理员会话",
empty: "没有需要清理的过期管理员会话",
failed: "清理过期管理员会话失败",
stopped: "管理员会话清理任务已关闭",
};
/// 管理员会话清理任务。
pub struct SessionCleanupTask {
store: SqliteAdminSessionStore,
interval_secs: u64,
task: CleanupTask<SqliteAdminSessionStore>,
}
impl SessionCleanupTask {
/// 构造管理员会话清理任务。
pub fn new(store: SqliteAdminSessionStore, interval_secs: u64) -> Self {
Self {
store,
interval_secs,
task: CleanupTask::new(store, interval_secs, MESSAGES),
}
}
/// 启动后台清理任务。
pub fn start(self) -> SchedulerHandle {
let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1);
tokio::spawn(async move {
let mut ticker = interval(Duration::from_secs(self.interval_secs));
info!(
interval_secs = self.interval_secs,
"管理员会话清理任务已启动"
);
loop {
tokio::select! {
_ = ticker.tick() => {
match self.cleanup_once().await {
Ok(deleted) if deleted > 0 => {
info!(deleted, "已清理过期管理员会话");
}
Ok(_) => {
debug!("没有需要清理的过期管理员会话");
}
Err(error) => {
warn!(error = %error, "清理过期管理员会话失败");
}
}
}
_ = shutdown_rx.recv() => {
info!("管理员会话清理任务已关闭");
break;
}
}
}
});
SchedulerHandle::new(shutdown_tx)
self.task.start()
}
/// 执行一次清理。
pub async fn cleanup_once(&self) -> Result<u64, sqlx::Error> {
self.cleanup_once_at(Utc::now()).await
self.task.cleanup_once().await
}
/// 在指定时间点执行一次清理。
pub async fn cleanup_once_at(&self, now: DateTime<Utc>) -> Result<u64, sqlx::Error> {
self.store.cleanup_expired_sessions(now).await
self.task.cleanup_once_at(now).await
}
}
impl ExpiredCleanupStore for SqliteAdminSessionStore {
type Error = sqlx::Error;
fn delete_expired_at(
&self,
now: DateTime<Utc>,
) -> impl Future<Output = Result<u64, Self::Error>> + Send + '_ {
self.cleanup_expired_sessions(now)
}
}
backend/src/upstream/accounts/mod.rs
+1
View File
@@ -16,3 +16,4 @@ pub mod service;
pub mod store;
/// 刷新租约存储与 JWT claims 解码。
pub mod token_refresh;
mod window;
backend/src/upstream/accounts/pool.rs
+9 -28
View File
@@ -11,9 +11,11 @@ use chrono::{DateTime, Duration, Utc};
use indexmap::IndexMap;
use thiserror::Error;
use crate::infra::time::parse_rfc3339_utc;
use crate::upstream::accounts::quota::{quota_snapshot_limit_reached, quota_snapshot_reset_at};
use crate::upstream::accounts::store::AccountStore;
use crate::upstream::accounts::token_refresh::{jwt_expiry, JwtExpiry};
use crate::upstream::accounts::window::should_reset_usage_window;
use crate::upstream::accounts::{
model::{Account, AccountModelUsageDelta, AccountStatus, AccountUsageDelta},
service::AccountService,
@@ -737,7 +739,12 @@ impl AccountPool {
let Some(account) = self.accounts.get_mut(account_id) else {
return false;
};
if should_reset_window_counters(account, new_reset_at, limit_window_seconds) {
if should_reset_usage_window(
account.window_reset_at,
account.limit_window_seconds,
new_reset_at,
limit_window_seconds,
) {
reset_window_counters(account);
account.window_started_at = Some(now);
}
@@ -1316,32 +1323,6 @@ fn refresh_quota_window(account: &mut Account, now: DateTime<Utc>) {
}
}
fn should_reset_window_counters(
account: &Account,
new_reset_at: DateTime<Utc>,
limit_window_seconds: Option<u64>,
) -> bool {
let Some(old_reset_at) = account.window_reset_at else {
return false;
};
if old_reset_at == new_reset_at {
return false;
}
let drift = old_reset_at
.signed_duration_since(new_reset_at)
.num_seconds()
.unsigned_abs();
let window_seconds = limit_window_seconds
.or(account.limit_window_seconds)
.unwrap_or(0);
let threshold = if window_seconds > 0 {
window_seconds / 2
} else {
3_600
};
drift >= threshold
}
fn reset_window_counters(account: &mut Account) {
account.window_request_count = 0;
account.window_input_tokens = 0;
@@ -1406,7 +1387,7 @@ fn compare_last_used(a: Option<&str>, b: Option<&str>) -> Ordering {
fn last_used_millis(value: Option<&str>) -> i64 {
value
.and_then(|value| DateTime::parse_from_rfc3339(value).ok())
.and_then(|value| parse_rfc3339_utc(value).ok())
.map(|datetime| datetime.timestamp_millis())
.unwrap_or(0)
}
backend/src/upstream/accounts/store.rs
+7 -41
View File
@@ -7,13 +7,18 @@ use serde_json::Value;
use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite, SqlitePool};
use thiserror::Error;
use crate::infra::json::{decode_cursor, page_offset, NumberedPage, Page};
use crate::infra::{
format::nonnegative_i64_to_u64,
json::{decode_cursor, page_offset, NumberedPage, Page},
time::parse_optional_rfc3339_utc as parse_optional_rfc3339,
};
use crate::upstream::accounts::model::{
Account, AccountModelUsageDelta, AccountStatus, AccountUsageDelta,
};
use crate::upstream::accounts::quota::{
quota_snapshot_limit_window_seconds, quota_snapshot_reset_at,
};
use crate::upstream::accounts::window::should_reset_usage_window;
// ============================================================================
// SQL 常量
@@ -1806,33 +1811,6 @@ fn sqlite_usage_delta(usage: AccountUsageDelta) -> UsageDelta {
}
}
fn should_reset_usage_window(
existing_reset_at: Option<DateTime<Utc>>,
existing_limit_window_seconds: Option<u64>,
new_reset_at: DateTime<Utc>,
new_limit_window_seconds: Option<u64>,
) -> bool {
let Some(existing_reset_at) = existing_reset_at else {
return false;
};
if existing_reset_at == new_reset_at {
return false;
}
let drift = existing_reset_at
.signed_duration_since(new_reset_at)
.num_seconds()
.unsigned_abs();
let window_seconds = new_limit_window_seconds
.or(existing_limit_window_seconds)
.unwrap_or(0);
let threshold = if window_seconds > 0 {
window_seconds / 2
} else {
3_600
};
drift >= threshold
}
fn quota_plan_type(quota_json: &str) -> Option<String> {
serde_json::from_str::<Value>(quota_json)
.ok()?
@@ -1902,18 +1880,6 @@ fn status_from_db(value: &str) -> SqliteAccountStoreResult<AccountStatus> {
}
}
fn parse_optional_rfc3339(value: Option<&str>) -> SqliteAccountStoreResult<Option<DateTime<Utc>>> {
value.map(parse_rfc3339).transpose()
}
fn parse_rfc3339(value: &str) -> SqliteAccountStoreResult<DateTime<Utc>> {
Ok(DateTime::parse_from_rfc3339(value)?.with_timezone(&Utc))
}
fn nonnegative_i64_to_u64(value: i64) -> u64 {
value.max(0).cast_unsigned()
}
fn u64_to_i64_saturating(value: u64) -> i64 {
value.min(i64::MAX as u64) as i64
}
@@ -1931,7 +1897,7 @@ async fn count_account_metadata(
let mut builder = QueryBuilder::<Sqlite>::new("select count(*) from accounts");
push_account_metadata_search(&mut builder, search);
let (total,): (i64,) = builder.build_query_as().fetch_one(pool).await?;
Ok(total.max(0).cast_unsigned())
Ok(nonnegative_i64_to_u64(total))
}
fn push_account_metadata_search(builder: &mut QueryBuilder<Sqlite>, search: Option<&str>) {
backend/src/upstream/accounts/window.rs
+31
View File
@@ -0,0 +1,31 @@
//! 账号用量窗口规则。
use chrono::{DateTime, Utc};
pub(crate) fn should_reset_usage_window(
existing_reset_at: Option<DateTime<Utc>>,
existing_limit_window_seconds: Option<u64>,
new_reset_at: DateTime<Utc>,
new_limit_window_seconds: Option<u64>,
) -> bool {
let Some(existing_reset_at) = existing_reset_at else {
return false;
};
if existing_reset_at == new_reset_at {
return false;
}
let drift = existing_reset_at
.signed_duration_since(new_reset_at)
.num_seconds()
.unsigned_abs();
let window_seconds = new_limit_window_seconds
.or(existing_limit_window_seconds)
.unwrap_or(0);
let threshold = if window_seconds > 0 {
window_seconds / 2
} else {
3_600
};
drift >= threshold
}
backend/src/upstream/protocol/responses.rs
+20 -6
View File
@@ -1,13 +1,16 @@
use std::collections::BTreeSet;
use std::{collections::BTreeSet, time::Instant};
use serde::{Deserialize, Serialize};
use serde_json::{json, Map, Value};
use crate::upstream::{
models::ParsedModelName,
protocol::{
schema::reconvert_tuple_values,
sse::{parse_sse_events, SseError, SseEvent},
use crate::{
infra::time::elapsed_millis_i64,
upstream::{
models::ParsedModelName,
protocol::{
schema::reconvert_tuple_values,
sse::{parse_sse_events, SseError, SseEvent},
},
},
};
@@ -150,6 +153,17 @@ pub fn response_body_has_first_event(body_bytes: &[u8]) -> bool {
.is_ok_and(|events| events.iter().any(response_sse_event_has_first_output))
}
/// 已收到首个有效 Responses SSE 输出事件时记录首 token 耗时。
pub fn update_first_response_event_ms(
started_at: Instant,
body_bytes: &[u8],
first_token_ms: &mut Option<i64>,
) {
if first_token_ms.is_none() && response_body_has_first_event(body_bytes) {
*first_token_ms = Some(elapsed_millis_i64(started_at).max(1));
}
}
fn response_sse_event_has_first_output(event: &SseEvent) -> bool {
let data = event.data.trim();
if data.is_empty() || data == "[DONE]" {
backend/src/upstream/protocol/sse.rs
+26
View File
@@ -27,6 +27,16 @@ pub fn sse_body_has_done(body: &str) -> bool {
.ends_with(DONE_SSE_FRAME.trim_end_matches(['\r', '\n']))
}
/// 返回下一个完整 SSE 帧分隔符位置与长度。
pub fn sse_frame_separator(input: &str) -> Option<(usize, usize)> {
sse_frame_separator_bytes(input.as_bytes())
}
/// 返回下一个完整 SSE 帧结束位置（含分隔符）。
pub fn sse_frame_end(bytes: &[u8]) -> Option<usize> {
sse_frame_separator_bytes(bytes).map(|(position, separator_len)| position + separator_len)
}
/// SSE 解析错误。
#[derive(Debug, Error, PartialEq, Eq)]
pub enum SseError {
@@ -170,6 +180,22 @@ fn track_event_buffer_bytes(current_bytes: &mut usize, line: &str) -> Result<(),
Ok(())
}
fn sse_frame_separator_bytes(bytes: &[u8]) -> Option<(usize, usize)> {
let lf = bytes
.windows(2)
.position(|window| window == b"\n\n")
.map(|position| (position, 2));
let crlf = bytes
.windows(4)
.position(|window| window == b"\r\n\r\n")
.map(|position| (position, 4));
match (lf, crlf) {
(Some(left), Some(right)) => Some(if left.0 <= right.0 { left } else { right }),
(Some(found), None) | (None, Some(found)) => Some(found),
(None, None) => None,
}
}
fn is_sse_metadata_line(line: &str) -> bool {
line.starts_with("event:")
|| line.starts_with("data:")
backend/src/upstream/protocol/websocket.rs
+1 -1
View File
@@ -140,7 +140,7 @@ pub fn websocket_event_to_sse_frame(raw: &str) -> Option<String> {
Some(encode_sse_event(&event, raw))
}
fn websocket_event_type(raw: &str) -> Option<String> {
pub(crate) fn websocket_event_type(raw: &str) -> Option<String> {
serde_json::from_str::<Value>(raw).ok().and_then(|value| {
value
.get("type")
backend/src/upstream/transport/client.rs
+13 -64
View File
@@ -13,7 +13,7 @@ use futures::{Stream, TryStreamExt};
use reqwest::{
header::{
HeaderMap, HeaderName, HeaderValue, ACCEPT, ACCEPT_ENCODING, AUTHORIZATION, CONTENT_TYPE,
COOKIE, RETRY_AFTER, SET_COOKIE, USER_AGENT,
COOKIE, RETRY_AFTER, USER_AGENT,
},
Client, Response as ReqwestResponse, StatusCode,
};
@@ -26,8 +26,8 @@ use crate::upstream::fingerprint::Fingerprint;
use crate::upstream::models::BackendModelEntry;
use crate::upstream::protocol::events::extract_sse_usage;
use crate::upstream::protocol::responses::{
http_sse_fallback_allowed, response_body_has_first_event, transport_for_request,
CodexCompactRequest, CodexResponsesRequest, CodexTransport,
http_sse_fallback_allowed, transport_for_request, CodexCompactRequest, CodexResponsesRequest,
CodexTransport,
};
use crate::upstream::protocol::sse::SseError;
use crate::upstream::protocol::websocket::{
@@ -39,6 +39,7 @@ use super::headers::{
build_ordered_codex_base_headers, insert_optional_header, insert_ordered_headers,
websocket_header_pairs,
};
use super::response_meta;
use super::tls::{build_reqwest_client_with_custom_ca, custom_ca_env_cache_key, CustomCaError};
use super::websocket::{
execute_response_create_request_stream_with_pool, execute_response_create_request_with_pool,
@@ -366,9 +367,9 @@ impl CodexBackendClient {
.send()
.await?;
let status = response.status();
let turn_state = turn_state(&response);
let set_cookie_headers = set_cookie_headers(&response);
let rate_limit_headers = rate_limit_headers(response.headers());
let turn_state = response_meta::turn_state(response.headers());
let set_cookie_headers = response_meta::set_cookie_headers(response.headers());
let rate_limit_headers = response_meta::rate_limit_headers(response.headers());
let retry_after_seconds = retry_after_seconds(response.headers(), None);
if !status.is_success() {
@@ -387,7 +388,7 @@ impl CodexBackendClient {
let mut stream = response.bytes_stream();
while let Some(chunk) = stream.try_next().await? {
body_bytes.extend_from_slice(&chunk);
update_first_token_ms(started_at, &body_bytes, &mut first_token_ms);
response_meta::update_first_token_ms(started_at, &body_bytes, &mut first_token_ms);
}
let body = String::from_utf8_lossy(&body_bytes).into_owned();
let usage = extract_sse_usage(&body).map_err(CodexClientError::InvalidSse)?;
@@ -418,9 +419,9 @@ impl CodexBackendClient {
.send()
.await?;
let status = response.status();
let turn_state = turn_state(&response);
let set_cookie_headers = set_cookie_headers(&response);
let rate_limit_headers = rate_limit_headers(response.headers());
let turn_state = response_meta::turn_state(response.headers());
let set_cookie_headers = response_meta::set_cookie_headers(response.headers());
let rate_limit_headers = response_meta::rate_limit_headers(response.headers());
let retry_after_seconds = retry_after_seconds(response.headers(), None);
if !status.is_success() {
@@ -704,8 +705,8 @@ impl CodexBackendClient {
.await?;
let status = response.status();
let set_cookie_headers = set_cookie_headers(&response);
let rate_limit_headers = rate_limit_headers(response.headers());
let set_cookie_headers = response_meta::set_cookie_headers(response.headers());
let rate_limit_headers = response_meta::rate_limit_headers(response.headers());
let retry_after_seconds = retry_after_seconds(response.headers(), None);
let body = response.text().await?;
@@ -953,20 +954,6 @@ impl CodexModelCatalogClient for CodexBackendClient {
// Response helpers
// ---------------------------------------------------------------------------
fn turn_state(response: &ReqwestResponse) -> Option<String> {
response
.headers()
.get("x-codex-turn-state")
.and_then(|value| value.to_str().ok())
.map(ToString::to_string)
}
fn update_first_token_ms(started_at: Instant, body_bytes: &[u8], first_token_ms: &mut Option<i64>) {
if first_token_ms.is_none() && response_body_has_first_event(body_bytes) {
*first_token_ms = Some(elapsed_millis_i64(started_at).max(1));
}
}
fn log_websocket_pool_decision(
request_id: &str,
account_id: Option<&str>,
@@ -1016,44 +1003,6 @@ impl WebSocketPoolLogContext {
}
}
fn elapsed_millis_i64(started_at: Instant) -> i64 {
started_at.elapsed().as_millis().min(i64::MAX as u128) as i64
}
fn set_cookie_headers(response: &ReqwestResponse) -> Vec<String> {
response
.headers()
.get_all(SET_COOKIE)
.iter()
.filter_map(|value| value.to_str().ok().map(ToString::to_string))
.collect()
}
fn rate_limit_headers(headers: &HeaderMap) -> Vec<(String, String)> {
headers
.iter()
.filter(|(name, _)| is_rate_limit_header(name.as_str()))
.filter_map(|(name, value)| {
value
.to_str()
.ok()
.map(|value| (name.as_str().to_string(), value.to_string()))
})
.collect()
}
fn is_rate_limit_header(name: &str) -> bool {
let name = name.to_ascii_lowercase();
name == "retry-after"
|| name.contains("ratelimit")
|| name.contains("rate-limit")
|| name.starts_with("x-codex-primary-")
|| name.starts_with("x-codex-secondary-")
|| name.starts_with("x-codex-code-review-")
|| name.starts_with("x-codex-review-")
|| name.starts_with("x-code-review-")
}
pub(super) fn retry_after_seconds(headers: &HeaderMap, body: Option<&str>) -> Option<u64> {
headers
.get(RETRY_AFTER)
backend/src/upstream/transport/mod.rs
+1
View File
@@ -3,6 +3,7 @@
pub mod client;
pub mod endpoints;
pub mod headers;
mod response_meta;
pub mod tls;
pub mod usage;
pub mod websocket;
backend/src/upstream/transport/response_meta.rs
+55
View File
@@ -0,0 +1,55 @@
//! 上游响应元数据提取辅助。
use std::time::Instant;
use reqwest::header::{HeaderMap, SET_COOKIE};
use crate::upstream::protocol::responses::update_first_response_event_ms;
pub(super) fn turn_state(headers: &HeaderMap) -> Option<String> {
headers
.get("x-codex-turn-state")
.and_then(|value| value.to_str().ok())
.map(ToString::to_string)
}
pub(super) fn set_cookie_headers(headers: &HeaderMap) -> Vec<String> {
headers
.get_all(SET_COOKIE)
.iter()
.filter_map(|value| value.to_str().ok().map(ToString::to_string))
.collect()
}
pub(super) fn rate_limit_headers(headers: &HeaderMap) -> Vec<(String, String)> {
headers
.iter()
.filter(|(name, _)| is_rate_limit_header(name.as_str()))
.filter_map(|(name, value)| {
value
.to_str()
.ok()
.map(|value| (name.as_str().to_string(), value.to_string()))
})
.collect()
}
pub(super) fn update_first_token_ms(
started_at: Instant,
body_bytes: &[u8],
first_token_ms: &mut Option<i64>,
) {
update_first_response_event_ms(started_at, body_bytes, first_token_ms);
}
fn is_rate_limit_header(name: &str) -> bool {
let name = name.to_ascii_lowercase();
name == "retry-after"
|| name.contains("ratelimit")
|| name.contains("rate-limit")
|| name.starts_with("x-codex-primary-")
|| name.starts_with("x-codex-secondary-")
|| name.starts_with("x-codex-code-review-")
|| name.starts_with("x-codex-review-")
|| name.starts_with("x-code-review-")
}
backend/src/upstream/transport/websocket.rs
+9 -68
View File
@@ -11,7 +11,6 @@ use std::{
use bytes::Bytes;
use chrono::Utc;
use futures::{channel::mpsc, SinkExt, Stream, StreamExt};
use serde_json::Value;
use thiserror::Error;
use tokio::{sync::Mutex, time::timeout};
use tokio_tungstenite::{connect_async_tls_with_config, Connector};
@@ -19,23 +18,24 @@ use tungstenite::{
self,
extensions::{compression::deflate::DeflateConfig, ExtensionsConfig},
handshake::client::Request as WsRequest,
http::{HeaderMap, Response as WsResponse},
http::Response as WsResponse,
protocol::WebSocketConfig,
Message,
};
use crate::infra::time::china_filename_timestamp_millis;
use crate::upstream::protocol::events::{self, TokenUsage};
use crate::upstream::protocol::responses::{response_body_has_first_event, CodexResponsesRequest};
use crate::upstream::protocol::responses::CodexResponsesRequest;
use crate::upstream::protocol::sse::SseError;
use crate::upstream::protocol::websocket::{
classify_websocket_error_frame, is_terminal_websocket_event,
retry_after_seconds_from_wrapped_error_headers, websocket_event_to_sse_frame,
websocket_incomplete_response_reason, websocket_metadata_turn_state,
websocket_event_type, websocket_incomplete_response_reason, websocket_metadata_turn_state,
websocket_response_completed_parse_error, websocket_response_create_payload_text,
OpeningAuditHeader, OpeningAuditSnapshot, WebSocketAuditArtifact,
};
use super::response_meta;
use super::websocket_pool::{
CodexWebSocketConnectionMetadata, CodexWebSocketPool, CodexWebSocketPoolKey, CodexWsStream,
PooledWebSocketConnection, WebSocketPoolAcquire, WebSocketPoolDecision,
@@ -373,7 +373,7 @@ fn websocket_opening_error(response: &WsResponse<Option<Vec<u8>>>) -> CodexWebSo
status_code,
retry_after_seconds,
body,
set_cookie_headers: websocket_set_cookie_headers(response.headers()),
set_cookie_headers: response_meta::set_cookie_headers(response.headers()),
}
}
@@ -795,7 +795,7 @@ async fn collect_websocket_response(
let forwarded = if let Some(frame) = websocket_event_to_sse_frame(&raw) {
body.push_str(&frame);
saw_response_frame = true;
update_first_token_ms(started_at, body.as_bytes(), &mut first_token_ms);
response_meta::update_first_token_ms(started_at, body.as_bytes(), &mut first_token_ms);
true
} else {
false
@@ -835,23 +835,13 @@ fn reused_connection_died_before_first_frame(
}
}
fn update_first_token_ms(started_at: Instant, body_bytes: &[u8], first_token_ms: &mut Option<i64>) {
if first_token_ms.is_none() && response_body_has_first_event(body_bytes) {
*first_token_ms = Some(elapsed_millis_i64(started_at).max(1));
}
}
fn elapsed_millis_i64(started_at: Instant) -> i64 {
started_at.elapsed().as_millis().min(i64::MAX as u128) as i64
}
fn websocket_connection_metadata(
response: &WsResponse<Option<Vec<u8>>>,
) -> CodexWebSocketConnectionMetadata {
CodexWebSocketConnectionMetadata {
turn_state: websocket_response_turn_state(response.headers()),
set_cookie_headers: websocket_set_cookie_headers(response.headers()),
rate_limit_headers: websocket_rate_limit_headers(response.headers()),
turn_state: response_meta::turn_state(response.headers()),
set_cookie_headers: response_meta::set_cookie_headers(response.headers()),
rate_limit_headers: response_meta::rate_limit_headers(response.headers()),
handshake_status: response.status().as_u16(),
}
}
@@ -1071,55 +1061,6 @@ async fn discard_stream_websocket(
// Header / metadata helpers
// ---------------------------------------------------------------------------
fn websocket_response_turn_state(headers: &HeaderMap) -> Option<String> {
headers
.get("x-codex-turn-state")
.and_then(|value| value.to_str().ok())
.map(ToString::to_string)
}
fn websocket_set_cookie_headers(headers: &HeaderMap) -> Vec<String> {
headers
.get_all("set-cookie")
.iter()
.filter_map(|value| value.to_str().ok().map(ToString::to_string))
.collect()
}
fn websocket_rate_limit_headers(headers: &HeaderMap) -> Vec<(String, String)> {
headers
.iter()
.filter(|(name, _)| is_rate_limit_header(name.as_str()))
.filter_map(|(name, value)| {
value
.to_str()
.ok()
.map(|value| (name.as_str().to_string(), value.to_string()))
})
.collect()
}
fn is_rate_limit_header(name: &str) -> bool {
let name = name.to_ascii_lowercase();
name == "retry-after"
|| name.contains("ratelimit")
|| name.contains("rate-limit")
|| name.starts_with("x-codex-primary-")
|| name.starts_with("x-codex-secondary-")
|| name.starts_with("x-codex-code-review-")
|| name.starts_with("x-codex-review-")
|| name.starts_with("x-code-review-")
}
fn websocket_event_type(raw: &str) -> Option<String> {
serde_json::from_str::<Value>(raw).ok().and_then(|value| {
value
.get("type")
.and_then(Value::as_str)
.map(ToString::to_string)
})
}
fn websocket_rate_limit_event_headers(raw: &str) -> Option<Vec<(String, String)>> {
events::parse_rate_limits_event_raw(raw)
.map(|parsed| events::rate_limits_to_header_pairs(&parsed))
backend/src/upstream/transport/websocket_pool.rs
+3 -2
View File
@@ -271,7 +271,7 @@ impl CodexWebSocketPool {
break;
};
let pool = CodexWebSocketPool { inner, config };
if pool.is_shutting_down().await {
if pool.is_shutdown().await {
break;
}
pool.maintain_idle_connections().await;
@@ -279,7 +279,8 @@ impl CodexWebSocketPool {
});
}
async fn is_shutting_down(&self) -> bool {
/// 返回连接池是否已进入关闭状态。
pub async fn is_shutdown(&self) -> bool {
self.inner.lock().await.shutting_down
}
backend/src/web/assets.rs
+5 -1
View File
@@ -37,7 +37,7 @@ async fn serve_index(State(state): State<AssetState>) -> Response {
}
async fn serve_spa_fallback(State(state): State<AssetState>, uri: Uri) -> Response {
if uri.path() == "/api" || uri.path().starts_with("/api/") {
if is_api_path(uri.path()) {
return (
StatusCode::NOT_FOUND,
Json(json!({
@@ -52,6 +52,10 @@ async fn serve_spa_fallback(State(state): State<AssetState>, uri: Uri) -> Respon
serve_file(&state.dist_dir.join("index.html"), "/").await
}
fn is_api_path(path: &str) -> bool {
path == "/api" || path.starts_with("/api/") || path == "/v1" || path.starts_with("/v1/")
}
async fn serve_asset(
State(state): State<AssetState>,
AxumPath(path): AxumPath<String>,
backend/src/web/headers.rs
+1 -1
View File
@@ -5,7 +5,7 @@ const ASSET_CACHE_CONTROL: &str = "public, max-age=31536000, immutable";
const CONTENT_SECURITY_POLICY: &str = concat!(
"default-src 'self'; ",
"script-src 'self'; ",
"style-src 'self' 'unsafe-inline'; ",
"style-src 'self'; ",
"img-src 'self' data:; ",
"font-src 'self' data:; ",
"connect-src 'self'; ",
backend/tests/admin/accounts/exporting.rs
+170 -16
View File
@@ -1,7 +1,10 @@
use super::*;
use sqlx::Row;
const EXPORT_CONFIRM: &str = "confirm=export_sensitive_accounts";
#[tokio::test]
async fn admin_accounts_export_should_return_cpr_payload_for_all_accounts() {
async fn admin_accounts_export_should_reject_empty_ids() {
let (app, _state, pool, _dir) =
admin_accounts_test_app("admin-accounts-export-all.sqlite", 131).await;
seed_account(
@@ -59,7 +62,7 @@ async fn admin_accounts_export_should_return_cpr_payload_for_all_accounts() {
.oneshot(
Request::builder()
.method("GET")
.uri("/api/admin/accounts/export")
.uri(format!("/api/admin/accounts/export?{EXPORT_CONFIRM}"))
.header("cookie", "cpr_admin_session=session_1")
.body(Body::empty())
.unwrap(),
@@ -68,18 +71,10 @@ async fn admin_accounts_export_should_return_cpr_payload_for_all_accounts() {
.unwrap();
let status = response.status();
let body = response_json(response).await;
let accounts = body["data"]["accounts"].as_array().unwrap();
assert_eq!(status, StatusCode::OK);
assert_eq!(body["data"]["sourceFormat"], "cpr");
assert_eq!(accounts.len(), 2);
assert!(accounts.iter().any(|account| {
account["id"] == "acct_export_a"
&& account["token"]
.as_str()
.is_some_and(|token| !token.is_empty())
&& account["refreshToken"] == "refresh-export-a"
}));
assert_eq!(status, StatusCode::BAD_REQUEST);
assert_eq!(body["code"], 40001);
assert_eq!(body["message"], "account ids are required");
}
#[tokio::test]
@@ -141,7 +136,9 @@ async fn admin_accounts_export_should_filter_by_ids() {
.oneshot(
Request::builder()
.method("GET")
.uri("/api/admin/accounts/export?ids=acct_export_selected")
.uri(format!(
"/api/admin/accounts/export?ids=acct_export_selected&{EXPORT_CONFIRM}"
))
.header("cookie", "cpr_admin_session=session_1")
.body(Body::empty())
.unwrap(),
@@ -164,7 +161,9 @@ async fn admin_accounts_export_should_fail_when_selected_id_is_missing() {
.oneshot(
Request::builder()
.method("GET")
.uri("/api/admin/accounts/export?ids=missing_account")
.uri(format!(
"/api/admin/accounts/export?ids=missing_account&{EXPORT_CONFIRM}"
))
.header("cookie", "cpr_admin_session=session_1")
.body(Body::empty())
.unwrap(),
@@ -209,7 +208,9 @@ async fn admin_accounts_export_payload_should_import_directly() {
.oneshot(
Request::builder()
.method("GET")
.uri("/api/admin/accounts/export?ids=acct_export_roundtrip")
.uri(format!(
"/api/admin/accounts/export?ids=acct_export_roundtrip&{EXPORT_CONFIRM}"
))
.header("cookie", "cpr_admin_session=session_1")
.body(Body::empty())
.unwrap(),
@@ -249,3 +250,156 @@ async fn admin_accounts_export_payload_should_import_directly() {
"refresh-roundtrip"
);
}
#[tokio::test]
async fn admin_accounts_export_should_require_explicit_sensitive_export_confirmation() {
let (app, _state, pool, _dir) =
admin_accounts_test_app("admin-accounts-export-confirm.sqlite", 136).await;
seed_account(
&pool,
NewAccount {
id: "acct_export_confirm".into(),
email: Some("confirm@example.com".into()),
account_id: Some("chatgpt_export_confirm".into()),
user_id: Some("user_export_confirm".into()),
label: None,
plan_type: Some("pro".into()),
access_token: SecretString::new(
test_jwt(
"chatgpt_export_confirm",
Some("user_export_confirm"),
Some("confirm@example.com"),
Some("pro"),
)
.into(),
),
refresh_token: Some(SecretString::new("refresh-confirm".into())),
access_token_expires_at: None,
status: AccountStatus::Active,
added_at: None,
},
)
.await;
let response = app
.oneshot(
Request::builder()
.method("GET")
.uri("/api/admin/accounts/export?ids=acct_export_confirm")
.header("cookie", "cpr_admin_session=session_1")
.header("x-request-id", "req_export_missing_confirmation")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
let status = response.status();
let body = response_json(response).await;
let audit = account_export_audit_row(&pool, "req_export_missing_confirmation").await;
assert_eq!(status, StatusCode::BAD_REQUEST);
assert_eq!(
body["message"],
"account export requires confirm=export_sensitive_accounts"
);
assert_eq!(audit.level, "warn");
assert_eq!(audit.status_code, Some(400));
assert_eq!(
audit.failure_class,
Some("missing_confirmation".to_string())
);
assert_eq!(audit.metadata["accountIds"][0], "acct_export_confirm");
assert!(audit.metadata.get("token").is_none());
assert!(audit.metadata.get("refreshToken").is_none());
}
#[tokio::test]
async fn admin_accounts_export_should_record_success_audit_event() {
let (app, _state, pool, _dir) =
admin_accounts_test_app("admin-accounts-export-audit.sqlite", 137).await;
seed_account(
&pool,
NewAccount {
id: "acct_export_audit".into(),
email: Some("audit@example.com".into()),
account_id: Some("chatgpt_export_audit".into()),
user_id: Some("user_export_audit".into()),
label: None,
plan_type: Some("team".into()),
access_token: SecretString::new(
test_jwt(
"chatgpt_export_audit",
Some("user_export_audit"),
Some("audit@example.com"),
Some("team"),
)
.into(),
),
refresh_token: Some(SecretString::new("refresh-audit".into())),
access_token_expires_at: None,
status: AccountStatus::Active,
added_at: None,
},
)
.await;
let response = app
.oneshot(
Request::builder()
.method("GET")
.uri(format!(
"/api/admin/accounts/export?ids=acct_export_audit&{EXPORT_CONFIRM}"
))
.header("cookie", "cpr_admin_session=session_1")
.header("x-request-id", "req_export_success_audit")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let audit = account_export_audit_row(&pool, "req_export_success_audit").await;
assert_eq!(audit.kind, "admin_account_export");
assert_eq!(audit.level, "warn");
assert_eq!(audit.status_code, Some(200));
assert_eq!(audit.failure_class, None);
assert_eq!(audit.message, "Admin account export succeeded");
assert_eq!(audit.metadata["accountCount"], 1);
assert_eq!(audit.metadata["accountIds"][0], "acct_export_audit");
assert!(audit.metadata.get("token").is_none());
assert!(audit.metadata.get("refreshToken").is_none());
}
#[derive(Debug)]
struct AccountExportAuditRow {
kind: String,
level: String,
status_code: Option<i64>,
failure_class: Option<String>,
message: String,
metadata: Value,
}
async fn account_export_audit_row(pool: &SqlitePool, request_id: &str) -> AccountExportAuditRow {
let row = sqlx::query(
"select kind, level, status_code, failure_class, message, metadata_json \
from usage_records \
where kind = 'admin_account_export' and request_id = ? \
order by created_at desc, id desc \
limit 1",
)
.bind(request_id)
.fetch_one(pool)
.await
.unwrap();
AccountExportAuditRow {
kind: row.get("kind"),
level: row.get("level"),
status_code: row.get("status_code"),
failure_class: row.get("failure_class"),
message: row.get("message"),
metadata: serde_json::from_str(&row.get::<String, _>("metadata_json")).unwrap(),
}
}
backend/tests/admin/auth/session.rs
+179 -54
View File
@@ -4,6 +4,8 @@ use axum::{
};
use chrono::Utc;
use codex_proxy_rs::infra::{database::connect_sqlite, identity::hash_admin_password};
use serde_json::Value;
use sqlx::Row;
use sqlx::SqlitePool;
use tower::util::ServiceExt;
@@ -11,45 +13,7 @@ use crate::support::{config::test_config, http::response_json};
#[tokio::test]
async fn admin_login_should_issue_http_only_session_cookie() {
let dir = tempfile::tempdir().unwrap();
let db = dir.path().join("admin-login.sqlite");
let url = format!("sqlite://{}", db.display());
let pool = connect_sqlite(&url).await.unwrap();
seed_admin_user(&pool, "correct-password").await;
let config = test_config(url);
let stores = codex_proxy_rs::runtime::services::BackgroundTaskStores {
accounts: codex_proxy_rs::upstream::accounts::store::SqliteAccountStore::new(pool.clone()),
admin_sessions: codex_proxy_rs::admin::auth::service::SqliteAdminSessionStore::new(
pool.clone(),
),
cookies: codex_proxy_rs::upstream::accounts::cookies::SqliteCookieStore::new(pool.clone()),
fingerprints: codex_proxy_rs::upstream::fingerprint::FingerprintRepository::new(
pool.clone(),
),
session_affinity:
codex_proxy_rs::proxy::dispatch::session_affinity::SqliteSessionAffinityStore::new(
pool.clone(),
),
refresh_leases: codex_proxy_rs::upstream::accounts::token_refresh::RefreshLeaseStore::new(
pool.clone(),
),
client_keys: codex_proxy_rs::admin::keys::service::SqliteClientKeyStore::new(pool.clone()),
usage_records:
codex_proxy_rs::admin::monitoring::usage_record_store::SqliteUsageRecordStore::new(
pool.clone(),
),
};
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = std::sync::Arc::new(codex_proxy_rs::runtime::services::Services::new(
&config,
stores,
fingerprint,
));
let state = codex_proxy_rs::runtime::state::AppState {
config,
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router().with_state(state);
let (app, _dir, _pool) = admin_login_test_app("admin-login.sqlite", "correct-password").await;
let response = app
.clone()
@@ -75,6 +39,7 @@ async fn admin_login_should_issue_http_only_session_cookie() {
let body = response_json(response).await;
let cookie = cookie.expect("login should set admin session cookie");
assert!(cookie.starts_with("cpr_admin_session="));
assert!(cookie.contains("Secure"));
assert!(cookie.contains("HttpOnly"));
assert!(cookie.contains("SameSite=Lax"));
assert_eq!(body["code"], 200);
@@ -96,11 +61,132 @@ async fn admin_login_should_issue_http_only_session_cookie() {
#[tokio::test]
async fn admin_login_should_reject_client_api_key_as_password_or_authorization() {
let (app, _dir, _pool) =
admin_login_test_app("admin-login-rejects-client-key.sqlite", "correct-password").await;
let response = app
.clone()
.oneshot(
Request::builder()
.method("POST")
.uri("/api/admin/login")
.header("content-type", "application/json")
.header("authorization", "Bearer cpr_not_an_admin_session")
.header("x-request-id", "req_login_bad")
.body(Body::from(r#"{"password":"cpr_not_an_admin_password"}"#))
.unwrap(),
)
.await
.unwrap();
let status = response.status();
let has_cookie = response.headers().get("set-cookie").is_some();
assert_eq!(status, StatusCode::UNAUTHORIZED);
let body = response_json(response).await;
assert!(!has_cookie);
assert_eq!(body["code"], 40102);
}
#[tokio::test]
async fn admin_login_should_throttle_repeated_failures_from_same_source() {
let (app, _dir, _pool) =
admin_login_test_app("admin-login-throttle.sqlite", "correct-password").await;
for _ in 0..5 {
let response = post_login(&app, "wrong-password").await;
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
}
let response = post_login(&app, "correct-password").await;
let status = response.status();
let body = response_json(response).await;
assert_eq!(status, StatusCode::TOO_MANY_REQUESTS);
assert_eq!(body["code"], 42901);
}
#[tokio::test]
async fn admin_login_should_record_success_audit_event() {
let (app, _dir, pool) =
admin_login_test_app("admin-login-success-audit.sqlite", "correct-password").await;
let response = post_login_with_request_id(&app, "correct-password", "req_login_audit_ok").await;
assert_eq!(response.status(), StatusCode::OK);
let audit = admin_login_audit_row(&pool, "req_login_audit_ok").await;
assert_eq!(audit.kind, "admin_auth");
assert_eq!(audit.level, "info");
assert_eq!(audit.route, Some("/api/admin/login".to_string()));
assert_eq!(audit.status_code, Some(200));
assert_eq!(audit.failure_class, None);
assert_eq!(audit.message, "Admin login succeeded");
assert_eq!(audit.metadata["source"], "unknown");
assert_eq!(audit.metadata["usernameProvided"], false);
assert!(audit.metadata.get("username").is_none());
}
#[tokio::test]
async fn admin_login_should_record_invalid_credentials_audit_event() {
let (app, _dir, pool) =
admin_login_test_app("admin-login-invalid-audit.sqlite", "correct-password").await;
let response = app
.clone()
.oneshot(
Request::builder()
.method("POST")
.uri("/api/admin/login")
.header("content-type", "application/json")
.header("x-request-id", "req_login_audit_bad")
.body(Body::from(
r#"{"username":"admin","password":"wrong-password"}"#,
))
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
let audit = admin_login_audit_row(&pool, "req_login_audit_bad").await;
assert_eq!(audit.level, "warn");
assert_eq!(audit.status_code, Some(401));
assert_eq!(audit.failure_class, Some("invalid_credentials".to_string()));
assert_eq!(audit.message, "Admin login failed");
assert_eq!(audit.metadata["usernameProvided"], true);
assert!(audit.metadata.get("username").is_none());
}
#[tokio::test]
async fn admin_login_should_record_throttled_audit_event() {
let (app, _dir, pool) =
admin_login_test_app("admin-login-throttled-audit.sqlite", "correct-password").await;
for index in 0..5 {
let response =
post_login_with_request_id(&app, "wrong-password", &format!("req_login_audit_{index}"))
.await;
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
}
let response =
post_login_with_request_id(&app, "correct-password", "req_login_audit_throttled").await;
assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
let audit = admin_login_audit_row(&pool, "req_login_audit_throttled").await;
assert_eq!(audit.level, "warn");
assert_eq!(audit.status_code, Some(429));
assert_eq!(audit.failure_class, Some("login_throttled".to_string()));
assert_eq!(audit.message, "Admin login throttled");
}
async fn admin_login_test_app(
db_name: &str,
password: &str,
) -> (axum::Router, tempfile::TempDir, SqlitePool) {
let dir = tempfile::tempdir().unwrap();
let db = dir.path().join("admin-login-rejects-client-key.sqlite");
let db = dir.path().join(db_name);
let url = format!("sqlite://{}", db.display());
let pool = connect_sqlite(&url).await.unwrap();
seed_admin_user(&pool, "correct-password").await;
seed_admin_user(&pool, password).await;
let config = test_config(url);
let stores = codex_proxy_rs::runtime::services::BackgroundTaskStores {
accounts: codex_proxy_rs::upstream::accounts::store::SqliteAccountStore::new(pool.clone()),
@@ -134,29 +220,68 @@ async fn admin_login_should_reject_client_api_key_as_password_or_authorization()
config,
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router().with_state(state);
(
codex_proxy_rs::http::router::router().with_state(state),
dir,
pool,
)
}
let response = app
.clone()
async fn post_login(app: &axum::Router, password: &str) -> axum::response::Response {
post_login_with_request_id(app, password, "req_login_throttle").await
}
async fn post_login_with_request_id(
app: &axum::Router,
password: &str,
request_id: &str,
) -> axum::response::Response {
app.clone()
.oneshot(
Request::builder()
.method("POST")
.uri("/api/admin/login")
.header("content-type", "application/json")
.header("authorization", "Bearer cpr_not_an_admin_session")
.header("x-request-id", "req_login_bad")
.body(Body::from(r#"{"password":"cpr_not_an_admin_password"}"#))
.header("x-request-id", request_id)
.body(Body::from(format!(r#"{{"password":"{password}"}}"#)))
.unwrap(),
)
.await
.unwrap();
let status = response.status();
let has_cookie = response.headers().get("set-cookie").is_some();
.unwrap()
}
assert_eq!(status, StatusCode::UNAUTHORIZED);
let body = response_json(response).await;
assert!(!has_cookie);
assert_eq!(body["code"], 40102);
#[derive(Debug)]
struct AdminLoginAuditRow {
kind: String,
level: String,
route: Option<String>,
status_code: Option<i64>,
failure_class: Option<String>,
message: String,
metadata: Value,
}
async fn admin_login_audit_row(pool: &SqlitePool, request_id: &str) -> AdminLoginAuditRow {
let row = sqlx::query(
"select kind, level, route, status_code, failure_class, message, metadata_json \
from usage_records \
where kind = 'admin_auth' and request_id = ? \
order by created_at desc, id desc \
limit 1",
)
.bind(request_id)
.fetch_one(pool)
.await
.unwrap();
AdminLoginAuditRow {
kind: row.get("kind"),
level: row.get("level"),
route: row.get("route"),
status_code: row.get("status_code"),
failure_class: row.get("failure_class"),
message: row.get("message"),
metadata: serde_json::from_str(&row.get::<String, _>("metadata_json")).unwrap(),
}
}
async fn seed_admin_user(pool: &SqlitePool, password: &str) {
backend/tests/admin/monitoring/dashboard.rs
+1 -1
View File
@@ -8,7 +8,7 @@ use codex_proxy_rs::{
auth::service::SqliteAdminSessionStore,
keys::service::SqliteClientKeyStore,
monitoring::{
usage_record::{UsageRecord, UsageRecordLevel},
usage_record_model::{UsageRecord, UsageRecordLevel},
usage_record_store::SqliteUsageRecordStore,
},
},
backend/tests/admin/monitoring/mod.rs
+1 -1
View File
@@ -1,3 +1,3 @@
pub(crate) mod dashboard;
pub(crate) mod usage_record_routes;
pub(crate) mod usage_record_store;
pub(crate) mod usage_records_routes;
backend/tests/admin/monitoring/usage_records_routes.rs → backend/tests/admin/monitoring/usage_record_routes.rs
+1 -1
View File
@@ -7,7 +7,7 @@ use codex_proxy_rs::{
admin::auth::service::SqliteAdminSessionStore,
admin::keys::service::SqliteClientKeyStore,
admin::monitoring::{
usage_record::{UsageRecord, UsageRecordLevel},
usage_record_model::{UsageRecord, UsageRecordLevel},
usage_record_store::SqliteUsageRecordStore,
},
config::types::AppConfig,
backend/tests/admin/monitoring/usage_record_store.rs
+1 -1
View File
@@ -1,6 +1,6 @@
use codex_proxy_rs::{
admin::monitoring::{
usage_record::{UsageRecord, UsageRecordLevel},
usage_record_model::{UsageRecord, UsageRecordLevel},
usage_record_store::{SqliteUsageRecordStore, UsageRecordFilter},
},
infra::database::connect_sqlite,
backend/tests/admin/system/mod.rs
+563 -12
View File
@@ -1,4 +1,9 @@
use std::{io::Write, path::Path, sync::Arc};
use std::{
io::Write,
path::Path,
sync::Arc,
time::{Duration, SystemTime},
};
use axum::{
body::Body,
@@ -21,7 +26,7 @@ use flate2::{write::GzEncoder, Compression};
use serde_json::{json, Value};
use sha2::{Digest, Sha256};
use sqlx::SqlitePool;
use tar::{Builder, Header};
use tar::{Builder, EntryType, Header};
use tokio::sync::Mutex;
use tower::util::ServiceExt;
use wiremock::{
@@ -141,6 +146,58 @@ async fn check_updates_should_report_source_build_as_unsupported() {
);
}
#[tokio::test]
async fn check_updates_should_reject_untrusted_github_api_base() {
let _guard = SYSTEM_ENV_LOCK.lock().await;
set_system_update_env(
"zyycn/codex-proxy-rs-untrusted-api-base-route",
"https://mirror.example",
);
let (app, _dir) = admin_system_test_app("system-untrusted-api-base.sqlite").await;
let info = get_check_updates(&app, true, "req_system_untrusted_api_base").await;
assert!(
info["hasUpdate"] == false
&& info["updateSupported"] == false
&& info["unsupportedReason"] == "GitHub API base must be https://api.github.com/repos"
);
}
#[tokio::test]
async fn restart_should_request_graceful_shutdown() {
let _guard = SYSTEM_ENV_LOCK.lock().await;
set_system_update_env("zyycn/codex-proxy-rs-restart-route", "http://127.0.0.1:9");
std::env::set_var("CPR_ENABLE_SELF_RESTART", "true");
let (app, _dir) = admin_system_test_app("system-restart.sqlite").await;
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/api/admin/system/restart")
.header("cookie", "cpr_admin_session=session_1")
.header("x-request-id", "req_system_restart")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
let body = response_json(response).await;
let shutdown = tokio::time::timeout(
Duration::from_secs(2),
codex_proxy_rs::runtime::shutdown::shutdown_signal(),
)
.await;
assert!(
body["data"]["operationId"]
.as_str()
.is_some_and(|id| id.starts_with("sysop-restart-"))
&& shutdown.is_ok()
);
}
#[tokio::test]
async fn update_should_replace_local_release_files_with_latest_asset() {
let _guard = SYSTEM_ENV_LOCK.lock().await;
@@ -190,6 +247,310 @@ async fn update_should_replace_local_release_files_with_latest_asset() {
);
}
#[tokio::test]
async fn update_should_reject_untrusted_github_api_base() {
let _guard = SYSTEM_ENV_LOCK.lock().await;
let deploy = tempfile::tempdir().unwrap();
let exe_path = deploy.path().join("codex-proxy-rs");
let web_dist = deploy.path().join("web/dist");
std::fs::create_dir_all(&web_dist).unwrap();
std::fs::write(&exe_path, "old-binary").unwrap();
std::fs::write(web_dist.join("index.html"), "old-web").unwrap();
set_system_update_env(
"zyycn/codex-proxy-rs-update-untrusted-api-base-route",
"https://mirror.example",
);
set_system_update_paths(deploy.path(), &exe_path, &web_dist);
let (app, _dir) = admin_system_test_app("system-update-untrusted-api-base.sqlite").await;
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/api/admin/system/update")
.header("cookie", "cpr_admin_session=session_1")
.header("x-request-id", "req_system_update_untrusted_api_base")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
let status = response.status();
let body = response_json(response).await;
assert!(
status == StatusCode::CONFLICT
&& body["message"] == "GitHub API base must be https://api.github.com/repos"
&& std::fs::read_to_string(&exe_path).unwrap() == "old-binary"
&& std::fs::read_to_string(web_dist.join("index.html")).unwrap() == "old-web"
);
}
#[tokio::test]
async fn update_should_fail_when_release_checksum_is_missing() {
let _guard = SYSTEM_ENV_LOCK.lock().await;
let repository = "zyycn/codex-proxy-rs-missing-checksum-route";
let github = MockServer::start().await;
let deploy = tempfile::tempdir().unwrap();
let exe_path = deploy.path().join("codex-proxy-rs");
let web_dist = deploy.path().join("web/dist");
std::fs::create_dir_all(&web_dist).unwrap();
std::fs::write(&exe_path, "old-binary").unwrap();
std::fs::write(web_dist.join("index.html"), "old-web").unwrap();
mount_latest_release_with_archive_without_checksum(
&github,
repository,
"0.4.0",
"new-binary",
"new-web",
)
.await;
set_system_update_env(repository, &github.uri());
set_system_update_paths(deploy.path(), &exe_path, &web_dist);
let (app, _dir) = admin_system_test_app("system-update-missing-checksum.sqlite").await;
assert_update_failure_preserves_files(
&app,
"req_system_update_missing_checksum",
StatusCode::BAD_GATEWAY,
"Release checksums.txt is required",
&exe_path,
&web_dist,
)
.await;
}
#[tokio::test]
async fn update_should_fail_when_release_checksum_mismatches() {
let _guard = SYSTEM_ENV_LOCK.lock().await;
let repository = "zyycn/codex-proxy-rs-bad-checksum-route";
let github = MockServer::start().await;
let deploy = tempfile::tempdir().unwrap();
let exe_path = deploy.path().join("codex-proxy-rs");
let web_dist = deploy.path().join("web/dist");
std::fs::create_dir_all(&web_dist).unwrap();
std::fs::write(&exe_path, "old-binary").unwrap();
std::fs::write(web_dist.join("index.html"), "old-web").unwrap();
mount_latest_release_with_archive_bad_checksum(
&github,
repository,
"0.4.0",
"new-binary",
"new-web",
)
.await;
set_system_update_env(repository, &github.uri());
set_system_update_paths(deploy.path(), &exe_path, &web_dist);
let (app, _dir) = admin_system_test_app("system-update-bad-checksum.sqlite").await;
assert_update_failure_preserves_files(
&app,
"req_system_update_bad_checksum",
StatusCode::BAD_GATEWAY,
"Checksum mismatch",
&exe_path,
&web_dist,
)
.await;
}
#[tokio::test]
async fn update_should_reject_release_archive_from_untrusted_host() {
let _guard = SYSTEM_ENV_LOCK.lock().await;
let repository = "zyycn/codex-proxy-rs-untrusted-download-route";
let github = MockServer::start().await;
let deploy = tempfile::tempdir().unwrap();
let exe_path = deploy.path().join("codex-proxy-rs");
let web_dist = deploy.path().join("web/dist");
std::fs::create_dir_all(&web_dist).unwrap();
std::fs::write(&exe_path, "old-binary").unwrap();
std::fs::write(web_dist.join("index.html"), "old-web").unwrap();
mount_latest_release_with_archive_urls(
&github,
repository,
"0.4.0",
"https://evil.example/download/codex-proxy-rs.tar.gz",
"https://github.com/zyycn/codex-proxy-rs/releases/download/v0.4.0/checksums.txt",
)
.await;
set_system_update_env(repository, &github.uri());
set_system_update_paths(deploy.path(), &exe_path, &web_dist);
let (app, _dir) = admin_system_test_app("system-update-untrusted-download.sqlite").await;
assert_update_failure_preserves_files(
&app,
"req_system_update_untrusted_download",
StatusCode::BAD_REQUEST,
"Download host is not allowed: evil.example",
&exe_path,
&web_dist,
)
.await;
}
#[tokio::test]
async fn update_should_reject_insecure_release_archive_even_when_env_flag_is_set() {
let _guard = SYSTEM_ENV_LOCK.lock().await;
let repository = "zyycn/codex-proxy-rs-insecure-download-route";
let github = MockServer::start().await;
let deploy = tempfile::tempdir().unwrap();
let exe_path = deploy.path().join("codex-proxy-rs");
let web_dist = deploy.path().join("web/dist");
std::fs::create_dir_all(&web_dist).unwrap();
std::fs::write(&exe_path, "old-binary").unwrap();
std::fs::write(web_dist.join("index.html"), "old-web").unwrap();
mount_latest_release_with_archive_urls(
&github,
repository,
"0.4.0",
"http://github.com/zyycn/codex-proxy-rs/releases/download/v0.4.0/codex-proxy-rs.tar.gz",
"https://github.com/zyycn/codex-proxy-rs/releases/download/v0.4.0/checksums.txt",
)
.await;
set_system_update_env(repository, &github.uri());
std::env::set_var("CPR_UPDATE_ALLOW_INSECURE_DOWNLOADS", "true");
set_system_update_paths(deploy.path(), &exe_path, &web_dist);
let (app, _dir) = admin_system_test_app("system-update-insecure-download.sqlite").await;
assert_update_failure_preserves_files(
&app,
"req_system_update_insecure_download",
StatusCode::BAD_REQUEST,
"Only HTTPS release downloads are allowed",
&exe_path,
&web_dist,
)
.await;
}
#[tokio::test]
async fn update_should_reject_release_archive_with_unsafe_path() {
let _guard = SYSTEM_ENV_LOCK.lock().await;
let repository = "zyycn/codex-proxy-rs-unsafe-archive-route";
let github = MockServer::start().await;
let deploy = tempfile::tempdir().unwrap();
let exe_path = deploy.path().join("codex-proxy-rs");
let web_dist = deploy.path().join("web/dist");
std::fs::create_dir_all(&web_dist).unwrap();
std::fs::write(&exe_path, "old-binary").unwrap();
std::fs::write(web_dist.join("index.html"), "old-web").unwrap();
mount_latest_release_with_unsafe_archive(&github, repository, "0.4.0", "new-binary", "new-web")
.await;
set_system_update_env(repository, &github.uri());
set_system_update_paths(deploy.path(), &exe_path, &web_dist);
let (app, _dir) = admin_system_test_app("system-update-unsafe-archive.sqlite").await;
assert_update_failure_preserves_files(
&app,
"req_system_update_unsafe_archive",
StatusCode::BAD_REQUEST,
"Unsafe archive path",
&exe_path,
&web_dist,
)
.await;
}
#[tokio::test]
async fn update_should_restore_web_assets_when_binary_backup_fails() {
let _guard = SYSTEM_ENV_LOCK.lock().await;
let repository = "zyycn/codex-proxy-rs-binary-backup-failure-route";
let github = MockServer::start().await;
let deploy = tempfile::tempdir().unwrap();
let exe_path = deploy.path().join("codex-proxy-rs");
let web_dist = deploy.path().join("web/dist");
std::fs::create_dir_all(&web_dist).unwrap();
std::fs::write(&exe_path, "old-binary").unwrap();
std::fs::write(web_dist.join("index.html"), "old-web").unwrap();
std::fs::create_dir(deploy.path().join("codex-proxy-rs.backup")).unwrap();
mount_latest_release_with_archive(&github, repository, "0.4.0", "new-binary", "new-web").await;
set_system_update_env(repository, &github.uri());
set_system_update_paths(deploy.path(), &exe_path, &web_dist);
let (app, _dir) = admin_system_test_app("system-update-binary-backup-failure.sqlite").await;
let response = app
.clone()
.oneshot(
Request::builder()
.method("POST")
.uri("/api/admin/system/update")
.header("cookie", "cpr_admin_session=session_1")
.header("x-request-id", "req_system_update_binary_backup_failure")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
let status_code = response.status();
let body = response_json(response).await;
let status = wait_for_operation_status(&app, "failed").await;
let message = body["message"].as_str().unwrap_or_default();
let state_error = status["operation"]["error"].as_str().unwrap_or_default();
assert!(
status_code == StatusCode::INTERNAL_SERVER_ERROR
&& message.starts_with("Failed to remove old binary backup:")
&& state_error.starts_with("Failed to remove old binary backup:")
&& std::fs::read_to_string(&exe_path).unwrap() == "old-binary"
&& std::fs::read_to_string(web_dist.join("index.html")).unwrap() == "old-web"
&& !deploy.path().join("web/dist.backup").exists(),
"status={status_code}, message={message}, state_error={state_error}, exe={}, web={}, web_backup_exists={}",
std::fs::read_to_string(&exe_path).unwrap_or_default(),
std::fs::read_to_string(web_dist.join("index.html")).unwrap_or_default(),
deploy.path().join("web/dist.backup").exists()
);
}
#[tokio::test]
async fn update_should_remove_stale_file_lock_and_continue() {
let _guard = SYSTEM_ENV_LOCK.lock().await;
let repository = "zyycn/codex-proxy-rs-stale-lock-route";
let github = MockServer::start().await;
let deploy = tempfile::tempdir().unwrap();
let exe_path = deploy.path().join("codex-proxy-rs");
let web_dist = deploy.path().join("web/dist");
let lock_path = deploy.path().join(".runtime/update.lock");
std::fs::create_dir_all(&web_dist).unwrap();
std::fs::create_dir_all(lock_path.parent().unwrap()).unwrap();
std::fs::write(&exe_path, "old-binary").unwrap();
std::fs::write(web_dist.join("index.html"), "old-web").unwrap();
std::fs::write(&lock_path, "pid=1\ncreated_at=stale\n").unwrap();
std::fs::File::options()
.write(true)
.open(&lock_path)
.unwrap()
.set_modified(SystemTime::now() - Duration::from_secs(31 * 60))
.unwrap();
mount_latest_release_with_archive(&github, repository, "0.4.0", "new-binary", "new-web").await;
set_system_update_env(repository, &github.uri());
set_system_update_paths(deploy.path(), &exe_path, &web_dist);
let (app, _dir) = admin_system_test_app("system-update-stale-lock.sqlite").await;
let response = app
.clone()
.oneshot(
Request::builder()
.method("POST")
.uri("/api/admin/system/update")
.header("cookie", "cpr_admin_session=session_1")
.header("x-request-id", "req_system_update_stale_lock")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
let body = response_json(response).await;
let status = wait_for_operation_status(&app, "succeeded").await;
assert!(
body["data"]["targetVersion"] == "0.4.0"
&& status["operation"]["status"] == "succeeded"
&& std::fs::read_to_string(&exe_path).unwrap() == "new-binary"
&& std::fs::read_to_string(web_dist.join("index.html")).unwrap() == "new-web"
&& !lock_path.exists()
);
}
#[tokio::test]
async fn update_should_replace_web_assets_across_filesystems() {
let _guard = SYSTEM_ENV_LOCK.lock().await;
@@ -339,6 +700,40 @@ async fn get_check_updates(app: &axum::Router, force: bool, request_id: &str) ->
response_json(response).await["data"].clone()
}
async fn assert_update_failure_preserves_files(
app: &axum::Router,
request_id: &str,
expected_status: StatusCode,
expected_message: &str,
exe_path: &Path,
web_dist: &Path,
) {
let response = app
.clone()
.oneshot(
Request::builder()
.method("POST")
.uri("/api/admin/system/update")
.header("cookie", "cpr_admin_session=session_1")
.header("x-request-id", request_id)
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
let status_code = response.status();
let body = response_json(response).await;
let status = wait_for_operation_status(app, "failed").await;
assert!(
status_code == expected_status
&& body["message"] == expected_message
&& status["operation"]["error"] == expected_message
&& std::fs::read_to_string(exe_path).unwrap() == "old-binary"
&& std::fs::read_to_string(web_dist.join("index.html")).unwrap() == "old-web"
);
}
async fn mount_latest_release(server: &MockServer, repository: &str, version: &str) {
Mock::given(method("GET"))
.and(path(format!("/repos/{repository}/releases/latest")))
@@ -371,6 +766,73 @@ async fn mount_latest_release_with_archive(
"{} {archive_name}\n",
hex::encode(Sha256::digest(&archive))
);
mount_latest_release_archive(
server,
repository,
version,
archive_name,
archive,
Some(checksum),
)
.await;
}
async fn mount_latest_release_with_archive_without_checksum(
server: &MockServer,
repository: &str,
version: &str,
binary: &str,
index: &str,
) {
let archive_name = format!(
"codex-proxy-rs_{version}_{}_{}.tar.gz",
std::env::consts::OS,
release_arch()
);
let archive = release_archive(binary, index);
mount_latest_release_archive(server, repository, version, archive_name, archive, None).await;
}
async fn mount_latest_release_with_archive_bad_checksum(
server: &MockServer,
repository: &str,
version: &str,
binary: &str,
index: &str,
) {
let archive_name = format!(
"codex-proxy-rs_{version}_{}_{}.tar.gz",
std::env::consts::OS,
release_arch()
);
let archive = release_archive(binary, index);
let checksum = format!(
"{} {archive_name}\n",
"0".repeat(Sha256::output_size() * 2)
);
mount_latest_release_archive(
server,
repository,
version,
archive_name,
archive,
Some(checksum),
)
.await;
}
async fn mount_latest_release_with_archive_urls(
server: &MockServer,
repository: &str,
version: &str,
archive_url: &str,
checksum_url: &str,
) {
let archive_name = format!(
"codex-proxy-rs_{version}_{}_{}.tar.gz",
std::env::consts::OS,
release_arch()
);
Mock::given(method("GET"))
.and(path(format!("/repos/{repository}/releases/latest")))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
@@ -382,28 +844,93 @@ async fn mount_latest_release_with_archive(
"assets": [
{
"name": archive_name,
"browser_download_url": format!("{}/download/{archive_name}", server.uri()),
"size": archive.len()
"browser_download_url": archive_url,
"size": 128
},
{
"name": "checksums.txt",
"browser_download_url": format!("{}/download/checksums.txt", server.uri()),
"size": checksum.len()
"browser_download_url": checksum_url,
"size": 128
}
]
})))
.mount(server)
.await;
}
async fn mount_latest_release_with_unsafe_archive(
server: &MockServer,
repository: &str,
version: &str,
binary: &str,
index: &str,
) {
let archive_name = format!(
"codex-proxy-rs_{version}_{}_{}.tar.gz",
std::env::consts::OS,
release_arch()
);
let archive = release_archive_with_extra_file(binary, index, "../escape", b"unsafe");
let checksum = format!(
"{} {archive_name}\n",
hex::encode(Sha256::digest(&archive))
);
mount_latest_release_archive(
server,
repository,
version,
archive_name,
archive,
Some(checksum),
)
.await;
}
async fn mount_latest_release_archive(
server: &MockServer,
repository: &str,
version: &str,
archive_name: String,
archive: Vec<u8>,
checksum: Option<String>,
) {
let mut assets = vec![json!({
"name": archive_name,
"browser_download_url": format!("{}/download/{archive_name}", server.uri()),
"size": archive.len()
})];
if let Some(checksum) = checksum.as_ref() {
assets.push(json!({
"name": "checksums.txt",
"browser_download_url": format!("{}/download/checksums.txt", server.uri()),
"size": checksum.len()
}));
}
Mock::given(method("GET"))
.and(path(format!("/repos/{repository}/releases/latest")))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"tag_name": format!("v{version}"),
"name": format!("v{version}"),
"body": "## Changes\n- Update package",
"html_url": format!("https://github.com/{repository}/releases/tag/v{version}"),
"prerelease": false,
"assets": assets
})))
.mount(server)
.await;
Mock::given(method("GET"))
.and(path(format!("/download/{archive_name}")))
.respond_with(ResponseTemplate::new(200).set_body_bytes(archive))
.mount(server)
.await;
Mock::given(method("GET"))
.and(path("/download/checksums.txt"))
.respond_with(ResponseTemplate::new(200).set_body_string(checksum))
.mount(server)
.await;
if let Some(checksum) = checksum {
Mock::given(method("GET"))
.and(path("/download/checksums.txt"))
.respond_with(ResponseTemplate::new(200).set_body_string(checksum))
.mount(server)
.await;
}
}
fn release_arch() -> &'static str {
@@ -415,12 +942,24 @@ fn release_arch() -> &'static str {
}
fn release_archive(binary: &str, index: &str) -> Vec<u8> {
release_archive_with_extra_file(binary, index, "", &[])
}
fn release_archive_with_extra_file(
binary: &str,
index: &str,
extra_path: &str,
extra_bytes: &[u8],
) -> Vec<u8> {
let mut data = Vec::new();
{
let encoder = GzEncoder::new(&mut data, Compression::default());
let mut builder = Builder::new(encoder);
append_tar_file(&mut builder, "codex-proxy-rs", binary.as_bytes());
append_tar_file(&mut builder, "web/dist/index.html", index.as_bytes());
if !extra_path.is_empty() {
append_raw_tar_file(&mut builder, extra_path, extra_bytes);
}
builder.finish().unwrap();
}
data
@@ -433,6 +972,17 @@ fn append_tar_file<W: Write>(builder: &mut Builder<W>, path: &str, bytes: &[u8])
builder.append_data(&mut header, path, bytes).unwrap();
}
fn append_raw_tar_file<W: Write>(builder: &mut Builder<W>, path: &str, bytes: &[u8]) {
let mut header = Header::new_old();
header.set_size(bytes.len() as u64);
header.set_mode(0o644);
header.set_entry_type(EntryType::Regular);
let name = path.as_bytes();
header.as_old_mut().name[..name.len()].copy_from_slice(name);
header.set_cksum();
builder.append(&header, bytes).unwrap();
}
fn set_system_update_env(repository: &str, github_api_base: &str) {
std::env::set_var("CPR_VERSION", "0.1.0");
std::env::set_var("CPR_DEPLOYMENT_MODE", "docker");
@@ -440,11 +990,12 @@ fn set_system_update_env(repository: &str, github_api_base: &str) {
std::env::set_var("CPR_UPDATE_CHANNEL", "stable");
std::env::set_var("CPR_UPDATE_REPOSITORY", repository);
std::env::set_var("CPR_GITHUB_API_BASE", format!("{github_api_base}/repos"));
std::env::set_var("CPR_UPDATE_ALLOW_INSECURE_DOWNLOADS", "true");
std::env::remove_var("CPR_UPDATE_ALLOW_INSECURE_DOWNLOADS");
std::env::remove_var("CPR_UPDATE_EXE_PATH");
std::env::remove_var("CPR_WEB_DIST_DIR");
std::env::remove_var("CPR_UPDATE_STATE_FILE");
std::env::remove_var("CPR_UPDATE_LOCK_FILE");
std::env::remove_var("CPR_ENABLE_SELF_RESTART");
}
fn set_system_update_paths(root: &Path, exe_path: &Path, web_dist: &Path) {
backend/tests/config/mod.rs
+89 -7
View File
@@ -1,12 +1,14 @@
use codex_proxy_rs::config::settings::account_pool_options_from_config;
use codex_proxy_rs::config::types::AppConfig;
use codex_proxy_rs::config::types::LoggingConfig;
use codex_proxy_rs::config::types::{AdminConfig, AdminConfigError, AppConfig};
use serde::de::DeserializeOwned;
use std::fs;
const DEFAULT_CONFIG_YAML: &str = r#"
server:
host: 0.0.0.0
port: 8080
trusted_proxies: []
api:
base_url: https://chatgpt.com/backend-api
database:
@@ -67,7 +69,7 @@ fingerprint:
admin:
session_ttl_minutes: 1440
default_username: admin
default_password: admin
default_password: test-admin-password
logging:
directory: .runtime/logs
retention_days: 14
@@ -76,8 +78,9 @@ logging:
#[test]
fn default_config_keeps_only_codex_backend() {
let cfg: AppConfig = serde_yml::from_str(DEFAULT_CONFIG_YAML).unwrap();
let cfg: AppConfig = parse_yaml_config(DEFAULT_CONFIG_YAML).unwrap();
assert_eq!(cfg.server.host, "0.0.0.0");
assert!(cfg.server.trusted_proxies.is_empty());
assert_eq!(cfg.api.base_url, "https://chatgpt.com/backend-api");
assert!(cfg.model_aliases.is_empty());
assert_eq!(cfg.auth.refresh_margin_seconds, 300);
@@ -99,7 +102,7 @@ fn default_config_keeps_only_codex_backend() {
#[test]
fn default_config_keeps_runtime_artifacts_under_runtime_directory() {
let cfg: AppConfig = serde_yml::from_str(DEFAULT_CONFIG_YAML).unwrap();
let cfg: AppConfig = parse_yaml_config(DEFAULT_CONFIG_YAML).unwrap();
assert_eq!(
cfg.database.url,
@@ -110,7 +113,7 @@ fn default_config_keeps_runtime_artifacts_under_runtime_directory() {
#[test]
fn account_pool_options_should_use_quota_skip_exhausted() {
let mut cfg: AppConfig = serde_yml::from_str(DEFAULT_CONFIG_YAML).unwrap();
let mut cfg: AppConfig = parse_yaml_config(DEFAULT_CONFIG_YAML).unwrap();
cfg.quota.skip_exhausted = false;
let options = account_pool_options_from_config(&cfg);
@@ -118,9 +121,32 @@ fn account_pool_options_should_use_quota_skip_exhausted() {
assert!(!options.skip_quota_limited);
}
#[test]
fn admin_config_should_reject_weak_default_password() {
let config = AdminConfig {
default_password: "123456".to_string(),
..AdminConfig::default()
};
assert_eq!(
config.validate_default_password(),
Err(AdminConfigError::WeakDefaultPassword)
);
}
#[test]
fn admin_config_should_accept_explicit_strong_default_password() {
let config = AdminConfig {
default_password: "correct-horse-battery-staple".to_string(),
..AdminConfig::default()
};
assert_eq!(config.validate_default_password(), Ok(()));
}
#[test]
fn config_should_reject_unknown_top_level_sections() {
let err = serde_yml::from_str::<AppConfig>(
let err = parse_yaml_config::<AppConfig>(
r"
server:
host: 127.0.0.1
@@ -233,6 +259,7 @@ ws_pool:
assert_eq!(cfg.server.host, "127.0.0.1");
assert_eq!(cfg.server.port, 8080);
assert!(cfg.server.trusted_proxies.is_empty());
assert_eq!(
cfg.database.url,
"sqlite://.runtime/data/codex-proxy-rs.sqlite"
@@ -253,9 +280,57 @@ ws_pool:
assert_eq!(cfg.logging.retention_days, 14);
}
#[test]
fn config_should_parse_trusted_proxy_entries() {
let cfg: AppConfig = parse_yaml_config(&DEFAULT_CONFIG_YAML.replace(
"trusted_proxies: []",
"trusted_proxies:\n - 127.0.0.1\n - 10.0.0.0/8",
))
.unwrap();
assert_eq!(
cfg.server.trusted_proxies,
vec!["127.0.0.1".to_string(), "10.0.0.0/8".to_string()]
);
}
#[test]
fn server_config_should_reject_unknown_fields() {
let err = parse_yaml_config::<AppConfig>(&DEFAULT_CONFIG_YAML.replace(
"trusted_proxies: []",
"trusted_proxies: []\n trusted_proxy: 127.0.0.1",
))
.unwrap_err();
assert!(
err.to_string().contains("trusted_proxy"),
"expected unknown server field to be rejected, got {err}"
);
}
#[test]
fn config_loader_should_read_explicit_config_file_from_env() {
let dir = tempfile::tempdir().unwrap();
let config_file = dir.path().join("custom-config.yaml");
fs::write(
&config_file,
DEFAULT_CONFIG_YAML
.replace("host: 0.0.0.0", "host: 127.0.0.2")
.replace("directory: .runtime/logs", "directory: .runtime/env-logs"),
)
.unwrap();
std::env::set_var("CPR_CONFIG_FILE", &config_file);
let cfg = AppConfig::load().unwrap();
std::env::remove_var("CPR_CONFIG_FILE");
assert_eq!(cfg.server.host, "127.0.0.2");
assert_eq!(cfg.logging.directory, ".runtime/env-logs");
}
#[test]
fn logging_config_should_reject_unknown_fields() {
let err = serde_yml::from_str::<LoggingConfig>(
let err = parse_yaml_config::<LoggingConfig>(
r"
directory: .runtime/logs
unexpected: true
@@ -270,3 +345,10 @@ enabled: false
"expected unknown logging field to be rejected, got {err}"
);
}
fn parse_yaml_config<T: DeserializeOwned>(yaml: &str) -> Result<T, config::ConfigError> {
config::Config::builder()
.add_source(config::File::from_str(yaml, config::FileFormat::Yaml))
.build()?
.try_deserialize()
}
backend/tests/http/trace_middleware/mod.rs
+199 -4
View File
@@ -9,11 +9,14 @@ use axum::{
extract::connect_info::ConnectInfo,
http::HeaderMap,
http::{Request, StatusCode},
middleware::from_fn,
middleware::{from_fn, from_fn_with_state},
routing::get,
Router,
};
use codex_proxy_rs::http::middleware::{request_id::attach_request_id, trace::http_trace_layer};
use codex_proxy_rs::http::middleware::{
request_id::{attach_request_id, attach_request_id_with_proxy_config, TrustedProxyConfig},
trace::http_trace_layer,
};
use tokio::time::{sleep, Duration};
use tower::ServiceExt;
use tracing_subscriber::fmt::MakeWriter;
@@ -108,6 +111,113 @@ async fn http_trace_should_include_request_id_and_completion_fields() {
);
}
#[tokio::test(flavor = "current_thread")]
async fn http_trace_should_redact_request_uri_query_values() {
let logs = SharedLogBuffer::default();
let subscriber = tracing_subscriber::fmt()
.json()
.with_writer(logs.clone())
.with_current_span(true)
.with_span_list(true)
.with_ansi(false)
.finish();
let app = Router::new()
.route("/trace-test", get(|| async { StatusCode::NO_CONTENT }))
.layer(http_trace_layer())
.layer(from_fn(attach_request_id));
let _guard = tracing::subscriber::set_default(subscriber);
let response = app
.oneshot(
Request::builder()
.uri("/trace-test?token=secret-token&code=secret-code")
.header("x-request-id", "req_trace_query")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
let _body = to_bytes(response.into_body(), 1024).await.unwrap();
let output = wait_for_trace_output(&logs).await;
assert!(
output.contains("/trace-test?token=<redacted>&code=<redacted>")
&& !output.contains("secret-token")
&& !output.contains("secret-code")
&& output.contains("token=<redacted>")
&& output.contains("code=<redacted>"),
"unexpected trace output: {output}"
);
}
#[tokio::test(flavor = "current_thread")]
async fn http_trace_should_redact_bare_query_segments() {
let logs = SharedLogBuffer::default();
let subscriber = tracing_subscriber::fmt()
.json()
.with_writer(logs.clone())
.with_current_span(true)
.with_span_list(true)
.with_ansi(false)
.finish();
let app = Router::new()
.route("/trace-test", get(|| async { StatusCode::NO_CONTENT }))
.layer(http_trace_layer())
.layer(from_fn(attach_request_id));
let _guard = tracing::subscriber::set_default(subscriber);
let response = app
.oneshot(
Request::builder()
.uri("/trace-test?bare-secret&empty=&page=2&page=3")
.header("x-request-id", "req_trace_query_bare")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
let _body = to_bytes(response.into_body(), 1024).await.unwrap();
let output = wait_for_trace_output(&logs).await;
assert!(
output.contains("/trace-test?<redacted>&empty=<redacted>&page=<redacted>&page=<redacted>")
&& !output.contains("bare-secret")
&& !output.contains("page=2")
&& !output.contains("page=3"),
"unexpected trace output: {output}"
);
}
#[tokio::test]
async fn request_context_should_generate_new_request_id_for_invalid_header() {
let app = Router::new()
.route("/request-id", get(|| async { StatusCode::NO_CONTENT }))
.layer(from_fn(attach_request_id));
let response = app
.oneshot(
Request::builder()
.uri("/request-id")
.header("x-request-id", "req invalid")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
let request_id = response
.headers()
.get("x-request-id")
.and_then(|value| value.to_str().ok())
.unwrap_or_default();
assert!(
request_id.starts_with("req_") && request_id != "req invalid",
"unexpected request id: {request_id}"
);
}
#[tokio::test]
async fn request_context_should_attach_real_ip_from_connection() {
let app = Router::new()
@@ -137,7 +247,7 @@ async fn request_context_should_attach_real_ip_from_connection() {
}
#[tokio::test]
async fn request_context_should_not_override_forwarded_ip_headers() {
async fn request_context_should_replace_untrusted_forwarded_ip_headers() {
let app = Router::new()
.route(
"/ip",
@@ -153,6 +263,8 @@ async fn request_context_should_not_override_forwarded_ip_headers() {
let mut request = Request::builder()
.uri("/ip")
.header("x-forwarded-for", "203.0.113.42, 10.0.0.2")
.header("cf-connecting-ip", "203.0.113.43")
.header("x-real-ip", "203.0.113.44")
.body(Body::empty())
.expect("request should build");
request
@@ -162,5 +274,88 @@ async fn request_context_should_not_override_forwarded_ip_headers() {
let response = app.oneshot(request).await.expect("response");
let bytes = to_bytes(response.into_body(), 1024).await.unwrap();
assert_eq!(std::str::from_utf8(&bytes).unwrap(), "missing");
assert_eq!(std::str::from_utf8(&bytes).unwrap(), "127.0.0.1");
}
#[tokio::test]
async fn request_context_should_accept_forwarded_ip_from_trusted_proxy() {
let trusted_proxies =
TrustedProxyConfig::from_entries(&["127.0.0.0/8".to_string()]).expect("trusted proxy");
let app = Router::new()
.route(
"/ip",
get(|headers: HeaderMap| async move {
let real_ip = headers
.get("x-real-ip")
.and_then(|value| value.to_str().ok())
.unwrap_or("missing");
let forwarded_present = headers.contains_key("x-forwarded-for");
format!("{real_ip}|forwarded={forwarded_present}")
}),
)
.layer(from_fn_with_state(
trusted_proxies,
attach_request_id_with_proxy_config,
));
let mut request = Request::builder()
.uri("/ip")
.header("x-forwarded-for", "203.0.113.42, 10.0.0.2")
.body(Body::empty())
.expect("request should build");
request
.extensions_mut()
.insert(ConnectInfo(SocketAddr::from(([127, 0, 0, 1], 61234))));
let response = app.oneshot(request).await.expect("response");
let bytes = to_bytes(response.into_body(), 1024).await.unwrap();
assert_eq!(
std::str::from_utf8(&bytes).unwrap(),
"203.0.113.42|forwarded=false"
);
}
#[tokio::test]
async fn request_context_should_prefer_cloudflare_ip_from_trusted_proxy() {
let trusted_proxies =
TrustedProxyConfig::from_entries(&["127.0.0.1".to_string()]).expect("trusted proxy");
let app = Router::new()
.route(
"/ip",
get(|headers: HeaderMap| async move {
headers
.get("x-real-ip")
.and_then(|value| value.to_str().ok())
.unwrap_or("missing")
.to_string()
}),
)
.layer(from_fn_with_state(
trusted_proxies,
attach_request_id_with_proxy_config,
));
let mut request = Request::builder()
.uri("/ip")
.header("cf-connecting-ip", "198.51.100.23")
.header("x-forwarded-for", "203.0.113.42, 10.0.0.2")
.body(Body::empty())
.expect("request should build");
request
.extensions_mut()
.insert(ConnectInfo(SocketAddr::from(([127, 0, 0, 1], 61234))));
let response = app.oneshot(request).await.expect("response");
let bytes = to_bytes(response.into_body(), 1024).await.unwrap();
assert_eq!(std::str::from_utf8(&bytes).unwrap(), "198.51.100.23");
}
#[test]
fn trusted_proxy_config_should_reject_invalid_entries() {
let error = TrustedProxyConfig::from_entries(&["not-an-ip".to_string()]).unwrap_err();
assert!(
error.to_string().contains("not-an-ip"),
"unexpected error: {error}"
);
}
backend/tests/http/web_assets/mod.rs
+126 -46
View File
@@ -1,9 +1,13 @@
use std::path::Path;
use axum::{
body::{to_bytes, Body},
http::{header, Request, StatusCode},
Router,
};
use codex_proxy_rs::{config::types::AppConfig, infra::database::connect_sqlite};
use serde_json::Value;
use sqlx::SqlitePool;
use tower::ServiceExt;
use crate::support::{config::test_config as base_test_config, http::response_json};
@@ -34,8 +38,124 @@ async fn server_router_should_serve_frontend_assets_without_shadowing_api_routes
.expect("asset should be written");
let db = dir.path().join("assets-routes.sqlite");
let database_url = format!("sqlite://{}", db.display());
let pool = connect_sqlite(&database_url).await.expect("sqlite pool");
let config = test_config(database_url);
let (app, _pool) = router_with_dist_and_database(&dist, &database_url).await;
let index = app.clone().oneshot(request("/")).await.expect("index");
let asset = app
.clone()
.oneshot(request("/assets/app.js"))
.await
.expect("asset");
let api = app
.clone()
.oneshot(request("/api/admin/settings"))
.await
.expect("api response");
let health = app
.oneshot(request("/healthz"))
.await
.expect("health response");
assert_eq!(index.status(), StatusCode::OK);
assert_eq!(asset.status(), StatusCode::OK);
assert_eq!(api.status(), StatusCode::UNAUTHORIZED);
assert_eq!(health.status(), StatusCode::NO_CONTENT);
let csp = index
.headers()
.get(header::CONTENT_SECURITY_POLICY)
.and_then(|value| value.to_str().ok())
.unwrap_or_default();
assert!(
csp.contains("style-src 'self';") && !csp.contains("unsafe-inline"),
"unexpected content security policy: {csp}"
);
assert_eq!(
asset.headers().get(header::CACHE_CONTROL),
Some(&header::HeaderValue::from_static(
"public, max-age=31536000, immutable"
))
);
assert_body_contains(index, "Codex Proxy Admin").await;
assert_body_contains(asset, "__asset").await;
assert_eq!(response_json(api).await["code"], Value::from(40101));
}
#[tokio::test]
async fn server_router_should_return_json_404_for_unknown_api_routes() {
let dir = tempfile::tempdir().expect("temp dir");
let dist = dir.path().join("dist");
std::fs::create_dir_all(&dist).expect("dist dir should be created");
std::fs::write(
dist.join("index.html"),
"<!doctype html><main id=\"app\">Codex Proxy Admin</main>",
)
.expect("index should be written");
let db = dir.path().join("unknown-api-routes.sqlite");
let database_url = format!("sqlite://{}", db.display());
let (app, _pool) = router_with_dist_and_database(&dist, &database_url).await;
let admin = app
.clone()
.oneshot(request("/api/admin/does-not-exist"))
.await
.expect("admin unknown response");
let openai = app
.oneshot(request("/v1/does-not-exist"))
.await
.expect("openai unknown response");
assert_unknown_api_route(admin).await;
assert_unknown_api_route(openai).await;
}
#[tokio::test]
async fn healthz_should_report_unavailable_when_sqlite_is_closed() {
let dir = tempfile::tempdir().expect("temp dir");
let dist = dir.path().join("dist");
std::fs::create_dir_all(&dist).expect("dist dir should be created");
std::fs::write(dist.join("index.html"), "<!doctype html>").expect("index should be written");
let db = dir.path().join("healthz.sqlite");
let database_url = format!("sqlite://{}", db.display());
let (app, pool) = router_with_dist_and_database(&dist, &database_url).await;
pool.close().await;
let response = app
.oneshot(request("/healthz"))
.await
.expect("health response");
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
}
fn request(uri: &str) -> Request<Body> {
Request::builder()
.uri(uri)
.body(Body::empty())
.expect("request should build")
}
async fn assert_body_contains(response: axum::response::Response, expected: &str) {
let bytes = to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should collect");
let text = std::str::from_utf8(&bytes).expect("body should be utf8");
assert!(
text.contains(expected),
"body should contain `{expected}`, got `{text}`"
);
}
async fn assert_unknown_api_route(response: axum::response::Response) {
assert_eq!(response.status(), StatusCode::NOT_FOUND);
let body = response_json(response).await;
assert_eq!(body["code"], Value::from(40401));
assert_eq!(body["message"], Value::from("API route not found"));
}
async fn router_with_dist_and_database(dist: &Path, database_url: &str) -> (Router, SqlitePool) {
let pool = connect_sqlite(database_url).await.expect("sqlite pool");
let config = test_config(database_url.to_string());
let stores = codex_proxy_rs::runtime::services::BackgroundTaskStores {
accounts: codex_proxy_rs::upstream::accounts::store::SqliteAccountStore::new(pool.clone()),
admin_sessions: codex_proxy_rs::admin::auth::service::SqliteAdminSessionStore::new(
@@ -68,50 +188,10 @@ async fn server_router_should_serve_frontend_assets_without_shadowing_api_routes
config,
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router_with_assets(&dist).with_state(state);
let index = app.clone().oneshot(request("/")).await.expect("index");
let asset = app
.clone()
.oneshot(request("/assets/app.js"))
.await
.expect("asset");
let api = app
.oneshot(request("/api/admin/settings"))
.await
.expect("api response");
assert_eq!(index.status(), StatusCode::OK);
assert_eq!(asset.status(), StatusCode::OK);
assert_eq!(api.status(), StatusCode::UNAUTHORIZED);
assert_eq!(
asset.headers().get(header::CACHE_CONTROL),
Some(&header::HeaderValue::from_static(
"public, max-age=31536000, immutable"
))
);
assert_body_contains(index, "Codex Proxy Admin").await;
assert_body_contains(asset, "__asset").await;
assert_eq!(response_json(api).await["code"], Value::from(40101));
}
fn request(uri: &str) -> Request<Body> {
Request::builder()
.uri(uri)
.body(Body::empty())
.expect("request should build")
}
async fn assert_body_contains(response: axum::response::Response, expected: &str) {
let bytes = to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should collect");
let text = std::str::from_utf8(&bytes).expect("body should be utf8");
assert!(
text.contains(expected),
"body should contain `{expected}`, got `{text}`"
);
(
codex_proxy_rs::http::router::router_with_assets(dist).with_state(state),
pool,
)
}
fn test_config(database_url: String) -> AppConfig {
backend/tests/infra/storage_schema/mod.rs
+57 -1
View File
@@ -1,3 +1,5 @@
use std::path::Path;
use codex_proxy_rs::infra::database::connect_sqlite;
#[tokio::test]
@@ -8,7 +10,7 @@ async fn sqlite_schema_should_create_current_tables() {
let pool = connect_sqlite(&url).await.unwrap();
let rows: Vec<(String,)> = sqlx::query_as(
"select name from sqlite_master where type = 'table' and name in ('admin_users', 'admin_sessions', 'client_api_keys', 'runtime_settings', 'accounts', 'account_refresh_leases', 'account_usage', 'account_model_usage', 'usage_time_buckets', 'model_account_routes', 'account_cookies', 'fingerprints', 'fingerprint_update_history', 'usage_records', 'model_plan_snapshots', 'session_affinities') order by name",
"select name from sqlite_master where type = 'table' and name in ('schema_migrations', 'admin_users', 'admin_sessions', 'client_api_keys', 'runtime_settings', 'accounts', 'account_refresh_leases', 'account_usage', 'account_model_usage', 'usage_time_buckets', 'model_account_routes', 'account_cookies', 'fingerprints', 'fingerprint_update_history', 'usage_records', 'model_plan_snapshots', 'session_affinities') order by name",
)
.fetch_all(&pool)
.await
@@ -29,6 +31,7 @@ async fn sqlite_schema_should_create_current_tables() {
"model_account_routes",
"model_plan_snapshots",
"runtime_settings",
"schema_migrations",
"session_affinities",
"usage_records",
"usage_time_buckets",
@@ -36,6 +39,59 @@ async fn sqlite_schema_should_create_current_tables() {
);
}
#[tokio::test]
async fn sqlite_schema_should_record_applied_migrations() {
let dir = tempfile::tempdir().unwrap();
let db = dir.path().join("schema-migrations.sqlite");
let url = format!("sqlite://{}", db.display());
let pool = connect_sqlite(&url).await.unwrap();
let rows: Vec<(i64, String)> =
sqlx::query_as("select version, name from schema_migrations order by version")
.fetch_all(&pool)
.await
.unwrap();
assert_eq!(rows, [(1, "initial".to_string())]);
}
#[test]
fn initial_migration_should_not_use_idempotent_business_schema_ddl() {
let migration_path = Path::new(env!("CARGO_MANIFEST_DIR"))
.join("src")
.join("infra")
.join("migrations")
.join("0001_initial.sql");
let migration = std::fs::read_to_string(migration_path).unwrap();
assert!(!migration.to_lowercase().contains("if not exists"));
}
#[tokio::test]
async fn connect_sqlite_should_reject_unversioned_existing_schema() {
let dir = tempfile::tempdir().unwrap();
let db = dir.path().join("unversioned.sqlite");
let url = format!("sqlite://{}", db.display());
std::fs::File::create(&db).unwrap();
let pool = sqlx::SqlitePool::connect(&url).await.unwrap();
sqlx::query("create table admin_users (id text primary key)")
.execute(&pool)
.await
.unwrap();
pool.close().await;
let err = match connect_sqlite(&url).await {
Ok(pool) => {
pool.close().await;
panic!("unversioned sqlite schema should be rejected");
}
Err(err) => err,
};
assert!(err.to_string().contains("unversioned sqlite schema"));
}
#[tokio::test]
async fn sqlite_schema_should_persist_runtime_settings_columns() {
let dir = tempfile::tempdir().unwrap();
backend/tests/proxy/auth.rs
+54
View File
@@ -0,0 +1,54 @@
use axum::http::{header::AUTHORIZATION, HeaderMap, HeaderValue};
use codex_proxy_rs::proxy::auth::{bearer_client_api_key, ClientApiKeyAuthError};
#[test]
fn bearer_client_api_key_should_report_missing_authorization() {
let headers = HeaderMap::new();
let result = bearer_client_api_key(&headers);
assert_eq!(result, Err(ClientApiKeyAuthError::MissingAuthorization));
}
#[test]
fn bearer_client_api_key_should_report_malformed_authorization() {
let mut headers = HeaderMap::new();
headers.insert(AUTHORIZATION, HeaderValue::from_static("Token sk_test"));
let result = bearer_client_api_key(&headers);
assert_eq!(result, Err(ClientApiKeyAuthError::MalformedAuthorization));
}
#[test]
fn bearer_client_api_key_should_report_empty_bearer_token() {
let mut headers = HeaderMap::new();
headers.insert(AUTHORIZATION, HeaderValue::from_static("Bearer "));
let result = bearer_client_api_key(&headers);
assert_eq!(result, Err(ClientApiKeyAuthError::MalformedAuthorization));
}
#[test]
fn bearer_client_api_key_should_report_invalid_key_format() {
let mut headers = HeaderMap::new();
headers.insert(
AUTHORIZATION,
HeaderValue::from_static("Bearer admin-secret"),
);
let result = bearer_client_api_key(&headers);
assert_eq!(result, Err(ClientApiKeyAuthError::InvalidKeyFormat));
}
#[test]
fn bearer_client_api_key_should_extract_trimmed_client_key() {
let mut headers = HeaderMap::new();
headers.insert(AUTHORIZATION, HeaderValue::from_static("Bearer sk_test "));
let result = bearer_client_api_key(&headers);
assert_eq!(result, Ok("sk_test"));
}
backend/tests/proxy/dispatch/chat_upstream/mod.rs
+1 -1
View File
@@ -15,7 +15,7 @@ use codex_proxy_rs::{
admin::auth::service::SqliteAdminSessionStore,
admin::keys::service::SqliteClientKeyStore,
admin::monitoring::{
usage_record::UsageRecordLevel,
usage_record_model::UsageRecordLevel,
usage_record_store::{
SqliteUsageRecordStore, UsageRecordFilter, DEFAULT_USAGE_RECORD_CAPACITY,
},
backend/tests/proxy/mod.rs
+1
View File
@@ -1,3 +1,4 @@
pub(crate) mod auth;
pub(crate) mod dispatch;
pub(crate) mod keys;
pub(crate) mod openai;
backend/tests/proxy/openai/models_routes.rs
+56 -41
View File
@@ -1,6 +1,7 @@
use axum::{
body::Body,
http::{Request, StatusCode},
Router,
};
use codex_proxy_rs::{
admin::auth::service::SqliteAdminSessionStore,
@@ -16,36 +17,13 @@ use codex_proxy_rs::{
};
use tower::util::ServiceExt;
use crate::support::{client_keys::insert_client_api_key, config::test_config};
use crate::support::{
client_keys::insert_client_api_key, config::test_config, http::response_json,
};
#[tokio::test]
async fn models_route_should_reject_unknown_client_api_key() {
let dir = tempfile::tempdir().unwrap();
let db = dir.path().join("openai-models-auth.sqlite");
let url = format!("sqlite://{}", db.display());
let pool = connect_sqlite(&url).await.unwrap();
let config = test_config(url);
let stores = BackgroundTaskStores {
accounts: SqliteAccountStore::new(pool.clone()),
admin_sessions: SqliteAdminSessionStore::new(pool.clone()),
cookies: SqliteCookieStore::new(pool.clone()),
fingerprints: FingerprintRepository::new(pool.clone()),
session_affinity: SqliteSessionAffinityStore::new(pool.clone()),
refresh_leases: RefreshLeaseStore::new(pool.clone()),
client_keys: SqliteClientKeyStore::new(pool.clone()),
usage_records: SqliteUsageRecordStore::new(pool.clone()),
};
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
services
.initialize_hot_path_state()
.await
.expect("hot path state should initialize");
let state = AppState {
config,
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router().with_state(state);
let (app, _key, _dir) = test_app("openai-models-auth.sqlite", false).await;
let response = app
.oneshot(
@@ -61,13 +39,61 @@ async fn models_route_should_reject_unknown_client_api_key() {
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
}
#[tokio::test]
async fn models_route_should_return_uniform_error_for_malformed_authorization() {
let (app, _key, _dir) = test_app("openai-models-auth-malformed.sqlite", false).await;
let response = app
.oneshot(
Request::builder()
.method("GET")
.uri("/v1/models")
.header("authorization", "Token sk_not_bearer")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
assert_eq!(
response_json(response).await["error"]["code"],
"invalid_api_key"
);
}
#[tokio::test]
async fn models_route_should_accept_stored_client_api_key() {
let (app, plaintext, _dir) = test_app("openai-models-auth-valid.sqlite", true).await;
let plaintext = plaintext.expect("client api key should be seeded");
let response = app
.oneshot(
Request::builder()
.method("GET")
.uri("/v1/models")
.header("authorization", format!("Bearer {plaintext}"))
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
}
async fn test_app(
db_name: &str,
seed_client_key: bool,
) -> (Router, Option<String>, tempfile::TempDir) {
let dir = tempfile::tempdir().unwrap();
let db = dir.path().join("openai-models-auth-valid.sqlite");
let db = dir.path().join(db_name);
let url = format!("sqlite://{}", db.display());
let pool = connect_sqlite(&url).await.unwrap();
let plaintext = insert_client_api_key(&pool).await;
let plaintext = if seed_client_key {
Some(insert_client_api_key(&pool).await)
} else {
None
};
let config = test_config(url);
let stores = BackgroundTaskStores {
accounts: SqliteAccountStore::new(pool.clone()),
@@ -91,16 +117,5 @@ async fn models_route_should_accept_stored_client_api_key() {
};
let app = codex_proxy_rs::http::router::router().with_state(state);
let response = app
.oneshot(
Request::builder()
.method("GET")
.uri("/v1/models")
.header("authorization", format!("Bearer {plaintext}"))
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
(app, plaintext, dir)
}
backend/tests/runtime/services.rs
+59 -13
View File
@@ -14,9 +14,26 @@ use codex_proxy_rs::{
transport::tls::CODEX_CA_CERT_ENV,
},
};
use sqlx::SqlitePool;
use crate::support::{config::test_config, sqlite::init_test_db};
fn stores(pool: SqlitePool) -> BackgroundTaskStores {
BackgroundTaskStores {
accounts: SqliteAccountStore::new(pool.clone()),
admin_sessions: SqliteAdminSessionStore::new(pool.clone()),
cookies: SqliteCookieStore::new(pool.clone()),
fingerprints: FingerprintRepository::new(pool.clone()),
session_affinity:
codex_proxy_rs::proxy::dispatch::session_affinity::SqliteSessionAffinityStore::new(
pool.clone(),
),
refresh_leases: RefreshLeaseStore::new(pool.clone()),
client_keys: SqliteClientKeyStore::new(pool.clone()),
usage_records: SqliteUsageRecordStore::new(pool),
}
}
#[tokio::test]
async fn services_try_new_should_use_configured_tls_transport_builder() {
const CASE_ENV: &str = "CODEX_PROXY_TEST_SERVICES_TLS_CASE";
@@ -27,19 +44,7 @@ async fn services_try_new_should_use_configured_tls_transport_builder() {
"sqlite://{}",
temp_dir.path().join("services-tls.sqlite").display()
));
let stores = BackgroundTaskStores {
accounts: SqliteAccountStore::new(pool.clone()),
admin_sessions: SqliteAdminSessionStore::new(pool.clone()),
cookies: SqliteCookieStore::new(pool.clone()),
fingerprints: FingerprintRepository::new(pool.clone()),
session_affinity:
codex_proxy_rs::proxy::dispatch::session_affinity::SqliteSessionAffinityStore::new(
pool.clone(),
),
refresh_leases: RefreshLeaseStore::new(pool.clone()),
client_keys: SqliteClientKeyStore::new(pool.clone()),
usage_records: SqliteUsageRecordStore::new(pool),
};
let stores = stores(pool);
let Err(error) = Services::try_new(
&config,
@@ -71,3 +76,44 @@ async fn services_try_new_should_use_configured_tls_transport_builder() {
String::from_utf8_lossy(&output.stderr)
);
}
#[tokio::test]
async fn services_try_new_should_expose_websocket_pool_when_enabled() {
let (pool, temp_dir) = init_test_db("services-ws-pool-enabled.sqlite").await;
let config = test_config(format!(
"sqlite://{}",
temp_dir
.path()
.join("services-ws-pool-enabled.sqlite")
.display()
));
let services = Services::try_new(
&config,
stores(pool),
crate::support::fingerprint::test_fingerprint(),
)
.expect("services should build");
assert!(services.websocket_pool.is_some());
}
#[tokio::test]
async fn services_try_new_should_not_expose_websocket_pool_when_disabled() {
let (pool, temp_dir) = init_test_db("services-ws-pool-disabled.sqlite").await;
let mut config = test_config(format!(
"sqlite://{}",
temp_dir
.path()
.join("services-ws-pool-disabled.sqlite")
.display()
));
config.ws_pool.enabled = false;
let services = Services::try_new(
&config,
stores(pool),
crate::support::fingerprint::test_fingerprint(),
)
.expect("services should build");
assert!(services.websocket_pool.is_none());
}
backend/tests/runtime/tasks/coordinator.rs
+26
View File
@@ -1,5 +1,31 @@
use std::{sync::Arc, time::Duration};
use codex_proxy_rs::{
runtime::tasks::coordinator::SchedulerHandle,
upstream::transport::{CodexWebSocketPool, CodexWebSocketPoolConfig},
};
#[tokio::test]
async fn task_coordinator_should_shutdown_without_panicking() {
let coordinator = codex_proxy_rs::runtime::tasks::coordinator::TaskCoordinator::default();
coordinator.shutdown().await;
}
#[tokio::test]
async fn scheduler_handle_should_shutdown_websocket_pool() {
let pool = Arc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig {
maintenance_interval: None,
ping_interval: None,
liveness_timeout: None,
max_age: Duration::from_secs(60),
max_per_account: 1,
enabled: true,
ping_timeout: Duration::ZERO,
}));
SchedulerHandle::from_websocket_pool(pool.clone())
.shutdown()
.await;
assert!(pool.is_shutdown().await);
}
backend/tests/support/config.rs
+2 -1
View File
@@ -10,6 +10,7 @@ pub(crate) fn test_config(database_url: String) -> AppConfig {
server: ServerConfig {
host: "127.0.0.1".to_string(),
port: 0,
trusted_proxies: Vec::new(),
},
api: ApiConfig {
base_url: "https://chatgpt.com/backend-api".to_string(),
@@ -41,7 +42,7 @@ pub(crate) fn test_config(database_url: String) -> AppConfig {
session_ttl_minutes: 1440,
session_cleanup_interval_secs: 3600,
default_username: "admin".to_string(),
default_password: "admin".to_string(),
default_password: "test-admin-password".to_string(),
},
logging: LoggingConfig {
directory: "logs".to_string(),
deploy/Dockerfile
+13 -5
View File
@@ -51,7 +51,10 @@ COPY --from=release-asset-builder /out /
FROM debian:bookworm-slim AS runtime-base
RUN apt-get update \
&& apt-get install -y --no-install-recommends ca-certificates \
&& apt-get install -y --no-install-recommends ca-certificates curl \
&& groupadd --system --gid 10001 cpr \
&& useradd --system --uid 10001 --gid cpr --home-dir /app --shell /usr/sbin/nologin cpr \
&& install -d -o cpr -g cpr /app/data /app/logs /app/web \
&& rm -rf /var/lib/apt/lists/*
WORKDIR /app
@@ -59,18 +62,23 @@ ENV CPR_DEPLOYMENT_MODE=docker
EXPOSE 8080
HEALTHCHECK --interval=30s --timeout=5s --start-period=10s --retries=3 \
CMD curl -fsS http://127.0.0.1:8080/healthz || exit 1
CMD ["codex-proxy-rs"]
FROM runtime-base AS runtime
COPY --from=backend-builder /tmp/codex-proxy-rs /usr/local/bin/codex-proxy-rs
COPY --from=frontend-builder /app/frontend/dist ./web/dist
COPY --from=backend-builder --chown=cpr:cpr /tmp/codex-proxy-rs /usr/local/bin/codex-proxy-rs
COPY --from=frontend-builder --chown=cpr:cpr /app/frontend/dist ./web/dist
USER cpr
FROM runtime-base AS runtime-prebuilt
ARG TARGETARCH
ARG CPR_PREBUILT_DIR=dist/docker
COPY ${CPR_PREBUILT_DIR}/linux-${TARGETARCH}/codex-proxy-rs /usr/local/bin/codex-proxy-rs
COPY ${CPR_PREBUILT_DIR}/web/dist ./web/dist
COPY --chown=cpr:cpr ${CPR_PREBUILT_DIR}/linux-${TARGETARCH}/codex-proxy-rs /usr/local/bin/codex-proxy-rs
COPY --chown=cpr:cpr ${CPR_PREBUILT_DIR}/web/dist ./web/dist
RUN chmod +x /usr/local/bin/codex-proxy-rs
USER cpr
deploy/Dockerfile.dockerignore
+5
View File
@@ -2,6 +2,10 @@
**/target/
frontend/node_modules/
frontend/dist/
frontend/.vite/
dist/*
!dist/docker/
!dist/docker/**
# Application runtime state.
.runtime/
@@ -14,6 +18,7 @@ deploy/logs/
*.db
*.db-shm
*.db-wal
*.sqlite-*
*.log
# Local configuration.
deploy/README.md
+7 -7
View File
@@ -10,7 +10,9 @@ Docker 部署入口集中在本目录。
cp deploy/config.example.yaml deploy/config.yaml
```
`deploy/config.example.yaml` 是 Docker 部署模板，默认将数据写入容器内 `/app/data`、日志写入 `/app/logs`，由 Compose 挂载到宿主机 `deploy/data` 和 `deploy/logs`。
`deploy/config.example.yaml` 是 Docker 部署模板，默认将数据写入容器内 `/app/data`、日志写入 `/app/logs`，由 Compose 命名卷持久化。
复制后必须设置 `admin.default_password`，空值或常见弱口令会被后端拒绝启动。
如果前面还有 Nginx、Caddy、Cloudflare Tunnel 等反向代理，只在该代理会覆盖或清理 `X-Forwarded-For` / `X-Real-IP` / `CF-Connecting-IP` 时，把代理的 IP 或 CIDR 写入 `server.trusted_proxies`；默认空列表会剥离这些客户端可伪造的头并使用 TCP peer IP。
构建并启动：
@@ -19,18 +21,16 @@ docker compose -f deploy/docker-compose.yml build
docker compose -f deploy/docker-compose.yml up -d
```
默认挂载：
默认挂载和卷：
- `deploy/config.yaml` -> `/app/config.yaml`
- `deploy/data` -> `/app/data`
- `deploy/logs` -> `/app/logs`
- `cpr-data` -> `/app/data`
- `cpr-logs` -> `/app/logs`
可以通过环境变量覆盖：
可以通过环境变量覆盖宿主机配置文件路径；容器内会通过 `CPR_CONFIG_FILE=/app/config.yaml` 显式读取挂载后的配置：
```bash
CPR_CONFIG_FILE=/path/to/config.yaml \
CPR_DATA_DIR=/path/to/data \
CPR_LOG_DIR=/path/to/logs \
docker compose -f deploy/docker-compose.yml up -d
```
deploy/config.example.yaml
+4 -2
View File
@@ -3,6 +3,8 @@ server:
host: '0.0.0.0'
# 容器内服务端口，需与 docker-compose.yml 的 target port 一致。
port: 8080
# 可信反向代理 IP 或 CIDR。为空时会忽略并剥离客户端传入的转发 IP 头。
trusted_proxies: []
api:
# ChatGPT / Codex 上游 API 地址。
@@ -82,9 +84,9 @@ admin:
session_ttl_minutes: 1440
# 过期会话清理间隔，单位秒。
session_cleanup_interval_secs: 3600
# 首次启动创建的默认管理员账号。生产部署前必须修改默认密码。
# 首次启动创建的默认管理员账号。default_password 必须改成至少 12 位的非默认密码。
default_username: admin@cpr.local
default_password: 123456
default_password: ''
logging:
# Docker 部署默认写入 /app/logs，由 Compose 挂载到宿主机 deploy/logs。
deploy/docker-compose.yml
+26 -3
View File
@@ -12,13 +12,32 @@ services:
CPR_BUILD_TYPE: ${CPR_BUILD_TYPE:-release}
container_name: codex-proxy-rs
restart: unless-stopped
user: "10001:10001"
init: true
stop_grace_period: 30s
pids_limit: 256
cpus: "2.0"
mem_limit: 1g
ports:
- "8080:8080"
- "127.0.0.1:8080:8080"
volumes:
- ${CPR_CONFIG_FILE:-./config.yaml}:/app/config.yaml:ro
- ${CPR_DATA_DIR:-./data}:/app/data
- ${CPR_LOG_DIR:-./logs}:/app/logs
- cpr-data:/app/data
- cpr-logs:/app/logs
security_opt:
- no-new-privileges:true
cap_drop:
- ALL
tmpfs:
- /tmp:size=64m,mode=1777
healthcheck:
test: ["CMD", "curl", "-fsS", "http://127.0.0.1:8080/healthz"]
interval: 30s
timeout: 5s
retries: 3
start_period: 10s
environment:
CPR_CONFIG_FILE: /app/config.yaml
CPR_DEPLOYMENT_MODE: docker
CPR_UPDATE_REPOSITORY: zyycn/codex-proxy-rs
CPR_UPDATE_CHANNEL: stable
@@ -26,3 +45,7 @@ services:
CPR_UPDATE_LOCK_FILE: /app/data/update.lock
CPR_WEB_DIST_DIR: /app/web/dist
CPR_ENABLE_SELF_RESTART: "true"
volumes:
cpr-data:
cpr-logs:

Some files were not shown because too many files have changed in this diff Show More