diff --git a/.cargo/config.toml b/.cargo/config.toml new file mode 100644 index 0000000..a9d3ab0 --- /dev/null +++ b/.cargo/config.toml @@ -0,0 +1,8 @@ +# The usage ledger builds SQLite in-tree (rusqlite "bundled"), which needs a +# C compiler. cargo-dist's release CI builds aarch64-unknown-linux-musl on +# GitHub's arm64 runners with only `musl-tools` installed, whose compiler is +# `musl-gcc`; the cc crate defaults to `aarch64-linux-musl-gcc` for that +# target, which musl-tools does not provide. Point it at musl-gcc, the same +# compiler the x86_64-unknown-linux-musl build already uses. +[env] +CC_aarch64_unknown_linux_musl = "musl-gcc" diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e95eb05..fefde3c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -42,5 +42,7 @@ jobs: with: toolchain: 1.96.1 - uses: Swatinem/rust-cache@v2 - - run: cargo install cargo-shear --locked + # Pinned: cargo-shear >= 1.14 pulls ra_ap_* crates that require + # rustc 1.98, above the 1.96.1 toolchain pinned above. Bump both together. + - run: cargo install cargo-shear --locked --version 1.13.0 - run: cargo shear diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 0000000..1dfcd0f --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,296 @@ +# This file was autogenerated by dist: https://axodotdev.github.io/cargo-dist +# +# Copyright 2022-2024, axodotdev +# SPDX-License-Identifier: MIT or Apache-2.0 +# +# CI that: +# +# * checks for a Git Tag that looks like a release +# * builds artifacts with dist (archives, installers, hashes) +# * uploads those artifacts to temporary workflow zip +# * on success, uploads the artifacts to a GitHub Release +# +# Note that the GitHub Release will be created with a generated +# title/body based on your changelogs. + +name: Release +permissions: + "contents": "write" + +# This task will run whenever you push a git tag that looks like a version +# like "1.0.0", "v0.1.0-prerelease.1", "my-app/0.1.0", "releases/v1.0.0", etc. +# Various formats will be parsed into a VERSION and an optional PACKAGE_NAME, where +# PACKAGE_NAME must be the name of a Cargo package in your workspace, and VERSION +# must be a Cargo-style SemVer Version (must have at least major.minor.patch). +# +# If PACKAGE_NAME is specified, then the announcement will be for that +# package (erroring out if it doesn't have the given version or isn't dist-able). +# +# If PACKAGE_NAME isn't specified, then the announcement will be for all +# (dist-able) packages in the workspace with that version (this mode is +# intended for workspaces with only one dist-able package, or with all dist-able +# packages versioned/released in lockstep). +# +# If you push multiple tags at once, separate instances of this workflow will +# spin up, creating an independent announcement for each one. However, GitHub +# will hard limit this to 3 tags per commit, as it will assume more tags is a +# mistake. +# +# If there's a prerelease-style suffix to the version, then the release(s) +# will be marked as a prerelease. +on: + pull_request: + push: + tags: + - '**[0-9]+.[0-9]+.[0-9]+*' + +jobs: + # Run 'dist plan' (or host) to determine what tasks we need to do + plan: + runs-on: "ubuntu-22.04" + outputs: + val: ${{ steps.plan.outputs.manifest }} + tag: ${{ !github.event.pull_request && github.ref_name || '' }} + tag-flag: ${{ !github.event.pull_request && format('--tag={0}', github.ref_name) || '' }} + publishing: ${{ !github.event.pull_request }} + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + steps: + - uses: actions/checkout@v6 + with: + persist-credentials: false + submodules: recursive + - name: Install dist + # we specify bash to get pipefail; it guards against the `curl` command + # failing. otherwise `sh` won't catch that `curl` returned non-0 + shell: bash + run: "curl --proto '=https' --tlsv1.2 -LsSf https://github.com/axodotdev/cargo-dist/releases/download/v0.32.0/cargo-dist-installer.sh | sh" + - name: Cache dist + uses: actions/upload-artifact@v7 + with: + name: cargo-dist-cache + path: ~/.cargo/bin/dist + # sure would be cool if github gave us proper conditionals... + # so here's a doubly-nested ternary-via-truthiness to try to provide the best possible + # functionality based on whether this is a pull_request, and whether it's from a fork. + # (PRs run on the *source* but secrets are usually on the *target* -- that's *good* + # but also really annoying to build CI around when it needs secrets to work right.) + - id: plan + run: | + dist ${{ (!github.event.pull_request && format('host --steps=create --tag={0}', github.ref_name)) || 'plan' }} --output-format=json > plan-dist-manifest.json + echo "dist ran successfully" + cat plan-dist-manifest.json + echo "manifest=$(jq -c "." plan-dist-manifest.json)" >> "$GITHUB_OUTPUT" + - name: "Upload dist-manifest.json" + uses: actions/upload-artifact@v7 + with: + name: artifacts-plan-dist-manifest + path: plan-dist-manifest.json + + # Build and packages all the platform-specific things + build-local-artifacts: + name: build-local-artifacts (${{ join(matrix.targets, ', ') }}) + # Let the initial task tell us to not run (currently very blunt) + needs: + - plan + if: ${{ fromJson(needs.plan.outputs.val).ci.github.artifacts_matrix.include != null && (needs.plan.outputs.publishing == 'true' || fromJson(needs.plan.outputs.val).ci.github.pr_run_mode == 'upload') }} + strategy: + fail-fast: false + # Target platforms/runners are computed by dist in create-release. + # Each member of the matrix has the following arguments: + # + # - runner: the github runner + # - dist-args: cli flags to pass to dist + # - install-dist: expression to run to install dist on the runner + # + # Typically there will be: + # - 1 "global" task that builds universal installers + # - N "local" tasks that build each platform's binaries and platform-specific installers + matrix: ${{ fromJson(needs.plan.outputs.val).ci.github.artifacts_matrix }} + runs-on: ${{ matrix.runner }} + container: ${{ matrix.container && matrix.container.image || null }} + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + BUILD_MANIFEST_NAME: target/distrib/${{ join(matrix.targets, '-') }}-dist-manifest.json + steps: + - name: enable windows longpaths + run: | + git config --global core.longpaths true + - uses: actions/checkout@v6 + with: + persist-credentials: false + submodules: recursive + - name: Install Rust non-interactively if not already installed + if: ${{ matrix.container }} + run: | + if ! command -v cargo > /dev/null 2>&1; then + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y + echo "$HOME/.cargo/bin" >> $GITHUB_PATH + fi + - name: Install dist + run: ${{ matrix.install_dist.run }} + # Get the dist-manifest + - name: Fetch local artifacts + uses: actions/download-artifact@v8 + with: + pattern: artifacts-* + path: target/distrib/ + merge-multiple: true + - name: Install dependencies + run: | + ${{ matrix.packages_install }} + - name: Build artifacts + run: | + # Actually do builds and make zips and whatnot + dist build ${{ needs.plan.outputs.tag-flag }} --print=linkage --output-format=json ${{ matrix.dist_args }} > dist-manifest.json + echo "dist ran successfully" + - id: cargo-dist + name: Post-build + # We force bash here just because github makes it really hard to get values up + # to "real" actions without writing to env-vars, and writing to env-vars has + # inconsistent syntax between shell and powershell. + shell: bash + run: | + # Parse out what we just built and upload it to scratch storage + echo "paths<> "$GITHUB_OUTPUT" + dist print-upload-files-from-manifest --manifest dist-manifest.json >> "$GITHUB_OUTPUT" + echo "EOF" >> "$GITHUB_OUTPUT" + + cp dist-manifest.json "$BUILD_MANIFEST_NAME" + - name: "Upload artifacts" + uses: actions/upload-artifact@v7 + with: + name: artifacts-build-local-${{ join(matrix.targets, '_') }} + path: | + ${{ steps.cargo-dist.outputs.paths }} + ${{ env.BUILD_MANIFEST_NAME }} + + # Build and package all the platform-agnostic(ish) things + build-global-artifacts: + needs: + - plan + - build-local-artifacts + runs-on: "ubuntu-22.04" + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + BUILD_MANIFEST_NAME: target/distrib/global-dist-manifest.json + steps: + - uses: actions/checkout@v6 + with: + persist-credentials: false + submodules: recursive + - name: Install cached dist + uses: actions/download-artifact@v8 + with: + name: cargo-dist-cache + path: ~/.cargo/bin/ + - run: chmod +x ~/.cargo/bin/dist + # Get all the local artifacts for the global tasks to use (for e.g. checksums) + - name: Fetch local artifacts + uses: actions/download-artifact@v8 + with: + pattern: artifacts-* + path: target/distrib/ + merge-multiple: true + - id: cargo-dist + shell: bash + run: | + dist build ${{ needs.plan.outputs.tag-flag }} --output-format=json "--artifacts=global" > dist-manifest.json + echo "dist ran successfully" + + # Parse out what we just built and upload it to scratch storage + echo "paths<> "$GITHUB_OUTPUT" + jq --raw-output ".upload_files[]" dist-manifest.json >> "$GITHUB_OUTPUT" + echo "EOF" >> "$GITHUB_OUTPUT" + + cp dist-manifest.json "$BUILD_MANIFEST_NAME" + - name: "Upload artifacts" + uses: actions/upload-artifact@v7 + with: + name: artifacts-build-global + path: | + ${{ steps.cargo-dist.outputs.paths }} + ${{ env.BUILD_MANIFEST_NAME }} + # Determines if we should publish/announce + host: + needs: + - plan + - build-local-artifacts + - build-global-artifacts + # Only run if we're "publishing", and only if plan, local and global didn't fail (skipped is fine) + if: ${{ always() && needs.plan.result == 'success' && needs.plan.outputs.publishing == 'true' && (needs.build-global-artifacts.result == 'skipped' || needs.build-global-artifacts.result == 'success') && (needs.build-local-artifacts.result == 'skipped' || needs.build-local-artifacts.result == 'success') }} + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + runs-on: "ubuntu-22.04" + outputs: + val: ${{ steps.host.outputs.manifest }} + steps: + - uses: actions/checkout@v6 + with: + persist-credentials: false + submodules: recursive + - name: Install cached dist + uses: actions/download-artifact@v8 + with: + name: cargo-dist-cache + path: ~/.cargo/bin/ + - run: chmod +x ~/.cargo/bin/dist + # Fetch artifacts from scratch-storage + - name: Fetch artifacts + uses: actions/download-artifact@v8 + with: + pattern: artifacts-* + path: target/distrib/ + merge-multiple: true + - id: host + shell: bash + run: | + dist host ${{ needs.plan.outputs.tag-flag }} --steps=upload --steps=release --output-format=json > dist-manifest.json + echo "artifacts uploaded and released successfully" + cat dist-manifest.json + echo "manifest=$(jq -c "." dist-manifest.json)" >> "$GITHUB_OUTPUT" + - name: "Upload dist-manifest.json" + uses: actions/upload-artifact@v7 + with: + # Overwrite the previous copy + name: artifacts-dist-manifest + path: dist-manifest.json + # Create a GitHub Release while uploading all files to it + - name: "Download GitHub Artifacts" + uses: actions/download-artifact@v8 + with: + pattern: artifacts-* + path: artifacts + merge-multiple: true + - name: Cleanup + run: | + # Remove the granular manifests + rm -f artifacts/*-dist-manifest.json + - name: Create GitHub Release + env: + PRERELEASE_FLAG: "${{ fromJson(steps.host.outputs.manifest).announcement_is_prerelease && '--prerelease' || '' }}" + ANNOUNCEMENT_TITLE: "${{ fromJson(steps.host.outputs.manifest).announcement_title }}" + ANNOUNCEMENT_BODY: "${{ fromJson(steps.host.outputs.manifest).announcement_github_body }}" + RELEASE_COMMIT: "${{ github.sha }}" + run: | + # Write and read notes from a file to avoid quoting breaking things + echo "$ANNOUNCEMENT_BODY" > $RUNNER_TEMP/notes.txt + + gh release create "${{ needs.plan.outputs.tag }}" --target "$RELEASE_COMMIT" $PRERELEASE_FLAG --title "$ANNOUNCEMENT_TITLE" --notes-file "$RUNNER_TEMP/notes.txt" artifacts/* + + announce: + needs: + - plan + - host + # use "always() && ..." to allow us to wait for all publish jobs while + # still allowing individual publish jobs to skip themselves (for prereleases). + # "host" however must run to completion, no skipping allowed! + if: ${{ always() && needs.host.result == 'success' }} + runs-on: "ubuntu-22.04" + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + steps: + - uses: actions/checkout@v6 + with: + persist-credentials: false + submodules: recursive diff --git a/.gitignore b/.gitignore index 96f35bd..69dae47 100644 --- a/.gitignore +++ b/.gitignore @@ -27,3 +27,6 @@ !/ref/REFS .playwright-mcp + +.z-proflow/ +.zcode/ diff --git a/Cargo.lock b/Cargo.lock index cb4e128..6feb413 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -111,6 +111,18 @@ dependencies = [ "rustversion", ] +[[package]] +name = "argon2" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "134c52ddac6d63c576bef8168db10c83c49c26444ecbc68060fef078925a901c" +dependencies = [ + "base64ct", + "blake2", + "cpufeatures 0.3.1", + "password-hash", +] + [[package]] name = "arrayvec" version = "0.7.6" @@ -205,6 +217,12 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "base64ct" +version = "1.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" + [[package]] name = "bit-set" version = "0.5.3" @@ -238,6 +256,15 @@ dependencies = [ "wyz", ] +[[package]] +name = "blake2" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b5d4d889834ee8ecfc0f8426ad30faf7cdcb10f741a8e6d7224d95325479f6f" +dependencies = [ + "digest 0.11.3", +] + [[package]] name = "block-buffer" version = "0.10.4" @@ -247,6 +274,15 @@ dependencies = [ "generic-array", ] +[[package]] +name = "block-buffer" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2f6c7dbe95a6ed67ad9f18e57daf93a2f034c524b99fd2b76d18fdfeb6660aa" +dependencies = [ + "hybrid-array", +] + [[package]] name = "borsh" version = "1.7.0" @@ -393,6 +429,12 @@ dependencies = [ "hashbrown 0.16.1", ] +[[package]] +name = "cmov" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" + [[package]] name = "colorchoice" version = "1.0.5" @@ -445,6 +487,15 @@ dependencies = [ "libc", ] +[[package]] +name = "cpufeatures" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ca28b0ae3115b884660db4118d803791fd6756b6e88f39c0f3f7859060d7566" +dependencies = [ + "libc", +] + [[package]] name = "crc32fast" version = "1.5.0" @@ -470,6 +521,24 @@ dependencies = [ "typenum", ] +[[package]] +name = "crypto-common" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce6e4c961d6cd6c9a86db418387425e8bdeaf05b3c8bc1411e6dca4c252f1453" +dependencies = [ + "hybrid-array", +] + +[[package]] +name = "ctutils" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03bb0e1cc970d482d121d9a1744999169b69a07470b3d644a7894e53fcaf4574" +dependencies = [ + "cmov", +] + [[package]] name = "darling" version = "0.20.11" @@ -566,8 +635,19 @@ version = "0.10.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" dependencies = [ - "block-buffer", - "crypto-common", + "block-buffer 0.10.4", + "crypto-common 0.1.7", +] + +[[package]] +name = "digest" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2" +dependencies = [ + "block-buffer 0.12.1", + "crypto-common 0.2.2", + "ctutils", ] [[package]] @@ -618,6 +698,18 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "fallible-iterator" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2acce4a10f12dc2fb14a218589d4f1f62ef011b2d0cc4b3cb1bba8e94da14649" + +[[package]] +name = "fallible-streaming-iterator" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" + [[package]] name = "fancy-regex" version = "0.13.0" @@ -1749,6 +1841,18 @@ name = "hashbrown" version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +dependencies = [ + "foldhash 0.2.0", +] + +[[package]] +name = "hashlink" +version = "0.12.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a596f1b20ed2cc5ecac41a164aaebc7258057060f06c0cf7a2ba3991ee7990fb" +dependencies = [ + "hashbrown 0.17.1", +] [[package]] name = "heapless" @@ -1827,6 +1931,15 @@ dependencies = [ "serde", ] +[[package]] +name = "hybrid-array" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "27f864f10dfb56725ce5ce5472bc52252c8f93a4ab86327122cebf62c5f59a17" +dependencies = [ + "typenum", +] + [[package]] name = "hyper" version = "1.10.1" @@ -2155,6 +2268,17 @@ version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" +[[package]] +name = "libsqlite3-sys" +version = "0.38.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1d20bef17f513b9b3004532233187769cd072d790971f4e4da0e346eb6401e8" +dependencies = [ + "cc", + "pkg-config", + "vcpkg", +] + [[package]] name = "linux-raw-sys" version = "0.4.15" @@ -2183,8 +2307,11 @@ dependencies = [ "insta", "libc", "llm-proxy-core", + "llm-proxy-protocol", "llm-proxy-provider", "llm-proxy-server", + "llm-proxy-storage", + "secrecy", "socket2", "tempfile", "time", @@ -2198,6 +2325,25 @@ dependencies = [ [[package]] name = "llm-proxy-api" version = "0.1.0" +dependencies = [ + "axum", + "humantime", + "llm-proxy-protocol", + "llm-proxy-storage", + "rusqlite", + "rust_decimal", + "serde", + "serde_json", + "sha2", + "tempfile", + "thiserror", + "time", + "tokio", + "tower", + "tower-http", + "tracing", + "tracing-subscriber", +] [[package]] name = "llm-proxy-core" @@ -2252,9 +2398,11 @@ dependencies = [ name = "llm-proxy-server" version = "0.1.0" dependencies = [ + "argon2", "axum", "bytes", "futures", + "llm-proxy-api", "llm-proxy-core", "llm-proxy-protocol", "llm-proxy-provider", @@ -2265,6 +2413,7 @@ dependencies = [ "serde", "serde_json", "sha2", + "tempfile", "thiserror", "time", "tokio", @@ -2274,6 +2423,7 @@ dependencies = [ "tower", "tower-http", "tracing", + "tracing-subscriber", "uuid", ] @@ -2282,9 +2432,15 @@ name = "llm-proxy-storage" version = "0.1.0" dependencies = [ "llm-proxy-protocol", + "rusqlite", + "rust_decimal", "serde", "serde_json", + "tempfile", + "thiserror", "time", + "tokio", + "tracing", ] [[package]] @@ -2504,12 +2660,33 @@ dependencies = [ "windows-link", ] +[[package]] +name = "password-hash" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aab41826031698d6ffcd9cff78ef56ef998e39dc7e5067cdfebe373842d4723b" +dependencies = [ + "getrandom 0.4.2", + "phc", +] + [[package]] name = "percent-encoding" version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" +[[package]] +name = "phc" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44dc769b75f93afdddd8c7fa12d685292ddeff1e66f7f0f3a234cf1818afe892" +dependencies = [ + "base64ct", + "ctutils", + "getrandom 0.4.2", +] + [[package]] name = "pin-project-lite" version = "0.2.17" @@ -2897,6 +3074,31 @@ dependencies = [ "syn 1.0.109", ] +[[package]] +name = "rsqlite-vfs" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c51c9ae4df8a7fba42103df5c621fa3c37eccf3a3c650879e90fc48b11cc192c" +dependencies = [ + "hashbrown 0.16.1", + "thiserror", +] + +[[package]] +name = "rusqlite" +version = "0.40.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23f2a97da3e3873c73cb2a2e71b35c40ff95e0b1eefa8d72d8499a6928c3b5b3" +dependencies = [ + "bitflags", + "fallible-iterator", + "fallible-streaming-iterator", + "hashlink", + "libsqlite3-sys", + "smallvec", + "sqlite-wasm-rs", +] + [[package]] name = "rust_decimal" version = "1.42.1" @@ -3150,8 +3352,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" dependencies = [ "cfg-if", - "cpufeatures", - "digest", + "cpufeatures 0.2.17", + "digest 0.10.7", ] [[package]] @@ -3160,7 +3362,7 @@ version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "89f599ac0c323ebb1c6082821a54962b839832b03984598375bff3975b804423" dependencies = [ - "digest", + "digest 0.10.7", "sha1", ] @@ -3171,8 +3373,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" dependencies = [ "cfg-if", - "cpufeatures", - "digest", + "cpufeatures 0.2.17", + "digest 0.10.7", ] [[package]] @@ -3256,6 +3458,18 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "sqlite-wasm-rs" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc3efc0da82635d7e1ced0053bbbfa8c7ab9645d0bf36ceb4f7127bb85315d75" +dependencies = [ + "cc", + "js-sys", + "rsqlite-vfs", + "wasm-bindgen", +] + [[package]] name = "stable_deref_trait" version = "1.2.1" diff --git a/Cargo.toml b/Cargo.toml index d454d12..3ca0631 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -25,7 +25,7 @@ llm-proxy-core = { path = "crates/llm-proxy-core" } llm-proxy-server = { path = "crates/llm-proxy-server" } # Async / runtime -tokio = { version = "1", features = ["macros", "rt-multi-thread", "signal", "time"] } +tokio = { version = "1", features = ["macros", "rt-multi-thread", "signal", "sync", "time"] } tokio-util = { version = "0.7", features = ["rt"] } # Web @@ -50,6 +50,8 @@ anyhow = "1" # Secrets secrecy = { version = "0.10", features = ["serde"] } +# Password hashing KDF for client-key verifiers (PHC strings) +argon2 = "0.6" # Tokenization tiktoken-rs = "0.6" @@ -57,6 +59,10 @@ tiktoken-rs = "0.6" # Money (cost/pricing) -- fixed-point, never float rust_decimal = "1" +# Persistence (usage ledger). "bundled" compiles SQLite in-tree so no +# system libsqlite3 is needed on deploy targets. +rusqlite = { version = "0.40", features = ["bundled"] } + # Observability tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt", "json"] } @@ -161,3 +167,8 @@ literal_string_with_formatting_args = "allow" lto = "thin" codegen-units = 1 strip = "symbols" + +# The profile that 'dist' will build with (cargo-dist passes --profile dist) +[profile.dist] +inherits = "release" +lto = "thin" diff --git a/apps/llm-proxy/Cargo.toml b/apps/llm-proxy/Cargo.toml index 0c595c7..25984f5 100644 --- a/apps/llm-proxy/Cargo.toml +++ b/apps/llm-proxy/Cargo.toml @@ -25,6 +25,7 @@ workspace = true llm-proxy-server = { workspace = true } llm-proxy-core = { workspace = true } llm-proxy-provider = { path = "../../crates/llm-proxy-provider" } +llm-proxy-storage = { path = "../../crates/llm-proxy-storage" } anyhow = { workspace = true } axum = { workspace = true } @@ -43,5 +44,7 @@ vergen-gix = { workspace = true } [dev-dependencies] insta = { version = "1.42", features = ["yaml", "filters"] } +llm-proxy-protocol = { path = "../../crates/llm-proxy-protocol" } +secrecy = { workspace = true } tempfile = "3" toml = { workspace = true } diff --git a/apps/llm-proxy/src/commands/serve.rs b/apps/llm-proxy/src/commands/serve.rs index 53fffc7..abd23d6 100644 --- a/apps/llm-proxy/src/commands/serve.rs +++ b/apps/llm-proxy/src/commands/serve.rs @@ -8,7 +8,7 @@ use anyhow::{Context, Result, bail}; use axum::serve::{Listener, ListenerExt}; use llm_proxy_core::LogFormat; use llm_proxy_provider::{ProviderAdapterRegistry, ProxyClient}; -use llm_proxy_server::{build_router, shutdown_signal}; +use llm_proxy_server::{admin_router, build_router, shutdown_signal}; use tokio::net::TcpListener; use tracing::info; @@ -61,6 +61,11 @@ pub async fn cmd_serve( // ensures a bad config does not leave a stale PID file on disk. let state = load_toml_state(&path, adapter_registry, proxy_client, port_override)?; + // Dedicated admin listener for the `/usage` surface (`[auth.admin]`), + // bound before the PID file and the main listener so a busy or unusable + // admin port is a clean startup error that leaves no stale PID file. + let admin = bind_admin_listener(&state).await?; + // Attempt atomic PID file creation. If it fails because the file already // exists, check whether the existing owner is still alive before removing // the stale file and retrying. This closes the TOCTOU window between the @@ -92,6 +97,11 @@ pub async fn cmd_serve( // by process::exit), so remove_pid() runs exactly once per shutdown. let bind_addr = state.bind_address(); let shutdown_deadline = state.shutdown_timeout(); + // Hold the event bus so the usage ledger can be flushed after the server + // stops. The router (which owns the state) keeps its own Arc; closing is + // watch-signalled, not sender-drop-signalled, so those clones do not + // hold the writer open. + let event_bus = std::sync::Arc::clone(state.event_bus()); let app = build_router(state); // `tap_io` runs on every accepted connection (audit MEDIUM-2 option A). @@ -127,31 +137,35 @@ pub async fn cmd_serve( ) .with_graceful_shutdown(shutdown_signal()); - let result = if shutdown_deadline.is_zero() { - serve.await - } else { - match tokio::time::timeout(shutdown_deadline, serve).await { - Ok(r) => r, - Err(_) => { - tracing::error!( - deadline = ?shutdown_deadline, - "graceful-shutdown deadline elapsed; forcing exit with in-flight connections dropped" - ); - // Run PID cleanup best-effort, then exit. The deadline exists - // precisely to escape stuck connections, so we do not wait for - // them; remove_pid is fast, synchronous filesystem I/O. - if let Err(e) = remove_pid() { - tracing::warn!(error = %e, "failed to remove PID file during shutdown"); - } - std::process::exit(1); + // The admin listener serves concurrently and drains on its own copy of + // the shutdown signal; its failure is logged, not fatal to the main + // server (the proxies' clients come first). + let admin_serve = async { + if let Some((listener, router)) = admin { + if let Err(error) = axum::serve(listener, router.into_make_service()) + .with_graceful_shutdown(shutdown_signal()) + .await + { + tracing::error!(%error, "admin usage listener failed"); } } }; + let result = serve_with_deadline(serve.into_future(), admin_serve, shutdown_deadline).await; + if let Err(e) = remove_pid() { tracing::warn!(error = %e, "failed to remove PID file during shutdown"); } + // Flush the usage ledger now that the graceful drain has finished and no + // new events can arrive: the writer drains its channel and exits. For + // the default NoopBus this is a no-op. close() blocks (bounded), so it + // runs on the blocking pool. (The hard-deadline branch above force-exits + // without flushing; queued events are lost there by design.) + if let Err(close_err) = tokio::task::spawn_blocking(move || event_bus.close()).await { + tracing::warn!(error = %close_err, "usage-ledger close task failed to join"); + } + match result { Ok(()) => { info!("server stopped cleanly"); @@ -164,6 +178,94 @@ pub async fn cmd_serve( } } +/// Serve the main and admin listeners concurrently under the graceful- +/// shutdown hard deadline. +/// +/// axum 0.8's `with_graceful_shutdown` waits indefinitely for in-flight +/// connections to drain. Because the proxy serves long-lived streaming +/// responses an upstream can hold open arbitrarily, one stuck stream could +/// otherwise pin the process forever on SIGTERM/SIGINT. Once the deadline +/// elapses after the signal, stuck connections are dropped and the process +/// force-exits so an orchestrator (Kubernetes/LB) can recycle the pod. +/// `shutdown_timeout == 0` disables the deadline (drain forever) for +/// operators who want that. +async fn serve_with_deadline( + serve: F, + admin_serve: A, + shutdown_deadline: std::time::Duration, +) -> std::io::Result<()> +where + F: std::future::Future>, + A: std::future::Future, +{ + if shutdown_deadline.is_zero() { + let (main, ()) = tokio::join!(serve, admin_serve); + return main; + } + match tokio::time::timeout(shutdown_deadline, async { + tokio::join!(serve, admin_serve) + }) + .await + { + Ok((main, ())) => main, + Err(_) => { + tracing::error!( + deadline = ?shutdown_deadline, + "graceful-shutdown deadline elapsed; forcing exit with in-flight connections dropped" + ); + // Run PID cleanup best-effort, then exit. The deadline exists + // precisely to escape stuck connections, so we do not wait for + // them; remove_pid is fast, synchronous filesystem I/O. + if let Err(e) = remove_pid() { + tracing::warn!(error = %e, "failed to remove PID file during shutdown"); + } + std::process::exit(1); + } + } +} + +/// Bind the dedicated admin `/usage` listener when `[auth.admin]` is +/// configured and the dedicated socket should exist. +/// +/// The default bind is loopback on the port after the main bind +/// (`ServerConfig::effective_admin_bind`). Under the +/// `server.admin_on_main_listener` opt-in with no explicit `admin_bind`, +/// NO dedicated socket is attempted — that is precisely the posture for +/// setups where a second port is not available, and `/usage` is served on +/// the main listener instead (handled inside `build_router`). An explicit +/// `admin_bind` always binds the dedicated listener, alongside any +/// main-listener mount. Returns `None` when there is nothing to bind. +async fn bind_admin_listener( + state: &llm_proxy_server::AppState, +) -> Result> { + let Some(posture) = state.app_config().admin_posture() else { + return Ok(None); + }; + if !state.app_config().server.dedicated_admin_listener() { + tracing::info!( + posture = posture.name(), + "admin usage API served on the main listener only; no dedicated socket" + ); + return Ok(None); + } + let Some(addr) = state.app_config().server.effective_admin_bind() else { + // Unreachable for validated configs: `load_app_config` rejects an + // admin API whose dedicated listener would bind a default port that + // cannot be derived. Refuse to guess a bind. + bail!("admin API enabled but no admin bind could be derived"); + }; + let Some(router) = admin_router(state) else { + bail!("admin API enabled but no /usage router could be built"); + }; + let listener = bind_listener(addr).await?; + info!( + addr = %addr, + posture = posture.name(), + "admin usage API listening (read-only, behind the admin token)" + ); + Ok(Some((listener, router))) +} + /// Bind the listening TCP socket with connection-hardening options. /// /// On Unix the socket is created via `socket2` so we can enable `SO_KEEPALIVE` diff --git a/apps/llm-proxy/src/platform.rs b/apps/llm-proxy/src/platform.rs index 2defb52..e66010a 100644 --- a/apps/llm-proxy/src/platform.rs +++ b/apps/llm-proxy/src/platform.rs @@ -3,6 +3,9 @@ //! Provides plist generation (macOS), desktop entry generation (Linux), //! and associated escaping utilities. +// Only the macOS plist builder needs the config dir (for the log path); the +// Linux desktop entry carries no absolute paths. +#[cfg(target_os = "macos")] use crate::paths::config_dir; /// Generate a macOS launchd plist for auto-start. @@ -139,6 +142,7 @@ mod tests { /// Used to build an `insta::Settings::add_filter` regex that matches the /// verbatim log path (which contains regex-special characters like `.`) /// so it can be redacted to `[log_path]` in snapshots. + #[cfg(target_os = "macos")] fn regex_escape(literal: &str) -> String { let mut out = String::with_capacity(literal.len()); for ch in literal.chars() { @@ -269,10 +273,16 @@ mod tests { #[test] fn desktop_exec_escape_shell_metacharacters() { let escaped = desktop_exec_escape("$HOME/`echo test`"); - assert!(!escaped.contains('$')); - assert!(!escaped.contains('`')); + // Every metacharacter must survive in backslash-escaped form... assert!(escaped.contains("\\$")); assert!(escaped.contains("\\`")); + // ...and none may remain unescaped: once the escaped forms are + // removed, no bare '$' or '`' may be left in the output. + let unescaped = escaped.replace("\\$", "").replace("\\`", ""); + assert!( + !unescaped.contains('$') && !unescaped.contains('`'), + "metacharacters must always be escaped, got: {escaped}" + ); } #[cfg(target_os = "linux")] diff --git a/apps/llm-proxy/src/snapshots/llm_proxy_app__platform__tests__desktop_entry_basic.snap b/apps/llm-proxy/src/snapshots/llm_proxy_app__platform__tests__desktop_entry_basic.snap new file mode 100644 index 0000000..67aeae1 --- /dev/null +++ b/apps/llm-proxy/src/snapshots/llm_proxy_app__platform__tests__desktop_entry_basic.snap @@ -0,0 +1,11 @@ +--- +source: apps/llm-proxy/src/platform.rs +assertion_line: 302 +expression: entry +--- +[Desktop Entry] +Type=Application +Name=LLM Proxy +Exec=/usr/bin/llm-proxy serve +X-GNOME-Autostart-enabled=true +Hidden=false diff --git a/apps/llm-proxy/src/state.rs b/apps/llm-proxy/src/state.rs index da90357..49f69e0 100644 --- a/apps/llm-proxy/src/state.rs +++ b/apps/llm-proxy/src/state.rs @@ -66,12 +66,72 @@ pub(crate) fn load_toml_state( } } + // Build the client-auth state here (before the listener binds) so an + // invalid verifier hash is a startup error, not a runtime fail-closed. + let client_auth = llm_proxy_server::ClientAuth::from_config(&app_config.auth) + .with_context(|| "building client auth from [auth] config")?; + + // Build per-provider egress clients from [provider.egress] entries here + // (also before the listener binds) so a proxy URL the transport cannot + // use is a startup error, never a silent fallback to direct egress. + // Providers without an egress section get no entry and keep using the + // shared default client, exactly as before the section existed. + let mut egress_clients = std::collections::HashMap::new(); + for provider in registry.iter() { + if let Some(egress) = &provider.egress { + let client = ProxyClient::try_with_egress_proxy(&egress.proxy).with_context(|| { + format!( + "building egress proxy client for provider \"{}\" from [provider.egress]", + provider.name + ) + })?; + // Log the provider and scheme, not the URL: egress proxy URLs may + // carry userinfo credentials. + info!( + provider = %provider.name, + scheme = egress.proxy.split("://").next().unwrap_or_default(), + "routing upstream traffic through the configured egress proxy" + ); + egress_clients.insert(provider.name.clone(), std::sync::Arc::new(client)); + } + } + + // Build the persistent event bus from `[storage]` (also before the + // listener binds) so an unusable ledger path is a startup error rather + // than a silent no-op. Absent `[storage]` keeps the NoopBus default: + // nothing is persisted, exactly as before the section existed. + let event_bus: std::sync::Arc = match &app_config.storage { + Some(storage) => { + let bus = llm_proxy_storage::SqliteEventBus::open( + &storage.path, + llm_proxy_storage::RetentionPolicy { + max_age: storage.max_age, + max_rows: storage.max_rows, + prune_every: storage.prune_interval, + }, + ) + .with_context(|| format!("opening usage ledger at {}", storage.path.display()))?; + info!( + path = %storage.path.display(), + "persisting usage events to the SQLite ledger" + ); + std::sync::Arc::new(bus) + } + None => std::sync::Arc::new(llm_proxy_storage::NoopBus), + }; + info!( config = %path.display(), providers = registry.len(), "loaded TOML config" ); + // Announce the admin-mount security posture loudly at startup so the + // operator sees which surface their config selected. The + // main-listener opt-in is a different, deliberately less-isolated + // posture than the default loopback-only dedicated listener. + warn_admin_posture(&app_config); + Ok(AppState::new_with_catalog_dir( app_config, registry, @@ -79,7 +139,48 @@ pub(crate) fn load_toml_state( proxy_client, build_info(), Some(cache_dir), - )) + ) + .with_client_auth(std::sync::Arc::new(client_auth)) + .with_egress_clients(egress_clients) + .with_event_bus(event_bus)) +} + +/// Emit the startup warnings naming the admin `/usage` mount posture. +/// +/// `admin_on_main_listener = true` deliberately exposes the admin surface on +/// the main listener, so it is called out by name. A dedicated admin +/// listener on a non-loopback address is likewise called out (the default is +/// loopback-only); the operator may have pointed it at a tunnel interface, +/// but the warning makes that choice visible in every startup log. +pub fn warn_admin_posture(config: &AppConfig) { + use llm_proxy_core::AdminMountPosture; + match config.admin_posture() { + Some(AdminMountPosture::MainListener) => { + tracing::warn!( + posture = AdminMountPosture::MainListener.name(), + main_bind = %config.server.bind, + "ADMIN SECURITY POSTURE: /usage is also served on the main listener, \ + protected only by the admin token — every network that reaches the \ + main listener reaches the admin API. Set server.admin_on_main_listener \ + = false to return to the loopback-only dedicated admin listener." + ); + } + Some(AdminMountPosture::DedicatedListener) => { + if let Some(bind) = config.server.effective_admin_bind() + && !bind.ip().is_loopback() + { + tracing::warn!( + posture = AdminMountPosture::DedicatedListener.name(), + admin_bind = %bind, + "ADMIN SECURITY POSTURE: the dedicated admin listener binds a \ + non-loopback address; any peer that reaches it can attempt the \ + admin token. Point server.admin_bind back at loopback unless a \ + tunnel requires otherwise." + ); + } + } + None => {} + } } /// Check if a provider with catalog enforcement enabled has no usable models. @@ -416,4 +517,328 @@ enforce = true "Plain log must still carry the emitted message, got: {captured}" ); } + + // ----------------------------------------------------------------------- + // [storage] wiring into load_toml_state + // ----------------------------------------------------------------------- + + /// Write a minimal valid main config, optionally with a [storage] section. + fn write_main_config( + dir: &std::path::Path, + storage_section: Option<&str>, + ) -> std::path::PathBuf { + let path = dir.join("config.toml"); + let storage = storage_section.unwrap_or_default(); + std::fs::write( + &path, + format!( + r#" +[server] +bind = "127.0.0.1:3456" +request_timeout = "300s" +log_level = "info" +hot_reload = false +server_name = "test" +{storage}"# + ), + ) + .unwrap(); + path + } + + #[tokio::test] + async fn load_toml_state_wires_sqlite_ledger_when_storage_configured() { + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("usage.db"); + let storage = format!( + " +[storage] +path = \"{}\" +", + db_path.display() + ); + let path = write_main_config(dir.path(), Some(&storage)); + + let state = load_toml_state( + &path, + llm_proxy_provider::ProviderAdapterRegistry::builtin(), + llm_proxy_provider::ProxyClient::try_new().unwrap(), + None, + ) + .expect("state with [storage]"); + + let debug = format!("{state:?}"); + assert!( + debug.contains("SqliteEventBus"), + "[storage] must wire the SQLite ledger, got: {debug}" + ); + + // The ledger accepts and persists an event end to end: emit through + // the wired bus, close (drain), then the database file must exist + // with content (the writer checkpoints its WAL on close). + let event = + llm_proxy_storage::ProxyEvent::RequestReceived(llm_proxy_storage::RequestReceived { + request_id: "wiring-1".to_owned(), + timestamp: time::OffsetDateTime::now_utc(), + caller: "laptop".to_owned(), + provider: "prov-a".to_owned(), + route_kind: "messages".to_owned(), + client_protocol: "anthropic".to_owned(), + model: llm_proxy_protocol::core::ModelRef { + requested: "claude-3".to_owned(), + upstream: None, + }, + streaming: false, + body_hash: "digest".to_owned(), + }); + let bus = std::sync::Arc::clone(state.event_bus()); + bus.emit(&event); + let closer = std::sync::Arc::clone(&bus); + tokio::task::spawn_blocking(move || llm_proxy_storage::EventBus::close(&*closer)) + .await + .expect("close"); + + let db = std::fs::read(&db_path).expect("database file exists after close"); + assert!( + !db.is_empty(), + "database file must hold the committed schema and rows" + ); + } + + #[tokio::test] + async fn load_toml_state_keeps_noop_bus_without_storage() { + let dir = tempfile::tempdir().unwrap(); + let path = write_main_config(dir.path(), None); + + let state = load_toml_state( + &path, + llm_proxy_provider::ProviderAdapterRegistry::builtin(), + llm_proxy_provider::ProxyClient::try_new().unwrap(), + None, + ) + .expect("state without [storage]"); + + let debug = format!("{state:?}"); + assert!( + debug.contains("NoopBus"), + "absent [storage] must keep the no-op bus (today's behavior), got: {debug}" + ); + // And close() on the default bus is a no-op that must not hang. + llm_proxy_storage::EventBus::close(state.event_bus().as_ref()); + } + + // ----------------------------------------------------------------------- + // [provider.egress] wiring into load_toml_state + // ----------------------------------------------------------------------- + + /// Write a providers/ dir next to the main config with one provider file, + /// optionally carrying an egress section. + fn write_provider_file(dir: &std::path::Path, egress: Option<&str>) { + let egress_section = egress.unwrap_or_default(); + std::fs::create_dir_all(dir.join("providers")).unwrap(); + std::fs::write( + dir.join("providers/egress-test.toml"), + format!( + r#" +[provider] +name = "egress-test" +api_key = "test-key" +auth_style = "bearer" + +[provider.adapters.chat] +protocol = "openai_chat_completions" +endpoint = "https://example.com/v1/chat/completions" + +[provider.routes] +chat_completions = "chat" +{egress_section}"# + ), + ) + .unwrap(); + } + + #[tokio::test] + async fn load_toml_state_wires_egress_client_only_for_configured_providers() { + let dir = tempfile::tempdir().unwrap(); + // The proxy URL carries userinfo: it must reach the transport as proxy + // authentication but must never be printed with the credentials. + write_provider_file( + dir.path(), + Some( + "[provider.egress]\nproxy = \"socks5h://egress-user:egress-pass@127.0.0.1:1080\"\n", + ), + ); + // A second provider with no egress section must keep the shared + // default client. + std::fs::write( + dir.path().join("providers/direct.toml"), + r#" +[provider] +name = "direct" +api_key = "test-key" +auth_style = "bearer" + +[provider.adapters.chat] +protocol = "openai_chat_completions" +endpoint = "https://example.com/v1/chat/completions" + +[provider.routes] +chat_completions = "chat" +"#, + ) + .unwrap(); + let path = write_main_config(dir.path(), None); + + let state = load_toml_state( + &path, + llm_proxy_provider::ProviderAdapterRegistry::builtin(), + llm_proxy_provider::ProxyClient::try_new().unwrap(), + None, + ) + .expect("state builds with an egress provider"); + + // Observable through the state Debug: the egress map holds exactly the + // provider that configured egress, and the proxy credentials never + // print (the credential-free host:port may). + let debug = format!("{state:?}"); + let egress_map = debug + .split("egress_clients: ") + .nth(1) + .and_then(|rest| rest.split(')').next()) + .unwrap_or_default(); + assert!( + egress_map.contains("egress-test"), + "the egress provider must have a wired client: {debug}" + ); + assert!( + !egress_map.contains("direct"), + "providers without egress must not appear in the egress map: {debug}" + ); + assert!( + !debug.contains("egress-user") && !debug.contains("egress-pass"), + "proxy credentials must never print in the state Debug: {debug}" + ); + } + + #[tokio::test] + async fn load_toml_state_rejects_egress_url_the_transport_cannot_use() { + let dir = tempfile::tempdir().unwrap(); + // Malformed IPv6 host: passes config validation (scheme http, host + // non-empty) but the transport cannot build a proxy from it, so the + // serve path must refuse to start rather than fall back to direct. + write_provider_file( + dir.path(), + Some("[provider.egress]\nproxy = \"http://[::1:8080\"\n"), + ); + let path = write_main_config(dir.path(), None); + + let err = load_toml_state( + &path, + llm_proxy_provider::ProviderAdapterRegistry::builtin(), + llm_proxy_provider::ProxyClient::try_new().unwrap(), + None, + ) + .expect_err("unusable egress proxy must be a startup error"); + + let message = format!("{err:#}"); + assert!( + message.contains("egress"), + "the error must point at the egress wiring: {message}" + ); + assert!( + message.contains("egress-test"), + "the error must name the provider: {message}" + ); + } + + // ----------------------------------------------------------------------- + // Admin posture startup warnings + // ----------------------------------------------------------------------- + + fn admin_app_config( + on_main: bool, + admin_bind: Option, + ) -> llm_proxy_core::AppConfig { + use llm_proxy_core::{AdminAuthConfig, ServerConfig, StorageConfig}; + llm_proxy_core::AppConfig { + server: ServerConfig { + bind: "127.0.0.1:3456".parse().unwrap(), + request_timeout: std::time::Duration::from_secs(300), + shutdown_timeout: std::time::Duration::from_secs(30), + log_level: "info".to_owned(), + hot_reload: false, + allowed_origins: None, + server_name: "posture-test".to_owned(), + rate_limit_rpm: 100, + trust_forwarded_headers: false, + dedup_window: std::time::Duration::ZERO, + log_format: Default::default(), + admin_bind, + admin_on_main_listener: on_main, + }, + auth: llm_proxy_core::AuthConfig { + mode: llm_proxy_core::AuthMode::None, + clients: std::collections::HashMap::new(), + admin: Some(AdminAuthConfig { + token: secrecy::SecretString::from("admin-secret"), + }), + }, + storage: Some(StorageConfig { + path: "usage.db".into(), + max_age: None, + max_rows: None, + prune_interval: std::time::Duration::from_secs(3600), + }), + } + } + + /// Capture the warnings [`warn_admin_posture`] emits for `config`. + fn captured_posture_warnings(config: &llm_proxy_core::AppConfig) -> Vec { + let buf = std::sync::Arc::new(std::sync::Mutex::new(Vec::::new())); + let layer = build_log_layer(&LogFormat::Json, BufferWriter(std::sync::Arc::clone(&buf))); + use tracing_subscriber::layer::SubscriberExt; + let subscriber = tracing_subscriber::registry().with(layer); + { + let _guard = tracing::subscriber::set_default(subscriber); + warn_admin_posture(config); + } + buf.lock().expect("buffer lock not poisoned").clone() + } + + #[test] + fn main_listener_optin_is_warned_about_loudly_by_name() { + let config = admin_app_config(true, None); + let captured = String::from_utf8(captured_posture_warnings(&config)) + .expect("captured bytes are UTF-8"); + assert!( + captured.contains("admin-on-main-listener"), + "the posture must be named for grepping: {captured}" + ); + assert!( + captured.contains("ADMIN SECURITY POSTURE"), + "the warning must be loud and identifiable: {captured}" + ); + } + + #[test] + fn dedicated_loopback_posture_stays_quiet() { + let config = admin_app_config(false, None); + let captured = String::from_utf8(captured_posture_warnings(&config)) + .expect("captured bytes are UTF-8"); + assert!( + !captured.contains("ADMIN SECURITY POSTURE"), + "the default loopback posture needs no warning: {captured}" + ); + } + + #[test] + fn dedicated_non_loopback_bind_is_warned_about() { + let config = admin_app_config(false, Some("10.0.0.5:9900".parse().unwrap())); + let captured = String::from_utf8(captured_posture_warnings(&config)) + .expect("captured bytes are UTF-8"); + assert!( + captured.contains("non-loopback"), + "a non-loopback admin listener must be called out: {captured}" + ); + } } diff --git a/apps/llm-proxy/tests/cli_validate.rs b/apps/llm-proxy/tests/cli_validate.rs index f2cbf48..abceeca 100644 --- a/apps/llm-proxy/tests/cli_validate.rs +++ b/apps/llm-proxy/tests/cli_validate.rs @@ -49,4 +49,191 @@ mod tests { "error should mention unsupported extension: {msg}" ); } + + #[test] + fn validate_rejects_non_loopback_bind_without_client_keys() { + // The client-identity safety rule, exercised through the command + // layer: an anonymous proxy exposed past loopback must refuse to + // start, so validation fails before any listener could bind. + let dir = tempfile::tempdir().unwrap(); + let config = dir.path().join("config.toml"); + std::fs::write( + &config, + r#" +[server] +bind = "0.0.0.0:3456" +request_timeout = "300s" +log_level = "info" +hot_reload = false +server_name = "test" +"#, + ) + .unwrap(); + + let result = cmd_validate(Some(&config)); + // anyhow's `to_string()` shows only the outer context; `{:#}` walks + // the source chain down to the CoreError validation message. + let msg = format!( + "{:#}", + result.expect_err("non-loopback anonymous bind must be rejected") + ); + assert!( + msg.contains("not loopback"), + "error should explain the loopback restriction: {msg}" + ); + } + + #[test] + fn validate_accepts_client_keys_on_non_loopback_bind() { + // With at least one hashed client key configured, binding off + // loopback is a legitimate deployment and must validate. + let dir = tempfile::tempdir().unwrap(); + let config = dir.path().join("config.toml"); + std::fs::write( + &config, + r#" +[server] +bind = "0.0.0.0:3456" +request_timeout = "300s" +log_level = "info" +hot_reload = false +server_name = "test" + +[auth] +mode = "client-keys" + +[auth.clients] +laptop = "$argon2id$v=19$m=8,t=1,p=1$MDEyMzQ1Njc4OWFiY2RlZg$eEnjdfBL8hfvU2QuGOpow8Yq2Jyrr8GXKYgtAAsscSs" +"#, + ) + .unwrap(); + + let result = cmd_validate(Some(&config)); + assert!( + result.is_ok(), + "client-keys on a non-loopback bind should validate: {result:?}" + ); + } + + #[test] + fn validate_accepts_admin_api_with_storage() { + let dir = tempfile::tempdir().unwrap(); + let config = dir.path().join("config.toml"); + std::fs::write( + &config, + r#" +[server] +bind = "127.0.0.1:3456" +request_timeout = "300s" +log_level = "info" +hot_reload = false +server_name = "test" + +[auth.admin] +token = "admin-secret" + +[storage] +path = "usage.db" +"#, + ) + .unwrap(); + + let result = cmd_validate(Some(&config)); + assert!( + result.is_ok(), + "admin + storage should validate: {result:?}" + ); + } + + #[test] + fn validate_rejects_admin_token_without_storage() { + let dir = tempfile::tempdir().unwrap(); + let config = dir.path().join("config.toml"); + std::fs::write( + &config, + r#" +[server] +bind = "127.0.0.1:3456" +request_timeout = "300s" +log_level = "info" +hot_reload = false +server_name = "test" + +[auth.admin] +token = "admin-secret" +"#, + ) + .unwrap(); + + let result = cmd_validate(Some(&config)); + let msg = format!( + "{:#}", + result.expect_err("admin without storage must be rejected") + ); + assert!( + msg.contains("[auth.admin] requires"), + "error should point at the missing ledger: {msg}" + ); + } + + #[test] + fn validate_rejects_admin_on_main_listener_without_admin_token() { + let dir = tempfile::tempdir().unwrap(); + let config = dir.path().join("config.toml"); + std::fs::write( + &config, + r#" +[server] +bind = "127.0.0.1:3456" +request_timeout = "300s" +log_level = "info" +hot_reload = false +server_name = "test" +admin_on_main_listener = true + +[storage] +path = "usage.db" +"#, + ) + .unwrap(); + + let result = cmd_validate(Some(&config)); + let msg = format!("{:#}", result.expect_err("dead opt-in must be rejected")); + assert!( + msg.contains("admin_on_main_listener"), + "error should name the dead flag: {msg}" + ); + } + + #[test] + fn validate_rejects_admin_bind_equal_to_main_bind() { + let dir = tempfile::tempdir().unwrap(); + let config = dir.path().join("config.toml"); + std::fs::write( + &config, + r#" +[server] +bind = "127.0.0.1:3456" +request_timeout = "300s" +log_level = "info" +hot_reload = false +server_name = "test" +admin_bind = "127.0.0.1:3456" + +[auth.admin] +token = "admin-secret" + +[storage] +path = "usage.db" +"#, + ) + .unwrap(); + + let result = cmd_validate(Some(&config)); + let msg = format!("{:#}", result.expect_err("same-socket bind rejected")); + assert!( + msg.contains("must differ from server.bind"), + "error should explain the socket conflict: {msg}" + ); + } } diff --git a/config.toml.example b/config.toml.example index 858ffcc..e79526f 100644 --- a/config.toml.example +++ b/config.toml.example @@ -20,3 +20,56 @@ trust_forwarded_headers = false # Window for rejecting identical path/body requests. "0ms" disables dedup. dedup_window = "0ms" server_name = "llm-proxy" +# Dedicated listener for the read-only admin /usage API. Only read when +# [auth.admin] is configured. Defaults to "127.0.0.1" on the port after +# `bind` (loopback-only); set it explicitly to point at a tunnel interface +# or a different port. Must differ from `bind`. +# admin_bind = "127.0.0.1:3457" +# Also serve /usage on the main listener behind the same admin token. +# Off by default; enabling it names a distinct security posture +# ("admin-on-main-listener") that startup warns about loudly: the admin +# surface then reaches every network the main listener does. With no +# explicit admin_bind set, this posture serves /usage on the main listener +# only (no dedicated socket) — the option for setups where a second port +# is not available. +# admin_on_main_listener = false + +# Client authentication for proxy routes. +# +# mode = "none" (default): no client authentication. Only permitted while +# [server].bind is a loopback address; the server refuses to start on a +# non-loopback bind without client keys. +[auth] +mode = "none" +# mode = "client-keys": every proxy request must present +# "Authorization: Bearer ". Named keys map key ids to argon2id +# password hashes in PHC format -- never plaintext secrets. Generate one with +# e.g. `printf 'my-secret' | argon2 'some-salt' -id -e`. +# The key id (not the secret) appears in usage events and rate-limit keys. +# +# [auth] +# mode = "client-keys" +# +# [auth.clients] +# laptop = "$argon2id$v=19$m=19456,t=2,p=1$c29tZXNhbHQ$..." + +# Usage ledger persistence. Absent (the default) disables persistence: +# usage events go to an in-memory no-op sink and nothing is written to disk. +# +# [storage] +# SQLite database file. Parent directories are created at startup. +# path = "${HOME}/.local/state/llm-proxy/usage.db" +# Delete ledger rows older than this age; absent keeps rows forever. +# max_age = "30d" +# Keep at most this many rows; the oldest rows are deleted when the cap is +# exceeded. Absent means unlimited. +# max_rows = 1000000 +# How often retention pruning runs. Default "1h"; must be at least "1s". +# prune_interval = "1h" + +# Admin API for the usage ledger. Absent (the default) disables the /usage +# surface entirely. Requires [storage]; every /usage route on every mount +# demands "Authorization: Bearer " (compared in constant time). +# +# [auth.admin] +# token = "${ADMIN_TOKEN}" diff --git a/crates/llm-proxy-api/Cargo.toml b/crates/llm-proxy-api/Cargo.toml index 906abd0..fbbec00 100644 --- a/crates/llm-proxy-api/Cargo.toml +++ b/crates/llm-proxy-api/Cargo.toml @@ -4,7 +4,7 @@ version = "0.1.0" edition.workspace = true rust-version.workspace = true license.workspace = true -description = "Reserved public HTTP API surface for llm-proxy; a placeholder until the handler layer lands." +description = "Read-only authenticated /usage query surface for the llm-proxy usage ledger." repository.workspace = true readme.workspace = true keywords.workspace = true @@ -12,3 +12,24 @@ categories.workspace = true [lints] workspace = true + +[dependencies] +axum = { workspace = true } +serde = { workspace = true } +thiserror = { workspace = true } +time = { workspace = true } +tokio = { workspace = true } +tracing = { workspace = true } +rusqlite = { workspace = true } +rust_decimal = { workspace = true } +sha2 = "0.10" +humantime = "2.3" +tower-http = { version = "0.6", features = ["trace", "catch-panic", "set-header"] } + +[dev-dependencies] +serde_json = { workspace = true } +tempfile = "3" +tower = { workspace = true, features = ["util"] } +tracing-subscriber = { workspace = true } +llm-proxy-storage = { path = "../llm-proxy-storage" } +llm-proxy-protocol = { path = "../llm-proxy-protocol" } diff --git a/crates/llm-proxy-api/src/auth.rs b/crates/llm-proxy-api/src/auth.rs new file mode 100644 index 0000000..35cb16c --- /dev/null +++ b/crates/llm-proxy-api/src/auth.rs @@ -0,0 +1,176 @@ +//! Admin bearer-token authentication for the `/usage` surface. + +use axum::extract::State; +use axum::http::{StatusCode, header}; +use axum::response::{IntoResponse, Response}; +use axum::{Json, middleware::Next}; +use serde::Serialize; +use sha2::{Digest, Sha256}; + +use crate::routes::UsageState; + +/// Guard for the admin token: stores only its SHA-256 digest. +/// +/// The configured token is hashed once at construction and never retained, +/// so neither the plaintext nor anything it could be compared against +/// directly survives in memory. Presented tokens are hashed the same way and +/// the two fixed-length digests are compared in constant time, which also +/// removes any length-dependence from the comparison. +#[derive(Clone)] +pub(crate) struct AdminAuth { + expected_digest: [u8; 32], +} + +impl std::fmt::Debug for AdminAuth { + // The digest is offline-brute-forceable for low-entropy tokens (SHA-256 + // preimage search), so it is treated as sensitive and never printed. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AdminAuth") + .field("token", &"[REDACTED]") + .finish() + } +} + +impl AdminAuth { + /// Build the guard from the configured admin token. + pub(crate) fn new(token: &str) -> Self { + Self { + expected_digest: Sha256::digest(token.as_bytes()).into(), + } + } + + /// Whether `presented` (a bearer token already extracted from the + /// Authorization header) is the configured admin token. + pub(crate) fn accepts(&self, presented: Option<&str>) -> bool { + let Some(token) = presented else { + return false; + }; + let digest: [u8; 32] = Sha256::digest(token.as_bytes()).into(); + constant_time_eq(&digest, &self.expected_digest) + } +} + +/// Constant-time equality of two equal-length byte strings. +/// +/// Compares all bytes unconditionally so the runtime does not depend on +/// where the first difference appears. Both inputs here are fixed 32-byte +/// digests, so the length branch is not secret-dependent. +fn constant_time_eq(a: &[u8], b: &[u8]) -> bool { + if a.len() != b.len() { + return false; + } + let mut diff = 0u8; + for (x, y) in a.iter().zip(b.iter()) { + diff |= x ^ y; + } + diff == 0 +} + +/// Extract the bearer token from `Authorization: Bearer `. +/// +/// The scheme is matched case-insensitively (RFC 7235). Returns `None` when +/// the header is absent, uses another scheme, or carries an empty token. +fn bearer_token(headers: &axum::http::HeaderMap) -> Option<&str> { + let value = headers.get(header::AUTHORIZATION)?; + let header_value = value.to_str().ok()?; + let mut parts = header_value.splitn(2, ' '); + match (parts.next(), parts.next()) { + (Some(scheme), Some(token)) if scheme.eq_ignore_ascii_case("bearer") => { + let token = token.trim(); + (!token.is_empty()).then_some(token) + } + _ => None, + } +} + +/// Body of every admin-auth rejection: static text that never echoes any +/// part of the presented credentials. +#[derive(Serialize)] +struct UnauthorizedBody { + error: UnauthorizedError, +} + +#[derive(Serialize)] +struct UnauthorizedError { + r#type: &'static str, + message: &'static str, +} + +/// Middleware requiring the admin token on every `/usage` route. +/// +/// Every rejection produces the same uniform 401 body, so the response does +/// not reveal whether a token was missing, malformed, or merely wrong. The +/// token is never logged and never appears in any response. +pub(crate) async fn require_admin( + State(state): State, + req: axum::extract::Request, + next: Next, +) -> Response { + if state.admin.accepts(bearer_token(req.headers())) { + return next.run(req).await; + } + ( + StatusCode::UNAUTHORIZED, + Json(UnauthorizedBody { + error: UnauthorizedError { + r#type: "unauthorized", + message: "missing or invalid admin token", + }, + }), + ) + .into_response() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn accepts_only_the_configured_token() { + let auth = AdminAuth::new("admin-secret"); + assert!(auth.accepts(Some("admin-secret"))); + assert!(!auth.accepts(Some("admin-secret2"))); + assert!(!auth.accepts(Some(""))); + assert!(!auth.accepts(None)); + } + + #[test] + fn debug_never_prints_the_digest() { + let auth = AdminAuth::new("admin-secret"); + let debug = format!("{auth:?}"); + assert!( + !debug.contains("admin-secret"), + "Debug must not expose the token, got: {debug}" + ); + } + + #[test] + fn bearer_token_extraction() { + let mut headers = axum::http::HeaderMap::new(); + assert_eq!(bearer_token(&headers), None); + headers.insert( + header::AUTHORIZATION, + "Bearer tok".parse().expect("valid header value"), + ); + assert_eq!(bearer_token(&headers), Some("tok")); + headers.insert( + header::AUTHORIZATION, + "bearer tok".parse().expect("valid header value"), + ); + assert_eq!( + bearer_token(&headers), + Some("tok"), + "scheme matching is case-insensitive" + ); + headers.insert( + header::AUTHORIZATION, + "Basic dXNlcjpwYXNz".parse().expect("valid header value"), + ); + assert_eq!(bearer_token(&headers), None); + headers.insert( + header::AUTHORIZATION, + "Bearer ".parse().expect("valid header value"), + ); + assert_eq!(bearer_token(&headers), None, "empty token is no token"); + } +} diff --git a/crates/llm-proxy-api/src/lib.rs b/crates/llm-proxy-api/src/lib.rs index 67de5c4..a7d7684 100644 --- a/crates/llm-proxy-api/src/lib.rs +++ b/crates/llm-proxy-api/src/lib.rs @@ -1,30 +1,25 @@ -//! llm-proxy-api: public API surface for the LLM proxy. +//! llm-proxy-api: the read-only, admin-authenticated `/usage` query surface +//! over the proxy's usage ledger. //! -//! This crate is intentionally empty in v1. It reserves the namespace for a -//! future public HTTP API layer. See `docs/protocol.md` for the implemented -//! protocol boundaries. +//! [`UsageApi`] builds a small axum router with three `GET` routes backed by +//! the SQLite ledger written by `llm-proxy-storage`'s `SqliteEventBus`: +//! +//! - `/usage/summary?window=1h&group_by=provider|model|caller` — token, +//! request, error, and exact-cost totals grouped over a time window. +//! - `/usage/requests?since=…&until=…&provider=…&limit=…&cursor=…` — +//! bounded, cursor-paged raw ledger rows. Never request bodies; the ledger +//! stores only a body digest and so do these responses. +//! - `/usage/keys` — caller list with request counts and last-seen times. +//! +//! Every route requires `Authorization: Bearer `; the token is +//! compared in constant time against a digest and never appears in logs or +//! responses. All responses are strictly read-only: queries run on +//! read-only SQLite connections and there is no write path. -/// Placeholder module to reserve the public API namespace. -/// -/// Remove this module and replace with real types once implementation -/// begins. Tracking: . -pub mod placeholder { - /// Sentinel type indicating this crate has not yet been implemented. - /// - /// This type exists so that downstream crates can reference - /// `llm_proxy_api::placeholder::NotImplemented` in feature-gated - /// code without pulling in a fully implemented crate. - #[derive(Debug, Clone, Copy, Default)] - pub struct NotImplemented; -} +#![deny(missing_docs)] -// Ensure the crate compiles as a valid library target. -#[cfg(test)] -mod tests { - use super::placeholder::NotImplemented; +mod auth; +mod reader; +mod routes; - #[test] - fn placeholder_type_exists() { - let _ = NotImplemented; - } -} +pub use routes::UsageApi; diff --git a/crates/llm-proxy-api/src/reader.rs b/crates/llm-proxy-api/src/reader.rs new file mode 100644 index 0000000..a793a60 --- /dev/null +++ b/crates/llm-proxy-api/src/reader.rs @@ -0,0 +1,615 @@ +//! Read-only queries over the SQLite usage ledger written by +//! `llm-proxy-storage`'s `SqliteEventBus`. +//! +//! Every query opens its own connection with +//! [`rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY`]: SQLite refuses any write +//! on such a connection, so the read-only guarantee comes from the database +//! handle itself, not from discipline in the query text. The ledger runs in +//! WAL mode (set by the writer), where readers do not block the writer and +//! the writer does not block readers. +//! +//! Costs are stored as exact decimal text and are summed with +//! [`rust_decimal::Decimal`] in Rust — never through SQLite's `SUM`, which +//! would coerce the text to a float. + +use std::collections::HashMap; +use std::path::{Path, PathBuf}; +use std::time::Duration; + +use rusqlite::Connection; +use rusqlite::OpenFlags; +use rusqlite::types::Value as SqlValue; +use serde::Serialize; +use time::OffsetDateTime; + +/// SQLite busy timeout for the read connections (matches the writer's). +const SQLITE_BUSY_TIMEOUT: Duration = Duration::from_secs(5); + +/// Upper bound accepted for the `limit` query parameter on +/// `/usage/requests`. Requests above it are rejected rather than clamped so +/// the caller learns the bound instead of silently receiving fewer rows +/// than asked. +pub(crate) const MAX_PAGE_LIMIT: i64 = 1000; + +/// Errors raised while querying the usage ledger. +#[derive(Debug, thiserror::Error)] +pub(crate) enum UsageQueryError { + /// The SQLite query or connection failed. + #[error("usage ledger query failed: {0}")] + Sqlite(#[from] rusqlite::Error), + /// A stored cost value is not parseable as an exact decimal. + #[error("usage ledger cost value is not a valid decimal: {0}")] + Cost(String), + /// The blocking query task failed to join (panicked or was cancelled). + #[error("usage ledger query task failed: {0}")] + Task(String), +} + +/// Field to group a `/usage/summary` request by. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum GroupBy { + /// The upstream provider name (`provider` column). + Provider, + /// The client-requested model (`model_requested` column). Upstream model + /// attribution is sparse (streaming completions may not carry it), so + /// the requested model is the grouping that is always present. + Model, + /// The caller identity (`caller` column: client key id or `anonymous`). + Caller, +} + +impl GroupBy { + /// Ledger column selected for this grouping. The value is fixed per + /// variant, never built from request input. + const fn column(self) -> &'static str { + match self { + Self::Provider => "provider", + Self::Model => "model_requested", + Self::Caller => "caller", + } + } + + /// Parse the `group_by` query parameter; `None` means unrecognized. + pub(crate) fn parse(value: &str) -> Option { + match value { + "provider" => Some(Self::Provider), + "model" => Some(Self::Model), + "caller" => Some(Self::Caller), + _ => None, + } + } + + /// Canonical name echoed back in responses. + pub(crate) const fn as_str(self) -> &'static str { + match self { + Self::Provider => "provider", + Self::Model => "model", + Self::Caller => "caller", + } + } +} + +/// Token counts for one summary group or raw row. Serialized with snake_case +/// keys; absent columns read as zero / `null` depending on the container. +#[derive(Debug, Clone, Default, Serialize)] +pub(crate) struct TokenCounts { + /// Input (prompt) tokens. + pub input: i64, + /// Output (completion) tokens. + pub output: i64, + /// Reasoning tokens (provider-reported thinking, when present). + pub reasoning: i64, + /// Cache-read input tokens. + pub cache_read: i64, + /// Cache-creation input tokens. + pub cache_creation: i64, +} + +/// Exact cost totals, serialized as decimal strings (never floats). +#[derive(Debug, Clone, Serialize)] +pub(crate) struct CostTotals { + /// Input-token cost. + pub input: String, + /// Output-token cost. + pub output: String, + /// Cache-read cost. + pub cache_read: String, + /// Cache-creation cost. + pub cache_creation: String, + /// Reasoning-token cost. + pub reasoning: String, + /// Total cost. + pub total: String, +} + +impl CostTotals { + /// All-zero totals for groups with no costed rows. + fn zeroed() -> Self { + Self { + input: "0".to_owned(), + output: "0".to_owned(), + cache_read: "0".to_owned(), + cache_creation: "0".to_owned(), + reasoning: "0".to_owned(), + total: "0".to_owned(), + } + } +} + +/// One aggregated group of a `/usage/summary` response. +#[derive(Debug, Serialize)] +pub(crate) struct SummaryGroup { + /// Grouping key (`null` when the underlying column is NULL, e.g. the + /// provider of a pre-routing failure). + pub key: Option, + /// Requests in the window. + pub requests: i64, + /// Failed requests in the window (`error_kind` present). + pub errors: i64, + /// Token totals. + pub tokens: TokenCounts, + /// Exact cost totals. + pub cost: CostTotals, +} + +/// One raw ledger row of a `/usage/requests` response. Columns mirror the +/// ledger schema; there is no request-body field anywhere in this type (the +/// ledger stores only the body digest, `body_hash`). +#[derive(Debug, Serialize)] +pub(crate) struct RequestRow { + /// Ledger request id (primary key). + pub request_id: String, + /// Unix epoch milliseconds the request was received (`null` for + /// failure-only rows). + pub received_at_ms: Option, + /// Unix epoch milliseconds the request completed or failed (`null` + /// while in flight). + pub completed_at_ms: Option, + /// Caller identity (client key id or `anonymous`). + pub caller: Option, + /// Upstream provider name. + pub provider: Option, + /// Route kind (`messages`, `chat`, ...). + pub route_kind: Option, + /// Client protocol (`anthropic`, `openai`, ...). + pub client_protocol: Option, + /// Whether the request was streaming. + pub streaming: Option, + /// Client-requested model. + pub model_requested: Option, + /// Upstream-resolved model, when routing resolved an alias. + pub model_upstream: Option, + /// Token usage, when the request completed. + pub tokens: Option, + /// Usage provenance label (`provider_reported`, ...), when present. + pub usage_provenance: Option, + /// Exact costs, when the request completed with cost data. + pub cost: Option, + /// Stop reason label, when the request completed. + pub stop_reason: Option, + /// Error kind, when the request failed. + pub error_kind: Option, + /// HTTP status of the failure, when applicable. + pub http_status: Option, + /// End-to-end latency in milliseconds, when the terminal event arrived. + pub latency_ms: Option, + /// Request-body digest. Never the body itself. + pub body_hash: Option, +} + +/// One page of `/usage/requests` results. +#[derive(Debug, Serialize)] +pub(crate) struct RequestsPage { + /// Rows, newest first. + pub rows: Vec, + /// Opaque cursor to pass as `?cursor=` for the next page; `None` when + /// this page is the last. + pub next_cursor: Option, +} + +/// One `/usage/keys` entry: a caller and its last-seen time. +#[derive(Debug, Serialize)] +pub(crate) struct CallerSummary { + /// Caller identity (client key id or `anonymous`); `null` only for + /// failure-only rows that carry no caller. + pub caller: Option, + /// Total ledgered requests for this caller. + pub requests: i64, + /// Failed requests for this caller. + pub errors: i64, + /// Unix epoch milliseconds of the caller's most recent ledgered request. + pub last_seen_ms: Option, +} + +/// Filters for `/usage/requests`. +#[derive(Debug, Default)] +pub(crate) struct RequestFilter { + /// Inclusive lower bound on the ledger time axis. + pub since_ms: Option, + /// Exclusive upper bound on the ledger time axis. + pub until_ms: Option, + /// Exact provider-name match. + pub provider: Option, + /// Page size (1..=[`MAX_PAGE_LIMIT`]). + pub limit: i64, + /// Keyset cursor: `(timestamp_ms, request_id)` of the last row of the + /// previous page. + pub cursor: Option<(i64, String)>, +} + +/// Read-only accessor for the usage ledger at `path`. +#[derive(Debug, Clone)] +pub(crate) struct UsageReader { + path: PathBuf, +} + +/// Unix epoch milliseconds for an [`OffsetDateTime`], saturating instead of +/// panicking on absurd values. +fn offset_to_unix_millis(ts: OffsetDateTime) -> i64 { + ts.unix_timestamp() + .saturating_mul(1000) + .saturating_add(i64::from(ts.millisecond())) +} + +/// Current Unix epoch milliseconds. +pub(crate) fn now_unix_millis() -> i64 { + offset_to_unix_millis(OffsetDateTime::now_utc()) +} + +impl UsageReader { + /// Reader for the ledger at `path`. The file is not touched here; each + /// query opens its own read-only connection, so construction cannot + /// fail and a missing/unusable file surfaces per query. + pub(crate) fn new(path: impl Into) -> Self { + Self { path: path.into() } + } + + /// Aggregated usage per group for rows at or after `since_ms`. + pub(crate) async fn summary( + &self, + since_ms: i64, + group: GroupBy, + ) -> Result, UsageQueryError> { + let path = self.path.clone(); + run_blocking(move || summary_sync(&path, since_ms, group)).await + } + + /// One bounded page of raw rows matching `filter`. + pub(crate) async fn requests( + &self, + filter: RequestFilter, + ) -> Result { + let path = self.path.clone(); + run_blocking(move || requests_sync(&path, &filter)).await + } + + /// Per-caller totals with last-seen times. + pub(crate) async fn callers(&self) -> Result, UsageQueryError> { + let path = self.path.clone(); + run_blocking(move || callers_sync(&path)).await + } + + /// Open one read-only connection with the busy timeout applied. + fn open(path: &Path) -> Result { + let conn = Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_ONLY)?; + conn.busy_timeout(SQLITE_BUSY_TIMEOUT)?; + Ok(conn) + } +} + +/// Run `f` on the blocking pool (SQLite calls are synchronous); a join +/// failure is reported rather than unwrapped. +async fn run_blocking(f: F) -> Result +where + F: FnOnce() -> Result + Send + 'static, + T: Send + 'static, +{ + tokio::task::spawn_blocking(f) + .await + .map_err(|error| UsageQueryError::Task(error.to_string()))? +} + +/// Parse a stored cost cell; `None` counts as zero, malformed text is an +/// error rather than a silently skipped amount. +fn parse_cost(value: Option) -> Result { + match value { + None => Ok(rust_decimal::Decimal::ZERO), + Some(text) => rust_decimal::Decimal::from_str_exact(&text) + .map_err(|_parse_error| UsageQueryError::Cost(text)), + } +} + +fn summary_sync( + path: &Path, + since_ms: i64, + group: GroupBy, +) -> Result, UsageQueryError> { + let conn = UsageReader::open(path)?; + let key = group.column(); + let sql = format!( + "SELECT {key} AS k, COUNT(*), \ + SUM(CASE WHEN error_kind IS NOT NULL THEN 1 ELSE 0 END), \ + COALESCE(SUM(COALESCE(input_tokens, 0)), 0), \ + COALESCE(SUM(COALESCE(output_tokens, 0)), 0), \ + COALESCE(SUM(COALESCE(reasoning_tokens, 0)), 0), \ + COALESCE(SUM(COALESCE(cache_read_tokens, 0)), 0), \ + COALESCE(SUM(COALESCE(cache_creation_tokens, 0)), 0) \ + FROM requests \ + WHERE COALESCE(received_at, completed_at) >= ?1 \ + GROUP BY {key} ORDER BY 2 DESC" + ); + let mut groups: Vec = Vec::new(); + { + let mut stmt = conn.prepare(&sql)?; + let rows = stmt.query_map([since_ms], |row| { + Ok(SummaryGroup { + key: row.get(0)?, + requests: row.get(1)?, + errors: row.get::<_, Option>(2)?.unwrap_or(0), + tokens: TokenCounts { + input: row.get(3)?, + output: row.get(4)?, + reasoning: row.get(5)?, + cache_read: row.get(6)?, + cache_creation: row.get(7)?, + }, + cost: CostTotals::zeroed(), + }) + })?; + for row in rows { + groups.push(row?); + } + } + let costs = sum_costs_by_group(&conn, key, since_ms)?; + for group_row in &mut groups { + if let Some(total) = costs.get(&group_row.key) { + group_row.cost = cost_totals_from(*total); + } + } + Ok(groups) +} + +/// Exact per-group cost sums. The aggregation runs in Rust over decimal +/// text because SQLite's `SUM` coerces text to float; [Decimal] keeps every +/// digit. Cost columns are only present together (written from one `Cost` +/// value), so `cost_total IS NOT NULL` selects exactly the costed rows. +fn sum_costs_by_group( + conn: &Connection, + key: &str, + since_ms: i64, +) -> Result, [rust_decimal::Decimal; 6]>, UsageQueryError> { + let sql = format!( + "SELECT {key} AS k, cost_input, cost_output, cost_cache_read, \ + cost_cache_creation, cost_reasoning, cost_total \ + FROM requests \ + WHERE COALESCE(received_at, completed_at) >= ?1 AND cost_total IS NOT NULL" + ); + let mut sums: HashMap, [rust_decimal::Decimal; 6]> = HashMap::new(); + let mut stmt = conn.prepare(&sql)?; + let rows = stmt.query_map([since_ms], |row| { + Ok(( + row.get::<_, Option>(0)?, + row.get::<_, Option>(1)?, + row.get::<_, Option>(2)?, + row.get::<_, Option>(3)?, + row.get::<_, Option>(4)?, + row.get::<_, Option>(5)?, + row.get::<_, Option>(6)?, + )) + })?; + for row in rows { + let (key, input, output, cache_read, cache_creation, reasoning, total) = row?; + let entry = sums.entry(key).or_insert([rust_decimal::Decimal::ZERO; 6]); + let parts = [ + parse_cost(input)?, + parse_cost(output)?, + parse_cost(cache_read)?, + parse_cost(cache_creation)?, + parse_cost(reasoning)?, + parse_cost(total)?, + ]; + for (slot, part) in entry.iter_mut().zip(parts) { + *slot += part; + } + } + Ok(sums) +} + +/// Render summed decimals as the string-shaped response type. +/// The array order matches [`sum_costs_by_group`]'s SELECT order. +fn cost_totals_from(sums: [rust_decimal::Decimal; 6]) -> CostTotals { + let show = |d: rust_decimal::Decimal| d.normalize().to_string(); + CostTotals { + input: show(sums[0]), + output: show(sums[1]), + cache_read: show(sums[2]), + cache_creation: show(sums[3]), + reasoning: show(sums[4]), + total: show(sums[5]), + } +} + +/// Map one `requests` row (plus its ledger timestamp, selected as the final +/// column) into the response shape. +fn map_request_row(row: &rusqlite::Row<'_>) -> rusqlite::Result<(RequestRow, i64)> { + let input: Option = row.get(10)?; + let output: Option = row.get(11)?; + let reasoning: Option = row.get(12)?; + let cache_read: Option = row.get(13)?; + let cache_creation: Option = row.get(14)?; + let cost_input: Option = row.get(16)?; + let cost_output: Option = row.get(17)?; + let cost_cache_creation: Option = row.get(18)?; + let cost_cache_read: Option = row.get(19)?; + let cost_reasoning: Option = row.get(20)?; + let cost_total: Option = row.get(21)?; + let streaming: Option = row.get(7)?; + let tokens = input.map(|input| TokenCounts { + input, + output: output.unwrap_or(0), + reasoning: reasoning.unwrap_or(0), + cache_read: cache_read.unwrap_or(0), + cache_creation: cache_creation.unwrap_or(0), + }); + let cost = cost_total.map(|total| CostTotals { + input: cost_input.unwrap_or_else(|| "0".to_owned()), + output: cost_output.unwrap_or_else(|| "0".to_owned()), + cache_read: cost_cache_read.unwrap_or_else(|| "0".to_owned()), + cache_creation: cost_cache_creation.unwrap_or_else(|| "0".to_owned()), + reasoning: cost_reasoning.unwrap_or_else(|| "0".to_owned()), + total, + }); + let row_out = RequestRow { + request_id: row.get(0)?, + received_at_ms: row.get(1)?, + completed_at_ms: row.get(2)?, + caller: row.get(3)?, + provider: row.get(4)?, + route_kind: row.get(5)?, + client_protocol: row.get(6)?, + streaming: streaming.map(|flag| flag != 0), + model_requested: row.get(8)?, + model_upstream: row.get(9)?, + tokens, + usage_provenance: row.get(15)?, + cost, + stop_reason: row.get(22)?, + error_kind: row.get(23)?, + http_status: row.get(24)?, + latency_ms: row.get(25)?, + body_hash: row.get(26)?, + }; + let ts: Option = row.get(27)?; + // The ledger time axis is COALESCE(received, completed); a row where + // both are NULL cannot be ordered, so it gets the smallest possible key. + Ok((row_out, ts.unwrap_or(i64::MIN))) +} + +fn requests_sync(path: &Path, filter: &RequestFilter) -> Result { + let conn = UsageReader::open(path)?; + let mut clauses: Vec<&'static str> = Vec::new(); + let mut params: Vec = Vec::new(); + if let Some(since) = filter.since_ms { + clauses.push("COALESCE(received_at, completed_at) >= ?"); + params.push(SqlValue::Integer(since)); + } + if let Some(until) = filter.until_ms { + clauses.push("COALESCE(received_at, completed_at) < ?"); + params.push(SqlValue::Integer(until)); + } + if let Some(provider) = &filter.provider { + clauses.push("provider = ?"); + params.push(SqlValue::Text(provider.clone())); + } + if let Some((cursor_ts, cursor_id)) = &filter.cursor { + clauses.push( + "(COALESCE(received_at, completed_at) < ? \ + OR (COALESCE(received_at, completed_at) = ? AND request_id < ?))", + ); + params.push(SqlValue::Integer(*cursor_ts)); + params.push(SqlValue::Integer(*cursor_ts)); + params.push(SqlValue::Text(cursor_id.clone())); + } + let where_sql = if clauses.is_empty() { + String::new() + } else { + format!("WHERE {}", clauses.join(" AND ")) + }; + let sql = format!( + "SELECT request_id, received_at, completed_at, caller, provider, route_kind, \ + client_protocol, streaming, model_requested, model_upstream, input_tokens, \ + output_tokens, reasoning_tokens, cache_read_tokens, cache_creation_tokens, \ + usage_provenance, cost_input, cost_output, cost_cache_creation, cost_cache_read, \ + cost_reasoning, cost_total, stop_reason, error_kind, http_status, latency_ms, \ + body_hash, COALESCE(received_at, completed_at) AS ts \ + FROM requests {where_sql} \ + ORDER BY ts DESC, request_id DESC \ + LIMIT ?" + ); + params.push(SqlValue::Integer(filter.limit + 1)); + let mut fetched: Vec<(RequestRow, i64)> = Vec::new(); + { + let mut stmt = conn.prepare(&sql)?; + let rows = stmt.query_map(rusqlite::params_from_iter(params), map_request_row)?; + for row in rows { + fetched.push(row?); + } + } + let has_more = fetched.len() > usize::try_from(filter.limit).unwrap_or(0); + if has_more { + fetched.truncate(usize::try_from(filter.limit).unwrap_or(0)); + } + let next_cursor = if has_more { + fetched + .last() + .map(|(row, ts)| format!("{ts}:{}", row.request_id)) + } else { + None + }; + Ok(RequestsPage { + rows: fetched.into_iter().map(|(row, _ts)| row).collect(), + next_cursor, + }) +} + +fn callers_sync(path: &Path) -> Result, UsageQueryError> { + let conn = UsageReader::open(path)?; + let mut stmt = conn.prepare( + "SELECT caller, COUNT(*), \ + SUM(CASE WHEN error_kind IS NOT NULL THEN 1 ELSE 0 END), \ + MAX(COALESCE(received_at, completed_at)) \ + FROM requests GROUP BY caller ORDER BY 4 DESC NULLS LAST", + )?; + let rows = stmt.query_map([], |row| { + Ok(CallerSummary { + caller: row.get(0)?, + requests: row.get(1)?, + errors: row.get::<_, Option>(2)?.unwrap_or(0), + last_seen_ms: row.get(3)?, + }) + })?; + let mut callers = Vec::new(); + for row in rows { + callers.push(row?); + } + Ok(callers) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn read_only_connections_refuse_writes() { + let dir = tempfile::tempdir().expect("tempdir"); + let db = dir.path().join("usage.db"); + { + let writer = Connection::open(&db).expect("create db"); + writer + .execute_batch("CREATE TABLE t (a INTEGER);") + .expect("seed table"); + } + let reader = UsageReader::open(&db).expect("open read-only"); + let attempt = reader.execute("INSERT INTO t (a) VALUES (1)", []); + assert!( + attempt.is_err(), + "SQLITE_OPEN_READ_ONLY must refuse writes at the connection level" + ); + // Reads still work on the same connection. + let count: i64 = reader + .query_row("SELECT COUNT(*) FROM t", [], |row| row.get(0)) + .expect("read on the read-only connection"); + assert_eq!(count, 0); + } + + #[test] + fn group_by_parses_known_values_and_rejects_unknown() { + assert_eq!(GroupBy::parse("provider"), Some(GroupBy::Provider)); + assert_eq!(GroupBy::parse("model"), Some(GroupBy::Model)); + assert_eq!(GroupBy::parse("caller"), Some(GroupBy::Caller)); + assert_eq!(GroupBy::parse("team"), None); + assert_eq!(GroupBy::Provider.as_str(), "provider"); + assert_eq!(GroupBy::Model.as_str(), "model"); + assert_eq!(GroupBy::Caller.as_str(), "caller"); + } +} diff --git a/crates/llm-proxy-api/src/routes.rs b/crates/llm-proxy-api/src/routes.rs new file mode 100644 index 0000000..aa474d0 --- /dev/null +++ b/crates/llm-proxy-api/src/routes.rs @@ -0,0 +1,301 @@ +//! Axum handlers and router for the `/usage` query surface. + +use std::path::PathBuf; + +use axum::Json; +use axum::extract::{Query, State}; +use axum::http::{HeaderValue, StatusCode, header}; +use axum::middleware::from_fn_with_state; +use axum::response::{IntoResponse, Response}; +use axum::routing::get; +use serde::Deserialize; +use serde::Serialize; +use tower_http::catch_panic::CatchPanicLayer; +use tower_http::set_header::SetResponseHeaderLayer; +use tower_http::trace::TraceLayer; + +use crate::auth::{AdminAuth, require_admin}; +use crate::reader::{GroupBy, MAX_PAGE_LIMIT, RequestFilter, UsageReader, now_unix_millis}; + +/// Default summary window when `?window=` is absent: 24 hours. +const DEFAULT_WINDOW: std::time::Duration = std::time::Duration::from_secs(24 * 3600); + +/// Shared state for the `/usage` routes: the ledger reader and the admin +/// token guard. +#[derive(Clone, Debug)] +pub(crate) struct UsageState { + pub(crate) reader: UsageReader, + pub(crate) admin: AdminAuth, +} + +/// The `/usage` query surface: a read-only, admin-authenticated router. +/// +/// Build one with [`UsageApi::new`] and either serve it on its own listener +/// (the dedicated admin listener) or merge it into another router (the +/// opt-in main-listener mount). There is no write path: every route is a +/// `GET` backed by read-only SQLite connections. +#[derive(Debug)] +pub struct UsageApi { + state: UsageState, +} + +impl UsageApi { + /// Build the `/usage` surface over the ledger at `db_path`, guarded by + /// `admin_token`. + #[must_use] + pub fn new(db_path: impl Into, admin_token: &str) -> Self { + Self { + state: UsageState { + reader: UsageReader::new(db_path), + admin: AdminAuth::new(admin_token), + }, + } + } + + /// The `/usage` router with state applied, ready to serve or merge. + /// + /// Every route sits behind the admin-token middleware. The outer layers + /// (`nosniff`, panic catching, tracing) are installed here so the same + /// protections apply on both the dedicated listener and any merged + /// mount, independent of the surrounding router's own layers. + #[must_use = "the Router must be used with a server or merged into one"] + pub fn router(self) -> axum::Router { + let state = self.state; + axum::Router::new() + .route("/usage/summary", get(summary)) + .route("/usage/requests", get(requests)) + .route("/usage/keys", get(keys)) + .layer(from_fn_with_state(state.clone(), require_admin)) + .layer(TraceLayer::new_for_http()) + .layer(SetResponseHeaderLayer::if_not_present( + header::X_CONTENT_TYPE_OPTIONS, + HeaderValue::from_static("nosniff"), + )) + .layer(CatchPanicLayer::new()) + .with_state(state) + } +} + +/// Handler error: invalid query parameters (400) or a ledger failure (500, +/// with the detail logged server-side only). +#[derive(Debug)] +pub(crate) enum UsageApiError { + /// A query parameter failed validation; the message is safe to return. + InvalidParam(String), + /// The ledger query failed; the response carries a generic message. + Query(crate::reader::UsageQueryError), +} + +#[derive(Serialize)] +struct ErrorBody { + error: ErrorDetail, +} + +#[derive(Serialize)] +struct ErrorDetail { + r#type: &'static str, + message: String, +} + +impl IntoResponse for UsageApiError { + fn into_response(self) -> Response { + match self { + Self::InvalidParam(message) => ( + StatusCode::BAD_REQUEST, + Json(ErrorBody { + error: ErrorDetail { + r#type: "invalid_request", + message, + }, + }), + ) + .into_response(), + Self::Query(error) => { + tracing::error!(error = %error, "usage ledger query failed"); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(ErrorBody { + error: ErrorDetail { + r#type: "internal_error", + message: "usage ledger query failed".to_owned(), + }, + }), + ) + .into_response() + } + } + } +} + +#[derive(Deserialize)] +struct SummaryParams { + window: Option, + group_by: Option, +} + +#[derive(Serialize)] +struct SummaryBody { + since_ms: i64, + until_ms: i64, + group_by: &'static str, + groups: Vec, +} + +/// `GET /usage/summary?window=1h&group_by=provider|model|caller`. +async fn summary( + State(state): State, + Query(params): Query, +) -> Result, UsageApiError> { + let until_ms = now_unix_millis(); + let window_ms = match ¶ms.window { + None => DEFAULT_WINDOW.as_millis(), + Some(raw) => { + let parsed = humantime::parse_duration(raw).map_err(|_parse_error| { + UsageApiError::InvalidParam(format!( + "window must be a duration like \"1h\" or \"30m\", got {raw:?}" + )) + })?; + if parsed.is_zero() { + return Err(UsageApiError::InvalidParam( + "window must be a positive duration".to_owned(), + )); + } + parsed.as_millis() + } + }; + let since_ms = until_ms.saturating_sub(i64::try_from(window_ms).unwrap_or(i64::MAX)); + let group = match params.group_by.as_deref() { + None | Some("provider") => GroupBy::Provider, + Some(other) => GroupBy::parse(other).ok_or_else(|| { + UsageApiError::InvalidParam(format!( + "group_by must be one of provider, model, caller; got {other:?}" + )) + })?, + }; + let groups = state + .reader + .summary(since_ms, group) + .await + .map_err(UsageApiError::Query)?; + Ok(Json(SummaryBody { + since_ms, + until_ms, + group_by: group.as_str(), + groups, + })) +} + +#[derive(Deserialize)] +struct RequestsParams { + since: Option, + until: Option, + provider: Option, + limit: Option, + cursor: Option, +} + +#[derive(Serialize)] +struct RequestsBody { + rows: Vec, + next_cursor: Option, +} + +/// Parse an RFC 3339 timestamp parameter into Unix epoch milliseconds. +fn parse_rfc3339_ms(name: &str, raw: &str) -> Result { + use time::format_description::well_known::Rfc3339; + time::OffsetDateTime::parse(raw, &Rfc3339) + .map(|ts| { + ts.unix_timestamp() + .saturating_mul(1000) + .saturating_add(i64::from(ts.millisecond())) + }) + .map_err(|_parse_error| { + UsageApiError::InvalidParam(format!( + "{name} must be an RFC 3339 timestamp like 2026-01-02T03:04:05Z, got {raw:?}" + )) + }) +} + +/// Parse an opaque `cursor` (`:`). +fn parse_cursor(raw: &str) -> Result<(i64, String), UsageApiError> { + let Some((ts, id)) = raw.split_once(':') else { + return Err(UsageApiError::InvalidParam( + "cursor is malformed; pass back the next_cursor value unchanged".to_owned(), + )); + }; + let ts: i64 = ts.parse().map_err(|_parse_error| { + UsageApiError::InvalidParam( + "cursor is malformed; pass back the next_cursor value unchanged".to_owned(), + ) + })?; + if id.is_empty() { + return Err(UsageApiError::InvalidParam( + "cursor is malformed; pass back the next_cursor value unchanged".to_owned(), + )); + } + Ok((ts, id.to_owned())) +} + +/// `GET /usage/requests?since=…&until=…&provider=…&limit=…&cursor=…`. +/// +/// Rows come back newest first, bounded by `limit` (default 100, maximum +/// 1000), paged via the returned `next_cursor`. `since` is inclusive, +/// `until` exclusive; both are RFC 3339. Responses never contain request +/// bodies — the ledger stores only the body digest. +async fn requests( + State(state): State, + Query(params): Query, +) -> Result, UsageApiError> { + let since_ms = match ¶ms.since { + Some(raw) => Some(parse_rfc3339_ms("since", raw)?), + None => None, + }; + let until_ms = match ¶ms.until { + Some(raw) => Some(parse_rfc3339_ms("until", raw)?), + None => None, + }; + if let (Some(since), Some(until)) = (since_ms, until_ms) + && since >= until + { + return Err(UsageApiError::InvalidParam( + "since must be before until".to_owned(), + )); + } + let limit = params.limit.unwrap_or(100); + if !(1..=MAX_PAGE_LIMIT).contains(&limit) { + return Err(UsageApiError::InvalidParam(format!( + "limit must be between 1 and {MAX_PAGE_LIMIT}, got {limit}" + ))); + } + let cursor = match ¶ms.cursor { + Some(raw) => Some(parse_cursor(raw)?), + None => None, + }; + let page = state + .reader + .requests(RequestFilter { + since_ms, + until_ms, + provider: params.provider, + limit, + cursor, + }) + .await + .map_err(UsageApiError::Query)?; + Ok(Json(RequestsBody { + rows: page.rows, + next_cursor: page.next_cursor, + })) +} + +#[derive(Serialize)] +struct KeysBody { + callers: Vec, +} + +/// `GET /usage/keys` — caller list with request counts, error counts, and +/// last-seen times. +async fn keys(State(state): State) -> Result, UsageApiError> { + let callers = state.reader.callers().await.map_err(UsageApiError::Query)?; + Ok(Json(KeysBody { callers })) +} diff --git a/crates/llm-proxy-api/tests/secret_abuse.rs b/crates/llm-proxy-api/tests/secret_abuse.rs new file mode 100644 index 0000000..4ea42c4 --- /dev/null +++ b/crates/llm-proxy-api/tests/secret_abuse.rs @@ -0,0 +1,216 @@ +//! Secret-abuse family (PRX-06) for the `/usage` surface itself: the admin +//! token never reaches the captured tracing output of usage traffic, and +//! neither the admin token nor any rejected credential ever lands in the raw +//! ledger file bytes. Negative assertions carry a positive control (seeded +//! ledger content that IS findable), so a passing scan proves the scan read +//! the real file rather than the wrong path. + +#[cfg(test)] +mod tests { + use std::path::{Path, PathBuf}; + use std::sync::{Arc, Mutex, MutexGuard}; + use std::time::Duration; + + use axum::Router; + use axum::body::Body; + use axum::http::{Request, StatusCode}; + use llm_proxy_storage::{EventBus, ProxyEvent, RequestReceived, SqliteEventBus}; + use tower::ServiceExt; + + const ADMIN_TOKEN: &str = "prx6-api-admin-token-0a3f5e91"; + const WRONG_TOKEN: &str = "prx6-api-attacker-guess-2c4d"; + const CLIENT_TOKEN: &str = "prx6-api-client-token-7e8f"; + /// Seeded caller: legitimately persisted, so the byte scan has a positive + /// control proving it reads real rows. + const CALLER_MARKER: &str = "prx6-seeded-caller-marker-4451"; + + /// Seed a one-row ledger and close the writer, returning the db path. + async fn seeded_ledger(dir: &tempfile::TempDir) -> PathBuf { + let path = dir.path().join("usage.db"); + let bus = SqliteEventBus::open( + &path, + llm_proxy_storage::RetentionPolicy { + max_age: None, + max_rows: None, + prune_every: Duration::from_secs(3600), + }, + ) + .expect("ledger opens"); + bus.emit(&ProxyEvent::RequestReceived(RequestReceived { + request_id: "prx6-seed".to_owned(), + timestamp: time::OffsetDateTime::now_utc(), + caller: CALLER_MARKER.to_owned(), + provider: "prov-seed".to_owned(), + route_kind: "messages".to_owned(), + client_protocol: "anthropic".to_owned(), + model: llm_proxy_protocol::core::ModelRef { + requested: "claude-3".to_owned(), + upstream: None, + }, + streaming: false, + body_hash: "prx6-seed-digest".to_owned(), + })); + let closer = Arc::new(bus); + tokio::task::spawn_blocking(move || closer.close()) + .await + .expect("close joins"); + path + } + + fn usage_app(path: &Path) -> Router { + llm_proxy_api::UsageApi::new(path.to_owned(), ADMIN_TOKEN).router() + } + + fn get(uri: &str, token: Option<&str>) -> Request { + let mut builder = Request::builder().uri(uri); + if let Some(token) = token { + builder = builder.header("authorization", format!("Bearer {token}")); + } + builder.body(Body::empty()).expect("request builds") + } + + async fn drive_all_outcomes(path: &Path) { + // One authenticated call per route, plus every rejection mode, so the + // whole surface's traffic flows before the assertions. + for uri in ["/usage/summary", "/usage/requests", "/usage/keys"] { + for token in [ + Some(ADMIN_TOKEN), + Some(WRONG_TOKEN), + Some(CLIENT_TOKEN), + None, + ] { + let (status, body) = send(usage_app(path).clone(), get(uri, token)).await; + let expected = if token == Some(ADMIN_TOKEN) { + StatusCode::OK + } else { + StatusCode::UNAUTHORIZED + }; + assert_eq!(status, expected, "{uri} with {token:?}: {body}"); + } + } + } + + async fn send(app: Router, request: Request) -> (StatusCode, String) { + let response = app.oneshot(request).await.expect("oneshot completes"); + let status = response.status(); + let bytes = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .expect("body reads"); + (status, String::from_utf8_lossy(&bytes).into_owned()) + } + + /// Concatenated raw bytes of the ledger database and any WAL/SHM + /// sidecars — everything readable on disk. + fn ledger_bytes(db: &Path) -> Vec { + let mut all = Vec::new(); + for suffix in ["", "-wal", "-shm"] { + let sidecar = PathBuf::from(format!("{}{suffix}", db.display())); + if let Ok(bytes) = std::fs::read(&sidecar) { + all.extend_from_slice(&bytes); + } + } + all + } + + /// Captured tracing output for the whole test binary. + /// + /// A process-global subscriber (installed exactly once) instead of a + /// thread-local default: with parallel tests, callsites first emitted on + /// a thread that has no subscriber can keep a stale "never" interest + /// cache, silently dropping later events from a thread-local capture — + /// observed as flaky partial captures during development. The global + /// subscriber receives every test's traffic on every thread, which is + /// strictly more content for the secret scans below. + struct TraceCapture { + buffer: Arc>>, + } + + static CAPTURE_BUFFER: std::sync::OnceLock>>> = std::sync::OnceLock::new(); + static CAPTURE_INSTALLED: std::sync::OnceLock<()> = std::sync::OnceLock::new(); + + impl TraceCapture { + fn install() -> Self { + use tracing_subscriber::layer::SubscriberExt; + let buffer = CAPTURE_BUFFER.get_or_init(|| Arc::new(Mutex::new(Vec::::new()))); + CAPTURE_INSTALLED.get_or_init(|| { + let layer = tracing_subscriber::fmt::layer() + .with_writer(BufferWriter(Arc::clone(buffer))) + .with_ansi(false); + tracing::subscriber::set_global_default(tracing_subscriber::registry().with(layer)) + .expect("the global capture subscriber installs exactly once"); + }); + Self { + buffer: Arc::clone(buffer), + } + } + + fn text(&self) -> String { + let bytes = self.buffer.lock().expect("buffer lock").clone(); + String::from_utf8_lossy(&bytes).into_owned() + } + } + + #[derive(Clone)] + struct BufferWriter(Arc>>); + + impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for BufferWriter { + type Writer = BufferWriterGuard<'a>; + + fn make_writer(&'a self) -> Self::Writer { + BufferWriterGuard(self.0.lock().expect("buffer lock")) + } + } + + struct BufferWriterGuard<'a>(MutexGuard<'a, Vec>); + + impl std::io::Write for BufferWriterGuard<'_> { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + self.0.extend_from_slice(bytes); + Ok(bytes.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + + #[tokio::test] + async fn no_credential_ever_lands_in_the_raw_ledger_file_bytes() { + let dir = tempfile::tempdir().expect("tempdir"); + let path = seeded_ledger(&dir).await; + drive_all_outcomes(&path).await; + + let bytes = ledger_bytes(&path); + let text = String::from_utf8_lossy(&bytes); + assert!( + text.contains(CALLER_MARKER) && text.contains("prx6-seed-digest"), + "control: seeded persisted values must be findable in the ledger bytes" + ); + for secret in [ADMIN_TOKEN, WRONG_TOKEN, CLIENT_TOKEN] { + assert!( + !text.contains(secret), + "credential {secret} leaked into the raw ledger file bytes" + ); + } + } + + #[tokio::test] + async fn admin_token_never_appears_in_captured_usage_tracing() { + let dir = tempfile::tempdir().expect("tempdir"); + let path = seeded_ledger(&dir).await; + + let capture = TraceCapture::install(); + drive_all_outcomes(&path).await; + let captured = capture.text(); + assert!( + captured.contains("started processing request"), + "control: usage request traffic must land in the capture, got: {captured}" + ); + for secret in [ADMIN_TOKEN, WRONG_TOKEN, CLIENT_TOKEN] { + assert!( + !captured.contains(secret), + "credential {secret} leaked into tracing output: {captured}" + ); + } + } +} diff --git a/crates/llm-proxy-api/tests/usage_api.rs b/crates/llm-proxy-api/tests/usage_api.rs new file mode 100644 index 0000000..bf0e95b --- /dev/null +++ b/crates/llm-proxy-api/tests/usage_api.rs @@ -0,0 +1,639 @@ +//! Integration tests for the `/usage` query surface over a real SQLite +//! ledger seeded through `SqliteEventBus`. + +#[cfg(test)] +mod tests { + use axum::body::Body; + use axum::http::Request; + use axum::http::StatusCode; + use llm_proxy_protocol::core::{Cost, ModelRef, StopReason, Usage, UsageProvenance}; + use llm_proxy_storage::{ + EventBus, ProxyEvent, RequestReceived, ResponseCompleted, ResponseFailed, SqliteEventBus, + }; + use rust_decimal::Decimal; + use serde_json::Value; + use time::OffsetDateTime; + use tower::ServiceExt; + + const ADMIN_TOKEN: &str = "test-admin-token"; + + /// Fixture: five ledgered requests with fixed, distinct timestamps so + /// window/order/cursor behavior is deterministic. + /// + /// - req-1 (laptop/prov-a/claude-3): completed, exact-decimal cost A. + /// - req-2 (laptop/prov-a/gpt-4): completed, exact-decimal cost B. + /// - req-3 (phone/prov-b/claude-3): failed upstream. + /// - orphan-1 (no caller/provider): failure-only row. + /// - old-1 (laptop/prov-old/claude-3): completed, 2 hours old. + struct Ledger { + _dir: tempfile::TempDir, + path: std::path::PathBuf, + t0: OffsetDateTime, + } + + fn received_event( + id: &str, + caller: &str, + provider: &str, + model: &str, + ts: OffsetDateTime, + ) -> ProxyEvent { + ProxyEvent::RequestReceived(RequestReceived { + request_id: id.to_owned(), + timestamp: ts, + caller: caller.to_owned(), + provider: provider.to_owned(), + route_kind: "messages".to_owned(), + client_protocol: "anthropic".to_owned(), + model: ModelRef { + requested: model.to_owned(), + upstream: None, + }, + streaming: true, + body_hash: format!("sha256-digest-{id}"), + }) + } + + fn completed_event( + id: &str, + provider: &str, + model: &str, + ts: OffsetDateTime, + total: &str, + ) -> ProxyEvent { + ProxyEvent::ResponseCompleted(ResponseCompleted { + request_id: id.to_owned(), + timestamp: ts, + provider: provider.to_owned(), + upstream_message_id: None, + model: ModelRef { + requested: model.to_owned(), + upstream: None, + }, + usage: Usage { + input_tokens: 100, + output_tokens: 50, + reasoning_tokens: Some(7), + cache_creation_input_tokens: Some(30), + cache_read_input_tokens: Some(200), + provenance: UsageProvenance::ProviderReported, + }, + cost: Some(Cost { + input: Decimal::from_str_exact(total).expect("decimal fixture"), + output: Decimal::ZERO, + cache_creation: Decimal::ZERO, + cache_read: Decimal::ZERO, + reasoning: Decimal::ZERO, + total: Decimal::from_str_exact(total).expect("decimal fixture"), + }), + stop_reason: StopReason::EndTurn, + latency_ms: 4321, + }) + } + + fn failed_event(id: &str, provider: Option<&str>, ts: OffsetDateTime) -> ProxyEvent { + ProxyEvent::ResponseFailed(ResponseFailed { + request_id: id.to_owned(), + timestamp: ts, + provider: provider.map(str::to_owned), + model: None, + error_kind: "UpstreamError".to_owned(), + message: "sanitized failure text".to_owned(), + http_status: 502, + latency_ms: 99, + }) + } + + async fn seed() -> Ledger { + let dir = tempfile::tempdir().expect("tempdir"); + let path = dir.path().join("usage.db"); + let bus = SqliteEventBus::open( + &path, + llm_proxy_storage::RetentionPolicy { + max_age: None, + max_rows: None, + prune_every: std::time::Duration::from_secs(3600), + }, + ) + .expect("open ledger"); + + let t0 = OffsetDateTime::now_utc().saturating_sub(time::Duration::seconds(60)); + let second = std::time::Duration::from_secs(1); + let events = [ + // req-1: laptop / prov-a / claude-3, cost A + received_event("req-1", "laptop", "prov-a", "claude-3", t0), + completed_event( + "req-1", + "prov-a", + "claude-3", + t0 + second * 4, + "0.1000000000000000000000001", + ), + // req-2: laptop / prov-a / gpt-4, cost B + received_event("req-2", "laptop", "prov-a", "gpt-4", t0 + second), + completed_event( + "req-2", + "prov-a", + "gpt-4", + t0 + second * 5, + "0.2000000000000000000000002", + ), + // req-3: phone / prov-b, failed + received_event("req-3", "phone", "prov-b", "claude-3", t0 + second * 2), + failed_event("req-3", Some("prov-b"), t0 + second * 6), + // orphan-1: failure-only row (no caller, no provider) + failed_event("orphan-1", None, t0 + second * 3), + // old-1: two hours old, its own provider so window tests stay clean + received_event( + "old-1", + "laptop", + "prov-old", + "claude-3", + t0.saturating_sub(time::Duration::seconds(7200)), + ), + completed_event( + "old-1", + "prov-old", + "claude-3", + t0.saturating_sub(time::Duration::seconds(7190)), + "0.5", + ), + ]; + for event in &events { + bus.emit(event); + } + let closer = std::sync::Arc::new(bus); + tokio::task::spawn_blocking(move || closer.close()) + .await + .expect("close joined"); + Ledger { + _dir: dir, + path, + t0, + } + } + + fn app_for(path: &std::path::Path) -> axum::Router { + llm_proxy_api::UsageApi::new(path.to_owned(), ADMIN_TOKEN).router() + } + + async fn get_json(app: axum::Router, uri: &str, token: Option<&str>) -> (StatusCode, Value) { + let mut builder = Request::builder().uri(uri); + if let Some(token) = token { + builder = builder.header("authorization", format!("Bearer {token}")); + } + let response = app + .oneshot(builder.body(Body::empty()).expect("request builds")) + .await + .expect("oneshot"); + let status = response.status(); + let bytes = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .expect("body"); + let json = serde_json::from_slice(&bytes).expect("body is JSON"); + (status, json) + } + + #[tokio::test] + async fn usage_routes_refuse_missing_and_wrong_tokens_without_echoing_them() { + let ledger = seed().await; + for path in ["/usage/summary", "/usage/requests", "/usage/keys"] { + let (status, body) = get_json(app_for(&ledger.path), path, None).await; + assert_eq!(status, StatusCode::UNAUTHORIZED, "{path} without token"); + assert_eq!( + body["error"]["message"], "missing or invalid admin token", + "uniform 401 body for {path}" + ); + + let (status, body_text) = { + let (status, body) = + get_json(app_for(&ledger.path), path, Some("wrong-token")).await; + (status, body.to_string()) + }; + assert_eq!(status, StatusCode::UNAUTHORIZED, "{path} with wrong token"); + assert!( + !body_text.contains("wrong-token") && !body_text.contains(ADMIN_TOKEN), + "the 401 body must not echo any presented or configured token: {body_text}" + ); + + let (status, _) = get_json(app_for(&ledger.path), path, Some(ADMIN_TOKEN)).await; + assert_eq!(status, StatusCode::OK, "{path} with the admin token"); + } + } + + #[tokio::test] + async fn summary_groups_by_provider_with_exact_decimal_costs() { + let ledger = seed().await; + let (status, body) = get_json( + app_for(&ledger.path), + "/usage/summary?group_by=provider", + Some(ADMIN_TOKEN), + ) + .await; + assert_eq!(status, StatusCode::OK); + + let groups = body["groups"].as_array().expect("groups array"); + let find = |key: &str| { + groups + .iter() + .find(|g| g["key"] == key) + .unwrap_or_else(|| panic!("group {key} missing: {groups:?}")) + }; + + let prov_a = find("prov-a"); + assert_eq!(prov_a["requests"], 2); + assert_eq!(prov_a["errors"], 0, "both prov-a requests completed"); + assert_eq!( + prov_a["tokens"]["input"], 200, + "100 input tokens per request" + ); + assert_eq!(prov_a["tokens"]["output"], 100); + assert_eq!(prov_a["tokens"]["reasoning"], 14); + assert_eq!(prov_a["tokens"]["cache_read"], 400); + assert_eq!(prov_a["tokens"]["cache_creation"], 60); + assert_eq!( + prov_a["cost"]["total"], "0.3000000000000000000000003", + "cost sums must be exact decimals — a float sum of these two values \ + yields 0.30000000000000004" + ); + assert_eq!( + prov_a["cost"]["input"], "0.3000000000000000000000003", + "input-cost column sums exactly too" + ); + + let prov_b = find("prov-b"); + assert_eq!(prov_b["requests"], 1); + assert_eq!(prov_b["errors"], 1); + + // The failure-only orphan row groups under a null key. + let null_group = groups + .iter() + .find(|g| g["key"].is_null()) + .expect("failure-only rows group under a null provider"); + assert_eq!(null_group["requests"], 1); + assert_eq!(null_group["errors"], 1); + + assert_eq!(body["group_by"], "provider"); + } + + #[tokio::test] + async fn summary_groups_by_caller_and_model() { + let ledger = seed().await; + let (_, body) = get_json( + app_for(&ledger.path), + "/usage/summary?group_by=caller", + Some(ADMIN_TOKEN), + ) + .await; + let groups = &body["groups"]; + let laptop = groups + .as_array() + .expect("groups") + .iter() + .find(|g| g["key"] == "laptop") + .expect("laptop group"); + assert_eq!( + laptop["requests"], 3, + "req-1, req-2 and the two-hour-old row" + ); + assert_eq!(laptop["errors"], 0); + + let (_, body) = get_json( + app_for(&ledger.path), + "/usage/summary?group_by=model", + Some(ADMIN_TOKEN), + ) + .await; + let model_groups = body["groups"].as_array().expect("groups"); + let claude = model_groups + .iter() + .find(|g| g["key"] == "claude-3") + .expect("claude-3 group"); + assert_eq!(claude["requests"], 3, "req-1, req-3 and old-1"); + let gpt = model_groups + .iter() + .find(|g| g["key"] == "gpt-4") + .expect("gpt-4 group"); + assert_eq!(gpt["requests"], 1); + } + + #[tokio::test] + async fn summary_window_excludes_rows_older_than_it() { + let ledger = seed().await; + let (_, body) = get_json( + app_for(&ledger.path), + "/usage/summary?window=1h", + Some(ADMIN_TOKEN), + ) + .await; + let groups = body["groups"].as_array().expect("groups"); + assert!( + groups.iter().all(|g| g["key"] != "prov-old"), + "the two-hour-old row must fall outside a 1h window: {groups:?}" + ); + let (_, body) = get_json( + app_for(&ledger.path), + "/usage/summary?window=4h", + Some(ADMIN_TOKEN), + ) + .await; + let groups = body["groups"].as_array().expect("groups"); + assert!( + groups.iter().any(|g| g["key"] == "prov-old"), + "a 4h window covers the two-hour-old row" + ); + + // Day-denominated windows are the natural operator input for + // retention horizons; they must parse through the same grammar. + let (status, _) = get_json( + app_for(&ledger.path), + "/usage/summary?window=7d", + Some(ADMIN_TOKEN), + ) + .await; + assert_eq!(status, StatusCode::OK, "\"7d\" must be an accepted window"); + } + + #[tokio::test] + async fn summary_rejects_bad_window_and_group_by() { + let ledger = seed().await; + let (status, body) = get_json( + app_for(&ledger.path), + "/usage/summary?window=abc", + Some(ADMIN_TOKEN), + ) + .await; + assert_eq!(status, StatusCode::BAD_REQUEST); + assert!(body["error"]["message"].to_string().contains("window")); + + let (status, _) = get_json( + app_for(&ledger.path), + "/usage/summary?window=0s", + Some(ADMIN_TOKEN), + ) + .await; + assert_eq!( + status, + StatusCode::BAD_REQUEST, + "zero window is not a window" + ); + + let (status, body) = get_json( + app_for(&ledger.path), + "/usage/summary?group_by=team", + Some(ADMIN_TOKEN), + ) + .await; + assert_eq!(status, StatusCode::BAD_REQUEST); + assert!(body["error"]["message"].to_string().contains("group_by")); + } + + #[tokio::test] + async fn requests_pages_newest_first_with_cursor_and_provider_filter() { + let ledger = seed().await; + let (status, page1) = get_json( + app_for(&ledger.path), + "/usage/requests?limit=2", + Some(ADMIN_TOKEN), + ) + .await; + assert_eq!(status, StatusCode::OK); + let rows1 = page1["rows"].as_array().expect("rows"); + assert_eq!(rows1.len(), 2, "limit is honored"); + let first_id = rows1[0]["request_id"].as_str().expect("id"); + assert_eq!( + first_id, "orphan-1", + "newest first: orphan-1 completed latest among the seeded rows" + ); + let cursor = page1["next_cursor"].as_str().expect("cursor present"); + assert_eq!(rows1[1]["request_id"], "req-3"); + + // Follow the cursor: the next page must not repeat page-1 rows. + let (status, page2) = get_json( + app_for(&ledger.path), + &format!("/usage/requests?limit=10&cursor={cursor}"), + Some(ADMIN_TOKEN), + ) + .await; + assert_eq!(status, StatusCode::OK); + let rows2 = page2["rows"].as_array().expect("rows"); + assert!( + rows2 + .iter() + .all(|r| r["request_id"] != "orphan-1" && r["request_id"] != "req-3"), + "cursor paging must not repeat rows: {rows2:?}" + ); + + // Provider filter narrows to that provider's rows only. + let (_, filtered) = get_json( + app_for(&ledger.path), + "/usage/requests?provider=prov-a", + Some(ADMIN_TOKEN), + ) + .await; + let filtered_rows = filtered["rows"].as_array().expect("rows"); + assert!( + !filtered_rows.is_empty() && filtered_rows.iter().all(|r| r["provider"] == "prov-a"), + "provider filter: {filtered_rows:?}" + ); + } + + #[tokio::test] + async fn requests_honors_since_and_until_bounds() { + let ledger = seed().await; + // until = t0+3s (exclusive): only rows whose ledger timestamp (received + // time for normal rows, failure time for orphan rows) is strictly older. + let until = (ledger.t0 + std::time::Duration::from_secs(3)) + .format(&time::format_description::well_known::Rfc3339) + .expect("rfc3339"); + let (_, body) = get_json( + app_for(&ledger.path), + &format!("/usage/requests?until={until}"), + Some(ADMIN_TOKEN), + ) + .await; + let ids: Vec<&str> = body["rows"] + .as_array() + .expect("rows") + .iter() + .filter_map(|r| r["request_id"].as_str()) + .collect(); + assert_eq!( + ids, + vec!["req-3", "req-2", "req-1", "old-1"], + "until is exclusive on the ledger time axis (orphan-1 at t0+3s is out)" + ); + + // since = t0+1s (inclusive): rows at or after t0+1s. + let since = (ledger.t0 + std::time::Duration::from_secs(1)) + .format(&time::format_description::well_known::Rfc3339) + .expect("rfc3339"); + let (_, body) = get_json( + app_for(&ledger.path), + &format!("/usage/requests?since={since}"), + Some(ADMIN_TOKEN), + ) + .await; + let ids: Vec<&str> = body["rows"] + .as_array() + .expect("rows") + .iter() + .filter_map(|r| r["request_id"].as_str()) + .collect(); + assert_eq!( + ids, + vec!["orphan-1", "req-3", "req-2"], + "since is inclusive" + ); + } + + #[tokio::test] + async fn requests_rejects_invalid_parameters() { + let ledger = seed().await; + let cases = [ + "/usage/requests?limit=0", + "/usage/requests?limit=1001", + "/usage/requests?limit=-5", + "/usage/requests?since=not-a-time", + "/usage/requests?until=not-a-time", + "/usage/requests?cursor=garbage", + "/usage/requests?cursor=123:", // empty id half + "/usage/requests?since=2026-01-02T03:04:05Z&until=2026-01-01T00:00:00Z", + ]; + for uri in cases { + let (status, body) = get_json(app_for(&ledger.path), uri, Some(ADMIN_TOKEN)).await; + assert_eq!(status, StatusCode::BAD_REQUEST, "{uri}: {body}"); + } + } + + #[tokio::test] + async fn requests_rows_carry_digests_but_never_bodies() { + let ledger = seed().await; + let (_, body) = get_json(app_for(&ledger.path), "/usage/requests", Some(ADMIN_TOKEN)).await; + let text = body.to_string(); + assert!( + text.contains("sha256-digest-"), + "rows expose the body digest for correlation" + ); + for row in body["rows"].as_array().expect("rows") { + let keys: Vec<&str> = row + .as_object() + .expect("row object") + .keys() + .map(String::as_str) + .collect(); + assert!( + !keys + .iter() + .any(|k| *k == "body" || *k == "request_body" || *k == "message"), + "raw rows must never carry request bodies or sanitized messages: {keys:?}" + ); + } + } + + #[tokio::test] + async fn keys_lists_callers_with_counts_and_last_seen() { + let ledger = seed().await; + let (status, body) = + get_json(app_for(&ledger.path), "/usage/keys", Some(ADMIN_TOKEN)).await; + assert_eq!(status, StatusCode::OK); + let callers = body["callers"].as_array().expect("callers"); + let find = |name: &str| { + callers + .iter() + .find(|c| c["caller"] == name) + .unwrap_or_else(|| panic!("caller {name} missing: {callers:?}")) + }; + let laptop = find("laptop"); + assert_eq!(laptop["requests"], 3, "req-1, req-2, old-1"); + assert_eq!(laptop["errors"], 0); + let phone = find("phone"); + assert_eq!(phone["requests"], 1); + assert_eq!(phone["errors"], 1); + assert!(laptop["last_seen_ms"].is_u64(), "last_seen_ms is populated"); + + // Ordering: most recently seen first. The last-seen value follows the + // ledger time axis — the request's received time for normal rows, the + // failure time for the caller-less orphan — so with this fixture: + // (t0+3s), phone (t0+2s), laptop (t0+1s). + let dated: Vec = callers + .iter() + .filter_map(|c| c["last_seen_ms"].as_i64()) + .collect(); + assert!( + !dated.is_empty(), + "callers with ledger activity must carry last_seen_ms: {callers:?}" + ); + assert_eq!( + dated, + { + let mut sorted = dated.clone(); + sorted.sort_by(|a, b| b.cmp(a)); + sorted + }, + "callers appear most-recently-seen first" + ); + let order: Vec<&str> = callers + .iter() + .map(|c| c["caller"].as_str().unwrap_or("")) + .collect(); + assert_eq!( + order, + vec!["", "phone", "laptop"], + "caller ordering follows last-seen timestamps, null caller included" + ); + } + + #[tokio::test] + async fn summary_serves_while_the_ledger_writer_is_active() { + // Read-only queries must work against a live WAL writer (the steady + // state of a running proxy), not only against a closed database. + let dir = tempfile::tempdir().expect("tempdir"); + let path = dir.path().join("usage.db"); + let bus = std::sync::Arc::new( + SqliteEventBus::open( + &path, + llm_proxy_storage::RetentionPolicy { + max_age: None, + max_rows: None, + prune_every: std::time::Duration::from_secs(3600), + }, + ) + .expect("open ledger"), + ); + bus.emit(&received_event( + "live-1", + "laptop", + "prov-live", + "claude-3", + OffsetDateTime::now_utc(), + )); + + let mut seen = false; + for _ in 0..200 { + let (status, body) = get_json( + app_for(&path), + "/usage/summary?group_by=provider", + Some(ADMIN_TOKEN), + ) + .await; + assert_eq!(status, StatusCode::OK, "query against live writer"); + let found = body["groups"] + .as_array() + .is_some_and(|g| g.iter().any(|x| x["key"] == "prov-live")); + if found { + seen = true; + break; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + assert!( + seen, + "the row must become visible while the writer is active" + ); + let closer = std::sync::Arc::clone(&bus); + tokio::task::spawn_blocking(move || closer.close()) + .await + .expect("close joined"); + } +} diff --git a/crates/llm-proxy-core/src/lib.rs b/crates/llm-proxy-core/src/lib.rs index bdcb108..a975ef4 100644 --- a/crates/llm-proxy-core/src/lib.rs +++ b/crates/llm-proxy-core/src/lib.rs @@ -56,10 +56,11 @@ pub use error::CoreError; pub use metrics::{Metrics, Snapshot}; pub use pid::{PidError, PidManager}; pub use provider_config::{ - AppConfig, AuthStyle, ConfigValidationError, LogFormat, ModelId, ModelPricing, - ProviderAdapterConfig, ProviderCatalogConfig, ProviderCatalogMode, ProviderConfig, - ProviderDiscoveryConfig, ProviderDiscoveryKind, ProviderFile, ProviderRouteKind, - ProviderRoutesConfig, ServerConfig, StaticModelCatalogEntry, load_app_config, + ANONYMOUS_CALLER, AdminAuthConfig, AdminMountPosture, AppConfig, AuthConfig, AuthMode, + AuthStyle, ConfigValidationError, LogFormat, ModelId, ModelPricing, ProviderAdapterConfig, + ProviderCatalogConfig, ProviderCatalogMode, ProviderConfig, ProviderDiscoveryConfig, + ProviderDiscoveryKind, ProviderEgressConfig, ProviderFile, ProviderRouteKind, + ProviderRoutesConfig, ServerConfig, StaticModelCatalogEntry, StorageConfig, load_app_config, load_provider_config, validate_provider_config, }; pub use provider_registry::{ diff --git a/crates/llm-proxy-core/src/provider_config.rs b/crates/llm-proxy-core/src/provider_config.rs index e4a4994..1846239 100644 --- a/crates/llm-proxy-core/src/provider_config.rs +++ b/crates/llm-proxy-core/src/provider_config.rs @@ -7,7 +7,7 @@ use std::collections::HashMap; use std::net::SocketAddr; -use std::path::Path; +use std::path::{Path, PathBuf}; use std::sync::Arc; use std::time::Duration; @@ -24,12 +24,62 @@ use crate::error::CoreError; /// Top-level TOML application configuration. /// -/// Loaded from the main `config.toml`, which contains server settings only. +/// Loaded from the main `config.toml`, which contains server and auth +/// settings. #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct AppConfig { /// Server bind and operational settings. pub server: ServerConfig, + /// Client authentication settings (`[auth]` section). + #[serde(default)] + pub auth: AuthConfig, + /// Usage-ledger persistence settings (`[storage]` section). Absent keeps + /// the in-memory no-op event bus (nothing is persisted). + #[serde(default)] + pub storage: Option, +} + +/// Named security postures for where the admin `/usage` API is reachable. +/// +/// Selected by config; the serve path logs the posture name loudly at startup +/// so operators can see which surface they configured. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AdminMountPosture { + /// Default: `/usage` is served only on the dedicated + /// [`ServerConfig::admin_bind`] listener (loopback unless configured + /// otherwise); the main listener has no admin surface at all. + DedicatedListener, + /// Opt-in (`server.admin_on_main_listener = true`): `/usage` is also + /// served on the main `bind` listener behind the same admin token, so + /// the admin surface reaches every network the main listener does. + MainListener, +} + +impl AdminMountPosture { + /// Stable name used in startup logs so operators can grep the posture + /// their config selected. + #[must_use] + pub const fn name(self) -> &'static str { + match self { + Self::DedicatedListener => "dedicated-admin-listener", + Self::MainListener => "admin-on-main-listener", + } + } +} + +impl AppConfig { + /// The admin `/usage` mount posture, or `None` when no + /// `[auth.admin]` token is configured (no admin surface anywhere). + #[must_use] + pub fn admin_posture(&self) -> Option { + let posture = if self.server.admin_on_main_listener { + AdminMountPosture::MainListener + } else { + AdminMountPosture::DedicatedListener + }; + self.auth.admin.as_ref().map(|_| posture) + } } // --------------------------------------------------------------------------- @@ -131,12 +181,128 @@ pub struct ServerConfig { /// header value are skipped (with a warning) when the router is built. #[serde(default)] pub allowed_origins: Option>, + /// Dedicated listener bind for the admin `/usage` API. Only read when + /// `[auth.admin]` is configured; the default is `127.0.0.1` on the port + /// after `bind`'s port (see [`ServerConfig::effective_admin_bind`]), so + /// the admin surface stays loopback-only unless deliberately pointed at + /// a tunnel interface or removed port. + #[serde(default)] + pub admin_bind: Option, + /// Also serve `/usage` on the main listener behind the same admin token. + /// + /// Opt-in; selects the [`AdminMountPosture::MainListener`] posture, which + /// the serve path announces with a loud startup warning: the admin + /// surface then reaches every network the main listener does. With no + /// explicit `admin_bind` this posture serves `/usage` on the main + /// listener ONLY — no dedicated socket is attempted, which is exactly + /// what setups without a second available port need. + #[serde(default)] + pub admin_on_main_listener: bool, +} + +impl ServerConfig { + /// Whether the dedicated admin listener binds at all. + /// + /// An explicit `admin_bind` always binds one. With none set, the + /// default loopback listener binds unless the main-listener opt-in + /// took over — `admin_on_main_listener = true` without an explicit + /// `admin_bind` is exactly the posture for setups where a second port + /// is not available, so no dedicated socket is attempted there. + #[must_use] + pub fn dedicated_admin_listener(&self) -> bool { + self.admin_bind.is_some() || !self.admin_on_main_listener + } + + /// Bind address for the dedicated admin listener. + /// + /// `admin_bind` when configured; otherwise the default: `127.0.0.1` on + /// the port after `bind`'s port. `None` when that default port cannot be + /// derived (`bind` on port 65535) — `load_app_config` rejects that + /// combination whenever the dedicated listener would bind the default, + /// so validated configs that consult this while + /// [`ServerConfig::dedicated_admin_listener`] is true always resolve to + /// an address. + #[must_use] + pub fn effective_admin_bind(&self) -> Option { + match self.admin_bind { + Some(addr) => Some(addr), + None => self + .bind + .port() + .checked_add(1) + .map(|port| SocketAddr::new(std::net::IpAddr::from([127, 0, 0, 1]), port)), + } + } } const fn default_rate_limit_rpm() -> u32 { 100 } +// --------------------------------------------------------------------------- +// AuthConfig +// --------------------------------------------------------------------------- + +/// Client authentication mode for proxy routes. +/// +/// Serialized in TOML as `"none"` / `"client-keys"`. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "kebab-case")] +pub enum AuthMode { + /// No client authentication: every request is attributed to the + /// `"anonymous"` caller. Only permitted while the server binds loopback; + /// `load_app_config` rejects a non-loopback bind in this mode. + #[default] + None, + /// Requests to proxy routes must present + /// `Authorization: Bearer ` matching a secret hashed under + /// `[auth.clients]`. + ClientKeys, +} + +/// Client authentication configuration (`[auth]` section of `config.toml`). +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct AuthConfig { + /// Authentication mode. Defaults to [`AuthMode::None`]. + #[serde(default)] + pub mode: AuthMode, + /// Key id to hashed client secret, under `[auth.clients]`. + /// + /// Each value is a PHC-format password-hashing string (argon2id), never a + /// plaintext secret. Wrapped in [`SecretString`] so the value is redacted + /// in `Debug` output. + /// + /// `skip_serializing` (mirroring `ProviderConfig::api_key`) prevents the + /// hashes from leaking through any `Serialize` path; deserializing a + /// config that was serialized will therefore fail on this field, which is + /// intentional. + #[serde(default, skip_serializing)] + pub clients: HashMap, + /// Admin API authentication (`[auth.admin]` section). When present, the + /// read-only `/usage` query surface is enabled and every `/usage` route + /// on every mount requires its bearer token. + #[serde(default)] + pub admin: Option, +} + +/// Admin API configuration (`[auth.admin]` section of `config.toml`). +/// +/// A single bearer token guarding every `/usage` route. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct AdminAuthConfig { + /// The admin bearer token, presented as `Authorization: Bearer `. + /// + /// Compared in constant time against a SHA-256 digest of the configured + /// token, so the plaintext never needs to be retained in memory. Wrapped + /// in [`SecretString`] so it is redacted in `Debug` output; + /// `skip_serializing` (mirroring `ProviderConfig::api_key`) keeps it out + /// of every `Serialize` path. + #[serde(skip_serializing)] + pub token: SecretString, +} + /// Default graceful-shutdown hard deadline (30 seconds). /// /// Chosen to comfortably exceed the maximum upstream request timeout while @@ -149,6 +315,88 @@ const fn default_dedup_window() -> Duration { Duration::ZERO } +// --------------------------------------------------------------------------- +// StorageConfig +// --------------------------------------------------------------------------- + +/// Default retention-pruning interval (1 hour). +const fn default_prune_interval() -> Duration { + Duration::from_secs(3600) +} + +/// Usage-ledger persistence configuration (`[storage]` section of +/// `config.toml`). +/// +/// When present, every proxied request is persisted to a SQLite ledger by the +/// storage crate's event bus. When absent, the in-memory no-op bus is used and +/// nothing is written to disk. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct StorageConfig { + /// Path to the SQLite database file. Parent directories are created on + /// startup. + pub path: PathBuf, + /// Delete ledger rows older than this age. Absent means rows are kept + /// forever. Must be positive when present. + #[serde(default, with = "humantime_serde")] + pub max_age: Option, + /// Keep at most this many ledger rows; the oldest rows are deleted when + /// the cap is exceeded. Absent means unlimited. Must be at least 1 when + /// present. + #[serde(default)] + pub max_rows: Option, + /// How often the retention pruning pass runs. Defaults to 1 hour. Must be + /// at least 1 second (the pruning timer rejects a zero period). + #[serde(default = "default_prune_interval", with = "humantime_serde")] + pub prune_interval: Duration, +} + +/// Validate `[storage]`. +/// +/// Rules: +/// - `path` must not be empty. +/// - `max_age`, when set, must be positive: a zero age would delete every row +/// on every pass and indicates a unit mistake rather than intent. +/// - `max_rows`, when set, must be at least 1 (zero would keep nothing). +/// - `prune_interval` must be at least 1 second; the pruning timer in the +/// storage crate panics on a zero period, and a sub-second prune cadence +/// serves no purpose. +fn validate_storage_config(storage: &StorageConfig) -> Result<(), CoreError> { + if storage.path.as_os_str().is_empty() { + return Err(CoreError::ConfigValidation { + message: "storage.path must not be empty".to_owned(), + source: None, + }); + } + if storage.max_age == Some(Duration::ZERO) { + return Err(CoreError::ConfigValidation { + message: "storage.max_age must be a positive duration; a zero age would \ + delete every ledger row on every pruning pass" + .to_owned(), + source: None, + }); + } + if storage.max_rows == Some(0) { + return Err(CoreError::ConfigValidation { + message: "storage.max_rows must be at least 1; a cap of zero would keep \ + no rows at all" + .to_owned(), + source: None, + }); + } + const MIN_PRUNE_INTERVAL: Duration = Duration::from_secs(1); + if storage.prune_interval < MIN_PRUNE_INTERVAL { + return Err(CoreError::ConfigValidation { + message: format!( + "storage.prune_interval must be at least 1 second, but is {:?}", + storage.prune_interval + ), + source: None, + }); + } + Ok(()) +} + // --------------------------------------------------------------------------- // ProviderFile // --------------------------------------------------------------------------- @@ -224,6 +472,59 @@ pub struct ProviderConfig { /// the key is the upstream model id. Defaults to empty (no cost accounting). #[serde(default)] pub pricing: HashMap, + /// Upstream egress (network exit) configuration. When present, this + /// provider's upstream requests are sent through the configured forward + /// proxy instead of a direct connection. Absent means direct egress, + /// exactly as before the section existed. + #[serde(default)] + pub egress: Option, +} + +// --------------------------------------------------------------------------- +// ProviderEgressConfig +// --------------------------------------------------------------------------- + +/// Per-provider forward-proxy egress configuration (`[provider.egress]`). +/// +/// The proxy URL scheme must be one of `socks5h://`, `socks5://`, `http://`, +/// or `https://`; anything else is rejected at config load, never silently +/// ignored. `socks5h://` resolves upstream hostnames at the proxy (the exit +/// point's DNS view); `socks5://` resolves them locally. +/// +/// The URL may carry userinfo (`http://user:pass@host:port`), which is passed +/// to the transport as proxy authentication. [`Debug`] redacts the userinfo; +/// only `scheme://host:port` is ever printed. +#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ProviderEgressConfig { + /// Forward proxy URL (`socks5h://host:port`, `socks5://`, `http://`, or + /// `https://`). Supports `${ENV_VAR}` interpolation like other string + /// fields. The port is optional and defaults per scheme (1080 for socks). + pub proxy: String, +} + +impl std::fmt::Debug for ProviderEgressConfig { + // SECURITY: the raw URL may carry proxy credentials as userinfo; print + // only the credential-free form. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ProviderEgressConfig") + .field("proxy", &redacted_proxy_url(&self.proxy)) + .finish() + } +} + +/// Render `scheme://host:port` from a proxy URL, dropping any userinfo. +/// +/// URLs without a `://` separator print as `` rather than echoing a +/// string that has not been through validation. +fn redacted_proxy_url(url: &str) -> String { + let Some((scheme, rest)) = url.split_once("://") else { + return "".to_owned(); + }; + let authority = rest.split(['/', '?', '#']).next().unwrap_or_default(); + let host = authority.rsplit_once('@').map_or(authority, |(_, h)| h); + let after_authority = rest.strip_prefix(authority).unwrap_or_default(); + format!("{scheme}://{host}{after_authority}") } // --------------------------------------------------------------------------- @@ -738,6 +1039,18 @@ pub enum ConfigValidationError { /// The endpoint URL that has an invalid scheme. endpoint: String, }, + /// The `[provider.egress]` proxy URL is not usable as a forward proxy. + #[error( + "provider \"{provider}\": [provider.egress] proxy URL is invalid ({reason}): \"{proxy}\"" + )] + InvalidEgressProxy { + /// Provider name. + provider: String, + /// The rejected proxy URL. + proxy: String, + /// What is wrong with it. + reason: &'static str, + }, /// An adapter header name or value contains CR or LF characters. #[error("provider \"{provider}\": adapter \"{adapter}\" header contains CR/LF characters")] HeaderContainsCrlf { @@ -880,6 +1193,7 @@ pub enum ConfigValidationError { /// api_key: SecretString::from("secret"), /// auth_style: AuthStyle::Bearer, /// passthrough_auth: false, +/// egress: None, /// adapters: HashMap::from([( /// "chat".to_owned(), /// ProviderAdapterConfig { @@ -912,6 +1226,7 @@ pub enum ConfigValidationError { /// api_key: SecretString::from("secret"), /// auth_style: AuthStyle::Bearer, /// passthrough_auth: false, +/// egress: None, /// adapters: HashMap::new(), /// routes: ProviderRoutesConfig::default(), /// model_aliases: HashMap::new(), @@ -935,6 +1250,7 @@ pub fn validate_provider_config( validate_pricing(provider, name)?; validate_adapters(provider, name, known_protocols)?; validate_discovery(provider, name)?; + validate_egress(provider, name)?; // Both the catalog and the model aliases reject unsafe model IDs, but only // when an adapter interpolates `{model}` into a Gemini endpoint URL. @@ -1247,6 +1563,70 @@ fn scheme_region(endpoint: &str) -> &str { } } +/// Egress proxy schemes accepted in `[provider.egress]`. +/// +/// `socks5h` resolves hostnames at the proxy; `socks5` resolves locally. +/// SOCKS support is the reason the provider crate enables reqwest's `socks` +/// cargo feature. Schemes outside this list are rejected at config load +/// rather than silently ignored. +const EGRESS_PROXY_SCHEMES: [&str; 4] = ["socks5h", "socks5", "http", "https"]; + +/// An egress proxy URL, when configured, must name an accepted scheme and a +/// non-empty host. +/// +/// reqwest itself is deliberately lenient (it accepts scheme-less URLs by +/// prefixing `http://`, and `socks4`/unknown schemes without complaint), so +/// this check is the real gate: anything outside the accepted schemes is a +/// config error, never a silently different proxy. +fn validate_egress(provider: &ProviderConfig, name: &str) -> Result<(), ConfigValidationError> { + let Some(egress) = &provider.egress else { + return Ok(()); + }; + let url = egress.proxy.trim(); + if url.is_empty() { + return Err(ConfigValidationError::InvalidEgressProxy { + provider: name.to_owned(), + proxy: egress.proxy.clone(), + reason: "empty proxy URL", + }); + } + let scheme = scheme_region(url); + if scheme.is_empty() { + return Err(ConfigValidationError::InvalidEgressProxy { + provider: name.to_owned(), + proxy: egress.proxy.clone(), + reason: "no URL scheme; expected socks5h://, socks5://, http://, or https://", + }); + } + if !EGRESS_PROXY_SCHEMES + .iter() + .any(|accepted| scheme.eq_ignore_ascii_case(accepted)) + { + return Err(ConfigValidationError::InvalidEgressProxy { + provider: name.to_owned(), + proxy: egress.proxy.clone(), + reason: "unsupported scheme; expected socks5h://, socks5://, http://, or https://", + }); + } + // Everything after the scheme must hold a host. Cut the authority region + // at the first path/query/fragment separator, then drop any userinfo + // (`user:pass@`) so only the host portion is checked. + let authority = url + .strip_prefix(scheme) + .and_then(|after_scheme| after_scheme.strip_prefix("://")) + .unwrap_or_default(); + let authority = authority.split(['/', '?', '#']).next().unwrap_or_default(); + let host = authority.rsplit_once('@').map_or(authority, |(_, h)| h); + if host.trim().is_empty() { + return Err(ConfigValidationError::InvalidEgressProxy { + provider: name.to_owned(), + proxy: egress.proxy.clone(), + reason: "no proxy host", + }); + } + Ok(()) +} + /// Validate a map of static headers for CRLF injection and forbidden header names. /// /// Shared between adapter header validation and discovery header validation @@ -1398,9 +1778,153 @@ pub fn load_app_config(path: impl AsRef) -> Result { source: None, }); } + validate_auth_config(&cfg.auth, cfg.server.bind)?; + if let Some(storage) = cfg.storage.as_ref() { + validate_storage_config(storage)?; + } + validate_admin_mount(&cfg)?; Ok(cfg) } +/// Validate the admin `/usage` mount against the rest of the config. +/// +/// Rules: +/// - `[auth.admin]` requires `[storage]`: without the usage ledger there is +/// nothing to query, so serving `/usage` would only produce errors. +/// - The admin token must not be empty (an empty token can never match a +/// bearer token, so the section would be dead config hiding a mistake). +/// - `server.admin_bind` must differ from `server.bind`: two listeners cannot +/// share one socket. +/// - Without `[auth.admin]`, `server.admin_bind` and +/// `server.admin_on_main_listener` are dead config and are rejected. +/// - The derived default admin port (`bind` port + 1) must exist: a `bind` +/// on port 65535 with no explicit `admin_bind` has no distinct port to +/// default to. +fn validate_admin_mount(cfg: &AppConfig) -> Result<(), CoreError> { + let Some(admin) = cfg.auth.admin.as_ref() else { + if cfg.server.admin_bind.is_some() { + return Err(CoreError::ConfigValidation { + message: "server.admin_bind is configured but [auth.admin] is not; \ + set [auth.admin] token or remove server.admin_bind" + .to_owned(), + source: None, + }); + } + if cfg.server.admin_on_main_listener { + return Err(CoreError::ConfigValidation { + message: "server.admin_on_main_listener = true but [auth.admin] is not \ + configured; there would be no admin token to guard /usage" + .to_owned(), + source: None, + }); + } + return Ok(()); + }; + if admin.token.expose_secret().is_empty() { + return Err(CoreError::ConfigValidation { + message: "[auth.admin] token must not be empty".to_owned(), + source: None, + }); + } + if cfg.storage.is_none() { + return Err(CoreError::ConfigValidation { + message: "[auth.admin] requires the [storage] usage ledger; configure \ + [storage] or remove [auth.admin]" + .to_owned(), + source: None, + }); + } + if let Some(admin_bind) = cfg.server.admin_bind { + if admin_bind == cfg.server.bind { + return Err(CoreError::ConfigValidation { + message: format!( + "server.admin_bind {admin_bind} must differ from server.bind; \ + the admin listener needs its own socket" + ), + source: None, + }); + } + } else if !cfg.server.admin_on_main_listener && cfg.server.bind.port() == u16::MAX { + // The default admin port (bind port + 1) is only consulted when the + // dedicated listener binds it; under the main-listener opt-in with + // no explicit admin_bind, no second socket is attempted. + return Err(CoreError::ConfigValidation { + message: "server.bind port 65535 leaves no distinct port for the default \ + admin listener; set server.admin_bind explicitly or enable \ + server.admin_on_main_listener" + .to_owned(), + source: None, + }); + } + Ok(()) +} + +/// Reserved caller name used for unauthenticated requests under +/// [`AuthMode::None`]; a client key id must not collide with it or event +/// attribution becomes ambiguous. +pub const ANONYMOUS_CALLER: &str = "anonymous"; + +/// Validate `[auth]` against itself and against the bind address. +/// +/// Rules: +/// - `mode = "none"` with a non-empty `[auth.clients]` is contradictory +/// (configured keys would be silently ignored) and is rejected. +/// - `mode = "client-keys"` requires at least one `[auth.clients]` entry. +/// - `mode = "none"` is only permitted on a loopback bind: a non-loopback +/// bind with no client keys would expose the proxy to every reachable +/// peer with no authentication, so startup refuses. +/// - Key ids must be non-empty and must not be [`ANONYMOUS_CALLER`]. +fn validate_auth_config(auth: &AuthConfig, bind: SocketAddr) -> Result<(), CoreError> { + let clients_empty = auth.clients.is_empty(); + match auth.mode { + AuthMode::None if !clients_empty => { + return Err(CoreError::ConfigValidation { + message: "auth.mode = \"none\" but [auth.clients] is non-empty; \ + set auth.mode = \"client-keys\" to enforce the configured keys" + .to_owned(), + source: None, + }); + } + AuthMode::ClientKeys if clients_empty => { + return Err(CoreError::ConfigValidation { + message: "auth.mode = \"client-keys\" requires at least one entry \ + under [auth.clients]" + .to_owned(), + source: None, + }); + } + AuthMode::None if !bind.ip().is_loopback() => { + return Err(CoreError::ConfigValidation { + message: format!( + "server.bind {bind} is not loopback and no client keys are \ + configured; anonymous auth is only permitted on loopback binds. \ + Configure [auth.clients] and set auth.mode = \"client-keys\"" + ), + source: None, + }); + } + AuthMode::None | AuthMode::ClientKeys => {} + } + for key_id in auth.clients.keys() { + if key_id.is_empty() { + return Err(CoreError::ConfigValidation { + message: "[auth.clients] contains an empty key id".to_owned(), + source: None, + }); + } + if key_id == ANONYMOUS_CALLER { + return Err(CoreError::ConfigValidation { + message: format!( + "[auth.clients] key id \"{ANONYMOUS_CALLER}\" is reserved for \ + unattributed requests; choose a different key id" + ), + source: None, + }); + } + } + Ok(()) +} + /// Load and parse a provider [`ProviderFile`] TOML file. /// /// The raw file content is first processed for `${ENV_VAR}` interpolation, @@ -1652,6 +2176,183 @@ server_name = "test" ); } + /// Write a config file with the given `bind` and `[auth]` body into a + /// fresh temp dir, returning the dir (the config lives at + /// `config.toml` inside it). The returned guard keeps the file alive; + /// callers run [`load_app_config`] on it directly. + fn config_dir_with_auth(bind: &str, auth_body: &str) -> tempfile::TempDir { + let directory = tempfile::tempdir().expect("tempdir"); + let path = directory.path().join("config.toml"); + std::fs::write( + &path, + format!( + r#" +[server] +bind = "{bind}" +request_timeout = "300s" +log_level = "info" +hot_reload = false +server_name = "test" + +{auth_body} +"# + ), + ) + .expect("write config file"); + directory + } + + #[test] + fn auth_absent_section_defaults_to_none_on_loopback() { + let dir = config_dir_with_auth("127.0.0.1:3456", ""); + let config = + load_app_config(dir.path().join("config.toml")).expect("default auth accepted"); + assert_eq!(config.auth.mode, AuthMode::None); + assert!(config.auth.clients.is_empty()); + } + + #[test] + fn auth_none_on_non_loopback_bind_is_rejected() { + // The core safety rule of client identity: an anonymous proxy exposed + // past loopback must refuse to start, before any listener binds. + for bind in ["0.0.0.0:3456", "192.168.1.4:3456", "[::]:3456"] { + let dir = config_dir_with_auth(bind, ""); + let error = load_app_config(dir.path().join("config.toml")) + .expect_err("non-loopback none must be rejected"); + assert!( + error.to_string().contains("not loopback"), + "bind {bind}: got: {error}" + ); + } + } + + #[test] + fn auth_none_on_ipv6_loopback_bind_is_accepted() { + // `::1` is loopback; the loopback exemption is about the address, not + // the address family. + let dir = config_dir_with_auth("[::1]:3456", ""); + load_app_config(dir.path().join("config.toml")).expect("ipv6 loopback accepted"); + } + + #[test] + fn auth_none_with_configured_clients_is_rejected() { + let dir = config_dir_with_auth( + "127.0.0.1:3456", + r#" +[auth] +mode = "none" + +[auth.clients] +laptop = "$argon2id$v=19$m=19456,t=2,p=1$c29tZXNhbHQ$hash" +"#, + ); + let error = load_app_config(dir.path().join("config.toml")) + .expect_err("none mode with clients is contradictory"); + assert!( + error.to_string().contains("[auth.clients] is non-empty"), + "got: {error}" + ); + } + + #[test] + fn auth_client_keys_without_clients_is_rejected() { + let dir = config_dir_with_auth( + "0.0.0.0:3456", + r#" +[auth] +mode = "client-keys" +"#, + ); + let error = load_app_config(dir.path().join("config.toml")) + .expect_err("client-keys mode needs at least one key"); + assert!( + error + .to_string() + .contains("requires at least one entry under [auth.clients]"), + "got: {error}" + ); + } + + #[test] + fn auth_client_keys_with_a_client_accepts_any_bind() { + // With at least one hashed key configured, a non-loopback bind is a + // legitimate deployment and must load. + let dir = config_dir_with_auth( + "0.0.0.0:3456", + r#" +[auth] +mode = "client-keys" + +[auth.clients] +laptop = "$argon2id$v=19$m=19456,t=2,p=1$c29tZXNhbHQ$hash" +"#, + ); + let config = load_app_config(dir.path().join("config.toml")) + .expect("client-keys on non-loopback accepted"); + assert_eq!(config.auth.mode, AuthMode::ClientKeys); + assert_eq!(config.auth.clients.len(), 1); + } + + #[test] + fn auth_rejects_reserved_and_empty_key_ids() { + let reserved_dir = config_dir_with_auth( + "127.0.0.1:3456", + r#" +[auth] +mode = "client-keys" + +[auth.clients] +anonymous = "$argon2id$v=19$m=19456,t=2,p=1$c29tZXNhbHQ$hash" +"#, + ); + let reserved = load_app_config(reserved_dir.path().join("config.toml")) + .expect_err("reserved key id must be rejected"); + assert!(reserved.to_string().contains("reserved"), "got: {reserved}"); + + let empty_dir = config_dir_with_auth( + "127.0.0.1:3456", + r#" +[auth] +mode = "client-keys" + +[auth.clients] +"" = "$argon2id$v=19$m=19456,t=2,p=1$c29tZXNhbHQ$hash" +"#, + ); + let empty = load_app_config(empty_dir.path().join("config.toml")) + .expect_err("empty key id must be rejected"); + assert!(empty.to_string().contains("empty key id"), "got: {empty}"); + } + + #[test] + fn auth_clients_do_not_leak_through_debug() { + // The hash (let alone a plaintext secret) must never surface in Debug + // output of the config, matching the api_key discipline. + let config: AppConfig = toml::from_str( + r#" +[server] +bind = "127.0.0.1:3456" +request_timeout = "300s" +log_level = "info" +hot_reload = false +server_name = "test" + +[auth] +mode = "client-keys" + +[auth.clients] +laptop = "$argon2id$v=19$m=19456,t=2,p=1$c29tZXNhbHQ$hashvalue" +"#, + ) + .expect("app config"); + let debug = format!("{config:?}"); + assert!( + !debug.contains("hashvalue"), + "client hash must be redacted in Debug, got: {debug}" + ); + assert!(debug.contains("[REDACTED]"), "got: {debug}"); + } + #[test] fn has_valid_http_scheme_is_case_insensitive() { // Schemes are case-insensitive per RFC 3986 and reqwest accepts mixed @@ -1684,18 +2385,204 @@ chat_completion = "chat" result.unwrap_err(); } + // ----------------------------------------------------------------------- + // [provider.egress] + // ----------------------------------------------------------------------- + #[test] - fn provider_debug_redacts_header_values_and_endpoint_query() { - let config: ProviderFile = toml::from_str( + fn egress_section_parses_and_absent_section_means_none() { + let with_egress: ProviderFile = toml::from_str( r#" [provider] name = "example" -api_key = "secret-api-key" +api_key = "key" auth_style = "bearer" -[provider.adapters.chat] -protocol = "openai_chat_completions" -endpoint = "https://example.com/v1/chat/completions?key=adapter-query-secret" +[provider.egress] +proxy = "socks5h://127.0.0.1:1080" +"#, + ) + .expect("provider with egress parses"); + assert_eq!( + with_egress.provider.egress, + Some(ProviderEgressConfig { + proxy: "socks5h://127.0.0.1:1080".to_owned(), + }) + ); + + let without_egress: ProviderFile = toml::from_str( + r#" +[provider] +name = "example" +api_key = "key" +auth_style = "bearer" +"#, + ) + .expect("provider without egress parses"); + assert_eq!(without_egress.provider.egress, None); + } + + #[test] + fn egress_section_rejects_unknown_fields() { + let result = toml::from_str::( + r#" +[provider] +name = "example" +api_key = "key" +auth_style = "bearer" + +[provider.egress] +proxy = "socks5h://127.0.0.1:1080" +bind_out = "198.51.100.10" +"#, + ); + result.unwrap_err(); + } + + #[test] + fn egress_debug_redacts_proxy_userinfo_but_keeps_the_host() { + let egress = ProviderEgressConfig { + proxy: "http://proxy-user:proxy-pass@proxy.internal:3128".to_owned(), + }; + let debug = format!("{egress:?}"); + assert!( + !debug.contains("proxy-user") && !debug.contains("proxy-pass"), + "proxy credentials must never print: {debug}" + ); + assert!( + debug.contains("proxy.internal:3128"), + "the credential-free host must stay visible for diagnostics: {debug}" + ); + // No-scheme strings (possible only pre-validation) print nothing raw. + let raw = ProviderEgressConfig { + proxy: "user:pass@host".to_owned(), + }; + assert!( + !format!("{raw:?}").contains("user:pass"), + "unvalidated URLs must not be echoed: {:?}", + raw + ); + } + + #[test] + fn validate_accepts_each_supported_egress_proxy_scheme() { + for url in [ + "socks5h://tunnel.example.com:1080", + "socks5://127.0.0.1:1080", + "http://proxy.internal:3128", + "https://proxy.internal:8443", + // Schemes are case-insensitive (RFC 3986), like endpoint schemes. + "SOCKS5H://tunnel.example.com:1080", + // Userinfo is proxy authentication, not part of the host check. + "http://user:pass@proxy.internal:3128", + // Portless socks: reqwest defaults to 1080. + "socks5h://tunnel.example.com", + ] { + let file: ProviderFile = toml::from_str(&format!( + r#" +[provider] +name = "example" +api_key = "key" +auth_style = "bearer" + +[provider.egress] +proxy = "{url}" +"# + )) + .expect("parses"); + assert_eq!( + validate_provider_config(&file.provider, None), + Ok(()), + "egress proxy {url} must validate" + ); + } + } + + #[test] + fn validate_rejects_unusable_egress_proxy_urls() { + // (url, reason fragment) pairs. Scheme-less URLs matter most: reqwest + // would silently prefix `http://` at client build, so config load is + // the only gate that can refuse them. + for (url, reason_fragment) in [ + ("ftp://proxy:21", "unsupported scheme"), + ("socks4://127.0.0.1:1080", "unsupported scheme"), + ("127.0.0.1:1080", "no URL scheme"), + ("proxy.internal:3128", "no URL scheme"), + ("socks5h://", "no proxy host"), + ("http:///path", "no proxy host"), + ("", "empty proxy URL"), + (" ", "empty proxy URL"), + ] { + let file: ProviderFile = toml::from_str(&format!( + r#" +[provider] +name = "example" +api_key = "key" +auth_style = "bearer" + +[provider.egress] +proxy = "{url}" +"# + )) + .expect("parses"); + let err = validate_provider_config(&file.provider, None) + .expect_err("egress proxy must be rejected"); + match &err { + ConfigValidationError::InvalidEgressProxy { + provider, reason, .. + } => { + assert_eq!(provider, "example"); + assert!( + reason.contains(reason_fragment), + "reason {reason:?} must name {reason_fragment:?} for {url:?}" + ); + } + other => panic!("expected InvalidEgressProxy for {url:?}, got {other:?}"), + } + // The operator-facing message names the provider and the URL. + let message = err.to_string(); + assert!(message.contains("example")); + assert!(message.contains("[provider.egress]")); + } + } + + #[test] + fn load_provider_config_rejects_invalid_egress_at_load_time() { + let dir = tempfile::tempdir().expect("temp dir"); + let path = dir.path().join("bad-egress.toml"); + std::fs::write( + &path, + r#" +[provider] +name = "example" +api_key = "key" +auth_style = "bearer" + +[provider.egress] +proxy = "ftp://proxy.example.com:21" +"#, + ) + .expect("write provider toml"); + + let err = load_provider_config(&path, None).expect_err("load must reject bad egress"); + assert!( + err.to_string().contains("egress"), + "error must point at the egress section: {err}" + ); + } + + #[test] + fn provider_debug_redacts_header_values_and_endpoint_query() { + let config: ProviderFile = toml::from_str( + r#" +[provider] +name = "example" +api_key = "secret-api-key" +auth_style = "bearer" + +[provider.adapters.chat] +protocol = "openai_chat_completions" +endpoint = "https://example.com/v1/chat/completions?key=adapter-query-secret" [provider.adapters.chat.headers] x-secret = "adapter-secret" @@ -2001,6 +2888,7 @@ endpoint = "https://example.com/v1" api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::from([( "chat".to_owned(), ProviderAdapterConfig { @@ -2028,6 +2916,7 @@ endpoint = "https://example.com/v1" api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::from([( "chat".to_owned(), ProviderAdapterConfig { @@ -2055,6 +2944,7 @@ endpoint = "https://example.com/v1" api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::from([( "chat".to_owned(), ProviderAdapterConfig { @@ -2082,6 +2972,7 @@ endpoint = "https://example.com/v1" api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -2102,6 +2993,7 @@ endpoint = "https://example.com/v1" api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -2167,6 +3059,7 @@ input = 0.0000025 api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -2193,6 +3086,7 @@ input = 0.0000025 api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -2217,6 +3111,7 @@ input = 0.0000025 api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -2251,6 +3146,7 @@ input = 0.0000025 api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -2297,6 +3193,7 @@ input = 0.0000025 api_key: SecretString::from(""), auth_style: AuthStyle::Bearer, passthrough_auth: true, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -2318,6 +3215,7 @@ input = 0.0000025 api_key: SecretString::from(""), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -2330,4 +3228,380 @@ input = 0.0000025 Err(ConfigValidationError::EmptyApiKey { .. }) )); } + + // -- [storage] ----------------------------------------------------------- + + /// Minimal valid main config with a `[storage]` section appended. + fn write_app_config(extra_storage: &str) -> (tempfile::TempDir, std::path::PathBuf) { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("config.toml"); + std::fs::write( + &path, + format!( + r#" +[server] +bind = "127.0.0.1:3456" +request_timeout = "300s" +log_level = "info" +hot_reload = false +server_name = "test" +{extra_storage}"# + ), + ) + .unwrap(); + (directory, path) + } + + #[test] + fn storage_section_defaults_to_none_when_absent() { + let (_dir, path) = write_app_config(""); + let config = load_app_config(&path).expect("valid config without [storage]"); + assert!( + config.storage.is_none(), + "absent [storage] must keep the no-op event bus" + ); + } + + #[test] + fn storage_section_parses_and_applies_defaults() { + let (_dir, path) = write_app_config( + r#" +[storage] +path = "/tmp/usage.db" +"#, + ); + let config = load_app_config(&path).expect("valid [storage]"); + let storage = config.storage.expect("storage present"); + assert_eq!(storage.path, std::path::PathBuf::from("/tmp/usage.db")); + assert_eq!(storage.max_age, None, "absent max_age keeps rows forever"); + assert_eq!(storage.max_rows, None, "absent max_rows is unlimited"); + assert_eq!( + storage.prune_interval, + Duration::from_secs(3600), + "absent prune_interval defaults to 1h" + ); + } + + #[test] + fn storage_section_parses_full_retention() { + let (_dir, path) = write_app_config( + r#" +[storage] +path = "/tmp/usage.db" +max_age = "30d" +max_rows = 500_000 +prune_interval = "15m" +"#, + ); + let config = load_app_config(&path).expect("valid [storage]"); + let storage = config.storage.expect("storage present"); + assert_eq!(storage.max_age, Some(Duration::from_secs(30 * 24 * 3600))); + assert_eq!(storage.max_rows, Some(500_000)); + assert_eq!(storage.prune_interval, Duration::from_secs(15 * 60)); + } + + #[test] + fn storage_zero_prune_interval_is_rejected() { + let (_dir, path) = write_app_config( + r#" +[storage] +path = "/tmp/usage.db" +prune_interval = "0s" +"#, + ); + let error = load_app_config(&path).expect_err("zero prune interval rejected"); + assert!( + error.to_string().contains("storage.prune_interval"), + "got: {error}" + ); + } + + #[test] + fn storage_subsecond_prune_interval_is_rejected() { + let (_dir, path) = write_app_config( + r#" +[storage] +path = "/tmp/usage.db" +max_age = "1d" +prune_interval = "500ms" +"#, + ); + assert!( + load_app_config(&path).is_err(), + "sub-second prune interval must be rejected (timer panics on zero/near-zero)" + ); + } + + #[test] + fn storage_zero_max_age_is_rejected() { + let (_dir, path) = write_app_config( + r#" +[storage] +path = "/tmp/usage.db" +max_age = "0s" +"#, + ); + let error = load_app_config(&path).expect_err("zero max_age rejected"); + assert!( + error.to_string().contains("storage.max_age"), + "got: {error}" + ); + } + + #[test] + fn storage_zero_max_rows_is_rejected() { + let (_dir, path) = write_app_config( + r#" +[storage] +path = "/tmp/usage.db" +max_rows = 0 +"#, + ); + let error = load_app_config(&path).expect_err("zero max_rows rejected"); + assert!( + error.to_string().contains("storage.max_rows"), + "got: {error}" + ); + } + + #[test] + fn storage_empty_path_is_rejected() { + let (_dir, path) = write_app_config( + r#" +[storage] +path = "" +"#, + ); + let error = load_app_config(&path).expect_err("empty path rejected"); + assert!(error.to_string().contains("storage.path"), "got: {error}"); + } + + #[test] + fn storage_unknown_field_is_rejected() { + let (_dir, path) = write_app_config( + r#" +[storage] +path = "/tmp/usage.db" +db_host = "postgres.example" +"#, + ); + assert!( + load_app_config(&path).is_err(), + "deny_unknown_fields must reject unrecognized [storage] keys" + ); + } + + // -- [auth.admin] / admin mount ------------------------------------------- + + /// Minimal valid `[storage]` section so admin tests vary only the admin + /// fields. + const STORAGE_SECTION: &str = "\n[storage]\npath = \"/tmp/usage.db\"\n"; + + #[test] + fn admin_section_parses_token_and_defaults_off() { + let (_dir, path) = write_app_config(&format!( + "{STORAGE_SECTION}\n[auth.admin]\ntoken = \"admin-secret\"\n" + )); + let config = load_app_config(&path).expect("valid [auth.admin]"); + assert!( + config + .auth + .admin + .as_ref() + .is_some_and(|a| a.token.expose_secret() == "admin-secret"), + "admin token must parse into the SecretString" + ); + assert!( + !config.server.admin_on_main_listener, + "main-listener mount must be opt-in" + ); + assert_eq!( + config.server.effective_admin_bind(), + Some("127.0.0.1:3457".parse().unwrap()), + "default admin bind is loopback on the port after the main bind" + ); + assert_eq!( + config.admin_posture(), + Some(AdminMountPosture::DedicatedListener), + "the default posture is the dedicated listener" + ); + } + + #[test] + fn admin_explicit_bind_and_main_listener_opt_in_parse() { + let (_dir, path) = write_app_config(&format!( + "\nadmin_bind = \"127.0.0.1:9900\"\nadmin_on_main_listener = true\n{STORAGE_SECTION}\n\ + [auth.admin]\ntoken = \"admin-secret\"\n" + )); + let config = load_app_config(&path).expect("valid admin config"); + assert_eq!( + config.server.effective_admin_bind(), + Some("127.0.0.1:9900".parse().unwrap()), + "explicit admin_bind wins over the derived default" + ); + assert_eq!( + config.admin_posture(), + Some(AdminMountPosture::MainListener), + "admin_on_main_listener selects the main-listener posture" + ); + assert_eq!( + AdminMountPosture::MainListener.name(), + "admin-on-main-listener", + "posture names are stable for log grepping" + ); + } + + #[test] + fn admin_without_storage_is_rejected() { + let (_dir, path) = write_app_config("\n[auth.admin]\ntoken = \"admin-secret\"\n"); + let error = load_app_config(&path).expect_err("admin requires [storage]"); + assert!( + error.to_string().contains("[auth.admin] requires"), + "got: {error}" + ); + } + + #[test] + fn admin_empty_token_is_rejected() { + let (_dir, path) = + write_app_config(&format!("{STORAGE_SECTION}\n[auth.admin]\ntoken = \"\"\n")); + let error = load_app_config(&path).expect_err("empty admin token rejected"); + assert!( + error.to_string().contains("token must not be empty"), + "got: {error}" + ); + } + + #[test] + fn admin_bind_equal_to_main_bind_is_rejected() { + let (_dir, path) = write_app_config(&format!( + "\nadmin_bind = \"127.0.0.1:3456\"\n{STORAGE_SECTION}\n[auth.admin]\n\ + token = \"admin-secret\"\n" + )); + let error = load_app_config(&path).expect_err("same-socket bind rejected"); + assert!( + error.to_string().contains("must differ from server.bind"), + "got: {error}" + ); + } + + #[test] + fn admin_default_port_underviable_on_65535_is_rejected() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("config.toml"); + std::fs::write( + &path, + format!( + "[server]\nbind = \"127.0.0.1:65535\"\nrequest_timeout = \"300s\"\n\ + log_level = \"info\"\nhot_reload = false\nserver_name = \"t\"\n{STORAGE_SECTION}\n\ + \n[auth.admin]\ntoken = \"admin-secret\"\n" + ), + ) + .unwrap(); + let error = load_app_config(&path).expect_err("no derivable admin port"); + assert!(error.to_string().contains("65535"), "got: {error}"); + } + + #[test] + fn admin_main_listener_optin_needs_no_derivable_port() { + // The "no second port available" posture: with the main-listener + // opt-in and no explicit admin_bind, no dedicated socket is + // attempted, so an underivable default port must not reject the + // config. + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("config.toml"); + std::fs::write( + &path, + format!( + "[server]\nbind = \"127.0.0.1:65535\"\nrequest_timeout = \"300s\"\n\ + log_level = \"info\"\nhot_reload = false\nserver_name = \"t\"\n\ + admin_on_main_listener = true\n{STORAGE_SECTION}\n\ + \n[auth.admin]\ntoken = \"admin-secret\"\n" + ), + ) + .unwrap(); + let config = load_app_config(&path).expect("main-listener posture needs no port"); + assert!( + !config.server.dedicated_admin_listener(), + "the opt-in without an explicit admin_bind must not bind a dedicated socket" + ); + } + + #[test] + fn explicit_admin_bind_keeps_the_dedicated_listener_under_the_optin() { + let (_dir, path) = write_app_config(&format!( + "\nadmin_bind = \"127.0.0.1:9900\"\nadmin_on_main_listener = true\n{STORAGE_SECTION}\n\ + [auth.admin]\ntoken = \"admin-secret\"\n" + )); + let config = load_app_config(&path).expect("both mounts are legitimate"); + assert!( + config.server.dedicated_admin_listener(), + "an explicit admin_bind always binds the dedicated listener" + ); + assert_eq!( + config.admin_posture(), + Some(AdminMountPosture::MainListener), + "and the main-listener mount is on top of it" + ); + } + + #[test] + fn admin_bind_without_admin_section_is_rejected() { + let (_dir, path) = write_app_config(&format!( + "\nadmin_bind = \"127.0.0.1:9900\"\n{STORAGE_SECTION}" + )); + let error = load_app_config(&path).expect_err("dead admin_bind rejected"); + assert!( + error + .to_string() + .contains("server.admin_bind is configured"), + "got: {error}" + ); + } + + #[test] + fn admin_main_listener_flag_without_admin_section_is_rejected() { + let (_dir, path) = write_app_config(&format!( + "\nadmin_on_main_listener = true\n{STORAGE_SECTION}" + )); + let error = load_app_config(&path).expect_err("dead flag rejected"); + assert!( + error.to_string().contains("admin_on_main_listener"), + "got: {error}" + ); + } + + #[test] + fn admin_unknown_field_is_rejected() { + let (_dir, path) = write_app_config(&format!( + "{STORAGE_SECTION}\n[auth.admin]\ntoken = \"admin-secret\"\nrotates = \"weekly\"\n" + )); + assert!( + load_app_config(&path).is_err(), + "deny_unknown_fields must reject unrecognized [auth.admin] keys" + ); + } + + #[test] + fn admin_posture_is_none_without_admin_section() { + let (_dir, path) = write_app_config(STORAGE_SECTION); + let config = load_app_config(&path).expect("valid config"); + assert_eq!( + config.admin_posture(), + None, + "no [auth.admin] means no admin surface anywhere" + ); + } + + #[test] + fn admin_auth_config_debug_redacts_token() { + let admin = AdminAuthConfig { + token: SecretString::from("super-secret-admin-token"), + }; + let debug = format!("{admin:?}"); + assert!( + !debug.contains("super-secret-admin-token"), + "Debug must redact the admin token, got: {debug}" + ); + } } diff --git a/crates/llm-proxy-core/src/provider_registry.rs b/crates/llm-proxy-core/src/provider_registry.rs index 977e0dd..902a5e2 100644 --- a/crates/llm-proxy-core/src/provider_registry.rs +++ b/crates/llm-proxy-core/src/provider_registry.rs @@ -486,6 +486,7 @@ mod tests { api_key: SecretString::from("api-secret"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::from([( "chat".to_owned(), ProviderAdapterConfig { @@ -525,6 +526,7 @@ mod tests { api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::from([( "chat".to_owned(), ProviderAdapterConfig { @@ -590,6 +592,7 @@ mod tests { api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -707,6 +710,7 @@ chat_completions = "chat" api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), // no adapters defined routes: ProviderRoutesConfig { chat_completions: Some("nonexistent_adapter".to_owned()), @@ -739,6 +743,7 @@ chat_completions = "chat" api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -775,6 +780,7 @@ chat_completions = "chat" api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -804,6 +810,7 @@ chat_completions = "chat" api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -823,6 +830,7 @@ chat_completions = "chat" api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -845,6 +853,7 @@ chat_completions = "chat" api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::from([( "chat".to_owned(), ProviderAdapterConfig { @@ -879,6 +888,7 @@ chat_completions = "chat" api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -919,6 +929,7 @@ chat_completions = "chat" api_key: SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters, routes: ProviderRoutesConfig { chat_completions: Some("nonexistent".to_owned()), diff --git a/crates/llm-proxy-protocol/src/client/anthropic.rs b/crates/llm-proxy-protocol/src/client/anthropic.rs index b8fbaba..d6eaeda 100644 --- a/crates/llm-proxy-protocol/src/client/anthropic.rs +++ b/crates/llm-proxy-protocol/src/client/anthropic.rs @@ -785,6 +785,13 @@ impl StreamEncoder { } /// Encode a single [`CoreEvent`] into zero or more Anthropic [`MessageEvent`]s. + #[expect( + clippy::too_many_lines, + reason = "one match arm per CoreEvent variant: the Anthropic Messages SSE \ + wire shape is a single translation table, and splitting the arms \ + into helpers would separate each event's field mapping from the \ + frame layout it produces" + )] pub fn encode_event(&mut self, event: CoreEvent) -> Result, ProtocolError> { if self.finished { tracing::trace!( diff --git a/crates/llm-proxy-protocol/src/client/openai_chat.rs b/crates/llm-proxy-protocol/src/client/openai_chat.rs index 78e82e9..f2f2623 100644 --- a/crates/llm-proxy-protocol/src/client/openai_chat.rs +++ b/crates/llm-proxy-protocol/src/client/openai_chat.rs @@ -91,6 +91,12 @@ fn parse_tool_call_arguments(arguments: &str) -> Option { // --------------------------------------------------------------------------- /// Decode an OpenAI [`ChatCompletionRequest`] into a normalised [`CoreRequest`]. +#[expect( + clippy::too_many_lines, + reason = "field-by-field wire decode of the OpenAI Chat request shape, kept \ + in one place so the mapping can be audited against the API schema \ + without jumping between helpers" +)] pub fn decode_request(req: ChatCompletionRequest) -> Result { if req.model.is_empty() { return Err(ProtocolError::InvalidRequest("model is required".into())); @@ -367,6 +373,12 @@ fn decode_tool_choice(value: serde_json::Value) -> CoreToolChoice { /// separator into a single string for `reasoning_content`. /// - `stop_sequence` is silently dropped since OpenAI has no native field for it. /// - `provider_meta` is silently dropped (for provider adapters only). +#[expect( + clippy::too_many_lines, + reason = "field-by-field wire encode of the OpenAI Chat response shape, kept \ + in one place so the mapping can be audited against the API schema \ + without jumping between helpers" +)] pub fn encode_response(resp: CoreResponse) -> Result { if !resp.provider_meta.is_empty() { tracing::warn!( @@ -625,6 +637,13 @@ impl StreamEncoder { } /// Encode a single [`CoreEvent`] into zero or more [`ChatCompletionChunk`]s. + #[expect( + clippy::too_many_lines, + reason = "one match arm per CoreEvent variant: the OpenAI Chat SSE chunk \ + wire shape is a single translation table, and splitting the arms \ + into helpers would separate each event's field mapping from the \ + chunk layout it produces" + )] pub fn encode_event( &mut self, event: CoreEvent, diff --git a/crates/llm-proxy-provider/Cargo.toml b/crates/llm-proxy-provider/Cargo.toml index f9a8b6a..3a0fb0e 100644 --- a/crates/llm-proxy-provider/Cargo.toml +++ b/crates/llm-proxy-provider/Cargo.toml @@ -18,7 +18,9 @@ llm-proxy-core = { workspace = true } llm-proxy-protocol = { path = "../llm-proxy-protocol" } bytes = "1" form_urlencoded = "1" -reqwest = { version = "0.12", features = ["json", "stream", "rustls-tls"] } +# "socks" enables socks5/socks5h forward-proxy support for [provider.egress] +# (verified against docs.rs/reqwest 0.12 and empirically in transport.rs tests). +reqwest = { version = "0.12", features = ["json", "stream", "rustls-tls", "socks"] } futures = "0.3" tokio = { workspace = true } serde_json = { workspace = true } @@ -30,4 +32,5 @@ secrecy = { workspace = true } [dev-dependencies] axum = { workspace = true, features = ["macros", "json"] } -tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } +# io-util: read_exact/write_all for the in-test SOCKS5 handshake server. +tokio = { workspace = true, features = ["macros", "rt-multi-thread", "io-util"] } diff --git a/crates/llm-proxy-provider/src/adapter/gemini.rs b/crates/llm-proxy-provider/src/adapter/gemini.rs index b2a0382..a7ba2b7 100644 --- a/crates/llm-proxy-provider/src/adapter/gemini.rs +++ b/crates/llm-proxy-provider/src/adapter/gemini.rs @@ -360,6 +360,12 @@ fn encode_tool_result_payload( impl GeminiAdapter { /// Encode a core request into a Gemini GenerateContent request. + #[expect( + clippy::too_many_lines, + reason = "field-by-field wire encode of the Gemini GenerateContent request \ + shape, kept in one place so the mapping can be audited against \ + the API schema without jumping between helpers" + )] pub fn encode_request( &self, core: &CoreRequest, diff --git a/crates/llm-proxy-provider/src/adapter/openai_chat.rs b/crates/llm-proxy-provider/src/adapter/openai_chat.rs index c547c70..1198fc1 100644 --- a/crates/llm-proxy-provider/src/adapter/openai_chat.rs +++ b/crates/llm-proxy-provider/src/adapter/openai_chat.rs @@ -411,6 +411,12 @@ fn push_tool_result_messages(content: &[CoreContent], messages: &mut Vec Result, ProviderError> { let data = frame.data.trim(); if data.is_empty() || data == "[DONE]" { @@ -534,6 +541,12 @@ fn push_message_text(output: &ResponsesOutput, content: &mut Vec) { impl ResponsesAdapter { /// Encode a core request into a Responses API request. + #[expect( + clippy::too_many_lines, + reason = "field-by-field wire encode of the Responses API request shape, \ + kept in one place so the mapping can be audited against the \ + API schema without jumping between helpers" + )] pub fn encode_request( &self, core: &CoreRequest, diff --git a/crates/llm-proxy-provider/src/discovery.rs b/crates/llm-proxy-provider/src/discovery.rs index 7f7f952..d859152 100644 --- a/crates/llm-proxy-provider/src/discovery.rs +++ b/crates/llm-proxy-provider/src/discovery.rs @@ -488,6 +488,7 @@ mod tests { api_key: secrecy::SecretString::from("test-key"), auth_style: AuthStyle::XApiKey, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -738,6 +739,7 @@ mod tests { api_key: secrecy::SecretString::from(api_key), auth_style, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), diff --git a/crates/llm-proxy-provider/src/transport.rs b/crates/llm-proxy-provider/src/transport.rs index 5aa29c4..f49e589 100644 --- a/crates/llm-proxy-provider/src/transport.rs +++ b/crates/llm-proxy-provider/src/transport.rs @@ -140,7 +140,7 @@ pub const DEFAULT_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(60); /// /// Owns a connection-pooled [`reqwest::Client`] and sends [`ProxyRequest`] /// instances without any knowledge of protocol names or model IDs. -#[derive(Debug, Clone)] +#[derive(Clone)] #[must_use = "ProxyClient does nothing until send/send_stream is called"] pub struct ProxyClient { http: reqwest::Client, @@ -157,6 +157,20 @@ pub struct ProxyClient { upstream_headers_timeout: Duration, } +impl fmt::Debug for ProxyClient { + // SECURITY: do not delegate to `reqwest::Client`'s Debug. It prints the + // configured forward-proxy URI whenever a proxy is set, and an egress + // proxy URL is operator infrastructure that does not belong in logs or + // `{:?}` output of containing types (e.g. `AppState`). + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ProxyClient") + .field("http", &"") + .field("stream_idle_timeout", &self.stream_idle_timeout) + .field("upstream_headers_timeout", &self.upstream_headers_timeout) + .finish() + } +} + impl Default for ProxyClient { fn default() -> Self { Self::new() @@ -200,11 +214,36 @@ impl ProxyClient { /// Use this when configuration may make client construction fallible /// (e.g. custom TLS backends, proxy environment validation). pub fn try_new() -> Result { - let http = reqwest::Client::builder() - .pool_max_idle_per_host(20) - .pool_idle_timeout(Duration::from_secs(90)) - .connect_timeout(Duration::from_secs(10)) - .build()?; + let http = pooled_http_builder().build()?; + + Ok(Self { + http, + stream_idle_timeout: DEFAULT_STREAM_IDLE_TIMEOUT, + upstream_headers_timeout: DEFAULT_UPSTREAM_HEADERS_TIMEOUT, + }) + } + + /// Creates a transport client that sends every upstream request through a + /// forward proxy (per-provider egress, `[provider.egress]`). + /// + /// The pool, connect-timeout, and default stream/header timeout settings + /// are identical to [`Self::try_new`]; only the proxy differs. Providers + /// without egress must keep using the shared default client — this + /// constructor exists for providers that configure one. + /// + /// `proxy_url` must be a `socks5h://`, `socks5://`, `http://`, or `https://` + /// URL with a host (userinfo allowed as proxy authentication). Anything + /// else is an error here as well as at config load, because reqwest would + /// otherwise silently accept it: it auto-prefixes `http://` onto scheme-less + /// URLs and builds proxies from schemes it cannot actually speak. + /// + /// SOCKS support requires reqwest's `socks` cargo feature, which this crate + /// enables. `socks5h://` resolves upstream hostnames at the proxy; + /// `socks5://` resolves them locally. + pub fn try_with_egress_proxy(proxy_url: &str) -> Result { + validate_egress_proxy_url(proxy_url)?; + let proxy = reqwest::Proxy::all(proxy_url)?; + let http = pooled_http_builder().proxy(proxy).build()?; Ok(Self { http, @@ -450,6 +489,60 @@ impl Stream for IdleTimeoutStream { // Helpers // --------------------------------------------------------------------------- +/// The connection-pool defaults shared by every [`ProxyClient`], egress or not. +/// +/// Extracted so the default client ([`ProxyClient::try_new`]) and the +/// per-provider egress client ([`ProxyClient::try_with_egress_proxy`]) cannot +/// drift apart: a provider without egress must behave exactly like today's +/// single shared client. +fn pooled_http_builder() -> reqwest::ClientBuilder { + reqwest::Client::builder() + .pool_max_idle_per_host(20) + .pool_idle_timeout(Duration::from_secs(90)) + .connect_timeout(Duration::from_secs(10)) +} + +/// Egress proxy schemes the transport will route through. +/// +/// Mirrors the config-level allowlist in `llm-proxy-core`'s +/// `validate_egress`; both are pinned by tests. reqwest itself accepts far +/// more (scheme-less URLs, `socks4`, unknown schemes), so this check — not +/// reqwest — defines what "a configured egress proxy" means. +const EGRESS_PROXY_SCHEMES: [&str; 4] = ["socks5h", "socks5", "http", "https"]; + +/// Reject egress proxy URLs outside the supported schemes or without a host. +/// +/// reqwest's `Proxy::all` silently coerces scheme-less URLs to `http://` and +/// builds proxies from schemes it cannot speak, so the transport validates +/// before handing the URL over. Config load applies the same rule earlier; +/// this is the second, defense-in-depth gate. +fn validate_egress_proxy_url(proxy_url: &str) -> Result<(), ProviderError> { + let url = proxy_url.trim(); + let (scheme, rest) = url.split_once("://").ok_or_else(|| { + ProviderError::InvalidConfig(format!( + "egress proxy URL has no scheme (expected socks5h://, socks5://, http://, or https://): {url:?}" + )) + })?; + if !EGRESS_PROXY_SCHEMES + .iter() + .any(|accepted| scheme.eq_ignore_ascii_case(accepted)) + { + return Err(ProviderError::InvalidConfig(format!( + "egress proxy scheme must be socks5h://, socks5://, http://, or https://, got {scheme}://" + ))); + } + // The authority region must hold a host; drop userinfo (`user:pass@`) + // before checking so proxy credentials do not count as a host. + let authority = rest.split(['/', '?', '#']).next().unwrap_or_default(); + let host = authority.rsplit_once('@').map_or(authority, |(_, h)| h); + if host.is_empty() { + return Err(ProviderError::InvalidConfig(format!( + "egress proxy URL has no host: {url:?}" + ))); + } + Ok(()) +} + /// Apply authentication headers to a request builder based on [`AuthStyle`]. /// /// # Security note on `AuthStyle::Both` @@ -1498,6 +1591,247 @@ mod tests { ); } + // ----------------------------------------------------------------------- + // Per-provider egress proxy client + // ----------------------------------------------------------------------- + + #[test] + fn egress_proxy_client_builds_for_each_supported_scheme() { + // This is also the empirical proof that reqwest's `socks` cargo + // feature (enabled in this crate) accepts socks5/socks5h proxy URLs: + // without the feature these two builds fail. + for url in [ + "socks5h://tunnel.example.com:1080", + "socks5://127.0.0.1:1080", + "http://proxy.internal:3128", + "https://proxy.internal:8443", + "SOCKS5H://tunnel.example.com:1080", + "http://user:pass@proxy.internal:3128", + "socks5h://tunnel.example.com", + ] { + let client = ProxyClient::try_with_egress_proxy(url) + .unwrap_or_else(|e| panic!("egress proxy {url} must build: {e}")); + // The egress client carries the same streaming timeouts as the + // default client: only the proxy differs. + assert_eq!( + client.stream_idle_timeout, DEFAULT_STREAM_IDLE_TIMEOUT, + "egress client must keep the default stream idle timeout" + ); + assert_eq!( + client.upstream_headers_timeout, DEFAULT_UPSTREAM_HEADERS_TIMEOUT, + "egress client must keep the default headers timeout" + ); + } + } + + #[test] + fn egress_proxy_client_rejects_unsupported_or_malformed_urls() { + for url in [ + "ftp://proxy:21", + "socks4://127.0.0.1:1080", + // Scheme-less: reqwest would silently prefix http:// — the + // transport must refuse instead. + "127.0.0.1:1080", + "socks5h://", + "", + ] { + let err = ProxyClient::try_with_egress_proxy(url) + .expect_err(&format!("egress proxy {url:?} must be rejected")); + assert!( + matches!(err, ProviderError::InvalidConfig(_)), + "egress proxy {url:?} must be InvalidConfig, got {err:?}" + ); + } + } + + /// Handler for the egress end-to-end test: echoes the request target. + /// + /// A forward proxy receives the request in absolute form + /// (`POST http://upstream/v1/test`), so echoing the URI proves the + /// connection landed on the proxy with the upstream target intact. + async fn absolute_target_echo_handler(uri: axum::http::Uri) -> String { + uri.to_string() + } + + #[test] + fn egress_proxy_client_debug_does_not_print_the_proxy() { + let client = ProxyClient::try_with_egress_proxy("socks5h://debug-secret-host.invalid:1080") + .expect("builds"); + let debug = format!("{client:?}"); + assert!( + !debug.contains("debug-secret-host.invalid"), + "ProxyClient Debug must not print the egress proxy: {debug}" + ); + } + + #[tokio::test] + async fn egress_proxy_client_sends_requests_through_the_proxy() { + // "Proxy" = an ordinary HTTP server on an ephemeral loopback port. It + // is the only listener in this test; the upstream host does not exist. + let app = Router::new().route("/v1/test", post(absolute_target_echo_handler)); + let proxy_base = start_test_server(app).await; + + let client = + ProxyClient::try_with_egress_proxy(&proxy_base).expect("http egress proxy builds"); + let req = ProxyRequest { + // Nonexistent host: only reachable via the proxy, which resolves + // nothing for plain-http absolute-form targets. + url: "http://egress-upstream.invalid/v1/test".to_owned(), + auth: AuthHeaders { + style: AuthStyle::Bearer, + api_key: SecretString::from("key"), + }, + body: br#"{"hello":"world"}"#.to_vec(), + stream: false, + extra_headers: std::collections::HashMap::new(), + }; + + let resp = client.send(req).await.expect("request via proxy succeeds"); + let text = String::from_utf8_lossy(&resp); + assert!( + text.contains("http://egress-upstream.invalid/v1/test"), + "the proxy must receive the absolute-form upstream target, got: {text}" + ); + } + + /// Read exactly `buf.len()` bytes (test-side helper for the SOCKS5 + /// handshake, where every field is fixed-width). + async fn read_exact(stream: &mut tokio::net::TcpStream, buf: &mut [u8]) -> std::io::Result<()> { + use tokio::io::AsyncReadExt; + let mut filled = 0; + while filled < buf.len() { + let n = stream.read(&mut buf[filled..]).await?; + if n == 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + "socks test client closed early", + )); + } + filled += n; + } + Ok(()) + } + + /// Serve exactly one no-auth SOCKS5 CONNECT, report the destination the + /// client asked to reach on `destination_tx`, then pipe the tunnel to + /// `upstream` (the copy loops stay alive until the test runtime drops + /// them — reqwest pools the tunneled connection, so waiting for it to + /// close would block for the whole pool idle timeout). + async fn serve_one_socks5_connect( + listener: tokio::net::TcpListener, + upstream: std::net::SocketAddr, + destination_tx: tokio::sync::oneshot::Sender, + ) { + use tokio::io::AsyncWriteExt; + let (mut sock, _peer) = listener.accept().await.expect("socks client connects"); + + // Client greeting: [VER, NMETHODS, METHODS..]; answer NO AUTH. + let mut head = [0u8; 2]; + read_exact(&mut sock, &mut head).await.expect("greeting"); + let mut methods = vec![0u8; head[1] as usize]; + read_exact(&mut sock, &mut methods).await.expect("methods"); + sock.write_all(&[0x05, 0x00]).await.expect("auth reply"); + + // CONNECT request: [VER, CMD=1, RSV, ATYP, ADDR..., PORT]. + let mut req_head = [0u8; 4]; + read_exact(&mut sock, &mut req_head) + .await + .expect("request head"); + assert_eq!(req_head[1], 0x01, "expected CONNECT"); + let addr_len = match req_head[3] { + 0x01 => 4, // IPv4 + 0x04 => 16, // IPv6 + 0x03 => { + // domain: length-prefixed + let mut len = [0u8; 1]; + read_exact(&mut sock, &mut len) + .await + .expect("domain length"); + len[0] as usize + } + other => panic!("unknown SOCKS address type {other}"), + }; + let mut addr_bytes = vec![0u8; addr_len]; + read_exact(&mut sock, &mut addr_bytes) + .await + .expect("address"); + let mut port = [0u8; 2]; + read_exact(&mut sock, &mut port).await.expect("port"); + let destination = if req_head[3] == 0x03 { + String::from_utf8(addr_bytes).expect("domain is UTF-8") + } else { + format!("{}-byte address", addr_bytes.len()) + }; + + // Reply: succeeded, bound to the real upstream. + let up_port = upstream.port().to_be_bytes(); + sock.write_all(&[0x05, 0x00, 0x00, 0x01, 127, 0, 0, 1, up_port[0], up_port[1]]) + .await + .expect("connect reply"); + + destination_tx + .send(destination) + .expect("test holds the receiver"); + + // Pipe the tunnel to the upstream. + let up = tokio::net::TcpStream::connect(upstream) + .await + .expect("connect to upstream"); + let (mut client_read, mut client_write) = sock.into_split(); + let (mut up_read, mut up_write) = up.into_split(); + let a = tokio::spawn(async move { tokio::io::copy(&mut client_read, &mut up_write).await }); + let b = tokio::spawn(async move { tokio::io::copy(&mut up_read, &mut client_write).await }); + let (_, _) = tokio::join!(a, b); + } + + #[tokio::test] + async fn socks5h_egress_delivers_the_hostname_to_the_proxy_tunnel() { + // Upstream: an ordinary echo server; only the SOCKS proxy can reach it. + let app = Router::new().route("/socks-test", post(echo_handler)); + let upstream_base = start_test_server(app).await; + let upstream_addr: std::net::SocketAddr = upstream_base + .trim_start_matches("http://") + .parse() + .expect("upstream addr"); + + let socks_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let socks_addr = socks_listener.local_addr().unwrap(); + let (destination_tx, destination_rx) = tokio::sync::oneshot::channel(); + tokio::spawn(serve_one_socks5_connect( + socks_listener, + upstream_addr, + destination_tx, + )); + + let client = ProxyClient::try_with_egress_proxy(&format!("socks5h://{socks_addr}")) + .expect("socks5h egress proxy builds"); + let req = ProxyRequest { + // Nonexistent hostname: socks5h must hand the NAME to the proxy + // (remote DNS), so the proxy observes it without resolving. + url: "http://egress-upstream.invalid/socks-test".to_owned(), + auth: AuthHeaders { + style: AuthStyle::Bearer, + api_key: SecretString::from("key"), + }, + body: br#"{"via":"socks"}"#.to_vec(), + stream: false, + extra_headers: std::collections::HashMap::new(), + }; + + let resp = client.send(req).await.expect("request via socks tunnel"); + let text = String::from_utf8_lossy(&resp); + assert!( + text.contains(r#""via":"socks""#), + "the echo upstream must see the body through the tunnel, got: {text}" + ); + + let destination = destination_rx.await.expect("destination"); + assert_eq!( + destination, "egress-upstream.invalid", + "socks5h must deliver the unresolved hostname to the proxy" + ); + } + // ----------------------------------------------------------------------- // Edge case tests // ----------------------------------------------------------------------- diff --git a/crates/llm-proxy-provider/tests/fixture_tests.rs b/crates/llm-proxy-provider/tests/fixture_tests.rs index 6e6c3ad..818c52b 100644 --- a/crates/llm-proxy-provider/tests/fixture_tests.rs +++ b/crates/llm-proxy-provider/tests/fixture_tests.rs @@ -562,6 +562,12 @@ mod tests { } } + #[expect( + clippy::too_many_lines, + reason = "walks one fixture's expected wire output field by field \ + (messages, tools, sampling); splitting the walk would scatter \ + the fixture's assertions across helpers" + )] fn run_encode_request_fixture(adapter_name: &str, case: &str) { let dir = Path::new(FIXTURE_ROOT).join(adapter_name).join(case); let core_json = read_fixture(&dir, "core.json"); @@ -705,6 +711,12 @@ mod tests { // --------------------------------------------------------------------------- /// Test that decoding a provider response produces the expected output.json fields. + #[expect( + clippy::too_many_lines, + reason = "walks one fixture's decoded response field by field (content \ + blocks, usage, stop reason); splitting the walk would scatter \ + the fixture's assertions across helpers" + )] fn run_decode_response_fixture(adapter_name: &str, case: &str) { let dir = Path::new(FIXTURE_ROOT).join(adapter_name).join(case); let input_json = read_fixture(&dir, "input.json"); @@ -959,6 +971,12 @@ mod tests { /// per event (e.g., `{"type": "message_start", ...}`). The test compares: /// 1. Event counts must match. /// 2. Each event's variant kind (from the `"type"` field) must match. + #[expect( + clippy::too_many_lines, + reason = "walks one fixture's expected SSE event sequence frame by \ + frame; splitting the walk would scatter the fixture's \ + assertions across helpers" + )] fn run_stream_decode_fixture(adapter_name: &str, case: &str) { let dir = Path::new(FIXTURE_ROOT).join(adapter_name).join(case); let sse_input = read_fixture_raw(&dir, "input.sse"); diff --git a/crates/llm-proxy-server/Cargo.toml b/crates/llm-proxy-server/Cargo.toml index 8547f84..71944e0 100644 --- a/crates/llm-proxy-server/Cargo.toml +++ b/crates/llm-proxy-server/Cargo.toml @@ -18,6 +18,7 @@ llm-proxy-core = { workspace = true } llm-proxy-protocol = { path = "../llm-proxy-protocol" } llm-proxy-provider = { path = "../llm-proxy-provider" } llm-proxy-storage = { path = "../llm-proxy-storage" } +llm-proxy-api = { path = "../llm-proxy-api" } axum = { workspace = true, features = ["macros", "json"] } tokio = { workspace = true, features = ["fs"] } @@ -34,6 +35,7 @@ bytes = "1" futures = "0.3" regex = { workspace = true } sha2 = "0.10" +argon2 = { workspace = true } toml = { workspace = true } time = { workspace = true } secrecy = { workspace = true } @@ -42,3 +44,5 @@ rust_decimal = { workspace = true } [dev-dependencies] tower = { workspace = true, features = ["util"] } llm-proxy-storage = { path = "../llm-proxy-storage", features = ["test-utils"] } +tempfile = "3" +tracing-subscriber = { workspace = true } diff --git a/crates/llm-proxy-server/src/catalog_service.rs b/crates/llm-proxy-server/src/catalog_service.rs index 7f2b411..36c9c79 100644 --- a/crates/llm-proxy-server/src/catalog_service.rs +++ b/crates/llm-proxy-server/src/catalog_service.rs @@ -623,6 +623,7 @@ mod tests { api_key: secrecy::SecretString::from("test-key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -791,6 +792,7 @@ mod tests { api_key: secrecy::SecretString::from(""), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -1131,6 +1133,7 @@ mod tests { api_key: secrecy::SecretString::from(""), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), diff --git a/crates/llm-proxy-server/src/lib.rs b/crates/llm-proxy-server/src/lib.rs index 642c4b7..946517a 100644 --- a/crates/llm-proxy-server/src/lib.rs +++ b/crates/llm-proxy-server/src/lib.rs @@ -17,11 +17,15 @@ pub mod shutdown; pub mod state; pub use catalog_service::{ModelCatalogService, write_catalog_atomic}; -pub use middleware::{RateLimiter, RequestDeduplicator, RequestId, RequestIdGenerator}; +pub use middleware::{ + CallerId, ClientAuth, ClientAuthError, RateLimiter, RequestDeduplicator, RequestId, + RequestIdGenerator, +}; pub use shutdown::shutdown_signal; pub use state::{AppState, BuildInfo}; use axum::Router; +use secrecy::ExposeSecret; /// Build the application router for the given state. /// @@ -36,3 +40,17 @@ use axum::Router; pub fn build_router(state: AppState) -> Router { routes::router(state) } + +/// Build the router served on the dedicated admin listener: the `/usage` +/// query surface behind the `[auth.admin]` token. +/// +/// Returns `None` when no `[auth.admin]` token is configured (there is no +/// admin surface to serve) or — impossible for configs that passed +/// `load_app_config` — `[storage]` is absent. The surface is strictly +/// read-only; every route is a `GET` over read-only SQLite connections. +#[must_use = "the Router must be used with a server"] +pub fn admin_router(state: &AppState) -> Option { + let admin = state.app_config.auth.admin.as_ref()?; + let storage = state.app_config.storage.as_ref()?; + Some(llm_proxy_api::UsageApi::new(storage.path.clone(), admin.token.expose_secret()).router()) +} diff --git a/crates/llm-proxy-server/src/middleware.rs b/crates/llm-proxy-server/src/middleware.rs index ceaf3a7..49f7a23 100644 --- a/crates/llm-proxy-server/src/middleware.rs +++ b/crates/llm-proxy-server/src/middleware.rs @@ -1,7 +1,8 @@ //! Middleware components for the LLM proxy server. //! //! Provides request deduplication, rate limiting, request ID generation, -//! and client IP extraction. +//! client IP extraction, and client-key authentication state +//! ([`ClientAuth`], [`CallerId`]). use std::collections::HashMap; use std::convert::Infallible; @@ -12,8 +13,12 @@ use std::time::Instant; use axum::extract::{ConnectInfo, FromRequestParts}; use axum::http::request::Parts; +use llm_proxy_core::{ANONYMOUS_CALLER, AuthConfig, AuthMode}; +use secrecy::{ExposeSecret, SecretString}; use sha2::{Digest, Sha256}; +use argon2::PasswordVerifier; + // --------------------------------------------------------------------------- // RequestDeduplicator // --------------------------------------------------------------------------- @@ -278,6 +283,37 @@ impl RateLimiter { bucket.try_consume(refill_rate) } + + /// Whether a request from `client_ip` would currently be allowed, + /// **without consuming a token**. + /// + /// Pure read of the peer's current refilled balance: a peer with no + /// bucket yet would start full, so it reports `true`. Used by the + /// client-auth gate to reject peers whose failure budget is already + /// exhausted *before* running the CPU-costly argon2 verification; the + /// actual token consumption still happens on the failure path via + /// [`Self::is_allowed`], so successful authentications are never charged. + #[must_use] + pub fn would_allow(&self, client_ip: &str) -> bool { + if self.max_requests_per_minute == 0.0 { + return true; + } + let buckets = self.buckets.lock().unwrap_or_else(|e| { + tracing::warn!("rate limiter mutex poisoned; recovering from panic"); + e.into_inner() + }); + let Some(bucket) = buckets.get(client_ip) else { + return true; + }; + // Same refill math as `try_consume`, but without mutating the bucket: + // compute the balance as of now and check it holds at least one token. + let elapsed = Instant::now() + .duration_since(bucket.last_refill) + .as_secs_f64(); + let refill_rate = self.max_requests_per_minute / 60.0; + let tokens = (bucket.tokens + elapsed * refill_rate).min(bucket.max_tokens); + tokens >= 1.0 + } } impl Default for RateLimiter { @@ -431,6 +467,395 @@ pub fn get_client_ip( "unknown".to_owned() } +// --------------------------------------------------------------------------- +// Client-key authentication +// --------------------------------------------------------------------------- + +/// One configured client key: its id and its PHC-format argon2 verifier. +/// +/// The verifier is the *hashed* secret only; the plaintext client key exists +/// solely in the operator's records and the client's `Authorization` header, +/// never in this process's long-lived state. +struct ClientVerifier { + /// Operator-chosen key id used for caller attribution and rate-limit keys. + key_id: String, + /// PHC-format argon2 hash of the client key (e.g. + /// `$argon2id$v=19$m=19456,t=2,p=1$...`). Parsed (and therefore + /// validated) whenever the verifier is checked; the raw string is kept + /// because `PasswordHash` borrows its input. + phc: SecretString, +} + +/// SHA-256 digest of an authenticated client key, with the key id it matched. +struct CachedDigest { + digest: [u8; 32], + key_id: String, +} + +/// Constant-time equality of two equal-length byte strings. +/// +/// Compares all bytes unconditionally so the runtime does not depend on where +/// the first difference appears. Length mismatch returns `false` immediately; +/// digest lengths are fixed at 32 bytes here, so that branch is not secret +/// dependent. +fn constant_time_eq(a: &[u8], b: &[u8]) -> bool { + if a.len() != b.len() { + return false; + } + let mut diff = 0u8; + for (x, y) in a.iter().zip(b.iter()) { + diff |= x ^ y; + } + diff == 0 +} + +/// Errors raised while building [`ClientAuth`] from configuration. +#[derive(Debug, thiserror::Error)] +pub enum ClientAuthError { + /// `auth.mode = "none"` was combined with configured client keys. + #[error("auth.mode is \"none\" but [auth.clients] is non-empty")] + ModeNoneWithClients, + /// `auth.mode = "client-keys"` had no client keys configured. + #[error("auth.mode is \"client-keys\" but [auth.clients] is empty")] + ClientKeysWithoutClients, + /// A key id was empty or the reserved `anonymous` name. + #[error("client key id \"{0}\" is empty or reserved")] + InvalidKeyId(String), + /// A verifier string is not a parseable PHC argon2 hash. + /// + /// The error carries the key id, never the hash string itself. + #[error("client key \"{key_id}\" has an invalid PHC verifier: {source}")] + InvalidVerifier { + /// The key id whose verifier failed to parse. + key_id: String, + /// The underlying parse error. + source: argon2::password_hash::Error, + }, +} + +/// Client-key authentication state: the configured verifiers plus a digest +/// cache for verified keys. +/// +/// Holds hashed verifiers only, never plaintext secrets. A presented bearer +/// token is verified with argon2 (constant-time tag comparison inside the +/// KDF); after a first successful verification the SHA-256 digest of the +/// token is cached and compared in constant time, so steady-state requests +/// skip the (CPU-heavy) KDF entirely. Only digests of tokens that already +/// verified are cached, so the cache holds at most one entry per configured +/// key id. +/// +/// Manual `Debug`: prints the mode and the key ids, never the verifier +/// strings (a leaked PHC hash enables offline brute-force of the key). +pub struct ClientAuth { + /// Whether proxy routes must present a valid client key. True whenever + /// verifiers were configured (`auth.mode = "client-keys"`). + enforced: bool, + /// Configured verifiers, sorted by key id for deterministic iteration. + verifiers: Vec, + /// Digest cache: SHA-256(token) -> key id for tokens that verified. + cache: Mutex>, +} + +impl std::fmt::Debug for ClientAuth { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ClientAuth") + .field("enforced", &self.enforced) + .field( + "client_key_ids", + &self + .verifiers + .iter() + .map(|v| v.key_id.as_str()) + .collect::>(), + ) + .finish() + } +} + +impl ClientAuth { + /// Build authentication state from the `[auth]` config section. + /// + /// Validates mode consistency, key ids, and that every verifier parses as + /// a PHC argon2 string. `auth.mode = "none"` builds a disabled + /// (not enforced) authenticator; the caller is responsible for the + /// loopback-bind rule that `load_app_config` enforces for that mode. + /// + /// # Errors + /// + /// Returns [`ClientAuthError`] on contradictory mode/client combinations, + /// reserved or empty key ids, or an unparseable verifier string. + pub fn from_config(config: &AuthConfig) -> Result { + match config.mode { + AuthMode::None if !config.clients.is_empty() => { + return Err(ClientAuthError::ModeNoneWithClients); + } + AuthMode::ClientKeys if config.clients.is_empty() => { + return Err(ClientAuthError::ClientKeysWithoutClients); + } + AuthMode::None | AuthMode::ClientKeys => {} + } + let mut verifiers = Vec::with_capacity(config.clients.len()); + for (key_id, phc) in &config.clients { + if key_id.is_empty() || key_id == ANONYMOUS_CALLER { + return Err(ClientAuthError::InvalidKeyId(key_id.clone())); + } + // Validate parseability now so a malformed verifier fails startup + // (via `load_toml_state`) instead of failing every request later. + if let Err(source) = argon2::password_hash::phc::PasswordHash::new(phc.expose_secret()) + { + return Err(ClientAuthError::InvalidVerifier { + key_id: key_id.clone(), + source: source.into(), + }); + } + verifiers.push(ClientVerifier { + key_id: key_id.clone(), + phc: SecretString::from(phc.expose_secret().to_owned()), + }); + } + verifiers.sort_by(|a, b| a.key_id.cmp(&b.key_id)); + Ok(Self { + enforced: matches!(config.mode, AuthMode::ClientKeys), + verifiers, + cache: Mutex::new(Vec::new()), + }) + } + + /// Fail-closed authenticator: enforcement on, nothing can authenticate. + /// + /// Used when config is invalid at construction time so every proxied + /// request is rejected rather than silently downgraded to no-auth. + #[must_use] + pub fn fail_closed() -> Self { + Self { + enforced: true, + verifiers: Vec::new(), + cache: Mutex::new(Vec::new()), + } + } + + /// Whether proxy routes require a valid client key. + #[must_use] + pub fn is_enforced(&self) -> bool { + self.enforced + } + + /// Number of configured client keys (0 when not enforced). + #[must_use] + pub fn client_count(&self) -> usize { + self.verifiers.len() + } + + /// Look up a cached match for `digest` in constant time. + fn cached_key_id(&self, digest: &[u8; 32]) -> Option { + let cache = self.cache.lock().unwrap_or_else(|e| { + tracing::warn!("client auth digest cache poisoned; recovering from panic"); + e.into_inner() + }); + cache + .iter() + .find(|entry| constant_time_eq(&entry.digest, digest)) + .map(|entry| entry.key_id.clone()) + } + + /// Verify a presented bearer token against every configured verifier. + /// + /// Runs the argon2 KDFs on the blocking pool (a single verification with + /// default parameters costs tens of milliseconds; running it on the async + /// executor would stall every other future on the worker). Returns the + /// matching key id, or `None` for unknown/invalid tokens. The token is + /// never logged. + async fn verify_uncached(&self, token: &str) -> Option { + // Digest before the move into the blocking task; needed for the + // cache insert when verification succeeds. + let digest: [u8; 32] = Sha256::digest(token.as_bytes()).into(); + let verifiers: Vec<(String, String)> = self + .verifiers + .iter() + .map(|v| (v.key_id.clone(), v.phc.expose_secret().to_owned())) + .collect(); + let token = token.to_owned(); + let matched = tokio::task::spawn_blocking(move || { + let argon2 = argon2::Argon2::default(); + for (key_id, phc) in &verifiers { + // `verify_password` derives the tag with the parameters + // embedded in the PHC string and compares in constant time. + if let Ok(hash) = argon2::password_hash::phc::PasswordHash::new(phc) + && argon2.verify_password(token.as_bytes(), &hash).is_ok() + { + return Some(key_id.clone()); + } + } + None + }) + .await + .map_err(|error| { + tracing::warn!(%error, "client auth verification task failed"); + error + }) + .ok() + .flatten(); + if let Some(key_id) = &matched { + let mut cache = self.cache.lock().unwrap_or_else(|e| { + tracing::warn!("client auth digest cache poisoned; recovering from panic"); + e.into_inner() + }); + if !cache + .iter() + .any(|entry| constant_time_eq(&entry.digest, &digest)) + { + cache.push(CachedDigest { + digest, + key_id: key_id.clone(), + }); + } + } + matched + } + + /// Authenticate a bearer token, returning the matching key id. + /// + /// Cache hits skip the KDF; misses verify against every configured key on + /// the blocking pool. `None` means the token is not a configured client + /// key (or auth is not enforced, in which case nothing authenticates). + pub(crate) async fn authenticate(&self, token: &str) -> Option { + if !self.enforced || self.verifiers.is_empty() { + return None; + } + let digest: [u8; 32] = Sha256::digest(token.as_bytes()).into(); + if let Some(key_id) = self.cached_key_id(&digest) { + return Some(key_id); + } + self.verify_uncached(token).await + } +} + +/// Extract the bearer token from `Authorization: Bearer `. +/// +/// The scheme is matched case-insensitively (RFC 7235). Returns `None` when +/// the header is absent, uses another scheme, or carries an empty token. The +/// returned secret is never logged; holders use it only for verification. +pub(crate) fn extract_bearer_token(headers: &axum::http::HeaderMap) -> Option { + let value = headers.get(axum::http::header::AUTHORIZATION)?; + let header = value.to_str().ok()?; + let mut parts = header.splitn(2, ' '); + match (parts.next(), parts.next()) { + (Some(scheme), Some(token)) if scheme.eq_ignore_ascii_case("bearer") => { + let token = token.trim(); + if token.is_empty() { + None + } else { + Some(SecretString::from(token.to_owned())) + } + } + _ => None, + } +} + +/// Typed caller identity carried in request extensions. +/// +/// Inserted by the client-auth middleware on proxy routes: the authenticated +/// key id when client keys are enforced, or [`CallerId::anonymous`] under the +/// loopback no-auth exemption. Handlers read it for caller attribution on +/// events and for rate-limit keying. +#[derive(Debug, Clone)] +pub struct CallerId { + /// Authenticated client key id; `None` when unauthenticated (no-auth + /// mode). + key_id: Option, +} + +impl CallerId { + /// The unauthenticated caller used under the loopback no-auth exemption. + #[must_use] + pub fn anonymous() -> Self { + Self { key_id: None } + } + + /// An authenticated caller identified by `key_id`. + #[must_use] + pub fn authenticated(key_id: impl Into) -> Self { + Self { + key_id: Some(key_id.into()), + } + } + + /// Whether a client key authenticated this request. + #[must_use] + pub fn is_authenticated(&self) -> bool { + self.key_id.is_some() + } + + /// The attribution string for events and logs: the key id, or + /// `"anonymous"` when unauthenticated. + #[must_use] + pub fn as_str(&self) -> &str { + self.key_id.as_deref().unwrap_or(ANONYMOUS_CALLER) + } +} + +impl std::fmt::Display for CallerId { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + +/// Extractor bundling the per-request identity stamped by the outer middleware +/// layers: the injected [`RequestId`], the [`CallerId`] from the client-auth +/// layer, and the peer address (when connection info is available). +/// +/// One extractor replaces three in every proxy handler signature, and the +/// owned struct is passed whole into inner handlers and +/// [`prepare_request`](crate::routes::core_pipeline::prepare_request). +/// +/// Absent extensions degrade the same way the individual extractors would: +/// a missing request id becomes `"unknown"` and a missing caller becomes +/// [`CallerId::anonymous`] (defense-in-depth for routers built without the +/// usual layers; through [`crate::routes::router`] both are always present). +#[derive(Debug, Clone)] +pub(crate) struct RequestMeta { + /// Request id minted by the outermost `inject_request_id` layer. + pub(crate) request_id: RequestId, + /// Caller identity stamped by the client-auth layer. + pub(crate) caller: CallerId, + /// Peer socket address, when the server was built with connection info. + pub(crate) connect_info: Option, +} + +impl FromRequestParts for RequestMeta +where + S: Send + Sync, +{ + type Rejection = Infallible; + + fn from_request_parts( + parts: &mut Parts, + _state: &S, + ) -> impl Future> + Send { + let request_id = parts + .extensions + .get::() + .cloned() + .unwrap_or_else(|| RequestId("unknown".to_owned())); + let caller = parts + .extensions + .get::() + .cloned() + .unwrap_or_else(CallerId::anonymous); + let connect_info = parts + .extensions + .get::>() + .map(|connect_info| connect_info.0); + async move { + Ok(Self { + request_id, + caller, + connect_info, + }) + } + } +} + // --------------------------------------------------------------------------- // Tests // --------------------------------------------------------------------------- @@ -439,6 +864,208 @@ pub fn get_client_ip( mod tests { use super::*; + // -- ClientAuth ------------------------------------------------------------ + + /// Hash a token with deliberately cheap argon2id parameters (8 KiB, one + /// pass) so tests run in microseconds; `ClientAuth` verification reads the + /// cost parameters from the PHC string, so verification is equally cheap. + fn cheap_phc(token: &str, salt: &[u8]) -> String { + let params = argon2::Params::new(8, 1, 1, None).expect("valid test params"); + let hasher = argon2::Argon2::new( + argon2::Algorithm::Argon2id, + argon2::Version::V0x13, + params.clone(), + ); + use argon2::CustomizedPasswordHasher; + hasher + .hash_password_customized( + token.as_bytes(), + salt, + Some(&argon2::ARGON2ID_IDENT), + None, + params, + ) + .expect("hash token") + .to_string() + } + + fn auth_config_with(phc: &str) -> llm_proxy_core::AuthConfig { + let mut clients = HashMap::new(); + clients.insert("laptop".to_owned(), SecretString::from(phc.to_owned())); + llm_proxy_core::AuthConfig { + mode: llm_proxy_core::AuthMode::ClientKeys, + clients, + admin: None, + } + } + + #[test] + fn client_auth_from_config_rejects_mode_contradictions() { + let none_with_clients = { + let mut clients = HashMap::new(); + clients.insert( + "laptop".to_owned(), + SecretString::from(cheap_phc("secret", b"0123456789abcdef")), + ); + llm_proxy_core::AuthConfig { + mode: llm_proxy_core::AuthMode::None, + clients, + admin: None, + } + }; + assert!(matches!( + ClientAuth::from_config(&none_with_clients), + Err(ClientAuthError::ModeNoneWithClients) + )); + + let keys_without_clients = llm_proxy_core::AuthConfig { + mode: llm_proxy_core::AuthMode::ClientKeys, + clients: HashMap::new(), + admin: None, + }; + assert!(matches!( + ClientAuth::from_config(&keys_without_clients), + Err(ClientAuthError::ClientKeysWithoutClients) + )); + } + + #[test] + fn client_auth_from_config_rejects_bad_key_ids_and_hashes() { + for bad_id in ["", "anonymous"] { + let mut config = auth_config_with(&cheap_phc("secret", b"0123456789abcdef")); + config.clients.remove("laptop"); + config + .clients + .insert(bad_id.to_owned(), SecretString::from("x".to_owned())); + assert!( + matches!( + ClientAuth::from_config(&config), + Err(ClientAuthError::InvalidKeyId(id)) if id == bad_id + ), + "key id {bad_id:?} must be rejected" + ); + } + + let config = auth_config_with("not-a-phc-string"); + let error = ClientAuth::from_config(&config).expect_err("malformed verifier rejected"); + assert!( + error.to_string().contains("laptop"), + "error names the key id, got: {error}" + ); + assert!( + !error.to_string().contains("not-a-phc-string"), + "error must not echo the verifier, got: {error}" + ); + } + + #[tokio::test] + async fn client_auth_authenticates_a_configured_key_and_rejects_others() { + let secret = cheap_phc("the-real-secret", b"0123456789abcdef"); + let auth = ClientAuth::from_config(&auth_config_with(&secret)).expect("valid config"); + + let matched = auth.authenticate("the-real-secret").await; + assert_eq!(matched.as_deref(), Some("laptop")); + // Repeat hits the digest cache and must still resolve. + assert_eq!( + auth.authenticate("the-real-secret").await.as_deref(), + Some("laptop") + ); + assert_eq!(auth.authenticate("wrong-secret").await, None); + assert_eq!(auth.authenticate("").await, None); + } + + #[tokio::test] + async fn client_auth_picks_the_matching_key_among_many() { + let mut clients = HashMap::new(); + for (id, token) in [("alpha", "one"), ("beta", "two"), ("gamma", "three")] { + clients.insert( + id.to_owned(), + SecretString::from(cheap_phc(token, b"0123456789abcdef")), + ); + } + let config = llm_proxy_core::AuthConfig { + mode: llm_proxy_core::AuthMode::ClientKeys, + clients, + admin: None, + }; + let auth = ClientAuth::from_config(&config).expect("valid config"); + assert_eq!(auth.authenticate("two").await.as_deref(), Some("beta")); + assert_eq!(auth.authenticate("three").await.as_deref(), Some("gamma")); + assert_eq!(auth.authenticate("four").await, None); + } + + #[test] + fn client_auth_none_mode_is_not_enforced() { + let auth = + ClientAuth::from_config(&llm_proxy_core::AuthConfig::default()).expect("valid config"); + assert!(!auth.is_enforced()); + } + + #[test] + fn client_auth_debug_shows_key_ids_but_never_verifiers() { + let secret = cheap_phc("the-real-secret", b"0123456789abcdef"); + let auth = ClientAuth::from_config(&auth_config_with(&secret)).expect("valid config"); + let debug = format!("{auth:?}"); + assert!(debug.contains("laptop"), "got: {debug}"); + assert!( + !debug.contains("argon2id") && !debug.contains(&secret), + "verifier string must not appear in Debug, got: {debug}" + ); + } + + #[test] + fn constant_time_eq_behaves_as_equality() { + assert!(constant_time_eq(b"abc", b"abc")); + assert!(!constant_time_eq(b"abc", b"abd")); + assert!(!constant_time_eq(b"abc", b"ab")); + } + + // -- extract_bearer_token -------------------------------------------------- + + #[test] + fn bearer_token_extracts_case_insensitive_scheme_and_trims() { + use axum::http::HeaderMap; + let mut headers = HeaderMap::new(); + headers.insert("authorization", "Bearer abc123".parse().unwrap()); + assert_eq!( + extract_bearer_token(&headers).map(|t| t.expose_secret().to_owned()), + Some("abc123".to_owned()) + ); + + let mut headers = HeaderMap::new(); + headers.insert("authorization", "bearer abc123 ".parse().unwrap()); + assert_eq!( + extract_bearer_token(&headers).map(|t| t.expose_secret().to_owned()), + Some("abc123".to_owned()) + ); + } + + #[test] + fn bearer_token_rejects_missing_other_scheme_and_empty() { + use axum::http::HeaderMap; + // SecretString has no PartialEq (deliberately); assert via Option::is_none. + assert!(extract_bearer_token(&HeaderMap::new()).is_none()); + + let mut headers = HeaderMap::new(); + headers.insert("authorization", "Basic dXNlcjpwYXNz".parse().unwrap()); + assert!(extract_bearer_token(&headers).is_none()); + + let mut headers = HeaderMap::new(); + headers.insert("authorization", "Bearer ".parse().unwrap()); + assert!(extract_bearer_token(&headers).is_none()); + } + + // -- CallerId --------------------------------------------------------------- + + #[test] + fn caller_id_attribution_strings() { + assert_eq!(CallerId::anonymous().as_str(), "anonymous"); + assert_eq!(CallerId::authenticated("laptop").as_str(), "laptop"); + assert!(!CallerId::anonymous().is_authenticated()); + assert!(CallerId::authenticated("laptop").is_authenticated()); + assert_eq!(CallerId::authenticated("laptop").to_string(), "laptop"); + } + // -- RequestDeduplicator --------------------------------------------------- #[test] @@ -580,6 +1207,54 @@ mod tests { } } + // -- RateLimiter::would_allow (auth flood gate) ----------------------------- + + #[test] + fn would_allow_is_false_once_the_peer_bucket_is_drained() { + let limiter = RateLimiter::new(1); + assert!( + limiter.is_allowed("10.0.0.9"), + "first request consumes the only token" + ); + // Refill is 1 token/minute, so within the test's lifetime the drained + // bucket stays empty and the peek must report denial. + assert!( + !limiter.would_allow("10.0.0.9"), + "drained peer must be denied by the peek" + ); + // Other peers are unaffected: a fresh bucket starts full. + assert!(limiter.would_allow("10.0.0.8")); + } + + #[test] + fn would_allow_reports_true_for_unknown_peer_and_disabled_limiter() { + let limiter = RateLimiter::new(100); + assert!( + limiter.would_allow("10.0.0.7"), + "unknown peer starts with a full bucket" + ); + + let disabled = RateLimiter::new(0); + for _ in 0..50 { + assert!(disabled.is_allowed("10.0.0.7")); + } + assert!( + disabled.would_allow("10.0.0.7"), + "a disabled limiter never denies" + ); + } + + #[test] + fn would_allow_does_not_consume_tokens() { + let limiter = RateLimiter::new(1); + assert!(limiter.is_allowed("10.0.0.6")); + // Repeated peeks must neither drain nor refill the bucket: the peer + // stays denied and a subsequent consume attempt is also denied. + assert!(!limiter.would_allow("10.0.0.6")); + assert!(!limiter.would_allow("10.0.0.6")); + assert!(!limiter.is_allowed("10.0.0.6")); + } + #[test] fn rate_limiter_rejects_requests_at_the_limit() { // With 1 RPM, the first request should be allowed and the second rejected. diff --git a/crates/llm-proxy-server/src/routes/chat.rs b/crates/llm-proxy-server/src/routes/chat.rs index 07d620e..d61edc0 100644 --- a/crates/llm-proxy-server/src/routes/chat.rs +++ b/crates/llm-proxy-server/src/routes/chat.rs @@ -12,14 +12,14 @@ use std::sync::Arc; use axum::Json; use axum::body::Body; -use axum::extract::{Extension, Path, State}; +use axum::extract::{Path, State}; use axum::http::{HeaderMap, Response}; use llm_proxy_core::ProviderRouteKind; use llm_proxy_protocol::client::openai_chat; use llm_proxy_protocol::openai::ChatCompletionRequest; use tracing::{info, warn}; -use crate::middleware::{OptionalConnectInfo, RequestId}; +use crate::middleware::RequestMeta; use crate::state::AppState; use super::core_pipeline; @@ -32,8 +32,7 @@ use super::error_response::{ClientProtocol, RouteError, route_error_response}; pub(crate) async fn handle_chat_completions( State(state): State, Path(provider): Path, - Extension(req_id): Extension, - OptionalConnectInfo(connect_info): OptionalConnectInfo, + meta: RequestMeta, headers: HeaderMap, body: axum::body::Bytes, ) -> Response { @@ -44,8 +43,7 @@ pub(crate) async fn handle_chat_completions( // blanket emit here would double-count every pipeline failure (the pipeline // emits AND returns Err via `?`), so the wrapper must stay emit-free // (audit route-responsefailed-gaps regression). - match handle_chat_completions_inner(state, req_id, provider, connect_info, headers, body).await - { + match handle_chat_completions_inner(state, meta, provider, headers, body).await { Ok(response) => response, Err(error) => { warn!(error = %error, "request failed"); @@ -63,9 +61,8 @@ pub(crate) async fn handle_chat_completions( /// function propagates without re-emitting. The outer wrapper renders only. async fn handle_chat_completions_inner( state: AppState, - req_id: RequestId, + meta: RequestMeta, provider: String, - connect_info: Option, headers: HeaderMap, body: axum::body::Bytes, ) -> Result, RouteError> { @@ -73,7 +70,7 @@ async fn handle_chat_completions_inner( let event_bus = Arc::clone(&state.event_bus); // The request id is consumed by prepare_request below; snapshot it once so // every early-gate emit references the same id. - let request_id = req_id.0.clone(); + let request_id = meta.request_id.0.clone(); // Every early gate emits exactly one ResponseFailed with the same // provider/request/start triple; bind it once instead of repeating the // six-argument call at each gate. @@ -102,15 +99,8 @@ async fn handle_chat_completions_inner( core_pipeline::validate_json_content_type(&headers).inspect_err(&emit_failed)?; let request_path = format!("/providers/{provider}/v1/chat/completions"); - let ctx = core_pipeline::prepare_request( - &state, - req_id.0, - &headers, - connect_info.as_ref(), - &body, - &request_path, - ) - .inspect_err(&emit_failed)?; + let ctx = core_pipeline::prepare_request(&state, &meta, &headers, &body, &request_path) + .inspect_err(&emit_failed)?; // Parse the OpenAI ChatCompletionRequest via axum's `Json` helper so the // `JsonRejection` taxonomy (syntax vs data error) is preserved and mapped @@ -145,28 +135,15 @@ async fn handle_chat_completions_inner( // Extract the inbound client auth token (used by passthrough-auth providers). let inbound_auth = core_pipeline::extract_inbound_auth(&headers); + let dispatch = core_pipeline::CoreDispatch { + provider_name: &provider, + route_kind: ProviderRouteKind::ChatCompletions, + client_protocol: ClientProtocol::OpenAiChat, + }; if is_streaming { - core_pipeline::handle_core_stream( - state, - ctx, - &provider, - ProviderRouteKind::ChatCompletions, - core, - ClientProtocol::OpenAiChat, - inbound_auth, - ) - .await + core_pipeline::handle_core_stream(state, ctx, dispatch, core, inbound_auth).await } else { - core_pipeline::handle_core_once( - state, - ctx, - &provider, - ProviderRouteKind::ChatCompletions, - core, - ClientProtocol::OpenAiChat, - inbound_auth, - ) - .await + core_pipeline::handle_core_once(state, ctx, dispatch, core, inbound_auth).await } } diff --git a/crates/llm-proxy-server/src/routes/core_pipeline.rs b/crates/llm-proxy-server/src/routes/core_pipeline.rs index 731bc20..a5688a3 100644 --- a/crates/llm-proxy-server/src/routes/core_pipeline.rs +++ b/crates/llm-proxy-server/src/routes/core_pipeline.rs @@ -215,6 +215,9 @@ pub(crate) struct RequestContext { pub(crate) request_id: String, /// Instant when the handler was entered (for latency metrics). pub(crate) start: Instant, + /// Authenticated client key id, or `"anonymous"` under the loopback + /// no-auth exemption. Stamped onto `RequestReceived` events. + pub(crate) caller: String, /// SHA-256 of the raw client request body, for the `RequestReceived.body_hash` /// correlation key. Computed in [`prepare_request`] over the raw wire bytes /// (where they first arrive), per the plan's locked decision, so it is shared @@ -232,17 +235,20 @@ pub(crate) struct RequestContext { /// the handler can immediately produce an error response without entering the /// core pipeline. /// -/// `request_id` is minted by the outermost `inject_request_id` layer and passed -/// in by the handler so the id attached to the `x-request-id` header matches -/// the id stamped on the request's tracing span (audit MEDIUM-1). +/// `meta.request_id` is minted by the outermost `inject_request_id` layer and +/// `meta.caller` is stamped by the client-auth middleware, so the id attached +/// to the `x-request-id` header matches the id stamped on the request's +/// tracing span (audit MEDIUM-1). An authenticated request is rate-limited +/// under its key id, an anonymous one under its client IP. pub(crate) fn prepare_request( state: &AppState, - request_id: String, + meta: &crate::middleware::RequestMeta, headers: &HeaderMap, - connect_info: Option<&std::net::SocketAddr>, body: &[u8], path: &str, ) -> Result { + let caller = &meta.caller; + let connect_info = meta.connect_info.as_ref(); let trust_forwarded_headers = state.trust_forwarded_headers(); let client_ip = get_client_ip(headers, connect_info, trust_forwarded_headers); // Security: loopback bypass is intentional for local development and testing. @@ -251,10 +257,21 @@ pub(crate) fn prepare_request( // processes on the same machine can bypass rate limiting. When // `trust_forwarded_headers` is true, the loopback check is disabled because // the peer IP may be spoofed by an intermediate reverse proxy. - let direct_loopback = - !trust_forwarded_headers && connect_info.is_some_and(|address| address.ip().is_loopback()); - - if !direct_loopback && !state.rate_limiter.is_allowed(&client_ip) { + // + // The bypass applies only to unauthenticated callers: anonymous requests + // exist solely under the loopback-only `auth.mode = "none"`, so this + // preserves today's local behavior, while a request authenticated with a + // client key is always limited under its key id. + let direct_loopback = !trust_forwarded_headers + && !caller.is_authenticated() + && connect_info.is_some_and(|address| address.ip().is_loopback()); + + let rate_limit_key = if caller.is_authenticated() { + caller.as_str().to_owned() + } else { + client_ip + }; + if !direct_loopback && !state.rate_limiter.is_allowed(&rate_limit_key) { state.metrics.record_rate_limited(); return Err(RouteError::RateLimited); } @@ -265,8 +282,9 @@ pub(crate) fn prepare_request( } Ok(RequestContext { - request_id, + request_id: meta.request_id.0.clone(), start: Instant::now(), + caller: caller.as_str().to_owned(), body_hash: body_hash_of_bytes(body), }) } @@ -684,20 +702,43 @@ fn compute_cost(usage: &Usage, pricing: &ModelPricing) -> Cost { } } +/// Resolved dispatch parameters shared by [`handle_core_once`] and +/// [`handle_core_stream`]: which provider the URL named, which route kind is +/// being served, and which client protocol the response must be encoded into. +#[derive(Debug, Clone, Copy)] +pub(crate) struct CoreDispatch<'a> { + /// Provider name from the URL path. + pub(crate) provider_name: &'a str, + /// Route kind (`chat_completions` or `messages`). + pub(crate) route_kind: ProviderRouteKind, + /// Client protocol for the response encoding. + pub(crate) client_protocol: ClientProtocol, +} + /// Non-streaming core pipeline. /// /// ```text /// CoreRequest -> encode -> send -> decode -> client encode -> HTTP response /// ``` +#[expect( + clippy::too_many_lines, + reason = "the pipeline stages run in request order and each stage's \ + metric + ResponseFailed emission sits next to the call that can \ + fail; extracting stages would separate the failure bookkeeping \ + from the step it belongs to" +)] pub(crate) async fn handle_core_once( state: AppState, ctx: RequestContext, - provider_name: &str, - route_kind: ProviderRouteKind, + dispatch: CoreDispatch<'_>, core: CoreRequest, - client_protocol: ClientProtocol, inbound_auth: Option, ) -> Result, RouteError> { + let CoreDispatch { + provider_name, + route_kind, + client_protocol, + } = dispatch; validate_provider_name_ref(provider_name)?; state.metrics.record_request(false); @@ -726,6 +767,7 @@ pub(crate) async fn handle_core_once( .emit(&ProxyEvent::RequestReceived(RequestReceived { request_id: ctx.request_id.clone(), timestamp: time::OffsetDateTime::now_utc(), + caller: ctx.caller.clone(), provider: provider_name.to_owned(), route_kind: route_kind_str(route_kind).to_owned(), client_protocol: client_protocol_str(client_protocol).to_owned(), @@ -779,19 +821,23 @@ pub(crate) async fn handle_core_once( err })?; - let response_bytes = state.proxy_client.send(proxy_req).await.map_err(|e| { - let err = map_provider_error(&e, passthrough_auth); - state.metrics.record_failure(); - emit_response_failed( - &state.event_bus, - &ctx.request_id, - Some(provider_name), - Some(&core.model), - &err, - ctx.start, - ); - err - })?; + let response_bytes = state + .upstream_client_for(&target.provider_name) + .send(proxy_req) + .await + .map_err(|e| { + let err = map_provider_error(&e, passthrough_auth); + state.metrics.record_failure(); + emit_response_failed( + &state.event_bus, + &ctx.request_id, + Some(provider_name), + Some(&core.model), + &err, + ctx.start, + ); + err + })?; // Decode the provider response into a CoreResponse. let mut core_resp = adapter @@ -992,15 +1038,24 @@ enum FirstByteResult { /// response: the spawned task sends the first event (or an error) back to /// the handler via a oneshot channel, and the handler decides whether to /// return an HTTP error or an SSE response. +#[expect( + clippy::too_many_lines, + reason = "mirrors handle_core_once for the streaming path: stage order is \ + the audit trail and each stage's failure bookkeeping stays next \ + to the failing call" +)] pub(crate) async fn handle_core_stream( state: AppState, ctx: RequestContext, - provider_name: &str, - route_kind: ProviderRouteKind, + dispatch: CoreDispatch<'_>, core: CoreRequest, - client_protocol: ClientProtocol, inbound_auth: Option, ) -> Result, RouteError> { + let CoreDispatch { + provider_name, + route_kind, + client_protocol, + } = dispatch; validate_provider_name_ref(provider_name)?; state.metrics.record_request(true); @@ -1033,6 +1088,7 @@ pub(crate) async fn handle_core_stream( .emit(&ProxyEvent::RequestReceived(RequestReceived { request_id: ctx.request_id.clone(), timestamp: time::OffsetDateTime::now_utc(), + caller: ctx.caller.clone(), provider: provider_name.to_owned(), route_kind: route_kind_str(route_kind).to_owned(), client_protocol: client_protocol_str(client_protocol).to_owned(), @@ -1087,7 +1143,7 @@ pub(crate) async fn handle_core_stream( // Open the streaming connection. let byte_stream = state - .proxy_client + .upstream_client_for(&target.provider_name) .send_stream(proxy_req) .await .map_err(|e| { @@ -1502,9 +1558,16 @@ async fn emit_encoder_finish( /// /// The `first_byte_tx` channel is used to signal the first-byte boundary: /// the spawned task sends either `FirstByteResult::FirstEvent` (once the -/// first SSE event is produced) or `FirstByteResult::PreStreamError` (if -/// the stream fails before any event). After the first event is sent, the +/// first SSE event is produced) or `FirstByteResult::PreStreamError` (if the +/// stream fails before any event). After the first event is sent, the /// channel is dropped and all subsequent errors become in-band SSE events. +#[expect( + clippy::too_many_lines, + reason = "one streaming state machine: the select-loop's cancellation, \ + decode, first-byte, and disconnect transitions read as a unit; \ + splitting the arms into helpers would scatter the state they \ + mutate" +)] fn build_sse_output_stream( byte_stream: std::pin::Pin< Box< @@ -2563,7 +2626,11 @@ mod tests { dedup_window: std::time::Duration::ZERO, server_name: "test".to_owned(), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }; let providers = llm_proxy_core::ProviderRegistry::from_providers(vec![]).expect("empty registry"); @@ -2612,59 +2679,113 @@ mod tests { fn direct_loopback_bypasses_rate_limit() { let state = state_with_operational_config(1, false); let headers = HeaderMap::new(); - let address = "127.0.0.1:12345".parse().unwrap(); + let meta = test_meta( + crate::middleware::CallerId::anonymous(), + Some("127.0.0.1:12345"), + ); - prepare_request( - &state, - "req-test".to_owned(), - &headers, - Some(&address), - b"one", - "/test", - ) - .expect("request under the rate limit should be admitted"); - prepare_request( - &state, - "req-test".to_owned(), - &headers, - Some(&address), - b"two", - "/test", - ) - .expect("request under the rate limit should be admitted"); + prepare_request(&state, &meta, &headers, b"one", "/test") + .expect("request under the rate limit should be admitted"); + prepare_request(&state, &meta, &headers, b"two", "/test") + .expect("request under the rate limit should be admitted"); + } + + /// Build a `RequestMeta` for direct `prepare_request` calls. + fn test_meta( + caller: crate::middleware::CallerId, + peer: Option<&str>, + ) -> crate::middleware::RequestMeta { + crate::middleware::RequestMeta { + request_id: crate::middleware::RequestId("req-test".to_owned()), + caller, + connect_info: peer.map(|addr| addr.parse().unwrap()), + } } #[test] fn untrusted_forwarded_headers_cannot_bypass_rate_limit() { let state = state_with_operational_config(1, false); - let address = "198.51.100.10:12345".parse().unwrap(); + let meta = test_meta( + crate::middleware::CallerId::anonymous(), + Some("198.51.100.10:12345"), + ); let mut first_headers = HeaderMap::new(); first_headers.insert("x-forwarded-for", "203.0.113.1".parse().unwrap()); let mut second_headers = HeaderMap::new(); second_headers.insert("x-forwarded-for", "203.0.113.2".parse().unwrap()); - prepare_request( - &state, - "req-test".to_owned(), - &first_headers, - Some(&address), - b"one", - "/test", - ) - .expect("request under the rate limit should be admitted"); + prepare_request(&state, &meta, &first_headers, b"one", "/test") + .expect("request under the rate limit should be admitted"); assert!(matches!( - prepare_request( - &state, - "req-test".to_owned(), - &second_headers, - Some(&address), - b"two", - "/test" - ), + prepare_request(&state, &meta, &second_headers, b"two", "/test"), Err(RouteError::RateLimited) )); } + #[test] + fn authenticated_caller_is_limited_by_key_id_across_addresses() { + // An authenticated caller's bucket is the key id, not the peer IP: a + // second request carrying the same key from a different address + // shares the budget (1 rpm here), so it is rejected. + let state = state_with_operational_config(1, false); + let caller = crate::middleware::CallerId::authenticated("laptop"); + let first_meta = test_meta(caller.clone(), Some("198.51.100.10:12345")); + let second_meta = test_meta(caller, Some("198.51.100.11:12345")); + let headers = HeaderMap::new(); + + prepare_request(&state, &first_meta, &headers, b"one", "/test") + .expect("first request under the rate limit should be admitted"); + assert!(matches!( + prepare_request(&state, &second_meta, &headers, b"two", "/test"), + Err(RouteError::RateLimited) + )); + } + + #[test] + fn authenticated_loopback_caller_is_still_rate_limited() { + // The loopback bypass exists for anonymous local clients (today's + // behavior under auth = "none"); an authenticated caller pays the + // limiter even from 127.0.0.1. + let state = state_with_operational_config(1, false); + let meta = test_meta( + crate::middleware::CallerId::authenticated("laptop"), + Some("127.0.0.1:12345"), + ); + let headers = HeaderMap::new(); + + prepare_request(&state, &meta, &headers, b"one", "/test") + .expect("first request under the rate limit should be admitted"); + assert!(matches!( + prepare_request(&state, &meta, &headers, b"two", "/test"), + Err(RouteError::RateLimited) + )); + } + + #[test] + fn request_context_carries_caller_attribution() { + let state = state_with_operational_config(100, false); + let headers = HeaderMap::new(); + let ctx = prepare_request( + &state, + &test_meta(crate::middleware::CallerId::anonymous(), None), + &headers, + b"", + "/test", + ) + .expect("request admitted"); + assert_eq!(ctx.caller, "anonymous"); + + let ctx = prepare_request( + &state, + &test_meta(crate::middleware::CallerId::authenticated("laptop"), None), + &headers, + b"", + "/test", + ) + .expect("request admitted"); + assert_eq!(ctx.caller, "laptop"); + } + // -- map_provider_error ----------------------------------------------------- #[test] @@ -2757,7 +2878,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: std::time::Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }; let providers = @@ -2827,7 +2952,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: std::time::Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }; let providers = diff --git a/crates/llm-proxy-server/src/routes/messages.rs b/crates/llm-proxy-server/src/routes/messages.rs index 5a98ea1..ac8ef12 100644 --- a/crates/llm-proxy-server/src/routes/messages.rs +++ b/crates/llm-proxy-server/src/routes/messages.rs @@ -12,14 +12,14 @@ use std::sync::Arc; use axum::Json; use axum::body::Body; -use axum::extract::{Extension, Path, State}; +use axum::extract::{Path, State}; use axum::http::{HeaderMap, Response}; use llm_proxy_core::ProviderRouteKind; use llm_proxy_protocol::anthropic::MessageRequest; use llm_proxy_protocol::client::anthropic; use tracing::{info, warn}; -use crate::middleware::{OptionalConnectInfo, RequestId}; +use crate::middleware::RequestMeta; use crate::state::AppState; use super::core_pipeline; @@ -32,8 +32,7 @@ use super::error_response::{ClientProtocol, RouteError, route_error_response}; pub(crate) async fn handle_messages( State(state): State, Path(provider): Path, - Extension(req_id): Extension, - OptionalConnectInfo(connect_info): OptionalConnectInfo, + meta: RequestMeta, headers: HeaderMap, body: axum::body::Bytes, ) -> Response { @@ -44,7 +43,7 @@ pub(crate) async fn handle_messages( // blanket emit here would double-count every pipeline failure (the pipeline // emits AND returns Err via `?`), so the wrapper must stay emit-free // (audit route-responsefailed-gaps regression). - match handle_messages_inner(state, req_id, provider, connect_info, headers, body).await { + match handle_messages_inner(state, meta, provider, headers, body).await { Ok(response) => response, Err(error) => { warn!(error = %error, "request failed"); @@ -62,9 +61,8 @@ pub(crate) async fn handle_messages( /// function propagates without re-emitting. The outer wrapper renders only. async fn handle_messages_inner( state: AppState, - req_id: RequestId, + meta: RequestMeta, provider: String, - connect_info: Option, headers: HeaderMap, body: axum::body::Bytes, ) -> Result, RouteError> { @@ -72,7 +70,7 @@ async fn handle_messages_inner( let event_bus = Arc::clone(&state.event_bus); // The request id is consumed by prepare_request below; snapshot it once so // every early-gate emit references the same id. - let request_id = req_id.0.clone(); + let request_id = meta.request_id.0.clone(); // Every early gate emits exactly one ResponseFailed with the same // provider/request/start triple; bind it once instead of repeating the // six-argument call at each gate. @@ -106,15 +104,8 @@ async fn handle_messages_inner( // Pre-flight: rate limit, dedup, request ID. let request_path = format!("/providers/{provider}/v1/messages"); - let ctx = core_pipeline::prepare_request( - &state, - req_id.0, - &headers, - connect_info.as_ref(), - &body, - &request_path, - ) - .inspect_err(&emit_failed)?; + let ctx = core_pipeline::prepare_request(&state, &meta, &headers, &body, &request_path) + .inspect_err(&emit_failed)?; // Parse the Anthropic MessageRequest via the axum Json helper so that // syntax-vs-data parse failures are distinguished into precise messages @@ -153,28 +144,15 @@ async fn handle_messages_inner( // Extract the inbound client auth token (used by passthrough-auth providers). let inbound_auth = core_pipeline::extract_inbound_auth(&headers); + let dispatch = core_pipeline::CoreDispatch { + provider_name: &provider, + route_kind: ProviderRouteKind::Messages, + client_protocol: ClientProtocol::Anthropic, + }; if is_streaming { - core_pipeline::handle_core_stream( - state, - ctx, - &provider, - ProviderRouteKind::Messages, - core, - ClientProtocol::Anthropic, - inbound_auth, - ) - .await + core_pipeline::handle_core_stream(state, ctx, dispatch, core, inbound_auth).await } else { - core_pipeline::handle_core_once( - state, - ctx, - &provider, - ProviderRouteKind::Messages, - core, - ClientProtocol::Anthropic, - inbound_auth, - ) - .await + core_pipeline::handle_core_once(state, ctx, dispatch, core, inbound_auth).await } } diff --git a/crates/llm-proxy-server/src/routes/mod.rs b/crates/llm-proxy-server/src/routes/mod.rs index 1b67671..fa4cfd5 100644 --- a/crates/llm-proxy-server/src/routes/mod.rs +++ b/crates/llm-proxy-server/src/routes/mod.rs @@ -16,9 +16,10 @@ use tower_http::{ }; use tracing::info_span; -use crate::middleware::{RequestId, RequestIdGenerator}; +use crate::middleware::{CallerId, RequestId, RequestIdGenerator, extract_bearer_token}; use crate::state::AppState; use llm_proxy_storage::EventBus; +use secrecy::ExposeSecret; mod chat; mod core_pipeline; @@ -86,6 +87,10 @@ pub fn router(state: AppState) -> Router { // rejections produced by the timeout/body-limit layers (audit finding: layer // fallbacks emit no terminal event). let event_bus = Arc::clone(&state.event_bus); + // The client-auth layer needs the full AppState (auth state, rate limiter, + // metrics), so the state is cloned once for it before being moved into + // `with_state` below. + let state_arc = state.clone(); // Span factory shared by both routers: it reads the id injected by the // outermost layer so every request span carries `request_id`, matching the @@ -103,26 +108,31 @@ pub fn router(state: AppState) -> Router { .layer(trace.clone()), ); - // API routes: request-id injection, error normalisation, body limit, - // tracing, timeout. + // API routes: request-id injection, client auth, error normalisation, + // body limit, tracing, timeout. // // Layer order (outermost first): // 1. `inject_request_id` -- generates the id and stamps it into extensions - // so the span (step 4) and the handler see the same id. - // 2. `normalize_error_responses` -- rewrites the 408 (timeout) and 413 + // so the span (step 5) and the handler see the same id. + // 2. `require_client_auth` -- under `auth.mode = "client-keys"`, rejects + // requests without a valid `Authorization: Bearer ` with a + // 401 before any body buffering. Under `auth.mode = "none"` it stamps + // the anonymous caller and passes through. + // 3. `normalize_error_responses` -- rewrites the 408 (timeout) and 413 // (body-limit) rejections emitted by layers below into protocol-shaped // JSON envelopes carrying `x-request-id` (audit MEDIUM-7). It runs // outside the body-limit/timeout layers so their rejections pass // through it on the way out. - // 3. `DefaultBodyLimit` -- caps the body the JSON extractors can buffer. - // 4. `TraceLayer` -- opens the request span (carrying the request id). - // 5. `TimeoutLayer` -- cancels handler work exceeding `request_timeout`. + // 4. `DefaultBodyLimit` -- caps the body the JSON extractors can buffer. + // 5. `TraceLayer` -- opens the request span (carrying the request id). + // 6. `TimeoutLayer` -- cancels handler work exceeding `request_timeout`. // // Because the timeout and body-limit layers emit non-JSON, header-less - // rejections, step 2 normalises their final responses into the same + // rejections, step 3 normalises their final responses into the same // Anthropic/OpenAI-shaped JSON schema every handler uses. let api_middleware = ServiceBuilder::new() .layer(from_fn_with_state(id_gen, inject_request_id)) + .layer(from_fn_with_state(state_arc, require_client_auth)) .layer(from_fn_with_state(event_bus, normalize_error_responses)) .layer(DefaultBodyLimit::max(MAX_BODY_BYTES)) .layer(trace) @@ -153,6 +163,23 @@ pub fn router(state: AppState) -> Router { // (audit LOW-27). .method_not_allowed_fallback(method_not_allowed); + // The admin `/usage` surface joins the MAIN listener only under the + // explicit `server.admin_on_main_listener` opt-in — a distinct, named + // security posture the serve path warns about loudly at startup. The + // default posture serves `/usage` on the dedicated admin listener + // ([`crate::admin_router`]) and leaves this listener without any admin + // surface. Config validation guarantees the flag implies `[auth.admin]`; + // `admin_router` returning `None` here keeps the default closed. The + // merge happens after `with_state` below because the usage router is + // already stateless. + let usage_on_main = match ( + crate::admin_router(&state), + state.app_config.server.admin_on_main_listener, + ) { + (Some(usage), true) => Some(usage), + _ => None, + }; + // Config-gated CORS (audit LOW-12). A `CorsLayer` is installed only when // the operator explicitly lists cross-origin callers, and then only for // those exact origins. This replaces a former `CorsLayer::very_permissive()` @@ -165,7 +192,7 @@ pub fn router(state: AppState) -> Router { None => router, }; - router + let router: Router = router // Global, outermost middleware (audit LOW-11, LOW-12) -- applied to the // merged router so they cover every route and both fallbacks. CORS sits // just inside these so a preflight still receives the headers below and @@ -182,7 +209,15 @@ pub fn router(state: AppState) -> Router { HeaderValue::from_static("nosniff"), )) .layer(CatchPanicLayer::new()) - .with_state(state) + .with_state(state); + + // Both routers are stateless here, so the opt-in usage surface merges + // cleanly (see `usage_on_main` above). It carries its own admin-auth, + // nosniff, panic-catch, and tracing layers. + match usage_on_main { + Some(usage) => router.merge(usage), + None => router, + } } /// Build a restrictive `CorsLayer` for an explicit allow-origin list, or return @@ -250,6 +285,117 @@ async fn inject_request_id( next.run(req).await } +/// Client-key authentication for proxy routes. +/// +/// Under `auth.mode = "client-keys"`, a request must carry +/// `Authorization: Bearer ` matching a configured +/// `[auth.clients]` verifier; the matching key id is stamped into extensions +/// as [`CallerId`] for event attribution and rate-limit keying. Anything +/// else — missing header, unknown scheme, wrong token — gets a 401 whose +/// body carries no secret material (the standard `Unauthorized` envelope). +/// +/// Under `auth.mode = "none"` (loopback binds only, enforced by config +/// validation at startup), every request is stamped +/// [`CallerId::anonymous`] and passes through unchanged. +/// +/// The bearer token is never logged and never echoed in the response. +/// +/// # Garbage-token flood bound +/// +/// Every failed attempt (missing or invalid token) is charged to the peer's +/// IP rate-limit bucket, and a peer whose bucket is empty is rejected +/// **before** the CPU-costly argon2 verification runs. A flood of garbage +/// tokens from one peer therefore pays at most `rate_limit_rpm` verifications +/// per minute (times the number of configured keys) instead of one per +/// request, and the blocking pool cannot be pinned by unauthenticated +/// traffic. Successful authentications are never charged to the IP bucket — +/// their requests are limited under the key id inside `prepare_request` — +/// and the response stays a uniform 401 on every rejection path, so the gate +/// leaks nothing about which check failed. When rate limiting is disabled +/// (`rate_limit_rpm = 0`) no bound applies; the operator turned the limiter +/// off. A failed attempt also emits a `ResponseFailed` so rejected requests +/// stay attributable in the event log, mirroring the 404/405/408/413 +/// layer-rejection pattern. +async fn require_client_auth(State(state): State, req: Request, next: Next) -> Response { + if !state.client_auth.is_enforced() { + let mut req = req; + req.extensions_mut().insert(CallerId::anonymous()); + return next.run(req).await; + } + + // Peer identity for the failure-budget gate; captured before the request + // is consumed by verification or the inner service. + let connect_info = req + .extensions() + .get::>() + .map(|ci| ci.0); + let client_ip = crate::middleware::get_client_ip( + req.headers(), + connect_info.as_ref(), + state.trust_forwarded_headers(), + ); + + let authenticated = match extract_bearer_token(req.headers()) { + Some(token) => { + // Flood gate BEFORE verification: a peer that already exhausted + // its failure budget is rejected without paying for argon2. The + // response is the same uniform 401 the verification failure + // produces, so the gate is not observable to the peer. + if !state.rate_limiter.would_allow(&client_ip) { + state.metrics.record_rate_limited(); + None + } else { + state.client_auth.authenticate(token.expose_secret()).await + } + } + None => None, + }; + + let mut req = req; + match authenticated { + Some(key_id) => { + tracing::debug!(caller = %key_id, "client authenticated"); + req.extensions_mut().insert(CallerId::authenticated(key_id)); + next.run(req).await + } + None => { + // Charge the failure to the peer's IP bucket so repeated failures + // drain it and trip the pre-verification gate above. When the + // bucket is already empty this consumes nothing and records the + // limiter metric. + if !state.rate_limiter.is_allowed(&client_ip) { + state.metrics.record_rate_limited(); + } + let request_id = req + .extensions() + .get::() + .cloned() + .unwrap_or_else(|| RequestId("unknown".to_owned())); + let path = req.uri().path().to_owned(); + // Drain the body so a rejected connection is not left half-read. + if let Err(e) = axum::body::to_bytes(req.into_body(), NOT_FOUND_BODY_DRAIN_LIMIT).await + { + tracing::trace!(error = %e, "draining body of auth-rejected request failed"); + } + let error = error_response::RouteError::Unauthorized; + core_pipeline::emit_response_failed( + &state.event_bus, + &request_id.0, + None, + None, + &error, + std::time::Instant::now(), + ); + let mut response = + error_response::route_error_response(protocol_for_path(&path), &error); + if let Ok(value) = request_id.0.parse() { + response.headers_mut().insert("x-request-id", value); + } + response + } + } +} + /// Build the per-request tracing span, carrying the injected request id. /// /// `make_span_with` runs after [`inject_request_id`] has stamped the id into @@ -491,3 +637,137 @@ mod tests { assert!(cors_layer_for(Some(&origins)).is_some()); } } + +/// Tests for the garbage-token flood gate in [`require_client_auth`]. +/// +/// The gate's purpose: a peer whose failure budget is exhausted is rejected +/// BEFORE the argon2 verification runs, so an unauthenticated flood cannot +/// pin the blocking pool. The peer-visible response is the same uniform 401 +/// either way, so the gate is proven here through its one observable +/// consequence: a peer that burned its budget cannot authenticate even with a +/// VALID token until the bucket refills. +#[cfg(test)] +mod auth_gate_tests { + use super::*; + + use axum::body::Body; + use axum::http::Request; + use tower::ServiceExt; + + /// Build an AppState with enforced client auth for `token` and the given + /// per-IP rate limit, plus a minimal router whose probe handler answers + /// 200 when a request makes it past the auth layer. + fn auth_router(rate_limit_rpm: u32) -> (Router, AppState) { + use argon2::{Algorithm, Argon2, CustomizedPasswordHasher, Params, Version}; + + let params = Params::new(8, 1, 1, None).expect("valid test params"); + let hasher = Argon2::new(Algorithm::Argon2id, Version::V0x13, params.clone()); + let phc = hasher + .hash_password_customized( + b"sekrit-token", + b"0123456789abcdef", + Some(&argon2::ARGON2ID_IDENT), + None, + params, + ) + .expect("hash token") + .to_string(); + + let mut clients = std::collections::HashMap::new(); + clients.insert("laptop".to_owned(), secrecy::SecretString::from(phc)); + let app_config = llm_proxy_core::AppConfig { + server: llm_proxy_core::ServerConfig { + bind: "127.0.0.1:3456".parse().unwrap(), + request_timeout: std::time::Duration::from_secs(300), + shutdown_timeout: std::time::Duration::from_secs(30), + log_level: "info".to_owned(), + hot_reload: false, + allowed_origins: None, + server_name: "gate-test".to_owned(), + rate_limit_rpm, + trust_forwarded_headers: false, + dedup_window: std::time::Duration::ZERO, + log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, + }, + auth: llm_proxy_core::AuthConfig { + mode: llm_proxy_core::AuthMode::ClientKeys, + clients, + admin: None, + }, + storage: None, + }; + + let state = crate::state::AppState::new( + app_config, + llm_proxy_core::ProviderRegistry::from_providers(vec![]).expect("empty registry"), + llm_proxy_provider::ProviderAdapterRegistry::builtin(), + llm_proxy_provider::ProxyClient::new(), + crate::state::BuildInfo { + name: "test", + version: "0.0.0", + target: "test", + git_sha: "test", + }, + ); + + let router = Router::new() + .route( + "/providers/x/v1/messages", + post(|| async { StatusCode::OK }), + ) + .layer(from_fn_with_state( + state.clone(), + super::require_client_auth, + )); + (router, state) + } + + fn bearer_request(token: &str) -> Request { + Request::builder() + .method("POST") + .uri("/providers/x/v1/messages") + .header("authorization", format!("Bearer {token}")) + .body(Body::empty()) + .unwrap() + } + + #[tokio::test] + async fn valid_token_passes_when_the_peer_has_budget() { + let (app, _state) = auth_router(100); + let response = app.oneshot(bearer_request("sekrit-token")).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + } + + #[tokio::test] + async fn exhausted_failure_budget_gates_verification_before_argon2() { + // rate_limit_rpm = 1: one failed attempt drains the peer's budget. + // `oneshot` consumes the router, so each request clones it; the clones + // share the same underlying state (one rate-limit bucket). + let (app, state) = auth_router(1); + + // Failure #1 pays for verification and charges the last token. + let first = app + .clone() + .oneshot(bearer_request("wrong-token")) + .await + .unwrap(); + assert_eq!(first.status(), StatusCode::UNAUTHORIZED); + assert_eq!(state.metrics.get_snapshot().rate_limited, 0); + + // Failure #2 carries a VALID token: the bucket is empty, so the gate + // rejects before verification. If the argon2 check ran first, this + // request would authenticate and the probe handler would answer 200. + let second = app.oneshot(bearer_request("sekrit-token")).await.unwrap(); + assert_eq!( + second.status(), + StatusCode::UNAUTHORIZED, + "a peer with an exhausted failure budget must be gated before verification" + ); + assert!( + state.metrics.get_snapshot().rate_limited >= 1, + "the gate records a limiter rejection" + ); + } +} diff --git a/crates/llm-proxy-server/src/routes/models.rs b/crates/llm-proxy-server/src/routes/models.rs index 0314d93..9ea0dc7 100644 --- a/crates/llm-proxy-server/src/routes/models.rs +++ b/crates/llm-proxy-server/src/routes/models.rs @@ -323,6 +323,7 @@ mod tests { api_key: secrecy::SecretString::from("sk-test"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -370,7 +371,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: std::time::Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }, registry, llm_proxy_provider::ProviderAdapterRegistry::builtin(), @@ -391,6 +396,7 @@ mod tests { api_key: secrecy::SecretString::from("sk-test"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -415,7 +421,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: std::time::Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }, registry, llm_proxy_provider::ProviderAdapterRegistry::builtin(), diff --git a/crates/llm-proxy-server/src/routes/token_count.rs b/crates/llm-proxy-server/src/routes/token_count.rs index c3f17e7..04b1539 100644 --- a/crates/llm-proxy-server/src/routes/token_count.rs +++ b/crates/llm-proxy-server/src/routes/token_count.rs @@ -8,7 +8,7 @@ use std::sync::Arc; use axum::body::Body; use axum::extract::rejection::JsonRejection; -use axum::extract::{Extension, Path, State}; +use axum::extract::{Path, State}; use axum::http::{HeaderMap, Response}; use axum::response::IntoResponse; use llm_proxy_core::MessageContent; @@ -18,7 +18,7 @@ use llm_proxy_protocol::core::{StopReason, Usage}; use llm_proxy_storage::{ProxyEvent, RequestReceived, ResponseCompleted}; use serde::Serialize; -use crate::middleware::{OptionalConnectInfo, RequestId}; +use crate::middleware::RequestMeta; use crate::state::AppState; use super::core_pipeline; @@ -54,29 +54,17 @@ pub(crate) struct TokenCountResponse { pub(crate) async fn count_tokens( State(state): State, Path(provider): Path, - Extension(req_id): Extension, - OptionalConnectInfo(connect_info): OptionalConnectInfo, + meta: RequestMeta, headers: HeaderMap, body: axum::body::Bytes, ) -> Response { - let start = std::time::Instant::now(); // The wrapper only renders: it does NOT emit ResponseFailed. token_count // has no core-pipeline dispatch, so each early-gate error return inside the // inner function emits exactly one ResponseFailed itself, and the success // path emits exactly one ResponseCompleted. A blanket emit here would // double-count, so the wrapper stays emit-free // (audit route-responsefailed-gaps regression). - match count_tokens_inner( - &state, - req_id, - &provider, - connect_info.as_ref(), - &headers, - body, - start, - ) - .await - { + match count_tokens_inner(&state, meta, &provider, &headers, body).await { Ok(response) => response, Err(error) => { tracing::warn!(error = %error, "request failed"); @@ -99,15 +87,100 @@ fn json_rejection_message(rejection: &JsonRejection) -> String { } } +/// Extract the text a token estimate counts from a decoded [`CoreRequest`]: +/// the concatenated system text and one [`MessageContent`] per message. +/// +/// NOTE: This estimate intentionally only counts text from `CoreContent::Text` +/// blocks. The following are NOT counted: +/// - Tool definitions (`core.tools`) -- their name, description, and +/// input_schema text fields are excluded. +/// - Non-text content blocks (images, documents, etc.). +/// - Tool result content blocks. +/// - System prompt non-text blocks. +/// +/// TODO(future): Extend counting to include tool definitions, tool_use blocks, +/// and tool_result blocks for a more accurate estimate. +fn extract_countable_text( + core: &llm_proxy_protocol::core::CoreRequest, +) -> (String, Vec) { + // Collect system text without intermediate Vec allocation. + let system_text: String = core + .system + .iter() + .filter_map(|block| match block { + llm_proxy_protocol::core::CoreContent::Text { text, .. } => Some(text.as_str()), + _ => None, + }) + .fold(String::new(), |mut acc, s| { + acc.push_str(s); + acc + }); + + let messages: Vec = core + .messages + .iter() + .map(|msg| { + let text: String = msg + .content + .iter() + .filter_map(|block| match block { + llm_proxy_protocol::core::CoreContent::Text { text, .. } => Some(text.as_str()), + _ => None, + }) + .collect(); + // Non-text content blocks (images, tool_use, thinking) produce + // empty text here. This is a known limitation of the heuristic + // counter. Log at debug level so operators can detect requests + // where the estimate may be significantly off. + if text.is_empty() && !msg.content.is_empty() { + tracing::debug!( + role = ?msg.role, + blocks = msg.content.len(), + "message has content blocks but no extractable text; token estimate will be zero for this message" + ); + } + let role_name = match msg.role { + llm_proxy_protocol::core::CoreRole::User => "user", + llm_proxy_protocol::core::CoreRole::Assistant => "assistant", + llm_proxy_protocol::core::CoreRole::System => "system", + // The Anthropic Messages API does not have a 'tool' role. + // Tool results are carried in 'user' role messages with + // tool_result content blocks. If a CoreRole::Tool reaches + // here, treat it as 'user' since tool_result content is + // typically nested under the user role in Anthropic. + llm_proxy_protocol::core::CoreRole::Tool => { + tracing::debug!( + "CoreRole::Tool mapped to 'user' for token counting; \ + tool_result content is counted under user role" + ); + "user" + } + // All known variants are handled above. This catch-all exists + // for forward compatibility. Unknown roles are treated as + // 'user' since that is the least-lossy mapping. + _ => { + tracing::warn!( + role = ?msg.role, + "unknown CoreRole variant in token count, mapping to 'user'" + ); + "user" + } + }; + MessageContent::new(role_name, text) + }) + .collect(); + + (system_text, messages) +} + async fn count_tokens_inner( state: &AppState, - req_id: RequestId, + meta: RequestMeta, provider: &str, - connect_info: Option<&std::net::SocketAddr>, headers: &HeaderMap, body: axum::body::Bytes, - start: std::time::Instant, ) -> Result, RouteError> { + let start = std::time::Instant::now(); // Pre-flight: rate limit, dedup, request ID. token_count has no core-pipeline // dispatch, so each EARLY-GATE error return below emits exactly one // ResponseFailed itself, and the success path emits exactly one @@ -115,7 +188,7 @@ async fn count_tokens_inner( let event_bus = Arc::clone(&state.event_bus); // The request id is consumed by prepare_request below; snapshot it once so // every early-gate emit references the same id. - let request_id = req_id.0.clone(); + let request_id = meta.request_id.0.clone(); // Every early gate emits exactly one ResponseFailed with the same // provider/request/start triple; bind it once instead of repeating the // six-argument call at each gate. @@ -139,15 +212,8 @@ async fn count_tokens_inner( // misreported as a JSON syntax error (audit LOW-30). core_pipeline::validate_json_content_type(headers).inspect_err(&emit_failed)?; let request_path = format!("/providers/{provider}/v1/messages/count_tokens"); - let ctx = core_pipeline::prepare_request( - state, - req_id.0, - headers, - connect_info, - &body, - &request_path, - ) - .inspect_err(&emit_failed)?; + let ctx = core_pipeline::prepare_request(state, &meta, headers, &body, &request_path) + .inspect_err(&emit_failed)?; // Parse and validate the Anthropic MessageRequest. // // Go through `axum::Json::from_bytes` (rather than `serde_json::from_slice`) @@ -183,6 +249,7 @@ async fn count_tokens_inner( .emit(&ProxyEvent::RequestReceived(RequestReceived { request_id: ctx.request_id.clone(), timestamp: time::OffsetDateTime::now_utc(), + caller: ctx.caller.clone(), provider: provider.to_owned(), route_kind: "count_tokens".to_owned(), client_protocol: "anthropic".to_owned(), @@ -192,83 +259,7 @@ async fn count_tokens_inner( })); // Extract text content from core messages for token counting. - // - // NOTE: This estimate intentionally only counts text from `CoreContent::Text` - // blocks. The following are NOT counted: - // - Tool definitions (`core.tools`) -- their name, description, and - // input_schema text fields are excluded. - // - Non-text content blocks (images, documents, etc.). - // - Tool result content blocks. - // - System prompt non-text blocks. - // - // TODO(future): Extend counting to include tool definitions, tool_use blocks, - // and tool_result blocks for a more accurate estimate. - // Collect system text without intermediate Vec allocation. - let system_text: String = core - .system - .iter() - .filter_map(|block| match block { - llm_proxy_protocol::core::CoreContent::Text { text, .. } => Some(text.as_str()), - _ => None, - }) - .fold(String::new(), |mut acc, s| { - acc.push_str(s); - acc - }); - - let messages: Vec = core - .messages - .iter() - .map(|msg| { - let text: String = msg - .content - .iter() - .filter_map(|block| match block { - llm_proxy_protocol::core::CoreContent::Text { text, .. } => Some(text.as_str()), - _ => None, - }) - .collect(); - // Non-text content blocks (images, tool_use, thinking) produce - // empty text here. This is a known limitation of the heuristic - // counter. Log at debug level so operators can detect requests - // where the estimate may be significantly off. - if text.is_empty() && !msg.content.is_empty() { - tracing::debug!( - role = ?msg.role, - blocks = msg.content.len(), - "message has content blocks but no extractable text; token estimate will be zero for this message" - ); - } - let role_name = match msg.role { - llm_proxy_protocol::core::CoreRole::User => "user", - llm_proxy_protocol::core::CoreRole::Assistant => "assistant", - llm_proxy_protocol::core::CoreRole::System => "system", - // The Anthropic Messages API does not have a 'tool' role. - // Tool results are carried in 'user' role messages with - // tool_result content blocks. If a CoreRole::Tool reaches - // here, treat it as 'user' since tool_result content is - // typically nested under the user role in Anthropic. - llm_proxy_protocol::core::CoreRole::Tool => { - tracing::debug!( - "CoreRole::Tool mapped to 'user' for token counting; \ - tool_result content is counted under user role" - ); - "user" - } - // All known variants are handled above. This catch-all exists - // for forward compatibility. Unknown roles are treated as - // 'user' since that is the least-lossy mapping. - _ => { - tracing::warn!( - role = ?msg.role, - "unknown CoreRole variant in token count, mapping to 'user'" - ); - "user" - } - }; - MessageContent::new(role_name, text) - }) - .collect(); + let (system_text, messages) = extract_countable_text(&core); // BPE tokenization is CPU-bound (fancy-regex driven) and the request body // is attacker-controlled up to MAX_BODY_BYTES (32 MiB). Running it on the diff --git a/crates/llm-proxy-server/src/state.rs b/crates/llm-proxy-server/src/state.rs index d439402..e30772f 100644 --- a/crates/llm-proxy-server/src/state.rs +++ b/crates/llm-proxy-server/src/state.rs @@ -1,3 +1,4 @@ +use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; @@ -6,7 +7,7 @@ use llm_proxy_provider::{ProviderAdapterRegistry, ProxyClient}; use llm_proxy_storage::{EventBus, NoopBus}; use crate::ModelCatalogService; -use crate::middleware::{RateLimiter, RequestDeduplicator, RequestIdGenerator}; +use crate::middleware::{ClientAuth, RateLimiter, RequestDeduplicator, RequestIdGenerator}; /// Static information about this binary. #[derive(Debug, Clone)] @@ -47,9 +48,15 @@ pub struct BuildInfo { /// `ProviderAdapter` types whose `api_key` is `SecretString`; the adapter /// config/target types keep a *manual* `Debug` that additionally redacts the /// endpoint query string and auth-header *values* (not just the key). Safe. -/// - `AppState` -> remaining fields: `BuildInfo`, `proxy_client`, `Counter`, -/// `Metrics`, `RateLimiter`, `RequestDeduplicator`, `RequestIdGenerator`, -/// `model_catalogs`, `event_bus` -- none hold secrets. Safe. +/// - `AppState` -> `client_auth` field: `ClientAuth` holds **hashed argon2id +/// verifiers only, never plaintext client secrets**; its manual `Debug` +/// prints the mode and key ids and never the verifier strings (a PHC hash +/// is offline-crackable, so it is treated as sensitive). Safe. +/// - `AppState` -> remaining fields: `BuildInfo`, `proxy_client`, +/// `egress_clients`, `Counter`, `Metrics`, `RateLimiter`, +/// `RequestDeduplicator`, `RequestIdGenerator`, `model_catalogs`, +/// `event_bus` -- none hold secrets (the egress map keys are provider +/// names and `ProxyClient`'s Debug prints no URLs). Safe. /// /// **Maintenance note:** Secret-backed fields use `secrecy::SecretString`, which /// redacts in `Debug`/`Display` regardless of whether the containing type @@ -69,6 +76,13 @@ pub struct AppState { pub(crate) provider_adapters: Arc, /// Protocol-neutral HTTP transport client. pub(crate) proxy_client: Arc, + /// Per-provider egress clients for providers with `[provider.egress]` + /// configured, keyed by provider name. Empty by default, in which case + /// every provider uses `proxy_client` -- exactly the behavior from before + /// egress support existed. The serve path builds this map from the + /// provider registry before the listener binds, so an unusable proxy URL + /// is a startup error, not a runtime fallback. + pub(crate) egress_clients: Arc>>, /// Build identifier reported by `/version`. pub(crate) build: Arc, /// Token counter for estimating token usage. @@ -92,6 +106,13 @@ pub struct AppState { /// Structured request/response event sink. Defaults to [`NoopBus`]; a real /// persisted backend lands in a later phase. pub(crate) event_bus: Arc, + /// Client-key authentication state: hashed argon2id verifiers plus a + /// digest cache. Built from `[auth]` config; construction from an invalid + /// auth config fails closed (enforced, nothing can authenticate) rather + /// than silently downgrading to no-auth, because the constructors are + /// infallible. `cmd_serve` surfaces invalid verifiers as a startup error + /// via `ClientAuth::from_config` before the state is built. + pub(crate) client_auth: Arc, } impl std::fmt::Debug for AppState { @@ -111,6 +132,7 @@ impl std::fmt::Debug for AppState { .field("providers", &self.providers) .field("provider_adapters", &self.provider_adapters) .field("proxy_client", &self.proxy_client) + .field("egress_clients", &self.egress_clients) .field("build", &self.build) .field("token_counter", &self.token_counter) .field("metrics", &self.metrics) @@ -119,6 +141,7 @@ impl std::fmt::Debug for AppState { .field("request_id_gen", &self.request_id_gen) .field("model_catalogs", &self.model_catalogs) .field("event_bus", &self.event_bus) + .field("client_auth", &self.client_auth) .finish() } } @@ -159,11 +182,28 @@ impl AppState { // matching the documented "0s disables deduplication" contract. let dedup_window_ms = u64::try_from(app_config.server.dedup_window.as_millis()).unwrap_or(0); + // Auth state from `[auth]` config. Construction is infallible, so an + // invalid auth config (e.g. an unparseable verifier) fails CLOSED -- + // enforced with zero verifiers, rejecting every proxied request -- and + // logs the error, rather than silently downgrading to no-auth. The + // serve path validates the auth config first (`ClientAuth::from_config` + // in `load_toml_state`) and turns this into a startup error. + let client_auth = match ClientAuth::from_config(&app_config.auth) { + Ok(auth) => Arc::new(auth), + Err(error) => { + tracing::error!( + error = %error, + "invalid [auth] config; failing closed (all proxy requests will be rejected)" + ); + Arc::new(ClientAuth::fail_closed()) + } + }; Self { app_config: Arc::new(app_config), providers: Arc::new(providers), provider_adapters: Arc::new(provider_adapters), proxy_client: Arc::new(proxy_client), + egress_clients: Arc::new(HashMap::new()), build: Arc::new(build), token_counter: Counter::new(), metrics: Arc::new(Metrics::new()), @@ -172,9 +212,36 @@ impl AppState { request_id_gen: Arc::new(RequestIdGenerator::new()), model_catalogs: Arc::new(ModelCatalogService::new(catalog_dir)), event_bus: Arc::new(NoopBus) as Arc, + client_auth, } } + /// Replace the client-auth state (builder-style). Used by `cmd_serve` + /// (after `ClientAuth::from_config` succeeds, so an invalid verifier is a + /// startup error rather than a runtime fail-closed) and by tests that + /// exercise authenticated routing. + #[must_use] + pub fn with_client_auth(mut self, client_auth: Arc) -> Self { + self.client_auth = client_auth; + self + } + + /// Replace the per-provider egress clients (builder-style). + /// + /// `cmd_serve` builds this map from `[provider.egress]` entries in the + /// provider registry before the listener binds, so a proxy URL reqwest + /// cannot use is a startup error; tests inject proxied clients the same + /// way. An empty map (the default) keeps every provider on the shared + /// `proxy_client`, exactly as before egress support existed. + #[must_use] + pub fn with_egress_clients( + mut self, + egress_clients: HashMap>, + ) -> Self { + self.egress_clients = Arc::new(egress_clients); + self + } + /// Replace the event bus (builder-style). Used by tests to inject a /// [`RecordingBus`](llm_proxy_storage::RecordingBus); production wires the /// real bus in `cmd_serve` (still `NoopBus` for now). @@ -226,11 +293,50 @@ impl AppState { &self.providers } + /// The transport client to use for an upstream request to `provider`. + /// + /// Providers with a wired `[provider.egress]` client get that client; + /// every other provider gets the shared default client. When no provider + /// has egress configured (the common case) this returns the shared client + /// for every lookup, preserving the single-client behavior from before + /// egress support existed. + /// + /// If a provider's config carries an egress section but no client was + /// wired for it, that is a wiring bug: the request still goes out, but a + /// loud warning names the provider on every affected request so the + /// mis-wiring cannot silently turn into direct egress. The serve path + /// builds the map from the same registry, so production cannot hit this. + pub fn upstream_client_for(&self, provider: &str) -> &ProxyClient { + if let Some(client) = self.egress_clients.get(provider) { + return client; + } + if self + .providers + .get(provider) + .is_some_and(|config| config.egress.is_some()) + { + tracing::warn!( + provider, + "provider has [provider.egress] configured but no egress client was wired; \ + its upstream traffic is leaving directly, not through the configured proxy" + ); + } + &self.proxy_client + } + /// Access the provider model catalog service. pub fn model_catalogs(&self) -> &ModelCatalogService { &self.model_catalogs } + /// Access the request/response event sink. + /// + /// Holds the bus so the serve path can [`EventBus::close`] it (flush the + /// usage ledger) after the HTTP server has stopped. + pub fn event_bus(&self) -> &Arc { + &self.event_bus + } + /// Return the bind address the server should listen on. pub fn bind_address(&self) -> std::net::SocketAddr { self.app_config.server.bind @@ -245,7 +351,6 @@ impl AppState { #[cfg(test)] mod tests { use super::*; - use std::collections::HashMap; use std::net::SocketAddr; /// Helper: build a minimal TOML AppConfig for tests. @@ -264,7 +369,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, } } @@ -298,6 +407,64 @@ mod tests { assert_eq!(state.shutdown_timeout(), Duration::from_secs(30)); } + // -- Per-provider egress client selection --------------------------------- + + #[test] + fn upstream_client_for_returns_the_shared_client_without_egress() { + let state = AppState::new( + make_app_config(), + make_provider_registry(), + ProviderAdapterRegistry::builtin(), + ProxyClient::new(), + make_build_info(), + ); + // Same object for every provider: the single-client behavior from + // before egress support existed. + for provider in ["opencode-go", "opencode-zen", "anything-else"] { + assert!( + std::ptr::eq( + state.upstream_client_for(provider), + state.proxy_client.as_ref() + ), + "provider {provider} without egress must use the shared client" + ); + } + } + + #[test] + fn upstream_client_for_returns_the_wired_client_only_for_that_provider() { + let egress_client = Arc::new( + ProxyClient::try_with_egress_proxy("socks5h://127.0.0.1:1080") + .expect("egress proxy client builds"), + ); + let state = AppState::new( + make_app_config(), + make_provider_registry(), + ProviderAdapterRegistry::builtin(), + ProxyClient::new(), + make_build_info(), + ) + .with_egress_clients(HashMap::from([( + "opencode-go".to_owned(), + Arc::clone(&egress_client), + )])); + + assert!( + std::ptr::eq( + state.upstream_client_for("opencode-go"), + egress_client.as_ref() + ), + "the provider with egress must use its wired client" + ); + assert!( + std::ptr::eq( + state.upstream_client_for("opencode-zen"), + state.proxy_client.as_ref() + ), + "providers without egress must keep the shared client" + ); + } + // -- Helper methods ------------------------------------------------------- #[test] @@ -323,6 +490,7 @@ mod tests { api_key: secrecy::SecretString::from("sk-secret-key-99999"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: HashMap::new(), routes: ProviderRoutesConfig::default(), model_aliases: HashMap::new(), @@ -376,6 +544,7 @@ mod tests { "request_id_gen", "model_catalogs", "event_bus", + "client_auth", ] { assert!( debug_output.contains(field), @@ -426,6 +595,7 @@ mod tests { api_key: secrecy::SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: { let mut m = HashMap::new(); m.insert( @@ -470,6 +640,7 @@ mod tests { api_key: secrecy::SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: { let mut m = HashMap::new(); m.insert( @@ -590,6 +761,7 @@ mod tests { api_key: secrecy::SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: { let mut m = HashMap::new(); m.insert( @@ -615,6 +787,7 @@ mod tests { api_key: secrecy::SecretString::from("key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: { let mut m = HashMap::new(); m.insert( @@ -709,4 +882,73 @@ mod tests { ); assert_eq!(state.server_name(), long_name); } + + // -- Client auth state ------------------------------------------------------ + + #[test] + fn invalid_auth_config_fails_closed() { + // An unparseable verifier cannot be allowed to downgrade the server to + // no-auth: the state must come up enforced with zero verifiers, so + // every proxied request is rejected. + let mut cfg = make_app_config(); + cfg.auth.mode = llm_proxy_core::AuthMode::ClientKeys; + cfg.auth.clients.insert( + "laptop".to_owned(), + secrecy::SecretString::from("not-a-phc-string"), + ); + let state = AppState::new( + cfg, + make_provider_registry(), + ProviderAdapterRegistry::builtin(), + ProxyClient::new(), + make_build_info(), + ); + assert!( + state.client_auth.is_enforced(), + "invalid auth config must fail closed, not open" + ); + assert_eq!(state.client_auth.client_count(), 0); + } + + #[test] + fn app_state_debug_does_not_leak_client_verifiers() { + // The client verifier hash must never surface through Debug, matching + // the api_key discipline. Use a real (cheap-parameter) argon2id hash + // so the verifier parses. + use argon2::{Algorithm, Argon2, CustomizedPasswordHasher, Params, Version}; + let params = Params::new(8, 1, 1, None).expect("valid test params"); + let hasher = Argon2::new(Algorithm::Argon2id, Version::V0x13, params.clone()); + let phc = hasher + .hash_password_customized( + b"leakable-secret-tail", + b"0123456789abcdef", + Some(&argon2::ARGON2ID_IDENT), + None, + params, + ) + .expect("hash token") + .to_string(); + + let mut cfg = make_app_config(); + cfg.auth.mode = llm_proxy_core::AuthMode::ClientKeys; + cfg.auth + .clients + .insert("laptop".to_owned(), secrecy::SecretString::from(phc)); + let client_auth = + crate::middleware::ClientAuth::from_config(&cfg.auth).expect("valid verifier"); + let state = AppState::new( + cfg, + make_provider_registry(), + ProviderAdapterRegistry::builtin(), + ProxyClient::new(), + make_build_info(), + ) + .with_client_auth(Arc::new(client_auth)); + let debug = format!("{state:?}"); + assert!( + !debug.contains("leakable-secret-tail"), + "client secret/verifier must not appear in Debug, got: {debug}" + ); + assert!(debug.contains("laptop"), "key id may appear: {debug}"); + } } diff --git a/crates/llm-proxy-server/tests/admin_usage.rs b/crates/llm-proxy-server/tests/admin_usage.rs new file mode 100644 index 0000000..c8741a5 --- /dev/null +++ b/crates/llm-proxy-server/tests/admin_usage.rs @@ -0,0 +1,212 @@ +//! Router-level tests for the admin `/usage` surface on both mounts: the +//! dedicated admin listener router and the opt-in main-listener mount. + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use axum::body::Body; + use axum::http::Request; + use axum::http::StatusCode; + use llm_proxy_core::{AppConfig, ServerConfig, StorageConfig}; + use llm_proxy_server::{AppState, BuildInfo, admin_router, build_router}; + use llm_proxy_storage::{EventBus, SqliteEventBus}; + use tower::ServiceExt; + + const ADMIN_TOKEN: &str = "integration-admin-token"; + + /// Seed a ledger with one completed laptop/prov-a request and return its + /// path. + async fn seeded_ledger_path(dir: &tempfile::TempDir) -> std::path::PathBuf { + let path = dir.path().join("usage.db"); + let bus = SqliteEventBus::open( + &path, + llm_proxy_storage::RetentionPolicy { + max_age: None, + max_rows: None, + prune_every: Duration::from_secs(3600), + }, + ) + .expect("open ledger"); + bus.emit(&llm_proxy_storage::ProxyEvent::RequestReceived( + llm_proxy_storage::RequestReceived { + request_id: "srv-1".to_owned(), + timestamp: time::OffsetDateTime::now_utc(), + caller: "laptop".to_owned(), + provider: "prov-a".to_owned(), + route_kind: "messages".to_owned(), + client_protocol: "anthropic".to_owned(), + model: llm_proxy_protocol::core::ModelRef { + requested: "claude-3".to_owned(), + upstream: None, + }, + streaming: false, + body_hash: "digest".to_owned(), + }, + )); + let closer = std::sync::Arc::new(bus); + tokio::task::spawn_blocking(move || closer.close()) + .await + .expect("close joined"); + path + } + + fn make_config( + db_path: std::path::PathBuf, + admin: Option, + on_main_listener: bool, + ) -> AppConfig { + AppConfig { + server: ServerConfig { + bind: "127.0.0.1:3456".parse().unwrap(), + request_timeout: Duration::from_secs(300), + shutdown_timeout: Duration::from_secs(30), + log_level: "info".to_owned(), + hot_reload: false, + allowed_origins: None, + server_name: "admin-usage-test".to_owned(), + rate_limit_rpm: 100, + trust_forwarded_headers: false, + dedup_window: Duration::ZERO, + log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: on_main_listener, + }, + auth: llm_proxy_core::AuthConfig { + mode: llm_proxy_core::AuthMode::None, + clients: std::collections::HashMap::new(), + admin, + }, + storage: Some(StorageConfig { + path: db_path, + max_age: None, + max_rows: None, + prune_interval: Duration::from_secs(3600), + }), + } + } + + fn make_state(config: AppConfig) -> AppState { + AppState::new( + config, + llm_proxy_core::ProviderRegistry::from_providers(vec![]).expect("empty registry"), + llm_proxy_provider::ProviderAdapterRegistry::builtin(), + llm_proxy_provider::ProxyClient::new(), + BuildInfo { + name: "test", + version: "0.0.0", + target: "test", + git_sha: "test", + }, + ) + } + + async fn get(app: axum::Router, uri: &str, token: Option<&str>) -> (StatusCode, String) { + let mut builder = Request::builder().uri(uri); + if let Some(token) = token { + builder = builder.header("authorization", format!("Bearer {token}")); + } + let response = app + .oneshot(builder.body(Body::empty()).expect("request builds")) + .await + .expect("oneshot"); + let status = response.status(); + let bytes = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .expect("body"); + (status, String::from_utf8_lossy(&bytes).into_owned()) + } + + #[tokio::test] + async fn main_listener_has_no_usage_surface_by_default() { + let dir = tempfile::tempdir().unwrap(); + let path = seeded_ledger_path(&dir).await; + let state = make_state(make_config( + path, + Some(llm_proxy_core::AdminAuthConfig { + token: secrecy::SecretString::from(ADMIN_TOKEN), + }), + false, + )); + let (status, _) = get(build_router(state), "/usage/summary", Some(ADMIN_TOKEN)).await; + assert_eq!( + status, + StatusCode::NOT_FOUND, + "without the opt-in, the main listener must not serve /usage even to the admin token" + ); + } + + #[tokio::test] + async fn main_listener_optin_serves_usage_behind_the_admin_token() { + let dir = tempfile::tempdir().unwrap(); + let path = seeded_ledger_path(&dir).await; + let state = make_state(make_config( + path, + Some(llm_proxy_core::AdminAuthConfig { + token: secrecy::SecretString::from(ADMIN_TOKEN), + }), + true, + )); + + let (status, body) = get(build_router(state.clone()), "/usage/summary", None).await; + assert_eq!(status, StatusCode::UNAUTHORIZED); + assert!( + !body.contains(ADMIN_TOKEN), + "401 must not echo the configured token" + ); + + let (status, body) = get( + build_router(state.clone()), + "/usage/summary", + Some("wrong-token"), + ) + .await; + assert_eq!(status, StatusCode::UNAUTHORIZED); + assert!( + !body.contains("wrong-token") && !body.contains(ADMIN_TOKEN), + "401 must not echo any token: {body}" + ); + + let (status, body) = get(build_router(state), "/usage/summary", Some(ADMIN_TOKEN)).await; + assert_eq!(status, StatusCode::OK); + assert!( + body.contains("prov-a"), + "the seeded provider group is served: {body}" + ); + } + + #[tokio::test] + async fn dedicated_admin_router_serves_usage_behind_the_admin_token() { + let dir = tempfile::tempdir().unwrap(); + let path = seeded_ledger_path(&dir).await; + let state = make_state(make_config( + path, + Some(llm_proxy_core::AdminAuthConfig { + token: secrecy::SecretString::from(ADMIN_TOKEN), + }), + false, + )); + + let admin = admin_router(&state).expect("admin router with [auth.admin]"); + let (status, _) = get(admin.clone(), "/usage/keys", None).await; + assert_eq!(status, StatusCode::UNAUTHORIZED); + + let (status, body) = get(admin, "/usage/keys", Some(ADMIN_TOKEN)).await; + assert_eq!(status, StatusCode::OK); + assert!( + body.contains("laptop"), + "the seeded caller is listed: {body}" + ); + } + + #[tokio::test] + async fn admin_router_is_none_without_admin_config() { + let dir = tempfile::tempdir().unwrap(); + let path = seeded_ledger_path(&dir).await; + let state = make_state(make_config(path, None, false)); + assert!( + admin_router(&state).is_none(), + "no [auth.admin] means no admin surface anywhere" + ); + } +} diff --git a/crates/llm-proxy-server/tests/chat_completions.rs b/crates/llm-proxy-server/tests/chat_completions.rs index 4eb7a82..b009dfa 100644 --- a/crates/llm-proxy-server/tests/chat_completions.rs +++ b/crates/llm-proxy-server/tests/chat_completions.rs @@ -187,6 +187,7 @@ mod tests { api_key: secrecy::SecretString::from("test-key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: { let mut m = HashMap::new(); m.insert( @@ -221,7 +222,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }; AppState::new( @@ -415,7 +420,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }; let registry = ProviderRegistry::from_providers(vec![]).expect("empty registry"); AppState::new( @@ -450,6 +459,7 @@ mod tests { api_key: secrecy::SecretString::from("test-key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: { let mut m = HashMap::new(); m.insert( @@ -485,7 +495,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }; AppState::new( app_config, @@ -954,6 +968,12 @@ mod tests { /// streaming tool-call start/delta/stop maps to choices[].delta.tool_calls /// Verifies the full start/delta/stop sequence: start has id + function.name, /// delta has function.arguments with partial JSON, and correct index values. + #[expect( + clippy::too_many_lines, + reason = "one end-to-end streaming scenario: the mock feed, the SSE \ + consumption, and the start/delta/stop assertions read as a \ + single narrative" + )] #[tokio::test] async fn streaming_tool_call_maps_to_delta_tool_calls() { // This test uses an Anthropic provider that returns tool call events. @@ -2008,6 +2028,12 @@ mod tests { } /// Same model name routed to different providers via URL path. + #[expect( + clippy::too_many_lines, + reason = "one end-to-end routing scenario: two provider registries, two \ + dispatches, and their divergent assertions read as a single \ + narrative" + )] #[tokio::test] async fn same_model_routes_to_different_providers() { // Create two providers both supporting chat_completions. @@ -2025,6 +2051,7 @@ mod tests { api_key: secrecy::SecretString::from("key-a"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: adapters_a, routes: llm_proxy_core::ProviderRoutesConfig { chat_completions: Some("chat".to_owned()), @@ -2050,6 +2077,7 @@ mod tests { api_key: secrecy::SecretString::from("key-b"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: adapters_b, routes: llm_proxy_core::ProviderRoutesConfig { chat_completions: Some("chat".to_owned()), @@ -2076,7 +2104,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }, registry, ProviderAdapterRegistry::builtin(), diff --git a/crates/llm-proxy-server/tests/core_pipeline.rs b/crates/llm-proxy-server/tests/core_pipeline.rs index dc5dfde..c6588a1 100644 --- a/crates/llm-proxy-server/tests/core_pipeline.rs +++ b/crates/llm-proxy-server/tests/core_pipeline.rs @@ -204,6 +204,7 @@ mod tests { api_key: secrecy::SecretString::from("test-key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: { let mut m = HashMap::new(); m.insert( @@ -238,7 +239,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }; AppState::new( @@ -1285,6 +1290,185 @@ mod tests { ); } + // =========================================================================== + // Client-key authentication + // =========================================================================== + + /// Hash a token with deliberately cheap argon2id parameters so router + /// tests run fast; verification reads the cost from the PHC string. + fn client_auth_for(key_id: &str, token: &str) -> llm_proxy_server::ClientAuth { + use argon2::{Algorithm, Argon2, CustomizedPasswordHasher, Params, Version}; + let params = Params::new(8, 1, 1, None).expect("valid test params"); + let hasher = Argon2::new(Algorithm::Argon2id, Version::V0x13, params.clone()); + let phc = hasher + .hash_password_customized( + token.as_bytes(), + b"0123456789abcdef", + Some(&argon2::ARGON2ID_IDENT), + None, + params, + ) + .expect("hash token") + .to_string(); + let mut clients = HashMap::new(); + clients.insert(key_id.to_owned(), secrecy::SecretString::from(phc)); + llm_proxy_server::ClientAuth::from_config(&llm_proxy_core::AuthConfig { + mode: llm_proxy_core::AuthMode::ClientKeys, + clients, + admin: None, + }) + .expect("valid client auth") + } + + /// Under enforced client keys, a request without a bearer token is + /// rejected with a 401 that carries no secret material and still has the + /// x-request-id header. + #[tokio::test] + async fn enforced_client_keys_reject_missing_bearer_with_sanitized_401() { + let mock_url = spawn_mock_anthropic_non_stream().await; + let state = state_with_anthropic_provider(&mock_url).with_client_auth(std::sync::Arc::new( + client_auth_for("laptop", "sekrit-token"), + )); + let app = build_router(state); + + let body = make_messages_body("claude-sonnet-4-6", false); + let resp = app.oneshot(messages_request(&body)).await.unwrap(); + assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); + assert!( + resp.headers().get("x-request-id").is_some(), + "401 must carry the request id" + ); + let body = axum::body::to_bytes(resp.into_body(), 64 * 1024) + .await + .unwrap(); + let text = String::from_utf8_lossy(&body); + assert!( + !text.contains("sekrit-token"), + "401 body must not echo any secret material, got: {text}" + ); + } + + /// A wrong bearer token is rejected exactly like a missing one: same 401 + /// body, and the attempted token never appears in the response. + #[tokio::test] + async fn enforced_client_keys_reject_wrong_bearer_without_echoing_it() { + let mock_url = spawn_mock_anthropic_non_stream().await; + let state = state_with_anthropic_provider(&mock_url).with_client_auth(std::sync::Arc::new( + client_auth_for("laptop", "sekrit-token"), + )); + let app = build_router(state); + + let request = Request::builder() + .method("POST") + .uri("/providers/mock-provider/v1/messages") + .header("content-type", "application/json") + .header("authorization", "Bearer attacker-guess") + .body(Body::from(make_messages_body("claude-sonnet-4-6", false))) + .unwrap(); + let resp = app.oneshot(request).await.unwrap(); + assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); + let body = axum::body::to_bytes(resp.into_body(), 64 * 1024) + .await + .unwrap(); + let text = String::from_utf8_lossy(&body); + assert!( + !text.contains("attacker-guess"), + "401 body must not echo the attempted token, got: {text}" + ); + } + + /// A valid bearer token passes the auth gate, and the RequestReceived + /// event is attributed to the key id (not "anonymous", not the secret). + #[tokio::test] + async fn valid_bearer_passes_and_attributes_caller_on_event() { + use std::sync::Arc; + + use llm_proxy_storage::{ProxyEvent, RecordingBus}; + + let mock_url = spawn_mock_anthropic_non_stream().await; + let bus = Arc::new(RecordingBus::new()); + let state = state_with_anthropic_provider(&mock_url) + .with_event_bus(Arc::::clone(&bus)) + .with_client_auth(std::sync::Arc::new(client_auth_for( + "laptop", + "sekrit-token", + ))); + let app = build_router(state); + + let request = Request::builder() + .method("POST") + .uri("/providers/mock-provider/v1/messages") + .header("content-type", "application/json") + .header("authorization", "bearer sekrit-token") + .body(Body::from(make_messages_body("claude-sonnet-4-6", false))) + .unwrap(); + let resp = app.oneshot(request).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let events = bus.snapshot(); + let received = events + .iter() + .find_map(|e| match e { + ProxyEvent::RequestReceived(r) => Some(r), + _ => None, + }) + .expect("RequestReceived emitted"); + assert_eq!(received.caller, "laptop"); + let event_json = + serde_json::to_string(&ProxyEvent::RequestReceived(received.clone())).unwrap(); + assert!( + !event_json.contains("sekrit-token"), + "client secret must never appear on the event, got: {event_json}" + ); + } + + /// Under the default no-auth mode, requests pass and events attribute the + /// caller as "anonymous". + #[tokio::test] + async fn no_auth_mode_attributes_anonymous_caller() { + use std::sync::Arc; + + use llm_proxy_storage::{ProxyEvent, RecordingBus}; + + let mock_url = spawn_mock_anthropic_non_stream().await; + let bus = Arc::new(RecordingBus::new()); + let state = state_with_anthropic_provider(&mock_url) + .with_event_bus(Arc::::clone(&bus)); + let app = build_router(state); + + let body = make_messages_body("claude-sonnet-4-6", false); + let resp = app.oneshot(messages_request(&body)).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let events = bus.snapshot(); + let received = events + .iter() + .find_map(|e| match e { + ProxyEvent::RequestReceived(r) => Some(r), + _ => None, + }) + .expect("RequestReceived emitted"); + assert_eq!(received.caller, "anonymous"); + } + + /// The auth gate is not applied to lightweight routes: /health stays + /// open even under enforced client keys (it exposes no model traffic). + #[tokio::test] + async fn enforced_client_keys_leave_health_open() { + let mock_url = spawn_mock_anthropic_non_stream().await; + let state = state_with_anthropic_provider(&mock_url).with_client_auth(std::sync::Arc::new( + client_auth_for("laptop", "sekrit-token"), + )); + let app = build_router(state); + + let request = Request::builder() + .uri("/health") + .body(Body::empty()) + .unwrap(); + let resp = app.oneshot(request).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + } + /// A request for an unknown provider emits `ResponseFailed` with the mapped /// HTTP status (404) and the provider name from the URL path. #[tokio::test] @@ -2007,6 +2191,7 @@ mod tests { api_key: secrecy::SecretString::from("test-key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: { let mut m = HashMap::new(); m.insert( @@ -2043,7 +2228,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }, registry, ProviderAdapterRegistry::builtin(), @@ -2116,6 +2305,7 @@ mod tests { api_key: secrecy::SecretString::from("test-key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: { let mut m = HashMap::new(); m.insert( @@ -2152,7 +2342,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_secs(60), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }, registry, ProviderAdapterRegistry::builtin(), @@ -2449,6 +2643,7 @@ mod tests { api_key: secrecy::SecretString::from("test-key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: { let mut m = HashMap::new(); m.insert( @@ -2485,7 +2680,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }, registry, ProviderAdapterRegistry::builtin(), @@ -2705,6 +2904,7 @@ mod tests { api_key: secrecy::SecretString::from("test-key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: { let mut m = HashMap::new(); m.insert( @@ -2741,7 +2941,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }; AppState::new( app_config, diff --git a/crates/llm-proxy-server/tests/integration.rs b/crates/llm-proxy-server/tests/integration.rs index 9f1ea4c..ed9cc5a 100644 --- a/crates/llm-proxy-server/tests/integration.rs +++ b/crates/llm-proxy-server/tests/integration.rs @@ -36,7 +36,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }; let registry = ProviderRegistry::from_providers(vec![]).expect("empty registry"); AppState::new( @@ -63,6 +67,7 @@ mod tests { api_key: secrecy::SecretString::from("test-key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: { let mut m = HashMap::new(); m.insert( @@ -99,7 +104,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }, registry, ProviderAdapterRegistry::builtin(), @@ -904,7 +913,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }; let registry = ProviderRegistry::from_providers(vec![]).expect("empty registry"); AppState::new( diff --git a/crates/llm-proxy-server/tests/openrouter_e2e.rs b/crates/llm-proxy-server/tests/openrouter_e2e.rs index c3b9e09..bdeca2d 100644 --- a/crates/llm-proxy-server/tests/openrouter_e2e.rs +++ b/crates/llm-proxy-server/tests/openrouter_e2e.rs @@ -322,6 +322,7 @@ mod tests { api_key: secrecy::SecretString::from("test-key"), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: { let mut m = HashMap::new(); m.insert( @@ -359,7 +360,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }; AppState::new( @@ -1068,6 +1073,7 @@ mod tests { api_key: secrecy::SecretString::from(api_key), auth_style: AuthStyle::Bearer, passthrough_auth: false, + egress: None, adapters: { let mut m = HashMap::new(); m.insert( @@ -1103,7 +1109,11 @@ mod tests { trust_forwarded_headers: false, dedup_window: Duration::from_millis(500), log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: false, }, + auth: Default::default(), + storage: None, }; AppState::new( app_config, diff --git a/crates/llm-proxy-server/tests/secret_abuse.rs b/crates/llm-proxy-server/tests/secret_abuse.rs new file mode 100644 index 0000000..5f1a3b3 --- /dev/null +++ b/crates/llm-proxy-server/tests/secret_abuse.rs @@ -0,0 +1,921 @@ +//! Secret-abuse family (PRX-06): secret material must never be observable. +//! +//! Every test drives the real surfaces — [`build_router`], [`admin_router`], +//! the client-auth layer, the SQLite ledger, the `/usage` endpoints — and +//! asserts that client tokens, the admin token, provider keys, and request +//! bodies never appear in events, tracing output, error responses, the raw +//! ledger file bytes, or anything another caller can observe. Negative +//! assertions are always paired with a positive control (content that *is* +//! expected to appear), so a passing scan proves the scan itself worked +//! rather than reading the wrong file. + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::path::{Path, PathBuf}; + use std::sync::{Arc, Mutex, MutexGuard}; + use std::time::Duration; + + use axum::Router; + use axum::body::Body; + use axum::http::{Request, StatusCode, header}; + use axum::routing::post; + use llm_proxy_core::{ + AdminAuthConfig, AppConfig, AuthMode, AuthStyle, ProviderAdapterConfig, ProviderConfig, + ProviderRegistry, ProviderRoutesConfig, ServerConfig, StorageConfig, + }; + use llm_proxy_provider::{ProviderAdapterRegistry, ProxyClient}; + use llm_proxy_server::{AppState, BuildInfo, ClientAuth, admin_router, build_router}; + use llm_proxy_storage::{EventBus, ProxyEvent, RecordingBus, SqliteEventBus}; + use tower::ServiceExt; + + /// Distinctive secret material, each unique to this family so a hit can + /// never be a coincidence with unrelated content. + const LAPTOP_TOKEN: &str = "prx6-laptop-client-token-4f1a9c73"; + const PHONE_TOKEN: &str = "prx6-phone-client-token-b27e03d1"; + const ADMIN_TOKEN: &str = "prx6-admin-token-9d5c48ea"; + const WRONG_TOKEN: &str = "prx6-attacker-guess-00000011"; + const STRAY_TOKEN: &str = "prx6-stray-bearer-under-noauth-77c1"; + const PROVIDER_KEY: &str = "prx6-provider-upstream-key-6b8f2e45"; + const BODY_MARKER: &str = "PRX6-BODY-MARKER-never-persist-me"; + /// Every string that must never appear in events, logs, responses, or the + /// ledger file bytes. + const SECRETS: &[&str] = &[ + LAPTOP_TOKEN, + PHONE_TOKEN, + ADMIN_TOKEN, + WRONG_TOKEN, + STRAY_TOKEN, + PROVIDER_KEY, + BODY_MARKER, + ]; + const KEY_ID_LAPTOP: &str = "laptop"; + const KEY_ID_PHONE: &str = "phone"; + + /// Canned Anthropic Messages success response for the mock upstream. + const ANTHROPIC_SUCCESS: &str = concat!( + r#"{"id":"msg_prx6","type":"message","role":"assistant","content":"#, + r#"[{"type":"text","text":"prx6-mock-reply"}],"model":"claude-sonnet-4-6","#, + r#""stop_reason":"end_turn","stop_sequence":null,"#, + r#""usage":{"input_tokens":10,"output_tokens":5}}"# + ); + + // ----------------------------------------------------------------------- + // Fixtures + // ----------------------------------------------------------------------- + + /// Hash `token` with deliberately cheap argon2id parameters so tests run + /// fast; verification reads the cost from the PHC string, so these + /// verifiers behave exactly like production-cost ones. + fn cheap_phc(token: &str, salt: &[u8]) -> String { + use argon2::{Algorithm, Argon2, CustomizedPasswordHasher, Params, Version}; + let params = Params::new(8, 1, 1, None).expect("valid cheap test params"); + Argon2::new(Algorithm::Argon2id, Version::V0x13, params.clone()) + .hash_password_customized( + token.as_bytes(), + salt, + Some(&argon2::ARGON2ID_IDENT), + None, + params, + ) + .expect("hash token") + .to_string() + } + + /// Enforced [`ClientAuth`] for the given (key id, token) pairs, plus the + /// verifier strings so tests can assert they never leak either. + fn client_auth(pairs: &[(&str, &str)]) -> (ClientAuth, Vec) { + let mut clients = HashMap::new(); + let mut phcs = Vec::new(); + for (index, (key_id, token)) in pairs.iter().enumerate() { + let salt = format!("{index:016x}"); + let phc = cheap_phc(token, salt.as_bytes()); + clients.insert( + (*key_id).to_owned(), + secrecy::SecretString::from(phc.clone()), + ); + phcs.push(phc); + } + let auth = ClientAuth::from_config(&llm_proxy_core::AuthConfig { + mode: AuthMode::ClientKeys, + clients, + admin: None, + }) + .expect("valid client auth"); + (auth, phcs) + } + + fn admin_config() -> AdminAuthConfig { + AdminAuthConfig { + token: secrecy::SecretString::from(ADMIN_TOKEN), + } + } + + /// AppState with one Anthropic-protocol provider routed at + /// `mock_endpoint`. `storage_path`/`admin` wire the `/usage` surface; + /// `on_main` also mounts it on the main listener (the opt-in posture). + fn base_state( + mock_endpoint: &str, + storage_path: Option, + admin: Option, + on_main: bool, + ) -> AppState { + let provider = ProviderConfig { + name: "mock-provider".to_owned(), + api_key: secrecy::SecretString::from(PROVIDER_KEY), + auth_style: AuthStyle::Bearer, + passthrough_auth: false, + egress: None, + adapters: { + let mut adapters = HashMap::new(); + adapters.insert( + "main".to_owned(), + ProviderAdapterConfig { + protocol: "anthropic_messages".to_owned(), + endpoint: mock_endpoint.to_owned(), + headers: Arc::new(HashMap::new()), + }, + ); + adapters + }, + routes: ProviderRoutesConfig { + messages: Some("main".to_owned()), + chat_completions: None, + }, + model_aliases: HashMap::new(), + discovery: None, + catalog: None, + pricing: HashMap::new(), + }; + + AppState::new( + AppConfig { + server: ServerConfig { + bind: "127.0.0.1:3456".parse().expect("bind parses"), + request_timeout: Duration::from_secs(300), + shutdown_timeout: Duration::from_secs(30), + log_level: "info".to_owned(), + hot_reload: false, + allowed_origins: None, + server_name: "secret-abuse-test".to_owned(), + rate_limit_rpm: 100, + trust_forwarded_headers: false, + dedup_window: Duration::ZERO, + log_format: Default::default(), + admin_bind: None, + admin_on_main_listener: on_main, + }, + auth: llm_proxy_core::AuthConfig { + mode: AuthMode::None, + clients: HashMap::new(), + admin, + }, + storage: storage_path.map(|path| StorageConfig { + path, + max_age: None, + max_rows: None, + prune_interval: Duration::from_secs(3600), + }), + }, + ProviderRegistry::from_providers(vec![provider]).expect("registry builds"), + ProviderAdapterRegistry::builtin(), + ProxyClient::new(), + BuildInfo { + name: "test", + version: "0.0.0", + target: "test", + git_sha: "test", + }, + ) + } + + /// A messages request body whose user content carries `marker`. + fn marked_body(marker: &str) -> String { + serde_json::json!({ + "model": "claude-sonnet-4-6", + "messages": [{ "role": "user", "content": format!("hello {marker}") }], + "max_tokens": 64, + "stream": false + }) + .to_string() + } + + fn post_messages(token: Option<&str>, body: &str) -> Request { + request_with_token("POST", "/providers/mock-provider/v1/messages", token, body) + } + + fn get_usage(uri: &str, token: Option<&str>) -> Request { + request_with_token("GET", uri, token, "") + } + + fn request_with_token( + method: &str, + uri: &str, + token: Option<&str>, + body: &str, + ) -> Request { + let mut builder = Request::builder() + .method(method) + .uri(uri) + .header("content-type", "application/json"); + if let Some(token) = token { + builder = builder.header("authorization", format!("Bearer {token}")); + } + builder + .body(Body::from(body.to_owned())) + .expect("request builds") + } + + /// Drive one request through the router, returning status, body text, and + /// the `x-request-id` response header. + async fn send(app: Router, request: Request) -> (StatusCode, String, Option) { + let response = app.oneshot(request).await.expect("oneshot completes"); + let status = response.status(); + let request_id = response + .headers() + .get("x-request-id") + .and_then(|value| value.to_str().ok()) + .map(str::to_owned); + let bytes = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .expect("body reads"); + ( + status, + String::from_utf8_lossy(&bytes).into_owned(), + request_id, + ) + } + + fn assert_free_of_secrets(context: &str, text: &str) { + for secret in SECRETS { + assert!( + !text.contains(secret), + "{context} leaked secret material ({secret}): {text}" + ); + } + } + + fn hex_sha256(bytes: &[u8]) -> String { + use sha2::{Digest, Sha256}; + let mut hasher = Sha256::new(); + hasher.update(bytes); + format!("{:x}", hasher.finalize()) + } + + /// Serialized forms of every recorded event that belongs to `request_id`. + fn events_for_request(bus: &RecordingBus, request_id: &str) -> Vec { + bus.snapshot() + .iter() + .filter_map(|event| { + let json = serde_json::to_string(event).expect("event serializes"); + json.contains(request_id).then_some(json) + }) + .collect() + } + + fn open_ledger(path: &Path) -> SqliteEventBus { + SqliteEventBus::open( + path, + llm_proxy_storage::RetentionPolicy { + max_age: None, + max_rows: None, + prune_every: Duration::from_secs(3600), + }, + ) + .expect("ledger opens") + } + + async fn close_ledger(bus: Arc) { + tokio::task::spawn_blocking(move || bus.close()) + .await + .expect("close joins"); + } + + /// Concatenated raw bytes of the ledger database and any WAL/SHM sidecar + /// files still on disk — everything an attacker with filesystem access + /// (or a copied backup) could read. + fn ledger_bytes(db: &Path) -> Vec { + let mut all = Vec::new(); + for suffix in ["", "-wal", "-shm"] { + let sidecar = PathBuf::from(format!("{}{suffix}", db.display())); + if let Ok(bytes) = std::fs::read(&sidecar) { + all.extend_from_slice(&bytes); + } + } + all + } + + async fn wait_ready(addr: std::net::SocketAddr) { + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + loop { + if tokio::net::TcpStream::connect(addr).await.is_ok() { + return; + } + assert!( + tokio::time::Instant::now() < deadline, + "mock upstream at {addr} never became ready" + ); + tokio::time::sleep(Duration::from_millis(5)).await; + } + } + + /// Mock upstream answering every POST with a canned Anthropic response. + async fn spawn_anthropic_mock() -> String { + let app = Router::new().route( + "/{*path}", + post(|| async { + ( + StatusCode::OK, + [(header::CONTENT_TYPE, "application/json")], + ANTHROPIC_SUCCESS, + ) + }), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("mock binds"); + let addr = listener.local_addr().expect("mock address"); + tokio::spawn(async move { + if let Err(error) = axum::serve(listener, app).await { + eprintln!("mock upstream serve error: {error}"); + } + }); + wait_ready(addr).await; + format!("http://{addr}/providers/mock-provider/v1/messages") + } + + /// Pathological mock upstream that reflects the request body verbatim: + /// the worst case for body confidentiality, since everything the client + /// sent comes back through the proxy. + async fn spawn_echo_mock() -> String { + let app = Router::new().route( + "/{*path}", + post(|body: axum::body::Bytes| async move { + ( + StatusCode::OK, + [(header::CONTENT_TYPE, "application/json")], + body, + ) + }), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("mock binds"); + let addr = listener.local_addr().expect("mock address"); + tokio::spawn(async move { + if let Err(error) = axum::serve(listener, app).await { + eprintln!("mock upstream serve error: {error}"); + } + }); + wait_ready(addr).await; + format!("http://{addr}/providers/mock-provider/v1/messages") + } + + // ----------------------------------------------------------------------- + // Tracing capture (process-global, installed once per test binary) + // ----------------------------------------------------------------------- + + /// Captured tracing output for the whole test binary. + /// + /// A process-global subscriber instead of a thread-local default: with + /// parallel tests, callsites first emitted on a thread that has no + /// subscriber can keep a stale "never" interest cache, silently dropping + /// later events from a thread-local capture — observed as flaky partial + /// captures during development. The shared buffer therefore contains + /// every test's traffic; tests that need per-caller isolation assert on + /// lines correlated with their own request ids (see the caller-isolation + /// test), and the secret scans simply scan more traffic. + struct TraceCapture { + buffer: Arc>>, + } + + static CAPTURE_BUFFER: std::sync::OnceLock>>> = std::sync::OnceLock::new(); + static CAPTURE_INSTALLED: std::sync::OnceLock<()> = std::sync::OnceLock::new(); + + impl TraceCapture { + fn install() -> Self { + use tracing_subscriber::layer::SubscriberExt; + let buffer = CAPTURE_BUFFER.get_or_init(|| Arc::new(Mutex::new(Vec::::new()))); + CAPTURE_INSTALLED.get_or_init(|| { + let layer = tracing_subscriber::fmt::layer() + .with_writer(BufferWriter(Arc::clone(buffer))) + .with_ansi(false); + tracing::subscriber::set_global_default(tracing_subscriber::registry().with(layer)) + .expect("the global capture subscriber installs exactly once"); + }); + Self { + buffer: Arc::clone(buffer), + } + } + + fn text(&self) -> String { + let bytes = self.buffer.lock().expect("buffer lock").clone(); + String::from_utf8_lossy(&bytes).into_owned() + } + } + + /// `MakeWriter` backed by the shared capture buffer. + #[derive(Clone)] + struct BufferWriter(Arc>>); + + impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for BufferWriter { + type Writer = BufferWriterGuard<'a>; + + fn make_writer(&'a self) -> Self::Writer { + BufferWriterGuard(self.0.lock().expect("buffer lock")) + } + } + + struct BufferWriterGuard<'a>(MutexGuard<'a, Vec>); + + impl std::io::Write for BufferWriterGuard<'_> { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + self.0.extend_from_slice(bytes); + Ok(bytes.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + + // ----------------------------------------------------------------------- + // 1. Tokens never appear in events, tracing output, or error responses + // ----------------------------------------------------------------------- + + #[tokio::test] + async fn client_and_admin_tokens_never_reach_events_tracing_or_responses() { + let mock_url = spawn_anthropic_mock().await; + let (auth, phcs) = + client_auth(&[(KEY_ID_LAPTOP, LAPTOP_TOKEN), (KEY_ID_PHONE, PHONE_TOKEN)]); + let bus = Arc::new(RecordingBus::new()); + let state = base_state(&mock_url, None, None, false) + .with_client_auth(Arc::new(auth)) + .with_event_bus(Arc::clone(&bus) as Arc); + let app = build_router(state.clone()); + + // The admin token is presented to the PROXY surface and a client key + // plus a garbage guess to the auth gate: both auth realms and both + // failure modes flow through the same observable surfaces. + let traffic: [(Option<&str>, StatusCode); 5] = [ + (Some(LAPTOP_TOKEN), StatusCode::OK), + (Some(PHONE_TOKEN), StatusCode::OK), + (Some(WRONG_TOKEN), StatusCode::UNAUTHORIZED), + (None, StatusCode::UNAUTHORIZED), + (Some(ADMIN_TOKEN), StatusCode::UNAUTHORIZED), + ]; + + let capture = TraceCapture::install(); + for (token, expected) in traffic { + let (status, body, _) = + send(app.clone(), post_messages(token, &marked_body(BODY_MARKER))).await; + assert_eq!( + status, expected, + "unexpected status for token {token:?}: {body}" + ); + assert_free_of_secrets(&format!("response body for {token:?}"), &body); + for phc in &phcs { + assert!( + !body.contains(phc.as_str()), + "response body leaked the argon2 verifier: {body}" + ); + } + } + + // Events: every recorded event serialized whole must be secret-free. + let events = bus.snapshot(); + assert!( + events.len() >= 6, + "traffic must produce events (got {})", + events.len() + ); + for event in &events { + let json = serde_json::to_string(event).expect("event serializes"); + assert_free_of_secrets("recorded event", &json); + for phc in &phcs { + assert!( + !json.contains(phc.as_str()), + "event leaked the argon2 verifier: {json}" + ); + } + } + // Control: attribution itself survived — both key ids are on events. + let callers: Vec<&str> = events + .iter() + .filter_map(|event| match event { + ProxyEvent::RequestReceived(received) => Some(received.caller.as_str()), + _ => None, + }) + .collect(); + assert!(callers.contains(&KEY_ID_LAPTOP), "callers: {callers:?}"); + assert!(callers.contains(&KEY_ID_PHONE), "callers: {callers:?}"); + + // Tracing: the capture must be live (control — the buffer receives + // debug-level request traffic from this binary's authenticated + // calls) and secret-free across everything it collected. + let captured = capture.text(); + assert!( + captured.contains("client authenticated"), + "control: debug-level request traffic must land in the capture, got: {captured}" + ); + assert_free_of_secrets("tracing output", &captured); + for phc in &phcs { + assert!( + !captured.contains(phc.as_str()), + "tracing output leaked the argon2 verifier: {captured}" + ); + } + + // The state's own Debug render (panic messages, error logs) is a leak + // vector for the verifiers it holds. + let state_debug = format!("{state:?}"); + assert!( + !state_debug.contains(LAPTOP_TOKEN) && !state_debug.contains(PHONE_TOKEN), + "AppState Debug leaked a client token: {state_debug}" + ); + for phc in &phcs { + assert!( + !state_debug.contains(phc.as_str()), + "AppState Debug leaked a verifier: {state_debug}" + ); + } + } + + // ----------------------------------------------------------------------- + // 2. Tokens, verifiers, and bodies never appear in the raw ledger bytes + // ----------------------------------------------------------------------- + + #[tokio::test] + async fn secrets_absent_from_raw_ledger_file_bytes() { + let dir = tempfile::tempdir().expect("tempdir"); + let db = dir.path().join("abuse.db"); + let mock_url = spawn_anthropic_mock().await; + let (auth, phcs) = + client_auth(&[(KEY_ID_LAPTOP, LAPTOP_TOKEN), (KEY_ID_PHONE, PHONE_TOKEN)]); + let bus = Arc::new(open_ledger(&db)); + let state = base_state(&mock_url, Some(db.clone()), Some(admin_config()), true) + .with_client_auth(Arc::new(auth)) + .with_event_bus(Arc::clone(&bus) as Arc); + let app = build_router(state); + + // Client traffic with both secrets and body content in flight, plus + // admin traffic (authenticated and rejected) against the same state. + let body = marked_body(BODY_MARKER); + let digest = hex_sha256(body.as_bytes()); + let (status, response, _) = + send(app.clone(), post_messages(Some(LAPTOP_TOKEN), &body)).await; + assert_eq!(status, StatusCode::OK, "valid client request: {response}"); + let (status, _, _) = send(app.clone(), post_messages(Some(WRONG_TOKEN), &body)).await; + assert_eq!(status, StatusCode::UNAUTHORIZED); + let (status, _, _) = + send(app.clone(), get_usage("/usage/requests", Some(ADMIN_TOKEN))).await; + assert_eq!(status, StatusCode::OK); + let (status, _, _) = send( + app.clone(), + get_usage("/usage/requests", Some(LAPTOP_TOKEN)), + ) + .await; + assert_eq!( + status, + StatusCode::UNAUTHORIZED, + "client keys are not admin" + ); + + close_ledger(Arc::clone(&bus)).await; + let raw = ledger_bytes(&db); + assert!(!raw.is_empty(), "the ledger file must exist on disk"); + + // Positive controls: content that IS persisted (key id, provider, + // body digest) proves the scan reads real rows, so the negative + // assertions below are not vacuous. + let text = String::from_utf8_lossy(&raw); + for expected in [KEY_ID_LAPTOP, "mock-provider", digest.as_str()] { + assert!( + text.contains(expected), + "control: persisted value {expected} must be findable in the ledger bytes" + ); + } + assert_free_of_secrets("raw ledger file bytes", &text); + for phc in &phcs { + assert!( + !text.contains(phc.as_str()), + "ledger bytes leaked an argon2 verifier" + ); + } + } + + // ----------------------------------------------------------------------- + // 3. /usage responses carry the digest but never the body content + // ----------------------------------------------------------------------- + + #[tokio::test] + async fn usage_serves_the_body_digest_but_never_the_body_content() { + let dir = tempfile::tempdir().expect("tempdir"); + let db = dir.path().join("abuse.db"); + // The upstream REFLECTS the request body: the marker round-trips + // through the proxy's response handling, so the assertions below + // hold even in the worst case for body confidentiality. + let mock_url = spawn_echo_mock().await; + let (auth, _phcs) = client_auth(&[(KEY_ID_LAPTOP, LAPTOP_TOKEN)]); + let bus = Arc::new(open_ledger(&db)); + let state = base_state(&mock_url, Some(db.clone()), Some(admin_config()), true) + .with_client_auth(Arc::new(auth)) + .with_event_bus(Arc::clone(&bus) as Arc); + let main_app = build_router(state.clone()); + let dedicated = admin_router(&state).expect("admin router for the dedicated mount"); + + let body = marked_body(BODY_MARKER); + let digest = hex_sha256(body.as_bytes()); + let (status, response, _) = + send(main_app.clone(), post_messages(Some(LAPTOP_TOKEN), &body)).await; + assert_ne!( + status, + StatusCode::OK, + "a reflected request-shaped body must not decode as a valid response" + ); + assert_free_of_secrets("client-facing error response", &response); + + close_ledger(Arc::clone(&bus)).await; + for (mount, app) in [("dedicated", dedicated), ("main-opt-in", main_app)] { + for uri in [ + "/usage/requests?limit=1000", + "/usage/summary?group_by=caller", + "/usage/keys", + ] { + let (status, body, _) = send(app.clone(), get_usage(uri, Some(ADMIN_TOKEN))).await; + assert_eq!(status, StatusCode::OK, "{mount} {uri}: {body}"); + assert_free_of_secrets(&format!("{mount} {uri} response"), &body); + assert!( + !body.contains(BODY_MARKER), + "{mount} {uri} must never carry request-body content: {body}" + ); + } + + let (_, body, _) = send( + app.clone(), + get_usage("/usage/requests?limit=1000", Some(ADMIN_TOKEN)), + ) + .await; + let rows = serde_json::from_str::(&body) + .expect("rows JSON") + .get("rows") + .cloned() + .unwrap_or_default(); + let digested: Vec<&str> = rows + .as_array() + .map(|rows| { + rows.iter() + .filter_map(|row| row["body_hash"].as_str()) + .collect() + }) + .unwrap_or_default(); + assert_eq!( + digested, + vec![digest.as_str()], + "{mount}: the request row must carry the sha256 of the actual body, not the body" + ); + assert!( + body.contains(KEY_ID_LAPTOP), + "{mount}: caller attribution must be visible to the admin" + ); + } + } + + // ----------------------------------------------------------------------- + // 4. Admin routes refuse non-admin credentials on BOTH mounts + // ----------------------------------------------------------------------- + + #[tokio::test] + async fn admin_surface_refuses_non_admin_credentials_on_both_mounts() { + let dir = tempfile::tempdir().expect("tempdir"); + let db = dir.path().join("abuse.db"); + // Seed one row so the authenticated calls have something to serve. + let bus = Arc::new(open_ledger(&db)); + bus.emit(&ProxyEvent::RequestReceived( + llm_proxy_storage::RequestReceived { + request_id: "prx6-seed".to_owned(), + timestamp: time::OffsetDateTime::now_utc(), + caller: "prx6-seed-caller".to_owned(), + provider: "mock-provider".to_owned(), + route_kind: "messages".to_owned(), + client_protocol: "anthropic".to_owned(), + model: llm_proxy_protocol::core::ModelRef { + requested: "claude-3".to_owned(), + upstream: None, + }, + streaming: false, + body_hash: "seed-digest".to_owned(), + }, + )); + close_ledger(Arc::clone(&bus)).await; + + let (auth, _phcs) = client_auth(&[(KEY_ID_LAPTOP, LAPTOP_TOKEN)]); + let state = base_state("http://127.0.0.1:9", Some(db), Some(admin_config()), true) + .with_client_auth(Arc::new(auth)); + let mounts = [ + ("dedicated", admin_router(&state).expect("admin router")), + ("main-opt-in", build_router(state)), + ]; + + for (mount, app) in mounts { + let (missing_status, missing_body, _) = + send(app.clone(), get_usage("/usage/summary", None)).await; + assert_eq!( + missing_status, + StatusCode::UNAUTHORIZED, + "{mount}: missing credentials must be 401 (not 403/5xx)" + ); + + for token in [WRONG_TOKEN, LAPTOP_TOKEN] { + let (status, body, _) = + send(app.clone(), get_usage("/usage/summary", Some(token))).await; + assert_eq!( + status, + StatusCode::UNAUTHORIZED, + "{mount}: {token:?} must not authenticate to the admin surface" + ); + assert_free_of_secrets(&format!("{mount} rejection for {token:?}"), &body); + assert_eq!( + body, missing_body, + "{mount}: every rejection carries the identical uniform body" + ); + } + + let (status, _, _) = + send(app.clone(), get_usage("/usage/summary", Some(ADMIN_TOKEN))).await; + assert_eq!(status, StatusCode::OK, "{mount}: the admin token serves"); + } + + // Realm separation in the other direction: the admin token is not a + // client key and must not ride the proxy surface. + let (proxy_auth, _phcs) = client_auth(&[(KEY_ID_LAPTOP, LAPTOP_TOKEN)]); + let proxy_state = base_state("http://127.0.0.1:9", None, None, false) + .with_client_auth(Arc::new(proxy_auth)); + let (status, body, _) = send( + build_router(proxy_state), + post_messages(Some(ADMIN_TOKEN), &marked_body(BODY_MARKER)), + ) + .await; + assert_eq!( + status, + StatusCode::UNAUTHORIZED, + "the admin token must not authenticate proxy routes" + ); + assert_free_of_secrets("proxy-surface rejection of the admin token", &body); + } + + // ----------------------------------------------------------------------- + // 5. One caller's key id is never observable by another caller + // ----------------------------------------------------------------------- + + #[tokio::test] + async fn one_callers_key_id_is_never_observable_by_another_caller() { + // Key ids unique to this test: the capture buffer is process-global + // (see TraceCapture), so ids shared with the other tests would make + // parallel traffic indistinguishable from a leak. + const ALPHA_ID: &str = "prx6-iso-alpha"; + const BETA_ID: &str = "prx6-iso-beta"; + const ALPHA_TOKEN: &str = "prx6-iso-alpha-token-3d5e"; + const BETA_TOKEN: &str = "prx6-iso-beta-token-8f1b"; + + let mock_url = spawn_anthropic_mock().await; + let (auth, _phcs) = client_auth(&[(ALPHA_ID, ALPHA_TOKEN), (BETA_ID, BETA_TOKEN)]); + let bus = Arc::new(RecordingBus::new()); + let state = base_state(&mock_url, None, None, false) + .with_client_auth(Arc::new(auth)) + .with_event_bus(Arc::clone(&bus) as Arc); + let app = build_router(state); + + // alpha's traffic first, so both callers' material is in flight in + // the process before beta's request is observed. + let (status, _, _) = send( + app.clone(), + post_messages(Some(ALPHA_TOKEN), &marked_body("alpha traffic")), + ) + .await; + assert_eq!(status, StatusCode::OK); + + // Everything beta can observe: its HTTP response, the events its + // request emits, and the log lines its request produces. None of + // them may name alpha. + let capture = TraceCapture::install(); + let (status, body, request_id) = send( + app.clone(), + post_messages(Some(BETA_TOKEN), &marked_body("beta traffic")), + ) + .await; + assert_eq!(status, StatusCode::OK); + assert!( + !body.contains(ALPHA_ID), + "beta's response names alpha's key id: {body}" + ); + + // Log lines correlated with beta's request (its request id in the + // span, or its own caller attribution) must never name alpha. + let request_id = request_id.expect("x-request-id header present"); + let captured = capture.text(); + let correlated: Vec<&str> = captured + .lines() + .filter(|line| line.contains(&request_id) || line.contains(BETA_ID)) + .collect(); + assert!( + !correlated.is_empty(), + "control: beta's traffic must produce log lines" + ); + for line in &correlated { + assert!( + !line.contains(ALPHA_ID), + "a log line for beta's request names alpha's key id: {line}" + ); + } + + let beta_event_json = events_for_request(&bus, &request_id); + assert!(beta_event_json.len() >= 2, "beta's request emitted events"); + for json in &beta_event_json { + assert!( + !json.contains(ALPHA_ID), + "an event for beta's request names alpha's key id: {json}" + ); + } + // Self-attribution intact: the received event names beta, and only + // beta, as the caller. + let beta_callers: Vec = bus + .snapshot() + .into_iter() + .filter_map(|event| match event { + ProxyEvent::RequestReceived(received) if received.request_id == request_id => { + Some(received.caller) + } + _ => None, + }) + .collect(); + assert_eq!( + beta_callers, + vec![BETA_ID.to_owned()], + "self-attribution intact" + ); + + // beta's client key must not open the attribution-bearing admin + // surface (both mounts refuse it — see + // `admin_surface_refuses_non_admin_credentials_on_both_mounts`). + } + + // ----------------------------------------------------------------------- + // 6. Attribution survives the loopback no-auth exemption + // ----------------------------------------------------------------------- + + #[tokio::test] + async fn anonymous_attribution_survives_the_loopback_exemption() { + let dir = tempfile::tempdir().expect("tempdir"); + let db = dir.path().join("abuse.db"); + let mock_url = spawn_anthropic_mock().await; + let bus = Arc::new(open_ledger(&db)); + let state = base_state(&mock_url, Some(db.clone()), Some(admin_config()), false) + .with_event_bus(Arc::clone(&bus) as Arc); + let app = build_router(state.clone()); + + // Under auth mode "none" the client still sends a bearer token (LLM + // SDKs do habitually). The proxy must neither attribute the request + // to that token nor let it leak anywhere. + let body = marked_body(BODY_MARKER); + let (status, response, _) = + send(app.clone(), post_messages(Some(STRAY_TOKEN), &body)).await; + assert_eq!(status, StatusCode::OK); + assert_free_of_secrets("response under the loopback exemption", &response); + + close_ledger(Arc::clone(&bus)).await; + let admin = admin_router(&state).expect("admin router"); + let (status, keys, _) = + send(admin.clone(), get_usage("/usage/keys", Some(ADMIN_TOKEN))).await; + assert_eq!(status, StatusCode::OK); + assert!( + keys.contains("anonymous"), + "attribution survives as the anonymous caller: {keys}" + ); + assert_free_of_secrets("/usage/keys under the loopback exemption", &keys); + assert!( + !keys.contains(STRAY_TOKEN), + "the stray bearer must not become the caller: {keys}" + ); + + let (status, rows, _) = send( + admin, + get_usage("/usage/requests?limit=1000", Some(ADMIN_TOKEN)), + ) + .await; + assert_eq!(status, StatusCode::OK); + assert!( + rows.contains("anonymous") && rows.contains(&hex_sha256(body.as_bytes())), + "the anonymous row carries the digest: {rows}" + ); + assert_free_of_secrets("/usage/requests under the loopback exemption", &rows); + + let bytes = ledger_bytes(&db); + let text = String::from_utf8_lossy(&bytes); + assert!( + text.contains("anonymous"), + "control: the anonymous caller is persisted" + ); + assert_free_of_secrets("ledger bytes under the loopback exemption", &text); + } +} diff --git a/crates/llm-proxy-storage/Cargo.toml b/crates/llm-proxy-storage/Cargo.toml index 0ce7f1a..ca1a376 100644 --- a/crates/llm-proxy-storage/Cargo.toml +++ b/crates/llm-proxy-storage/Cargo.toml @@ -20,8 +20,14 @@ test-utils = [] [dependencies] llm-proxy-protocol = { path = "../llm-proxy-protocol" } +rusqlite = { workspace = true } serde = { workspace = true } +thiserror = { workspace = true } time = { workspace = true } +tokio = { workspace = true } +tracing = { workspace = true } [dev-dependencies] +rust_decimal = { workspace = true } serde_json = { workspace = true } +tempfile = "3" diff --git a/crates/llm-proxy-storage/src/lib.rs b/crates/llm-proxy-storage/src/lib.rs index a6b47c9..8525bce 100644 --- a/crates/llm-proxy-storage/src/lib.rs +++ b/crates/llm-proxy-storage/src/lib.rs @@ -1,13 +1,14 @@ -//! Structured request/response event log surface. +//! Structured request/response event log surface and persistence. //! -//! The persistence backend (SQLite, etc.) lands in a later phase. This crate -//! defines the event types and the [`EventBus`] trait that backends implement, -//! plus a [`NoopBus`] default and a [`RecordingBus`] for tests. +//! This crate defines the event types and the [`EventBus`] trait that +//! backends implement: [`NoopBus`] (the default when no `[storage]` config +//! is present), [`RecordingBus`] for tests, and [`SqliteEventBus`], the +//! persistent per-request usage ledger. //! -//! Events are emitted at request/response boundaries in the server crate. A -//! real persisted backend will implement [`EventBus`]; for now [`NoopBus`] is -//! the production sink and [`RecordingBus`] (behind the `test-utils` cargo -//! feature) lets integration tests assert on the event sequence. +//! Events are emitted at request/response boundaries in the server crate. +//! [`SqliteEventBus`] (behind the storage config) persists them to SQLite; +//! [`RecordingBus`] (behind the `test-utils` cargo feature) lets integration +//! tests assert on the event sequence. #![deny(missing_docs)] @@ -15,6 +16,10 @@ use llm_proxy_protocol::core::{Cost, ModelRef, StopReason, Usage}; use serde::{Deserialize, Serialize}; use time::OffsetDateTime; +mod sqlite; + +pub use sqlite::{RetentionPolicy, SqliteEventBus, StorageError, run_retention_pass}; + // --------------------------------------------------------------------------- // ProxyEvent // --------------------------------------------------------------------------- @@ -46,6 +51,10 @@ pub struct RequestReceived { /// When the request was received (RFC 3339). #[serde(with = "time::serde::rfc3339")] pub timestamp: OffsetDateTime, + /// Authenticated client key id, or `"anonymous"` when the request was + /// served under the loopback no-auth exemption. Key ids are operator + /// chosen; client secrets never appear on events. + pub caller: String, /// Provider name from the URL path. pub provider: String, /// Route kind (`"chat_completions"` or `"messages"`). @@ -133,8 +142,8 @@ pub struct ResponseFailed { /// In-process event sink. /// -/// Implementations: [`NoopBus`], [`RecordingBus`] (tests), and the future -/// SQLite backend. +/// Implementations: [`NoopBus`], [`RecordingBus`] (tests), and +/// [`SqliteEventBus`] (persistent ledger). /// /// `Debug` is a super-trait so that `Arc` can be formatted in /// `AppState`'s manual `Debug` impl. @@ -142,6 +151,16 @@ pub trait EventBus: Send + Sync + std::fmt::Debug { /// Emit one event. Must not block the calling task; backends do IO on a /// dedicated writer task. fn emit(&self, event: &ProxyEvent); + + /// Signal that no further events will be emitted, flush anything still + /// queued, and release the writer. + /// + /// Called once at process shutdown, after the HTTP server has stopped + /// serving (so no new events can race the flush). The default is a + /// no-op, which is correct for stateless sinks. Implementations must + /// stay bounded: give up after a finite wait rather than hang shutdown. + /// May block briefly; wrap in `spawn_blocking` from async code. + fn close(&self) {} } /// Default no-op sink. @@ -204,6 +223,7 @@ mod tests { RequestReceived { request_id: "req-1".to_owned(), timestamp: OffsetDateTime::now_utc(), + caller: "anonymous".to_owned(), provider: "p".to_owned(), route_kind: "messages".to_owned(), client_protocol: "anthropic".to_owned(), @@ -275,6 +295,7 @@ mod tests { let event = ProxyEvent::RequestReceived(RequestReceived { request_id: "req-1".to_owned(), timestamp, + caller: "anonymous".to_owned(), provider: "p".to_owned(), route_kind: "messages".to_owned(), client_protocol: "anthropic".to_owned(), @@ -413,6 +434,7 @@ mod tests { let req = ProxyEvent::RequestReceived(RequestReceived { request_id: SENTINEL.to_owned(), timestamp: OffsetDateTime::now_utc(), + caller: SENTINEL.to_owned(), provider: SENTINEL.to_owned(), route_kind: SENTINEL.to_owned(), client_protocol: SENTINEL.to_owned(), diff --git a/crates/llm-proxy-storage/src/sqlite.rs b/crates/llm-proxy-storage/src/sqlite.rs new file mode 100644 index 0000000..b11d614 --- /dev/null +++ b/crates/llm-proxy-storage/src/sqlite.rs @@ -0,0 +1,1370 @@ +//! SQLite-backed [`EventBus`]: the persistent per-request usage ledger. +//! +//! One `requests` row per request id, filled in as the request's events +//! arrive: `RequestReceived` inserts the row, and the terminal event +//! (`ResponseCompleted` or `ResponseFailed`) completes it. Requests that fail +//! before routing (client-auth rejections, unknown routes) produce only a +//! terminal event, so those rows carry `received_at IS NULL` and the failure +//! timestamp in `completed_at`. +//! +//! Writes are group-committed by a single dedicated tokio task behind a +//! bounded channel; [`EventBus::emit`] never blocks the calling task and +//! drops (and counts) events when the channel is full. + +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Condvar, Mutex}; +use std::time::{Duration, Instant}; + +use llm_proxy_protocol::core::{StopReason, UsageProvenance}; +use rusqlite::Connection; +use time::OffsetDateTime; +use tokio::sync::mpsc::{self, Sender}; +use tokio::sync::watch; + +use crate::{EventBus, ProxyEvent}; + +/// Bounded-channel capacity for the default constructor. +const DEFAULT_CHANNEL_CAPACITY: usize = 1024; + +/// Maximum events pulled from the channel and committed per transaction. +const WRITE_BATCH_MAX: usize = 256; + +/// How long [`EventBus::close`] waits for the writer to drain before giving +/// up (shutdown must stay bounded). +const CLOSE_DRAIN_TIMEOUT: Duration = Duration::from_secs(10); + +/// SQLite busy timeout for concurrent connections (tests, admin readers). +const SQLITE_BUSY_TIMEOUT: Duration = Duration::from_secs(5); + +// --------------------------------------------------------------------------- +// Errors +// --------------------------------------------------------------------------- + +/// Errors from the SQLite usage ledger. +#[derive(Debug, thiserror::Error)] +pub enum StorageError { + /// Failed to create the ledger file's parent directory. + #[error("creating usage-ledger parent directory {path}: {source}")] + CreateDir { + /// Parent directory that could not be created. + path: String, + /// Underlying I/O error. + source: std::io::Error, + }, + /// Failed to open or initialize the SQLite database. + #[error("opening usage ledger at {path}: {source}")] + Open { + /// Database path. + path: String, + /// Underlying rusqlite error. + source: rusqlite::Error, + }, + /// A write or retention operation on an open ledger failed. + #[error("usage ledger operation failed: {0}")] + Sqlite(#[from] rusqlite::Error), +} + +// --------------------------------------------------------------------------- +// Retention +// --------------------------------------------------------------------------- + +/// Retention bounds applied to the ledger on a timer. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct RetentionPolicy { + /// Delete rows older than this age. `None` keeps rows forever. + pub max_age: Option, + /// Keep at most this many rows; the oldest rows are deleted when the cap + /// is exceeded. `None` is unlimited. + pub max_rows: Option, + /// How often the pruning pass runs. Must be at least 1 second (the + /// pruning timer rejects a zero period; `load_app_config` enforces this + /// for config-file users). + pub prune_every: Duration, +} + +impl RetentionPolicy { + /// Whether any retention bound is configured. + fn is_bounded(&self) -> bool { + self.max_age.is_some() || self.max_rows.is_some() + } +} + +/// Run one retention pass, deleting rows past `policy`'s bounds. +/// +/// Public so operators/tests can prune deterministically; the writer task +/// calls it on the [`RetentionPolicy::prune_every`] timer. Age pruning uses +/// `COALESCE(received_at, completed_at)` so failure-only rows (no +/// `received_at`) are still subject to retention. +/// +/// Returns the number of rows deleted. +/// +/// # Errors +/// +/// Returns [`StorageError::Sqlite`] if any statement fails. +pub fn run_retention_pass( + conn: &mut Connection, + policy: &RetentionPolicy, +) -> Result { + if !policy.is_bounded() { + return Ok(0); + } + let mut deleted: u64 = 0; + if let Some(max_age) = policy.max_age { + // A max_age whose millis do not fit i64 (Duration::MAX) saturates to + // "the beginning of time", deleting nothing — the safe direction for + // an absurd value. + let cutoff = unix_millis(OffsetDateTime::now_utc()) + .saturating_sub(i64::try_from(max_age.as_millis()).unwrap_or(i64::MAX)); + let n = conn.execute( + "DELETE FROM requests WHERE COALESCE(received_at, completed_at) < ?1", + rusqlite::params![cutoff], + )?; + deleted += u64::try_from(n).unwrap_or(u64::MAX); + } + if let Some(max_rows) = policy.max_rows { + let count: i64 = conn.query_row("SELECT COUNT(*) FROM requests", [], |row| row.get(0))?; + let excess = u64::try_from(count) + .unwrap_or(u64::MAX) + .saturating_sub(max_rows); + if excess > 0 { + let limit = i64::try_from(excess).unwrap_or(i64::MAX); + let n = conn.execute( + "DELETE FROM requests WHERE request_id IN ( + SELECT request_id FROM requests + ORDER BY COALESCE(received_at, completed_at) ASC + LIMIT ?1 + )", + rusqlite::params![limit], + )?; + deleted += u64::try_from(n).unwrap_or(u64::MAX); + } + } + Ok(deleted) +} + +// --------------------------------------------------------------------------- +// Schema +// --------------------------------------------------------------------------- + +/// Ledger schema. Costs are stored as decimal strings (never floats or +/// reals); `body_hash` is the request-body *digest* from +/// [`crate::RequestReceived`], never the body itself. +const SCHEMA_SQL: &str = "\ +CREATE TABLE IF NOT EXISTS requests ( + request_id TEXT PRIMARY KEY, + received_at INTEGER, + completed_at INTEGER, + caller TEXT, + provider TEXT, + route_kind TEXT, + client_protocol TEXT, + streaming INTEGER, + model_requested TEXT, + model_upstream TEXT, + input_tokens INTEGER, + output_tokens INTEGER, + reasoning_tokens INTEGER, + cache_read_tokens INTEGER, + cache_creation_tokens INTEGER, + usage_provenance TEXT, + cost_input TEXT, + cost_output TEXT, + cost_cache_creation TEXT, + cost_cache_read TEXT, + cost_reasoning TEXT, + cost_total TEXT, + stop_reason TEXT, + error_kind TEXT, + http_status INTEGER, + latency_ms INTEGER, + body_hash TEXT +); +CREATE INDEX IF NOT EXISTS idx_requests_timestamp + ON requests (received_at); +CREATE INDEX IF NOT EXISTS idx_requests_provider_timestamp + ON requests (provider, received_at); +CREATE INDEX IF NOT EXISTS idx_requests_caller_timestamp + ON requests (caller, received_at); +"; + +/// Open a WAL-mode connection and make sure the schema exists. +fn open_connection(path: &Path) -> Result { + let open_err = |source: rusqlite::Error| StorageError::Open { + path: path.display().to_string(), + source, + }; + let conn = Connection::open(path).map_err(open_err)?; + // WAL keeps readers (tests, the future admin API) lock-free while the + // single writer commits. `pragma_update` errors if the engine cannot + // honor the mode (e.g. a filesystem without shared memory), which is a + // startup error rather than a silent downgrade. + conn.pragma_update(None, "journal_mode", "WAL") + .map_err(open_err)?; + // NORMAL is the recommended WAL pairing: commits do not fsync the WAL on + // every transaction; durability is preserved across process crashes and + // only a power loss can lose the last commits (the database never + // corrupts). + conn.pragma_update(None, "synchronous", "NORMAL") + .map_err(open_err)?; + conn.busy_timeout(SQLITE_BUSY_TIMEOUT).map_err(open_err)?; + conn.execute_batch(SCHEMA_SQL).map_err(open_err)?; + Ok(conn) +} + +// --------------------------------------------------------------------------- +// Event → row mapping +// --------------------------------------------------------------------------- + +/// OffsetDateTime → Unix epoch milliseconds. +fn unix_millis(ts: OffsetDateTime) -> i64 { + ts.unix_timestamp() + .saturating_mul(1000) + .saturating_add(i64::from(ts.millisecond())) +} + +/// Stable, human-readable label for a stop reason (used verbatim by ledger +/// consumers grouping on this column). +fn stop_reason_label(reason: &StopReason) -> String { + match reason { + StopReason::EndTurn => "end_turn".to_owned(), + StopReason::MaxTokens => "max_tokens".to_owned(), + StopReason::ToolUse => "tool_use".to_owned(), + StopReason::StopSequence => "stop_sequence".to_owned(), + StopReason::Refusal => "refusal".to_owned(), + StopReason::Error => "error".to_owned(), + StopReason::Unknown(inner) => format!("unknown:{inner}"), + _ => "other".to_owned(), + } +} + +/// Stable label for usage provenance (`Usage`'s companion field; lets ledger +/// consumers distinguish provider-reported usage from synthetic zeros). +fn provenance_label(provenance: &UsageProvenance) -> &'static str { + match provenance { + UsageProvenance::ProviderReported => "provider_reported", + UsageProvenance::SyntheticZero => "synthetic_zero", + UsageProvenance::Unknown => "unknown", + _ => "other", + } +} + +/// Write one event's row effects. Terminal events upsert, so a failure for a +/// request that was never received still leaves a row. +fn write_event(conn: &Connection, event: &ProxyEvent) -> Result<(), StorageError> { + match event { + ProxyEvent::RequestReceived(e) => insert_request_received(conn, e), + ProxyEvent::ResponseCompleted(e) => upsert_response_completed(conn, e), + ProxyEvent::ResponseFailed(e) => upsert_response_failed(conn, e), + } +} + +/// Insert the request row. A duplicate request id (retry, bug) must not +/// clobber the original row. +fn insert_request_received( + conn: &Connection, + e: &crate::RequestReceived, +) -> Result<(), StorageError> { + conn.execute( + "INSERT OR IGNORE INTO requests ( + request_id, received_at, caller, provider, route_kind, + client_protocol, streaming, model_requested, model_upstream, + body_hash + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)", + rusqlite::params![ + e.request_id, + unix_millis(e.timestamp), + e.caller, + e.provider, + e.route_kind, + e.client_protocol, + e.streaming, + e.model.requested, + e.model.upstream, + e.body_hash, + ], + )?; + Ok(()) +} + +/// Complete the request row with usage, cost, and stop reason (or insert it, +/// if the request event was never seen). +fn upsert_response_completed( + conn: &Connection, + e: &crate::ResponseCompleted, +) -> Result<(), StorageError> { + conn.execute( + "INSERT INTO requests ( + request_id, completed_at, provider, model_requested, + model_upstream, input_tokens, output_tokens, reasoning_tokens, + cache_read_tokens, cache_creation_tokens, usage_provenance, + cost_input, cost_output, cost_cache_creation, cost_cache_read, + cost_reasoning, cost_total, stop_reason, latency_ms + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, + ?13, ?14, ?15, ?16, ?17, ?18, ?19) + ON CONFLICT(request_id) DO UPDATE SET + completed_at = excluded.completed_at, + provider = excluded.provider, + model_requested = excluded.model_requested, + model_upstream = excluded.model_upstream, + input_tokens = excluded.input_tokens, + output_tokens = excluded.output_tokens, + reasoning_tokens = excluded.reasoning_tokens, + cache_read_tokens = excluded.cache_read_tokens, + cache_creation_tokens = excluded.cache_creation_tokens, + usage_provenance = excluded.usage_provenance, + cost_input = excluded.cost_input, + cost_output = excluded.cost_output, + cost_cache_creation = excluded.cost_cache_creation, + cost_cache_read = excluded.cost_cache_read, + cost_reasoning = excluded.cost_reasoning, + cost_total = excluded.cost_total, + stop_reason = excluded.stop_reason, + latency_ms = excluded.latency_ms", + rusqlite::params![ + e.request_id, + unix_millis(e.timestamp), + e.provider, + e.model.requested, + e.model.upstream, + e.usage.input_tokens, + e.usage.output_tokens, + e.usage.reasoning_tokens, + e.usage.cache_read_input_tokens, + e.usage.cache_creation_input_tokens, + provenance_label(&e.usage.provenance), + e.cost.as_ref().map(|c| c.input.to_string()), + e.cost.as_ref().map(|c| c.output.to_string()), + e.cost.as_ref().map(|c| c.cache_creation.to_string()), + e.cost.as_ref().map(|c| c.cache_read.to_string()), + e.cost.as_ref().map(|c| c.reasoning.to_string()), + e.cost.as_ref().map(|c| c.total.to_string()), + stop_reason_label(&e.stop_reason), + i64::try_from(e.latency_ms).unwrap_or(i64::MAX), + ], + )?; + Ok(()) +} + +/// Complete the request row with failure facts (or insert it, if the request +/// event was never seen). provider/model are optional on failures; COALESCE +/// keeps values already stored by the request's RequestReceived row instead +/// of overwriting them with NULL. +fn upsert_response_failed( + conn: &Connection, + e: &crate::ResponseFailed, +) -> Result<(), StorageError> { + conn.execute( + "INSERT INTO requests ( + request_id, completed_at, provider, model_requested, + model_upstream, error_kind, http_status, latency_ms + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8) + ON CONFLICT(request_id) DO UPDATE SET + completed_at = excluded.completed_at, + provider = COALESCE(excluded.provider, requests.provider), + model_requested = COALESCE( + excluded.model_requested, requests.model_requested), + model_upstream = COALESCE( + excluded.model_upstream, requests.model_upstream), + error_kind = excluded.error_kind, + http_status = excluded.http_status, + latency_ms = excluded.latency_ms", + rusqlite::params![ + e.request_id, + unix_millis(e.timestamp), + e.provider, + e.model.as_ref().map(|m| m.requested.clone()), + e.model.as_ref().and_then(|m| m.upstream.clone()), + e.error_kind, + i64::from(e.http_status), + i64::try_from(e.latency_ms).unwrap_or(i64::MAX), + ], + )?; + Ok(()) +} + +// --------------------------------------------------------------------------- +// Writer task +// --------------------------------------------------------------------------- + +/// One unit of blocking work for the writer's connection. +enum WriterOp { + /// Commit a batch of events in one transaction (group commit). + Write { + /// Events to persist. + events: Vec, + }, + /// Apply the retention policy. + Prune { + /// Bounds to enforce. + policy: RetentionPolicy, + }, +} + +fn run_op(conn: &mut Connection, op: WriterOp) -> Result { + match op { + WriterOp::Write { events } => { + let tx = conn.transaction()?; + for event in &events { + write_event(&tx, event)?; + } + tx.commit()?; + Ok(u64::try_from(events.len()).unwrap_or(u64::MAX)) + } + WriterOp::Prune { policy } => run_retention_pass(conn, &policy), + } +} + +/// Run `op` on the writer's connection on the blocking pool, then hand the +/// connection back. If the blocking closure panics the original connection is +/// lost inside it; a fresh one is opened so the writer can continue. Returns +/// `None` only when the ledger cannot be reopened (the writer stops and +/// subsequent emits are counted as dropped). +async fn exec_on_conn( + path: &Path, + conn: Connection, + op: WriterOp, +) -> Option<(Connection, Result)> { + let path = path.to_path_buf(); + let joined = tokio::task::spawn_blocking(move || { + let mut conn = conn; + let result = run_op(&mut conn, op); + (conn, result) + }) + .await; + match joined { + Ok((conn, result)) => Some((conn, result)), + Err(join_error) => { + tracing::error!( + error = %join_error, + "usage-ledger writer panicked; reopening the database" + ); + let reopen_path = path.clone(); + let reopened = tokio::task::spawn_blocking(move || open_connection(&reopen_path)).await; + match reopened { + Ok(Ok(conn)) => Some((conn, Ok(0))), + Ok(Err(error)) => { + tracing::error!( + error = %error, + path = %path.display(), + "failed to reopen usage ledger; writer stopping" + ); + None + } + Err(join_error) => { + tracing::error!( + error = %join_error, + "reopening the usage ledger panicked; writer stopping" + ); + None + } + } + } + } +} + +/// Flag the writer sets once it has exited (drained or failed). +type DoneFlag = Arc<(Mutex, Condvar)>; + +fn signal_done(done: &DoneFlag) { + let mut guard = done.0.lock().unwrap_or_else(|e| e.into_inner()); + *guard = true; + drop(guard); + done.1.notify_all(); +} + +/// The single ledger writer: drains the bounded channel in batches, applies +/// retention on a timer, and drains-and-exits when closed. +async fn writer_task( + path: PathBuf, + conn: Connection, + mut rx: mpsc::Receiver, + mut shutdown: watch::Receiver, + retention: RetentionPolicy, + done: DoneFlag, +) { + let pruning = retention.is_bounded(); + // Defensively clamp: a zero period would panic the interval timer. + // Config validation rejects it, direct construction might not. + let mut prune = tokio::time::interval(retention.prune_every.max(Duration::from_secs(1))); + prune.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + // The connection moves into each blocking call and back out; the Option + // slot lets a path that loses the connection (unrecoverable failure) + // break the loop without the borrow checker seeing a moved value on the + // paths that continue. + let mut conn_slot = Some(conn); + let mut buf: Vec = Vec::with_capacity(WRITE_BATCH_MAX); + + loop { + // Poll the shutdown value each iteration so a signal that arrived + // while a batch was being written is never missed. + if *shutdown.borrow_and_update() { + break; + } + tokio::select! { + biased; + _ = shutdown.changed() => { + if *shutdown.borrow() { + break; + } + } + received = rx.recv_many(&mut buf, WRITE_BATCH_MAX) => { + if received == 0 { + // Every sender dropped; nothing more can arrive. + break; + } + let Some(conn) = conn_slot.take() else { break }; + match exec_on_conn( + &path, + conn, + WriterOp::Write { events: std::mem::take(&mut buf) }, + ) + .await + { + Some((c, result)) => { + conn_slot = Some(c); + if let Err(error) = result { + tracing::error!( + error = %error, + "usage-ledger batch write failed; batch events lost" + ); + } + } + None => break, + } + } + _ = prune.tick(), if pruning => { + if let Some(conn) = conn_slot.take() { + match exec_on_conn( + &path, + conn, + WriterOp::Prune { policy: retention }, + ) + .await + { + Some((c, result)) => { + conn_slot = Some(c); + match result { + Ok(rows) if rows > 0 => tracing::info!( + rows, + "pruned usage-ledger rows past retention bounds" + ), + Ok(_) => {} + Err(error) => tracing::error!( + error = %error, + "usage-ledger retention pass failed" + ), + } + } + None => break, + } + } + } + } + } + + // Final drain so close() is a flush: everything queued behind the + // shutdown signal is persisted before the writer exits. + while let Ok(event) = rx.try_recv() { + buf.push(event); + } + if !buf.is_empty() { + if let Some(conn) = conn_slot.take() { + match exec_on_conn(&path, conn, WriterOp::Write { events: buf }).await { + Some((_conn, Ok(_))) => {} + Some((_conn, Err(error))) => { + tracing::error!(error = %error, "usage-ledger final drain failed"); + } + None => {} + } + } + } + + signal_done(&done); +} + +// --------------------------------------------------------------------------- +// SqliteEventBus +// --------------------------------------------------------------------------- + +/// SQLite-backed [`EventBus`]: persists every event to one `requests` ledger. +/// +/// `emit` hands the event to a dedicated writer task through a bounded +/// channel and never blocks: a full channel drops the event and increments a +/// dropped-events counter ([`SqliteEventBus::dropped_events`]) rather than +/// back-pressuring the response path. +/// +/// # Runtime requirement +/// +/// [`SqliteEventBus::open`] spawns the writer task and must be called from +/// within a tokio runtime. +#[derive(Debug)] +pub struct SqliteEventBus { + tx: Sender, + shutdown_tx: Mutex>, + done: DoneFlag, + dropped: AtomicU64, + path: PathBuf, +} + +impl SqliteEventBus { + /// Open (or create) the ledger at `path` and start the writer task. + /// + /// The database is opened in WAL mode with the `requests` schema and its + /// indexes created if absent; parent directories of `path` are created. + /// Retention bounds, when configured, are enforced on a timer. + /// + /// # Errors + /// + /// Returns [`StorageError`] if the directory cannot be created, the + /// database cannot be opened (or WAL mode refused), or the schema cannot + /// be initialized. Blocking filesystem work happens here; call at + /// startup, not on the hot path. + pub fn open(path: impl AsRef, retention: RetentionPolicy) -> Result { + Self::open_with_capacity(path, retention, DEFAULT_CHANNEL_CAPACITY) + } + + /// Like [`SqliteEventBus::open`] with an explicit channel capacity. + fn open_with_capacity( + path: impl AsRef, + retention: RetentionPolicy, + capacity: usize, + ) -> Result { + let path = path.as_ref(); + if let Some(parent) = path.parent() { + if !parent.as_os_str().is_empty() { + std::fs::create_dir_all(parent).map_err(|source| StorageError::CreateDir { + path: parent.display().to_string(), + source, + })?; + } + } + let conn = open_connection(path)?; + let (tx, rx) = mpsc::channel(capacity.max(1)); + let (shutdown_tx, shutdown_rx) = watch::channel(false); + let done: DoneFlag = Arc::new((Mutex::new(false), Condvar::new())); + tokio::spawn(writer_task( + path.to_path_buf(), + conn, + rx, + shutdown_rx, + retention, + Arc::clone(&done), + )); + Ok(Self { + tx, + shutdown_tx: Mutex::new(shutdown_tx), + done, + dropped: AtomicU64::new(0), + path: path.to_path_buf(), + }) + } + + /// Test-only constructor with a live channel but no writer task and no + /// database: the channel fills and never drains, so drop-on-full behavior + /// is deterministic. + #[cfg(test)] + fn unconnected(capacity: usize) -> Self { + // Leak the receiver so it outlives this constructor (a local binding + // would drop at return, making try_send report Closed instead of + // Full). The process-exit reclaim makes the leak irrelevant. + let (tx, rx) = mpsc::channel(capacity.max(1)); + let (shutdown_tx, _shutdown_rx) = watch::channel(false); + Box::leak(Box::new(rx)); + Self { + tx, + shutdown_tx: Mutex::new(shutdown_tx), + done: Arc::new((Mutex::new(false), Condvar::new())), + dropped: AtomicU64::new(0), + path: PathBuf::new(), + } + } + + /// Number of events dropped so far (channel full, or writer gone). + #[must_use] + pub fn dropped_events(&self) -> u64 { + self.dropped.load(Ordering::Relaxed) + } + + /// Path of the backing database file. + #[must_use] + pub fn path(&self) -> &Path { + &self.path + } +} + +impl EventBus for SqliteEventBus { + fn emit(&self, event: &ProxyEvent) { + match self.tx.try_send(event.clone()) { + Ok(()) => {} + Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { + let dropped = self.dropped.fetch_add(1, Ordering::Relaxed) + 1; + tracing::debug!( + dropped, + "usage-ledger channel full; event dropped (never backpressure a response)" + ); + } + Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => { + let dropped = self.dropped.fetch_add(1, Ordering::Relaxed) + 1; + tracing::debug!(dropped, "usage-ledger writer is gone; event dropped"); + } + } + } + + fn close(&self) { + // Signal first; an error means the writer already exited and the done + // flag is already set. + if let Ok(shutdown_tx) = self.shutdown_tx.lock() { + if shutdown_tx.send(true).is_err() { + tracing::debug!("usage-ledger writer had already exited"); + } + } + // Bounded wait for the writer's final drain. + let deadline = Instant::now() + CLOSE_DRAIN_TIMEOUT; + let mut guard = self.done.0.lock().unwrap_or_else(|e| e.into_inner()); + while !*guard { + let now = Instant::now(); + if now >= deadline { + tracing::warn!( + timeout_secs = CLOSE_DRAIN_TIMEOUT.as_secs(), + "usage-ledger close timed out; queued events may not be flushed" + ); + return; + } + let (next, timeout) = self + .done + .1 + .wait_timeout(guard, deadline - now) + .unwrap_or_else(|e| e.into_inner()); + guard = next; + if timeout.timed_out() { + continue; + } + } + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + use crate::{RequestReceived, ResponseCompleted, ResponseFailed}; + use llm_proxy_protocol::core::{Cost, ModelRef, Usage}; + use rust_decimal::Decimal; + + const RECEIVED_MS: i64 = 1_700_000_000_000; + const COMPLETED_MS: i64 = 1_700_000_004_000; + + fn unbounded() -> RetentionPolicy { + RetentionPolicy { + max_age: None, + max_rows: None, + prune_every: Duration::from_secs(3600), + } + } + + fn received_event(id: &str, caller: &str) -> ProxyEvent { + ProxyEvent::RequestReceived(RequestReceived { + request_id: id.to_owned(), + timestamp: OffsetDateTime::from_unix_timestamp(RECEIVED_MS / 1000) + .expect("valid unix timestamp"), + caller: caller.to_owned(), + provider: "prov-a".to_owned(), + route_kind: "messages".to_owned(), + client_protocol: "anthropic".to_owned(), + model: ModelRef { + requested: "claude-3".to_owned(), + upstream: Some("claude-3-upstream".to_owned()), + }, + streaming: true, + body_hash: "sha256-digest-only".to_owned(), + }) + } + + fn completed_event(id: &str) -> ProxyEvent { + ProxyEvent::ResponseCompleted(ResponseCompleted { + request_id: id.to_owned(), + timestamp: OffsetDateTime::from_unix_timestamp(COMPLETED_MS / 1000) + .expect("valid unix timestamp"), + provider: "prov-a".to_owned(), + upstream_message_id: Some("msg_7".to_owned()), + model: ModelRef { + requested: "claude-3".to_owned(), + upstream: None, + }, + usage: Usage { + input_tokens: 100, + output_tokens: 50, + reasoning_tokens: Some(7), + cache_creation_input_tokens: Some(30), + cache_read_input_tokens: Some(200), + provenance: UsageProvenance::ProviderReported, + }, + cost: Some(Cost { + input: Decimal::from_str_exact("0.000123456789012345678901234") + .expect("decimal fits"), + output: Decimal::from_str_exact("2.25").expect("decimal fits"), + cache_creation: Decimal::ZERO, + cache_read: Decimal::from_str_exact("0.0000005").expect("decimal fits"), + reasoning: Decimal::ZERO, + total: Decimal::from_str_exact("2.250123956789012345678901234") + .expect("decimal fits"), + }), + stop_reason: StopReason::Unknown("weird-finish".to_owned()), + latency_ms: 4321, + }) + } + + fn failed_event(id: &str, provider: Option<&str>) -> ProxyEvent { + ProxyEvent::ResponseFailed(ResponseFailed { + request_id: id.to_owned(), + timestamp: OffsetDateTime::from_unix_timestamp(COMPLETED_MS / 1000) + .expect("valid unix timestamp"), + provider: provider.map(str::to_owned), + model: None, + error_kind: "UpstreamError".to_owned(), + message: "sanitized failure text".to_owned(), + http_status: 502, + latency_ms: 99, + }) + } + + /// Poll until `cond` holds (2s budget at 10ms cadence). + async fn wait_until(description: &str, mut cond: impl FnMut() -> bool) { + for _ in 0..200 { + if cond() { + return; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + panic!("condition not met within 2s: {description}"); + } + + /// Open a second connection to the ledger for assertions. + fn reader(path: &Path) -> Connection { + Connection::open(path).expect("open ledger for assertions") + } + + fn scalar(conn: &Connection, sql: &str) -> i64 { + conn.query_row(sql, [], |row| row.get(0)) + .expect("scalar query") + } + + fn text_column(conn: &Connection, sql: &str) -> Option { + conn.query_row(sql, [], |row| row.get::<_, Option>(0)) + .expect("text query") + } + + #[tokio::test] + async fn open_enables_wal_and_creates_schema_and_indexes() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("usage.db"); + let bus = SqliteEventBus::open(&db, unbounded()).expect("open ledger"); + let bus = Arc::new(bus); + let closer = Arc::clone(&bus); + tokio::task::spawn_blocking(move || closer.close()) + .await + .expect("close"); + + let conn = reader(&db); + let mode: String = conn + .query_row("PRAGMA journal_mode", [], |row| row.get(0)) + .expect("journal mode"); + assert_eq!(mode, "wal", "ledger must run in WAL mode"); + + let indexes: Vec = { + let mut stmt = conn + .prepare("SELECT name FROM sqlite_master WHERE type = 'index'") + .expect("prepare index query"); + let rows = stmt + .query_map([], |row| row.get::<_, String>(0)) + .expect("query indexes"); + let mut names = Vec::new(); + for row in rows { + names.push(row.expect("index row")); + } + names + }; + for expected in [ + "idx_requests_timestamp", + "idx_requests_provider_timestamp", + "idx_requests_caller_timestamp", + ] { + assert!( + indexes.contains(&expected.to_owned()), + "missing index {expected}; got {indexes:?}" + ); + } + assert_eq!( + scalar(&conn, "SELECT COUNT(*) FROM requests"), + 0, + "a fresh ledger starts empty" + ); + } + + #[tokio::test] + async fn request_received_persists_row_with_caller_and_body_hash() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("usage.db"); + let bus = SqliteEventBus::open(&db, unbounded()).expect("open ledger"); + bus.emit(&received_event("req-1", "laptop")); + let conn = reader(&db); + wait_until("request row appears", || { + scalar( + &conn, + "SELECT COUNT(*) FROM requests WHERE request_id = 'req-1'", + ) == 1 + }) + .await; + + let row = conn + .query_row( + "SELECT received_at, caller, provider, route_kind, client_protocol, \ + streaming, model_requested, model_upstream, body_hash \ + FROM requests WHERE request_id = 'req-1'", + [], + |row| { + Ok(( + row.get::<_, i64>(0)?, + row.get::<_, String>(1)?, + row.get::<_, String>(2)?, + row.get::<_, String>(3)?, + row.get::<_, String>(4)?, + row.get::<_, i64>(5)?, + row.get::<_, String>(6)?, + row.get::<_, Option>(7)?, + row.get::<_, String>(8)?, + )) + }, + ) + .expect("fetch row"); + assert_eq!(row.0, RECEIVED_MS); + assert_eq!(row.1, "laptop", "caller (PRX-01 key id) must persist"); + assert_eq!(row.2, "prov-a"); + assert_eq!(row.3, "messages"); + assert_eq!(row.4, "anthropic"); + assert_eq!(row.5, 1, "streaming flag stored as 1"); + assert_eq!(row.6, "claude-3"); + assert_eq!(row.7.as_deref(), Some("claude-3-upstream")); + assert_eq!(row.8, "sha256-digest-only"); + // No body column exists at all: only the digest is ever stored. + let has_body_column: i64 = scalar( + &conn, + "SELECT COUNT(*) FROM pragma_table_info('requests') WHERE name = 'body'", + ); + assert_eq!(has_body_column, 0, "no request-body column may exist"); + + let closer = Arc::new(bus); + tokio::task::spawn_blocking(move || closer.close()) + .await + .expect("close"); + } + + #[tokio::test] + async fn completed_event_fills_the_same_row_with_usage() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("usage.db"); + let bus = SqliteEventBus::open(&db, unbounded()).expect("open ledger"); + bus.emit(&received_event("req-2", "laptop")); + bus.emit(&completed_event("req-2")); + let conn = reader(&db); + wait_until("completed row is filled", || { + scalar( + &conn, + "SELECT COUNT(*) FROM requests \ + WHERE request_id = 'req-2' AND completed_at IS NOT NULL", + ) == 1 + }) + .await; + + assert_eq!( + scalar(&conn, "SELECT COUNT(*) FROM requests"), + 1, + "one request, one row: the terminal event updates in place" + ); + assert_eq!( + scalar( + &conn, + "SELECT input_tokens FROM requests WHERE request_id = 'req-2'" + ), + 100 + ); + assert_eq!( + scalar( + &conn, + "SELECT output_tokens FROM requests WHERE request_id = 'req-2'" + ), + 50 + ); + assert_eq!( + scalar( + &conn, + "SELECT reasoning_tokens FROM requests WHERE request_id = 'req-2'" + ), + 7 + ); + assert_eq!( + scalar( + &conn, + "SELECT cache_read_tokens FROM requests WHERE request_id = 'req-2'" + ), + 200 + ); + assert_eq!( + scalar( + &conn, + "SELECT cache_creation_tokens FROM requests WHERE request_id = 'req-2'" + ), + 30 + ); + assert_eq!( + text_column( + &conn, + "SELECT usage_provenance FROM requests WHERE request_id = 'req-2'" + ) + .as_deref(), + Some("provider_reported") + ); + assert_eq!( + scalar( + &conn, + "SELECT received_at FROM requests WHERE request_id = 'req-2'" + ), + RECEIVED_MS, + "the terminal event must not clobber received_at" + ); + assert_eq!( + scalar( + &conn, + "SELECT completed_at FROM requests WHERE request_id = 'req-2'" + ), + COMPLETED_MS + ); + + let closer = Arc::new(bus); + tokio::task::spawn_blocking(move || closer.close()) + .await + .expect("close"); + } + + #[tokio::test] + async fn completed_event_persists_exact_cost_and_stop_reason() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("usage.db"); + let bus = SqliteEventBus::open(&db, unbounded()).expect("open ledger"); + bus.emit(&received_event("req-cost", "laptop")); + bus.emit(&completed_event("req-cost")); + let conn = reader(&db); + wait_until("completed row carries its cost", || { + scalar( + &conn, + "SELECT COUNT(*) FROM requests \ + WHERE request_id = 'req-cost' AND cost_total IS NOT NULL", + ) == 1 + }) + .await; + + // Costs are stored as exact decimal text and must round-trip. + let cost_input_text = text_column( + &conn, + "SELECT cost_input FROM requests WHERE request_id = 'req-cost'", + ) + .expect("cost_input stored"); + assert_eq!( + Decimal::from_str_exact(&cost_input_text).expect("cost parses back"), + Decimal::from_str_exact("0.000123456789012345678901234").expect("original"), + "cost must round-trip exactly through its stored text" + ); + assert_eq!( + text_column( + &conn, + "SELECT cost_total FROM requests WHERE request_id = 'req-cost'" + ) + .as_deref(), + Some("2.250123956789012345678901234") + ); + assert_eq!( + text_column( + &conn, + "SELECT stop_reason FROM requests WHERE request_id = 'req-cost'" + ) + .as_deref(), + Some("unknown:weird-finish"), + "unknown stop reasons keep their provider value" + ); + assert_eq!( + scalar( + &conn, + "SELECT latency_ms FROM requests WHERE request_id = 'req-cost'" + ), + 4321 + ); + + let closer = Arc::new(bus); + tokio::task::spawn_blocking(move || closer.close()) + .await + .expect("close"); + } + + #[tokio::test] + async fn failed_event_upserts_and_preserves_known_fields() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("usage.db"); + let bus = SqliteEventBus::open(&db, unbounded()).expect("open ledger"); + + // A failure with no provider must not erase the provider stored by + // the request event. + bus.emit(&received_event("req-3", "laptop")); + bus.emit(&failed_event("req-3", None)); + // A failure for a request that was never received still gets a row. + bus.emit(&failed_event("orphan-1", Some("prov-b"))); + + let conn = reader(&db); + wait_until("both failure rows exist", || { + scalar( + &conn, + "SELECT COUNT(*) FROM requests \ + WHERE request_id IN ('req-3', 'orphan-1') AND error_kind IS NOT NULL", + ) == 2 + }) + .await; + + assert_eq!( + text_column( + &conn, + "SELECT provider FROM requests WHERE request_id = 'req-3'" + ) + .as_deref(), + Some("prov-a"), + "a None provider on failure must not NULL out the stored provider" + ); + assert_eq!( + text_column( + &conn, + "SELECT error_kind FROM requests WHERE request_id = 'req-3'" + ) + .as_deref(), + Some("UpstreamError") + ); + assert_eq!( + scalar( + &conn, + "SELECT http_status FROM requests WHERE request_id = 'req-3'" + ), + 502 + ); + + let orphan_received: Option = conn + .query_row( + "SELECT received_at FROM requests WHERE request_id = 'orphan-1'", + [], + |row| row.get(0), + ) + .expect("orphan row"); + assert_eq!( + orphan_received, None, + "a failure-only row has no received timestamp to report" + ); + assert_eq!( + scalar( + &conn, + "SELECT completed_at FROM requests WHERE request_id = 'orphan-1'" + ), + COMPLETED_MS + ); + assert_eq!( + text_column( + &conn, + "SELECT provider FROM requests WHERE request_id = 'orphan-1'" + ) + .as_deref(), + Some("prov-b") + ); + // The sanitized message field is not a ledger column by design; the + // error kind and status are the persisted failure facts. + let has_message_column: i64 = scalar( + &conn, + "SELECT COUNT(*) FROM pragma_table_info('requests') WHERE name = 'message'", + ); + assert_eq!(has_message_column, 0); + + let closer = Arc::new(bus); + tokio::task::spawn_blocking(move || closer.close()) + .await + .expect("close"); + } + + #[test] + fn emit_counts_drops_when_channel_is_full_and_never_blocks() { + // No writer task: the channel fills and stays full, so every emit + // past capacity must return immediately and be counted. + let bus = SqliteEventBus::unconnected(8); + for i in 0..20 { + bus.emit(&received_event(&format!("req-{i}"), "laptop")); + } + assert_eq!(bus.dropped_events(), 12, "20 emits into capacity 8 drop 12"); + } + + #[tokio::test] + async fn every_emitted_event_is_persisted_or_counted_never_lost() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("usage.db"); + let bus = Arc::new( + SqliteEventBus::open_with_capacity(&db, unbounded(), 32).expect("open ledger"), + ); + for i in 0..300 { + bus.emit(&received_event(&format!("r{i}"), "laptop")); + } + let closer = Arc::clone(&bus); + tokio::task::spawn_blocking(move || closer.close()) + .await + .expect("close"); + + let conn = reader(&db); + let persisted = scalar(&conn, "SELECT COUNT(*) FROM requests"); + assert_eq!( + persisted + i64::try_from(bus.dropped_events()).expect("u64 fits i64"), + 300, + "every event is either persisted or counted as dropped — never lost silently" + ); + } + + #[tokio::test] + async fn close_flushes_everything_queued_before_returning() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("usage.db"); + let bus = Arc::new(SqliteEventBus::open(&db, unbounded()).expect("open ledger")); + for i in 0..25 { + bus.emit(&received_event(&format!("f{i}"), "laptop")); + } + let closer = Arc::clone(&bus); + tokio::task::spawn_blocking(move || closer.close()) + .await + .expect("close"); + + let conn = reader(&db); + assert_eq!( + scalar(&conn, "SELECT COUNT(*) FROM requests"), + 25, + "close() returns only after the queue is fully drained" + ); + assert_eq!(bus.dropped_events(), 0); + } + + #[tokio::test] + async fn retention_pass_prunes_by_age_then_row_cap() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("usage.db"); + let bus = SqliteEventBus::open(&db, unbounded()).expect("open ledger"); + for i in 0..5 { + bus.emit(&received_event(&format!("p{i}"), "laptop")); + } + let closer = Arc::new(bus); + tokio::task::spawn_blocking(move || closer.close()) + .await + .expect("close"); + + let mut conn = Connection::open(&db).expect("open ledger"); + wait_until("all five rows visible", || { + scalar(&conn, "SELECT COUNT(*) FROM requests") == 5 + }) + .await; + // Two ancient rows, three recent (just under an hour old). + let now_ms = unix_millis(OffsetDateTime::now_utc()); + for (id, ts) in [ + ("p0", 1000_i64), + ("p1", 2000), + ("p2", now_ms - 60_000), + ("p3", now_ms - 50_000), + ("p4", now_ms - 40_000), + ] { + conn.execute( + "UPDATE requests SET received_at = ?1 WHERE request_id = ?2", + rusqlite::params![ts, id], + ) + .expect("backdate row"); + } + + let deleted = run_retention_pass( + &mut conn, + &RetentionPolicy { + max_age: Some(Duration::from_secs(3600)), + max_rows: None, + prune_every: Duration::from_secs(3600), + }, + ) + .expect("age pass"); + assert_eq!(deleted, 2, "the two ancient rows are pruned by age"); + assert_eq!(scalar(&conn, "SELECT COUNT(*) FROM requests"), 3); + + let deleted = run_retention_pass( + &mut conn, + &RetentionPolicy { + max_age: None, + max_rows: Some(1), + prune_every: Duration::from_secs(3600), + }, + ) + .expect("row-cap pass"); + assert_eq!(deleted, 2, "the two oldest of the remaining rows go"); + assert_eq!( + text_column(&conn, "SELECT request_id FROM requests"), + Some("p4".to_owned()), + "the newest row survives the cap" + ); + + // An unbounded policy prunes nothing, even when called directly. + let deleted = run_retention_pass(&mut conn, &unbounded()).expect("no-op pass"); + assert_eq!(deleted, 0); + } + + #[tokio::test] + async fn retention_timer_prunes_old_rows_without_manual_invocation() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("usage.db"); + let _bus = SqliteEventBus::open( + &db, + RetentionPolicy { + max_age: Some(Duration::from_secs(3600)), + max_rows: None, + prune_every: Duration::from_secs(1), + }, + ) + .expect("open ledger"); + // The row must carry a current timestamp: an ancient one would be + // deleted by the timer's immediate first tick before this test can + // observe it. The 1h age bound lets it survive until the test + // backdates it below. + let ProxyEvent::RequestReceived(mut fresh) = received_event("stale", "laptop") else { + panic!("builder returns RequestReceived"); + }; + fresh.timestamp = OffsetDateTime::now_utc(); + _bus.emit(&ProxyEvent::RequestReceived(fresh)); + + let conn = reader(&db); + wait_until("row appears", || { + scalar( + &conn, + "SELECT COUNT(*) FROM requests WHERE request_id = 'stale'", + ) == 1 + }) + .await; + conn.execute( + "UPDATE requests SET received_at = 1000 WHERE request_id = 'stale'", + [], + ) + .expect("backdate row"); + // The writer's own timer (1s cadence) must delete the stale row with + // no manual pass. + wait_until("timer prunes the stale row", || { + scalar( + &conn, + "SELECT COUNT(*) FROM requests WHERE request_id = 'stale'", + ) == 0 + }) + .await; + } +} diff --git a/crates/llm-proxy-storage/tests/secret_abuse.rs b/crates/llm-proxy-storage/tests/secret_abuse.rs new file mode 100644 index 0000000..bf3e299 --- /dev/null +++ b/crates/llm-proxy-storage/tests/secret_abuse.rs @@ -0,0 +1,152 @@ +//! Secret-abuse family (PRX-06) for the SQLite ledger: only the documented +//! columns ever reach the raw file bytes. The `ResponseFailed.message` field +//! can carry sanitized upstream text (truncated error bodies), so the abuse +//! property is that whatever an event's message holds, it is never persisted +//! — a raw-bytes assertion, stronger than checking the column list, because +//! it scans every byte of the database and its WAL sidecars, exactly what an +//! attacker with filesystem access would read. + +#[cfg(test)] +mod tests { + use std::path::{Path, PathBuf}; + use std::time::Duration; + + use llm_proxy_protocol::core::{Cost, ModelRef, StopReason, Usage, UsageProvenance}; + use llm_proxy_storage::{ + EventBus, ProxyEvent, RequestReceived, ResponseCompleted, ResponseFailed, SqliteEventBus, + }; + use rust_decimal::Decimal; + use time::OffsetDateTime; + + /// Distinctive markers, unique to this family so a hit can never be a + /// coincidence with unrelated file content. + const CALLER_MARKER: &str = "prx6-storage-caller-77aa01"; + const DIGEST_MARKER: &str = "prx6-storage-digest-88bb02"; + const STOP_LABEL_CONTROL: &str = "end_turn"; + /// Marker placed in `ResponseFailed.message` — event-carried text that + /// must never be persisted. + const MESSAGE_MARKER: &str = "prx6-storage-message-99cc03"; + + fn received(ts: OffsetDateTime) -> ProxyEvent { + ProxyEvent::RequestReceived(RequestReceived { + request_id: "prx6-row-1".to_owned(), + timestamp: ts, + caller: CALLER_MARKER.to_owned(), + provider: "prov".to_owned(), + route_kind: "messages".to_owned(), + client_protocol: "anthropic".to_owned(), + model: ModelRef { + requested: "claude-3".to_owned(), + upstream: None, + }, + streaming: false, + body_hash: DIGEST_MARKER.to_owned(), + }) + } + + fn completed(ts: OffsetDateTime) -> ProxyEvent { + ProxyEvent::ResponseCompleted(ResponseCompleted { + request_id: "prx6-row-1".to_owned(), + timestamp: ts, + provider: "prov".to_owned(), + upstream_message_id: Some("msg_upstream".to_owned()), + model: ModelRef { + requested: "claude-3".to_owned(), + upstream: None, + }, + usage: Usage { + input_tokens: 11, + output_tokens: 7, + reasoning_tokens: None, + cache_read_input_tokens: None, + cache_creation_input_tokens: None, + provenance: UsageProvenance::ProviderReported, + }, + cost: Some(Cost { + input: Decimal::from(11), + output: Decimal::ZERO, + cache_creation: Decimal::ZERO, + cache_read: Decimal::ZERO, + reasoning: Decimal::ZERO, + total: Decimal::from(11), + }), + stop_reason: StopReason::EndTurn, + latency_ms: 42, + }) + } + + fn failed(request_id: &str, ts: OffsetDateTime) -> ProxyEvent { + ProxyEvent::ResponseFailed(ResponseFailed { + request_id: request_id.to_owned(), + timestamp: ts, + provider: Some("prov".to_owned()), + model: None, + error_kind: "UpstreamError".to_owned(), + // The sanitized message stands in for any event-carried text: + // upstream error fragments, echoes, anything. None of it may be + // persisted, even for a failure whose request row exists. + message: format!("upstream failed after echoing {MESSAGE_MARKER}"), + http_status: 502, + latency_ms: 5, + }) + } + + /// Concatenated raw bytes of the ledger database and any WAL/SHM + /// sidecars still on disk. + fn ledger_bytes(db: &Path) -> Vec { + let mut all = Vec::new(); + for suffix in ["", "-wal", "-shm"] { + let sidecar = PathBuf::from(format!("{}{suffix}", db.display())); + if let Ok(bytes) = std::fs::read(&sidecar) { + all.extend_from_slice(&bytes); + } + } + all + } + + #[tokio::test] + async fn raw_ledger_bytes_persist_columns_but_never_event_messages() { + let dir = tempfile::tempdir().expect("tempdir"); + let db = dir.path().join("abuse.db"); + let bus = SqliteEventBus::open( + &db, + llm_proxy_storage::RetentionPolicy { + max_age: None, + max_rows: None, + prune_every: Duration::from_secs(3600), + }, + ) + .expect("ledger opens"); + + let t0 = OffsetDateTime::now_utc(); + // A completed request whose terminal failure text carries the marker + // (a paired row that gets both INSERT and UPDATE), plus a + // failure-only orphan row whose message also carries it. + bus.emit(&received(t0)); + bus.emit(&failed("prx6-row-1", t0)); + bus.emit(&completed(t0)); + bus.emit(&failed("prx6-orphan", t0)); + + let closer = std::sync::Arc::new(bus); + tokio::task::spawn_blocking(move || closer.close()) + .await + .expect("close joins"); + + let bytes = ledger_bytes(&db); + let text = String::from_utf8_lossy(&bytes); + // Positive controls: the scan finds values that ARE persisted — the + // caller key id, the body digest, and a stop-reason label — proving + // the negative assertion below reads the real rows. + for control in [CALLER_MARKER, DIGEST_MARKER, STOP_LABEL_CONTROL] { + assert!( + text.contains(control), + "control: persisted value {control} must be findable in the raw bytes" + ); + } + assert!( + !text.contains(MESSAGE_MARKER), + "event message text (upstream echoes, error fragments) leaked into the \ + raw ledger file bytes" + ); + } +} diff --git a/dist-workspace.toml b/dist-workspace.toml new file mode 100644 index 0000000..8e2573d --- /dev/null +++ b/dist-workspace.toml @@ -0,0 +1,19 @@ +[workspace] +members = ["cargo:."] + +# Config for 'dist' (https://github.com/axodotdev/cargo-dist) +[dist] +# The preferred dist version to use in CI (Cargo.toml SemVer syntax) +cargo-dist-version = "0.32.0" +# CI backends to support +ci = "github" +# The installers to generate for each app +installers = ["shell"] +# Target platforms to build apps for (Rust target-triple syntax) +targets = ["x86_64-unknown-linux-musl", "aarch64-unknown-linux-musl"] +# Path that installers should place binaries in +install-path = "CARGO_HOME" +# Where to host releases +hosting = "github" +# Whether to install an updater program +install-updater = false diff --git a/providers/opencode-go.toml.example b/providers/opencode-go.toml.example index f670984..6a6e0ce 100644 --- a/providers/opencode-go.toml.example +++ b/providers/opencode-go.toml.example @@ -45,3 +45,11 @@ messages = "anthropic" # cache_creation = "0.000003" # cache_read = "0.00000125" # reasoning = "0.000010" + +# Route this provider's upstream traffic through a forward proxy. +# Allowed schemes: socks5h:// (hostnames resolved at the proxy), socks5://, +# http://, https://. The URL may carry userinfo (proxy authentication). +# Absent [provider.egress] = direct connection, as before. +# +# [provider.egress] +# proxy = "socks5h://127.0.0.1:1080" diff --git a/providers/opencode-zen.toml.example b/providers/opencode-zen.toml.example index ce481eb..ec4d66a 100644 --- a/providers/opencode-zen.toml.example +++ b/providers/opencode-zen.toml.example @@ -53,3 +53,11 @@ messages = "anthropic" # cache_creation = "0.000003" # cache_read = "0.00000125" # reasoning = "0.000010" + +# Route this provider's upstream traffic through a forward proxy. +# Allowed schemes: socks5h:// (hostnames resolved at the proxy), socks5://, +# http://, https://. The URL may carry userinfo (proxy authentication). +# Absent [provider.egress] = direct connection, as before. +# +# [provider.egress] +# proxy = "socks5h://127.0.0.1:1080"