From 806293396c1e1afbb7f7933a4ebbe1fbee51c887 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Tue, 18 Aug 2026 04:48:55 +0800 Subject: [PATCH 01/74] Add scoped temporary credential authentication --- .github/workflows/docker-publish.yml | 14 + .github/workflows/release.yml | 4 + CHANGELOG.md | 7 + Cargo.lock | 9 +- Cargo.toml | 4 +- DOCKER_README.md | 37 +- README.md | 20 +- README.zh-CN.md | 20 +- docker/docker-compose.yml | 7 +- docker/pb-mapper.dockerfile | 3 +- docs/authentication-v2.md | 297 +++ docs/authentication-v2.zh-CN.md | 232 ++ docs/pb-mapper-intro.zh-CN.md | 9 +- docs/user-guide.md | 80 +- docs/user-guide.zh-CN.md | 73 +- examples/pb_local_server.rs | 2 + scripts/install-server-gitee.sh | 20 +- scripts/install-server-github.sh | 20 +- scripts/release/entrypoint/pb-mapper.sh | 19 +- services/pb-mapper-server.service | 4 +- services/readme.md | 21 +- skills/pb-mapper-connect-deploy/SKILL.md | 33 +- skills/pb-mapper-server-deploy/SKILL.md | 27 +- src/bin/pb-mapper.rs | 746 ++++++- src/common/auth.rs | 2548 ++++++++++++++++++++++ src/common/checksum.rs | 191 +- src/common/error.rs | 2 + src/common/message/command.rs | 184 +- src/common/message/mod.rs | 11 +- src/common/message/secure.rs | 1063 +++++++++ src/common/mod.rs | 1 + src/local/client/mod.rs | 95 +- src/local/client/status.rs | 56 +- src/local/client/stream.rs | 34 +- src/local/server/mod.rs | 97 +- src/local/server/stream.rs | 40 +- src/pb_server/admin.rs | 220 ++ src/pb_server/client.rs | 126 +- src/pb_server/error.rs | 4 + src/pb_server/mod.rs | 792 ++++++- src/pb_server/server.rs | 61 +- src/pb_server/status.rs | 9 +- src/utils/codec.rs | 2 +- tests/regression.rs | 482 +++- tests/test_delay.rs | 23 +- ui/lib/l10n/app_en.arb | 10 +- ui/lib/l10n/app_zh.arb | 10 +- ui/lib/src/views/configuration_view.dart | 18 +- ui/lib/src/views/setup_wizard_view.dart | 2 +- ui/native/pb_mapper_ffi/src/state.rs | 12 +- ui/test/widget_test.dart | 7 +- 51 files changed, 7434 insertions(+), 374 deletions(-) create mode 100644 docs/authentication-v2.md create mode 100644 docs/authentication-v2.zh-CN.md create mode 100644 src/common/auth.rs create mode 100644 src/common/message/secure.rs create mode 100644 src/pb_server/admin.rs diff --git a/.github/workflows/docker-publish.yml b/.github/workflows/docker-publish.yml index b5d2ffd..cf43090 100644 --- a/.github/workflows/docker-publish.yml +++ b/.github/workflows/docker-publish.yml @@ -36,6 +36,20 @@ jobs: - name: Set up Docker Buildx uses: docker/setup-buildx-action@v3 + - name: Verify release tag matches Cargo version + shell: bash + run: | + set -euo pipefail + package_version="$(awk ' + /^\[workspace.package\]$/ { in_workspace = 1; next } + in_workspace && /^version = / { + gsub(/"/, "", $3) + print $3 + exit + } + ' Cargo.toml)" + test "${GITHUB_REF_NAME}" = "v${package_version}" + - name: Log in to Docker Hub uses: docker/login-action@v3 with: diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 518b522..7c12da9 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -77,3 +77,7 @@ jobs: LICENSE README.md README.zh-CN.md + docs/authentication-v2.md + docs/authentication-v2.zh-CN.md + docs/user-guide.md + docs/user-guide.zh-CN.md diff --git a/CHANGELOG.md b/CHANGELOG.md index 14aa463..652db8a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,13 @@ All notable changes to this project will be documented in this file. +## [0.4.0] - 2026-08-18 +- Added a sole administrator credential plus renewable, expiring, and immediately revocable `pbmt1_` temporary credentials with fixed-slot O(1) lookup and isolated per-key service namespaces. +- Added single-flight protocol-v2 authentication with directional AES-256-GCM keys, monotonic frame counters, authenticated routing metadata, replay detection, stable structured errors, and optional legacy framing during migration. +- Added encrypted snapshot/WAL authentication state, lifecycle audit records, hierarchical timing-wheel expiry, hard closure of revoked live connections, safe-mode recovery, root-key rotation, and explicit auth-state reset. +- Extended the unified CLI with temporary-key lifecycle, service/connection inventory, auth status, protocol policy, root rotation, namespace targeting, and human/JSON/NDJSON output. +- Replaced insecure default-key fallback with first-start random administrator-key generation, retained machine-derived keys only for explicit compatibility, and updated Flutter, installers, systemd, Docker, release metadata, and bilingual documentation. + ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. - Consolidated release archives into one cross-platform binary artifact per target and updated Docker, installers, systemd templates, build scripts, deployment skills, and documentation to use it. diff --git a/Cargo.lock b/Cargo.lock index d0d1077..b829fa7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -84,6 +84,12 @@ dependencies = [ "syn", ] +[[package]] +name = "base64" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" + [[package]] name = "better_mimalloc_rs" version = "0.1.2" @@ -771,8 +777,9 @@ dependencies = [ [[package]] name = "pb-mapper" -version = "0.3.0" +version = "0.4.0" dependencies = [ + "base64", "better_mimalloc_rs", "bytes", "clap", diff --git a/Cargo.toml b/Cargo.toml index 68584db..a127ba0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -26,6 +26,7 @@ ring.workspace = true once_cell.workspace = true uni-stream.workspace = true kanal.workspace = true +base64.workspace = true [dev-dependencies] dotenvy = "0.15.7" @@ -38,7 +39,7 @@ members = ["ui/native/pb_mapper_ffi"] exclude = ["deps/uni-stream", "deps/kanal"] [workspace.package] -version = "0.3.0" +version = "0.4.0" authors = ["L_B__"] edition = "2021" @@ -64,6 +65,7 @@ bytes = "1.11" trust-dns-resolver = { version = "0.23.2" } ring = "0.17.14" once_cell = "1.20.2" +base64 = "0.22.1" uni-stream = { git = "https://github.com/acking-you/uni-stream.git", branch = "master" } kanal = { git = "https://github.com/acking-you/kanal.git", branch = "dev/pb-mapper" } diff --git a/DOCKER_README.md b/DOCKER_README.md index 7cc7bdb..da26103 100644 --- a/DOCKER_README.md +++ b/DOCKER_README.md @@ -13,8 +13,8 @@ docker run -d \ --name pb-mapper \ -p 7666:7666 \ -e PB_MAPPER_PORT=7666 \ - -e USE_MACHINE_MSG_HEADER_KEY=true \ -e RUST_LOG=error \ + -v pb-mapper-auth:/var/lib/pb-mapper/auth \ ackingliu/pb-mapper:latest-x86_64_musl ``` @@ -28,11 +28,16 @@ services: environment: - PB_MAPPER_PORT=7666 - USE_IPV6=false - - USE_MACHINE_MSG_HEADER_KEY=true + - USE_MACHINE_MSG_HEADER_KEY=false - RUST_LOG=error + volumes: + - pb-mapper-auth:/var/lib/pb-mapper/auth ports: - "7666:7666" restart: unless-stopped + +volumes: + pb-mapper-auth: ``` Save as `docker-compose.yml` and run: @@ -46,10 +51,12 @@ docker-compose up -d |----------|---------|-------------| | `PB_MAPPER_PORT` | `7666` | **Required** - Port for the pb-mapper server to listen on | | `USE_IPV6` | `false` | Enable IPv6 support (`true`/`false`) | -| `USE_MACHINE_MSG_HEADER_KEY` | `true` | Derive `MSG_HEADER_KEY` from hostname + MAC and persist to `/var/lib/pb-mapper-server/msg_header_key` | +| `MSG_HEADER_KEY` | unset | Optional 32-character administrator key used only to initialize a new persistent auth volume | +| `USE_MACHINE_MSG_HEADER_KEY` | `false` | Legacy compatibility: derive the administrator key from hostname + MAC | +| `PB_MAPPER_AUTH_STATE_DIR` | `/var/lib/pb-mapper/auth` | Persistent encrypted authentication state | | `RUST_LOG` | `error` | Logging level (`error`, `warn`, `info`, `debug`, `trace`) | -⚠️ **Important**: `PB_MAPPER_PORT` must be set or the container will exit with an error. +⚠️ **Important**: `PB_MAPPER_PORT` must be set and `/var/lib/pb-mapper/auth` must be persistent. The first start creates a random administrator key at `admin.key`; losing the volume changes the root credential and loses temporary-key state. ## 📋 Ubuntu Deployment Guide @@ -85,11 +92,16 @@ services: environment: PB_MAPPER_PORT: 7666 USE_IPV6: false - USE_MACHINE_MSG_HEADER_KEY: true + USE_MACHINE_MSG_HEADER_KEY: false RUST_LOG: error + volumes: + - pb-mapper-auth:/var/lib/pb-mapper/auth ports: - "7666:7666" restart: unless-stopped + +volumes: + pb-mapper-auth: EOF ``` @@ -112,6 +124,9 @@ docker-compose ps # View logs docker-compose logs -f pb-mapper + +# Read the administrator key on the Docker host +docker exec pb-mapper cat /var/lib/pb-mapper/auth/admin.key ``` ### Step 5: Verify Installation @@ -154,18 +169,18 @@ For other architectures, you can build the image yourself using the provided Doc |-----|-------------| | `latest-x86_64_musl` | Latest stable x86_64 build (recommended) | | `latest-aarch64_musl` | Latest stable ARM64 build | -| `v0.3.0-x86_64_musl` | Tagged-release x86_64 build | -| `v0.3.0-aarch64_musl` | Tagged-release ARM64 build | -| `0.3.0-x86_64_musl` | Semver x86_64 alias | -| `0.3.0-aarch64_musl` | Semver ARM64 alias | +| `v0.4.0-x86_64_musl` | Tagged-release x86_64 build | +| `v0.4.0-aarch64_musl` | Tagged-release ARM64 build | +| `0.4.0-x86_64_musl` | Semver x86_64 alias | +| `0.4.0-aarch64_musl` | Semver ARM64 alias | **Recommendation**: Use `latest-x86_64_musl` for x86_64 systems or `latest-aarch64_musl` for ARM64 systems for best compatibility. ## 🛡️ Security Considerations - **Firewall**: Only expose port 7666 to trusted networks -- **Encryption**: Use the encryption features in client/server tools -- **Access Control**: Implement service key management strategy +- **Authentication**: Keep `admin.key` on the relay and distribute expiring `pbmt1_` temporary credentials to workloads +- **Forwarded payload encryption**: Use `register --codec` when the inner application protocol is plaintext - **Updates**: Regularly update to the latest version for security patches ## 📊 Monitoring and Logs diff --git a/README.md b/README.md index 3745f7c..24b76fb 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,8 @@ ## Highlights - **One binary, one public port** — `pb-mapper` provides every runtime role, while a service-key registry replaces per-service port planning. +- **Scoped temporary credentials** — the administrator key can issue renewable, expiring `pbmt1_` credentials; each credential gets an isolated service namespace and can only inspect, register, and connect inside it. +- **Authenticated protocol v2** — directional AES-256-GCM control frames authenticate in the first request without adding a handshake round trip. New clients use v2; the server can temporarily allow legacy clients during migration. - **Optional encryption** — AES-256-GCM (via `ring`) on forwarded traffic, enabled with `--codec` at registration. - **Proven in production** — on real workloads (e.g. a Palworld UDP server), latency matches frp with a directly exposed port. @@ -41,18 +43,21 @@ With an AI coding agent (Claude Code, Cursor, Kiro), the built-in skills handle ### Alternative — one-liner install script -If the remote host can reach GitHub directly, this installs the unified `pb-mapper` binary and runs its `server` command as a systemd service on Linux (x86_64, musl) — port `7666`, `--use-machine-msg-header-key` on, key stored at `/var/lib/pb-mapper-server/msg_header_key`. +If the remote host can reach GitHub directly, this installs the unified `pb-mapper` binary and runs its `server` command as a systemd service on Linux (x86_64, musl). The relay listens on port `7666` and creates a random administrator key at `/var/lib/pb-mapper/auth/admin.key` on first start. ```bash curl -fsSL https://raw.githubusercontent.com/acking-you/pb-mapper/master/scripts/install-server-github.sh | bash ``` -After install, load the same key before running `pb-mapper register` or `pb-mapper connect`: +Use the administrator key only for management and issue a temporary credential for a workload: ```bash -export MSG_HEADER_KEY="$(cat /var/lib/pb-mapper-server/msg_header_key)" +export MSG_HEADER_KEY="$(sudo cat /var/lib/pb-mapper/auth/admin.key)" +pb-mapper admin --server :7666 key issue --ttl 24h --label home-web ``` +Copy the printed `pbmt1_...` credential to the register and connect machines as their `MSG_HEADER_KEY`. They may use the same service name without colliding with another temporary credential's namespace. + ## Architecture ![pb-mapper architecture](docs/assets/architecture-flow.svg) @@ -80,10 +85,13 @@ Your web server runs on `localhost:8080` at home. # 1. on the public server — start the central router pb-mapper server --port 7666 -# 2. at home — register the web server under key 'web' +# 2. issue a temporary credential, then export it on both endpoint machines +export MSG_HEADER_KEY='' + +# 3. at home — register the web server under key 'web' pb-mapper register tcp --server :7666 --key web --addr 127.0.0.1:8080 -# 3. at the coffee shop — subscribe and expose it locally +# 4. at the coffee shop — subscribe and expose it locally pb-mapper connect tcp --server :7666 --key web --addr 127.0.0.1:3000 ``` @@ -97,6 +105,7 @@ Open `http://localhost:3000` in the coffee-shop browser — traffic flows throug | `pb-mapper register tcp\|udp` | Registers a local TCP/UDP service with the server | | `pb-mapper connect tcp\|udp` | Subscribes to a registered service and exposes a local port | | `pb-mapper status keys\|remote-id` | Queries the central router | +| `pb-mapper admin ...` | Issues/renews/revokes credentials and inspects auth, services, and connections | | **Flutter UI** (`ui/`) | GUI for server, register, connect, and status workflows | ## Developer view @@ -107,6 +116,7 @@ Open `http://localhost:3000` in the coffee-shop browser — traffic flows throug ## Documentation - User guide (build / run / use): [`docs/user-guide.md`](docs/user-guide.md) +- Authentication and protocol v2: [`docs/authentication-v2.md`](docs/authentication-v2.md) - Docker server guide: [`DOCKER_README.md`](DOCKER_README.md) - 中文文档: [`README.zh-CN.md`](README.zh-CN.md), [`docs/user-guide.zh-CN.md`](docs/user-guide.zh-CN.md) diff --git a/README.zh-CN.md b/README.zh-CN.md index 58fe928..e4d37dd 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -27,6 +27,8 @@ ## 亮点 - **单二进制、单公网端口**：统一的 `pb-mapper` 命令覆盖所有运行角色，服务 key 注册表取代逐个服务规划端口。 +- **临时凭据与命名空间隔离**：管理员密钥可签发可续期、自动过期的 `pbmt1_` 凭据；每把临时凭据只能查看、注册和连接自己的命名空间。 +- **V2 首帧鉴权**：控制帧使用按方向派生的 AES-256-GCM 密钥，在第一个请求内完成鉴权，不增加额外握手往返；新客户端固定使用 V2，服务端可在迁移期兼容旧协议。 - **可选加密**：转发流量可启用 AES-256-GCM（基于 `ring`），注册服务时用 `--codec` 开启。 - **生产可用**：真实负载下（例如 Palworld UDP 服务器），延迟与 frp 直暴端口相当。 @@ -41,18 +43,21 @@ ### 备选方式：一键安装脚本 -远程主机能直连 GitHub 时，一条命令即可在 Linux（x86_64，musl）上安装统一的 `pb-mapper` 二进制，并以 `server` 子命令启动 systemd 服务：端口 `7666`，启用 `--use-machine-msg-header-key`，key 落盘在 `/var/lib/pb-mapper-server/msg_header_key`。 +远程主机能直连 GitHub 时，一条命令即可在 Linux（x86_64，musl）上安装统一的 `pb-mapper` 二进制，并以 `server` 子命令启动 systemd 服务。中继监听 `7666`，首次启动时会在 `/var/lib/pb-mapper/auth/admin.key` 创建随机管理员密钥。 ```bash curl -fsSL https://raw.githubusercontent.com/acking-you/pb-mapper/master/scripts/install-server-github.sh | bash ``` -安装完成后，在运行 `pb-mapper register` 或 `pb-mapper connect` 前加载同一把 key： +管理员密钥只用于管理；先为一项业务签发临时凭据： ```bash -export MSG_HEADER_KEY="$(cat /var/lib/pb-mapper-server/msg_header_key)" +export MSG_HEADER_KEY="$(sudo cat /var/lib/pb-mapper/auth/admin.key)" +pb-mapper admin --server :7666 key issue --ttl 24h --label home-web ``` +把输出的 `pbmt1_...` 凭据作为 register 与 connect 机器上的 `MSG_HEADER_KEY`。不同临时凭据即使使用相同的 service name，也不会相互冲突。 + ## 架构 ![pb-mapper architecture](docs/assets/architecture-flow.svg) @@ -80,10 +85,13 @@ register 与 connect 工作流也可以通过 Flutter UI 操作。 # 1. 公网服务器：启动中心路由 pb-mapper server --port 7666 -# 2. 家中机器：以 key 'web' 注册服务 +# 2. 签发临时凭据，并在两端机器导入 +export MSG_HEADER_KEY='' + +# 3. 家中机器：以 key 'web' 注册服务 pb-mapper register tcp --server :7666 --key web --addr 127.0.0.1:8080 -# 3. 咖啡店机器：订阅并在本地暴露 +# 4. 咖啡店机器：订阅并在本地暴露 pb-mapper connect tcp --server :7666 --key web --addr 127.0.0.1:3000 ``` @@ -97,6 +105,7 @@ pb-mapper connect tcp --server :7666 --key web --addr 127.0.0.1:3000 | `pb-mapper register tcp\|udp` | 将本地 TCP/UDP 服务注册到服务器 | | `pb-mapper connect tcp\|udp` | 订阅已注册的服务并在本地暴露端口 | | `pb-mapper status keys\|remote-id` | 查询中心路由状态 | +| `pb-mapper admin ...` | 签发/续期/吊销临时凭据并查看认证、服务与连接状态 | | **Flutter UI**（`ui/`） | server、register、connect、status 的图形化界面 | ## 开发者视角 @@ -107,6 +116,7 @@ pb-mapper connect tcp --server :7666 --key web --addr 127.0.0.1:3000 ## 文档 - 使用手册（编译/运行/使用）：[`docs/user-guide.zh-CN.md`](docs/user-guide.zh-CN.md) +- 认证与 V2 协议：[`docs/authentication-v2.zh-CN.md`](docs/authentication-v2.zh-CN.md) - Docker 服务器指南：[`DOCKER_README.md`](DOCKER_README.md) - English docs: [`README.md`](README.md)、[`docs/user-guide.md`](docs/user-guide.md) diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml index c172cb2..d25b769 100644 --- a/docker/docker-compose.yml +++ b/docker/docker-compose.yml @@ -5,8 +5,13 @@ services: image: ackingliu/pb-mapper:x86_64_musl environment: PB_MAPPER_PORT: 7666 - USE_MACHINE_MSG_HEADER_KEY: true + USE_MACHINE_MSG_HEADER_KEY: false RUST_LOG: error + volumes: + - pb-mapper-auth:/var/lib/pb-mapper/auth ports: - "7666:7666" restart: unless-stopped + +volumes: + pb-mapper-auth: diff --git a/docker/pb-mapper.dockerfile b/docker/pb-mapper.dockerfile index 4e9f1cd..0cf2f6d 100644 --- a/docker/pb-mapper.dockerfile +++ b/docker/pb-mapper.dockerfile @@ -9,7 +9,8 @@ RUN chmod +x ./pb-mapper ./pb-mapper.sh ENV PB_MAPPER_PORT=7666 ENV USE_IPV6=false -ENV USE_MACHINE_MSG_HEADER_KEY=true +ENV USE_MACHINE_MSG_HEADER_KEY=false +VOLUME ["/var/lib/pb-mapper/auth"] EXPOSE $PB_MAPPER_PORT ENTRYPOINT [ "./pb-mapper.sh" ] diff --git a/docs/authentication-v2.md b/docs/authentication-v2.md new file mode 100644 index 0000000..8ce2c83 --- /dev/null +++ b/docs/authentication-v2.md @@ -0,0 +1,297 @@ +# Authentication and Protocol V2 + +## Background and goals + +pb-mapper uses one public relay port for registration, subscription, status, and +administration. Version 0.4 keeps that transport model and adds two credential +levels without adding a TLS-style handshake: + +- one 32-byte administrator key owns the relay; +- renewable `pbmt1_` temporary credentials can inspect, register, and connect + only inside their own namespace; +- the first protocol-v2 frame authenticates and carries the request in one TCP + flight; +- revocation, expiry, and root-key rotation close affected live control and data + connections. + +TLS is still appropriate when endpoint identity, certificate trust, or traffic +analysis resistance is required. Protocol v2 protects pb-mapper frames with a +pre-shared credential; it is not a replacement for a public-key PKI. + +## Model and terminology + +| Term | Meaning | +| --- | --- | +| Administrator key | The sole 32-byte root credential. It can manage keys and inspect every namespace. | +| Temporary credential | A printable `pbmt1_...` value containing a key ID and derived 32-byte secret. | +| Key ID | A 64-bit `generation:u32 | slot:u32` identifier used for direct slot lookup. | +| Namespace | `0` for the administrator, otherwise the temporary key ID. | +| Service name | A user-facing name within one namespace. Equal names in different temporary namespaces do not collide. | +| Credential lease | The cancellation object shared by connections authenticated with one credential. | + +The administrator key is never copied into a temporary credential. A temporary +secret is derived with HKDF-SHA256 from the administrator key, the persistent +server instance ID, and the key ID. The fixed slot table stores lifecycle +metadata and a weak lease reference, not the temporary secret. + +## End-to-end architecture + +```mermaid +sequenceDiagram + participant C as register/connect/admin CLI + participant R as pb-mapper relay + participant A as auth actor + participant M as connection manager + + C->>R: V2 prefix + encrypted first request + R->>R: derive directional keys and authenticate frame + R->>A: validate key ID, generation, state, expiry + A-->>R: namespace + weak credential lease + alt administrator operation + R->>A: issue/renew/revoke/status + A-->>R: durable result after WAL fsync + else register/connect/status + R->>M: namespace-scoped service operation + M-->>R: scoped result or stable error + end + R-->>C: encrypted response on the same connection +``` + +Only the long-lived registration control connection and each independently +opened subscribe/data connection carry a V2 first frame. The relay does not add +an extra authentication exchange when the register process opens a data TCP +connection for a request. + +## Protocol-v2 framing + +### Initial prefix + +Every new client writes this 32-byte clear-text routing prefix: + +| Bytes | Field | +| ---: | --- | +| 4 | Magic `PBM2` | +| 1 | Version `2` | +| 1 | Flags, currently `0` | +| 2 | Reserved, currently `0` | +| 8 | Big-endian key ID; `0` means administrator | +| 16 | Random connection salt | + +The prefix is not secret. It is authenticated as associated data on every +encrypted frame. Unsupported flags, versions, and non-zero reserved bytes are +rejected before request dispatch. + +### Directional frame keys + +HKDF-SHA256 uses the connection salt as salt and the credential's 32-byte secret +as input key material. Two independent outputs are expanded with +`pb-mapper-v2-c2s` and `pb-mapper-v2-s2c`. This prevents nonce reuse across +directions even though both directions begin with counter zero. + +### Encrypted frames + +Each frame is encoded as: + +| Bytes | Field | +| ---: | --- | +| 8 | Big-endian monotonically increasing counter | +| 4 | Big-endian ciphertext length, including the 16-byte GCM tag | +| variable | AES-256-GCM ciphertext and tag | + +The 96-bit AES-GCM nonce is four zero bytes followed by the 64-bit counter. AAD +contains the complete initial prefix, one direction byte, the counter, and the +ciphertext length. Counter mismatch, authentication failure, oversized frames, +and counter exhaustion close the connection. + +The first client request uses client-to-server counter `0`. The first response +uses server-to-client counter `0`. Later control frames continue from counter +`1` through one stateful reader/writer per direction. + +### Replay resistance + +The relay fingerprints `(key_id, connection_salt)` and checks two rotating +1 MiB Bloom filters covering the current and previous 60-second windows. A +probable duplicate returns the stable retryable error +`connection_salt_replayed`; one-shot administrator CLI operations retry once +with a fresh salt. + +## Credential lifecycle + +### Issuance and renewal + +`key issue` allocates a free fixed-table slot, increments its generation, +derives the secret, appends an encrypted WAL mutation, calls `fsync`, and only +then exposes the credential. `key renew` keeps the same credential and key ID, +updates its absolute expiry, and inserts a new versioned timing-wheel entry. +Stale wheel entries are ignored. + +```bash +export MSG_HEADER_KEY="$(sudo cat /var/lib/pb-mapper/auth/admin.key)" + +pb-mapper admin --server relay.example.com:7666 \ + key issue --ttl 24h --label home-web + +pb-mapper admin --server relay.example.com:7666 \ + key renew 4294967296 --ttl 7d +``` + +Temporary TTLs are at least 10 seconds and at most 30 days by default. The +server maximum is configurable. + +### Expiry, revocation, and garbage collection + +A four-level hierarchical timing wheel owns the strong `Arc` for every active +temporary lease. Foreground authentication state contains only `Weak` +references. Expiry or explicit revocation cancels the lease, immediately +causing authenticated control and data tasks to drop their TCP streams. +Tombstones remain briefly for stable diagnostics, then become reusable slots. + +```bash +pb-mapper admin --server relay.example.com:7666 key revoke 4294967296 +pb-mapper admin --server relay.example.com:7666 key gc +``` + +### Root rotation and state reset + +Root rotation writes an empty snapshot encrypted with the new key, appends the +audit record, persists `admin.key`, and then switches live state. It invalidates +all temporary credentials and closes connections authenticated with the old +administrator or temporary keys. The CLI stages the candidate key before the +request and verifies the new key with an authenticated status call. + +An explicit auth-state reset also invalidates all temporary credentials. It +rotates the server instance ID so credentials from a corrupted or lost slot +table cannot become valid again if a key ID is later reused. + +## Namespace authorization + +Temporary credentials may perform `register`, `connect`, and `status` only in +their own namespace. They cannot issue keys, reveal credentials, inspect other +namespaces, alter protocol policy, reset auth state, or rotate the root key. + +The administrator defaults to namespace `0`. It may inspect or connect to a +temporary namespace with `--namespace `. Registering into another +namespace additionally requires `--force` to avoid accidental ownership +confusion. + +Temporary-key service names are 1-128 ASCII bytes from +`[A-Za-z0-9._:-]`. The relay enforces per-namespace caps for services, +registration connections, active streams, and new-stream rate. + +| Approach | Memory and lookup | Revocation | Namespace isolation | Wire cost | +| --- | --- | --- | --- | --- | +| Stateless signed token | Minimal server state | Requires a deny list | Token claim based | One request | +| Stateful hash map | Proportional allocations and hashing | Direct | Direct | One request | +| Fixed slots plus derived secrets | Fixed hot memory and O(1) lookup | Direct slot cancellation | Key ID is namespace | One request | + +The fixed-slot design deliberately accepts bounded server state to make early +revocation and hard connection closure deterministic. + +## Persistence and safe mode + +The default state directory is `/var/lib/pb-mapper/auth`: + +| File | Purpose | +| --- | --- | +| `admin.key` | Root credential, mode `0600` | +| `server-instance-id` | 16-byte persistent derivation identity | +| `auth.snapshot` | AES-256-GCM encrypted compact slot state | +| `auth.wal` | Length-prefixed, individually encrypted mutations and audit records | + +The directory is mode `0700`. Mutating operations acknowledge only after the +WAL record is synced. The actor compacts state every five minutes with an atomic +snapshot replacement and WAL truncation. + +Invalid authentication-state headers, failed integrity checks, truncated WAL +records, schema mismatch, and failed compaction place temporary authentication +in safe mode. Administrator authentication stays available for inspection and +explicit reset; temporary authentication fails closed. + +## Administration and output contracts + +```bash +pb-mapper admin --server relay.example.com:7666 status +pb-mapper admin --server relay.example.com:7666 key list --page-size 100 +pb-mapper admin --server relay.example.com:7666 key show 4294967296 +pb-mapper admin --server relay.example.com:7666 key reveal 4294967296 +pb-mapper admin --server relay.example.com:7666 service list --key-id 4294967296 +pb-mapper admin --server relay.example.com:7666 connection list --all +pb-mapper admin --server relay.example.com:7666 legacy-protocol set deny +pb-mapper admin --server relay.example.com:7666 auth-state reset --confirm +pb-mapper admin --server relay.example.com:7666 root-key rotate +``` + +`--output human|json|ndjson` controls rendering. Pages default to 100 and are +capped at 1000. `--all` follows pages and emits one NDJSON object per item so a +large inventory does not need to be buffered by the CLI. + +Stable structured errors contain `code`, `message`, `retryable`, and +`server_time`. Authentication failure logs include stage, key ID, peer, and +reason but never credential material. Repeated failures are emitted five times +per minute per `(peer IP, key ID, reason)`, followed by a suppression summary. + +## Migration and compatibility + +New clients always emit protocol v2. A v0.4 server defaults to accepting legacy +framing so older clients can be upgraded without an outage. Operators can view +legacy connection counters, upgrade all clients, and then set the policy to +`deny`. Upgrade the relay before any client because v0.3 relays do not understand +the v2 first-frame magic. + +Fresh servers generate a random administrator key. Both the relay and install +scripts preserve an existing `/var/lib/pb-mapper-server/msg_header_key` by +copying it to the new `admin.key` path when no new key or environment credential +is configured. `--use-machine-msg-header-key` remains available only as an +explicit legacy compatibility option. + +Docker deployments must persist `/var/lib/pb-mapper/auth`; otherwise a recreated +container generates a different root key and cannot decrypt previous auth +state. + +## Operations playbook + +### Temporary credential rejected after renewal + +1. Run `pb-mapper admin status` and confirm `safe_mode=false`. +2. Run `key show ` and verify the key is `active` and its absolute expiry. +3. Check structured logs for `temporary_key_generation_mismatch`, + `temporary_key_expired`, or `protocol_v2_decrypt_failed`. +4. If the credential text was lost or copied incorrectly, run `key reveal ` + and replace the endpoint configuration. Renewal does not change the value. + +### Relay starts in safe mode + +1. Preserve the entire auth directory for diagnosis. +2. Confirm `admin.key`, `server-instance-id`, snapshot, and WAL belong to the + same server instance and were not partially restored. +3. Use `pb-mapper admin status`; administrator access remains available. +4. If recovery is impossible, run `auth-state reset --confirm`, then issue new + temporary credentials. Reset rotates the server instance ID and closes old + workloads. + +### Legacy clients stop connecting + +1. Check `admin status` for the current legacy policy and active legacy count. +2. If policy is `deny`, upgrade the client or temporarily set it to `allow`. +3. New clients should log protocol `V2`; a continuing legacy count identifies + an old binary or integration that still needs replacement. + +## Code index + +- Credential format and process configuration: `src/common/checksum.rs` +- Slot table, timing wheel, encrypted WAL, and lifecycle actor: + `src/common/auth.rs` +- V2 framing, key derivation, counters, and replay filter: + `src/common/message/secure.rs` +- Namespace dispatch and relay resource limits: `src/pb_server/mod.rs` +- Administrator request execution: `src/pb_server/admin.rs` +- Unified command surface: `src/bin/pb-mapper.rs` + +## Summary + +Version 0.4 retains pb-mapper's one-port, long-lived-control-connection model +while separating root administration from scoped workload access. Temporary +keys are renewable but revocable, namespace collisions are eliminated, and +authentication remains part of the first request. The explicit operational +boundary is unchanged: protocol v2 is symmetric pre-shared-key security, while +TLS remains the layer for certificate-based endpoint identity. diff --git a/docs/authentication-v2.zh-CN.md b/docs/authentication-v2.zh-CN.md new file mode 100644 index 0000000..9faa921 --- /dev/null +++ b/docs/authentication-v2.zh-CN.md @@ -0,0 +1,232 @@ +# 认证体系与 V2 协议 + +## 背景与目标 + +pb-mapper 的注册、订阅、状态与管理流量共用一个公网端口。0.4 版本不改变这一 +连接模型，也不引入类似 TLS 的额外握手，而是在原有对称密钥体系内增加两级权限： + +- 一把 32 字节管理员密钥拥有中继的全部权限； +- 可续期、可提前吊销、自动过期的 `pbmt1_` 临时凭据只能查看、注册和连接自己的 + 命名空间； +- V2 第一个加密帧同时完成鉴权与请求传输，不增加一次网络往返； +- 过期、吊销与根密钥轮换会主动关闭受影响的控制连接和数据连接。 + +V2 是预共享密钥协议。当系统还需要证书身份、公开信任链或对流量分析的额外防护时， +TLS 仍然有独立价值，V2 不替代公钥 PKI。 + +## 核心模型 + +| 概念 | 含义 | +| --- | --- | +| 管理员密钥 | 唯一的 32 字节根凭据，可管理密钥并查看全部命名空间。 | +| 临时凭据 | `pbmt1_...` 字符串，携带 key ID 与派生后的 32 字节 secret。 | +| Key ID | 64 位 `generation:u32 | slot:u32`，用于直接定位固定槽位。 | +| 命名空间 | 管理员默认为 `0`；临时凭据的命名空间就是自己的 key ID。 | +| 凭据租约 | 同一凭据认证出的连接共同观察的取消对象。 | + +临时 secret 由管理员密钥、持久化 server instance ID 与 key ID 通过 +HKDF-SHA256 派生。服务端固定槽位只存生命周期元数据与弱引用，不存临时 secret。 + +## 端到端数据流 + +```mermaid +sequenceDiagram + participant C as register/connect/admin CLI + participant R as pb-mapper relay + participant A as auth actor + participant M as connection manager + + C->>R: V2 前缀 + 加密后的首个请求 + R->>R: 派生双向密钥并验证加密帧 + R->>A: 校验 key ID、generation、状态与过期时间 + A-->>R: 命名空间 + 凭据租约弱引用 + alt 管理操作 + R->>A: 签发/续期/吊销/状态查询 + A-->>R: WAL fsync 后的结果 + else 业务操作 + R->>M: 命名空间内的注册/订阅/查询 + M-->>R: 结果或稳定错误码 + end + R-->>C: 同一连接上的加密响应 +``` + +register 进程长期保持的控制连接只在建立时认证一次。后续每次业务请求拉起的新数据 +TCP 连接仍有自己的 V2 首帧，但不会在其上再做多轮鉴权交换。 + +## V2 帧结构 + +### 首帧前缀 + +新客户端先写入 32 字节明文路由前缀： + +| 字节数 | 字段 | +| ---: | --- | +| 4 | Magic `PBM2` | +| 1 | 版本 `2` | +| 1 | Flags，当前必须为 `0` | +| 2 | Reserved，当前必须为 `0` | +| 8 | 大端 key ID；`0` 表示管理员 | +| 16 | 随机 connection salt | + +前缀不承担保密作用，但会作为每个加密帧的 AAD 被完整认证。未知版本、flags 或 +reserved 值会在请求分发前被拒绝。 + +### 双向密钥与计数器 + +HKDF-SHA256 使用 connection salt 作为 salt，凭据的 32 字节 secret 作为 IKM， +分别以 `pb-mapper-v2-c2s` 与 `pb-mapper-v2-s2c` 派生两个 AES-256-GCM 密钥。 +因此两个方向都从计数器 0 开始，也不会重复使用同一密钥与 nonce 组合。 + +每个加密帧由 8 字节大端计数器、4 字节密文长度、密文与 16 字节 GCM tag 组成。 +96 位 nonce 是四个零字节加 64 位计数器。AAD 包含完整首帧前缀、方向字节、计数器 +与密文长度。计数器不连续、认证失败、帧过大或计数器耗尽都会关闭连接。 + +首个请求使用 C2S counter 0，首个响应使用 S2C counter 0；后续控制消息由同一组 +有状态 reader/writer 从 counter 1 继续。 + +### 重放检测 + +服务端对 `(key_id, connection_salt)` 做指纹，并使用两个轮换的 1 MiB Bloom filter +覆盖当前与上一个 60 秒窗口。疑似重复会返回可重试错误 +`connection_salt_replayed`；一次性 admin CLI 会自动换 salt 重试一次。 + +## 临时凭据生命周期 + +### 签发与续期 + +签发会寻找空槽位、递增 generation、派生 secret、写入加密 WAL 并 `fsync`，随后才 +把凭据返回给管理员。续期不换 key ID 与凭据文本，只更新绝对过期时间并把新版本任务 +放入时间轮，旧任务到期时因版本不匹配而被忽略。 + +```bash +export MSG_HEADER_KEY="$(sudo cat /var/lib/pb-mapper/auth/admin.key)" +pb-mapper admin --server relay.example.com:7666 \ + key issue --ttl 24h --label home-web +pb-mapper admin --server relay.example.com:7666 \ + key renew 4294967296 --ttl 7d +``` + +默认最短 TTL 为 10 秒，最长为 30 天，服务端可调整最大值。 + +### 到期、吊销与 GC + +四层层级时间轮持有每个活动凭据租约的强 `Arc`；前台认证状态只持有 `Weak`。 +到期或管理员吊销会取消租约，相关控制任务和数据转发任务立即释放 TCP 连接。短暂保留 +tombstone 以给出稳定错误后，槽位可以复用。显式 `key gc` 可立即清理非活动槽位。 + +### 根密钥轮换与状态重置 + +根密钥轮换先用新密钥写空 snapshot、追加审计、持久化 `admin.key`，再切换内存状态。 +它会使全部临时凭据失效，并关闭旧管理员或临时凭据建立的连接。CLI 在发请求前保存 +候选 key，完成后再用新 key 执行一次 `admin status` 验证。 + +`auth-state reset --confirm` 同样会清空临时凭据，并轮换 server instance ID。这样即使 +原槽位表损坏或丢失，旧凭据也不会因为未来复用了相同 key ID 而重新有效。 + +## 命名空间与权限边界 + +临时凭据只能在自己的命名空间执行 `register`、`connect` 与 `status`。它不能签发或 +查看其他 key、进入其他命名空间、修改 legacy 策略、重置状态或轮换管理员密钥。 + +管理员默认使用命名空间 0；通过 `--namespace ` 可以查看或连接临时命名空间。 +管理员要在临时命名空间内注册服务时还必须显式使用 `--force`，避免误把业务服务挂到 +错误租户。 + +临时凭据的 service name 限制为 1 到 128 个 ASCII 字节，字符集为 +`[A-Za-z0-9._:-]`。服务端分别按命名空间限制 service 数、单 service 注册连接数、 +活动 stream 数与新建 stream 速率。 + +| 方案 | 内存与查询 | 提前吊销 | 命名空间隔离 | 网络成本 | +| --- | --- | --- | --- | --- | +| 无状态签名 token | 服务端状态少 | 仍需 deny list | 依赖 token claim | 一个请求 | +| 通用 HashMap | 动态分配与哈希 | 直接删除 | 直接 | 一个请求 | +| 固定槽位 + 派生 secret | 固定热内存、O(1) 查找 | 直接取消槽位租约 | key ID 即 namespace | 一个请求 | + +当前方案明确接受有上限的服务端状态，以换取确定的提前吊销与活动连接硬关闭。 + +## 持久化与安全模式 + +默认目录 `/var/lib/pb-mapper/auth` 权限为 `0700`： + +| 文件 | 用途 | +| --- | --- | +| `admin.key` | 根凭据，权限 `0600` | +| `server-instance-id` | 16 字节持久派生身份 | +| `auth.snapshot` | AES-256-GCM 加密的紧凑槽位快照 | +| `auth.wal` | 带长度前缀、逐条加密的 mutation 与 audit | + +变更只有在 WAL 同步成功后才对外确认。后台 actor 每五分钟原子替换 snapshot 并截断 +WAL。无效文件头、完整性验证失败、WAL 截断、schema 不匹配或 compact 失败都会进入 +safe mode：临时凭据全部 fail closed，管理员仍可查看状态并执行显式 reset。 + +## 管理命令与输出 + +```bash +pb-mapper admin --server relay.example.com:7666 status +pb-mapper admin --server relay.example.com:7666 key list --page-size 100 +pb-mapper admin --server relay.example.com:7666 key show 4294967296 +pb-mapper admin --server relay.example.com:7666 key reveal 4294967296 +pb-mapper admin --server relay.example.com:7666 service list --key-id 4294967296 +pb-mapper admin --server relay.example.com:7666 connection list --all +pb-mapper admin --server relay.example.com:7666 legacy-protocol set deny +pb-mapper admin --server relay.example.com:7666 auth-state reset --confirm +pb-mapper admin --server relay.example.com:7666 root-key rotate +``` + +`--output human|json|ndjson` 控制展示格式。默认每页 100，最大 1000；`--all` 自动翻页 +并逐行输出 NDJSON，避免 CLI 一次缓存完整列表。稳定错误结构包含 `code`、`message`、 +`retryable` 与 `server_time`。 + +日志记录 auth stage、key ID、peer 与 reason，但不记录凭据。相同 +`(peer IP, key ID, reason)` 每分钟最多直接输出 5 次，下一窗口汇总被抑制的数量。 + +## 迁移与兼容性 + +新客户端固定发送 V2。0.4 服务端默认暂时接受旧帧，方便滚动升级；确认 +`active_legacy_connections` 归零后，可执行 `legacy-protocol set deny`。必须先升级中继、 +再升级客户端，因为 0.3 中继无法识别 V2 首帧 magic。 + +新安装会随机生成管理员密钥。中继自身与安装脚本在未配置新 key 或环境变量时，如果 +发现旧的 `/var/lib/pb-mapper-server/msg_header_key`，会将其复制到新路径，保留现有 +业务连接。`--use-machine-msg-header-key` 只作为明确的兼容选项继续存在。 + +Docker 必须持久化 `/var/lib/pb-mapper/auth`；否则重建容器会产生新管理员密钥，并且 +无法读取先前认证状态。 + +## 运维排障 + +### 续期后临时凭据仍被拒绝 + +1. 执行 `admin status`，确认 `safe_mode=false`。 +2. 执行 `key show `，确认状态为 active 并核对绝对过期时间。 +3. 从结构化日志区分 generation 不匹配、已过期与 V2 解密失败。 +4. 如果只是凭据文本复制错误，执行 `key reveal ` 重新配置；续期本身不会换凭据。 + +### 服务端进入 safe mode + +1. 先完整保留 auth 目录用于诊断。 +2. 确认 key、instance ID、snapshot 与 WAL 是否来自同一份服务器状态。 +3. 使用管理员 key 查询 `admin status`；管理员通道仍然可用。 +4. 无法恢复时执行 `auth-state reset --confirm`，再重新签发业务凭据。该操作会轮换 + instance ID，并断开旧业务。 + +### 禁用 legacy 后仍有旧客户端 + +1. 在 `admin status` 查看 legacy policy、当前连接数与最后连接时间。 +2. 如果业务尚未升级，可短暂改回 `allow`，但应尽快升级客户端。 +3. 新客户端的服务端日志应显示协议 `V2`；仍增长的 legacy 计数可以定位旧 binary。 + +## 代码索引 + +- 凭据格式与进程配置：`src/common/checksum.rs` +- 固定槽位、时间轮、加密 WAL 与生命周期 actor：`src/common/auth.rs` +- V2 帧、双向派生、计数器与 replay filter：`src/common/message/secure.rs` +- 命名空间分发与资源限制：`src/pb_server/mod.rs` +- 管理请求执行：`src/pb_server/admin.rs` +- 统一 CLI：`src/bin/pb-mapper.rs` + +## 总结 + +0.4 在保持单端口与长控制连接模型的同时，把根管理权限与业务访问权限拆开。临时凭据 +可续期、可吊销、按命名空间隔离，鉴权仍然包含在第一个业务请求内。需要明确保留的 +边界是：V2 解决预共享密钥下的帧认证与权限控制，证书身份仍属于 TLS 层。 diff --git a/docs/pb-mapper-intro.zh-CN.md b/docs/pb-mapper-intro.zh-CN.md index e8c6606..0497b0f 100644 --- a/docs/pb-mapper-intro.zh-CN.md +++ b/docs/pb-mapper-intro.zh-CN.md @@ -87,11 +87,11 @@ VPS 能直连 GitHub 的话： curl -fsSL https://raw.githubusercontent.com/acking-you/pb-mapper/master/scripts/install-server-github.sh | bash ``` -装完默认监听 `7666`，开启 `--use-machine-msg-header-key`，密钥写到 `/var/lib/pb-mapper-server/msg_header_key`。client 侧 `export MSG_HEADER_KEY="$(cat /var/lib/pb-mapper-server/msg_header_key)"` 就能对上。 +装完默认监听 `7666`，首次启动会在 `/var/lib/pb-mapper/auth/admin.key` 生成随机管理员密钥。管理员用它签发带过期时间的 `pbmt1_` 临时凭据，再把临时凭据交给 register/connect 两端；不同临时凭据的同名 service 不会撞名。 ### 方式三：手动跑 CLI 或者用 Flutter UI -三个二进制，名字就是功能： +同一个 `pb-mapper` 二进制通过子命令切换功能： - `pb-mapper server`：公网中继 - `pb-mapper register`：跑在本地服务那一侧，把 `127.0.0.1:xxx` 注册成一个 service key @@ -102,6 +102,11 @@ curl -fsSL https://raw.githubusercontent.com/acking-you/pb-mapper/master/scripts ```bash # VPS 上 pb-mapper server --port 7666 +export MSG_HEADER_KEY="$(sudo cat /var/lib/pb-mapper/auth/admin.key)" +pb-mapper admin --server 127.0.0.1:7666 key issue --ttl 24h --label web + +# 家里和咖啡店都先导入上一步输出的 pbmt1_ 临时凭据 +export MSG_HEADER_KEY='' # 家里 pb-mapper register tcp --server :7666 --key web --addr 127.0.0.1:8080 diff --git a/docs/user-guide.md b/docs/user-guide.md index 67ae576..1165378 100644 --- a/docs/user-guide.md +++ b/docs/user-guide.md @@ -4,7 +4,7 @@ ## Overview -pb-mapper exposes local TCP/UDP services through a public relay using a service key. One `pb-mapper` binary provides the `server`, `register`, `connect`, and `status` commands, alongside an optional Flutter GUI. +pb-mapper exposes local TCP/UDP services through a public relay using a service key. One `pb-mapper` binary provides the `server`, `register`, `connect`, `status`, and `admin` commands, alongside an optional Flutter GUI. ## How it works @@ -112,31 +112,49 @@ Optional flags: - `--ipv6`: enable IPv6 listening - `--keep-alive`: enable TCP keep-alive -- `--use-machine-msg-header-key`: derive `MSG_HEADER_KEY` from current machine hostname + MAC, - and write it to `/var/lib/pb-mapper-server/msg_header_key` +- `--auth-state-dir`: authentication state directory (default `/var/lib/pb-mapper/auth`) +- `--max-temporary-keys`: fixed temporary-key slot capacity (default `65536`) +- `--max-temporary-key-ttl`: maximum issued TTL (default `30d`) +- `--legacy-protocol allow|deny`: initial legacy-client policy +- `--use-machine-msg-header-key`: explicit legacy compatibility mode -### Machine-derived `MSG_HEADER_KEY` (optional) +### Administrator and temporary credentials -When you want each deployed server to use a host-specific key (instead of the built-in default), -start server with: +On first start, the relay creates a random administrator key at +`/var/lib/pb-mapper/auth/admin.key`. There is no built-in default credential. +Keep the administrator key on the relay host and use it to issue a temporary +credential for a workload: ```bash -pb-mapper server --port 7666 --use-machine-msg-header-key +export MSG_HEADER_KEY="$(sudo cat /var/lib/pb-mapper/auth/admin.key)" +pb-mapper admin --server "your-server:7666" \ + key issue --ttl 24h --label my-service ``` -This will: +Export the printed `pbmt1_...` credential on both the register and connect +machines. The temporary key can see and use only its own namespace: + +```bash +export MSG_HEADER_KEY='pbmt1_...' +pb-mapper register tcp --server "your-server:7666" --key "my-service" --addr "127.0.0.1:8080" +``` -- derive a stable 32-byte key from hostname + MAC addresses -- set server process `MSG_HEADER_KEY` automatically -- persist the key to `/var/lib/pb-mapper-server/msg_header_key` +Renewing a key preserves the credential text. Revocation or expiry immediately +closes its active control and data connections. See +[`authentication-v2.md`](authentication-v2.md) for the full lifecycle, +namespace model, protocol framing, and migration procedure. -Then use the same key for the `register` and `connect` commands: +The machine-derived option remains available for an existing deployment, but +it is not recommended for new installations: ```bash -export MSG_HEADER_KEY="$(cat /var/lib/pb-mapper-server/msg_header_key)" -pb-mapper register tcp --server "your-server:7666" --key "my-service" --addr "127.0.0.1:8080" +pb-mapper server --port 7666 --use-machine-msg-header-key ``` +On upgrade, if no new administrator key or `MSG_HEADER_KEY` is present, the +relay automatically imports `/var/lib/pb-mapper-server/msg_header_key` so +legacy clients keep working. + ### 2) Register a local service Register a TCP service: @@ -177,6 +195,30 @@ pb-mapper status remote-id --server "your-server:7666" pb-mapper status keys --server "your-server:7666" ``` +An administrator can explicitly inspect or connect to a temporary namespace: + +```bash +pb-mapper status keys --server "your-server:7666" --namespace 4294967296 +pb-mapper connect tcp --server "your-server:7666" --namespace 4294967296 \ + --key "my-service" --addr "127.0.0.1:9090" +``` + +Registering as administrator inside a temporary namespace additionally requires +`--force`. + +### Administrator commands + +```bash +pb-mapper admin --server "your-server:7666" status +pb-mapper admin --server "your-server:7666" key list +pb-mapper admin --server "your-server:7666" key reveal 4294967296 +pb-mapper admin --server "your-server:7666" service list --all +pb-mapper admin --server "your-server:7666" connection list --all +``` + +Use `--output json` for one JSON document or `--output ndjson` for streaming +automation. Page size defaults to 100 and is capped at 1000. + ## Run (GUI) The Flutter UI can start the server, register services, and connect clients through a graphical workflow. Start it from `ui/`: @@ -189,6 +231,16 @@ flutter run ## Environment variables - `PB_MAPPER_SERVER`: default server address for the CLI +- `MSG_HEADER_KEY`: 32-character administrator key or a `pbmt1_` temporary credential +- `PB_MAPPER_AUTH_STATE_DIR`: relay auth-state directory, default `/var/lib/pb-mapper/auth` +- `PB_MAPPER_AUTH_MAX_TEMP_KEYS`: fixed temporary-key capacity, default `65536` +- `PB_MAPPER_AUTH_MAX_TEMP_TTL_SECS`: maximum temporary-key TTL, default 30 days +- `PB_MAPPER_LEGACY_PROTOCOL`: `allow` or `deny`, default `allow` +- `PB_MAPPER_MAX_SERVICES_PER_NAMESPACE`: service names per namespace, default `256` +- `PB_MAPPER_MAX_REGISTER_CONNECTIONS_PER_SERVICE`: control connections per service, default `16` +- `PB_MAPPER_MAX_STREAMS_PER_NAMESPACE`: active streams per namespace, default `1024` +- `PB_MAPPER_NEW_STREAMS_PER_SECOND`: sustained new-stream rate per namespace, default `100` +- `PB_MAPPER_NEW_STREAMS_BURST`: new-stream burst per namespace, default `200` - `PB_MAPPER_KEEP_ALIVE`: enable TCP keep-alive (set to `ON`) - `PB_MAPPER_LOG_FORMAT`: tracing output format, one of `pretty` (default), `compact`, or `json` - `PB_MAPPER_CONTROL_IO_TIMEOUT`: close stalled control-plane handshakes after this duration, default `30s` diff --git a/docs/user-guide.zh-CN.md b/docs/user-guide.zh-CN.md index 4278365..aea60e2 100644 --- a/docs/user-guide.zh-CN.md +++ b/docs/user-guide.zh-CN.md @@ -4,7 +4,7 @@ ## 概览 -pb-mapper 通过“服务 key”将本地 TCP/UDP 服务暴露到公网中继。统一的 `pb-mapper` 二进制提供 `server`、`register`、`connect`、`status` 四类命令，并保留可选的 Flutter GUI。 +pb-mapper 通过“服务 key”将本地 TCP/UDP 服务暴露到公网中继。统一的 `pb-mapper` 二进制提供 `server`、`register`、`connect`、`status`、`admin` 五类命令，并保留可选的 Flutter GUI。 ## 运转机制 @@ -112,30 +112,44 @@ pb-mapper server --port 7666 - `--ipv6`：开启 IPv6 监听 - `--keep-alive`：开启 TCP keep-alive -- `--use-machine-msg-header-key`：基于当前机器 hostname + MAC 派生 `MSG_HEADER_KEY`， - 并写入 `/var/lib/pb-mapper-server/msg_header_key` +- `--auth-state-dir`：认证状态目录，默认 `/var/lib/pb-mapper/auth` +- `--max-temporary-keys`：临时 key 固定槽位容量，默认 `65536` +- `--max-temporary-key-ttl`：临时 key 最大 TTL，默认 `30d` +- `--legacy-protocol allow|deny`：旧协议初始接入策略 +- `--use-machine-msg-header-key`：明确启用旧版机器派生 key 兼容模式 -### 基于机器信息派生 `MSG_HEADER_KEY`（可选） +### 管理员密钥与临时凭据 -如果你希望每台部署机器都使用各自唯一的 key（而不是内置默认 key），可以这样启动服务端： +中继首次启动时会在 `/var/lib/pb-mapper/auth/admin.key` 生成随机管理员密钥，系统不再 +提供内置默认 key。管理员密钥留在中继机器上，用它为业务签发临时凭据： ```bash -pb-mapper server --port 7666 --use-machine-msg-header-key +export MSG_HEADER_KEY="$(sudo cat /var/lib/pb-mapper/auth/admin.key)" +pb-mapper admin --server "your-server:7666" \ + key issue --ttl 24h --label my-service ``` -该参数会完成： +把输出的 `pbmt1_...` 凭据导入 register 与 connect 两端。临时 key 只能看到并操作 +自己的命名空间： + +```bash +export MSG_HEADER_KEY='pbmt1_...' +pb-mapper register tcp --server "your-server:7666" --key "my-service" --addr "127.0.0.1:8080" +``` -- 基于 hostname + MAC 地址派生稳定的 32 字节 key -- 自动设置当前服务端进程的 `MSG_HEADER_KEY` -- 将 key 持久化到 `/var/lib/pb-mapper-server/msg_header_key` +续期不会改变凭据文本；吊销或到期会立即关闭对应的控制连接与数据连接。完整生命周期、 +命名空间、V2 帧格式和迁移步骤见 +[`authentication-v2.zh-CN.md`](authentication-v2.zh-CN.md)。 -随后在 `register` 与 `connect` 命令中使用同一 key： +已有部署仍可显式启用机器派生 key，但不建议新安装继续使用： ```bash -export MSG_HEADER_KEY="$(cat /var/lib/pb-mapper-server/msg_header_key)" -pb-mapper register tcp --server "your-server:7666" --key "my-service" --addr "127.0.0.1:8080" +pb-mapper server --port 7666 --use-machine-msg-header-key ``` +升级时，如果新管理员密钥和 `MSG_HEADER_KEY` 都不存在，中继会自动导入 +`/var/lib/pb-mapper-server/msg_header_key`，旧客户端无需立刻换 key。 + ### 2）注册本地服务 注册 TCP 服务： @@ -176,6 +190,29 @@ pb-mapper status remote-id --server "your-server:7666" pb-mapper status keys --server "your-server:7666" ``` +管理员可明确查看或连接某个临时命名空间： + +```bash +pb-mapper status keys --server "your-server:7666" --namespace 4294967296 +pb-mapper connect tcp --server "your-server:7666" --namespace 4294967296 \ + --key "my-service" --addr "127.0.0.1:9090" +``` + +管理员要在临时命名空间注册服务，还必须增加 `--force`。 + +### 管理命令 + +```bash +pb-mapper admin --server "your-server:7666" status +pb-mapper admin --server "your-server:7666" key list +pb-mapper admin --server "your-server:7666" key reveal 4294967296 +pb-mapper admin --server "your-server:7666" service list --all +pb-mapper admin --server "your-server:7666" connection list --all +``` + +自动化场景可使用 `--output json` 输出单个 JSON 文档，或用 `--output ndjson` 流式输出。 +默认每页 100 条，最大 1000 条。 + ## 运行（GUI） Flutter UI 可用于启动服务器、注册服务与建立连接。启动方式： @@ -188,6 +225,16 @@ flutter run ## 环境变量 - `PB_MAPPER_SERVER`：CLI 默认服务器地址 +- `MSG_HEADER_KEY`：32 字符管理员密钥或 `pbmt1_` 临时凭据 +- `PB_MAPPER_AUTH_STATE_DIR`：中继认证状态目录，默认 `/var/lib/pb-mapper/auth` +- `PB_MAPPER_AUTH_MAX_TEMP_KEYS`：临时 key 固定容量，默认 `65536` +- `PB_MAPPER_AUTH_MAX_TEMP_TTL_SECS`：临时 key 最大 TTL，默认 30 天 +- `PB_MAPPER_LEGACY_PROTOCOL`：`allow` 或 `deny`，默认 `allow` +- `PB_MAPPER_MAX_SERVICES_PER_NAMESPACE`：每命名空间 service 数，默认 `256` +- `PB_MAPPER_MAX_REGISTER_CONNECTIONS_PER_SERVICE`：每 service 控制连接数，默认 `16` +- `PB_MAPPER_MAX_STREAMS_PER_NAMESPACE`：每命名空间活动 stream 数，默认 `1024` +- `PB_MAPPER_NEW_STREAMS_PER_SECOND`：每命名空间持续新建 stream 速率，默认 `100` +- `PB_MAPPER_NEW_STREAMS_BURST`：每命名空间新建 stream 突发量，默认 `200` - `PB_MAPPER_KEEP_ALIVE`：启用 TCP keep-alive（设置为 `ON`） - `PB_MAPPER_LOG_FORMAT`：tracing 输出格式，可选 `pretty`（默认）、`compact` 或 `json` - `PB_MAPPER_CONTROL_IO_TIMEOUT`：控制面握手卡住后的关闭时间，默认 `30s` diff --git a/examples/pb_local_server.rs b/examples/pb_local_server.rs index 1ae9e97..a282436 100644 --- a/examples/pb_local_server.rs +++ b/examples/pb_local_server.rs @@ -13,6 +13,8 @@ async fn main() { need_codec: false, is_datagram: false, keep_alive: false, + namespace: None, + force_namespace: false, }, ) .await; diff --git a/scripts/install-server-gitee.sh b/scripts/install-server-gitee.sh index 51cd1c8..abc6506 100755 --- a/scripts/install-server-gitee.sh +++ b/scripts/install-server-gitee.sh @@ -2,7 +2,7 @@ set -euo pipefail # Configuration -VERSION="${PB_MAPPER_VERSION:-0.3.0}" +VERSION="${PB_MAPPER_VERSION:-0.4.0}" ARCH="${PB_MAPPER_ARCH:-x86_64-unknown-linux-musl}" TARBALL="pb-mapper-${ARCH}.tar.gz" DOWNLOAD_URL="https://gitee.com/acking-you/pb-mapper/releases/download/v${VERSION}/${TARBALL}" @@ -10,6 +10,9 @@ INSTALL_DIR="/usr/local/bin" SERVICE_NAME="pb-mapper-server" SERVICE_PATH="/etc/systemd/system/${SERVICE_NAME}.service" PORT="${PB_MAPPER_PORT:-7666}" +AUTH_DIR="/var/lib/pb-mapper/auth" +ADMIN_KEY_PATH="${AUTH_DIR}/admin.key" +LEGACY_KEY_PATH="/var/lib/pb-mapper-server/msg_header_key" # Re-run with sudo if needed if [ "${EUID:-$(id -u)}" -ne 0 ]; then @@ -67,6 +70,14 @@ fi mkdir -p "$INSTALL_DIR" install -m 0755 "$BIN_PATH" "${INSTALL_DIR}/pb-mapper" +# Preserve the former machine-derived credential on upgrade. New installations let +# pb-mapper create a random administrator key on first start. +install -d -m 0700 "$AUTH_DIR" +if [ ! -s "$ADMIN_KEY_PATH" ] && [ -s "$LEGACY_KEY_PATH" ]; then + install -m 0600 "$LEGACY_KEY_PATH" "$ADMIN_KEY_PATH" + echo "Migrated the legacy machine-derived key into $ADMIN_KEY_PATH" +fi + # Stop and remove existing service if present if systemctl is-active --quiet "${SERVICE_NAME}.service"; then systemctl stop "${SERVICE_NAME}.service" @@ -86,8 +97,10 @@ After=network.target [Service] Type=simple -ExecStart=${INSTALL_DIR}/pb-mapper server --port ${PORT} --use-machine-msg-header-key +ExecStart=${INSTALL_DIR}/pb-mapper server --port ${PORT} Environment=RUST_LOG=info +StateDirectory=pb-mapper +StateDirectoryMode=0700 Restart=on-failure RestartSec=3 LimitNOFILE=65535 @@ -102,4 +115,5 @@ systemctl enable --now "${SERVICE_NAME}.service" echo "pb-mapper server is installed and running." echo "Service name: ${SERVICE_NAME}.service" -echo "Machine-derived key file: /var/lib/pb-mapper-server/msg_header_key" +echo "Administrator key file: /var/lib/pb-mapper/auth/admin.key" +echo "Read it locally as root and issue temporary credentials for register/connect clients." diff --git a/scripts/install-server-github.sh b/scripts/install-server-github.sh index beb7d9b..5a11466 100755 --- a/scripts/install-server-github.sh +++ b/scripts/install-server-github.sh @@ -2,7 +2,7 @@ set -euo pipefail # Configuration -VERSION="${PB_MAPPER_VERSION:-0.3.0}" +VERSION="${PB_MAPPER_VERSION:-0.4.0}" ARCH="${PB_MAPPER_ARCH:-x86_64-unknown-linux-musl}" TARBALL="pb-mapper-${ARCH}.tar.gz" DOWNLOAD_URL="https://github.com/acking-you/pb-mapper/releases/download/v${VERSION}/${TARBALL}" @@ -10,6 +10,9 @@ INSTALL_DIR="/usr/local/bin" SERVICE_NAME="pb-mapper-server" SERVICE_PATH="/etc/systemd/system/${SERVICE_NAME}.service" PORT="${PB_MAPPER_PORT:-7666}" +AUTH_DIR="/var/lib/pb-mapper/auth" +ADMIN_KEY_PATH="${AUTH_DIR}/admin.key" +LEGACY_KEY_PATH="/var/lib/pb-mapper-server/msg_header_key" # Re-run with sudo if needed if [ "${EUID:-$(id -u)}" -ne 0 ]; then @@ -67,6 +70,14 @@ fi mkdir -p "$INSTALL_DIR" install -m 0755 "$BIN_PATH" "${INSTALL_DIR}/pb-mapper" +# Preserve the former machine-derived credential on upgrade. New installations let +# pb-mapper create a random administrator key on first start. +install -d -m 0700 "$AUTH_DIR" +if [ ! -s "$ADMIN_KEY_PATH" ] && [ -s "$LEGACY_KEY_PATH" ]; then + install -m 0600 "$LEGACY_KEY_PATH" "$ADMIN_KEY_PATH" + echo "Migrated the legacy machine-derived key into $ADMIN_KEY_PATH" +fi + # Stop and remove existing service if present if systemctl is-active --quiet "${SERVICE_NAME}.service"; then systemctl stop "${SERVICE_NAME}.service" @@ -86,8 +97,10 @@ After=network.target [Service] Type=simple -ExecStart=${INSTALL_DIR}/pb-mapper server --port ${PORT} --use-machine-msg-header-key +ExecStart=${INSTALL_DIR}/pb-mapper server --port ${PORT} Environment=RUST_LOG=info +StateDirectory=pb-mapper +StateDirectoryMode=0700 Restart=on-failure RestartSec=3 LimitNOFILE=65535 @@ -102,4 +115,5 @@ systemctl enable --now "${SERVICE_NAME}.service" echo "pb-mapper server is installed and running." echo "Service name: ${SERVICE_NAME}.service" -echo "Machine-derived key file: /var/lib/pb-mapper-server/msg_header_key" +echo "Administrator key file: /var/lib/pb-mapper/auth/admin.key" +echo "Read it locally as root and issue temporary credentials for register/connect clients." diff --git a/scripts/release/entrypoint/pb-mapper.sh b/scripts/release/entrypoint/pb-mapper.sh index b750c00..54c3399 100644 --- a/scripts/release/entrypoint/pb-mapper.sh +++ b/scripts/release/entrypoint/pb-mapper.sh @@ -7,9 +7,18 @@ if [ -z "${PB_MAPPER_PORT:-}" ]; then exit 1 fi -USE_MACHINE_MSG_HEADER_KEY=${USE_MACHINE_MSG_HEADER_KEY:-true} +USE_MACHINE_MSG_HEADER_KEY=${USE_MACHINE_MSG_HEADER_KEY:-false} ARGS=(-p "$PB_MAPPER_PORT") +AUTH_DIR="${PB_MAPPER_AUTH_STATE_DIR:-/var/lib/pb-mapper/auth}" +ADMIN_KEY_PATH="$AUTH_DIR/admin.key" +LEGACY_KEY_PATH="/var/lib/pb-mapper-server/msg_header_key" + +install -d -m 0700 "$AUTH_DIR" +if [ ! -s "$ADMIN_KEY_PATH" ] && [ -s "$LEGACY_KEY_PATH" ]; then + install -m 0600 "$LEGACY_KEY_PATH" "$ADMIN_KEY_PATH" + echo "Migrated the legacy machine-derived key into $ADMIN_KEY_PATH" +fi if [ "${USE_IPV6:-false}" = "true" ]; then echo "USE_IPV6 is set to true" @@ -19,10 +28,16 @@ else fi if [ "$USE_MACHINE_MSG_HEADER_KEY" = "true" ]; then - echo "USE_MACHINE_MSG_HEADER_KEY is set to true" + echo "WARNING: USE_MACHINE_MSG_HEADER_KEY is a legacy compatibility mode" ARGS+=(--use-machine-msg-header-key) else echo "USE_MACHINE_MSG_HEADER_KEY is set to false" fi +if [ -n "${MSG_HEADER_KEY:-}" ]; then + echo "Using the configured administrator credential" +else + echo "Using or initializing $ADMIN_KEY_PATH" +fi + exec ./pb-mapper server "${ARGS[@]}" diff --git a/services/pb-mapper-server.service b/services/pb-mapper-server.service index 822a251..9f13c92 100644 --- a/services/pb-mapper-server.service +++ b/services/pb-mapper-server.service @@ -7,7 +7,9 @@ Wants=network-online.target Type=simple Environment=RUST_LOG=info EnvironmentFile=-/etc/pb-mapper/server.env -ExecStart=/usr/local/bin/pb-mapper server --port 7666 --use-machine-msg-header-key +ExecStart=/usr/local/bin/pb-mapper server --port 7666 +StateDirectory=pb-mapper +StateDirectoryMode=0700 Restart=on-failure RestartSec=3 LimitNOFILE=65535 diff --git a/services/readme.md b/services/readme.md index e4dda8e..49d6942 100644 --- a/services/readme.md +++ b/services/readme.md @@ -9,9 +9,20 @@ sudo install -m 0644 services/pb-mapper-register@.service /etc/systemd/system/ sudo install -m 0644 services/pb-mapper-connect@.service /etc/systemd/system/ ``` -The relay unit runs `pb-mapper server` directly. Override it with a systemd -drop-in if the default `7666` port or machine-derived key behavior is not -appropriate. +The relay unit runs `pb-mapper server` directly. On first start it creates a +random administrator key at `/var/lib/pb-mapper/auth/admin.key` with mode +`0600`. Keep that directory persistent. If upgrading from the old +machine-derived mode, the first v0.4 start automatically copies +`/var/lib/pb-mapper-server/msg_header_key` to the new path when neither an +administrator key file nor `MSG_HEADER_KEY` is already configured. + +Use the administrator key locally for management, then issue a scoped temporary +credential for each tenant or workload: + +```bash +export MSG_HEADER_KEY="$(sudo cat /var/lib/pb-mapper/auth/admin.key)" +pb-mapper admin --server relay.example.com:7666 key issue --ttl 24h --label home-web +``` Registration instances read `/etc/pb-mapper/register/.env`: @@ -21,7 +32,7 @@ SERVICE_KEY=home-web LOCAL_ADDR=127.0.0.1:8080 TRANSPORT=tcp REGISTER_EXTRA_ARGS=--codec --keep-alive -MSG_HEADER_KEY=replace-with-the-shared-32-byte-key +MSG_HEADER_KEY=pbmt1_replace-with-an-issued-temporary-credential ``` Connect instances read `/etc/pb-mapper/connect/.env`: @@ -32,7 +43,7 @@ SERVICE_KEY=home-web LOCAL_ADDR=127.0.0.1:9090 TRANSPORT=tcp CONNECT_EXTRA_ARGS=--keep-alive -MSG_HEADER_KEY=replace-with-the-shared-32-byte-key +MSG_HEADER_KEY=pbmt1_replace-with-the-same-temporary-credential ``` Create the matching directory and env file, then enable the instance: diff --git a/skills/pb-mapper-connect-deploy/SKILL.md b/skills/pb-mapper-connect-deploy/SKILL.md index 5c6f7b7..15141a4 100644 --- a/skills/pb-mapper-connect-deploy/SKILL.md +++ b/skills/pb-mapper-connect-deploy/SKILL.md @@ -38,7 +38,7 @@ Prompt the user for each value below. Do NOT assume or hardcode any value. | `SERVICE_KEY` | pb-mapper service name to subscribe | — | Required | | `LISTEN_IP` | Listen on localhost only (`127.0.0.1`) or all interfaces (`0.0.0.0`)? | `127.0.0.1` | Required | | `LISTEN_PORT` | Local listening port on remote host | — | Required | -| `MSG_HEADER_KEY` | Encryption key (exactly 32 chars) | *(empty)* | Optional, **confidential** — never log or echo | +| `MSG_HEADER_KEY` | Administrator key or `pbmt1_...` temporary credential | *(empty)* | Required for authenticated v2 connections; **confidential** — never log or echo | | `PUBLIC_CHECK_URL` | URL for external validation | *(empty)* | Optional | | `VERSION` | Release version (without `v` prefix) | Latest release | Auto-detect or user-specified | | `TARGET_TRIPLE` | Build target | `x86_64-unknown-linux-musl` | User can override | @@ -143,7 +143,10 @@ REMOTE_SYSTEMD ### 4. Write instance env file and start service -`MSG_HEADER_KEY` must be omitted when empty; never write an empty value to env file. +Prefer a temporary credential issued by the relay administrator. It can register, +connect, and inspect only its own namespace, and it expires automatically. Use the +administrator key only for relay administration or an intentional namespace-0 +deployment. Never write an empty credential to the environment file. ```bash ssh ${SSH_PORT_OPT} "${SSH_TARGET}" \ @@ -159,14 +162,21 @@ RUST_LOG=info PB_MAPPER_KEEP_ALIVE=ON EOF -if [ -n "${MSG_HEADER_KEY}" ]; then - CLEAN_KEY="$(printf '%s' "${MSG_HEADER_KEY}" | tr -d '\r\n')" - if [ "${#CLEAN_KEY}" -ne 32 ]; then - echo "MSG_HEADER_KEY must be exactly 32 characters" >&2 - exit 1 - fi - echo "MSG_HEADER_KEY=${CLEAN_KEY}" | sudo tee -a "${ENV_FILE}" >/dev/null +if [ -z "${MSG_HEADER_KEY}" ]; then + echo "MSG_HEADER_KEY is required" >&2 + exit 1 fi +CLEAN_KEY="$(printf '%s' "${MSG_HEADER_KEY}" | tr -d '\r\n')" +case "${CLEAN_KEY}" in + pbmt1_*) ;; + *) + if [ "${#CLEAN_KEY}" -ne 32 ]; then + echo "MSG_HEADER_KEY must be a 32-byte administrator key or pbmt1_ temporary credential" >&2 + exit 1 + fi + ;; +esac +echo "MSG_HEADER_KEY=${CLEAN_KEY}" | sudo tee -a "${ENV_FILE}" >/dev/null sudo systemctl daemon-reload sudo systemctl enable --now "pb-mapper-connect@${INSTANCE_NAME}.service" @@ -191,7 +201,10 @@ If `jq` is available, pipe through `jq .` for formatted JSON output. Use this quick triage when startup or forwarding fails: -- `datalen not valid`: likely `MSG_HEADER_KEY` mismatch or hidden newline; verify both sides use the same 32-byte key. +- `protocol_v2_decrypt_failed`: credential does not belong to this relay or was corrupted in transit. +- `temporary_key_expired` / `temporary_key_revoked`: ask an administrator to renew the same key ID or issue a replacement. +- `namespace_access_denied`: a temporary credential can only use its own namespace; remove an incorrect `--namespace` value. +- Legacy `datalen not valid`: administrator key mismatch or a hidden newline in an old protocol-v1 client. - Service restarts immediately: inspect logs with `journalctl -u pb-mapper-connect@${INSTANCE_NAME} -n 200 --no-pager`. - Remote port not listening: confirm `LOCAL_ADDR` host/port and no port conflict. - Public URL fails but localhost works: investigate reverse proxy (for example, Caddy route/TLS config). diff --git a/skills/pb-mapper-server-deploy/SKILL.md b/skills/pb-mapper-server-deploy/SKILL.md index d3ddcd7..ec24086 100644 --- a/skills/pb-mapper-server-deploy/SKILL.md +++ b/skills/pb-mapper-server-deploy/SKILL.md @@ -39,7 +39,6 @@ Prompt the user for each value below. Do NOT assume or hardcode any value. | `SERVER_PORT` | `pb-mapper server` listening port | `7666` | User can override | | `USE_IPV6` | Listen on IPv6 (`::`) instead of IPv4? | `false` | Optional | | `ENABLE_KEEP_ALIVE` | Enable TCP keep-alive? | `false` | Optional | -| `USE_MACHINE_KEY` | Use machine-derived msg header key? | `true` | Recommended; generates and persists a 32-char key | | `VERSION` | Release version (without `v` prefix) | Latest release | Auto-detect or user-specified | | `TARGET_TRIPLE` | Build target | `x86_64-unknown-linux-musl` | User can override | @@ -67,9 +66,6 @@ fi if [ "${ENABLE_KEEP_ALIVE}" = "true" ]; then EXTRA_FLAGS="${EXTRA_FLAGS} --keep-alive" fi -if [ "${USE_MACHINE_KEY}" = "true" ]; then - EXTRA_FLAGS="${EXTRA_FLAGS} --use-machine-msg-header-key" -fi export EXTRA_FLAGS ``` @@ -138,6 +134,8 @@ After=network.target Type=simple ExecStart=/usr/local/bin/pb-mapper server --port ${SERVER_PORT} ${EXTRA_FLAGS} Environment=RUST_LOG=info +StateDirectory=pb-mapper +StateDirectoryMode=0700 Restart=on-failure RestartSec=3 LimitNOFILE=65535 @@ -165,13 +163,25 @@ ssh ${SSH_PORT_OPT} "${SSH_TARGET}" "sudo systemctl --no-pager --full status pb- ssh ${SSH_PORT_OPT} "${SSH_TARGET}" "ss -tlnp | grep ':${SERVER_PORT}'" ``` -If `USE_MACHINE_KEY=true`, retrieve the generated key for use with `register` and `connect`: +On a fresh installation, the relay creates a random administrator key with mode +`0600`. Retrieve it once to initialize the administrator CLI: + +```bash +ssh ${SSH_PORT_OPT} "${SSH_TARGET}" "sudo cat /var/lib/pb-mapper/auth/admin.key" +``` + +Store it securely. Do not distribute it to ordinary register/connect instances. +Instead, load it only in the administrator shell and issue a scoped temporary +credential: ```bash -ssh ${SSH_PORT_OPT} "${SSH_TARGET}" "sudo cat /var/lib/pb-mapper-server/msg_header_key" +export MSG_HEADER_KEY='' +pb-mapper admin --server "${SSH_HOST}:${SERVER_PORT}" key issue --ttl 30d --label '' ``` -Store this key securely — it is needed by `pb-mapper register` and `pb-mapper connect` when `--codec` or `MSG_HEADER_KEY` is used. +Copy the returned `pbmt1_...` credential to the business instance. Use +`pb-mapper admin key renew --ttl 30d` to extend it in place without +changing deployed credentials. ## Troubleshooting Checklist @@ -180,7 +190,8 @@ Use this quick triage when the server fails to start or accept connections: - Port already in use: check with `ss -tlnp | grep :${SERVER_PORT}` and stop the conflicting process. - Firewall blocking: ensure the server port is open (`ufw allow ${SERVER_PORT}/tcp` or equivalent). - Service crashes on start: inspect logs with `journalctl -u pb-mapper-server -n 200 --no-pager`. -- Machine key not generated: verify `/var/lib/pb-mapper-server/` directory exists and is writable; check logs for key derivation errors. +- Administrator key missing: verify `/var/lib/pb-mapper/auth/` is writable and inspect startup logs for `administrator_key_initialized` or an `auth_stage` error. +- Temporary clients rejected: run `pb-mapper admin status`, then `pb-mapper admin key show ` and correlate the stable error code in the server log. - Permission denied: binary must be owned by root with `0755` permissions. ## Safe Update Procedure diff --git a/src/bin/pb-mapper.rs b/src/bin/pb-mapper.rs index 686e445..f738bac 100644 --- a/src/bin/pb-mapper.rs +++ b/src/bin/pb-mapper.rs @@ -1,16 +1,29 @@ use std::error::Error; use std::net::{IpAddr, Ipv4Addr, Ipv6Addr}; +use std::path::PathBuf; +use std::time::Duration; use better_mimalloc_rs::MiMalloc; use clap::{Args, Parser, Subcommand, ValueEnum}; +use pb_mapper::common::auth::{ + generate_admin_key, initialize_admin_key, write_admin_key_file, LegacyProtocolPolicy, + DEFAULT_AUTH_STATE_DIR, +}; +use pb_mapper::common::checksum::set_process_msg_header_key; use pb_mapper::common::checksum::{setup_machine_msg_header_key, MACHINE_MSG_HEADER_KEY_PATH}; use pb_mapper::common::config::{ get_pb_mapper_server_async, get_sockaddr_async, init_tracing, keep_alive_from_env, StatusOp, }; +use pb_mapper::common::message::command::{ + AdminRequest, AdminResponse, MessageSerializer, PbConnRequest, PbConnResponse, +}; use pb_mapper::common::message::forward::StreamForward; -use pb_mapper::local::client::{handle_status_cli, run_client_side_cli}; +use pb_mapper::common::message::secure::ClientHeaderSession; +use pb_mapper::common::message::MessageReader; +use pb_mapper::local::client::{handle_status_cli_scoped, run_client_side_cli_scoped}; use pb_mapper::local::server::{run_server_side_cli, ServerTunnelOptions}; use pb_mapper::pb_server::run_server_with_shutdown; +use tokio::net::TcpStream; use tokio_util::sync::CancellationToken; use uni_stream::stream::{ StreamProvider, TcpListenerProvider, TcpStreamProvider, UdpListenerProvider, UdpStreamProvider, @@ -42,6 +55,8 @@ enum Command { Connect(ConnectArgs), /// Query relay status. Status(StatusArgs), + /// Manage temporary credentials and inspect relay authentication state. + Admin(AdminArgs), } #[derive(Debug, Args)] @@ -58,6 +73,162 @@ struct ServerArgs { /// Derive MSG_HEADER_KEY from this machine and persist it for other roles. #[arg(long, default_value_t = false)] use_machine_msg_header_key: bool, + /// Directory containing encrypted authentication state and the administrator key file. + #[arg(long, default_value = DEFAULT_AUTH_STATE_DIR)] + auth_state_dir: PathBuf, + /// Create a random administrator key before starting the relay. + #[arg(long, default_value_t = false)] + init_admin_key: bool, + /// Replace an existing administrator key when used with --init-admin-key. + #[arg(long, requires = "init_admin_key", default_value_t = false)] + force_init_admin_key: bool, + /// Maximum temporary-key slots allocated by the relay. + #[arg(long, default_value_t = 65_536)] + max_temporary_keys: usize, + /// Maximum accepted temporary-key TTL. + #[arg(long, default_value = "30d", value_parser = parse_duration)] + max_temporary_key_ttl: Duration, + /// Allow or deny the legacy encrypted framing protocol. + #[arg(long, value_enum, default_value_t = LegacyProtocolArg::Allow)] + legacy_protocol: LegacyProtocolArg, +} + +#[derive(Debug, Args)] +struct AdminArgs { + /// Relay address. Falls back to PB_MAPPER_SERVER. + #[arg(short, long, visible_alias = "pb-mapper-server", value_name = "ADDR")] + server: Option, + /// Machine-readable output mode. + #[arg(long, value_enum, default_value_t = OutputFormat::Human)] + output: OutputFormat, + #[command(subcommand)] + command: AdminCommand, +} + +#[derive(Debug, Subcommand)] +enum AdminCommand { + /// Issue, inspect, renew, reveal, revoke, or collect temporary keys. + Key(AdminKeyArgs), + /// List relay connections across namespaces. + Connection(AdminConnectionArgs), + /// List registered services across namespaces. + Service(AdminServiceArgs), + /// Show authentication state and protocol counters. + Status, + /// Repair or reset encrypted temporary-key state. + AuthState(AdminAuthStateArgs), + /// Rotate the sole administrator key and invalidate every existing credential. + RootKey(AdminRootKeyArgs), + /// Change legacy protocol acceptance at runtime. + LegacyProtocol(AdminLegacyProtocolArgs), +} + +#[derive(Debug, Args)] +struct AdminKeyArgs { + #[command(subcommand)] + command: AdminKeyCommand, +} + +#[derive(Debug, Subcommand)] +enum AdminKeyCommand { + Issue { + #[arg(long, value_parser = parse_duration)] + ttl: Duration, + #[arg(long)] + label: Option, + }, + List { + #[arg(long, default_value_t = 0)] + page: u32, + #[arg(long, default_value_t = 100, value_parser = clap::value_parser!(u16).range(1..=1000))] + page_size: u16, + #[arg(long, default_value_t = false)] + all: bool, + }, + Show { + key_id: u64, + }, + Reveal { + key_id: u64, + }, + Renew { + key_id: u64, + #[arg(long, value_parser = parse_duration)] + ttl: Duration, + }, + Revoke { + key_id: u64, + }, + Gc, +} + +#[derive(Debug, Args)] +struct AdminConnectionArgs { + #[command(subcommand)] + command: AdminListCommand, +} + +#[derive(Debug, Args)] +struct AdminServiceArgs { + #[command(subcommand)] + command: AdminListCommand, +} + +#[derive(Debug, Clone, Subcommand)] +enum AdminListCommand { + List { + #[arg(long)] + key_id: Option, + #[arg(long, default_value_t = 0)] + page: u32, + #[arg(long, default_value_t = 100, value_parser = clap::value_parser!(u16).range(1..=1000))] + page_size: u16, + #[arg(long, default_value_t = false)] + all: bool, + }, +} + +#[derive(Debug, Args)] +struct AdminAuthStateArgs { + #[command(subcommand)] + command: AdminAuthStateCommand, +} + +#[derive(Debug, Subcommand)] +enum AdminAuthStateCommand { + Reset { + #[arg(long, default_value_t = false)] + confirm: bool, + }, +} + +#[derive(Debug, Args)] +struct AdminRootKeyArgs { + #[command(subcommand)] + command: AdminRootKeyCommand, +} + +#[derive(Debug, Subcommand)] +enum AdminRootKeyCommand { + Rotate { + /// New 32-byte administrator key. A cryptographically random printable key is generated when omitted. + #[arg(long)] + new_key: Option, + /// Save the new key here before asking the relay to rotate. + #[arg(long, default_value = "/var/lib/pb-mapper/auth/admin.key")] + key_file: PathBuf, + }, +} + +#[derive(Debug, Args)] +struct AdminLegacyProtocolArgs { + #[command(subcommand)] + command: AdminLegacyProtocolCommand, +} + +#[derive(Debug, Subcommand)] +enum AdminLegacyProtocolCommand { + Set { policy: LegacyProtocolArg }, } #[derive(Debug, Args)] @@ -76,6 +247,12 @@ struct RegisterArgs { /// Encrypt forwarded traffic with the configured MSG_HEADER_KEY. #[arg(short, long, default_value_t = false)] codec: bool, + /// Administrator-only target namespace. Temporary credentials always use their own key id. + #[arg(long)] + namespace: Option, + /// Required when an administrator registers a service inside a temporary-key namespace. + #[arg(long, requires = "namespace", default_value_t = false)] + force: bool, } #[derive(Debug, Args)] @@ -91,6 +268,9 @@ struct ConnectArgs { addr: String, #[command(flatten)] relay: RelayArgs, + /// Administrator-only target namespace. Temporary credentials always use their own key id. + #[arg(long)] + namespace: Option, } #[derive(Debug, Args)] @@ -101,6 +281,9 @@ struct StatusArgs { /// Relay address. Falls back to PB_MAPPER_SERVER. #[arg(short, long, visible_alias = "pb-mapper-server", value_name = "ADDR")] server: Option, + /// Administrator-only namespace to inspect. + #[arg(long)] + namespace: Option, } #[derive(Debug, Args)] @@ -119,6 +302,28 @@ enum Transport { Udp, } +#[derive(Clone, Copy, Debug, Eq, PartialEq, ValueEnum)] +enum OutputFormat { + Human, + Json, + Ndjson, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq, ValueEnum)] +enum LegacyProtocolArg { + Allow, + Deny, +} + +impl From for LegacyProtocolPolicy { + fn from(value: LegacyProtocolArg) -> Self { + match value { + LegacyProtocolArg::Allow => Self::Allow, + LegacyProtocolArg::Deny => Self::Deny, + } + } +} + #[tokio::main] async fn main() { MiMalloc::init(); @@ -137,12 +342,38 @@ async fn run(cli: Cli) -> Result<(), Box> { Command::Register(args) => run_register(args).await?, Command::Connect(args) => run_connect(args).await?, Command::Status(args) => run_status(args).await?, + Command::Admin(args) => run_admin(args).await?, } Ok(()) } async fn run_server(args: ServerArgs) -> Result<(), Box> { + std::env::set_var("PB_MAPPER_AUTH_STATE_DIR", &args.auth_state_dir); + std::env::set_var( + "PB_MAPPER_AUTH_MAX_TEMP_KEYS", + args.max_temporary_keys.to_string(), + ); + std::env::set_var( + "PB_MAPPER_AUTH_MAX_TEMP_TTL_SECS", + args.max_temporary_key_ttl.as_secs().to_string(), + ); + std::env::set_var( + "PB_MAPPER_LEGACY_PROTOCOL", + match args.legacy_protocol { + LegacyProtocolArg::Allow => "allow", + LegacyProtocolArg::Deny => "deny", + }, + ); + if args.init_admin_key { + let key_path = args.auth_state_dir.join("admin.key"); + let key = initialize_admin_key(&key_path, args.force_init_admin_key)?; + set_process_msg_header_key(Some(&key))?; + eprintln!("administrator key initialized at {}", key_path.display()); + } if args.use_machine_msg_header_key { + tracing::warn!( + "--use-machine-msg-header-key is a legacy compatibility option; prefer a random administrator key" + ); setup_machine_msg_header_key()?; tracing::info!( path = MACHINE_MSG_HEADER_KEY_PATH, @@ -172,6 +403,8 @@ async fn run_register(args: RegisterArgs) -> Result<(), Box> { need_codec: args.codec, is_datagram: args.transport == Transport::Udp, keep_alive: args.relay.keep_alive || keep_alive_from_env(), + namespace: args.namespace, + force_namespace: args.force, }; match args.transport { @@ -204,12 +437,24 @@ async fn run_connect(args: ConnectArgs) -> Result<(), Box> { match args.transport { Transport::Tcp => { - run_client_side_cli::(local_addr, remote_addr, key, keep_alive) - .await; + run_client_side_cli_scoped::( + local_addr, + remote_addr, + key, + keep_alive, + args.namespace, + ) + .await; } Transport::Udp => { - run_client_side_cli::(local_addr, remote_addr, key, keep_alive) - .await; + run_client_side_cli_scoped::( + local_addr, + remote_addr, + key, + keep_alive, + args.namespace, + ) + .await; } } Ok(()) @@ -217,10 +462,443 @@ async fn run_connect(args: ConnectArgs) -> Result<(), Box> { async fn run_status(args: StatusArgs) -> Result<(), Box> { let remote_addr = get_pb_mapper_server_async(args.server.as_deref()).await?; - handle_status_cli(args.op, remote_addr).await; + handle_status_cli_scoped(args.op, remote_addr, args.namespace).await; Ok(()) } +async fn run_admin(args: AdminArgs) -> Result<(), Box> { + let remote_addr = get_pb_mapper_server_async(args.server.as_deref()).await?; + match args.command { + AdminCommand::Key(AdminKeyArgs { command }) => match command { + AdminKeyCommand::Issue { ttl, label } => { + let response = send_admin_request( + remote_addr, + AdminRequest::KeyIssue { + ttl_seconds: ttl.as_secs(), + label, + }, + ) + .await?; + print_admin_response(args.output, &response)?; + } + AdminKeyCommand::List { + page, + page_size, + all, + } => { + stream_key_pages(remote_addr, args.output, page, page_size, all).await?; + } + AdminKeyCommand::Show { key_id } => { + let response = + send_admin_request(remote_addr, AdminRequest::KeyShow { key_id }).await?; + print_admin_response(args.output, &response)?; + } + AdminKeyCommand::Reveal { key_id } => { + let response = + send_admin_request(remote_addr, AdminRequest::KeyReveal { key_id }).await?; + print_admin_response(args.output, &response)?; + } + AdminKeyCommand::Renew { key_id, ttl } => { + let response = send_admin_request( + remote_addr, + AdminRequest::KeyRenew { + key_id, + ttl_seconds: ttl.as_secs(), + }, + ) + .await?; + print_admin_response(args.output, &response)?; + } + AdminKeyCommand::Revoke { key_id } => { + let response = + send_admin_request(remote_addr, AdminRequest::KeyRevoke { key_id }).await?; + print_admin_response(args.output, &response)?; + } + AdminKeyCommand::Gc => { + let response = send_admin_request(remote_addr, AdminRequest::KeyGc).await?; + print_admin_response(args.output, &response)?; + } + }, + AdminCommand::Connection(AdminConnectionArgs { command }) => { + let AdminListCommand::List { + key_id, + page, + page_size, + all, + } = command; + stream_connection_pages(remote_addr, args.output, key_id, page, page_size, all).await?; + } + AdminCommand::Service(AdminServiceArgs { command }) => { + let AdminListCommand::List { + key_id, + page, + page_size, + all, + } = command; + stream_service_pages(remote_addr, args.output, key_id, page, page_size, all).await?; + } + AdminCommand::Status => { + let response = send_admin_request(remote_addr, AdminRequest::AuthStatus).await?; + print_admin_response(args.output, &response)?; + } + AdminCommand::AuthState(AdminAuthStateArgs { + command: AdminAuthStateCommand::Reset { confirm }, + }) => { + let response = + send_admin_request(remote_addr, AdminRequest::AuthStateReset { confirm }).await?; + print_admin_response(args.output, &response)?; + } + AdminCommand::RootKey(AdminRootKeyArgs { + command: AdminRootKeyCommand::Rotate { new_key, key_file }, + }) => { + let new_key = new_key.unwrap_or_else(generate_admin_key); + let staged_key_file = key_file.with_file_name(format!( + ".{}.next", + key_file + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or("admin.key") + )); + write_admin_key_file(&staged_key_file, &new_key, true)?; + let response = send_admin_request( + remote_addr, + AdminRequest::RootKeyRotate { + new_admin_key: new_key.clone(), + }, + ) + .await + .map_err(|error| { + std::io::Error::other(format!( + "root rotation request failed; the candidate key remains at `{}`: {error}", + staged_key_file.display() + )) + })?; + set_process_msg_header_key(Some(&new_key))?; + let verification = send_admin_request(remote_addr, AdminRequest::AuthStatus).await?; + if !matches!(verification, AdminResponse::AuthStatus(_)) { + return Err(std::io::Error::other( + "new administrator key did not pass the post-rotation status check", + ) + .into()); + } + write_admin_key_file(&key_file, &new_key, true).map_err(|error| { + std::io::Error::other(format!( + "administrator key rotated and verified, but `{}` could not be updated; recover the key from `{}`: {error}", + key_file.display(), + staged_key_file.display() + )) + })?; + if let Err(error) = std::fs::remove_file(&staged_key_file) { + tracing::warn!( + path = %staged_key_file.display(), + %error, + "administrator key was rotated, but the staged key file could not be removed" + ); + } + if args.output == OutputFormat::Human { + println!("administrator key rotated and verified"); + println!("key file: {}", key_file.display()); + } else { + print_admin_response(args.output, &response)?; + } + } + AdminCommand::LegacyProtocol(AdminLegacyProtocolArgs { + command: AdminLegacyProtocolCommand::Set { policy }, + }) => { + let response = send_admin_request( + remote_addr, + AdminRequest::LegacyProtocolSet { + policy: policy.into(), + }, + ) + .await?; + print_admin_response(args.output, &response)?; + } + } + Ok(()) +} + +async fn send_admin_request( + remote_addr: std::net::SocketAddr, + request: AdminRequest, +) -> Result> { + let encoded = PbConnRequest::Admin(request).encode()?; + for attempt in 0..2 { + let mut stream = TcpStream::connect(remote_addr).await?; + let session = ClientHeaderSession::from_process()?; + session.write_initial(&mut stream, &encoded).await?; + let mut reader = session.response_reader(&mut stream)?; + let message = reader.read_msg().await?; + match PbConnResponse::decode(message)? { + PbConnResponse::Admin(response) => return Ok(response), + PbConnResponse::Error(error) + if error.code == "connection_salt_replayed" && error.retryable && attempt == 0 => + { + continue; + } + PbConnResponse::Error(error) => { + return Err(std::io::Error::other(format!( + "{}: {} (retryable={})", + error.code, error.message, error.retryable + )) + .into()); + } + response => { + return Err(std::io::Error::other(format!( + "unexpected administrator response: {response:?}" + )) + .into()); + } + } + } + Err(std::io::Error::other("connection salt replay retry was exhausted").into()) +} + +async fn stream_key_pages( + remote_addr: std::net::SocketAddr, + output: OutputFormat, + mut page: u32, + page_size: u16, + all: bool, +) -> Result<(), Box> { + loop { + let response = + send_admin_request(remote_addr, AdminRequest::KeyList { page, page_size }).await?; + let AdminResponse::KeyList(key_page) = &response else { + return Err(std::io::Error::other("unexpected key-list response").into()); + }; + if all { + for item in &key_page.items { + println!("{}", serde_json::to_string(item)?); + } + } else { + print_admin_response(output, &response)?; + } + let Some(next_page) = key_page.next_page else { + break; + }; + if !all { + break; + } + page = next_page; + } + Ok(()) +} + +async fn stream_service_pages( + remote_addr: std::net::SocketAddr, + output: OutputFormat, + key_id: Option, + mut page: u32, + page_size: u16, + all: bool, +) -> Result<(), Box> { + loop { + let response = send_admin_request( + remote_addr, + AdminRequest::ServiceList { + key_id, + page, + page_size, + }, + ) + .await?; + let AdminResponse::Services(service_page) = &response else { + return Err(std::io::Error::other("unexpected service-list response").into()); + }; + if all { + for item in &service_page.items { + println!("{}", serde_json::to_string(item)?); + } + } else { + print_admin_response(output, &response)?; + } + let Some(next_page) = service_page.next_page else { + break; + }; + if !all { + break; + } + page = next_page; + } + Ok(()) +} + +async fn stream_connection_pages( + remote_addr: std::net::SocketAddr, + output: OutputFormat, + key_id: Option, + mut page: u32, + page_size: u16, + all: bool, +) -> Result<(), Box> { + loop { + let response = send_admin_request( + remote_addr, + AdminRequest::ConnectionList { + key_id, + page, + page_size, + }, + ) + .await?; + let AdminResponse::Connections(connection_page) = &response else { + return Err(std::io::Error::other("unexpected connection-list response").into()); + }; + if all { + for item in &connection_page.items { + println!("{}", serde_json::to_string(item)?); + } + } else { + print_admin_response(output, &response)?; + } + let Some(next_page) = connection_page.next_page else { + break; + }; + if !all { + break; + } + page = next_page; + } + Ok(()) +} + +fn print_admin_response( + output: OutputFormat, + response: &AdminResponse, +) -> Result<(), Box> { + match output { + OutputFormat::Json => println!( + "{}", + serde_json::to_string_pretty(&serde_json::json!({ + "schema_version": 1, + "data": response, + }))? + ), + OutputFormat::Ndjson => println!("{}", serde_json::to_string(response)?), + OutputFormat::Human => print_human_admin_response(response), + } + Ok(()) +} + +fn print_human_admin_response(response: &AdminResponse) { + match response { + AdminResponse::KeyIssued(key) + | AdminResponse::KeyShown(key) + | AdminResponse::KeyRenewed(key) => { + println!("key id: {}", key.metadata.key_id); + println!("state: {}", key.metadata.state); + println!("expires at: {}", key.metadata.expires_at); + if let Some(label) = &key.metadata.label { + println!("label: {label}"); + } + if !key.credential.is_empty() { + println!("credential: {}", key.credential); + } + } + AdminResponse::KeyRevoked(key) => { + println!("key {}: {}", key.key_id, key.state); + } + AdminResponse::KeyList(page) => { + println!("KEY ID\tSTATE\tEXPIRES\tLABEL"); + for key in &page.items { + println!( + "{}\t{}\t{}\t{}", + key.key_id, + key.state, + key.expires_at, + key.label.as_deref().unwrap_or("") + ); + } + if let Some(next) = page.next_page { + println!("next page: {next}"); + } + } + AdminResponse::KeyGc { removed } => println!("removed {removed} inactive keys"), + AdminResponse::AuthStatus(status) => { + println!("safe mode: {}", status.safe_mode); + println!( + "keys: {} active / {} expired / {} revoked / {} capacity", + status.active_keys, status.expired_keys, status.revoked_keys, status.capacity + ); + println!("legacy protocol: {:?}", status.legacy_protocol); + println!( + "active legacy connections: {}", + status.active_legacy_connections + ); + println!( + "last legacy connection: {}", + status + .last_legacy_connection_at + .map(|value| value.to_string()) + .unwrap_or_else(|| "never".to_string()) + ); + println!( + "authentication: {} succeeded / {} failed", + status.auth_successes, status.auth_failures + ); + println!("server instance: {}", status.server_instance_id); + } + AdminResponse::Services(page) => { + println!("KEY ID\tSERVICE\tTRANSPORT\tCODEC\tCONNECTIONS"); + for service in &page.items { + println!( + "{}\t{}\t{}\t{}\t{}", + service.key_id, + service.service_name, + service.transport, + service.codec_enabled, + service.connection_count + ); + } + } + AdminResponse::Connections(page) => { + println!("KEY ID\tSERVICE\tCONN ID\tHEALTHY\tTRANSPORT\tCODEC"); + for connection in &page.items { + println!( + "{}\t{}\t{}\t{}\t{}\t{}", + connection.key_id, + connection.service_name, + connection.conn_id, + connection.healthy, + connection.transport, + connection.codec_enabled + ); + } + } + AdminResponse::Ok { action } => println!("ok: {action}"), + } +} + +fn parse_duration(raw: &str) -> Result { + let raw = raw.trim(); + if raw.is_empty() { + return Err("duration must not be empty".to_string()); + } + let split = raw + .find(|character: char| !character.is_ascii_digit()) + .unwrap_or(raw.len()); + let (number, unit) = raw.split_at(split); + let value = number + .parse::() + .map_err(|_| format!("invalid duration `{raw}`"))?; + let multiplier = match unit { + "" | "s" => 1, + "m" => 60, + "h" => 60 * 60, + "d" => 24 * 60 * 60, + _ => { + return Err(format!( + "unsupported duration unit `{unit}`; use s, m, h, or d" + )) + } + }; + value + .checked_mul(multiplier) + .map(Duration::from_secs) + .ok_or_else(|| "duration is too large".to_string()) +} + #[cfg(test)] mod tests { use super::*; @@ -253,6 +931,37 @@ mod tests { "relay:7666", ], vec!["pb-mapper", "status", "keys", "--server", "relay:7666"], + vec![ + "pb-mapper", + "admin", + "--server", + "relay:7666", + "key", + "issue", + "--ttl", + "30d", + "--label", + "build-agent", + ], + vec![ + "pb-mapper", + "admin", + "--output", + "ndjson", + "connection", + "list", + "--page-size", + "1000", + "--all", + ], + vec![ + "pb-mapper", + "admin", + "root-key", + "rotate", + "--key-file", + "/tmp/pb-mapper-admin.key", + ], ]; for args in cases { @@ -282,5 +991,30 @@ mod tests { "relay:7666", ]) .expect("relay and local aliases should parse"); + Cli::try_parse_from([ + "pb-mapper", + "register", + "tcp", + "--key", + "web", + "--addr", + "127.0.0.1:8080", + "--namespace", + "4294967296", + "--force", + ]) + .expect("administrator namespace registration flags should parse"); + } + + #[test] + fn rejects_invalid_admin_paging_and_duration() { + assert!( + Cli::try_parse_from(["pb-mapper", "admin", "key", "list", "--page-size", "1001",]) + .is_err() + ); + assert!(Cli::try_parse_from( + ["pb-mapper", "admin", "key", "issue", "--ttl", "1fortnight",] + ) + .is_err()); } } diff --git a/src/common/auth.rs b/src/common/auth.rs new file mode 100644 index 0000000..1b3b66e --- /dev/null +++ b/src/common/auth.rs @@ -0,0 +1,2548 @@ +//! Authentication state for protocol-v2 connections and administrator operations. +//! +//! The root administrator key is never copied into a temporary credential. Temporary +//! keys are derived from `(root key, server instance id, key id)` and the hot slot table +//! stores only lifecycle metadata plus a weak lease reference. The background actor owns +//! the strong leases through a hierarchical timing wheel. + +use std::collections::{HashMap, VecDeque}; +use std::fmt; +use std::fs::{File, OpenOptions}; +use std::io::{Read, Write}; +#[cfg(unix)] +use std::os::unix::fs::PermissionsExt; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::{Arc, RwLock, Weak}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use rand::RngExt; +use ring::aead::{Aad, LessSafeKey, Nonce, UnboundKey, AES_256_GCM}; +use ring::hkdf::{Salt, HKDF_SHA256}; +use serde::{Deserialize, Serialize}; +use tokio::sync::{mpsc, oneshot}; +use tokio_util::sync::CancellationToken; + +use super::checksum::{ + encode_temporary_credential, get_process_credential, parse_credential, + set_process_msg_header_key, AesKeyType, Credential, MACHINE_MSG_HEADER_KEY_PATH, +}; + +pub const ADMIN_NAMESPACE: u64 = 0; +pub const DEFAULT_AUTH_STATE_DIR: &str = "/var/lib/pb-mapper/auth"; +pub const DEFAULT_TEMP_KEY_CAPACITY: usize = 65_536; +pub const DEFAULT_MAX_TEMP_KEY_TTL: Duration = Duration::from_secs(30 * 24 * 60 * 60); +pub const MIN_TEMP_KEY_TTL: Duration = Duration::from_secs(10); +const TOMBSTONE_RETENTION: Duration = Duration::from_secs(60); +const SNAPSHOT_COMPACTION_INTERVAL: Duration = Duration::from_secs(5 * 60); +const SNAPSHOT_SCHEMA_VERSION: u16 = 1; +const STATE_BLOB_MAGIC: &[u8; 5] = b"PBAS1"; +const STATE_AAD: &[u8] = b"pb-mapper-auth-state-v1"; +const INSTANCE_ID_LEN: usize = 16; + +#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum LegacyProtocolPolicy { + Allow, + Deny, +} + +impl LegacyProtocolPolicy { + pub fn is_allowed(self) -> bool { + matches!(self, Self::Allow) + } +} + +#[derive(Clone, Debug)] +pub struct AuthConfig { + pub state_dir: PathBuf, + pub max_temporary_keys: usize, + pub max_temporary_key_ttl: Duration, + pub legacy_protocol: LegacyProtocolPolicy, +} + +impl Default for AuthConfig { + fn default() -> Self { + Self { + state_dir: std::env::var_os("PB_MAPPER_AUTH_STATE_DIR") + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from(DEFAULT_AUTH_STATE_DIR)), + max_temporary_keys: env_usize( + "PB_MAPPER_AUTH_MAX_TEMP_KEYS", + DEFAULT_TEMP_KEY_CAPACITY, + 1, + 1_048_576, + ), + max_temporary_key_ttl: Duration::from_secs(env_u64( + "PB_MAPPER_AUTH_MAX_TEMP_TTL_SECS", + DEFAULT_MAX_TEMP_KEY_TTL.as_secs(), + MIN_TEMP_KEY_TTL.as_secs(), + 365 * 24 * 60 * 60, + )), + legacy_protocol: match std::env::var("PB_MAPPER_LEGACY_PROTOCOL") + .unwrap_or_else(|_| "allow".to_string()) + .to_ascii_lowercase() + .as_str() + { + "deny" => LegacyProtocolPolicy::Deny, + _ => LegacyProtocolPolicy::Allow, + }, + } + } +} + +fn env_usize(name: &str, default: usize, min: usize, max: usize) -> usize { + std::env::var(name) + .ok() + .and_then(|value| value.parse::().ok()) + .filter(|value| (*value >= min) && (*value <= max)) + .unwrap_or(default) +} + +fn env_u64(name: &str, default: u64, min: u64, max: u64) -> u64 { + std::env::var(name) + .ok() + .and_then(|value| value.parse::().ok()) + .filter(|value| (*value >= min) && (*value <= max)) + .unwrap_or(default) +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct AuthFailure { + pub code: String, + pub message: String, + pub retryable: bool, +} + +impl AuthFailure { + pub fn new(code: impl Into, message: impl Into, retryable: bool) -> Self { + Self { + code: code.into(), + message: message.into(), + retryable, + } + } + + pub fn internal(message: impl Into) -> Self { + Self::new("auth_internal_error", message, false) + } +} + +impl fmt::Display for AuthFailure { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(formatter, "{}: {}", self.code, self.message) + } +} + +impl std::error::Error for AuthFailure {} + +#[derive(Debug)] +pub struct AuthLease { + key_id: u64, + expires_at: AtomicU64, + wheel_version: AtomicU64, + cancellation: CancellationToken, +} + +impl AuthLease { + fn new(key_id: u64, expires_at: u64) -> Self { + Self { + key_id, + expires_at: AtomicU64::new(expires_at), + wheel_version: AtomicU64::new(1), + cancellation: CancellationToken::new(), + } + } + + pub fn key_id(&self) -> u64 { + self.key_id + } + + pub fn expires_at(&self) -> u64 { + self.expires_at.load(Ordering::Acquire) + } + + pub fn cancellation_token(&self) -> CancellationToken { + self.cancellation.clone() + } +} + +#[derive(Clone, Debug)] +pub struct AuthContext { + pub key_id: u64, + pub namespace: u64, + pub is_admin: bool, + lease: Weak, +} + +impl AuthContext { + fn from_lease(key_id: u64, is_admin: bool, lease: &Arc) -> Self { + Self { + key_id, + namespace: if is_admin { ADMIN_NAMESPACE } else { key_id }, + is_admin, + lease: Arc::downgrade(lease), + } + } + + pub fn ensure_active(&self) -> Result, AuthFailure> { + let lease = self.lease.upgrade().ok_or_else(|| { + AuthFailure::new( + if self.is_admin { + "administrator_key_rotated" + } else { + "temporary_key_inactive" + }, + "credential lease is no longer active", + false, + ) + })?; + if lease.cancellation.is_cancelled() { + return Err(AuthFailure::new( + if self.is_admin { + "administrator_key_rotated" + } else { + "temporary_key_revoked" + }, + "credential lease has been cancelled", + false, + )); + } + if !self.is_admin && lease.expires_at() <= unix_seconds() { + lease.cancellation.cancel(); + return Err(AuthFailure::new( + "temporary_key_expired", + "temporary key has expired", + false, + )); + } + Ok(lease) + } + + pub fn cancellation_token(&self) -> Result { + Ok(self.ensure_active()?.cancellation_token()) + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +enum SlotState { + Free, + Active, + Expired, + Revoked, +} + +#[derive(Debug)] +struct SlotHot { + generation: u32, + state: SlotState, + expires_at: u64, + lease: Weak, +} + +impl Default for SlotHot { + fn default() -> Self { + Self { + generation: 0, + state: SlotState::Free, + expires_at: 0, + lease: Weak::new(), + } + } +} + +#[derive(Debug)] +struct AuthStateInner { + admin_key: RwLock, + admin_lease: RwLock>, + instance_id: RwLock<[u8; INSTANCE_ID_LEN]>, + slots: RwLock>, + safe_mode: AtomicBool, + legacy_protocol_allowed: AtomicBool, + active_legacy_connections: AtomicU64, + last_legacy_connection_at: AtomicU64, + auth_successes: AtomicU64, + auth_failures: AtomicU64, +} + +impl AuthStateInner { + fn admin_key(&self) -> AesKeyType { + *self + .admin_key + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + } + + fn instance_id(&self) -> [u8; INSTANCE_ID_LEN] { + *self + .instance_id + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + } +} + +#[derive(Clone)] +pub struct AuthRuntime { + inner: Weak, + command_tx: mpsc::Sender, + config: AuthConfig, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct TemporaryKeyMetadata { + pub key_id: u64, + pub state: String, + pub issued_at: u64, + pub expires_at: u64, + pub label: Option, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct IssuedTemporaryKey { + #[serde(flatten)] + pub metadata: TemporaryKeyMetadata, + pub credential: String, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct KeyPage { + pub schema_version: u16, + pub items: Vec, + pub next_page: Option, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct AuthStatus { + pub schema_version: u16, + pub safe_mode: bool, + pub capacity: usize, + pub active_keys: usize, + pub expired_keys: usize, + pub revoked_keys: usize, + pub legacy_protocol: LegacyProtocolPolicy, + pub active_legacy_connections: u64, + pub last_legacy_connection_at: Option, + pub auth_successes: u64, + pub auth_failures: u64, + pub server_instance_id: String, +} + +#[derive(Clone, Debug)] +struct ColdMetadata { + issued_at: u64, + label: Option, +} + +enum AuthCommand { + Issue { + ttl: Duration, + label: Option, + response: oneshot::Sender>, + }, + List { + page: u32, + page_size: u16, + response: oneshot::Sender>, + }, + Show { + key_id: u64, + reveal: bool, + response: oneshot::Sender>, + }, + Renew { + key_id: u64, + ttl: Duration, + response: oneshot::Sender>, + }, + Revoke { + key_id: u64, + response: oneshot::Sender>, + }, + Gc { + response: oneshot::Sender>, + }, + Reset { + response: oneshot::Sender>, + }, + RotateRoot { + new_key: AesKeyType, + response: oneshot::Sender>, + }, + SetLegacyProtocol { + policy: LegacyProtocolPolicy, + response: oneshot::Sender>, + }, + Status { + response: oneshot::Sender>, + }, + Audit { + action: String, + key_id: Option, + detail: Option, + response: oneshot::Sender>, + }, +} + +impl AuthRuntime { + pub async fn from_process(config: AuthConfig) -> Result { + prepare_state_dir(&config.state_dir)?; + let credential = load_server_admin_credential(&config.state_dir)?; + let Credential::Admin(admin_key) = credential else { + return Err(AuthFailure::new( + "administrator_key_required", + "the relay server must start with the administrator credential", + false, + )); + }; + Self::start(admin_key, config).await + } + + pub async fn start(admin_key: AesKeyType, config: AuthConfig) -> Result { + prepare_state_dir(&config.state_dir)?; + let instance_id = load_or_create_instance_id(&config.state_dir)?; + let (loaded, safe_mode) = load_persisted_state(&config, &admin_key, instance_id); + let mut slots = (0..config.max_temporary_keys) + .map(|_| SlotHot::default()) + .collect::>() + .into_boxed_slice(); + let mut cold = HashMap::new(); + let mut wheel = TimingWheel::new(unix_seconds()); + let now = unix_seconds(); + + let admin_lease = Arc::new(AuthLease::new(0, u64::MAX)); + if let Some(state) = loaded.as_ref() { + for (index, generation) in state.generations.iter().copied().enumerate() { + if let Some(slot) = slots.get_mut(index) { + slot.generation = generation; + } + } + for entry in &state.entries { + let index = key_slot(entry.key_id) as usize; + let Some(slot) = slots.get_mut(index) else { + continue; + }; + if slot.generation != key_generation(entry.key_id) { + continue; + } + let state = if entry.state == SlotState::Active && entry.expires_at <= now { + SlotState::Expired + } else { + entry.state + }; + slot.state = state; + slot.expires_at = entry.expires_at; + cold.insert( + entry.key_id, + ColdMetadata { + issued_at: entry.issued_at, + label: entry.label.clone(), + }, + ); + if state == SlotState::Active { + let lease = Arc::new(AuthLease::new(entry.key_id, entry.expires_at)); + slot.lease = Arc::downgrade(&lease); + wheel.insert(lease); + } + } + } + + let legacy_protocol = loaded + .as_ref() + .map(|state| state.legacy_protocol) + .unwrap_or(config.legacy_protocol); + let inner = Arc::new(AuthStateInner { + admin_key: RwLock::new(admin_key), + admin_lease: RwLock::new(Arc::downgrade(&admin_lease)), + instance_id: RwLock::new(instance_id), + slots: RwLock::new(slots), + safe_mode: AtomicBool::new(safe_mode), + legacy_protocol_allowed: AtomicBool::new(legacy_protocol.is_allowed()), + active_legacy_connections: AtomicU64::new(0), + last_legacy_connection_at: AtomicU64::new(0), + auth_successes: AtomicU64::new(0), + auth_failures: AtomicU64::new(0), + }); + let (command_tx, command_rx) = mpsc::channel(256); + let runtime = Self { + inner: Arc::downgrade(&inner), + command_tx, + config: config.clone(), + }; + + tokio::spawn(run_auth_actor( + inner, + admin_lease, + command_rx, + config, + cold, + wheel, + )); + Ok(runtime) + } + + pub fn config(&self) -> &AuthConfig { + &self.config + } + + fn inner(&self) -> Result, AuthFailure> { + self.inner.upgrade().ok_or_else(|| { + AuthFailure::new( + "auth_state_unavailable", + "authentication state manager is not running", + true, + ) + }) + } + + pub fn admin_key(&self) -> Result { + Ok(self.inner()?.admin_key()) + } + + pub fn derive_key(&self, key_id: u64) -> Result { + let inner = self.inner()?; + if key_id == 0 { + return Ok(inner.admin_key()); + } + derive_temporary_key(&inner.admin_key(), &inner.instance_id(), key_id) + } + + pub fn authenticate(&self, key_id: u64) -> Result { + let inner = self.inner()?; + if key_id == 0 { + let lease = inner + .admin_lease + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .upgrade() + .ok_or_else(|| { + AuthFailure::new( + "administrator_key_rotated", + "administrator credential was rotated", + false, + ) + })?; + inner.auth_successes.fetch_add(1, Ordering::Relaxed); + return Ok(AuthContext::from_lease(0, true, &lease)); + } + if inner.safe_mode.load(Ordering::Acquire) { + inner.auth_failures.fetch_add(1, Ordering::Relaxed); + return Err(AuthFailure::new( + "temporary_key_store_unavailable", + "temporary key state is unavailable; administrator reset is required", + false, + )); + } + + let index = key_slot(key_id) as usize; + let generation = key_generation(key_id); + let slots = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let Some(slot) = slots.get(index) else { + inner.auth_failures.fetch_add(1, Ordering::Relaxed); + return Err(AuthFailure::new( + "temporary_key_not_found", + "temporary key id is outside the configured slot table", + false, + )); + }; + if slot.generation != generation { + inner.auth_failures.fetch_add(1, Ordering::Relaxed); + return Err(AuthFailure::new( + "temporary_key_generation_mismatch", + "temporary key generation does not match the current slot", + false, + )); + } + let failure = match slot.state { + SlotState::Free => Some(AuthFailure::new( + "temporary_key_not_found", + "temporary key does not exist", + false, + )), + SlotState::Expired => Some(AuthFailure::new( + "temporary_key_expired", + "temporary key has expired", + false, + )), + SlotState::Revoked => Some(AuthFailure::new( + "temporary_key_revoked", + "temporary key was revoked", + false, + )), + SlotState::Active if slot.expires_at <= unix_seconds() => { + if let Some(lease) = slot.lease.upgrade() { + lease.cancellation.cancel(); + } + Some(AuthFailure::new( + "temporary_key_expired", + "temporary key has expired", + false, + )) + } + SlotState::Active => None, + }; + if let Some(failure) = failure { + inner.auth_failures.fetch_add(1, Ordering::Relaxed); + return Err(failure); + } + let lease = slot.lease.upgrade().ok_or_else(|| { + inner.auth_failures.fetch_add(1, Ordering::Relaxed); + AuthFailure::new( + "temporary_key_inactive", + "temporary key lease is no longer active", + true, + ) + })?; + inner.auth_successes.fetch_add(1, Ordering::Relaxed); + Ok(AuthContext::from_lease(key_id, false, &lease)) + } + + pub fn legacy_protocol_allowed(&self) -> Result { + Ok(self + .inner()? + .legacy_protocol_allowed + .load(Ordering::Acquire)) + } + + pub fn record_legacy_connection(&self) -> Result { + let inner = self.inner()?; + inner + .active_legacy_connections + .fetch_add(1, Ordering::AcqRel); + inner + .last_legacy_connection_at + .store(unix_seconds(), Ordering::Release); + Ok(LegacyConnectionGuard { + inner: Arc::downgrade(&inner), + }) + } + + async fn request( + &self, + build: impl FnOnce(oneshot::Sender>) -> AuthCommand, + ) -> Result { + let (response, receiver) = oneshot::channel(); + self.command_tx.send(build(response)).await.map_err(|_| { + AuthFailure::new( + "auth_state_unavailable", + "authentication state manager is not running", + true, + ) + })?; + receiver.await.map_err(|_| { + AuthFailure::new( + "auth_state_unavailable", + "authentication state manager dropped the response", + true, + ) + })? + } + + pub async fn issue( + &self, + ttl: Duration, + label: Option, + ) -> Result { + self.request(|response| AuthCommand::Issue { + ttl, + label, + response, + }) + .await + } + + pub async fn list(&self, page: u32, page_size: u16) -> Result { + self.request(|response| AuthCommand::List { + page, + page_size, + response, + }) + .await + } + + pub async fn show(&self, key_id: u64, reveal: bool) -> Result { + self.request(|response| AuthCommand::Show { + key_id, + reveal, + response, + }) + .await + } + + pub async fn renew( + &self, + key_id: u64, + ttl: Duration, + ) -> Result { + self.request(|response| AuthCommand::Renew { + key_id, + ttl, + response, + }) + .await + } + + pub async fn revoke(&self, key_id: u64) -> Result { + self.request(|response| AuthCommand::Revoke { key_id, response }) + .await + } + + pub async fn gc(&self) -> Result { + self.request(|response| AuthCommand::Gc { response }).await + } + + pub async fn reset(&self) -> Result<(), AuthFailure> { + self.request(|response| AuthCommand::Reset { response }) + .await + } + + pub async fn rotate_root(&self, new_key: AesKeyType) -> Result<(), AuthFailure> { + self.request(|response| AuthCommand::RotateRoot { new_key, response }) + .await + } + + pub async fn set_legacy_protocol( + &self, + policy: LegacyProtocolPolicy, + ) -> Result<(), AuthFailure> { + self.request(|response| AuthCommand::SetLegacyProtocol { policy, response }) + .await + } + + pub async fn status(&self) -> Result { + self.request(|response| AuthCommand::Status { response }) + .await + } + + pub async fn audit_admin( + &self, + action: impl Into, + key_id: Option, + detail: Option, + ) -> Result<(), AuthFailure> { + let action = action.into(); + self.request(|response| AuthCommand::Audit { + action, + key_id, + detail, + response, + }) + .await + } +} + +pub struct LegacyConnectionGuard { + inner: Weak, +} + +impl Drop for LegacyConnectionGuard { + fn drop(&mut self) { + if let Some(inner) = self.inner.upgrade() { + inner + .active_legacy_connections + .fetch_sub(1, Ordering::AcqRel); + } + } +} + +fn load_server_admin_credential(state_dir: &Path) -> Result { + let path = state_dir.join("admin.key"); + let raw = if path.exists() { + #[cfg(unix)] + { + let metadata = std::fs::metadata(&path).map_err(|error| { + AuthFailure::new( + "administrator_key_required", + format!( + "administrator key file `{}` metadata could not be read: {error}", + path.display() + ), + false, + ) + })?; + if metadata.permissions().mode() & 0o077 != 0 { + std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600)).map_err( + |error| { + AuthFailure::new( + "administrator_key_required", + format!( + "administrator key file `{}` permissions could not be secured: {error}", + path.display() + ), + false, + ) + }, + )?; + tracing::warn!( + event = "administrator_key_permissions_repaired", + path = %path.display(), + "restricted administrator key file permissions to 0600" + ); + } + } + std::fs::read_to_string(&path).map_err(|error| { + AuthFailure::new( + "administrator_key_required", + format!( + "administrator key file `{}` could not be read: {error}", + path.display() + ), + false, + ) + })? + } else if let Ok(credential) = get_process_credential() { + let Credential::Admin(key) = credential else { + return Err(AuthFailure::new( + "administrator_key_required", + "the relay server cannot start with a temporary credential", + false, + )); + }; + let key = String::from_utf8(key.to_vec()).map_err(|_| { + AuthFailure::new( + "administrator_key_invalid", + "the relay administrator key must be printable UTF-8 so it can be persisted", + false, + ) + })?; + write_admin_key(state_dir, &key)?; + key + } else if Path::new(MACHINE_MSG_HEADER_KEY_PATH).is_file() { + let key = std::fs::read_to_string(MACHINE_MSG_HEADER_KEY_PATH).map_err(|error| { + AuthFailure::new( + "administrator_key_required", + format!( + "legacy administrator key file `{MACHINE_MSG_HEADER_KEY_PATH}` could not be read: {error}" + ), + false, + ) + })?; + let Credential::Admin(_) = parse_credential(key.trim()) + .map_err(|error| AuthFailure::new("administrator_key_invalid", error, false))? + else { + return Err(AuthFailure::new( + "administrator_key_required", + "the legacy server key file contains a temporary credential", + false, + )); + }; + write_admin_key(state_dir, key.trim())?; + tracing::warn!( + event = "administrator_key_migrated", + source = MACHINE_MSG_HEADER_KEY_PATH, + destination = %path.display(), + "migrated the legacy administrator key into the v0.4 authentication state directory" + ); + key + } else { + let key = initialize_admin_key(&path, false)?; + tracing::warn!( + event = "administrator_key_initialized", + path = %path.display(), + "no administrator credential was configured; generated a random key file" + ); + key + }; + let credential = parse_credential(raw.trim()) + .map_err(|error| AuthFailure::new("administrator_key_invalid", error, false))?; + if !credential.is_admin() { + return Err(AuthFailure::new( + "administrator_key_required", + "the server key file contains a temporary credential", + false, + )); + } + set_process_msg_header_key(Some(raw.trim())).map_err(AuthFailure::internal)?; + Ok(credential) +} + +pub fn make_key_id(generation: u32, slot: u32) -> u64 { + (u64::from(generation) << 32) | u64::from(slot) +} + +pub fn key_generation(key_id: u64) -> u32 { + (key_id >> 32) as u32 +} + +pub fn key_slot(key_id: u64) -> u32 { + key_id as u32 +} + +pub fn derive_temporary_key( + admin_key: &AesKeyType, + instance_id: &[u8; INSTANCE_ID_LEN], + key_id: u64, +) -> Result { + let salt = Salt::new(HKDF_SHA256, instance_id); + let pseudo_random_key = salt.extract(admin_key); + let key_id_bytes = key_id.to_be_bytes(); + let info = [b"pb-mapper-temp-key-v1".as_slice(), key_id_bytes.as_slice()]; + let output = pseudo_random_key + .expand(&info, HkdfLen(32)) + .map_err(|_| AuthFailure::internal("failed to expand temporary key"))?; + let mut key = [0_u8; 32]; + output + .fill(&mut key) + .map_err(|_| AuthFailure::internal("failed to fill temporary key"))?; + Ok(key) +} + +struct HkdfLen(usize); + +impl ring::hkdf::KeyType for HkdfLen { + fn len(&self) -> usize { + self.0 + } +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +struct PersistedEntry { + key_id: u64, + state: SlotState, + issued_at: u64, + expires_at: u64, + label: Option, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +struct PersistedSnapshot { + schema_version: u16, + instance_id: [u8; INSTANCE_ID_LEN], + generations: Vec, + entries: Vec, + legacy_protocol: LegacyProtocolPolicy, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +enum StateMutation { + Issue(PersistedEntry), + Renew { key_id: u64, expires_at: u64 }, + Revoke { key_id: u64, at: u64 }, + LegacyProtocol(LegacyProtocolPolicy), +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +struct AuditRecord { + at: u64, + action: String, + key_id: Option, + label: Option, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +enum WalRecord { + Mutation { + mutation: StateMutation, + audit: AuditRecord, + }, + Audit(AuditRecord), +} + +async fn run_auth_actor( + inner: Arc, + mut admin_lease: Arc, + mut command_rx: mpsc::Receiver, + config: AuthConfig, + mut cold: HashMap, + mut wheel: TimingWheel, +) { + let now = unix_seconds(); + let mut tombstones = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .iter() + .enumerate() + .filter_map(|(index, slot)| { + matches!(slot.state, SlotState::Expired | SlotState::Revoked).then_some(( + now.saturating_add(TOMBSTONE_RETENTION.as_secs()), + make_key_id(slot.generation, index as u32), + )) + }) + .collect::>(); + let mut last_snapshot_at = unix_seconds(); + let mut tick = tokio::time::interval(Duration::from_secs(1)); + tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + tokio::select! { + _ = tick.tick() => { + let now = unix_seconds(); + for lease in wheel.advance(now) { + let key_id = lease.key_id(); + let version = lease.wheel_version.load(Ordering::Acquire); + if lease.expires_at() > now { + wheel.insert_with_version(lease, version); + continue; + } + let mut slots = inner.slots.write().unwrap_or_else(|poisoned| poisoned.into_inner()); + if let Some(slot) = slots.get_mut(key_slot(key_id) as usize) { + if slot.generation == key_generation(key_id) && slot.state == SlotState::Active { + slot.state = SlotState::Expired; + lease.cancellation.cancel(); + tombstones.push_back((now.saturating_add(TOMBSTONE_RETENTION.as_secs()), key_id)); + tracing::info!( + event = "temporary_key_expired", + auth_stage = "expiry", + key_id, + expires_at = lease.expires_at(), + "temporary key expired and active work was cancelled" + ); + } + } + } + while let Some((cleanup_at, key_id)) = tombstones.front().copied() { + if cleanup_at > now { + break; + } + tombstones.pop_front(); + let mut slots = inner.slots.write().unwrap_or_else(|poisoned| poisoned.into_inner()); + if let Some(slot) = slots.get_mut(key_slot(key_id) as usize) { + if slot.generation == key_generation(key_id) && matches!(slot.state, SlotState::Expired | SlotState::Revoked) { + slot.state = SlotState::Free; + slot.expires_at = 0; + slot.lease = Weak::new(); + cold.remove(&key_id); + } + } + } + if now.saturating_sub(last_snapshot_at) >= SNAPSHOT_COMPACTION_INTERVAL.as_secs() { + let snapshot = build_snapshot(&inner, &cold); + if let Err(error) = write_snapshot_and_truncate_wal( + &config, + &inner.admin_key(), + &snapshot, + ) { + inner.safe_mode.store(true, Ordering::Release); + cancel_all_temporary_leases(&inner); + tracing::error!( + event = "auth_state_safe_mode", + auth_stage = "snapshot_compaction", + reason = %error.code, + error = %error, + "authentication state compaction failed closed" + ); + } else { + last_snapshot_at = now; + } + } + } + command = command_rx.recv() => { + let Some(command) = command else { + admin_lease.cancellation.cancel(); + cancel_all_temporary_leases(&inner); + break; + }; + match command { + AuthCommand::Issue { ttl, label, response } => { + let result = actor_issue(&inner, &config, &mut cold, &mut wheel, ttl, label); + let _ = response.send(result); + } + AuthCommand::List { page, page_size, response } => { + let _ = response.send(actor_list(&inner, &cold, page, page_size)); + } + AuthCommand::Show { key_id, reveal, response } => { + let result = actor_show(&inner, &config, &cold, key_id, reveal); + let _ = response.send(result); + } + AuthCommand::Renew { key_id, ttl, response } => { + let result = actor_renew(&inner, &config, &cold, &mut wheel, key_id, ttl); + let _ = response.send(result); + } + AuthCommand::Revoke { key_id, response } => { + let result = actor_revoke(&inner, &config, &cold, &mut tombstones, key_id); + let _ = response.send(result); + } + AuthCommand::Gc { response } => { + let result = actor_gc(&inner, &config, &mut cold, &mut tombstones); + let _ = response.send(result); + } + AuthCommand::Reset { response } => { + let result = actor_reset(&inner, &config, &mut cold, &mut wheel, "auth_state_reset"); + let _ = response.send(result); + } + AuthCommand::RotateRoot { new_key, response } => { + let result = actor_rotate_root(&inner, &config, &mut cold, &mut wheel, &mut admin_lease, new_key); + let _ = response.send(result); + } + AuthCommand::SetLegacyProtocol { policy, response } => { + let result = actor_set_legacy_protocol(&inner, &config, policy); + let _ = response.send(result); + } + AuthCommand::Status { response } => { + let _ = response.send(Ok(actor_status(&inner))); + } + AuthCommand::Audit { action, key_id, detail, response } => { + let result = append_audit( + &config, + &inner.admin_key(), + audit(&action, key_id, detail), + ); + let _ = response.send(result); + } + } + } + } + } +} + +fn validate_ttl(config: &AuthConfig, ttl: Duration) -> Result { + if ttl < MIN_TEMP_KEY_TTL { + return Err(AuthFailure::new( + "temporary_key_ttl_too_short", + format!( + "temporary key TTL must be at least {} seconds", + MIN_TEMP_KEY_TTL.as_secs() + ), + false, + )); + } + if ttl > config.max_temporary_key_ttl { + return Err(AuthFailure::new( + "temporary_key_ttl_too_long", + format!( + "temporary key TTL exceeds the configured maximum of {} seconds", + config.max_temporary_key_ttl.as_secs() + ), + false, + )); + } + Ok(unix_seconds().saturating_add(ttl.as_secs())) +} + +fn validate_label(label: Option) -> Result, AuthFailure> { + let label = label + .map(|label| label.trim().to_string()) + .filter(|label| !label.is_empty()); + if label.as_ref().is_some_and(|label| label.len() > 64) { + return Err(AuthFailure::new( + "temporary_key_label_too_long", + "temporary key label must not exceed 64 UTF-8 bytes", + false, + )); + } + Ok(label) +} + +fn actor_issue( + inner: &Arc, + config: &AuthConfig, + cold: &mut HashMap, + wheel: &mut TimingWheel, + ttl: Duration, + label: Option, +) -> Result { + ensure_store_available(inner)?; + let expires_at = validate_ttl(config, ttl)?; + let label = validate_label(label)?; + let issued_at = unix_seconds(); + let mut slots = inner + .slots + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let Some((index, slot)) = slots + .iter_mut() + .enumerate() + .find(|(_, slot)| slot.state == SlotState::Free && slot.generation < u32::MAX) + else { + return Err(AuthFailure::new( + "temporary_key_capacity_exhausted", + "temporary key slot table is full", + true, + )); + }; + let generation = slot.generation + 1; + let key_id = make_key_id(generation, index as u32); + let entry = PersistedEntry { + key_id, + state: SlotState::Active, + issued_at, + expires_at, + label: label.clone(), + }; + append_mutation( + config, + &inner.admin_key(), + StateMutation::Issue(entry.clone()), + audit("temporary_key_issue", Some(key_id), label.clone()), + )?; + let lease = Arc::new(AuthLease::new(key_id, expires_at)); + slot.generation = generation; + slot.state = SlotState::Active; + slot.expires_at = expires_at; + slot.lease = Arc::downgrade(&lease); + cold.insert(key_id, ColdMetadata { issued_at, label }); + wheel.insert(lease); + drop(slots); + metadata_with_credential(inner, cold, key_id, true) +} + +fn actor_list( + inner: &Arc, + cold: &HashMap, + page: u32, + page_size: u16, +) -> Result { + let page_size = page_size.clamp(1, 1000) as usize; + let start = (page as usize).saturating_mul(page_size); + let slots = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let mut all = slots + .iter() + .enumerate() + .filter_map(|(index, slot)| { + if slot.state == SlotState::Free { + return None; + } + let key_id = make_key_id(slot.generation, index as u32); + let cold = cold.get(&key_id)?; + Some(TemporaryKeyMetadata { + key_id, + state: slot_state_name(slot.state).to_string(), + issued_at: cold.issued_at, + expires_at: slot.expires_at, + label: cold.label.clone(), + }) + }) + .collect::>(); + all.sort_by_key(|item| std::cmp::Reverse(item.issued_at)); + let items = all.iter().skip(start).take(page_size).cloned().collect(); + let next_page = (start.saturating_add(page_size) < all.len()).then_some(page.saturating_add(1)); + Ok(KeyPage { + schema_version: 1, + items, + next_page, + }) +} + +fn actor_show( + inner: &Arc, + config: &AuthConfig, + cold: &HashMap, + key_id: u64, + reveal: bool, +) -> Result { + let result = metadata_with_credential(inner, cold, key_id, reveal)?; + append_audit( + config, + &inner.admin_key(), + audit( + if reveal { + "temporary_key_reveal" + } else { + "temporary_key_show" + }, + Some(key_id), + result.metadata.label.clone(), + ), + )?; + Ok(result) +} + +fn actor_renew( + inner: &Arc, + config: &AuthConfig, + cold: &HashMap, + wheel: &mut TimingWheel, + key_id: u64, + ttl: Duration, +) -> Result { + ensure_store_available(inner)?; + let expires_at = validate_ttl(config, ttl)?; + let index = key_slot(key_id) as usize; + let mut slots = inner + .slots + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let slot = slots.get_mut(index).ok_or_else(|| key_not_found(key_id))?; + validate_slot_identity(slot, key_id)?; + if slot.state != SlotState::Active || slot.expires_at <= unix_seconds() { + return Err(AuthFailure::new( + "temporary_key_not_renewable", + "only an active, unexpired temporary key can be renewed", + false, + )); + } + let label = cold + .get(&key_id) + .and_then(|metadata| metadata.label.clone()); + append_mutation( + config, + &inner.admin_key(), + StateMutation::Renew { key_id, expires_at }, + audit("temporary_key_renew", Some(key_id), label), + )?; + let lease = slot.lease.upgrade().ok_or_else(|| { + AuthFailure::new( + "temporary_key_inactive", + "temporary key lease is no longer active", + true, + ) + })?; + slot.expires_at = expires_at; + lease.expires_at.store(expires_at, Ordering::Release); + lease.wheel_version.fetch_add(1, Ordering::AcqRel); + wheel.insert(lease); + drop(slots); + metadata_with_credential(inner, cold, key_id, true) +} + +fn actor_revoke( + inner: &Arc, + config: &AuthConfig, + cold: &HashMap, + tombstones: &mut VecDeque<(u64, u64)>, + key_id: u64, +) -> Result { + ensure_store_available(inner)?; + let now = unix_seconds(); + let index = key_slot(key_id) as usize; + let mut slots = inner + .slots + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let slot = slots.get_mut(index).ok_or_else(|| key_not_found(key_id))?; + validate_slot_identity(slot, key_id)?; + if slot.state != SlotState::Active { + return Err(AuthFailure::new( + "temporary_key_not_active", + "temporary key is not active", + false, + )); + } + let cold_metadata = cold.get(&key_id).ok_or_else(|| key_not_found(key_id))?; + append_mutation( + config, + &inner.admin_key(), + StateMutation::Revoke { key_id, at: now }, + audit( + "temporary_key_revoke", + Some(key_id), + cold_metadata.label.clone(), + ), + )?; + slot.state = SlotState::Revoked; + if let Some(lease) = slot.lease.upgrade() { + lease.cancellation.cancel(); + } + tombstones.push_back((now.saturating_add(TOMBSTONE_RETENTION.as_secs()), key_id)); + Ok(TemporaryKeyMetadata { + key_id, + state: slot_state_name(slot.state).to_string(), + issued_at: cold_metadata.issued_at, + expires_at: slot.expires_at, + label: cold_metadata.label.clone(), + }) +} + +fn actor_gc( + inner: &Arc, + config: &AuthConfig, + cold: &mut HashMap, + tombstones: &mut VecDeque<(u64, u64)>, +) -> Result { + ensure_store_available(inner)?; + let now = unix_seconds(); + let mut removed = 0_u64; + let mut slots = inner + .slots + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + for (index, slot) in slots.iter_mut().enumerate() { + if matches!(slot.state, SlotState::Expired | SlotState::Revoked) + || (slot.state == SlotState::Active && slot.expires_at <= now) + { + let key_id = make_key_id(slot.generation, index as u32); + if let Some(lease) = slot.lease.upgrade() { + lease.cancellation.cancel(); + } + slot.state = SlotState::Free; + slot.expires_at = 0; + slot.lease = Weak::new(); + cold.remove(&key_id); + removed = removed.saturating_add(1); + } + } + tombstones.clear(); + drop(slots); + let snapshot = build_snapshot(inner, cold); + let admin_key = inner.admin_key(); + if let Err(error) = + write_snapshot_and_truncate_wal(config, &admin_key, &snapshot).and_then(|()| { + append_audit( + config, + &admin_key, + audit("temporary_key_gc", None, Some(format!("removed={removed}"))), + ) + }) + { + inner.safe_mode.store(true, Ordering::Release); + cancel_all_temporary_leases(inner); + return Err(error); + } + Ok(removed) +} + +fn actor_reset( + inner: &Arc, + config: &AuthConfig, + cold: &mut HashMap, + wheel: &mut TimingWheel, + action: &str, +) -> Result<(), AuthFailure> { + let new_instance_id = random_instance_id(); + let snapshot = empty_snapshot(inner, new_instance_id); + let admin_key = inner.admin_key(); + if let Err(error) = write_snapshot_and_truncate_wal(config, &admin_key, &snapshot) + .and_then(|()| append_audit(config, &admin_key, audit(action, None, None))) + { + inner.safe_mode.store(true, Ordering::Release); + cancel_all_temporary_leases(inner); + return Err(error); + } + if let Err(error) = atomic_write( + &config.state_dir.join("server-instance-id"), + &new_instance_id, + 0o600, + ) { + inner.safe_mode.store(true, Ordering::Release); + cancel_all_temporary_leases(inner); + return Err(error); + } + + cancel_all_temporary_leases(inner); + { + let mut slots = inner + .slots + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + for slot in slots.iter_mut() { + slot.state = SlotState::Free; + slot.expires_at = 0; + slot.lease = Weak::new(); + } + } + *inner + .instance_id + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = new_instance_id; + cold.clear(); + wheel.clear(unix_seconds()); + inner.safe_mode.store(false, Ordering::Release); + Ok(()) +} + +fn actor_rotate_root( + inner: &Arc, + config: &AuthConfig, + cold: &mut HashMap, + wheel: &mut TimingWheel, + admin_lease: &mut Arc, + new_key: AesKeyType, +) -> Result<(), AuthFailure> { + if new_key == inner.admin_key() { + return Err(AuthFailure::new( + "administrator_key_unchanged", + "new administrator key must differ from the current key", + false, + )); + } + let new_key_string = String::from_utf8(new_key.to_vec()).map_err(|_| { + AuthFailure::new( + "administrator_key_invalid", + "administrator key must be 32 UTF-8 bytes for MSG_HEADER_KEY compatibility", + false, + ) + })?; + if new_key_string.chars().any(char::is_whitespace) { + return Err(AuthFailure::new( + "administrator_key_invalid", + "administrator key must not contain whitespace", + false, + )); + } + + let snapshot = empty_snapshot(inner, inner.instance_id()); + if let Err(error) = write_snapshot_and_truncate_wal(config, &new_key, &snapshot) + .and_then(|()| { + append_audit( + config, + &new_key, + audit("administrator_key_rotate", None, None), + ) + }) + .and_then(|()| write_admin_key(&config.state_dir, &new_key_string)) + { + inner.safe_mode.store(true, Ordering::Release); + cancel_all_temporary_leases(inner); + return Err(error); + } + + cancel_all_temporary_leases(inner); + let old_admin_lease = admin_lease.clone(); + { + let mut slots = inner + .slots + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + for slot in slots.iter_mut() { + slot.state = SlotState::Free; + slot.expires_at = 0; + slot.lease = Weak::new(); + } + } + cold.clear(); + wheel.clear(unix_seconds()); + let new_admin_lease = Arc::new(AuthLease::new(0, u64::MAX)); + *inner + .admin_key + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = new_key; + *inner + .admin_lease + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = Arc::downgrade(&new_admin_lease); + set_process_msg_header_key(Some(&new_key_string)).map_err(AuthFailure::internal)?; + inner.safe_mode.store(false, Ordering::Release); + old_admin_lease.cancellation.cancel(); + *admin_lease = new_admin_lease; + Ok(()) +} + +fn actor_set_legacy_protocol( + inner: &Arc, + config: &AuthConfig, + policy: LegacyProtocolPolicy, +) -> Result<(), AuthFailure> { + ensure_store_available(inner)?; + append_mutation( + config, + &inner.admin_key(), + StateMutation::LegacyProtocol(policy), + audit("legacy_protocol_update", None, Some(format!("{policy:?}"))), + )?; + inner + .legacy_protocol_allowed + .store(policy.is_allowed(), Ordering::Release); + Ok(()) +} + +fn actor_status(inner: &Arc) -> AuthStatus { + let slots = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let active_keys = slots + .iter() + .filter(|slot| slot.state == SlotState::Active) + .count(); + let expired_keys = slots + .iter() + .filter(|slot| slot.state == SlotState::Expired) + .count(); + let revoked_keys = slots + .iter() + .filter(|slot| slot.state == SlotState::Revoked) + .count(); + let last_legacy_connection_at = inner.last_legacy_connection_at.load(Ordering::Acquire); + AuthStatus { + schema_version: 1, + safe_mode: inner.safe_mode.load(Ordering::Acquire), + capacity: slots.len(), + active_keys, + expired_keys, + revoked_keys, + legacy_protocol: if inner.legacy_protocol_allowed.load(Ordering::Acquire) { + LegacyProtocolPolicy::Allow + } else { + LegacyProtocolPolicy::Deny + }, + active_legacy_connections: inner.active_legacy_connections.load(Ordering::Acquire), + last_legacy_connection_at: (last_legacy_connection_at != 0) + .then_some(last_legacy_connection_at), + auth_successes: inner.auth_successes.load(Ordering::Relaxed), + auth_failures: inner.auth_failures.load(Ordering::Relaxed), + server_instance_id: hex(&inner.instance_id()), + } +} + +fn ensure_store_available(inner: &AuthStateInner) -> Result<(), AuthFailure> { + if inner.safe_mode.load(Ordering::Acquire) { + Err(AuthFailure::new( + "temporary_key_store_unavailable", + "temporary key store is in administrator safe mode", + false, + )) + } else { + Ok(()) + } +} + +fn validate_slot_identity(slot: &SlotHot, key_id: u64) -> Result<(), AuthFailure> { + if slot.generation != key_generation(key_id) || slot.state == SlotState::Free { + Err(key_not_found(key_id)) + } else { + Ok(()) + } +} + +fn key_not_found(key_id: u64) -> AuthFailure { + AuthFailure::new( + "temporary_key_not_found", + format!("temporary key {key_id} does not exist"), + false, + ) +} + +fn metadata_with_credential( + inner: &Arc, + cold: &HashMap, + key_id: u64, + reveal: bool, +) -> Result { + let slots = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let slot = slots + .get(key_slot(key_id) as usize) + .ok_or_else(|| key_not_found(key_id))?; + validate_slot_identity(slot, key_id)?; + let cold = cold.get(&key_id).ok_or_else(|| key_not_found(key_id))?; + let credential = if reveal { + let key = derive_temporary_key(&inner.admin_key(), &inner.instance_id(), key_id)?; + encode_temporary_credential(key_id, &key) + } else { + String::new() + }; + Ok(IssuedTemporaryKey { + metadata: TemporaryKeyMetadata { + key_id, + state: slot_state_name(slot.state).to_string(), + issued_at: cold.issued_at, + expires_at: slot.expires_at, + label: cold.label.clone(), + }, + credential, + }) +} + +fn slot_state_name(state: SlotState) -> &'static str { + match state { + SlotState::Free => "free", + SlotState::Active => "active", + SlotState::Expired => "expired", + SlotState::Revoked => "revoked", + } +} + +fn audit(action: &str, key_id: Option, label: Option) -> AuditRecord { + AuditRecord { + at: unix_seconds(), + action: action.to_string(), + key_id, + label, + } +} + +fn cancel_all_temporary_leases(inner: &AuthStateInner) { + let slots = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + for lease in slots.iter().filter_map(|slot| slot.lease.upgrade()) { + lease.cancellation.cancel(); + } +} + +fn build_snapshot(inner: &AuthStateInner, cold: &HashMap) -> PersistedSnapshot { + let slots = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let generations = slots.iter().map(|slot| slot.generation).collect(); + let entries = slots + .iter() + .enumerate() + .filter_map(|(index, slot)| { + if slot.state == SlotState::Free { + return None; + } + let key_id = make_key_id(slot.generation, index as u32); + let cold = cold.get(&key_id)?; + Some(PersistedEntry { + key_id, + state: slot.state, + issued_at: cold.issued_at, + expires_at: slot.expires_at, + label: cold.label.clone(), + }) + }) + .collect(); + PersistedSnapshot { + schema_version: SNAPSHOT_SCHEMA_VERSION, + instance_id: inner.instance_id(), + generations, + entries, + legacy_protocol: if inner.legacy_protocol_allowed.load(Ordering::Acquire) { + LegacyProtocolPolicy::Allow + } else { + LegacyProtocolPolicy::Deny + }, + } +} + +fn empty_snapshot(inner: &AuthStateInner, instance_id: [u8; INSTANCE_ID_LEN]) -> PersistedSnapshot { + let slots = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + PersistedSnapshot { + schema_version: SNAPSHOT_SCHEMA_VERSION, + instance_id, + generations: slots.iter().map(|slot| slot.generation).collect(), + entries: Vec::new(), + legacy_protocol: if inner.legacy_protocol_allowed.load(Ordering::Acquire) { + LegacyProtocolPolicy::Allow + } else { + LegacyProtocolPolicy::Deny + }, + } +} + +fn load_persisted_state( + config: &AuthConfig, + admin_key: &AesKeyType, + instance_id: [u8; INSTANCE_ID_LEN], +) -> (Option, bool) { + match try_load_persisted_state(config, admin_key, instance_id) { + Ok(state) => (Some(state), false), + Err(error) => { + tracing::error!( + event = "auth_state_safe_mode", + auth_stage = "state_load", + reason = %error.code, + error = %error, + "temporary key store failed closed in administrator safe mode" + ); + (None, true) + } + } +} + +fn try_load_persisted_state( + config: &AuthConfig, + admin_key: &AesKeyType, + instance_id: [u8; INSTANCE_ID_LEN], +) -> Result { + let snapshot_path = config.state_dir.join("auth.snapshot"); + let mut snapshot = if snapshot_path.exists() { + let bytes = std::fs::read(&snapshot_path).map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to read `{}`: {error}", snapshot_path.display()), + false, + ) + })?; + let plain = open_blob(admin_key, &bytes)?; + serde_json::from_slice::(&plain).map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to decode auth snapshot: {error}"), + false, + ) + })? + } else { + PersistedSnapshot { + schema_version: SNAPSHOT_SCHEMA_VERSION, + instance_id, + generations: vec![0; config.max_temporary_keys], + entries: Vec::new(), + legacy_protocol: config.legacy_protocol, + } + }; + if snapshot.schema_version != SNAPSHOT_SCHEMA_VERSION || snapshot.instance_id != instance_id { + return Err(AuthFailure::new( + "temporary_key_store_unavailable", + "auth snapshot schema or server instance id does not match", + false, + )); + } + snapshot.generations.resize(config.max_temporary_keys, 0); + snapshot.generations.truncate(config.max_temporary_keys); + + let wal_path = config.state_dir.join("auth.wal"); + if wal_path.exists() { + for record in read_wal(&wal_path, admin_key)? { + if let WalRecord::Mutation { mutation, .. } = record { + apply_persisted_mutation(&mut snapshot, mutation, config.max_temporary_keys)?; + } + } + } + Ok(snapshot) +} + +fn apply_persisted_mutation( + snapshot: &mut PersistedSnapshot, + mutation: StateMutation, + capacity: usize, +) -> Result<(), AuthFailure> { + match mutation { + StateMutation::Issue(entry) => { + let index = key_slot(entry.key_id) as usize; + if index >= capacity { + return Err(AuthFailure::new( + "temporary_key_store_unavailable", + "WAL issue record references a slot outside the configured capacity", + false, + )); + } + snapshot.generations[index] = key_generation(entry.key_id); + snapshot + .entries + .retain(|current| key_slot(current.key_id) as usize != index); + snapshot.entries.push(entry); + } + StateMutation::Renew { key_id, expires_at } => { + let entry = snapshot + .entries + .iter_mut() + .find(|entry| entry.key_id == key_id) + .ok_or_else(|| { + AuthFailure::new( + "temporary_key_store_unavailable", + "WAL renew record references an unknown key", + false, + ) + })?; + entry.expires_at = expires_at; + entry.state = SlotState::Active; + } + StateMutation::Revoke { key_id, .. } => { + let entry = snapshot + .entries + .iter_mut() + .find(|entry| entry.key_id == key_id) + .ok_or_else(|| { + AuthFailure::new( + "temporary_key_store_unavailable", + "WAL revoke record references an unknown key", + false, + ) + })?; + entry.state = SlotState::Revoked; + } + StateMutation::LegacyProtocol(policy) => snapshot.legacy_protocol = policy, + } + Ok(()) +} + +fn append_mutation( + config: &AuthConfig, + admin_key: &AesKeyType, + mutation: StateMutation, + audit: AuditRecord, +) -> Result<(), AuthFailure> { + append_wal(config, admin_key, &WalRecord::Mutation { mutation, audit }) +} + +fn append_audit( + config: &AuthConfig, + admin_key: &AesKeyType, + audit: AuditRecord, +) -> Result<(), AuthFailure> { + append_wal(config, admin_key, &WalRecord::Audit(audit)) +} + +fn append_wal( + config: &AuthConfig, + admin_key: &AesKeyType, + record: &WalRecord, +) -> Result<(), AuthFailure> { + let plain = serde_json::to_vec(record).map_err(|error| { + AuthFailure::internal(format!("failed to encode auth WAL record: {error}")) + })?; + let sealed = seal_blob(admin_key, &plain)?; + let length = u32::try_from(sealed.len()) + .map_err(|_| AuthFailure::internal("auth WAL record is too large"))?; + let path = config.state_dir.join("auth.wal"); + let mut file = OpenOptions::new() + .create(true) + .append(true) + .open(&path) + .map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to open `{}`: {error}", path.display()), + true, + ) + })?; + #[cfg(unix)] + file.set_permissions(std::fs::Permissions::from_mode(0o600)) + .map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to secure `{}`: {error}", path.display()), + false, + ) + })?; + file.write_all(&length.to_be_bytes()) + .and_then(|()| file.write_all(&sealed)) + .and_then(|()| file.sync_data()) + .map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to durably append `{}`: {error}", path.display()), + true, + ) + }) +} + +fn read_wal(path: &Path, admin_key: &AesKeyType) -> Result, AuthFailure> { + let mut file = File::open(path).map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to open `{}`: {error}", path.display()), + false, + ) + })?; + let mut records = Vec::new(); + loop { + let mut length = [0_u8; 4]; + match file.read(&mut length[..1]) { + Ok(0) => break, + Ok(1) => {} + Ok(_) => unreachable!("single-byte WAL prefix read"), + Err(error) => { + return Err(AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to read auth WAL length: {error}"), + false, + )); + } + } + file.read_exact(&mut length[1..]).map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("truncated auth WAL length: {error}"), + false, + ) + })?; + let length = u32::from_be_bytes(length) as usize; + if length > 1024 * 1024 { + return Err(AuthFailure::new( + "temporary_key_store_unavailable", + "auth WAL record exceeds 1 MiB", + false, + )); + } + let mut sealed = vec![0_u8; length]; + file.read_exact(&mut sealed).map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("truncated auth WAL record: {error}"), + false, + ) + })?; + let plain = open_blob(admin_key, &sealed)?; + records.push(serde_json::from_slice(&plain).map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to decode auth WAL record: {error}"), + false, + ) + })?); + } + Ok(records) +} + +fn write_snapshot_and_truncate_wal( + config: &AuthConfig, + admin_key: &AesKeyType, + snapshot: &PersistedSnapshot, +) -> Result<(), AuthFailure> { + let plain = serde_json::to_vec(snapshot).map_err(|error| { + AuthFailure::internal(format!("failed to encode auth snapshot: {error}")) + })?; + let sealed = seal_blob(admin_key, &plain)?; + let snapshot_path = config.state_dir.join("auth.snapshot"); + atomic_write(&snapshot_path, &sealed, 0o600)?; + let wal_path = config.state_dir.join("auth.wal"); + let wal = OpenOptions::new() + .create(true) + .write(true) + .truncate(true) + .open(&wal_path) + .map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to truncate `{}`: {error}", wal_path.display()), + true, + ) + })?; + wal.sync_all().map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to sync `{}`: {error}", wal_path.display()), + true, + ) + }) +} + +fn seal_blob(admin_key: &AesKeyType, plain: &[u8]) -> Result, AuthFailure> { + let key = LessSafeKey::new( + UnboundKey::new(&AES_256_GCM, admin_key) + .map_err(|_| AuthFailure::internal("failed to initialize state encryption key"))?, + ); + let mut nonce_bytes = [0_u8; 12]; + let mut rng = rand::rng(); + for byte in &mut nonce_bytes { + *byte = rng.random(); + } + let mut output = plain.to_vec(); + key.seal_in_place_append_tag( + Nonce::assume_unique_for_key(nonce_bytes), + Aad::from(STATE_AAD), + &mut output, + ) + .map_err(|_| AuthFailure::internal("failed to encrypt authentication state"))?; + let mut sealed = Vec::with_capacity(STATE_BLOB_MAGIC.len() + nonce_bytes.len() + output.len()); + sealed.extend_from_slice(STATE_BLOB_MAGIC); + sealed.extend_from_slice(&nonce_bytes); + sealed.extend_from_slice(&output); + Ok(sealed) +} + +fn open_blob(admin_key: &AesKeyType, sealed: &[u8]) -> Result, AuthFailure> { + if sealed.len() < STATE_BLOB_MAGIC.len() + 12 + AES_256_GCM.tag_len() + || &sealed[..STATE_BLOB_MAGIC.len()] != STATE_BLOB_MAGIC + { + return Err(AuthFailure::new( + "temporary_key_store_unavailable", + "authentication state blob has an invalid header", + false, + )); + } + let nonce_start = STATE_BLOB_MAGIC.len(); + let nonce_end = nonce_start + 12; + let nonce_bytes: [u8; 12] = sealed[nonce_start..nonce_end] + .try_into() + .expect("validated nonce width"); + let mut plain = sealed[nonce_end..].to_vec(); + let key = LessSafeKey::new(UnboundKey::new(&AES_256_GCM, admin_key).map_err(|_| { + AuthFailure::new( + "temporary_key_store_unavailable", + "failed to initialize state decryption key", + false, + ) + })?); + let opened = key + .open_in_place( + Nonce::assume_unique_for_key(nonce_bytes), + Aad::from(STATE_AAD), + &mut plain, + ) + .map_err(|_| { + AuthFailure::new( + "temporary_key_store_unavailable", + "authentication state integrity check failed", + false, + ) + })?; + let len = opened.len(); + plain.truncate(len); + Ok(plain) +} + +fn prepare_state_dir(path: &Path) -> Result<(), AuthFailure> { + std::fs::create_dir_all(path).map_err(|error| { + AuthFailure::new( + "auth_state_unavailable", + format!( + "failed to create auth state directory `{}`: {error}", + path.display() + ), + false, + ) + })?; + #[cfg(unix)] + std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o700)).map_err(|error| { + AuthFailure::new( + "auth_state_unavailable", + format!( + "failed to secure auth state directory `{}`: {error}", + path.display() + ), + false, + ) + })?; + Ok(()) +} + +fn load_or_create_instance_id(path: &Path) -> Result<[u8; INSTANCE_ID_LEN], AuthFailure> { + let instance_path = path.join("server-instance-id"); + if instance_path.exists() { + let bytes = std::fs::read(&instance_path).map_err(|error| { + AuthFailure::new( + "auth_state_unavailable", + format!("failed to read `{}`: {error}", instance_path.display()), + false, + ) + })?; + return bytes.try_into().map_err(|_| { + AuthFailure::new( + "auth_state_unavailable", + "server instance id must be exactly 16 bytes", + false, + ) + }); + } + let instance_id = random_instance_id(); + atomic_write(&instance_path, &instance_id, 0o600)?; + Ok(instance_id) +} + +fn random_instance_id() -> [u8; INSTANCE_ID_LEN] { + let mut instance_id = [0_u8; INSTANCE_ID_LEN]; + let mut rng = rand::rng(); + for byte in &mut instance_id { + *byte = rng.random(); + } + instance_id +} + +fn write_admin_key(state_dir: &Path, key: &str) -> Result<(), AuthFailure> { + atomic_write( + &state_dir.join("admin.key"), + format!("{key}\n").as_bytes(), + 0o600, + ) +} + +pub fn generate_admin_key() -> String { + const CHARSET: &[u8] = b"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"; + let mut rng = rand::rng(); + (0..32) + .map(|_| CHARSET[rng.random_range(0..CHARSET.len())] as char) + .collect() +} + +pub fn initialize_admin_key(path: &Path, force: bool) -> Result { + if path.exists() && !force { + return Err(AuthFailure::new( + "administrator_key_exists", + format!("administrator key file `{}` already exists", path.display()), + false, + )); + } + let key = generate_admin_key(); + atomic_write(path, format!("{key}\n").as_bytes(), 0o600)?; + Ok(key) +} + +pub fn write_admin_key_file(path: &Path, key: &str, force: bool) -> Result<(), AuthFailure> { + let Credential::Admin(_) = parse_credential(key) + .map_err(|error| AuthFailure::new("administrator_key_invalid", error, false))? + else { + return Err(AuthFailure::new( + "administrator_key_invalid", + "administrator key file requires a 32-byte administrator key", + false, + )); + }; + if path.exists() && !force { + return Err(AuthFailure::new( + "administrator_key_exists", + format!( + "administrator key file `{}` already exists; pass --force to replace it", + path.display() + ), + false, + )); + } + atomic_write(path, format!("{key}\n").as_bytes(), 0o600) +} + +fn atomic_write(path: &Path, data: &[u8], mode: u32) -> Result<(), AuthFailure> { + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent).map_err(|error| { + AuthFailure::new( + "auth_state_unavailable", + format!("failed to create `{}`: {error}", parent.display()), + false, + ) + })?; + } + let file_name = path + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or("auth-state"); + let mut random_suffix = [0_u8; 8]; + let mut rng = rand::rng(); + for byte in &mut random_suffix { + *byte = rng.random(); + } + let temporary = path.with_file_name(format!( + ".{file_name}.tmp-{}-{}", + std::process::id(), + hex(&random_suffix) + )); + let mut file = OpenOptions::new() + .create_new(true) + .write(true) + .open(&temporary) + .map_err(|error| { + AuthFailure::new( + "auth_state_unavailable", + format!("failed to open `{}`: {error}", temporary.display()), + false, + ) + })?; + let result = (|| { + #[cfg(unix)] + file.set_permissions(std::fs::Permissions::from_mode(mode)) + .map_err(|error| { + AuthFailure::internal(format!("failed to set key permissions: {error}")) + })?; + #[cfg(not(unix))] + let _ = mode; + file.write_all(data) + .and_then(|()| file.sync_all()) + .map_err(|error| { + AuthFailure::new( + "auth_state_unavailable", + format!("failed to write `{}`: {error}", temporary.display()), + false, + ) + })?; + drop(file); + std::fs::rename(&temporary, path).map_err(|error| { + AuthFailure::new( + "auth_state_unavailable", + format!("failed to replace `{}`: {error}", path.display()), + false, + ) + })?; + #[cfg(unix)] + if let Some(parent) = path.parent() { + File::open(parent) + .and_then(|directory| directory.sync_all()) + .map_err(|error| { + AuthFailure::new( + "auth_state_unavailable", + format!("failed to sync `{}`: {error}", parent.display()), + false, + ) + })?; + } + Ok(()) + })(); + if result.is_err() { + let _ = std::fs::remove_file(&temporary); + } + result +} + +fn unix_seconds() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() +} + +fn hex(bytes: &[u8]) -> String { + const DIGITS: &[u8; 16] = b"0123456789abcdef"; + let mut output = String::with_capacity(bytes.len() * 2); + for byte in bytes { + output.push(DIGITS[(byte >> 4) as usize] as char); + output.push(DIGITS[(byte & 0x0f) as usize] as char); + } + output +} + +struct WheelEntry { + lease: Arc, + version: u64, +} + +struct TimingWheel { + now: u64, + level0: Vec>, + level1: Vec>, + level2: Vec>, + level3: Vec>, +} + +impl TimingWheel { + fn new(now: u64) -> Self { + Self { + now, + level0: empty_buckets(256), + level1: empty_buckets(64), + level2: empty_buckets(64), + level3: empty_buckets(64), + } + } + + fn insert(&mut self, lease: Arc) { + let version = lease.wheel_version.load(Ordering::Acquire); + self.insert_with_version(lease, version); + } + + fn insert_with_version(&mut self, lease: Arc, version: u64) { + let expires_at = lease.expires_at(); + let delta = expires_at.saturating_sub(self.now); + let entry = WheelEntry { lease, version }; + if delta < 1 << 8 { + self.level0[(expires_at & 0xff) as usize].push(entry); + } else if delta < 1 << 14 { + self.level1[((expires_at >> 8) & 0x3f) as usize].push(entry); + } else if delta < 1 << 20 { + self.level2[((expires_at >> 14) & 0x3f) as usize].push(entry); + } else { + self.level3[((expires_at >> 20) & 0x3f) as usize].push(entry); + } + } + + fn advance(&mut self, target: u64) -> Vec> { + let mut due = Vec::new(); + while self.now < target { + self.now = self.now.saturating_add(1); + if self.now & 0xff == 0 { + self.cascade(1); + if (self.now >> 8) & 0x3f == 0 { + self.cascade(2); + if (self.now >> 14) & 0x3f == 0 { + self.cascade(3); + } + } + } + let index = (self.now & 0xff) as usize; + for entry in std::mem::take(&mut self.level0[index]) { + if entry.version == entry.lease.wheel_version.load(Ordering::Acquire) { + if entry.lease.expires_at() <= self.now { + due.push(entry.lease); + } else { + self.insert(entry.lease); + } + } + } + } + due + } + + fn cascade(&mut self, level: u8) { + let entries = match level { + 1 => { + let index = ((self.now >> 8) & 0x3f) as usize; + std::mem::take(&mut self.level1[index]) + } + 2 => { + let index = ((self.now >> 14) & 0x3f) as usize; + std::mem::take(&mut self.level2[index]) + } + 3 => { + let index = ((self.now >> 20) & 0x3f) as usize; + std::mem::take(&mut self.level3[index]) + } + _ => Vec::new(), + }; + for entry in entries { + if entry.version == entry.lease.wheel_version.load(Ordering::Acquire) { + self.insert_with_version(entry.lease, entry.version); + } + } + } + + fn clear(&mut self, now: u64) { + *self = Self::new(now); + } +} + +fn empty_buckets(count: usize) -> Vec> { + std::iter::repeat_with(Vec::new).take(count).collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + fn temp_state_dir(name: &str) -> PathBuf { + let mut suffix = [0_u8; 8]; + let mut rng = rand::rng(); + for byte in &mut suffix { + *byte = rng.random(); + } + std::env::temp_dir().join(format!("pb-mapper-{name}-{}", hex(&suffix))) + } + + #[test] + fn key_id_round_trip() { + let key_id = make_key_id(42, 65_535); + assert_eq!(key_generation(key_id), 42); + assert_eq!(key_slot(key_id), 65_535); + } + + #[test] + fn derived_key_is_bound_to_instance_and_key_id() { + let admin = *b"0123456789abcdefghijklmnopqrstuv"; + let instance_a = [1_u8; INSTANCE_ID_LEN]; + let instance_b = [2_u8; INSTANCE_ID_LEN]; + let key = derive_temporary_key(&admin, &instance_a, make_key_id(1, 7)).unwrap(); + assert_eq!( + key, + derive_temporary_key(&admin, &instance_a, make_key_id(1, 7)).unwrap() + ); + assert_ne!( + key, + derive_temporary_key(&admin, &instance_b, make_key_id(1, 7)).unwrap() + ); + assert_ne!( + key, + derive_temporary_key(&admin, &instance_a, make_key_id(2, 7)).unwrap() + ); + } + + #[tokio::test] + async fn issue_renew_revoke_and_persist() { + let state_dir = temp_state_dir("auth-lifecycle"); + let admin = *b"0123456789abcdefghijklmnopqrstuv"; + let config = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 8, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + let runtime = AuthRuntime::start(admin, config.clone()).await.unwrap(); + let issued = runtime + .issue(Duration::from_secs(60), Some("demo".to_string())) + .await + .unwrap(); + assert!(issued.credential.starts_with("pbmt1_")); + let context = runtime.authenticate(issued.metadata.key_id).unwrap(); + assert!(!context.is_admin); + let cancellation = context.cancellation_token().unwrap(); + let renewed = runtime + .renew(issued.metadata.key_id, Duration::from_secs(120)) + .await + .unwrap(); + assert_eq!(renewed.metadata.key_id, issued.metadata.key_id); + assert_eq!(renewed.credential, issued.credential); + assert!(renewed.metadata.expires_at > issued.metadata.expires_at); + runtime.revoke(issued.metadata.key_id).await.unwrap(); + assert!(cancellation.is_cancelled()); + assert_eq!( + context.ensure_active().unwrap_err().code, + "temporary_key_revoked" + ); + assert_eq!( + runtime + .authenticate(issued.metadata.key_id) + .unwrap_err() + .code, + "temporary_key_revoked" + ); + drop(runtime); + + tokio::time::sleep(Duration::from_millis(20)).await; + let restored = AuthRuntime::start(admin, config).await.unwrap(); + assert_eq!( + restored + .authenticate(issued.metadata.key_id) + .unwrap_err() + .code, + "temporary_key_revoked" + ); + let _ = std::fs::remove_dir_all(state_dir); + } + + #[tokio::test] + async fn reset_rotates_instance_and_prevents_old_key_id_reuse() { + let state_dir = temp_state_dir("auth-reset"); + let admin = *b"0123456789abcdefghijklmnopqrstuv"; + let config = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 1, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + let runtime = AuthRuntime::start(admin, config).await.unwrap(); + let before = runtime.status().await.unwrap().server_instance_id; + let old = runtime + .issue(Duration::from_secs(60), Some("before-reset".to_string())) + .await + .unwrap(); + let old_context = runtime.authenticate(old.metadata.key_id).unwrap(); + let old_cancellation = old_context.cancellation_token().unwrap(); + + runtime.reset().await.unwrap(); + + let after = runtime.status().await.unwrap().server_instance_id; + assert_ne!(after, before); + assert!(old_cancellation.is_cancelled()); + assert!(runtime.authenticate(old.metadata.key_id).is_err()); + let replacement = runtime + .issue(Duration::from_secs(60), Some("after-reset".to_string())) + .await + .unwrap(); + assert_ne!(replacement.metadata.key_id, old.metadata.key_id); + assert_ne!(replacement.credential, old.credential); + + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + let _ = std::fs::remove_dir_all(state_dir); + } + + #[tokio::test] + async fn corrupt_wal_fails_temporary_keys_closed_until_admin_reset() { + let state_dir = temp_state_dir("auth-safe-mode"); + let admin = *b"0123456789abcdefghijklmnopqrstuv"; + let config = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 4, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + let runtime = AuthRuntime::start(admin, config.clone()).await.unwrap(); + let issued = runtime + .issue(Duration::from_secs(60), Some("corrupt-me".to_string())) + .await + .unwrap(); + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + std::fs::write(state_dir.join("auth.wal"), b"broken-wal").unwrap(); + + let recovered = AuthRuntime::start(admin, config).await.unwrap(); + assert!(recovered.status().await.unwrap().safe_mode); + assert_eq!( + recovered + .authenticate(issued.metadata.key_id) + .unwrap_err() + .code, + "temporary_key_store_unavailable" + ); + recovered.reset().await.unwrap(); + assert!(!recovered.status().await.unwrap().safe_mode); + + drop(recovered); + tokio::time::sleep(Duration::from_millis(20)).await; + let _ = std::fs::remove_dir_all(state_dir); + } + + #[test] + fn timing_wheel_ignores_stale_renewal_entry() { + let now = 1_000; + let lease = Arc::new(AuthLease::new(make_key_id(1, 0), now + 5)); + let mut wheel = TimingWheel::new(now); + wheel.insert(lease.clone()); + lease.expires_at.store(now + 20, Ordering::Release); + lease.wheel_version.fetch_add(1, Ordering::AcqRel); + wheel.insert(lease.clone()); + assert!(wheel.advance(now + 6).is_empty()); + assert_eq!(wheel.advance(now + 20).len(), 1); + } +} diff --git a/src/common/checksum.rs b/src/common/checksum.rs index 419c1fe..dc07241 100644 --- a/src/common/checksum.rs +++ b/src/common/checksum.rs @@ -6,6 +6,8 @@ use std::process::Command; use std::sync::atomic::{AtomicU32, Ordering}; use std::sync::{LazyLock, RwLock}; +use base64::engine::general_purpose::URL_SAFE_NO_PAD; +use base64::Engine; use rand::RngExt; use ring::digest::{digest, SHA256}; @@ -13,52 +15,75 @@ use super::message::DataLenType; pub type ChecksumType = u32; -const DEFAULT_KEY: &str = "abcdefghijklmnopqlsn123456789j01"; /// Environment variable used by server/client processes to carry the 32-byte header key. pub const ENV_MSG_HEADER_KEY: &str = "MSG_HEADER_KEY"; /// Fixed file path used to persist a machine-derived key for operators to reuse. pub const MACHINE_MSG_HEADER_KEY_PATH: &str = "/var/lib/pb-mapper-server/msg_header_key"; +pub const ADMIN_KEY_PATH: &str = "/var/lib/pb-mapper/auth/admin.key"; +pub const TEMP_CREDENTIAL_PREFIX: &str = "pbmt1_"; const DERIVE_MSG_HEADER_KEY_TAG: &str = "pb-mapper-msg-header-key-v1"; const DERIVE_MSG_HEADER_KEY_CHARSET: &[u8] = b"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"; struct MsgHeaderKeyState { - key: RwLock>, + credential: RwLock>, hash: AtomicU32, } -fn key_len_error(input: &str) -> String { - format!("`{ENV_MSG_HEADER_KEY}` must have 256 bit(32 byte)!. current input key:{input}") +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Credential { + Admin(AesKeyType), + Temporary { key_id: u64, key: AesKeyType }, } -fn load_msg_header_key_from_env_or_default() -> Vec { - let key = match std::env::var(ENV_MSG_HEADER_KEY) { - Ok(k) => { - let key = k.as_bytes(); - if key.len() != 32 { - tracing::warn!("{}", key_len_error(&k)); - std::process::exit(1); - } - key.to_vec() +impl Credential { + pub fn key_id(&self) -> u64 { + match self { + Self::Admin(_) => 0, + Self::Temporary { key_id, .. } => *key_id, + } + } + + pub fn key(&self) -> &AesKeyType { + match self { + Self::Admin(key) | Self::Temporary { key, .. } => key, } - Err(_) => { - tracing::warn!( - "No ENV:`{ENV_MSG_HEADER_KEY}` provided,we use default key:{DEFAULT_KEY}" - ); - DEFAULT_KEY.as_bytes().to_vec() + } + + pub fn is_admin(&self) -> bool { + matches!(self, Self::Admin(_)) + } +} + +fn key_len_error(input: &str) -> String { + format!( + "`{ENV_MSG_HEADER_KEY}` administrator key must be exactly 32 bytes; received {} bytes", + input.len() + ) +} + +fn load_credential_from_env() -> Option { + let raw = std::env::var(ENV_MSG_HEADER_KEY).ok()?; + match parse_credential(raw.trim()) { + Ok(credential) => Some(credential), + Err(error) => { + tracing::error!(reason = "credential_invalid", %error, "invalid MSG_HEADER_KEY"); + None } - }; - key + } } -fn update_runtime_msg_header_key(key: Vec) { - let hash = gen_checksum_by_key(&key); +fn update_runtime_credential(credential: Option) { + let hash = credential + .as_ref() + .map(|credential| gen_checksum_by_key(credential.key())) + .unwrap_or_default(); let mut guard = MSG_HEADER_KEY_STATE - .key + .credential .write() .unwrap_or_else(|poisoned| poisoned.into_inner()); - *guard = key; + *guard = credential; MSG_HEADER_KEY_STATE.hash.store(hash, Ordering::Release); } @@ -67,45 +92,108 @@ fn update_runtime_msg_header_key(key: Vec) { /// This state is mutable so FFI/UI can update `MSG_HEADER_KEY` at runtime /// without restarting the process. static MSG_HEADER_KEY_STATE: LazyLock = LazyLock::new(|| { - let key = load_msg_header_key_from_env_or_default(); - let hash = gen_checksum_by_key(&key); + let credential = load_credential_from_env(); + let hash = credential + .as_ref() + .map(|credential| gen_checksum_by_key(credential.key())) + .unwrap_or_default(); MsgHeaderKeyState { - key: RwLock::new(key), + credential: RwLock::new(credential), hash: AtomicU32::new(hash), } }); -/// Get current message header key bytes. -pub fn get_msg_header_key() -> Vec { +/// Return the configured process credential, failing closed when none exists. +pub fn get_process_credential() -> Result { MSG_HEADER_KEY_STATE - .key + .credential .read() .unwrap_or_else(|poisoned| poisoned.into_inner()) .clone() + .ok_or_else(|| { + format!( + "`{ENV_MSG_HEADER_KEY}` is required; no insecure default credential is available" + ) + }) +} + +/// Get current message header key bytes. +pub fn get_msg_header_key() -> Result, String> { + get_process_credential().map(|credential| credential.key().to_vec()) } /// Set process `MSG_HEADER_KEY` and update runtime checksum/key state. /// -/// - `Some(non-empty)` => validate length 32, set env, apply immediately. -/// - `None` or empty => remove env and reset to default key. +/// - `Some(non-empty)` => validate an admin or temporary credential and apply it immediately. +/// - `None` or empty => remove the credential. Subsequent network operations fail closed. pub fn set_process_msg_header_key(msg_header_key: Option<&str>) -> Result<(), String> { let normalized = msg_header_key.map(str::trim).unwrap_or(""); if normalized.is_empty() { std::env::remove_var(ENV_MSG_HEADER_KEY); - update_runtime_msg_header_key(DEFAULT_KEY.as_bytes().to_vec()); + update_runtime_credential(None); return Ok(()); } - let key = normalized.as_bytes(); - if key.len() != 32 { - return Err(key_len_error(normalized)); - } + let credential = parse_credential(normalized)?; std::env::set_var(ENV_MSG_HEADER_KEY, normalized); - update_runtime_msg_header_key(key.to_vec()); + update_runtime_credential(Some(credential)); Ok(()) } +pub fn parse_credential(raw: &str) -> Result { + if let Some(encoded) = raw.strip_prefix(TEMP_CREDENTIAL_PREFIX) { + let payload = URL_SAFE_NO_PAD + .decode(encoded) + .map_err(|_| "temporary credential is not valid base64url".to_string())?; + if payload.len() != 45 { + return Err(format!( + "temporary credential payload must be 45 bytes, got {}", + payload.len() + )); + } + if payload[0] != 1 { + return Err(format!( + "unsupported temporary credential version {}", + payload[0] + )); + } + let expected = digest(&SHA256, &payload[..41]); + if expected.as_ref()[..4] != payload[41..45] { + return Err("temporary credential checksum mismatch".to_string()); + } + let key_id = u64::from_be_bytes(payload[1..9].try_into().expect("fixed key id width")); + if key_id == 0 { + return Err("temporary credential key id must not be zero".to_string()); + } + let key = payload[9..41] + .try_into() + .expect("fixed temporary key width"); + return Ok(Credential::Temporary { key_id, key }); + } + + let bytes = raw.as_bytes(); + if bytes.len() != 32 { + return Err(key_len_error(raw)); + } + Ok(Credential::Admin( + bytes.try_into().expect("validated admin key width"), + )) +} + +pub fn encode_temporary_credential(key_id: u64, key: &AesKeyType) -> String { + let mut payload = Vec::with_capacity(45); + payload.push(1); + payload.extend_from_slice(&key_id.to_be_bytes()); + payload.extend_from_slice(key); + let checksum = digest(&SHA256, &payload); + payload.extend_from_slice(&checksum.as_ref()[..4]); + format!( + "{TEMP_CREDENTIAL_PREFIX}{}", + URL_SAFE_NO_PAD.encode(payload) + ) +} + /// Derive a stable machine-specific `MSG_HEADER_KEY` and persist it. /// /// The derivation seed is built from normalized hostname + normalized MAC list, @@ -363,7 +451,7 @@ fn write_machine_msg_header_key(key: &str) -> io::Result<()> { ) })?; #[cfg(unix)] - std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o644))?; + std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600))?; Ok(()) } @@ -408,7 +496,10 @@ mod tests { #[test] fn test_random_checksum() { use super::*; - println!("{}", gen_checksum_by_key(DEFAULT_KEY.as_bytes())); + println!( + "{}", + gen_checksum_by_key(b"0123456789abcdefghijklmnopqrstuv") + ); } #[test] @@ -424,4 +515,24 @@ mod tests { assert_eq!(key1.len(), 32); assert!(key1.chars().all(|ch| ch.is_ascii_alphanumeric())); } + + #[test] + fn temporary_credential_round_trip_and_checksum() { + use super::*; + + let key = [7_u8; 32]; + let encoded = encode_temporary_credential(0x0000_0007_0000_002a, &key); + assert_eq!( + parse_credential(&encoded).unwrap(), + Credential::Temporary { + key_id: 0x0000_0007_0000_002a, + key + } + ); + + let mut corrupted = encoded.into_bytes(); + let last = corrupted.last_mut().unwrap(); + *last = if *last == b'A' { b'B' } else { b'A' }; + assert!(parse_credential(std::str::from_utf8(&corrupted).unwrap()).is_err()); + } } diff --git a/src/common/error.rs b/src/common/error.rs index 66bb572..9717721 100644 --- a/src/common/error.rs +++ b/src/common/error.rs @@ -59,6 +59,8 @@ pub enum Error { // specific error explanation detail: String, }, + #[snafu(display("protocol-v2 error: {detail}"))] + MsgProtocol { detail: String }, #[snafu(display("`{action}` forward message failed: {source}"))] MsgForward { // must be "read" or "write" diff --git a/src/common/message/command.rs b/src/common/message/command.rs index e4feef0..5cc6df8 100644 --- a/src/common/message/command.rs +++ b/src/common/message/command.rs @@ -2,6 +2,9 @@ use serde::{Deserialize, Serialize}; use snafu::ResultExt; use super::super::error::{MsgSerializeSnafu, Result}; +use crate::common::auth::{ + AuthStatus, IssuedTemporaryKey, KeyPage, LegacyProtocolPolicy, TemporaryKeyMetadata, +}; use crate::common::checksum::AesKeyType; pub const CONTROL_PROTOCOL_V2: u16 = 2; @@ -58,16 +61,166 @@ pub enum PbConnRequest { #[serde(default, skip_serializing_if = "Option::is_none")] heartbeat_tolerance_ms: Option, }, + RegisterScoped { + need_codec: bool, + is_datagram: bool, + key: String, + namespace: u64, + force_namespace: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + protocol_version: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + client_instance_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + heartbeat_interval_ms: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + heartbeat_tolerance_ms: Option, + }, Subcribe { key: String, }, + SubcribeScoped { + key: String, + namespace: u64, + }, Status(PbConnStatusReq), + StatusScoped { + status: PbConnStatusReq, + namespace: u64, + }, Stream { key: String, dst_id: u32, #[serde(default)] server_generation: u64, }, + StreamScoped { + key: String, + namespace: u64, + dst_id: u32, + #[serde(default)] + server_generation: u64, + }, + Admin(AdminRequest), +} + +#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)] +pub enum AdminRequest { + KeyIssue { + ttl_seconds: u64, + #[serde(default, skip_serializing_if = "Option::is_none")] + label: Option, + }, + KeyList { + #[serde(default)] + page: u32, + #[serde(default = "default_page_size")] + page_size: u16, + }, + KeyShow { + key_id: u64, + }, + KeyReveal { + key_id: u64, + }, + KeyRenew { + key_id: u64, + ttl_seconds: u64, + }, + KeyRevoke { + key_id: u64, + }, + KeyGc, + AuthStatus, + AuthStateReset { + confirm: bool, + }, + RootKeyRotate { + new_admin_key: String, + }, + LegacyProtocolSet { + policy: LegacyProtocolPolicy, + }, + ConnectionList { + #[serde(default, skip_serializing_if = "Option::is_none")] + key_id: Option, + #[serde(default)] + page: u32, + #[serde(default = "default_page_size")] + page_size: u16, + }, + ServiceList { + #[serde(default, skip_serializing_if = "Option::is_none")] + key_id: Option, + #[serde(default)] + page: u32, + #[serde(default = "default_page_size")] + page_size: u16, + }, +} + +const fn default_page_size() -> u16 { + 100 +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct PbErrorResponse { + pub code: String, + pub message: String, + pub retryable: bool, + pub server_time: u64, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct AdminServiceInfo { + pub key_id: u64, + pub namespace: u64, + pub service_name: String, + pub transport: String, + pub codec_enabled: bool, + pub connection_count: u32, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct AdminConnectionInfo { + pub key_id: u64, + pub namespace: u64, + pub service_name: String, + pub conn_id: u32, + pub generation: u64, + pub protocol_version: u16, + pub healthy: bool, + pub transport: String, + pub codec_enabled: bool, + pub last_rx_age_ms: u64, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct AdminServicePage { + pub schema_version: u16, + pub items: Vec, + pub next_page: Option, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct AdminConnectionPage { + pub schema_version: u16, + pub items: Vec, + pub next_page: Option, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub enum AdminResponse { + KeyIssued(IssuedTemporaryKey), + KeyList(KeyPage), + KeyShown(IssuedTemporaryKey), + KeyRenewed(IssuedTemporaryKey), + KeyRevoked(TemporaryKeyMetadata), + KeyGc { removed: u64 }, + AuthStatus(AuthStatus), + Services(AdminServicePage), + Connections(AdminConnectionPage), + Ok { action: String }, } #[derive(Debug, Clone, Deserialize, Serialize)] @@ -87,6 +240,22 @@ pub enum PbConnResponse { codec_key: Option, }, Status(PbConnStatusResp), + Admin(AdminResponse), + Error(PbErrorResponse), +} + +impl PbConnResponse { + pub fn error(code: impl Into, message: impl Into, retryable: bool) -> Self { + Self::Error(PbErrorResponse { + code: code.into(), + message: message.into(), + retryable, + server_time: std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs(), + }) + } } #[derive(Debug, Clone, Deserialize, Serialize)] @@ -123,17 +292,6 @@ pub enum LocalServer { }, } -const CONTENT_SHOW_LIMIT_SIZE: usize = 1024; - -#[inline] -fn get_content(raw_content: String) -> String { - if raw_content.len() > CONTENT_SHOW_LIMIT_SIZE { - raw_content[0..CONTENT_SHOW_LIMIT_SIZE].to_string() - } else { - raw_content - } -} - macro_rules! gen_impl_msg_serializer { ($struct_name:ident) => { impl MessageSerializer for $struct_name { @@ -141,7 +299,7 @@ macro_rules! gen_impl_msg_serializer { serde_json::to_vec(self).with_context(|_| MsgSerializeSnafu { action: "encode", struct_name: stringify!($struct_name), - content: get_content(format!("{self:?}")), + content: "payload redacted".to_string(), }) } @@ -149,7 +307,7 @@ macro_rules! gen_impl_msg_serializer { serde_json::from_slice(msg).with_context(|_| MsgSerializeSnafu { action: "decode", struct_name: stringify!($struct_name), - content: get_content(format!("{}", String::from_utf8_lossy(msg))), + content: format!("{}-byte payload redacted", msg.len()), }) } } diff --git a/src/common/message/mod.rs b/src/common/message/mod.rs index 0477ed8..6379532 100644 --- a/src/common/message/mod.rs +++ b/src/common/message/mod.rs @@ -2,6 +2,7 @@ //! messages pub mod command; pub mod forward; +pub mod secure; use snafu::{ensure, ResultExt}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; @@ -281,7 +282,10 @@ pub fn get_header_msg_writer( #[inline] pub fn get_default_encodec() -> Result { - let key = get_msg_header_key(); + let key = get_msg_header_key().map_err(|detail| error::Error::MsgCodec { + action: "load configured credential", + detail, + })?; Aes256GcmEnCodec::try_new(&key).map_err(|e| error::Error::MsgCodec { action: "create default encodec", detail: format!("{e}"), @@ -290,7 +294,10 @@ pub fn get_default_encodec() -> Result { #[inline] pub fn get_default_decodec() -> Result { - let key = get_msg_header_key(); + let key = get_msg_header_key().map_err(|detail| error::Error::MsgCodec { + action: "load configured credential", + detail, + })?; Aes256GcmDeCodec::try_new(&key).map_err(|e| error::Error::MsgCodec { action: "create default decodec", detail: format!("{e}"), diff --git a/src/common/message/secure.rs b/src/common/message/secure.rs new file mode 100644 index 0000000..ec3d1f2 --- /dev/null +++ b/src/common/message/secure.rs @@ -0,0 +1,1063 @@ +//! Protocol-v2 single-flight authentication framing. +//! +//! The first client frame carries a clear-text routing prefix and an authenticated encrypted +//! request. It does not add a handshake or round trip. All following control messages on the +//! same TCP connection use independently derived directional keys and monotonically increasing +//! 64-bit counters. + +use std::sync::{Arc, Mutex}; +use std::time::{SystemTime, UNIX_EPOCH}; + +use rand::RngExt; +use ring::aead::{Aad, LessSafeKey, Nonce, UnboundKey, AES_256_GCM}; +use ring::digest::{digest, SHA256}; +use ring::hkdf::{Salt, HKDF_SHA256}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; + +use super::{ + CodecMessageReader, CodecMessageWriter, DataLenType, MessageReader, MessageWriter, MAX_MSG_LEN, +}; +use crate::common::auth::{AuthContext, AuthFailure, AuthRuntime, LegacyConnectionGuard}; +use crate::common::checksum::{get_process_credential, valid_checksum, AesKeyType, Credential}; +use crate::common::error::{Error, Result}; +use crate::utils::codec::{Aes256GcmDeCodec, Aes256GcmEnCodec, Decryptor}; + +pub const PROTOCOL_V2_MAGIC: [u8; 4] = *b"PBM2"; +pub const PROTOCOL_V2_VERSION: u8 = 2; +const CONNECTION_SALT_LEN: usize = 16; +const FIRST_PREFIX_REMAINDER_LEN: usize = 28; +const FRAME_HEADER_LEN: usize = 12; +const DIRECTION_CLIENT_TO_SERVER: u8 = 0; +const DIRECTION_SERVER_TO_CLIENT: u8 = 1; +const DEFAULT_REPLAY_WINDOW_SECONDS: u64 = 60; +const DEFAULT_REPLAY_FILTER_BYTES: usize = 1024 * 1024; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum HeaderProtocol { + Legacy, + V2, +} + +#[derive(Clone)] +struct V2Material { + key_id: u64, + flags: u8, + salt: [u8; CONNECTION_SALT_LEN], + client_to_server: AesKeyType, + server_to_client: AesKeyType, +} + +pub struct ClientHeaderSession { + protocol: HeaderProtocol, + legacy_key: AesKeyType, + v2: Option, +} + +impl ClientHeaderSession { + /// New clients always use protocol v2, for both administrator and temporary credentials. + pub fn from_process() -> Result { + let credential = get_process_credential().map_err(protocol_error)?; + Self::new_v2(&credential) + } + + pub fn new_v2(credential: &Credential) -> Result { + let mut salt = [0_u8; CONNECTION_SALT_LEN]; + let mut rng = rand::rng(); + for byte in &mut salt { + *byte = rng.random(); + } + let material = derive_material(credential.key_id(), credential.key(), salt)?; + Ok(Self { + protocol: HeaderProtocol::V2, + legacy_key: *credential.key(), + v2: Some(material), + }) + } + + #[cfg(test)] + pub fn new_legacy(key: AesKeyType) -> Self { + Self { + protocol: HeaderProtocol::Legacy, + legacy_key: key, + v2: None, + } + } + + pub fn protocol(&self) -> HeaderProtocol { + self.protocol + } + + pub async fn write_initial( + &self, + writer: &mut T, + message: &[u8], + ) -> Result<()> { + match self.protocol { + HeaderProtocol::Legacy => { + let codec = Aes256GcmEnCodec::try_new(&self.legacy_key) + .map_err(|_| protocol_error("failed to initialize legacy writer"))?; + CodecMessageWriter::new(writer, codec) + .write_msg(message) + .await + } + HeaderProtocol::V2 => { + let material = self.v2.as_ref().expect("v2 session material"); + writer + .write_all(&first_prefix(material)) + .await + .map_err(|error| { + protocol_error(format!("failed to write v2 prefix: {error}")) + })?; + V2MessageWriter::new(writer, material.clone(), DIRECTION_CLIENT_TO_SERVER, 0)? + .write_msg(message) + .await + } + } + } + + pub fn response_reader<'a, T: AsyncReadExt + Unpin>( + &self, + reader: &'a mut T, + ) -> Result> { + match self.protocol { + HeaderProtocol::Legacy => Ok(HeaderMessageReader::Legacy(CodecMessageReader::new( + reader, + Aes256GcmDeCodec::try_new(&self.legacy_key) + .map_err(|_| protocol_error("failed to initialize legacy reader"))?, + ))), + HeaderProtocol::V2 => Ok(HeaderMessageReader::V2(V2MessageReader::new( + reader, + self.v2.as_ref().expect("v2 session material").clone(), + DIRECTION_SERVER_TO_CLIENT, + 0, + )?)), + } + } + + pub fn continuation_writer<'a, T: AsyncWriteExt + Unpin>( + &self, + writer: &'a mut T, + ) -> Result> { + match self.protocol { + HeaderProtocol::Legacy => Ok(HeaderMessageWriter::Legacy(CodecMessageWriter::new( + writer, + Aes256GcmEnCodec::try_new(&self.legacy_key) + .map_err(|_| protocol_error("failed to initialize legacy writer"))?, + ))), + HeaderProtocol::V2 => Ok(HeaderMessageWriter::V2(V2MessageWriter::new( + writer, + self.v2.as_ref().expect("v2 session material").clone(), + DIRECTION_CLIENT_TO_SERVER, + 1, + )?)), + } + } +} + +pub struct ServerHeaderSession { + protocol: HeaderProtocol, + legacy_key: AesKeyType, + v2: Option, + context: Option, + _legacy_guard: Option, +} + +impl fmt::Debug for ServerHeaderSession { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ServerHeaderSession") + .field("protocol", &self.protocol) + .field("key_id", &self.key_id()) + .field("authenticated", &self.context.is_some()) + .finish() + } +} + +impl ServerHeaderSession { + pub fn protocol(&self) -> HeaderProtocol { + self.protocol + } + + pub fn key_id(&self) -> u64 { + self.context + .as_ref() + .map(|context| context.key_id) + .unwrap_or_else(|| { + self.v2 + .as_ref() + .map(|material| material.key_id) + .unwrap_or_default() + }) + } + + pub fn context(&self) -> Result<&AuthContext> { + self.context + .as_ref() + .ok_or_else(|| protocol_error("server session was not authenticated")) + } + + pub fn take_context(&mut self) -> Result { + self.context + .take() + .ok_or_else(|| protocol_error("server session was not authenticated")) + } + + pub fn response_writer<'a, T: AsyncWriteExt + Unpin>( + &self, + writer: &'a mut T, + ) -> Result> { + match self.protocol { + HeaderProtocol::Legacy => Ok(HeaderMessageWriter::Legacy(CodecMessageWriter::new( + writer, + Aes256GcmEnCodec::try_new(&self.legacy_key) + .map_err(|_| protocol_error("failed to initialize legacy response writer"))?, + ))), + HeaderProtocol::V2 => Ok(HeaderMessageWriter::V2(V2MessageWriter::new( + writer, + self.v2.as_ref().expect("v2 session material").clone(), + DIRECTION_SERVER_TO_CLIENT, + 0, + )?)), + } + } + + pub fn continuation_reader<'a, T: AsyncReadExt + Unpin>( + &self, + reader: &'a mut T, + ) -> Result> { + match self.protocol { + HeaderProtocol::Legacy => Ok(HeaderMessageReader::Legacy(CodecMessageReader::new( + reader, + Aes256GcmDeCodec::try_new(&self.legacy_key) + .map_err(|_| protocol_error("failed to initialize legacy reader"))?, + ))), + HeaderProtocol::V2 => Ok(HeaderMessageReader::V2(V2MessageReader::new( + reader, + self.v2.as_ref().expect("v2 session material").clone(), + DIRECTION_CLIENT_TO_SERVER, + 1, + )?)), + } + } +} + +pub struct ServerInitialMessage { + pub payload: Vec, + pub session: ServerHeaderSession, +} + +pub struct ServerInitialError { + pub failure: AuthFailure, + pub response_session: Option, +} + +impl fmt::Debug for ServerInitialError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ServerInitialError") + .field("failure", &self.failure) + .field("has_response_session", &self.response_session.is_some()) + .finish() + } +} + +impl fmt::Display for ServerInitialError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + self.failure.fmt(formatter) + } +} + +impl std::error::Error for ServerInitialError {} + +use std::fmt; + +#[derive(Clone)] +pub struct ServerSecurity { + auth: AuthRuntime, + replay: Arc>, + failure_logs: Arc>, +} + +impl ServerSecurity { + pub fn new(auth: AuthRuntime) -> Self { + Self { + auth, + replay: Arc::new(Mutex::new(RotatingBloom::new( + DEFAULT_REPLAY_FILTER_BYTES, + DEFAULT_REPLAY_WINDOW_SECONDS, + ))), + failure_logs: Arc::new(Mutex::new(FailureLogLimiter::default())), + } + } + + pub fn auth(&self) -> &AuthRuntime { + &self.auth + } + + pub fn record_failure_log( + &self, + peer_ip: std::net::IpAddr, + key_id: u64, + reason: &str, + ) -> FailureLogDecision { + self.failure_logs + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .record(peer_ip, key_id, reason, unix_seconds()) + } + + pub async fn read_initial( + &self, + reader: &mut T, + ) -> std::result::Result { + let mut first = [0_u8; 4]; + reader + .read_exact(&mut first) + .await + .map_err(|error| ServerInitialError { + failure: AuthFailure::new( + "protocol_header_read_failed", + format!("failed to read initial protocol header: {error}"), + true, + ), + response_session: None, + })?; + if first == PROTOCOL_V2_MAGIC { + self.read_v2_initial(reader).await + } else { + self.read_legacy_initial(reader, first).await + } + } + + async fn read_legacy_initial( + &self, + reader: &mut T, + checksum_bytes: [u8; 4], + ) -> std::result::Result { + if !self.auth.legacy_protocol_allowed().unwrap_or(false) { + return Err(ServerInitialError { + failure: AuthFailure::new( + "legacy_protocol_disabled", + "legacy protocol is disabled by the administrator", + false, + ), + response_session: None, + }); + } + let key = self + .auth + .admin_key() + .map_err(|failure| ServerInitialError { + failure, + response_session: None, + })?; + let context = self + .auth + .authenticate(0) + .map_err(|failure| ServerInitialError { + failure, + response_session: None, + })?; + let checksum = u32::from_be_bytes(checksum_bytes); + let datalen = reader + .read_u32() + .await + .map_err(|error| ServerInitialError { + failure: AuthFailure::new( + "legacy_frame_invalid", + format!("failed to read legacy frame length: {error}"), + true, + ), + response_session: None, + })?; + if !valid_checksum(datalen, checksum) || datalen > MAX_MSG_LEN { + return Err(ServerInitialError { + failure: AuthFailure::new( + "legacy_frame_invalid", + "legacy frame checksum or length is invalid", + false, + ), + response_session: None, + }); + } + let mut encrypted = vec![0_u8; datalen as usize]; + reader + .read_exact(&mut encrypted) + .await + .map_err(|error| ServerInitialError { + failure: AuthFailure::new( + "legacy_frame_invalid", + format!("failed to read legacy frame body: {error}"), + true, + ), + response_session: None, + })?; + let mut codec = Aes256GcmDeCodec::try_new(&key).map_err(|_| ServerInitialError { + failure: AuthFailure::new( + "legacy_decrypt_failed", + "failed to initialize legacy decryption", + false, + ), + response_session: None, + })?; + let plain = codec + .decrypt(&mut encrypted) + .map_err(|_| ServerInitialError { + failure: AuthFailure::new( + "legacy_decrypt_failed", + "legacy credential or encrypted frame is invalid", + false, + ), + response_session: None, + })?; + let legacy_guard = + self.auth + .record_legacy_connection() + .map_err(|failure| ServerInitialError { + failure, + response_session: None, + })?; + Ok(ServerInitialMessage { + payload: plain.to_vec(), + session: ServerHeaderSession { + protocol: HeaderProtocol::Legacy, + legacy_key: key, + v2: None, + context: Some(context), + _legacy_guard: Some(legacy_guard), + }, + }) + } + + async fn read_v2_initial( + &self, + reader: &mut T, + ) -> std::result::Result { + let mut remainder = [0_u8; FIRST_PREFIX_REMAINDER_LEN]; + reader + .read_exact(&mut remainder) + .await + .map_err(|error| ServerInitialError { + failure: AuthFailure::new( + "protocol_v2_header_invalid", + format!("failed to read protocol-v2 header: {error}"), + true, + ), + response_session: None, + })?; + let version = remainder[0]; + let flags = remainder[1]; + let reserved = u16::from_be_bytes([remainder[2], remainder[3]]); + if version != PROTOCOL_V2_VERSION || flags != 0 || reserved != 0 { + return Err(ServerInitialError { + failure: AuthFailure::new( + if version != PROTOCOL_V2_VERSION { + "protocol_version_unsupported" + } else { + "protocol_v2_header_invalid" + }, + format!( + "unsupported protocol header version={version} flags={flags} reserved={reserved}" + ), + false, + ), + response_session: None, + }); + } + let key_id = u64::from_be_bytes(remainder[4..12].try_into().expect("fixed key id")); + let salt: [u8; CONNECTION_SALT_LEN] = + remainder[12..28].try_into().expect("fixed connection salt"); + let key = self + .auth + .derive_key(key_id) + .map_err(|failure| ServerInitialError { + failure, + response_session: None, + })?; + let material = derive_material(key_id, &key, salt).map_err(|error| ServerInitialError { + failure: AuthFailure::new( + "protocol_v2_key_derivation_failed", + error.to_string(), + false, + ), + response_session: None, + })?; + let mut session = ServerHeaderSession { + protocol: HeaderProtocol::V2, + legacy_key: key, + v2: Some(material.clone()), + context: None, + _legacy_guard: None, + }; + let mut message_reader = V2MessageReader::new( + reader, + material, + DIRECTION_CLIENT_TO_SERVER, + 0, + ) + .map_err(|error| ServerInitialError { + failure: AuthFailure::new("protocol_v2_decrypt_failed", error.to_string(), false), + response_session: Some(session_without_context(&session)), + })?; + let payload = message_reader + .read_msg() + .await + .map_err(|error| ServerInitialError { + failure: AuthFailure::new("protocol_v2_decrypt_failed", error.to_string(), false), + response_session: None, + })? + .to_vec(); + + let fingerprint = replay_fingerprint(key_id, &salt); + let replayed = self + .replay + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .contains(&fingerprint, unix_seconds()); + if replayed { + return Err(ServerInitialError { + failure: AuthFailure::new( + "connection_salt_replayed", + "protocol-v2 connection salt was already accepted", + true, + ), + response_session: Some(session), + }); + } + + let context = self + .auth + .authenticate(key_id) + .map_err(|failure| ServerInitialError { + failure, + response_session: Some(session_without_context(&session)), + })?; + self.replay + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .insert(&fingerprint, unix_seconds()); + session.context = Some(context); + Ok(ServerInitialMessage { payload, session }) + } +} + +#[derive(Clone, Copy, Debug)] +pub struct FailureLogDecision { + pub emit: bool, + pub suppressed: u64, +} + +struct FailureLogEntry { + window_started_at: u64, + emitted: u8, + suppressed: u64, +} + +#[derive(Default)] +struct FailureLogLimiter { + entries: std::collections::HashMap<(std::net::IpAddr, u64, String), FailureLogEntry>, + overflow: Option, +} + +impl FailureLogLimiter { + fn record( + &mut self, + peer_ip: std::net::IpAddr, + key_id: u64, + reason: &str, + now: u64, + ) -> FailureLogDecision { + let key = (peer_ip, key_id, reason.to_string()); + if !self.entries.contains_key(&key) && self.entries.len() >= 4096 { + self.entries + .retain(|_, entry| now.saturating_sub(entry.window_started_at) < 120); + if self.entries.len() >= 4096 { + let entry = self.overflow.get_or_insert(FailureLogEntry { + window_started_at: now, + emitted: 0, + suppressed: 0, + }); + return record_failure_entry(entry, now); + } + } + let entry = self.entries.entry(key).or_insert(FailureLogEntry { + window_started_at: now, + emitted: 0, + suppressed: 0, + }); + record_failure_entry(entry, now) + } +} + +fn record_failure_entry(entry: &mut FailureLogEntry, now: u64) -> FailureLogDecision { + if now.saturating_sub(entry.window_started_at) >= 60 { + let suppressed = entry.suppressed; + *entry = FailureLogEntry { + window_started_at: now, + emitted: 1, + suppressed: 0, + }; + return FailureLogDecision { + emit: true, + suppressed, + }; + } + if entry.emitted < 5 { + entry.emitted += 1; + FailureLogDecision { + emit: true, + suppressed: 0, + } + } else { + entry.suppressed = entry.suppressed.saturating_add(1); + FailureLogDecision { + emit: false, + suppressed: 0, + } + } +} + +fn session_without_context(session: &ServerHeaderSession) -> ServerHeaderSession { + ServerHeaderSession { + protocol: session.protocol, + legacy_key: session.legacy_key, + v2: session.v2.clone(), + context: None, + _legacy_guard: None, + } +} + +pub enum HeaderMessageReader<'a, T: AsyncReadExt + Unpin> { + Legacy(CodecMessageReader<'a, T, Aes256GcmDeCodec>), + V2(V2MessageReader<'a, T>), +} + +impl MessageReader for HeaderMessageReader<'_, T> { + async fn read_msg(&mut self) -> Result<&'_ [u8]> { + match self { + Self::Legacy(reader) => reader.read_msg().await, + Self::V2(reader) => reader.read_msg().await, + } + } +} + +pub enum HeaderMessageWriter<'a, T: AsyncWriteExt + Unpin> { + Legacy(CodecMessageWriter<'a, T, Aes256GcmEnCodec>), + V2(V2MessageWriter<'a, T>), +} + +impl MessageWriter for HeaderMessageWriter<'_, T> { + async fn write_msg(&mut self, message: &[u8]) -> Result<()> { + match self { + Self::Legacy(writer) => writer.write_msg(message).await, + Self::V2(writer) => writer.write_msg(message).await, + } + } +} + +pub struct V2MessageReader<'a, T: AsyncReadExt + Unpin> { + reader: &'a mut T, + material: V2Material, + key: LessSafeKey, + direction: u8, + expected_counter: u64, + buffer: Vec, +} + +impl<'a, T: AsyncReadExt + Unpin> V2MessageReader<'a, T> { + fn new( + reader: &'a mut T, + material: V2Material, + direction: u8, + expected_counter: u64, + ) -> Result { + let key_bytes = direction_key(&material, direction); + let key = LessSafeKey::new( + UnboundKey::new(&AES_256_GCM, key_bytes) + .map_err(|_| protocol_error("invalid protocol-v2 read key"))?, + ); + Ok(Self { + reader, + material, + key, + direction, + expected_counter, + buffer: Vec::new(), + }) + } +} + +impl MessageReader for V2MessageReader<'_, T> { + async fn read_msg(&mut self) -> Result<&'_ [u8]> { + let counter = self + .reader + .read_u64() + .await + .map_err(|error| protocol_error(format!("failed to read v2 counter: {error}")))?; + if counter != self.expected_counter { + return Err(protocol_error(format!( + "protocol-v2 counter mismatch: expected {}, got {counter}", + self.expected_counter + ))); + } + let datalen = self + .reader + .read_u32() + .await + .map_err(|error| protocol_error(format!("failed to read v2 length: {error}")))?; + if datalen < AES_256_GCM.tag_len() as u32 || datalen > MAX_MSG_LEN { + return Err(protocol_error(format!( + "protocol-v2 payload length {datalen} is invalid" + ))); + } + self.buffer.resize(datalen as usize, 0); + self.reader + .read_exact(&mut self.buffer) + .await + .map_err(|error| protocol_error(format!("failed to read v2 payload: {error}")))?; + let aad = frame_aad(&self.material, self.direction, counter, datalen); + let plain = self + .key + .open_in_place(nonce(counter), Aad::from(aad.as_slice()), &mut self.buffer) + .map_err(|_| protocol_error("protocol-v2 payload authentication failed"))?; + let plain_len = plain.len(); + self.buffer.truncate(plain_len); + self.expected_counter = self + .expected_counter + .checked_add(1) + .ok_or_else(|| protocol_error("protocol-v2 receive counter exhausted"))?; + Ok(&self.buffer) + } +} + +pub struct V2MessageWriter<'a, T: AsyncWriteExt + Unpin> { + writer: &'a mut T, + material: V2Material, + key: LessSafeKey, + direction: u8, + counter: u64, +} + +impl<'a, T: AsyncWriteExt + Unpin> V2MessageWriter<'a, T> { + fn new(writer: &'a mut T, material: V2Material, direction: u8, counter: u64) -> Result { + let key_bytes = direction_key(&material, direction); + let key = LessSafeKey::new( + UnboundKey::new(&AES_256_GCM, key_bytes) + .map_err(|_| protocol_error("invalid protocol-v2 write key"))?, + ); + Ok(Self { + writer, + material, + key, + direction, + counter, + }) + } +} + +impl MessageWriter for V2MessageWriter<'_, T> { + async fn write_msg(&mut self, message: &[u8]) -> Result<()> { + let encrypted_len = message + .len() + .checked_add(AES_256_GCM.tag_len()) + .and_then(|len| DataLenType::try_from(len).ok()) + .ok_or_else(|| protocol_error("protocol-v2 message is too large"))?; + if encrypted_len > MAX_MSG_LEN { + return Err(protocol_error( + "protocol-v2 message exceeds the maximum length", + )); + } + let counter = self.counter; + let aad = frame_aad(&self.material, self.direction, counter, encrypted_len); + let mut encrypted = message.to_vec(); + self.key + .seal_in_place_append_tag(nonce(counter), Aad::from(aad.as_slice()), &mut encrypted) + .map_err(|_| protocol_error("failed to encrypt protocol-v2 message"))?; + self.writer + .write_u64(counter) + .await + .map_err(|error| protocol_error(format!("failed to write v2 frame header: {error}")))?; + self.writer + .write_u32(encrypted_len) + .await + .map_err(|error| protocol_error(format!("failed to write v2 frame header: {error}")))?; + self.writer + .write_all(&encrypted) + .await + .map_err(|error| protocol_error(format!("failed to write v2 frame body: {error}")))?; + self.counter = self + .counter + .checked_add(1) + .ok_or_else(|| protocol_error("protocol-v2 send counter exhausted"))?; + Ok(()) + } +} + +fn derive_material( + key_id: u64, + credential_key: &AesKeyType, + salt_bytes: [u8; CONNECTION_SALT_LEN], +) -> Result { + let salt = Salt::new(HKDF_SHA256, &salt_bytes); + let pseudo_random_key = salt.extract(credential_key); + let client_to_server = expand_direction(&pseudo_random_key, b"pb-mapper-v2-c2s")?; + let server_to_client = expand_direction(&pseudo_random_key, b"pb-mapper-v2-s2c")?; + Ok(V2Material { + key_id, + flags: 0, + salt: salt_bytes, + client_to_server, + server_to_client, + }) +} + +fn expand_direction( + pseudo_random_key: &ring::hkdf::Prk, + label: &'static [u8], +) -> Result { + let info = [label]; + let output = pseudo_random_key + .expand(&info, HkdfLen(32)) + .map_err(|_| protocol_error("failed to derive protocol-v2 direction key"))?; + let mut key = [0_u8; 32]; + output + .fill(&mut key) + .map_err(|_| protocol_error("failed to fill protocol-v2 direction key"))?; + Ok(key) +} + +struct HkdfLen(usize); + +impl ring::hkdf::KeyType for HkdfLen { + fn len(&self) -> usize { + self.0 + } +} + +fn direction_key(material: &V2Material, direction: u8) -> &AesKeyType { + if direction == DIRECTION_CLIENT_TO_SERVER { + &material.client_to_server + } else { + &material.server_to_client + } +} + +fn first_prefix(material: &V2Material) -> Vec { + let mut prefix = Vec::with_capacity(PROTOCOL_V2_MAGIC.len() + FIRST_PREFIX_REMAINDER_LEN); + prefix.extend_from_slice(&PROTOCOL_V2_MAGIC); + prefix.push(PROTOCOL_V2_VERSION); + prefix.push(material.flags); + prefix.extend_from_slice(&0_u16.to_be_bytes()); + prefix.extend_from_slice(&material.key_id.to_be_bytes()); + prefix.extend_from_slice(&material.salt); + prefix +} + +fn frame_aad(material: &V2Material, direction: u8, counter: u64, datalen: u32) -> Vec { + let mut aad = Vec::with_capacity( + PROTOCOL_V2_MAGIC.len() + FIRST_PREFIX_REMAINDER_LEN + 1 + FRAME_HEADER_LEN, + ); + aad.extend_from_slice(&first_prefix(material)); + aad.push(direction); + aad.extend_from_slice(&counter.to_be_bytes()); + aad.extend_from_slice(&datalen.to_be_bytes()); + aad +} + +fn nonce(counter: u64) -> Nonce { + let mut bytes = [0_u8; 12]; + bytes[4..].copy_from_slice(&counter.to_be_bytes()); + Nonce::assume_unique_for_key(bytes) +} + +fn replay_fingerprint(key_id: u64, salt: &[u8; CONNECTION_SALT_LEN]) -> [u8; 32] { + let mut input = [0_u8; 8 + CONNECTION_SALT_LEN]; + input[..8].copy_from_slice(&key_id.to_be_bytes()); + input[8..].copy_from_slice(salt); + digest(&SHA256, &input) + .as_ref() + .try_into() + .expect("SHA-256 width") +} + +struct RotatingBloom { + current: Vec, + previous: Vec, + current_started_at: u64, + window_seconds: u64, +} + +impl RotatingBloom { + fn new(bytes: usize, window_seconds: u64) -> Self { + Self { + current: vec![0; bytes], + previous: vec![0; bytes], + current_started_at: unix_seconds(), + window_seconds, + } + } + + fn contains(&mut self, fingerprint: &[u8; 32], now: u64) -> bool { + self.rotate(now); + bloom_contains(&self.current, fingerprint) || bloom_contains(&self.previous, fingerprint) + } + + fn insert(&mut self, fingerprint: &[u8; 32], now: u64) { + self.rotate(now); + bloom_insert(&mut self.current, fingerprint); + } + + fn rotate(&mut self, now: u64) { + let elapsed = now.saturating_sub(self.current_started_at); + if elapsed < self.window_seconds { + return; + } + if elapsed >= self.window_seconds.saturating_mul(2) { + self.current.fill(0); + self.previous.fill(0); + } else { + std::mem::swap(&mut self.current, &mut self.previous); + self.current.fill(0); + } + self.current_started_at = now; + } +} + +fn bloom_positions(filter_len: usize, fingerprint: &[u8; 32]) -> [usize; 4] { + let bits = filter_len * 8; + std::array::from_fn(|index| { + let offset = index * 8; + let hash = u64::from_be_bytes( + fingerprint[offset..offset + 8] + .try_into() + .expect("fingerprint chunk"), + ); + hash as usize % bits + }) +} + +fn bloom_contains(filter: &[u8], fingerprint: &[u8; 32]) -> bool { + bloom_positions(filter.len(), fingerprint) + .into_iter() + .all(|position| filter[position / 8] & (1 << (position % 8)) != 0) +} + +fn bloom_insert(filter: &mut [u8], fingerprint: &[u8; 32]) { + for position in bloom_positions(filter.len(), fingerprint) { + filter[position / 8] |= 1 << (position % 8); + } +} + +fn protocol_error(detail: impl Into) -> Error { + Error::MsgProtocol { + detail: detail.into(), + } +} + +fn unix_seconds() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::common::auth::{AuthConfig, LegacyProtocolPolicy}; + use crate::common::checksum::encode_temporary_credential; + + fn temp_config() -> AuthConfig { + let mut random = [0_u8; 8]; + let mut rng = rand::rng(); + for byte in &mut random { + *byte = rng.random(); + } + AuthConfig { + state_dir: std::env::temp_dir() + .join(format!("pb-mapper-v2-{}", u64::from_be_bytes(random))), + max_temporary_keys: 8, + max_temporary_key_ttl: std::time::Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + } + } + + #[tokio::test] + async fn v2_round_trip_uses_directional_counters() { + let credential = Credential::Admin(*b"0123456789abcdefghijklmnopqrstuv"); + let client = ClientHeaderSession::new_v2(&credential).unwrap(); + let config = temp_config(); + let auth = AuthRuntime::start(*credential.key(), config.clone()) + .await + .unwrap(); + let security = ServerSecurity::new(auth); + let (mut client_io, mut server_io) = tokio::io::duplex(4096); + + let client_task = async { + client + .write_initial(&mut client_io, b"request") + .await + .unwrap(); + let mut reader = client.response_reader(&mut client_io).unwrap(); + assert_eq!(reader.read_msg().await.unwrap(), b"response"); + }; + let server_task = async { + let initial = security.read_initial(&mut server_io).await.unwrap(); + assert_eq!(initial.payload, b"request"); + let mut writer = initial.session.response_writer(&mut server_io).unwrap(); + writer.write_msg(b"response").await.unwrap(); + }; + tokio::join!(client_task, server_task); + let _ = std::fs::remove_dir_all(config.state_dir); + } + + #[tokio::test] + async fn temporary_credential_authenticates_without_storing_secret() { + let admin = *b"0123456789abcdefghijklmnopqrstuv"; + let config = temp_config(); + let auth = AuthRuntime::start(admin, config.clone()).await.unwrap(); + let issued = auth + .issue(std::time::Duration::from_secs(60), None) + .await + .unwrap(); + let Credential::Temporary { key_id, key } = + crate::common::checksum::parse_credential(&issued.credential).unwrap() + else { + panic!("expected temporary credential") + }; + assert_eq!(issued.credential, encode_temporary_credential(key_id, &key)); + let client = ClientHeaderSession::new_v2(&Credential::Temporary { key_id, key }).unwrap(); + let security = ServerSecurity::new(auth); + let (mut client_io, mut server_io) = tokio::io::duplex(4096); + let client_task = client.write_initial(&mut client_io, b"temporary"); + let server_task = security.read_initial(&mut server_io); + let (client_result, server_result) = tokio::join!(client_task, server_task); + client_result.unwrap(); + let initial = server_result.unwrap(); + assert_eq!(initial.payload, b"temporary"); + assert_eq!(initial.session.context().unwrap().namespace, key_id); + let _ = std::fs::remove_dir_all(config.state_dir); + } + + #[test] + fn rotating_bloom_covers_current_and_previous_window() { + let mut bloom = RotatingBloom::new(1024, 60); + let value = [7_u8; 32]; + let start = bloom.current_started_at; + assert!(!bloom.contains(&value, start)); + bloom.insert(&value, start); + assert!(bloom.contains(&value, start + 60)); + assert!(!bloom.contains(&value, start + 121)); + } + + #[test] + fn failure_log_limiter_has_a_hard_cardinality_bound() { + let mut limiter = FailureLogLimiter::default(); + let peer = "127.0.0.1".parse().unwrap(); + for key_id in 0..10_000 { + limiter.record(peer, key_id, "invalid", 1_000); + } + assert_eq!(limiter.entries.len(), 4096); + assert!(limiter.overflow.is_some()); + } +} diff --git a/src/common/mod.rs b/src/common/mod.rs index 7dad4ac..4d7513e 100644 --- a/src/common/mod.rs +++ b/src/common/mod.rs @@ -1,3 +1,4 @@ +pub mod auth; pub mod buffer; pub mod checksum; pub mod config; diff --git a/src/local/client/mod.rs b/src/local/client/mod.rs index 697cd40..b46b4c5 100644 --- a/src/local/client/mod.rs +++ b/src/local/client/mod.rs @@ -13,7 +13,7 @@ use tokio::time::MissedTickBehavior; use uni_stream::udp::set_custom_timeout; use self::error::{AcceptLocalStreamSnafu, BindLocalListenerSnafu}; -use self::status::get_status; +use self::status::{get_status, get_status_scoped}; use self::stream::handle_local_stream; use crate::common::config::{ client_health_check_interval, client_health_check_timeout, client_health_failure_threshold, @@ -48,6 +48,26 @@ pub async fn run_client_side_cli( + local_addr: A, + remote_addr: A, + key: Arc, + keep_alive: bool, + namespace: Option, +) where + ::Item: StreamForward, +{ + run_client_side_cli_with_callback_scoped::( + local_addr, + remote_addr, + key, + keep_alive, + namespace, + None, + ) + .await +} + pub async fn run_client_side_cli_with_callback( local_addr: A, remote_addr: A, @@ -56,6 +76,30 @@ pub async fn run_client_side_cli_with_callback, ) where ::Item: StreamForward, +{ + run_client_side_cli_with_callback_scoped::( + local_addr, + remote_addr, + key, + keep_alive, + None, + status_callback, + ) + .await +} + +pub async fn run_client_side_cli_with_callback_scoped< + LocalListener: ListenerProvider, + A: ToSocketAddrs, +>( + local_addr: A, + remote_addr: A, + key: Arc, + keep_alive: bool, + namespace: Option, + status_callback: Option, +) where + ::Item: StreamForward, { set_custom_timeout(Duration::from_secs(120)); @@ -92,7 +136,7 @@ pub async fn run_client_side_cli_with_callback { - if let Err(reason) = probe_remote_key(remote_addr, key.as_ref()).await { + if let Err(reason) = probe_remote_key(remote_addr, key.as_ref(), namespace).await { consecutive_health_failures = consecutive_health_failures.saturating_add(1); if consecutive_health_failures < health_failure_threshold { tracing::warn!( @@ -234,7 +278,7 @@ pub async fn run_client_side_cli_with_callback std::result::Result<(), String> { +async fn probe_remote_key( + remote_addr: SocketAddr, + key: &str, + namespace: Option, +) -> std::result::Result<(), String> { let timeout = client_health_check_timeout(); - match tokio::time::timeout(timeout, probe_remote_key_once(remote_addr, key)).await { + match tokio::time::timeout(timeout, probe_remote_key_once(remote_addr, key, namespace)).await { Ok(result) => result, Err(_) => Err(format!("remote key probe timed out after {timeout:?}")), } @@ -279,12 +327,14 @@ async fn probe_remote_key(remote_addr: SocketAddr, key: &str) -> std::result::Re async fn probe_remote_key_once( remote_addr: SocketAddr, key: &str, + namespace: Option, ) -> std::result::Result<(), String> { match fetch_remote_status( remote_addr, PbConnStatusReq::Service { key: key.to_string(), }, + namespace, ) .await { @@ -312,7 +362,7 @@ async fn probe_remote_key_once( } } - let status_resp = fetch_remote_status(remote_addr, PbConnStatusReq::Keys).await?; + let status_resp = fetch_remote_status(remote_addr, PbConnStatusReq::Keys, namespace).await?; let PbConnStatusResp::Keys(keys) = status_resp else { return Err(format!( "expected keys status response, got {status_resp:?}" @@ -330,11 +380,12 @@ async fn probe_remote_key_once( async fn fetch_remote_status( remote_addr: SocketAddr, req: PbConnStatusReq, + namespace: Option, ) -> std::result::Result { let mut stream = each_addr(remote_addr, TcpStream::connect) .await .map_err(|e| format!("connect remote stream failed: {e}"))?; - get_status(&mut stream, req) + get_status_scoped(&mut stream, req, namespace) .await .map_err(|e| format!("get status failed: {}", snafu::Report::from_error(e))) } @@ -356,9 +407,31 @@ pub async fn show_status( pub async fn handle_status_cli( op: StatusOp, addr: A, +) { + handle_status_cli_scoped(op, addr, None).await +} + +pub async fn handle_status_cli_scoped( + op: StatusOp, + addr: A, + namespace: Option, ) { match op { - StatusOp::RemoteId => show_status(addr, PbConnStatusReq::RemoteId).await, - StatusOp::Keys => show_status(addr, PbConnStatusReq::Keys).await, + StatusOp::RemoteId => show_status_scoped(addr, PbConnStatusReq::RemoteId, namespace).await, + StatusOp::Keys => show_status_scoped(addr, PbConnStatusReq::Keys, namespace).await, } } + +pub async fn show_status_scoped( + remote_addr: A, + req: PbConnStatusReq, + namespace: Option, +) { + let mut stream = snafu_error_get_or_return!( + each_addr(remote_addr, TcpStream::connect).await, + "get status stream" + ); + let status = snafu_error_get_or_return!(get_status_scoped(&mut stream, req, namespace).await); + let status = snafu_error_get_or_return!(serde_json::to_string_pretty(&status)); + println!("Status:{status}"); +} diff --git a/src/local/client/status.rs b/src/local/client/status.rs index ea47acc..1902889 100644 --- a/src/local/client/status.rs +++ b/src/local/client/status.rs @@ -9,36 +9,46 @@ use crate::common::config::control_io_timeout; use crate::common::message::command::{ MessageSerializer, PbConnRequest, PbConnResponse, PbConnStatusReq, PbConnStatusResp, }; -use crate::common::message::{ - get_header_msg_reader, get_header_msg_writer, MessageReader, MessageWriter, -}; +use crate::common::message::secure::ClientHeaderSession; +use crate::common::message::MessageReader; pub async fn get_status( remote_stream: &mut S, req: PbConnStatusReq, +) -> super::error::Result { + get_status_scoped(remote_stream, req, None).await +} + +pub async fn get_status_scoped( + remote_stream: &mut S, + req: PbConnStatusReq, + namespace: Option, ) -> super::error::Result { let timeout = control_io_timeout(); - let msg = PbConnRequest::Status(req) - .encode() - .context(EncodeStatusReqSnafu)?; + let request = match namespace { + Some(namespace) => PbConnRequest::StatusScoped { + status: req, + namespace, + }, + None => PbConnRequest::Status(req), + }; + let msg = request.encode().context(EncodeStatusReqSnafu)?; - // send status request - { - let mut msg_writer = get_header_msg_writer(remote_stream) - .context(CreateHeaderToolSnafu { action: "writer" })?; - match tokio::time::timeout(timeout, msg_writer.write_msg(&msg)).await { - Ok(result) => result.context(WriteStatusReqSnafu)?, - Err(_) => ControlIoTimeoutSnafu { - action: "write status request", - timeout, - } - .fail()?, + let session = + ClientHeaderSession::from_process().context(CreateHeaderToolSnafu { action: "session" })?; + match tokio::time::timeout(timeout, session.write_initial(remote_stream, &msg)).await { + Ok(result) => result.context(WriteStatusReqSnafu)?, + Err(_) => ControlIoTimeoutSnafu { + action: "write status request", + timeout, } + .fail()?, } // get status - let mut msg_reader = - get_header_msg_reader(remote_stream).context(CreateHeaderToolSnafu { action: "reader" })?; + let mut msg_reader = session + .response_reader(remote_stream) + .context(CreateHeaderToolSnafu { action: "reader" })?; let msg = match tokio::time::timeout(timeout, msg_reader.read_msg()).await { Ok(result) => result.context(ReadStatusRespSnafu)?, Err(_) => ControlIoTimeoutSnafu { @@ -50,8 +60,12 @@ pub async fn get_status( let resp = PbConnResponse::decode(msg).context(DecodeStatusRespSnafu)?; match resp { PbConnResponse::Status(status) => Ok(status), - _ => StatusRespNotMatchSnafu { - resp: String::from_utf8_lossy(msg), + PbConnResponse::Error(error) => StatusRespNotMatchSnafu { + resp: format!("{}: {}", error.code, error.message), + } + .fail(), + other => StatusRespNotMatchSnafu { + resp: format!("{other:?}"), } .fail(), } diff --git a/src/local/client/stream.rs b/src/local/client/stream.rs index e978536..cf9781c 100644 --- a/src/local/client/stream.rs +++ b/src/local/client/stream.rs @@ -13,9 +13,8 @@ use super::error::{ use crate::common::config::control_io_timeout; use crate::common::message::command::{MessageSerializer, PbConnRequest, PbConnResponse}; use crate::common::message::forward::StreamForward; -use crate::common::message::{ - get_header_msg_reader, get_header_msg_writer, MessageReader, MessageWriter, -}; +use crate::common::message::secure::ClientHeaderSession; +use crate::common::message::MessageReader; use crate::local::client::error::CreateHeaderToolSnafu; use crate::snafu_error_handle; use uni_stream::addr::{each_addr, ToSocketAddrs}; @@ -30,6 +29,7 @@ pub async fn handle_local_stream< key: Arc, remote_addr: A, keep_alive: bool, + namespace: Option, ) -> Result<()> { let mut remote_stream = each_addr(remote_addr, TcpStream::connect) .await @@ -47,14 +47,19 @@ pub async fn handle_local_stream< let (codec_key, client_id, server_id) = { let timeout = control_io_timeout(); // handle request - let msg = PbConnRequest::Subcribe { - key: key.to_string(), - } - .encode() - .context(EncodeSubcribeReqSnafu)?; - let mut msg_writer = get_header_msg_writer(&mut remote_stream) - .context(CreateHeaderToolSnafu { action: "writer" })?; - match tokio::time::timeout(timeout, msg_writer.write_msg(&msg)).await { + let request = match namespace { + Some(namespace) => PbConnRequest::SubcribeScoped { + key: key.to_string(), + namespace, + }, + None => PbConnRequest::Subcribe { + key: key.to_string(), + }, + }; + let msg = request.encode().context(EncodeSubcribeReqSnafu)?; + let session = ClientHeaderSession::from_process() + .context(CreateHeaderToolSnafu { action: "session" })?; + match tokio::time::timeout(timeout, session.write_initial(&mut remote_stream, &msg)).await { Ok(result) => result.context(WriteSubcribeReqSnafu)?, Err(_) => ControlIoTimeoutSnafu { action: "write subcribe request", @@ -63,7 +68,8 @@ pub async fn handle_local_stream< .fail()?, } // handle response - let mut msg_reader = get_header_msg_reader(&mut remote_stream) + let mut msg_reader = session + .response_reader(&mut remote_stream) .context(CreateHeaderToolSnafu { action: "reader" })?; let msg = match tokio::time::timeout(timeout, msg_reader.read_msg()).await { Ok(result) => result.context(ReadSubcribeRespSnafu)?, @@ -80,6 +86,10 @@ pub async fn handle_local_stream< client_id, server_id, } => (codec_key, client_id, server_id), + PbConnResponse::Error(error) => SubcribeRespNotMatchSnafu { + resp: format!("{}: {}", error.code, error.message), + } + .fail()?, resp => SubcribeRespNotMatchSnafu { resp: format!("{resp:?}"), } diff --git a/src/local/server/mod.rs b/src/local/server/mod.rs index 9fca9df..891272a 100644 --- a/src/local/server/mod.rs +++ b/src/local/server/mod.rs @@ -26,9 +26,8 @@ use crate::common::message::command::{ PbConnStatusResp, PbServerRequest, CONTROL_PROTOCOL_V2, }; use crate::common::message::forward::StreamForward; -use crate::common::message::{ - get_header_msg_reader, get_header_msg_writer, MessageReader, MessageWriter, -}; +use crate::common::message::secure::ClientHeaderSession; +use crate::common::message::{MessageReader, MessageWriter}; use crate::utils::timeout::RetryBackoff; use crate::{ snafu_error_get_or_continue, snafu_error_get_or_return, snafu_error_get_or_return_ok, @@ -131,6 +130,8 @@ pub struct ServerTunnelOptions { pub need_codec: bool, pub is_datagram: bool, pub keep_alive: bool, + pub namespace: Option, + pub force_namespace: bool, } #[derive(Clone, Debug)] @@ -148,6 +149,7 @@ struct StreamTarget { local_addr: A, remote_addr: A, keep_alive: bool, + namespace: Option, } fn duration_to_millis(duration: Duration) -> u64 { @@ -166,17 +168,19 @@ async fn probe_remote_registration( remote_addr: SocketAddr, key: Arc, registration: ControlRegistration, + namespace: Option, ) -> RegistrationProbeResult { let timeout = registration_probe_timeout(); let result = tokio::time::timeout(timeout, async { let mut stream = each_addr(remote_addr, TcpStream::connect) .await .map_err(|e| format!("connect remote status stream failed: {e}"))?; - crate::local::client::status::get_status( + crate::local::client::status::get_status_scoped( &mut stream, PbConnStatusReq::Service { key: key.to_string(), }, + namespace, ) .await .map_err(|e| { @@ -392,6 +396,8 @@ where need_codec, is_datagram, keep_alive, + namespace, + force_namespace, }, worker_index, } = config; @@ -437,40 +443,51 @@ where "manager stream set tcp nodelay" ); - // start register server with key - { - let timeout = control_io_timeout(); - let heartbeat_interval = control_heartbeat_interval(); - let heartbeat_tolerance = control_heartbeat_tolerance(); - let msg = snafu_error_get_or_return_ok!(PbConnRequest::Register { + // Start registration with a protocol-v2 first frame. The session is reused for all + // subsequent control messages on this TCP connection. + let session = match ClientHeaderSession::from_process() { + Ok(session) => session, + Err(error) => { + tracing::error!("create manager protocol-v2 session failed: {error}"); + return Err(Status::ConnectRemote); + } + }; + let timeout = control_io_timeout(); + let heartbeat_interval = control_heartbeat_interval(); + let heartbeat_tolerance = control_heartbeat_tolerance(); + let request = match namespace { + Some(namespace) => PbConnRequest::RegisterScoped { key: key.to_string(), + namespace, + force_namespace, need_codec, is_datagram, protocol_version: Some(CONTROL_PROTOCOL_V2), client_instance_id: Some(new_client_instance_id(worker_index)), heartbeat_interval_ms: Some(duration_to_millis(heartbeat_interval)), heartbeat_tolerance_ms: Some(duration_to_millis(heartbeat_tolerance)), + }, + None => PbConnRequest::Register { + key: key.to_string(), + need_codec, + is_datagram, + protocol_version: Some(CONTROL_PROTOCOL_V2), + client_instance_id: Some(new_client_instance_id(worker_index)), + heartbeat_interval_ms: Some(duration_to_millis(heartbeat_interval)), + heartbeat_tolerance_ms: Some(duration_to_millis(heartbeat_tolerance)), + }, + }; + let msg = snafu_error_get_or_return_ok!(request.encode().context(EncodeRegisterReqSnafu)); + match tokio::time::timeout(timeout, session.write_initial(&mut manager_stream, &msg)).await { + Ok(result) => snafu_error_get_or_return_ok!(result.context(SendRegisterReqSnafu)), + Err(_) => snafu_error_get_or_return_ok!(ControlIoTimeoutSnafu { + action: "send register request", + timeout, } - .encode() - .context(EncodeRegisterReqSnafu)); - let mut msg_writer = match get_header_msg_writer(&mut manager_stream) { - Ok(writer) => writer, - Err(e) => { - tracing::error!("create manager header writer failed: {e}"); - return Err(Status::ConnectRemote); - } - }; - match tokio::time::timeout(timeout, msg_writer.write_msg(&msg)).await { - Ok(result) => snafu_error_get_or_return_ok!(result.context(SendRegisterReqSnafu)), - Err(_) => snafu_error_get_or_return_ok!(ControlIoTimeoutSnafu { - action: "send register request", - timeout, - } - .fail()), - } + .fail()), } let (mut reader, mut writer) = manager_stream.into_split(); - let mut msg_reader = match get_header_msg_reader(&mut reader) { + let mut msg_reader = match session.response_reader(&mut reader) { Ok(reader) => reader, Err(e) => { tracing::error!("create manager header reader failed: {e}"); @@ -508,6 +525,16 @@ where protocol_version: 1, lease_ttl_ms: 0, }, + PbConnResponse::Error(error) => { + tracing::error!( + event = "local_server_registration_rejected", + reason = %error.code, + retryable = error.retryable, + message = %error.message, + "pb server rejected service registration" + ); + snafu_error_get_or_return_ok!(RegisterRespNotMatchSnafu {}.fail()) + } _ => snafu_error_get_or_return_ok!(RegisterRespNotMatchSnafu {}.fail()), }; tracing::info!( @@ -536,7 +563,7 @@ where let writer_key = key.clone(); let writer_registration = registration; let mut writer_handle = tokio::spawn(async move { - let mut msg_writer = match get_header_msg_writer(&mut writer) { + let mut msg_writer = match session.continuation_writer(&mut writer) { Ok(writer) => writer, Err(e) => { tracing::error!("create manager header writer failed: {e}"); @@ -621,6 +648,7 @@ where local_addr, remote_addr, keep_alive, + namespace, }, key.clone(), registration.conn_id, @@ -671,7 +699,13 @@ where let probe_tx = probe_tx.clone(); let probe_key = key.clone(); tokio::spawn(async move { - let result = probe_remote_registration(remote_addr, probe_key, registration).await; + let result = probe_remote_registration( + remote_addr, + probe_key, + registration, + namespace, + ) + .await; let _ = probe_tx.send(result); }); } @@ -817,7 +851,8 @@ where key, client_id, server_generation, - target.keep_alive + target.keep_alive, + target.namespace, ) .await ) diff --git a/src/local/server/stream.rs b/src/local/server/stream.rs index 0acdcd7..741b7ca 100644 --- a/src/local/server/stream.rs +++ b/src/local/server/stream.rs @@ -13,9 +13,8 @@ use super::error::{ use crate::common::config::control_io_timeout; use crate::common::message::command::{MessageSerializer, PbConnRequest, PbConnResponse}; use crate::common::message::forward::StreamForward; -use crate::common::message::{ - get_header_msg_reader, get_header_msg_writer, MessageReader, MessageWriter, -}; +use crate::common::message::secure::ClientHeaderSession; +use crate::common::message::MessageReader; use crate::local::server::error::CreateHeaderToolSnafu; use crate::snafu_error_handle; use uni_stream::addr::{each_addr, ToSocketAddrs}; @@ -34,6 +33,7 @@ pub async fn handle_stream< client_id: u32, server_generation: u64, keep_alive: bool, + namespace: Option, ) -> Result<()> where LocalStream::Item: StreamForward, @@ -42,13 +42,20 @@ where let client_id_span = info_span!("client_id", key_ref, client_id); let _enter = client_id_span.enter(); - let msg = PbConnRequest::Stream { - key: key.to_string(), - dst_id: client_id, - server_generation, - } - .encode() - .context(EncodePbConnStreamReqSnafu)?; + let request = match namespace { + Some(namespace) => PbConnRequest::StreamScoped { + key: key.to_string(), + namespace, + dst_id: client_id, + server_generation, + }, + None => PbConnRequest::Stream { + key: key.to_string(), + dst_id: client_id, + server_generation, + }, + }; + let msg = request.encode().context(EncodePbConnStreamReqSnafu)?; let timeout = control_io_timeout(); let mut remote_stream = @@ -70,9 +77,9 @@ where // write stream request and read response let codec_key = { - let mut msg_writer = get_header_msg_writer(&mut remote_stream) - .context(CreateHeaderToolSnafu { action: "writer" })?; - match tokio::time::timeout(timeout, msg_writer.write_msg(&msg)).await { + let session = ClientHeaderSession::from_process() + .context(CreateHeaderToolSnafu { action: "session" })?; + match tokio::time::timeout(timeout, session.write_initial(&mut remote_stream, &msg)).await { Ok(result) => result.context(WritePbConnStreamReqSnafu)?, Err(_) => ControlIoTimeoutSnafu { action: "write pb conn stream request", @@ -80,7 +87,8 @@ where } .fail()?, } - let mut msg_reader = get_header_msg_reader(&mut remote_stream) + let mut msg_reader = session + .response_reader(&mut remote_stream) .context(CreateHeaderToolSnafu { action: "reader" })?; let msg = match tokio::time::timeout(timeout, msg_reader.read_msg()).await { Ok(result) => result.context(ReadPbConnStreamRespSnafu)?, @@ -93,6 +101,10 @@ where let resp = PbConnResponse::decode(msg).context(DecodePbConnStreamRespSnafu)?; match resp { PbConnResponse::Stream { codec_key } => codec_key, + PbConnResponse::Error(error) => PbConnStreamRespNotMatchSnafu { + resp: format!("{}: {}", error.code, error.message), + } + .fail()?, _ => PbConnStreamRespNotMatchSnafu { resp: format!("{resp:?}"), } diff --git a/src/pb_server/admin.rs b/src/pb_server/admin.rs new file mode 100644 index 0000000..964fcef --- /dev/null +++ b/src/pb_server/admin.rs @@ -0,0 +1,220 @@ +use std::time::Duration; + +use tokio::net::TcpStream; + +use super::error::Error; +use super::{ManagerTask, ManagerTaskSender, Result}; +use crate::common::auth::{AuthFailure, AuthRuntime}; +use crate::common::checksum::{parse_credential, Credential}; +use crate::common::conn_id::RemoteConnId; +use crate::common::message::command::{ + AdminRequest, AdminResponse, MessageSerializer, PbConnResponse, +}; +use crate::common::message::secure::ServerHeaderSession; +use crate::common::message::MessageWriter; + +pub async fn handle_admin_request( + request: AdminRequest, + auth: AuthRuntime, + manager: ManagerTaskSender, + conn_id: RemoteConnId, + mut conn: TcpStream, + session: ServerHeaderSession, +) -> Result<()> { + let result = execute(request, auth, manager).await; + let response = match result { + Ok(response) => PbConnResponse::Admin(response), + Err(failure) => { + tracing::warn!( + event = "admin_operation_failed", + auth_stage = "permission_or_state", + conn_id = %conn_id, + reason = %failure.code, + retryable = failure.retryable, + error = %failure.message, + "administrator operation failed" + ); + PbConnResponse::error(failure.code, failure.message, failure.retryable) + } + }; + let message = response.encode().map_err(|error| Error::AdminOperation { + detail: format!("failed to encode response: {error}"), + })?; + let mut writer = session + .response_writer(&mut conn) + .map_err(|error| Error::AdminOperation { + detail: format!("failed to create response writer: {error}"), + })?; + writer + .write_msg(&message) + .await + .map_err(|error| Error::AdminOperation { + detail: format!("failed to write response: {error}"), + }) +} + +async fn execute( + request: AdminRequest, + auth: AuthRuntime, + manager: ManagerTaskSender, +) -> std::result::Result { + match request { + AdminRequest::KeyIssue { ttl_seconds, label } => auth + .issue(Duration::from_secs(ttl_seconds), label) + .await + .map(AdminResponse::KeyIssued), + AdminRequest::KeyList { page, page_size } => { + audit_read( + &auth, + "temporary_key_list", + None, + Some(format!("page={page},page_size={page_size}")), + ) + .await; + auth.list(page, page_size).await.map(AdminResponse::KeyList) + } + AdminRequest::KeyShow { key_id } => { + auth.show(key_id, false).await.map(AdminResponse::KeyShown) + } + AdminRequest::KeyReveal { key_id } => { + auth.show(key_id, true).await.map(AdminResponse::KeyShown) + } + AdminRequest::KeyRenew { + key_id, + ttl_seconds, + } => auth + .renew(key_id, Duration::from_secs(ttl_seconds)) + .await + .map(AdminResponse::KeyRenewed), + AdminRequest::KeyRevoke { key_id } => { + auth.revoke(key_id).await.map(AdminResponse::KeyRevoked) + } + AdminRequest::KeyGc => auth + .gc() + .await + .map(|removed| AdminResponse::KeyGc { removed }), + AdminRequest::AuthStatus => { + audit_read(&auth, "auth_status", None, None).await; + auth.status().await.map(AdminResponse::AuthStatus) + } + AdminRequest::AuthStateReset { confirm } => { + if !confirm { + return Err(AuthFailure::new( + "confirmation_required", + "auth-state reset requires explicit confirmation", + false, + )); + } + auth.reset().await?; + Ok(AdminResponse::Ok { + action: "auth_state_reset".to_string(), + }) + } + AdminRequest::RootKeyRotate { new_admin_key } => { + let Credential::Admin(new_key) = parse_credential(new_admin_key.trim()) + .map_err(|error| AuthFailure::new("administrator_key_invalid", error, false))? + else { + return Err(AuthFailure::new( + "administrator_key_invalid", + "root rotation requires a 32-byte administrator key", + false, + )); + }; + auth.rotate_root(new_key).await?; + Ok(AdminResponse::Ok { + action: "administrator_key_rotated".to_string(), + }) + } + AdminRequest::LegacyProtocolSet { policy } => { + auth.set_legacy_protocol(policy).await?; + Ok(AdminResponse::Ok { + action: "legacy_protocol_updated".to_string(), + }) + } + AdminRequest::ServiceList { + key_id, + page, + page_size, + } => { + audit_read( + &auth, + "service_list", + key_id, + Some(format!("page={page},page_size={page_size}")), + ) + .await; + let (response_sender, receiver) = tokio::sync::oneshot::channel(); + manager + .send(ManagerTask::AdminServiceList { + key_id, + page, + page_size, + response_sender, + }) + .await + .map_err(|_| { + AuthFailure::new( + "server_state_unavailable", + "relay connection manager is unavailable", + true, + ) + })?; + receiver.await.map(AdminResponse::Services).map_err(|_| { + AuthFailure::new( + "server_state_unavailable", + "relay connection manager dropped the service query", + true, + ) + }) + } + AdminRequest::ConnectionList { + key_id, + page, + page_size, + } => { + audit_read( + &auth, + "connection_list", + key_id, + Some(format!("page={page},page_size={page_size}")), + ) + .await; + let (response_sender, receiver) = tokio::sync::oneshot::channel(); + manager + .send(ManagerTask::AdminConnectionList { + key_id, + page, + page_size, + response_sender, + }) + .await + .map_err(|_| { + AuthFailure::new( + "server_state_unavailable", + "relay connection manager is unavailable", + true, + ) + })?; + receiver.await.map(AdminResponse::Connections).map_err(|_| { + AuthFailure::new( + "server_state_unavailable", + "relay connection manager dropped the connection query", + true, + ) + }) + } + } +} + +async fn audit_read(auth: &AuthRuntime, action: &str, key_id: Option, detail: Option) { + if let Err(error) = auth.audit_admin(action, key_id, detail).await { + tracing::warn!( + event = "admin_audit_failed", + auth_stage = "audit", + action, + reason = %error.code, + error = %error.message, + "administrator read operation could not be audited" + ); + } +} diff --git a/src/pb_server/client.rs b/src/pb_server/client.rs index 317412a..faf8e70 100644 --- a/src/pb_server/client.rs +++ b/src/pb_server/client.rs @@ -1,3 +1,4 @@ +use std::sync::Arc; use std::time::Duration; use snafu::ResultExt; @@ -21,7 +22,8 @@ use crate::common::message::forward::{ CodecForwardReader, CodecForwardWriter, NormalDatagramReader, NormalDatagramWriter, NormalForwardReader, NormalForwardWriter, }; -use crate::common::message::{get_decodec, get_encodec, get_header_msg_writer, MessageWriter}; +use crate::common::message::secure::ServerHeaderSession; +use crate::common::message::{get_decodec, get_encodec, MessageWriter}; use crate::pb_server::error::{ ClientConnCreateHeaderToolSnafu, ClientConnEncodeStreamRespSnafu, ClientConnWriteStreamRespSnafu, @@ -137,17 +139,26 @@ const CLIENT_CONN_CONTROL_TIMEOUT: Duration = Duration::from_secs(30); /// 1. Request server stream /// 2. Forward the traffic between client stream and server stream -#[instrument(skip(task_sender, conn))] +#[instrument(skip(task_sender, conn, session))] pub async fn handle_client_conn( key: ImutableKey, conn_id: RemoteConnId, task_sender: ManagerTaskSender, mut conn: TcpStream, + session: ServerHeaderSession, ) -> Result<()> { let prev_time = Instant::now(); let mut guard = ClientConnGuard::new(conn_id, None, task_sender.clone(), key.clone()); - let (mut server_stream, server_id, codec_key, is_datagram) = - match get_server_stream(&mut conn, key.clone(), conn_id, task_sender.clone()).await { + let (mut server_stream, server_session, server_id, codec_key, is_datagram) = + match get_server_stream( + &mut conn, + &session, + key.clone(), + conn_id, + task_sender.clone(), + ) + .await + { Ok(res) => res, Err(e) => { tracing::warn!( @@ -162,8 +173,17 @@ pub async fn handle_client_conn( } }; guard.set_server_id(server_id); + let server_cancellation = server_session + .context() + .map_err(|error| super::error::Error::ClientConnAuthInactive { + detail: error.to_string(), + })? + .cancellation_token() + .map_err(|error| super::error::Error::ClientConnAuthInactive { + detail: error.to_string(), + })?; - let result = async { + let forwarding = async { let duration = Instant::now() - prev_time; tracing::info!( @@ -182,7 +202,8 @@ pub async fn handle_client_conn( // response message to server to indicate that stream handling has finished { - let mut msg_writer = get_header_msg_writer(&mut server_writer) + let mut msg_writer = server_session + .response_writer(&mut server_writer) .context(ClientConnCreateHeaderToolSnafu { tool: "writer" })?; let msg = PbConnResponse::Stream { codec_key }.encode().context( ClientConnEncodeStreamRespSnafu { @@ -247,8 +268,20 @@ pub async fn handle_client_conn( } Ok(()) - } - .await; + }; + let result = tokio::select! { + result = forwarding => result, + _ = server_cancellation.cancelled() => { + tracing::info!( + event = "connection_auth_expired", + key = %key, + client_conn_id = %conn_id, + server_conn_id = %server_id, + "closing active data stream because its registering credential expired or was revoked" + ); + Ok(()) + } + }; match &result { Ok(()) => tracing::info!( event = "client_forward_finished", @@ -297,14 +330,16 @@ async fn retire_server_conn( async fn write_subscribe_response( conn: &mut TcpStream, + session: &ServerHeaderSession, key: &ImutableKey, conn_id: RemoteConnId, server_id: RemoteConnId, codec_key: Option, is_datagram: bool, ) -> Result<()> { - let mut msg_writer = - get_header_msg_writer(conn).context(ClientConnCreateHeaderToolSnafu { tool: "writer" })?; + let mut msg_writer = session + .response_writer(conn) + .context(ClientConnCreateHeaderToolSnafu { tool: "writer" })?; let msg = PbConnResponse::Subcribe { codec_key, client_id: conn_id.into(), @@ -336,10 +371,17 @@ async fn write_subscribe_response( async fn get_server_stream( conn: &mut TcpStream, + session: &ServerHeaderSession, key: ImutableKey, conn_id: RemoteConnId, task_sender: ManagerTaskSender, -) -> Result<(TcpStream, RemoteConnId, Option, bool)> { +) -> Result<( + TcpStream, + ServerHeaderSession, + RemoteConnId, + Option, + bool, +)> { let (tx, rx) = kanal::bounded_async(DEFAULT_CLIENT_CHAN_CAP); let ack_timeout = stream_ack_timeout(); let ready_timeout = stream_ready_timeout(); @@ -407,7 +449,11 @@ async fn get_server_stream( ); (codec_key, is_datagram, server_conn_id, server_generation) } - ConnTask::SubcribeFailed { reason } => { + ConnTask::SubcribeFailed { + code, + reason, + retryable, + } => { tracing::warn!( event = "subscribe_failed", key = %key, @@ -415,6 +461,7 @@ async fn get_server_stream( reason = %reason, "subscribe failed before stream forwarding" ); + write_subscribe_error(conn, session, &code, &reason, retryable).await?; ClientConnSubcribeFailedSnafu { key: key.clone(), conn_id, @@ -492,10 +539,19 @@ async fn get_server_stream( server_id, server_generation: response_generation, stream, + session: stream_session, } if response_generation == server_generation => { - write_subscribe_response(conn, &key, conn_id, server_id, codec_key, is_datagram) - .await?; - return Ok((stream, server_id, codec_key, is_datagram)); + write_subscribe_response( + conn, + session, + &key, + conn_id, + server_id, + codec_key, + is_datagram, + ) + .await?; + return Ok((stream, stream_session, server_id, codec_key, is_datagram)); } ConnTask::StreamAck { server_id, @@ -557,12 +613,21 @@ async fn get_server_stream( server_id, server_generation: response_generation, stream, + session: stream_session, } = resp { if response_generation == server_generation { - write_subscribe_response(conn, &key, conn_id, server_id, codec_key, is_datagram) - .await?; - return Ok((stream, server_id, codec_key, is_datagram)); + write_subscribe_response( + conn, + session, + &key, + conn_id, + server_id, + codec_key, + is_datagram, + ) + .await?; + return Ok((stream, stream_session, server_id, codec_key, is_datagram)); } tracing::warn!( event = "server_stream_generation_mismatch", @@ -585,3 +650,28 @@ async fn get_server_stream( .fail()? } } + +async fn write_subscribe_error( + conn: &mut TcpStream, + session: &ServerHeaderSession, + code: &str, + reason: &str, + retryable: bool, +) -> Result<()> { + let message = PbConnResponse::error(code, reason, retryable) + .encode() + .context(ClientConnEncodeSubcribeRespSnafu { + key: Arc::from(""), + conn_id: RemoteConnId::default(), + })?; + let mut writer = session + .response_writer(conn) + .context(ClientConnCreateHeaderToolSnafu { tool: "writer" })?; + writer + .write_msg(&message) + .await + .context(ClientConnWriteSubcribeRespSnafu { + key: Arc::from(""), + conn_id: RemoteConnId::default(), + }) +} diff --git a/src/pb_server/error.rs b/src/pb_server/error.rs index dc896da..552c68e 100644 --- a/src/pb_server/error.rs +++ b/src/pb_server/error.rs @@ -8,6 +8,8 @@ use crate::common::{self}; #[derive(Debug, Snafu)] #[snafu(visibility(pub(super)))] pub enum Error { + #[snafu(display("administrator operation failed: {detail}"))] + AdminOperation { detail: String }, /// server task center error #[snafu(display("read pb conn init request with `conn_id:{conn_id}`"))] TaskCenterReadInitRequest { @@ -201,6 +203,8 @@ pub enum Error { tool: &'static str, source: common::error::Error, }, + #[snafu(display("client data stream credential is inactive: {detail}"))] + ClientConnAuthInactive { detail: String }, #[snafu(display( "send deregister client task error with `key:{key}` `server:{server_id:?}` <-> \ `client:{client_id}`, type:{source:?} detail:{source}" diff --git a/src/pb_server/mod.rs b/src/pb_server/mod.rs index dd622cd..0abbe23 100644 --- a/src/pb_server/mod.rs +++ b/src/pb_server/mod.rs @@ -1,3 +1,4 @@ +mod admin; mod client; mod error; mod server; @@ -13,22 +14,26 @@ use tokio::net::{TcpListener, TcpStream, ToSocketAddrs}; use tokio_util::sync::CancellationToken; use tracing::instrument; +use self::admin::handle_admin_request; use self::client::handle_client_conn; use self::error::{ TaskCenterDecodeInitRequestSnafu, TaskCenterInitRequestTimeoutSnafu, TaskCenterReadInitRequestSnafu, TaskCenterSendListenerSnafu, TaskCenterSendStatusRespSnafu, TaskCenterSendStreamRespToManagerSnafu, TaskCenterSetKeepAliveSnafu, }; -use self::server::handle_server_conn; +use self::server::{handle_server_conn, ServerRegistration}; use self::status::handle_show_status; +use crate::common::auth::{AuthConfig, AuthContext, AuthRuntime}; use crate::common::config::{control_io_timeout, keep_alive_from_env, server_lease_timeout}; use crate::common::conn_id::{ConnIdProvider, RemoteConnId}; use crate::common::manager::{ForwardMessage, SenderChan, TaskManager}; use crate::common::message::command::{ + AdminConnectionInfo, AdminConnectionPage, AdminServiceInfo, AdminServicePage, MessageSerializer, PbConnRequest, PbConnResponse, PbConnStatusReq, PbConnStatusResp, PbServiceConnStatus, }; -use crate::common::message::{get_header_msg_reader, MessageReader}; +use crate::common::message::secure::{ServerHeaderSession, ServerSecurity}; +use crate::common::message::{get_header_msg_reader, MessageReader, MessageWriter}; use crate::pb_server::error::{ ServerListenSnafu, TaskCenterClientSendStreamSnafu, TaskCenterSendRegisterRespSnafu, TaskCenterSendStreamRespToClientSnafu, TaskCenterSendSubcribeRespSnafu, @@ -61,7 +66,9 @@ pub enum ManagerTask { excluded_server_conns: Vec<(RemoteConnId, u64)>, }, Stream { + key: ImutableKey, stream: TcpStream, + session: ServerHeaderSession, server_id: RemoteConnId, client_id: RemoteConnId, server_generation: u64, @@ -74,11 +81,24 @@ pub enum ManagerTask { Status { conn_sender: ConnTaskSender, status: PbConnStatusReq, + namespace: u64, conn_id: RemoteConnId, }, StatusQuery { response_sender: tokio::sync::oneshot::Sender, }, + AdminServiceList { + key_id: Option, + page: u32, + page_size: u16, + response_sender: tokio::sync::oneshot::Sender, + }, + AdminConnectionList { + key_id: Option, + page: u32, + page_size: u16, + response_sender: tokio::sync::oneshot::Sender, + }, DeRegisterServerConn { key: ImutableKey, conn_id: RemoteConnId, @@ -103,6 +123,11 @@ pub enum ConnTask { protocol_version: u16, lease_ttl_ms: u64, }, + RegisterFailed { + code: String, + reason: String, + retryable: bool, + }, SubcribeResp { server_conn_id: RemoteConnId, server_generation: u64, @@ -110,7 +135,9 @@ pub enum ConnTask { is_datagram: bool, }, SubcribeFailed { + code: String, reason: String, + retryable: bool, }, SubcribeRetry { reason: String, @@ -130,6 +157,7 @@ pub enum ConnTask { server_id: RemoteConnId, server_generation: u64, stream: TcpStream, + session: ServerHeaderSession, }, StatusResp(PbConnResponse), } @@ -158,6 +186,46 @@ pub struct ServerConnInfo { pub type ServerConnMap = hashbrown::HashMap>; +struct NamespaceRateLimit { + tokens: f64, + last_refill: Instant, + rate_per_second: f64, + burst: f64, +} + +impl NamespaceRateLimit { + fn new(rate_per_second: usize, burst: usize) -> Self { + Self { + tokens: burst as f64, + last_refill: Instant::now(), + rate_per_second: rate_per_second as f64, + burst: burst as f64, + } + } + + fn allow(&mut self) -> bool { + let now = Instant::now(); + self.tokens = (self.tokens + + now.duration_since(self.last_refill).as_secs_f64() * self.rate_per_second) + .min(self.burst); + self.last_refill = now; + if self.tokens < 1.0 { + false + } else { + self.tokens -= 1.0; + true + } + } +} + +fn env_limit(name: &str, default: usize) -> usize { + std::env::var(name) + .ok() + .and_then(|value| value.parse::().ok()) + .filter(|value| *value > 0) + .unwrap_or(default) +} + #[derive(Debug, Clone)] pub struct ServerStatusInfo { pub active_connections: u32, @@ -253,7 +321,9 @@ async fn send_subcribe_failed( let reason = reason.into(); if conn_sender .send(ConnTask::SubcribeFailed { + code: "service_not_available".to_string(), reason: reason.clone(), + retryable: true, }) .await .is_err() @@ -333,10 +403,42 @@ pub async fn run_server_with_shutdown( >, keep_alive: bool, ) -> std::io::Result<()> { + run_server_with_auth_config( + addr, + shutdown_token, + status_channel, + keep_alive, + AuthConfig::default(), + ) + .await +} + +pub async fn run_server_with_auth_config( + addr: A, + shutdown_token: CancellationToken, + status_channel: Option< + tokio::sync::mpsc::UnboundedReceiver>, + >, + keep_alive: bool, + auth_config: AuthConfig, +) -> std::io::Result<()> { + let auth = AuthRuntime::from_process(auth_config) + .await + .map_err(|error| std::io::Error::other(error.to_string()))?; + let security = ServerSecurity::new(auth); let mut manager = ServerMananger::new(RemoteIdProvider::new()); // represent the mapping of the `key` to the id of the server-side conn let mut server_conn_map = ServerConnMap::new(); - let mut pending_streams = hashbrown::HashMap::::new(); + let mut pending_streams = + hashbrown::HashMap::::new(); + let mut namespace_stream_counts = hashbrown::HashMap::::new(); + let mut namespace_rate_limits = hashbrown::HashMap::::new(); + let max_services_per_namespace = env_limit("PB_MAPPER_MAX_SERVICES_PER_NAMESPACE", 256); + let max_register_connections_per_service = + env_limit("PB_MAPPER_MAX_REGISTER_CONNECTIONS_PER_SERVICE", 16); + let max_streams_per_namespace = env_limit("PB_MAPPER_MAX_STREAMS_PER_NAMESPACE", 1024); + let new_streams_per_second = env_limit("PB_MAPPER_NEW_STREAMS_PER_SECOND", 100); + let new_streams_burst = env_limit("PB_MAPPER_NEW_STREAMS_BURST", 200); let mut next_server_generation = 1_u64; let listener = TcpListener::bind(addr).await?; @@ -399,6 +501,97 @@ pub async fn run_server_with_shutdown( }; match task { + ManagerTask::AdminServiceList { + key_id, + page, + page_size, + response_sender, + } => { + let page_size = page_size.clamp(1, 1000) as usize; + let start = (page as usize).saturating_mul(page_size); + let mut all = server_conn_map + .iter() + .filter_map(|(key, connections)| { + let (namespace, service_name) = split_scoped_service_key(key); + if key_id.is_some_and(|key_id| key_id != namespace) { + return None; + } + let first = connections.first()?; + Some(AdminServiceInfo { + key_id: namespace, + namespace, + service_name: service_name.to_string(), + transport: if first.is_datagram { "udp" } else { "tcp" }.to_string(), + codec_enabled: first.need_codec, + connection_count: connections.len() as u32, + }) + }) + .collect::>(); + all.sort_by(|left, right| { + left.namespace + .cmp(&right.namespace) + .then_with(|| left.service_name.cmp(&right.service_name)) + }); + let items = all.iter().skip(start).take(page_size).cloned().collect(); + let next_page = + (start.saturating_add(page_size) < all.len()).then_some(page.saturating_add(1)); + let _ = response_sender.send(AdminServicePage { + schema_version: 1, + items, + next_page, + }); + } + ManagerTask::AdminConnectionList { + key_id, + page, + page_size, + response_sender, + } => { + let now = Instant::now(); + let page_size = page_size.clamp(1, 1000) as usize; + let start = (page as usize).saturating_mul(page_size); + let mut all = server_conn_map + .iter() + .flat_map(|(key, connections)| { + let (namespace, service_name) = split_scoped_service_key(key); + connections.iter().filter_map(move |connection| { + if key_id.is_some_and(|key_id| key_id != namespace) { + return None; + } + Some(AdminConnectionInfo { + key_id: namespace, + namespace, + service_name: service_name.to_string(), + conn_id: connection.conn_id.into(), + generation: connection.generation, + protocol_version: connection.protocol_version, + healthy: connection.health == ServerConnHealth::Healthy, + transport: if connection.is_datagram { "udp" } else { "tcp" } + .to_string(), + codec_enabled: connection.need_codec, + last_rx_age_ms: now + .duration_since(connection.last_rx_at) + .as_millis() + as u64, + }) + }) + }) + .collect::>(); + all.sort_by(|left, right| { + left.namespace + .cmp(&right.namespace) + .then_with(|| left.service_name.cmp(&right.service_name)) + .then_with(|| left.conn_id.cmp(&right.conn_id)) + }); + let items = all.iter().skip(start).take(page_size).cloned().collect(); + let next_page = + (start.saturating_add(page_size) < all.len()).then_some(page.saturating_add(1)); + let _ = response_sender.send(AdminConnectionPage { + schema_version: 1, + items, + next_page, + }); + } ManagerTask::StatusQuery { response_sender } => { let total_connections = server_conn_map .values() @@ -425,23 +618,55 @@ pub async fn run_server_with_shutdown( ManagerTask::Status { conn_sender, status, + namespace, conn_id, } => { let resp = match status { PbConnStatusReq::RemoteId => { + let scoped = server_conn_map + .iter() + .filter(|(key, _)| split_scoped_service_key(key).0 == namespace) + .map(|(key, value)| (split_scoped_service_key(key).1, value)) + .collect::>(); + let registered_ids = scoped + .iter() + .flat_map(|(_, connections)| { + connections.iter().map(|connection| connection.conn_id) + }) + .collect::>(); + let client_ids = pending_streams + .iter() + .filter_map(|(client_id, (_, _, key))| { + (split_scoped_service_key(key).0 == namespace).then_some(*client_id) + }) + .collect::>(); PbConnResponse::Status(PbConnStatusResp::RemoteId { - server_map: format!("{server_conn_map:?}"), - active: manager.active_conn_id_msg(), - idle: manager.idle_conn_id_msg(), + server_map: format!("{scoped:?}"), + active: format!( + "registered={registered_ids:?}, clients={client_ids:?}" + ), + idle: "namespace scoped; use `pb-mapper admin connection list` for global inspection" + .to_string(), }) } PbConnStatusReq::Keys => PbConnResponse::Status(PbConnStatusResp::Keys( - server_conn_map.keys().map(|k| k.to_string()).collect(), + server_conn_map + .keys() + .filter_map(|key| { + let (key_namespace, service_name) = split_scoped_service_key(key); + (key_namespace == namespace).then(|| service_name.to_string()) + }) + .collect(), )), PbConnStatusReq::Service { key } => { - let key: ImutableKey = key.into(); + let display_key = key.clone(); + let key: ImutableKey = if namespace == 0 { + key.into() + } else { + Arc::from(format!("@{namespace:016x}\u{0}{key}")) + }; PbConnResponse::Status(PbConnStatusResp::Service { - key: key.to_string(), + key: display_key, connections: service_status_connections(&server_conn_map, &key), }) } @@ -469,9 +694,11 @@ pub async fn run_server_with_shutdown( "accepted pb connection" ); let manager_task_sender = manager.get_task_sender(); + let security = security.clone(); tokio::spawn(async move { snafu_error_handle!( - handle_conn(conn_id, peer_addr, manager_task_sender, stream).await + handle_conn(conn_id, peer_addr, manager_task_sender, stream, security) + .await ); }); } @@ -479,13 +706,24 @@ pub async fn run_server_with_shutdown( let removed_from_service_map = remove_server_conn(&mut server_conn_map, &key, conn_id); let removed_from_active_map = manager.deregister_conn(conn_id); - pending_streams.retain(|_, (server_id, _)| *server_id != conn_id); + let removed_pending_streams = remove_pending_streams_for_server( + &mut pending_streams, + &mut namespace_stream_counts, + conn_id, + ); + release_namespace_rate_limit_if_idle( + split_scoped_service_key(&key).0, + &server_conn_map, + &pending_streams, + &mut namespace_rate_limits, + ); tracing::info!( event = "server_conn_deregistered", key = %key, conn_id = %conn_id, removed_from_service_map, removed_from_active_map, + removed_pending_streams, registered_services = server_conn_map.len(), server_connections = registered_server_conn_count(&server_conn_map), active_connections = manager.active_conn_count(), @@ -512,14 +750,17 @@ pub async fn run_server_with_shutdown( let removed_from_service_map = remove_server_conn(&mut server_conn_map, &key, conn_id); let removed_from_active_map = manager.deregister_conn(conn_id); - let mut removed_pending_streams = 0usize; - pending_streams.retain(|_, (server_id, _)| { - let keep = *server_id != conn_id; - if !keep { - removed_pending_streams += 1; - } - keep - }); + let removed_pending_streams = remove_pending_streams_for_server( + &mut pending_streams, + &mut namespace_stream_counts, + conn_id, + ); + release_namespace_rate_limit_if_idle( + split_scoped_service_key(&key).0, + &server_conn_map, + &pending_streams, + &mut namespace_rate_limits, + ); let retire_notified = conn_sender .as_ref() .and_then(|sender| { @@ -550,13 +791,25 @@ pub async fn run_server_with_shutdown( server_id, client_id, } => { - pending_streams.remove(&client_id); + let removed_namespace = pending_streams.remove(&client_id).map(|(_, _, key)| { + let namespace = split_scoped_service_key(&key).0; + decrement_namespace_stream_count(&mut namespace_stream_counts, namespace); + namespace + }); let removed_server_conn = if let Some(server_id) = server_id { manager.deregister_conn(server_id) } else { false }; let removed_client_conn = manager.deregister_conn(client_id); + if let Some(namespace) = removed_namespace { + release_namespace_rate_limit_if_idle( + namespace, + &server_conn_map, + &pending_streams, + &mut namespace_rate_limits, + ); + } if removed_server_conn || removed_client_conn { tracing::info!( event = "client_conn_deregistered", @@ -591,6 +844,51 @@ pub async fn run_server_with_shutdown( is_datagram, protocol_version, } => { + let namespace = split_scoped_service_key(&key).0; + let existing = server_conn_map.get(&key); + let failure = if existing.is_some_and(|connections| { + connections + .first() + .is_some_and(|connection| connection.is_datagram != is_datagram) + }) { + Some(( + "service_transport_mismatch", + "the service name is already registered with a different transport", + false, + )) + } else if existing.is_some_and(|connections| { + connections.len() >= max_register_connections_per_service + }) { + Some(( + "service_connection_limit_exceeded", + "the service has reached its register connection limit", + true, + )) + } else if existing.is_none() + && server_conn_map + .keys() + .filter(|registered| split_scoped_service_key(registered).0 == namespace) + .count() + >= max_services_per_namespace + { + Some(( + "namespace_service_limit_exceeded", + "the namespace has reached its service name limit", + true, + )) + } else { + None + }; + if let Some((code, reason, retryable)) = failure { + let _ = conn_sender + .send(ConnTask::RegisterFailed { + code: code.to_string(), + reason: reason.to_string(), + retryable, + }) + .await; + continue; + } let generation = next_server_generation; next_server_generation = next_server_generation.saturating_add(1).max(1); let now = Instant::now(); @@ -648,13 +946,15 @@ pub async fn run_server_with_shutdown( .context(TaskCenterSendRegisterRespSnafu { key, conn_id })); } ManagerTask::Stream { + key, stream, + session, server_id, client_id, server_generation, } => { - let Some((expected_control_conn_id, expected_generation)) = - pending_streams.get(&client_id).copied() + let Some((expected_control_conn_id, expected_generation, expected_key)) = + pending_streams.get(&client_id).cloned() else { tracing::warn!( event = "stale_stream_without_pending_client", @@ -665,6 +965,17 @@ pub async fn run_server_with_shutdown( ); continue; }; + if key != expected_key { + tracing::warn!( + event = "stream_namespace_mismatch", + stream_conn_id = %server_id, + client_conn_id = %client_id, + expected_key = %expected_key, + actual_key = %key, + "dropping stream that does not belong to the pending namespace and service" + ); + continue; + } if server_generation != 0 && expected_generation != server_generation { tracing::warn!( event = "stale_stream_generation_mismatch", @@ -703,7 +1014,8 @@ pub async fn run_server_with_shutdown( .send(ConnTask::StreamResp { server_id, server_generation: expected_generation, - stream + stream, + session, }) .await .map_err(|_| kanal::SendError(())) @@ -716,8 +1028,8 @@ pub async fn run_server_with_shutdown( } => { let recorded_activity = record_server_conn_activity_by_conn_id(&mut server_conn_map, server_id); - let Some((expected_server_id, expected_generation)) = - pending_streams.get(&client_id).copied() + let Some((expected_server_id, expected_generation, _)) = + pending_streams.get(&client_id).cloned() else { tracing::warn!( event = "stale_stream_ack_without_pending_client", @@ -767,6 +1079,22 @@ pub async fn run_server_with_shutdown( conn_sender, excluded_server_conns, } => { + let namespace = split_scoped_service_key(&key).0; + if namespace_stream_counts + .get(&namespace) + .copied() + .unwrap_or_default() + >= max_streams_per_namespace + { + let _ = conn_sender + .send(ConnTask::SubcribeFailed { + code: "namespace_stream_limit_exceeded".to_string(), + reason: "the namespace has reached its active stream limit".to_string(), + retryable: true, + }) + .await; + continue; + } let Some(server_conn_id_list) = server_conn_map.get(&key).cloned() else { let reason = format!("server key `{key}` is not registered"); tracing::warn!( @@ -785,6 +1113,22 @@ pub async fn run_server_with_shutdown( } continue; }; + if !namespace_rate_limits + .entry(namespace) + .or_insert_with(|| { + NamespaceRateLimit::new(new_streams_per_second, new_streams_burst) + }) + .allow() + { + let _ = conn_sender + .send(ConnTask::SubcribeFailed { + code: "namespace_stream_rate_exceeded".to_string(), + reason: "the namespace new-stream rate limit was exceeded".to_string(), + retryable: true, + }) + .await; + continue; + } let mut selected = false; let mut candidates = Vec::new(); candidates.extend(server_conn_id_list.iter().rev().copied().filter(|info| { @@ -845,7 +1189,12 @@ pub async fn run_server_with_shutdown( if manager.get_conn_sender_chan(&conn_id).is_none() { manager.sign_up_conn_sender(conn_id, conn_sender.clone()); } - pending_streams.insert(conn_id, (server_conn_id, server_generation)); + let is_new_stream = pending_streams + .insert(conn_id, (server_conn_id, server_generation, key.clone())) + .is_none(); + if is_new_stream { + *namespace_stream_counts.entry(namespace).or_default() += 1; + } // 2. Response subcribe ok if let Err(e) = conn_sender .send(ConnTask::SubcribeResp { @@ -953,15 +1302,158 @@ async fn handle_listener( } } -#[instrument(skip(manager_task_sender, conn), fields(conn_id = %conn_id, peer_addr = %peer_addr))] +#[instrument(skip(manager_task_sender, conn, security), fields(conn_id = %conn_id, peer_addr = %peer_addr))] async fn handle_conn( conn_id: RemoteConnId, peer_addr: SocketAddr, manager_task_sender: ManagerTaskSender, mut conn: TcpStream, + security: ServerSecurity, ) -> Result<()> { - // handle by action - let init_request = get_init_request(&mut conn, conn_id).await?; + let timeout = control_io_timeout(); + let initial = match tokio::time::timeout(timeout, security.read_initial(&mut conn)).await { + Err(_) => TaskCenterInitRequestTimeoutSnafu { conn_id, timeout }.fail()?, + Ok(Err(error)) => { + let key_id = error + .response_session + .as_ref() + .map(|session| session.key_id()) + .unwrap_or_default(); + let decision = security.record_failure_log(peer_addr.ip(), key_id, &error.failure.code); + if decision.suppressed > 0 { + tracing::warn!( + event = "auth_failures_suppressed", + peer_ip = %peer_addr.ip(), + key_id, + reason = %error.failure.code, + suppressed = decision.suppressed, + "suppressed repeated authentication failures in the previous window" + ); + } + if decision.emit { + tracing::warn!( + event = "auth_failed", + auth_stage = "initial_frame", + conn_id = %conn_id, + peer_addr = %peer_addr, + key_id, + reason = %error.failure.code, + retryable = error.failure.retryable, + error = %error.failure.message, + "connection authentication failed" + ); + } + if let Some(session) = error.response_session { + write_protocol_error(&mut conn, &session, &error.failure).await; + } + return Ok(()); + } + Ok(Ok(initial)) => initial, + }; + let init_request = match PbConnRequest::decode(&initial.payload) { + Ok(request) => request, + Err(error) => { + tracing::warn!( + event = "auth_failed", + auth_stage = "request_decode", + conn_id = %conn_id, + peer_addr = %peer_addr, + error = %error, + "authenticated request could not be decoded" + ); + write_protocol_error( + &mut conn, + &initial.session, + &crate::common::auth::AuthFailure::new( + "request_decode_failed", + "authenticated request payload is malformed", + false, + ), + ) + .await; + return Ok(()); + } + }; + let mut requested_namespace = None; + let mut force_register_namespace = false; + let init_request = match init_request { + PbConnRequest::RegisterScoped { + need_codec, + is_datagram, + key, + namespace, + force_namespace, + protocol_version, + client_instance_id, + heartbeat_interval_ms, + heartbeat_tolerance_ms, + } => { + requested_namespace = Some(namespace); + force_register_namespace = force_namespace; + PbConnRequest::Register { + need_codec, + is_datagram, + key, + protocol_version, + client_instance_id, + heartbeat_interval_ms, + heartbeat_tolerance_ms, + } + } + PbConnRequest::SubcribeScoped { key, namespace } => { + requested_namespace = Some(namespace); + PbConnRequest::Subcribe { key } + } + PbConnRequest::StatusScoped { status, namespace } => { + requested_namespace = Some(namespace); + PbConnRequest::Status(status) + } + PbConnRequest::StreamScoped { + key, + namespace, + dst_id, + server_generation, + } => { + requested_namespace = Some(namespace); + PbConnRequest::Stream { + key, + dst_id, + server_generation, + } + } + request => request, + }; + let session = initial.session; + let auth_context = match session.context() { + Ok(context) => context.clone(), + Err(error) => { + tracing::warn!(conn_id = %conn_id, peer_addr = %peer_addr, %error, "missing auth context"); + return Ok(()); + } + }; + tracing::info!( + event = "auth_succeeded", + auth_stage = "session", + conn_id = %conn_id, + peer_addr = %peer_addr, + key_id = auth_context.key_id, + namespace = auth_context.namespace, + protocol = ?session.protocol(), + is_admin = auth_context.is_admin, + "connection authentication succeeded" + ); + let effective_namespace = match resolve_namespace( + &auth_context, + requested_namespace, + force_register_namespace, + matches!(&init_request, PbConnRequest::Register { .. }), + ) { + Ok(namespace) => namespace, + Err(failure) => { + write_protocol_error(&mut conn, &session, &failure).await; + return Ok(()); + } + }; match init_request { PbConnRequest::Register { key, @@ -987,16 +1479,37 @@ async fn handle_conn( is_datagram, "received pb init request" ); - handle_server_conn( - key.into(), - need_codec, - is_datagram, - protocol_version, - conn_id, + let key = match scoped_service_key(&auth_context, effective_namespace, &key) { + Ok(key) => key, + Err(failure) => { + write_protocol_error(&mut conn, &session, &failure).await; + return Ok(()); + } + }; + let cancellation = match auth_context.cancellation_token() { + Ok(token) => token, + Err(failure) => { + write_protocol_error(&mut conn, &session, &failure).await; + return Ok(()); + } + }; + tokio::select! { + result = handle_server_conn( + ServerRegistration { + key, + need_codec, + is_datagram, + protocol_version, + conn_id, + }, manager_task_sender, conn, - ) - .await?; + session, + ) => result?, + _ = cancellation.cancelled() => { + tracing::info!(event = "connection_auth_expired", key_id = auth_context.key_id, conn_id = %conn_id, "closing registered service connection"); + } + } } PbConnRequest::Subcribe { key } => { tracing::info!( @@ -1007,7 +1520,26 @@ async fn handle_conn( key = %key, "received pb init request" ); - handle_client_conn(key.into(), conn_id, manager_task_sender, conn).await?; + let key = match scoped_service_key(&auth_context, effective_namespace, &key) { + Ok(key) => key, + Err(failure) => { + write_protocol_error(&mut conn, &session, &failure).await; + return Ok(()); + } + }; + let cancellation = match auth_context.cancellation_token() { + Ok(token) => token, + Err(failure) => { + write_protocol_error(&mut conn, &session, &failure).await; + return Ok(()); + } + }; + tokio::select! { + result = handle_client_conn(key, conn_id, manager_task_sender, conn, session) => result?, + _ = cancellation.cancelled() => { + tracing::info!(event = "connection_auth_expired", key_id = auth_context.key_id, conn_id = %conn_id, "closing subscribed data connection"); + } + } } PbConnRequest::Stream { key, @@ -1024,10 +1556,18 @@ async fn handle_conn( server_generation, "received pb init request" ); - let key = ImutableKey::from(key); + let key = match scoped_service_key(&auth_context, effective_namespace, &key) { + Ok(key) => key, + Err(failure) => { + write_protocol_error(&mut conn, &session, &failure).await; + return Ok(()); + } + }; manager_task_sender .send(ManagerTask::Stream { + key: key.clone(), stream: conn, + session, server_id: conn_id, client_id: dst_id.into(), server_generation, @@ -1045,12 +1585,184 @@ async fn handle_conn( status = ?status, "received pb init request" ); - handle_show_status(status, manager_task_sender, conn_id, conn).await?; + handle_show_status( + status, + effective_namespace, + manager_task_sender, + conn_id, + conn, + session, + ) + .await?; } + PbConnRequest::Admin(request) => { + if !auth_context.is_admin { + write_protocol_error( + &mut conn, + &session, + &crate::common::auth::AuthFailure::new( + "admin_permission_required", + "administrator credential is required for this operation", + false, + ), + ) + .await; + return Ok(()); + } + handle_admin_request( + request, + security.auth().clone(), + manager_task_sender, + conn_id, + conn, + session, + ) + .await?; + } + PbConnRequest::RegisterScoped { .. } + | PbConnRequest::SubcribeScoped { .. } + | PbConnRequest::StatusScoped { .. } + | PbConnRequest::StreamScoped { .. } => unreachable!("scoped request was normalized"), } Ok(()) } +async fn write_protocol_error( + conn: &mut TcpStream, + session: &ServerHeaderSession, + failure: &crate::common::auth::AuthFailure, +) { + let response = PbConnResponse::error( + failure.code.clone(), + failure.message.clone(), + failure.retryable, + ); + let Ok(message) = response.encode() else { + return; + }; + let Ok(mut writer) = session.response_writer(conn) else { + return; + }; + if let Err(error) = writer.write_msg(&message).await { + tracing::debug!(%error, reason = %failure.code, "failed to write structured protocol error"); + } +} + +fn resolve_namespace( + context: &AuthContext, + requested: Option, + force_register_namespace: bool, + is_register: bool, +) -> std::result::Result { + let namespace = requested.unwrap_or(context.namespace); + if !context.is_admin && namespace != context.namespace { + return Err(crate::common::auth::AuthFailure::new( + "namespace_access_denied", + "temporary credentials can only access their own namespace", + false, + )); + } + if context.is_admin && is_register && namespace != 0 && !force_register_namespace { + return Err(crate::common::auth::AuthFailure::new( + "namespace_force_required", + "administrator registration in a temporary namespace requires --force", + false, + )); + } + Ok(namespace) +} + +fn scoped_service_key( + context: &AuthContext, + namespace: u64, + service_name: &str, +) -> std::result::Result { + if service_name.is_empty() || service_name.len() > 1024 || service_name.contains('\0') { + return Err(crate::common::auth::AuthFailure::new( + "service_name_invalid", + "service names must be 1-1024 bytes and must not contain NUL", + false, + )); + } + if !context.is_admin + && (service_name.len() > 128 + || !service_name + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || b"._:-".contains(&byte))) + { + return Err(crate::common::auth::AuthFailure::new( + "service_name_invalid", + "temporary-key service names must be 1-128 ASCII bytes from [A-Za-z0-9._:-]", + false, + )); + } + if namespace == 0 { + Ok(Arc::from(service_name)) + } else { + Ok(Arc::from(format!("@{namespace:016x}\u{0}{service_name}"))) + } +} + +fn split_scoped_service_key(key: &str) -> (u64, &str) { + let Some((prefix, name)) = key.split_once('\0') else { + return (0, key); + }; + let Some(hex) = prefix.strip_prefix('@') else { + return (0, key); + }; + match u64::from_str_radix(hex, 16) { + Ok(namespace) => (namespace, name), + Err(_) => (0, key), + } +} + +fn decrement_namespace_stream_count( + namespace_stream_counts: &mut hashbrown::HashMap, + namespace: u64, +) { + let Some(count) = namespace_stream_counts.get_mut(&namespace) else { + return; + }; + *count = count.saturating_sub(1); + if *count == 0 { + namespace_stream_counts.remove(&namespace); + } +} + +fn release_namespace_rate_limit_if_idle( + namespace: u64, + server_conn_map: &ServerConnMap, + pending_streams: &hashbrown::HashMap, + namespace_rate_limits: &mut hashbrown::HashMap, +) { + let has_registered_service = server_conn_map + .keys() + .any(|key| split_scoped_service_key(key).0 == namespace); + let has_pending_stream = pending_streams + .values() + .any(|(_, _, key)| split_scoped_service_key(key).0 == namespace); + if !has_registered_service && !has_pending_stream { + namespace_rate_limits.remove(&namespace); + } +} + +fn remove_pending_streams_for_server( + pending_streams: &mut hashbrown::HashMap, + namespace_stream_counts: &mut hashbrown::HashMap, + server_id_to_remove: RemoteConnId, +) -> usize { + let mut removed = 0; + pending_streams.retain(|_, (server_id, _, key)| { + if *server_id != server_id_to_remove { + return true; + } + decrement_namespace_stream_count(namespace_stream_counts, split_scoped_service_key(key).0); + removed += 1; + false + }); + removed +} + pub async fn get_init_request( conn: &mut TcpStream, conn_id: RemoteConnId, diff --git a/src/pb_server/server.rs b/src/pb_server/server.rs index f330e4f..b379057 100644 --- a/src/pb_server/server.rs +++ b/src/pb_server/server.rs @@ -17,9 +17,8 @@ use crate::common::conn_id::RemoteConnId; use crate::common::message::command::{ LocalServer, MessageSerializer, PbConnResponse, PbServerRequest, CONTROL_PROTOCOL_V2, }; -use crate::common::message::{ - get_header_msg_reader, get_header_msg_writer, MessageReader, MessageWriter, -}; +use crate::common::message::secure::ServerHeaderSession; +use crate::common::message::{MessageReader, MessageWriter}; /// Ensure that server-side connections are properly deregistered before a normal connection is /// disconnected or an exception occurs @@ -132,18 +131,30 @@ enum ServerControlWrite { Pong(Vec), } +pub struct ServerRegistration { + pub key: ImutableKey, + pub need_codec: bool, + pub is_datagram: bool, + pub protocol_version: u16, + pub conn_id: RemoteConnId, +} + /// Maintaining a connection to the server. /// This connection is used to send channel request -#[instrument(skip(task_sender))] +#[instrument(skip(registration, task_sender, session))] pub async fn handle_server_conn( - key: ImutableKey, - need_codec: bool, - is_datagram: bool, - protocol_version: u16, - conn_id: RemoteConnId, + registration: ServerRegistration, task_sender: ManagerTaskSender, - conn: TcpStream, + mut conn: TcpStream, + session: ServerHeaderSession, ) -> Result<()> { + let ServerRegistration { + key, + need_codec, + is_datagram, + protocol_version, + conn_id, + } = registration; let (tx, rx) = kanal::bounded_async(DEFAULT_SERVER_CHAN_CAP); // register metadate @@ -187,6 +198,30 @@ pub async fn handle_server_conn( lease_ttl_ms, } = response else { + if let ConnTask::RegisterFailed { + code, + reason, + retryable, + } = response + { + let response = PbConnResponse::error(code, reason, retryable) + .encode() + .context(ServerConnEncodeRegisterRespSnafu { + key: key.clone(), + conn_id, + })?; + let mut writer = session + .response_writer(&mut conn) + .context(ServerConnCreateHeaderToolSnafu { tool: "writer" })?; + writer + .write_msg(&response) + .await + .context(ServerConnWriteRegisteredOkSnafu { + key: key.clone(), + conn_id, + })?; + return Ok(()); + } ServerConnRegisteredRespNotMatchSnafu { key: key.clone(), conn_id, @@ -203,7 +238,8 @@ pub async fn handle_server_conn( ); let (mut reader, mut writer) = conn.into_split(); - let mut msg_reader = get_header_msg_reader(&mut reader) + let mut msg_reader = session + .continuation_reader(&mut reader) .context(ServerConnCreateHeaderToolSnafu { tool: "reader" })?; // Keep one header writer for the register response and all later control frames. The // encrypted header codec is stateful; recreating it between frames breaks peer decoding. @@ -224,7 +260,8 @@ pub async fn handle_server_conn( let (write_tx, mut write_rx) = tokio::sync::mpsc::unbounded_channel::(); let writer_key = key.clone(); let mut writer_handle = tokio::spawn(async move { - let mut msg_writer = get_header_msg_writer(&mut writer) + let mut msg_writer = session + .response_writer(&mut writer) .context(ServerConnCreateHeaderToolSnafu { tool: "writer" })?; msg_writer.write_msg(®ister_response).await.context( ServerConnWriteRegisteredOkSnafu { diff --git a/src/pb_server/status.rs b/src/pb_server/status.rs index 789a7c3..77c48ef 100644 --- a/src/pb_server/status.rs +++ b/src/pb_server/status.rs @@ -9,7 +9,8 @@ use super::error::{ use super::{ConnTask, ManagerTask, ManagerTaskSender}; use crate::common::conn_id::RemoteConnId; use crate::common::message::command::{MessageSerializer, PbConnStatusReq}; -use crate::common::message::{get_header_msg_writer, MessageWriter}; +use crate::common::message::secure::ServerHeaderSession; +use crate::common::message::MessageWriter; struct StatusConnGuard { conn_id: RemoteConnId, @@ -85,9 +86,11 @@ impl Drop for StatusConnGuard { pub async fn handle_show_status( status: PbConnStatusReq, + namespace: u64, manager_sender: ManagerTaskSender, conn_id: RemoteConnId, mut conn: TcpStream, + session: ServerHeaderSession, ) -> Result<()> { let info_span = info_span!("show status", "{status:?},{conn_id:?}"); let mut guard = StatusConnGuard::new(conn_id, manager_sender.clone()); @@ -97,6 +100,7 @@ pub async fn handle_show_status( let req = ManagerTask::Status { conn_sender: tx, status, + namespace, conn_id, }; manager_sender @@ -108,7 +112,8 @@ pub async fn handle_show_status( let resp = rx.recv().await.context(StatusRecvConnTaskSnafu)?; if let ConnTask::StatusResp(resp) = resp { let msg = resp.encode().context(StatusEncodeRespSnafu)?; - let mut msg_writer = get_header_msg_writer(&mut conn) + let mut msg_writer = session + .response_writer(&mut conn) .context(StatusCreateHeaderToolSnafu { tool: "writer" })?; msg_writer .write_msg(&msg) diff --git a/src/utils/codec.rs b/src/utils/codec.rs index b375470..c5a00dd 100644 --- a/src/utils/codec.rs +++ b/src/utils/codec.rs @@ -56,7 +56,7 @@ impl Aes256GcmCodec { } pub fn try_new_with_default_key() -> RingResult { - let key = get_msg_header_key(); + let key = get_msg_header_key().map_err(|_| ring::error::Unspecified)?; Aes256GcmCodec::try_new(key.as_ref()) } diff --git a/tests/regression.rs b/tests/regression.rs index 145b3ce..bbe1d1b 100644 --- a/tests/regression.rs +++ b/tests/regression.rs @@ -3,17 +3,25 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use std::time::Duration; +use pb_mapper::common::auth::{ + write_admin_key_file, AuthConfig, AuthRuntime, LegacyProtocolPolicy, +}; +use pb_mapper::common::checksum::{parse_credential, set_process_msg_header_key, Credential}; use pb_mapper::common::message::command::{ - LocalServer, MessageSerializer, PbConnRequest, PbConnResponse, PbConnStatusReq, - PbConnStatusResp, PbServerRequest, PbServiceConnStatus, + AdminRequest, AdminResponse, LocalServer, MessageSerializer, PbConnRequest, PbConnResponse, + PbConnStatusReq, PbConnStatusResp, PbServerRequest, PbServiceConnStatus, +}; +use pb_mapper::common::message::secure::{ + ClientHeaderSession, ServerHeaderSession, ServerSecurity, }; use pb_mapper::common::message::{ get_header_msg_reader, get_header_msg_writer, MessageReader, MessageWriter, }; use pb_mapper::local::client::run_client_side_cli_with_callback; use pb_mapper::local::server::{run_server_side_cli_with_callback, ServerTunnelOptions}; -use pb_mapper::pb_server::{get_init_request, run_server_with_shutdown}; +use pb_mapper::pb_server::run_server_with_auth_config; use tokio::io::AsyncReadExt; +use tokio::io::AsyncWriteExt; use tokio::net::{TcpListener, TcpStream}; use tokio::time::timeout; use tokio_util::sync::CancellationToken; @@ -55,6 +63,33 @@ async fn wait_for_server(server_addr: SocketAddr) -> TcpStream { .expect("server did not start") } +async fn read_secure_request( + security: &ServerSecurity, + stream: &mut TcpStream, +) -> (PbConnRequest, ServerHeaderSession) { + let initial = security.read_initial(stream).await.unwrap(); + ( + PbConnRequest::decode(&initial.payload).unwrap(), + initial.session, + ) +} + +fn auth_config(server_addr: SocketAddr) -> AuthConfig { + set_process_msg_header_key(Some(TEST_ADMIN_KEY)).unwrap(); + AuthConfig { + state_dir: std::env::temp_dir().join(format!( + "pb-mapper-regression-{}-{}", + std::process::id(), + server_addr.port() + )), + max_temporary_keys: 64, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + } +} + +const TEST_ADMIN_KEY: &str = "0123456789abcdefghijklmnopqrstuv"; + async fn register_control_conn_parts( reader: &mut impl MessageReader, writer: &mut impl MessageWriter, @@ -139,18 +174,348 @@ async fn read_status_keys(server_addr: SocketAddr) -> Vec { keys } +async fn send_v2_request( + server_addr: SocketAddr, + credential: &Credential, + request: PbConnRequest, +) -> (TcpStream, ClientHeaderSession, PbConnResponse) { + let mut stream = wait_for_server(server_addr).await; + let session = ClientHeaderSession::new_v2(credential).unwrap(); + session + .write_initial(&mut stream, &request.encode().unwrap()) + .await + .unwrap(); + let response = { + let mut reader = session.response_reader(&mut stream).unwrap(); + let message = timeout(Duration::from_secs(1), reader.read_msg()) + .await + .expect("v2 response timed out") + .unwrap(); + PbConnResponse::decode(message).unwrap() + }; + (stream, session, response) +} + #[tokio::test] -async fn status_service_reports_registered_v2_control_connection() { +async fn temporary_credentials_are_isolated_denied_admin_and_revoked_live() { + let probe_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let server_addr = probe_listener.local_addr().unwrap(); + drop(probe_listener); + + let config = auth_config(server_addr); + let _ = std::fs::remove_dir_all(&config.state_dir); + write_admin_key_file(&config.state_dir.join("admin.key"), TEST_ADMIN_KEY, true).unwrap(); + let runtime = AuthRuntime::start( + *TEST_ADMIN_KEY.as_bytes().first_chunk::<32>().unwrap(), + config.clone(), + ) + .await + .unwrap(); + let first = runtime + .issue(Duration::from_secs(120), Some("first".to_string())) + .await + .unwrap(); + let second = runtime + .issue(Duration::from_secs(120), Some("second".to_string())) + .await + .unwrap(); + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + + let shutdown_token = CancellationToken::new(); + let server_shutdown = shutdown_token.clone(); + let server_config = config.clone(); + let server = tokio::spawn(async move { + run_server_with_auth_config(server_addr, server_shutdown, None, false, server_config) + .await + .unwrap(); + }); + + let first_credential = parse_credential(&first.credential).unwrap(); + let second_credential = parse_credential(&second.credential).unwrap(); + let register_request = |service: &str| PbConnRequest::Register { + need_codec: false, + is_datagram: false, + key: service.to_string(), + protocol_version: Some(2), + client_instance_id: Some("temporary-auth-regression".to_string()), + heartbeat_interval_ms: Some(50), + heartbeat_tolerance_ms: Some(150), + }; + let (mut first_control, _, first_register) = send_v2_request( + server_addr, + &first_credential, + register_request("same-name"), + ) + .await; + let (_second_control, _, second_register) = send_v2_request( + server_addr, + &second_credential, + register_request("same-name"), + ) + .await; + let first_conn_id = match first_register { + PbConnResponse::RegisterV2 { conn_id, .. } => conn_id, + response => panic!("unexpected first register response: {response:?}"), + }; + let second_conn_id = match second_register { + PbConnResponse::RegisterV2 { conn_id, .. } => conn_id, + response => panic!("unexpected second register response: {response:?}"), + }; + assert_ne!(first_conn_id, second_conn_id); + + let status_request = PbConnRequest::Status(PbConnStatusReq::Service { + key: "same-name".to_string(), + }); + for (credential, expected_conn_id) in [ + (&first_credential, first_conn_id), + (&second_credential, second_conn_id), + ] { + let (_, _, response) = + send_v2_request(server_addr, credential, status_request.clone()).await; + let PbConnResponse::Status(PbConnStatusResp::Service { connections, .. }) = response else { + panic!("unexpected scoped status response: {response:?}"); + }; + assert_eq!(connections.len(), 1); + assert_eq!(connections[0].conn_id, expected_conn_id); + } + + let (_, _, denied) = send_v2_request( + server_addr, + &first_credential, + PbConnRequest::Admin(AdminRequest::AuthStatus), + ) + .await; + let PbConnResponse::Error(denied) = denied else { + panic!("temporary credential unexpectedly received an admin response"); + }; + assert_eq!(denied.code, "admin_permission_required"); + + let admin_credential = + Credential::Admin(*TEST_ADMIN_KEY.as_bytes().first_chunk::<32>().unwrap()); + let (_, _, revoked) = send_v2_request( + server_addr, + &admin_credential, + PbConnRequest::Admin(AdminRequest::KeyRevoke { + key_id: first.metadata.key_id, + }), + ) + .await; + assert!(matches!( + revoked, + PbConnResponse::Admin(AdminResponse::KeyRevoked(_)) + )); + let mut byte = [0_u8; 1]; + let closed = timeout(Duration::from_secs(1), first_control.read(&mut byte)) + .await + .expect("revoked control connection was not closed") + .unwrap(); + assert_eq!(closed, 0); + + let (_, _, second_still_active) = + send_v2_request(server_addr, &second_credential, status_request).await; + assert!(matches!( + second_still_active, + PbConnResponse::Status(PbConnStatusResp::Service { .. }) + )); + + shutdown_token.cancel(); + server.await.unwrap(); + let _ = std::fs::remove_dir_all(config.state_dir); +} + +#[tokio::test] +async fn revoking_temporary_credential_closes_active_data_stream() { let probe_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let server_addr = probe_listener.local_addr().unwrap(); drop(probe_listener); + let config = auth_config(server_addr); + let _ = std::fs::remove_dir_all(&config.state_dir); + write_admin_key_file(&config.state_dir.join("admin.key"), TEST_ADMIN_KEY, true).unwrap(); + let runtime = AuthRuntime::start( + *TEST_ADMIN_KEY.as_bytes().first_chunk::<32>().unwrap(), + config.clone(), + ) + .await + .unwrap(); + let issued = runtime + .issue(Duration::from_secs(120), Some("active-stream".to_string())) + .await + .unwrap(); + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + let shutdown_token = CancellationToken::new(); let server_shutdown = shutdown_token.clone(); + let server_config = config.clone(); let server = tokio::spawn(async move { - run_server_with_shutdown(server_addr, server_shutdown, None, false) + run_server_with_auth_config(server_addr, server_shutdown, None, false, server_config) + .await + .unwrap(); + }); + + let credential = parse_credential(&issued.credential).unwrap(); + let service = "revoked-stream"; + let mut control = wait_for_server(server_addr).await; + let control_session = ClientHeaderSession::new_v2(&credential).unwrap(); + let register = PbConnRequest::Register { + need_codec: false, + is_datagram: false, + key: service.to_string(), + protocol_version: Some(2), + client_instance_id: Some("active-stream-test".to_string()), + heartbeat_interval_ms: Some(50), + heartbeat_tolerance_ms: Some(150), + }; + control_session + .write_initial(&mut control, ®ister.encode().unwrap()) + .await + .unwrap(); + let (mut control_read, mut control_write) = control.into_split(); + let mut control_reader = control_session.response_reader(&mut control_read).unwrap(); + let register_response = timeout(Duration::from_secs(1), control_reader.read_msg()) + .await + .expect("register response timed out") + .unwrap(); + assert!(matches!( + PbConnResponse::decode(register_response).unwrap(), + PbConnResponse::RegisterV2 { .. } + )); + + let mut subscriber = wait_for_server(server_addr).await; + let subscriber_session = ClientHeaderSession::new_v2(&credential).unwrap(); + subscriber_session + .write_initial( + &mut subscriber, + &PbConnRequest::Subcribe { + key: service.to_string(), + } + .encode() + .unwrap(), + ) + .await + .unwrap(); + + let stream_request = timeout(Duration::from_secs(1), control_reader.read_msg()) + .await + .expect("stream request timed out") + .unwrap(); + let LocalServer::Stream { + client_id, + server_generation, + } = LocalServer::decode(stream_request).unwrap() + else { + panic!("unexpected local server stream request"); + }; + let mut control_writer = control_session + .continuation_writer(&mut control_write) + .unwrap(); + control_writer + .write_msg( + &PbServerRequest::StreamAck { + client_id, + server_generation, + } + .encode() + .unwrap(), + ) + .await + .unwrap(); + + let mut provider = wait_for_server(server_addr).await; + let provider_session = ClientHeaderSession::new_v2(&credential).unwrap(); + provider_session + .write_initial( + &mut provider, + &PbConnRequest::Stream { + key: service.to_string(), + dst_id: client_id, + server_generation, + } + .encode() + .unwrap(), + ) + .await + .unwrap(); + { + let mut subscriber_reader = subscriber_session.response_reader(&mut subscriber).unwrap(); + let response = timeout(Duration::from_secs(1), subscriber_reader.read_msg()) .await + .expect("subscribe response timed out") .unwrap(); + assert!(matches!( + PbConnResponse::decode(response).unwrap(), + PbConnResponse::Subcribe { .. } + )); + } + { + let mut provider_reader = provider_session.response_reader(&mut provider).unwrap(); + let response = timeout(Duration::from_secs(1), provider_reader.read_msg()) + .await + .expect("provider stream response timed out") + .unwrap(); + assert!(matches!( + PbConnResponse::decode(response).unwrap(), + PbConnResponse::Stream { .. } + )); + } + + subscriber.write_all(b"ready").await.unwrap(); + let mut ready = [0_u8; 5]; + timeout(Duration::from_secs(1), provider.read_exact(&mut ready)) + .await + .expect("active data stream did not forward") + .unwrap(); + assert_eq!(&ready, b"ready"); + + let admin_credential = + Credential::Admin(*TEST_ADMIN_KEY.as_bytes().first_chunk::<32>().unwrap()); + let (_, _, revoked) = send_v2_request( + server_addr, + &admin_credential, + PbConnRequest::Admin(AdminRequest::KeyRevoke { + key_id: issued.metadata.key_id, + }), + ) + .await; + assert!(matches!( + revoked, + PbConnResponse::Admin(AdminResponse::KeyRevoked(_)) + )); + + let mut byte = [0_u8; 1]; + for (name, stream) in [("subscriber", &mut subscriber), ("provider", &mut provider)] { + let read = timeout(Duration::from_secs(1), stream.read(&mut byte)) + .await + .unwrap_or_else(|_| panic!("revoked {name} data stream was not closed")) + .unwrap(); + assert_eq!(read, 0, "revoked {name} data stream remained open"); + } + + shutdown_token.cancel(); + server.await.unwrap(); + let _ = std::fs::remove_dir_all(config.state_dir); +} + +#[tokio::test] +async fn status_service_reports_registered_v2_control_connection() { + let probe_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let server_addr = probe_listener.local_addr().unwrap(); + drop(probe_listener); + + let shutdown_token = CancellationToken::new(); + let server_shutdown = shutdown_token.clone(); + let server = tokio::spawn(async move { + run_server_with_auth_config( + server_addr, + server_shutdown, + None, + false, + auth_config(server_addr), + ) + .await + .unwrap(); }); let key = "sf-backend"; @@ -210,15 +575,19 @@ async fn local_server_reconnects_when_registered_conn_is_missing_from_remote_sta let fake_register_count = register_count.clone(); let fake_second_register_tx = second_register_tx.clone(); + let fake_security = ServerSecurity::new( + AuthRuntime::from_process(auth_config(remote_addr)) + .await + .unwrap(), + ); let fake_server = tokio::spawn(async move { loop { let (mut stream, _) = remote_listener.accept().await.unwrap(); let register_count = fake_register_count.clone(); let second_register_tx = fake_second_register_tx.clone(); + let security = fake_security.clone(); tokio::spawn(async move { - let Ok(request) = get_init_request(&mut stream, 0.into()).await else { - return; - }; + let (request, session) = read_secure_request(&security, &mut stream).await; match request { PbConnRequest::Register { key, .. } => { let count = register_count.fetch_add(1, Ordering::SeqCst) + 1; @@ -229,7 +598,7 @@ async fn local_server_reconnects_when_registered_conn_is_missing_from_remote_sta } .encode() .unwrap(); - let mut writer = get_header_msg_writer(&mut stream).unwrap(); + let mut writer = session.response_writer(&mut stream).unwrap(); writer.write_msg(&response).await.unwrap(); if count == 2 { if let Some(tx) = second_register_tx.lock().await.take() { @@ -246,14 +615,14 @@ async fn local_server_reconnects_when_registered_conn_is_missing_from_remote_sta }) .encode() .unwrap(); - let mut writer = get_header_msg_writer(&mut stream).unwrap(); + let mut writer = session.response_writer(&mut stream).unwrap(); writer.write_msg(&response).await.unwrap(); } PbConnRequest::Status(PbConnStatusReq::Keys) => { let response = PbConnResponse::Status(PbConnStatusResp::Keys(Vec::new())) .encode() .unwrap(); - let mut writer = get_header_msg_writer(&mut stream).unwrap(); + let mut writer = session.response_writer(&mut stream).unwrap(); writer.write_msg(&response).await.unwrap(); } _ => {} @@ -271,6 +640,8 @@ async fn local_server_reconnects_when_registered_conn_is_missing_from_remote_sta need_codec: false, is_datagram: false, keep_alive: false, + namespace: None, + force_namespace: false, }, None, )); @@ -288,10 +659,15 @@ async fn local_server_reconnects_when_registered_conn_is_missing_from_remote_sta async fn client_closes_initial_status_probe_after_key_check() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let remote_addr = listener.local_addr().unwrap(); + let security = ServerSecurity::new( + AuthRuntime::from_process(auth_config(remote_addr)) + .await + .unwrap(), + ); let fake_server = tokio::spawn(async move { let (mut stream, _) = listener.accept().await.unwrap(); - let request = get_init_request(&mut stream, 0.into()).await.unwrap(); + let (request, session) = read_secure_request(&security, &mut stream).await; let PbConnRequest::Status(PbConnStatusReq::Service { key }) = request else { panic!("client did not use service status for initial key check"); }; @@ -309,7 +685,7 @@ async fn client_closes_initial_status_probe_after_key_check() { .encode() .unwrap(); { - let mut writer = get_header_msg_writer(&mut stream).unwrap(); + let mut writer = session.response_writer(&mut stream).unwrap(); writer.write_msg(&response).await.unwrap(); } @@ -351,15 +727,19 @@ async fn client_tolerates_one_failed_health_check_while_listener_is_active() { let fake_failed_status_responses = failed_status_responses.clone(); let fake_status_count = status_count.clone(); + let fake_security = ServerSecurity::new( + AuthRuntime::from_process(auth_config(remote_addr)) + .await + .unwrap(), + ); let fake_server = tokio::spawn(async move { loop { let (mut stream, _) = remote_listener.accept().await.unwrap(); let failed_status_responses = fake_failed_status_responses.clone(); let status_count = fake_status_count.clone(); + let security = fake_security.clone(); tokio::spawn(async move { - let Ok(request) = get_init_request(&mut stream, 0.into()).await else { - return; - }; + let (request, session) = read_secure_request(&security, &mut stream).await; match request { PbConnRequest::Status(PbConnStatusReq::Service { key }) => { status_count.fetch_add(1, Ordering::SeqCst); @@ -383,7 +763,7 @@ async fn client_tolerates_one_failed_health_check_while_listener_is_active() { PbConnResponse::Status(PbConnStatusResp::Service { key, connections }) .encode() .unwrap(); - let mut writer = get_header_msg_writer(&mut stream).unwrap(); + let mut writer = session.response_writer(&mut stream).unwrap(); writer.write_msg(&response).await.unwrap(); } PbConnRequest::Status(PbConnStatusReq::Keys) => { @@ -401,7 +781,7 @@ async fn client_tolerates_one_failed_health_check_while_listener_is_active() { let response = PbConnResponse::Status(PbConnStatusResp::Keys(keys)) .encode() .unwrap(); - let mut writer = get_header_msg_writer(&mut stream).unwrap(); + let mut writer = session.response_writer(&mut stream).unwrap(); writer.write_msg(&response).await.unwrap(); } PbConnRequest::Subcribe { .. } => std::future::pending().await, @@ -484,9 +864,15 @@ async fn subscribe_retires_unacked_control_connection() { let shutdown_token = CancellationToken::new(); let server_shutdown = shutdown_token.clone(); let server = tokio::spawn(async move { - run_server_with_shutdown(server_addr, server_shutdown, None, false) - .await - .unwrap(); + run_server_with_auth_config( + server_addr, + server_shutdown, + None, + false, + auth_config(server_addr), + ) + .await + .unwrap(); }); let key = "sf-backend"; @@ -548,9 +934,15 @@ async fn subscribe_waits_for_replacement_after_retiring_stale_control_connection let shutdown_token = CancellationToken::new(); let server_shutdown = shutdown_token.clone(); let server = tokio::spawn(async move { - run_server_with_shutdown(server_addr, server_shutdown, None, false) - .await - .unwrap(); + run_server_with_auth_config( + server_addr, + server_shutdown, + None, + false, + auth_config(server_addr), + ) + .await + .unwrap(); }); let key = "sf-backend"; @@ -667,9 +1059,15 @@ async fn subscribe_missing_key_closes_without_hanging() { let shutdown_token = CancellationToken::new(); let server_shutdown = shutdown_token.clone(); let server = tokio::spawn(async move { - run_server_with_shutdown(server_addr, server_shutdown, None, false) - .await - .unwrap(); + run_server_with_auth_config( + server_addr, + server_shutdown, + None, + false, + auth_config(server_addr), + ) + .await + .unwrap(); }); let mut stream = timeout(Duration::from_secs(2), async { @@ -697,7 +1095,11 @@ async fn subscribe_missing_key_closes_without_hanging() { let result = timeout(Duration::from_millis(200), reader.read_msg()) .await .expect("missing-key subscribe hung instead of closing"); - assert!(result.is_err()); + let response = PbConnResponse::decode(result.unwrap()).unwrap(); + let PbConnResponse::Error(error) = response else { + panic!("expected structured missing-service error"); + }; + assert_eq!(error.code, "service_not_available"); shutdown_token.cancel(); server.await.unwrap(); @@ -712,9 +1114,15 @@ async fn subscribe_bypasses_unacked_stale_control_connection() { let shutdown_token = CancellationToken::new(); let server_shutdown = shutdown_token.clone(); let server = tokio::spawn(async move { - run_server_with_shutdown(server_addr, server_shutdown, None, false) - .await - .unwrap(); + run_server_with_auth_config( + server_addr, + server_shutdown, + None, + false, + auth_config(server_addr), + ) + .await + .unwrap(); }); let key = "sf-backend"; @@ -817,9 +1225,15 @@ async fn subscribe_bypasses_acked_control_connection_without_stream() { let shutdown_token = CancellationToken::new(); let server_shutdown = shutdown_token.clone(); let server = tokio::spawn(async move { - run_server_with_shutdown(server_addr, server_shutdown, None, false) - .await - .unwrap(); + run_server_with_auth_config( + server_addr, + server_shutdown, + None, + false, + auth_config(server_addr), + ) + .await + .unwrap(); }); let key = "sf-backend"; diff --git a/tests/test_delay.rs b/tests/test_delay.rs index 533f330..b80b87d 100644 --- a/tests/test_delay.rs +++ b/tests/test_delay.rs @@ -2,17 +2,19 @@ use std::env; use std::sync::LazyLock; use std::time::Duration; +use pb_mapper::common::auth::{AuthConfig, LegacyProtocolPolicy}; use pb_mapper::common::config::init_tracing; use pb_mapper::common::message::{ MessageReader, MessageWriter, NormalMessageReader, NormalMessageWriter, }; use pb_mapper::local::client::run_client_side_cli; use pb_mapper::local::server::{run_server_side_cli, ServerTunnelOptions}; -use pb_mapper::pb_server::run_server; +use pb_mapper::pb_server::run_server_with_auth_config; use rand::RngExt; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::UdpSocket; use tokio::time::{timeout, Instant}; +use tokio_util::sync::CancellationToken; use uni_stream::addr::ToSocketAddrs; use uni_stream::stream::{ListenerProvider, TcpListenerProvider, UdpListenerProvider}; use uni_stream::stream::{StreamProvider, StreamSplit, TcpStreamProvider, UdpStreamProvider}; @@ -109,7 +111,20 @@ async fn run_udp_echo_server(addr: &str) -> Result<(), Box().ok()) + .unwrap_or_default(); + let auth_config = AuthConfig { + state_dir: std::env::temp_dir() + .join(format!("pb-mapper-delay-{}-{port}", std::process::id())), + max_temporary_keys: 64, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + if let Err(e) = + run_server_with_auth_config(addr, CancellationToken::new(), None, false, auth_config).await + { eprintln!("pb-mapper server failed to start: {e}"); } } @@ -131,6 +146,8 @@ async fn run_pb_mapper_server_cli( need_codec, is_datagram: true, keep_alive: false, + namespace: None, + force_namespace: false, }, ) .await @@ -144,6 +161,8 @@ async fn run_pb_mapper_server_cli( need_codec, is_datagram: false, keep_alive: false, + namespace: None, + force_namespace: false, }, ) .await diff --git a/ui/lib/l10n/app_en.arb b/ui/lib/l10n/app_en.arb index b58cf4f..19245f4 100644 --- a/ui/lib/l10n/app_en.arb +++ b/ui/lib/l10n/app_en.arb @@ -41,9 +41,9 @@ "setupServerBody": "Host and port. All traffic goes through it.", "setupServerHint": "example.com:7666", "setupServerInvalid": "Use host:port", - "setupKeyLabel": "Header key (optional)", - "setupKeyBody": "Only if your server was started with one.", - "setupKeyInvalid": "Must be exactly 32 characters", + "setupKeyLabel": "Credential", + "setupKeyBody": "Use a temporary credential issued by the relay administrator. The administrator key also works.", + "setupKeyInvalid": "Use a 32-character administrator key or a pbmt1_ temporary credential", "setupCheckingServer": "Checking…", "setupServerOk": "Server reached", "setupServerFailed": "Not reachable. You can continue anyway.", @@ -161,7 +161,7 @@ "saving": "Saving…", "checkServer": "Check Server Connectivity", "serverAddressHelp": "Address of the pb-mapper server to connect to", - "msgHeaderKeyHelp": "Used for the message checksum and encryption handshake. 32 characters when set.", + "msgHeaderKeyHelp": "Required. Prefer a pbmt1_ temporary credential issued by the relay; the 32-character administrator key also works.", "keepAliveHelp": "Enable TCP keep-alive for connections", "configServerAddress": "Server Address: {value}", "@configServerAddress": { @@ -206,7 +206,7 @@ }, "saveFailed": "Failed to save the configuration", "serverCheckFailed": "Server check failed", - "keyLengthInvalid": "The header key must be exactly 32 characters", + "keyLengthInvalid": "Use a 32-character administrator key or a pbmt1_ temporary credential", "registerTitle": "Register Service", "registerAction": "Register & Start", "registeredList": "Registered Services ({count})", diff --git a/ui/lib/l10n/app_zh.arb b/ui/lib/l10n/app_zh.arb index cf946cd..25e5ddc 100644 --- a/ui/lib/l10n/app_zh.arb +++ b/ui/lib/l10n/app_zh.arb @@ -41,9 +41,9 @@ "setupServerBody": "地址和端口，所有流量都经过它。", "setupServerHint": "example.com:7666", "setupServerInvalid": "请用 host:port 格式", - "setupKeyLabel": "消息头密钥（可选）", - "setupKeyBody": "仅当服务器启动时带了密钥。", - "setupKeyInvalid": "必须是 32 个字符", + "setupKeyLabel": "连接凭据", + "setupKeyBody": "请使用中继管理员签发的临时凭据；管理员密钥也可以连接。", + "setupKeyInvalid": "请输入 32 字符管理员密钥或 pbmt1_ 临时凭据", "setupCheckingServer": "检测中…", "setupServerOk": "已连上服务器", "setupServerFailed": "连不上，也可以先继续。", @@ -161,7 +161,7 @@ "saving": "保存中…", "checkServer": "检测服务器连通性", "serverAddressHelp": "要连接的 pb-mapper 服务器地址", - "msgHeaderKeyHelp": "用于消息校验与加密握手，填写时必须是 32 个字符。", + "msgHeaderKeyHelp": "必填。优先使用中继签发的 pbmt1_ 临时凭据，也可使用 32 字符管理员密钥。", "keepAliveHelp": "为连接启用 TCP keep-alive", "configServerAddress": "服务器地址：{value}", "@configServerAddress": { @@ -206,7 +206,7 @@ }, "saveFailed": "保存配置失败", "serverCheckFailed": "服务器检测失败", - "keyLengthInvalid": "消息头密钥必须是 32 个字符", + "keyLengthInvalid": "请输入 32 字符管理员密钥或 pbmt1_ 临时凭据", "registerTitle": "注册服务", "registerAction": "注册并启动", "registeredList": "已注册服务（{count}）", diff --git a/ui/lib/src/views/configuration_view.dart b/ui/lib/src/views/configuration_view.dart index de7059a..23c3bd0 100644 --- a/ui/lib/src/views/configuration_view.dart +++ b/ui/lib/src/views/configuration_view.dart @@ -35,11 +35,8 @@ class _ConfigurationViewState extends State { ChangeSubscription? _changes; - @override - void initState() { - super.initState(); // Reload when anything changes this list, including a change made @@ -47,20 +44,19 @@ class _ConfigurationViewState extends State { // from a terminal while this window was open. _changes = ChangeSubscription.listen( - PbMapperService.changeStream, {StateChangeKind.config}, - (_) { if (mounted) _loadConfig(); }, - + (_) { + if (mounted) _loadConfig(); + }, ); _loadConfig(); } @override void dispose() { - _changes?.cancel(); _serverAddressController.dispose(); _msgHeaderKeyController.dispose(); @@ -90,7 +86,7 @@ class _ConfigurationViewState extends State { Future _saveConfiguration() async { if (_isSaving) return; // Prevent multiple simultaneous saves final msgHeaderKey = _msgHeaderKeyController.text.trim(); - if (msgHeaderKey.isNotEmpty && msgHeaderKey.length != 32) { + if (msgHeaderKey.length != 32 && !msgHeaderKey.startsWith('pbmt1_')) { showToast(context, context.l10n.keyLengthInvalid, kind: ToastKind.error); return; } @@ -308,9 +304,9 @@ class _ConfigurationViewState extends State { if (serverAddress.isEmpty) { throw const FormatException('serverAddress is required'); } - if (msgHeaderKey.isNotEmpty && msgHeaderKey.length != 32) { + if (msgHeaderKey.length != 32 && !msgHeaderKey.startsWith('pbmt1_')) { throw const FormatException( - 'MSG_HEADER_KEY must be exactly 32 characters', + 'MSG_HEADER_KEY must be a 32-character administrator key or a pbmt1_ temporary credential', ); } @@ -364,7 +360,7 @@ class _ConfigurationViewState extends State { controller: _msgHeaderKeyController, decoration: InputDecoration( labelText: 'MSG_HEADER_KEY', - hintText: '32 characters, or empty', + hintText: '32-character admin key or pbmt1_ credential', border: OutlineInputBorder(), helperText: context.l10n.msgHeaderKeyHelp, ), diff --git a/ui/lib/src/views/setup_wizard_view.dart b/ui/lib/src/views/setup_wizard_view.dart index 364c3e8..538d5be 100644 --- a/ui/lib/src/views/setup_wizard_view.dart +++ b/ui/lib/src/views/setup_wizard_view.dart @@ -158,7 +158,7 @@ class _SetupWizardViewState extends State { setState(() => _error = l10n.setupServerInvalid); return; } - if (key.isNotEmpty && key.length != 32) { + if (key.length != 32 && !key.startsWith('pbmt1_')) { setState(() => _error = l10n.setupKeyInvalid); return; } diff --git a/ui/native/pb_mapper_ffi/src/state.rs b/ui/native/pb_mapper_ffi/src/state.rs index 71d4bf5..632f487 100644 --- a/ui/native/pb_mapper_ffi/src/state.rs +++ b/ui/native/pb_mapper_ffi/src/state.rs @@ -12,7 +12,7 @@ use tokio::sync::{Mutex, RwLock}; use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; -use pb_mapper::common::checksum::set_process_msg_header_key; +use pb_mapper::common::checksum::{parse_credential, set_process_msg_header_key}; use pb_mapper::common::config::{get_pb_mapper_server_async, get_sockaddr_async}; use pb_mapper::common::message::command::{PbConnStatusReq, PbConnStatusResp}; use pb_mapper::local::client::status::get_status; @@ -205,11 +205,7 @@ fn normalize_msg_header_key(msg_header_key: String) -> Result if normalized.is_empty() { return Ok(normalized); } - if normalized.len() != 32 { - return Err(CtlError::invalid_argument( - "MSG_HEADER_KEY must be exactly 32 bytes (256-bit) when provided", - )); - } + parse_credential(&normalized).map_err(CtlError::invalid_argument)?; Ok(normalized) } @@ -944,6 +940,8 @@ impl PbMapperState { need_codec: enable_encryption, is_datagram: false, keep_alive: enable_keep_alive, + namespace: None, + force_namespace: false, }, Some(callback), ) @@ -959,6 +957,8 @@ impl PbMapperState { need_codec: enable_encryption, is_datagram: true, keep_alive: enable_keep_alive, + namespace: None, + force_namespace: false, }, Some(callback), ) diff --git a/ui/test/widget_test.dart b/ui/test/widget_test.dart index 6c2e1ea..c7e2dda 100644 --- a/ui/test/widget_test.dart +++ b/ui/test/widget_test.dart @@ -818,7 +818,12 @@ void main() { await tester.tap(find.text('Next')); await tester.pump(); - expect(find.text('Must be exactly 32 characters'), findsOneWidget); + expect( + find.text( + 'Use a 32-character administrator key or a pbmt1_ temporary credential', + ), + findsOneWidget, + ); }); testWidgets('server-only mode starts at the server question', (tester) async { From 1884717c709f27390721eb32c18cbc2ee70f0dde Mon Sep 17 00:00:00 2001 From: LB7666 Date: Tue, 18 Aug 2026 04:55:28 +0800 Subject: [PATCH 02/74] Make codec test configuration-independent --- src/utils/codec.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/utils/codec.rs b/src/utils/codec.rs index c5a00dd..59be873 100644 --- a/src/utils/codec.rs +++ b/src/utils/codec.rs @@ -197,8 +197,9 @@ mod tests { #[test] fn test_codec() { + const TEST_KEY: [u8; 32] = [0x42; 32]; let data = String::from("fdafas反对fdasfasfasfsdafdasfsdfasd范德萨发顺🤣❤️😁😍👍👍丰十大大师傅士大夫大撒发射点发士大夫大师傅大师傅士大夫士大夫阿斯蒂芬大师傅阿斯顿法大师傅看叫阿三的发就可是大家发开始打客服开始大幅喀什的开发点卡收费就开始打客服就是的咖啡肯定撒法开始打客服就是的咖啡就开始大幅扣税的急啊看发叫阿三的发生的开发就是大家可是大家发看大数据开发大数据开发大家ask发就是的咖啡的萨芬就卡死的房价开始打家开发商的JFK上的飞机卡上的纠纷开始打飞机宽带技术开发就开始大家开发建设的卡JFK大数据风控静安寺的看法角度看萨芬卡上的纠纷看静安寺的看法角度思考积分可是大家发卡是大家看法就大肆砍伐尽快打算减肥肯定是积分开始大幅技术大咖积分开始打飞机扣税的急啊看发的技术开发就是JFK十大福克斯大家开发大撒发射点幅度萨芬撒旦发发收范德萨发顺丰士大夫十大阿斯蒂芬大师傅阿斯顿附件是的客服对接撒巨大石块积分的课时费阿斯蒂芬法大师傅大师傅十大法大师傅阿斯蒂芬阿斯顿法大师傅阿斯蒂芬大师傅阿斯顿法大师傅大师傅阿斯蒂芬阿斯蒂芬士大夫阿斯蒂芬大师傅的萨芬打算减肥上岛咖啡加快速度大数据开发就是打客服看大数据开发就开始减肥卡萨丁JFK是大家看法加快速度JFK技术大咖积分喀什的开发独守空房技术大咖积分空手道解放扣税的开发商的开发接口是大家看法角度看是否扣税的急啊看发生的开发的快速减肥开始大幅就是打客服卡上的纠纷啊撒旦解放扣税的急啊看发加快速度点卡JFK啥的但是法大师傅技术大咖积分卡萨丁就反馈是大家看法啊是大家看法卡上的纠纷可是大家发喀什的开发大卡司喀什的开发就是打客服法大师傅士大夫的式咖啡机上岛咖啡就是的咖啡艰苦大师傅看上雕刻技法喀什的开发上岛咖啡就喀什的开发就是打客服卡上的纠纷技术的咖啡机肯定撒开发啊十大科技开发速度加啊反馈就是的咖啡开始大幅大师傅似的十大放假啊上岛咖啡就可是大家发空间的是否撒旦士大夫的撒娇开发是大家看法大肆砍伐就喀什的开发氨基酸的考虑非军事对抗疗法金克拉撒旦发艰苦拉萨的飞机喀什打开发就可是大家发可是大家看附件卡上的纠纷卡刷点卡技术的咖啡机可是大家发卡是大家看法静安寺的看法就可是大家发卡萨丁就开发商的急啊看飞机迪斯科发技术的咖啡机可是大家发看电视剧开发商大开始打到发大水发大水"); - let mut cryption = Aes256GcmCodec::try_new_with_default_key().unwrap(); + let mut cryption = Aes256GcmCodec::try_new(&TEST_KEY).unwrap(); let mut out_buf = data.as_bytes().to_vec(); let tag = { let _timer = Timer::new_with_hint("Encrypt".into()); From 2fa4c46613031e5176720e34c888729562223f12 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Tue, 18 Aug 2026 05:38:13 +0800 Subject: [PATCH 03/74] Harden temporary credential authentication --- Cargo.lock | 7 + Cargo.toml | 2 + docs/authentication-v2.md | 62 +- docs/authentication-v2.zh-CN.md | 49 +- src/bin/pb-mapper.rs | 559 +---- src/bin/pb-mapper/admin.rs | 602 +++++ src/common/auth.rs | 2028 +---------------- src/common/auth/actor.rs | 872 +++++++ src/common/auth/persistence.rs | 674 ++++++ src/common/auth/runtime.rs | 483 ++++ src/common/auth/tests.rs | 307 +++ src/common/auth/timing_wheel.rs | 104 + src/common/checksum.rs | 39 +- src/common/message/command.rs | 15 + src/common/message/secure.rs | 526 +---- src/common/message/secure/frame.rs | 236 ++ src/common/message/secure/limiter.rs | 75 + src/common/message/secure/replay.rs | 87 + src/common/message/secure/tests.rs | 171 ++ src/pb_server/admin.rs | 65 +- src/pb_server/connection.rs | 522 +++++ src/pb_server/mod.rs | 1413 +----------- src/pb_server/runtime.rs | 925 ++++++++ tests/regression.rs | 147 +- ui/native/pb_mapper_ffi/src/state.rs | 1235 +--------- .../pb_mapper_ffi/src/state/configuration.rs | 320 +++ ui/native/pb_mapper_ffi/src/state/runtime.rs | 436 ++++ ui/native/pb_mapper_ffi/src/state/status.rs | 466 ++++ 28 files changed, 6733 insertions(+), 5694 deletions(-) create mode 100644 src/bin/pb-mapper/admin.rs create mode 100644 src/common/auth/actor.rs create mode 100644 src/common/auth/persistence.rs create mode 100644 src/common/auth/runtime.rs create mode 100644 src/common/auth/tests.rs create mode 100644 src/common/auth/timing_wheel.rs create mode 100644 src/common/message/secure/frame.rs create mode 100644 src/common/message/secure/limiter.rs create mode 100644 src/common/message/secure/replay.rs create mode 100644 src/common/message/secure/tests.rs create mode 100644 src/pb_server/connection.rs create mode 100644 src/pb_server/runtime.rs create mode 100644 ui/native/pb_mapper_ffi/src/state/configuration.rs create mode 100644 ui/native/pb_mapper_ffi/src/state/runtime.rs create mode 100644 ui/native/pb_mapper_ffi/src/state/status.rs diff --git a/Cargo.lock b/Cargo.lock index b829fa7..bc6876c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -794,6 +794,7 @@ dependencies = [ "serde_json", "snafu", "socket2 0.6.1", + "subtle", "tokio", "tokio-util", "tracing", @@ -1147,6 +1148,12 @@ version = "0.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" +[[package]] +name = "subtle" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" + [[package]] name = "syn" version = "2.0.114" diff --git a/Cargo.toml b/Cargo.toml index a127ba0..47d0b77 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,6 +27,7 @@ once_cell.workspace = true uni-stream.workspace = true kanal.workspace = true base64.workspace = true +subtle.workspace = true [dev-dependencies] dotenvy = "0.15.7" @@ -66,6 +67,7 @@ trust-dns-resolver = { version = "0.23.2" } ring = "0.17.14" once_cell = "1.20.2" base64 = "0.22.1" +subtle = "2.6.1" uni-stream = { git = "https://github.com/acking-you/uni-stream.git", branch = "master" } kanal = { git = "https://github.com/acking-you/kanal.git", branch = "dev/pb-mapper" } diff --git a/docs/authentication-v2.md b/docs/authentication-v2.md index 8ce2c83..96967c6 100644 --- a/docs/authentication-v2.md +++ b/docs/authentication-v2.md @@ -75,11 +75,14 @@ Every new client writes this 32-byte clear-text routing prefix: | 1 | Flags, currently `0` | | 2 | Reserved, currently `0` | | 8 | Big-endian key ID; `0` means administrator | -| 16 | Random connection salt | +| 16 | Connection salt: 8-byte Unix timestamp plus 8 random bytes | The prefix is not secret. It is authenticated as associated data on every -encrypted frame. Unsupported flags, versions, and non-zero reserved bytes are -rejected before request dispatch. +encrypted frame. Unsupported flags, versions, non-zero reserved bytes, and +timestamps outside the five-minute clock-skew window are rejected before +request dispatch. The encrypted first request is capped at 64 KiB before +authentication, while authenticated continuation frames retain the normal +protocol limit. ### Directional frame keys @@ -109,11 +112,14 @@ uses server-to-client counter `0`. Later control frames continue from counter ### Replay resistance -The relay fingerprints `(key_id, connection_salt)` and checks two rotating -1 MiB Bloom filters covering the current and previous 60-second windows. A -probable duplicate returns the stable retryable error +The relay fingerprints `(key_id, connection_salt)` and atomically checks and +inserts it in two rotating 1 MiB Bloom filters covering the current and previous +60-second windows. A probable duplicate returns the stable retryable error `connection_salt_replayed`; one-shot administrator CLI operations retry once -with a fresh salt. +with a fresh salt. Mutating administrator requests additionally claim their +exact fingerprint in the encrypted WAL before dispatch. Those claims survive +restart and compaction for ten minutes, so an old captured mutation cannot be +replayed after the Bloom window or a process restart. ## Credential lifecycle @@ -153,11 +159,14 @@ pb-mapper admin --server relay.example.com:7666 key gc ### Root rotation and state reset -Root rotation writes an empty snapshot encrypted with the new key, appends the -audit record, persists `admin.key`, and then switches live state. It invalidates -all temporary credentials and closes connections authenticated with the old -administrator or temporary keys. The CLI stages the candidate key before the -request and verifies the new key with an authenticated status call. +Root rotation writes an empty snapshot encrypted with the new key, preserves +the bounded audit history, persists `admin.key`, and then switches the key and +administrator lease as one state transition. It invalidates all temporary +credentials and closes connections authenticated with the old administrator or +temporary keys. The CLI stages the candidate key before the request and verifies +the new key with an authenticated status call. When `--key-file` is omitted, +the recovery copy is written below `$XDG_CONFIG_HOME/pb-mapper` (or +`$HOME/.config/pb-mapper`) rather than requiring local `/var/lib` access. An explicit auth-state reset also invalidates all temporary credentials. It rotates the server instance ID so credentials from a corrupted or lost slot @@ -200,7 +209,14 @@ The default state directory is `/var/lib/pb-mapper/auth`: The directory is mode `0700`. Mutating operations acknowledge only after the WAL record is synced. The actor compacts state every five minutes with an atomic -snapshot replacement and WAL truncation. +snapshot replacement and WAL truncation. The snapshot carries the bounded +audit history and active administrator replay claims, so compaction does not +discard either security record. + +The Flutter server uses its application config directory's `auth/` child and +does not report itself running until both the TCP listener and authentication +state have initialized successfully. This keeps desktop/mobile starts writable +without pretending that a failed `/var/lib` initialization succeeded. Invalid authentication-state headers, failed integrity checks, truncated WAL records, schema mismatch, and failed compaction place temporary authentication @@ -222,8 +238,9 @@ pb-mapper admin --server relay.example.com:7666 root-key rotate ``` `--output human|json|ndjson` controls rendering. Pages default to 100 and are -capped at 1000. `--all` follows pages and emits one NDJSON object per item so a -large inventory does not need to be buffered by the CLI. +capped at 1000. `--all` follows every page while preserving the selected output +format. NDJSON is the streaming choice for large inventories; JSON emits one +combined document and human output emits one combined table. Stable structured errors contain `code`, `message`, `retryable`, and `server_time`. Authentication failure logs include stage, key ID, peer, and @@ -279,13 +296,16 @@ state. ## Code index - Credential format and process configuration: `src/common/checksum.rs` -- Slot table, timing wheel, encrypted WAL, and lifecycle actor: - `src/common/auth.rs` -- V2 framing, key derivation, counters, and replay filter: - `src/common/message/secure.rs` -- Namespace dispatch and relay resource limits: `src/pb_server/mod.rs` +- Authentication facade and shared model: `src/common/auth.rs` +- Lifecycle actor, persistence, runtime, and timing wheel: + `src/common/auth/{actor,persistence,runtime,timing_wheel}.rs` +- V2 session facade plus frame, limiter, and replay modules: + `src/common/message/secure.rs` and `src/common/message/secure/` +- Relay state, runtime loop, and connection dispatch: + `src/pb_server/{mod,runtime,connection}.rs` - Administrator request execution: `src/pb_server/admin.rs` -- Unified command surface: `src/bin/pb-mapper.rs` +- Unified CLI and administrator command module: `src/bin/pb-mapper.rs` and + `src/bin/pb-mapper/admin.rs` ## Summary diff --git a/docs/authentication-v2.zh-CN.md b/docs/authentication-v2.zh-CN.md index 9faa921..b713334 100644 --- a/docs/authentication-v2.zh-CN.md +++ b/docs/authentication-v2.zh-CN.md @@ -66,10 +66,11 @@ TCP 连接仍有自己的 V2 首帧，但不会在其上再做多轮鉴权交换 | 1 | Flags，当前必须为 `0` | | 2 | Reserved，当前必须为 `0` | | 8 | 大端 key ID；`0` 表示管理员 | -| 16 | 随机 connection salt | +| 16 | connection salt：8 字节 Unix 时间戳 + 8 字节随机数 | -前缀不承担保密作用，但会作为每个加密帧的 AAD 被完整认证。未知版本、flags 或 -reserved 值会在请求分发前被拒绝。 +前缀不承担保密作用，但会作为每个加密帧的 AAD 被完整认证。未知版本、flags、 +reserved 值或超出五分钟时钟偏差窗口的时间戳会在请求分发前被拒绝。未认证的首个 +加密请求上限为 64 KiB；鉴权完成后的后续帧仍沿用正常协议上限。 ### 双向密钥与计数器 @@ -86,9 +87,11 @@ HKDF-SHA256 使用 connection salt 作为 salt，凭据的 32 字节 secret 作 ### 重放检测 -服务端对 `(key_id, connection_salt)` 做指纹，并使用两个轮换的 1 MiB Bloom filter -覆盖当前与上一个 60 秒窗口。疑似重复会返回可重试错误 -`connection_salt_replayed`；一次性 admin CLI 会自动换 salt 重试一次。 +服务端对 `(key_id, connection_salt)` 做指纹，并在同一个临界区内完成两个轮换的 +1 MiB Bloom filter 的检查与写入，覆盖当前与上一个 60 秒窗口。疑似重复会返回 +可重试错误 `connection_salt_replayed`；一次性 admin CLI 会自动换 salt 重试一次。 +会修改状态的管理员请求还会在分发前把精确指纹写入加密 WAL；该记录在十分钟内跨 +重启、跨 compact 保留，不能通过等待 Bloom 窗口结束或重启进程来重放旧操作。 ## 临时凭据生命周期 @@ -116,9 +119,12 @@ tombstone 以给出稳定错误后，槽位可以复用。显式 `key gc` 可立 ### 根密钥轮换与状态重置 -根密钥轮换先用新密钥写空 snapshot、追加审计、持久化 `admin.key`，再切换内存状态。 -它会使全部临时凭据失效，并关闭旧管理员或临时凭据建立的连接。CLI 在发请求前保存 -候选 key，完成后再用新 key 执行一次 `admin status` 验证。 +根密钥轮换先用新密钥写空 snapshot，同时保留有上限的审计历史，持久化 +`admin.key`，再把密钥与管理员 lease 作为一次状态变更切换。它会使全部临时凭据失效， +并关闭旧管理员或临时凭据建立的连接。CLI 在发请求前保存候选 key，完成后再用新 key +执行一次 `admin status` 验证。未指定 `--key-file` 时，恢复副本默认写到 +`$XDG_CONFIG_HOME/pb-mapper`（或 `$HOME/.config/pb-mapper`），不要求本机能写 +`/var/lib`。 `auth-state reset --confirm` 同样会清空临时凭据，并轮换 server instance ID。这样即使 原槽位表损坏或丢失，旧凭据也不会因为未来复用了相同 key ID 而重新有效。 @@ -156,8 +162,13 @@ tombstone 以给出稳定错误后，槽位可以复用。显式 `key gc` 可立 | `auth.wal` | 带长度前缀、逐条加密的 mutation 与 audit | 变更只有在 WAL 同步成功后才对外确认。后台 actor 每五分钟原子替换 snapshot 并截断 -WAL。无效文件头、完整性验证失败、WAL 截断、schema 不匹配或 compact 失败都会进入 -safe mode：临时凭据全部 fail closed，管理员仍可查看状态并执行显式 reset。 +WAL；snapshot 同时保存有上限的审计历史和仍有效的管理员重放声明，compact 不会丢弃 +这些安全记录。无效文件头、完整性验证失败、WAL 截断、schema 不匹配或 compact 失败 +都会进入 safe mode：临时凭据全部 fail closed，管理员仍可查看状态并执行显式 reset。 + +Flutter 启动服务端时使用应用配置目录下的 `auth/` 子目录，并且只有 TCP listener 与 +认证状态都初始化成功后才会报告 running；桌面和移动端无需写 `/var/lib`，初始化失败 +时也不会出现虚假的运行状态。 ## 管理命令与输出 @@ -173,8 +184,9 @@ pb-mapper admin --server relay.example.com:7666 auth-state reset --confirm pb-mapper admin --server relay.example.com:7666 root-key rotate ``` -`--output human|json|ndjson` 控制展示格式。默认每页 100，最大 1000；`--all` 自动翻页 -并逐行输出 NDJSON，避免 CLI 一次缓存完整列表。稳定错误结构包含 `code`、`message`、 +`--output human|json|ndjson` 控制展示格式。默认每页 100，最大 1000；`--all` 自动翻完 +所有页面并保留选定的输出格式。大列表应选择 NDJSON 流式输出；JSON 输出单个合并文档， +human 输出单个合并表格。稳定错误结构包含 `code`、`message`、 `retryable` 与 `server_time`。 日志记录 auth stage、key ID、peer 与 reason，但不记录凭据。相同 @@ -219,11 +231,14 @@ Docker 必须持久化 `/var/lib/pb-mapper/auth`；否则重建容器会产生 ## 代码索引 - 凭据格式与进程配置：`src/common/checksum.rs` -- 固定槽位、时间轮、加密 WAL 与生命周期 actor：`src/common/auth.rs` -- V2 帧、双向派生、计数器与 replay filter：`src/common/message/secure.rs` -- 命名空间分发与资源限制：`src/pb_server/mod.rs` +- 认证 facade 与共享模型：`src/common/auth.rs` +- 生命周期 actor、持久化、runtime 与时间轮： + `src/common/auth/{actor,persistence,runtime,timing_wheel}.rs` +- V2 session facade、frame、限流与 replay 模块： + `src/common/message/secure.rs` 与 `src/common/message/secure/` +- 中继状态、runtime loop 与连接分发：`src/pb_server/{mod,runtime,connection}.rs` - 管理请求执行：`src/pb_server/admin.rs` -- 统一 CLI：`src/bin/pb-mapper.rs` +- 统一 CLI 与管理员命令模块：`src/bin/pb-mapper.rs`、`src/bin/pb-mapper/admin.rs` ## 总结 diff --git a/src/bin/pb-mapper.rs b/src/bin/pb-mapper.rs index f738bac..8a26084 100644 --- a/src/bin/pb-mapper.rs +++ b/src/bin/pb-mapper.rs @@ -6,7 +6,7 @@ use std::time::Duration; use better_mimalloc_rs::MiMalloc; use clap::{Args, Parser, Subcommand, ValueEnum}; use pb_mapper::common::auth::{ - generate_admin_key, initialize_admin_key, write_admin_key_file, LegacyProtocolPolicy, + generate_admin_key, initialize_admin_key, write_admin_key_file, KeyPage, LegacyProtocolPolicy, DEFAULT_AUTH_STATE_DIR, }; use pb_mapper::common::checksum::set_process_msg_header_key; @@ -15,7 +15,8 @@ use pb_mapper::common::config::{ get_pb_mapper_server_async, get_sockaddr_async, init_tracing, keep_alive_from_env, StatusOp, }; use pb_mapper::common::message::command::{ - AdminRequest, AdminResponse, MessageSerializer, PbConnRequest, PbConnResponse, + AdminConnectionPage, AdminRequest, AdminResponse, AdminServicePage, MessageSerializer, + PbConnRequest, PbConnResponse, }; use pb_mapper::common::message::forward::StreamForward; use pb_mapper::common::message::secure::ClientHeaderSession; @@ -93,144 +94,9 @@ struct ServerArgs { legacy_protocol: LegacyProtocolArg, } -#[derive(Debug, Args)] -struct AdminArgs { - /// Relay address. Falls back to PB_MAPPER_SERVER. - #[arg(short, long, visible_alias = "pb-mapper-server", value_name = "ADDR")] - server: Option, - /// Machine-readable output mode. - #[arg(long, value_enum, default_value_t = OutputFormat::Human)] - output: OutputFormat, - #[command(subcommand)] - command: AdminCommand, -} - -#[derive(Debug, Subcommand)] -enum AdminCommand { - /// Issue, inspect, renew, reveal, revoke, or collect temporary keys. - Key(AdminKeyArgs), - /// List relay connections across namespaces. - Connection(AdminConnectionArgs), - /// List registered services across namespaces. - Service(AdminServiceArgs), - /// Show authentication state and protocol counters. - Status, - /// Repair or reset encrypted temporary-key state. - AuthState(AdminAuthStateArgs), - /// Rotate the sole administrator key and invalidate every existing credential. - RootKey(AdminRootKeyArgs), - /// Change legacy protocol acceptance at runtime. - LegacyProtocol(AdminLegacyProtocolArgs), -} - -#[derive(Debug, Args)] -struct AdminKeyArgs { - #[command(subcommand)] - command: AdminKeyCommand, -} - -#[derive(Debug, Subcommand)] -enum AdminKeyCommand { - Issue { - #[arg(long, value_parser = parse_duration)] - ttl: Duration, - #[arg(long)] - label: Option, - }, - List { - #[arg(long, default_value_t = 0)] - page: u32, - #[arg(long, default_value_t = 100, value_parser = clap::value_parser!(u16).range(1..=1000))] - page_size: u16, - #[arg(long, default_value_t = false)] - all: bool, - }, - Show { - key_id: u64, - }, - Reveal { - key_id: u64, - }, - Renew { - key_id: u64, - #[arg(long, value_parser = parse_duration)] - ttl: Duration, - }, - Revoke { - key_id: u64, - }, - Gc, -} - -#[derive(Debug, Args)] -struct AdminConnectionArgs { - #[command(subcommand)] - command: AdminListCommand, -} - -#[derive(Debug, Args)] -struct AdminServiceArgs { - #[command(subcommand)] - command: AdminListCommand, -} - -#[derive(Debug, Clone, Subcommand)] -enum AdminListCommand { - List { - #[arg(long)] - key_id: Option, - #[arg(long, default_value_t = 0)] - page: u32, - #[arg(long, default_value_t = 100, value_parser = clap::value_parser!(u16).range(1..=1000))] - page_size: u16, - #[arg(long, default_value_t = false)] - all: bool, - }, -} - -#[derive(Debug, Args)] -struct AdminAuthStateArgs { - #[command(subcommand)] - command: AdminAuthStateCommand, -} - -#[derive(Debug, Subcommand)] -enum AdminAuthStateCommand { - Reset { - #[arg(long, default_value_t = false)] - confirm: bool, - }, -} - -#[derive(Debug, Args)] -struct AdminRootKeyArgs { - #[command(subcommand)] - command: AdminRootKeyCommand, -} - -#[derive(Debug, Subcommand)] -enum AdminRootKeyCommand { - Rotate { - /// New 32-byte administrator key. A cryptographically random printable key is generated when omitted. - #[arg(long)] - new_key: Option, - /// Save the new key here before asking the relay to rotate. - #[arg(long, default_value = "/var/lib/pb-mapper/auth/admin.key")] - key_file: PathBuf, - }, -} - -#[derive(Debug, Args)] -struct AdminLegacyProtocolArgs { - #[command(subcommand)] - command: AdminLegacyProtocolCommand, -} - -#[derive(Debug, Subcommand)] -enum AdminLegacyProtocolCommand { - Set { policy: LegacyProtocolArg }, -} - +#[path = "pb-mapper/admin.rs"] +mod admin; +use admin::AdminArgs; #[derive(Debug, Args)] struct RegisterArgs { /// Transport used by the local service. @@ -302,13 +168,6 @@ enum Transport { Udp, } -#[derive(Clone, Copy, Debug, Eq, PartialEq, ValueEnum)] -enum OutputFormat { - Human, - Json, - Ndjson, -} - #[derive(Clone, Copy, Debug, Eq, PartialEq, ValueEnum)] enum LegacyProtocolArg { Allow, @@ -342,7 +201,7 @@ async fn run(cli: Cli) -> Result<(), Box> { Command::Register(args) => run_register(args).await?, Command::Connect(args) => run_connect(args).await?, Command::Status(args) => run_status(args).await?, - Command::Admin(args) => run_admin(args).await?, + Command::Admin(args) => admin::run_admin(args).await?, } Ok(()) } @@ -466,410 +325,6 @@ async fn run_status(args: StatusArgs) -> Result<(), Box> { Ok(()) } -async fn run_admin(args: AdminArgs) -> Result<(), Box> { - let remote_addr = get_pb_mapper_server_async(args.server.as_deref()).await?; - match args.command { - AdminCommand::Key(AdminKeyArgs { command }) => match command { - AdminKeyCommand::Issue { ttl, label } => { - let response = send_admin_request( - remote_addr, - AdminRequest::KeyIssue { - ttl_seconds: ttl.as_secs(), - label, - }, - ) - .await?; - print_admin_response(args.output, &response)?; - } - AdminKeyCommand::List { - page, - page_size, - all, - } => { - stream_key_pages(remote_addr, args.output, page, page_size, all).await?; - } - AdminKeyCommand::Show { key_id } => { - let response = - send_admin_request(remote_addr, AdminRequest::KeyShow { key_id }).await?; - print_admin_response(args.output, &response)?; - } - AdminKeyCommand::Reveal { key_id } => { - let response = - send_admin_request(remote_addr, AdminRequest::KeyReveal { key_id }).await?; - print_admin_response(args.output, &response)?; - } - AdminKeyCommand::Renew { key_id, ttl } => { - let response = send_admin_request( - remote_addr, - AdminRequest::KeyRenew { - key_id, - ttl_seconds: ttl.as_secs(), - }, - ) - .await?; - print_admin_response(args.output, &response)?; - } - AdminKeyCommand::Revoke { key_id } => { - let response = - send_admin_request(remote_addr, AdminRequest::KeyRevoke { key_id }).await?; - print_admin_response(args.output, &response)?; - } - AdminKeyCommand::Gc => { - let response = send_admin_request(remote_addr, AdminRequest::KeyGc).await?; - print_admin_response(args.output, &response)?; - } - }, - AdminCommand::Connection(AdminConnectionArgs { command }) => { - let AdminListCommand::List { - key_id, - page, - page_size, - all, - } = command; - stream_connection_pages(remote_addr, args.output, key_id, page, page_size, all).await?; - } - AdminCommand::Service(AdminServiceArgs { command }) => { - let AdminListCommand::List { - key_id, - page, - page_size, - all, - } = command; - stream_service_pages(remote_addr, args.output, key_id, page, page_size, all).await?; - } - AdminCommand::Status => { - let response = send_admin_request(remote_addr, AdminRequest::AuthStatus).await?; - print_admin_response(args.output, &response)?; - } - AdminCommand::AuthState(AdminAuthStateArgs { - command: AdminAuthStateCommand::Reset { confirm }, - }) => { - let response = - send_admin_request(remote_addr, AdminRequest::AuthStateReset { confirm }).await?; - print_admin_response(args.output, &response)?; - } - AdminCommand::RootKey(AdminRootKeyArgs { - command: AdminRootKeyCommand::Rotate { new_key, key_file }, - }) => { - let new_key = new_key.unwrap_or_else(generate_admin_key); - let staged_key_file = key_file.with_file_name(format!( - ".{}.next", - key_file - .file_name() - .and_then(|name| name.to_str()) - .unwrap_or("admin.key") - )); - write_admin_key_file(&staged_key_file, &new_key, true)?; - let response = send_admin_request( - remote_addr, - AdminRequest::RootKeyRotate { - new_admin_key: new_key.clone(), - }, - ) - .await - .map_err(|error| { - std::io::Error::other(format!( - "root rotation request failed; the candidate key remains at `{}`: {error}", - staged_key_file.display() - )) - })?; - set_process_msg_header_key(Some(&new_key))?; - let verification = send_admin_request(remote_addr, AdminRequest::AuthStatus).await?; - if !matches!(verification, AdminResponse::AuthStatus(_)) { - return Err(std::io::Error::other( - "new administrator key did not pass the post-rotation status check", - ) - .into()); - } - write_admin_key_file(&key_file, &new_key, true).map_err(|error| { - std::io::Error::other(format!( - "administrator key rotated and verified, but `{}` could not be updated; recover the key from `{}`: {error}", - key_file.display(), - staged_key_file.display() - )) - })?; - if let Err(error) = std::fs::remove_file(&staged_key_file) { - tracing::warn!( - path = %staged_key_file.display(), - %error, - "administrator key was rotated, but the staged key file could not be removed" - ); - } - if args.output == OutputFormat::Human { - println!("administrator key rotated and verified"); - println!("key file: {}", key_file.display()); - } else { - print_admin_response(args.output, &response)?; - } - } - AdminCommand::LegacyProtocol(AdminLegacyProtocolArgs { - command: AdminLegacyProtocolCommand::Set { policy }, - }) => { - let response = send_admin_request( - remote_addr, - AdminRequest::LegacyProtocolSet { - policy: policy.into(), - }, - ) - .await?; - print_admin_response(args.output, &response)?; - } - } - Ok(()) -} - -async fn send_admin_request( - remote_addr: std::net::SocketAddr, - request: AdminRequest, -) -> Result> { - let encoded = PbConnRequest::Admin(request).encode()?; - for attempt in 0..2 { - let mut stream = TcpStream::connect(remote_addr).await?; - let session = ClientHeaderSession::from_process()?; - session.write_initial(&mut stream, &encoded).await?; - let mut reader = session.response_reader(&mut stream)?; - let message = reader.read_msg().await?; - match PbConnResponse::decode(message)? { - PbConnResponse::Admin(response) => return Ok(response), - PbConnResponse::Error(error) - if error.code == "connection_salt_replayed" && error.retryable && attempt == 0 => - { - continue; - } - PbConnResponse::Error(error) => { - return Err(std::io::Error::other(format!( - "{}: {} (retryable={})", - error.code, error.message, error.retryable - )) - .into()); - } - response => { - return Err(std::io::Error::other(format!( - "unexpected administrator response: {response:?}" - )) - .into()); - } - } - } - Err(std::io::Error::other("connection salt replay retry was exhausted").into()) -} - -async fn stream_key_pages( - remote_addr: std::net::SocketAddr, - output: OutputFormat, - mut page: u32, - page_size: u16, - all: bool, -) -> Result<(), Box> { - loop { - let response = - send_admin_request(remote_addr, AdminRequest::KeyList { page, page_size }).await?; - let AdminResponse::KeyList(key_page) = &response else { - return Err(std::io::Error::other("unexpected key-list response").into()); - }; - if all { - for item in &key_page.items { - println!("{}", serde_json::to_string(item)?); - } - } else { - print_admin_response(output, &response)?; - } - let Some(next_page) = key_page.next_page else { - break; - }; - if !all { - break; - } - page = next_page; - } - Ok(()) -} - -async fn stream_service_pages( - remote_addr: std::net::SocketAddr, - output: OutputFormat, - key_id: Option, - mut page: u32, - page_size: u16, - all: bool, -) -> Result<(), Box> { - loop { - let response = send_admin_request( - remote_addr, - AdminRequest::ServiceList { - key_id, - page, - page_size, - }, - ) - .await?; - let AdminResponse::Services(service_page) = &response else { - return Err(std::io::Error::other("unexpected service-list response").into()); - }; - if all { - for item in &service_page.items { - println!("{}", serde_json::to_string(item)?); - } - } else { - print_admin_response(output, &response)?; - } - let Some(next_page) = service_page.next_page else { - break; - }; - if !all { - break; - } - page = next_page; - } - Ok(()) -} - -async fn stream_connection_pages( - remote_addr: std::net::SocketAddr, - output: OutputFormat, - key_id: Option, - mut page: u32, - page_size: u16, - all: bool, -) -> Result<(), Box> { - loop { - let response = send_admin_request( - remote_addr, - AdminRequest::ConnectionList { - key_id, - page, - page_size, - }, - ) - .await?; - let AdminResponse::Connections(connection_page) = &response else { - return Err(std::io::Error::other("unexpected connection-list response").into()); - }; - if all { - for item in &connection_page.items { - println!("{}", serde_json::to_string(item)?); - } - } else { - print_admin_response(output, &response)?; - } - let Some(next_page) = connection_page.next_page else { - break; - }; - if !all { - break; - } - page = next_page; - } - Ok(()) -} - -fn print_admin_response( - output: OutputFormat, - response: &AdminResponse, -) -> Result<(), Box> { - match output { - OutputFormat::Json => println!( - "{}", - serde_json::to_string_pretty(&serde_json::json!({ - "schema_version": 1, - "data": response, - }))? - ), - OutputFormat::Ndjson => println!("{}", serde_json::to_string(response)?), - OutputFormat::Human => print_human_admin_response(response), - } - Ok(()) -} - -fn print_human_admin_response(response: &AdminResponse) { - match response { - AdminResponse::KeyIssued(key) - | AdminResponse::KeyShown(key) - | AdminResponse::KeyRenewed(key) => { - println!("key id: {}", key.metadata.key_id); - println!("state: {}", key.metadata.state); - println!("expires at: {}", key.metadata.expires_at); - if let Some(label) = &key.metadata.label { - println!("label: {label}"); - } - if !key.credential.is_empty() { - println!("credential: {}", key.credential); - } - } - AdminResponse::KeyRevoked(key) => { - println!("key {}: {}", key.key_id, key.state); - } - AdminResponse::KeyList(page) => { - println!("KEY ID\tSTATE\tEXPIRES\tLABEL"); - for key in &page.items { - println!( - "{}\t{}\t{}\t{}", - key.key_id, - key.state, - key.expires_at, - key.label.as_deref().unwrap_or("") - ); - } - if let Some(next) = page.next_page { - println!("next page: {next}"); - } - } - AdminResponse::KeyGc { removed } => println!("removed {removed} inactive keys"), - AdminResponse::AuthStatus(status) => { - println!("safe mode: {}", status.safe_mode); - println!( - "keys: {} active / {} expired / {} revoked / {} capacity", - status.active_keys, status.expired_keys, status.revoked_keys, status.capacity - ); - println!("legacy protocol: {:?}", status.legacy_protocol); - println!( - "active legacy connections: {}", - status.active_legacy_connections - ); - println!( - "last legacy connection: {}", - status - .last_legacy_connection_at - .map(|value| value.to_string()) - .unwrap_or_else(|| "never".to_string()) - ); - println!( - "authentication: {} succeeded / {} failed", - status.auth_successes, status.auth_failures - ); - println!("server instance: {}", status.server_instance_id); - } - AdminResponse::Services(page) => { - println!("KEY ID\tSERVICE\tTRANSPORT\tCODEC\tCONNECTIONS"); - for service in &page.items { - println!( - "{}\t{}\t{}\t{}\t{}", - service.key_id, - service.service_name, - service.transport, - service.codec_enabled, - service.connection_count - ); - } - } - AdminResponse::Connections(page) => { - println!("KEY ID\tSERVICE\tCONN ID\tHEALTHY\tTRANSPORT\tCODEC"); - for connection in &page.items { - println!( - "{}\t{}\t{}\t{}\t{}\t{}", - connection.key_id, - connection.service_name, - connection.conn_id, - connection.healthy, - connection.transport, - connection.codec_enabled - ); - } - } - AdminResponse::Ok { action } => println!("ok: {action}"), - } -} - fn parse_duration(raw: &str) -> Result { let raw = raw.trim(); if raw.is_empty() { diff --git a/src/bin/pb-mapper/admin.rs b/src/bin/pb-mapper/admin.rs new file mode 100644 index 0000000..bff21a1 --- /dev/null +++ b/src/bin/pb-mapper/admin.rs @@ -0,0 +1,602 @@ +use super::*; + +#[derive(Debug, Args)] +pub(super) struct AdminArgs { + /// Relay address. Falls back to PB_MAPPER_SERVER. + #[arg(short, long, visible_alias = "pb-mapper-server", value_name = "ADDR")] + server: Option, + /// Machine-readable output mode. + #[arg(long, value_enum, default_value_t = OutputFormat::Human)] + output: OutputFormat, + #[command(subcommand)] + command: AdminCommand, +} + +#[derive(Debug, Subcommand)] +enum AdminCommand { + /// Issue, inspect, renew, reveal, revoke, or collect temporary keys. + Key(AdminKeyArgs), + /// List relay connections across namespaces. + Connection(AdminConnectionArgs), + /// List registered services across namespaces. + Service(AdminServiceArgs), + /// Show authentication state and protocol counters. + Status, + /// Repair or reset encrypted temporary-key state. + AuthState(AdminAuthStateArgs), + /// Rotate the sole administrator key and invalidate every existing credential. + RootKey(AdminRootKeyArgs), + /// Change legacy protocol acceptance at runtime. + LegacyProtocol(AdminLegacyProtocolArgs), +} + +#[derive(Debug, Args)] +struct AdminKeyArgs { + #[command(subcommand)] + command: AdminKeyCommand, +} + +#[derive(Debug, Subcommand)] +enum AdminKeyCommand { + Issue { + #[arg(long, value_parser = parse_duration)] + ttl: Duration, + #[arg(long)] + label: Option, + }, + List { + #[arg(long, default_value_t = 0)] + page: u32, + #[arg(long, default_value_t = 100, value_parser = clap::value_parser!(u16).range(1..=1000))] + page_size: u16, + #[arg(long, default_value_t = false)] + all: bool, + }, + Show { + key_id: u64, + }, + Reveal { + key_id: u64, + }, + Renew { + key_id: u64, + #[arg(long, value_parser = parse_duration)] + ttl: Duration, + }, + Revoke { + key_id: u64, + }, + Gc, +} + +#[derive(Debug, Args)] +struct AdminConnectionArgs { + #[command(subcommand)] + command: AdminListCommand, +} + +#[derive(Debug, Args)] +struct AdminServiceArgs { + #[command(subcommand)] + command: AdminListCommand, +} + +#[derive(Debug, Clone, Subcommand)] +enum AdminListCommand { + List { + #[arg(long)] + key_id: Option, + #[arg(long, default_value_t = 0)] + page: u32, + #[arg(long, default_value_t = 100, value_parser = clap::value_parser!(u16).range(1..=1000))] + page_size: u16, + #[arg(long, default_value_t = false)] + all: bool, + }, +} + +#[derive(Debug, Args)] +struct AdminAuthStateArgs { + #[command(subcommand)] + command: AdminAuthStateCommand, +} + +#[derive(Debug, Subcommand)] +enum AdminAuthStateCommand { + Reset { + #[arg(long, default_value_t = false)] + confirm: bool, + }, +} + +#[derive(Debug, Args)] +struct AdminRootKeyArgs { + #[command(subcommand)] + command: AdminRootKeyCommand, +} + +#[derive(Debug, Subcommand)] +enum AdminRootKeyCommand { + Rotate { + /// New 32-byte administrator key. A cryptographically random printable key is generated when omitted. + #[arg(long)] + new_key: Option, + /// Save the new key here before asking the relay to rotate. + #[arg(long)] + key_file: Option, + }, +} + +#[derive(Debug, Args)] +struct AdminLegacyProtocolArgs { + #[command(subcommand)] + command: AdminLegacyProtocolCommand, +} + +#[derive(Debug, Subcommand)] +enum AdminLegacyProtocolCommand { + Set { policy: LegacyProtocolArg }, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq, ValueEnum)] +enum OutputFormat { + Human, + Json, + Ndjson, +} + +pub(super) async fn run_admin(args: AdminArgs) -> Result<(), Box> { + let remote_addr = get_pb_mapper_server_async(args.server.as_deref()).await?; + match args.command { + AdminCommand::Key(AdminKeyArgs { command }) => match command { + AdminKeyCommand::Issue { ttl, label } => { + let response = send_admin_request( + remote_addr, + AdminRequest::KeyIssue { + ttl_seconds: ttl.as_secs(), + label, + }, + ) + .await?; + print_admin_response(args.output, &response)?; + } + AdminKeyCommand::List { + page, + page_size, + all, + } => { + stream_key_pages(remote_addr, args.output, page, page_size, all).await?; + } + AdminKeyCommand::Show { key_id } => { + let response = + send_admin_request(remote_addr, AdminRequest::KeyShow { key_id }).await?; + print_admin_response(args.output, &response)?; + } + AdminKeyCommand::Reveal { key_id } => { + let response = + send_admin_request(remote_addr, AdminRequest::KeyReveal { key_id }).await?; + print_admin_response(args.output, &response)?; + } + AdminKeyCommand::Renew { key_id, ttl } => { + let response = send_admin_request( + remote_addr, + AdminRequest::KeyRenew { + key_id, + ttl_seconds: ttl.as_secs(), + }, + ) + .await?; + print_admin_response(args.output, &response)?; + } + AdminKeyCommand::Revoke { key_id } => { + let response = + send_admin_request(remote_addr, AdminRequest::KeyRevoke { key_id }).await?; + print_admin_response(args.output, &response)?; + } + AdminKeyCommand::Gc => { + let response = send_admin_request(remote_addr, AdminRequest::KeyGc).await?; + print_admin_response(args.output, &response)?; + } + }, + AdminCommand::Connection(AdminConnectionArgs { command }) => { + let AdminListCommand::List { + key_id, + page, + page_size, + all, + } = command; + stream_connection_pages(remote_addr, args.output, key_id, page, page_size, all).await?; + } + AdminCommand::Service(AdminServiceArgs { command }) => { + let AdminListCommand::List { + key_id, + page, + page_size, + all, + } = command; + stream_service_pages(remote_addr, args.output, key_id, page, page_size, all).await?; + } + AdminCommand::Status => { + let response = send_admin_request(remote_addr, AdminRequest::AuthStatus).await?; + print_admin_response(args.output, &response)?; + } + AdminCommand::AuthState(AdminAuthStateArgs { + command: AdminAuthStateCommand::Reset { confirm }, + }) => { + let response = + send_admin_request(remote_addr, AdminRequest::AuthStateReset { confirm }).await?; + print_admin_response(args.output, &response)?; + } + AdminCommand::RootKey(AdminRootKeyArgs { + command: AdminRootKeyCommand::Rotate { new_key, key_file }, + }) => { + let key_file = key_file.unwrap_or_else(default_admin_recovery_key_file); + let new_key = new_key.unwrap_or_else(generate_admin_key); + let staged_key_file = key_file.with_file_name(format!( + ".{}.next", + key_file + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or("admin.key") + )); + write_admin_key_file(&staged_key_file, &new_key, true)?; + let response = send_admin_request( + remote_addr, + AdminRequest::RootKeyRotate { + new_admin_key: new_key.clone(), + }, + ) + .await + .map_err(|error| { + std::io::Error::other(format!( + "root rotation request failed; the candidate key remains at `{}`: {error}", + staged_key_file.display() + )) + })?; + set_process_msg_header_key(Some(&new_key))?; + let verification = send_admin_request(remote_addr, AdminRequest::AuthStatus).await?; + if !matches!(verification, AdminResponse::AuthStatus(_)) { + return Err(std::io::Error::other( + "new administrator key did not pass the post-rotation status check", + ) + .into()); + } + write_admin_key_file(&key_file, &new_key, true).map_err(|error| { + std::io::Error::other(format!( + "administrator key rotated and verified, but `{}` could not be updated; recover the key from `{}`: {error}", + key_file.display(), + staged_key_file.display() + )) + })?; + if let Err(error) = std::fs::remove_file(&staged_key_file) { + tracing::warn!( + path = %staged_key_file.display(), + %error, + "administrator key was rotated, but the staged key file could not be removed" + ); + } + if args.output == OutputFormat::Human { + println!("administrator key rotated and verified"); + println!("key file: {}", key_file.display()); + } else { + print_admin_response(args.output, &response)?; + } + } + AdminCommand::LegacyProtocol(AdminLegacyProtocolArgs { + command: AdminLegacyProtocolCommand::Set { policy }, + }) => { + let response = send_admin_request( + remote_addr, + AdminRequest::LegacyProtocolSet { + policy: policy.into(), + }, + ) + .await?; + print_admin_response(args.output, &response)?; + } + } + Ok(()) +} + +async fn send_admin_request( + remote_addr: std::net::SocketAddr, + request: AdminRequest, +) -> Result> { + let encoded = PbConnRequest::Admin(request).encode()?; + for attempt in 0..2 { + let mut stream = TcpStream::connect(remote_addr).await?; + let session = ClientHeaderSession::from_process()?; + session.write_initial(&mut stream, &encoded).await?; + let mut reader = session.response_reader(&mut stream)?; + let message = reader.read_msg().await?; + match PbConnResponse::decode(message)? { + PbConnResponse::Admin(response) => return Ok(response), + PbConnResponse::Error(error) + if error.code == "connection_salt_replayed" && error.retryable && attempt == 0 => + { + continue; + } + PbConnResponse::Error(error) => { + return Err(std::io::Error::other(format!( + "{}: {} (retryable={})", + error.code, error.message, error.retryable + )) + .into()); + } + response => { + return Err(std::io::Error::other(format!( + "unexpected administrator response: {response:?}" + )) + .into()); + } + } + } + Err(std::io::Error::other("connection salt replay retry was exhausted").into()) +} + +async fn stream_key_pages( + remote_addr: std::net::SocketAddr, + output: OutputFormat, + mut page: u32, + page_size: u16, + all: bool, +) -> Result<(), Box> { + let mut combined: Option = None; + loop { + let response = + send_admin_request(remote_addr, AdminRequest::KeyList { page, page_size }).await?; + let AdminResponse::KeyList(key_page) = &response else { + return Err(std::io::Error::other("unexpected key-list response").into()); + }; + if all { + if output == OutputFormat::Ndjson { + for item in &key_page.items { + println!("{}", serde_json::to_string(item)?); + } + } else { + let page = combined.get_or_insert_with(|| { + let mut page = key_page.clone(); + page.items.clear(); + page.next_page = None; + page + }); + page.items.extend(key_page.items.iter().cloned()); + } + } else { + print_admin_response(output, &response)?; + } + let Some(next_page) = key_page.next_page else { + break; + }; + if !all { + break; + } + page = next_page; + } + if let Some(page) = combined { + print_admin_response(output, &AdminResponse::KeyList(page))?; + } + Ok(()) +} + +async fn stream_service_pages( + remote_addr: std::net::SocketAddr, + output: OutputFormat, + key_id: Option, + mut page: u32, + page_size: u16, + all: bool, +) -> Result<(), Box> { + let mut combined: Option = None; + loop { + let response = send_admin_request( + remote_addr, + AdminRequest::ServiceList { + key_id, + page, + page_size, + }, + ) + .await?; + let AdminResponse::Services(service_page) = &response else { + return Err(std::io::Error::other("unexpected service-list response").into()); + }; + if all { + if output == OutputFormat::Ndjson { + for item in &service_page.items { + println!("{}", serde_json::to_string(item)?); + } + } else { + let page = combined.get_or_insert_with(|| { + let mut page = service_page.clone(); + page.items.clear(); + page.next_page = None; + page + }); + page.items.extend(service_page.items.iter().cloned()); + } + } else { + print_admin_response(output, &response)?; + } + let Some(next_page) = service_page.next_page else { + break; + }; + if !all { + break; + } + page = next_page; + } + if let Some(page) = combined { + print_admin_response(output, &AdminResponse::Services(page))?; + } + Ok(()) +} + +async fn stream_connection_pages( + remote_addr: std::net::SocketAddr, + output: OutputFormat, + key_id: Option, + mut page: u32, + page_size: u16, + all: bool, +) -> Result<(), Box> { + let mut combined: Option = None; + loop { + let response = send_admin_request( + remote_addr, + AdminRequest::ConnectionList { + key_id, + page, + page_size, + }, + ) + .await?; + let AdminResponse::Connections(connection_page) = &response else { + return Err(std::io::Error::other("unexpected connection-list response").into()); + }; + if all { + if output == OutputFormat::Ndjson { + for item in &connection_page.items { + println!("{}", serde_json::to_string(item)?); + } + } else { + let page = combined.get_or_insert_with(|| { + let mut page = connection_page.clone(); + page.items.clear(); + page.next_page = None; + page + }); + page.items.extend(connection_page.items.iter().cloned()); + } + } else { + print_admin_response(output, &response)?; + } + let Some(next_page) = connection_page.next_page else { + break; + }; + if !all { + break; + } + page = next_page; + } + if let Some(page) = combined { + print_admin_response(output, &AdminResponse::Connections(page))?; + } + Ok(()) +} + +fn print_admin_response( + output: OutputFormat, + response: &AdminResponse, +) -> Result<(), Box> { + match output { + OutputFormat::Json => println!( + "{}", + serde_json::to_string_pretty(&serde_json::json!({ + "schema_version": 1, + "data": response, + }))? + ), + OutputFormat::Ndjson => println!("{}", serde_json::to_string(response)?), + OutputFormat::Human => print_human_admin_response(response), + } + Ok(()) +} + +fn print_human_admin_response(response: &AdminResponse) { + match response { + AdminResponse::KeyIssued(key) + | AdminResponse::KeyShown(key) + | AdminResponse::KeyRenewed(key) => { + println!("key id: {}", key.metadata.key_id); + println!("state: {}", key.metadata.state); + println!("expires at: {}", key.metadata.expires_at); + if let Some(label) = &key.metadata.label { + println!("label: {label}"); + } + if !key.credential.is_empty() { + println!("credential: {}", key.credential); + } + } + AdminResponse::KeyRevoked(key) => { + println!("key {}: {}", key.key_id, key.state); + } + AdminResponse::KeyList(page) => { + println!("KEY ID\tSTATE\tEXPIRES\tLABEL"); + for key in &page.items { + println!( + "{}\t{}\t{}\t{}", + key.key_id, + key.state, + key.expires_at, + key.label.as_deref().unwrap_or("") + ); + } + if let Some(next) = page.next_page { + println!("next page: {next}"); + } + } + AdminResponse::KeyGc { removed } => println!("removed {removed} inactive keys"), + AdminResponse::AuthStatus(status) => { + println!("safe mode: {}", status.safe_mode); + println!( + "keys: {} active / {} expired / {} revoked / {} capacity", + status.active_keys, status.expired_keys, status.revoked_keys, status.capacity + ); + println!("legacy protocol: {:?}", status.legacy_protocol); + println!( + "active legacy connections: {}", + status.active_legacy_connections + ); + println!( + "last legacy connection: {}", + status + .last_legacy_connection_at + .map(|value| value.to_string()) + .unwrap_or_else(|| "never".to_string()) + ); + println!( + "authentication: {} succeeded / {} failed", + status.auth_successes, status.auth_failures + ); + println!("server instance: {}", status.server_instance_id); + } + AdminResponse::Services(page) => { + println!("KEY ID\tSERVICE\tTRANSPORT\tCODEC\tCONNECTIONS"); + for service in &page.items { + println!( + "{}\t{}\t{}\t{}\t{}", + service.key_id, + service.service_name, + service.transport, + service.codec_enabled, + service.connection_count + ); + } + } + AdminResponse::Connections(page) => { + println!("KEY ID\tSERVICE\tCONN ID\tHEALTHY\tTRANSPORT\tCODEC"); + for connection in &page.items { + println!( + "{}\t{}\t{}\t{}\t{}\t{}", + connection.key_id, + connection.service_name, + connection.conn_id, + connection.healthy, + connection.transport, + connection.codec_enabled + ); + } + } + AdminResponse::Ok { action } => println!("ok: {action}"), + } +} + +fn default_admin_recovery_key_file() -> PathBuf { + std::env::var_os("XDG_CONFIG_HOME") + .map(PathBuf::from) + .or_else(|| std::env::var_os("HOME").map(|home| PathBuf::from(home).join(".config"))) + .unwrap_or_else(|| PathBuf::from(".")) + .join("pb-mapper") + .join("admin.key") +} diff --git a/src/common/auth.rs b/src/common/auth.rs index 1b3b66e..39fdbe1 100644 --- a/src/common/auth.rs +++ b/src/common/auth.rs @@ -5,7 +5,7 @@ //! stores only lifecycle metadata plus a weak lease reference. The background actor owns //! the strong leases through a hierarchical timing wheel. -use std::collections::{HashMap, VecDeque}; +use std::collections::{HashMap, HashSet, VecDeque}; use std::fmt; use std::fs::{File, OpenOptions}; use std::io::{Read, Write}; @@ -20,12 +20,14 @@ use rand::RngExt; use ring::aead::{Aad, LessSafeKey, Nonce, UnboundKey, AES_256_GCM}; use ring::hkdf::{Salt, HKDF_SHA256}; use serde::{Deserialize, Serialize}; +use subtle::ConstantTimeEq; use tokio::sync::{mpsc, oneshot}; use tokio_util::sync::CancellationToken; use super::checksum::{ encode_temporary_credential, get_process_credential, parse_credential, - set_process_msg_header_key, AesKeyType, Credential, MACHINE_MSG_HEADER_KEY_PATH, + set_process_msg_header_key, AesKeyType, Credential, ENV_MSG_HEADER_KEY, + MACHINE_MSG_HEADER_KEY_PATH, }; pub const ADMIN_NAMESPACE: u64 = 0; @@ -39,6 +41,9 @@ const SNAPSHOT_SCHEMA_VERSION: u16 = 1; const STATE_BLOB_MAGIC: &[u8; 5] = b"PBAS1"; const STATE_AAD: &[u8] = b"pb-mapper-auth-state-v1"; const INSTANCE_ID_LEN: usize = 16; +const ADMIN_REPLAY_RETENTION: Duration = Duration::from_secs(10 * 60); +const ADMIN_REPLAY_CAPACITY: usize = 65_536; +const AUDIT_RECORD_CAPACITY: usize = 4096; #[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] @@ -222,6 +227,18 @@ impl AuthContext { pub fn cancellation_token(&self) -> Result { Ok(self.ensure_active()?.cancellation_token()) } + + fn admin_authority(&self) -> Result, AuthFailure> { + if !self.is_admin { + return Err(AuthFailure::new( + "admin_permission_required", + "administrator credential is required for this operation", + false, + )); + } + self.ensure_active()?; + Ok(self.lease.clone()) + } } #[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)] @@ -252,10 +269,15 @@ impl Default for SlotHot { } } +#[derive(Debug)] +struct AdminState { + key: AesKeyType, + lease: Weak, +} + #[derive(Debug)] struct AuthStateInner { - admin_key: RwLock, - admin_lease: RwLock>, + admin: RwLock, instance_id: RwLock<[u8; INSTANCE_ID_LEN]>, slots: RwLock>, safe_mode: AtomicBool, @@ -264,14 +286,15 @@ struct AuthStateInner { last_legacy_connection_at: AtomicU64, auth_successes: AtomicU64, auth_failures: AtomicU64, + audit_records: RwLock>, } impl AuthStateInner { fn admin_key(&self) -> AesKeyType { - *self - .admin_key + self.admin .read() .unwrap_or_else(|poisoned| poisoned.into_inner()) + .key } fn instance_id(&self) -> [u8; INSTANCE_ID_LEN] { @@ -335,48 +358,65 @@ struct ColdMetadata { } enum AuthCommand { + ClaimAdminMutation { + authority: Weak, + fingerprint: [u8; 32], + client_timestamp: u64, + response: oneshot::Sender>, + }, Issue { + authority: Weak, ttl: Duration, label: Option, response: oneshot::Sender>, }, List { + authority: Weak, page: u32, page_size: u16, response: oneshot::Sender>, }, Show { + authority: Weak, key_id: u64, reveal: bool, response: oneshot::Sender>, }, Renew { + authority: Weak, key_id: u64, ttl: Duration, response: oneshot::Sender>, }, Revoke { + authority: Weak, key_id: u64, response: oneshot::Sender>, }, Gc { + authority: Weak, response: oneshot::Sender>, }, Reset { + authority: Weak, response: oneshot::Sender>, }, RotateRoot { + authority: Weak, new_key: AesKeyType, response: oneshot::Sender>, }, SetLegacyProtocol { + authority: Weak, policy: LegacyProtocolPolicy, response: oneshot::Sender>, }, Status { + authority: Weak, response: oneshot::Sender>, }, Audit { + authority: Weak, action: String, key_id: Option, detail: Option, @@ -384,355 +424,7 @@ enum AuthCommand { }, } -impl AuthRuntime { - pub async fn from_process(config: AuthConfig) -> Result { - prepare_state_dir(&config.state_dir)?; - let credential = load_server_admin_credential(&config.state_dir)?; - let Credential::Admin(admin_key) = credential else { - return Err(AuthFailure::new( - "administrator_key_required", - "the relay server must start with the administrator credential", - false, - )); - }; - Self::start(admin_key, config).await - } - - pub async fn start(admin_key: AesKeyType, config: AuthConfig) -> Result { - prepare_state_dir(&config.state_dir)?; - let instance_id = load_or_create_instance_id(&config.state_dir)?; - let (loaded, safe_mode) = load_persisted_state(&config, &admin_key, instance_id); - let mut slots = (0..config.max_temporary_keys) - .map(|_| SlotHot::default()) - .collect::>() - .into_boxed_slice(); - let mut cold = HashMap::new(); - let mut wheel = TimingWheel::new(unix_seconds()); - let now = unix_seconds(); - - let admin_lease = Arc::new(AuthLease::new(0, u64::MAX)); - if let Some(state) = loaded.as_ref() { - for (index, generation) in state.generations.iter().copied().enumerate() { - if let Some(slot) = slots.get_mut(index) { - slot.generation = generation; - } - } - for entry in &state.entries { - let index = key_slot(entry.key_id) as usize; - let Some(slot) = slots.get_mut(index) else { - continue; - }; - if slot.generation != key_generation(entry.key_id) { - continue; - } - let state = if entry.state == SlotState::Active && entry.expires_at <= now { - SlotState::Expired - } else { - entry.state - }; - slot.state = state; - slot.expires_at = entry.expires_at; - cold.insert( - entry.key_id, - ColdMetadata { - issued_at: entry.issued_at, - label: entry.label.clone(), - }, - ); - if state == SlotState::Active { - let lease = Arc::new(AuthLease::new(entry.key_id, entry.expires_at)); - slot.lease = Arc::downgrade(&lease); - wheel.insert(lease); - } - } - } - - let legacy_protocol = loaded - .as_ref() - .map(|state| state.legacy_protocol) - .unwrap_or(config.legacy_protocol); - let inner = Arc::new(AuthStateInner { - admin_key: RwLock::new(admin_key), - admin_lease: RwLock::new(Arc::downgrade(&admin_lease)), - instance_id: RwLock::new(instance_id), - slots: RwLock::new(slots), - safe_mode: AtomicBool::new(safe_mode), - legacy_protocol_allowed: AtomicBool::new(legacy_protocol.is_allowed()), - active_legacy_connections: AtomicU64::new(0), - last_legacy_connection_at: AtomicU64::new(0), - auth_successes: AtomicU64::new(0), - auth_failures: AtomicU64::new(0), - }); - let (command_tx, command_rx) = mpsc::channel(256); - let runtime = Self { - inner: Arc::downgrade(&inner), - command_tx, - config: config.clone(), - }; - - tokio::spawn(run_auth_actor( - inner, - admin_lease, - command_rx, - config, - cold, - wheel, - )); - Ok(runtime) - } - - pub fn config(&self) -> &AuthConfig { - &self.config - } - - fn inner(&self) -> Result, AuthFailure> { - self.inner.upgrade().ok_or_else(|| { - AuthFailure::new( - "auth_state_unavailable", - "authentication state manager is not running", - true, - ) - }) - } - - pub fn admin_key(&self) -> Result { - Ok(self.inner()?.admin_key()) - } - - pub fn derive_key(&self, key_id: u64) -> Result { - let inner = self.inner()?; - if key_id == 0 { - return Ok(inner.admin_key()); - } - derive_temporary_key(&inner.admin_key(), &inner.instance_id(), key_id) - } - - pub fn authenticate(&self, key_id: u64) -> Result { - let inner = self.inner()?; - if key_id == 0 { - let lease = inner - .admin_lease - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .upgrade() - .ok_or_else(|| { - AuthFailure::new( - "administrator_key_rotated", - "administrator credential was rotated", - false, - ) - })?; - inner.auth_successes.fetch_add(1, Ordering::Relaxed); - return Ok(AuthContext::from_lease(0, true, &lease)); - } - if inner.safe_mode.load(Ordering::Acquire) { - inner.auth_failures.fetch_add(1, Ordering::Relaxed); - return Err(AuthFailure::new( - "temporary_key_store_unavailable", - "temporary key state is unavailable; administrator reset is required", - false, - )); - } - - let index = key_slot(key_id) as usize; - let generation = key_generation(key_id); - let slots = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let Some(slot) = slots.get(index) else { - inner.auth_failures.fetch_add(1, Ordering::Relaxed); - return Err(AuthFailure::new( - "temporary_key_not_found", - "temporary key id is outside the configured slot table", - false, - )); - }; - if slot.generation != generation { - inner.auth_failures.fetch_add(1, Ordering::Relaxed); - return Err(AuthFailure::new( - "temporary_key_generation_mismatch", - "temporary key generation does not match the current slot", - false, - )); - } - let failure = match slot.state { - SlotState::Free => Some(AuthFailure::new( - "temporary_key_not_found", - "temporary key does not exist", - false, - )), - SlotState::Expired => Some(AuthFailure::new( - "temporary_key_expired", - "temporary key has expired", - false, - )), - SlotState::Revoked => Some(AuthFailure::new( - "temporary_key_revoked", - "temporary key was revoked", - false, - )), - SlotState::Active if slot.expires_at <= unix_seconds() => { - if let Some(lease) = slot.lease.upgrade() { - lease.cancellation.cancel(); - } - Some(AuthFailure::new( - "temporary_key_expired", - "temporary key has expired", - false, - )) - } - SlotState::Active => None, - }; - if let Some(failure) = failure { - inner.auth_failures.fetch_add(1, Ordering::Relaxed); - return Err(failure); - } - let lease = slot.lease.upgrade().ok_or_else(|| { - inner.auth_failures.fetch_add(1, Ordering::Relaxed); - AuthFailure::new( - "temporary_key_inactive", - "temporary key lease is no longer active", - true, - ) - })?; - inner.auth_successes.fetch_add(1, Ordering::Relaxed); - Ok(AuthContext::from_lease(key_id, false, &lease)) - } - - pub fn legacy_protocol_allowed(&self) -> Result { - Ok(self - .inner()? - .legacy_protocol_allowed - .load(Ordering::Acquire)) - } - - pub fn record_legacy_connection(&self) -> Result { - let inner = self.inner()?; - inner - .active_legacy_connections - .fetch_add(1, Ordering::AcqRel); - inner - .last_legacy_connection_at - .store(unix_seconds(), Ordering::Release); - Ok(LegacyConnectionGuard { - inner: Arc::downgrade(&inner), - }) - } - - async fn request( - &self, - build: impl FnOnce(oneshot::Sender>) -> AuthCommand, - ) -> Result { - let (response, receiver) = oneshot::channel(); - self.command_tx.send(build(response)).await.map_err(|_| { - AuthFailure::new( - "auth_state_unavailable", - "authentication state manager is not running", - true, - ) - })?; - receiver.await.map_err(|_| { - AuthFailure::new( - "auth_state_unavailable", - "authentication state manager dropped the response", - true, - ) - })? - } - - pub async fn issue( - &self, - ttl: Duration, - label: Option, - ) -> Result { - self.request(|response| AuthCommand::Issue { - ttl, - label, - response, - }) - .await - } - - pub async fn list(&self, page: u32, page_size: u16) -> Result { - self.request(|response| AuthCommand::List { - page, - page_size, - response, - }) - .await - } - - pub async fn show(&self, key_id: u64, reveal: bool) -> Result { - self.request(|response| AuthCommand::Show { - key_id, - reveal, - response, - }) - .await - } - - pub async fn renew( - &self, - key_id: u64, - ttl: Duration, - ) -> Result { - self.request(|response| AuthCommand::Renew { - key_id, - ttl, - response, - }) - .await - } - - pub async fn revoke(&self, key_id: u64) -> Result { - self.request(|response| AuthCommand::Revoke { key_id, response }) - .await - } - - pub async fn gc(&self) -> Result { - self.request(|response| AuthCommand::Gc { response }).await - } - - pub async fn reset(&self) -> Result<(), AuthFailure> { - self.request(|response| AuthCommand::Reset { response }) - .await - } - - pub async fn rotate_root(&self, new_key: AesKeyType) -> Result<(), AuthFailure> { - self.request(|response| AuthCommand::RotateRoot { new_key, response }) - .await - } - - pub async fn set_legacy_protocol( - &self, - policy: LegacyProtocolPolicy, - ) -> Result<(), AuthFailure> { - self.request(|response| AuthCommand::SetLegacyProtocol { policy, response }) - .await - } - - pub async fn status(&self) -> Result { - self.request(|response| AuthCommand::Status { response }) - .await - } - - pub async fn audit_admin( - &self, - action: impl Into, - key_id: Option, - detail: Option, - ) -> Result<(), AuthFailure> { - let action = action.into(); - self.request(|response| AuthCommand::Audit { - action, - key_id, - detail, - response, - }) - .await - } -} +mod runtime; pub struct LegacyConnectionGuard { inner: Weak, @@ -793,7 +485,9 @@ fn load_server_admin_credential(state_dir: &Path) -> Result, entries: Vec, legacy_protocol: LegacyProtocolPolicy, + #[serde(default)] + admin_replays: Vec, + #[serde(default)] + audit_records: VecDeque, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +struct AdminReplayRecord { + fingerprint: [u8; 32], + client_timestamp: u64, } #[derive(Clone, Debug, Serialize, Deserialize)] @@ -939,1610 +643,14 @@ enum WalRecord { audit: AuditRecord, }, Audit(AuditRecord), + AdminReplay(AdminReplayRecord), } -async fn run_auth_actor( - inner: Arc, - mut admin_lease: Arc, - mut command_rx: mpsc::Receiver, - config: AuthConfig, - mut cold: HashMap, - mut wheel: TimingWheel, -) { - let now = unix_seconds(); - let mut tombstones = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .iter() - .enumerate() - .filter_map(|(index, slot)| { - matches!(slot.state, SlotState::Expired | SlotState::Revoked).then_some(( - now.saturating_add(TOMBSTONE_RETENTION.as_secs()), - make_key_id(slot.generation, index as u32), - )) - }) - .collect::>(); - let mut last_snapshot_at = unix_seconds(); - let mut tick = tokio::time::interval(Duration::from_secs(1)); - tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - loop { - tokio::select! { - _ = tick.tick() => { - let now = unix_seconds(); - for lease in wheel.advance(now) { - let key_id = lease.key_id(); - let version = lease.wheel_version.load(Ordering::Acquire); - if lease.expires_at() > now { - wheel.insert_with_version(lease, version); - continue; - } - let mut slots = inner.slots.write().unwrap_or_else(|poisoned| poisoned.into_inner()); - if let Some(slot) = slots.get_mut(key_slot(key_id) as usize) { - if slot.generation == key_generation(key_id) && slot.state == SlotState::Active { - slot.state = SlotState::Expired; - lease.cancellation.cancel(); - tombstones.push_back((now.saturating_add(TOMBSTONE_RETENTION.as_secs()), key_id)); - tracing::info!( - event = "temporary_key_expired", - auth_stage = "expiry", - key_id, - expires_at = lease.expires_at(), - "temporary key expired and active work was cancelled" - ); - } - } - } - while let Some((cleanup_at, key_id)) = tombstones.front().copied() { - if cleanup_at > now { - break; - } - tombstones.pop_front(); - let mut slots = inner.slots.write().unwrap_or_else(|poisoned| poisoned.into_inner()); - if let Some(slot) = slots.get_mut(key_slot(key_id) as usize) { - if slot.generation == key_generation(key_id) && matches!(slot.state, SlotState::Expired | SlotState::Revoked) { - slot.state = SlotState::Free; - slot.expires_at = 0; - slot.lease = Weak::new(); - cold.remove(&key_id); - } - } - } - if now.saturating_sub(last_snapshot_at) >= SNAPSHOT_COMPACTION_INTERVAL.as_secs() { - let snapshot = build_snapshot(&inner, &cold); - if let Err(error) = write_snapshot_and_truncate_wal( - &config, - &inner.admin_key(), - &snapshot, - ) { - inner.safe_mode.store(true, Ordering::Release); - cancel_all_temporary_leases(&inner); - tracing::error!( - event = "auth_state_safe_mode", - auth_stage = "snapshot_compaction", - reason = %error.code, - error = %error, - "authentication state compaction failed closed" - ); - } else { - last_snapshot_at = now; - } - } - } - command = command_rx.recv() => { - let Some(command) = command else { - admin_lease.cancellation.cancel(); - cancel_all_temporary_leases(&inner); - break; - }; - match command { - AuthCommand::Issue { ttl, label, response } => { - let result = actor_issue(&inner, &config, &mut cold, &mut wheel, ttl, label); - let _ = response.send(result); - } - AuthCommand::List { page, page_size, response } => { - let _ = response.send(actor_list(&inner, &cold, page, page_size)); - } - AuthCommand::Show { key_id, reveal, response } => { - let result = actor_show(&inner, &config, &cold, key_id, reveal); - let _ = response.send(result); - } - AuthCommand::Renew { key_id, ttl, response } => { - let result = actor_renew(&inner, &config, &cold, &mut wheel, key_id, ttl); - let _ = response.send(result); - } - AuthCommand::Revoke { key_id, response } => { - let result = actor_revoke(&inner, &config, &cold, &mut tombstones, key_id); - let _ = response.send(result); - } - AuthCommand::Gc { response } => { - let result = actor_gc(&inner, &config, &mut cold, &mut tombstones); - let _ = response.send(result); - } - AuthCommand::Reset { response } => { - let result = actor_reset(&inner, &config, &mut cold, &mut wheel, "auth_state_reset"); - let _ = response.send(result); - } - AuthCommand::RotateRoot { new_key, response } => { - let result = actor_rotate_root(&inner, &config, &mut cold, &mut wheel, &mut admin_lease, new_key); - let _ = response.send(result); - } - AuthCommand::SetLegacyProtocol { policy, response } => { - let result = actor_set_legacy_protocol(&inner, &config, policy); - let _ = response.send(result); - } - AuthCommand::Status { response } => { - let _ = response.send(Ok(actor_status(&inner))); - } - AuthCommand::Audit { action, key_id, detail, response } => { - let result = append_audit( - &config, - &inner.admin_key(), - audit(&action, key_id, detail), - ); - let _ = response.send(result); - } - } - } - } - } -} - -fn validate_ttl(config: &AuthConfig, ttl: Duration) -> Result { - if ttl < MIN_TEMP_KEY_TTL { - return Err(AuthFailure::new( - "temporary_key_ttl_too_short", - format!( - "temporary key TTL must be at least {} seconds", - MIN_TEMP_KEY_TTL.as_secs() - ), - false, - )); - } - if ttl > config.max_temporary_key_ttl { - return Err(AuthFailure::new( - "temporary_key_ttl_too_long", - format!( - "temporary key TTL exceeds the configured maximum of {} seconds", - config.max_temporary_key_ttl.as_secs() - ), - false, - )); - } - Ok(unix_seconds().saturating_add(ttl.as_secs())) -} - -fn validate_label(label: Option) -> Result, AuthFailure> { - let label = label - .map(|label| label.trim().to_string()) - .filter(|label| !label.is_empty()); - if label.as_ref().is_some_and(|label| label.len() > 64) { - return Err(AuthFailure::new( - "temporary_key_label_too_long", - "temporary key label must not exceed 64 UTF-8 bytes", - false, - )); - } - Ok(label) -} - -fn actor_issue( - inner: &Arc, - config: &AuthConfig, - cold: &mut HashMap, - wheel: &mut TimingWheel, - ttl: Duration, - label: Option, -) -> Result { - ensure_store_available(inner)?; - let expires_at = validate_ttl(config, ttl)?; - let label = validate_label(label)?; - let issued_at = unix_seconds(); - let mut slots = inner - .slots - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let Some((index, slot)) = slots - .iter_mut() - .enumerate() - .find(|(_, slot)| slot.state == SlotState::Free && slot.generation < u32::MAX) - else { - return Err(AuthFailure::new( - "temporary_key_capacity_exhausted", - "temporary key slot table is full", - true, - )); - }; - let generation = slot.generation + 1; - let key_id = make_key_id(generation, index as u32); - let entry = PersistedEntry { - key_id, - state: SlotState::Active, - issued_at, - expires_at, - label: label.clone(), - }; - append_mutation( - config, - &inner.admin_key(), - StateMutation::Issue(entry.clone()), - audit("temporary_key_issue", Some(key_id), label.clone()), - )?; - let lease = Arc::new(AuthLease::new(key_id, expires_at)); - slot.generation = generation; - slot.state = SlotState::Active; - slot.expires_at = expires_at; - slot.lease = Arc::downgrade(&lease); - cold.insert(key_id, ColdMetadata { issued_at, label }); - wheel.insert(lease); - drop(slots); - metadata_with_credential(inner, cold, key_id, true) -} - -fn actor_list( - inner: &Arc, - cold: &HashMap, - page: u32, - page_size: u16, -) -> Result { - let page_size = page_size.clamp(1, 1000) as usize; - let start = (page as usize).saturating_mul(page_size); - let slots = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let mut all = slots - .iter() - .enumerate() - .filter_map(|(index, slot)| { - if slot.state == SlotState::Free { - return None; - } - let key_id = make_key_id(slot.generation, index as u32); - let cold = cold.get(&key_id)?; - Some(TemporaryKeyMetadata { - key_id, - state: slot_state_name(slot.state).to_string(), - issued_at: cold.issued_at, - expires_at: slot.expires_at, - label: cold.label.clone(), - }) - }) - .collect::>(); - all.sort_by_key(|item| std::cmp::Reverse(item.issued_at)); - let items = all.iter().skip(start).take(page_size).cloned().collect(); - let next_page = (start.saturating_add(page_size) < all.len()).then_some(page.saturating_add(1)); - Ok(KeyPage { - schema_version: 1, - items, - next_page, - }) -} - -fn actor_show( - inner: &Arc, - config: &AuthConfig, - cold: &HashMap, - key_id: u64, - reveal: bool, -) -> Result { - let result = metadata_with_credential(inner, cold, key_id, reveal)?; - append_audit( - config, - &inner.admin_key(), - audit( - if reveal { - "temporary_key_reveal" - } else { - "temporary_key_show" - }, - Some(key_id), - result.metadata.label.clone(), - ), - )?; - Ok(result) -} - -fn actor_renew( - inner: &Arc, - config: &AuthConfig, - cold: &HashMap, - wheel: &mut TimingWheel, - key_id: u64, - ttl: Duration, -) -> Result { - ensure_store_available(inner)?; - let expires_at = validate_ttl(config, ttl)?; - let index = key_slot(key_id) as usize; - let mut slots = inner - .slots - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let slot = slots.get_mut(index).ok_or_else(|| key_not_found(key_id))?; - validate_slot_identity(slot, key_id)?; - if slot.state != SlotState::Active || slot.expires_at <= unix_seconds() { - return Err(AuthFailure::new( - "temporary_key_not_renewable", - "only an active, unexpired temporary key can be renewed", - false, - )); - } - let label = cold - .get(&key_id) - .and_then(|metadata| metadata.label.clone()); - append_mutation( - config, - &inner.admin_key(), - StateMutation::Renew { key_id, expires_at }, - audit("temporary_key_renew", Some(key_id), label), - )?; - let lease = slot.lease.upgrade().ok_or_else(|| { - AuthFailure::new( - "temporary_key_inactive", - "temporary key lease is no longer active", - true, - ) - })?; - slot.expires_at = expires_at; - lease.expires_at.store(expires_at, Ordering::Release); - lease.wheel_version.fetch_add(1, Ordering::AcqRel); - wheel.insert(lease); - drop(slots); - metadata_with_credential(inner, cold, key_id, true) -} - -fn actor_revoke( - inner: &Arc, - config: &AuthConfig, - cold: &HashMap, - tombstones: &mut VecDeque<(u64, u64)>, - key_id: u64, -) -> Result { - ensure_store_available(inner)?; - let now = unix_seconds(); - let index = key_slot(key_id) as usize; - let mut slots = inner - .slots - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let slot = slots.get_mut(index).ok_or_else(|| key_not_found(key_id))?; - validate_slot_identity(slot, key_id)?; - if slot.state != SlotState::Active { - return Err(AuthFailure::new( - "temporary_key_not_active", - "temporary key is not active", - false, - )); - } - let cold_metadata = cold.get(&key_id).ok_or_else(|| key_not_found(key_id))?; - append_mutation( - config, - &inner.admin_key(), - StateMutation::Revoke { key_id, at: now }, - audit( - "temporary_key_revoke", - Some(key_id), - cold_metadata.label.clone(), - ), - )?; - slot.state = SlotState::Revoked; - if let Some(lease) = slot.lease.upgrade() { - lease.cancellation.cancel(); - } - tombstones.push_back((now.saturating_add(TOMBSTONE_RETENTION.as_secs()), key_id)); - Ok(TemporaryKeyMetadata { - key_id, - state: slot_state_name(slot.state).to_string(), - issued_at: cold_metadata.issued_at, - expires_at: slot.expires_at, - label: cold_metadata.label.clone(), - }) -} - -fn actor_gc( - inner: &Arc, - config: &AuthConfig, - cold: &mut HashMap, - tombstones: &mut VecDeque<(u64, u64)>, -) -> Result { - ensure_store_available(inner)?; - let now = unix_seconds(); - let mut removed = 0_u64; - let mut slots = inner - .slots - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - for (index, slot) in slots.iter_mut().enumerate() { - if matches!(slot.state, SlotState::Expired | SlotState::Revoked) - || (slot.state == SlotState::Active && slot.expires_at <= now) - { - let key_id = make_key_id(slot.generation, index as u32); - if let Some(lease) = slot.lease.upgrade() { - lease.cancellation.cancel(); - } - slot.state = SlotState::Free; - slot.expires_at = 0; - slot.lease = Weak::new(); - cold.remove(&key_id); - removed = removed.saturating_add(1); - } - } - tombstones.clear(); - drop(slots); - let snapshot = build_snapshot(inner, cold); - let admin_key = inner.admin_key(); - if let Err(error) = - write_snapshot_and_truncate_wal(config, &admin_key, &snapshot).and_then(|()| { - append_audit( - config, - &admin_key, - audit("temporary_key_gc", None, Some(format!("removed={removed}"))), - ) - }) - { - inner.safe_mode.store(true, Ordering::Release); - cancel_all_temporary_leases(inner); - return Err(error); - } - Ok(removed) -} - -fn actor_reset( - inner: &Arc, - config: &AuthConfig, - cold: &mut HashMap, - wheel: &mut TimingWheel, - action: &str, -) -> Result<(), AuthFailure> { - let new_instance_id = random_instance_id(); - let snapshot = empty_snapshot(inner, new_instance_id); - let admin_key = inner.admin_key(); - if let Err(error) = write_snapshot_and_truncate_wal(config, &admin_key, &snapshot) - .and_then(|()| append_audit(config, &admin_key, audit(action, None, None))) - { - inner.safe_mode.store(true, Ordering::Release); - cancel_all_temporary_leases(inner); - return Err(error); - } - if let Err(error) = atomic_write( - &config.state_dir.join("server-instance-id"), - &new_instance_id, - 0o600, - ) { - inner.safe_mode.store(true, Ordering::Release); - cancel_all_temporary_leases(inner); - return Err(error); - } - - cancel_all_temporary_leases(inner); - { - let mut slots = inner - .slots - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - for slot in slots.iter_mut() { - slot.state = SlotState::Free; - slot.expires_at = 0; - slot.lease = Weak::new(); - } - } - *inner - .instance_id - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()) = new_instance_id; - cold.clear(); - wheel.clear(unix_seconds()); - inner.safe_mode.store(false, Ordering::Release); - Ok(()) -} - -fn actor_rotate_root( - inner: &Arc, - config: &AuthConfig, - cold: &mut HashMap, - wheel: &mut TimingWheel, - admin_lease: &mut Arc, - new_key: AesKeyType, -) -> Result<(), AuthFailure> { - if new_key == inner.admin_key() { - return Err(AuthFailure::new( - "administrator_key_unchanged", - "new administrator key must differ from the current key", - false, - )); - } - let new_key_string = String::from_utf8(new_key.to_vec()).map_err(|_| { - AuthFailure::new( - "administrator_key_invalid", - "administrator key must be 32 UTF-8 bytes for MSG_HEADER_KEY compatibility", - false, - ) - })?; - if new_key_string.chars().any(char::is_whitespace) { - return Err(AuthFailure::new( - "administrator_key_invalid", - "administrator key must not contain whitespace", - false, - )); - } - - let snapshot = empty_snapshot(inner, inner.instance_id()); - if let Err(error) = write_snapshot_and_truncate_wal(config, &new_key, &snapshot) - .and_then(|()| { - append_audit( - config, - &new_key, - audit("administrator_key_rotate", None, None), - ) - }) - .and_then(|()| write_admin_key(&config.state_dir, &new_key_string)) - { - inner.safe_mode.store(true, Ordering::Release); - cancel_all_temporary_leases(inner); - return Err(error); - } - - cancel_all_temporary_leases(inner); - let old_admin_lease = admin_lease.clone(); - { - let mut slots = inner - .slots - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - for slot in slots.iter_mut() { - slot.state = SlotState::Free; - slot.expires_at = 0; - slot.lease = Weak::new(); - } - } - cold.clear(); - wheel.clear(unix_seconds()); - let new_admin_lease = Arc::new(AuthLease::new(0, u64::MAX)); - *inner - .admin_key - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()) = new_key; - *inner - .admin_lease - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()) = Arc::downgrade(&new_admin_lease); - set_process_msg_header_key(Some(&new_key_string)).map_err(AuthFailure::internal)?; - inner.safe_mode.store(false, Ordering::Release); - old_admin_lease.cancellation.cancel(); - *admin_lease = new_admin_lease; - Ok(()) -} - -fn actor_set_legacy_protocol( - inner: &Arc, - config: &AuthConfig, - policy: LegacyProtocolPolicy, -) -> Result<(), AuthFailure> { - ensure_store_available(inner)?; - append_mutation( - config, - &inner.admin_key(), - StateMutation::LegacyProtocol(policy), - audit("legacy_protocol_update", None, Some(format!("{policy:?}"))), - )?; - inner - .legacy_protocol_allowed - .store(policy.is_allowed(), Ordering::Release); - Ok(()) -} - -fn actor_status(inner: &Arc) -> AuthStatus { - let slots = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let active_keys = slots - .iter() - .filter(|slot| slot.state == SlotState::Active) - .count(); - let expired_keys = slots - .iter() - .filter(|slot| slot.state == SlotState::Expired) - .count(); - let revoked_keys = slots - .iter() - .filter(|slot| slot.state == SlotState::Revoked) - .count(); - let last_legacy_connection_at = inner.last_legacy_connection_at.load(Ordering::Acquire); - AuthStatus { - schema_version: 1, - safe_mode: inner.safe_mode.load(Ordering::Acquire), - capacity: slots.len(), - active_keys, - expired_keys, - revoked_keys, - legacy_protocol: if inner.legacy_protocol_allowed.load(Ordering::Acquire) { - LegacyProtocolPolicy::Allow - } else { - LegacyProtocolPolicy::Deny - }, - active_legacy_connections: inner.active_legacy_connections.load(Ordering::Acquire), - last_legacy_connection_at: (last_legacy_connection_at != 0) - .then_some(last_legacy_connection_at), - auth_successes: inner.auth_successes.load(Ordering::Relaxed), - auth_failures: inner.auth_failures.load(Ordering::Relaxed), - server_instance_id: hex(&inner.instance_id()), - } -} - -fn ensure_store_available(inner: &AuthStateInner) -> Result<(), AuthFailure> { - if inner.safe_mode.load(Ordering::Acquire) { - Err(AuthFailure::new( - "temporary_key_store_unavailable", - "temporary key store is in administrator safe mode", - false, - )) - } else { - Ok(()) - } -} - -fn validate_slot_identity(slot: &SlotHot, key_id: u64) -> Result<(), AuthFailure> { - if slot.generation != key_generation(key_id) || slot.state == SlotState::Free { - Err(key_not_found(key_id)) - } else { - Ok(()) - } -} - -fn key_not_found(key_id: u64) -> AuthFailure { - AuthFailure::new( - "temporary_key_not_found", - format!("temporary key {key_id} does not exist"), - false, - ) -} - -fn metadata_with_credential( - inner: &Arc, - cold: &HashMap, - key_id: u64, - reveal: bool, -) -> Result { - let slots = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let slot = slots - .get(key_slot(key_id) as usize) - .ok_or_else(|| key_not_found(key_id))?; - validate_slot_identity(slot, key_id)?; - let cold = cold.get(&key_id).ok_or_else(|| key_not_found(key_id))?; - let credential = if reveal { - let key = derive_temporary_key(&inner.admin_key(), &inner.instance_id(), key_id)?; - encode_temporary_credential(key_id, &key) - } else { - String::new() - }; - Ok(IssuedTemporaryKey { - metadata: TemporaryKeyMetadata { - key_id, - state: slot_state_name(slot.state).to_string(), - issued_at: cold.issued_at, - expires_at: slot.expires_at, - label: cold.label.clone(), - }, - credential, - }) -} - -fn slot_state_name(state: SlotState) -> &'static str { - match state { - SlotState::Free => "free", - SlotState::Active => "active", - SlotState::Expired => "expired", - SlotState::Revoked => "revoked", - } -} - -fn audit(action: &str, key_id: Option, label: Option) -> AuditRecord { - AuditRecord { - at: unix_seconds(), - action: action.to_string(), - key_id, - label, - } -} - -fn cancel_all_temporary_leases(inner: &AuthStateInner) { - let slots = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - for lease in slots.iter().filter_map(|slot| slot.lease.upgrade()) { - lease.cancellation.cancel(); - } -} - -fn build_snapshot(inner: &AuthStateInner, cold: &HashMap) -> PersistedSnapshot { - let slots = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let generations = slots.iter().map(|slot| slot.generation).collect(); - let entries = slots - .iter() - .enumerate() - .filter_map(|(index, slot)| { - if slot.state == SlotState::Free { - return None; - } - let key_id = make_key_id(slot.generation, index as u32); - let cold = cold.get(&key_id)?; - Some(PersistedEntry { - key_id, - state: slot.state, - issued_at: cold.issued_at, - expires_at: slot.expires_at, - label: cold.label.clone(), - }) - }) - .collect(); - PersistedSnapshot { - schema_version: SNAPSHOT_SCHEMA_VERSION, - instance_id: inner.instance_id(), - generations, - entries, - legacy_protocol: if inner.legacy_protocol_allowed.load(Ordering::Acquire) { - LegacyProtocolPolicy::Allow - } else { - LegacyProtocolPolicy::Deny - }, - } -} - -fn empty_snapshot(inner: &AuthStateInner, instance_id: [u8; INSTANCE_ID_LEN]) -> PersistedSnapshot { - let slots = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - PersistedSnapshot { - schema_version: SNAPSHOT_SCHEMA_VERSION, - instance_id, - generations: slots.iter().map(|slot| slot.generation).collect(), - entries: Vec::new(), - legacy_protocol: if inner.legacy_protocol_allowed.load(Ordering::Acquire) { - LegacyProtocolPolicy::Allow - } else { - LegacyProtocolPolicy::Deny - }, - } -} - -fn load_persisted_state( - config: &AuthConfig, - admin_key: &AesKeyType, - instance_id: [u8; INSTANCE_ID_LEN], -) -> (Option, bool) { - match try_load_persisted_state(config, admin_key, instance_id) { - Ok(state) => (Some(state), false), - Err(error) => { - tracing::error!( - event = "auth_state_safe_mode", - auth_stage = "state_load", - reason = %error.code, - error = %error, - "temporary key store failed closed in administrator safe mode" - ); - (None, true) - } - } -} - -fn try_load_persisted_state( - config: &AuthConfig, - admin_key: &AesKeyType, - instance_id: [u8; INSTANCE_ID_LEN], -) -> Result { - let snapshot_path = config.state_dir.join("auth.snapshot"); - let mut snapshot = if snapshot_path.exists() { - let bytes = std::fs::read(&snapshot_path).map_err(|error| { - AuthFailure::new( - "temporary_key_store_unavailable", - format!("failed to read `{}`: {error}", snapshot_path.display()), - false, - ) - })?; - let plain = open_blob(admin_key, &bytes)?; - serde_json::from_slice::(&plain).map_err(|error| { - AuthFailure::new( - "temporary_key_store_unavailable", - format!("failed to decode auth snapshot: {error}"), - false, - ) - })? - } else { - PersistedSnapshot { - schema_version: SNAPSHOT_SCHEMA_VERSION, - instance_id, - generations: vec![0; config.max_temporary_keys], - entries: Vec::new(), - legacy_protocol: config.legacy_protocol, - } - }; - if snapshot.schema_version != SNAPSHOT_SCHEMA_VERSION || snapshot.instance_id != instance_id { - return Err(AuthFailure::new( - "temporary_key_store_unavailable", - "auth snapshot schema or server instance id does not match", - false, - )); - } - snapshot.generations.resize(config.max_temporary_keys, 0); - snapshot.generations.truncate(config.max_temporary_keys); - - let wal_path = config.state_dir.join("auth.wal"); - if wal_path.exists() { - for record in read_wal(&wal_path, admin_key)? { - if let WalRecord::Mutation { mutation, .. } = record { - apply_persisted_mutation(&mut snapshot, mutation, config.max_temporary_keys)?; - } - } - } - Ok(snapshot) -} - -fn apply_persisted_mutation( - snapshot: &mut PersistedSnapshot, - mutation: StateMutation, - capacity: usize, -) -> Result<(), AuthFailure> { - match mutation { - StateMutation::Issue(entry) => { - let index = key_slot(entry.key_id) as usize; - if index >= capacity { - return Err(AuthFailure::new( - "temporary_key_store_unavailable", - "WAL issue record references a slot outside the configured capacity", - false, - )); - } - snapshot.generations[index] = key_generation(entry.key_id); - snapshot - .entries - .retain(|current| key_slot(current.key_id) as usize != index); - snapshot.entries.push(entry); - } - StateMutation::Renew { key_id, expires_at } => { - let entry = snapshot - .entries - .iter_mut() - .find(|entry| entry.key_id == key_id) - .ok_or_else(|| { - AuthFailure::new( - "temporary_key_store_unavailable", - "WAL renew record references an unknown key", - false, - ) - })?; - entry.expires_at = expires_at; - entry.state = SlotState::Active; - } - StateMutation::Revoke { key_id, .. } => { - let entry = snapshot - .entries - .iter_mut() - .find(|entry| entry.key_id == key_id) - .ok_or_else(|| { - AuthFailure::new( - "temporary_key_store_unavailable", - "WAL revoke record references an unknown key", - false, - ) - })?; - entry.state = SlotState::Revoked; - } - StateMutation::LegacyProtocol(policy) => snapshot.legacy_protocol = policy, - } - Ok(()) -} - -fn append_mutation( - config: &AuthConfig, - admin_key: &AesKeyType, - mutation: StateMutation, - audit: AuditRecord, -) -> Result<(), AuthFailure> { - append_wal(config, admin_key, &WalRecord::Mutation { mutation, audit }) -} - -fn append_audit( - config: &AuthConfig, - admin_key: &AesKeyType, - audit: AuditRecord, -) -> Result<(), AuthFailure> { - append_wal(config, admin_key, &WalRecord::Audit(audit)) -} - -fn append_wal( - config: &AuthConfig, - admin_key: &AesKeyType, - record: &WalRecord, -) -> Result<(), AuthFailure> { - let plain = serde_json::to_vec(record).map_err(|error| { - AuthFailure::internal(format!("failed to encode auth WAL record: {error}")) - })?; - let sealed = seal_blob(admin_key, &plain)?; - let length = u32::try_from(sealed.len()) - .map_err(|_| AuthFailure::internal("auth WAL record is too large"))?; - let path = config.state_dir.join("auth.wal"); - let mut file = OpenOptions::new() - .create(true) - .append(true) - .open(&path) - .map_err(|error| { - AuthFailure::new( - "temporary_key_store_unavailable", - format!("failed to open `{}`: {error}", path.display()), - true, - ) - })?; - #[cfg(unix)] - file.set_permissions(std::fs::Permissions::from_mode(0o600)) - .map_err(|error| { - AuthFailure::new( - "temporary_key_store_unavailable", - format!("failed to secure `{}`: {error}", path.display()), - false, - ) - })?; - file.write_all(&length.to_be_bytes()) - .and_then(|()| file.write_all(&sealed)) - .and_then(|()| file.sync_data()) - .map_err(|error| { - AuthFailure::new( - "temporary_key_store_unavailable", - format!("failed to durably append `{}`: {error}", path.display()), - true, - ) - }) -} - -fn read_wal(path: &Path, admin_key: &AesKeyType) -> Result, AuthFailure> { - let mut file = File::open(path).map_err(|error| { - AuthFailure::new( - "temporary_key_store_unavailable", - format!("failed to open `{}`: {error}", path.display()), - false, - ) - })?; - let mut records = Vec::new(); - loop { - let mut length = [0_u8; 4]; - match file.read(&mut length[..1]) { - Ok(0) => break, - Ok(1) => {} - Ok(_) => unreachable!("single-byte WAL prefix read"), - Err(error) => { - return Err(AuthFailure::new( - "temporary_key_store_unavailable", - format!("failed to read auth WAL length: {error}"), - false, - )); - } - } - file.read_exact(&mut length[1..]).map_err(|error| { - AuthFailure::new( - "temporary_key_store_unavailable", - format!("truncated auth WAL length: {error}"), - false, - ) - })?; - let length = u32::from_be_bytes(length) as usize; - if length > 1024 * 1024 { - return Err(AuthFailure::new( - "temporary_key_store_unavailable", - "auth WAL record exceeds 1 MiB", - false, - )); - } - let mut sealed = vec![0_u8; length]; - file.read_exact(&mut sealed).map_err(|error| { - AuthFailure::new( - "temporary_key_store_unavailable", - format!("truncated auth WAL record: {error}"), - false, - ) - })?; - let plain = open_blob(admin_key, &sealed)?; - records.push(serde_json::from_slice(&plain).map_err(|error| { - AuthFailure::new( - "temporary_key_store_unavailable", - format!("failed to decode auth WAL record: {error}"), - false, - ) - })?); - } - Ok(records) -} - -fn write_snapshot_and_truncate_wal( - config: &AuthConfig, - admin_key: &AesKeyType, - snapshot: &PersistedSnapshot, -) -> Result<(), AuthFailure> { - let plain = serde_json::to_vec(snapshot).map_err(|error| { - AuthFailure::internal(format!("failed to encode auth snapshot: {error}")) - })?; - let sealed = seal_blob(admin_key, &plain)?; - let snapshot_path = config.state_dir.join("auth.snapshot"); - atomic_write(&snapshot_path, &sealed, 0o600)?; - let wal_path = config.state_dir.join("auth.wal"); - let wal = OpenOptions::new() - .create(true) - .write(true) - .truncate(true) - .open(&wal_path) - .map_err(|error| { - AuthFailure::new( - "temporary_key_store_unavailable", - format!("failed to truncate `{}`: {error}", wal_path.display()), - true, - ) - })?; - wal.sync_all().map_err(|error| { - AuthFailure::new( - "temporary_key_store_unavailable", - format!("failed to sync `{}`: {error}", wal_path.display()), - true, - ) - }) -} - -fn seal_blob(admin_key: &AesKeyType, plain: &[u8]) -> Result, AuthFailure> { - let key = LessSafeKey::new( - UnboundKey::new(&AES_256_GCM, admin_key) - .map_err(|_| AuthFailure::internal("failed to initialize state encryption key"))?, - ); - let mut nonce_bytes = [0_u8; 12]; - let mut rng = rand::rng(); - for byte in &mut nonce_bytes { - *byte = rng.random(); - } - let mut output = plain.to_vec(); - key.seal_in_place_append_tag( - Nonce::assume_unique_for_key(nonce_bytes), - Aad::from(STATE_AAD), - &mut output, - ) - .map_err(|_| AuthFailure::internal("failed to encrypt authentication state"))?; - let mut sealed = Vec::with_capacity(STATE_BLOB_MAGIC.len() + nonce_bytes.len() + output.len()); - sealed.extend_from_slice(STATE_BLOB_MAGIC); - sealed.extend_from_slice(&nonce_bytes); - sealed.extend_from_slice(&output); - Ok(sealed) -} - -fn open_blob(admin_key: &AesKeyType, sealed: &[u8]) -> Result, AuthFailure> { - if sealed.len() < STATE_BLOB_MAGIC.len() + 12 + AES_256_GCM.tag_len() - || &sealed[..STATE_BLOB_MAGIC.len()] != STATE_BLOB_MAGIC - { - return Err(AuthFailure::new( - "temporary_key_store_unavailable", - "authentication state blob has an invalid header", - false, - )); - } - let nonce_start = STATE_BLOB_MAGIC.len(); - let nonce_end = nonce_start + 12; - let nonce_bytes: [u8; 12] = sealed[nonce_start..nonce_end] - .try_into() - .expect("validated nonce width"); - let mut plain = sealed[nonce_end..].to_vec(); - let key = LessSafeKey::new(UnboundKey::new(&AES_256_GCM, admin_key).map_err(|_| { - AuthFailure::new( - "temporary_key_store_unavailable", - "failed to initialize state decryption key", - false, - ) - })?); - let opened = key - .open_in_place( - Nonce::assume_unique_for_key(nonce_bytes), - Aad::from(STATE_AAD), - &mut plain, - ) - .map_err(|_| { - AuthFailure::new( - "temporary_key_store_unavailable", - "authentication state integrity check failed", - false, - ) - })?; - let len = opened.len(); - plain.truncate(len); - Ok(plain) -} - -fn prepare_state_dir(path: &Path) -> Result<(), AuthFailure> { - std::fs::create_dir_all(path).map_err(|error| { - AuthFailure::new( - "auth_state_unavailable", - format!( - "failed to create auth state directory `{}`: {error}", - path.display() - ), - false, - ) - })?; - #[cfg(unix)] - std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o700)).map_err(|error| { - AuthFailure::new( - "auth_state_unavailable", - format!( - "failed to secure auth state directory `{}`: {error}", - path.display() - ), - false, - ) - })?; - Ok(()) -} - -fn load_or_create_instance_id(path: &Path) -> Result<[u8; INSTANCE_ID_LEN], AuthFailure> { - let instance_path = path.join("server-instance-id"); - if instance_path.exists() { - let bytes = std::fs::read(&instance_path).map_err(|error| { - AuthFailure::new( - "auth_state_unavailable", - format!("failed to read `{}`: {error}", instance_path.display()), - false, - ) - })?; - return bytes.try_into().map_err(|_| { - AuthFailure::new( - "auth_state_unavailable", - "server instance id must be exactly 16 bytes", - false, - ) - }); - } - let instance_id = random_instance_id(); - atomic_write(&instance_path, &instance_id, 0o600)?; - Ok(instance_id) -} - -fn random_instance_id() -> [u8; INSTANCE_ID_LEN] { - let mut instance_id = [0_u8; INSTANCE_ID_LEN]; - let mut rng = rand::rng(); - for byte in &mut instance_id { - *byte = rng.random(); - } - instance_id -} - -fn write_admin_key(state_dir: &Path, key: &str) -> Result<(), AuthFailure> { - atomic_write( - &state_dir.join("admin.key"), - format!("{key}\n").as_bytes(), - 0o600, - ) -} - -pub fn generate_admin_key() -> String { - const CHARSET: &[u8] = b"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"; - let mut rng = rand::rng(); - (0..32) - .map(|_| CHARSET[rng.random_range(0..CHARSET.len())] as char) - .collect() -} - -pub fn initialize_admin_key(path: &Path, force: bool) -> Result { - if path.exists() && !force { - return Err(AuthFailure::new( - "administrator_key_exists", - format!("administrator key file `{}` already exists", path.display()), - false, - )); - } - let key = generate_admin_key(); - atomic_write(path, format!("{key}\n").as_bytes(), 0o600)?; - Ok(key) -} - -pub fn write_admin_key_file(path: &Path, key: &str, force: bool) -> Result<(), AuthFailure> { - let Credential::Admin(_) = parse_credential(key) - .map_err(|error| AuthFailure::new("administrator_key_invalid", error, false))? - else { - return Err(AuthFailure::new( - "administrator_key_invalid", - "administrator key file requires a 32-byte administrator key", - false, - )); - }; - if path.exists() && !force { - return Err(AuthFailure::new( - "administrator_key_exists", - format!( - "administrator key file `{}` already exists; pass --force to replace it", - path.display() - ), - false, - )); - } - atomic_write(path, format!("{key}\n").as_bytes(), 0o600) -} - -fn atomic_write(path: &Path, data: &[u8], mode: u32) -> Result<(), AuthFailure> { - if let Some(parent) = path.parent() { - std::fs::create_dir_all(parent).map_err(|error| { - AuthFailure::new( - "auth_state_unavailable", - format!("failed to create `{}`: {error}", parent.display()), - false, - ) - })?; - } - let file_name = path - .file_name() - .and_then(|name| name.to_str()) - .unwrap_or("auth-state"); - let mut random_suffix = [0_u8; 8]; - let mut rng = rand::rng(); - for byte in &mut random_suffix { - *byte = rng.random(); - } - let temporary = path.with_file_name(format!( - ".{file_name}.tmp-{}-{}", - std::process::id(), - hex(&random_suffix) - )); - let mut file = OpenOptions::new() - .create_new(true) - .write(true) - .open(&temporary) - .map_err(|error| { - AuthFailure::new( - "auth_state_unavailable", - format!("failed to open `{}`: {error}", temporary.display()), - false, - ) - })?; - let result = (|| { - #[cfg(unix)] - file.set_permissions(std::fs::Permissions::from_mode(mode)) - .map_err(|error| { - AuthFailure::internal(format!("failed to set key permissions: {error}")) - })?; - #[cfg(not(unix))] - let _ = mode; - file.write_all(data) - .and_then(|()| file.sync_all()) - .map_err(|error| { - AuthFailure::new( - "auth_state_unavailable", - format!("failed to write `{}`: {error}", temporary.display()), - false, - ) - })?; - drop(file); - std::fs::rename(&temporary, path).map_err(|error| { - AuthFailure::new( - "auth_state_unavailable", - format!("failed to replace `{}`: {error}", path.display()), - false, - ) - })?; - #[cfg(unix)] - if let Some(parent) = path.parent() { - File::open(parent) - .and_then(|directory| directory.sync_all()) - .map_err(|error| { - AuthFailure::new( - "auth_state_unavailable", - format!("failed to sync `{}`: {error}", parent.display()), - false, - ) - })?; - } - Ok(()) - })(); - if result.is_err() { - let _ = std::fs::remove_file(&temporary); - } - result -} - -fn unix_seconds() -> u64 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_secs() -} - -fn hex(bytes: &[u8]) -> String { - const DIGITS: &[u8; 16] = b"0123456789abcdef"; - let mut output = String::with_capacity(bytes.len() * 2); - for byte in bytes { - output.push(DIGITS[(byte >> 4) as usize] as char); - output.push(DIGITS[(byte & 0x0f) as usize] as char); - } - output -} - -struct WheelEntry { - lease: Arc, - version: u64, -} - -struct TimingWheel { - now: u64, - level0: Vec>, - level1: Vec>, - level2: Vec>, - level3: Vec>, -} - -impl TimingWheel { - fn new(now: u64) -> Self { - Self { - now, - level0: empty_buckets(256), - level1: empty_buckets(64), - level2: empty_buckets(64), - level3: empty_buckets(64), - } - } - - fn insert(&mut self, lease: Arc) { - let version = lease.wheel_version.load(Ordering::Acquire); - self.insert_with_version(lease, version); - } - - fn insert_with_version(&mut self, lease: Arc, version: u64) { - let expires_at = lease.expires_at(); - let delta = expires_at.saturating_sub(self.now); - let entry = WheelEntry { lease, version }; - if delta < 1 << 8 { - self.level0[(expires_at & 0xff) as usize].push(entry); - } else if delta < 1 << 14 { - self.level1[((expires_at >> 8) & 0x3f) as usize].push(entry); - } else if delta < 1 << 20 { - self.level2[((expires_at >> 14) & 0x3f) as usize].push(entry); - } else { - self.level3[((expires_at >> 20) & 0x3f) as usize].push(entry); - } - } - - fn advance(&mut self, target: u64) -> Vec> { - let mut due = Vec::new(); - while self.now < target { - self.now = self.now.saturating_add(1); - if self.now & 0xff == 0 { - self.cascade(1); - if (self.now >> 8) & 0x3f == 0 { - self.cascade(2); - if (self.now >> 14) & 0x3f == 0 { - self.cascade(3); - } - } - } - let index = (self.now & 0xff) as usize; - for entry in std::mem::take(&mut self.level0[index]) { - if entry.version == entry.lease.wheel_version.load(Ordering::Acquire) { - if entry.lease.expires_at() <= self.now { - due.push(entry.lease); - } else { - self.insert(entry.lease); - } - } - } - } - due - } - - fn cascade(&mut self, level: u8) { - let entries = match level { - 1 => { - let index = ((self.now >> 8) & 0x3f) as usize; - std::mem::take(&mut self.level1[index]) - } - 2 => { - let index = ((self.now >> 14) & 0x3f) as usize; - std::mem::take(&mut self.level2[index]) - } - 3 => { - let index = ((self.now >> 20) & 0x3f) as usize; - std::mem::take(&mut self.level3[index]) - } - _ => Vec::new(), - }; - for entry in entries { - if entry.version == entry.lease.wheel_version.load(Ordering::Acquire) { - self.insert_with_version(entry.lease, entry.version); - } - } - } - - fn clear(&mut self, now: u64) { - *self = Self::new(now); - } -} - -fn empty_buckets(count: usize) -> Vec> { - std::iter::repeat_with(Vec::new).take(count).collect() -} - +mod actor; +use actor::{run_auth_actor, AuthActorState}; +mod persistence; +pub use persistence::*; +mod timing_wheel; +use timing_wheel::TimingWheel; #[cfg(test)] -mod tests { - use super::*; - - fn temp_state_dir(name: &str) -> PathBuf { - let mut suffix = [0_u8; 8]; - let mut rng = rand::rng(); - for byte in &mut suffix { - *byte = rng.random(); - } - std::env::temp_dir().join(format!("pb-mapper-{name}-{}", hex(&suffix))) - } - - #[test] - fn key_id_round_trip() { - let key_id = make_key_id(42, 65_535); - assert_eq!(key_generation(key_id), 42); - assert_eq!(key_slot(key_id), 65_535); - } - - #[test] - fn derived_key_is_bound_to_instance_and_key_id() { - let admin = *b"0123456789abcdefghijklmnopqrstuv"; - let instance_a = [1_u8; INSTANCE_ID_LEN]; - let instance_b = [2_u8; INSTANCE_ID_LEN]; - let key = derive_temporary_key(&admin, &instance_a, make_key_id(1, 7)).unwrap(); - assert_eq!( - key, - derive_temporary_key(&admin, &instance_a, make_key_id(1, 7)).unwrap() - ); - assert_ne!( - key, - derive_temporary_key(&admin, &instance_b, make_key_id(1, 7)).unwrap() - ); - assert_ne!( - key, - derive_temporary_key(&admin, &instance_a, make_key_id(2, 7)).unwrap() - ); - } - - #[tokio::test] - async fn issue_renew_revoke_and_persist() { - let state_dir = temp_state_dir("auth-lifecycle"); - let admin = *b"0123456789abcdefghijklmnopqrstuv"; - let config = AuthConfig { - state_dir: state_dir.clone(), - max_temporary_keys: 8, - max_temporary_key_ttl: Duration::from_secs(3600), - legacy_protocol: LegacyProtocolPolicy::Allow, - }; - let runtime = AuthRuntime::start(admin, config.clone()).await.unwrap(); - let issued = runtime - .issue(Duration::from_secs(60), Some("demo".to_string())) - .await - .unwrap(); - assert!(issued.credential.starts_with("pbmt1_")); - let context = runtime.authenticate(issued.metadata.key_id).unwrap(); - assert!(!context.is_admin); - let cancellation = context.cancellation_token().unwrap(); - let renewed = runtime - .renew(issued.metadata.key_id, Duration::from_secs(120)) - .await - .unwrap(); - assert_eq!(renewed.metadata.key_id, issued.metadata.key_id); - assert_eq!(renewed.credential, issued.credential); - assert!(renewed.metadata.expires_at > issued.metadata.expires_at); - runtime.revoke(issued.metadata.key_id).await.unwrap(); - assert!(cancellation.is_cancelled()); - assert_eq!( - context.ensure_active().unwrap_err().code, - "temporary_key_revoked" - ); - assert_eq!( - runtime - .authenticate(issued.metadata.key_id) - .unwrap_err() - .code, - "temporary_key_revoked" - ); - drop(runtime); - - tokio::time::sleep(Duration::from_millis(20)).await; - let restored = AuthRuntime::start(admin, config).await.unwrap(); - assert_eq!( - restored - .authenticate(issued.metadata.key_id) - .unwrap_err() - .code, - "temporary_key_revoked" - ); - let _ = std::fs::remove_dir_all(state_dir); - } - - #[tokio::test] - async fn reset_rotates_instance_and_prevents_old_key_id_reuse() { - let state_dir = temp_state_dir("auth-reset"); - let admin = *b"0123456789abcdefghijklmnopqrstuv"; - let config = AuthConfig { - state_dir: state_dir.clone(), - max_temporary_keys: 1, - max_temporary_key_ttl: Duration::from_secs(3600), - legacy_protocol: LegacyProtocolPolicy::Allow, - }; - let runtime = AuthRuntime::start(admin, config).await.unwrap(); - let before = runtime.status().await.unwrap().server_instance_id; - let old = runtime - .issue(Duration::from_secs(60), Some("before-reset".to_string())) - .await - .unwrap(); - let old_context = runtime.authenticate(old.metadata.key_id).unwrap(); - let old_cancellation = old_context.cancellation_token().unwrap(); - - runtime.reset().await.unwrap(); - - let after = runtime.status().await.unwrap().server_instance_id; - assert_ne!(after, before); - assert!(old_cancellation.is_cancelled()); - assert!(runtime.authenticate(old.metadata.key_id).is_err()); - let replacement = runtime - .issue(Duration::from_secs(60), Some("after-reset".to_string())) - .await - .unwrap(); - assert_ne!(replacement.metadata.key_id, old.metadata.key_id); - assert_ne!(replacement.credential, old.credential); - - drop(runtime); - tokio::time::sleep(Duration::from_millis(20)).await; - let _ = std::fs::remove_dir_all(state_dir); - } - - #[tokio::test] - async fn corrupt_wal_fails_temporary_keys_closed_until_admin_reset() { - let state_dir = temp_state_dir("auth-safe-mode"); - let admin = *b"0123456789abcdefghijklmnopqrstuv"; - let config = AuthConfig { - state_dir: state_dir.clone(), - max_temporary_keys: 4, - max_temporary_key_ttl: Duration::from_secs(3600), - legacy_protocol: LegacyProtocolPolicy::Allow, - }; - let runtime = AuthRuntime::start(admin, config.clone()).await.unwrap(); - let issued = runtime - .issue(Duration::from_secs(60), Some("corrupt-me".to_string())) - .await - .unwrap(); - drop(runtime); - tokio::time::sleep(Duration::from_millis(20)).await; - std::fs::write(state_dir.join("auth.wal"), b"broken-wal").unwrap(); - - let recovered = AuthRuntime::start(admin, config).await.unwrap(); - assert!(recovered.status().await.unwrap().safe_mode); - assert_eq!( - recovered - .authenticate(issued.metadata.key_id) - .unwrap_err() - .code, - "temporary_key_store_unavailable" - ); - recovered.reset().await.unwrap(); - assert!(!recovered.status().await.unwrap().safe_mode); - - drop(recovered); - tokio::time::sleep(Duration::from_millis(20)).await; - let _ = std::fs::remove_dir_all(state_dir); - } - - #[test] - fn timing_wheel_ignores_stale_renewal_entry() { - let now = 1_000; - let lease = Arc::new(AuthLease::new(make_key_id(1, 0), now + 5)); - let mut wheel = TimingWheel::new(now); - wheel.insert(lease.clone()); - lease.expires_at.store(now + 20, Ordering::Release); - lease.wheel_version.fetch_add(1, Ordering::AcqRel); - wheel.insert(lease.clone()); - assert!(wheel.advance(now + 6).is_empty()); - assert_eq!(wheel.advance(now + 20).len(), 1); - } -} +mod tests; diff --git a/src/common/auth/actor.rs b/src/common/auth/actor.rs new file mode 100644 index 0000000..5ab4967 --- /dev/null +++ b/src/common/auth/actor.rs @@ -0,0 +1,872 @@ +use super::*; + +pub(super) struct AuthActorState { + cold: HashMap, + wheel: TimingWheel, + admin_replays: HashSet<[u8; 32]>, + admin_replay_order: VecDeque, +} + +impl AuthActorState { + pub(super) fn new( + cold: HashMap, + wheel: TimingWheel, + admin_replays: HashSet<[u8; 32]>, + admin_replay_order: VecDeque, + ) -> Self { + Self { + cold, + wheel, + admin_replays, + admin_replay_order, + } + } +} + +pub(super) async fn run_auth_actor( + inner: Arc, + mut admin_lease: Arc, + mut command_rx: mpsc::Receiver, + config: AuthConfig, + state: AuthActorState, +) { + let AuthActorState { + mut cold, + mut wheel, + mut admin_replays, + mut admin_replay_order, + } = state; + let now = unix_seconds(); + let mut tombstones = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .iter() + .enumerate() + .filter_map(|(index, slot)| { + matches!(slot.state, SlotState::Expired | SlotState::Revoked).then_some(( + now.saturating_add(TOMBSTONE_RETENTION.as_secs()), + make_key_id(slot.generation, index as u32), + )) + }) + .collect::>(); + let mut last_snapshot_at = unix_seconds(); + let mut tick = tokio::time::interval(Duration::from_secs(1)); + tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + tokio::select! { + _ = tick.tick() => { + let now = unix_seconds(); + for lease in wheel.advance(now) { + let key_id = lease.key_id(); + let version = lease.wheel_version.load(Ordering::Acquire); + if lease.expires_at() > now { + wheel.insert_with_version(lease, version); + continue; + } + let mut slots = inner.slots.write().unwrap_or_else(|poisoned| poisoned.into_inner()); + if let Some(slot) = slots.get_mut(key_slot(key_id) as usize) { + if slot.generation == key_generation(key_id) && slot.state == SlotState::Active { + slot.state = SlotState::Expired; + lease.cancellation.cancel(); + tombstones.push_back((now.saturating_add(TOMBSTONE_RETENTION.as_secs()), key_id)); + tracing::info!( + event = "temporary_key_expired", + auth_stage = "expiry", + key_id, + expires_at = lease.expires_at(), + "temporary key expired and active work was cancelled" + ); + } + } + } + while let Some((cleanup_at, key_id)) = tombstones.front().copied() { + if cleanup_at > now { + break; + } + tombstones.pop_front(); + let mut slots = inner.slots.write().unwrap_or_else(|poisoned| poisoned.into_inner()); + if let Some(slot) = slots.get_mut(key_slot(key_id) as usize) { + if slot.generation == key_generation(key_id) && matches!(slot.state, SlotState::Expired | SlotState::Revoked) { + slot.state = SlotState::Free; + slot.expires_at = 0; + slot.lease = Weak::new(); + cold.remove(&key_id); + } + } + } + admin_replay_order.retain(|record| { + let keep = now.saturating_sub(record.client_timestamp) + <= ADMIN_REPLAY_RETENTION.as_secs(); + if !keep { + admin_replays.remove(&record.fingerprint); + } + keep + }); + if now.saturating_sub(last_snapshot_at) >= SNAPSHOT_COMPACTION_INTERVAL.as_secs() { + let snapshot = build_snapshot(&inner, &cold, &admin_replay_order); + if let Err(error) = write_snapshot_and_truncate_wal( + &config, + &inner.admin_key(), + &snapshot, + ) { + inner.safe_mode.store(true, Ordering::Release); + cancel_all_temporary_leases(&inner); + tracing::error!( + event = "auth_state_safe_mode", + auth_stage = "snapshot_compaction", + reason = %error.code, + error = %error, + "authentication state compaction failed closed" + ); + } else { + last_snapshot_at = now; + } + } + } + command = command_rx.recv() => { + let Some(command) = command else { + admin_lease.cancellation.cancel(); + cancel_all_temporary_leases(&inner); + break; + }; + match command { + AuthCommand::ClaimAdminMutation { + authority, + fingerprint, + client_timestamp, + response, + } => { + let result = validate_admin_authority(&inner, &authority).and_then(|()| { + actor_claim_admin_mutation( + &inner, + &config, + &mut admin_replays, + &mut admin_replay_order, + fingerprint, + client_timestamp, + ) + }); + let _ = response.send(result); + } + AuthCommand::Issue { authority, ttl, label, response } => { + let result = validate_admin_authority(&inner, &authority) + .and_then(|()| actor_issue(&inner, &config, &mut cold, &mut wheel, ttl, label)); + let _ = response.send(result); + } + AuthCommand::List { authority, page, page_size, response } => { + let result = validate_admin_authority(&inner, &authority) + .and_then(|()| actor_list(&inner, &cold, page, page_size)); + let _ = response.send(result); + } + AuthCommand::Show { authority, key_id, reveal, response } => { + let result = validate_admin_authority(&inner, &authority) + .and_then(|()| actor_show(&inner, &config, &cold, key_id, reveal)); + let _ = response.send(result); + } + AuthCommand::Renew { authority, key_id, ttl, response } => { + let result = validate_admin_authority(&inner, &authority) + .and_then(|()| actor_renew(&inner, &config, &cold, &mut wheel, key_id, ttl)); + let _ = response.send(result); + } + AuthCommand::Revoke { authority, key_id, response } => { + let result = validate_admin_authority(&inner, &authority) + .and_then(|()| actor_revoke(&inner, &config, &cold, &mut tombstones, key_id)); + let _ = response.send(result); + } + AuthCommand::Gc { authority, response } => { + let result = validate_admin_authority(&inner, &authority) + .and_then(|()| actor_gc( + &inner, + &config, + &mut cold, + &mut tombstones, + &admin_replay_order, + )); + let _ = response.send(result); + } + AuthCommand::Reset { authority, response } => { + let result = validate_admin_authority(&inner, &authority) + .and_then(|()| actor_reset( + &inner, + &config, + &mut cold, + &mut wheel, + &admin_replay_order, + "auth_state_reset", + )); + let _ = response.send(result); + } + AuthCommand::RotateRoot { authority, new_key, response } => { + let result = validate_admin_authority(&inner, &authority) + .and_then(|()| actor_rotate_root(&inner, &config, &mut cold, &mut wheel, &mut admin_lease, new_key)); + if result.is_ok() { + admin_replays.clear(); + admin_replay_order.clear(); + } + let _ = response.send(result); + } + AuthCommand::SetLegacyProtocol { authority, policy, response } => { + let result = validate_admin_authority(&inner, &authority) + .and_then(|()| actor_set_legacy_protocol(&inner, &config, policy)); + let _ = response.send(result); + } + AuthCommand::Status { authority, response } => { + let result = validate_admin_authority(&inner, &authority) + .map(|()| actor_status(&inner)); + let _ = response.send(result); + } + AuthCommand::Audit { authority, action, key_id, detail, response } => { + let result = validate_admin_authority(&inner, &authority).and_then(|()| { + append_audit( + &config, + &inner, + audit(&action, key_id, detail), + ) + }); + let _ = response.send(result); + } + } + } + } + } +} + +fn actor_claim_admin_mutation( + inner: &AuthStateInner, + config: &AuthConfig, + admin_replays: &mut HashSet<[u8; 32]>, + admin_replay_order: &mut VecDeque, + fingerprint: [u8; 32], + client_timestamp: u64, +) -> Result<(), AuthFailure> { + if admin_replays.contains(&fingerprint) { + return Err(AuthFailure::new( + "admin_request_replayed", + "administrator mutation was already admitted", + false, + )); + } + if admin_replays.len() >= ADMIN_REPLAY_CAPACITY { + return Err(AuthFailure::new( + "admin_replay_capacity_exhausted", + "administrator mutation replay window is full; retry after older claims expire", + true, + )); + } + let now = unix_seconds(); + if now.abs_diff(client_timestamp) > ADMIN_REPLAY_RETENTION.as_secs() / 2 { + return Err(AuthFailure::new( + "admin_request_timestamp_invalid", + "administrator mutation timestamp is outside the accepted window", + false, + )); + } + let record = AdminReplayRecord { + fingerprint, + client_timestamp, + }; + append_wal( + config, + &inner.admin_key(), + &WalRecord::AdminReplay(record.clone()), + )?; + admin_replays.insert(fingerprint); + admin_replay_order.push_back(record); + Ok(()) +} + +fn validate_ttl(config: &AuthConfig, ttl: Duration) -> Result { + if ttl < MIN_TEMP_KEY_TTL { + return Err(AuthFailure::new( + "temporary_key_ttl_too_short", + format!( + "temporary key TTL must be at least {} seconds", + MIN_TEMP_KEY_TTL.as_secs() + ), + false, + )); + } + if ttl > config.max_temporary_key_ttl { + return Err(AuthFailure::new( + "temporary_key_ttl_too_long", + format!( + "temporary key TTL exceeds the configured maximum of {} seconds", + config.max_temporary_key_ttl.as_secs() + ), + false, + )); + } + Ok(unix_seconds().saturating_add(ttl.as_secs())) +} + +fn validate_label(label: Option) -> Result, AuthFailure> { + let label = label + .map(|label| label.trim().to_string()) + .filter(|label| !label.is_empty()); + if label.as_ref().is_some_and(|label| label.len() > 64) { + return Err(AuthFailure::new( + "temporary_key_label_too_long", + "temporary key label must not exceed 64 UTF-8 bytes", + false, + )); + } + Ok(label) +} + +fn actor_issue( + inner: &Arc, + config: &AuthConfig, + cold: &mut HashMap, + wheel: &mut TimingWheel, + ttl: Duration, + label: Option, +) -> Result { + ensure_store_available(inner)?; + let expires_at = validate_ttl(config, ttl)?; + let label = validate_label(label)?; + let issued_at = unix_seconds(); + let mut slots = inner + .slots + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let Some((index, slot)) = slots + .iter_mut() + .enumerate() + .find(|(_, slot)| slot.state == SlotState::Free && slot.generation < u32::MAX) + else { + return Err(AuthFailure::new( + "temporary_key_capacity_exhausted", + "temporary key slot table is full", + true, + )); + }; + let generation = slot.generation + 1; + let key_id = make_key_id(generation, index as u32); + let entry = PersistedEntry { + key_id, + state: SlotState::Active, + issued_at, + expires_at, + label: label.clone(), + }; + append_mutation( + config, + inner, + StateMutation::Issue(entry.clone()), + audit("temporary_key_issue", Some(key_id), label.clone()), + )?; + let lease = Arc::new(AuthLease::new(key_id, expires_at)); + slot.generation = generation; + slot.state = SlotState::Active; + slot.expires_at = expires_at; + slot.lease = Arc::downgrade(&lease); + cold.insert(key_id, ColdMetadata { issued_at, label }); + wheel.insert(lease); + drop(slots); + metadata_with_credential(inner, cold, key_id, true) +} + +fn actor_list( + inner: &Arc, + cold: &HashMap, + page: u32, + page_size: u16, +) -> Result { + let page_size = page_size.clamp(1, 1000) as usize; + let start = (page as usize).saturating_mul(page_size); + let slots = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let mut all = slots + .iter() + .enumerate() + .filter_map(|(index, slot)| { + if slot.state == SlotState::Free { + return None; + } + let key_id = make_key_id(slot.generation, index as u32); + let cold = cold.get(&key_id)?; + Some(TemporaryKeyMetadata { + key_id, + state: slot_state_name(slot.state).to_string(), + issued_at: cold.issued_at, + expires_at: slot.expires_at, + label: cold.label.clone(), + }) + }) + .collect::>(); + all.sort_by_key(|item| std::cmp::Reverse(item.issued_at)); + let items = all.iter().skip(start).take(page_size).cloned().collect(); + let next_page = (start.saturating_add(page_size) < all.len()).then_some(page.saturating_add(1)); + Ok(KeyPage { + schema_version: 1, + items, + next_page, + }) +} + +fn actor_show( + inner: &Arc, + config: &AuthConfig, + cold: &HashMap, + key_id: u64, + reveal: bool, +) -> Result { + let result = metadata_with_credential(inner, cold, key_id, reveal)?; + append_audit( + config, + inner, + audit( + if reveal { + "temporary_key_reveal" + } else { + "temporary_key_show" + }, + Some(key_id), + result.metadata.label.clone(), + ), + )?; + Ok(result) +} + +fn actor_renew( + inner: &Arc, + config: &AuthConfig, + cold: &HashMap, + wheel: &mut TimingWheel, + key_id: u64, + ttl: Duration, +) -> Result { + ensure_store_available(inner)?; + let expires_at = validate_ttl(config, ttl)?; + let index = key_slot(key_id) as usize; + let mut slots = inner + .slots + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let slot = slots.get_mut(index).ok_or_else(|| key_not_found(key_id))?; + validate_slot_identity(slot, key_id)?; + if slot.state != SlotState::Active || slot.expires_at <= unix_seconds() { + return Err(AuthFailure::new( + "temporary_key_not_renewable", + "only an active, unexpired temporary key can be renewed", + false, + )); + } + let label = cold + .get(&key_id) + .and_then(|metadata| metadata.label.clone()); + append_mutation( + config, + inner, + StateMutation::Renew { key_id, expires_at }, + audit("temporary_key_renew", Some(key_id), label), + )?; + let lease = slot.lease.upgrade().ok_or_else(|| { + AuthFailure::new( + "temporary_key_inactive", + "temporary key lease is no longer active", + true, + ) + })?; + slot.expires_at = expires_at; + lease.expires_at.store(expires_at, Ordering::Release); + lease.wheel_version.fetch_add(1, Ordering::AcqRel); + wheel.insert(lease); + drop(slots); + metadata_with_credential(inner, cold, key_id, true) +} + +fn actor_revoke( + inner: &Arc, + config: &AuthConfig, + cold: &HashMap, + tombstones: &mut VecDeque<(u64, u64)>, + key_id: u64, +) -> Result { + ensure_store_available(inner)?; + let now = unix_seconds(); + let index = key_slot(key_id) as usize; + let mut slots = inner + .slots + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let slot = slots.get_mut(index).ok_or_else(|| key_not_found(key_id))?; + validate_slot_identity(slot, key_id)?; + if slot.state != SlotState::Active { + return Err(AuthFailure::new( + "temporary_key_not_active", + "temporary key is not active", + false, + )); + } + let cold_metadata = cold.get(&key_id).ok_or_else(|| key_not_found(key_id))?; + append_mutation( + config, + inner, + StateMutation::Revoke { key_id, at: now }, + audit( + "temporary_key_revoke", + Some(key_id), + cold_metadata.label.clone(), + ), + )?; + slot.state = SlotState::Revoked; + if let Some(lease) = slot.lease.upgrade() { + lease.cancellation.cancel(); + } + tombstones.push_back((now.saturating_add(TOMBSTONE_RETENTION.as_secs()), key_id)); + Ok(TemporaryKeyMetadata { + key_id, + state: slot_state_name(slot.state).to_string(), + issued_at: cold_metadata.issued_at, + expires_at: slot.expires_at, + label: cold_metadata.label.clone(), + }) +} + +fn actor_gc( + inner: &Arc, + config: &AuthConfig, + cold: &mut HashMap, + tombstones: &mut VecDeque<(u64, u64)>, + admin_replays: &VecDeque, +) -> Result { + ensure_store_available(inner)?; + let now = unix_seconds(); + let mut removed = 0_u64; + let mut slots = inner + .slots + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + for (index, slot) in slots.iter_mut().enumerate() { + if matches!(slot.state, SlotState::Expired | SlotState::Revoked) + || (slot.state == SlotState::Active && slot.expires_at <= now) + { + let key_id = make_key_id(slot.generation, index as u32); + if let Some(lease) = slot.lease.upgrade() { + lease.cancellation.cancel(); + } + slot.state = SlotState::Free; + slot.expires_at = 0; + slot.lease = Weak::new(); + cold.remove(&key_id); + removed = removed.saturating_add(1); + } + } + tombstones.clear(); + drop(slots); + let gc_audit = audit("temporary_key_gc", None, Some(format!("removed={removed}"))); + let mut snapshot = build_snapshot(inner, cold, admin_replays); + push_persisted_audit(&mut snapshot.audit_records, gc_audit.clone()); + let admin_key = inner.admin_key(); + if let Err(error) = write_snapshot_and_truncate_wal(config, &admin_key, &snapshot) { + inner.safe_mode.store(true, Ordering::Release); + cancel_all_temporary_leases(inner); + return Err(error); + } + push_audit_record(inner, gc_audit); + Ok(removed) +} + +fn actor_reset( + inner: &Arc, + config: &AuthConfig, + cold: &mut HashMap, + wheel: &mut TimingWheel, + admin_replays: &VecDeque, + action: &str, +) -> Result<(), AuthFailure> { + let new_instance_id = random_instance_id(); + let reset_audit = audit(action, None, None); + let mut snapshot = empty_snapshot(inner, new_instance_id, admin_replays); + push_persisted_audit(&mut snapshot.audit_records, reset_audit.clone()); + let admin_key = inner.admin_key(); + if let Err(error) = write_snapshot_and_truncate_wal(config, &admin_key, &snapshot) { + inner.safe_mode.store(true, Ordering::Release); + cancel_all_temporary_leases(inner); + return Err(error); + } + push_audit_record(inner, reset_audit); + if let Err(error) = atomic_write( + &config.state_dir.join("server-instance-id"), + &new_instance_id, + 0o600, + ) { + inner.safe_mode.store(true, Ordering::Release); + cancel_all_temporary_leases(inner); + return Err(error); + } + + cancel_all_temporary_leases(inner); + { + let mut slots = inner + .slots + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + for slot in slots.iter_mut() { + slot.state = SlotState::Free; + slot.expires_at = 0; + slot.lease = Weak::new(); + } + } + *inner + .instance_id + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = new_instance_id; + cold.clear(); + wheel.clear(unix_seconds()); + inner.safe_mode.store(false, Ordering::Release); + Ok(()) +} + +fn actor_rotate_root( + inner: &Arc, + config: &AuthConfig, + cold: &mut HashMap, + wheel: &mut TimingWheel, + admin_lease: &mut Arc, + new_key: AesKeyType, +) -> Result<(), AuthFailure> { + if new_key == inner.admin_key() { + return Err(AuthFailure::new( + "administrator_key_unchanged", + "new administrator key must differ from the current key", + false, + )); + } + let new_key_string = String::from_utf8(new_key.to_vec()).map_err(|_| { + AuthFailure::new( + "administrator_key_invalid", + "administrator key must be 32 UTF-8 bytes for MSG_HEADER_KEY compatibility", + false, + ) + })?; + if new_key_string.chars().any(char::is_whitespace) { + return Err(AuthFailure::new( + "administrator_key_invalid", + "administrator key must not contain whitespace", + false, + )); + } + + let rotate_audit = audit("administrator_key_rotate", None, None); + let mut snapshot = empty_snapshot(inner, inner.instance_id(), &VecDeque::new()); + push_persisted_audit(&mut snapshot.audit_records, rotate_audit.clone()); + if let Err(error) = write_snapshot_and_truncate_wal(config, &new_key, &snapshot) + .and_then(|()| write_admin_key(&config.state_dir, &new_key_string)) + { + inner.safe_mode.store(true, Ordering::Release); + cancel_all_temporary_leases(inner); + return Err(error); + } + push_audit_record(inner, rotate_audit); + + cancel_all_temporary_leases(inner); + let old_admin_lease = admin_lease.clone(); + { + let mut slots = inner + .slots + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + for slot in slots.iter_mut() { + slot.state = SlotState::Free; + slot.expires_at = 0; + slot.lease = Weak::new(); + } + } + cold.clear(); + wheel.clear(unix_seconds()); + let new_admin_lease = Arc::new(AuthLease::new(0, u64::MAX)); + *inner + .admin + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = AdminState { + key: new_key, + lease: Arc::downgrade(&new_admin_lease), + }; + set_process_msg_header_key(Some(&new_key_string)).map_err(AuthFailure::internal)?; + inner.safe_mode.store(false, Ordering::Release); + old_admin_lease.cancellation.cancel(); + *admin_lease = new_admin_lease; + Ok(()) +} + +fn actor_set_legacy_protocol( + inner: &Arc, + config: &AuthConfig, + policy: LegacyProtocolPolicy, +) -> Result<(), AuthFailure> { + ensure_store_available(inner)?; + append_mutation( + config, + inner, + StateMutation::LegacyProtocol(policy), + audit("legacy_protocol_update", None, Some(format!("{policy:?}"))), + )?; + inner + .legacy_protocol_allowed + .store(policy.is_allowed(), Ordering::Release); + Ok(()) +} + +fn actor_status(inner: &Arc) -> AuthStatus { + let slots = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let active_keys = slots + .iter() + .filter(|slot| slot.state == SlotState::Active) + .count(); + let expired_keys = slots + .iter() + .filter(|slot| slot.state == SlotState::Expired) + .count(); + let revoked_keys = slots + .iter() + .filter(|slot| slot.state == SlotState::Revoked) + .count(); + let last_legacy_connection_at = inner.last_legacy_connection_at.load(Ordering::Acquire); + AuthStatus { + schema_version: 1, + safe_mode: inner.safe_mode.load(Ordering::Acquire), + capacity: slots.len(), + active_keys, + expired_keys, + revoked_keys, + legacy_protocol: if inner.legacy_protocol_allowed.load(Ordering::Acquire) { + LegacyProtocolPolicy::Allow + } else { + LegacyProtocolPolicy::Deny + }, + active_legacy_connections: inner.active_legacy_connections.load(Ordering::Acquire), + last_legacy_connection_at: (last_legacy_connection_at != 0) + .then_some(last_legacy_connection_at), + auth_successes: inner.auth_successes.load(Ordering::Relaxed), + auth_failures: inner.auth_failures.load(Ordering::Relaxed), + server_instance_id: hex(&inner.instance_id()), + } +} + +fn validate_admin_authority( + inner: &AuthStateInner, + authority: &Weak, +) -> Result<(), AuthFailure> { + let presented = authority.upgrade().ok_or_else(|| { + AuthFailure::new( + "administrator_key_rotated", + "administrator credential lease is no longer active", + false, + ) + })?; + if presented.cancellation.is_cancelled() { + return Err(AuthFailure::new( + "administrator_key_rotated", + "administrator credential lease has been cancelled", + false, + )); + } + let current = inner + .admin + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .lease + .upgrade() + .ok_or_else(|| { + AuthFailure::new( + "administrator_key_rotated", + "active administrator credential lease is unavailable", + false, + ) + })?; + if !Arc::ptr_eq(&presented, ¤t) { + return Err(AuthFailure::new( + "administrator_key_rotated", + "administrator request was authenticated before the latest root-key rotation", + false, + )); + } + Ok(()) +} + +fn ensure_store_available(inner: &AuthStateInner) -> Result<(), AuthFailure> { + if inner.safe_mode.load(Ordering::Acquire) { + Err(AuthFailure::new( + "temporary_key_store_unavailable", + "temporary key store is in administrator safe mode", + false, + )) + } else { + Ok(()) + } +} + +fn validate_slot_identity(slot: &SlotHot, key_id: u64) -> Result<(), AuthFailure> { + if slot.generation != key_generation(key_id) || slot.state == SlotState::Free { + Err(key_not_found(key_id)) + } else { + Ok(()) + } +} + +fn key_not_found(key_id: u64) -> AuthFailure { + AuthFailure::new( + "temporary_key_not_found", + format!("temporary key {key_id} does not exist"), + false, + ) +} + +fn metadata_with_credential( + inner: &Arc, + cold: &HashMap, + key_id: u64, + reveal: bool, +) -> Result { + let slots = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let slot = slots + .get(key_slot(key_id) as usize) + .ok_or_else(|| key_not_found(key_id))?; + validate_slot_identity(slot, key_id)?; + let cold = cold.get(&key_id).ok_or_else(|| key_not_found(key_id))?; + let credential = if reveal { + let key = derive_temporary_key(&inner.admin_key(), &inner.instance_id(), key_id)?; + encode_temporary_credential(key_id, &key) + } else { + String::new() + }; + Ok(IssuedTemporaryKey { + metadata: TemporaryKeyMetadata { + key_id, + state: slot_state_name(slot.state).to_string(), + issued_at: cold.issued_at, + expires_at: slot.expires_at, + label: cold.label.clone(), + }, + credential, + }) +} + +fn slot_state_name(state: SlotState) -> &'static str { + match state { + SlotState::Free => "free", + SlotState::Active => "active", + SlotState::Expired => "expired", + SlotState::Revoked => "revoked", + } +} + +fn audit(action: &str, key_id: Option, label: Option) -> AuditRecord { + AuditRecord { + at: unix_seconds(), + action: action.to_string(), + key_id, + label, + } +} diff --git a/src/common/auth/persistence.rs b/src/common/auth/persistence.rs new file mode 100644 index 0000000..b2aea88 --- /dev/null +++ b/src/common/auth/persistence.rs @@ -0,0 +1,674 @@ +use super::*; + +pub(super) fn push_audit_record(inner: &AuthStateInner, record: AuditRecord) { + let mut records = inner + .audit_records + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + while records.len() >= AUDIT_RECORD_CAPACITY { + records.pop_front(); + } + records.push_back(record); +} + +pub(super) fn cancel_all_temporary_leases(inner: &AuthStateInner) { + let slots = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + for lease in slots.iter().filter_map(|slot| slot.lease.upgrade()) { + lease.cancellation.cancel(); + } +} + +pub(super) fn build_snapshot( + inner: &AuthStateInner, + cold: &HashMap, + admin_replays: &VecDeque, +) -> PersistedSnapshot { + let slots = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let generations = slots.iter().map(|slot| slot.generation).collect(); + let entries = slots + .iter() + .enumerate() + .filter_map(|(index, slot)| { + if slot.state == SlotState::Free { + return None; + } + let key_id = make_key_id(slot.generation, index as u32); + let cold = cold.get(&key_id)?; + Some(PersistedEntry { + key_id, + state: slot.state, + issued_at: cold.issued_at, + expires_at: slot.expires_at, + label: cold.label.clone(), + }) + }) + .collect(); + PersistedSnapshot { + schema_version: SNAPSHOT_SCHEMA_VERSION, + instance_id: inner.instance_id(), + generations, + entries, + legacy_protocol: if inner.legacy_protocol_allowed.load(Ordering::Acquire) { + LegacyProtocolPolicy::Allow + } else { + LegacyProtocolPolicy::Deny + }, + admin_replays: admin_replays.iter().cloned().collect(), + audit_records: inner + .audit_records + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clone(), + } +} + +pub(super) fn empty_snapshot( + inner: &AuthStateInner, + instance_id: [u8; INSTANCE_ID_LEN], + admin_replays: &VecDeque, +) -> PersistedSnapshot { + let slots = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + PersistedSnapshot { + schema_version: SNAPSHOT_SCHEMA_VERSION, + instance_id, + generations: slots.iter().map(|slot| slot.generation).collect(), + entries: Vec::new(), + legacy_protocol: if inner.legacy_protocol_allowed.load(Ordering::Acquire) { + LegacyProtocolPolicy::Allow + } else { + LegacyProtocolPolicy::Deny + }, + admin_replays: admin_replays.iter().cloned().collect(), + audit_records: inner + .audit_records + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clone(), + } +} + +pub(super) fn load_persisted_state( + config: &AuthConfig, + admin_key: &AesKeyType, + instance_id: [u8; INSTANCE_ID_LEN], +) -> (Option, bool) { + match try_load_persisted_state(config, admin_key, instance_id) { + Ok(state) => (Some(state), false), + Err(error) => { + tracing::error!( + event = "auth_state_safe_mode", + auth_stage = "state_load", + reason = %error.code, + error = %error, + "temporary key store failed closed in administrator safe mode" + ); + (None, true) + } + } +} + +pub(super) fn try_load_persisted_state( + config: &AuthConfig, + admin_key: &AesKeyType, + instance_id: [u8; INSTANCE_ID_LEN], +) -> Result { + let snapshot_path = config.state_dir.join("auth.snapshot"); + let mut snapshot = if snapshot_path.exists() { + let bytes = std::fs::read(&snapshot_path).map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to read `{}`: {error}", snapshot_path.display()), + false, + ) + })?; + let plain = open_blob(admin_key, &bytes)?; + serde_json::from_slice::(&plain).map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to decode auth snapshot: {error}"), + false, + ) + })? + } else { + PersistedSnapshot { + schema_version: SNAPSHOT_SCHEMA_VERSION, + instance_id, + generations: vec![0; config.max_temporary_keys], + entries: Vec::new(), + legacy_protocol: config.legacy_protocol, + admin_replays: Vec::new(), + audit_records: VecDeque::new(), + } + }; + if snapshot.schema_version != SNAPSHOT_SCHEMA_VERSION || snapshot.instance_id != instance_id { + return Err(AuthFailure::new( + "temporary_key_store_unavailable", + "auth snapshot schema or server instance id does not match", + false, + )); + } + snapshot.generations.resize(config.max_temporary_keys, 0); + snapshot.generations.truncate(config.max_temporary_keys); + + let wal_path = config.state_dir.join("auth.wal"); + if wal_path.exists() { + for record in read_wal(&wal_path, admin_key)? { + match record { + WalRecord::Mutation { mutation, audit } => { + apply_persisted_mutation(&mut snapshot, mutation, config.max_temporary_keys)?; + push_persisted_audit(&mut snapshot.audit_records, audit); + } + WalRecord::AdminReplay(record) => snapshot.admin_replays.push(record), + WalRecord::Audit(audit) => push_persisted_audit(&mut snapshot.audit_records, audit), + } + } + } + Ok(snapshot) +} + +pub(super) fn apply_persisted_mutation( + snapshot: &mut PersistedSnapshot, + mutation: StateMutation, + capacity: usize, +) -> Result<(), AuthFailure> { + match mutation { + StateMutation::Issue(entry) => { + let index = key_slot(entry.key_id) as usize; + if index >= capacity { + return Err(AuthFailure::new( + "temporary_key_store_unavailable", + "WAL issue record references a slot outside the configured capacity", + false, + )); + } + snapshot.generations[index] = key_generation(entry.key_id); + snapshot + .entries + .retain(|current| key_slot(current.key_id) as usize != index); + snapshot.entries.push(entry); + } + StateMutation::Renew { key_id, expires_at } => { + let entry = snapshot + .entries + .iter_mut() + .find(|entry| entry.key_id == key_id) + .ok_or_else(|| { + AuthFailure::new( + "temporary_key_store_unavailable", + "WAL renew record references an unknown key", + false, + ) + })?; + entry.expires_at = expires_at; + entry.state = SlotState::Active; + } + StateMutation::Revoke { key_id, .. } => { + let entry = snapshot + .entries + .iter_mut() + .find(|entry| entry.key_id == key_id) + .ok_or_else(|| { + AuthFailure::new( + "temporary_key_store_unavailable", + "WAL revoke record references an unknown key", + false, + ) + })?; + entry.state = SlotState::Revoked; + } + StateMutation::LegacyProtocol(policy) => snapshot.legacy_protocol = policy, + } + Ok(()) +} + +pub(super) fn append_mutation( + config: &AuthConfig, + inner: &AuthStateInner, + mutation: StateMutation, + audit: AuditRecord, +) -> Result<(), AuthFailure> { + append_wal( + config, + &inner.admin_key(), + &WalRecord::Mutation { + mutation, + audit: audit.clone(), + }, + )?; + push_audit_record(inner, audit); + Ok(()) +} + +pub(super) fn append_audit( + config: &AuthConfig, + inner: &AuthStateInner, + audit: AuditRecord, +) -> Result<(), AuthFailure> { + append_wal(config, &inner.admin_key(), &WalRecord::Audit(audit.clone()))?; + push_audit_record(inner, audit); + Ok(()) +} + +pub(super) fn push_persisted_audit(records: &mut VecDeque, record: AuditRecord) { + while records.len() >= AUDIT_RECORD_CAPACITY { + records.pop_front(); + } + records.push_back(record); +} + +pub(super) fn append_wal( + config: &AuthConfig, + admin_key: &AesKeyType, + record: &WalRecord, +) -> Result<(), AuthFailure> { + let plain = serde_json::to_vec(record).map_err(|error| { + AuthFailure::internal(format!("failed to encode auth WAL record: {error}")) + })?; + let sealed = seal_blob(admin_key, &plain)?; + let length = u32::try_from(sealed.len()) + .map_err(|_| AuthFailure::internal("auth WAL record is too large"))?; + let path = config.state_dir.join("auth.wal"); + let mut file = OpenOptions::new() + .create(true) + .append(true) + .open(&path) + .map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to open `{}`: {error}", path.display()), + true, + ) + })?; + #[cfg(unix)] + file.set_permissions(std::fs::Permissions::from_mode(0o600)) + .map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to secure `{}`: {error}", path.display()), + false, + ) + })?; + file.write_all(&length.to_be_bytes()) + .and_then(|()| file.write_all(&sealed)) + .and_then(|()| file.sync_data()) + .map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to durably append `{}`: {error}", path.display()), + true, + ) + }) +} + +pub(super) fn read_wal(path: &Path, admin_key: &AesKeyType) -> Result, AuthFailure> { + let mut file = File::open(path).map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to open `{}`: {error}", path.display()), + false, + ) + })?; + let mut records = Vec::new(); + loop { + let mut length = [0_u8; 4]; + match file.read(&mut length[..1]) { + Ok(0) => break, + Ok(1) => {} + Ok(_) => unreachable!("single-byte WAL prefix read"), + Err(error) => { + return Err(AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to read auth WAL length: {error}"), + false, + )); + } + } + file.read_exact(&mut length[1..]).map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("truncated auth WAL length: {error}"), + false, + ) + })?; + let length = u32::from_be_bytes(length) as usize; + if length > 1024 * 1024 { + return Err(AuthFailure::new( + "temporary_key_store_unavailable", + "auth WAL record exceeds 1 MiB", + false, + )); + } + let mut sealed = vec![0_u8; length]; + file.read_exact(&mut sealed).map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("truncated auth WAL record: {error}"), + false, + ) + })?; + let plain = open_blob(admin_key, &sealed)?; + records.push(serde_json::from_slice(&plain).map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to decode auth WAL record: {error}"), + false, + ) + })?); + } + Ok(records) +} + +pub(super) fn write_snapshot_and_truncate_wal( + config: &AuthConfig, + admin_key: &AesKeyType, + snapshot: &PersistedSnapshot, +) -> Result<(), AuthFailure> { + let plain = serde_json::to_vec(snapshot).map_err(|error| { + AuthFailure::internal(format!("failed to encode auth snapshot: {error}")) + })?; + let sealed = seal_blob(admin_key, &plain)?; + let snapshot_path = config.state_dir.join("auth.snapshot"); + atomic_write(&snapshot_path, &sealed, 0o600)?; + let wal_path = config.state_dir.join("auth.wal"); + let wal = OpenOptions::new() + .create(true) + .write(true) + .truncate(true) + .open(&wal_path) + .map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to truncate `{}`: {error}", wal_path.display()), + true, + ) + })?; + wal.sync_all().map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to sync `{}`: {error}", wal_path.display()), + true, + ) + }) +} + +pub(super) fn seal_blob(admin_key: &AesKeyType, plain: &[u8]) -> Result, AuthFailure> { + let key = LessSafeKey::new( + UnboundKey::new(&AES_256_GCM, admin_key) + .map_err(|_| AuthFailure::internal("failed to initialize state encryption key"))?, + ); + let mut nonce_bytes = [0_u8; 12]; + let mut rng = rand::rng(); + for byte in &mut nonce_bytes { + *byte = rng.random(); + } + let mut output = plain.to_vec(); + key.seal_in_place_append_tag( + Nonce::assume_unique_for_key(nonce_bytes), + Aad::from(STATE_AAD), + &mut output, + ) + .map_err(|_| AuthFailure::internal("failed to encrypt authentication state"))?; + let mut sealed = Vec::with_capacity(STATE_BLOB_MAGIC.len() + nonce_bytes.len() + output.len()); + sealed.extend_from_slice(STATE_BLOB_MAGIC); + sealed.extend_from_slice(&nonce_bytes); + sealed.extend_from_slice(&output); + Ok(sealed) +} + +pub(super) fn open_blob(admin_key: &AesKeyType, sealed: &[u8]) -> Result, AuthFailure> { + if sealed.len() < STATE_BLOB_MAGIC.len() + 12 + AES_256_GCM.tag_len() + || &sealed[..STATE_BLOB_MAGIC.len()] != STATE_BLOB_MAGIC + { + return Err(AuthFailure::new( + "temporary_key_store_unavailable", + "authentication state blob has an invalid header", + false, + )); + } + let nonce_start = STATE_BLOB_MAGIC.len(); + let nonce_end = nonce_start + 12; + let nonce_bytes: [u8; 12] = sealed[nonce_start..nonce_end] + .try_into() + .expect("validated nonce width"); + let mut plain = sealed[nonce_end..].to_vec(); + let key = LessSafeKey::new(UnboundKey::new(&AES_256_GCM, admin_key).map_err(|_| { + AuthFailure::new( + "temporary_key_store_unavailable", + "failed to initialize state decryption key", + false, + ) + })?); + let opened = key + .open_in_place( + Nonce::assume_unique_for_key(nonce_bytes), + Aad::from(STATE_AAD), + &mut plain, + ) + .map_err(|_| { + AuthFailure::new( + "temporary_key_store_unavailable", + "authentication state integrity check failed", + false, + ) + })?; + let len = opened.len(); + plain.truncate(len); + Ok(plain) +} + +pub(super) fn prepare_state_dir(path: &Path) -> Result<(), AuthFailure> { + std::fs::create_dir_all(path).map_err(|error| { + AuthFailure::new( + "auth_state_unavailable", + format!( + "failed to create auth state directory `{}`: {error}", + path.display() + ), + false, + ) + })?; + #[cfg(unix)] + std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o700)).map_err(|error| { + AuthFailure::new( + "auth_state_unavailable", + format!( + "failed to secure auth state directory `{}`: {error}", + path.display() + ), + false, + ) + })?; + Ok(()) +} + +pub(super) fn load_or_create_instance_id( + path: &Path, +) -> Result<[u8; INSTANCE_ID_LEN], AuthFailure> { + let instance_path = path.join("server-instance-id"); + if instance_path.exists() { + let bytes = std::fs::read(&instance_path).map_err(|error| { + AuthFailure::new( + "auth_state_unavailable", + format!("failed to read `{}`: {error}", instance_path.display()), + false, + ) + })?; + return bytes.try_into().map_err(|_| { + AuthFailure::new( + "auth_state_unavailable", + "server instance id must be exactly 16 bytes", + false, + ) + }); + } + let instance_id = random_instance_id(); + atomic_write(&instance_path, &instance_id, 0o600)?; + Ok(instance_id) +} + +pub(super) fn random_instance_id() -> [u8; INSTANCE_ID_LEN] { + let mut instance_id = [0_u8; INSTANCE_ID_LEN]; + let mut rng = rand::rng(); + for byte in &mut instance_id { + *byte = rng.random(); + } + instance_id +} + +pub(super) fn write_admin_key(state_dir: &Path, key: &str) -> Result<(), AuthFailure> { + atomic_write( + &state_dir.join("admin.key"), + format!("{key}\n").as_bytes(), + 0o600, + ) +} + +pub fn generate_admin_key() -> String { + const CHARSET: &[u8] = b"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"; + let mut rng = rand::rng(); + (0..32) + .map(|_| CHARSET[rng.random_range(0..CHARSET.len())] as char) + .collect() +} + +pub fn initialize_admin_key(path: &Path, force: bool) -> Result { + if path.exists() && !force { + return Err(AuthFailure::new( + "administrator_key_exists", + format!("administrator key file `{}` already exists", path.display()), + false, + )); + } + let key = generate_admin_key(); + atomic_write(path, format!("{key}\n").as_bytes(), 0o600)?; + Ok(key) +} + +pub fn write_admin_key_file(path: &Path, key: &str, force: bool) -> Result<(), AuthFailure> { + let Credential::Admin(_) = parse_credential(key) + .map_err(|error| AuthFailure::new("administrator_key_invalid", error, false))? + else { + return Err(AuthFailure::new( + "administrator_key_invalid", + "administrator key file requires a 32-byte administrator key", + false, + )); + }; + if path.exists() && !force { + return Err(AuthFailure::new( + "administrator_key_exists", + format!( + "administrator key file `{}` already exists; pass --force to replace it", + path.display() + ), + false, + )); + } + atomic_write(path, format!("{key}\n").as_bytes(), 0o600) +} + +pub(super) fn atomic_write(path: &Path, data: &[u8], mode: u32) -> Result<(), AuthFailure> { + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent).map_err(|error| { + AuthFailure::new( + "auth_state_unavailable", + format!("failed to create `{}`: {error}", parent.display()), + false, + ) + })?; + } + let file_name = path + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or("auth-state"); + let mut random_suffix = [0_u8; 8]; + let mut rng = rand::rng(); + for byte in &mut random_suffix { + *byte = rng.random(); + } + let temporary = path.with_file_name(format!( + ".{file_name}.tmp-{}-{}", + std::process::id(), + hex(&random_suffix) + )); + let mut file = OpenOptions::new() + .create_new(true) + .write(true) + .open(&temporary) + .map_err(|error| { + AuthFailure::new( + "auth_state_unavailable", + format!("failed to open `{}`: {error}", temporary.display()), + false, + ) + })?; + let result = (|| { + #[cfg(unix)] + file.set_permissions(std::fs::Permissions::from_mode(mode)) + .map_err(|error| { + AuthFailure::internal(format!("failed to set key permissions: {error}")) + })?; + #[cfg(not(unix))] + let _ = mode; + file.write_all(data) + .and_then(|()| file.sync_all()) + .map_err(|error| { + AuthFailure::new( + "auth_state_unavailable", + format!("failed to write `{}`: {error}", temporary.display()), + false, + ) + })?; + drop(file); + std::fs::rename(&temporary, path).map_err(|error| { + AuthFailure::new( + "auth_state_unavailable", + format!("failed to replace `{}`: {error}", path.display()), + false, + ) + })?; + #[cfg(unix)] + if let Some(parent) = path.parent() { + File::open(parent) + .and_then(|directory| directory.sync_all()) + .map_err(|error| { + AuthFailure::new( + "auth_state_unavailable", + format!("failed to sync `{}`: {error}", parent.display()), + false, + ) + })?; + } + Ok(()) + })(); + if result.is_err() { + let _ = std::fs::remove_file(&temporary); + } + result +} + +pub(super) fn unix_seconds() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() +} + +pub(super) fn hex(bytes: &[u8]) -> String { + const DIGITS: &[u8; 16] = b"0123456789abcdef"; + let mut output = String::with_capacity(bytes.len() * 2); + for byte in bytes { + output.push(DIGITS[(byte >> 4) as usize] as char); + output.push(DIGITS[(byte & 0x0f) as usize] as char); + } + output +} diff --git a/src/common/auth/runtime.rs b/src/common/auth/runtime.rs new file mode 100644 index 0000000..4f236f7 --- /dev/null +++ b/src/common/auth/runtime.rs @@ -0,0 +1,483 @@ +use super::*; + +impl AuthRuntime { + pub async fn from_process(config: AuthConfig) -> Result { + prepare_state_dir(&config.state_dir)?; + let credential = load_server_admin_credential(&config.state_dir)?; + let Credential::Admin(admin_key) = credential else { + return Err(AuthFailure::new( + "administrator_key_required", + "the relay server must start with the administrator credential", + false, + )); + }; + Self::start(admin_key, config).await + } + + pub async fn start(admin_key: AesKeyType, config: AuthConfig) -> Result { + prepare_state_dir(&config.state_dir)?; + let instance_id = load_or_create_instance_id(&config.state_dir)?; + let (loaded, safe_mode) = load_persisted_state(&config, &admin_key, instance_id); + let mut slots = (0..config.max_temporary_keys) + .map(|_| SlotHot::default()) + .collect::>() + .into_boxed_slice(); + let mut cold = HashMap::new(); + let mut wheel = TimingWheel::new(unix_seconds()); + let now = unix_seconds(); + + let admin_lease = Arc::new(AuthLease::new(0, u64::MAX)); + if let Some(state) = loaded.as_ref() { + for (index, generation) in state.generations.iter().copied().enumerate() { + if let Some(slot) = slots.get_mut(index) { + slot.generation = generation; + } + } + for entry in &state.entries { + let index = key_slot(entry.key_id) as usize; + let Some(slot) = slots.get_mut(index) else { + continue; + }; + if slot.generation != key_generation(entry.key_id) { + continue; + } + let state = if entry.state == SlotState::Active && entry.expires_at <= now { + SlotState::Expired + } else { + entry.state + }; + slot.state = state; + slot.expires_at = entry.expires_at; + cold.insert( + entry.key_id, + ColdMetadata { + issued_at: entry.issued_at, + label: entry.label.clone(), + }, + ); + if state == SlotState::Active { + let lease = Arc::new(AuthLease::new(entry.key_id, entry.expires_at)); + slot.lease = Arc::downgrade(&lease); + wheel.insert(lease); + } + } + } + + let legacy_protocol = loaded + .as_ref() + .map(|state| state.legacy_protocol) + .unwrap_or(config.legacy_protocol); + let mut admin_replay_order = loaded + .as_ref() + .map(|state| { + state + .admin_replays + .iter() + .filter(|record| { + now.saturating_sub(record.client_timestamp) + <= ADMIN_REPLAY_RETENTION.as_secs() + }) + .cloned() + .collect::>() + }) + .unwrap_or_default(); + while admin_replay_order.len() > ADMIN_REPLAY_CAPACITY { + admin_replay_order.pop_front(); + } + let admin_replays = admin_replay_order + .iter() + .map(|record| record.fingerprint) + .collect::>(); + let mut audit_records: VecDeque = loaded + .as_ref() + .map(|state| state.audit_records.iter().cloned().collect()) + .unwrap_or_default(); + while audit_records.len() > AUDIT_RECORD_CAPACITY { + audit_records.pop_front(); + } + let inner = Arc::new(AuthStateInner { + admin: RwLock::new(AdminState { + key: admin_key, + lease: Arc::downgrade(&admin_lease), + }), + instance_id: RwLock::new(instance_id), + slots: RwLock::new(slots), + safe_mode: AtomicBool::new(safe_mode), + legacy_protocol_allowed: AtomicBool::new(legacy_protocol.is_allowed()), + active_legacy_connections: AtomicU64::new(0), + last_legacy_connection_at: AtomicU64::new(0), + auth_successes: AtomicU64::new(0), + auth_failures: AtomicU64::new(0), + audit_records: RwLock::new(audit_records), + }); + let (command_tx, command_rx) = mpsc::channel(256); + let runtime = Self { + inner: Arc::downgrade(&inner), + command_tx, + config: config.clone(), + }; + + tokio::spawn(run_auth_actor( + inner, + admin_lease, + command_rx, + config, + AuthActorState::new(cold, wheel, admin_replays, admin_replay_order), + )); + Ok(runtime) + } + + pub fn config(&self) -> &AuthConfig { + &self.config + } + + fn inner(&self) -> Result, AuthFailure> { + self.inner.upgrade().ok_or_else(|| { + AuthFailure::new( + "auth_state_unavailable", + "authentication state manager is not running", + true, + ) + }) + } + + pub fn admin_key(&self) -> Result { + Ok(self.inner()?.admin_key()) + } + + pub fn derive_key(&self, key_id: u64) -> Result { + let inner = self.inner()?; + if key_id == 0 { + return Ok(inner.admin_key()); + } + derive_temporary_key(&inner.admin_key(), &inner.instance_id(), key_id) + } + + pub fn authenticate(&self, key_id: u64) -> Result { + let presented_key = self.derive_key(key_id)?; + self.authenticate_presented(key_id, &presented_key) + } + + pub fn authenticate_presented( + &self, + key_id: u64, + presented_key: &AesKeyType, + ) -> Result { + let inner = self.inner()?; + if key_id == 0 { + let admin = inner + .admin + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if !bool::from(presented_key.ct_eq(&admin.key)) { + inner.auth_failures.fetch_add(1, Ordering::Relaxed); + return Err(AuthFailure::new( + "administrator_key_rotated", + "administrator credential no longer matches the active root key", + false, + )); + } + let lease = admin.lease.upgrade().ok_or_else(|| { + AuthFailure::new( + "administrator_key_rotated", + "administrator credential was rotated", + false, + ) + })?; + inner.auth_successes.fetch_add(1, Ordering::Relaxed); + return Ok(AuthContext::from_lease(0, true, &lease)); + } + if inner.safe_mode.load(Ordering::Acquire) { + inner.auth_failures.fetch_add(1, Ordering::Relaxed); + return Err(AuthFailure::new( + "temporary_key_store_unavailable", + "temporary key state is unavailable; administrator reset is required", + false, + )); + } + + let expected_key = derive_temporary_key(&inner.admin_key(), &inner.instance_id(), key_id)?; + if !bool::from(presented_key.ct_eq(&expected_key)) { + inner.auth_failures.fetch_add(1, Ordering::Relaxed); + return Err(AuthFailure::new( + "temporary_key_invalid", + "temporary credential does not match the active relay key material", + false, + )); + } + + let index = key_slot(key_id) as usize; + let generation = key_generation(key_id); + let slots = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let Some(slot) = slots.get(index) else { + inner.auth_failures.fetch_add(1, Ordering::Relaxed); + return Err(AuthFailure::new( + "temporary_key_not_found", + "temporary key id is outside the configured slot table", + false, + )); + }; + if slot.generation != generation { + inner.auth_failures.fetch_add(1, Ordering::Relaxed); + return Err(AuthFailure::new( + "temporary_key_generation_mismatch", + "temporary key generation does not match the current slot", + false, + )); + } + let failure = match slot.state { + SlotState::Free => Some(AuthFailure::new( + "temporary_key_not_found", + "temporary key does not exist", + false, + )), + SlotState::Expired => Some(AuthFailure::new( + "temporary_key_expired", + "temporary key has expired", + false, + )), + SlotState::Revoked => Some(AuthFailure::new( + "temporary_key_revoked", + "temporary key was revoked", + false, + )), + SlotState::Active if slot.expires_at <= unix_seconds() => { + if let Some(lease) = slot.lease.upgrade() { + lease.cancellation.cancel(); + } + Some(AuthFailure::new( + "temporary_key_expired", + "temporary key has expired", + false, + )) + } + SlotState::Active => None, + }; + if let Some(failure) = failure { + inner.auth_failures.fetch_add(1, Ordering::Relaxed); + return Err(failure); + } + let lease = slot.lease.upgrade().ok_or_else(|| { + inner.auth_failures.fetch_add(1, Ordering::Relaxed); + AuthFailure::new( + "temporary_key_inactive", + "temporary key lease is no longer active", + true, + ) + })?; + inner.auth_successes.fetch_add(1, Ordering::Relaxed); + Ok(AuthContext::from_lease(key_id, false, &lease)) + } + + pub fn legacy_protocol_allowed(&self) -> Result { + Ok(self + .inner()? + .legacy_protocol_allowed + .load(Ordering::Acquire)) + } + + pub fn record_legacy_connection(&self) -> Result { + let inner = self.inner()?; + inner + .active_legacy_connections + .fetch_add(1, Ordering::AcqRel); + inner + .last_legacy_connection_at + .store(unix_seconds(), Ordering::Release); + Ok(LegacyConnectionGuard { + inner: Arc::downgrade(&inner), + }) + } + + async fn request( + &self, + build: impl FnOnce(oneshot::Sender>) -> AuthCommand, + ) -> Result { + let (response, receiver) = oneshot::channel(); + self.command_tx.send(build(response)).await.map_err(|_| { + AuthFailure::new( + "auth_state_unavailable", + "authentication state manager is not running", + true, + ) + })?; + receiver.await.map_err(|_| { + AuthFailure::new( + "auth_state_unavailable", + "authentication state manager dropped the response", + true, + ) + })? + } + + pub async fn claim_admin_mutation( + &self, + authorization: &AuthContext, + fingerprint: [u8; 32], + client_timestamp: u64, + ) -> Result<(), AuthFailure> { + let authority = authorization.admin_authority()?; + self.request(|response| AuthCommand::ClaimAdminMutation { + authority, + fingerprint, + client_timestamp, + response, + }) + .await + } + + pub async fn issue( + &self, + authorization: &AuthContext, + ttl: Duration, + label: Option, + ) -> Result { + let authority = authorization.admin_authority()?; + self.request(|response| AuthCommand::Issue { + authority, + ttl, + label, + response, + }) + .await + } + + pub async fn list( + &self, + authorization: &AuthContext, + page: u32, + page_size: u16, + ) -> Result { + let authority = authorization.admin_authority()?; + self.request(|response| AuthCommand::List { + authority, + page, + page_size, + response, + }) + .await + } + + pub async fn show( + &self, + authorization: &AuthContext, + key_id: u64, + reveal: bool, + ) -> Result { + let authority = authorization.admin_authority()?; + self.request(|response| AuthCommand::Show { + authority, + key_id, + reveal, + response, + }) + .await + } + + pub async fn renew( + &self, + authorization: &AuthContext, + key_id: u64, + ttl: Duration, + ) -> Result { + let authority = authorization.admin_authority()?; + self.request(|response| AuthCommand::Renew { + authority, + key_id, + ttl, + response, + }) + .await + } + + pub async fn revoke( + &self, + authorization: &AuthContext, + key_id: u64, + ) -> Result { + let authority = authorization.admin_authority()?; + self.request(|response| AuthCommand::Revoke { + authority, + key_id, + response, + }) + .await + } + + pub async fn gc(&self, authorization: &AuthContext) -> Result { + let authority = authorization.admin_authority()?; + self.request(|response| AuthCommand::Gc { + authority, + response, + }) + .await + } + + pub async fn reset(&self, authorization: &AuthContext) -> Result<(), AuthFailure> { + let authority = authorization.admin_authority()?; + self.request(|response| AuthCommand::Reset { + authority, + response, + }) + .await + } + + pub async fn rotate_root( + &self, + authorization: &AuthContext, + new_key: AesKeyType, + ) -> Result<(), AuthFailure> { + let authority = authorization.admin_authority()?; + self.request(|response| AuthCommand::RotateRoot { + authority, + new_key, + response, + }) + .await + } + + pub async fn set_legacy_protocol( + &self, + authorization: &AuthContext, + policy: LegacyProtocolPolicy, + ) -> Result<(), AuthFailure> { + let authority = authorization.admin_authority()?; + self.request(|response| AuthCommand::SetLegacyProtocol { + authority, + policy, + response, + }) + .await + } + + pub async fn status(&self, authorization: &AuthContext) -> Result { + let authority = authorization.admin_authority()?; + self.request(|response| AuthCommand::Status { + authority, + response, + }) + .await + } + + pub async fn audit_admin( + &self, + authorization: &AuthContext, + action: impl Into, + key_id: Option, + detail: Option, + ) -> Result<(), AuthFailure> { + let authority = authorization.admin_authority()?; + let action = action.into(); + self.request(|response| AuthCommand::Audit { + authority, + action, + key_id, + detail, + response, + }) + .await + } +} diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs new file mode 100644 index 0000000..84dde37 --- /dev/null +++ b/src/common/auth/tests.rs @@ -0,0 +1,307 @@ +use super::*; + +fn temp_state_dir(name: &str) -> PathBuf { + let mut suffix = [0_u8; 8]; + let mut rng = rand::rng(); + for byte in &mut suffix { + *byte = rng.random(); + } + std::env::temp_dir().join(format!("pb-mapper-{name}-{}", hex(&suffix))) +} + +#[test] +fn key_id_round_trip() { + let key_id = make_key_id(42, 65_535); + assert_eq!(key_generation(key_id), 42); + assert_eq!(key_slot(key_id), 65_535); +} + +#[test] +fn derived_key_is_bound_to_instance_and_key_id() { + let admin_key = *b"0123456789abcdefghijklmnopqrstuv"; + let instance_a = [1_u8; INSTANCE_ID_LEN]; + let instance_b = [2_u8; INSTANCE_ID_LEN]; + let key = derive_temporary_key(&admin_key, &instance_a, make_key_id(1, 7)).unwrap(); + assert_eq!( + key, + derive_temporary_key(&admin_key, &instance_a, make_key_id(1, 7)).unwrap() + ); + assert_ne!( + key, + derive_temporary_key(&admin_key, &instance_b, make_key_id(1, 7)).unwrap() + ); + assert_ne!( + key, + derive_temporary_key(&admin_key, &instance_a, make_key_id(2, 7)).unwrap() + ); +} + +#[tokio::test] +async fn issue_renew_revoke_and_persist() { + let state_dir = temp_state_dir("auth-lifecycle"); + let admin_key = *b"0123456789abcdefghijklmnopqrstuv"; + let config = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 8, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + let runtime = AuthRuntime::start(admin_key, config.clone()).await.unwrap(); + let admin = runtime.authenticate(0).unwrap(); + let issued = runtime + .issue(&admin, Duration::from_secs(60), Some("demo".to_string())) + .await + .unwrap(); + assert!(issued.credential.starts_with("pbmt1_")); + let context = runtime.authenticate(issued.metadata.key_id).unwrap(); + assert!(!context.is_admin); + let cancellation = context.cancellation_token().unwrap(); + let renewed = runtime + .renew(&admin, issued.metadata.key_id, Duration::from_secs(120)) + .await + .unwrap(); + assert_eq!(renewed.metadata.key_id, issued.metadata.key_id); + assert_eq!(renewed.credential, issued.credential); + assert!(renewed.metadata.expires_at > issued.metadata.expires_at); + runtime + .revoke(&admin, issued.metadata.key_id) + .await + .unwrap(); + assert!(cancellation.is_cancelled()); + assert_eq!( + context.ensure_active().unwrap_err().code, + "temporary_key_revoked" + ); + assert_eq!( + runtime + .authenticate(issued.metadata.key_id) + .unwrap_err() + .code, + "temporary_key_revoked" + ); + drop(runtime); + + tokio::time::sleep(Duration::from_millis(20)).await; + let restored = AuthRuntime::start(admin_key, config).await.unwrap(); + assert_eq!( + restored + .authenticate(issued.metadata.key_id) + .unwrap_err() + .code, + "temporary_key_revoked" + ); + let _ = std::fs::remove_dir_all(state_dir); +} + +#[tokio::test] +async fn reset_rotates_instance_and_prevents_old_key_id_reuse() { + let state_dir = temp_state_dir("auth-reset"); + let admin_key = *b"0123456789abcdefghijklmnopqrstuv"; + let config = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 1, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + let runtime = AuthRuntime::start(admin_key, config).await.unwrap(); + let admin = runtime.authenticate(0).unwrap(); + let before = runtime.status(&admin).await.unwrap().server_instance_id; + let old = runtime + .issue( + &admin, + Duration::from_secs(60), + Some("before-reset".to_string()), + ) + .await + .unwrap(); + let old_context = runtime.authenticate(old.metadata.key_id).unwrap(); + let old_cancellation = old_context.cancellation_token().unwrap(); + + runtime.reset(&admin).await.unwrap(); + + let after = runtime.status(&admin).await.unwrap().server_instance_id; + assert_ne!(after, before); + assert!(old_cancellation.is_cancelled()); + assert!(runtime.authenticate(old.metadata.key_id).is_err()); + let replacement = runtime + .issue( + &admin, + Duration::from_secs(60), + Some("after-reset".to_string()), + ) + .await + .unwrap(); + assert_ne!(replacement.metadata.key_id, old.metadata.key_id); + assert_ne!(replacement.credential, old.credential); + + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + let _ = std::fs::remove_dir_all(state_dir); +} + +#[tokio::test] +async fn corrupt_wal_fails_temporary_keys_closed_until_admin_reset() { + let state_dir = temp_state_dir("auth-safe-mode"); + let admin_key = *b"0123456789abcdefghijklmnopqrstuv"; + let config = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 4, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + let runtime = AuthRuntime::start(admin_key, config.clone()).await.unwrap(); + let admin = runtime.authenticate(0).unwrap(); + let issued = runtime + .issue( + &admin, + Duration::from_secs(60), + Some("corrupt-me".to_string()), + ) + .await + .unwrap(); + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + std::fs::write(state_dir.join("auth.wal"), b"broken-wal").unwrap(); + + let recovered = AuthRuntime::start(admin_key, config).await.unwrap(); + let recovered_admin = recovered.authenticate(0).unwrap(); + assert!(recovered.status(&recovered_admin).await.unwrap().safe_mode); + assert_eq!( + recovered + .authenticate(issued.metadata.key_id) + .unwrap_err() + .code, + "temporary_key_store_unavailable" + ); + recovered.reset(&recovered_admin).await.unwrap(); + assert!(!recovered.status(&recovered_admin).await.unwrap().safe_mode); + + drop(recovered); + tokio::time::sleep(Duration::from_millis(20)).await; + let _ = std::fs::remove_dir_all(state_dir); +} + +#[tokio::test] +async fn root_rotation_rejects_old_key_and_in_flight_admin_context() { + let state_dir = temp_state_dir("auth-root-rotation"); + let old_key = *b"0123456789abcdefghijklmnopqrstuv"; + let new_key = *b"abcdefghijklmnopqrstuvwxyz012345"; + let config = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 4, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + let runtime = AuthRuntime::start(old_key, config).await.unwrap(); + let old_admin = runtime.authenticate_presented(0, &old_key).unwrap(); + + runtime + .rotate_root(&old_admin, new_key) + .await + .expect("root rotation should succeed"); + + assert_eq!( + runtime + .authenticate_presented(0, &old_key) + .unwrap_err() + .code, + "administrator_key_rotated" + ); + assert_eq!( + runtime + .issue(&old_admin, Duration::from_secs(60), None) + .await + .unwrap_err() + .code, + "administrator_key_rotated" + ); + let new_admin = runtime.authenticate_presented(0, &new_key).unwrap(); + assert!(runtime.status(&new_admin).await.is_ok()); + + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + let _ = std::fs::remove_dir_all(state_dir); +} + +#[tokio::test] +async fn admitted_admin_mutation_replay_survives_restart() { + let state_dir = temp_state_dir("admin-replay-restart"); + let admin_key = *b"0123456789abcdefghijklmnopqrstuv"; + let config = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 4, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + let fingerprint = [0x5a; 32]; + let timestamp = unix_seconds(); + let runtime = AuthRuntime::start(admin_key, config.clone()).await.unwrap(); + let admin = runtime.authenticate(0).unwrap(); + runtime + .claim_admin_mutation(&admin, fingerprint, timestamp) + .await + .unwrap(); + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + + let restored = AuthRuntime::start(admin_key, config).await.unwrap(); + let restored_admin = restored.authenticate(0).unwrap(); + assert_eq!( + restored + .claim_admin_mutation(&restored_admin, fingerprint, timestamp) + .await + .unwrap_err() + .code, + "admin_request_replayed" + ); + + drop(restored); + tokio::time::sleep(Duration::from_millis(20)).await; + let _ = std::fs::remove_dir_all(state_dir); +} + +#[tokio::test] +async fn snapshot_compaction_preserves_audit_records() { + let state_dir = temp_state_dir("audit-compaction"); + let admin_key = *b"0123456789abcdefghijklmnopqrstuv"; + let config = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 4, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + let runtime = AuthRuntime::start(admin_key, config.clone()).await.unwrap(); + let admin = runtime.authenticate(0).unwrap(); + runtime + .issue(&admin, Duration::from_secs(60), Some("audited".to_string())) + .await + .unwrap(); + runtime.gc(&admin).await.unwrap(); + + let instance_id = load_or_create_instance_id(&state_dir).unwrap(); + let persisted = try_load_persisted_state(&config, &admin_key, instance_id).unwrap(); + let actions = persisted + .audit_records + .iter() + .map(|record| record.action.as_str()) + .collect::>(); + assert!(actions.contains(&"temporary_key_issue")); + assert!(actions.contains(&"temporary_key_gc")); + + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + let _ = std::fs::remove_dir_all(state_dir); +} + +#[test] +fn timing_wheel_ignores_stale_renewal_entry() { + let now = 1_000; + let lease = Arc::new(AuthLease::new(make_key_id(1, 0), now + 5)); + let mut wheel = TimingWheel::new(now); + wheel.insert(lease.clone()); + lease.expires_at.store(now + 20, Ordering::Release); + lease.wheel_version.fetch_add(1, Ordering::AcqRel); + wheel.insert(lease.clone()); + assert!(wheel.advance(now + 6).is_empty()); + assert_eq!(wheel.advance(now + 20).len(), 1); +} diff --git a/src/common/auth/timing_wheel.rs b/src/common/auth/timing_wheel.rs new file mode 100644 index 0000000..d96dafd --- /dev/null +++ b/src/common/auth/timing_wheel.rs @@ -0,0 +1,104 @@ +use super::*; + +struct WheelEntry { + lease: Arc, + version: u64, +} + +pub(super) struct TimingWheel { + now: u64, + level0: Vec>, + level1: Vec>, + level2: Vec>, + level3: Vec>, +} + +impl TimingWheel { + pub(super) fn new(now: u64) -> Self { + Self { + now, + level0: empty_buckets(256), + level1: empty_buckets(64), + level2: empty_buckets(64), + level3: empty_buckets(64), + } + } + + pub(super) fn insert(&mut self, lease: Arc) { + let version = lease.wheel_version.load(Ordering::Acquire); + self.insert_with_version(lease, version); + } + + pub(super) fn insert_with_version(&mut self, lease: Arc, version: u64) { + let expires_at = lease.expires_at(); + let delta = expires_at.saturating_sub(self.now); + let entry = WheelEntry { lease, version }; + if delta < 1 << 8 { + self.level0[(expires_at & 0xff) as usize].push(entry); + } else if delta < 1 << 14 { + self.level1[((expires_at >> 8) & 0x3f) as usize].push(entry); + } else if delta < 1 << 20 { + self.level2[((expires_at >> 14) & 0x3f) as usize].push(entry); + } else { + self.level3[((expires_at >> 20) & 0x3f) as usize].push(entry); + } + } + + pub(super) fn advance(&mut self, target: u64) -> Vec> { + let mut due = Vec::new(); + while self.now < target { + self.now = self.now.saturating_add(1); + if self.now & 0xff == 0 { + self.cascade(1); + if (self.now >> 8) & 0x3f == 0 { + self.cascade(2); + if (self.now >> 14) & 0x3f == 0 { + self.cascade(3); + } + } + } + let index = (self.now & 0xff) as usize; + for entry in std::mem::take(&mut self.level0[index]) { + if entry.version == entry.lease.wheel_version.load(Ordering::Acquire) { + if entry.lease.expires_at() <= self.now { + due.push(entry.lease); + } else { + self.insert(entry.lease); + } + } + } + } + due + } + + fn cascade(&mut self, level: u8) { + let entries = match level { + 1 => { + let index = ((self.now >> 8) & 0x3f) as usize; + std::mem::take(&mut self.level1[index]) + } + 2 => { + let index = ((self.now >> 14) & 0x3f) as usize; + std::mem::take(&mut self.level2[index]) + } + 3 => { + let index = ((self.now >> 20) & 0x3f) as usize; + std::mem::take(&mut self.level3[index]) + } + _ => Vec::new(), + }; + for entry in entries { + if entry.version == entry.lease.wheel_version.load(Ordering::Acquire) { + self.insert_with_version(entry.lease, entry.version); + } + } + } + + pub(super) fn clear(&mut self, now: u64) { + *self = Self::new(now); + } +} + +fn empty_buckets(count: usize) -> Vec> { + std::iter::repeat_with(Vec::new).take(count).collect() +} diff --git a/src/common/checksum.rs b/src/common/checksum.rs index dc07241..3823ec4 100644 --- a/src/common/checksum.rs +++ b/src/common/checksum.rs @@ -28,6 +28,7 @@ const DERIVE_MSG_HEADER_KEY_CHARSET: &[u8] = struct MsgHeaderKeyState { credential: RwLock>, + load_error: RwLock>, hash: AtomicU32, } @@ -63,15 +64,14 @@ fn key_len_error(input: &str) -> String { ) } -fn load_credential_from_env() -> Option { - let raw = std::env::var(ENV_MSG_HEADER_KEY).ok()?; - match parse_credential(raw.trim()) { - Ok(credential) => Some(credential), - Err(error) => { - tracing::error!(reason = "credential_invalid", %error, "invalid MSG_HEADER_KEY"); - None - } - } +fn load_credential_from_env() -> Result, String> { + let Some(raw) = std::env::var_os(ENV_MSG_HEADER_KEY) else { + return Ok(None); + }; + let raw = raw + .into_string() + .map_err(|_| format!("`{ENV_MSG_HEADER_KEY}` must contain valid UTF-8 credential text"))?; + parse_credential(raw.trim()).map(Some) } fn update_runtime_credential(credential: Option) { @@ -84,6 +84,10 @@ fn update_runtime_credential(credential: Option) { .write() .unwrap_or_else(|poisoned| poisoned.into_inner()); *guard = credential; + *MSG_HEADER_KEY_STATE + .load_error + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = None; MSG_HEADER_KEY_STATE.hash.store(hash, Ordering::Release); } @@ -92,19 +96,34 @@ fn update_runtime_credential(credential: Option) { /// This state is mutable so FFI/UI can update `MSG_HEADER_KEY` at runtime /// without restarting the process. static MSG_HEADER_KEY_STATE: LazyLock = LazyLock::new(|| { - let credential = load_credential_from_env(); + let (credential, load_error) = match load_credential_from_env() { + Ok(credential) => (credential, None), + Err(error) => { + tracing::error!(reason = "credential_invalid", %error, "invalid MSG_HEADER_KEY"); + (None, Some(error)) + } + }; let hash = credential .as_ref() .map(|credential| gen_checksum_by_key(credential.key())) .unwrap_or_default(); MsgHeaderKeyState { credential: RwLock::new(credential), + load_error: RwLock::new(load_error), hash: AtomicU32::new(hash), } }); /// Return the configured process credential, failing closed when none exists. pub fn get_process_credential() -> Result { + if let Some(error) = MSG_HEADER_KEY_STATE + .load_error + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clone() + { + return Err(error); + } MSG_HEADER_KEY_STATE .credential .read() diff --git a/src/common/message/command.rs b/src/common/message/command.rs index 5cc6df8..f660c6e 100644 --- a/src/common/message/command.rs +++ b/src/common/message/command.rs @@ -159,6 +159,21 @@ pub enum AdminRequest { }, } +impl AdminRequest { + pub fn is_mutating(&self) -> bool { + matches!( + self, + Self::KeyIssue { .. } + | Self::KeyRenew { .. } + | Self::KeyRevoke { .. } + | Self::KeyGc + | Self::AuthStateReset { .. } + | Self::RootKeyRotate { .. } + | Self::LegacyProtocolSet { .. } + ) + } +} + const fn default_page_size() -> u16 { 100 } diff --git a/src/common/message/secure.rs b/src/common/message/secure.rs index ec3d1f2..346ec2e 100644 --- a/src/common/message/secure.rs +++ b/src/common/message/secure.rs @@ -31,6 +31,8 @@ const DIRECTION_CLIENT_TO_SERVER: u8 = 0; const DIRECTION_SERVER_TO_CLIENT: u8 = 1; const DEFAULT_REPLAY_WINDOW_SECONDS: u64 = 60; const DEFAULT_REPLAY_FILTER_BYTES: usize = 1024 * 1024; +const MAX_INITIAL_PLAINTEXT_LEN: u32 = 64 * 1024; +const MAX_CONNECTION_CLOCK_SKEW_SECONDS: u64 = 5 * 60; #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum HeaderProtocol { @@ -38,15 +40,6 @@ pub enum HeaderProtocol { V2, } -#[derive(Clone)] -struct V2Material { - key_id: u64, - flags: u8, - salt: [u8; CONNECTION_SALT_LEN], - client_to_server: AesKeyType, - server_to_client: AesKeyType, -} - pub struct ClientHeaderSession { protocol: HeaderProtocol, legacy_key: AesKeyType, @@ -62,8 +55,9 @@ impl ClientHeaderSession { pub fn new_v2(credential: &Credential) -> Result { let mut salt = [0_u8; CONNECTION_SALT_LEN]; + salt[..8].copy_from_slice(&unix_seconds().to_be_bytes()); let mut rng = rand::rng(); - for byte in &mut salt { + for byte in &mut salt[8..] { *byte = rng.random(); } let material = derive_material(credential.key_id(), credential.key(), salt)?; @@ -244,6 +238,8 @@ impl ServerHeaderSession { pub struct ServerInitialMessage { pub payload: Vec, pub session: ServerHeaderSession, + pub replay_fingerprint: Option<[u8; 32]>, + pub client_timestamp: Option, } pub struct ServerInitialError { @@ -426,6 +422,8 @@ impl ServerSecurity { context: Some(context), _legacy_guard: Some(legacy_guard), }, + replay_fingerprint: None, + client_timestamp: None, }) } @@ -467,6 +465,18 @@ impl ServerSecurity { let key_id = u64::from_be_bytes(remainder[4..12].try_into().expect("fixed key id")); let salt: [u8; CONNECTION_SALT_LEN] = remainder[12..28].try_into().expect("fixed connection salt"); + let client_timestamp = u64::from_be_bytes(salt[..8].try_into().expect("fixed timestamp")); + let now = unix_seconds(); + if now.abs_diff(client_timestamp) > MAX_CONNECTION_CLOCK_SKEW_SECONDS { + return Err(ServerInitialError { + failure: AuthFailure::new( + "connection_timestamp_invalid", + "protocol-v2 connection timestamp is outside the accepted clock-skew window", + false, + ), + response_session: None, + }); + } let key = self .auth .derive_key(key_id) @@ -500,11 +510,11 @@ impl ServerSecurity { response_session: Some(session_without_context(&session)), })?; let payload = message_reader - .read_msg() + .read_msg_with_limit(MAX_INITIAL_PLAINTEXT_LEN) .await .map_err(|error| ServerInitialError { failure: AuthFailure::new("protocol_v2_decrypt_failed", error.to_string(), false), - response_session: None, + response_session: Some(session_without_context(&session)), })? .to_vec(); @@ -513,7 +523,7 @@ impl ServerSecurity { .replay .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()) - .contains(&fingerprint, unix_seconds()); + .check_and_insert(&fingerprint, unix_seconds()); if replayed { return Err(ServerInitialError { failure: AuthFailure::new( @@ -527,96 +537,24 @@ impl ServerSecurity { let context = self .auth - .authenticate(key_id) + .authenticate_presented(key_id, &key) .map_err(|failure| ServerInitialError { failure, response_session: Some(session_without_context(&session)), })?; - self.replay - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .insert(&fingerprint, unix_seconds()); session.context = Some(context); - Ok(ServerInitialMessage { payload, session }) - } -} - -#[derive(Clone, Copy, Debug)] -pub struct FailureLogDecision { - pub emit: bool, - pub suppressed: u64, -} - -struct FailureLogEntry { - window_started_at: u64, - emitted: u8, - suppressed: u64, -} - -#[derive(Default)] -struct FailureLogLimiter { - entries: std::collections::HashMap<(std::net::IpAddr, u64, String), FailureLogEntry>, - overflow: Option, -} - -impl FailureLogLimiter { - fn record( - &mut self, - peer_ip: std::net::IpAddr, - key_id: u64, - reason: &str, - now: u64, - ) -> FailureLogDecision { - let key = (peer_ip, key_id, reason.to_string()); - if !self.entries.contains_key(&key) && self.entries.len() >= 4096 { - self.entries - .retain(|_, entry| now.saturating_sub(entry.window_started_at) < 120); - if self.entries.len() >= 4096 { - let entry = self.overflow.get_or_insert(FailureLogEntry { - window_started_at: now, - emitted: 0, - suppressed: 0, - }); - return record_failure_entry(entry, now); - } - } - let entry = self.entries.entry(key).or_insert(FailureLogEntry { - window_started_at: now, - emitted: 0, - suppressed: 0, - }); - record_failure_entry(entry, now) - } -} - -fn record_failure_entry(entry: &mut FailureLogEntry, now: u64) -> FailureLogDecision { - if now.saturating_sub(entry.window_started_at) >= 60 { - let suppressed = entry.suppressed; - *entry = FailureLogEntry { - window_started_at: now, - emitted: 1, - suppressed: 0, - }; - return FailureLogDecision { - emit: true, - suppressed, - }; - } - if entry.emitted < 5 { - entry.emitted += 1; - FailureLogDecision { - emit: true, - suppressed: 0, - } - } else { - entry.suppressed = entry.suppressed.saturating_add(1); - FailureLogDecision { - emit: false, - suppressed: 0, - } + Ok(ServerInitialMessage { + payload, + session, + replay_fingerprint: Some(fingerprint), + client_timestamp: Some(client_timestamp), + }) } } +mod limiter; +pub use limiter::FailureLogDecision; +use limiter::FailureLogLimiter; fn session_without_context(session: &ServerHeaderSession) -> ServerHeaderSession { ServerHeaderSession { protocol: session.protocol, @@ -655,299 +593,11 @@ impl MessageWriter for HeaderMessageWriter<'_, T> { } } -pub struct V2MessageReader<'a, T: AsyncReadExt + Unpin> { - reader: &'a mut T, - material: V2Material, - key: LessSafeKey, - direction: u8, - expected_counter: u64, - buffer: Vec, -} - -impl<'a, T: AsyncReadExt + Unpin> V2MessageReader<'a, T> { - fn new( - reader: &'a mut T, - material: V2Material, - direction: u8, - expected_counter: u64, - ) -> Result { - let key_bytes = direction_key(&material, direction); - let key = LessSafeKey::new( - UnboundKey::new(&AES_256_GCM, key_bytes) - .map_err(|_| protocol_error("invalid protocol-v2 read key"))?, - ); - Ok(Self { - reader, - material, - key, - direction, - expected_counter, - buffer: Vec::new(), - }) - } -} - -impl MessageReader for V2MessageReader<'_, T> { - async fn read_msg(&mut self) -> Result<&'_ [u8]> { - let counter = self - .reader - .read_u64() - .await - .map_err(|error| protocol_error(format!("failed to read v2 counter: {error}")))?; - if counter != self.expected_counter { - return Err(protocol_error(format!( - "protocol-v2 counter mismatch: expected {}, got {counter}", - self.expected_counter - ))); - } - let datalen = self - .reader - .read_u32() - .await - .map_err(|error| protocol_error(format!("failed to read v2 length: {error}")))?; - if datalen < AES_256_GCM.tag_len() as u32 || datalen > MAX_MSG_LEN { - return Err(protocol_error(format!( - "protocol-v2 payload length {datalen} is invalid" - ))); - } - self.buffer.resize(datalen as usize, 0); - self.reader - .read_exact(&mut self.buffer) - .await - .map_err(|error| protocol_error(format!("failed to read v2 payload: {error}")))?; - let aad = frame_aad(&self.material, self.direction, counter, datalen); - let plain = self - .key - .open_in_place(nonce(counter), Aad::from(aad.as_slice()), &mut self.buffer) - .map_err(|_| protocol_error("protocol-v2 payload authentication failed"))?; - let plain_len = plain.len(); - self.buffer.truncate(plain_len); - self.expected_counter = self - .expected_counter - .checked_add(1) - .ok_or_else(|| protocol_error("protocol-v2 receive counter exhausted"))?; - Ok(&self.buffer) - } -} - -pub struct V2MessageWriter<'a, T: AsyncWriteExt + Unpin> { - writer: &'a mut T, - material: V2Material, - key: LessSafeKey, - direction: u8, - counter: u64, -} - -impl<'a, T: AsyncWriteExt + Unpin> V2MessageWriter<'a, T> { - fn new(writer: &'a mut T, material: V2Material, direction: u8, counter: u64) -> Result { - let key_bytes = direction_key(&material, direction); - let key = LessSafeKey::new( - UnboundKey::new(&AES_256_GCM, key_bytes) - .map_err(|_| protocol_error("invalid protocol-v2 write key"))?, - ); - Ok(Self { - writer, - material, - key, - direction, - counter, - }) - } -} - -impl MessageWriter for V2MessageWriter<'_, T> { - async fn write_msg(&mut self, message: &[u8]) -> Result<()> { - let encrypted_len = message - .len() - .checked_add(AES_256_GCM.tag_len()) - .and_then(|len| DataLenType::try_from(len).ok()) - .ok_or_else(|| protocol_error("protocol-v2 message is too large"))?; - if encrypted_len > MAX_MSG_LEN { - return Err(protocol_error( - "protocol-v2 message exceeds the maximum length", - )); - } - let counter = self.counter; - let aad = frame_aad(&self.material, self.direction, counter, encrypted_len); - let mut encrypted = message.to_vec(); - self.key - .seal_in_place_append_tag(nonce(counter), Aad::from(aad.as_slice()), &mut encrypted) - .map_err(|_| protocol_error("failed to encrypt protocol-v2 message"))?; - self.writer - .write_u64(counter) - .await - .map_err(|error| protocol_error(format!("failed to write v2 frame header: {error}")))?; - self.writer - .write_u32(encrypted_len) - .await - .map_err(|error| protocol_error(format!("failed to write v2 frame header: {error}")))?; - self.writer - .write_all(&encrypted) - .await - .map_err(|error| protocol_error(format!("failed to write v2 frame body: {error}")))?; - self.counter = self - .counter - .checked_add(1) - .ok_or_else(|| protocol_error("protocol-v2 send counter exhausted"))?; - Ok(()) - } -} - -fn derive_material( - key_id: u64, - credential_key: &AesKeyType, - salt_bytes: [u8; CONNECTION_SALT_LEN], -) -> Result { - let salt = Salt::new(HKDF_SHA256, &salt_bytes); - let pseudo_random_key = salt.extract(credential_key); - let client_to_server = expand_direction(&pseudo_random_key, b"pb-mapper-v2-c2s")?; - let server_to_client = expand_direction(&pseudo_random_key, b"pb-mapper-v2-s2c")?; - Ok(V2Material { - key_id, - flags: 0, - salt: salt_bytes, - client_to_server, - server_to_client, - }) -} - -fn expand_direction( - pseudo_random_key: &ring::hkdf::Prk, - label: &'static [u8], -) -> Result { - let info = [label]; - let output = pseudo_random_key - .expand(&info, HkdfLen(32)) - .map_err(|_| protocol_error("failed to derive protocol-v2 direction key"))?; - let mut key = [0_u8; 32]; - output - .fill(&mut key) - .map_err(|_| protocol_error("failed to fill protocol-v2 direction key"))?; - Ok(key) -} - -struct HkdfLen(usize); - -impl ring::hkdf::KeyType for HkdfLen { - fn len(&self) -> usize { - self.0 - } -} - -fn direction_key(material: &V2Material, direction: u8) -> &AesKeyType { - if direction == DIRECTION_CLIENT_TO_SERVER { - &material.client_to_server - } else { - &material.server_to_client - } -} - -fn first_prefix(material: &V2Material) -> Vec { - let mut prefix = Vec::with_capacity(PROTOCOL_V2_MAGIC.len() + FIRST_PREFIX_REMAINDER_LEN); - prefix.extend_from_slice(&PROTOCOL_V2_MAGIC); - prefix.push(PROTOCOL_V2_VERSION); - prefix.push(material.flags); - prefix.extend_from_slice(&0_u16.to_be_bytes()); - prefix.extend_from_slice(&material.key_id.to_be_bytes()); - prefix.extend_from_slice(&material.salt); - prefix -} - -fn frame_aad(material: &V2Material, direction: u8, counter: u64, datalen: u32) -> Vec { - let mut aad = Vec::with_capacity( - PROTOCOL_V2_MAGIC.len() + FIRST_PREFIX_REMAINDER_LEN + 1 + FRAME_HEADER_LEN, - ); - aad.extend_from_slice(&first_prefix(material)); - aad.push(direction); - aad.extend_from_slice(&counter.to_be_bytes()); - aad.extend_from_slice(&datalen.to_be_bytes()); - aad -} - -fn nonce(counter: u64) -> Nonce { - let mut bytes = [0_u8; 12]; - bytes[4..].copy_from_slice(&counter.to_be_bytes()); - Nonce::assume_unique_for_key(bytes) -} - -fn replay_fingerprint(key_id: u64, salt: &[u8; CONNECTION_SALT_LEN]) -> [u8; 32] { - let mut input = [0_u8; 8 + CONNECTION_SALT_LEN]; - input[..8].copy_from_slice(&key_id.to_be_bytes()); - input[8..].copy_from_slice(salt); - digest(&SHA256, &input) - .as_ref() - .try_into() - .expect("SHA-256 width") -} - -struct RotatingBloom { - current: Vec, - previous: Vec, - current_started_at: u64, - window_seconds: u64, -} - -impl RotatingBloom { - fn new(bytes: usize, window_seconds: u64) -> Self { - Self { - current: vec![0; bytes], - previous: vec![0; bytes], - current_started_at: unix_seconds(), - window_seconds, - } - } - - fn contains(&mut self, fingerprint: &[u8; 32], now: u64) -> bool { - self.rotate(now); - bloom_contains(&self.current, fingerprint) || bloom_contains(&self.previous, fingerprint) - } - - fn insert(&mut self, fingerprint: &[u8; 32], now: u64) { - self.rotate(now); - bloom_insert(&mut self.current, fingerprint); - } - - fn rotate(&mut self, now: u64) { - let elapsed = now.saturating_sub(self.current_started_at); - if elapsed < self.window_seconds { - return; - } - if elapsed >= self.window_seconds.saturating_mul(2) { - self.current.fill(0); - self.previous.fill(0); - } else { - std::mem::swap(&mut self.current, &mut self.previous); - self.current.fill(0); - } - self.current_started_at = now; - } -} - -fn bloom_positions(filter_len: usize, fingerprint: &[u8; 32]) -> [usize; 4] { - let bits = filter_len * 8; - std::array::from_fn(|index| { - let offset = index * 8; - let hash = u64::from_be_bytes( - fingerprint[offset..offset + 8] - .try_into() - .expect("fingerprint chunk"), - ); - hash as usize % bits - }) -} - -fn bloom_contains(filter: &[u8], fingerprint: &[u8; 32]) -> bool { - bloom_positions(filter.len(), fingerprint) - .into_iter() - .all(|position| filter[position / 8] & (1 << (position % 8)) != 0) -} - -fn bloom_insert(filter: &mut [u8], fingerprint: &[u8; 32]) { - for position in bloom_positions(filter.len(), fingerprint) { - filter[position / 8] |= 1 << (position % 8); - } -} - +mod frame; +use frame::{derive_material, first_prefix, V2Material}; +pub use frame::{V2MessageReader, V2MessageWriter}; +mod replay; +use replay::{replay_fingerprint, RotatingBloom}; fn protocol_error(detail: impl Into) -> Error { Error::MsgProtocol { detail: detail.into(), @@ -962,102 +612,4 @@ fn unix_seconds() -> u64 { } #[cfg(test)] -mod tests { - use super::*; - use crate::common::auth::{AuthConfig, LegacyProtocolPolicy}; - use crate::common::checksum::encode_temporary_credential; - - fn temp_config() -> AuthConfig { - let mut random = [0_u8; 8]; - let mut rng = rand::rng(); - for byte in &mut random { - *byte = rng.random(); - } - AuthConfig { - state_dir: std::env::temp_dir() - .join(format!("pb-mapper-v2-{}", u64::from_be_bytes(random))), - max_temporary_keys: 8, - max_temporary_key_ttl: std::time::Duration::from_secs(3600), - legacy_protocol: LegacyProtocolPolicy::Allow, - } - } - - #[tokio::test] - async fn v2_round_trip_uses_directional_counters() { - let credential = Credential::Admin(*b"0123456789abcdefghijklmnopqrstuv"); - let client = ClientHeaderSession::new_v2(&credential).unwrap(); - let config = temp_config(); - let auth = AuthRuntime::start(*credential.key(), config.clone()) - .await - .unwrap(); - let security = ServerSecurity::new(auth); - let (mut client_io, mut server_io) = tokio::io::duplex(4096); - - let client_task = async { - client - .write_initial(&mut client_io, b"request") - .await - .unwrap(); - let mut reader = client.response_reader(&mut client_io).unwrap(); - assert_eq!(reader.read_msg().await.unwrap(), b"response"); - }; - let server_task = async { - let initial = security.read_initial(&mut server_io).await.unwrap(); - assert_eq!(initial.payload, b"request"); - let mut writer = initial.session.response_writer(&mut server_io).unwrap(); - writer.write_msg(b"response").await.unwrap(); - }; - tokio::join!(client_task, server_task); - let _ = std::fs::remove_dir_all(config.state_dir); - } - - #[tokio::test] - async fn temporary_credential_authenticates_without_storing_secret() { - let admin = *b"0123456789abcdefghijklmnopqrstuv"; - let config = temp_config(); - let auth = AuthRuntime::start(admin, config.clone()).await.unwrap(); - let issued = auth - .issue(std::time::Duration::from_secs(60), None) - .await - .unwrap(); - let Credential::Temporary { key_id, key } = - crate::common::checksum::parse_credential(&issued.credential).unwrap() - else { - panic!("expected temporary credential") - }; - assert_eq!(issued.credential, encode_temporary_credential(key_id, &key)); - let client = ClientHeaderSession::new_v2(&Credential::Temporary { key_id, key }).unwrap(); - let security = ServerSecurity::new(auth); - let (mut client_io, mut server_io) = tokio::io::duplex(4096); - let client_task = client.write_initial(&mut client_io, b"temporary"); - let server_task = security.read_initial(&mut server_io); - let (client_result, server_result) = tokio::join!(client_task, server_task); - client_result.unwrap(); - let initial = server_result.unwrap(); - assert_eq!(initial.payload, b"temporary"); - assert_eq!(initial.session.context().unwrap().namespace, key_id); - let _ = std::fs::remove_dir_all(config.state_dir); - } - - #[test] - fn rotating_bloom_covers_current_and_previous_window() { - let mut bloom = RotatingBloom::new(1024, 60); - let value = [7_u8; 32]; - let start = bloom.current_started_at; - assert!(!bloom.contains(&value, start)); - bloom.insert(&value, start); - assert!(bloom.contains(&value, start + 60)); - assert!(!bloom.contains(&value, start + 121)); - } - - #[test] - fn failure_log_limiter_has_a_hard_cardinality_bound() { - let mut limiter = FailureLogLimiter::default(); - let peer = "127.0.0.1".parse().unwrap(); - for key_id in 0..10_000 { - limiter.record(peer, key_id, "invalid", 1_000); - } - assert_eq!(limiter.entries.len(), 4096); - assert!(limiter.overflow.is_some()); - } -} +mod tests; diff --git a/src/common/message/secure/frame.rs b/src/common/message/secure/frame.rs new file mode 100644 index 0000000..23222af --- /dev/null +++ b/src/common/message/secure/frame.rs @@ -0,0 +1,236 @@ +use super::*; + +#[derive(Clone)] +pub(super) struct V2Material { + pub(super) key_id: u64, + pub(super) flags: u8, + pub(super) salt: [u8; CONNECTION_SALT_LEN], + pub(super) client_to_server: AesKeyType, + pub(super) server_to_client: AesKeyType, +} + +pub struct V2MessageReader<'a, T: AsyncReadExt + Unpin> { + reader: &'a mut T, + material: V2Material, + key: LessSafeKey, + direction: u8, + expected_counter: u64, + buffer: Vec, +} + +impl<'a, T: AsyncReadExt + Unpin> V2MessageReader<'a, T> { + pub(super) fn new( + reader: &'a mut T, + material: V2Material, + direction: u8, + expected_counter: u64, + ) -> Result { + let key_bytes = direction_key(&material, direction); + let key = LessSafeKey::new( + UnboundKey::new(&AES_256_GCM, key_bytes) + .map_err(|_| protocol_error("invalid protocol-v2 read key"))?, + ); + Ok(Self { + reader, + material, + key, + direction, + expected_counter, + buffer: Vec::new(), + }) + } + + pub(super) async fn read_msg_with_limit(&mut self, max_plaintext_len: u32) -> Result<&'_ [u8]> { + let counter = self + .reader + .read_u64() + .await + .map_err(|error| protocol_error(format!("failed to read v2 counter: {error}")))?; + if counter != self.expected_counter { + return Err(protocol_error(format!( + "protocol-v2 counter mismatch: expected {}, got {counter}", + self.expected_counter + ))); + } + let datalen = self + .reader + .read_u32() + .await + .map_err(|error| protocol_error(format!("failed to read v2 length: {error}")))?; + let max_encrypted_len = max_plaintext_len.saturating_add(AES_256_GCM.tag_len() as u32); + if datalen < AES_256_GCM.tag_len() as u32 || datalen > max_encrypted_len { + return Err(protocol_error(format!( + "protocol-v2 payload length {datalen} exceeds the {max_plaintext_len}-byte limit" + ))); + } + self.buffer.resize(datalen as usize, 0); + self.reader + .read_exact(&mut self.buffer) + .await + .map_err(|error| protocol_error(format!("failed to read v2 payload: {error}")))?; + let aad = frame_aad(&self.material, self.direction, counter, datalen); + let plain = self + .key + .open_in_place(nonce(counter), Aad::from(aad.as_slice()), &mut self.buffer) + .map_err(|_| protocol_error("protocol-v2 payload authentication failed"))?; + let plain_len = plain.len(); + self.buffer.truncate(plain_len); + self.expected_counter = self + .expected_counter + .checked_add(1) + .ok_or_else(|| protocol_error("protocol-v2 receive counter exhausted"))?; + Ok(&self.buffer) + } +} + +impl MessageReader for V2MessageReader<'_, T> { + async fn read_msg(&mut self) -> Result<&'_ [u8]> { + self.read_msg_with_limit(MAX_MSG_LEN - AES_256_GCM.tag_len() as u32) + .await + } +} + +pub struct V2MessageWriter<'a, T: AsyncWriteExt + Unpin> { + writer: &'a mut T, + material: V2Material, + key: LessSafeKey, + direction: u8, + counter: u64, +} + +impl<'a, T: AsyncWriteExt + Unpin> V2MessageWriter<'a, T> { + pub(super) fn new( + writer: &'a mut T, + material: V2Material, + direction: u8, + counter: u64, + ) -> Result { + let key_bytes = direction_key(&material, direction); + let key = LessSafeKey::new( + UnboundKey::new(&AES_256_GCM, key_bytes) + .map_err(|_| protocol_error("invalid protocol-v2 write key"))?, + ); + Ok(Self { + writer, + material, + key, + direction, + counter, + }) + } +} + +impl MessageWriter for V2MessageWriter<'_, T> { + async fn write_msg(&mut self, message: &[u8]) -> Result<()> { + let encrypted_len = message + .len() + .checked_add(AES_256_GCM.tag_len()) + .and_then(|len| DataLenType::try_from(len).ok()) + .ok_or_else(|| protocol_error("protocol-v2 message is too large"))?; + if encrypted_len > MAX_MSG_LEN { + return Err(protocol_error( + "protocol-v2 message exceeds the maximum length", + )); + } + let counter = self.counter; + let aad = frame_aad(&self.material, self.direction, counter, encrypted_len); + let mut encrypted = message.to_vec(); + self.key + .seal_in_place_append_tag(nonce(counter), Aad::from(aad.as_slice()), &mut encrypted) + .map_err(|_| protocol_error("failed to encrypt protocol-v2 message"))?; + self.writer + .write_u64(counter) + .await + .map_err(|error| protocol_error(format!("failed to write v2 frame header: {error}")))?; + self.writer + .write_u32(encrypted_len) + .await + .map_err(|error| protocol_error(format!("failed to write v2 frame header: {error}")))?; + self.writer + .write_all(&encrypted) + .await + .map_err(|error| protocol_error(format!("failed to write v2 frame body: {error}")))?; + self.counter = self + .counter + .checked_add(1) + .ok_or_else(|| protocol_error("protocol-v2 send counter exhausted"))?; + Ok(()) + } +} + +pub(super) fn derive_material( + key_id: u64, + credential_key: &AesKeyType, + salt_bytes: [u8; CONNECTION_SALT_LEN], +) -> Result { + let salt = Salt::new(HKDF_SHA256, &salt_bytes); + let pseudo_random_key = salt.extract(credential_key); + let client_to_server = expand_direction(&pseudo_random_key, b"pb-mapper-v2-c2s")?; + let server_to_client = expand_direction(&pseudo_random_key, b"pb-mapper-v2-s2c")?; + Ok(V2Material { + key_id, + flags: 0, + salt: salt_bytes, + client_to_server, + server_to_client, + }) +} + +fn expand_direction( + pseudo_random_key: &ring::hkdf::Prk, + label: &'static [u8], +) -> Result { + let info = [label]; + let output = pseudo_random_key + .expand(&info, HkdfLen(32)) + .map_err(|_| protocol_error("failed to derive protocol-v2 direction key"))?; + let mut key = [0_u8; 32]; + output + .fill(&mut key) + .map_err(|_| protocol_error("failed to fill protocol-v2 direction key"))?; + Ok(key) +} + +struct HkdfLen(usize); + +impl ring::hkdf::KeyType for HkdfLen { + fn len(&self) -> usize { + self.0 + } +} + +fn direction_key(material: &V2Material, direction: u8) -> &AesKeyType { + if direction == DIRECTION_CLIENT_TO_SERVER { + &material.client_to_server + } else { + &material.server_to_client + } +} + +pub(super) fn first_prefix(material: &V2Material) -> Vec { + let mut prefix = Vec::with_capacity(PROTOCOL_V2_MAGIC.len() + FIRST_PREFIX_REMAINDER_LEN); + prefix.extend_from_slice(&PROTOCOL_V2_MAGIC); + prefix.push(PROTOCOL_V2_VERSION); + prefix.push(material.flags); + prefix.extend_from_slice(&0_u16.to_be_bytes()); + prefix.extend_from_slice(&material.key_id.to_be_bytes()); + prefix.extend_from_slice(&material.salt); + prefix +} + +fn frame_aad(material: &V2Material, direction: u8, counter: u64, datalen: u32) -> Vec { + let mut aad = Vec::with_capacity( + PROTOCOL_V2_MAGIC.len() + FIRST_PREFIX_REMAINDER_LEN + 1 + FRAME_HEADER_LEN, + ); + aad.extend_from_slice(&first_prefix(material)); + aad.push(direction); + aad.extend_from_slice(&counter.to_be_bytes()); + aad.extend_from_slice(&datalen.to_be_bytes()); + aad +} + +fn nonce(counter: u64) -> Nonce { + let mut bytes = [0_u8; 12]; + bytes[4..].copy_from_slice(&counter.to_be_bytes()); + Nonce::assume_unique_for_key(bytes) +} diff --git a/src/common/message/secure/limiter.rs b/src/common/message/secure/limiter.rs new file mode 100644 index 0000000..59177ce --- /dev/null +++ b/src/common/message/secure/limiter.rs @@ -0,0 +1,75 @@ +#[derive(Clone, Copy, Debug)] +pub struct FailureLogDecision { + pub emit: bool, + pub suppressed: u64, +} + +pub(super) struct FailureLogEntry { + window_started_at: u64, + emitted: u8, + suppressed: u64, +} + +#[derive(Default)] +pub(super) struct FailureLogLimiter { + pub(super) entries: std::collections::HashMap<(std::net::IpAddr, u64, String), FailureLogEntry>, + pub(super) overflow: Option, +} + +impl FailureLogLimiter { + pub(super) fn record( + &mut self, + peer_ip: std::net::IpAddr, + key_id: u64, + reason: &str, + now: u64, + ) -> FailureLogDecision { + let key = (peer_ip, key_id, reason.to_string()); + if !self.entries.contains_key(&key) && self.entries.len() >= 4096 { + self.entries + .retain(|_, entry| now.saturating_sub(entry.window_started_at) < 120); + if self.entries.len() >= 4096 { + let entry = self.overflow.get_or_insert(FailureLogEntry { + window_started_at: now, + emitted: 0, + suppressed: 0, + }); + return record_failure_entry(entry, now); + } + } + let entry = self.entries.entry(key).or_insert(FailureLogEntry { + window_started_at: now, + emitted: 0, + suppressed: 0, + }); + record_failure_entry(entry, now) + } +} + +fn record_failure_entry(entry: &mut FailureLogEntry, now: u64) -> FailureLogDecision { + if now.saturating_sub(entry.window_started_at) >= 60 { + let suppressed = entry.suppressed; + *entry = FailureLogEntry { + window_started_at: now, + emitted: 1, + suppressed: 0, + }; + return FailureLogDecision { + emit: true, + suppressed, + }; + } + if entry.emitted < 5 { + entry.emitted += 1; + FailureLogDecision { + emit: true, + suppressed: 0, + } + } else { + entry.suppressed = entry.suppressed.saturating_add(1); + FailureLogDecision { + emit: false, + suppressed: 0, + } + } +} diff --git a/src/common/message/secure/replay.rs b/src/common/message/secure/replay.rs new file mode 100644 index 0000000..0ebd0d2 --- /dev/null +++ b/src/common/message/secure/replay.rs @@ -0,0 +1,87 @@ +use super::*; + +pub(super) fn replay_fingerprint(key_id: u64, salt: &[u8; CONNECTION_SALT_LEN]) -> [u8; 32] { + let mut input = [0_u8; 8 + CONNECTION_SALT_LEN]; + input[..8].copy_from_slice(&key_id.to_be_bytes()); + input[8..].copy_from_slice(salt); + digest(&SHA256, &input) + .as_ref() + .try_into() + .expect("SHA-256 width") +} + +pub(super) struct RotatingBloom { + current: Vec, + previous: Vec, + pub(super) current_started_at: u64, + window_seconds: u64, +} + +impl RotatingBloom { + pub(super) fn new(bytes: usize, window_seconds: u64) -> Self { + Self { + current: vec![0; bytes], + previous: vec![0; bytes], + current_started_at: unix_seconds(), + window_seconds, + } + } + + pub(super) fn contains(&mut self, fingerprint: &[u8; 32], now: u64) -> bool { + self.rotate(now); + bloom_contains(&self.current, fingerprint) || bloom_contains(&self.previous, fingerprint) + } + + pub(super) fn insert(&mut self, fingerprint: &[u8; 32], now: u64) { + self.rotate(now); + bloom_insert(&mut self.current, fingerprint); + } + + pub(super) fn check_and_insert(&mut self, fingerprint: &[u8; 32], now: u64) -> bool { + if self.contains(fingerprint, now) { + return true; + } + self.insert(fingerprint, now); + false + } + + fn rotate(&mut self, now: u64) { + let elapsed = now.saturating_sub(self.current_started_at); + if elapsed < self.window_seconds { + return; + } + if elapsed >= self.window_seconds.saturating_mul(2) { + self.current.fill(0); + self.previous.fill(0); + } else { + std::mem::swap(&mut self.current, &mut self.previous); + self.current.fill(0); + } + self.current_started_at = now; + } +} + +fn bloom_positions(filter_len: usize, fingerprint: &[u8; 32]) -> [usize; 4] { + let bits = filter_len * 8; + std::array::from_fn(|index| { + let offset = index * 8; + let hash = u64::from_be_bytes( + fingerprint[offset..offset + 8] + .try_into() + .expect("fingerprint chunk"), + ); + hash as usize % bits + }) +} + +fn bloom_contains(filter: &[u8], fingerprint: &[u8; 32]) -> bool { + bloom_positions(filter.len(), fingerprint) + .into_iter() + .all(|position| filter[position / 8] & (1 << (position % 8)) != 0) +} + +fn bloom_insert(filter: &mut [u8], fingerprint: &[u8; 32]) { + for position in bloom_positions(filter.len(), fingerprint) { + filter[position / 8] |= 1 << (position % 8); + } +} diff --git a/src/common/message/secure/tests.rs b/src/common/message/secure/tests.rs new file mode 100644 index 0000000..e8cda72 --- /dev/null +++ b/src/common/message/secure/tests.rs @@ -0,0 +1,171 @@ +use super::*; +use crate::common::auth::{AuthConfig, LegacyProtocolPolicy}; +use crate::common::checksum::encode_temporary_credential; + +fn temp_config() -> AuthConfig { + let mut random = [0_u8; 8]; + let mut rng = rand::rng(); + for byte in &mut random { + *byte = rng.random(); + } + AuthConfig { + state_dir: std::env::temp_dir() + .join(format!("pb-mapper-v2-{}", u64::from_be_bytes(random))), + max_temporary_keys: 8, + max_temporary_key_ttl: std::time::Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + } +} + +#[tokio::test] +async fn v2_round_trip_uses_directional_counters() { + let credential = Credential::Admin(*b"0123456789abcdefghijklmnopqrstuv"); + let client = ClientHeaderSession::new_v2(&credential).unwrap(); + let config = temp_config(); + let auth = AuthRuntime::start(*credential.key(), config.clone()) + .await + .unwrap(); + let security = ServerSecurity::new(auth); + let (mut client_io, mut server_io) = tokio::io::duplex(4096); + + let client_task = async { + client + .write_initial(&mut client_io, b"request") + .await + .unwrap(); + let mut reader = client.response_reader(&mut client_io).unwrap(); + assert_eq!(reader.read_msg().await.unwrap(), b"response"); + }; + let server_task = async { + let initial = security.read_initial(&mut server_io).await.unwrap(); + assert_eq!(initial.payload, b"request"); + let mut writer = initial.session.response_writer(&mut server_io).unwrap(); + writer.write_msg(b"response").await.unwrap(); + }; + tokio::join!(client_task, server_task); + let _ = std::fs::remove_dir_all(config.state_dir); +} + +#[tokio::test] +async fn temporary_credential_authenticates_without_storing_secret() { + let admin = *b"0123456789abcdefghijklmnopqrstuv"; + let config = temp_config(); + let auth = AuthRuntime::start(admin, config.clone()).await.unwrap(); + let admin_context = auth.authenticate(0).unwrap(); + let issued = auth + .issue(&admin_context, std::time::Duration::from_secs(60), None) + .await + .unwrap(); + let Credential::Temporary { key_id, key } = + crate::common::checksum::parse_credential(&issued.credential).unwrap() + else { + panic!("expected temporary credential") + }; + assert_eq!(issued.credential, encode_temporary_credential(key_id, &key)); + let client = ClientHeaderSession::new_v2(&Credential::Temporary { key_id, key }).unwrap(); + let security = ServerSecurity::new(auth); + let (mut client_io, mut server_io) = tokio::io::duplex(4096); + let client_task = client.write_initial(&mut client_io, b"temporary"); + let server_task = security.read_initial(&mut server_io); + let (client_result, server_result) = tokio::join!(client_task, server_task); + client_result.unwrap(); + let initial = server_result.unwrap(); + assert_eq!(initial.payload, b"temporary"); + assert_eq!(initial.session.context().unwrap().namespace, key_id); + let _ = std::fs::remove_dir_all(config.state_dir); +} + +async fn encode_initial(session: &ClientHeaderSession, payload: &[u8]) -> Vec { + let (mut writer, mut reader) = tokio::io::duplex(128 * 1024); + let write = async { + session.write_initial(&mut writer, payload).await.unwrap(); + drop(writer); + }; + let read = async { + let mut bytes = Vec::new(); + reader.read_to_end(&mut bytes).await.unwrap(); + bytes + }; + let (_, bytes) = tokio::join!(write, read); + bytes +} + +#[tokio::test] +async fn identical_initial_frames_are_admitted_only_once() { + let credential = Credential::Admin(*b"0123456789abcdefghijklmnopqrstuv"); + let session = ClientHeaderSession::new_v2(&credential).unwrap(); + let bytes = encode_initial(&session, b"same-request").await; + let config = temp_config(); + let auth = AuthRuntime::start(*credential.key(), config.clone()) + .await + .unwrap(); + let security = ServerSecurity::new(auth); + let mut first = std::io::Cursor::new(bytes.clone()); + let mut second = std::io::Cursor::new(bytes); + + let (first, second) = tokio::join!( + security.read_initial(&mut first), + security.read_initial(&mut second) + ); + let results = [first, second]; + assert_eq!(results.iter().filter(|result| result.is_ok()).count(), 1); + assert_eq!( + results + .iter() + .filter_map(|result| result.as_ref().err()) + .next() + .unwrap() + .failure + .code, + "connection_salt_replayed" + ); + + let _ = std::fs::remove_dir_all(config.state_dir); +} + +#[tokio::test] +async fn oversized_initial_frame_is_rejected_before_reading_its_body() { + let credential = Credential::Admin(*b"0123456789abcdefghijklmnopqrstuv"); + let session = ClientHeaderSession::new_v2(&credential).unwrap(); + let material = session.v2.as_ref().unwrap(); + let mut bytes = first_prefix(material); + bytes.extend_from_slice(&0_u64.to_be_bytes()); + bytes.extend_from_slice(&(MAX_INITIAL_PLAINTEXT_LEN + 17).to_be_bytes()); + let config = temp_config(); + let auth = AuthRuntime::start(*credential.key(), config.clone()) + .await + .unwrap(); + let security = ServerSecurity::new(auth); + let mut input = std::io::Cursor::new(bytes); + + let error = match security.read_initial(&mut input).await { + Ok(_) => panic!("oversized initial frame was accepted"), + Err(error) => error, + }; + assert_eq!(error.failure.code, "protocol_v2_decrypt_failed"); + assert!(error.failure.message.contains("65536-byte limit")); + + let _ = std::fs::remove_dir_all(config.state_dir); +} + +#[test] +fn rotating_bloom_covers_current_and_previous_window() { + let mut bloom = RotatingBloom::new(1024, 60); + let value = [7_u8; 32]; + let start = bloom.current_started_at; + assert!(!bloom.contains(&value, start)); + bloom.insert(&value, start); + assert!(bloom.contains(&value, start + 60)); + assert!(!bloom.contains(&value, start + 121)); +} + +#[test] +fn failure_log_limiter_has_a_hard_cardinality_bound() { + let mut limiter = FailureLogLimiter::default(); + let peer = "127.0.0.1".parse().unwrap(); + for key_id in 0..10_000 { + limiter.record(peer, key_id, "invalid", 1_000); + } + assert_eq!(limiter.entries.len(), 4096); + assert!(limiter.overflow.is_some()); +} diff --git a/src/pb_server/admin.rs b/src/pb_server/admin.rs index 964fcef..46e103f 100644 --- a/src/pb_server/admin.rs +++ b/src/pb_server/admin.rs @@ -4,7 +4,7 @@ use tokio::net::TcpStream; use super::error::Error; use super::{ManagerTask, ManagerTaskSender, Result}; -use crate::common::auth::{AuthFailure, AuthRuntime}; +use crate::common::auth::{AuthContext, AuthFailure, AuthRuntime}; use crate::common::checksum::{parse_credential, Credential}; use crate::common::conn_id::RemoteConnId; use crate::common::message::command::{ @@ -15,13 +15,14 @@ use crate::common::message::MessageWriter; pub async fn handle_admin_request( request: AdminRequest, + authorization: AuthContext, auth: AuthRuntime, manager: ManagerTaskSender, conn_id: RemoteConnId, mut conn: TcpStream, session: ServerHeaderSession, ) -> Result<()> { - let result = execute(request, auth, manager).await; + let result = execute(request, &authorization, auth, manager).await; let response = match result { Ok(response) => PbConnResponse::Admin(response), Err(failure) => { @@ -55,47 +56,56 @@ pub async fn handle_admin_request( async fn execute( request: AdminRequest, + authorization: &AuthContext, auth: AuthRuntime, manager: ManagerTaskSender, ) -> std::result::Result { match request { AdminRequest::KeyIssue { ttl_seconds, label } => auth - .issue(Duration::from_secs(ttl_seconds), label) + .issue(authorization, Duration::from_secs(ttl_seconds), label) .await .map(AdminResponse::KeyIssued), AdminRequest::KeyList { page, page_size } => { audit_read( &auth, + authorization, "temporary_key_list", None, Some(format!("page={page},page_size={page_size}")), ) .await; - auth.list(page, page_size).await.map(AdminResponse::KeyList) - } - AdminRequest::KeyShow { key_id } => { - auth.show(key_id, false).await.map(AdminResponse::KeyShown) - } - AdminRequest::KeyReveal { key_id } => { - auth.show(key_id, true).await.map(AdminResponse::KeyShown) + auth.list(authorization, page, page_size) + .await + .map(AdminResponse::KeyList) } + AdminRequest::KeyShow { key_id } => auth + .show(authorization, key_id, false) + .await + .map(AdminResponse::KeyShown), + AdminRequest::KeyReveal { key_id } => auth + .show(authorization, key_id, true) + .await + .map(AdminResponse::KeyShown), AdminRequest::KeyRenew { key_id, ttl_seconds, } => auth - .renew(key_id, Duration::from_secs(ttl_seconds)) + .renew(authorization, key_id, Duration::from_secs(ttl_seconds)) .await .map(AdminResponse::KeyRenewed), - AdminRequest::KeyRevoke { key_id } => { - auth.revoke(key_id).await.map(AdminResponse::KeyRevoked) - } + AdminRequest::KeyRevoke { key_id } => auth + .revoke(authorization, key_id) + .await + .map(AdminResponse::KeyRevoked), AdminRequest::KeyGc => auth - .gc() + .gc(authorization) .await .map(|removed| AdminResponse::KeyGc { removed }), AdminRequest::AuthStatus => { - audit_read(&auth, "auth_status", None, None).await; - auth.status().await.map(AdminResponse::AuthStatus) + audit_read(&auth, authorization, "auth_status", None, None).await; + auth.status(authorization) + .await + .map(AdminResponse::AuthStatus) } AdminRequest::AuthStateReset { confirm } => { if !confirm { @@ -105,7 +115,7 @@ async fn execute( false, )); } - auth.reset().await?; + auth.reset(authorization).await?; Ok(AdminResponse::Ok { action: "auth_state_reset".to_string(), }) @@ -120,13 +130,13 @@ async fn execute( false, )); }; - auth.rotate_root(new_key).await?; + auth.rotate_root(authorization, new_key).await?; Ok(AdminResponse::Ok { action: "administrator_key_rotated".to_string(), }) } AdminRequest::LegacyProtocolSet { policy } => { - auth.set_legacy_protocol(policy).await?; + auth.set_legacy_protocol(authorization, policy).await?; Ok(AdminResponse::Ok { action: "legacy_protocol_updated".to_string(), }) @@ -138,6 +148,7 @@ async fn execute( } => { audit_read( &auth, + authorization, "service_list", key_id, Some(format!("page={page},page_size={page_size}")), @@ -174,6 +185,7 @@ async fn execute( } => { audit_read( &auth, + authorization, "connection_list", key_id, Some(format!("page={page},page_size={page_size}")), @@ -206,8 +218,17 @@ async fn execute( } } -async fn audit_read(auth: &AuthRuntime, action: &str, key_id: Option, detail: Option) { - if let Err(error) = auth.audit_admin(action, key_id, detail).await { +async fn audit_read( + auth: &AuthRuntime, + authorization: &AuthContext, + action: &str, + key_id: Option, + detail: Option, +) { + if let Err(error) = auth + .audit_admin(authorization, action, key_id, detail) + .await + { tracing::warn!( event = "admin_audit_failed", auth_stage = "audit", diff --git a/src/pb_server/connection.rs b/src/pb_server/connection.rs new file mode 100644 index 0000000..dff693c --- /dev/null +++ b/src/pb_server/connection.rs @@ -0,0 +1,522 @@ +use super::*; + +pub(super) async fn handle_listener( + task_sender: ManagerTaskSender, + listener: TcpListener, + keep_alive: bool, +) -> Result<()> { + loop { + let (stream, addr) = listener.accept().await.context(ServerListenSnafu)?; + tracing::debug!( + event = "tcp_conn_accepted", + peer_addr = %addr, + "accepted tcp connection" + ); + // set keepalive (optional) and nodelay + if keep_alive { + snafu_error_handle!(set_tcp_keep_alive(&stream).context(TaskCenterSetKeepAliveSnafu)); + } + snafu_error_handle!(set_tcp_nodelay(&stream), "remote stream set nodelay"); + task_sender + .send(ManagerTask::Accept { + stream, + peer_addr: addr, + }) + .await + .map_err(|_| kanal::SendError(())) + .context(TaskCenterSendListenerSnafu)? + } +} + +#[instrument(skip(manager_task_sender, conn, security), fields(conn_id = %conn_id, peer_addr = %peer_addr))] +pub(super) async fn handle_conn( + conn_id: RemoteConnId, + peer_addr: SocketAddr, + manager_task_sender: ManagerTaskSender, + mut conn: TcpStream, + security: ServerSecurity, +) -> Result<()> { + let timeout = control_io_timeout(); + let initial = match tokio::time::timeout(timeout, security.read_initial(&mut conn)).await { + Err(_) => TaskCenterInitRequestTimeoutSnafu { conn_id, timeout }.fail()?, + Ok(Err(error)) => { + let key_id = error + .response_session + .as_ref() + .map(|session| session.key_id()) + .unwrap_or_default(); + let decision = security.record_failure_log(peer_addr.ip(), key_id, &error.failure.code); + if decision.suppressed > 0 { + tracing::warn!( + event = "auth_failures_suppressed", + peer_ip = %peer_addr.ip(), + key_id, + reason = %error.failure.code, + suppressed = decision.suppressed, + "suppressed repeated authentication failures in the previous window" + ); + } + if decision.emit { + tracing::warn!( + event = "auth_failed", + auth_stage = "initial_frame", + conn_id = %conn_id, + peer_addr = %peer_addr, + key_id, + reason = %error.failure.code, + retryable = error.failure.retryable, + error = %error.failure.message, + "connection authentication failed" + ); + } + if let Some(session) = error.response_session { + write_protocol_error(&mut conn, &session, &error.failure).await; + } + return Ok(()); + } + Ok(Ok(initial)) => initial, + }; + let replay_fingerprint = initial.replay_fingerprint; + let client_timestamp = initial.client_timestamp; + let init_request = match PbConnRequest::decode(&initial.payload) { + Ok(request) => request, + Err(error) => { + tracing::warn!( + event = "auth_failed", + auth_stage = "request_decode", + conn_id = %conn_id, + peer_addr = %peer_addr, + error = %error, + "authenticated request could not be decoded" + ); + write_protocol_error( + &mut conn, + &initial.session, + &crate::common::auth::AuthFailure::new( + "request_decode_failed", + "authenticated request payload is malformed", + false, + ), + ) + .await; + return Ok(()); + } + }; + let mut requested_namespace = None; + let mut force_register_namespace = false; + let init_request = match init_request { + PbConnRequest::RegisterScoped { + need_codec, + is_datagram, + key, + namespace, + force_namespace, + protocol_version, + client_instance_id, + heartbeat_interval_ms, + heartbeat_tolerance_ms, + } => { + requested_namespace = Some(namespace); + force_register_namespace = force_namespace; + PbConnRequest::Register { + need_codec, + is_datagram, + key, + protocol_version, + client_instance_id, + heartbeat_interval_ms, + heartbeat_tolerance_ms, + } + } + PbConnRequest::SubcribeScoped { key, namespace } => { + requested_namespace = Some(namespace); + PbConnRequest::Subcribe { key } + } + PbConnRequest::StatusScoped { status, namespace } => { + requested_namespace = Some(namespace); + PbConnRequest::Status(status) + } + PbConnRequest::StreamScoped { + key, + namespace, + dst_id, + server_generation, + } => { + requested_namespace = Some(namespace); + PbConnRequest::Stream { + key, + dst_id, + server_generation, + } + } + request => request, + }; + let session = initial.session; + let auth_context = match session.context() { + Ok(context) => context.clone(), + Err(error) => { + tracing::warn!(conn_id = %conn_id, peer_addr = %peer_addr, %error, "missing auth context"); + return Ok(()); + } + }; + tracing::info!( + event = "auth_succeeded", + auth_stage = "session", + conn_id = %conn_id, + peer_addr = %peer_addr, + key_id = auth_context.key_id, + namespace = auth_context.namespace, + protocol = ?session.protocol(), + is_admin = auth_context.is_admin, + "connection authentication succeeded" + ); + let effective_namespace = match resolve_namespace( + &auth_context, + requested_namespace, + force_register_namespace, + matches!(&init_request, PbConnRequest::Register { .. }), + ) { + Ok(namespace) => namespace, + Err(failure) => { + write_protocol_error(&mut conn, &session, &failure).await; + return Ok(()); + } + }; + match init_request { + PbConnRequest::Register { + key, + need_codec, + is_datagram, + protocol_version, + client_instance_id, + heartbeat_interval_ms, + heartbeat_tolerance_ms, + } => { + let protocol_version = protocol_version.unwrap_or(1); + tracing::info!( + event = "init_request", + request = "register", + conn_id = %conn_id, + peer_addr = %peer_addr, + key = %key, + protocol_version, + client_instance_id = ?client_instance_id, + heartbeat_interval_ms = ?heartbeat_interval_ms, + heartbeat_tolerance_ms = ?heartbeat_tolerance_ms, + need_codec, + is_datagram, + "received pb init request" + ); + let key = match scoped_service_key(&auth_context, effective_namespace, &key) { + Ok(key) => key, + Err(failure) => { + write_protocol_error(&mut conn, &session, &failure).await; + return Ok(()); + } + }; + let cancellation = match auth_context.cancellation_token() { + Ok(token) => token, + Err(failure) => { + write_protocol_error(&mut conn, &session, &failure).await; + return Ok(()); + } + }; + tokio::select! { + result = handle_server_conn( + ServerRegistration { + key, + need_codec, + is_datagram, + protocol_version, + conn_id, + }, + manager_task_sender, + conn, + session, + ) => result?, + _ = cancellation.cancelled() => { + tracing::info!(event = "connection_auth_expired", key_id = auth_context.key_id, conn_id = %conn_id, "closing registered service connection"); + } + } + } + PbConnRequest::Subcribe { key } => { + tracing::info!( + event = "init_request", + request = "subscribe", + conn_id = %conn_id, + peer_addr = %peer_addr, + key = %key, + "received pb init request" + ); + let key = match scoped_service_key(&auth_context, effective_namespace, &key) { + Ok(key) => key, + Err(failure) => { + write_protocol_error(&mut conn, &session, &failure).await; + return Ok(()); + } + }; + let cancellation = match auth_context.cancellation_token() { + Ok(token) => token, + Err(failure) => { + write_protocol_error(&mut conn, &session, &failure).await; + return Ok(()); + } + }; + tokio::select! { + result = handle_client_conn(key, conn_id, manager_task_sender, conn, session) => result?, + _ = cancellation.cancelled() => { + tracing::info!(event = "connection_auth_expired", key_id = auth_context.key_id, conn_id = %conn_id, "closing subscribed data connection"); + } + } + } + PbConnRequest::Stream { + key, + dst_id, + server_generation, + } => { + tracing::debug!( + event = "init_request", + request = "stream", + conn_id = %conn_id, + peer_addr = %peer_addr, + key = %key, + client_conn_id = dst_id, + server_generation, + "received pb init request" + ); + let key = match scoped_service_key(&auth_context, effective_namespace, &key) { + Ok(key) => key, + Err(failure) => { + write_protocol_error(&mut conn, &session, &failure).await; + return Ok(()); + } + }; + manager_task_sender + .send(ManagerTask::Stream { + key: key.clone(), + stream: conn, + session, + server_id: conn_id, + client_id: dst_id.into(), + server_generation, + }) + .await + .map_err(|_| kanal::SendError(())) + .context(TaskCenterSendStreamRespToManagerSnafu { key, conn_id })?; + } + PbConnRequest::Status(status) => { + tracing::debug!( + event = "init_request", + request = "status", + conn_id = %conn_id, + peer_addr = %peer_addr, + status = ?status, + "received pb init request" + ); + handle_show_status( + status, + effective_namespace, + manager_task_sender, + conn_id, + conn, + session, + ) + .await?; + } + PbConnRequest::Admin(request) => { + if !auth_context.is_admin { + write_protocol_error( + &mut conn, + &session, + &crate::common::auth::AuthFailure::new( + "admin_permission_required", + "administrator credential is required for this operation", + false, + ), + ) + .await; + return Ok(()); + } + if session.protocol() != HeaderProtocol::V2 { + write_protocol_error( + &mut conn, + &session, + &crate::common::auth::AuthFailure::new( + "admin_protocol_v2_required", + "administrator operations require protocol v2", + false, + ), + ) + .await; + return Ok(()); + } + if request.is_mutating() { + let Some(fingerprint) = replay_fingerprint else { + unreachable!("protocol-v2 sessions always carry a replay fingerprint"); + }; + let Some(client_timestamp) = client_timestamp else { + unreachable!("protocol-v2 sessions always carry a client timestamp"); + }; + if let Err(failure) = security + .auth() + .claim_admin_mutation(&auth_context, fingerprint, client_timestamp) + .await + { + write_protocol_error(&mut conn, &session, &failure).await; + return Ok(()); + } + } + handle_admin_request( + request, + auth_context, + security.auth().clone(), + manager_task_sender, + conn_id, + conn, + session, + ) + .await?; + } + PbConnRequest::RegisterScoped { .. } + | PbConnRequest::SubcribeScoped { .. } + | PbConnRequest::StatusScoped { .. } + | PbConnRequest::StreamScoped { .. } => unreachable!("scoped request was normalized"), + } + Ok(()) +} + +async fn write_protocol_error( + conn: &mut TcpStream, + session: &ServerHeaderSession, + failure: &crate::common::auth::AuthFailure, +) { + let response = PbConnResponse::error( + failure.code.clone(), + failure.message.clone(), + failure.retryable, + ); + let Ok(message) = response.encode() else { + return; + }; + let Ok(mut writer) = session.response_writer(conn) else { + return; + }; + if let Err(error) = writer.write_msg(&message).await { + tracing::debug!(%error, reason = %failure.code, "failed to write structured protocol error"); + } +} + +fn resolve_namespace( + context: &AuthContext, + requested: Option, + force_register_namespace: bool, + is_register: bool, +) -> std::result::Result { + let namespace = requested.unwrap_or(context.namespace); + if !context.is_admin && namespace != context.namespace { + return Err(crate::common::auth::AuthFailure::new( + "namespace_access_denied", + "temporary credentials can only access their own namespace", + false, + )); + } + if context.is_admin && is_register && namespace != 0 && !force_register_namespace { + return Err(crate::common::auth::AuthFailure::new( + "namespace_force_required", + "administrator registration in a temporary namespace requires --force", + false, + )); + } + Ok(namespace) +} + +fn scoped_service_key( + context: &AuthContext, + namespace: u64, + service_name: &str, +) -> std::result::Result { + if service_name.is_empty() || service_name.len() > 1024 || service_name.contains('\0') { + return Err(crate::common::auth::AuthFailure::new( + "service_name_invalid", + "service names must be 1-1024 bytes and must not contain NUL", + false, + )); + } + if !context.is_admin + && (service_name.len() > 128 + || !service_name + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || b"._:-".contains(&byte))) + { + return Err(crate::common::auth::AuthFailure::new( + "service_name_invalid", + "temporary-key service names must be 1-128 ASCII bytes from [A-Za-z0-9._:-]", + false, + )); + } + if namespace == 0 { + Ok(Arc::from(service_name)) + } else { + Ok(Arc::from(format!("@{namespace:016x}\u{0}{service_name}"))) + } +} + +pub(super) fn split_scoped_service_key(key: &str) -> (u64, &str) { + let Some((prefix, name)) = key.split_once('\0') else { + return (0, key); + }; + let Some(hex) = prefix.strip_prefix('@') else { + return (0, key); + }; + match u64::from_str_radix(hex, 16) { + Ok(namespace) => (namespace, name), + Err(_) => (0, key), + } +} + +pub(super) fn decrement_namespace_stream_count( + namespace_stream_counts: &mut hashbrown::HashMap, + namespace: u64, +) { + let Some(count) = namespace_stream_counts.get_mut(&namespace) else { + return; + }; + *count = count.saturating_sub(1); + if *count == 0 { + namespace_stream_counts.remove(&namespace); + } +} + +pub(super) fn release_namespace_rate_limit_if_idle( + namespace: u64, + server_conn_map: &ServerConnMap, + pending_streams: &hashbrown::HashMap, + namespace_rate_limits: &mut hashbrown::HashMap, +) { + let has_registered_service = server_conn_map + .keys() + .any(|key| split_scoped_service_key(key).0 == namespace); + let has_pending_stream = pending_streams + .values() + .any(|(_, _, key)| split_scoped_service_key(key).0 == namespace); + if !has_registered_service && !has_pending_stream { + namespace_rate_limits.remove(&namespace); + } +} + +pub(super) fn remove_pending_streams_for_server( + pending_streams: &mut hashbrown::HashMap, + namespace_stream_counts: &mut hashbrown::HashMap, + server_id_to_remove: RemoteConnId, +) -> usize { + let mut removed = 0; + pending_streams.retain(|_, (server_id, _, key)| { + if *server_id != server_id_to_remove { + return true; + } + decrement_namespace_stream_count(namespace_stream_counts, split_scoped_service_key(key).0); + removed += 1; + false + }); + removed +} diff --git a/src/pb_server/mod.rs b/src/pb_server/mod.rs index 0abbe23..37238f4 100644 --- a/src/pb_server/mod.rs +++ b/src/pb_server/mod.rs @@ -32,7 +32,7 @@ use crate::common::message::command::{ MessageSerializer, PbConnRequest, PbConnResponse, PbConnStatusReq, PbConnStatusResp, PbServiceConnStatus, }; -use crate::common::message::secure::{ServerHeaderSession, ServerSecurity}; +use crate::common::message::secure::{HeaderProtocol, ServerHeaderSession, ServerSecurity}; use crate::common::message::{get_header_msg_reader, MessageReader, MessageWriter}; use crate::pb_server::error::{ ServerListenSnafu, TaskCenterClientSendStreamSnafu, TaskCenterSendRegisterRespSnafu, @@ -362,1407 +362,16 @@ async fn send_subcribe_retry( } } -struct RemoteIdProvider { - next_id: RemoteConnId, -} - -impl RemoteIdProvider { - fn new() -> Self { - Self { - next_id: RemoteConnId::default(), - } - } -} - -impl ConnIdProvider for RemoteIdProvider { - fn get_next_id(&mut self) -> RemoteConnId { - let ret = self.next_id; - self.next_id += 1; - ret - } - - fn is_valid_id(&self, id: &RemoteConnId) -> bool { - id < &self.next_id - } -} -type ServerMananger = TaskManager; - -/// Run a server that takes its keep-alive setting from the environment. -/// -/// Callers that own the setting — the binary, and the UI, which has a toggle for -/// it — should use [`run_server_with_shutdown`] and pass it explicitly. -pub async fn run_server(addr: A) -> std::io::Result<()> { - run_server_with_shutdown(addr, CancellationToken::new(), None, keep_alive_from_env()).await -} - -pub async fn run_server_with_shutdown( - addr: A, - shutdown_token: CancellationToken, - status_channel: Option< - tokio::sync::mpsc::UnboundedReceiver>, - >, - keep_alive: bool, -) -> std::io::Result<()> { - run_server_with_auth_config( - addr, - shutdown_token, - status_channel, - keep_alive, - AuthConfig::default(), - ) - .await -} - -pub async fn run_server_with_auth_config( - addr: A, - shutdown_token: CancellationToken, - status_channel: Option< - tokio::sync::mpsc::UnboundedReceiver>, - >, - keep_alive: bool, - auth_config: AuthConfig, -) -> std::io::Result<()> { - let auth = AuthRuntime::from_process(auth_config) - .await - .map_err(|error| std::io::Error::other(error.to_string()))?; - let security = ServerSecurity::new(auth); - let mut manager = ServerMananger::new(RemoteIdProvider::new()); - // represent the mapping of the `key` to the id of the server-side conn - let mut server_conn_map = ServerConnMap::new(); - let mut pending_streams = - hashbrown::HashMap::::new(); - let mut namespace_stream_counts = hashbrown::HashMap::::new(); - let mut namespace_rate_limits = hashbrown::HashMap::::new(); - let max_services_per_namespace = env_limit("PB_MAPPER_MAX_SERVICES_PER_NAMESPACE", 256); - let max_register_connections_per_service = - env_limit("PB_MAPPER_MAX_REGISTER_CONNECTIONS_PER_SERVICE", 16); - let max_streams_per_namespace = env_limit("PB_MAPPER_MAX_STREAMS_PER_NAMESPACE", 1024); - let new_streams_per_second = env_limit("PB_MAPPER_NEW_STREAMS_PER_SECOND", 100); - let new_streams_burst = env_limit("PB_MAPPER_NEW_STREAMS_BURST", 200); - let mut next_server_generation = 1_u64; - - let listener = TcpListener::bind(addr).await?; - let listen_addr = listener.local_addr()?; - tracing::info!( - event = "pb_server_listening", - listen_addr = %listen_addr, - control_timeout = ?control_io_timeout(), - "pb-mapper server is listening" - ); - - let task_sender = manager.get_task_sender(); - let shutdown_token_clone = shutdown_token.clone(); - - let listener_handle = tokio::spawn(async move { - tokio::select! { - result = handle_listener(task_sender, listener, keep_alive) => { - if let Err(e) = result { - tracing::error!("Listener error: {}", e); - } - } - _ = shutdown_token_clone.cancelled() => { - tracing::info!("Listener shutdown requested"); - } - } - }); - - let start_time = std::time::Instant::now(); - - let status_forward_handle = status_channel.map(|mut receiver| { - let status_sender = manager.get_task_sender(); - tokio::spawn(async move { - while let Some(response_sender) = receiver.recv().await { - if status_sender - .send(ManagerTask::StatusQuery { response_sender }) - .await - .is_err() - { - break; - } - } - }) - }); - - let shutdown_handle = { - let shutdown_sender = manager.get_task_sender(); - tokio::spawn(async move { - shutdown_token.cancelled().await; - let _ = shutdown_sender.send(ManagerTask::Shutdown).await; - }) - }; - - loop { - let task = match manager.wait_for_task().await { - Ok(task) => task, - Err(e) => { - tracing::error!("Manager task error: {}", e); - break; - } - }; - - match task { - ManagerTask::AdminServiceList { - key_id, - page, - page_size, - response_sender, - } => { - let page_size = page_size.clamp(1, 1000) as usize; - let start = (page as usize).saturating_mul(page_size); - let mut all = server_conn_map - .iter() - .filter_map(|(key, connections)| { - let (namespace, service_name) = split_scoped_service_key(key); - if key_id.is_some_and(|key_id| key_id != namespace) { - return None; - } - let first = connections.first()?; - Some(AdminServiceInfo { - key_id: namespace, - namespace, - service_name: service_name.to_string(), - transport: if first.is_datagram { "udp" } else { "tcp" }.to_string(), - codec_enabled: first.need_codec, - connection_count: connections.len() as u32, - }) - }) - .collect::>(); - all.sort_by(|left, right| { - left.namespace - .cmp(&right.namespace) - .then_with(|| left.service_name.cmp(&right.service_name)) - }); - let items = all.iter().skip(start).take(page_size).cloned().collect(); - let next_page = - (start.saturating_add(page_size) < all.len()).then_some(page.saturating_add(1)); - let _ = response_sender.send(AdminServicePage { - schema_version: 1, - items, - next_page, - }); - } - ManagerTask::AdminConnectionList { - key_id, - page, - page_size, - response_sender, - } => { - let now = Instant::now(); - let page_size = page_size.clamp(1, 1000) as usize; - let start = (page as usize).saturating_mul(page_size); - let mut all = server_conn_map - .iter() - .flat_map(|(key, connections)| { - let (namespace, service_name) = split_scoped_service_key(key); - connections.iter().filter_map(move |connection| { - if key_id.is_some_and(|key_id| key_id != namespace) { - return None; - } - Some(AdminConnectionInfo { - key_id: namespace, - namespace, - service_name: service_name.to_string(), - conn_id: connection.conn_id.into(), - generation: connection.generation, - protocol_version: connection.protocol_version, - healthy: connection.health == ServerConnHealth::Healthy, - transport: if connection.is_datagram { "udp" } else { "tcp" } - .to_string(), - codec_enabled: connection.need_codec, - last_rx_age_ms: now - .duration_since(connection.last_rx_at) - .as_millis() - as u64, - }) - }) - }) - .collect::>(); - all.sort_by(|left, right| { - left.namespace - .cmp(&right.namespace) - .then_with(|| left.service_name.cmp(&right.service_name)) - .then_with(|| left.conn_id.cmp(&right.conn_id)) - }); - let items = all.iter().skip(start).take(page_size).cloned().collect(); - let next_page = - (start.saturating_add(page_size) < all.len()).then_some(page.saturating_add(1)); - let _ = response_sender.send(AdminConnectionPage { - schema_version: 1, - items, - next_page, - }); - } - ManagerTask::StatusQuery { response_sender } => { - let total_connections = server_conn_map - .values() - .map(|conns| conns.len() as u32) - .sum(); - - let status_info = ServerStatusInfo { - active_connections: total_connections, - registered_services: server_conn_map.len() as u32, - uptime_seconds: start_time.elapsed().as_secs(), - }; - - // Send response back (ignore if receiver dropped) - let _ = response_sender.send(status_info); - tracing::debug!( - event = "status_query_served", - registered_services = server_conn_map.len(), - server_connections = total_connections, - active_connections = manager.active_conn_count(), - idle_connections = manager.idle_conn_count(), - "server status query served" - ); - } - ManagerTask::Status { - conn_sender, - status, - namespace, - conn_id, - } => { - let resp = match status { - PbConnStatusReq::RemoteId => { - let scoped = server_conn_map - .iter() - .filter(|(key, _)| split_scoped_service_key(key).0 == namespace) - .map(|(key, value)| (split_scoped_service_key(key).1, value)) - .collect::>(); - let registered_ids = scoped - .iter() - .flat_map(|(_, connections)| { - connections.iter().map(|connection| connection.conn_id) - }) - .collect::>(); - let client_ids = pending_streams - .iter() - .filter_map(|(client_id, (_, _, key))| { - (split_scoped_service_key(key).0 == namespace).then_some(*client_id) - }) - .collect::>(); - PbConnResponse::Status(PbConnStatusResp::RemoteId { - server_map: format!("{scoped:?}"), - active: format!( - "registered={registered_ids:?}, clients={client_ids:?}" - ), - idle: "namespace scoped; use `pb-mapper admin connection list` for global inspection" - .to_string(), - }) - } - PbConnStatusReq::Keys => PbConnResponse::Status(PbConnStatusResp::Keys( - server_conn_map - .keys() - .filter_map(|key| { - let (key_namespace, service_name) = split_scoped_service_key(key); - (key_namespace == namespace).then(|| service_name.to_string()) - }) - .collect(), - )), - PbConnStatusReq::Service { key } => { - let display_key = key.clone(); - let key: ImutableKey = if namespace == 0 { - key.into() - } else { - Arc::from(format!("@{namespace:016x}\u{0}{key}")) - }; - PbConnResponse::Status(PbConnStatusResp::Service { - key: display_key, - connections: service_status_connections(&server_conn_map, &key), - }) - } - }; - snafu_error_get_or_continue!(conn_sender - .send(ConnTask::StatusResp(resp)) - .await - .map_err(|_| kanal::SendError(())) - .context(TaskCenterSendStatusRespSnafu { conn_id })); - } - ManagerTask::Accept { stream, peer_addr } => { - let conn_id = manager.get_conn_id( - server_conn_map - .iter() - .flat_map(|(_, ids)| ids.iter().map(|v| v.conn_id)), - ); - tracing::info!( - event = "conn_accepted", - conn_id = %conn_id, - peer_addr = %peer_addr, - registered_services = server_conn_map.len(), - server_connections = registered_server_conn_count(&server_conn_map), - active_connections = manager.active_conn_count(), - idle_connections = manager.idle_conn_count(), - "accepted pb connection" - ); - let manager_task_sender = manager.get_task_sender(); - let security = security.clone(); - tokio::spawn(async move { - snafu_error_handle!( - handle_conn(conn_id, peer_addr, manager_task_sender, stream, security) - .await - ); - }); - } - ManagerTask::DeRegisterServerConn { key, conn_id } => { - let removed_from_service_map = - remove_server_conn(&mut server_conn_map, &key, conn_id); - let removed_from_active_map = manager.deregister_conn(conn_id); - let removed_pending_streams = remove_pending_streams_for_server( - &mut pending_streams, - &mut namespace_stream_counts, - conn_id, - ); - release_namespace_rate_limit_if_idle( - split_scoped_service_key(&key).0, - &server_conn_map, - &pending_streams, - &mut namespace_rate_limits, - ); - tracing::info!( - event = "server_conn_deregistered", - key = %key, - conn_id = %conn_id, - removed_from_service_map, - removed_from_active_map, - removed_pending_streams, - registered_services = server_conn_map.len(), - server_connections = registered_server_conn_count(&server_conn_map), - active_connections = manager.active_conn_count(), - idle_connections = manager.idle_conn_count(), - "server connection deregistered" - ); - } - ManagerTask::ServerConnActivity { key, conn_id } => { - let recorded = record_server_conn_activity(&mut server_conn_map, &key, conn_id); - tracing::debug!( - event = "server_conn_lease_renewed", - key = %key, - conn_id = %conn_id, - recorded, - "server control connection activity recorded" - ); - } - ManagerTask::RetireServerConn { - key, - conn_id, - reason, - } => { - let conn_sender = manager.get_conn_sender_chan(&conn_id); - let removed_from_service_map = - remove_server_conn(&mut server_conn_map, &key, conn_id); - let removed_from_active_map = manager.deregister_conn(conn_id); - let removed_pending_streams = remove_pending_streams_for_server( - &mut pending_streams, - &mut namespace_stream_counts, - conn_id, - ); - release_namespace_rate_limit_if_idle( - split_scoped_service_key(&key).0, - &server_conn_map, - &pending_streams, - &mut namespace_rate_limits, - ); - let retire_notified = conn_sender - .as_ref() - .and_then(|sender| { - sender - .try_send(ConnTask::Retire { - reason: reason.clone(), - }) - .ok() - }) - .is_some(); - tracing::warn!( - event = "server_conn_retired", - key = %key, - conn_id = %conn_id, - reason = %reason, - removed_from_service_map, - removed_from_active_map, - removed_pending_streams, - retire_notified, - registered_services = server_conn_map.len(), - server_connections = registered_server_conn_count(&server_conn_map), - active_connections = manager.active_conn_count(), - idle_connections = manager.idle_conn_count(), - "server connection retired" - ); - } - ManagerTask::DeRegisterClientConn { - server_id, - client_id, - } => { - let removed_namespace = pending_streams.remove(&client_id).map(|(_, _, key)| { - let namespace = split_scoped_service_key(&key).0; - decrement_namespace_stream_count(&mut namespace_stream_counts, namespace); - namespace - }); - let removed_server_conn = if let Some(server_id) = server_id { - manager.deregister_conn(server_id) - } else { - false - }; - let removed_client_conn = manager.deregister_conn(client_id); - if let Some(namespace) = removed_namespace { - release_namespace_rate_limit_if_idle( - namespace, - &server_conn_map, - &pending_streams, - &mut namespace_rate_limits, - ); - } - if removed_server_conn || removed_client_conn { - tracing::info!( - event = "client_conn_deregistered", - server_conn_id = ?server_id, - client_conn_id = %client_id, - removed_server_conn, - removed_client_conn, - registered_services = server_conn_map.len(), - server_connections = registered_server_conn_count(&server_conn_map), - active_connections = manager.active_conn_count(), - idle_connections = manager.idle_conn_count(), - "client connection deregistered" - ); - } else { - tracing::debug!( - event = "client_conn_deregister_skipped", - server_conn_id = ?server_id, - client_conn_id = %client_id, - registered_services = server_conn_map.len(), - server_connections = registered_server_conn_count(&server_conn_map), - active_connections = manager.active_conn_count(), - idle_connections = manager.idle_conn_count(), - "client connection was already inactive" - ); - } - } - ManagerTask::Register { - key, - conn_id, - conn_sender, - need_codec, - is_datagram, - protocol_version, - } => { - let namespace = split_scoped_service_key(&key).0; - let existing = server_conn_map.get(&key); - let failure = if existing.is_some_and(|connections| { - connections - .first() - .is_some_and(|connection| connection.is_datagram != is_datagram) - }) { - Some(( - "service_transport_mismatch", - "the service name is already registered with a different transport", - false, - )) - } else if existing.is_some_and(|connections| { - connections.len() >= max_register_connections_per_service - }) { - Some(( - "service_connection_limit_exceeded", - "the service has reached its register connection limit", - true, - )) - } else if existing.is_none() - && server_conn_map - .keys() - .filter(|registered| split_scoped_service_key(registered).0 == namespace) - .count() - >= max_services_per_namespace - { - Some(( - "namespace_service_limit_exceeded", - "the namespace has reached its service name limit", - true, - )) - } else { - None - }; - if let Some((code, reason, retryable)) = failure { - let _ = conn_sender - .send(ConnTask::RegisterFailed { - code: code.to_string(), - reason: reason.to_string(), - retryable, - }) - .await; - continue; - } - let generation = next_server_generation; - next_server_generation = next_server_generation.saturating_add(1).max(1); - let now = Instant::now(); - // sign up server connection - manager.sign_up_conn_sender(conn_id, conn_sender.clone()); - match server_conn_map.entry(key.clone()) { - hashbrown::hash_map::Entry::Occupied(mut o) => { - o.get_mut().push(ServerConnInfo { - conn_id, - generation, - health: ServerConnHealth::Healthy, - need_codec, - is_datagram, - protocol_version, - last_rx_at: now, - }); - } - hashbrown::hash_map::Entry::Vacant(v) => { - v.insert(vec![ServerConnInfo { - conn_id, - generation, - health: ServerConnHealth::Healthy, - need_codec, - is_datagram, - protocol_version, - last_rx_at: now, - }]); - } - } - - // response registered ok - tracing::info!( - event = "server_conn_registered", - key = %key, - conn_id = %conn_id, - generation, - protocol_version, - need_codec, - is_datagram, - service_connections = service_conn_count(&server_conn_map, &key), - registered_services = server_conn_map.len(), - server_connections = registered_server_conn_count(&server_conn_map), - active_connections = manager.active_conn_count(), - idle_connections = manager.idle_conn_count(), - "server connection registered" - ); - snafu_error_get_or_continue!(conn_sender - .send(ConnTask::RegisterResp { - generation, - protocol_version, - lease_ttl_ms: server_lease_timeout().as_millis() as u64, - }) - .await - .map_err(|_| kanal::SendError(())) - .context(TaskCenterSendRegisterRespSnafu { key, conn_id })); - } - ManagerTask::Stream { - key, - stream, - session, - server_id, - client_id, - server_generation, - } => { - let Some((expected_control_conn_id, expected_generation, expected_key)) = - pending_streams.get(&client_id).cloned() - else { - tracing::warn!( - event = "stale_stream_without_pending_client", - server_conn_id = %server_id, - client_conn_id = %client_id, - server_generation, - "dropping stream for client without pending subscribe" - ); - continue; - }; - if key != expected_key { - tracing::warn!( - event = "stream_namespace_mismatch", - stream_conn_id = %server_id, - client_conn_id = %client_id, - expected_key = %expected_key, - actual_key = %key, - "dropping stream that does not belong to the pending namespace and service" - ); - continue; - } - if server_generation != 0 && expected_generation != server_generation { - tracing::warn!( - event = "stale_stream_generation_mismatch", - stream_conn_id = %server_id, - client_conn_id = %client_id, - expected_control_conn_id = %expected_control_conn_id, - expected_generation, - server_generation, - "dropping stale stream for a previous subscribe attempt" - ); - continue; - } - if let Some(info) = server_conn_map - .values_mut() - .flat_map(|infos| infos.iter_mut()) - .find(|info| { - info.conn_id == expected_control_conn_id - && info.generation == expected_generation - }) - { - info.health = ServerConnHealth::Healthy; - } - tracing::debug!( - event = "stream_ready_for_client", - stream_conn_id = %server_id, - control_conn_id = %expected_control_conn_id, - client_conn_id = %client_id, - server_generation = expected_generation, - active_connections = manager.active_conn_count(), - "server stream ready for client" - ); - let client_sender = snafu_error_get_or_continue!(manager - .get_conn_sender_chan(&client_id) - .context(TaskCenterStreamConnIdNotExistSnafu { conn_id: client_id })); - snafu_error_handle!(client_sender - .send(ConnTask::StreamResp { - server_id, - server_generation: expected_generation, - stream, - session, - }) - .await - .map_err(|_| kanal::SendError(())) - .context(TaskCenterSendStreamRespToClientSnafu { conn_id: client_id })); - } - ManagerTask::StreamAck { - server_id, - client_id, - server_generation, - } => { - let recorded_activity = - record_server_conn_activity_by_conn_id(&mut server_conn_map, server_id); - let Some((expected_server_id, expected_generation, _)) = - pending_streams.get(&client_id).cloned() - else { - tracing::warn!( - event = "stale_stream_ack_without_pending_client", - server_conn_id = %server_id, - client_conn_id = %client_id, - server_generation, - recorded_activity, - "dropping stream ack for client without pending subscribe" - ); - continue; - }; - if expected_server_id != server_id || expected_generation != server_generation { - tracing::warn!( - event = "stale_stream_ack_generation_mismatch", - server_conn_id = %server_id, - client_conn_id = %client_id, - expected_server_conn_id = %expected_server_id, - expected_generation, - server_generation, - recorded_activity, - "dropping stale stream ack for a previous subscribe attempt" - ); - continue; - } - if let Some(info) = server_conn_map - .values_mut() - .flat_map(|infos| infos.iter_mut()) - .find(|info| info.conn_id == server_id && info.generation == server_generation) - { - info.health = ServerConnHealth::Healthy; - } - let client_sender = snafu_error_get_or_continue!(manager - .get_conn_sender_chan(&client_id) - .context(TaskCenterStreamConnIdNotExistSnafu { conn_id: client_id })); - snafu_error_handle!(client_sender - .send(ConnTask::StreamAck { - server_id, - server_generation, - }) - .await - .map_err(|_| kanal::SendError(())) - .context(TaskCenterSendStreamRespToClientSnafu { conn_id: client_id })); - } - ManagerTask::Subcribe { - key, - conn_id, - conn_sender, - excluded_server_conns, - } => { - let namespace = split_scoped_service_key(&key).0; - if namespace_stream_counts - .get(&namespace) - .copied() - .unwrap_or_default() - >= max_streams_per_namespace - { - let _ = conn_sender - .send(ConnTask::SubcribeFailed { - code: "namespace_stream_limit_exceeded".to_string(), - reason: "the namespace has reached its active stream limit".to_string(), - retryable: true, - }) - .await; - continue; - } - let Some(server_conn_id_list) = server_conn_map.get(&key).cloned() else { - let reason = format!("server key `{key}` is not registered"); - tracing::warn!( - event = "subscribe_key_missing", - key = %key, - client_conn_id = %conn_id, - excluded_server_conns = ?excluded_server_conns, - registered_services = server_conn_map.len(), - server_connections = registered_server_conn_count(&server_conn_map), - "subscribe key is not registered" - ); - if excluded_server_conns.is_empty() { - send_subcribe_failed(&conn_sender, &key, conn_id, reason).await; - } else { - send_subcribe_retry(&conn_sender, &key, conn_id, reason).await; - } - continue; - }; - if !namespace_rate_limits - .entry(namespace) - .or_insert_with(|| { - NamespaceRateLimit::new(new_streams_per_second, new_streams_burst) - }) - .allow() - { - let _ = conn_sender - .send(ConnTask::SubcribeFailed { - code: "namespace_stream_rate_exceeded".to_string(), - reason: "the namespace new-stream rate limit was exceeded".to_string(), - retryable: true, - }) - .await; - continue; - } - let mut selected = false; - let mut candidates = Vec::new(); - candidates.extend(server_conn_id_list.iter().rev().copied().filter(|info| { - info.health == ServerConnHealth::Healthy - && !excluded_server_conns.contains(&(info.conn_id, info.generation)) - })); - for server_info in candidates { - let ServerConnInfo { - conn_id: server_conn_id, - generation: server_generation, - health, - need_codec, - is_datagram, - protocol_version: _, - last_rx_at: _, - } = server_info; - let Some(server_conn_sender) = manager.get_conn_sender_chan(&server_conn_id) - else { - tracing::warn!( - event = "subscribe_stale_server_conn", - key = %key, - client_conn_id = %conn_id, - server_conn_id = %server_conn_id, - reason = "sender_not_found", - "subscribe skipped stale server connection" - ); - remove_server_conn(&mut server_conn_map, &key, server_conn_id); - let _ = manager.deregister_conn(server_conn_id); - continue; - }; - // 1. Send a request to get server stream - if let Err(e) = server_conn_sender - .send(ConnTask::StreamReq { - client_id: conn_id, - server_generation, - }) - .await - .map_err(|_| kanal::SendError(())) - .context(TaskCenterClientSendStreamSnafu { - key: key.clone(), - conn_id, - }) - { - let report = snafu::Report::from_error(e); - tracing::error!( - event = "subscribe_stream_request_failed", - key = %key, - client_conn_id = %conn_id, - server_conn_id = %server_conn_id, - error = %report, - "failed to send stream request to registered server" - ); - remove_server_conn(&mut server_conn_map, &key, server_conn_id); - let _ = manager.deregister_conn(server_conn_id); - continue; - } - // sign up client connection after a server accepted the stream request - if manager.get_conn_sender_chan(&conn_id).is_none() { - manager.sign_up_conn_sender(conn_id, conn_sender.clone()); - } - let is_new_stream = pending_streams - .insert(conn_id, (server_conn_id, server_generation, key.clone())) - .is_none(); - if is_new_stream { - *namespace_stream_counts.entry(namespace).or_default() += 1; - } - // 2. Response subcribe ok - if let Err(e) = conn_sender - .send(ConnTask::SubcribeResp { - server_conn_id, - server_generation, - need_codec, - is_datagram, - }) - .await - .map_err(|_| kanal::SendError(())) - .context(TaskCenterSendSubcribeRespSnafu { - key: key.clone(), - conn_id, - }) - { - let report = snafu::Report::from_error(e); - tracing::error!( - event = "subscribe_response_failed", - key = %key, - client_conn_id = %conn_id, - server_conn_id = %server_conn_id, - error = %report, - "failed to send subscribe response to client" - ); - manager.deregister_conn(conn_id); - selected = true; - break; - } - tracing::info!( - event = "subscribe_server_selected", - key = %key, - client_conn_id = %conn_id, - server_conn_id = %server_conn_id, - server_generation, - health = ?health, - need_codec, - is_datagram, - service_connections = service_conn_count(&server_conn_map, &key), - active_connections = manager.active_conn_count(), - "selected server connection for client subscribe" - ); - selected = true; - break; - } - if !selected { - let reason = format!("no usable server connection for key `{key}`"); - tracing::warn!( - event = "subscribe_no_usable_server_conn", - key = %key, - client_conn_id = %conn_id, - excluded_server_conns = ?excluded_server_conns, - registered_services = server_conn_map.len(), - server_connections = registered_server_conn_count(&server_conn_map), - "no usable server connection for subscribe" - ); - if excluded_server_conns.is_empty() { - send_subcribe_failed(&conn_sender, &key, conn_id, reason).await; - } else { - send_subcribe_retry(&conn_sender, &key, conn_id, reason).await; - } - } - } - ManagerTask::Shutdown => { - tracing::info!("Server shutdown requested, stopping main loop"); - break; - } - } - } - - // Gracefully shutdown the listener - listener_handle.abort(); - shutdown_handle.abort(); - if let Some(handle) = status_forward_handle { - handle.abort(); - } - tracing::info!("Server shutdown completed"); - Ok(()) -} - -async fn handle_listener( - task_sender: ManagerTaskSender, - listener: TcpListener, - keep_alive: bool, -) -> Result<()> { - loop { - let (stream, addr) = listener.accept().await.context(ServerListenSnafu)?; - tracing::debug!( - event = "tcp_conn_accepted", - peer_addr = %addr, - "accepted tcp connection" - ); - // set keepalive (optional) and nodelay - if keep_alive { - snafu_error_handle!(set_tcp_keep_alive(&stream).context(TaskCenterSetKeepAliveSnafu)); - } - snafu_error_handle!(set_tcp_nodelay(&stream), "remote stream set nodelay"); - task_sender - .send(ManagerTask::Accept { - stream, - peer_addr: addr, - }) - .await - .map_err(|_| kanal::SendError(())) - .context(TaskCenterSendListenerSnafu)? - } -} - -#[instrument(skip(manager_task_sender, conn, security), fields(conn_id = %conn_id, peer_addr = %peer_addr))] -async fn handle_conn( - conn_id: RemoteConnId, - peer_addr: SocketAddr, - manager_task_sender: ManagerTaskSender, - mut conn: TcpStream, - security: ServerSecurity, -) -> Result<()> { - let timeout = control_io_timeout(); - let initial = match tokio::time::timeout(timeout, security.read_initial(&mut conn)).await { - Err(_) => TaskCenterInitRequestTimeoutSnafu { conn_id, timeout }.fail()?, - Ok(Err(error)) => { - let key_id = error - .response_session - .as_ref() - .map(|session| session.key_id()) - .unwrap_or_default(); - let decision = security.record_failure_log(peer_addr.ip(), key_id, &error.failure.code); - if decision.suppressed > 0 { - tracing::warn!( - event = "auth_failures_suppressed", - peer_ip = %peer_addr.ip(), - key_id, - reason = %error.failure.code, - suppressed = decision.suppressed, - "suppressed repeated authentication failures in the previous window" - ); - } - if decision.emit { - tracing::warn!( - event = "auth_failed", - auth_stage = "initial_frame", - conn_id = %conn_id, - peer_addr = %peer_addr, - key_id, - reason = %error.failure.code, - retryable = error.failure.retryable, - error = %error.failure.message, - "connection authentication failed" - ); - } - if let Some(session) = error.response_session { - write_protocol_error(&mut conn, &session, &error.failure).await; - } - return Ok(()); - } - Ok(Ok(initial)) => initial, - }; - let init_request = match PbConnRequest::decode(&initial.payload) { - Ok(request) => request, - Err(error) => { - tracing::warn!( - event = "auth_failed", - auth_stage = "request_decode", - conn_id = %conn_id, - peer_addr = %peer_addr, - error = %error, - "authenticated request could not be decoded" - ); - write_protocol_error( - &mut conn, - &initial.session, - &crate::common::auth::AuthFailure::new( - "request_decode_failed", - "authenticated request payload is malformed", - false, - ), - ) - .await; - return Ok(()); - } - }; - let mut requested_namespace = None; - let mut force_register_namespace = false; - let init_request = match init_request { - PbConnRequest::RegisterScoped { - need_codec, - is_datagram, - key, - namespace, - force_namespace, - protocol_version, - client_instance_id, - heartbeat_interval_ms, - heartbeat_tolerance_ms, - } => { - requested_namespace = Some(namespace); - force_register_namespace = force_namespace; - PbConnRequest::Register { - need_codec, - is_datagram, - key, - protocol_version, - client_instance_id, - heartbeat_interval_ms, - heartbeat_tolerance_ms, - } - } - PbConnRequest::SubcribeScoped { key, namespace } => { - requested_namespace = Some(namespace); - PbConnRequest::Subcribe { key } - } - PbConnRequest::StatusScoped { status, namespace } => { - requested_namespace = Some(namespace); - PbConnRequest::Status(status) - } - PbConnRequest::StreamScoped { - key, - namespace, - dst_id, - server_generation, - } => { - requested_namespace = Some(namespace); - PbConnRequest::Stream { - key, - dst_id, - server_generation, - } - } - request => request, - }; - let session = initial.session; - let auth_context = match session.context() { - Ok(context) => context.clone(), - Err(error) => { - tracing::warn!(conn_id = %conn_id, peer_addr = %peer_addr, %error, "missing auth context"); - return Ok(()); - } - }; - tracing::info!( - event = "auth_succeeded", - auth_stage = "session", - conn_id = %conn_id, - peer_addr = %peer_addr, - key_id = auth_context.key_id, - namespace = auth_context.namespace, - protocol = ?session.protocol(), - is_admin = auth_context.is_admin, - "connection authentication succeeded" - ); - let effective_namespace = match resolve_namespace( - &auth_context, - requested_namespace, - force_register_namespace, - matches!(&init_request, PbConnRequest::Register { .. }), - ) { - Ok(namespace) => namespace, - Err(failure) => { - write_protocol_error(&mut conn, &session, &failure).await; - return Ok(()); - } - }; - match init_request { - PbConnRequest::Register { - key, - need_codec, - is_datagram, - protocol_version, - client_instance_id, - heartbeat_interval_ms, - heartbeat_tolerance_ms, - } => { - let protocol_version = protocol_version.unwrap_or(1); - tracing::info!( - event = "init_request", - request = "register", - conn_id = %conn_id, - peer_addr = %peer_addr, - key = %key, - protocol_version, - client_instance_id = ?client_instance_id, - heartbeat_interval_ms = ?heartbeat_interval_ms, - heartbeat_tolerance_ms = ?heartbeat_tolerance_ms, - need_codec, - is_datagram, - "received pb init request" - ); - let key = match scoped_service_key(&auth_context, effective_namespace, &key) { - Ok(key) => key, - Err(failure) => { - write_protocol_error(&mut conn, &session, &failure).await; - return Ok(()); - } - }; - let cancellation = match auth_context.cancellation_token() { - Ok(token) => token, - Err(failure) => { - write_protocol_error(&mut conn, &session, &failure).await; - return Ok(()); - } - }; - tokio::select! { - result = handle_server_conn( - ServerRegistration { - key, - need_codec, - is_datagram, - protocol_version, - conn_id, - }, - manager_task_sender, - conn, - session, - ) => result?, - _ = cancellation.cancelled() => { - tracing::info!(event = "connection_auth_expired", key_id = auth_context.key_id, conn_id = %conn_id, "closing registered service connection"); - } - } - } - PbConnRequest::Subcribe { key } => { - tracing::info!( - event = "init_request", - request = "subscribe", - conn_id = %conn_id, - peer_addr = %peer_addr, - key = %key, - "received pb init request" - ); - let key = match scoped_service_key(&auth_context, effective_namespace, &key) { - Ok(key) => key, - Err(failure) => { - write_protocol_error(&mut conn, &session, &failure).await; - return Ok(()); - } - }; - let cancellation = match auth_context.cancellation_token() { - Ok(token) => token, - Err(failure) => { - write_protocol_error(&mut conn, &session, &failure).await; - return Ok(()); - } - }; - tokio::select! { - result = handle_client_conn(key, conn_id, manager_task_sender, conn, session) => result?, - _ = cancellation.cancelled() => { - tracing::info!(event = "connection_auth_expired", key_id = auth_context.key_id, conn_id = %conn_id, "closing subscribed data connection"); - } - } - } - PbConnRequest::Stream { - key, - dst_id, - server_generation, - } => { - tracing::debug!( - event = "init_request", - request = "stream", - conn_id = %conn_id, - peer_addr = %peer_addr, - key = %key, - client_conn_id = dst_id, - server_generation, - "received pb init request" - ); - let key = match scoped_service_key(&auth_context, effective_namespace, &key) { - Ok(key) => key, - Err(failure) => { - write_protocol_error(&mut conn, &session, &failure).await; - return Ok(()); - } - }; - manager_task_sender - .send(ManagerTask::Stream { - key: key.clone(), - stream: conn, - session, - server_id: conn_id, - client_id: dst_id.into(), - server_generation, - }) - .await - .map_err(|_| kanal::SendError(())) - .context(TaskCenterSendStreamRespToManagerSnafu { key, conn_id })?; - } - PbConnRequest::Status(status) => { - tracing::debug!( - event = "init_request", - request = "status", - conn_id = %conn_id, - peer_addr = %peer_addr, - status = ?status, - "received pb init request" - ); - handle_show_status( - status, - effective_namespace, - manager_task_sender, - conn_id, - conn, - session, - ) - .await?; - } - PbConnRequest::Admin(request) => { - if !auth_context.is_admin { - write_protocol_error( - &mut conn, - &session, - &crate::common::auth::AuthFailure::new( - "admin_permission_required", - "administrator credential is required for this operation", - false, - ), - ) - .await; - return Ok(()); - } - handle_admin_request( - request, - security.auth().clone(), - manager_task_sender, - conn_id, - conn, - session, - ) - .await?; - } - PbConnRequest::RegisterScoped { .. } - | PbConnRequest::SubcribeScoped { .. } - | PbConnRequest::StatusScoped { .. } - | PbConnRequest::StreamScoped { .. } => unreachable!("scoped request was normalized"), - } - Ok(()) -} - -async fn write_protocol_error( - conn: &mut TcpStream, - session: &ServerHeaderSession, - failure: &crate::common::auth::AuthFailure, -) { - let response = PbConnResponse::error( - failure.code.clone(), - failure.message.clone(), - failure.retryable, - ); - let Ok(message) = response.encode() else { - return; - }; - let Ok(mut writer) = session.response_writer(conn) else { - return; - }; - if let Err(error) = writer.write_msg(&message).await { - tracing::debug!(%error, reason = %failure.code, "failed to write structured protocol error"); - } -} - -fn resolve_namespace( - context: &AuthContext, - requested: Option, - force_register_namespace: bool, - is_register: bool, -) -> std::result::Result { - let namespace = requested.unwrap_or(context.namespace); - if !context.is_admin && namespace != context.namespace { - return Err(crate::common::auth::AuthFailure::new( - "namespace_access_denied", - "temporary credentials can only access their own namespace", - false, - )); - } - if context.is_admin && is_register && namespace != 0 && !force_register_namespace { - return Err(crate::common::auth::AuthFailure::new( - "namespace_force_required", - "administrator registration in a temporary namespace requires --force", - false, - )); - } - Ok(namespace) -} - -fn scoped_service_key( - context: &AuthContext, - namespace: u64, - service_name: &str, -) -> std::result::Result { - if service_name.is_empty() || service_name.len() > 1024 || service_name.contains('\0') { - return Err(crate::common::auth::AuthFailure::new( - "service_name_invalid", - "service names must be 1-1024 bytes and must not contain NUL", - false, - )); - } - if !context.is_admin - && (service_name.len() > 128 - || !service_name - .bytes() - .all(|byte| byte.is_ascii_alphanumeric() || b"._:-".contains(&byte))) - { - return Err(crate::common::auth::AuthFailure::new( - "service_name_invalid", - "temporary-key service names must be 1-128 ASCII bytes from [A-Za-z0-9._:-]", - false, - )); - } - if namespace == 0 { - Ok(Arc::from(service_name)) - } else { - Ok(Arc::from(format!("@{namespace:016x}\u{0}{service_name}"))) - } -} - -fn split_scoped_service_key(key: &str) -> (u64, &str) { - let Some((prefix, name)) = key.split_once('\0') else { - return (0, key); - }; - let Some(hex) = prefix.strip_prefix('@') else { - return (0, key); - }; - match u64::from_str_radix(hex, 16) { - Ok(namespace) => (namespace, name), - Err(_) => (0, key), - } -} - -fn decrement_namespace_stream_count( - namespace_stream_counts: &mut hashbrown::HashMap, - namespace: u64, -) { - let Some(count) = namespace_stream_counts.get_mut(&namespace) else { - return; - }; - *count = count.saturating_sub(1); - if *count == 0 { - namespace_stream_counts.remove(&namespace); - } -} - -fn release_namespace_rate_limit_if_idle( - namespace: u64, - server_conn_map: &ServerConnMap, - pending_streams: &hashbrown::HashMap, - namespace_rate_limits: &mut hashbrown::HashMap, -) { - let has_registered_service = server_conn_map - .keys() - .any(|key| split_scoped_service_key(key).0 == namespace); - let has_pending_stream = pending_streams - .values() - .any(|(_, _, key)| split_scoped_service_key(key).0 == namespace); - if !has_registered_service && !has_pending_stream { - namespace_rate_limits.remove(&namespace); - } -} - -fn remove_pending_streams_for_server( - pending_streams: &mut hashbrown::HashMap, - namespace_stream_counts: &mut hashbrown::HashMap, - server_id_to_remove: RemoteConnId, -) -> usize { - let mut removed = 0; - pending_streams.retain(|_, (server_id, _, key)| { - if *server_id != server_id_to_remove { - return true; - } - decrement_namespace_stream_count(namespace_stream_counts, split_scoped_service_key(key).0); - removed += 1; - false - }); - removed -} - +mod runtime; +pub use runtime::{ + run_server, run_server_on_listener, run_server_with_auth_config, run_server_with_shutdown, +}; +mod connection; +use connection::{ + decrement_namespace_stream_count, handle_conn, handle_listener, + release_namespace_rate_limit_if_idle, remove_pending_streams_for_server, + split_scoped_service_key, +}; pub async fn get_init_request( conn: &mut TcpStream, conn_id: RemoteConnId, diff --git a/src/pb_server/runtime.rs b/src/pb_server/runtime.rs new file mode 100644 index 0000000..a89a382 --- /dev/null +++ b/src/pb_server/runtime.rs @@ -0,0 +1,925 @@ +use super::*; + +struct RemoteIdProvider { + next_id: RemoteConnId, +} + +impl RemoteIdProvider { + fn new() -> Self { + Self { + next_id: RemoteConnId::default(), + } + } +} + +impl ConnIdProvider for RemoteIdProvider { + fn get_next_id(&mut self) -> RemoteConnId { + let ret = self.next_id; + self.next_id += 1; + ret + } + + fn is_valid_id(&self, id: &RemoteConnId) -> bool { + id < &self.next_id + } +} +type ServerMananger = TaskManager; + +/// Run a server that takes its keep-alive setting from the environment. +/// +/// Callers that own the setting — the binary, and the UI, which has a toggle for +/// it — should use [`run_server_with_shutdown`] and pass it explicitly. +pub async fn run_server(addr: A) -> std::io::Result<()> { + run_server_with_shutdown(addr, CancellationToken::new(), None, keep_alive_from_env()).await +} + +pub async fn run_server_with_shutdown( + addr: A, + shutdown_token: CancellationToken, + status_channel: Option< + tokio::sync::mpsc::UnboundedReceiver>, + >, + keep_alive: bool, +) -> std::io::Result<()> { + run_server_with_auth_config( + addr, + shutdown_token, + status_channel, + keep_alive, + AuthConfig::default(), + ) + .await +} + +pub async fn run_server_with_auth_config( + addr: A, + shutdown_token: CancellationToken, + status_channel: Option< + tokio::sync::mpsc::UnboundedReceiver>, + >, + keep_alive: bool, + auth_config: AuthConfig, +) -> std::io::Result<()> { + let auth = AuthRuntime::from_process(auth_config) + .await + .map_err(|error| std::io::Error::other(error.to_string()))?; + let listener = TcpListener::bind(addr).await?; + run_server_on_listener(listener, shutdown_token, status_channel, keep_alive, auth).await +} + +pub async fn run_server_on_listener( + listener: TcpListener, + shutdown_token: CancellationToken, + status_channel: Option< + tokio::sync::mpsc::UnboundedReceiver>, + >, + keep_alive: bool, + auth: AuthRuntime, +) -> std::io::Result<()> { + let security = ServerSecurity::new(auth); + let mut manager = ServerMananger::new(RemoteIdProvider::new()); + // represent the mapping of the `key` to the id of the server-side conn + let mut server_conn_map = ServerConnMap::new(); + let mut pending_streams = + hashbrown::HashMap::::new(); + let mut namespace_stream_counts = hashbrown::HashMap::::new(); + let mut namespace_rate_limits = hashbrown::HashMap::::new(); + let max_services_per_namespace = env_limit("PB_MAPPER_MAX_SERVICES_PER_NAMESPACE", 256); + let max_register_connections_per_service = + env_limit("PB_MAPPER_MAX_REGISTER_CONNECTIONS_PER_SERVICE", 16); + let max_streams_per_namespace = env_limit("PB_MAPPER_MAX_STREAMS_PER_NAMESPACE", 1024); + let new_streams_per_second = env_limit("PB_MAPPER_NEW_STREAMS_PER_SECOND", 100); + let new_streams_burst = env_limit("PB_MAPPER_NEW_STREAMS_BURST", 200); + let mut next_server_generation = 1_u64; + + let listen_addr = listener.local_addr()?; + tracing::info!( + event = "pb_server_listening", + listen_addr = %listen_addr, + control_timeout = ?control_io_timeout(), + "pb-mapper server is listening" + ); + + let task_sender = manager.get_task_sender(); + let shutdown_token_clone = shutdown_token.clone(); + + let listener_handle = tokio::spawn(async move { + tokio::select! { + result = handle_listener(task_sender, listener, keep_alive) => { + if let Err(e) = result { + tracing::error!("Listener error: {}", e); + } + } + _ = shutdown_token_clone.cancelled() => { + tracing::info!("Listener shutdown requested"); + } + } + }); + + let start_time = std::time::Instant::now(); + + let status_forward_handle = status_channel.map(|mut receiver| { + let status_sender = manager.get_task_sender(); + tokio::spawn(async move { + while let Some(response_sender) = receiver.recv().await { + if status_sender + .send(ManagerTask::StatusQuery { response_sender }) + .await + .is_err() + { + break; + } + } + }) + }); + + let shutdown_handle = { + let shutdown_sender = manager.get_task_sender(); + tokio::spawn(async move { + shutdown_token.cancelled().await; + let _ = shutdown_sender.send(ManagerTask::Shutdown).await; + }) + }; + + loop { + let task = match manager.wait_for_task().await { + Ok(task) => task, + Err(e) => { + tracing::error!("Manager task error: {}", e); + break; + } + }; + + match task { + ManagerTask::AdminServiceList { + key_id, + page, + page_size, + response_sender, + } => { + let page_size = page_size.clamp(1, 1000) as usize; + let start = (page as usize).saturating_mul(page_size); + let mut all = server_conn_map + .iter() + .filter_map(|(key, connections)| { + let (namespace, service_name) = split_scoped_service_key(key); + if key_id.is_some_and(|key_id| key_id != namespace) { + return None; + } + let first = connections.first()?; + Some(AdminServiceInfo { + key_id: namespace, + namespace, + service_name: service_name.to_string(), + transport: if first.is_datagram { "udp" } else { "tcp" }.to_string(), + codec_enabled: first.need_codec, + connection_count: connections.len() as u32, + }) + }) + .collect::>(); + all.sort_by(|left, right| { + left.namespace + .cmp(&right.namespace) + .then_with(|| left.service_name.cmp(&right.service_name)) + }); + let items = all.iter().skip(start).take(page_size).cloned().collect(); + let next_page = + (start.saturating_add(page_size) < all.len()).then_some(page.saturating_add(1)); + let _ = response_sender.send(AdminServicePage { + schema_version: 1, + items, + next_page, + }); + } + ManagerTask::AdminConnectionList { + key_id, + page, + page_size, + response_sender, + } => { + let now = Instant::now(); + let page_size = page_size.clamp(1, 1000) as usize; + let start = (page as usize).saturating_mul(page_size); + let mut all = server_conn_map + .iter() + .flat_map(|(key, connections)| { + let (namespace, service_name) = split_scoped_service_key(key); + connections.iter().filter_map(move |connection| { + if key_id.is_some_and(|key_id| key_id != namespace) { + return None; + } + Some(AdminConnectionInfo { + key_id: namespace, + namespace, + service_name: service_name.to_string(), + conn_id: connection.conn_id.into(), + generation: connection.generation, + protocol_version: connection.protocol_version, + healthy: connection.health == ServerConnHealth::Healthy, + transport: if connection.is_datagram { "udp" } else { "tcp" } + .to_string(), + codec_enabled: connection.need_codec, + last_rx_age_ms: now + .duration_since(connection.last_rx_at) + .as_millis() + as u64, + }) + }) + }) + .collect::>(); + all.sort_by(|left, right| { + left.namespace + .cmp(&right.namespace) + .then_with(|| left.service_name.cmp(&right.service_name)) + .then_with(|| left.conn_id.cmp(&right.conn_id)) + }); + let items = all.iter().skip(start).take(page_size).cloned().collect(); + let next_page = + (start.saturating_add(page_size) < all.len()).then_some(page.saturating_add(1)); + let _ = response_sender.send(AdminConnectionPage { + schema_version: 1, + items, + next_page, + }); + } + ManagerTask::StatusQuery { response_sender } => { + let total_connections = server_conn_map + .values() + .map(|conns| conns.len() as u32) + .sum(); + + let status_info = ServerStatusInfo { + active_connections: total_connections, + registered_services: server_conn_map.len() as u32, + uptime_seconds: start_time.elapsed().as_secs(), + }; + + // Send response back (ignore if receiver dropped) + let _ = response_sender.send(status_info); + tracing::debug!( + event = "status_query_served", + registered_services = server_conn_map.len(), + server_connections = total_connections, + active_connections = manager.active_conn_count(), + idle_connections = manager.idle_conn_count(), + "server status query served" + ); + } + ManagerTask::Status { + conn_sender, + status, + namespace, + conn_id, + } => { + let resp = match status { + PbConnStatusReq::RemoteId => { + let scoped = server_conn_map + .iter() + .filter(|(key, _)| split_scoped_service_key(key).0 == namespace) + .map(|(key, value)| (split_scoped_service_key(key).1, value)) + .collect::>(); + let registered_ids = scoped + .iter() + .flat_map(|(_, connections)| { + connections.iter().map(|connection| connection.conn_id) + }) + .collect::>(); + let client_ids = pending_streams + .iter() + .filter_map(|(client_id, (_, _, key))| { + (split_scoped_service_key(key).0 == namespace).then_some(*client_id) + }) + .collect::>(); + PbConnResponse::Status(PbConnStatusResp::RemoteId { + server_map: format!("{scoped:?}"), + active: format!( + "registered={registered_ids:?}, clients={client_ids:?}" + ), + idle: "namespace scoped; use `pb-mapper admin connection list` for global inspection" + .to_string(), + }) + } + PbConnStatusReq::Keys => PbConnResponse::Status(PbConnStatusResp::Keys( + server_conn_map + .keys() + .filter_map(|key| { + let (key_namespace, service_name) = split_scoped_service_key(key); + (key_namespace == namespace).then(|| service_name.to_string()) + }) + .collect(), + )), + PbConnStatusReq::Service { key } => { + let display_key = key.clone(); + let key: ImutableKey = if namespace == 0 { + key.into() + } else { + Arc::from(format!("@{namespace:016x}\u{0}{key}")) + }; + PbConnResponse::Status(PbConnStatusResp::Service { + key: display_key, + connections: service_status_connections(&server_conn_map, &key), + }) + } + }; + snafu_error_get_or_continue!(conn_sender + .send(ConnTask::StatusResp(resp)) + .await + .map_err(|_| kanal::SendError(())) + .context(TaskCenterSendStatusRespSnafu { conn_id })); + } + ManagerTask::Accept { stream, peer_addr } => { + let conn_id = manager.get_conn_id( + server_conn_map + .iter() + .flat_map(|(_, ids)| ids.iter().map(|v| v.conn_id)), + ); + tracing::info!( + event = "conn_accepted", + conn_id = %conn_id, + peer_addr = %peer_addr, + registered_services = server_conn_map.len(), + server_connections = registered_server_conn_count(&server_conn_map), + active_connections = manager.active_conn_count(), + idle_connections = manager.idle_conn_count(), + "accepted pb connection" + ); + let manager_task_sender = manager.get_task_sender(); + let security = security.clone(); + tokio::spawn(async move { + snafu_error_handle!( + handle_conn(conn_id, peer_addr, manager_task_sender, stream, security) + .await + ); + }); + } + ManagerTask::DeRegisterServerConn { key, conn_id } => { + let removed_from_service_map = + remove_server_conn(&mut server_conn_map, &key, conn_id); + let removed_from_active_map = manager.deregister_conn(conn_id); + let removed_pending_streams = remove_pending_streams_for_server( + &mut pending_streams, + &mut namespace_stream_counts, + conn_id, + ); + release_namespace_rate_limit_if_idle( + split_scoped_service_key(&key).0, + &server_conn_map, + &pending_streams, + &mut namespace_rate_limits, + ); + tracing::info!( + event = "server_conn_deregistered", + key = %key, + conn_id = %conn_id, + removed_from_service_map, + removed_from_active_map, + removed_pending_streams, + registered_services = server_conn_map.len(), + server_connections = registered_server_conn_count(&server_conn_map), + active_connections = manager.active_conn_count(), + idle_connections = manager.idle_conn_count(), + "server connection deregistered" + ); + } + ManagerTask::ServerConnActivity { key, conn_id } => { + let recorded = record_server_conn_activity(&mut server_conn_map, &key, conn_id); + tracing::debug!( + event = "server_conn_lease_renewed", + key = %key, + conn_id = %conn_id, + recorded, + "server control connection activity recorded" + ); + } + ManagerTask::RetireServerConn { + key, + conn_id, + reason, + } => { + let conn_sender = manager.get_conn_sender_chan(&conn_id); + let removed_from_service_map = + remove_server_conn(&mut server_conn_map, &key, conn_id); + let removed_from_active_map = manager.deregister_conn(conn_id); + let removed_pending_streams = remove_pending_streams_for_server( + &mut pending_streams, + &mut namespace_stream_counts, + conn_id, + ); + release_namespace_rate_limit_if_idle( + split_scoped_service_key(&key).0, + &server_conn_map, + &pending_streams, + &mut namespace_rate_limits, + ); + let retire_notified = conn_sender + .as_ref() + .and_then(|sender| { + sender + .try_send(ConnTask::Retire { + reason: reason.clone(), + }) + .ok() + }) + .is_some(); + tracing::warn!( + event = "server_conn_retired", + key = %key, + conn_id = %conn_id, + reason = %reason, + removed_from_service_map, + removed_from_active_map, + removed_pending_streams, + retire_notified, + registered_services = server_conn_map.len(), + server_connections = registered_server_conn_count(&server_conn_map), + active_connections = manager.active_conn_count(), + idle_connections = manager.idle_conn_count(), + "server connection retired" + ); + } + ManagerTask::DeRegisterClientConn { + server_id, + client_id, + } => { + let removed_namespace = pending_streams.remove(&client_id).map(|(_, _, key)| { + let namespace = split_scoped_service_key(&key).0; + decrement_namespace_stream_count(&mut namespace_stream_counts, namespace); + namespace + }); + let removed_server_conn = if let Some(server_id) = server_id { + manager.deregister_conn(server_id) + } else { + false + }; + let removed_client_conn = manager.deregister_conn(client_id); + if let Some(namespace) = removed_namespace { + release_namespace_rate_limit_if_idle( + namespace, + &server_conn_map, + &pending_streams, + &mut namespace_rate_limits, + ); + } + if removed_server_conn || removed_client_conn { + tracing::info!( + event = "client_conn_deregistered", + server_conn_id = ?server_id, + client_conn_id = %client_id, + removed_server_conn, + removed_client_conn, + registered_services = server_conn_map.len(), + server_connections = registered_server_conn_count(&server_conn_map), + active_connections = manager.active_conn_count(), + idle_connections = manager.idle_conn_count(), + "client connection deregistered" + ); + } else { + tracing::debug!( + event = "client_conn_deregister_skipped", + server_conn_id = ?server_id, + client_conn_id = %client_id, + registered_services = server_conn_map.len(), + server_connections = registered_server_conn_count(&server_conn_map), + active_connections = manager.active_conn_count(), + idle_connections = manager.idle_conn_count(), + "client connection was already inactive" + ); + } + } + ManagerTask::Register { + key, + conn_id, + conn_sender, + need_codec, + is_datagram, + protocol_version, + } => { + let namespace = split_scoped_service_key(&key).0; + let existing = server_conn_map.get(&key); + let failure = if existing.is_some_and(|connections| { + connections + .first() + .is_some_and(|connection| connection.is_datagram != is_datagram) + }) { + Some(( + "service_transport_mismatch", + "the service name is already registered with a different transport", + false, + )) + } else if existing.is_some_and(|connections| { + connections.len() >= max_register_connections_per_service + }) { + Some(( + "service_connection_limit_exceeded", + "the service has reached its register connection limit", + true, + )) + } else if existing.is_none() + && server_conn_map + .keys() + .filter(|registered| split_scoped_service_key(registered).0 == namespace) + .count() + >= max_services_per_namespace + { + Some(( + "namespace_service_limit_exceeded", + "the namespace has reached its service name limit", + true, + )) + } else { + None + }; + if let Some((code, reason, retryable)) = failure { + let _ = conn_sender + .send(ConnTask::RegisterFailed { + code: code.to_string(), + reason: reason.to_string(), + retryable, + }) + .await; + continue; + } + let generation = next_server_generation; + next_server_generation = next_server_generation.saturating_add(1).max(1); + let now = Instant::now(); + // sign up server connection + manager.sign_up_conn_sender(conn_id, conn_sender.clone()); + match server_conn_map.entry(key.clone()) { + hashbrown::hash_map::Entry::Occupied(mut o) => { + o.get_mut().push(ServerConnInfo { + conn_id, + generation, + health: ServerConnHealth::Healthy, + need_codec, + is_datagram, + protocol_version, + last_rx_at: now, + }); + } + hashbrown::hash_map::Entry::Vacant(v) => { + v.insert(vec![ServerConnInfo { + conn_id, + generation, + health: ServerConnHealth::Healthy, + need_codec, + is_datagram, + protocol_version, + last_rx_at: now, + }]); + } + } + + // response registered ok + tracing::info!( + event = "server_conn_registered", + key = %key, + conn_id = %conn_id, + generation, + protocol_version, + need_codec, + is_datagram, + service_connections = service_conn_count(&server_conn_map, &key), + registered_services = server_conn_map.len(), + server_connections = registered_server_conn_count(&server_conn_map), + active_connections = manager.active_conn_count(), + idle_connections = manager.idle_conn_count(), + "server connection registered" + ); + snafu_error_get_or_continue!(conn_sender + .send(ConnTask::RegisterResp { + generation, + protocol_version, + lease_ttl_ms: server_lease_timeout().as_millis() as u64, + }) + .await + .map_err(|_| kanal::SendError(())) + .context(TaskCenterSendRegisterRespSnafu { key, conn_id })); + } + ManagerTask::Stream { + key, + stream, + session, + server_id, + client_id, + server_generation, + } => { + let Some((expected_control_conn_id, expected_generation, expected_key)) = + pending_streams.get(&client_id).cloned() + else { + tracing::warn!( + event = "stale_stream_without_pending_client", + server_conn_id = %server_id, + client_conn_id = %client_id, + server_generation, + "dropping stream for client without pending subscribe" + ); + continue; + }; + if key != expected_key { + tracing::warn!( + event = "stream_namespace_mismatch", + stream_conn_id = %server_id, + client_conn_id = %client_id, + expected_key = %expected_key, + actual_key = %key, + "dropping stream that does not belong to the pending namespace and service" + ); + continue; + } + if server_generation != 0 && expected_generation != server_generation { + tracing::warn!( + event = "stale_stream_generation_mismatch", + stream_conn_id = %server_id, + client_conn_id = %client_id, + expected_control_conn_id = %expected_control_conn_id, + expected_generation, + server_generation, + "dropping stale stream for a previous subscribe attempt" + ); + continue; + } + if let Some(info) = server_conn_map + .values_mut() + .flat_map(|infos| infos.iter_mut()) + .find(|info| { + info.conn_id == expected_control_conn_id + && info.generation == expected_generation + }) + { + info.health = ServerConnHealth::Healthy; + } + tracing::debug!( + event = "stream_ready_for_client", + stream_conn_id = %server_id, + control_conn_id = %expected_control_conn_id, + client_conn_id = %client_id, + server_generation = expected_generation, + active_connections = manager.active_conn_count(), + "server stream ready for client" + ); + let client_sender = snafu_error_get_or_continue!(manager + .get_conn_sender_chan(&client_id) + .context(TaskCenterStreamConnIdNotExistSnafu { conn_id: client_id })); + snafu_error_handle!(client_sender + .send(ConnTask::StreamResp { + server_id, + server_generation: expected_generation, + stream, + session, + }) + .await + .map_err(|_| kanal::SendError(())) + .context(TaskCenterSendStreamRespToClientSnafu { conn_id: client_id })); + } + ManagerTask::StreamAck { + server_id, + client_id, + server_generation, + } => { + let recorded_activity = + record_server_conn_activity_by_conn_id(&mut server_conn_map, server_id); + let Some((expected_server_id, expected_generation, _)) = + pending_streams.get(&client_id).cloned() + else { + tracing::warn!( + event = "stale_stream_ack_without_pending_client", + server_conn_id = %server_id, + client_conn_id = %client_id, + server_generation, + recorded_activity, + "dropping stream ack for client without pending subscribe" + ); + continue; + }; + if expected_server_id != server_id || expected_generation != server_generation { + tracing::warn!( + event = "stale_stream_ack_generation_mismatch", + server_conn_id = %server_id, + client_conn_id = %client_id, + expected_server_conn_id = %expected_server_id, + expected_generation, + server_generation, + recorded_activity, + "dropping stale stream ack for a previous subscribe attempt" + ); + continue; + } + if let Some(info) = server_conn_map + .values_mut() + .flat_map(|infos| infos.iter_mut()) + .find(|info| info.conn_id == server_id && info.generation == server_generation) + { + info.health = ServerConnHealth::Healthy; + } + let client_sender = snafu_error_get_or_continue!(manager + .get_conn_sender_chan(&client_id) + .context(TaskCenterStreamConnIdNotExistSnafu { conn_id: client_id })); + snafu_error_handle!(client_sender + .send(ConnTask::StreamAck { + server_id, + server_generation, + }) + .await + .map_err(|_| kanal::SendError(())) + .context(TaskCenterSendStreamRespToClientSnafu { conn_id: client_id })); + } + ManagerTask::Subcribe { + key, + conn_id, + conn_sender, + excluded_server_conns, + } => { + let namespace = split_scoped_service_key(&key).0; + if namespace_stream_counts + .get(&namespace) + .copied() + .unwrap_or_default() + >= max_streams_per_namespace + { + let _ = conn_sender + .send(ConnTask::SubcribeFailed { + code: "namespace_stream_limit_exceeded".to_string(), + reason: "the namespace has reached its active stream limit".to_string(), + retryable: true, + }) + .await; + continue; + } + let Some(server_conn_id_list) = server_conn_map.get(&key).cloned() else { + let reason = format!("server key `{key}` is not registered"); + tracing::warn!( + event = "subscribe_key_missing", + key = %key, + client_conn_id = %conn_id, + excluded_server_conns = ?excluded_server_conns, + registered_services = server_conn_map.len(), + server_connections = registered_server_conn_count(&server_conn_map), + "subscribe key is not registered" + ); + if excluded_server_conns.is_empty() { + send_subcribe_failed(&conn_sender, &key, conn_id, reason).await; + } else { + send_subcribe_retry(&conn_sender, &key, conn_id, reason).await; + } + continue; + }; + if !namespace_rate_limits + .entry(namespace) + .or_insert_with(|| { + NamespaceRateLimit::new(new_streams_per_second, new_streams_burst) + }) + .allow() + { + let _ = conn_sender + .send(ConnTask::SubcribeFailed { + code: "namespace_stream_rate_exceeded".to_string(), + reason: "the namespace new-stream rate limit was exceeded".to_string(), + retryable: true, + }) + .await; + continue; + } + let mut selected = false; + let mut candidates = Vec::new(); + candidates.extend(server_conn_id_list.iter().rev().copied().filter(|info| { + info.health == ServerConnHealth::Healthy + && !excluded_server_conns.contains(&(info.conn_id, info.generation)) + })); + for server_info in candidates { + let ServerConnInfo { + conn_id: server_conn_id, + generation: server_generation, + health, + need_codec, + is_datagram, + protocol_version: _, + last_rx_at: _, + } = server_info; + let Some(server_conn_sender) = manager.get_conn_sender_chan(&server_conn_id) + else { + tracing::warn!( + event = "subscribe_stale_server_conn", + key = %key, + client_conn_id = %conn_id, + server_conn_id = %server_conn_id, + reason = "sender_not_found", + "subscribe skipped stale server connection" + ); + remove_server_conn(&mut server_conn_map, &key, server_conn_id); + let _ = manager.deregister_conn(server_conn_id); + continue; + }; + // 1. Send a request to get server stream + if let Err(e) = server_conn_sender + .send(ConnTask::StreamReq { + client_id: conn_id, + server_generation, + }) + .await + .map_err(|_| kanal::SendError(())) + .context(TaskCenterClientSendStreamSnafu { + key: key.clone(), + conn_id, + }) + { + let report = snafu::Report::from_error(e); + tracing::error!( + event = "subscribe_stream_request_failed", + key = %key, + client_conn_id = %conn_id, + server_conn_id = %server_conn_id, + error = %report, + "failed to send stream request to registered server" + ); + remove_server_conn(&mut server_conn_map, &key, server_conn_id); + let _ = manager.deregister_conn(server_conn_id); + continue; + } + // sign up client connection after a server accepted the stream request + if manager.get_conn_sender_chan(&conn_id).is_none() { + manager.sign_up_conn_sender(conn_id, conn_sender.clone()); + } + let is_new_stream = pending_streams + .insert(conn_id, (server_conn_id, server_generation, key.clone())) + .is_none(); + if is_new_stream { + *namespace_stream_counts.entry(namespace).or_default() += 1; + } + // 2. Response subcribe ok + if let Err(e) = conn_sender + .send(ConnTask::SubcribeResp { + server_conn_id, + server_generation, + need_codec, + is_datagram, + }) + .await + .map_err(|_| kanal::SendError(())) + .context(TaskCenterSendSubcribeRespSnafu { + key: key.clone(), + conn_id, + }) + { + let report = snafu::Report::from_error(e); + tracing::error!( + event = "subscribe_response_failed", + key = %key, + client_conn_id = %conn_id, + server_conn_id = %server_conn_id, + error = %report, + "failed to send subscribe response to client" + ); + manager.deregister_conn(conn_id); + selected = true; + break; + } + tracing::info!( + event = "subscribe_server_selected", + key = %key, + client_conn_id = %conn_id, + server_conn_id = %server_conn_id, + server_generation, + health = ?health, + need_codec, + is_datagram, + service_connections = service_conn_count(&server_conn_map, &key), + active_connections = manager.active_conn_count(), + "selected server connection for client subscribe" + ); + selected = true; + break; + } + if !selected { + let reason = format!("no usable server connection for key `{key}`"); + tracing::warn!( + event = "subscribe_no_usable_server_conn", + key = %key, + client_conn_id = %conn_id, + excluded_server_conns = ?excluded_server_conns, + registered_services = server_conn_map.len(), + server_connections = registered_server_conn_count(&server_conn_map), + "no usable server connection for subscribe" + ); + if excluded_server_conns.is_empty() { + send_subcribe_failed(&conn_sender, &key, conn_id, reason).await; + } else { + send_subcribe_retry(&conn_sender, &key, conn_id, reason).await; + } + } + } + ManagerTask::Shutdown => { + tracing::info!("Server shutdown requested, stopping main loop"); + break; + } + } + } + + // Gracefully shutdown the listener + listener_handle.abort(); + shutdown_handle.abort(); + if let Some(handle) = status_forward_handle { + handle.abort(); + } + tracing::info!("Server shutdown completed"); + Ok(()) +} diff --git a/tests/regression.rs b/tests/regression.rs index bbe1d1b..4f0158a 100644 --- a/tests/regression.rs +++ b/tests/regression.rs @@ -63,6 +63,105 @@ async fn wait_for_server(server_addr: SocketAddr) -> TcpStream { .expect("server did not start") } +#[tokio::test] +async fn explicit_invalid_msg_header_key_fails_server_startup() { + let state_dir = + std::env::temp_dir().join(format!("pb-mapper-invalid-env-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&state_dir); + let mut command = tokio::process::Command::new(env!("CARGO_BIN_EXE_pb-mapper")); + command + .arg("server") + .arg("--port") + .arg("0") + .arg("--auth-state-dir") + .arg(&state_dir) + .env("MSG_HEADER_KEY", "invalid") + .kill_on_drop(true); + + let output = timeout(Duration::from_secs(3), command.output()) + .await + .expect("invalid explicit key must fail instead of starting the server") + .unwrap(); + assert!(!output.status.success()); + let logs = format!( + "{}{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + assert!(logs.contains("administrator_key_invalid"), "logs: {logs}"); + assert!(!state_dir.join("admin.key").exists()); + let _ = std::fs::remove_dir_all(state_dir); +} + +#[tokio::test] +async fn admin_all_preserves_json_output_mode() { + let probe_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let server_addr = probe_listener.local_addr().unwrap(); + drop(probe_listener); + let config = auth_config(server_addr); + let _ = std::fs::remove_dir_all(&config.state_dir); + write_admin_key_file(&config.state_dir.join("admin.key"), TEST_ADMIN_KEY, true).unwrap(); + let runtime = AuthRuntime::start( + *TEST_ADMIN_KEY.as_bytes().first_chunk::<32>().unwrap(), + config.clone(), + ) + .await + .unwrap(); + let admin = runtime.authenticate(0).unwrap(); + runtime + .issue( + &admin, + Duration::from_secs(120), + Some("json-output".to_string()), + ) + .await + .unwrap(); + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + + let shutdown = CancellationToken::new(); + let server_shutdown = shutdown.clone(); + let server_config = config.clone(); + let server = tokio::spawn(async move { + run_server_with_auth_config(server_addr, server_shutdown, None, false, server_config) + .await + .unwrap(); + }); + drop(wait_for_server(server_addr).await); + + let mut command = tokio::process::Command::new(env!("CARGO_BIN_EXE_pb-mapper")); + command + .arg("admin") + .arg("--server") + .arg(server_addr.to_string()) + .arg("--output") + .arg("json") + .arg("key") + .arg("list") + .arg("--all") + .env("MSG_HEADER_KEY", TEST_ADMIN_KEY) + .env("RUST_LOG", "off") + .kill_on_drop(true); + let output = timeout(Duration::from_secs(3), command.output()) + .await + .expect("admin JSON request timed out") + .unwrap(); + assert!(output.status.success()); + let document: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); + assert_eq!(document["schema_version"], 1); + assert_eq!( + document["data"]["KeyList"]["items"] + .as_array() + .unwrap() + .len(), + 1 + ); + + shutdown.cancel(); + server.await.unwrap(); + let _ = std::fs::remove_dir_all(config.state_dir); +} + async fn read_secure_request( security: &ServerSecurity, stream: &mut TcpStream, @@ -211,12 +310,13 @@ async fn temporary_credentials_are_isolated_denied_admin_and_revoked_live() { ) .await .unwrap(); + let admin = runtime.authenticate(0).unwrap(); let first = runtime - .issue(Duration::from_secs(120), Some("first".to_string())) + .issue(&admin, Duration::from_secs(120), Some("first".to_string())) .await .unwrap(); let second = runtime - .issue(Duration::from_secs(120), Some("second".to_string())) + .issue(&admin, Duration::from_secs(120), Some("second".to_string())) .await .unwrap(); drop(runtime); @@ -325,7 +425,7 @@ async fn temporary_credentials_are_isolated_denied_admin_and_revoked_live() { } #[tokio::test] -async fn revoking_temporary_credential_closes_active_data_stream() { +async fn revoking_subscriber_credential_closes_cross_credential_data_stream() { let probe_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let server_addr = probe_listener.local_addr().unwrap(); drop(probe_listener); @@ -339,8 +439,13 @@ async fn revoking_temporary_credential_closes_active_data_stream() { ) .await .unwrap(); + let admin = runtime.authenticate(0).unwrap(); let issued = runtime - .issue(Duration::from_secs(120), Some("active-stream".to_string())) + .issue( + &admin, + Duration::from_secs(120), + Some("active-stream".to_string()), + ) .await .unwrap(); drop(runtime); @@ -356,17 +461,21 @@ async fn revoking_temporary_credential_closes_active_data_stream() { }); let credential = parse_credential(&issued.credential).unwrap(); + let admin_credential = + Credential::Admin(*TEST_ADMIN_KEY.as_bytes().first_chunk::<32>().unwrap()); let service = "revoked-stream"; let mut control = wait_for_server(server_addr).await; - let control_session = ClientHeaderSession::new_v2(&credential).unwrap(); - let register = PbConnRequest::Register { + let control_session = ClientHeaderSession::new_v2(&admin_credential).unwrap(); + let register = PbConnRequest::RegisterScoped { need_codec: false, is_datagram: false, key: service.to_string(), + namespace: issued.metadata.key_id, + force_namespace: true, protocol_version: Some(2), client_instance_id: Some("active-stream-test".to_string()), - heartbeat_interval_ms: Some(50), - heartbeat_tolerance_ms: Some(150), + heartbeat_interval_ms: Some(5_000), + heartbeat_tolerance_ms: Some(15_000), }; control_session .write_initial(&mut control, ®ister.encode().unwrap()) @@ -424,12 +533,13 @@ async fn revoking_temporary_credential_closes_active_data_stream() { .unwrap(); let mut provider = wait_for_server(server_addr).await; - let provider_session = ClientHeaderSession::new_v2(&credential).unwrap(); + let provider_session = ClientHeaderSession::new_v2(&admin_credential).unwrap(); provider_session .write_initial( &mut provider, - &PbConnRequest::Stream { + &PbConnRequest::StreamScoped { key: service.to_string(), + namespace: issued.metadata.key_id, dst_id: client_id, server_generation, } @@ -469,8 +579,6 @@ async fn revoking_temporary_credential_closes_active_data_stream() { .unwrap(); assert_eq!(&ready, b"ready"); - let admin_credential = - Credential::Admin(*TEST_ADMIN_KEY.as_bytes().first_chunk::<32>().unwrap()); let (_, _, revoked) = send_v2_request( server_addr, &admin_credential, @@ -485,12 +593,15 @@ async fn revoking_temporary_credential_closes_active_data_stream() { )); let mut byte = [0_u8; 1]; - for (name, stream) in [("subscriber", &mut subscriber), ("provider", &mut provider)] { - let read = timeout(Duration::from_secs(1), stream.read(&mut byte)) - .await - .unwrap_or_else(|_| panic!("revoked {name} data stream was not closed")) - .unwrap(); - assert_eq!(read, 0, "revoked {name} data stream remained open"); + let read = timeout(Duration::from_secs(1), subscriber.read(&mut byte)) + .await + .expect("revoked subscriber data stream was not closed") + .unwrap(); + assert_eq!(read, 0, "revoked subscriber data stream remained open"); + + match timeout(Duration::from_millis(200), control_reader.read_msg()).await { + Err(_) | Ok(Ok(_)) => {} + Ok(Err(error)) => panic!("administrator registration was cancelled too: {error}"), } shutdown_token.cancel(); diff --git a/ui/native/pb_mapper_ffi/src/state.rs b/ui/native/pb_mapper_ffi/src/state.rs index 632f487..fc5f26c 100644 --- a/ui/native/pb_mapper_ffi/src/state.rs +++ b/ui/native/pb_mapper_ffi/src/state.rs @@ -12,6 +12,7 @@ use tokio::sync::{Mutex, RwLock}; use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; +use pb_mapper::common::auth::{AuthConfig, AuthRuntime}; use pb_mapper::common::checksum::{parse_credential, set_process_msg_header_key}; use pb_mapper::common::config::{get_pb_mapper_server_async, get_sockaddr_async}; use pb_mapper::common::message::command::{PbConnStatusReq, PbConnStatusResp}; @@ -20,7 +21,7 @@ use pb_mapper::local::client::{run_client_side_cli_with_callback, ClientStatusCa use pb_mapper::local::server::{ run_server_side_cli_with_callback, ServerTunnelOptions, StatusCallback, }; -use pb_mapper::pb_server::{run_server_with_shutdown, ServerStatusInfo}; +use pb_mapper::pb_server::{run_server_on_listener, ServerStatusInfo}; use pb_mapper::utils::addr::each_addr; use uni_stream::stream::got_one_socket_addr; use uni_stream::stream::{ @@ -449,1210 +450,9 @@ pub struct PbMapperState { connecting: Arc>>, } -impl PbMapperState { - async fn reset_status_caches(&self) { - { - let mut cache = self.local_server_status_cache.write().await; - *cache = LocalServerStatus { - is_running: false, - active_connections: 0, - registered_services: 0, - uptime_seconds: 0, - }; - } - { - let mut last_update = self.local_server_status_last_update.write().await; - *last_update = None; - } - self.local_server_status_refreshing - .store(false, Ordering::Release); - - self.service_status_cache.write().await.clear(); - self.client_status_cache.write().await.clear(); - self.service_status_refreshing.write().await.clear(); - self.client_status_refreshing.write().await.clear(); - } - pub fn new(app_directory_path: Option) -> Self { - let config_dir = Self::get_config_dir(&app_directory_path); - tracing::info!("Using config directory: {:?}", config_dir); - - let local_server_status_cache = Arc::new(RwLock::new(LocalServerStatus { - is_running: false, - active_connections: 0, - registered_services: 0, - uptime_seconds: 0, - })); - - let temp_state = Self { - server_handle: None, - server_shutdown_token: None, - server_status_sender: None, - server_start_time: None, - registered_services: Arc::new(RwLock::new(HashMap::new())), - active_connections: Arc::new(RwLock::new(HashMap::new())), - service_handles: HashMap::new(), - client_handles: HashMap::new(), - config: AppConfig::default(), - config_dir: config_dir.clone(), - app_directory_path: app_directory_path.clone(), - local_server_status_cache: local_server_status_cache.clone(), - local_server_status_last_update: Arc::new(RwLock::new(None)), - local_server_status_refreshing: Arc::new(AtomicBool::new(false)), - service_status_cache: Arc::new(RwLock::new(HashMap::new())), - client_status_cache: Arc::new(RwLock::new(HashMap::new())), - service_status_refreshing: Arc::new(RwLock::new(HashSet::new())), - client_status_refreshing: Arc::new(RwLock::new(HashSet::new())), - registering: Arc::new(StdMutex::new(HashSet::new())), - connecting: Arc::new(StdMutex::new(HashSet::new())), - }; - - let config = temp_state.load_config().unwrap_or_else(|e| { - tracing::warn!("Could not load config: {}, using defaults", e); - AppConfig::default() - }); - - tracing::info!( - "Loaded configuration: server_address={}, keep_alive={}, msg_header_key_set={}", - config.server_address, - config.keep_alive_enabled, - !config.msg_header_key.is_empty() - ); - - let state = Self { - server_handle: None, - server_shutdown_token: None, - server_status_sender: None, - server_start_time: None, - registered_services: Arc::new(RwLock::new(HashMap::new())), - active_connections: Arc::new(RwLock::new(HashMap::new())), - service_handles: HashMap::new(), - client_handles: HashMap::new(), - config, - config_dir, - app_directory_path, - local_server_status_cache, - local_server_status_last_update: Arc::new(RwLock::new(None)), - local_server_status_refreshing: Arc::new(AtomicBool::new(false)), - service_status_cache: Arc::new(RwLock::new(HashMap::new())), - client_status_cache: Arc::new(RwLock::new(HashMap::new())), - service_status_refreshing: Arc::new(RwLock::new(HashSet::new())), - client_status_refreshing: Arc::new(RwLock::new(HashSet::new())), - registering: Arc::new(StdMutex::new(HashSet::new())), - connecting: Arc::new(StdMutex::new(HashSet::new())), - }; - if let Err(e) = state.apply_msg_header_key_env() { - tracing::error!("Failed to apply MSG_HEADER_KEY during init: {}", e); - } - state - } - - pub fn set_app_directory_path(&mut self, path: Option) -> Result<(), CtlError> { - self.app_directory_path = path; - self.config_dir = Self::get_config_dir(&self.app_directory_path); - - // Reload config from new location if exists - match self.load_config() { - Ok(config) => self.config = config, - Err(e) => { - tracing::warn!("Failed to reload config after setting app dir: {}", e); - } - } - self.apply_msg_header_key_env()?; - - Ok(()) - } - - fn apply_msg_header_key_env(&self) -> Result<(), CtlError> { - let key = (!self.config.msg_header_key.is_empty()).then_some(&*self.config.msg_header_key); - // The library validates the key's length and shape; a rejection here is - // the stored setting being wrong, not something going wrong. - set_process_msg_header_key(key).map_err(CtlError::invalid_argument) - } - - #[allow(unused_variables)] - fn get_config_dir(app_directory_path: &Option) -> PathBuf { - // An explicit path wins everywhere. Mobile is where it normally comes - // from — Flutter hands it over, because there is no OS config dir to - // discover — but honouring it on desktop too is what lets a test point - // a state at a temporary directory instead of the user's real config. - if let Some(app_dir) = app_directory_path { - let path = PathBuf::from(app_dir).join("pb-mapper-ui"); - tracing::info!("Using caller-provided app directory: {:?}", path); - return path; - } - #[cfg(any(target_os = "android", target_os = "ios"))] - { - tracing::warn!("No app directory provided for mobile platform, using relative path"); - PathBuf::from("pb-mapper-ui") - } - #[cfg(not(any(target_os = "android", target_os = "ios")))] - { - if let Some(config_dir) = dirs::config_dir() { - config_dir.join("pb-mapper-ui") - } else if let Some(home_dir) = dirs::home_dir() { - home_dir.join(".config").join("pb-mapper-ui") - } else { - tracing::warn!("Could not determine home directory, using current directory"); - std::env::current_dir() - .unwrap_or_else(|_| PathBuf::from(".")) - .join("pb-mapper-ui-config") - } - } - } - - fn get_config_file_path(&self) -> PathBuf { - let config_dir = Self::get_config_dir(&self.app_directory_path); - - if let Err(e) = std::fs::create_dir_all(&config_dir) { - tracing::warn!( - "Failed to create config directory {:?}: {}, using current directory", - config_dir, - e - ); - return PathBuf::from("pb_mapper_config.json"); - } - - let config_file = config_dir.join("config.json"); - tracing::info!("Using config file path: {:?}", config_file); - config_file - } - - pub fn load_config(&self) -> Result { - let config_path = self.get_config_file_path(); - if config_path.exists() { - let contents = - fs::read_to_string(config_path).map_err(|e| CtlError::io(e.to_string()))?; - let mut config: AppConfig = - serde_json::from_str(&contents).map_err(|e| CtlError::io(e.to_string()))?; - config.msg_header_key = normalize_msg_header_key(config.msg_header_key)?; - Ok(config) - } else { - Ok(AppConfig::default()) - } - } - - pub fn save_config(&self) -> Result<(), CtlError> { - let config_path = self.get_config_file_path(); - let contents = - serde_json::to_string_pretty(&self.config).map_err(|e| CtlError::io(e.to_string()))?; - fs::write(config_path, contents).map_err(|e| CtlError::io(e.to_string()))?; - Ok(()) - } - - fn get_service_config_path(&self) -> PathBuf { - self.config_dir.join("services.json") - } - - fn get_client_config_path(&self) -> PathBuf { - self.config_dir.join("clients.json") - } - - pub fn load_service_configs(&self) -> ServiceConfigStore { - let path = self.get_service_config_path(); - match fs::read_to_string(&path) { - Ok(content) => serde_json::from_str(&content).unwrap_or_else(|_| ServiceConfigStore { - services: HashMap::new(), - }), - Err(_) => ServiceConfigStore { - services: HashMap::new(), - }, - } - } - - pub fn save_service_configs(&self, store: &ServiceConfigStore) -> Result<(), CtlError> { - let path = self.get_service_config_path(); - if let Some(parent) = path.parent() { - fs::create_dir_all(parent) - .map_err(|e| CtlError::io(format!("Failed to create config dir: {e}")))?; - } - - let content = serde_json::to_string_pretty(store) - .map_err(|e| CtlError::io(format!("Failed to serialize config: {e}")))?; - - fs::write(&path, content) - .map_err(|e| CtlError::io(format!("Failed to write config file: {e}")))?; - Ok(()) - } - - pub fn save_service_config( - &self, - service_key: &str, - local_address: &str, - protocol: &str, - enable_encryption: bool, - enable_keep_alive: bool, - ) -> Result<(), CtlError> { - let mut store = self.load_service_configs(); - let now = SystemTime::now(); - - let config = ServiceConfigData { - service_key: service_key.to_string(), - local_address: local_address.to_string(), - protocol: protocol.to_string(), - enable_encryption, - enable_keep_alive, - created_at: if store.services.contains_key(service_key) { - store.services[service_key].created_at - } else { - now - }, - }; - - store.services.insert(service_key.to_string(), config); - self.save_service_configs(&store) - } - - pub fn delete_service_config(&self, service_key: &str) -> Result<(), CtlError> { - let mut store = self.load_service_configs(); - store.services.remove(service_key); - self.save_service_configs(&store) - } - - pub fn load_client_configs(&self) -> ClientConfigStore { - let path = self.get_client_config_path(); - match fs::read_to_string(&path) { - Ok(content) => serde_json::from_str(&content).unwrap_or_else(|_| ClientConfigStore { - clients: HashMap::new(), - }), - Err(_) => ClientConfigStore { - clients: HashMap::new(), - }, - } - } - - pub fn save_client_configs(&self, store: &ClientConfigStore) -> Result<(), CtlError> { - let path = self.get_client_config_path(); - if let Some(parent) = path.parent() { - fs::create_dir_all(parent) - .map_err(|e| CtlError::io(format!("Failed to create config dir: {e}")))?; - } - - let content = serde_json::to_string_pretty(store) - .map_err(|e| CtlError::io(format!("Failed to serialize client config: {e}")))?; - - fs::write(&path, content) - .map_err(|e| CtlError::io(format!("Failed to write client config file: {e}")))?; - Ok(()) - } - - pub fn save_client_config( - &self, - service_key: &str, - local_address: &str, - protocol: &str, - enable_keep_alive: bool, - ) -> Result<(), CtlError> { - let mut store = self.load_client_configs(); - let now = SystemTime::now(); - - let config = ClientConfigData { - service_key: service_key.to_string(), - local_address: local_address.to_string(), - protocol: protocol.to_string(), - enable_keep_alive, - created_at: if store.clients.contains_key(service_key) { - store.clients[service_key].created_at - } else { - now - }, - }; - - store.clients.insert(service_key.to_string(), config); - self.save_client_configs(&store) - } - - pub fn delete_client_config(&self, service_key: &str) -> Result<(), CtlError> { - let mut store = self.load_client_configs(); - store.clients.remove(service_key); - self.save_client_configs(&store) - } - - pub async fn start_server( - &mut self, - port: u16, - enable_keep_alive: bool, - ) -> Result<(), CtlError> { - if self.server_handle.is_some() { - return Err(CtlError::already_exists("Server is already running")); - } - - let ip_addr = IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)); - let bind_addr = std::net::SocketAddr::new(ip_addr, port); - - // Preflight bind to surface "port already in use" errors before spawning. - let listener = TcpListener::bind(bind_addr).await.map_err(|e| { - CtlError::address_in_use(format!("Failed to bind server on {bind_addr}: {e}")) - })?; - drop(listener); - - tracing::info!("Starting pb-mapper server on {}:{}", ip_addr, port); - - let shutdown_token = CancellationToken::new(); - let shutdown_token_clone = shutdown_token.clone(); - - let (status_sender, status_receiver) = tokio::sync::mpsc::unbounded_channel(); - - let handle = tokio::spawn(async move { - if let Err(e) = run_server_with_shutdown( - (ip_addr, port), - shutdown_token_clone, - Some(status_receiver), - enable_keep_alive, - ) - .await - { - tracing::error!("pb-mapper server stopped with error: {e}"); - } - }); - - self.server_handle = Some(handle); - self.server_shutdown_token = Some(shutdown_token); - self.server_status_sender = Some(status_sender); - self.server_start_time = Some(SystemTime::now()); - - { - let mut cache = self.local_server_status_cache.write().await; - *cache = LocalServerStatus { - is_running: true, - active_connections: 0, - registered_services: 0, - uptime_seconds: 0, - }; - } - { - let mut last_update = self.local_server_status_last_update.write().await; - *last_update = Some(Instant::now()); - } - - tracing::info!("pb-mapper server started successfully"); - Ok(()) - } - - pub async fn stop_server(&mut self) -> Result<(), CtlError> { - if let (Some(handle), Some(shutdown_token)) = - (self.server_handle.take(), self.server_shutdown_token.take()) - { - self.server_status_sender = None; - - shutdown_token.cancel(); - - let shutdown_timeout = tokio::time::Duration::from_secs(5); - - match tokio::time::timeout(shutdown_timeout, handle).await { - Ok(_) => { - tracing::info!("Server shutdown gracefully"); - } - Err(_) => { - tracing::warn!("Server shutdown timed out, may not have closed gracefully"); - } - } - - self.server_start_time = None; - - { - let mut cache = self.local_server_status_cache.write().await; - *cache = LocalServerStatus { - is_running: false, - active_connections: 0, - registered_services: 0, - uptime_seconds: 0, - }; - } - { - let mut last_update = self.local_server_status_last_update.write().await; - *last_update = Some(Instant::now()); - } - - for (_, handle) in self.service_handles.drain() { - handle.abort(); - } - - for (_, handle) in self.client_handles.drain() { - handle.abort(); - } - - self.registered_services.write().await.clear(); - self.active_connections.write().await.clear(); - - tracing::info!("pb-mapper server stopped, all services and connections terminated"); - Ok(()) - } else { - Err(CtlError::not_found("Server is not running")) - } - } - - async fn finish_register(&mut self, commit: RegisterCommit) -> Result<(), CtlError> { - let RegisterCommit { - service_key, - local_address, - protocol, - enable_encryption, - enable_keep_alive, - local_sock_addr, - remote_sock_addr, - } = commit; - - if let Some(previous) = self.service_handles.remove(&service_key) { - tracing::warn!( - "Service '{service_key}' is already registered, replacing existing handle" - ); - // Dropping a `JoinHandle` does not stop the task. Without this the - // replaced tunnel kept running and retrying, with nothing left - // holding a handle able to abort it. - previous.abort(); - } - - tracing::info!( - "Registering service '{}' with protocol {}, local address {}, server address {}", - service_key, - protocol, - local_address, - self.config.server_address - ); - - self.save_service_config( - &service_key, - &local_address, - &protocol, - enable_encryption, - enable_keep_alive, - ) - .map_err(|e| CtlError::io(format!("Failed to save service configuration: {e}")))?; - - let key_clone = service_key.clone(); - let service_key_for_status = service_key.clone(); - - let callback: StatusCallback = Box::new(move |status: &str| { - tracing::info!( - "Service {} status update: {}", - service_key_for_status, - status - ); - }); - - let handle = if protocol.to_uppercase() == "TCP" { - tokio::spawn(async move { - let _ = run_server_side_cli_with_callback::( - local_sock_addr, - remote_sock_addr, - key_clone.into(), - ServerTunnelOptions { - need_codec: enable_encryption, - is_datagram: false, - keep_alive: enable_keep_alive, - namespace: None, - force_namespace: false, - }, - Some(callback), - ) - .await; - }) - } else { - tokio::spawn(async move { - let _ = run_server_side_cli_with_callback::( - local_sock_addr, - remote_sock_addr, - key_clone.into(), - ServerTunnelOptions { - need_codec: enable_encryption, - is_datagram: true, - keep_alive: enable_keep_alive, - namespace: None, - force_namespace: false, - }, - Some(callback), - ) - .await; - }) - }; - - self.service_handles.insert(service_key.clone(), handle); - - { - let mut cache = self.service_status_cache.write().await; - cache.insert( - service_key.clone(), - StatusCacheEntry { - status: "retrying".to_string(), - message: "Connecting to pb-mapper server...".to_string(), - updated_at: Instant::now(), - }, - ); - } - self.schedule_service_status_refresh(&service_key).await; - - let service_info = ServiceInfo { - service_key: service_key.clone(), - protocol, - local_address, - status: "Registering".to_string(), - }; - - self.registered_services - .write() - .await - .insert(service_key.clone(), service_info); - - tracing::info!("Service '{}' registration initiated", service_key); - Ok(()) - } - - pub async fn unregister_service(&mut self, service_key: String) -> Result<(), CtlError> { - if let Some(handle) = self.service_handles.remove(&service_key) { - handle.abort(); - } - - if self - .registered_services - .write() - .await - .remove(&service_key) - .is_some() - { - tracing::info!("Service '{}' unregistered successfully", service_key); - Ok(()) - } else { - Err(CtlError::not_found(format!( - "Service '{service_key}' is not registered" - ))) - } - } - - pub async fn delete_service_config_and_stop( - &mut self, - service_key: String, - ) -> Result<(), CtlError> { - if let Some(handle) = self.service_handles.remove(&service_key) { - handle.abort(); - } - - self.registered_services.write().await.remove(&service_key); - - self.delete_service_config(&service_key) - } - - async fn finish_connect(&mut self, commit: ConnectCommit) -> Result<(), CtlError> { - let ConnectCommit { - service_key, - local_address, - protocol, - enable_keep_alive, - local_sock_addr, - remote_sock_addr, - } = commit; - - if let Some(previous) = self.client_handles.remove(&service_key) { - tracing::warn!( - "Client for service '{service_key}' is already connected, replacing handle" - ); - // As in `finish_register`: dropping the handle leaves the old - // client's retry loop running with nothing able to stop it. - previous.abort(); - } - - let protocol_upper = protocol.to_uppercase(); - - tracing::info!( - "Connecting to service '{}' with protocol {}, local address {}, server address {}", - service_key, - protocol, - local_address, - self.config.server_address - ); - - let key_clone = service_key.clone(); - - let status_callback: ClientStatusCallback = { - let service_key_for_callback = service_key.clone(); - Box::new(move |status: &str| { - tracing::info!("Client {} status: {}", service_key_for_callback, status); - }) - }; - - let handle = if protocol_upper == "TCP" { - tokio::spawn(async move { - run_client_side_cli_with_callback::( - local_sock_addr, - remote_sock_addr, - key_clone.into(), - enable_keep_alive, - Some(status_callback), - ) - .await; - }) - } else { - tokio::spawn(async move { - run_client_side_cli_with_callback::( - local_sock_addr, - remote_sock_addr, - key_clone.into(), - enable_keep_alive, - Some(status_callback), - ) - .await; - }) - }; - - self.client_handles.insert(service_key.clone(), handle); - - { - let mut cache = self.client_status_cache.write().await; - cache.insert( - service_key.clone(), - StatusCacheEntry { - status: "retrying".to_string(), - message: "Connecting to pb-mapper server...".to_string(), - updated_at: Instant::now(), - }, - ); - } - self.schedule_client_status_refresh(&service_key).await; - - let connection_info = ConnectionInfo { - service_key: service_key.clone(), - client_id: format!("client-{service_key}"), - status: "Connected".to_string(), - }; - - self.active_connections - .write() - .await - .insert(service_key.clone(), connection_info); - - // Persist here rather than at the FFI boundary, so a connection made - // from a terminal is remembered exactly like one made from the window. - // `finish_register` has always done this; leaving it out here meant a - // CLI `connect` started a client that never appeared in the list. - if let Err(e) = - self.save_client_config(&service_key, &local_address, &protocol, enable_keep_alive) - { - // The client is up either way, so this is a warning and not a - // failure: losing the config costs the entry after a restart. - tracing::warn!("Failed to save client config for '{service_key}': {e}"); - } - - tracing::info!("Connected to service '{}' successfully", service_key); - Ok(()) - } - - /// Claims a service key for a registration. See [`KeyClaim`]. - fn claim_registering(&self, service_key: &str) -> Result { - claim_key(&self.registering, service_key, "being registered") - } - - /// Claims a service key for a client connection. See [`KeyClaim`]. - fn claim_connecting(&self, service_key: &str) -> Result { - claim_key(&self.connecting, service_key, "being connected") - } - - pub async fn disconnect_service(&mut self, service_key: String) -> Result<(), CtlError> { - // Aborting the task is the part that matters: it is what stops the - // retry loop still dialling in the background. - let aborted = match self.client_handles.remove(&service_key) { - Some(handle) => { - handle.abort(); - true - } - None => false, - }; - - let was_listed = self - .active_connections - .write() - .await - .remove(&service_key) - .is_some(); - - // Reported failure only when there was nothing to stop. It used to key - // off the bookkeeping map alone, so a client whose task had been - // aborted could still be reported as "not connected" — an error for an - // operation that had in fact just done its job. - if aborted || was_listed { - tracing::info!("Disconnected from service '{}'", service_key); - Ok(()) - } else { - Err(CtlError::not_found(format!( - "Service '{service_key}' is not connected" - ))) - } - } - - pub async fn delete_client_config_and_stop( - &mut self, - service_key: String, - ) -> Result<(), CtlError> { - if let Some(handle) = self.client_handles.remove(&service_key) { - handle.abort(); - } - - self.active_connections.write().await.remove(&service_key); - - self.delete_client_config(&service_key) - } - - pub async fn get_config_status(&self) -> AppConfig { - self.config.clone() - } - - pub async fn update_config( - &mut self, - server_address: String, - keep_alive: bool, - msg_header_key: String, - ) -> Result<(), CtlError> { - let msg_header_key = normalize_msg_header_key(msg_header_key)?; - self.config.server_address = server_address; - self.config.keep_alive_enabled = keep_alive; - self.config.msg_header_key = msg_header_key; - self.apply_msg_header_key_env()?; - self.save_config()?; - self.reset_status_caches().await; - Ok(()) - } - - pub async fn get_service_configs(&self) -> Vec { - let store = self.load_service_configs(); - let mut services = Vec::new(); - - let mut sorted_configs: Vec<_> = store.services.values().collect(); - sorted_configs.sort_by_key(|config| config.created_at); - - for config in sorted_configs { - let (status, message) = self.calculate_service_status(&config.service_key).await; - - services.push(ServiceConfigInfo { - service_key: config.service_key.clone(), - local_address: config.local_address.clone(), - protocol: config.protocol.clone(), - enable_encryption: config.enable_encryption, - enable_keep_alive: config.enable_keep_alive, - status, - status_message: message, - created_at_ms: config - .created_at - .duration_since(SystemTime::UNIX_EPOCH) - .unwrap_or_default() - .as_millis() as u64, - updated_at_ms: SystemTime::now() - .duration_since(SystemTime::UNIX_EPOCH) - .unwrap_or_default() - .as_millis() as u64, - }); - } - - services - } - - pub async fn get_service_status(&self, service_key: String) -> ServiceStatusResponse { - let (status, message) = self.calculate_service_status(&service_key).await; - ServiceStatusResponse { - service_key, - status, - message, - } - } - - pub async fn get_client_configs(&self) -> Vec { - let store = self.load_client_configs(); - let mut client_infos = Vec::new(); - - for (service_key, config) in store.clients.iter() { - let (status, status_message) = self.calculate_client_status(service_key).await; - - client_infos.push(ClientConfigInfo { - service_key: config.service_key.clone(), - local_address: config.local_address.clone(), - protocol: config.protocol.clone(), - enable_keep_alive: config.enable_keep_alive, - status, - status_message, - created_at_ms: config - .created_at - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_millis() as u64, - updated_at_ms: config - .created_at - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_millis() as u64, - }); - } - - client_infos.sort_by_key(|info| info.created_at_ms); - client_infos - } - - pub async fn get_client_status(&self, service_key: String) -> ClientStatusResponse { - let (status, message) = self.calculate_client_status(&service_key).await; - ClientStatusResponse { - service_key, - status, - message, - } - } - - pub async fn get_local_server_status(&self) -> LocalServerStatus { - let is_running = self.server_handle.is_some(); - if !is_running { - let status = LocalServerStatus { - is_running: false, - active_connections: 0, - registered_services: 0, - uptime_seconds: 0, - }; - { - let mut cache = self.local_server_status_cache.write().await; - *cache = status.clone(); - } - { - let mut last_update = self.local_server_status_last_update.write().await; - *last_update = Some(Instant::now()); - } - return status; - } - - let should_refresh = { - let last_update = self.local_server_status_last_update.read().await; - cache_is_stale(*last_update, STATUS_CACHE_TTL) - }; - - if should_refresh { - self.schedule_local_server_status_refresh(); - } - - let cache = self.local_server_status_cache.read().await; - cache.clone() - } - - fn schedule_local_server_status_refresh(&self) { - if self - .local_server_status_refreshing - .swap(true, Ordering::AcqRel) - { - return; - } - - let sender = self.server_status_sender.clone(); - let cache = self.local_server_status_cache.clone(); - let last_update = self.local_server_status_last_update.clone(); - let refreshing = self.local_server_status_refreshing.clone(); - let start_time = self.server_start_time; - - tokio::spawn(async move { - let mut status = LocalServerStatus { - is_running: true, - active_connections: 0, - registered_services: 0, - uptime_seconds: start_time - .and_then(|ts| SystemTime::now().duration_since(ts).ok()) - .map(|d| d.as_secs()) - .unwrap_or(0), - }; - - if let Some(sender) = sender { - let (response_sender, response_receiver) = tokio::sync::oneshot::channel(); - if sender.send(response_sender).is_ok() { - if let Ok(Ok(info)) = - tokio::time::timeout(Duration::from_millis(200), response_receiver).await - { - status.active_connections = info.active_connections; - status.registered_services = info.registered_services; - status.uptime_seconds = info.uptime_seconds; - } - } - } - - { - let mut cache = cache.write().await; - *cache = status; - } - { - let mut last_update = last_update.write().await; - *last_update = Some(Instant::now()); - } - refreshing.store(false, Ordering::Release); - }); - } - - pub async fn get_server_status_detail(&self) -> Result { - self.force_refresh_server_status().await - } - - /// The connections the server holds for one key, from the protocol's own - /// structured query rather than the Debug dump in `server_map`. - pub async fn get_service_conns( - &self, - service_key: String, - ) -> Result, CtlError> { - let server_addr = self.config.server_address.clone(); - match tokio::time::timeout( - FORCE_REFRESH_TIMEOUT, - get_service_conns_with_addr(&server_addr, &service_key), - ) - .await - { - Ok(result) => result, - Err(_) => Err(CtlError::timeout(format!( - "Timed out asking {server_addr} about {service_key}" - ))), - } - } - - /// Perform a blocking status refresh — waits for the actual network result - /// instead of returning stale cache. - pub async fn force_refresh_server_status(&self) -> Result { - let server_addr = self.config.server_address.clone(); - - let detail = match tokio::time::timeout( - FORCE_REFRESH_TIMEOUT, - fetch_real_status_with_addr(&server_addr), - ) - .await - { - Ok(Ok((services, remote_id_data))) => ServerStatusDetail { - server_available: true, - registered_services: services, - server_map: remote_id_data.server_map, - active_connections: remote_id_data.active, - idle_connections: remote_id_data.idle, - }, - Ok(Err(e)) => { - tracing::warn!("Force refresh failed: {}", e); - ServerStatusDetail { - server_available: false, - registered_services: Vec::new(), - server_map: String::new(), - active_connections: String::new(), - idle_connections: String::new(), - } - } - Err(_) => { - tracing::warn!("Force refresh timed out after {:?}", FORCE_REFRESH_TIMEOUT); - ServerStatusDetail { - server_available: false, - registered_services: Vec::new(), - server_map: String::new(), - active_connections: String::new(), - idle_connections: String::new(), - } - } - }; - - Ok(detail) - } - - // Cache service status to avoid blocking UI with network checks on every paint. - async fn get_cached_service_status(&self, service_key: &str) -> (String, String) { - if let Some(handle) = self.service_handles.get(service_key) { - if handle.is_finished() { - return ( - "failed".to_string(), - "Service connection terminated".to_string(), - ); - } - - let cached = { - let cache = self.service_status_cache.read().await; - cache.get(service_key).cloned() - }; - - let should_refresh = cached - .as_ref() - .map(|entry| entry.updated_at.elapsed() > STATUS_CACHE_TTL) - .unwrap_or(true); - - if should_refresh { - self.schedule_service_status_refresh(service_key).await; - } - - if let Some(entry) = cached { - return (entry.status, entry.message); - } - - return ( - "retrying".to_string(), - "Checking service status...".to_string(), - ); - } - - ( - "stopped".to_string(), - "Service is not registered".to_string(), - ) - } - - // Cache client status to avoid blocking UI with network checks on every paint. - async fn get_cached_client_status(&self, service_key: &str) -> (String, String) { - if let Some(handle) = self.client_handles.get(service_key) { - if handle.is_finished() { - return ( - "failed".to_string(), - "Client connection terminated".to_string(), - ); - } - - let cached = { - let cache = self.client_status_cache.read().await; - cache.get(service_key).cloned() - }; - - let should_refresh = cached - .as_ref() - .map(|entry| entry.updated_at.elapsed() > STATUS_CACHE_TTL) - .unwrap_or(true); - - if should_refresh { - self.schedule_client_status_refresh(service_key).await; - } - - if let Some(entry) = cached { - return (entry.status, entry.message); - } - - return ( - "retrying".to_string(), - "Checking client status...".to_string(), - ); - } - - ("stopped".to_string(), "Client is not connected".to_string()) - } - - async fn schedule_service_status_refresh(&self, service_key: &str) { - { - let mut refreshing = self.service_status_refreshing.write().await; - if refreshing.contains(service_key) { - return; - } - refreshing.insert(service_key.to_string()); - } - - let server_addr = self.config.server_address.clone(); - let cache = self.service_status_cache.clone(); - let refreshing = self.service_status_refreshing.clone(); - let key = service_key.to_string(); - - tokio::spawn(async move { - let result = tokio::time::timeout( - STATUS_REFRESH_TIMEOUT, - check_service_with_get_status(&server_addr, &key), - ) - .await; - - let (status, message) = match result { - Ok(Ok(true)) => ( - "running".to_string(), - "Service is running normally".to_string(), - ), - Ok(Ok(false)) => ( - "retrying".to_string(), - "Service is in retry connection loop".to_string(), - ), - Ok(Err(_)) | Err(_) => ( - "failed".to_string(), - "Cannot connect to pb-server".to_string(), - ), - }; - - let changed = { - let mut cache = cache.write().await; - let changed = cache - .get(&key) - .is_none_or(|entry| entry.status != status || entry.message != message); - cache.insert( - key.clone(), - StatusCacheEntry { - status, - message, - updated_at: Instant::now(), - }, - ); - changed - }; - // Only transitions the user can perceive. These run on a timer for - // every configured entry, so emitting on every refresh would reload - // the list several times a second for no visible reason. - if changed { - events::emit(events::ChangeKind::Services, Some(&key), Origin::Internal); - } - - let mut refreshing = refreshing.write().await; - refreshing.remove(&key); - }); - } - - async fn schedule_client_status_refresh(&self, service_key: &str) { - { - let mut refreshing = self.client_status_refreshing.write().await; - if refreshing.contains(service_key) { - return; - } - refreshing.insert(service_key.to_string()); - } - - let server_addr = self.config.server_address.clone(); - let cache = self.client_status_cache.clone(); - let refreshing = self.client_status_refreshing.clone(); - let key = service_key.to_string(); - - tokio::spawn(async move { - let result = tokio::time::timeout( - STATUS_REFRESH_TIMEOUT, - check_service_with_get_status(&server_addr, &key), - ) - .await; - - let (status, message) = match result { - Ok(Ok(true)) => ( - "running".to_string(), - "Client is connected normally".to_string(), - ), - Ok(Ok(false)) => ( - "retrying".to_string(), - "Client is in retry connection loop".to_string(), - ), - Ok(Err(_)) | Err(_) => ( - "failed".to_string(), - "Cannot connect to pb-server".to_string(), - ), - }; - - let changed = { - let mut cache = cache.write().await; - let changed = cache - .get(&key) - .is_none_or(|entry| entry.status != status || entry.message != message); - cache.insert( - key.clone(), - StatusCacheEntry { - status, - message, - updated_at: Instant::now(), - }, - ); - changed - }; - // Only transitions the user can perceive. These run on a timer for - // every configured entry, so emitting on every refresh would reload - // the list several times a second for no visible reason. - if changed { - events::emit(events::ChangeKind::Clients, Some(&key), Origin::Internal); - } - - let mut refreshing = refreshing.write().await; - refreshing.remove(&key); - }); - } - - async fn calculate_service_status(&self, service_key: &str) -> (String, String) { - self.get_cached_service_status(service_key).await - } - - async fn calculate_client_status(&self, service_key: &str) -> (String, String) { - self.get_cached_client_status(service_key).await - } -} +mod configuration; +mod runtime; +mod status; /// Registers a service, holding the state lock only for the bookkeeping. /// @@ -1784,6 +584,31 @@ mod tests { (Arc::new(Mutex::new(state)), root) } + #[tokio::test] + async fn ui_server_uses_its_writable_config_directory_and_reports_readiness() { + let (state, root) = temp_state("server-auth-path"); + let auth_dir = { + let mut state = state.lock().await; + let auth_dir = state.config_dir.join("auth"); + state + .start_server(0, false) + .await + .expect("UI server should bind and initialize authentication"); + assert!(state.server_handle.is_some()); + assert!(state.get_local_server_status().await.is_running); + auth_dir + }; + + assert!(auth_dir.join("admin.key").is_file()); + state + .lock() + .await + .stop_server() + .await + .expect("UI server should stop cleanly"); + let _ = std::fs::remove_dir_all(&root); + } + /// The claim is what stands in for the lock that registration no longer /// holds across its slow phase. Without it, two callers — the window and a /// terminal, say — could both finish the preflight and both insert, and the diff --git a/ui/native/pb_mapper_ffi/src/state/configuration.rs b/ui/native/pb_mapper_ffi/src/state/configuration.rs new file mode 100644 index 0000000..2d0e279 --- /dev/null +++ b/ui/native/pb_mapper_ffi/src/state/configuration.rs @@ -0,0 +1,320 @@ +use super::*; + +impl PbMapperState { + pub(super) async fn reset_status_caches(&self) { + { + let mut cache = self.local_server_status_cache.write().await; + *cache = LocalServerStatus { + is_running: false, + active_connections: 0, + registered_services: 0, + uptime_seconds: 0, + }; + } + { + let mut last_update = self.local_server_status_last_update.write().await; + *last_update = None; + } + self.local_server_status_refreshing + .store(false, Ordering::Release); + + self.service_status_cache.write().await.clear(); + self.client_status_cache.write().await.clear(); + self.service_status_refreshing.write().await.clear(); + self.client_status_refreshing.write().await.clear(); + } + pub fn new(app_directory_path: Option) -> Self { + let config_dir = Self::get_config_dir(&app_directory_path); + tracing::info!("Using config directory: {:?}", config_dir); + + let local_server_status_cache = Arc::new(RwLock::new(LocalServerStatus { + is_running: false, + active_connections: 0, + registered_services: 0, + uptime_seconds: 0, + })); + + let temp_state = Self { + server_handle: None, + server_shutdown_token: None, + server_status_sender: None, + server_start_time: None, + registered_services: Arc::new(RwLock::new(HashMap::new())), + active_connections: Arc::new(RwLock::new(HashMap::new())), + service_handles: HashMap::new(), + client_handles: HashMap::new(), + config: AppConfig::default(), + config_dir: config_dir.clone(), + app_directory_path: app_directory_path.clone(), + local_server_status_cache: local_server_status_cache.clone(), + local_server_status_last_update: Arc::new(RwLock::new(None)), + local_server_status_refreshing: Arc::new(AtomicBool::new(false)), + service_status_cache: Arc::new(RwLock::new(HashMap::new())), + client_status_cache: Arc::new(RwLock::new(HashMap::new())), + service_status_refreshing: Arc::new(RwLock::new(HashSet::new())), + client_status_refreshing: Arc::new(RwLock::new(HashSet::new())), + registering: Arc::new(StdMutex::new(HashSet::new())), + connecting: Arc::new(StdMutex::new(HashSet::new())), + }; + + let config = temp_state.load_config().unwrap_or_else(|e| { + tracing::warn!("Could not load config: {}, using defaults", e); + AppConfig::default() + }); + + tracing::info!( + "Loaded configuration: server_address={}, keep_alive={}, msg_header_key_set={}", + config.server_address, + config.keep_alive_enabled, + !config.msg_header_key.is_empty() + ); + + let state = Self { + server_handle: None, + server_shutdown_token: None, + server_status_sender: None, + server_start_time: None, + registered_services: Arc::new(RwLock::new(HashMap::new())), + active_connections: Arc::new(RwLock::new(HashMap::new())), + service_handles: HashMap::new(), + client_handles: HashMap::new(), + config, + config_dir, + app_directory_path, + local_server_status_cache, + local_server_status_last_update: Arc::new(RwLock::new(None)), + local_server_status_refreshing: Arc::new(AtomicBool::new(false)), + service_status_cache: Arc::new(RwLock::new(HashMap::new())), + client_status_cache: Arc::new(RwLock::new(HashMap::new())), + service_status_refreshing: Arc::new(RwLock::new(HashSet::new())), + client_status_refreshing: Arc::new(RwLock::new(HashSet::new())), + registering: Arc::new(StdMutex::new(HashSet::new())), + connecting: Arc::new(StdMutex::new(HashSet::new())), + }; + if let Err(e) = state.apply_msg_header_key_env() { + tracing::error!("Failed to apply MSG_HEADER_KEY during init: {}", e); + } + state + } + + pub fn set_app_directory_path(&mut self, path: Option) -> Result<(), CtlError> { + self.app_directory_path = path; + self.config_dir = Self::get_config_dir(&self.app_directory_path); + + // Reload config from new location if exists + match self.load_config() { + Ok(config) => self.config = config, + Err(e) => { + tracing::warn!("Failed to reload config after setting app dir: {}", e); + } + } + self.apply_msg_header_key_env()?; + + Ok(()) + } + + pub(super) fn apply_msg_header_key_env(&self) -> Result<(), CtlError> { + let key = (!self.config.msg_header_key.is_empty()).then_some(&*self.config.msg_header_key); + // The library validates the key's length and shape; a rejection here is + // the stored setting being wrong, not something going wrong. + set_process_msg_header_key(key).map_err(CtlError::invalid_argument) + } + + #[allow(unused_variables)] + fn get_config_dir(app_directory_path: &Option) -> PathBuf { + // An explicit path wins everywhere. Mobile is where it normally comes + // from — Flutter hands it over, because there is no OS config dir to + // discover — but honouring it on desktop too is what lets a test point + // a state at a temporary directory instead of the user's real config. + if let Some(app_dir) = app_directory_path { + let path = PathBuf::from(app_dir).join("pb-mapper-ui"); + tracing::info!("Using caller-provided app directory: {:?}", path); + return path; + } + #[cfg(any(target_os = "android", target_os = "ios"))] + { + tracing::warn!("No app directory provided for mobile platform, using relative path"); + PathBuf::from("pb-mapper-ui") + } + #[cfg(not(any(target_os = "android", target_os = "ios")))] + { + if let Some(config_dir) = dirs::config_dir() { + config_dir.join("pb-mapper-ui") + } else if let Some(home_dir) = dirs::home_dir() { + home_dir.join(".config").join("pb-mapper-ui") + } else { + tracing::warn!("Could not determine home directory, using current directory"); + std::env::current_dir() + .unwrap_or_else(|_| PathBuf::from(".")) + .join("pb-mapper-ui-config") + } + } + } + + fn get_config_file_path(&self) -> PathBuf { + let config_dir = Self::get_config_dir(&self.app_directory_path); + + if let Err(e) = std::fs::create_dir_all(&config_dir) { + tracing::warn!( + "Failed to create config directory {:?}: {}, using current directory", + config_dir, + e + ); + return PathBuf::from("pb_mapper_config.json"); + } + + let config_file = config_dir.join("config.json"); + tracing::info!("Using config file path: {:?}", config_file); + config_file + } + + pub fn load_config(&self) -> Result { + let config_path = self.get_config_file_path(); + if config_path.exists() { + let contents = + fs::read_to_string(config_path).map_err(|e| CtlError::io(e.to_string()))?; + let mut config: AppConfig = + serde_json::from_str(&contents).map_err(|e| CtlError::io(e.to_string()))?; + config.msg_header_key = normalize_msg_header_key(config.msg_header_key)?; + Ok(config) + } else { + Ok(AppConfig::default()) + } + } + + pub fn save_config(&self) -> Result<(), CtlError> { + let config_path = self.get_config_file_path(); + let contents = + serde_json::to_string_pretty(&self.config).map_err(|e| CtlError::io(e.to_string()))?; + fs::write(config_path, contents).map_err(|e| CtlError::io(e.to_string()))?; + Ok(()) + } + + fn get_service_config_path(&self) -> PathBuf { + self.config_dir.join("services.json") + } + + fn get_client_config_path(&self) -> PathBuf { + self.config_dir.join("clients.json") + } + + pub fn load_service_configs(&self) -> ServiceConfigStore { + let path = self.get_service_config_path(); + match fs::read_to_string(&path) { + Ok(content) => serde_json::from_str(&content).unwrap_or_else(|_| ServiceConfigStore { + services: HashMap::new(), + }), + Err(_) => ServiceConfigStore { + services: HashMap::new(), + }, + } + } + + pub fn save_service_configs(&self, store: &ServiceConfigStore) -> Result<(), CtlError> { + let path = self.get_service_config_path(); + if let Some(parent) = path.parent() { + fs::create_dir_all(parent) + .map_err(|e| CtlError::io(format!("Failed to create config dir: {e}")))?; + } + + let content = serde_json::to_string_pretty(store) + .map_err(|e| CtlError::io(format!("Failed to serialize config: {e}")))?; + + fs::write(&path, content) + .map_err(|e| CtlError::io(format!("Failed to write config file: {e}")))?; + Ok(()) + } + + pub fn save_service_config( + &self, + service_key: &str, + local_address: &str, + protocol: &str, + enable_encryption: bool, + enable_keep_alive: bool, + ) -> Result<(), CtlError> { + let mut store = self.load_service_configs(); + let now = SystemTime::now(); + + let config = ServiceConfigData { + service_key: service_key.to_string(), + local_address: local_address.to_string(), + protocol: protocol.to_string(), + enable_encryption, + enable_keep_alive, + created_at: if store.services.contains_key(service_key) { + store.services[service_key].created_at + } else { + now + }, + }; + + store.services.insert(service_key.to_string(), config); + self.save_service_configs(&store) + } + + pub fn delete_service_config(&self, service_key: &str) -> Result<(), CtlError> { + let mut store = self.load_service_configs(); + store.services.remove(service_key); + self.save_service_configs(&store) + } + + pub fn load_client_configs(&self) -> ClientConfigStore { + let path = self.get_client_config_path(); + match fs::read_to_string(&path) { + Ok(content) => serde_json::from_str(&content).unwrap_or_else(|_| ClientConfigStore { + clients: HashMap::new(), + }), + Err(_) => ClientConfigStore { + clients: HashMap::new(), + }, + } + } + + pub fn save_client_configs(&self, store: &ClientConfigStore) -> Result<(), CtlError> { + let path = self.get_client_config_path(); + if let Some(parent) = path.parent() { + fs::create_dir_all(parent) + .map_err(|e| CtlError::io(format!("Failed to create config dir: {e}")))?; + } + + let content = serde_json::to_string_pretty(store) + .map_err(|e| CtlError::io(format!("Failed to serialize client config: {e}")))?; + + fs::write(&path, content) + .map_err(|e| CtlError::io(format!("Failed to write client config file: {e}")))?; + Ok(()) + } + + pub fn save_client_config( + &self, + service_key: &str, + local_address: &str, + protocol: &str, + enable_keep_alive: bool, + ) -> Result<(), CtlError> { + let mut store = self.load_client_configs(); + let now = SystemTime::now(); + + let config = ClientConfigData { + service_key: service_key.to_string(), + local_address: local_address.to_string(), + protocol: protocol.to_string(), + enable_keep_alive, + created_at: if store.clients.contains_key(service_key) { + store.clients[service_key].created_at + } else { + now + }, + }; + + store.clients.insert(service_key.to_string(), config); + self.save_client_configs(&store) + } + + pub fn delete_client_config(&self, service_key: &str) -> Result<(), CtlError> { + let mut store = self.load_client_configs(); + store.clients.remove(service_key); + self.save_client_configs(&store) + } +} diff --git a/ui/native/pb_mapper_ffi/src/state/runtime.rs b/ui/native/pb_mapper_ffi/src/state/runtime.rs new file mode 100644 index 0000000..44e2444 --- /dev/null +++ b/ui/native/pb_mapper_ffi/src/state/runtime.rs @@ -0,0 +1,436 @@ +use super::*; + +impl PbMapperState { + pub async fn start_server( + &mut self, + port: u16, + enable_keep_alive: bool, + ) -> Result<(), CtlError> { + if self.server_handle.is_some() { + return Err(CtlError::already_exists("Server is already running")); + } + + let ip_addr = IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)); + let bind_addr = std::net::SocketAddr::new(ip_addr, port); + + let listener = TcpListener::bind(bind_addr).await.map_err(|e| { + CtlError::address_in_use(format!("Failed to bind server on {bind_addr}: {e}")) + })?; + let auth_config = AuthConfig { + state_dir: self.config_dir.join("auth"), + ..AuthConfig::default() + }; + let auth = AuthRuntime::from_process(auth_config) + .await + .map_err(|error| { + CtlError::io(format!( + "Failed to initialize relay authentication: {error}" + )) + })?; + + tracing::info!("Starting pb-mapper server on {}:{}", ip_addr, port); + + let shutdown_token = CancellationToken::new(); + let shutdown_token_clone = shutdown_token.clone(); + + let (status_sender, status_receiver) = tokio::sync::mpsc::unbounded_channel(); + + let handle = tokio::spawn(async move { + if let Err(e) = run_server_on_listener( + listener, + shutdown_token_clone, + Some(status_receiver), + enable_keep_alive, + auth, + ) + .await + { + tracing::error!("pb-mapper server stopped with error: {e}"); + } + }); + + self.server_handle = Some(handle); + self.server_shutdown_token = Some(shutdown_token); + self.server_status_sender = Some(status_sender); + self.server_start_time = Some(SystemTime::now()); + + { + let mut cache = self.local_server_status_cache.write().await; + *cache = LocalServerStatus { + is_running: true, + active_connections: 0, + registered_services: 0, + uptime_seconds: 0, + }; + } + { + let mut last_update = self.local_server_status_last_update.write().await; + *last_update = Some(Instant::now()); + } + + tracing::info!("pb-mapper server started successfully"); + Ok(()) + } + + pub async fn stop_server(&mut self) -> Result<(), CtlError> { + if let (Some(handle), Some(shutdown_token)) = + (self.server_handle.take(), self.server_shutdown_token.take()) + { + self.server_status_sender = None; + + shutdown_token.cancel(); + + let shutdown_timeout = tokio::time::Duration::from_secs(5); + + match tokio::time::timeout(shutdown_timeout, handle).await { + Ok(_) => { + tracing::info!("Server shutdown gracefully"); + } + Err(_) => { + tracing::warn!("Server shutdown timed out, may not have closed gracefully"); + } + } + + self.server_start_time = None; + + { + let mut cache = self.local_server_status_cache.write().await; + *cache = LocalServerStatus { + is_running: false, + active_connections: 0, + registered_services: 0, + uptime_seconds: 0, + }; + } + { + let mut last_update = self.local_server_status_last_update.write().await; + *last_update = Some(Instant::now()); + } + + for (_, handle) in self.service_handles.drain() { + handle.abort(); + } + + for (_, handle) in self.client_handles.drain() { + handle.abort(); + } + + self.registered_services.write().await.clear(); + self.active_connections.write().await.clear(); + + tracing::info!("pb-mapper server stopped, all services and connections terminated"); + Ok(()) + } else { + Err(CtlError::not_found("Server is not running")) + } + } + + pub(super) async fn finish_register(&mut self, commit: RegisterCommit) -> Result<(), CtlError> { + let RegisterCommit { + service_key, + local_address, + protocol, + enable_encryption, + enable_keep_alive, + local_sock_addr, + remote_sock_addr, + } = commit; + + if let Some(previous) = self.service_handles.remove(&service_key) { + tracing::warn!( + "Service '{service_key}' is already registered, replacing existing handle" + ); + // Dropping a `JoinHandle` does not stop the task. Without this the + // replaced tunnel kept running and retrying, with nothing left + // holding a handle able to abort it. + previous.abort(); + } + + tracing::info!( + "Registering service '{}' with protocol {}, local address {}, server address {}", + service_key, + protocol, + local_address, + self.config.server_address + ); + + self.save_service_config( + &service_key, + &local_address, + &protocol, + enable_encryption, + enable_keep_alive, + ) + .map_err(|e| CtlError::io(format!("Failed to save service configuration: {e}")))?; + + let key_clone = service_key.clone(); + let service_key_for_status = service_key.clone(); + + let callback: StatusCallback = Box::new(move |status: &str| { + tracing::info!( + "Service {} status update: {}", + service_key_for_status, + status + ); + }); + + let handle = if protocol.to_uppercase() == "TCP" { + tokio::spawn(async move { + let _ = run_server_side_cli_with_callback::( + local_sock_addr, + remote_sock_addr, + key_clone.into(), + ServerTunnelOptions { + need_codec: enable_encryption, + is_datagram: false, + keep_alive: enable_keep_alive, + namespace: None, + force_namespace: false, + }, + Some(callback), + ) + .await; + }) + } else { + tokio::spawn(async move { + let _ = run_server_side_cli_with_callback::( + local_sock_addr, + remote_sock_addr, + key_clone.into(), + ServerTunnelOptions { + need_codec: enable_encryption, + is_datagram: true, + keep_alive: enable_keep_alive, + namespace: None, + force_namespace: false, + }, + Some(callback), + ) + .await; + }) + }; + + self.service_handles.insert(service_key.clone(), handle); + + { + let mut cache = self.service_status_cache.write().await; + cache.insert( + service_key.clone(), + StatusCacheEntry { + status: "retrying".to_string(), + message: "Connecting to pb-mapper server...".to_string(), + updated_at: Instant::now(), + }, + ); + } + self.schedule_service_status_refresh(&service_key).await; + + let service_info = ServiceInfo { + service_key: service_key.clone(), + protocol, + local_address, + status: "Registering".to_string(), + }; + + self.registered_services + .write() + .await + .insert(service_key.clone(), service_info); + + tracing::info!("Service '{}' registration initiated", service_key); + Ok(()) + } + + pub async fn unregister_service(&mut self, service_key: String) -> Result<(), CtlError> { + if let Some(handle) = self.service_handles.remove(&service_key) { + handle.abort(); + } + + if self + .registered_services + .write() + .await + .remove(&service_key) + .is_some() + { + tracing::info!("Service '{}' unregistered successfully", service_key); + Ok(()) + } else { + Err(CtlError::not_found(format!( + "Service '{service_key}' is not registered" + ))) + } + } + + pub async fn delete_service_config_and_stop( + &mut self, + service_key: String, + ) -> Result<(), CtlError> { + if let Some(handle) = self.service_handles.remove(&service_key) { + handle.abort(); + } + + self.registered_services.write().await.remove(&service_key); + + self.delete_service_config(&service_key) + } + + pub(super) async fn finish_connect(&mut self, commit: ConnectCommit) -> Result<(), CtlError> { + let ConnectCommit { + service_key, + local_address, + protocol, + enable_keep_alive, + local_sock_addr, + remote_sock_addr, + } = commit; + + if let Some(previous) = self.client_handles.remove(&service_key) { + tracing::warn!( + "Client for service '{service_key}' is already connected, replacing handle" + ); + // As in `finish_register`: dropping the handle leaves the old + // client's retry loop running with nothing able to stop it. + previous.abort(); + } + + let protocol_upper = protocol.to_uppercase(); + + tracing::info!( + "Connecting to service '{}' with protocol {}, local address {}, server address {}", + service_key, + protocol, + local_address, + self.config.server_address + ); + + let key_clone = service_key.clone(); + + let status_callback: ClientStatusCallback = { + let service_key_for_callback = service_key.clone(); + Box::new(move |status: &str| { + tracing::info!("Client {} status: {}", service_key_for_callback, status); + }) + }; + + let handle = if protocol_upper == "TCP" { + tokio::spawn(async move { + run_client_side_cli_with_callback::( + local_sock_addr, + remote_sock_addr, + key_clone.into(), + enable_keep_alive, + Some(status_callback), + ) + .await; + }) + } else { + tokio::spawn(async move { + run_client_side_cli_with_callback::( + local_sock_addr, + remote_sock_addr, + key_clone.into(), + enable_keep_alive, + Some(status_callback), + ) + .await; + }) + }; + + self.client_handles.insert(service_key.clone(), handle); + + { + let mut cache = self.client_status_cache.write().await; + cache.insert( + service_key.clone(), + StatusCacheEntry { + status: "retrying".to_string(), + message: "Connecting to pb-mapper server...".to_string(), + updated_at: Instant::now(), + }, + ); + } + self.schedule_client_status_refresh(&service_key).await; + + let connection_info = ConnectionInfo { + service_key: service_key.clone(), + client_id: format!("client-{service_key}"), + status: "Connected".to_string(), + }; + + self.active_connections + .write() + .await + .insert(service_key.clone(), connection_info); + + // Persist here rather than at the FFI boundary, so a connection made + // from a terminal is remembered exactly like one made from the window. + // `finish_register` has always done this; leaving it out here meant a + // CLI `connect` started a client that never appeared in the list. + if let Err(e) = + self.save_client_config(&service_key, &local_address, &protocol, enable_keep_alive) + { + // The client is up either way, so this is a warning and not a + // failure: losing the config costs the entry after a restart. + tracing::warn!("Failed to save client config for '{service_key}': {e}"); + } + + tracing::info!("Connected to service '{}' successfully", service_key); + Ok(()) + } + + /// Claims a service key for a registration. See [`KeyClaim`]. + pub(super) fn claim_registering(&self, service_key: &str) -> Result { + claim_key(&self.registering, service_key, "being registered") + } + + /// Claims a service key for a client connection. See [`KeyClaim`]. + pub(super) fn claim_connecting(&self, service_key: &str) -> Result { + claim_key(&self.connecting, service_key, "being connected") + } + + pub async fn disconnect_service(&mut self, service_key: String) -> Result<(), CtlError> { + // Aborting the task is the part that matters: it is what stops the + // retry loop still dialling in the background. + let aborted = match self.client_handles.remove(&service_key) { + Some(handle) => { + handle.abort(); + true + } + None => false, + }; + + let was_listed = self + .active_connections + .write() + .await + .remove(&service_key) + .is_some(); + + // Reported failure only when there was nothing to stop. It used to key + // off the bookkeeping map alone, so a client whose task had been + // aborted could still be reported as "not connected" — an error for an + // operation that had in fact just done its job. + if aborted || was_listed { + tracing::info!("Disconnected from service '{}'", service_key); + Ok(()) + } else { + Err(CtlError::not_found(format!( + "Service '{service_key}' is not connected" + ))) + } + } + + pub async fn delete_client_config_and_stop( + &mut self, + service_key: String, + ) -> Result<(), CtlError> { + if let Some(handle) = self.client_handles.remove(&service_key) { + handle.abort(); + } + + self.active_connections.write().await.remove(&service_key); + + self.delete_client_config(&service_key) + } +} diff --git a/ui/native/pb_mapper_ffi/src/state/status.rs b/ui/native/pb_mapper_ffi/src/state/status.rs new file mode 100644 index 0000000..8ae5f36 --- /dev/null +++ b/ui/native/pb_mapper_ffi/src/state/status.rs @@ -0,0 +1,466 @@ +use super::*; + +impl PbMapperState { + pub async fn get_config_status(&self) -> AppConfig { + self.config.clone() + } + + pub async fn update_config( + &mut self, + server_address: String, + keep_alive: bool, + msg_header_key: String, + ) -> Result<(), CtlError> { + let msg_header_key = normalize_msg_header_key(msg_header_key)?; + self.config.server_address = server_address; + self.config.keep_alive_enabled = keep_alive; + self.config.msg_header_key = msg_header_key; + self.apply_msg_header_key_env()?; + self.save_config()?; + self.reset_status_caches().await; + Ok(()) + } + + pub async fn get_service_configs(&self) -> Vec { + let store = self.load_service_configs(); + let mut services = Vec::new(); + + let mut sorted_configs: Vec<_> = store.services.values().collect(); + sorted_configs.sort_by_key(|config| config.created_at); + + for config in sorted_configs { + let (status, message) = self.calculate_service_status(&config.service_key).await; + + services.push(ServiceConfigInfo { + service_key: config.service_key.clone(), + local_address: config.local_address.clone(), + protocol: config.protocol.clone(), + enable_encryption: config.enable_encryption, + enable_keep_alive: config.enable_keep_alive, + status, + status_message: message, + created_at_ms: config + .created_at + .duration_since(SystemTime::UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64, + updated_at_ms: SystemTime::now() + .duration_since(SystemTime::UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64, + }); + } + + services + } + + pub async fn get_service_status(&self, service_key: String) -> ServiceStatusResponse { + let (status, message) = self.calculate_service_status(&service_key).await; + ServiceStatusResponse { + service_key, + status, + message, + } + } + + pub async fn get_client_configs(&self) -> Vec { + let store = self.load_client_configs(); + let mut client_infos = Vec::new(); + + for (service_key, config) in store.clients.iter() { + let (status, status_message) = self.calculate_client_status(service_key).await; + + client_infos.push(ClientConfigInfo { + service_key: config.service_key.clone(), + local_address: config.local_address.clone(), + protocol: config.protocol.clone(), + enable_keep_alive: config.enable_keep_alive, + status, + status_message, + created_at_ms: config + .created_at + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64, + updated_at_ms: config + .created_at + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64, + }); + } + + client_infos.sort_by_key(|info| info.created_at_ms); + client_infos + } + + pub async fn get_client_status(&self, service_key: String) -> ClientStatusResponse { + let (status, message) = self.calculate_client_status(&service_key).await; + ClientStatusResponse { + service_key, + status, + message, + } + } + + pub async fn get_local_server_status(&self) -> LocalServerStatus { + let is_running = self.server_handle.is_some(); + if !is_running { + let status = LocalServerStatus { + is_running: false, + active_connections: 0, + registered_services: 0, + uptime_seconds: 0, + }; + { + let mut cache = self.local_server_status_cache.write().await; + *cache = status.clone(); + } + { + let mut last_update = self.local_server_status_last_update.write().await; + *last_update = Some(Instant::now()); + } + return status; + } + + let should_refresh = { + let last_update = self.local_server_status_last_update.read().await; + cache_is_stale(*last_update, STATUS_CACHE_TTL) + }; + + if should_refresh { + self.schedule_local_server_status_refresh(); + } + + let cache = self.local_server_status_cache.read().await; + cache.clone() + } + + fn schedule_local_server_status_refresh(&self) { + if self + .local_server_status_refreshing + .swap(true, Ordering::AcqRel) + { + return; + } + + let sender = self.server_status_sender.clone(); + let cache = self.local_server_status_cache.clone(); + let last_update = self.local_server_status_last_update.clone(); + let refreshing = self.local_server_status_refreshing.clone(); + let start_time = self.server_start_time; + + tokio::spawn(async move { + let mut status = LocalServerStatus { + is_running: true, + active_connections: 0, + registered_services: 0, + uptime_seconds: start_time + .and_then(|ts| SystemTime::now().duration_since(ts).ok()) + .map(|d| d.as_secs()) + .unwrap_or(0), + }; + + if let Some(sender) = sender { + let (response_sender, response_receiver) = tokio::sync::oneshot::channel(); + if sender.send(response_sender).is_ok() { + if let Ok(Ok(info)) = + tokio::time::timeout(Duration::from_millis(200), response_receiver).await + { + status.active_connections = info.active_connections; + status.registered_services = info.registered_services; + status.uptime_seconds = info.uptime_seconds; + } + } + } + + { + let mut cache = cache.write().await; + *cache = status; + } + { + let mut last_update = last_update.write().await; + *last_update = Some(Instant::now()); + } + refreshing.store(false, Ordering::Release); + }); + } + + pub async fn get_server_status_detail(&self) -> Result { + self.force_refresh_server_status().await + } + + /// The connections the server holds for one key, from the protocol's own + /// structured query rather than the Debug dump in `server_map`. + pub async fn get_service_conns( + &self, + service_key: String, + ) -> Result, CtlError> { + let server_addr = self.config.server_address.clone(); + match tokio::time::timeout( + FORCE_REFRESH_TIMEOUT, + get_service_conns_with_addr(&server_addr, &service_key), + ) + .await + { + Ok(result) => result, + Err(_) => Err(CtlError::timeout(format!( + "Timed out asking {server_addr} about {service_key}" + ))), + } + } + + /// Perform a blocking status refresh — waits for the actual network result + /// instead of returning stale cache. + pub async fn force_refresh_server_status(&self) -> Result { + let server_addr = self.config.server_address.clone(); + + let detail = match tokio::time::timeout( + FORCE_REFRESH_TIMEOUT, + fetch_real_status_with_addr(&server_addr), + ) + .await + { + Ok(Ok((services, remote_id_data))) => ServerStatusDetail { + server_available: true, + registered_services: services, + server_map: remote_id_data.server_map, + active_connections: remote_id_data.active, + idle_connections: remote_id_data.idle, + }, + Ok(Err(e)) => { + tracing::warn!("Force refresh failed: {}", e); + ServerStatusDetail { + server_available: false, + registered_services: Vec::new(), + server_map: String::new(), + active_connections: String::new(), + idle_connections: String::new(), + } + } + Err(_) => { + tracing::warn!("Force refresh timed out after {:?}", FORCE_REFRESH_TIMEOUT); + ServerStatusDetail { + server_available: false, + registered_services: Vec::new(), + server_map: String::new(), + active_connections: String::new(), + idle_connections: String::new(), + } + } + }; + + Ok(detail) + } + + // Cache service status to avoid blocking UI with network checks on every paint. + async fn get_cached_service_status(&self, service_key: &str) -> (String, String) { + if let Some(handle) = self.service_handles.get(service_key) { + if handle.is_finished() { + return ( + "failed".to_string(), + "Service connection terminated".to_string(), + ); + } + + let cached = { + let cache = self.service_status_cache.read().await; + cache.get(service_key).cloned() + }; + + let should_refresh = cached + .as_ref() + .map(|entry| entry.updated_at.elapsed() > STATUS_CACHE_TTL) + .unwrap_or(true); + + if should_refresh { + self.schedule_service_status_refresh(service_key).await; + } + + if let Some(entry) = cached { + return (entry.status, entry.message); + } + + return ( + "retrying".to_string(), + "Checking service status...".to_string(), + ); + } + + ( + "stopped".to_string(), + "Service is not registered".to_string(), + ) + } + + // Cache client status to avoid blocking UI with network checks on every paint. + async fn get_cached_client_status(&self, service_key: &str) -> (String, String) { + if let Some(handle) = self.client_handles.get(service_key) { + if handle.is_finished() { + return ( + "failed".to_string(), + "Client connection terminated".to_string(), + ); + } + + let cached = { + let cache = self.client_status_cache.read().await; + cache.get(service_key).cloned() + }; + + let should_refresh = cached + .as_ref() + .map(|entry| entry.updated_at.elapsed() > STATUS_CACHE_TTL) + .unwrap_or(true); + + if should_refresh { + self.schedule_client_status_refresh(service_key).await; + } + + if let Some(entry) = cached { + return (entry.status, entry.message); + } + + return ( + "retrying".to_string(), + "Checking client status...".to_string(), + ); + } + + ("stopped".to_string(), "Client is not connected".to_string()) + } + + pub(super) async fn schedule_service_status_refresh(&self, service_key: &str) { + { + let mut refreshing = self.service_status_refreshing.write().await; + if refreshing.contains(service_key) { + return; + } + refreshing.insert(service_key.to_string()); + } + + let server_addr = self.config.server_address.clone(); + let cache = self.service_status_cache.clone(); + let refreshing = self.service_status_refreshing.clone(); + let key = service_key.to_string(); + + tokio::spawn(async move { + let result = tokio::time::timeout( + STATUS_REFRESH_TIMEOUT, + check_service_with_get_status(&server_addr, &key), + ) + .await; + + let (status, message) = match result { + Ok(Ok(true)) => ( + "running".to_string(), + "Service is running normally".to_string(), + ), + Ok(Ok(false)) => ( + "retrying".to_string(), + "Service is in retry connection loop".to_string(), + ), + Ok(Err(_)) | Err(_) => ( + "failed".to_string(), + "Cannot connect to pb-server".to_string(), + ), + }; + + let changed = { + let mut cache = cache.write().await; + let changed = cache + .get(&key) + .is_none_or(|entry| entry.status != status || entry.message != message); + cache.insert( + key.clone(), + StatusCacheEntry { + status, + message, + updated_at: Instant::now(), + }, + ); + changed + }; + // Only transitions the user can perceive. These run on a timer for + // every configured entry, so emitting on every refresh would reload + // the list several times a second for no visible reason. + if changed { + events::emit(events::ChangeKind::Services, Some(&key), Origin::Internal); + } + + let mut refreshing = refreshing.write().await; + refreshing.remove(&key); + }); + } + + pub(super) async fn schedule_client_status_refresh(&self, service_key: &str) { + { + let mut refreshing = self.client_status_refreshing.write().await; + if refreshing.contains(service_key) { + return; + } + refreshing.insert(service_key.to_string()); + } + + let server_addr = self.config.server_address.clone(); + let cache = self.client_status_cache.clone(); + let refreshing = self.client_status_refreshing.clone(); + let key = service_key.to_string(); + + tokio::spawn(async move { + let result = tokio::time::timeout( + STATUS_REFRESH_TIMEOUT, + check_service_with_get_status(&server_addr, &key), + ) + .await; + + let (status, message) = match result { + Ok(Ok(true)) => ( + "running".to_string(), + "Client is connected normally".to_string(), + ), + Ok(Ok(false)) => ( + "retrying".to_string(), + "Client is in retry connection loop".to_string(), + ), + Ok(Err(_)) | Err(_) => ( + "failed".to_string(), + "Cannot connect to pb-server".to_string(), + ), + }; + + let changed = { + let mut cache = cache.write().await; + let changed = cache + .get(&key) + .is_none_or(|entry| entry.status != status || entry.message != message); + cache.insert( + key.clone(), + StatusCacheEntry { + status, + message, + updated_at: Instant::now(), + }, + ); + changed + }; + // Only transitions the user can perceive. These run on a timer for + // every configured entry, so emitting on every refresh would reload + // the list several times a second for no visible reason. + if changed { + events::emit(events::ChangeKind::Clients, Some(&key), Origin::Internal); + } + + let mut refreshing = refreshing.write().await; + refreshing.remove(&key); + }); + } + + async fn calculate_service_status(&self, service_key: &str) -> (String, String) { + self.get_cached_service_status(service_key).await + } + + async fn calculate_client_status(&self, service_key: &str) -> (String, String) { + self.get_cached_client_status(service_key).await + } +} From 781637fddceea9ab9c9db0c0de8004b58e64d3bf Mon Sep 17 00:00:00 2001 From: LB7666 Date: Tue, 18 Aug 2026 05:43:43 +0800 Subject: [PATCH 04/74] Document module boundaries --- src/bin/pb-mapper.rs | 13 +++++++++++++ src/bin/pb-mapper/admin.rs | 12 ++++++++++++ src/common/auth.rs | 12 ++++++++++++ src/common/auth/actor.rs | 18 ++++++++++++++++++ src/common/auth/persistence.rs | 12 ++++++++++++ src/common/auth/runtime.rs | 15 +++++++++++++++ src/common/auth/tests.rs | 12 ++++++++++++ src/common/auth/timing_wheel.rs | 11 +++++++++++ src/common/message/secure.rs | 12 ++++++++++++ src/common/message/secure/frame.rs | 12 ++++++++++++ src/common/message/secure/limiter.rs | 11 +++++++++++ src/common/message/secure/replay.rs | 11 +++++++++++ src/common/message/secure/tests.rs | 11 +++++++++++ src/pb_server/admin.rs | 11 +++++++++++ src/pb_server/connection.rs | 16 ++++++++++++++++ src/pb_server/mod.rs | 12 ++++++++++++ src/pb_server/runtime.rs | 13 +++++++++++++ ui/native/pb_mapper_ffi/src/state.rs | 11 +++++++++++ .../pb_mapper_ffi/src/state/configuration.rs | 11 +++++++++++ ui/native/pb_mapper_ffi/src/state/runtime.rs | 12 ++++++++++++ ui/native/pb_mapper_ffi/src/state/status.rs | 13 +++++++++++++ 21 files changed, 261 insertions(+) diff --git a/src/bin/pb-mapper.rs b/src/bin/pb-mapper.rs index 8a26084..425a4a7 100644 --- a/src/bin/pb-mapper.rs +++ b/src/bin/pb-mapper.rs @@ -1,3 +1,16 @@ +//! Unified command-line entry point for every pb-mapper role. +//! +//! ```text +//! +-> server (relay) +//! process args -> clap -+-> register (publish a local service) +//! +-> connect (open a local listener) +//! +-> status (namespace-scoped inspection) +//! +-> admin (credential/control plane) +//! ``` +//! +//! Role-specific execution stays below this dispatch layer. Administrator parsing, +//! pagination, wire requests, and output rendering live in the `admin` module. + use std::error::Error; use std::net::{IpAddr, Ipv4Addr, Ipv6Addr}; use std::path::PathBuf; diff --git a/src/bin/pb-mapper/admin.rs b/src/bin/pb-mapper/admin.rs index bff21a1..10450b4 100644 --- a/src/bin/pb-mapper/admin.rs +++ b/src/bin/pb-mapper/admin.rs @@ -1,3 +1,15 @@ +//! Administrator CLI: command parsing, one-shot V2 requests, pagination, and rendering. +//! +//! ```text +//! admin args -> AdminRequest -> authenticated V2 connection -> relay +//! ^ | +//! +--- human / JSON / NDJSON <- AdminResponse <--------+ +//! ``` +//! +//! `--all` keeps the selected output contract: human and JSON aggregate pages, +//! while NDJSON deliberately streams one item at a time. Root-key rotation stages +//! a recovery copy before contacting the relay, then verifies the new credential. + use super::*; #[derive(Debug, Args)] diff --git a/src/common/auth.rs b/src/common/auth.rs index 39fdbe1..7e707d2 100644 --- a/src/common/auth.rs +++ b/src/common/auth.rs @@ -4,6 +4,18 @@ //! keys are derived from `(root key, server instance id, key id)` and the hot slot table //! stores only lifecycle metadata plus a weak lease reference. The background actor owns //! the strong leases through a hierarchical timing wheel. +//! +//! ```text +//! administrator key + instance id + key id -> derived temporary credential +//! | +//! request -> AuthContext -> Weak lease -+-> actor-owned Arc lease -> timing wheel +//! +-> cancel on expiry/revoke/reset/rotation +//! +//! AuthRuntime facade -> serialized actor -> encrypted snapshot + WAL +//! ``` +//! +//! The facade/model types remain in this root module; runtime checks, actor mutations, +//! persistence, expiry scheduling, and focused tests live in their respective children. use std::collections::{HashMap, HashSet, VecDeque}; use std::fmt; diff --git a/src/common/auth/actor.rs b/src/common/auth/actor.rs index 5ab4967..c464d50 100644 --- a/src/common/auth/actor.rs +++ b/src/common/auth/actor.rs @@ -1,3 +1,21 @@ +//! Serialized owner of mutable authentication lifecycle state. +//! +//! ```text +//! authenticated admin command +//! | +//! v +//! validate current admin lease +//! | +//! v +//! append encrypted WAL -> mutate slots / leases / timing wheel +//! | +//! +-> periodic snapshot + bounded replay/audit retention +//! ``` +//! +//! Keeping authorization revalidation and mutations in one actor prevents a request +//! authenticated before root rotation from executing against the new administrator +//! state. The actor is also the sole strong owner of temporary-key leases. + use super::*; pub(super) struct AuthActorState { diff --git a/src/common/auth/persistence.rs b/src/common/auth/persistence.rs index b2aea88..1b66732 100644 --- a/src/common/auth/persistence.rs +++ b/src/common/auth/persistence.rs @@ -1,3 +1,15 @@ +//! Durable, encrypted authentication state and audit/replay retention. +//! +//! ```text +//! startup: admin.key -> decrypt snapshot -> replay WAL -> in-memory state +//! mutation: command -> fsync encrypted WAL -> publish hot-state change +//! compact: hot state + audit + replay set -> snapshot -> truncate WAL +//! ``` +//! +//! Snapshot replacement and administrator-key files use atomic rename. Bounded audit +//! and replay collections are carried through compaction so security history does not +//! disappear when the WAL is truncated. + use super::*; pub(super) fn push_audit_record(inner: &AuthStateInner, record: AuditRecord) { diff --git a/src/common/auth/runtime.rs b/src/common/auth/runtime.rs index 4f236f7..5d51faa 100644 --- a/src/common/auth/runtime.rs +++ b/src/common/auth/runtime.rs @@ -1,3 +1,18 @@ +//! Public authentication runtime facade and hot-path credential checks. +//! +//! ```text +//! process credential + persisted state +//! | +//! v +//! hot slot table (Weak leases) <---- request authentication +//! | +//! +----> lifecycle actor (strong leases + time wheel) +//! ``` +//! +//! Read-only authentication stays synchronous and allocation-light. Every administrator +//! API captures a weak authority lease and sends it to the actor, where it is compared +//! with the current lease immediately before the operation executes. + use super::*; impl AuthRuntime { diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index 84dde37..2603026 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -1,3 +1,15 @@ +//! Authentication invariants exercised at the state-machine boundary. +//! +//! ```text +//! issue -> renew -> expire/revoke -> persist/restart +//! | | +//! +-> lease cancellation +-> encrypted recovery +//! root rotate -> reject old key + reject already-authenticated old context +//! ``` +//! +//! Protocol framing has its own tests under `common::message::secure::tests`; this +//! module focuses on lifecycle, persistence, audit, replay, and timing-wheel behavior. + use super::*; fn temp_state_dir(name: &str) -> PathBuf { diff --git a/src/common/auth/timing_wheel.rs b/src/common/auth/timing_wheel.rs index d96dafd..ab56e80 100644 --- a/src/common/auth/timing_wheel.rs +++ b/src/common/auth/timing_wheel.rs @@ -1,3 +1,14 @@ +//! Hierarchical expiry scheduler for temporary credential leases. +//! +//! ```text +//! lease(expires_at) -> level/slot bucket -> one-second actor tick -> expired leases +//! renew ----> version bump ------^ stale bucket entries are ignored +//! ``` +//! +//! The wheel owns strong `Arc` references. Request-facing structures retain +//! only `Weak` references, so expiry, revoke, reset, and root rotation have one clear +//! cancellation owner without keeping dead credentials alive indefinitely. + use super::*; struct WheelEntry { diff --git a/src/common/message/secure.rs b/src/common/message/secure.rs index 346ec2e..8f84031 100644 --- a/src/common/message/secure.rs +++ b/src/common/message/secure.rs @@ -4,6 +4,18 @@ //! request. It does not add a handshake or round trip. All following control messages on the //! same TCP connection use independently derived directional keys and monotonically increasing //! 64-bit counters. +//! +//! ```text +//! first flight: PBM2 | version | key id | timestamp+salt | counter | len | ciphertext +//! | | | +//! | +-> replay/time checks +-> bounded AEAD open +//! +-> derive directional session keys +//! +//! continuation: counter(n+1) | len | ciphertext -> same authenticated session +//! ``` +//! +//! This root module coordinates client/server sessions. Frame mechanics, replay admission, +//! log suppression, and protocol tests are isolated in focused child modules. use std::sync::{Arc, Mutex}; use std::time::{SystemTime, UNIX_EPOCH}; diff --git a/src/common/message/secure/frame.rs b/src/common/message/secure/frame.rs index 23222af..659a7cd 100644 --- a/src/common/message/secure/frame.rs +++ b/src/common/message/secure/frame.rs @@ -1,3 +1,15 @@ +//! Protocol-v2 key derivation and directional authenticated frame codecs. +//! +//! ```text +//! credential + connection salt -> HKDF -> c2s key / s2c key +//! plaintext -> counter + length + AEAD(AAD) -> encrypted frame +//! encrypted frame -> bound length -> verify counter/tag -> plaintext +//! ``` +//! +//! Counters are monotonic per direction and are included in both the nonce and AAD. +//! The initial reader can impose a smaller pre-authentication limit before allocating +//! a body; continuation frames retain the normal protocol maximum. + use super::*; #[derive(Clone)] diff --git a/src/common/message/secure/limiter.rs b/src/common/message/secure/limiter.rs index 59177ce..a28912a 100644 --- a/src/common/message/secure/limiter.rs +++ b/src/common/message/secure/limiter.rs @@ -1,3 +1,14 @@ +//! Cardinality-bounded suppression for repeated authentication failure logs. +//! +//! ```text +//! (peer IP, key id, reason) -> per-window counter -> emit first / suppress repeats +//! too many distinct keys ---------> shared overflow bucket +//! ``` +//! +//! This limits log amplification from the public relay port without changing protocol +//! decisions: every authentication failure is still rejected, only duplicate logging +//! is coalesced. + #[derive(Clone, Copy, Debug)] pub struct FailureLogDecision { pub emit: bool, diff --git a/src/common/message/secure/replay.rs b/src/common/message/secure/replay.rs index 0ebd0d2..e12c1db 100644 --- a/src/common/message/secure/replay.rs +++ b/src/common/message/secure/replay.rs @@ -1,3 +1,14 @@ +//! Fast process-local duplicate admission guard for protocol-v2 first flights. +//! +//! ```text +//! key id + salt -> SHA-256 fingerprint -> current Bloom window +//! -> previous Bloom window +//! ``` +//! +//! `check_and_insert` is called while one mutex is held, making concurrent admission +//! atomic. This Bloom filter protects all connection types from immediate duplicates; +//! administrator mutations additionally use the exact durable replay set in `auth`. + use super::*; pub(super) fn replay_fingerprint(key_id: u64, salt: &[u8; CONNECTION_SALT_LEN]) -> [u8; 32] { diff --git a/src/common/message/secure/tests.rs b/src/common/message/secure/tests.rs index e8cda72..409a06b 100644 --- a/src/common/message/secure/tests.rs +++ b/src/common/message/secure/tests.rs @@ -1,3 +1,14 @@ +//! End-to-end protocol-v2 framing and admission invariants. +//! +//! ```text +//! client session -> duplex transport -> ServerSecurity -> authenticated context +//! captured frame -- concurrent replay -----------------> exactly one admission +//! oversized first header ------------------------------> reject before body read +//! ``` +//! +//! These tests intentionally exercise both administrator and derived temporary +//! credentials, while lifecycle persistence remains covered by `common::auth::tests`. + use super::*; use crate::common::auth::{AuthConfig, LegacyProtocolPolicy}; use crate::common::checksum::encode_temporary_credential; diff --git a/src/pb_server/admin.rs b/src/pb_server/admin.rs index 46e103f..ba25768 100644 --- a/src/pb_server/admin.rs +++ b/src/pb_server/admin.rs @@ -1,3 +1,14 @@ +//! Server-side administrator request execution. +//! +//! ```text +//! authenticated AdminRequest -> revalidated AuthContext -> auth actor / manager +//! -> AdminResponse +//! ``` +//! +//! Credential lifecycle operations go to `AuthRuntime`; service and connection +//! inventory requests go to the routing manager. Read operations are audited without +//! weakening the primary response when only audit emission fails. + use std::time::Duration; use tokio::net::TcpStream; diff --git a/src/pb_server/connection.rs b/src/pb_server/connection.rs index dff693c..a80168e 100644 --- a/src/pb_server/connection.rs +++ b/src/pb_server/connection.rs @@ -1,3 +1,19 @@ +//! Per-connection admission, authentication, namespace resolution, and role dispatch. +//! +//! ```text +//! accepted TCP socket +//! | +//! v +//! bounded V2/legacy first frame -> AuthContext -> namespace policy +//! | | +//! +-> structured auth error +-> register / subscribe / stream +//! +-> status / administrator request +//! ``` +//! +//! Long-lived register and subscribe futures are raced against the credential's +//! cancellation token here. This outer guard closes a subscriber even when the paired +//! service stream belongs to a different credential. + use super::*; pub(super) async fn handle_listener( diff --git a/src/pb_server/mod.rs b/src/pb_server/mod.rs index 37238f4..5bcbf01 100644 --- a/src/pb_server/mod.rs +++ b/src/pb_server/mod.rs @@ -1,3 +1,15 @@ +//! Relay server domain model and module wiring. +//! +//! ```text +//! TCP listener -> connection authentication/dispatch -> ManagerTask queue +//! -> routing runtime +//! registered control connection <------ ConnTask ----+------> subscriber +//! ``` +//! +//! `connection` owns per-socket protocol/authentication concerns, while `runtime` +//! serializes global routing maps and quotas. Service-side and client-side tunnel loops +//! remain isolated in `server` and `client`. + mod admin; mod client; mod error; diff --git a/src/pb_server/runtime.rs b/src/pb_server/runtime.rs index a89a382..8b8a655 100644 --- a/src/pb_server/runtime.rs +++ b/src/pb_server/runtime.rs @@ -1,3 +1,16 @@ +//! Relay orchestration and the serialized routing-manager event loop. +//! +//! ```text +//! listener task ---- Accept -------+ +//! control tasks ---- Register -----+-> ManagerTask loop -> routing maps / quotas +//! subscriber ------- Subcribe -----+ -> ConnTask responses +//! provider stream -- Stream/Ack ---+ +//! ``` +//! +//! The manager loop is the single writer for connection IDs, registrations, pending +//! streams, per-namespace counts, and rate limits. Socket I/O runs in spawned connection +//! tasks and communicates with this state only through typed tasks. + use super::*; struct RemoteIdProvider { diff --git a/ui/native/pb_mapper_ffi/src/state.rs b/ui/native/pb_mapper_ffi/src/state.rs index fc5f26c..c9fed75 100644 --- a/ui/native/pb_mapper_ffi/src/state.rs +++ b/ui/native/pb_mapper_ffi/src/state.rs @@ -1,3 +1,14 @@ +//! Shared Flutter-FFI application state and module boundaries. +//! +//! ```text +//! Flutter command -> Arc> -> configuration / runtime / status +//! -> change events back to Flutter +//! ``` +//! +//! Slow DNS, bind, and connectivity work is deliberately performed outside the global +//! state lock. Per-key claims prevent duplicate setup while keeping unrelated UI reads +//! and operations responsive. + use std::collections::{HashMap, HashSet}; use std::fs; use std::net::{IpAddr, Ipv4Addr, SocketAddr}; diff --git a/ui/native/pb_mapper_ffi/src/state/configuration.rs b/ui/native/pb_mapper_ffi/src/state/configuration.rs index 2d0e279..2ad22b8 100644 --- a/ui/native/pb_mapper_ffi/src/state/configuration.rs +++ b/ui/native/pb_mapper_ffi/src/state/configuration.rs @@ -1,3 +1,14 @@ +//! User-writable configuration storage and in-memory state initialization. +//! +//! ```text +//! app directory -> pb-mapper-ui/config.json -> AppConfig -> process credential +//! -> services.json / clients.json -> remembered tunnel definitions +//! ``` +//! +//! An explicit Flutter app directory wins on every platform. This same directory is +//! the root for relay authentication state, so desktop and mobile UI processes never +//! depend on root-owned `/var/lib` paths. + use super::*; impl PbMapperState { diff --git a/ui/native/pb_mapper_ffi/src/state/runtime.rs b/ui/native/pb_mapper_ffi/src/state/runtime.rs index 44e2444..91e848b 100644 --- a/ui/native/pb_mapper_ffi/src/state/runtime.rs +++ b/ui/native/pb_mapper_ffi/src/state/runtime.rs @@ -1,3 +1,15 @@ +//! Runtime lifecycle for the embedded relay, registered services, and local clients. +//! +//! ```text +//! start relay: bind listener -> initialize app-local auth -> spawn -> mark running +//! register: resolved addresses -> spawn control pool -> retain JoinHandle +//! connect: preflight local bind -> spawn listener ----> retain JoinHandle +//! stop: cancel relay + abort owned tunnel tasks + clear runtime maps +//! ``` +//! +//! Readiness is published only after both listener binding and authentication +//! initialization succeed, preventing the UI from displaying a phantom running relay. + use super::*; impl PbMapperState { diff --git a/ui/native/pb_mapper_ffi/src/state/status.rs b/ui/native/pb_mapper_ffi/src/state/status.rs index 8ae5f36..8146a06 100644 --- a/ui/native/pb_mapper_ffi/src/state/status.rs +++ b/ui/native/pb_mapper_ffi/src/state/status.rs @@ -1,3 +1,16 @@ +//! Non-blocking status views and bounded asynchronous refresh scheduling for Flutter. +//! +//! ```text +//! UI read -> cached snapshot -> immediate response +//! | +//! +-- stale? -> one deduplicated network refresh -> cache + change event +//! force refresh -----------------------------------------> awaited network result +//! ``` +//! +//! Service, client, and embedded-relay caches are independent. Refresh markers prevent +//! duplicate probes, while visible events are emitted only when the displayed state +//! actually changes. + use super::*; impl PbMapperState { From 93b501f30fa1e5e34e614a0a09b5440ef8fbc875 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Tue, 18 Aug 2026 05:47:05 +0800 Subject: [PATCH 05/74] Reject conflicting auth configuration --- docs/authentication-v2.md | 3 ++- docs/authentication-v2.zh-CN.md | 4 +++- src/bin/pb-mapper.rs | 13 ++++++++++- src/common/auth.rs | 38 ++++++++++++++++++++++++++------- src/common/auth/tests.rs | 14 ++++++++++++ 5 files changed, 61 insertions(+), 11 deletions(-) diff --git a/docs/authentication-v2.md b/docs/authentication-v2.md index 96967c6..032fab8 100644 --- a/docs/authentication-v2.md +++ b/docs/authentication-v2.md @@ -253,7 +253,8 @@ New clients always emit protocol v2. A v0.4 server defaults to accepting legacy framing so older clients can be upgraded without an outage. Operators can view legacy connection counters, upgrade all clients, and then set the policy to `deny`. Upgrade the relay before any client because v0.3 relays do not understand -the v2 first-frame magic. +the v2 first-frame magic. An explicitly configured `PB_MAPPER_LEGACY_PROTOCOL` +is trimmed and must be `allow` or `deny`; malformed values fail closed to `deny`. Fresh servers generate a random administrator key. Both the relay and install scripts preserve an existing `/var/lib/pb-mapper-server/msg_header_key` by diff --git a/docs/authentication-v2.zh-CN.md b/docs/authentication-v2.zh-CN.md index b713334..80af9e1 100644 --- a/docs/authentication-v2.zh-CN.md +++ b/docs/authentication-v2.zh-CN.md @@ -196,7 +196,9 @@ human 输出单个合并表格。稳定错误结构包含 `code`、`message`、 新客户端固定发送 V2。0.4 服务端默认暂时接受旧帧，方便滚动升级；确认 `active_legacy_connections` 归零后，可执行 `legacy-protocol set deny`。必须先升级中继、 -再升级客户端，因为 0.3 中继无法识别 V2 首帧 magic。 +再升级客户端，因为 0.3 中继无法识别 V2 首帧 magic。显式配置的 +`PB_MAPPER_LEGACY_PROTOCOL` 会先去除首尾空白，并且只能是 `allow` 或 `deny`； +无效值会 fail closed 为 `deny`。 新安装会随机生成管理员密钥。中继自身与安装脚本在未配置新 key 或环境变量时，如果 发现旧的 `/var/lib/pb-mapper-server/msg_header_key`，会将其复制到新路径，保留现有 diff --git a/src/bin/pb-mapper.rs b/src/bin/pb-mapper.rs index 425a4a7..b73cd02 100644 --- a/src/bin/pb-mapper.rs +++ b/src/bin/pb-mapper.rs @@ -91,7 +91,11 @@ struct ServerArgs { #[arg(long, default_value = DEFAULT_AUTH_STATE_DIR)] auth_state_dir: PathBuf, /// Create a random administrator key before starting the relay. - #[arg(long, default_value_t = false)] + #[arg( + long, + conflicts_with = "use_machine_msg_header_key", + default_value_t = false + )] init_admin_key: bool, /// Replace an existing administrator key when used with --init-admin-key. #[arg(long, requires = "init_admin_key", default_value_t = false)] @@ -484,5 +488,12 @@ mod tests { ["pb-mapper", "admin", "key", "issue", "--ttl", "1fortnight",] ) .is_err()); + assert!(Cli::try_parse_from([ + "pb-mapper", + "server", + "--init-admin-key", + "--use-machine-msg-header-key", + ]) + .is_err()); } } diff --git a/src/common/auth.rs b/src/common/auth.rs index 7e707d2..872f62f 100644 --- a/src/common/auth.rs +++ b/src/common/auth.rs @@ -96,18 +96,40 @@ impl Default for AuthConfig { MIN_TEMP_KEY_TTL.as_secs(), 365 * 24 * 60 * 60, )), - legacy_protocol: match std::env::var("PB_MAPPER_LEGACY_PROTOCOL") - .unwrap_or_else(|_| "allow".to_string()) - .to_ascii_lowercase() - .as_str() - { - "deny" => LegacyProtocolPolicy::Deny, - _ => LegacyProtocolPolicy::Allow, - }, + legacy_protocol: legacy_protocol_from_env(), } } } +fn legacy_protocol_from_env() -> LegacyProtocolPolicy { + match std::env::var("PB_MAPPER_LEGACY_PROTOCOL") { + Err(std::env::VarError::NotPresent) => LegacyProtocolPolicy::Allow, + Err(std::env::VarError::NotUnicode(_)) => { + tracing::error!( + event = "legacy_protocol_config_invalid", + "PB_MAPPER_LEGACY_PROTOCOL is not UTF-8; denying legacy framing" + ); + LegacyProtocolPolicy::Deny + } + Ok(value) => parse_legacy_protocol_policy(&value).unwrap_or_else(|| { + tracing::error!( + event = "legacy_protocol_config_invalid", + value, + "PB_MAPPER_LEGACY_PROTOCOL must be `allow` or `deny`; denying legacy framing" + ); + LegacyProtocolPolicy::Deny + }), + } +} + +fn parse_legacy_protocol_policy(value: &str) -> Option { + match value.trim().to_ascii_lowercase().as_str() { + "allow" => Some(LegacyProtocolPolicy::Allow), + "deny" => Some(LegacyProtocolPolicy::Deny), + _ => None, + } +} + fn env_usize(name: &str, default: usize, min: usize, max: usize) -> usize { std::env::var(name) .ok() diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index 2603026..f003ab8 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -21,6 +21,20 @@ fn temp_state_dir(name: &str) -> PathBuf { std::env::temp_dir().join(format!("pb-mapper-{name}-{}", hex(&suffix))) } +#[test] +fn legacy_protocol_policy_trims_valid_values_and_rejects_unknown_values() { + assert_eq!( + parse_legacy_protocol_policy(" allow\n"), + Some(LegacyProtocolPolicy::Allow) + ); + assert_eq!( + parse_legacy_protocol_policy(" DENY "), + Some(LegacyProtocolPolicy::Deny) + ); + assert_eq!(parse_legacy_protocol_policy("enabled"), None); + assert_eq!(parse_legacy_protocol_policy(""), None); +} + #[test] fn key_id_round_trip() { let key_id = make_key_id(42, 65_535); From 5cec089dc13d4ccf5e9658d928c9e00910e13828 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Tue, 18 Aug 2026 05:49:59 +0800 Subject: [PATCH 06/74] Distinguish invalid administrator keys --- src/common/auth/runtime.rs | 4 ++-- src/common/auth/tests.rs | 10 +++++++++- tests/regression.rs | 5 ++++- 3 files changed, 15 insertions(+), 4 deletions(-) diff --git a/src/common/auth/runtime.rs b/src/common/auth/runtime.rs index 5d51faa..2b7f274 100644 --- a/src/common/auth/runtime.rs +++ b/src/common/auth/runtime.rs @@ -187,8 +187,8 @@ impl AuthRuntime { if !bool::from(presented_key.ct_eq(&admin.key)) { inner.auth_failures.fetch_add(1, Ordering::Relaxed); return Err(AuthFailure::new( - "administrator_key_rotated", - "administrator credential no longer matches the active root key", + "administrator_key_invalid", + "administrator credential does not match the active root key", false, )); } diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index f003ab8..5ae3b31 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -220,6 +220,14 @@ async fn root_rotation_rejects_old_key_and_in_flight_admin_context() { }; let runtime = AuthRuntime::start(old_key, config).await.unwrap(); let old_admin = runtime.authenticate_presented(0, &old_key).unwrap(); + let mistyped_key = *b"1123456789abcdefghijklmnopqrstuv"; + assert_eq!( + runtime + .authenticate_presented(0, &mistyped_key) + .unwrap_err() + .code, + "administrator_key_invalid" + ); runtime .rotate_root(&old_admin, new_key) @@ -231,7 +239,7 @@ async fn root_rotation_rejects_old_key_and_in_flight_admin_context() { .authenticate_presented(0, &old_key) .unwrap_err() .code, - "administrator_key_rotated" + "administrator_key_invalid" ); assert_eq!( runtime diff --git a/tests/regression.rs b/tests/regression.rs index 4f0158a..23a2cff 100644 --- a/tests/regression.rs +++ b/tests/regression.rs @@ -174,10 +174,13 @@ async fn read_secure_request( } fn auth_config(server_addr: SocketAddr) -> AuthConfig { + static CONFIG_SEQUENCE: AtomicUsize = AtomicUsize::new(0); + set_process_msg_header_key(Some(TEST_ADMIN_KEY)).unwrap(); + let sequence = CONFIG_SEQUENCE.fetch_add(1, Ordering::Relaxed); AuthConfig { state_dir: std::env::temp_dir().join(format!( - "pb-mapper-regression-{}-{}", + "pb-mapper-regression-{}-{}-{sequence}", std::process::id(), server_addr.port() )), From 13c5df285b597874795e86b2b8c48e5f53680160 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Tue, 18 Aug 2026 05:57:58 +0800 Subject: [PATCH 07/74] Bound timing wheel catch-up work --- src/common/auth/tests.rs | 17 ++++++++++++++++ src/common/auth/timing_wheel.rs | 36 +++++++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+) diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index 5ae3b31..f30d5f8 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -339,3 +339,20 @@ fn timing_wheel_ignores_stale_renewal_entry() { assert!(wheel.advance(now + 6).is_empty()); assert_eq!(wheel.advance(now + 20).len(), 1); } + +#[test] +fn timing_wheel_fast_forwards_large_clock_jumps() { + let now = 1_000; + let target = now + 7 * 24 * 60 * 60; + let expired = Arc::new(AuthLease::new(make_key_id(1, 0), now + 5)); + let future = Arc::new(AuthLease::new(make_key_id(1, 1), target + 20)); + let mut wheel = TimingWheel::new(now); + wheel.insert(expired.clone()); + wheel.insert(future.clone()); + + let due = wheel.advance(target); + assert_eq!(due.len(), 1); + assert_eq!(due[0].key_id(), expired.key_id()); + assert!(wheel.advance(target + 19).is_empty()); + assert_eq!(wheel.advance(target + 20).len(), 1); +} diff --git a/src/common/auth/timing_wheel.rs b/src/common/auth/timing_wheel.rs index ab56e80..c64a0b0 100644 --- a/src/common/auth/timing_wheel.rs +++ b/src/common/auth/timing_wheel.rs @@ -3,6 +3,8 @@ //! ```text //! lease(expires_at) -> level/slot bucket -> one-second actor tick -> expired leases //! renew ----> version bump ------^ stale bucket entries are ignored +//! | +//! large clock jump -> bounded bucket scan + rebuild (never second-by-second catch-up) //! ``` //! //! The wheel owns strong `Arc` references. Request-facing structures retain @@ -11,6 +13,8 @@ use super::*; +const MAX_INCREMENTAL_ADVANCE_SECONDS: u64 = 256; + struct WheelEntry { lease: Arc, version: u64, @@ -56,6 +60,10 @@ impl TimingWheel { } pub(super) fn advance(&mut self, target: u64) -> Vec> { + if target.saturating_sub(self.now) > MAX_INCREMENTAL_ADVANCE_SECONDS { + return self.fast_forward(target); + } + let mut due = Vec::new(); while self.now < target { self.now = self.now.saturating_add(1); @@ -82,6 +90,28 @@ impl TimingWheel { due } + fn fast_forward(&mut self, target: u64) -> Vec> { + self.now = target; + let mut entries = Vec::new(); + take_all_entries(&mut self.level0, &mut entries); + take_all_entries(&mut self.level1, &mut entries); + take_all_entries(&mut self.level2, &mut entries); + take_all_entries(&mut self.level3, &mut entries); + + let mut due = Vec::new(); + for entry in entries { + if entry.version != entry.lease.wheel_version.load(Ordering::Acquire) { + continue; + } + if entry.lease.expires_at() <= target { + due.push(entry.lease); + } else { + self.insert_with_version(entry.lease, entry.version); + } + } + due + } + fn cascade(&mut self, level: u8) { let entries = match level { 1 => { @@ -113,3 +143,9 @@ impl TimingWheel { fn empty_buckets(count: usize) -> Vec> { std::iter::repeat_with(Vec::new).take(count).collect() } + +fn take_all_entries(buckets: &mut [Vec], entries: &mut Vec) { + for bucket in buckets { + entries.append(bucket); + } +} From 8e3a54c4d1afb89da04a55cf107510b18636c964 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Tue, 18 Aug 2026 06:04:35 +0800 Subject: [PATCH 08/74] Warn on invalid auth limits --- src/common/auth.rs | 47 ++++++++++++++++++++++++++++++++++++---------- 1 file changed, 37 insertions(+), 10 deletions(-) diff --git a/src/common/auth.rs b/src/common/auth.rs index 872f62f..1fe725d 100644 --- a/src/common/auth.rs +++ b/src/common/auth.rs @@ -131,19 +131,46 @@ fn parse_legacy_protocol_policy(value: &str) -> Option { } fn env_usize(name: &str, default: usize, min: usize, max: usize) -> usize { - std::env::var(name) - .ok() - .and_then(|value| value.parse::().ok()) - .filter(|value| (*value >= min) && (*value <= max)) - .unwrap_or(default) + env_bounded(name, default, min, max) } fn env_u64(name: &str, default: u64, min: u64, max: u64) -> u64 { - std::env::var(name) - .ok() - .and_then(|value| value.parse::().ok()) - .filter(|value| (*value >= min) && (*value <= max)) - .unwrap_or(default) + env_bounded(name, default, min, max) +} + +fn env_bounded(name: &str, default: T, min: T, max: T) -> T +where + T: std::str::FromStr + PartialOrd + Copy + fmt::Display, +{ + match std::env::var(name) { + Err(std::env::VarError::NotPresent) => default, + Ok(raw) => match raw.parse::() { + Ok(value) if value >= min && value <= max => value, + _ => { + tracing::warn!( + event = "auth_config_value_invalid", + variable = name, + value = raw, + min = %min, + max = %max, + fallback = %default, + "invalid authentication configuration value; using the default" + ); + default + } + }, + Err(std::env::VarError::NotUnicode(_)) => { + tracing::warn!( + event = "auth_config_value_invalid", + variable = name, + min = %min, + max = %max, + fallback = %default, + "authentication configuration value is not UTF-8; using the default" + ); + default + } + } } #[derive(Clone, Debug, Serialize, Deserialize)] From 6a5857a852794af336a7ea3e7bd3f3f811ded69e Mon Sep 17 00:00:00 2001 From: LB7666 Date: Tue, 18 Aug 2026 06:07:56 +0800 Subject: [PATCH 09/74] Expire overdue timing wheel entries promptly --- src/common/auth/tests.rs | 25 +++++++++++++++++++++++++ src/common/auth/timing_wheel.rs | 27 ++++++++++++++++++++++++--- 2 files changed, 49 insertions(+), 3 deletions(-) diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index f30d5f8..65c4978 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -356,3 +356,28 @@ fn timing_wheel_fast_forwards_large_clock_jumps() { assert!(wheel.advance(target + 19).is_empty()); assert_eq!(wheel.advance(target + 20).len(), 1); } + +#[test] +fn timing_wheel_returns_already_expired_insert_without_wrapping() { + let now = 1_000; + let expired = Arc::new(AuthLease::new(make_key_id(1, 0), now - 1)); + let mut wheel = TimingWheel::new(now); + wheel.insert(expired.clone()); + + let due = wheel.advance(now); + assert_eq!(due.len(), 1); + assert_eq!(due[0].key_id(), expired.key_id()); +} + +#[test] +fn timing_wheel_expires_cascaded_boundary_entry_without_an_extra_tick() { + let now = 700; + let expires_at = 1_024; + let lease = Arc::new(AuthLease::new(make_key_id(1, 0), expires_at)); + let mut wheel = TimingWheel::new(now); + wheel.insert(lease.clone()); + + let due = wheel.advance(expires_at); + assert_eq!(due.len(), 1); + assert_eq!(due[0].key_id(), lease.key_id()); +} diff --git a/src/common/auth/timing_wheel.rs b/src/common/auth/timing_wheel.rs index c64a0b0..03e0a6b 100644 --- a/src/common/auth/timing_wheel.rs +++ b/src/common/auth/timing_wheel.rs @@ -5,6 +5,7 @@ //! renew ----> version bump ------^ stale bucket entries are ignored //! | //! large clock jump -> bounded bucket scan + rebuild (never second-by-second catch-up) +//! overdue insert --> immediate-due queue --> next advance, without a wheel revolution //! ``` //! //! The wheel owns strong `Arc` references. Request-facing structures retain @@ -22,6 +23,7 @@ struct WheelEntry { pub(super) struct TimingWheel { now: u64, + immediate_due: Vec, level0: Vec>, level1: Vec>, level2: Vec>, @@ -32,6 +34,7 @@ impl TimingWheel { pub(super) fn new(now: u64) -> Self { Self { now, + immediate_due: Vec::new(), level0: empty_buckets(256), level1: empty_buckets(64), level2: empty_buckets(64), @@ -48,7 +51,9 @@ impl TimingWheel { let expires_at = lease.expires_at(); let delta = expires_at.saturating_sub(self.now); let entry = WheelEntry { lease, version }; - if delta < 1 << 8 { + if expires_at <= self.now { + self.immediate_due.push(entry); + } else if delta < 1 << 8 { self.level0[(expires_at & 0xff) as usize].push(entry); } else if delta < 1 << 14 { self.level1[((expires_at >> 8) & 0x3f) as usize].push(entry); @@ -64,7 +69,7 @@ impl TimingWheel { return self.fast_forward(target); } - let mut due = Vec::new(); + let mut due = self.take_immediate_due(target); while self.now < target { self.now = self.now.saturating_add(1); if self.now & 0xff == 0 { @@ -76,6 +81,7 @@ impl TimingWheel { } } } + due.extend(self.take_immediate_due(self.now)); let index = (self.now & 0xff) as usize; for entry in std::mem::take(&mut self.level0[index]) { if entry.version == entry.lease.wheel_version.load(Ordering::Acquire) { @@ -92,7 +98,7 @@ impl TimingWheel { fn fast_forward(&mut self, target: u64) -> Vec> { self.now = target; - let mut entries = Vec::new(); + let mut entries = std::mem::take(&mut self.immediate_due); take_all_entries(&mut self.level0, &mut entries); take_all_entries(&mut self.level1, &mut entries); take_all_entries(&mut self.level2, &mut entries); @@ -112,6 +118,21 @@ impl TimingWheel { due } + fn take_immediate_due(&mut self, target: u64) -> Vec> { + let mut due = Vec::new(); + for entry in std::mem::take(&mut self.immediate_due) { + if entry.version != entry.lease.wheel_version.load(Ordering::Acquire) { + continue; + } + if entry.lease.expires_at() <= target { + due.push(entry.lease); + } else { + self.insert_with_version(entry.lease, entry.version); + } + } + due + } + fn cascade(&mut self, level: u8) { let entries = match level { 1 => { From 866d20012fdef863caa1519b890bdd992a96217d Mon Sep 17 00:00:00 2001 From: LB7666 Date: Tue, 18 Aug 2026 06:14:22 +0800 Subject: [PATCH 10/74] Harden auth lifecycle cleanup --- src/common/auth/actor.rs | 30 +++++++++++++++++++++--------- src/common/auth/tests.rs | 33 +++++++++++++++++++++++++++++++++ src/common/auth/timing_wheel.rs | 9 +++++++++ 3 files changed, 63 insertions(+), 9 deletions(-) diff --git a/src/common/auth/actor.rs b/src/common/auth/actor.rs index c464d50..78e7ea0 100644 --- a/src/common/auth/actor.rs +++ b/src/common/auth/actor.rs @@ -113,14 +113,11 @@ pub(super) async fn run_auth_actor( } } } - admin_replay_order.retain(|record| { - let keep = now.saturating_sub(record.client_timestamp) - <= ADMIN_REPLAY_RETENTION.as_secs(); - if !keep { - admin_replays.remove(&record.fingerprint); - } - keep - }); + prune_expired_admin_replays( + now, + &mut admin_replays, + &mut admin_replay_order, + ); if now.saturating_sub(last_snapshot_at) >= SNAPSHOT_COMPACTION_INTERVAL.as_secs() { let snapshot = build_snapshot(&inner, &cold, &admin_replay_order); if let Err(error) = write_snapshot_and_truncate_wal( @@ -258,6 +255,8 @@ fn actor_claim_admin_mutation( fingerprint: [u8; 32], client_timestamp: u64, ) -> Result<(), AuthFailure> { + let now = unix_seconds(); + prune_expired_admin_replays(now, admin_replays, admin_replay_order); if admin_replays.contains(&fingerprint) { return Err(AuthFailure::new( "admin_request_replayed", @@ -272,7 +271,6 @@ fn actor_claim_admin_mutation( true, )); } - let now = unix_seconds(); if now.abs_diff(client_timestamp) > ADMIN_REPLAY_RETENTION.as_secs() / 2 { return Err(AuthFailure::new( "admin_request_timestamp_invalid", @@ -294,6 +292,20 @@ fn actor_claim_admin_mutation( Ok(()) } +pub(super) fn prune_expired_admin_replays( + now: u64, + admin_replays: &mut HashSet<[u8; 32]>, + admin_replay_order: &mut VecDeque, +) { + admin_replay_order.retain(|record| { + let keep = now.saturating_sub(record.client_timestamp) <= ADMIN_REPLAY_RETENTION.as_secs(); + if !keep { + admin_replays.remove(&record.fingerprint); + } + keep + }); +} + fn validate_ttl(config: &AuthConfig, ttl: Duration) -> Result { if ttl < MIN_TEMP_KEY_TTL { return Err(AuthFailure::new( diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index 65c4978..2741087 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -381,3 +381,36 @@ fn timing_wheel_expires_cascaded_boundary_entry_without_an_extra_tick() { assert_eq!(due.len(), 1); assert_eq!(due[0].key_id(), lease.key_id()); } + +#[test] +fn timing_wheel_clear_cancels_owned_leases() { + let now = 1_000; + let lease = Arc::new(AuthLease::new(make_key_id(1, 0), now + 60)); + let cancellation = lease.cancellation_token(); + let mut wheel = TimingWheel::new(now); + wheel.insert(lease); + + wheel.clear(now + 1); + assert!(cancellation.is_cancelled()); +} + +#[test] +fn replay_pruning_removes_only_records_outside_the_retention_window() { + let now = 10_000; + let expired = AdminReplayRecord { + fingerprint: [1; 32], + client_timestamp: now - ADMIN_REPLAY_RETENTION.as_secs() - 1, + }; + let current = AdminReplayRecord { + fingerprint: [2; 32], + client_timestamp: now, + }; + let mut replay_set = HashSet::from([expired.fingerprint, current.fingerprint]); + let mut replay_order = VecDeque::from([expired, current.clone()]); + + super::actor::prune_expired_admin_replays(now, &mut replay_set, &mut replay_order); + + assert_eq!(replay_set, HashSet::from([current.fingerprint])); + assert_eq!(replay_order.len(), 1); + assert_eq!(replay_order[0].fingerprint, current.fingerprint); +} diff --git a/src/common/auth/timing_wheel.rs b/src/common/auth/timing_wheel.rs index 03e0a6b..32208a7 100644 --- a/src/common/auth/timing_wheel.rs +++ b/src/common/auth/timing_wheel.rs @@ -6,6 +6,7 @@ //! | //! large clock jump -> bounded bucket scan + rebuild (never second-by-second catch-up) //! overdue insert --> immediate-due queue --> next advance, without a wheel revolution +//! reset/rotation -> cancel every wheel-owned lease -> clear all buckets //! ``` //! //! The wheel owns strong `Arc` references. Request-facing structures retain @@ -157,6 +158,14 @@ impl TimingWheel { } pub(super) fn clear(&mut self, now: u64) { + let mut entries = std::mem::take(&mut self.immediate_due); + take_all_entries(&mut self.level0, &mut entries); + take_all_entries(&mut self.level1, &mut entries); + take_all_entries(&mut self.level2, &mut entries); + take_all_entries(&mut self.level3, &mut entries); + for entry in entries { + entry.lease.cancellation.cancel(); + } *self = Self::new(now); } } From aa694992caa32ad344d1efcf76cf1f563e59b8fe Mon Sep 17 00:00:00 2001 From: LB7666 Date: Tue, 18 Aug 2026 06:22:28 +0800 Subject: [PATCH 11/74] Persist temporary key tombstone times --- src/common/auth.rs | 3 ++ src/common/auth/actor.rs | 65 +++++++++++++++++++++++++--------- src/common/auth/persistence.rs | 5 ++- src/common/auth/runtime.rs | 5 +++ src/common/auth/tests.rs | 9 +++++ 5 files changed, 70 insertions(+), 17 deletions(-) diff --git a/src/common/auth.rs b/src/common/auth.rs index 1fe725d..4d0e29d 100644 --- a/src/common/auth.rs +++ b/src/common/auth.rs @@ -416,6 +416,7 @@ pub struct AuthStatus { struct ColdMetadata { issued_at: u64, label: Option, + tombstoned_at: u64, } enum AuthCommand { @@ -660,6 +661,8 @@ struct PersistedEntry { issued_at: u64, expires_at: u64, label: Option, + #[serde(default)] + tombstoned_at: Option, } #[derive(Clone, Debug, Serialize, Deserialize)] diff --git a/src/common/auth/actor.rs b/src/common/auth/actor.rs index 78e7ea0..324256a 100644 --- a/src/common/auth/actor.rs +++ b/src/common/auth/actor.rs @@ -54,7 +54,6 @@ pub(super) async fn run_auth_actor( mut admin_replays, mut admin_replay_order, } = state; - let now = unix_seconds(); let mut tombstones = inner .slots .read() @@ -62,12 +61,22 @@ pub(super) async fn run_auth_actor( .iter() .enumerate() .filter_map(|(index, slot)| { - matches!(slot.state, SlotState::Expired | SlotState::Revoked).then_some(( - now.saturating_add(TOMBSTONE_RETENTION.as_secs()), - make_key_id(slot.generation, index as u32), + if !matches!(slot.state, SlotState::Expired | SlotState::Revoked) { + return None; + } + let key_id = make_key_id(slot.generation, index as u32); + let tombstoned_at = cold + .get(&key_id) + .map(|metadata| metadata.tombstoned_at) + .unwrap_or(0); + Some(( + tombstoned_at.saturating_add(TOMBSTONE_RETENTION.as_secs()), + key_id, )) }) - .collect::>(); + .collect::>(); + tombstones.sort_unstable_by_key(|(cleanup_at, _)| *cleanup_at); + let mut tombstones = VecDeque::from(tombstones); let mut last_snapshot_at = unix_seconds(); let mut tick = tokio::time::interval(Duration::from_secs(1)); tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); @@ -87,7 +96,15 @@ pub(super) async fn run_auth_actor( if slot.generation == key_generation(key_id) && slot.state == SlotState::Active { slot.state = SlotState::Expired; lease.cancellation.cancel(); - tombstones.push_back((now.saturating_add(TOMBSTONE_RETENTION.as_secs()), key_id)); + let tombstoned_at = slot.expires_at; + if let Some(metadata) = cold.get_mut(&key_id) { + metadata.tombstoned_at = tombstoned_at; + } + push_tombstone( + &mut tombstones, + tombstoned_at, + key_id, + ); tracing::info!( event = "temporary_key_expired", auth_stage = "expiry", @@ -186,7 +203,7 @@ pub(super) async fn run_auth_actor( } AuthCommand::Revoke { authority, key_id, response } => { let result = validate_admin_authority(&inner, &authority) - .and_then(|()| actor_revoke(&inner, &config, &cold, &mut tombstones, key_id)); + .and_then(|()| actor_revoke(&inner, &config, &mut cold, &mut tombstones, key_id)); let _ = response.send(result); } AuthCommand::Gc { authority, response } => { @@ -379,6 +396,7 @@ fn actor_issue( issued_at, expires_at, label: label.clone(), + tombstoned_at: None, }; append_mutation( config, @@ -391,7 +409,14 @@ fn actor_issue( slot.state = SlotState::Active; slot.expires_at = expires_at; slot.lease = Arc::downgrade(&lease); - cold.insert(key_id, ColdMetadata { issued_at, label }); + cold.insert( + key_id, + ColdMetadata { + issued_at, + label, + tombstoned_at: 0, + }, + ); wheel.insert(lease); drop(slots); metadata_with_credential(inner, cold, key_id, true) @@ -512,7 +537,7 @@ fn actor_renew( fn actor_revoke( inner: &Arc, config: &AuthConfig, - cold: &HashMap, + cold: &mut HashMap, tombstones: &mut VecDeque<(u64, u64)>, key_id: u64, ) -> Result { @@ -532,22 +557,24 @@ fn actor_revoke( false, )); } - let cold_metadata = cold.get(&key_id).ok_or_else(|| key_not_found(key_id))?; + let label = cold + .get(&key_id) + .ok_or_else(|| key_not_found(key_id))? + .label + .clone(); append_mutation( config, inner, StateMutation::Revoke { key_id, at: now }, - audit( - "temporary_key_revoke", - Some(key_id), - cold_metadata.label.clone(), - ), + audit("temporary_key_revoke", Some(key_id), label.clone()), )?; slot.state = SlotState::Revoked; if let Some(lease) = slot.lease.upgrade() { lease.cancellation.cancel(); } - tombstones.push_back((now.saturating_add(TOMBSTONE_RETENTION.as_secs()), key_id)); + let cold_metadata = cold.get_mut(&key_id).ok_or_else(|| key_not_found(key_id))?; + cold_metadata.tombstoned_at = now; + push_tombstone(tombstones, now, key_id); Ok(TemporaryKeyMetadata { key_id, state: slot_state_name(slot.state).to_string(), @@ -557,6 +584,12 @@ fn actor_revoke( }) } +fn push_tombstone(tombstones: &mut VecDeque<(u64, u64)>, tombstoned_at: u64, key_id: u64) { + let cleanup_at = tombstoned_at.saturating_add(TOMBSTONE_RETENTION.as_secs()); + let index = tombstones.partition_point(|(current, _)| *current <= cleanup_at); + tombstones.insert(index, (cleanup_at, key_id)); +} + fn actor_gc( inner: &Arc, config: &AuthConfig, diff --git a/src/common/auth/persistence.rs b/src/common/auth/persistence.rs index 1b66732..f4a737e 100644 --- a/src/common/auth/persistence.rs +++ b/src/common/auth/persistence.rs @@ -58,6 +58,7 @@ pub(super) fn build_snapshot( issued_at: cold.issued_at, expires_at: slot.expires_at, label: cold.label.clone(), + tombstoned_at: (cold.tombstoned_at != 0).then_some(cold.tombstoned_at), }) }) .collect(); @@ -222,8 +223,9 @@ pub(super) fn apply_persisted_mutation( })?; entry.expires_at = expires_at; entry.state = SlotState::Active; + entry.tombstoned_at = None; } - StateMutation::Revoke { key_id, .. } => { + StateMutation::Revoke { key_id, at } => { let entry = snapshot .entries .iter_mut() @@ -236,6 +238,7 @@ pub(super) fn apply_persisted_mutation( ) })?; entry.state = SlotState::Revoked; + entry.tombstoned_at = Some(at); } StateMutation::LegacyProtocol(policy) => snapshot.legacy_protocol = policy, } diff --git a/src/common/auth/runtime.rs b/src/common/auth/runtime.rs index 2b7f274..da8c0cc 100644 --- a/src/common/auth/runtime.rs +++ b/src/common/auth/runtime.rs @@ -68,6 +68,11 @@ impl AuthRuntime { ColdMetadata { issued_at: entry.issued_at, label: entry.label.clone(), + tombstoned_at: match state { + SlotState::Expired => entry.tombstoned_at.unwrap_or(entry.expires_at), + SlotState::Revoked => entry.tombstoned_at.unwrap_or(0), + SlotState::Free | SlotState::Active => 0, + }, }, ); if state == SlotState::Active { diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index 2741087..1a719d1 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -105,6 +105,15 @@ async fn issue_renew_revoke_and_persist() { .code, "temporary_key_revoked" ); + let instance_id = load_or_create_instance_id(&state_dir).unwrap(); + let persisted = try_load_persisted_state(&config, &admin_key, instance_id).unwrap(); + let revoked = persisted + .entries + .iter() + .find(|entry| entry.key_id == issued.metadata.key_id) + .unwrap(); + assert_eq!(revoked.state, SlotState::Revoked); + assert!(revoked.tombstoned_at.is_some()); drop(runtime); tokio::time::sleep(Duration::from_millis(20)).await; From 5d485f301d2e058e5c51d7fb9114cf66fd93feab Mon Sep 17 00:00:00 2001 From: LB7666 Date: Tue, 18 Aug 2026 06:29:48 +0800 Subject: [PATCH 12/74] Remove implicit credential authentication --- src/common/auth/runtime.rs | 9 ++------ src/common/auth/tests.rs | 34 ++++++++++++++++-------------- src/common/message/secure.rs | 14 ++++++------ src/common/message/secure/tests.rs | 2 +- tests/regression.rs | 12 ++++++++--- 5 files changed, 37 insertions(+), 34 deletions(-) diff --git a/src/common/auth/runtime.rs b/src/common/auth/runtime.rs index da8c0cc..ee3e787 100644 --- a/src/common/auth/runtime.rs +++ b/src/common/auth/runtime.rs @@ -161,11 +161,11 @@ impl AuthRuntime { }) } - pub fn admin_key(&self) -> Result { + pub(crate) fn admin_key(&self) -> Result { Ok(self.inner()?.admin_key()) } - pub fn derive_key(&self, key_id: u64) -> Result { + pub(crate) fn derive_key(&self, key_id: u64) -> Result { let inner = self.inner()?; if key_id == 0 { return Ok(inner.admin_key()); @@ -173,11 +173,6 @@ impl AuthRuntime { derive_temporary_key(&inner.admin_key(), &inner.instance_id(), key_id) } - pub fn authenticate(&self, key_id: u64) -> Result { - let presented_key = self.derive_key(key_id)?; - self.authenticate_presented(key_id, &presented_key) - } - pub fn authenticate_presented( &self, key_id: u64, diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index 1a719d1..543d566 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -21,6 +21,11 @@ fn temp_state_dir(name: &str) -> PathBuf { std::env::temp_dir().join(format!("pb-mapper-{name}-{}", hex(&suffix))) } +fn authenticate_for_test(runtime: &AuthRuntime, key_id: u64) -> Result { + let key = runtime.derive_key(key_id)?; + runtime.authenticate_presented(key_id, &key) +} + #[test] fn legacy_protocol_policy_trims_valid_values_and_rejects_unknown_values() { assert_eq!( @@ -73,13 +78,13 @@ async fn issue_renew_revoke_and_persist() { legacy_protocol: LegacyProtocolPolicy::Allow, }; let runtime = AuthRuntime::start(admin_key, config.clone()).await.unwrap(); - let admin = runtime.authenticate(0).unwrap(); + let admin = authenticate_for_test(&runtime, 0).unwrap(); let issued = runtime .issue(&admin, Duration::from_secs(60), Some("demo".to_string())) .await .unwrap(); assert!(issued.credential.starts_with("pbmt1_")); - let context = runtime.authenticate(issued.metadata.key_id).unwrap(); + let context = authenticate_for_test(&runtime, issued.metadata.key_id).unwrap(); assert!(!context.is_admin); let cancellation = context.cancellation_token().unwrap(); let renewed = runtime @@ -99,8 +104,7 @@ async fn issue_renew_revoke_and_persist() { "temporary_key_revoked" ); assert_eq!( - runtime - .authenticate(issued.metadata.key_id) + authenticate_for_test(&runtime, issued.metadata.key_id) .unwrap_err() .code, "temporary_key_revoked" @@ -119,8 +123,7 @@ async fn issue_renew_revoke_and_persist() { tokio::time::sleep(Duration::from_millis(20)).await; let restored = AuthRuntime::start(admin_key, config).await.unwrap(); assert_eq!( - restored - .authenticate(issued.metadata.key_id) + authenticate_for_test(&restored, issued.metadata.key_id) .unwrap_err() .code, "temporary_key_revoked" @@ -139,7 +142,7 @@ async fn reset_rotates_instance_and_prevents_old_key_id_reuse() { legacy_protocol: LegacyProtocolPolicy::Allow, }; let runtime = AuthRuntime::start(admin_key, config).await.unwrap(); - let admin = runtime.authenticate(0).unwrap(); + let admin = authenticate_for_test(&runtime, 0).unwrap(); let before = runtime.status(&admin).await.unwrap().server_instance_id; let old = runtime .issue( @@ -149,7 +152,7 @@ async fn reset_rotates_instance_and_prevents_old_key_id_reuse() { ) .await .unwrap(); - let old_context = runtime.authenticate(old.metadata.key_id).unwrap(); + let old_context = authenticate_for_test(&runtime, old.metadata.key_id).unwrap(); let old_cancellation = old_context.cancellation_token().unwrap(); runtime.reset(&admin).await.unwrap(); @@ -157,7 +160,7 @@ async fn reset_rotates_instance_and_prevents_old_key_id_reuse() { let after = runtime.status(&admin).await.unwrap().server_instance_id; assert_ne!(after, before); assert!(old_cancellation.is_cancelled()); - assert!(runtime.authenticate(old.metadata.key_id).is_err()); + assert!(authenticate_for_test(&runtime, old.metadata.key_id).is_err()); let replacement = runtime .issue( &admin, @@ -185,7 +188,7 @@ async fn corrupt_wal_fails_temporary_keys_closed_until_admin_reset() { legacy_protocol: LegacyProtocolPolicy::Allow, }; let runtime = AuthRuntime::start(admin_key, config.clone()).await.unwrap(); - let admin = runtime.authenticate(0).unwrap(); + let admin = authenticate_for_test(&runtime, 0).unwrap(); let issued = runtime .issue( &admin, @@ -199,11 +202,10 @@ async fn corrupt_wal_fails_temporary_keys_closed_until_admin_reset() { std::fs::write(state_dir.join("auth.wal"), b"broken-wal").unwrap(); let recovered = AuthRuntime::start(admin_key, config).await.unwrap(); - let recovered_admin = recovered.authenticate(0).unwrap(); + let recovered_admin = authenticate_for_test(&recovered, 0).unwrap(); assert!(recovered.status(&recovered_admin).await.unwrap().safe_mode); assert_eq!( - recovered - .authenticate(issued.metadata.key_id) + authenticate_for_test(&recovered, issued.metadata.key_id) .unwrap_err() .code, "temporary_key_store_unavailable" @@ -279,7 +281,7 @@ async fn admitted_admin_mutation_replay_survives_restart() { let fingerprint = [0x5a; 32]; let timestamp = unix_seconds(); let runtime = AuthRuntime::start(admin_key, config.clone()).await.unwrap(); - let admin = runtime.authenticate(0).unwrap(); + let admin = authenticate_for_test(&runtime, 0).unwrap(); runtime .claim_admin_mutation(&admin, fingerprint, timestamp) .await @@ -288,7 +290,7 @@ async fn admitted_admin_mutation_replay_survives_restart() { tokio::time::sleep(Duration::from_millis(20)).await; let restored = AuthRuntime::start(admin_key, config).await.unwrap(); - let restored_admin = restored.authenticate(0).unwrap(); + let restored_admin = authenticate_for_test(&restored, 0).unwrap(); assert_eq!( restored .claim_admin_mutation(&restored_admin, fingerprint, timestamp) @@ -314,7 +316,7 @@ async fn snapshot_compaction_preserves_audit_records() { legacy_protocol: LegacyProtocolPolicy::Allow, }; let runtime = AuthRuntime::start(admin_key, config.clone()).await.unwrap(); - let admin = runtime.authenticate(0).unwrap(); + let admin = authenticate_for_test(&runtime, 0).unwrap(); runtime .issue(&admin, Duration::from_secs(60), Some("audited".to_string())) .await diff --git a/src/common/message/secure.rs b/src/common/message/secure.rs index 8f84031..41c52a1 100644 --- a/src/common/message/secure.rs +++ b/src/common/message/secure.rs @@ -359,13 +359,6 @@ impl ServerSecurity { failure, response_session: None, })?; - let context = self - .auth - .authenticate(0) - .map_err(|failure| ServerInitialError { - failure, - response_session: None, - })?; let checksum = u32::from_be_bytes(checksum_bytes); let datalen = reader .read_u32() @@ -418,6 +411,13 @@ impl ServerSecurity { ), response_session: None, })?; + let context = self + .auth + .authenticate_presented(0, &key) + .map_err(|failure| ServerInitialError { + failure, + response_session: None, + })?; let legacy_guard = self.auth .record_legacy_connection() diff --git a/src/common/message/secure/tests.rs b/src/common/message/secure/tests.rs index 409a06b..eecc18c 100644 --- a/src/common/message/secure/tests.rs +++ b/src/common/message/secure/tests.rs @@ -62,7 +62,7 @@ async fn temporary_credential_authenticates_without_storing_secret() { let admin = *b"0123456789abcdefghijklmnopqrstuv"; let config = temp_config(); let auth = AuthRuntime::start(admin, config.clone()).await.unwrap(); - let admin_context = auth.authenticate(0).unwrap(); + let admin_context = auth.authenticate_presented(0, &admin).unwrap(); let issued = auth .issue(&admin_context, std::time::Duration::from_secs(60), None) .await diff --git a/tests/regression.rs b/tests/regression.rs index 23a2cff..bf066c3 100644 --- a/tests/regression.rs +++ b/tests/regression.rs @@ -107,7 +107,9 @@ async fn admin_all_preserves_json_output_mode() { ) .await .unwrap(); - let admin = runtime.authenticate(0).unwrap(); + let admin = runtime + .authenticate_presented(0, TEST_ADMIN_KEY.as_bytes().first_chunk::<32>().unwrap()) + .unwrap(); runtime .issue( &admin, @@ -313,7 +315,9 @@ async fn temporary_credentials_are_isolated_denied_admin_and_revoked_live() { ) .await .unwrap(); - let admin = runtime.authenticate(0).unwrap(); + let admin = runtime + .authenticate_presented(0, TEST_ADMIN_KEY.as_bytes().first_chunk::<32>().unwrap()) + .unwrap(); let first = runtime .issue(&admin, Duration::from_secs(120), Some("first".to_string())) .await @@ -442,7 +446,9 @@ async fn revoking_subscriber_credential_closes_cross_credential_data_stream() { ) .await .unwrap(); - let admin = runtime.authenticate(0).unwrap(); + let admin = runtime + .authenticate_presented(0, TEST_ADMIN_KEY.as_bytes().first_chunk::<32>().unwrap()) + .unwrap(); let issued = runtime .issue( &admin, From 86c2a1cc1963bc11cc1eb53a5cbd63eeb2d93fdc Mon Sep 17 00:00:00 2001 From: LB7666 Date: Tue, 18 Aug 2026 06:34:06 +0800 Subject: [PATCH 13/74] Defend against conflicting server key modes --- src/bin/pb-mapper.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/bin/pb-mapper.rs b/src/bin/pb-mapper.rs index b73cd02..b6c0616 100644 --- a/src/bin/pb-mapper.rs +++ b/src/bin/pb-mapper.rs @@ -245,8 +245,7 @@ async fn run_server(args: ServerArgs) -> Result<(), Box> { let key = initialize_admin_key(&key_path, args.force_init_admin_key)?; set_process_msg_header_key(Some(&key))?; eprintln!("administrator key initialized at {}", key_path.display()); - } - if args.use_machine_msg_header_key { + } else if args.use_machine_msg_header_key { tracing::warn!( "--use-machine-msg-header-key is a legacy compatibility option; prefer a random administrator key" ); From a3b1c4b457bb856d7aa51b7ba09f2de2e40e79e4 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Tue, 18 Aug 2026 06:40:25 +0800 Subject: [PATCH 14/74] Migrate missing tombstone timestamps safely --- src/common/auth/actor.rs | 3 ++- src/common/auth/persistence.rs | 29 ++++++++++++++++++++++- src/common/auth/runtime.rs | 13 +++++++---- src/common/auth/tests.rs | 42 ++++++++++++++++++++++++++++++++++ 4 files changed, 81 insertions(+), 6 deletions(-) diff --git a/src/common/auth/actor.rs b/src/common/auth/actor.rs index 324256a..8f8a758 100644 --- a/src/common/auth/actor.rs +++ b/src/common/auth/actor.rs @@ -54,6 +54,7 @@ pub(super) async fn run_auth_actor( mut admin_replays, mut admin_replay_order, } = state; + let now = unix_seconds(); let mut tombstones = inner .slots .read() @@ -68,7 +69,7 @@ pub(super) async fn run_auth_actor( let tombstoned_at = cold .get(&key_id) .map(|metadata| metadata.tombstoned_at) - .unwrap_or(0); + .unwrap_or(now); Some(( tombstoned_at.saturating_add(TOMBSTONE_RETENTION.as_secs()), key_id, diff --git a/src/common/auth/persistence.rs b/src/common/auth/persistence.rs index f4a737e..198f792 100644 --- a/src/common/auth/persistence.rs +++ b/src/common/auth/persistence.rs @@ -1,7 +1,7 @@ //! Durable, encrypted authentication state and audit/replay retention. //! //! ```text -//! startup: admin.key -> decrypt snapshot -> replay WAL -> in-memory state +//! startup: admin.key -> decrypt snapshot -> replay WAL -> normalize -> in-memory state //! mutation: command -> fsync encrypted WAL -> publish hot-state change //! compact: hot state + audit + replay set -> snapshot -> truncate WAL //! ``` @@ -81,6 +81,33 @@ pub(super) fn build_snapshot( } } +pub(super) fn normalize_tombstone_times(snapshot: &mut PersistedSnapshot, now: u64) -> bool { + let mut changed = false; + for entry in &mut snapshot.entries { + if entry.tombstoned_at.is_some() { + continue; + } + let tombstoned_at = match entry.state { + SlotState::Expired => Some(entry.expires_at), + SlotState::Revoked => snapshot + .audit_records + .iter() + .rev() + .find(|record| { + record.action == "temporary_key_revoke" && record.key_id == Some(entry.key_id) + }) + .map(|record| record.at) + .or(Some(now)), + SlotState::Free | SlotState::Active => None, + }; + if tombstoned_at.is_some() { + entry.tombstoned_at = tombstoned_at; + changed = true; + } + } + changed +} + pub(super) fn empty_snapshot( inner: &AuthStateInner, instance_id: [u8; INSTANCE_ID_LEN], diff --git a/src/common/auth/runtime.rs b/src/common/auth/runtime.rs index ee3e787..66a52f9 100644 --- a/src/common/auth/runtime.rs +++ b/src/common/auth/runtime.rs @@ -32,14 +32,19 @@ impl AuthRuntime { pub async fn start(admin_key: AesKeyType, config: AuthConfig) -> Result { prepare_state_dir(&config.state_dir)?; let instance_id = load_or_create_instance_id(&config.state_dir)?; - let (loaded, safe_mode) = load_persisted_state(&config, &admin_key, instance_id); + let (mut loaded, safe_mode) = load_persisted_state(&config, &admin_key, instance_id); + let now = unix_seconds(); + if let Some(state) = loaded.as_mut() { + if normalize_tombstone_times(state, now) { + write_snapshot_and_truncate_wal(&config, &admin_key, state)?; + } + } let mut slots = (0..config.max_temporary_keys) .map(|_| SlotHot::default()) .collect::>() .into_boxed_slice(); let mut cold = HashMap::new(); - let mut wheel = TimingWheel::new(unix_seconds()); - let now = unix_seconds(); + let mut wheel = TimingWheel::new(now); let admin_lease = Arc::new(AuthLease::new(0, u64::MAX)); if let Some(state) = loaded.as_ref() { @@ -70,7 +75,7 @@ impl AuthRuntime { label: entry.label.clone(), tombstoned_at: match state { SlotState::Expired => entry.tombstoned_at.unwrap_or(entry.expires_at), - SlotState::Revoked => entry.tombstoned_at.unwrap_or(0), + SlotState::Revoked => entry.tombstoned_at.unwrap_or(now), SlotState::Free | SlotState::Active => 0, }, }, diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index 543d566..7772208 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -425,3 +425,45 @@ fn replay_pruning_removes_only_records_outside_the_retention_window() { assert_eq!(replay_order.len(), 1); assert_eq!(replay_order[0].fingerprint, current.fingerprint); } + +#[test] +fn tombstone_migration_prefers_audit_time_and_persists_fail_closed_fallback() { + let now = 10_000; + let revoked_with_audit = PersistedEntry { + key_id: make_key_id(1, 0), + state: SlotState::Revoked, + issued_at: 100, + expires_at: 20_000, + label: None, + tombstoned_at: None, + }; + let revoked_without_audit = PersistedEntry { + key_id: make_key_id(1, 1), + state: SlotState::Revoked, + issued_at: 100, + expires_at: 20_000, + label: None, + tombstoned_at: None, + }; + let audit_at = now - 30; + let mut snapshot = PersistedSnapshot { + schema_version: SNAPSHOT_SCHEMA_VERSION, + instance_id: [1; INSTANCE_ID_LEN], + generations: vec![1, 1], + entries: vec![revoked_with_audit, revoked_without_audit], + legacy_protocol: LegacyProtocolPolicy::Deny, + admin_replays: Vec::new(), + audit_records: VecDeque::from([AuditRecord { + at: audit_at, + action: "temporary_key_revoke".to_string(), + key_id: Some(make_key_id(1, 0)), + label: None, + }]), + }; + + assert!(normalize_tombstone_times(&mut snapshot, now)); + assert_eq!(snapshot.entries[0].tombstoned_at, Some(audit_at)); + assert_eq!(snapshot.entries[1].tombstoned_at, Some(now)); + assert!(!normalize_tombstone_times(&mut snapshot, now + 1)); + assert_eq!(snapshot.entries[1].tombstoned_at, Some(now)); +} From 7466b4cfc5e3a38d202e7ca6c98f4b00ec0fd6b8 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Tue, 18 Aug 2026 13:49:05 +0800 Subject: [PATCH 15/74] Fix authentication review findings --- scripts/release/entrypoint/pb-mapper.sh | 2 +- src/bin/pb-mapper.rs | 106 ++++++--- src/bin/pb-mapper/admin.rs | 70 +++++- src/common/auth.rs | 153 ++++++++----- src/common/auth/actor.rs | 4 +- src/common/auth/runtime.rs | 28 ++- src/common/auth/tests.rs | 49 +++++ src/pb_server/admin.rs | 220 ++++++++++++++++--- tests/regression.rs | 34 ++- ui/native/pb_mapper_ffi/src/state/runtime.rs | 2 +- 10 files changed, 528 insertions(+), 140 deletions(-) diff --git a/scripts/release/entrypoint/pb-mapper.sh b/scripts/release/entrypoint/pb-mapper.sh index 54c3399..15124f7 100644 --- a/scripts/release/entrypoint/pb-mapper.sh +++ b/scripts/release/entrypoint/pb-mapper.sh @@ -15,7 +15,7 @@ ADMIN_KEY_PATH="$AUTH_DIR/admin.key" LEGACY_KEY_PATH="/var/lib/pb-mapper-server/msg_header_key" install -d -m 0700 "$AUTH_DIR" -if [ ! -s "$ADMIN_KEY_PATH" ] && [ -s "$LEGACY_KEY_PATH" ]; then +if [ -z "${MSG_HEADER_KEY:-}" ] && [ ! -s "$ADMIN_KEY_PATH" ] && [ -s "$LEGACY_KEY_PATH" ]; then install -m 0600 "$LEGACY_KEY_PATH" "$ADMIN_KEY_PATH" echo "Migrated the legacy machine-derived key into $ADMIN_KEY_PATH" fi diff --git a/src/bin/pb-mapper.rs b/src/bin/pb-mapper.rs index b6c0616..0c9a3dd 100644 --- a/src/bin/pb-mapper.rs +++ b/src/bin/pb-mapper.rs @@ -19,13 +19,14 @@ use std::time::Duration; use better_mimalloc_rs::MiMalloc; use clap::{Args, Parser, Subcommand, ValueEnum}; use pb_mapper::common::auth::{ - generate_admin_key, initialize_admin_key, write_admin_key_file, KeyPage, LegacyProtocolPolicy, - DEFAULT_AUTH_STATE_DIR, + generate_admin_key, initialize_admin_key, write_admin_key_file, AuthConfig, KeyPage, + LegacyProtocolPolicy, }; use pb_mapper::common::checksum::set_process_msg_header_key; use pb_mapper::common::checksum::{setup_machine_msg_header_key, MACHINE_MSG_HEADER_KEY_PATH}; use pb_mapper::common::config::{ - get_pb_mapper_server_async, get_sockaddr_async, init_tracing, keep_alive_from_env, StatusOp, + control_io_timeout, get_pb_mapper_server_async, get_sockaddr_async, init_tracing, + keep_alive_from_env, StatusOp, }; use pb_mapper::common::message::command::{ AdminConnectionPage, AdminRequest, AdminResponse, AdminServicePage, MessageSerializer, @@ -88,8 +89,8 @@ struct ServerArgs { #[arg(long, default_value_t = false)] use_machine_msg_header_key: bool, /// Directory containing encrypted authentication state and the administrator key file. - #[arg(long, default_value = DEFAULT_AUTH_STATE_DIR)] - auth_state_dir: PathBuf, + #[arg(long)] + auth_state_dir: Option, /// Create a random administrator key before starting the relay. #[arg( long, @@ -101,14 +102,14 @@ struct ServerArgs { #[arg(long, requires = "init_admin_key", default_value_t = false)] force_init_admin_key: bool, /// Maximum temporary-key slots allocated by the relay. - #[arg(long, default_value_t = 65_536)] - max_temporary_keys: usize, + #[arg(long)] + max_temporary_keys: Option, /// Maximum accepted temporary-key TTL. - #[arg(long, default_value = "30d", value_parser = parse_duration)] - max_temporary_key_ttl: Duration, + #[arg(long, value_parser = parse_duration)] + max_temporary_key_ttl: Option, /// Allow or deny the legacy encrypted framing protocol. - #[arg(long, value_enum, default_value_t = LegacyProtocolArg::Allow)] - legacy_protocol: LegacyProtocolArg, + #[arg(long, value_enum)] + legacy_protocol: Option, } #[path = "pb-mapper/admin.rs"] @@ -224,24 +225,33 @@ async fn run(cli: Cli) -> Result<(), Box> { } async fn run_server(args: ServerArgs) -> Result<(), Box> { - std::env::set_var("PB_MAPPER_AUTH_STATE_DIR", &args.auth_state_dir); - std::env::set_var( - "PB_MAPPER_AUTH_MAX_TEMP_KEYS", - args.max_temporary_keys.to_string(), - ); - std::env::set_var( - "PB_MAPPER_AUTH_MAX_TEMP_TTL_SECS", - args.max_temporary_key_ttl.as_secs().to_string(), - ); - std::env::set_var( - "PB_MAPPER_LEGACY_PROTOCOL", - match args.legacy_protocol { - LegacyProtocolArg::Allow => "allow", - LegacyProtocolArg::Deny => "deny", - }, - ); + if let Some(auth_state_dir) = &args.auth_state_dir { + std::env::set_var("PB_MAPPER_AUTH_STATE_DIR", auth_state_dir); + } + if let Some(max_temporary_keys) = args.max_temporary_keys { + std::env::set_var( + "PB_MAPPER_AUTH_MAX_TEMP_KEYS", + max_temporary_keys.to_string(), + ); + } + if let Some(max_temporary_key_ttl) = args.max_temporary_key_ttl { + std::env::set_var( + "PB_MAPPER_AUTH_MAX_TEMP_TTL_SECS", + max_temporary_key_ttl.as_secs().to_string(), + ); + } + if let Some(legacy_protocol) = args.legacy_protocol { + std::env::set_var( + "PB_MAPPER_LEGACY_PROTOCOL", + match legacy_protocol { + LegacyProtocolArg::Allow => "allow", + LegacyProtocolArg::Deny => "deny", + }, + ); + } + let auth_config = AuthConfig::default(); if args.init_admin_key { - let key_path = args.auth_state_dir.join("admin.key"); + let key_path = auth_config.state_dir.join("admin.key"); let key = initialize_admin_key(&key_path, args.force_init_admin_key)?; set_process_msg_header_key(Some(&key))?; eprintln!("administrator key initialized at {}", key_path.display()); @@ -495,4 +505,44 @@ mod tests { ]) .is_err()); } + + #[test] + fn server_auth_options_only_override_environment_when_explicit() { + let cli = + Cli::try_parse_from(["pb-mapper", "server"]).expect("server defaults should parse"); + let Command::Server(defaults) = cli.command else { + panic!("expected server command"); + }; + assert_eq!(defaults.auth_state_dir, None); + assert_eq!(defaults.max_temporary_keys, None); + assert_eq!(defaults.max_temporary_key_ttl, None); + assert_eq!(defaults.legacy_protocol, None); + + let cli = Cli::try_parse_from([ + "pb-mapper", + "server", + "--auth-state-dir", + "/tmp/pb-mapper-auth", + "--max-temporary-keys", + "1024", + "--max-temporary-key-ttl", + "2h", + "--legacy-protocol", + "deny", + ]) + .expect("explicit server authentication options should parse"); + let Command::Server(explicit) = cli.command else { + panic!("expected server command"); + }; + assert_eq!( + explicit.auth_state_dir, + Some(PathBuf::from("/tmp/pb-mapper-auth")) + ); + assert_eq!(explicit.max_temporary_keys, Some(1024)); + assert_eq!( + explicit.max_temporary_key_ttl, + Some(Duration::from_secs(2 * 60 * 60)) + ); + assert_eq!(explicit.legacy_protocol, Some(LegacyProtocolArg::Deny)); + } } diff --git a/src/bin/pb-mapper/admin.rs b/src/bin/pb-mapper/admin.rs index 10450b4..2fa74c1 100644 --- a/src/bin/pb-mapper/admin.rs +++ b/src/bin/pb-mapper/admin.rs @@ -313,15 +313,38 @@ pub(super) async fn run_admin(args: AdminArgs) -> Result<(), Box> { async fn send_admin_request( remote_addr: std::net::SocketAddr, request: AdminRequest, +) -> Result> { + send_admin_request_with_timeout(remote_addr, request, control_io_timeout()).await +} + +async fn send_admin_request_with_timeout( + remote_addr: std::net::SocketAddr, + request: AdminRequest, + io_timeout: Duration, ) -> Result> { let encoded = PbConnRequest::Admin(request).encode()?; for attempt in 0..2 { - let mut stream = TcpStream::connect(remote_addr).await?; - let session = ClientHeaderSession::from_process()?; - session.write_initial(&mut stream, &encoded).await?; - let mut reader = session.response_reader(&mut stream)?; - let message = reader.read_msg().await?; - match PbConnResponse::decode(message)? { + let response = tokio::time::timeout(io_timeout, async { + let mut stream = TcpStream::connect(remote_addr) + .await + .map_err(|error| -> Box { Box::new(error) })?; + let session = ClientHeaderSession::from_process()?; + session.write_initial(&mut stream, &encoded).await?; + let mut reader = session.response_reader(&mut stream)?; + let message = reader.read_msg().await?; + Ok::<_, Box>(PbConnResponse::decode(message)?) + }) + .await + .map_err(|_| { + std::io::Error::new( + std::io::ErrorKind::TimedOut, + format!( + "administrator request attempt timed out after {} ms", + io_timeout.as_millis() + ), + ) + })??; + match response { PbConnResponse::Admin(response) => return Ok(response), PbConnResponse::Error(error) if error.code == "connection_salt_replayed" && error.retryable && attempt == 0 => @@ -612,3 +635,38 @@ fn default_admin_recovery_key_file() -> PathBuf { .join("pb-mapper") .join("admin.key") } + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn administrator_request_times_out_when_peer_stalls() { + set_process_msg_header_key(Some("0123456789abcdefghijklmnopqrstuv")) + .expect("test administrator credential should be valid"); + let listener = tokio::net::TcpListener::bind((std::net::Ipv4Addr::LOCALHOST, 0)) + .await + .expect("test listener should bind"); + let remote_addr = listener + .local_addr() + .expect("listener should have an address"); + let stalled_peer = tokio::spawn(async move { + let (_stream, _) = listener.accept().await.expect("test peer should connect"); + std::future::pending::<()>().await; + }); + + let error = send_admin_request_with_timeout( + remote_addr, + AdminRequest::AuthStatus, + Duration::from_millis(50), + ) + .await + .expect_err("a stalled administrator request should time out"); + + let io_error = error + .downcast_ref::() + .expect("timeout should be reported as an I/O error"); + assert_eq!(io_error.kind(), std::io::ErrorKind::TimedOut); + stalled_peer.abort(); + } +} diff --git a/src/common/auth.rs b/src/common/auth.rs index 4d0e29d..a4f1303 100644 --- a/src/common/auth.rs +++ b/src/common/auth.rs @@ -57,6 +57,10 @@ const ADMIN_REPLAY_RETENTION: Duration = Duration::from_secs(10 * 60); const ADMIN_REPLAY_CAPACITY: usize = 65_536; const AUDIT_RECORD_CAPACITY: usize = 4096; +#[cfg(test)] +pub(crate) static PROCESS_CREDENTIAL_TEST_LOCK: std::sync::LazyLock> = + std::sync::LazyLock::new(|| tokio::sync::Mutex::new(())); + #[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum LegacyProtocolPolicy { @@ -289,6 +293,17 @@ impl AuthContext { Ok(self.ensure_active()?.cancellation_token()) } + pub(crate) fn admin_cancellation_token(&self) -> Result { + if !self.is_admin { + return Err(AuthFailure::new( + "admin_permission_required", + "administrator credential is required for this operation", + false, + )); + } + self.cancellation_token() + } + fn admin_authority(&self) -> Result, AuthFailure> { if !self.is_admin { return Err(AuthFailure::new( @@ -339,6 +354,7 @@ struct AdminState { #[derive(Debug)] struct AuthStateInner { admin: RwLock, + sync_process_credential: bool, instance_id: RwLock<[u8; INSTANCE_ID_LEN]>, slots: RwLock>, safe_mode: AtomicBool, @@ -502,51 +518,71 @@ impl Drop for LegacyConnectionGuard { } } -fn load_server_admin_credential(state_dir: &Path) -> Result { - let path = state_dir.join("admin.key"); - let raw = if path.exists() { - #[cfg(unix)] - { - let metadata = std::fs::metadata(&path).map_err(|error| { - AuthFailure::new( - "administrator_key_required", - format!( - "administrator key file `{}` metadata could not be read: {error}", - path.display() - ), - false, - ) - })?; - if metadata.permissions().mode() & 0o077 != 0 { - std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600)).map_err( - |error| { - AuthFailure::new( - "administrator_key_required", - format!( - "administrator key file `{}` permissions could not be secured: {error}", - path.display() - ), - false, - ) - }, - )?; - tracing::warn!( - event = "administrator_key_permissions_repaired", - path = %path.display(), - "restricted administrator key file permissions to 0600" - ); - } - } - std::fs::read_to_string(&path).map_err(|error| { +fn read_admin_key(path: &Path) -> Result, AuthFailure> { + if !path.exists() { + return Ok(None); + } + #[cfg(unix)] + { + let metadata = std::fs::metadata(path).map_err(|error| { AuthFailure::new( "administrator_key_required", format!( - "administrator key file `{}` could not be read: {error}", + "administrator key file `{}` metadata could not be read: {error}", path.display() ), false, ) - })? + })?; + if metadata.permissions().mode() & 0o077 != 0 { + std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600)).map_err( + |error| { + AuthFailure::new( + "administrator_key_required", + format!( + "administrator key file `{}` permissions could not be secured: {error}", + path.display() + ), + false, + ) + }, + )?; + tracing::warn!( + event = "administrator_key_permissions_repaired", + path = %path.display(), + "restricted administrator key file permissions to 0600" + ); + } + } + std::fs::read_to_string(path).map(Some).map_err(|error| { + AuthFailure::new( + "administrator_key_required", + format!( + "administrator key file `{}` could not be read: {error}", + path.display() + ), + false, + ) + }) +} + +fn validate_admin_credential(raw: &str) -> Result { + let credential = parse_credential(raw.trim()) + .map_err(|error| AuthFailure::new("administrator_key_invalid", error, false))?; + if !credential.is_admin() { + return Err(AuthFailure::new( + "administrator_key_required", + "the server key file contains a temporary credential", + false, + )); + } + Ok(credential) +} + +fn load_server_admin_credential(state_dir: &Path) -> Result { + let path = state_dir.join("admin.key"); + let raw = if let Some(raw) = read_admin_key(&path)? { + raw } else if std::env::var_os(ENV_MSG_HEADER_KEY).is_some() { let credential = get_process_credential() .map_err(|error| AuthFailure::new("administrator_key_invalid", error, false))?; @@ -576,15 +612,7 @@ fn load_server_admin_credential(state_dir: &Path) -> Result Result Result { + let path = state_dir.join("admin.key"); + let raw = match read_admin_key(&path)? { + Some(raw) => raw, + None => { + let key = initialize_admin_key(&path, false)?; + tracing::warn!( + event = "isolated_administrator_key_initialized", + path = %path.display(), + "generated an administrator key for an embedded relay" + ); + key + } + }; + validate_admin_credential(&raw) +} + pub fn make_key_id(generation: u32, slot: u32) -> u64 { (u64::from(generation) << 32) | u64::from(slot) } diff --git a/src/common/auth/actor.rs b/src/common/auth/actor.rs index 8f8a758..1cf9fc2 100644 --- a/src/common/auth/actor.rs +++ b/src/common/auth/actor.rs @@ -751,7 +751,9 @@ fn actor_rotate_root( key: new_key, lease: Arc::downgrade(&new_admin_lease), }; - set_process_msg_header_key(Some(&new_key_string)).map_err(AuthFailure::internal)?; + if inner.sync_process_credential { + set_process_msg_header_key(Some(&new_key_string)).map_err(AuthFailure::internal)?; + } inner.safe_mode.store(false, Ordering::Release); old_admin_lease.cancellation.cancel(); *admin_lease = new_admin_lease; diff --git a/src/common/auth/runtime.rs b/src/common/auth/runtime.rs index 66a52f9..9680680 100644 --- a/src/common/auth/runtime.rs +++ b/src/common/auth/runtime.rs @@ -26,10 +26,35 @@ impl AuthRuntime { false, )); }; - Self::start(admin_key, config).await + Self::start_with_process_sync(admin_key, config, true).await + } + + /// Start an embedded relay with an administrator key owned only by its state directory. + /// + /// This deliberately leaves the process credential untouched because the containing UI uses + /// that credential for its outbound register, connect, status, and stream connections. + pub async fn from_isolated_state(config: AuthConfig) -> Result { + prepare_state_dir(&config.state_dir)?; + let credential = load_isolated_server_admin_credential(&config.state_dir)?; + let Credential::Admin(admin_key) = credential else { + return Err(AuthFailure::new( + "administrator_key_required", + "the embedded relay must start with an administrator credential", + false, + )); + }; + Self::start_with_process_sync(admin_key, config, false).await } pub async fn start(admin_key: AesKeyType, config: AuthConfig) -> Result { + Self::start_with_process_sync(admin_key, config, true).await + } + + async fn start_with_process_sync( + admin_key: AesKeyType, + config: AuthConfig, + sync_process_credential: bool, + ) -> Result { prepare_state_dir(&config.state_dir)?; let instance_id = load_or_create_instance_id(&config.state_dir)?; let (mut loaded, safe_mode) = load_persisted_state(&config, &admin_key, instance_id); @@ -125,6 +150,7 @@ impl AuthRuntime { key: admin_key, lease: Arc::downgrade(&admin_lease), }), + sync_process_credential, instance_id: RwLock::new(instance_id), slots: RwLock::new(slots), safe_mode: AtomicBool::new(safe_mode), diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index 7772208..6fe5835 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -67,6 +67,54 @@ fn derived_key_is_bound_to_instance_and_key_id() { ); } +#[tokio::test] +async fn isolated_runtime_preserves_remote_temporary_process_credential() { + let _process_credential_guard = PROCESS_CREDENTIAL_TEST_LOCK.lock().await; + let state_dir = temp_state_dir("isolated-relay"); + let temporary_key_id = make_key_id(1, 0); + let temporary_key = *b"temporary-remote-key-0123456789a"; + let temporary_credential = encode_temporary_credential(temporary_key_id, &temporary_key); + set_process_msg_header_key(Some(&temporary_credential)).unwrap(); + let config = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 4, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + + let runtime = AuthRuntime::from_isolated_state(config).await.unwrap(); + assert_eq!( + get_process_credential().unwrap(), + Credential::Temporary { + key_id: temporary_key_id, + key: temporary_key, + } + ); + + let local_admin_raw = std::fs::read_to_string(state_dir.join("admin.key")).unwrap(); + let Credential::Admin(local_admin_key) = parse_credential(local_admin_raw.trim()).unwrap() + else { + panic!("isolated relay key should be an administrator credential"); + }; + let local_admin = runtime.authenticate_presented(0, &local_admin_key).unwrap(); + runtime + .rotate_root(&local_admin, *b"isolated-new-admin-key-012345678") + .await + .unwrap(); + assert_eq!( + get_process_credential().unwrap(), + Credential::Temporary { + key_id: temporary_key_id, + key: temporary_key, + } + ); + + set_process_msg_header_key(None).unwrap(); + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + let _ = std::fs::remove_dir_all(state_dir); +} + #[tokio::test] async fn issue_renew_revoke_and_persist() { let state_dir = temp_state_dir("auth-lifecycle"); @@ -220,6 +268,7 @@ async fn corrupt_wal_fails_temporary_keys_closed_until_admin_reset() { #[tokio::test] async fn root_rotation_rejects_old_key_and_in_flight_admin_context() { + let _process_credential_guard = PROCESS_CREDENTIAL_TEST_LOCK.lock().await; let state_dir = temp_state_dir("auth-root-rotation"); let old_key = *b"0123456789abcdefghijklmnopqrstuv"; let new_key = *b"abcdefghijklmnopqrstuvwxyz012345"; diff --git a/src/pb_server/admin.rs b/src/pb_server/admin.rs index ba25768..fa057dd 100644 --- a/src/pb_server/admin.rs +++ b/src/pb_server/admin.rs @@ -166,28 +166,20 @@ async fn execute( ) .await; let (response_sender, receiver) = tokio::sync::oneshot::channel(); - manager - .send(ManagerTask::AdminServiceList { + let response = query_inventory( + authorization, + &manager, + ManagerTask::AdminServiceList { key_id, page, page_size, response_sender, - }) - .await - .map_err(|_| { - AuthFailure::new( - "server_state_unavailable", - "relay connection manager is unavailable", - true, - ) - })?; - receiver.await.map(AdminResponse::Services).map_err(|_| { - AuthFailure::new( - "server_state_unavailable", - "relay connection manager dropped the service query", - true, - ) - }) + }, + receiver, + "service", + ) + .await?; + Ok(AdminResponse::Services(response)) } AdminRequest::ConnectionList { key_id, @@ -203,32 +195,75 @@ async fn execute( ) .await; let (response_sender, receiver) = tokio::sync::oneshot::channel(); - manager - .send(ManagerTask::AdminConnectionList { + let response = query_inventory( + authorization, + &manager, + ManagerTask::AdminConnectionList { key_id, page, page_size, response_sender, - }) - .await - .map_err(|_| { - AuthFailure::new( - "server_state_unavailable", - "relay connection manager is unavailable", - true, - ) - })?; - receiver.await.map(AdminResponse::Connections).map_err(|_| { - AuthFailure::new( - "server_state_unavailable", - "relay connection manager dropped the connection query", - true, - ) - }) + }, + receiver, + "connection", + ) + .await?; + Ok(AdminResponse::Connections(response)) } } } +/// Dispatch an inventory read only while the administrator lease remains current. +/// +/// Root rotation cancels the old lease. Racing both channel operations against that +/// cancellation prevents a request authenticated under the old root from waiting for or +/// returning relay inventory after the rotation has taken effect. The final revalidation +/// establishes the successful read's authorization point after the manager produced its page. +async fn query_inventory( + authorization: &AuthContext, + manager: &ManagerTaskSender, + task: ManagerTask, + receiver: tokio::sync::oneshot::Receiver, + inventory: &'static str, +) -> std::result::Result { + let cancellation = authorization.admin_cancellation_token()?; + tokio::select! { + biased; + _ = cancellation.cancelled() => return Err(cancelled_authorization(authorization)), + result = manager.send(task) => result.map_err(|_| { + AuthFailure::new( + "server_state_unavailable", + "relay connection manager is unavailable", + true, + ) + })?, + } + let response = tokio::select! { + biased; + _ = cancellation.cancelled() => return Err(cancelled_authorization(authorization)), + result = receiver => result.map_err(|_| { + AuthFailure::new( + "server_state_unavailable", + format!("relay connection manager dropped the {inventory} query"), + true, + ) + })?, + }; + authorization.ensure_active()?; + Ok(response) +} + +fn cancelled_authorization(authorization: &AuthContext) -> AuthFailure { + match authorization.ensure_active() { + Err(error) => error, + Ok(_) => AuthFailure::new( + "administrator_key_rotated", + "administrator credential lease was cancelled during the inventory query", + false, + ), + } +} + async fn audit_read( auth: &AuthRuntime, authorization: &AuthContext, @@ -250,3 +285,116 @@ async fn audit_read( ); } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::common::auth::{AuthConfig, LegacyProtocolPolicy}; + + fn temp_state_dir(name: &str) -> std::path::PathBuf { + std::env::temp_dir().join(format!( + "pb-mapper-admin-{name}-{}", + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|duration| duration.as_nanos()) + .unwrap_or_default() + )) + } + + async fn inventory_query_rejects_rotation(connection_query: bool) { + let _process_credential_guard = crate::common::auth::PROCESS_CREDENTIAL_TEST_LOCK + .lock() + .await; + let state_dir = temp_state_dir(if connection_query { + "connections" + } else { + "services" + }); + let old_key = *b"0123456789abcdefghijklmnopqrstuv"; + let new_key = *b"abcdefghijklmnopqrstuvwxyz012345"; + let runtime = AuthRuntime::start( + old_key, + AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 4, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }, + ) + .await + .expect("authentication runtime should start"); + let admin = runtime + .authenticate_presented(0, &old_key) + .expect("old administrator key should authenticate"); + let request = if connection_query { + AdminRequest::ConnectionList { + key_id: None, + page: 0, + page_size: 100, + } + } else { + AdminRequest::ServiceList { + key_id: None, + page: 0, + page_size: 100, + } + }; + let (manager, receiver) = kanal::unbounded_async(); + let request_admin = admin.clone(); + let request_runtime = runtime.clone(); + let pending = tokio::spawn(async move { + execute(request, &request_admin, request_runtime, manager).await + }); + + let manager_task = receiver + .recv() + .await + .expect("inventory request should reach the manager"); + runtime + .rotate_root(&admin, new_key) + .await + .expect("root rotation should succeed"); + match manager_task { + ManagerTask::AdminServiceList { + response_sender, .. + } => { + let _ = response_sender.send(crate::common::message::command::AdminServicePage { + schema_version: 1, + items: Vec::new(), + next_page: None, + }); + } + ManagerTask::AdminConnectionList { + response_sender, .. + } => { + let _ = + response_sender.send(crate::common::message::command::AdminConnectionPage { + schema_version: 1, + items: Vec::new(), + next_page: None, + }); + } + _ => panic!("expected an administrator inventory manager task"), + } + + let failure = pending + .await + .expect("inventory task should not panic") + .expect_err("rotated administrator must not receive inventory"); + assert_eq!(failure.code, "administrator_key_rotated"); + + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + let _ = std::fs::remove_dir_all(state_dir); + } + + #[tokio::test] + async fn service_inventory_rejects_root_rotation_after_dispatch() { + inventory_query_rejects_rotation(false).await; + } + + #[tokio::test] + async fn connection_inventory_rejects_root_rotation_after_dispatch() { + inventory_query_rejects_rotation(true).await; + } +} diff --git a/tests/regression.rs b/tests/regression.rs index bf066c3..4e3e62f 100644 --- a/tests/regression.rs +++ b/tests/regression.rs @@ -387,16 +387,30 @@ async fn temporary_credentials_are_isolated_denied_admin_and_revoked_live() { assert_eq!(connections[0].conn_id, expected_conn_id); } - let (_, _, denied) = send_v2_request( - server_addr, - &first_credential, - PbConnRequest::Admin(AdminRequest::AuthStatus), - ) - .await; - let PbConnResponse::Error(denied) = denied else { - panic!("temporary credential unexpectedly received an admin response"); - }; - assert_eq!(denied.code, "admin_permission_required"); + for admin_request in [ + AdminRequest::AuthStatus, + AdminRequest::ServiceList { + key_id: None, + page: 0, + page_size: 100, + }, + AdminRequest::ConnectionList { + key_id: None, + page: 0, + page_size: 100, + }, + ] { + let (_, _, denied) = send_v2_request( + server_addr, + &first_credential, + PbConnRequest::Admin(admin_request), + ) + .await; + let PbConnResponse::Error(denied) = denied else { + panic!("temporary credential unexpectedly received an admin response"); + }; + assert_eq!(denied.code, "admin_permission_required"); + } let admin_credential = Credential::Admin(*TEST_ADMIN_KEY.as_bytes().first_chunk::<32>().unwrap()); diff --git a/ui/native/pb_mapper_ffi/src/state/runtime.rs b/ui/native/pb_mapper_ffi/src/state/runtime.rs index 91e848b..b7d6a83 100644 --- a/ui/native/pb_mapper_ffi/src/state/runtime.rs +++ b/ui/native/pb_mapper_ffi/src/state/runtime.rs @@ -32,7 +32,7 @@ impl PbMapperState { state_dir: self.config_dir.join("auth"), ..AuthConfig::default() }; - let auth = AuthRuntime::from_process(auth_config) + let auth = AuthRuntime::from_isolated_state(auth_config) .await .map_err(|error| { CtlError::io(format!( From 03799cd1ea79cdf0007631d83159e9e02fc449e4 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 00:43:26 +0800 Subject: [PATCH 16/74] Fix remaining authentication review findings Honor an explicit MSG_HEADER_KEY before installer legacy-key migration, validate isolated-relay legacy frames with the relay key, retain first-flight replay fingerprints for the full clock-skew window, and use a user-writable auth directory on macOS and Windows. --- CHANGELOG.md | 1 + docs/authentication-v2.md | 6 ++- docs/authentication-v2.zh-CN.md | 6 ++- docs/user-guide.md | 11 ++++-- docs/user-guide.zh-CN.md | 10 +++-- scripts/install-server-gitee.sh | 29 +++++++++++++-- scripts/install-server-github.sh | 29 +++++++++++++-- src/bin/pb-mapper.rs | 2 + src/common/auth.rs | 36 ++++++++++++++++-- src/common/auth/tests.rs | 27 ++++++++++++++ src/common/checksum.rs | 26 +++++++++++++ src/common/message/secure.rs | 10 +++-- src/common/message/secure/tests.rs | 59 +++++++++++++++++++++++++++--- 13 files changed, 222 insertions(+), 30 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 652db8a..141c4c0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ All notable changes to this project will be documented in this file. - Added encrypted snapshot/WAL authentication state, lifecycle audit records, hierarchical timing-wheel expiry, hard closure of revoked live connections, safe-mode recovery, root-key rotation, and explicit auth-state reset. - Extended the unified CLI with temporary-key lifecycle, service/connection inventory, auth status, protocol policy, root rotation, namespace targeting, and human/JSON/NDJSON output. - Replaced insecure default-key fallback with first-start random administrator-key generation, retained machine-derived keys only for explicit compatibility, and updated Flutter, installers, systemd, Docker, release metadata, and bilingual documentation. +- Fixed remaining review findings: installer migration now honors `MSG_HEADER_KEY` from `/etc/pb-mapper/server.env`, isolated relays validate legacy frames with their own administrator key, first-flight replay retention covers the full clock-skew window, and desktop macOS/Windows servers use a user-writable auth directory. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/docs/authentication-v2.md b/docs/authentication-v2.md index 032fab8..aa45acc 100644 --- a/docs/authentication-v2.md +++ b/docs/authentication-v2.md @@ -114,7 +114,7 @@ uses server-to-client counter `0`. Later control frames continue from counter The relay fingerprints `(key_id, connection_salt)` and atomically checks and inserts it in two rotating 1 MiB Bloom filters covering the current and previous -60-second windows. A probable duplicate returns the stable retryable error +300-second windows, matching the accepted first-flight clock-skew interval. A probable duplicate returns the stable retryable error `connection_salt_replayed`; one-shot administrator CLI operations retry once with a fresh salt. Mutating administrator requests additionally claim their exact fingerprint in the encrypted WAL before dispatch. Those claims survive @@ -198,7 +198,9 @@ revocation and hard connection closure deterministic. ## Persistence and safe mode -The default state directory is `/var/lib/pb-mapper/auth`: +The Linux system-service state directory is `/var/lib/pb-mapper/auth`. macOS +and Windows desktop binaries default to a user-writable application directory +instead of `/var/lib`: | File | Purpose | | --- | --- | diff --git a/docs/authentication-v2.zh-CN.md b/docs/authentication-v2.zh-CN.md index 80af9e1..26b5295 100644 --- a/docs/authentication-v2.zh-CN.md +++ b/docs/authentication-v2.zh-CN.md @@ -88,7 +88,8 @@ HKDF-SHA256 使用 connection salt 作为 salt，凭据的 32 字节 secret 作 ### 重放检测 服务端对 `(key_id, connection_salt)` 做指纹，并在同一个临界区内完成两个轮换的 -1 MiB Bloom filter 的检查与写入，覆盖当前与上一个 60 秒窗口。疑似重复会返回 +1 MiB Bloom filter 的检查与写入，覆盖当前与上一个 300 秒窗口，与首帧可接受的 +时钟偏差窗口一致。疑似重复会返回 可重试错误 `connection_salt_replayed`；一次性 admin CLI 会自动换 salt 重试一次。 会修改状态的管理员请求还会在分发前把精确指纹写入加密 WAL；该记录在十分钟内跨 重启、跨 compact 保留，不能通过等待 Bloom 窗口结束或重启进程来重放旧操作。 @@ -152,7 +153,8 @@ tombstone 以给出稳定错误后，槽位可以复用。显式 `key gc` 可立 ## 持久化与安全模式 -默认目录 `/var/lib/pb-mapper/auth` 权限为 `0700`： +Linux 系统服务默认目录是 `/var/lib/pb-mapper/auth`，权限为 `0700`。macOS 与 +Windows 桌面二进制默认写到用户可写的应用目录，而不是 `/var/lib`： | 文件 | 用途 | | --- | --- | diff --git a/docs/user-guide.md b/docs/user-guide.md index 1165378..f8209f7 100644 --- a/docs/user-guide.md +++ b/docs/user-guide.md @@ -112,7 +112,7 @@ Optional flags: - `--ipv6`: enable IPv6 listening - `--keep-alive`: enable TCP keep-alive -- `--auth-state-dir`: authentication state directory (default `/var/lib/pb-mapper/auth`) +- `--auth-state-dir`: authentication state directory (Linux system default `/var/lib/pb-mapper/auth`; macOS and Windows use a user-writable application directory) - `--max-temporary-keys`: fixed temporary-key slot capacity (default `65536`) - `--max-temporary-key-ttl`: maximum issued TTL (default `30d`) - `--legacy-protocol allow|deny`: initial legacy-client policy @@ -120,8 +120,11 @@ Optional flags: ### Administrator and temporary credentials -On first start, the relay creates a random administrator key at -`/var/lib/pb-mapper/auth/admin.key`. There is no built-in default credential. +On first start, the relay creates a random administrator key in its +authentication state directory (`admin.key`). On Linux system services that +is `/var/lib/pb-mapper/auth/admin.key`. Desktop macOS and Windows builds use +a user-writable application directory instead. There is no built-in default +credential. Keep the administrator key on the relay host and use it to issue a temporary credential for a workload: @@ -232,7 +235,7 @@ flutter run - `PB_MAPPER_SERVER`: default server address for the CLI - `MSG_HEADER_KEY`: 32-character administrator key or a `pbmt1_` temporary credential -- `PB_MAPPER_AUTH_STATE_DIR`: relay auth-state directory, default `/var/lib/pb-mapper/auth` +- `PB_MAPPER_AUTH_STATE_DIR`: relay auth-state directory (Linux system default `/var/lib/pb-mapper/auth`; macOS and Windows use a user-writable application directory) - `PB_MAPPER_AUTH_MAX_TEMP_KEYS`: fixed temporary-key capacity, default `65536` - `PB_MAPPER_AUTH_MAX_TEMP_TTL_SECS`: maximum temporary-key TTL, default 30 days - `PB_MAPPER_LEGACY_PROTOCOL`: `allow` or `deny`, default `allow` diff --git a/docs/user-guide.zh-CN.md b/docs/user-guide.zh-CN.md index aea60e2..9cde552 100644 --- a/docs/user-guide.zh-CN.md +++ b/docs/user-guide.zh-CN.md @@ -112,7 +112,7 @@ pb-mapper server --port 7666 - `--ipv6`：开启 IPv6 监听 - `--keep-alive`：开启 TCP keep-alive -- `--auth-state-dir`：认证状态目录，默认 `/var/lib/pb-mapper/auth` +- `--auth-state-dir`：认证状态目录（Linux 系统服务默认 `/var/lib/pb-mapper/auth`；macOS 与 Windows 使用当前用户可写的应用目录） - `--max-temporary-keys`：临时 key 固定槽位容量，默认 `65536` - `--max-temporary-key-ttl`：临时 key 最大 TTL，默认 `30d` - `--legacy-protocol allow|deny`：旧协议初始接入策略 @@ -120,8 +120,10 @@ pb-mapper server --port 7666 ### 管理员密钥与临时凭据 -中继首次启动时会在 `/var/lib/pb-mapper/auth/admin.key` 生成随机管理员密钥，系统不再 -提供内置默认 key。管理员密钥留在中继机器上，用它为业务签发临时凭据： +中继首次启动时会在认证状态目录生成随机管理员密钥（`admin.key`）。Linux 系统服务 +默认写到 `/var/lib/pb-mapper/auth/admin.key`；macOS 与 Windows 桌面构建则使用当前 +用户可写的应用目录。系统不再提供内置默认 key。管理员密钥留在中继机器上，用它为 +业务签发临时凭据： ```bash export MSG_HEADER_KEY="$(sudo cat /var/lib/pb-mapper/auth/admin.key)" @@ -226,7 +228,7 @@ flutter run - `PB_MAPPER_SERVER`：CLI 默认服务器地址 - `MSG_HEADER_KEY`：32 字符管理员密钥或 `pbmt1_` 临时凭据 -- `PB_MAPPER_AUTH_STATE_DIR`：中继认证状态目录，默认 `/var/lib/pb-mapper/auth` +- `PB_MAPPER_AUTH_STATE_DIR`：中继认证状态目录（Linux 系统服务默认 `/var/lib/pb-mapper/auth`；macOS 与 Windows 使用当前用户可写的应用目录） - `PB_MAPPER_AUTH_MAX_TEMP_KEYS`：临时 key 固定容量，默认 `65536` - `PB_MAPPER_AUTH_MAX_TEMP_TTL_SECS`：临时 key 最大 TTL，默认 30 天 - `PB_MAPPER_LEGACY_PROTOCOL`：`allow` 或 `deny`，默认 `allow` diff --git a/scripts/install-server-gitee.sh b/scripts/install-server-gitee.sh index abc6506..d4cd398 100755 --- a/scripts/install-server-gitee.sh +++ b/scripts/install-server-gitee.sh @@ -13,6 +13,26 @@ PORT="${PB_MAPPER_PORT:-7666}" AUTH_DIR="/var/lib/pb-mapper/auth" ADMIN_KEY_PATH="${AUTH_DIR}/admin.key" LEGACY_KEY_PATH="/var/lib/pb-mapper-server/msg_header_key" +SERVER_ENV_FILE="/etc/pb-mapper/server.env" + +configured_msg_header_key() { + if [ -n "${MSG_HEADER_KEY:-}" ]; then + printf '%s' "$MSG_HEADER_KEY" + return 0 + fi + if [ ! -f "$SERVER_ENV_FILE" ]; then + return 0 + fi + awk -F= ' + $1 ~ /^[[:space:]]*#/ { next } + $1 ~ /^[[:space:]]*MSG_HEADER_KEY[[:space:]]*$/ { + val = substr($0, index($0, "=") + 1) + sub(/\r$/, "", val) + key = val + } + END { printf "%s", key } + ' "$SERVER_ENV_FILE" +} # Re-run with sudo if needed if [ "${EUID:-$(id -u)}" -ne 0 ]; then @@ -70,10 +90,12 @@ fi mkdir -p "$INSTALL_DIR" install -m 0755 "$BIN_PATH" "${INSTALL_DIR}/pb-mapper" -# Preserve the former machine-derived credential on upgrade. New installations let -# pb-mapper create a random administrator key on first start. +# Preserve the former machine-derived credential on upgrade only when neither +# admin.key nor an explicit MSG_HEADER_KEY is already configured. An explicit +# key in the environment or /etc/pb-mapper/server.env must win; otherwise the +# runtime would prefer the newly copied admin.key and lock operators out. install -d -m 0700 "$AUTH_DIR" -if [ ! -s "$ADMIN_KEY_PATH" ] && [ -s "$LEGACY_KEY_PATH" ]; then +if [ -z "$(configured_msg_header_key)" ] && [ ! -s "$ADMIN_KEY_PATH" ] && [ -s "$LEGACY_KEY_PATH" ]; then install -m 0600 "$LEGACY_KEY_PATH" "$ADMIN_KEY_PATH" echo "Migrated the legacy machine-derived key into $ADMIN_KEY_PATH" fi @@ -99,6 +121,7 @@ After=network.target Type=simple ExecStart=${INSTALL_DIR}/pb-mapper server --port ${PORT} Environment=RUST_LOG=info +EnvironmentFile=-/etc/pb-mapper/server.env StateDirectory=pb-mapper StateDirectoryMode=0700 Restart=on-failure diff --git a/scripts/install-server-github.sh b/scripts/install-server-github.sh index 5a11466..038d491 100755 --- a/scripts/install-server-github.sh +++ b/scripts/install-server-github.sh @@ -13,6 +13,26 @@ PORT="${PB_MAPPER_PORT:-7666}" AUTH_DIR="/var/lib/pb-mapper/auth" ADMIN_KEY_PATH="${AUTH_DIR}/admin.key" LEGACY_KEY_PATH="/var/lib/pb-mapper-server/msg_header_key" +SERVER_ENV_FILE="/etc/pb-mapper/server.env" + +configured_msg_header_key() { + if [ -n "${MSG_HEADER_KEY:-}" ]; then + printf '%s' "$MSG_HEADER_KEY" + return 0 + fi + if [ ! -f "$SERVER_ENV_FILE" ]; then + return 0 + fi + awk -F= ' + $1 ~ /^[[:space:]]*#/ { next } + $1 ~ /^[[:space:]]*MSG_HEADER_KEY[[:space:]]*$/ { + val = substr($0, index($0, "=") + 1) + sub(/\r$/, "", val) + key = val + } + END { printf "%s", key } + ' "$SERVER_ENV_FILE" +} # Re-run with sudo if needed if [ "${EUID:-$(id -u)}" -ne 0 ]; then @@ -70,10 +90,12 @@ fi mkdir -p "$INSTALL_DIR" install -m 0755 "$BIN_PATH" "${INSTALL_DIR}/pb-mapper" -# Preserve the former machine-derived credential on upgrade. New installations let -# pb-mapper create a random administrator key on first start. +# Preserve the former machine-derived credential on upgrade only when neither +# admin.key nor an explicit MSG_HEADER_KEY is already configured. An explicit +# key in the environment or /etc/pb-mapper/server.env must win; otherwise the +# runtime would prefer the newly copied admin.key and lock operators out. install -d -m 0700 "$AUTH_DIR" -if [ ! -s "$ADMIN_KEY_PATH" ] && [ -s "$LEGACY_KEY_PATH" ]; then +if [ -z "$(configured_msg_header_key)" ] && [ ! -s "$ADMIN_KEY_PATH" ] && [ -s "$LEGACY_KEY_PATH" ]; then install -m 0600 "$LEGACY_KEY_PATH" "$ADMIN_KEY_PATH" echo "Migrated the legacy machine-derived key into $ADMIN_KEY_PATH" fi @@ -99,6 +121,7 @@ After=network.target Type=simple ExecStart=${INSTALL_DIR}/pb-mapper server --port ${PORT} Environment=RUST_LOG=info +EnvironmentFile=-/etc/pb-mapper/server.env StateDirectory=pb-mapper StateDirectoryMode=0700 Restart=on-failure diff --git a/src/bin/pb-mapper.rs b/src/bin/pb-mapper.rs index 0c9a3dd..2c8efdc 100644 --- a/src/bin/pb-mapper.rs +++ b/src/bin/pb-mapper.rs @@ -89,6 +89,8 @@ struct ServerArgs { #[arg(long, default_value_t = false)] use_machine_msg_header_key: bool, /// Directory containing encrypted authentication state and the administrator key file. + /// Defaults to /var/lib/pb-mapper/auth on Linux, or a user-writable application + /// directory on macOS and Windows. #[arg(long)] auth_state_dir: Option, /// Create a random administrator key before starting the relay. diff --git a/src/common/auth.rs b/src/common/auth.rs index a4f1303..b01bcdf 100644 --- a/src/common/auth.rs +++ b/src/common/auth.rs @@ -82,12 +82,42 @@ pub struct AuthConfig { pub legacy_protocol: LegacyProtocolPolicy, } +pub fn default_auth_state_dir() -> PathBuf { + std::env::var_os("PB_MAPPER_AUTH_STATE_DIR") + .map(PathBuf::from) + .unwrap_or_else(platform_default_auth_state_dir) +} + +pub(crate) fn platform_default_auth_state_dir() -> PathBuf { + #[cfg(windows)] + { + let base = std::env::var_os("LOCALAPPDATA") + .or_else(|| std::env::var_os("APPDATA")) + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from(r"C:\ProgramData")); + base.join("pb-mapper").join("auth") + } + #[cfg(target_os = "macos")] + { + match std::env::var_os("HOME") { + Some(home) => PathBuf::from(home) + .join("Library") + .join("Application Support") + .join("pb-mapper") + .join("auth"), + None => PathBuf::from("/Library/Application Support/pb-mapper/auth"), + } + } + #[cfg(not(any(windows, target_os = "macos")))] + { + PathBuf::from(DEFAULT_AUTH_STATE_DIR) + } +} + impl Default for AuthConfig { fn default() -> Self { Self { - state_dir: std::env::var_os("PB_MAPPER_AUTH_STATE_DIR") - .map(PathBuf::from) - .unwrap_or_else(|| PathBuf::from(DEFAULT_AUTH_STATE_DIR)), + state_dir: default_auth_state_dir(), max_temporary_keys: env_usize( "PB_MAPPER_AUTH_MAX_TEMP_KEYS", DEFAULT_TEMP_KEY_CAPACITY, diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index 6fe5835..a73cf18 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -26,6 +26,33 @@ fn authenticate_for_test(runtime: &AuthRuntime, key_id: u64) -> Result ChecksumType { }) } +#[inline] +/// Compute frame checksum from payload length and an explicit header key. +pub fn get_checksum_for_key(datalen: DataLenType, key: &[u8]) -> ChecksumType { + datalen ^ gen_checksum_by_key(key) +} + #[inline] /// Compute frame checksum from payload length and the current header key hash. pub fn get_checksum(datalen: DataLenType) -> ChecksumType { datalen ^ MSG_HEADER_KEY_STATE.hash.load(Ordering::Acquire) } +#[inline] +/// Validate a frame checksum against an explicit header key. +pub fn valid_checksum_for_key(datalen: DataLenType, checksum: ChecksumType, key: &[u8]) -> bool { + checksum == get_checksum_for_key(datalen, key) +} + #[inline] /// Validate frame checksum generated by [`get_checksum`]. pub fn valid_checksum(datalen: DataLenType, checksum: ChecksumType) -> bool { @@ -521,6 +533,20 @@ mod tests { ); } + #[test] + fn checksum_for_an_explicit_key_is_independent_of_process_state() { + use super::*; + let key = b"0123456789abcdefghijklmnopqrstuv"; + let datalen = 32; + let checksum = get_checksum_for_key(datalen, key); + assert!(valid_checksum_for_key(datalen, checksum, key)); + assert!(!valid_checksum_for_key( + datalen, + checksum, + b"abcdefghijklmnopqrstuvwxyz012345" + )); + } + #[test] fn test_derive_msg_header_key_is_stable() { use super::*; diff --git a/src/common/message/secure.rs b/src/common/message/secure.rs index 41c52a1..2ea085e 100644 --- a/src/common/message/secure.rs +++ b/src/common/message/secure.rs @@ -30,7 +30,9 @@ use super::{ CodecMessageReader, CodecMessageWriter, DataLenType, MessageReader, MessageWriter, MAX_MSG_LEN, }; use crate::common::auth::{AuthContext, AuthFailure, AuthRuntime, LegacyConnectionGuard}; -use crate::common::checksum::{get_process_credential, valid_checksum, AesKeyType, Credential}; +use crate::common::checksum::{ + get_process_credential, valid_checksum_for_key, AesKeyType, Credential, +}; use crate::common::error::{Error, Result}; use crate::utils::codec::{Aes256GcmDeCodec, Aes256GcmEnCodec, Decryptor}; @@ -41,10 +43,10 @@ const FIRST_PREFIX_REMAINDER_LEN: usize = 28; const FRAME_HEADER_LEN: usize = 12; const DIRECTION_CLIENT_TO_SERVER: u8 = 0; const DIRECTION_SERVER_TO_CLIENT: u8 = 1; -const DEFAULT_REPLAY_WINDOW_SECONDS: u64 = 60; +const MAX_CONNECTION_CLOCK_SKEW_SECONDS: u64 = 5 * 60; +const DEFAULT_REPLAY_WINDOW_SECONDS: u64 = MAX_CONNECTION_CLOCK_SKEW_SECONDS; const DEFAULT_REPLAY_FILTER_BYTES: usize = 1024 * 1024; const MAX_INITIAL_PLAINTEXT_LEN: u32 = 64 * 1024; -const MAX_CONNECTION_CLOCK_SKEW_SECONDS: u64 = 5 * 60; #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum HeaderProtocol { @@ -371,7 +373,7 @@ impl ServerSecurity { ), response_session: None, })?; - if !valid_checksum(datalen, checksum) || datalen > MAX_MSG_LEN { + if !valid_checksum_for_key(datalen, checksum, &key) || datalen > MAX_MSG_LEN { return Err(ServerInitialError { failure: AuthFailure::new( "legacy_frame_invalid", diff --git a/src/common/message/secure/tests.rs b/src/common/message/secure/tests.rs index eecc18c..f824946 100644 --- a/src/common/message/secure/tests.rs +++ b/src/common/message/secure/tests.rs @@ -10,8 +10,8 @@ //! credentials, while lifecycle persistence remains covered by `common::auth::tests`. use super::*; -use crate::common::auth::{AuthConfig, LegacyProtocolPolicy}; -use crate::common::checksum::encode_temporary_credential; +use crate::common::auth::{AuthConfig, LegacyProtocolPolicy, PROCESS_CREDENTIAL_TEST_LOCK}; +use crate::common::checksum::{encode_temporary_credential, set_process_msg_header_key}; fn temp_config() -> AuthConfig { let mut random = [0_u8; 8]; @@ -161,13 +161,62 @@ async fn oversized_initial_frame_is_rejected_before_reading_its_body() { #[test] fn rotating_bloom_covers_current_and_previous_window() { - let mut bloom = RotatingBloom::new(1024, 60); + let mut bloom = RotatingBloom::new(1024, DEFAULT_REPLAY_WINDOW_SECONDS); let value = [7_u8; 32]; let start = bloom.current_started_at; assert!(!bloom.contains(&value, start)); bloom.insert(&value, start); - assert!(bloom.contains(&value, start + 60)); - assert!(!bloom.contains(&value, start + 121)); + assert!(bloom.contains(&value, start + DEFAULT_REPLAY_WINDOW_SECONDS)); + assert!(!bloom.contains( + &value, + start + DEFAULT_REPLAY_WINDOW_SECONDS.saturating_mul(2) + 1 + )); +} + +#[test] +fn rotating_bloom_retains_fingerprints_for_the_clock_skew_window() { + let mut bloom = RotatingBloom::new(1024, DEFAULT_REPLAY_WINDOW_SECONDS); + let value = [9_u8; 32]; + let start = bloom.current_started_at; + bloom.insert(&value, start); + assert!(bloom.contains(&value, start + 120)); + assert!(bloom.contains(&value, start + MAX_CONNECTION_CLOCK_SKEW_SECONDS)); +} + +#[tokio::test] +async fn legacy_initial_frame_validates_against_isolated_relay_key() { + let _process_credential_guard = PROCESS_CREDENTIAL_TEST_LOCK.lock().await; + let config = temp_config(); + let isolated_admin = *b"isolated-admin-key-0123456789abc"; + std::fs::create_dir_all(&config.state_dir).unwrap(); + std::fs::write(config.state_dir.join("admin.key"), isolated_admin).unwrap(); + let auth = AuthRuntime::from_isolated_state(config.clone()) + .await + .unwrap(); + + set_process_msg_header_key(Some( + std::str::from_utf8(&isolated_admin).expect("printable isolated key"), + )) + .unwrap(); + let client = ClientHeaderSession::new_legacy(isolated_admin); + let bytes = encode_initial(&client, b"legacy-isolated").await; + + let temporary_key_id = 1; + let temporary_key = *b"temporary-remote-key-0123456789a"; + set_process_msg_header_key(Some(&encode_temporary_credential( + temporary_key_id, + &temporary_key, + ))) + .unwrap(); + + let security = ServerSecurity::new(auth); + let mut input = std::io::Cursor::new(bytes); + let initial = security.read_initial(&mut input).await.unwrap(); + assert_eq!(initial.payload, b"legacy-isolated"); + assert_eq!(initial.session.protocol(), HeaderProtocol::Legacy); + + set_process_msg_header_key(None).unwrap(); + let _ = std::fs::remove_dir_all(config.state_dir); } #[test] From 91e9e7e55c2085325b248f4be1083721bb905f75 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 01:25:41 +0800 Subject: [PATCH 17/74] Fix remaining isolated-relay, replay, and key-lifecycle review findings Bind isolated-relay legacy continuation checksums to the session key, keep first-flight Bloom generations for twice the clock-skew window, reject NUL/non-printable rotated administrator keys, cancel in-flight status reads on credential revocation, and refuse --force-init-admin-key when encrypted auth state already exists. --- CHANGELOG.md | 1 + docs/authentication-v2.md | 3 +- docs/authentication-v2.zh-CN.md | 4 +- src/common/auth/actor.rs | 8 ++- src/common/auth/persistence.rs | 16 ++++++ src/common/auth/tests.rs | 33 ++++++++++++ src/common/checksum.rs | 14 +++++ src/common/message/mod.rs | 86 ++++++++++++++++++++++++++---- src/common/message/secure.rs | 56 +++++++++++-------- src/common/message/secure/tests.rs | 43 +++++++++++---- src/pb_server/connection.rs | 29 ++++++---- 11 files changed, 237 insertions(+), 56 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 141c4c0..a27019e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ All notable changes to this project will be documented in this file. - Extended the unified CLI with temporary-key lifecycle, service/connection inventory, auth status, protocol policy, root rotation, namespace targeting, and human/JSON/NDJSON output. - Replaced insecure default-key fallback with first-start random administrator-key generation, retained machine-derived keys only for explicit compatibility, and updated Flutter, installers, systemd, Docker, release metadata, and bilingual documentation. - Fixed remaining review findings: installer migration now honors `MSG_HEADER_KEY` from `/etc/pb-mapper/server.env`, isolated relays validate legacy frames with their own administrator key, first-flight replay retention covers the full clock-skew window, and desktop macOS/Windows servers use a user-writable auth directory. +- Rejected NUL/non-printable rotated administrator keys, cancelled in-flight status reads on credential revocation, refused `--force-init-admin-key` when encrypted auth state already exists, bound isolated-relay legacy continuation checksums to the relay key, and doubled first-flight Bloom retention so a max-future timestamp cannot outlive the filter. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/docs/authentication-v2.md b/docs/authentication-v2.md index aa45acc..a35236a 100644 --- a/docs/authentication-v2.md +++ b/docs/authentication-v2.md @@ -114,7 +114,8 @@ uses server-to-client counter `0`. Later control frames continue from counter The relay fingerprints `(key_id, connection_salt)` and atomically checks and inserts it in two rotating 1 MiB Bloom filters covering the current and previous -300-second windows, matching the accepted first-flight clock-skew interval. A probable duplicate returns the stable retryable error +600-second windows, so a max-future first-flight timestamp cannot outlive replay +retention. A probable duplicate returns the stable retryable error `connection_salt_replayed`; one-shot administrator CLI operations retry once with a fresh salt. Mutating administrator requests additionally claim their exact fingerprint in the encrypted WAL before dispatch. Those claims survive diff --git a/docs/authentication-v2.zh-CN.md b/docs/authentication-v2.zh-CN.md index 26b5295..afb4cd0 100644 --- a/docs/authentication-v2.zh-CN.md +++ b/docs/authentication-v2.zh-CN.md @@ -88,8 +88,8 @@ HKDF-SHA256 使用 connection salt 作为 salt，凭据的 32 字节 secret 作 ### 重放检测 服务端对 `(key_id, connection_salt)` 做指纹，并在同一个临界区内完成两个轮换的 -1 MiB Bloom filter 的检查与写入，覆盖当前与上一个 300 秒窗口，与首帧可接受的 -时钟偏差窗口一致。疑似重复会返回 +1 MiB Bloom filter 的检查与写入，覆盖当前与上一个 600 秒窗口，使首帧允许的 +最大未来时间戳无法在过滤器遗忘后继续重放。疑似重复会返回 可重试错误 `connection_salt_replayed`；一次性 admin CLI 会自动换 salt 重试一次。 会修改状态的管理员请求还会在分发前把精确指纹写入加密 WAL；该记录在十分钟内跨 重启、跨 compact 保留，不能通过等待 Bloom 窗口结束或重启进程来重放旧操作。 diff --git a/src/common/auth/actor.rs b/src/common/auth/actor.rs index 1cf9fc2..874f177 100644 --- a/src/common/auth/actor.rs +++ b/src/common/auth/actor.rs @@ -708,10 +708,14 @@ fn actor_rotate_root( false, ) })?; - if new_key_string.chars().any(char::is_whitespace) { + if !new_key_string + .as_bytes() + .iter() + .all(|byte| byte.is_ascii_graphic()) + { return Err(AuthFailure::new( "administrator_key_invalid", - "administrator key must not contain whitespace", + "administrator key must be 32 printable ASCII bytes without whitespace or NUL", false, )); } diff --git a/src/common/auth/persistence.rs b/src/common/auth/persistence.rs index 198f792..216e7f5 100644 --- a/src/common/auth/persistence.rs +++ b/src/common/auth/persistence.rs @@ -590,6 +590,22 @@ pub fn initialize_admin_key(path: &Path, force: bool) -> Result Result Result { if bytes.len() != 32 { return Err(key_len_error(raw)); } + if !bytes.iter().all(|byte| byte.is_ascii_graphic()) { + return Err(format!( + "`{ENV_MSG_HEADER_KEY}` administrator key must be 32 printable ASCII bytes without whitespace or NUL" + )); + } Ok(Credential::Admin( bytes.try_into().expect("validated admin key width"), )) @@ -547,6 +552,15 @@ mod tests { )); } + #[test] + fn administrator_credentials_reject_nul_and_whitespace() { + use super::*; + let mut with_nul = *b"0123456789abcdefghijklmnopqrstuv"; + with_nul[8] = 0; + assert!(parse_credential(std::str::from_utf8(&with_nul).unwrap()).is_err()); + assert!(parse_credential("0123456789abcdefghijklmnopq rstuv").is_err()); + } + #[test] fn test_derive_msg_header_key_is_stable() { use super::*; diff --git a/src/common/message/mod.rs b/src/common/message/mod.rs index 6379532..782d8f4 100644 --- a/src/common/message/mod.rs +++ b/src/common/message/mod.rs @@ -7,7 +7,10 @@ use snafu::{ensure, ResultExt}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use super::buffer::{BufferGetter, CommonBuffer, FixedSizeBuffer}; -use super::checksum::{get_checksum, get_msg_header_key, valid_checksum}; +use super::checksum::{ + get_checksum, get_checksum_for_key, get_msg_header_key, valid_checksum, valid_checksum_for_key, + AesKeyType, +}; use super::error::{ self, MsgDatalenValidateSnafu, MsgNetworkReadBodySnafu, MsgNetworkReadCheckSumSnafu, MsgNetworkReadDatalenSnafu, MsgNetworkWriteBodySnafu, MsgNetworkWriteCheckSumSnafu, @@ -131,10 +134,29 @@ gen_write_network_with_error!( ); #[inline] -async fn get_msg_len(reader: &mut T) -> Result { +fn checksum_matches(datalen: DataLenType, checksum: u32, key: Option<&[u8]>) -> bool { + match key { + Some(key) => valid_checksum_for_key(datalen, checksum, key), + None => valid_checksum(datalen, checksum), + } +} + +#[inline] +fn checksum_for(len: DataLenType, key: Option<&[u8]>) -> u32 { + match key { + Some(key) => get_checksum_for_key(len, key), + None => get_checksum(len), + } +} + +#[inline] +async fn get_msg_len( + reader: &mut T, + checksum_key: Option<&[u8]>, +) -> Result { let checksum = read_checksum(reader).await?; let datalen = read_datalen(reader).await?; - if valid_checksum(datalen, checksum) { + if checksum_matches(datalen, checksum, checksum_key) { ensure!( datalen <= MAX_MSG_LEN, MsgDatalenExceededSnafu { @@ -149,14 +171,19 @@ async fn get_msg_len(reader: &mut T) -> Result(writer: &mut T, len: DataLenType) -> Result<()> { - write_checksum(writer, get_checksum(len)).await?; +async fn set_msg_len( + writer: &mut T, + len: DataLenType, + checksum_key: Option<&[u8]>, +) -> Result<()> { + write_checksum(writer, checksum_for(len, checksum_key)).await?; write_datalen(writer, len).await } pub struct NormalMessageReader<'a, T: AsyncReadExt + Unpin> { reader: &'a mut T, buffer: CommonBuffer, + checksum_key: Option, } impl<'a, T: AsyncReadExt + Unpin> NormalMessageReader<'a, T> { @@ -164,11 +191,21 @@ impl<'a, T: AsyncReadExt + Unpin> NormalMessageReader<'a, T> { Self { reader, buffer: CommonBuffer::new(), + checksum_key: None, } } + pub fn with_checksum_key(mut self, key: AesKeyType) -> Self { + self.checksum_key = Some(key); + self + } + async fn read_msg_inner(&mut self) -> Result<&'_ [u8]> { - let datalen = get_msg_len(&mut self.reader).await?; + let datalen = get_msg_len( + &mut self.reader, + self.checksum_key.as_ref().map(|key| key.as_slice()), + ) + .await?; self.buffer.fixed_resize(datalen as usize); let n = read_msg_body(&mut self.reader, self.buffer.buffer_mut()).await?; Ok(&self.buffer.buffer()[0..n]) @@ -183,15 +220,24 @@ impl<'a, T: AsyncReadExt + Unpin> MessageReader for NormalMessageReader<'a, T> { pub struct NormalMessageWriter<'a, T: AsyncWriteExt> { writer: &'a mut T, + checksum_key: Option, } impl<'a, T: AsyncWriteExt + Unpin> NormalMessageWriter<'a, T> { pub fn new(writer: &'a mut T) -> Self { - Self { writer } + Self { + writer, + checksum_key: None, + } } async fn write_msg_inner(&mut self, msg: &[u8]) -> Result<()> { - set_msg_len(&mut self.writer, msg.len() as u32).await?; + set_msg_len( + &mut self.writer, + msg.len() as u32, + self.checksum_key.as_ref().map(|key| key.as_slice()), + ) + .await?; write_msg_body(&mut self.writer, msg).await } @@ -215,6 +261,11 @@ impl<'a, T: AsyncReadExt + Unpin, D: Decryptor> CodecMessageReader<'a, T, D> { decryptor, } } + + pub fn with_checksum_key(mut self, key: AesKeyType) -> Self { + self.reader.checksum_key = Some(key); + self + } } impl<'a, T: AsyncReadExt + Unpin, D: Decryptor> MessageReader for CodecMessageReader<'a, T, D> { @@ -236,11 +287,21 @@ impl<'a, T: AsyncReadExt + Unpin, D: Decryptor> MessageReader for CodecMessageRe pub struct CodecMessageWriter<'a, T: AsyncWriteExt + Unpin, E: Encryptor> { writer: &'a mut T, encryptor: E, + checksum_key: Option, } impl<'a, T: AsyncWriteExt + Unpin, E: Encryptor> CodecMessageWriter<'a, T, E> { pub fn new(writer: &'a mut T, encryptor: E) -> Self { - Self { writer, encryptor } + Self { + writer, + encryptor, + checksum_key: None, + } + } + + pub fn with_checksum_key(mut self, key: AesKeyType) -> Self { + self.checksum_key = Some(key); + self } pub async fn shutdown(&mut self) -> std::io::Result<()> { @@ -260,7 +321,12 @@ impl<'a, T: AsyncWriteExt + Unpin, E: Encryptor> MessageWriter for CodecMessageW })?; let msg_len = (buf.len() + tag.as_ref().len()) as DataLenType; - set_msg_len(self.writer, msg_len).await?; + set_msg_len( + self.writer, + msg_len, + self.checksum_key.as_ref().map(|key| key.as_slice()), + ) + .await?; write_codec_msg(self.writer, &buf).await?; write_codec_tag(self.writer, tag.as_ref()).await } diff --git a/src/common/message/secure.rs b/src/common/message/secure.rs index 2ea085e..4388458 100644 --- a/src/common/message/secure.rs +++ b/src/common/message/secure.rs @@ -44,7 +44,7 @@ const FRAME_HEADER_LEN: usize = 12; const DIRECTION_CLIENT_TO_SERVER: u8 = 0; const DIRECTION_SERVER_TO_CLIENT: u8 = 1; const MAX_CONNECTION_CLOCK_SKEW_SECONDS: u64 = 5 * 60; -const DEFAULT_REPLAY_WINDOW_SECONDS: u64 = MAX_CONNECTION_CLOCK_SKEW_SECONDS; +const DEFAULT_REPLAY_WINDOW_SECONDS: u64 = MAX_CONNECTION_CLOCK_SKEW_SECONDS.saturating_mul(2); const DEFAULT_REPLAY_FILTER_BYTES: usize = 1024 * 1024; const MAX_INITIAL_PLAINTEXT_LEN: u32 = 64 * 1024; @@ -105,6 +105,7 @@ impl ClientHeaderSession { let codec = Aes256GcmEnCodec::try_new(&self.legacy_key) .map_err(|_| protocol_error("failed to initialize legacy writer"))?; CodecMessageWriter::new(writer, codec) + .with_checksum_key(self.legacy_key) .write_msg(message) .await } @@ -128,11 +129,14 @@ impl ClientHeaderSession { reader: &'a mut T, ) -> Result> { match self.protocol { - HeaderProtocol::Legacy => Ok(HeaderMessageReader::Legacy(CodecMessageReader::new( - reader, - Aes256GcmDeCodec::try_new(&self.legacy_key) - .map_err(|_| protocol_error("failed to initialize legacy reader"))?, - ))), + HeaderProtocol::Legacy => Ok(HeaderMessageReader::Legacy( + CodecMessageReader::new( + reader, + Aes256GcmDeCodec::try_new(&self.legacy_key) + .map_err(|_| protocol_error("failed to initialize legacy reader"))?, + ) + .with_checksum_key(self.legacy_key), + )), HeaderProtocol::V2 => Ok(HeaderMessageReader::V2(V2MessageReader::new( reader, self.v2.as_ref().expect("v2 session material").clone(), @@ -147,11 +151,14 @@ impl ClientHeaderSession { writer: &'a mut T, ) -> Result> { match self.protocol { - HeaderProtocol::Legacy => Ok(HeaderMessageWriter::Legacy(CodecMessageWriter::new( - writer, - Aes256GcmEnCodec::try_new(&self.legacy_key) - .map_err(|_| protocol_error("failed to initialize legacy writer"))?, - ))), + HeaderProtocol::Legacy => Ok(HeaderMessageWriter::Legacy( + CodecMessageWriter::new( + writer, + Aes256GcmEnCodec::try_new(&self.legacy_key) + .map_err(|_| protocol_error("failed to initialize legacy writer"))?, + ) + .with_checksum_key(self.legacy_key), + )), HeaderProtocol::V2 => Ok(HeaderMessageWriter::V2(V2MessageWriter::new( writer, self.v2.as_ref().expect("v2 session material").clone(), @@ -215,11 +222,15 @@ impl ServerHeaderSession { writer: &'a mut T, ) -> Result> { match self.protocol { - HeaderProtocol::Legacy => Ok(HeaderMessageWriter::Legacy(CodecMessageWriter::new( - writer, - Aes256GcmEnCodec::try_new(&self.legacy_key) - .map_err(|_| protocol_error("failed to initialize legacy response writer"))?, - ))), + HeaderProtocol::Legacy => Ok(HeaderMessageWriter::Legacy( + CodecMessageWriter::new( + writer, + Aes256GcmEnCodec::try_new(&self.legacy_key).map_err(|_| { + protocol_error("failed to initialize legacy response writer") + })?, + ) + .with_checksum_key(self.legacy_key), + )), HeaderProtocol::V2 => Ok(HeaderMessageWriter::V2(V2MessageWriter::new( writer, self.v2.as_ref().expect("v2 session material").clone(), @@ -234,11 +245,14 @@ impl ServerHeaderSession { reader: &'a mut T, ) -> Result> { match self.protocol { - HeaderProtocol::Legacy => Ok(HeaderMessageReader::Legacy(CodecMessageReader::new( - reader, - Aes256GcmDeCodec::try_new(&self.legacy_key) - .map_err(|_| protocol_error("failed to initialize legacy reader"))?, - ))), + HeaderProtocol::Legacy => Ok(HeaderMessageReader::Legacy( + CodecMessageReader::new( + reader, + Aes256GcmDeCodec::try_new(&self.legacy_key) + .map_err(|_| protocol_error("failed to initialize legacy reader"))?, + ) + .with_checksum_key(self.legacy_key), + )), HeaderProtocol::V2 => Ok(HeaderMessageReader::V2(V2MessageReader::new( reader, self.v2.as_ref().expect("v2 session material").clone(), diff --git a/src/common/message/secure/tests.rs b/src/common/message/secure/tests.rs index f824946..e6c5edd 100644 --- a/src/common/message/secure/tests.rs +++ b/src/common/message/secure/tests.rs @@ -183,6 +183,19 @@ fn rotating_bloom_retains_fingerprints_for_the_clock_skew_window() { assert!(bloom.contains(&value, start + MAX_CONNECTION_CLOCK_SKEW_SECONDS)); } +#[test] +fn rotating_bloom_retains_a_max_future_timestamp_past_the_next_rotation() { + let mut bloom = RotatingBloom::new(1024, DEFAULT_REPLAY_WINDOW_SECONDS); + let value = [11_u8; 32]; + let start = bloom.current_started_at; + let insert_at = start + DEFAULT_REPLAY_WINDOW_SECONDS - 1; + bloom.insert(&value, insert_at); + assert!(bloom.contains( + &value, + insert_at + MAX_CONNECTION_CLOCK_SKEW_SECONDS.saturating_mul(2) - 1 + )); +} + #[tokio::test] async fn legacy_initial_frame_validates_against_isolated_relay_key() { let _process_credential_guard = PROCESS_CREDENTIAL_TEST_LOCK.lock().await; @@ -194,13 +207,6 @@ async fn legacy_initial_frame_validates_against_isolated_relay_key() { .await .unwrap(); - set_process_msg_header_key(Some( - std::str::from_utf8(&isolated_admin).expect("printable isolated key"), - )) - .unwrap(); - let client = ClientHeaderSession::new_legacy(isolated_admin); - let bytes = encode_initial(&client, b"legacy-isolated").await; - let temporary_key_id = 1; let temporary_key = *b"temporary-remote-key-0123456789a"; set_process_msg_header_key(Some(&encode_temporary_credential( @@ -210,10 +216,25 @@ async fn legacy_initial_frame_validates_against_isolated_relay_key() { .unwrap(); let security = ServerSecurity::new(auth); - let mut input = std::io::Cursor::new(bytes); - let initial = security.read_initial(&mut input).await.unwrap(); - assert_eq!(initial.payload, b"legacy-isolated"); - assert_eq!(initial.session.protocol(), HeaderProtocol::Legacy); + let (mut client_io, mut server_io) = tokio::io::duplex(4096); + let client = ClientHeaderSession::new_legacy(isolated_admin); + let client_task = async { + client + .write_initial(&mut client_io, b"legacy-isolated") + .await + .unwrap(); + let mut reader = client.response_reader(&mut client_io).unwrap(); + reader.read_msg().await.unwrap().to_vec() + }; + let server_task = async { + let initial = security.read_initial(&mut server_io).await.unwrap(); + assert_eq!(initial.payload, b"legacy-isolated"); + assert_eq!(initial.session.protocol(), HeaderProtocol::Legacy); + let mut writer = initial.session.response_writer(&mut server_io).unwrap(); + writer.write_msg(b"legacy-response").await.unwrap(); + }; + let (response, _) = tokio::join!(client_task, server_task); + assert_eq!(response, b"legacy-response"); set_process_msg_header_key(None).unwrap(); let _ = std::fs::remove_dir_all(config.state_dir); diff --git a/src/pb_server/connection.rs b/src/pb_server/connection.rs index a80168e..c05ff08 100644 --- a/src/pb_server/connection.rs +++ b/src/pb_server/connection.rs @@ -329,15 +329,26 @@ pub(super) async fn handle_conn( status = ?status, "received pb init request" ); - handle_show_status( - status, - effective_namespace, - manager_task_sender, - conn_id, - conn, - session, - ) - .await?; + let cancellation = match auth_context.cancellation_token() { + Ok(token) => token, + Err(failure) => { + write_protocol_error(&mut conn, &session, &failure).await; + return Ok(()); + } + }; + tokio::select! { + result = handle_show_status( + status, + effective_namespace, + manager_task_sender, + conn_id, + conn, + session, + ) => result?, + _ = cancellation.cancelled() => { + tracing::info!(event = "connection_auth_expired", key_id = auth_context.key_id, conn_id = %conn_id, "closing status request"); + } + } } PbConnRequest::Admin(request) => { if !auth_context.is_admin { From 460c540969ace402eb784c55d1b7185e244c0b75 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 01:47:03 +0800 Subject: [PATCH 18/74] Share auth and protocol helpers for later reuse Centralize env-safe administrator-key checks, isolated-relay legacy codec construction, credential-cancellation races, and snapshot/WAL paths. Comments explain the setenv, clock-skew, and force-init invariants that those helpers encode. --- CHANGELOG.md | 1 + src/common/auth.rs | 9 +- src/common/auth/actor.rs | 17 +-- src/common/auth/persistence.rs | 31 +++-- src/common/checksum.rs | 31 ++++- src/common/message/mod.rs | 32 +++-- src/common/message/secure.rs | 87 +++++++----- src/common/message/secure/replay.rs | 3 + src/pb_server/connection.rs | 207 ++++++++++++++++++---------- 9 files changed, 264 insertions(+), 154 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a27019e..c1e34d4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ All notable changes to this project will be documented in this file. - Replaced insecure default-key fallback with first-start random administrator-key generation, retained machine-derived keys only for explicit compatibility, and updated Flutter, installers, systemd, Docker, release metadata, and bilingual documentation. - Fixed remaining review findings: installer migration now honors `MSG_HEADER_KEY` from `/etc/pb-mapper/server.env`, isolated relays validate legacy frames with their own administrator key, first-flight replay retention covers the full clock-skew window, and desktop macOS/Windows servers use a user-writable auth directory. - Rejected NUL/non-printable rotated administrator keys, cancelled in-flight status reads on credential revocation, refused `--force-init-admin-key` when encrypted auth state already exists, bound isolated-relay legacy continuation checksums to the relay key, and doubled first-flight Bloom retention so a max-future timestamp cannot outlive the filter. +- Centralized env-safe administrator-key checks, isolated-relay legacy codec construction, credential-cancellation races, and auth snapshot/WAL paths so later protocol and lifecycle changes reuse one implementation. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/src/common/auth.rs b/src/common/auth.rs index b01bcdf..ce5d9f8 100644 --- a/src/common/auth.rs +++ b/src/common/auth.rs @@ -37,9 +37,9 @@ use tokio::sync::{mpsc, oneshot}; use tokio_util::sync::CancellationToken; use super::checksum::{ - encode_temporary_credential, get_process_credential, parse_credential, - set_process_msg_header_key, AesKeyType, Credential, ENV_MSG_HEADER_KEY, - MACHINE_MSG_HEADER_KEY_PATH, + encode_temporary_credential, env_safe_admin_key_error, get_process_credential, + is_env_safe_admin_key, parse_credential, set_process_msg_header_key, AesKeyType, Credential, + ENV_MSG_HEADER_KEY, MACHINE_MSG_HEADER_KEY_PATH, }; pub const ADMIN_NAMESPACE: u64 = 0; @@ -88,6 +88,9 @@ pub fn default_auth_state_dir() -> PathBuf { .unwrap_or_else(platform_default_auth_state_dir) } +/// Linux systemd/Docker keep `/var/lib/pb-mapper/auth`. Desktop macOS and +/// Windows binaries run as a normal user, so they need an application data +/// directory instead of a root-owned system path. pub(crate) fn platform_default_auth_state_dir() -> PathBuf { #[cfg(windows)] { diff --git a/src/common/auth/actor.rs b/src/common/auth/actor.rs index 874f177..30e81ee 100644 --- a/src/common/auth/actor.rs +++ b/src/common/auth/actor.rs @@ -701,24 +701,15 @@ fn actor_rotate_root( false, )); } - let new_key_string = String::from_utf8(new_key.to_vec()).map_err(|_| { - AuthFailure::new( - "administrator_key_invalid", - "administrator key must be 32 UTF-8 bytes for MSG_HEADER_KEY compatibility", - false, - ) - })?; - if !new_key_string - .as_bytes() - .iter() - .all(|byte| byte.is_ascii_graphic()) - { + if !is_env_safe_admin_key(&new_key) { return Err(AuthFailure::new( "administrator_key_invalid", - "administrator key must be 32 printable ASCII bytes without whitespace or NUL", + env_safe_admin_key_error(), false, )); } + let new_key_string = + String::from_utf8(new_key.to_vec()).expect("printable ASCII is valid UTF-8"); let rotate_audit = audit("administrator_key_rotate", None, None); let mut snapshot = empty_snapshot(inner, inner.instance_id(), &VecDeque::new()); diff --git a/src/common/auth/persistence.rs b/src/common/auth/persistence.rs index 216e7f5..ac41e6d 100644 --- a/src/common/auth/persistence.rs +++ b/src/common/auth/persistence.rs @@ -12,6 +12,21 @@ use super::*; +pub(super) const AUTH_SNAPSHOT_FILE: &str = "auth.snapshot"; +pub(super) const AUTH_WAL_FILE: &str = "auth.wal"; + +pub(super) fn auth_snapshot_path(state_dir: &Path) -> PathBuf { + state_dir.join(AUTH_SNAPSHOT_FILE) +} + +pub(super) fn auth_wal_path(state_dir: &Path) -> PathBuf { + state_dir.join(AUTH_WAL_FILE) +} + +pub fn encrypted_auth_state_exists(state_dir: &Path) -> bool { + auth_snapshot_path(state_dir).exists() || auth_wal_path(state_dir).exists() +} + pub(super) fn push_audit_record(inner: &AuthStateInner, record: AuditRecord) { let mut records = inner .audit_records @@ -161,7 +176,7 @@ pub(super) fn try_load_persisted_state( admin_key: &AesKeyType, instance_id: [u8; INSTANCE_ID_LEN], ) -> Result { - let snapshot_path = config.state_dir.join("auth.snapshot"); + let snapshot_path = auth_snapshot_path(&config.state_dir); let mut snapshot = if snapshot_path.exists() { let bytes = std::fs::read(&snapshot_path).map_err(|error| { AuthFailure::new( @@ -199,7 +214,7 @@ pub(super) fn try_load_persisted_state( snapshot.generations.resize(config.max_temporary_keys, 0); snapshot.generations.truncate(config.max_temporary_keys); - let wal_path = config.state_dir.join("auth.wal"); + let wal_path = auth_wal_path(&config.state_dir); if wal_path.exists() { for record in read_wal(&wal_path, admin_key)? { match record { @@ -318,7 +333,7 @@ pub(super) fn append_wal( let sealed = seal_blob(admin_key, &plain)?; let length = u32::try_from(sealed.len()) .map_err(|_| AuthFailure::internal("auth WAL record is too large"))?; - let path = config.state_dir.join("auth.wal"); + let path = auth_wal_path(&config.state_dir); let mut file = OpenOptions::new() .create(true) .append(true) @@ -418,9 +433,9 @@ pub(super) fn write_snapshot_and_truncate_wal( AuthFailure::internal(format!("failed to encode auth snapshot: {error}")) })?; let sealed = seal_blob(admin_key, &plain)?; - let snapshot_path = config.state_dir.join("auth.snapshot"); + let snapshot_path = auth_snapshot_path(&config.state_dir); atomic_write(&snapshot_path, &sealed, 0o600)?; - let wal_path = config.state_dir.join("auth.wal"); + let wal_path = auth_wal_path(&config.state_dir); let wal = OpenOptions::new() .create(true) .write(true) @@ -590,11 +605,11 @@ pub fn initialize_admin_key(path: &Path, force: bool) -> Result bool { + bytes.len() == ADMIN_KEY_LEN && bytes.iter().all(|byte| byte.is_ascii_graphic()) +} + +pub fn env_safe_admin_key_error() -> String { + format!( + "`{ENV_MSG_HEADER_KEY}` administrator key must be 32 printable ASCII bytes without whitespace or NUL" + ) +} const DERIVE_MSG_HEADER_KEY_TAG: &str = "pb-mapper-msg-header-key-v1"; const DERIVE_MSG_HEADER_KEY_CHARSET: &[u8] = @@ -192,13 +205,11 @@ pub fn parse_credential(raw: &str) -> Result { } let bytes = raw.as_bytes(); - if bytes.len() != 32 { + if bytes.len() != ADMIN_KEY_LEN { return Err(key_len_error(raw)); } - if !bytes.iter().all(|byte| byte.is_ascii_graphic()) { - return Err(format!( - "`{ENV_MSG_HEADER_KEY}` administrator key must be 32 printable ASCII bytes without whitespace or NUL" - )); + if !is_env_safe_admin_key(bytes) { + return Err(env_safe_admin_key_error()); } Ok(Credential::Admin( bytes.try_into().expect("validated admin key width"), @@ -552,6 +563,16 @@ mod tests { )); } + #[test] + fn env_safe_admin_key_rejects_nul_and_accepts_printable_ascii() { + use super::*; + assert!(is_env_safe_admin_key(b"0123456789abcdefghijklmnopqrstuv")); + let mut with_nul = *b"0123456789abcdefghijklmnopqrstuv"; + with_nul[8] = 0; + assert!(!is_env_safe_admin_key(&with_nul)); + assert!(!is_env_safe_admin_key(b"short")); + } + #[test] fn administrator_credentials_reject_nul_and_whitespace() { use super::*; diff --git a/src/common/message/mod.rs b/src/common/message/mod.rs index 782d8f4..8a4ca9f 100644 --- a/src/common/message/mod.rs +++ b/src/common/message/mod.rs @@ -133,6 +133,10 @@ gen_write_network_with_error!( &[u8] ); +fn checksum_key_bytes(key: &Option) -> Option<&[u8]> { + key.as_ref().map(|key| key.as_slice()) +} + #[inline] fn checksum_matches(datalen: DataLenType, checksum: u32, key: Option<&[u8]>) -> bool { match key { @@ -201,11 +205,7 @@ impl<'a, T: AsyncReadExt + Unpin> NormalMessageReader<'a, T> { } async fn read_msg_inner(&mut self) -> Result<&'_ [u8]> { - let datalen = get_msg_len( - &mut self.reader, - self.checksum_key.as_ref().map(|key| key.as_slice()), - ) - .await?; + let datalen = get_msg_len(&mut self.reader, checksum_key_bytes(&self.checksum_key)).await?; self.buffer.fixed_resize(datalen as usize); let n = read_msg_body(&mut self.reader, self.buffer.buffer_mut()).await?; Ok(&self.buffer.buffer()[0..n]) @@ -235,7 +235,7 @@ impl<'a, T: AsyncWriteExt + Unpin> NormalMessageWriter<'a, T> { set_msg_len( &mut self.writer, msg.len() as u32, - self.checksum_key.as_ref().map(|key| key.as_slice()), + checksum_key_bytes(&self.checksum_key), ) .await?; @@ -262,6 +262,13 @@ impl<'a, T: AsyncReadExt + Unpin, D: Decryptor> CodecMessageReader<'a, T, D> { } } + /// Bind the length checksum to `key` instead of the process credential. + /// Isolated relays keep a remote `MSG_HEADER_KEY` while speaking with a + /// different local administrator key. + pub fn for_session_key(reader: &'a mut T, decryptor: D, key: AesKeyType) -> Self { + Self::new(reader, decryptor).with_checksum_key(key) + } + pub fn with_checksum_key(mut self, key: AesKeyType) -> Self { self.reader.checksum_key = Some(key); self @@ -287,6 +294,8 @@ impl<'a, T: AsyncReadExt + Unpin, D: Decryptor> MessageReader for CodecMessageRe pub struct CodecMessageWriter<'a, T: AsyncWriteExt + Unpin, E: Encryptor> { writer: &'a mut T, encryptor: E, + /// `None` uses the process `MSG_HEADER_KEY` hash. Isolated relays must set + /// this to the session key so continuation frames stay decryptable. checksum_key: Option, } @@ -299,6 +308,10 @@ impl<'a, T: AsyncWriteExt + Unpin, E: Encryptor> CodecMessageWriter<'a, T, E> { } } + pub fn for_session_key(writer: &'a mut T, encryptor: E, key: AesKeyType) -> Self { + Self::new(writer, encryptor).with_checksum_key(key) + } + pub fn with_checksum_key(mut self, key: AesKeyType) -> Self { self.checksum_key = Some(key); self @@ -321,12 +334,7 @@ impl<'a, T: AsyncWriteExt + Unpin, E: Encryptor> MessageWriter for CodecMessageW })?; let msg_len = (buf.len() + tag.as_ref().len()) as DataLenType; - set_msg_len( - self.writer, - msg_len, - self.checksum_key.as_ref().map(|key| key.as_slice()), - ) - .await?; + set_msg_len(self.writer, msg_len, checksum_key_bytes(&self.checksum_key)).await?; write_codec_msg(self.writer, &buf).await?; write_codec_tag(self.writer, tag.as_ref()).await } diff --git a/src/common/message/secure.rs b/src/common/message/secure.rs index 4388458..f697b79 100644 --- a/src/common/message/secure.rs +++ b/src/common/message/secure.rs @@ -44,6 +44,9 @@ const FRAME_HEADER_LEN: usize = 12; const DIRECTION_CLIENT_TO_SERVER: u8 = 0; const DIRECTION_SERVER_TO_CLIENT: u8 = 1; const MAX_CONNECTION_CLOCK_SKEW_SECONDS: u64 = 5 * 60; +/// Each Bloom generation must outlive the accepted clock-skew interval. A +/// salt inserted at the end of a window with `ts = now + skew` stays valid +/// until `insert + 2*skew`, so one generation is `2 * skew`. const DEFAULT_REPLAY_WINDOW_SECONDS: u64 = MAX_CONNECTION_CLOCK_SKEW_SECONDS.saturating_mul(2); const DEFAULT_REPLAY_FILTER_BYTES: usize = 1024 * 1024; const MAX_INITIAL_PLAINTEXT_LEN: u32 = 64 * 1024; @@ -102,10 +105,7 @@ impl ClientHeaderSession { ) -> Result<()> { match self.protocol { HeaderProtocol::Legacy => { - let codec = Aes256GcmEnCodec::try_new(&self.legacy_key) - .map_err(|_| protocol_error("failed to initialize legacy writer"))?; - CodecMessageWriter::new(writer, codec) - .with_checksum_key(self.legacy_key) + legacy_message_writer(writer, &self.legacy_key, "legacy writer")? .write_msg(message) .await } @@ -129,14 +129,11 @@ impl ClientHeaderSession { reader: &'a mut T, ) -> Result> { match self.protocol { - HeaderProtocol::Legacy => Ok(HeaderMessageReader::Legacy( - CodecMessageReader::new( - reader, - Aes256GcmDeCodec::try_new(&self.legacy_key) - .map_err(|_| protocol_error("failed to initialize legacy reader"))?, - ) - .with_checksum_key(self.legacy_key), - )), + HeaderProtocol::Legacy => Ok(HeaderMessageReader::Legacy(legacy_message_reader( + reader, + &self.legacy_key, + "legacy reader", + )?)), HeaderProtocol::V2 => Ok(HeaderMessageReader::V2(V2MessageReader::new( reader, self.v2.as_ref().expect("v2 session material").clone(), @@ -151,14 +148,11 @@ impl ClientHeaderSession { writer: &'a mut T, ) -> Result> { match self.protocol { - HeaderProtocol::Legacy => Ok(HeaderMessageWriter::Legacy( - CodecMessageWriter::new( - writer, - Aes256GcmEnCodec::try_new(&self.legacy_key) - .map_err(|_| protocol_error("failed to initialize legacy writer"))?, - ) - .with_checksum_key(self.legacy_key), - )), + HeaderProtocol::Legacy => Ok(HeaderMessageWriter::Legacy(legacy_message_writer( + writer, + &self.legacy_key, + "legacy writer", + )?)), HeaderProtocol::V2 => Ok(HeaderMessageWriter::V2(V2MessageWriter::new( writer, self.v2.as_ref().expect("v2 session material").clone(), @@ -222,15 +216,11 @@ impl ServerHeaderSession { writer: &'a mut T, ) -> Result> { match self.protocol { - HeaderProtocol::Legacy => Ok(HeaderMessageWriter::Legacy( - CodecMessageWriter::new( - writer, - Aes256GcmEnCodec::try_new(&self.legacy_key).map_err(|_| { - protocol_error("failed to initialize legacy response writer") - })?, - ) - .with_checksum_key(self.legacy_key), - )), + HeaderProtocol::Legacy => Ok(HeaderMessageWriter::Legacy(legacy_message_writer( + writer, + &self.legacy_key, + "legacy response writer", + )?)), HeaderProtocol::V2 => Ok(HeaderMessageWriter::V2(V2MessageWriter::new( writer, self.v2.as_ref().expect("v2 session material").clone(), @@ -245,14 +235,11 @@ impl ServerHeaderSession { reader: &'a mut T, ) -> Result> { match self.protocol { - HeaderProtocol::Legacy => Ok(HeaderMessageReader::Legacy( - CodecMessageReader::new( - reader, - Aes256GcmDeCodec::try_new(&self.legacy_key) - .map_err(|_| protocol_error("failed to initialize legacy reader"))?, - ) - .with_checksum_key(self.legacy_key), - )), + HeaderProtocol::Legacy => Ok(HeaderMessageReader::Legacy(legacy_message_reader( + reader, + &self.legacy_key, + "legacy reader", + )?)), HeaderProtocol::V2 => Ok(HeaderMessageReader::V2(V2MessageReader::new( reader, self.v2.as_ref().expect("v2 session material").clone(), @@ -626,6 +613,32 @@ use frame::{derive_material, first_prefix, V2Material}; pub use frame::{V2MessageReader, V2MessageWriter}; mod replay; use replay::{replay_fingerprint, RotatingBloom}; +fn legacy_message_reader<'a, T: AsyncReadExt + Unpin>( + reader: &'a mut T, + key: &AesKeyType, + action: &str, +) -> Result> { + Ok(CodecMessageReader::for_session_key( + reader, + Aes256GcmDeCodec::try_new(key) + .map_err(|_| protocol_error(format!("failed to initialize {action}")))?, + *key, + )) +} + +fn legacy_message_writer<'a, T: AsyncWriteExt + Unpin>( + writer: &'a mut T, + key: &AesKeyType, + action: &str, +) -> Result> { + Ok(CodecMessageWriter::for_session_key( + writer, + Aes256GcmEnCodec::try_new(key) + .map_err(|_| protocol_error(format!("failed to initialize {action}")))?, + *key, + )) +} + fn protocol_error(detail: impl Into) -> Error { Error::MsgProtocol { detail: detail.into(), diff --git a/src/common/message/secure/replay.rs b/src/common/message/secure/replay.rs index e12c1db..d33c882 100644 --- a/src/common/message/secure/replay.rs +++ b/src/common/message/secure/replay.rs @@ -8,6 +8,9 @@ //! `check_and_insert` is called while one mutex is held, making concurrent admission //! atomic. This Bloom filter protects all connection types from immediate duplicates; //! administrator mutations additionally use the exact durable replay set in `auth`. +//! +//! Each generation lasts `2 *` the accepted clock-skew so a salt inserted at the +//! end of a window with a max-future timestamp cannot be replayed after rotation. use super::*; diff --git a/src/pb_server/connection.rs b/src/pb_server/connection.rs index c05ff08..7502ebe 100644 --- a/src/pb_server/connection.rs +++ b/src/pb_server/connection.rs @@ -10,9 +10,9 @@ //! +-> status / administrator request //! ``` //! -//! Long-lived register and subscribe futures are raced against the credential's -//! cancellation token here. This outer guard closes a subscriber even when the paired -//! service stream belongs to a different credential. +//! Long-lived register, subscribe, and status futures are raced against the +//! credential's cancellation token here. This outer guard closes a subscriber even +//! when the paired service stream belongs to a different credential. use super::*; @@ -223,37 +223,39 @@ pub(super) async fn handle_conn( is_datagram, "received pb init request" ); - let key = match scoped_service_key(&auth_context, effective_namespace, &key) { - Ok(key) => key, - Err(failure) => { - write_protocol_error(&mut conn, &session, &failure).await; - return Ok(()); - } - }; - let cancellation = match auth_context.cancellation_token() { - Ok(token) => token, - Err(failure) => { - write_protocol_error(&mut conn, &session, &failure).await; - return Ok(()); - } + let Some(key) = scope_service_or_reject( + &mut conn, + &session, + &auth_context, + effective_namespace, + &key, + ) + .await? + else { + return Ok(()); }; - tokio::select! { - result = handle_server_conn( - ServerRegistration { - key, - need_codec, - is_datagram, - protocol_version, - conn_id, - }, - manager_task_sender, + run_while_credential_active( conn, session, - ) => result?, - _ = cancellation.cancelled() => { - tracing::info!(event = "connection_auth_expired", key_id = auth_context.key_id, conn_id = %conn_id, "closing registered service connection"); - } - } + &auth_context, + conn_id, + "registered service connection", + |conn, session| { + handle_server_conn( + ServerRegistration { + key, + need_codec, + is_datagram, + protocol_version, + conn_id, + }, + manager_task_sender, + conn, + session, + ) + }, + ) + .await?; } PbConnRequest::Subcribe { key } => { tracing::info!( @@ -264,26 +266,28 @@ pub(super) async fn handle_conn( key = %key, "received pb init request" ); - let key = match scoped_service_key(&auth_context, effective_namespace, &key) { - Ok(key) => key, - Err(failure) => { - write_protocol_error(&mut conn, &session, &failure).await; - return Ok(()); - } - }; - let cancellation = match auth_context.cancellation_token() { - Ok(token) => token, - Err(failure) => { - write_protocol_error(&mut conn, &session, &failure).await; - return Ok(()); - } + let Some(key) = scope_service_or_reject( + &mut conn, + &session, + &auth_context, + effective_namespace, + &key, + ) + .await? + else { + return Ok(()); }; - tokio::select! { - result = handle_client_conn(key, conn_id, manager_task_sender, conn, session) => result?, - _ = cancellation.cancelled() => { - tracing::info!(event = "connection_auth_expired", key_id = auth_context.key_id, conn_id = %conn_id, "closing subscribed data connection"); - } - } + run_while_credential_active( + conn, + session, + &auth_context, + conn_id, + "subscribed data connection", + |conn, session| { + handle_client_conn(key, conn_id, manager_task_sender, conn, session) + }, + ) + .await?; } PbConnRequest::Stream { key, @@ -300,12 +304,16 @@ pub(super) async fn handle_conn( server_generation, "received pb init request" ); - let key = match scoped_service_key(&auth_context, effective_namespace, &key) { - Ok(key) => key, - Err(failure) => { - write_protocol_error(&mut conn, &session, &failure).await; - return Ok(()); - } + let Some(key) = scope_service_or_reject( + &mut conn, + &session, + &auth_context, + effective_namespace, + &key, + ) + .await? + else { + return Ok(()); }; manager_task_sender .send(ManagerTask::Stream { @@ -329,26 +337,24 @@ pub(super) async fn handle_conn( status = ?status, "received pb init request" ); - let cancellation = match auth_context.cancellation_token() { - Ok(token) => token, - Err(failure) => { - write_protocol_error(&mut conn, &session, &failure).await; - return Ok(()); - } - }; - tokio::select! { - result = handle_show_status( - status, - effective_namespace, - manager_task_sender, - conn_id, - conn, - session, - ) => result?, - _ = cancellation.cancelled() => { - tracing::info!(event = "connection_auth_expired", key_id = auth_context.key_id, conn_id = %conn_id, "closing status request"); - } - } + run_while_credential_active( + conn, + session, + &auth_context, + conn_id, + "status request", + |conn, session| { + handle_show_status( + status, + effective_namespace, + manager_task_sender, + conn_id, + conn, + session, + ) + }, + ) + .await?; } PbConnRequest::Admin(request) => { if !auth_context.is_admin { @@ -412,6 +418,55 @@ pub(super) async fn handle_conn( Ok(()) } +async fn scope_service_or_reject( + conn: &mut TcpStream, + session: &ServerHeaderSession, + auth_context: &AuthContext, + namespace: u64, + service_name: &str, +) -> Result> { + match scoped_service_key(auth_context, namespace, service_name) { + Ok(key) => Ok(Some(key)), + Err(failure) => { + write_protocol_error(conn, session, &failure).await; + Ok(None) + } + } +} + +async fn run_while_credential_active( + mut conn: TcpStream, + session: ServerHeaderSession, + auth_context: &AuthContext, + conn_id: RemoteConnId, + closed_what: &'static str, + work: F, +) -> Result<()> +where + F: FnOnce(TcpStream, ServerHeaderSession) -> Fut, + Fut: std::future::Future>, +{ + let cancellation = match auth_context.cancellation_token() { + Ok(token) => token, + Err(failure) => { + write_protocol_error(&mut conn, &session, &failure).await; + return Ok(()); + } + }; + tokio::select! { + result = work(conn, session) => result?, + _ = cancellation.cancelled() => { + tracing::info!( + event = "connection_auth_expired", + key_id = auth_context.key_id, + conn_id = %conn_id, + "closing {closed_what}" + ); + } + } + Ok(()) +} + async fn write_protocol_error( conn: &mut TcpStream, session: &ServerHeaderSession, From 4d68a2c7551dc1e48964b4e90dcd85f136a39bc8 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 02:01:18 +0800 Subject: [PATCH 19/74] Harden key initialization, slot generations, WAL, and stream auth Refuse admin-key creation whenever encrypted auth state remains, keep discarded slot generations across capacity changes, roll back or fail-closed uncertain WAL appends, authenticate first flights before consuming the replay filter, and reuse the registration credential for provider streams. --- CHANGELOG.md | 1 + src/common/auth.rs | 4 + src/common/auth/actor.rs | 11 +- src/common/auth/persistence.rs | 164 +++++++++++++++++++++-------- src/common/auth/runtime.rs | 6 ++ src/common/auth/tests.rs | 76 +++++++++++++ src/common/checksum.rs | 3 +- src/common/message/secure.rs | 15 ++- src/common/message/secure/tests.rs | 43 +++++++- src/local/server/mod.rs | 25 +++-- src/local/server/stream.rs | 26 +++-- 11 files changed, 304 insertions(+), 70 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c1e34d4..80f52ae 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,7 @@ All notable changes to this project will be documented in this file. - Fixed remaining review findings: installer migration now honors `MSG_HEADER_KEY` from `/etc/pb-mapper/server.env`, isolated relays validate legacy frames with their own administrator key, first-flight replay retention covers the full clock-skew window, and desktop macOS/Windows servers use a user-writable auth directory. - Rejected NUL/non-printable rotated administrator keys, cancelled in-flight status reads on credential revocation, refused `--force-init-admin-key` when encrypted auth state already exists, bound isolated-relay legacy continuation checksums to the relay key, and doubled first-flight Bloom retention so a max-future timestamp cannot outlive the filter. - Centralized env-safe administrator-key checks, isolated-relay legacy codec construction, credential-cancellation races, and auth snapshot/WAL paths so later protocol and lifecycle changes reuse one implementation. +- Refused administrator-key initialization whenever encrypted auth state is present, preserved discarded slot generations across capacity changes, rolled back or fail-closed uncertain WAL appends, authenticated first flights before consuming the replay filter, and reused the registration credential for provider streams. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/src/common/auth.rs b/src/common/auth.rs index ce5d9f8..f1c26e6 100644 --- a/src/common/auth.rs +++ b/src/common/auth.rs @@ -390,6 +390,10 @@ struct AuthStateInner { sync_process_credential: bool, instance_id: RwLock<[u8; INSTANCE_ID_LEN]>, slots: RwLock>, + /// Generations and entries for slots above the current capacity. Kept so a + /// later capacity increase cannot reuse a discarded slot's key id. + high_slot_generations: RwLock>, + high_slot_entries: RwLock>, safe_mode: AtomicBool, legacy_protocol_allowed: AtomicBool, active_legacy_connections: AtomicU64, diff --git a/src/common/auth/actor.rs b/src/common/auth/actor.rs index 30e81ee..47d7f30 100644 --- a/src/common/auth/actor.rs +++ b/src/common/auth/actor.rs @@ -300,10 +300,13 @@ fn actor_claim_admin_mutation( fingerprint, client_timestamp, }; - append_wal( - config, - &inner.admin_key(), - &WalRecord::AdminReplay(record.clone()), + fail_closed_on_uncertain_wal( + inner, + append_wal( + config, + &inner.admin_key(), + &WalRecord::AdminReplay(record.clone()), + ), )?; admin_replays.insert(fingerprint); admin_replay_order.push_back(record); diff --git a/src/common/auth/persistence.rs b/src/common/auth/persistence.rs index ac41e6d..44b8eb1 100644 --- a/src/common/auth/persistence.rs +++ b/src/common/auth/persistence.rs @@ -48,6 +48,34 @@ pub(super) fn cancel_all_temporary_leases(inner: &AuthStateInner) { } } +fn snapshot_generations(inner: &AuthStateInner) -> Vec { + let slots = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let extra = inner + .high_slot_generations + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let mut generations = slots.iter().map(|slot| slot.generation).collect::>(); + generations.extend_from_slice(&extra); + generations +} + +pub(super) fn split_high_slot_state( + snapshot: &PersistedSnapshot, + capacity: usize, +) -> (Vec, Vec) { + let high_generations = snapshot.generations.get(capacity..).unwrap_or(&[]).to_vec(); + let high_entries = snapshot + .entries + .iter() + .filter(|entry| key_slot(entry.key_id) as usize >= capacity) + .cloned() + .collect(); + (high_generations, high_entries) +} + pub(super) fn build_snapshot( inner: &AuthStateInner, cold: &HashMap, @@ -57,8 +85,8 @@ pub(super) fn build_snapshot( .slots .read() .unwrap_or_else(|poisoned| poisoned.into_inner()); - let generations = slots.iter().map(|slot| slot.generation).collect(); - let entries = slots + let generations = snapshot_generations(inner); + let mut entries = slots .iter() .enumerate() .filter_map(|(index, slot)| { @@ -76,7 +104,15 @@ pub(super) fn build_snapshot( tombstoned_at: (cold.tombstoned_at != 0).then_some(cold.tombstoned_at), }) }) - .collect(); + .collect::>(); + entries.extend( + inner + .high_slot_entries + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .iter() + .cloned(), + ); PersistedSnapshot { schema_version: SNAPSHOT_SCHEMA_VERSION, instance_id: inner.instance_id(), @@ -128,14 +164,10 @@ pub(super) fn empty_snapshot( instance_id: [u8; INSTANCE_ID_LEN], admin_replays: &VecDeque, ) -> PersistedSnapshot { - let slots = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); PersistedSnapshot { schema_version: SNAPSHOT_SCHEMA_VERSION, instance_id, - generations: slots.iter().map(|slot| slot.generation).collect(), + generations: snapshot_generations(inner), entries: Vec::new(), legacy_protocol: if inner.legacy_protocol_allowed.load(Ordering::Acquire) { LegacyProtocolPolicy::Allow @@ -211,8 +243,9 @@ pub(super) fn try_load_persisted_state( false, )); } - snapshot.generations.resize(config.max_temporary_keys, 0); - snapshot.generations.truncate(config.max_temporary_keys); + if snapshot.generations.len() < config.max_temporary_keys { + snapshot.generations.resize(config.max_temporary_keys, 0); + } let wal_path = auth_wal_path(&config.state_dir); if wal_path.exists() { @@ -238,14 +271,17 @@ pub(super) fn apply_persisted_mutation( match mutation { StateMutation::Issue(entry) => { let index = key_slot(entry.key_id) as usize; - if index >= capacity { - return Err(AuthFailure::new( - "temporary_key_store_unavailable", - "WAL issue record references a slot outside the configured capacity", - false, - )); + if snapshot.generations.len() <= index { + snapshot.generations.resize(index + 1, 0); } snapshot.generations[index] = key_generation(entry.key_id); + if index >= capacity { + snapshot + .entries + .retain(|current| key_slot(current.key_id) as usize != index); + snapshot.entries.push(entry); + return Ok(()); + } snapshot .entries .retain(|current| key_slot(current.key_id) as usize != index); @@ -287,19 +323,35 @@ pub(super) fn apply_persisted_mutation( Ok(()) } +pub(super) fn fail_closed_on_uncertain_wal( + inner: &AuthStateInner, + result: Result<(), AuthFailure>, +) -> Result<(), AuthFailure> { + if let Err(error) = &result { + if !error.retryable { + inner.safe_mode.store(true, Ordering::Release); + cancel_all_temporary_leases(inner); + } + } + result +} + pub(super) fn append_mutation( config: &AuthConfig, inner: &AuthStateInner, mutation: StateMutation, audit: AuditRecord, ) -> Result<(), AuthFailure> { - append_wal( - config, - &inner.admin_key(), - &WalRecord::Mutation { - mutation, - audit: audit.clone(), - }, + fail_closed_on_uncertain_wal( + inner, + append_wal( + config, + &inner.admin_key(), + &WalRecord::Mutation { + mutation, + audit: audit.clone(), + }, + ), )?; push_audit_record(inner, audit); Ok(()) @@ -310,7 +362,10 @@ pub(super) fn append_audit( inner: &AuthStateInner, audit: AuditRecord, ) -> Result<(), AuthFailure> { - append_wal(config, &inner.admin_key(), &WalRecord::Audit(audit.clone()))?; + fail_closed_on_uncertain_wal( + inner, + append_wal(config, &inner.admin_key(), &WalRecord::Audit(audit.clone())), + )?; push_audit_record(inner, audit); Ok(()) } @@ -354,16 +409,39 @@ pub(super) fn append_wal( false, ) })?; - file.write_all(&length.to_be_bytes()) - .and_then(|()| file.write_all(&sealed)) - .and_then(|()| file.sync_data()) + let start_len = file + .metadata() + .map(|metadata| metadata.len()) .map_err(|error| { AuthFailure::new( "temporary_key_store_unavailable", - format!("failed to durably append `{}`: {error}", path.display()), + format!("failed to inspect `{}`: {error}", path.display()), true, ) - }) + })?; + if let Err(error) = file + .write_all(&length.to_be_bytes()) + .and_then(|()| file.write_all(&sealed)) + .and_then(|()| file.sync_data()) + { + let rolled_back = file + .set_len(start_len) + .and_then(|()| file.sync_data()) + .is_ok(); + return Err(AuthFailure::new( + "temporary_key_store_unavailable", + if rolled_back { + format!("failed to durably append `{}`: {error}", path.display()) + } else { + format!( + "failed to durably append `{}` and could not restore the previous WAL length: {error}", + path.display() + ) + }, + rolled_back, + )); + } + Ok(()) } pub(super) fn read_wal(path: &Path, admin_key: &AesKeyType) -> Result, AuthFailure> { @@ -605,20 +683,20 @@ pub fn initialize_admin_key(path: &Path, force: bool) -> Result AUDIT_RECORD_CAPACITY { audit_records.pop_front(); } + let (high_slot_generations, high_slot_entries) = loaded + .as_ref() + .map(|state| split_high_slot_state(state, config.max_temporary_keys)) + .unwrap_or_default(); let inner = Arc::new(AuthStateInner { admin: RwLock::new(AdminState { key: admin_key, @@ -153,6 +157,8 @@ impl AuthRuntime { sync_process_credential, instance_id: RwLock::new(instance_id), slots: RwLock::new(slots), + high_slot_generations: RwLock::new(high_slot_generations), + high_slot_entries: RwLock::new(high_slot_entries), safe_mode: AtomicBool::new(safe_mode), legacy_protocol_allowed: AtomicBool::new(legacy_protocol.is_allowed()), active_legacy_connections: AtomicU64::new(0), diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index c9edfa5..c3a73cb 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -35,6 +35,82 @@ fn initialize_admin_key_refuses_to_replace_a_key_when_encrypted_state_exists() { std::fs::write(state_dir.join("auth.snapshot"), b"encrypted").unwrap(); let error = initialize_admin_key(&key_path, true).unwrap_err(); assert_eq!(error.code, "administrator_key_state_exists"); + let missing = state_dir.join("missing-admin.key"); + let error = initialize_admin_key(&missing, false).unwrap_err(); + assert_eq!(error.code, "administrator_key_state_exists"); + let _ = std::fs::remove_dir_all(state_dir); +} + +#[tokio::test] +async fn shrinking_then_expanding_capacity_does_not_reuse_old_key_ids() { + let state_dir = temp_state_dir("capacity-shrink"); + let admin_key = *b"0123456789abcdefghijklmnopqrstuv"; + let config_two = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 2, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + let runtime = AuthRuntime::start(admin_key, config_two.clone()) + .await + .unwrap(); + let admin = authenticate_for_test(&runtime, 0).unwrap(); + let first = runtime + .issue(&admin, Duration::from_secs(60), Some("first".to_string())) + .await + .unwrap(); + let second = runtime + .issue(&admin, Duration::from_secs(60), Some("second".to_string())) + .await + .unwrap(); + let Credential::Temporary { + key_id: first_id, .. + } = parse_credential(&first.credential).unwrap() + else { + panic!("expected temporary credential"); + }; + let Credential::Temporary { + key_id: second_id, .. + } = parse_credential(&second.credential).unwrap() + else { + panic!("expected temporary credential"); + }; + runtime.revoke(&admin, first_id).await.unwrap(); + runtime.revoke(&admin, second_id).await.unwrap(); + runtime.gc(&admin).await.unwrap(); + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + + let config_one = AuthConfig { + max_temporary_keys: 1, + ..config_two.clone() + }; + let runtime = AuthRuntime::start(admin_key, config_one).await.unwrap(); + let admin = authenticate_for_test(&runtime, 0).unwrap(); + assert!(!runtime.status(&admin).await.unwrap().safe_mode); + let _third = runtime + .issue(&admin, Duration::from_secs(60), Some("third".to_string())) + .await + .unwrap(); + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + + let runtime = AuthRuntime::start(admin_key, config_two).await.unwrap(); + let admin = authenticate_for_test(&runtime, 0).unwrap(); + assert!(!runtime.status(&admin).await.unwrap().safe_mode); + let fourth = runtime + .issue(&admin, Duration::from_secs(60), Some("fourth".to_string())) + .await + .unwrap(); + let Credential::Temporary { + key_id: fourth_id, .. + } = parse_credential(&fourth.credential).unwrap() + else { + panic!("expected temporary credential"); + }; + assert_ne!(fourth_id, second_id); + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; let _ = std::fs::remove_dir_all(state_dir); } diff --git a/src/common/checksum.rs b/src/common/checksum.rs index 7dc8a4a..302349a 100644 --- a/src/common/checksum.rs +++ b/src/common/checksum.rs @@ -45,7 +45,7 @@ struct MsgHeaderKeyState { hash: AtomicU32, } -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum Credential { Admin(AesKeyType), Temporary { key_id: u64, key: AesKeyType }, @@ -141,7 +141,6 @@ pub fn get_process_credential() -> Result { .credential .read() .unwrap_or_else(|poisoned| poisoned.into_inner()) - .clone() .ok_or_else(|| { format!( "`{ENV_MSG_HEADER_KEY}` is required; no insecure default credential is available" diff --git a/src/common/message/secure.rs b/src/common/message/secure.rs index f697b79..45f7f0c 100644 --- a/src/common/message/secure.rs +++ b/src/common/message/secure.rs @@ -533,6 +533,13 @@ impl ServerSecurity { })? .to_vec(); + let context = self + .auth + .authenticate_presented(key_id, &key) + .map_err(|failure| ServerInitialError { + failure, + response_session: Some(session_without_context(&session)), + })?; let fingerprint = replay_fingerprint(key_id, &salt); let replayed = self .replay @@ -549,14 +556,6 @@ impl ServerSecurity { response_session: Some(session), }); } - - let context = self - .auth - .authenticate_presented(key_id, &key) - .map_err(|failure| ServerInitialError { - failure, - response_session: Some(session_without_context(&session)), - })?; session.context = Some(context); Ok(ServerInitialMessage { payload, diff --git a/src/common/message/secure/tests.rs b/src/common/message/secure/tests.rs index e6c5edd..c605e14 100644 --- a/src/common/message/secure/tests.rs +++ b/src/common/message/secure/tests.rs @@ -11,7 +11,9 @@ use super::*; use crate::common::auth::{AuthConfig, LegacyProtocolPolicy, PROCESS_CREDENTIAL_TEST_LOCK}; -use crate::common::checksum::{encode_temporary_credential, set_process_msg_header_key}; +use crate::common::checksum::{ + encode_temporary_credential, parse_credential, set_process_msg_header_key, +}; fn temp_config() -> AuthConfig { let mut random = [0_u8; 8]; @@ -134,6 +136,45 @@ async fn identical_initial_frames_are_admitted_only_once() { let _ = std::fs::remove_dir_all(config.state_dir); } +#[tokio::test] +async fn revoked_first_flights_do_not_consume_the_replay_filter() { + let admin = *b"0123456789abcdefghijklmnopqrstuv"; + let config = temp_config(); + let auth = AuthRuntime::start(admin, config.clone()).await.unwrap(); + let admin_context = auth.authenticate_presented(0, &admin).unwrap(); + let issued = auth + .issue(&admin_context, std::time::Duration::from_secs(60), None) + .await + .unwrap(); + let Credential::Temporary { key_id, key } = parse_credential(&issued.credential).unwrap() + else { + panic!("expected temporary credential"); + }; + let client = ClientHeaderSession::new_v2(&Credential::Temporary { key_id, key }).unwrap(); + let bytes = encode_initial(&client, b"revoked").await; + auth.revoke(&admin_context, key_id).await.unwrap(); + let security = ServerSecurity::new(auth); + + let first = match security + .read_initial(&mut std::io::Cursor::new(bytes.clone())) + .await + { + Ok(_) => panic!("revoked credential should fail"), + Err(error) => error, + }; + let second = match security + .read_initial(&mut std::io::Cursor::new(bytes)) + .await + { + Ok(_) => panic!("revoked credential should fail again"), + Err(error) => error, + }; + assert_eq!(first.failure.code, "temporary_key_revoked"); + assert_eq!(second.failure.code, "temporary_key_revoked"); + + let _ = std::fs::remove_dir_all(config.state_dir); +} + #[tokio::test] async fn oversized_initial_frame_is_rejected_before_reading_its_body() { let credential = Credential::Admin(*b"0123456789abcdefghijklmnopqrstuv"); diff --git a/src/local/server/mod.rs b/src/local/server/mod.rs index 891272a..1273de2 100644 --- a/src/local/server/mod.rs +++ b/src/local/server/mod.rs @@ -16,7 +16,8 @@ use self::error::{ EncodeRegisterReqSnafu, EncodeStreamAckMsgSnafu, ReadRegisterRespSnafu, ReadStreamReqSnafu, RegisterRespNotMatchSnafu, SendRegisterReqSnafu, WritePingMsgSnafu, WriteStreamAckMsgSnafu, }; -use self::stream::handle_stream; +use self::stream::{handle_stream, StreamConnect}; +use crate::common::checksum::{get_process_credential, Credential}; use crate::common::config::{ control_conn_pool_size, control_heartbeat_interval, control_heartbeat_tolerance, control_io_timeout, control_suspect_grace, registration_probe_timeout, @@ -150,6 +151,7 @@ struct StreamTarget { remote_addr: A, keep_alive: bool, namespace: Option, + credential: Credential, } fn duration_to_millis(duration: Duration) -> u64 { @@ -445,7 +447,14 @@ where // Start registration with a protocol-v2 first frame. The session is reused for all // subsequent control messages on this TCP connection. - let session = match ClientHeaderSession::from_process() { + let credential = match get_process_credential() { + Ok(credential) => credential, + Err(error) => { + tracing::error!("load registration credential failed: {error}"); + return Err(Status::ConnectRemote); + } + }; + let session = match ClientHeaderSession::new_v2(&credential) { Ok(session) => session, Err(error) => { tracing::error!("create manager protocol-v2 session failed: {error}"); @@ -649,6 +658,7 @@ where remote_addr, keep_alive, namespace, + credential, }, key.clone(), registration.conn_id, @@ -846,13 +856,16 @@ where tokio::spawn(async move { snafu_error_handle!( handle_stream::( - target.local_addr, - target.remote_addr, key, client_id, server_generation, - target.keep_alive, - target.namespace, + StreamConnect { + local_addr: target.local_addr, + remote_addr: target.remote_addr, + keep_alive: target.keep_alive, + namespace: target.namespace, + credential: target.credential, + }, ) .await ) diff --git a/src/local/server/stream.rs b/src/local/server/stream.rs index 741b7ca..5c9186c 100644 --- a/src/local/server/stream.rs +++ b/src/local/server/stream.rs @@ -10,6 +10,7 @@ use super::error::{ DecodePbConnStreamRespSnafu, EncodePbConnStreamReqSnafu, PbConnStreamRespNotMatchSnafu, ReadPbConnStreamRespSnafu, Result, WritePbConnStreamReqSnafu, }; +use crate::common::checksum::Credential; use crate::common::config::control_io_timeout; use crate::common::message::command::{MessageSerializer, PbConnRequest, PbConnResponse}; use crate::common::message::forward::StreamForward; @@ -20,6 +21,14 @@ use crate::snafu_error_handle; use uni_stream::addr::{each_addr, ToSocketAddrs}; use uni_stream::stream::{set_tcp_keep_alive, set_tcp_nodelay, StreamProvider, StreamSplit}; +pub struct StreamConnect { + pub local_addr: A, + pub remote_addr: A, + pub keep_alive: bool, + pub namespace: Option, + pub credential: Credential, +} + /// Handle a stream connection and establish a forward network traffic forwarding. /// This function handles both local and remote streams, sets up message writers and readers, /// and starts forwarding network traffic between the two endpoints. @@ -27,17 +36,21 @@ pub async fn handle_stream< LocalStream: StreamProvider, A: ToSocketAddrs + Debug + Copy + Clone + Send, >( - local_addr: A, - remote_addr: A, key: Arc, client_id: u32, server_generation: u64, - keep_alive: bool, - namespace: Option, + connect: StreamConnect, ) -> Result<()> where LocalStream::Item: StreamForward, { + let StreamConnect { + local_addr, + remote_addr, + keep_alive, + namespace, + credential, + } = connect; let key_ref = key.as_ref(); let client_id_span = info_span!("client_id", key_ref, client_id); let _enter = client_id_span.enter(); @@ -75,9 +88,10 @@ where } snafu_error_handle!(set_tcp_nodelay(&remote_stream), "remote stream set nodelay"); - // write stream request and read response + // Use the credential captured at registration. A later UI/process key + // change must not move these streams into another namespace. let codec_key = { - let session = ClientHeaderSession::from_process() + let session = ClientHeaderSession::new_v2(&credential) .context(CreateHeaderToolSnafu { action: "session" })?; match tokio::time::timeout(timeout, session.write_initial(&mut remote_stream, &msg)).await { Ok(result) => result.context(WritePbConnStreamReqSnafu)?, From e07aa5c70464e05854a22276e7978e839aafc28f Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 02:03:50 +0800 Subject: [PATCH 20/74] Deny legacy framing in safe mode and cap legacy first flights Safe-mode recovery no longer falls back to the permissive migration default, and legacy initial frames are rejected before allocating more than the control-frame ciphertext limit. --- CHANGELOG.md | 1 + src/common/auth/runtime.rs | 12 ++++++++---- src/common/auth/tests.rs | 30 ++++++++++++++++++++++++++++++ src/common/message/secure.rs | 4 +++- src/common/message/secure/tests.rs | 24 ++++++++++++++++++++++++ 5 files changed, 66 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 80f52ae..dc6a7f7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ All notable changes to this project will be documented in this file. - Rejected NUL/non-printable rotated administrator keys, cancelled in-flight status reads on credential revocation, refused `--force-init-admin-key` when encrypted auth state already exists, bound isolated-relay legacy continuation checksums to the relay key, and doubled first-flight Bloom retention so a max-future timestamp cannot outlive the filter. - Centralized env-safe administrator-key checks, isolated-relay legacy codec construction, credential-cancellation races, and auth snapshot/WAL paths so later protocol and lifecycle changes reuse one implementation. - Refused administrator-key initialization whenever encrypted auth state is present, preserved discarded slot generations across capacity changes, rolled back or fail-closed uncertain WAL appends, authenticated first flights before consuming the replay filter, and reused the registration credential for provider streams. +- Capped legacy first-flight allocations and kept legacy framing denied when authentication state enters safe mode. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/src/common/auth/runtime.rs b/src/common/auth/runtime.rs index 476c01e..81c969c 100644 --- a/src/common/auth/runtime.rs +++ b/src/common/auth/runtime.rs @@ -113,10 +113,14 @@ impl AuthRuntime { } } - let legacy_protocol = loaded - .as_ref() - .map(|state| state.legacy_protocol) - .unwrap_or(config.legacy_protocol); + let legacy_protocol = if safe_mode { + LegacyProtocolPolicy::Deny + } else { + loaded + .as_ref() + .map(|state| state.legacy_protocol) + .unwrap_or(config.legacy_protocol) + }; let mut admin_replay_order = loaded .as_ref() .map(|state| { diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index c3a73cb..237393d 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -114,6 +114,36 @@ async fn shrinking_then_expanding_capacity_does_not_reuse_old_key_ids() { let _ = std::fs::remove_dir_all(state_dir); } +#[tokio::test] +async fn safe_mode_denies_legacy_protocol_instead_of_restoring_the_default() { + let state_dir = temp_state_dir("safe-mode-legacy"); + let admin_key = *b"0123456789abcdefghijklmnopqrstuv"; + let config = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 4, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + let runtime = AuthRuntime::start(admin_key, config.clone()).await.unwrap(); + let admin = authenticate_for_test(&runtime, 0).unwrap(); + runtime + .set_legacy_protocol(&admin, LegacyProtocolPolicy::Deny) + .await + .unwrap(); + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + std::fs::write(state_dir.join("auth.wal"), b"broken-wal").unwrap(); + + let runtime = AuthRuntime::start(admin_key, config).await.unwrap(); + let admin = authenticate_for_test(&runtime, 0).unwrap(); + let status = runtime.status(&admin).await.unwrap(); + assert!(status.safe_mode); + assert_eq!(status.legacy_protocol, LegacyProtocolPolicy::Deny); + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + let _ = std::fs::remove_dir_all(state_dir); +} + #[tokio::test] async fn rotate_root_rejects_a_nul_containing_key() { let state_dir = temp_state_dir("rotate-nul"); diff --git a/src/common/message/secure.rs b/src/common/message/secure.rs index 45f7f0c..e5593e1 100644 --- a/src/common/message/secure.rs +++ b/src/common/message/secure.rs @@ -50,6 +50,7 @@ const MAX_CONNECTION_CLOCK_SKEW_SECONDS: u64 = 5 * 60; const DEFAULT_REPLAY_WINDOW_SECONDS: u64 = MAX_CONNECTION_CLOCK_SKEW_SECONDS.saturating_mul(2); const DEFAULT_REPLAY_FILTER_BYTES: usize = 1024 * 1024; const MAX_INITIAL_PLAINTEXT_LEN: u32 = 64 * 1024; +const MAX_INITIAL_CIPHERTEXT_LEN: u32 = MAX_INITIAL_PLAINTEXT_LEN + 16; #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum HeaderProtocol { @@ -374,7 +375,8 @@ impl ServerSecurity { ), response_session: None, })?; - if !valid_checksum_for_key(datalen, checksum, &key) || datalen > MAX_MSG_LEN { + if !valid_checksum_for_key(datalen, checksum, &key) || datalen > MAX_INITIAL_CIPHERTEXT_LEN + { return Err(ServerInitialError { failure: AuthFailure::new( "legacy_frame_invalid", diff --git a/src/common/message/secure/tests.rs b/src/common/message/secure/tests.rs index c605e14..dc47efc 100644 --- a/src/common/message/secure/tests.rs +++ b/src/common/message/secure/tests.rs @@ -200,6 +200,30 @@ async fn oversized_initial_frame_is_rejected_before_reading_its_body() { let _ = std::fs::remove_dir_all(config.state_dir); } +#[tokio::test] +async fn oversized_legacy_initial_frame_is_rejected_before_reading_its_body() { + use crate::common::checksum::get_checksum_for_key; + + let admin = *b"0123456789abcdefghijklmnopqrstuv"; + let config = temp_config(); + let auth = AuthRuntime::start(admin, config.clone()).await.unwrap(); + let datalen = MAX_INITIAL_CIPHERTEXT_LEN + 1; + let checksum = get_checksum_for_key(datalen, &admin); + let mut bytes = checksum.to_be_bytes().to_vec(); + bytes.extend_from_slice(&datalen.to_be_bytes()); + let security = ServerSecurity::new(auth); + let error = match security + .read_initial(&mut std::io::Cursor::new(bytes)) + .await + { + Ok(_) => panic!("oversized legacy frame was accepted"), + Err(error) => error, + }; + assert_eq!(error.failure.code, "legacy_frame_invalid"); + + let _ = std::fs::remove_dir_all(config.state_dir); +} + #[test] fn rotating_bloom_covers_current_and_previous_window() { let mut bloom = RotatingBloom::new(1024, DEFAULT_REPLAY_WINDOW_SECONDS); From e5801bb7f3dede4041c0cd27d6137e483d0e37da Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 02:11:30 +0800 Subject: [PATCH 21/74] Document WAL rollback retryable contract Retryable append errors mean the file was restored to its pre-append length and hot state was never published. Fail closed only when that rollback itself failed. --- src/common/auth/persistence.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/common/auth/persistence.rs b/src/common/auth/persistence.rs index 44b8eb1..3fb211e 100644 --- a/src/common/auth/persistence.rs +++ b/src/common/auth/persistence.rs @@ -323,6 +323,10 @@ pub(super) fn apply_persisted_mutation( Ok(()) } +// WHY: `append_wal` sets `retryable` only after it restores the file to the +// pre-append length. Callers persist first and publish hot state only after +// that returns, so a retryable error is a clean no-op. Fail closed only when +// that rollback itself failed and the WAL length is unknown. pub(super) fn fail_closed_on_uncertain_wal( inner: &AuthStateInner, result: Result<(), AuthFailure>, @@ -424,6 +428,8 @@ pub(super) fn append_wal( .and_then(|()| file.write_all(&sealed)) .and_then(|()| file.sync_data()) { + // retryable == rolled_back. A later append can then start at a known + // good offset. If truncation fails, the next record would be unreadable. let rolled_back = file .set_len(start_len) .and_then(|()| file.sync_data()) From e3d4b3fc18bfd12ef29a4c09200d238143a7a6a8 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 02:56:12 +0800 Subject: [PATCH 22/74] Harden safe-mode compaction, WAL durability, and first-flight admission Skip snapshot compaction after a failed load so empty in-memory state cannot replace damaged files. Fsync the auth directory when creating the WAL. Clear retained high-slot entries on reset and rotation. Drop the slot write lock before WAL append so fail-closed cancellation cannot deadlock. Fail closed when the process credential is cleared instead of accepting an unkeyed checksum. Rate-limit first flights per credential before inserting into the shared replay filter. --- CHANGELOG.md | 1 + src/common/auth/actor.rs | 148 ++++++++++++++++++---------- src/common/auth/persistence.rs | 53 +++++++--- src/common/auth/tests.rs | 52 ++++++++++ src/common/checksum.rs | 26 ++++- src/common/message/mod.rs | 22 +++-- src/common/message/secure.rs | 43 +++++--- src/common/message/secure/replay.rs | 77 +++++++++++++-- src/common/message/secure/tests.rs | 11 +++ 9 files changed, 340 insertions(+), 93 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index dc6a7f7..5a4bc16 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ All notable changes to this project will be documented in this file. - Centralized env-safe administrator-key checks, isolated-relay legacy codec construction, credential-cancellation races, and auth snapshot/WAL paths so later protocol and lifecycle changes reuse one implementation. - Refused administrator-key initialization whenever encrypted auth state is present, preserved discarded slot generations across capacity changes, rolled back or fail-closed uncertain WAL appends, authenticated first flights before consuming the replay filter, and reused the registration credential for provider streams. - Capped legacy first-flight allocations and kept legacy framing denied when authentication state enters safe mode. +- Skipped snapshot compaction while startup is in safe mode, fsynced the auth directory when creating `auth.wal`, cleared retained high-slot entries on reset/rotate, dropped the slot write lock before WAL fail-closed cancellation, fail-closed process checksums after the credential is cleared, and rate-limited first flights per credential before consuming the shared replay filter. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/src/common/auth/actor.rs b/src/common/auth/actor.rs index 47d7f30..e2d25ac 100644 --- a/src/common/auth/actor.rs +++ b/src/common/auth/actor.rs @@ -136,7 +136,14 @@ pub(super) async fn run_auth_actor( &mut admin_replays, &mut admin_replay_order, ); - if now.saturating_sub(last_snapshot_at) >= SNAPSHOT_COMPACTION_INTERVAL.as_secs() { + // WHY: A failed load starts safe mode with empty in-memory + // generations. Compacting that reconstruction would replace the + // damaged snapshot, truncate the WAL, and let the next start + // exit safe mode without rotating the instance id. + if compaction_is_allowed(inner.safe_mode.load(Ordering::Acquire)) + && now.saturating_sub(last_snapshot_at) + >= SNAPSHOT_COMPACTION_INTERVAL.as_secs() + { let snapshot = build_snapshot(&inner, &cold, &admin_replay_order); if let Err(error) = write_snapshot_and_truncate_wal( &config, @@ -377,37 +384,53 @@ fn actor_issue( let expires_at = validate_ttl(config, ttl)?; let label = validate_label(label)?; let issued_at = unix_seconds(); - let mut slots = inner - .slots - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let Some((index, slot)) = slots - .iter_mut() - .enumerate() - .find(|(_, slot)| slot.state == SlotState::Free && slot.generation < u32::MAX) - else { - return Err(AuthFailure::new( - "temporary_key_capacity_exhausted", - "temporary key slot table is full", - true, - )); - }; - let generation = slot.generation + 1; - let key_id = make_key_id(generation, index as u32); - let entry = PersistedEntry { - key_id, - state: SlotState::Active, - issued_at, - expires_at, - label: label.clone(), - tombstoned_at: None, + let (index, generation, key_id, entry) = { + let slots = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let Some((index, slot)) = slots + .iter() + .enumerate() + .find(|(_, slot)| slot.state == SlotState::Free && slot.generation < u32::MAX) + else { + return Err(AuthFailure::new( + "temporary_key_capacity_exhausted", + "temporary key slot table is full", + true, + )); + }; + let generation = slot.generation + 1; + let key_id = make_key_id(generation, index as u32); + ( + index, + generation, + key_id, + PersistedEntry { + key_id, + state: SlotState::Active, + issued_at, + expires_at, + label: label.clone(), + tombstoned_at: None, + }, + ) }; + // Persist before taking the slot write lock. A fail-closed WAL error + // cancels leases via slots.read() and must not nest under slots.write(). append_mutation( config, inner, - StateMutation::Issue(entry.clone()), + StateMutation::Issue(entry), audit("temporary_key_issue", Some(key_id), label.clone()), )?; + let mut slots = inner + .slots + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let slot = slots + .get_mut(index) + .ok_or_else(|| AuthFailure::internal("issued slot disappeared"))?; let lease = Arc::new(AuthLease::new(key_id, expires_at)); slot.generation = generation; slot.state = SlotState::Active; @@ -501,18 +524,20 @@ fn actor_renew( ensure_store_available(inner)?; let expires_at = validate_ttl(config, ttl)?; let index = key_slot(key_id) as usize; - let mut slots = inner - .slots - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let slot = slots.get_mut(index).ok_or_else(|| key_not_found(key_id))?; - validate_slot_identity(slot, key_id)?; - if slot.state != SlotState::Active || slot.expires_at <= unix_seconds() { - return Err(AuthFailure::new( - "temporary_key_not_renewable", - "only an active, unexpired temporary key can be renewed", - false, - )); + { + let slots = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let slot = slots.get(index).ok_or_else(|| key_not_found(key_id))?; + validate_slot_identity(slot, key_id)?; + if slot.state != SlotState::Active || slot.expires_at <= unix_seconds() { + return Err(AuthFailure::new( + "temporary_key_not_renewable", + "only an active, unexpired temporary key can be renewed", + false, + )); + } } let label = cold .get(&key_id) @@ -523,6 +548,19 @@ fn actor_renew( StateMutation::Renew { key_id, expires_at }, audit("temporary_key_renew", Some(key_id), label), )?; + let mut slots = inner + .slots + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let slot = slots.get_mut(index).ok_or_else(|| key_not_found(key_id))?; + validate_slot_identity(slot, key_id)?; + if slot.state != SlotState::Active { + return Err(AuthFailure::new( + "temporary_key_inactive", + "temporary key lease is no longer active", + true, + )); + } let lease = slot.lease.upgrade().ok_or_else(|| { AuthFailure::new( "temporary_key_inactive", @@ -548,18 +586,20 @@ fn actor_revoke( ensure_store_available(inner)?; let now = unix_seconds(); let index = key_slot(key_id) as usize; - let mut slots = inner - .slots - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let slot = slots.get_mut(index).ok_or_else(|| key_not_found(key_id))?; - validate_slot_identity(slot, key_id)?; - if slot.state != SlotState::Active { - return Err(AuthFailure::new( - "temporary_key_not_active", - "temporary key is not active", - false, - )); + { + let slots = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let slot = slots.get(index).ok_or_else(|| key_not_found(key_id))?; + validate_slot_identity(slot, key_id)?; + if slot.state != SlotState::Active { + return Err(AuthFailure::new( + "temporary_key_not_active", + "temporary key is not active", + false, + )); + } } let label = cold .get(&key_id) @@ -572,6 +612,12 @@ fn actor_revoke( StateMutation::Revoke { key_id, at: now }, audit("temporary_key_revoke", Some(key_id), label.clone()), )?; + let mut slots = inner + .slots + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let slot = slots.get_mut(index).ok_or_else(|| key_not_found(key_id))?; + validate_slot_identity(slot, key_id)?; slot.state = SlotState::Revoked; if let Some(lease) = slot.lease.upgrade() { lease.cancellation.cancel(); @@ -685,6 +731,7 @@ fn actor_reset( .unwrap_or_else(|poisoned| poisoned.into_inner()) = new_instance_id; cold.clear(); wheel.clear(unix_seconds()); + clear_retained_high_slot_entries(inner); inner.safe_mode.store(false, Ordering::Release); Ok(()) } @@ -741,6 +788,7 @@ fn actor_rotate_root( } cold.clear(); wheel.clear(unix_seconds()); + clear_retained_high_slot_entries(inner); let new_admin_lease = Arc::new(AuthLease::new(0, u64::MAX)); *inner .admin diff --git a/src/common/auth/persistence.rs b/src/common/auth/persistence.rs index 3fb211e..568e020 100644 --- a/src/common/auth/persistence.rs +++ b/src/common/auth/persistence.rs @@ -27,6 +27,35 @@ pub fn encrypted_auth_state_exists(state_dir: &Path) -> bool { auth_snapshot_path(state_dir).exists() || auth_wal_path(state_dir).exists() } +pub(super) fn compaction_is_allowed(safe_mode: bool) -> bool { + !safe_mode +} + +pub(super) fn clear_retained_high_slot_entries(inner: &AuthStateInner) { + inner + .high_slot_entries + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clear(); +} + +fn sync_parent_directory(path: &Path) -> Result<(), AuthFailure> { + #[cfg(unix)] + if let Some(parent) = path.parent() { + File::open(parent) + .and_then(|directory| directory.sync_all()) + .map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to sync `{}`: {error}", parent.display()), + false, + ) + })?; + } + let _ = path; + Ok(()) +} + pub(super) fn push_audit_record(inner: &AuthStateInner, record: AuditRecord) { let mut records = inner .audit_records @@ -393,6 +422,7 @@ pub(super) fn append_wal( let length = u32::try_from(sealed.len()) .map_err(|_| AuthFailure::internal("auth WAL record is too large"))?; let path = auth_wal_path(&config.state_dir); + let created = !path.exists(); let mut file = OpenOptions::new() .create(true) .append(true) @@ -447,6 +477,9 @@ pub(super) fn append_wal( rolled_back, )); } + if created { + sync_parent_directory(&path)?; + } Ok(()) } @@ -520,6 +553,7 @@ pub(super) fn write_snapshot_and_truncate_wal( let snapshot_path = auth_snapshot_path(&config.state_dir); atomic_write(&snapshot_path, &sealed, 0o600)?; let wal_path = auth_wal_path(&config.state_dir); + let created = !wal_path.exists(); let wal = OpenOptions::new() .create(true) .write(true) @@ -538,7 +572,11 @@ pub(super) fn write_snapshot_and_truncate_wal( format!("failed to sync `{}`: {error}", wal_path.display()), true, ) - }) + })?; + if created { + sync_parent_directory(&wal_path)?; + } + Ok(()) } pub(super) fn seal_blob(admin_key: &AesKeyType, plain: &[u8]) -> Result, AuthFailure> { @@ -793,18 +831,7 @@ pub(super) fn atomic_write(path: &Path, data: &[u8], mode: u32) -> Result<(), Au false, ) })?; - #[cfg(unix)] - if let Some(parent) = path.parent() { - File::open(parent) - .and_then(|directory| directory.sync_all()) - .map_err(|error| { - AuthFailure::new( - "auth_state_unavailable", - format!("failed to sync `{}`: {error}", parent.display()), - false, - ) - })?; - } + sync_parent_directory(path)?; Ok(()) })(); if result.is_err() { diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index 237393d..5d9041a 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -144,6 +144,58 @@ async fn safe_mode_denies_legacy_protocol_instead_of_restoring_the_default() { let _ = std::fs::remove_dir_all(state_dir); } +#[test] +fn safe_mode_startup_does_not_allow_compaction() { + assert!(!compaction_is_allowed(true)); + assert!(compaction_is_allowed(false)); +} + +#[tokio::test] +async fn reset_clears_retained_high_slot_entries() { + let state_dir = temp_state_dir("reset-high-slots"); + let admin_key = *b"0123456789abcdefghijklmnopqrstuv"; + let config_two = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 2, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + let runtime = AuthRuntime::start(admin_key, config_two.clone()) + .await + .unwrap(); + let admin = authenticate_for_test(&runtime, 0).unwrap(); + let first = runtime + .issue(&admin, Duration::from_secs(60), Some("first".to_string())) + .await + .unwrap(); + let second = runtime + .issue(&admin, Duration::from_secs(60), Some("second".to_string())) + .await + .unwrap(); + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + + let config_one = AuthConfig { + max_temporary_keys: 1, + ..config_two.clone() + }; + let runtime = AuthRuntime::start(admin_key, config_one).await.unwrap(); + let admin = authenticate_for_test(&runtime, 0).unwrap(); + runtime.reset(&admin).await.unwrap(); + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + + let runtime = AuthRuntime::start(admin_key, config_two).await.unwrap(); + let admin = authenticate_for_test(&runtime, 0).unwrap(); + let page = runtime.list(&admin, 0, 100).await.unwrap(); + assert!(page.items.is_empty()); + assert!(authenticate_for_test(&runtime, first.metadata.key_id).is_err()); + assert!(authenticate_for_test(&runtime, second.metadata.key_id).is_err()); + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + let _ = std::fs::remove_dir_all(state_dir); +} + #[tokio::test] async fn rotate_root_rejects_a_nul_containing_key() { let state_dir = temp_state_dir("rotate-nul"); diff --git a/src/common/checksum.rs b/src/common/checksum.rs index 302349a..91576d2 100644 --- a/src/common/checksum.rs +++ b/src/common/checksum.rs @@ -501,6 +501,11 @@ pub fn get_checksum_for_key(datalen: DataLenType, key: &[u8]) -> ChecksumType { datalen ^ gen_checksum_by_key(key) } +/// `true` when the process credential can key a length checksum. +pub fn process_checksum_is_ready() -> bool { + get_process_credential().is_ok() +} + #[inline] /// Compute frame checksum from payload length and the current header key hash. pub fn get_checksum(datalen: DataLenType) -> ChecksumType { @@ -515,8 +520,12 @@ pub fn valid_checksum_for_key(datalen: DataLenType, checksum: ChecksumType, key: #[inline] /// Validate frame checksum generated by [`get_checksum`]. +/// +/// Missing or invalid process credentials fail closed instead of accepting +/// an unkeyed `datalen` checksum (`hash == 0`). pub fn valid_checksum(datalen: DataLenType, checksum: ChecksumType) -> bool { - datalen == (checksum ^ MSG_HEADER_KEY_STATE.hash.load(Ordering::Acquire)) + process_checksum_is_ready() + && datalen == (checksum ^ MSG_HEADER_KEY_STATE.hash.load(Ordering::Acquire)) } pub type AesKeyType = [u8; 32]; @@ -572,6 +581,21 @@ mod tests { assert!(!is_env_safe_admin_key(b"short")); } + #[tokio::test] + async fn clearing_the_process_credential_fails_closed_for_checksums() { + use super::*; + use crate::common::auth::PROCESS_CREDENTIAL_TEST_LOCK; + + let _guard = PROCESS_CREDENTIAL_TEST_LOCK.lock().await; + set_process_msg_header_key(Some("0123456789abcdefghijklmnopqrstuv")).unwrap(); + let checksum = get_checksum(32); + assert!(valid_checksum(32, checksum)); + set_process_msg_header_key(None).unwrap(); + assert!(!valid_checksum(32, checksum)); + assert!(!valid_checksum(32, 32)); + assert!(get_process_credential().is_err()); + } + #[test] fn administrator_credentials_reject_nul_and_whitespace() { use super::*; diff --git a/src/common/message/mod.rs b/src/common/message/mod.rs index 8a4ca9f..f41867d 100644 --- a/src/common/message/mod.rs +++ b/src/common/message/mod.rs @@ -8,8 +8,8 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt}; use super::buffer::{BufferGetter, CommonBuffer, FixedSizeBuffer}; use super::checksum::{ - get_checksum, get_checksum_for_key, get_msg_header_key, valid_checksum, valid_checksum_for_key, - AesKeyType, + get_checksum, get_checksum_for_key, get_msg_header_key, process_checksum_is_ready, + valid_checksum, valid_checksum_for_key, AesKeyType, }; use super::error::{ self, MsgDatalenValidateSnafu, MsgNetworkReadBodySnafu, MsgNetworkReadCheckSumSnafu, @@ -146,10 +146,20 @@ fn checksum_matches(datalen: DataLenType, checksum: u32, key: Option<&[u8]>) -> } #[inline] -fn checksum_for(len: DataLenType, key: Option<&[u8]>) -> u32 { +fn checksum_for(len: DataLenType, key: Option<&[u8]>) -> Result { match key { - Some(key) => get_checksum_for_key(len, key), - None => get_checksum(len), + Some(key) => Ok(get_checksum_for_key(len, key)), + None => { + if !process_checksum_is_ready() { + return Err(error::Error::MsgCodec { + action: "load configured credential", + detail: + "`MSG_HEADER_KEY` is required; no insecure default checksum is available" + .to_string(), + }); + } + Ok(get_checksum(len)) + } } } @@ -180,7 +190,7 @@ async fn set_msg_len( len: DataLenType, checksum_key: Option<&[u8]>, ) -> Result<()> { - write_checksum(writer, checksum_for(len, checksum_key)).await?; + write_checksum(writer, checksum_for(len, checksum_key)?).await?; write_datalen(writer, len).await } diff --git a/src/common/message/secure.rs b/src/common/message/secure.rs index e5593e1..37f7037 100644 --- a/src/common/message/secure.rs +++ b/src/common/message/secure.rs @@ -286,7 +286,7 @@ use std::fmt; #[derive(Clone)] pub struct ServerSecurity { auth: AuthRuntime, - replay: Arc>, + replay: Arc>, failure_logs: Arc>, } @@ -294,7 +294,7 @@ impl ServerSecurity { pub fn new(auth: AuthRuntime) -> Self { Self { auth, - replay: Arc::new(Mutex::new(RotatingBloom::new( + replay: Arc::new(Mutex::new(ReplayGuard::new( DEFAULT_REPLAY_FILTER_BYTES, DEFAULT_REPLAY_WINDOW_SECONDS, ))), @@ -543,20 +543,33 @@ impl ServerSecurity { response_session: Some(session_without_context(&session)), })?; let fingerprint = replay_fingerprint(key_id, &salt); - let replayed = self + match self .replay .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()) - .check_and_insert(&fingerprint, unix_seconds()); - if replayed { - return Err(ServerInitialError { - failure: AuthFailure::new( - "connection_salt_replayed", - "protocol-v2 connection salt was already accepted", - true, - ), - response_session: Some(session), - }); + .admit(key_id, &fingerprint, unix_seconds()) + { + FirstFlightAdmit::Replayed => { + return Err(ServerInitialError { + failure: AuthFailure::new( + "connection_salt_replayed", + "protocol-v2 connection salt was already accepted", + true, + ), + response_session: Some(session), + }); + } + FirstFlightAdmit::Limited => { + return Err(ServerInitialError { + failure: AuthFailure::new( + "connection_admission_limited", + "this credential has opened too many new connections in the current window", + true, + ), + response_session: Some(session), + }); + } + FirstFlightAdmit::Fresh => {} } session.context = Some(context); Ok(ServerInitialMessage { @@ -613,7 +626,9 @@ mod frame; use frame::{derive_material, first_prefix, V2Material}; pub use frame::{V2MessageReader, V2MessageWriter}; mod replay; -use replay::{replay_fingerprint, RotatingBloom}; +#[cfg(test)] +use replay::RotatingBloom; +use replay::{replay_fingerprint, FirstFlightAdmit, ReplayGuard}; fn legacy_message_reader<'a, T: AsyncReadExt + Unpin>( reader: &'a mut T, key: &AesKeyType, diff --git a/src/common/message/secure/replay.rs b/src/common/message/secure/replay.rs index d33c882..1c33ef4 100644 --- a/src/common/message/secure/replay.rs +++ b/src/common/message/secure/replay.rs @@ -3,17 +3,31 @@ //! ```text //! key id + salt -> SHA-256 fingerprint -> current Bloom window //! -> previous Bloom window +//! key id ---------> per-credential first-flight count (before Bloom insert) //! ``` //! -//! `check_and_insert` is called while one mutex is held, making concurrent admission +//! `admit` is called while one mutex is held, making concurrent admission //! atomic. This Bloom filter protects all connection types from immediate duplicates; //! administrator mutations additionally use the exact durable replay set in `auth`. //! //! Each generation lasts `2 *` the accepted clock-skew so a salt inserted at the //! end of a window with a max-future timestamp cannot be replayed after rotation. +//! Per-credential counts stop one tenant from filling the shared filter with +//! unique salts before the request payload is decoded. + +use std::collections::HashMap; use super::*; +pub(super) const MAX_FIRST_FLIGHTS_PER_KEY: u32 = 8_192; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum FirstFlightAdmit { + Fresh, + Replayed, + Limited, +} + pub(super) fn replay_fingerprint(key_id: u64, salt: &[u8; CONNECTION_SALT_LEN]) -> [u8; 32] { let mut input = [0_u8; 8 + CONNECTION_SALT_LEN]; input[..8].copy_from_slice(&key_id.to_be_bytes()); @@ -51,14 +65,6 @@ impl RotatingBloom { bloom_insert(&mut self.current, fingerprint); } - pub(super) fn check_and_insert(&mut self, fingerprint: &[u8; 32], now: u64) -> bool { - if self.contains(fingerprint, now) { - return true; - } - self.insert(fingerprint, now); - false - } - fn rotate(&mut self, now: u64) { let elapsed = now.saturating_sub(self.current_started_at); if elapsed < self.window_seconds { @@ -75,6 +81,59 @@ impl RotatingBloom { } } +pub(super) struct ReplayGuard { + bloom: RotatingBloom, + counts: HashMap, + counts_started_at: u64, + window_seconds: u64, + max_per_key: u32, +} + +impl ReplayGuard { + pub(super) fn new(bytes: usize, window_seconds: u64) -> Self { + Self { + bloom: RotatingBloom::new(bytes, window_seconds), + counts: HashMap::new(), + counts_started_at: unix_seconds(), + window_seconds, + max_per_key: MAX_FIRST_FLIGHTS_PER_KEY, + } + } + + #[cfg(test)] + pub(super) fn with_max_per_key(mut self, max_per_key: u32) -> Self { + self.max_per_key = max_per_key; + self + } + + pub(super) fn admit( + &mut self, + key_id: u64, + fingerprint: &[u8; 32], + now: u64, + ) -> FirstFlightAdmit { + self.rotate_counts(now); + if self.bloom.contains(fingerprint, now) { + return FirstFlightAdmit::Replayed; + } + let count = self.counts.entry(key_id).or_insert(0); + if *count >= self.max_per_key { + return FirstFlightAdmit::Limited; + } + self.bloom.insert(fingerprint, now); + *count = count.saturating_add(1); + FirstFlightAdmit::Fresh + } + + fn rotate_counts(&mut self, now: u64) { + if now.saturating_sub(self.counts_started_at) < self.window_seconds { + return; + } + self.counts.clear(); + self.counts_started_at = now; + } +} + fn bloom_positions(filter_len: usize, fingerprint: &[u8; 32]) -> [usize; 4] { let bits = filter_len * 8; std::array::from_fn(|index| { diff --git a/src/common/message/secure/tests.rs b/src/common/message/secure/tests.rs index dc47efc..4fd3625 100644 --- a/src/common/message/secure/tests.rs +++ b/src/common/message/secure/tests.rs @@ -261,6 +261,17 @@ fn rotating_bloom_retains_a_max_future_timestamp_past_the_next_rotation() { )); } +#[test] +fn per_credential_admission_limit_does_not_consume_other_keys() { + let now = unix_seconds(); + let mut guard = ReplayGuard::new(1024, DEFAULT_REPLAY_WINDOW_SECONDS).with_max_per_key(2); + assert_eq!(guard.admit(1, &[1_u8; 32], now), FirstFlightAdmit::Fresh); + assert_eq!(guard.admit(1, &[2_u8; 32], now), FirstFlightAdmit::Fresh); + assert_eq!(guard.admit(1, &[3_u8; 32], now), FirstFlightAdmit::Limited); + assert_eq!(guard.admit(2, &[3_u8; 32], now), FirstFlightAdmit::Fresh); + assert_eq!(guard.admit(1, &[1_u8; 32], now), FirstFlightAdmit::Replayed); +} + #[tokio::test] async fn legacy_initial_frame_validates_against_isolated_relay_key() { let _process_credential_guard = PROCESS_CREDENTIAL_TEST_LOCK.lock().await; From d3b23d7810cb7816b61e9049ae951e0c7e7aa807 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 03:12:21 +0800 Subject: [PATCH 23/74] Persist first-flight replay and close remaining auth gaps Store admitted first-flight fingerprints so a restart cannot replay a captured subscribe. Replace existing auth files with MoveFileEx on Windows. Reject explicit out-of-range --max-temporary-keys and TTL flags instead of silently restoring defaults. Expose the embedded relay administrator key through FFI and the configuration page. --- CHANGELOG.md | 1 + src/bin/pb-mapper.rs | 42 +++++++++++- src/common/auth.rs | 6 +- src/common/auth/persistence.rs | 41 +++++++++++- src/common/message/secure.rs | 14 +++- src/common/message/secure/replay.rs | 73 +++++++++++++++++++-- src/common/message/secure/tests.rs | 26 +++++++- ui/lib/l10n/app_en.arb | 2 + ui/lib/l10n/app_zh.arb | 2 + ui/lib/src/ffi/pb_mapper_api.dart | 3 + ui/lib/src/views/configuration_view.dart | 14 ++++ ui/native/pb_mapper_ffi/src/config.rs | 8 +-- ui/native/pb_mapper_ffi/src/ctl/mod.rs | 6 +- ui/native/pb_mapper_ffi/src/state.rs | 6 ++ ui/native/pb_mapper_ffi/src/state/status.rs | 8 +++ 15 files changed, 234 insertions(+), 18 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5a4bc16..566f5bf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,7 @@ All notable changes to this project will be documented in this file. - Refused administrator-key initialization whenever encrypted auth state is present, preserved discarded slot generations across capacity changes, rolled back or fail-closed uncertain WAL appends, authenticated first flights before consuming the replay filter, and reused the registration credential for provider streams. - Capped legacy first-flight allocations and kept legacy framing denied when authentication state enters safe mode. - Skipped snapshot compaction while startup is in safe mode, fsynced the auth directory when creating `auth.wal`, cleared retained high-slot entries on reset/rotate, dropped the slot write lock before WAL fail-closed cancellation, fail-closed process checksums after the credential is cleared, and rate-limited first flights per credential before consuming the shared replay filter. +- Persisted first-flight replay admissions across restarts, replaced existing auth files atomically on Windows, rejected explicit out-of-range server auth flags, and exposed the embedded relay's isolated administrator key through FFI/UI. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/src/bin/pb-mapper.rs b/src/bin/pb-mapper.rs index 2c8efdc..3b9c6c2 100644 --- a/src/bin/pb-mapper.rs +++ b/src/bin/pb-mapper.rs @@ -20,7 +20,7 @@ use better_mimalloc_rs::MiMalloc; use clap::{Args, Parser, Subcommand, ValueEnum}; use pb_mapper::common::auth::{ generate_admin_key, initialize_admin_key, write_admin_key_file, AuthConfig, KeyPage, - LegacyProtocolPolicy, + LegacyProtocolPolicy, MAX_TEMP_KEY_CAPACITY, MAX_TEMP_KEY_TTL, MIN_TEMP_KEY_TTL, }; use pb_mapper::common::checksum::set_process_msg_header_key; use pb_mapper::common::checksum::{setup_machine_msg_header_key, MACHINE_MSG_HEADER_KEY_PATH}; @@ -226,22 +226,41 @@ async fn run(cli: Cli) -> Result<(), Box> { Ok(()) } -async fn run_server(args: ServerArgs) -> Result<(), Box> { +fn apply_server_auth_overrides(args: &ServerArgs) -> Result<(), Box> { if let Some(auth_state_dir) = &args.auth_state_dir { std::env::set_var("PB_MAPPER_AUTH_STATE_DIR", auth_state_dir); } if let Some(max_temporary_keys) = args.max_temporary_keys { + if !(1..=MAX_TEMP_KEY_CAPACITY).contains(&max_temporary_keys) { + return Err(format!( + "`--max-temporary-keys` must be between 1 and {MAX_TEMP_KEY_CAPACITY}" + ) + .into()); + } std::env::set_var( "PB_MAPPER_AUTH_MAX_TEMP_KEYS", max_temporary_keys.to_string(), ); } if let Some(max_temporary_key_ttl) = args.max_temporary_key_ttl { + if max_temporary_key_ttl < MIN_TEMP_KEY_TTL || max_temporary_key_ttl > MAX_TEMP_KEY_TTL { + return Err(format!( + "`--max-temporary-key-ttl` must be between {}s and {}d", + MIN_TEMP_KEY_TTL.as_secs(), + MAX_TEMP_KEY_TTL.as_secs() / 86_400 + ) + .into()); + } std::env::set_var( "PB_MAPPER_AUTH_MAX_TEMP_TTL_SECS", max_temporary_key_ttl.as_secs().to_string(), ); } + Ok(()) +} + +async fn run_server(args: ServerArgs) -> Result<(), Box> { + apply_server_auth_overrides(&args)?; if let Some(legacy_protocol) = args.legacy_protocol { std::env::set_var( "PB_MAPPER_LEGACY_PROTOCOL", @@ -547,4 +566,23 @@ mod tests { ); assert_eq!(explicit.legacy_protocol, Some(LegacyProtocolArg::Deny)); } + + #[test] + fn explicit_out_of_range_server_auth_flags_are_rejected() { + let cli = Cli::try_parse_from(["pb-mapper", "server", "--max-temporary-keys", "0"]) + .expect("clap should accept the token before bounds checking"); + let Command::Server(args) = cli.command else { + panic!("expected server command"); + }; + let error = apply_server_auth_overrides(&args).unwrap_err(); + assert!(error.to_string().contains("--max-temporary-keys")); + + let cli = Cli::try_parse_from(["pb-mapper", "server", "--max-temporary-key-ttl", "5s"]) + .expect("clap should accept the token before bounds checking"); + let Command::Server(args) = cli.command else { + panic!("expected server command"); + }; + let error = apply_server_auth_overrides(&args).unwrap_err(); + assert!(error.to_string().contains("--max-temporary-key-ttl")); + } } diff --git a/src/common/auth.rs b/src/common/auth.rs index f1c26e6..330a975 100644 --- a/src/common/auth.rs +++ b/src/common/auth.rs @@ -45,8 +45,10 @@ use super::checksum::{ pub const ADMIN_NAMESPACE: u64 = 0; pub const DEFAULT_AUTH_STATE_DIR: &str = "/var/lib/pb-mapper/auth"; pub const DEFAULT_TEMP_KEY_CAPACITY: usize = 65_536; +pub const MAX_TEMP_KEY_CAPACITY: usize = 1_048_576; pub const DEFAULT_MAX_TEMP_KEY_TTL: Duration = Duration::from_secs(30 * 24 * 60 * 60); pub const MIN_TEMP_KEY_TTL: Duration = Duration::from_secs(10); +pub const MAX_TEMP_KEY_TTL: Duration = Duration::from_secs(365 * 24 * 60 * 60); const TOMBSTONE_RETENTION: Duration = Duration::from_secs(60); const SNAPSHOT_COMPACTION_INTERVAL: Duration = Duration::from_secs(5 * 60); const SNAPSHOT_SCHEMA_VERSION: u16 = 1; @@ -125,13 +127,13 @@ impl Default for AuthConfig { "PB_MAPPER_AUTH_MAX_TEMP_KEYS", DEFAULT_TEMP_KEY_CAPACITY, 1, - 1_048_576, + MAX_TEMP_KEY_CAPACITY, ), max_temporary_key_ttl: Duration::from_secs(env_u64( "PB_MAPPER_AUTH_MAX_TEMP_TTL_SECS", DEFAULT_MAX_TEMP_KEY_TTL.as_secs(), MIN_TEMP_KEY_TTL.as_secs(), - 365 * 24 * 60 * 60, + MAX_TEMP_KEY_TTL.as_secs(), )), legacy_protocol: legacy_protocol_from_env(), } diff --git a/src/common/auth/persistence.rs b/src/common/auth/persistence.rs index 568e020..70f5c14 100644 --- a/src/common/auth/persistence.rs +++ b/src/common/auth/persistence.rs @@ -39,6 +39,45 @@ pub(super) fn clear_retained_high_slot_entries(inner: &AuthStateInner) { .clear(); } +fn replace_file(from: &Path, to: &Path) -> std::io::Result<()> { + #[cfg(windows)] + { + use std::os::windows::ffi::OsStrExt; + + const MOVEFILE_REPLACE_EXISTING: u32 = 0x1; + const MOVEFILE_WRITE_THROUGH: u32 = 0x8; + extern "system" { + fn MoveFileExW( + lp_existing_file_name: *const u16, + lp_new_file_name: *const u16, + dw_flags: u32, + ) -> i32; + } + fn wide(path: &Path) -> Vec { + path.as_os_str() + .encode_wide() + .chain(std::iter::once(0)) + .collect() + } + let from_w = wide(from); + let to_w = wide(to); + let ok = unsafe { + MoveFileExW( + from_w.as_ptr(), + to_w.as_ptr(), + MOVEFILE_REPLACE_EXISTING | MOVEFILE_WRITE_THROUGH, + ) + }; + if ok == 0 { + Err(std::io::Error::last_os_error()) + } else { + Ok(()) + } + } + #[cfg(not(windows))] + std::fs::rename(from, to) +} + fn sync_parent_directory(path: &Path) -> Result<(), AuthFailure> { #[cfg(unix)] if let Some(parent) = path.parent() { @@ -824,7 +863,7 @@ pub(super) fn atomic_write(path: &Path, data: &[u8], mode: u32) -> Result<(), Au ) })?; drop(file); - std::fs::rename(&temporary, path).map_err(|error| { + replace_file(&temporary, path).map_err(|error| { AuthFailure::new( "auth_state_unavailable", format!("failed to replace `{}`: {error}", path.display()), diff --git a/src/common/message/secure.rs b/src/common/message/secure.rs index 37f7037..6cb7da9 100644 --- a/src/common/message/secure.rs +++ b/src/common/message/secure.rs @@ -292,9 +292,11 @@ pub struct ServerSecurity { impl ServerSecurity { pub fn new(auth: AuthRuntime) -> Self { + let replay_path = auth.config().state_dir.join("connection.replay"); Self { auth, - replay: Arc::new(Mutex::new(ReplayGuard::new( + replay: Arc::new(Mutex::new(ReplayGuard::open( + Some(replay_path), DEFAULT_REPLAY_FILTER_BYTES, DEFAULT_REPLAY_WINDOW_SECONDS, ))), @@ -569,6 +571,16 @@ impl ServerSecurity { response_session: Some(session), }); } + FirstFlightAdmit::Unavailable => { + return Err(ServerInitialError { + failure: AuthFailure::new( + "connection_replay_store_unavailable", + "failed to persist first-flight replay admission", + true, + ), + response_session: Some(session), + }); + } FirstFlightAdmit::Fresh => {} } session.context = Some(context); diff --git a/src/common/message/secure/replay.rs b/src/common/message/secure/replay.rs index 1c33ef4..ece6461 100644 --- a/src/common/message/secure/replay.rs +++ b/src/common/message/secure/replay.rs @@ -16,16 +16,21 @@ //! unique salts before the request payload is decoded. use std::collections::HashMap; +use std::fs::{File, OpenOptions}; +use std::io::{Read, Write}; +use std::path::PathBuf; use super::*; pub(super) const MAX_FIRST_FLIGHTS_PER_KEY: u32 = 8_192; +const REPLAY_RECORD_LEN: usize = 40; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(super) enum FirstFlightAdmit { Fresh, Replayed, Limited, + Unavailable, } pub(super) fn replay_fingerprint(key_id: u64, salt: &[u8; CONNECTION_SALT_LEN]) -> [u8; 32] { @@ -87,17 +92,21 @@ pub(super) struct ReplayGuard { counts_started_at: u64, window_seconds: u64, max_per_key: u32, + log_path: Option, } impl ReplayGuard { - pub(super) fn new(bytes: usize, window_seconds: u64) -> Self { - Self { + pub(super) fn open(log_path: Option, bytes: usize, window_seconds: u64) -> Self { + let mut guard = Self { bloom: RotatingBloom::new(bytes, window_seconds), counts: HashMap::new(), counts_started_at: unix_seconds(), window_seconds, max_per_key: MAX_FIRST_FLIGHTS_PER_KEY, - } + log_path, + }; + guard.load_persisted(); + guard } #[cfg(test)] @@ -116,12 +125,14 @@ impl ReplayGuard { if self.bloom.contains(fingerprint, now) { return FirstFlightAdmit::Replayed; } - let count = self.counts.entry(key_id).or_insert(0); - if *count >= self.max_per_key { + if self.counts.get(&key_id).copied().unwrap_or(0) >= self.max_per_key { return FirstFlightAdmit::Limited; } + if self.persist(fingerprint, now).is_err() { + return FirstFlightAdmit::Unavailable; + } self.bloom.insert(fingerprint, now); - *count = count.saturating_add(1); + *self.counts.entry(key_id).or_insert(0) += 1; FirstFlightAdmit::Fresh } @@ -132,6 +143,56 @@ impl ReplayGuard { self.counts.clear(); self.counts_started_at = now; } + + fn persist(&self, fingerprint: &[u8; 32], now: u64) -> std::io::Result<()> { + let Some(path) = &self.log_path else { + return Ok(()); + }; + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + let mut file = OpenOptions::new().create(true).append(true).open(path)?; + let mut record = [0_u8; REPLAY_RECORD_LEN]; + record[..32].copy_from_slice(fingerprint); + record[32..].copy_from_slice(&now.to_be_bytes()); + file.write_all(&record)?; + file.sync_data() + } + + fn load_persisted(&mut self) { + let Some(path) = self.log_path.clone() else { + return; + }; + let Ok(mut file) = File::open(&path) else { + return; + }; + let now = unix_seconds(); + let mut live = Vec::new(); + let mut record = [0_u8; REPLAY_RECORD_LEN]; + loop { + match file.read(&mut record) { + Ok(0) => break, + Ok(n) if n == REPLAY_RECORD_LEN => {} + _ => break, + } + let timestamp = u64::from_be_bytes(record[32..].try_into().expect("timestamp width")); + if now.saturating_sub(timestamp) > self.window_seconds { + continue; + } + let fingerprint: [u8; 32] = record[..32].try_into().expect("fingerprint width"); + self.bloom.insert(&fingerprint, timestamp); + live.push(record); + } + if let Ok(mut rewritten) = OpenOptions::new() + .create(true) + .write(true) + .truncate(true) + .open(&path) + { + let _ = rewritten.write_all(&live.concat()); + let _ = rewritten.sync_data(); + } + } } fn bloom_positions(filter_len: usize, fingerprint: &[u8; 32]) -> [usize; 4] { diff --git a/src/common/message/secure/tests.rs b/src/common/message/secure/tests.rs index 4fd3625..0c88c3c 100644 --- a/src/common/message/secure/tests.rs +++ b/src/common/message/secure/tests.rs @@ -264,7 +264,8 @@ fn rotating_bloom_retains_a_max_future_timestamp_past_the_next_rotation() { #[test] fn per_credential_admission_limit_does_not_consume_other_keys() { let now = unix_seconds(); - let mut guard = ReplayGuard::new(1024, DEFAULT_REPLAY_WINDOW_SECONDS).with_max_per_key(2); + let mut guard = + ReplayGuard::open(None, 1024, DEFAULT_REPLAY_WINDOW_SECONDS).with_max_per_key(2); assert_eq!(guard.admit(1, &[1_u8; 32], now), FirstFlightAdmit::Fresh); assert_eq!(guard.admit(1, &[2_u8; 32], now), FirstFlightAdmit::Fresh); assert_eq!(guard.admit(1, &[3_u8; 32], now), FirstFlightAdmit::Limited); @@ -272,6 +273,29 @@ fn per_credential_admission_limit_does_not_consume_other_keys() { assert_eq!(guard.admit(1, &[1_u8; 32], now), FirstFlightAdmit::Replayed); } +#[test] +fn persisted_first_flights_survive_replay_guard_restart() { + let mut random = [0_u8; 8]; + let mut rng = rand::rng(); + for byte in &mut random { + *byte = rng.random(); + } + let path = + std::env::temp_dir().join(format!("pb-mapper-replay-{}", u64::from_be_bytes(random))); + let now = unix_seconds(); + let fingerprint = [13_u8; 32]; + { + let mut guard = ReplayGuard::open(Some(path.clone()), 1024, DEFAULT_REPLAY_WINDOW_SECONDS); + assert_eq!(guard.admit(7, &fingerprint, now), FirstFlightAdmit::Fresh); + } + let mut restored = ReplayGuard::open(Some(path.clone()), 1024, DEFAULT_REPLAY_WINDOW_SECONDS); + assert_eq!( + restored.admit(7, &fingerprint, now), + FirstFlightAdmit::Replayed + ); + let _ = std::fs::remove_file(path); +} + #[tokio::test] async fn legacy_initial_frame_validates_against_isolated_relay_key() { let _process_credential_guard = PROCESS_CREDENTIAL_TEST_LOCK.lock().await; diff --git a/ui/lib/l10n/app_en.arb b/ui/lib/l10n/app_en.arb index 19245f4..b2cc307 100644 --- a/ui/lib/l10n/app_en.arb +++ b/ui/lib/l10n/app_en.arb @@ -162,6 +162,8 @@ "checkServer": "Check Server Connectivity", "serverAddressHelp": "Address of the pb-mapper server to connect to", "msgHeaderKeyHelp": "Required. Prefer a pbmt1_ temporary credential issued by the relay; the 32-character administrator key also works.", + "isolatedRelayAdminKey": "Embedded relay administrator key", + "isolatedRelayAdminKeyHelp": "Generated for this app's local relay only. Use it to register or connect to that relay; it is not the outbound MSG_HEADER_KEY.", "keepAliveHelp": "Enable TCP keep-alive for connections", "configServerAddress": "Server Address: {value}", "@configServerAddress": { diff --git a/ui/lib/l10n/app_zh.arb b/ui/lib/l10n/app_zh.arb index 25e5ddc..d59aba0 100644 --- a/ui/lib/l10n/app_zh.arb +++ b/ui/lib/l10n/app_zh.arb @@ -162,6 +162,8 @@ "checkServer": "检测服务器连通性", "serverAddressHelp": "要连接的 pb-mapper 服务器地址", "msgHeaderKeyHelp": "必填。优先使用中继签发的 pbmt1_ 临时凭据，也可使用 32 字符管理员密钥。", + "isolatedRelayAdminKey": "内嵌中继管理员密钥", + "isolatedRelayAdminKeyHelp": "仅用于本应用内嵌中继。用它向本地中继注册或连接，不要把它当成对外的 MSG_HEADER_KEY。", "keepAliveHelp": "为连接启用 TCP keep-alive", "configServerAddress": "服务器地址：{value}", "@configServerAddress": { diff --git a/ui/lib/src/ffi/pb_mapper_api.dart b/ui/lib/src/ffi/pb_mapper_api.dart index 454e4a6..85284d9 100644 --- a/ui/lib/src/ffi/pb_mapper_api.dart +++ b/ui/lib/src/ffi/pb_mapper_api.dart @@ -16,11 +16,13 @@ class ConfigStatus { final String serverAddress; final bool keepAliveEnabled; final String msgHeaderKey; + final String isolatedRelayAdminKey; const ConfigStatus({ required this.serverAddress, required this.keepAliveEnabled, required this.msgHeaderKey, + this.isolatedRelayAdminKey = '', }); factory ConfigStatus.fromMap(Map map) { @@ -31,6 +33,7 @@ class ConfigStatus { ), keepAliveEnabled: _asBool(map['keepAliveEnabled'], fallback: true), msgHeaderKey: _asString(map['msgHeaderKey']), + isolatedRelayAdminKey: _asString(map['isolatedRelayAdminKey']), ); } } diff --git a/ui/lib/src/views/configuration_view.dart b/ui/lib/src/views/configuration_view.dart index 23c3bd0..e61d71f 100644 --- a/ui/lib/src/views/configuration_view.dart +++ b/ui/lib/src/views/configuration_view.dart @@ -365,6 +365,20 @@ class _ConfigurationViewState extends State { helperText: context.l10n.msgHeaderKeyHelp, ), ), + if (_currentConfig?.isolatedRelayAdminKey.isNotEmpty == + true) ...[ + const SizedBox(height: 16), + InputDecorator( + decoration: InputDecoration( + labelText: context.l10n.isolatedRelayAdminKey, + border: const OutlineInputBorder(), + helperText: context.l10n.isolatedRelayAdminKeyHelp, + ), + child: SelectableText( + _currentConfig!.isolatedRelayAdminKey, + ), + ), + ], const SizedBox(height: 16), SwitchListTile( title: const Text('PB_MAPPER_KEEP_ALIVE'), diff --git a/ui/native/pb_mapper_ffi/src/config.rs b/ui/native/pb_mapper_ffi/src/config.rs index 2ba14ed..2efb6d9 100644 --- a/ui/native/pb_mapper_ffi/src/config.rs +++ b/ui/native/pb_mapper_ffi/src/config.rs @@ -9,7 +9,6 @@ use crate::ctl::Origin; use crate::events; use crate::handle::PbMapperHandle; use crate::response::{err_ctl, err_null_handle, ok_data, ok_message, parse_c_string}; -use crate::state::AppConfig; /// Get current app config. #[no_mangle] @@ -20,15 +19,16 @@ pub unsafe extern "C" fn pb_mapper_get_config_json(handle: *mut PbMapperHandle) let handle = unsafe { &mut *handle }; let state = handle.state.clone(); - let config: AppConfig = handle.runtime.block_on(async move { + let (config, isolated_admin_key) = handle.runtime.block_on(async move { let state = state.lock().await; - state.get_config_status().await + (state.get_config_status().await, state.isolated_admin_key()) }); ok_data(json!({ "serverAddress": config.server_address, "keepAliveEnabled": config.keep_alive_enabled, - "msgHeaderKey": config.msg_header_key + "msgHeaderKey": config.msg_header_key, + "isolatedRelayAdminKey": isolated_admin_key.unwrap_or_default(), })) } diff --git a/ui/native/pb_mapper_ffi/src/ctl/mod.rs b/ui/native/pb_mapper_ffi/src/ctl/mod.rs index 1e4866e..d01ecbf 100644 --- a/ui/native/pb_mapper_ffi/src/ctl/mod.rs +++ b/ui/native/pb_mapper_ffi/src/ctl/mod.rs @@ -206,12 +206,16 @@ async fn run( } Command::ConfigGet => { - let config = state.lock().await.get_config_status().await; + let guard = state.lock().await; + let config = guard.get_config_status().await; + let isolated_admin_key = guard.isolated_admin_key(); Ok(proto::Response::ok( Some(json!({ "serverAddress": config.server_address, "keepAliveEnabled": config.keep_alive_enabled, "msgHeaderKeySet": !config.msg_header_key.is_empty(), + "isolatedRelayAdminKeySet": isolated_admin_key.is_some(), + "isolatedRelayAdminKey": isolated_admin_key.unwrap_or_default(), })), None, )) diff --git a/ui/native/pb_mapper_ffi/src/state.rs b/ui/native/pb_mapper_ffi/src/state.rs index c9fed75..a3dc1fd 100644 --- a/ui/native/pb_mapper_ffi/src/state.rs +++ b/ui/native/pb_mapper_ffi/src/state.rs @@ -611,6 +611,12 @@ mod tests { }; assert!(auth_dir.join("admin.key").is_file()); + let isolated = state + .lock() + .await + .isolated_admin_key() + .expect("embedded relay should expose its administrator key"); + assert_eq!(isolated.len(), 32); state .lock() .await diff --git a/ui/native/pb_mapper_ffi/src/state/status.rs b/ui/native/pb_mapper_ffi/src/state/status.rs index 8146a06..489e567 100644 --- a/ui/native/pb_mapper_ffi/src/state/status.rs +++ b/ui/native/pb_mapper_ffi/src/state/status.rs @@ -18,6 +18,14 @@ impl PbMapperState { self.config.clone() } + pub fn isolated_admin_key(&self) -> Option { + let path = self.config_dir.join("auth").join("admin.key"); + std::fs::read_to_string(path).ok().and_then(|raw| { + let key = raw.trim().to_string(); + (!key.is_empty()).then_some(key) + }) + } + pub async fn update_config( &mut self, server_address: String, From ba02af78f006d7bfc78e7265dbca13cdfee66d8b Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 03:17:18 +0800 Subject: [PATCH 24/74] Surface the exhausted salt-replay error on the second admin attempt The first-flight retry loop returned a generic protocol error on the second connection_salt_replayed response, so the dedicated exhausted message was unreachable. --- CHANGELOG.md | 1 + src/bin/pb-mapper/admin.rs | 6 ++++-- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 566f5bf..4cc975f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,7 @@ All notable changes to this project will be documented in this file. - Capped legacy first-flight allocations and kept legacy framing denied when authentication state enters safe mode. - Skipped snapshot compaction while startup is in safe mode, fsynced the auth directory when creating `auth.wal`, cleared retained high-slot entries on reset/rotate, dropped the slot write lock before WAL fail-closed cancellation, fail-closed process checksums after the credential is cleared, and rate-limited first flights per credential before consuming the shared replay filter. - Persisted first-flight replay admissions across restarts, replaced existing auth files atomically on Windows, rejected explicit out-of-range server auth flags, and exposed the embedded relay's isolated administrator key through FFI/UI. +- Made a second administrator first-flight salt replay surface the dedicated retry-exhausted error instead of leaving that path unreachable. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/src/bin/pb-mapper/admin.rs b/src/bin/pb-mapper/admin.rs index 2fa74c1..bf026eb 100644 --- a/src/bin/pb-mapper/admin.rs +++ b/src/bin/pb-mapper/admin.rs @@ -347,9 +347,11 @@ async fn send_admin_request_with_timeout( match response { PbConnResponse::Admin(response) => return Ok(response), PbConnResponse::Error(error) - if error.code == "connection_salt_replayed" && error.retryable && attempt == 0 => + if error.code == "connection_salt_replayed" && error.retryable => { - continue; + if attempt == 0 { + continue; + } } PbConnResponse::Error(error) => { return Err(std::io::Error::other(format!( From df7dcf5afca83959f8e0f74a71ef725e02cec9dc Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 03:30:33 +0800 Subject: [PATCH 25/74] Harden the durable replay log and exclusive auth-state ownership Compact connection.replay while the relay is running, roll back torn appends, and replace the pruned log atomically. Size first-flight admission from PB_MAPPER_NEW_STREAMS_PER_SECOND so the documented stream rate is not clipped. Take an exclusive lock on the authentication state directory so two relays cannot share one store. --- CHANGELOG.md | 1 + src/common/auth.rs | 1 + src/common/auth/persistence.rs | 97 ++++++++++++++++++++++++++- src/common/auth/runtime.rs | 2 + src/common/auth/tests.rs | 24 +++++++ src/common/message/secure/replay.rs | 100 ++++++++++++++++++++++++---- 6 files changed, 211 insertions(+), 14 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4cc975f..fb7fbaf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,7 @@ All notable changes to this project will be documented in this file. - Skipped snapshot compaction while startup is in safe mode, fsynced the auth directory when creating `auth.wal`, cleared retained high-slot entries on reset/rotate, dropped the slot write lock before WAL fail-closed cancellation, fail-closed process checksums after the credential is cleared, and rate-limited first flights per credential before consuming the shared replay filter. - Persisted first-flight replay admissions across restarts, replaced existing auth files atomically on Windows, rejected explicit out-of-range server auth flags, and exposed the embedded relay's isolated administrator key through FFI/UI. - Made a second administrator first-flight salt replay surface the dedicated retry-exhausted error instead of leaving that path unreachable. +- Compacted the durable first-flight replay log while the relay is running, rolled back torn replay-log appends, rewrote that log atomically, sized first-flight admission from `PB_MAPPER_NEW_STREAMS_PER_SECOND`, and took an exclusive lock on the authentication state directory. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/src/common/auth.rs b/src/common/auth.rs index 330a975..1e71fc3 100644 --- a/src/common/auth.rs +++ b/src/common/auth.rs @@ -426,6 +426,7 @@ pub struct AuthRuntime { inner: Weak, command_tx: mpsc::Sender, config: AuthConfig, + _state_lock: Arc, } #[derive(Clone, Debug, Serialize, Deserialize)] diff --git a/src/common/auth/persistence.rs b/src/common/auth/persistence.rs index 70f5c14..d73f7a6 100644 --- a/src/common/auth/persistence.rs +++ b/src/common/auth/persistence.rs @@ -27,6 +27,101 @@ pub fn encrypted_auth_state_exists(state_dir: &Path) -> bool { auth_snapshot_path(state_dir).exists() || auth_wal_path(state_dir).exists() } +pub fn acquire_state_dir_lock(state_dir: &Path) -> Result { + let path = state_dir.join("auth.lock"); + let file = OpenOptions::new() + .create(true) + .read(true) + .write(true) + .truncate(false) + .open(&path) + .map_err(|error| { + AuthFailure::new( + "auth_state_unavailable", + format!("failed to open `{}`: {error}", path.display()), + false, + ) + })?; + lock_exclusive_nonblock(&file).map_err(|error| { + AuthFailure::new( + "auth_state_locked", + format!( + "authentication state directory `{}` is already in use: {error}", + state_dir.display() + ), + false, + ) + })?; + Ok(file) +} + +fn lock_exclusive_nonblock(file: &File) -> std::io::Result<()> { + #[cfg(unix)] + { + extern "C" { + fn flock(fd: i32, operation: i32) -> i32; + } + const LOCK_EX: i32 = 2; + const LOCK_NB: i32 = 4; + use std::os::unix::io::AsRawFd; + if unsafe { flock(file.as_raw_fd(), LOCK_EX | LOCK_NB) } != 0 { + return Err(std::io::Error::last_os_error()); + } + Ok(()) + } + #[cfg(windows)] + { + use std::os::windows::io::AsRawHandle; + const LOCKFILE_FAIL_IMMEDIATELY: u32 = 0x1; + const LOCKFILE_EXCLUSIVE_LOCK: u32 = 0x2; + #[repr(C)] + struct Overlapped { + internal: usize, + internal_high: usize, + offset: u32, + offset_high: u32, + event: *mut core::ffi::c_void, + } + extern "system" { + fn LockFileEx( + file: *mut core::ffi::c_void, + flags: u32, + reserved: u32, + bytes_low: u32, + bytes_high: u32, + overlapped: *mut Overlapped, + ) -> i32; + } + let mut overlapped = Overlapped { + internal: 0, + internal_high: 0, + offset: 0, + offset_high: 0, + event: core::ptr::null_mut(), + }; + let ok = unsafe { + LockFileEx( + file.as_raw_handle(), + LOCKFILE_FAIL_IMMEDIATELY | LOCKFILE_EXCLUSIVE_LOCK, + 0, + 1, + 0, + &mut overlapped, + ) + }; + if ok == 0 { + Err(std::io::Error::last_os_error()) + } else { + Ok(()) + } + } + #[cfg(not(any(unix, windows)))] + { + let _ = file; + Ok(()) + } +} + pub(super) fn compaction_is_allowed(safe_mode: bool) -> bool { !safe_mode } @@ -39,7 +134,7 @@ pub(super) fn clear_retained_high_slot_entries(inner: &AuthStateInner) { .clear(); } -fn replace_file(from: &Path, to: &Path) -> std::io::Result<()> { +pub(crate) fn replace_file(from: &Path, to: &Path) -> std::io::Result<()> { #[cfg(windows)] { use std::os::windows::ffi::OsStrExt; diff --git a/src/common/auth/runtime.rs b/src/common/auth/runtime.rs index 81c969c..bd5a0e0 100644 --- a/src/common/auth/runtime.rs +++ b/src/common/auth/runtime.rs @@ -56,6 +56,7 @@ impl AuthRuntime { sync_process_credential: bool, ) -> Result { prepare_state_dir(&config.state_dir)?; + let state_lock = Arc::new(acquire_state_dir_lock(&config.state_dir)?); let instance_id = load_or_create_instance_id(&config.state_dir)?; let (mut loaded, safe_mode) = load_persisted_state(&config, &admin_key, instance_id); let now = unix_seconds(); @@ -176,6 +177,7 @@ impl AuthRuntime { inner: Arc::downgrade(&inner), command_tx, config: config.clone(), + _state_lock: state_lock, }; tokio::spawn(run_auth_actor( diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index 5d9041a..2339c9b 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -144,6 +144,30 @@ async fn safe_mode_denies_legacy_protocol_instead_of_restoring_the_default() { let _ = std::fs::remove_dir_all(state_dir); } +#[tokio::test] +async fn overlapping_runtimes_cannot_share_an_auth_state_directory() { + let state_dir = temp_state_dir("auth-dir-lock"); + let admin_key = *b"0123456789abcdefghijklmnopqrstuv"; + let config = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 4, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + let first = AuthRuntime::start(admin_key, config.clone()).await.unwrap(); + let error = match AuthRuntime::start(admin_key, config.clone()).await { + Ok(_) => panic!("second runtime should not share the auth directory"), + Err(error) => error, + }; + assert_eq!(error.code, "auth_state_locked"); + drop(first); + tokio::time::sleep(Duration::from_millis(20)).await; + let recovered = AuthRuntime::start(admin_key, config).await.unwrap(); + drop(recovered); + tokio::time::sleep(Duration::from_millis(20)).await; + let _ = std::fs::remove_dir_all(state_dir); +} + #[test] fn safe_mode_startup_does_not_allow_compaction() { assert!(!compaction_is_allowed(true)); diff --git a/src/common/message/secure/replay.rs b/src/common/message/secure/replay.rs index ece6461..3ee859a 100644 --- a/src/common/message/secure/replay.rs +++ b/src/common/message/secure/replay.rs @@ -22,8 +22,23 @@ use std::path::PathBuf; use super::*; -pub(super) const MAX_FIRST_FLIGHTS_PER_KEY: u32 = 8_192; +const DEFAULT_NEW_STREAMS_PER_SECOND: u32 = 100; const REPLAY_RECORD_LEN: usize = 40; +const REPLAY_COMPACT_INTERVAL_SECONDS: u64 = 60; + +fn first_flight_budget(window_seconds: u64) -> u32 { + let streams_per_sec = std::env::var("PB_MAPPER_NEW_STREAMS_PER_SECOND") + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(DEFAULT_NEW_STREAMS_PER_SECOND) + .clamp(1, 1_000_000); + let window = u32::try_from(window_seconds).unwrap_or(u32::MAX); + streams_per_sec + .saturating_mul(2) + .saturating_mul(window) + .saturating_mul(2) + .max(8_192) +} #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(super) enum FirstFlightAdmit { @@ -93,17 +108,20 @@ pub(super) struct ReplayGuard { window_seconds: u64, max_per_key: u32, log_path: Option, + last_compact_at: u64, } impl ReplayGuard { pub(super) fn open(log_path: Option, bytes: usize, window_seconds: u64) -> Self { + let now = unix_seconds(); let mut guard = Self { bloom: RotatingBloom::new(bytes, window_seconds), counts: HashMap::new(), - counts_started_at: unix_seconds(), + counts_started_at: now, window_seconds, - max_per_key: MAX_FIRST_FLIGHTS_PER_KEY, + max_per_key: first_flight_budget(window_seconds), log_path, + last_compact_at: now, }; guard.load_persisted(); guard @@ -133,6 +151,9 @@ impl ReplayGuard { } self.bloom.insert(fingerprint, now); *self.counts.entry(key_id).or_insert(0) += 1; + if now.saturating_sub(self.last_compact_at) >= REPLAY_COMPACT_INTERVAL_SECONDS { + self.compact(now); + } FirstFlightAdmit::Fresh } @@ -152,11 +173,16 @@ impl ReplayGuard { std::fs::create_dir_all(parent)?; } let mut file = OpenOptions::new().create(true).append(true).open(path)?; + let start_len = file.metadata()?.len(); let mut record = [0_u8; REPLAY_RECORD_LEN]; record[..32].copy_from_slice(fingerprint); record[32..].copy_from_slice(&now.to_be_bytes()); - file.write_all(&record)?; - file.sync_data() + if let Err(error) = file.write_all(&record).and_then(|()| file.sync_data()) { + let _ = file.set_len(start_len); + let _ = file.sync_data(); + return Err(error); + } + Ok(()) } fn load_persisted(&mut self) { @@ -183,15 +209,63 @@ impl ReplayGuard { self.bloom.insert(&fingerprint, timestamp); live.push(record); } - if let Ok(mut rewritten) = OpenOptions::new() - .create(true) - .write(true) - .truncate(true) - .open(&path) - { - let _ = rewritten.write_all(&live.concat()); - let _ = rewritten.sync_data(); + if self.rewrite_live(&live).is_ok() { + self.last_compact_at = now; + } + } + + fn compact(&mut self, now: u64) { + let Some(path) = &self.log_path else { + self.last_compact_at = now; + return; + }; + let Ok(mut file) = File::open(path) else { + self.last_compact_at = now; + return; + }; + let mut live = Vec::new(); + let mut record = [0_u8; REPLAY_RECORD_LEN]; + loop { + match file.read(&mut record) { + Ok(0) => break, + Ok(n) if n == REPLAY_RECORD_LEN => {} + _ => break, + } + let timestamp = u64::from_be_bytes(record[32..].try_into().expect("timestamp width")); + if now.saturating_sub(timestamp) <= self.window_seconds { + live.push(record); + } + } + if self.rewrite_live(&live).is_ok() { + self.last_compact_at = now; + } + } + + fn rewrite_live(&self, live: &[[u8; REPLAY_RECORD_LEN]]) -> std::io::Result<()> { + let Some(path) = &self.log_path else { + return Ok(()); + }; + let temporary = path.with_file_name(format!( + ".{}.tmp-{}", + path.file_name() + .and_then(|name| name.to_str()) + .unwrap_or("connection.replay"), + std::process::id() + )); + let result = (|| { + let mut file = OpenOptions::new() + .create_new(true) + .write(true) + .open(&temporary)?; + file.write_all(&live.concat())?; + file.sync_all()?; + drop(file); + crate::common::auth::replace_file(&temporary, path) + })(); + if result.is_err() { + let _ = std::fs::remove_file(&temporary); } + result } } From b8a1e5d48452028caad1ed14272561f1b49a5a1e Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 03:44:28 +0800 Subject: [PATCH 26/74] Recover interrupted root rotation and release auth locks on shutdown Stage admin.key.next before rewriting the snapshot so a crash cannot leave state encrypted under an unpublished key. Fsync the replay-log directory on first creation. Take the state lock before --init-admin-key. Abort accepted connection tasks on relay shutdown so auth.lock is released. --- CHANGELOG.md | 1 + src/bin/pb-mapper.rs | 8 +++++-- src/common/auth.rs | 36 +++++++++++++++++++++++++++++ src/common/auth/actor.rs | 5 +++- src/common/auth/persistence.rs | 2 +- src/common/message/secure/replay.rs | 5 ++++ src/pb_server/runtime.rs | 13 +++++++++-- 7 files changed, 64 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index fb7fbaf..bc1c043 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,7 @@ All notable changes to this project will be documented in this file. - Persisted first-flight replay admissions across restarts, replaced existing auth files atomically on Windows, rejected explicit out-of-range server auth flags, and exposed the embedded relay's isolated administrator key through FFI/UI. - Made a second administrator first-flight salt replay surface the dedicated retry-exhausted error instead of leaving that path unreachable. - Compacted the durable first-flight replay log while the relay is running, rolled back torn replay-log appends, rewrote that log atomically, sized first-flight admission from `PB_MAPPER_NEW_STREAMS_PER_SECOND`, and took an exclusive lock on the authentication state directory. +- Fsynced the replay-log directory on first creation, took the state lock before `--init-admin-key`, aborted accepted connection tasks on relay shutdown, and staged `admin.key.next` so an interrupted root rotation can recover a matching key and snapshot. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/src/bin/pb-mapper.rs b/src/bin/pb-mapper.rs index 3b9c6c2..40325d9 100644 --- a/src/bin/pb-mapper.rs +++ b/src/bin/pb-mapper.rs @@ -19,8 +19,9 @@ use std::time::Duration; use better_mimalloc_rs::MiMalloc; use clap::{Args, Parser, Subcommand, ValueEnum}; use pb_mapper::common::auth::{ - generate_admin_key, initialize_admin_key, write_admin_key_file, AuthConfig, KeyPage, - LegacyProtocolPolicy, MAX_TEMP_KEY_CAPACITY, MAX_TEMP_KEY_TTL, MIN_TEMP_KEY_TTL, + acquire_state_dir_lock, generate_admin_key, initialize_admin_key, write_admin_key_file, + AuthConfig, KeyPage, LegacyProtocolPolicy, MAX_TEMP_KEY_CAPACITY, MAX_TEMP_KEY_TTL, + MIN_TEMP_KEY_TTL, }; use pb_mapper::common::checksum::set_process_msg_header_key; use pb_mapper::common::checksum::{setup_machine_msg_header_key, MACHINE_MSG_HEADER_KEY_PATH}; @@ -272,8 +273,11 @@ async fn run_server(args: ServerArgs) -> Result<(), Box> { } let auth_config = AuthConfig::default(); if args.init_admin_key { + std::fs::create_dir_all(&auth_config.state_dir)?; + let _lock = acquire_state_dir_lock(&auth_config.state_dir)?; let key_path = auth_config.state_dir.join("admin.key"); let key = initialize_admin_key(&key_path, args.force_init_admin_key)?; + drop(_lock); set_process_msg_header_key(Some(&key))?; eprintln!("administrator key initialized at {}", key_path.display()); } else if args.use_machine_msg_header_key { diff --git a/src/common/auth.rs b/src/common/auth.rs index 1e71fc3..9bd69de 100644 --- a/src/common/auth.rs +++ b/src/common/auth.rs @@ -619,6 +619,40 @@ fn validate_admin_credential(raw: &str) -> Result { Ok(credential) } +fn recover_admin_key_after_rotation( + state_dir: &Path, + current: &str, +) -> Result { + let snapshot_path = auth_snapshot_path(state_dir); + if !snapshot_path.exists() { + return Ok(current.to_string()); + } + let bytes = std::fs::read(&snapshot_path).map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to read `{}`: {error}", snapshot_path.display()), + false, + ) + })?; + if let Ok(Credential::Admin(current_key)) = parse_credential(current.trim()) { + if open_blob(¤t_key, &bytes).is_ok() { + return Ok(current.to_string()); + } + } + let Some(next) = read_admin_key(&state_dir.join("admin.key.next"))? else { + return Ok(current.to_string()); + }; + let Ok(Credential::Admin(next_key)) = parse_credential(next.trim()) else { + return Ok(current.to_string()); + }; + if open_blob(&next_key, &bytes).is_err() { + return Ok(current.to_string()); + } + write_admin_key(state_dir, next.trim())?; + let _ = std::fs::remove_file(state_dir.join("admin.key.next")); + Ok(next) +} + fn load_server_admin_credential(state_dir: &Path) -> Result { let path = state_dir.join("admin.key"); let raw = if let Some(raw) = read_admin_key(&path)? { @@ -670,6 +704,7 @@ fn load_server_admin_credential(state_dir: &Path) -> Result Result std::io::Result<()> { std::fs::rename(from, to) } -fn sync_parent_directory(path: &Path) -> Result<(), AuthFailure> { +pub(crate) fn sync_parent_directory(path: &Path) -> Result<(), AuthFailure> { #[cfg(unix)] if let Some(parent) = path.parent() { File::open(parent) diff --git a/src/common/message/secure/replay.rs b/src/common/message/secure/replay.rs index 3ee859a..e46dbb4 100644 --- a/src/common/message/secure/replay.rs +++ b/src/common/message/secure/replay.rs @@ -172,6 +172,7 @@ impl ReplayGuard { if let Some(parent) = path.parent() { std::fs::create_dir_all(parent)?; } + let created = !path.exists(); let mut file = OpenOptions::new().create(true).append(true).open(path)?; let start_len = file.metadata()?.len(); let mut record = [0_u8; REPLAY_RECORD_LEN]; @@ -182,6 +183,10 @@ impl ReplayGuard { let _ = file.sync_data(); return Err(error); } + if created { + crate::common::auth::sync_parent_directory(path) + .map_err(|error| std::io::Error::other(error.to_string()))?; + } Ok(()) } diff --git a/src/pb_server/runtime.rs b/src/pb_server/runtime.rs index 8b8a655..ec29ce5 100644 --- a/src/pb_server/runtime.rs +++ b/src/pb_server/runtime.rs @@ -104,6 +104,7 @@ pub async fn run_server_on_listener( let new_streams_per_second = env_limit("PB_MAPPER_NEW_STREAMS_PER_SECOND", 100); let new_streams_burst = env_limit("PB_MAPPER_NEW_STREAMS_BURST", 200); let mut next_server_generation = 1_u64; + let mut connection_tasks = Vec::new(); let listen_addr = listener.local_addr()?; tracing::info!( @@ -358,12 +359,14 @@ pub async fn run_server_on_listener( ); let manager_task_sender = manager.get_task_sender(); let security = security.clone(); - tokio::spawn(async move { + connection_tasks + .retain(|handle: &tokio::task::JoinHandle<()>| !handle.is_finished()); + connection_tasks.push(tokio::spawn(async move { snafu_error_handle!( handle_conn(conn_id, peer_addr, manager_task_sender, stream, security) .await ); - }); + })); } ManagerTask::DeRegisterServerConn { key, conn_id } => { let removed_from_service_map = @@ -922,6 +925,9 @@ pub async fn run_server_on_listener( } ManagerTask::Shutdown => { tracing::info!("Server shutdown requested, stopping main loop"); + for handle in connection_tasks.drain(..) { + handle.abort(); + } break; } } @@ -933,6 +939,9 @@ pub async fn run_server_on_listener( if let Some(handle) = status_forward_handle { handle.abort(); } + for handle in connection_tasks { + handle.abort(); + } tracing::info!("Server shutdown completed"); Ok(()) } From 60b7582cf18251b17f601a7c1de430f7760bb303 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 03:51:42 +0800 Subject: [PATCH 27/74] Gate the embedded relay administrator key behind an explicit reveal Routine config fetches now only report whether the isolated key exists. The plaintext secret is returned only from a dedicated reveal call, and the configuration page shows it after the user asks. --- CHANGELOG.md | 1 + ui/lib/l10n/app_en.arb | 1 + ui/lib/l10n/app_zh.arb | 1 + ui/lib/src/ffi/pb_mapper_api.dart | 16 +++++++++++ ui/lib/src/ffi/pb_mapper_ffi.dart | 5 ++++ ui/lib/src/ffi/pb_mapper_service.dart | 7 +++++ ui/lib/src/views/configuration_view.dart | 34 ++++++++++++++++-------- ui/native/pb_mapper_ffi/src/config.rs | 31 ++++++++++++++++++--- ui/native/pb_mapper_ffi/src/ctl/mod.rs | 1 - ui/native/pb_mapper_ffi/src/lib.rs | 4 ++- ui/test/fake_pb_mapper_api.dart | 6 +++++ 11 files changed, 91 insertions(+), 16 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index bc1c043..20bea9b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,7 @@ All notable changes to this project will be documented in this file. - Made a second administrator first-flight salt replay surface the dedicated retry-exhausted error instead of leaving that path unreachable. - Compacted the durable first-flight replay log while the relay is running, rolled back torn replay-log appends, rewrote that log atomically, sized first-flight admission from `PB_MAPPER_NEW_STREAMS_PER_SECOND`, and took an exclusive lock on the authentication state directory. - Fsynced the replay-log directory on first creation, took the state lock before `--init-admin-key`, aborted accepted connection tasks on relay shutdown, and staged `admin.key.next` so an interrupted root rotation can recover a matching key and snapshot. +- Stopped returning the embedded relay administrator key from routine config fetches; revealing it now requires an explicit FFI/UI action. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/ui/lib/l10n/app_en.arb b/ui/lib/l10n/app_en.arb index b2cc307..6d56cf7 100644 --- a/ui/lib/l10n/app_en.arb +++ b/ui/lib/l10n/app_en.arb @@ -164,6 +164,7 @@ "msgHeaderKeyHelp": "Required. Prefer a pbmt1_ temporary credential issued by the relay; the 32-character administrator key also works.", "isolatedRelayAdminKey": "Embedded relay administrator key", "isolatedRelayAdminKeyHelp": "Generated for this app's local relay only. Use it to register or connect to that relay; it is not the outbound MSG_HEADER_KEY.", + "isolatedRelayReveal": "Reveal embedded relay key", "keepAliveHelp": "Enable TCP keep-alive for connections", "configServerAddress": "Server Address: {value}", "@configServerAddress": { diff --git a/ui/lib/l10n/app_zh.arb b/ui/lib/l10n/app_zh.arb index d59aba0..b938b9a 100644 --- a/ui/lib/l10n/app_zh.arb +++ b/ui/lib/l10n/app_zh.arb @@ -164,6 +164,7 @@ "msgHeaderKeyHelp": "必填。优先使用中继签发的 pbmt1_ 临时凭据，也可使用 32 字符管理员密钥。", "isolatedRelayAdminKey": "内嵌中继管理员密钥", "isolatedRelayAdminKeyHelp": "仅用于本应用内嵌中继。用它向本地中继注册或连接，不要把它当成对外的 MSG_HEADER_KEY。", + "isolatedRelayReveal": "显示内嵌中继密钥", "keepAliveHelp": "为连接启用 TCP keep-alive", "configServerAddress": "服务器地址：{value}", "@configServerAddress": { diff --git a/ui/lib/src/ffi/pb_mapper_api.dart b/ui/lib/src/ffi/pb_mapper_api.dart index 85284d9..fb099b7 100644 --- a/ui/lib/src/ffi/pb_mapper_api.dart +++ b/ui/lib/src/ffi/pb_mapper_api.dart @@ -16,12 +16,14 @@ class ConfigStatus { final String serverAddress; final bool keepAliveEnabled; final String msgHeaderKey; + final bool isolatedRelayAdminKeySet; final String isolatedRelayAdminKey; const ConfigStatus({ required this.serverAddress, required this.keepAliveEnabled, required this.msgHeaderKey, + this.isolatedRelayAdminKeySet = false, this.isolatedRelayAdminKey = '', }); @@ -33,6 +35,10 @@ class ConfigStatus { ), keepAliveEnabled: _asBool(map['keepAliveEnabled'], fallback: true), msgHeaderKey: _asString(map['msgHeaderKey']), + isolatedRelayAdminKeySet: _asBool( + map['isolatedRelayAdminKeySet'], + fallback: false, + ), isolatedRelayAdminKey: _asString(map['isolatedRelayAdminKey']), ); } @@ -244,6 +250,7 @@ abstract interface class PbMapperApiClient { Future setAppDirectoryPath(String path); Future fetchConfig(); + Future revealIsolatedRelayAdminKey(); Future updateConfig({ required String serverAddress, required bool keepAlive, @@ -313,6 +320,15 @@ class PbMapperApi implements PbMapperApiClient { ); } + @override + Future revealIsolatedRelayAdminKey() async { + final result = await _service.revealIsolatedRelayAdminKey(); + if (result['success'] == true) { + return _asString(_asMap(result['data'])['isolatedRelayAdminKey']); + } + return ''; + } + @override Future updateConfig({ required String serverAddress, diff --git a/ui/lib/src/ffi/pb_mapper_ffi.dart b/ui/lib/src/ffi/pb_mapper_ffi.dart index 30b9529..b0b38f5 100644 --- a/ui/lib/src/ffi/pb_mapper_ffi.dart +++ b/ui/lib/src/ffi/pb_mapper_ffi.dart @@ -296,6 +296,11 @@ class PbMapperFFI { 'pb_mapper_get_config_json', ); + late final pbMapperRevealIsolatedAdminKey = lib + .lookupFunction<_PbMapperGetConfigNative, _PbMapperGetConfigDart>( + 'pb_mapper_reveal_isolated_admin_key', + ); + late final pbMapperUpdateConfig = lib .lookupFunction<_PbMapperUpdateConfigNative, _PbMapperUpdateConfigDart>( 'pb_mapper_update_config', diff --git a/ui/lib/src/ffi/pb_mapper_service.dart b/ui/lib/src/ffi/pb_mapper_service.dart index 2054d34..fe25a8f 100644 --- a/ui/lib/src/ffi/pb_mapper_service.dart +++ b/ui/lib/src/ffi/pb_mapper_service.dart @@ -240,6 +240,10 @@ class PbMapperService { return _runJsonOnWorker('getConfig', {}); } + Future> revealIsolatedRelayAdminKey() { + return _runJsonOnWorker('revealIsolatedRelayAdminKey', {}); + } + Future> updateConfig({ required String serverAddress, required bool keepAlive, @@ -383,6 +387,9 @@ Map _callJsonIsolate(Map params) { case 'getConfig': result = ffi.pbMapperGetConfig(handle); break; + case 'revealIsolatedRelayAdminKey': + result = ffi.pbMapperRevealIsolatedAdminKey(handle); + break; case 'updateConfig': arg1 = (params['serverAddress'] as String).toNativeUtf8(); arg2 = (params['msgHeaderKey'] as String).toNativeUtf8(); diff --git a/ui/lib/src/views/configuration_view.dart b/ui/lib/src/views/configuration_view.dart index e61d71f..e74fa3c 100644 --- a/ui/lib/src/views/configuration_view.dart +++ b/ui/lib/src/views/configuration_view.dart @@ -32,6 +32,7 @@ class _ConfigurationViewState extends State { bool? _serverReachable; String _serverCheckMessage = ''; ConfigStatus? _currentConfig; + String _revealedIsolatedRelayAdminKey = ''; ChangeSubscription? _changes; @@ -123,6 +124,14 @@ class _ConfigurationViewState extends State { } } + Future _revealIsolatedRelayAdminKey() async { + final key = await _api.revealIsolatedRelayAdminKey(); + if (!mounted) return; + setState(() { + _revealedIsolatedRelayAdminKey = key; + }); + } + Future _checkServerConnection() async { if (_isCheckingServer) return; setState(() { @@ -365,19 +374,22 @@ class _ConfigurationViewState extends State { helperText: context.l10n.msgHeaderKeyHelp, ), ), - if (_currentConfig?.isolatedRelayAdminKey.isNotEmpty == - true) ...[ + if (_currentConfig?.isolatedRelayAdminKeySet == true) ...[ const SizedBox(height: 16), - InputDecorator( - decoration: InputDecoration( - labelText: context.l10n.isolatedRelayAdminKey, - border: const OutlineInputBorder(), - helperText: context.l10n.isolatedRelayAdminKeyHelp, - ), - child: SelectableText( - _currentConfig!.isolatedRelayAdminKey, + if (_revealedIsolatedRelayAdminKey.isEmpty) + OutlinedButton( + onPressed: _revealIsolatedRelayAdminKey, + child: Text(context.l10n.isolatedRelayReveal), + ) + else + InputDecorator( + decoration: InputDecoration( + labelText: context.l10n.isolatedRelayAdminKey, + border: const OutlineInputBorder(), + helperText: context.l10n.isolatedRelayAdminKeyHelp, + ), + child: SelectableText(_revealedIsolatedRelayAdminKey), ), - ), ], const SizedBox(height: 16), SwitchListTile( diff --git a/ui/native/pb_mapper_ffi/src/config.rs b/ui/native/pb_mapper_ffi/src/config.rs index 2efb6d9..3ddda8a 100644 --- a/ui/native/pb_mapper_ffi/src/config.rs +++ b/ui/native/pb_mapper_ffi/src/config.rs @@ -19,16 +19,41 @@ pub unsafe extern "C" fn pb_mapper_get_config_json(handle: *mut PbMapperHandle) let handle = unsafe { &mut *handle }; let state = handle.state.clone(); - let (config, isolated_admin_key) = handle.runtime.block_on(async move { + let (config, isolated_admin_key_set) = handle.runtime.block_on(async move { let state = state.lock().await; - (state.get_config_status().await, state.isolated_admin_key()) + ( + state.get_config_status().await, + state.isolated_admin_key().is_some(), + ) }); ok_data(json!({ "serverAddress": config.server_address, "keepAliveEnabled": config.keep_alive_enabled, "msgHeaderKey": config.msg_header_key, - "isolatedRelayAdminKey": isolated_admin_key.unwrap_or_default(), + "isolatedRelayAdminKeySet": isolated_admin_key_set, + })) +} + +/// Reveal the embedded relay administrator key. This is a separate call so +/// routine config fetches cannot leak the root secret. +#[no_mangle] +pub unsafe extern "C" fn pb_mapper_reveal_isolated_admin_key( + handle: *mut PbMapperHandle, +) -> *mut c_char { + if handle.is_null() { + return err_null_handle(); + } + + let handle = unsafe { &mut *handle }; + let state = handle.state.clone(); + let key = handle.runtime.block_on(async move { + let state = state.lock().await; + state.isolated_admin_key() + }); + + ok_data(json!({ + "isolatedRelayAdminKey": key.unwrap_or_default(), })) } diff --git a/ui/native/pb_mapper_ffi/src/ctl/mod.rs b/ui/native/pb_mapper_ffi/src/ctl/mod.rs index d01ecbf..034b9a7 100644 --- a/ui/native/pb_mapper_ffi/src/ctl/mod.rs +++ b/ui/native/pb_mapper_ffi/src/ctl/mod.rs @@ -215,7 +215,6 @@ async fn run( "keepAliveEnabled": config.keep_alive_enabled, "msgHeaderKeySet": !config.msg_header_key.is_empty(), "isolatedRelayAdminKeySet": isolated_admin_key.is_some(), - "isolatedRelayAdminKey": isolated_admin_key.unwrap_or_default(), })), None, )) diff --git a/ui/native/pb_mapper_ffi/src/lib.rs b/ui/native/pb_mapper_ffi/src/lib.rs index c3db9e8..4635488 100644 --- a/ui/native/pb_mapper_ffi/src/lib.rs +++ b/ui/native/pb_mapper_ffi/src/lib.rs @@ -22,7 +22,9 @@ pub use client::{ pb_mapper_connect_service, pb_mapper_delete_client_config, pb_mapper_disconnect_service, pb_mapper_get_client_configs_json, pb_mapper_get_client_status_json, }; -pub use config::{pb_mapper_get_config_json, pb_mapper_update_config}; +pub use config::{ + pb_mapper_get_config_json, pb_mapper_reveal_isolated_admin_key, pb_mapper_update_config, +}; pub use events::pb_mapper_set_change_callback; pub use handle::{ pb_mapper_create, pb_mapper_destroy, pb_mapper_set_app_dir, pb_mapper_start_control_server, diff --git a/ui/test/fake_pb_mapper_api.dart b/ui/test/fake_pb_mapper_api.dart index fa4220d..58b384c 100644 --- a/ui/test/fake_pb_mapper_api.dart +++ b/ui/test/fake_pb_mapper_api.dart @@ -36,6 +36,12 @@ class FakePbMapperApi implements PbMapperApiClient { @override Future fetchConfig() async => config; + @override + Future revealIsolatedRelayAdminKey() async { + calls.add('revealIsolatedRelayAdminKey'); + return config.isolatedRelayAdminKey; + } + @override Future updateConfig({ required String serverAddress, From 9ab8f0aaabf366853c5ddbfcae82aa1f5f00e616 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 04:11:27 +0800 Subject: [PATCH 28/74] Recover interrupted resets and lock before first-start key creation Take auth.lock before loading or creating admin.key, stage server-instance-id.next so a crash mid-reset can promote a matching snapshot, and use a user-writable Linux auth directory when the system path is not usable. --- CHANGELOG.md | 1 + docs/authentication-v2.md | 8 +- docs/authentication-v2.zh-CN.md | 7 +- docs/user-guide.md | 10 +- docs/user-guide.zh-CN.md | 10 +- src/bin/pb-mapper.rs | 4 +- src/common/auth.rs | 66 ++++++++++- src/common/auth/actor.rs | 22 ++-- src/common/auth/persistence.rs | 92 ++++++++++++--- src/common/auth/runtime.rs | 18 +-- src/common/auth/tests.rs | 193 +++++++++++++++++++++++++++++++- 11 files changed, 375 insertions(+), 56 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 20bea9b..3ac17d1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,6 +19,7 @@ All notable changes to this project will be documented in this file. - Compacted the durable first-flight replay log while the relay is running, rolled back torn replay-log appends, rewrote that log atomically, sized first-flight admission from `PB_MAPPER_NEW_STREAMS_PER_SECOND`, and took an exclusive lock on the authentication state directory. - Fsynced the replay-log directory on first creation, took the state lock before `--init-admin-key`, aborted accepted connection tasks on relay shutdown, and staged `admin.key.next` so an interrupted root rotation can recover a matching key and snapshot. - Stopped returning the embedded relay administrator key from routine config fetches; revealing it now requires an explicit FFI/UI action. +- Took `auth.lock` before loading or creating `admin.key`, staged `server-instance-id.next` so an interrupted reset can recover a matching snapshot, and used a user-writable Linux auth directory when `/var/lib/pb-mapper/auth` is not usable. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/docs/authentication-v2.md b/docs/authentication-v2.md index a35236a..fcf0dbb 100644 --- a/docs/authentication-v2.md +++ b/docs/authentication-v2.md @@ -199,14 +199,16 @@ revocation and hard connection closure deterministic. ## Persistence and safe mode -The Linux system-service state directory is `/var/lib/pb-mapper/auth`. macOS -and Windows desktop binaries default to a user-writable application directory -instead of `/var/lib`: +The Linux system-service state directory is `/var/lib/pb-mapper/auth`. +Unprivileged Linux, macOS, and Windows desktop binaries default to a +user-writable application directory instead of `/var/lib`: | File | Purpose | | --- | --- | | `admin.key` | Root credential, mode `0600` | +| `admin.key.next` | Staged root key for an in-flight rotation | | `server-instance-id` | 16-byte persistent derivation identity | +| `server-instance-id.next` | Staged instance id for an in-flight reset | | `auth.snapshot` | AES-256-GCM encrypted compact slot state | | `auth.wal` | Length-prefixed, individually encrypted mutations and audit records | diff --git a/docs/authentication-v2.zh-CN.md b/docs/authentication-v2.zh-CN.md index afb4cd0..eb7b7ff 100644 --- a/docs/authentication-v2.zh-CN.md +++ b/docs/authentication-v2.zh-CN.md @@ -153,13 +153,16 @@ tombstone 以给出稳定错误后，槽位可以复用。显式 `key gc` 可立 ## 持久化与安全模式 -Linux 系统服务默认目录是 `/var/lib/pb-mapper/auth`，权限为 `0700`。macOS 与 -Windows 桌面二进制默认写到用户可写的应用目录，而不是 `/var/lib`： +Linux 系统服务默认目录是 `/var/lib/pb-mapper/auth`，权限为 `0700`。无特权 +Linux、macOS 与 Windows 桌面二进制默认写到用户可写的应用目录，而不是 +`/var/lib`： | 文件 | 用途 | | --- | --- | | `admin.key` | 根凭据，权限 `0600` | +| `admin.key.next` | 轮换进行中的暂存根密钥 | | `server-instance-id` | 16 字节持久派生身份 | +| `server-instance-id.next` | reset 进行中的暂存实例 ID | | `auth.snapshot` | AES-256-GCM 加密的紧凑槽位快照 | | `auth.wal` | 带长度前缀、逐条加密的 mutation 与 audit | diff --git a/docs/user-guide.md b/docs/user-guide.md index f8209f7..dc4a4b1 100644 --- a/docs/user-guide.md +++ b/docs/user-guide.md @@ -112,7 +112,7 @@ Optional flags: - `--ipv6`: enable IPv6 listening - `--keep-alive`: enable TCP keep-alive -- `--auth-state-dir`: authentication state directory (Linux system default `/var/lib/pb-mapper/auth`; macOS and Windows use a user-writable application directory) +- `--auth-state-dir`: authentication state directory (Linux services default `/var/lib/pb-mapper/auth`; unprivileged Linux, macOS, and Windows use a user-writable application directory) - `--max-temporary-keys`: fixed temporary-key slot capacity (default `65536`) - `--max-temporary-key-ttl`: maximum issued TTL (default `30d`) - `--legacy-protocol allow|deny`: initial legacy-client policy @@ -122,9 +122,9 @@ Optional flags: On first start, the relay creates a random administrator key in its authentication state directory (`admin.key`). On Linux system services that -is `/var/lib/pb-mapper/auth/admin.key`. Desktop macOS and Windows builds use -a user-writable application directory instead. There is no built-in default -credential. +is `/var/lib/pb-mapper/auth/admin.key`. Unprivileged Linux, macOS, and +Windows builds use a user-writable application directory instead. There is +no built-in default credential. Keep the administrator key on the relay host and use it to issue a temporary credential for a workload: @@ -235,7 +235,7 @@ flutter run - `PB_MAPPER_SERVER`: default server address for the CLI - `MSG_HEADER_KEY`: 32-character administrator key or a `pbmt1_` temporary credential -- `PB_MAPPER_AUTH_STATE_DIR`: relay auth-state directory (Linux system default `/var/lib/pb-mapper/auth`; macOS and Windows use a user-writable application directory) +- `PB_MAPPER_AUTH_STATE_DIR`: relay auth-state directory (Linux services default `/var/lib/pb-mapper/auth`; unprivileged Linux, macOS, and Windows use a user-writable application directory) - `PB_MAPPER_AUTH_MAX_TEMP_KEYS`: fixed temporary-key capacity, default `65536` - `PB_MAPPER_AUTH_MAX_TEMP_TTL_SECS`: maximum temporary-key TTL, default 30 days - `PB_MAPPER_LEGACY_PROTOCOL`: `allow` or `deny`, default `allow` diff --git a/docs/user-guide.zh-CN.md b/docs/user-guide.zh-CN.md index 9cde552..01a80e8 100644 --- a/docs/user-guide.zh-CN.md +++ b/docs/user-guide.zh-CN.md @@ -112,7 +112,7 @@ pb-mapper server --port 7666 - `--ipv6`：开启 IPv6 监听 - `--keep-alive`：开启 TCP keep-alive -- `--auth-state-dir`：认证状态目录（Linux 系统服务默认 `/var/lib/pb-mapper/auth`；macOS 与 Windows 使用当前用户可写的应用目录） +- `--auth-state-dir`：认证状态目录（Linux 系统服务默认 `/var/lib/pb-mapper/auth`；无特权 Linux、macOS 与 Windows 使用当前用户可写的应用目录） - `--max-temporary-keys`：临时 key 固定槽位容量，默认 `65536` - `--max-temporary-key-ttl`：临时 key 最大 TTL，默认 `30d` - `--legacy-protocol allow|deny`：旧协议初始接入策略 @@ -121,9 +121,9 @@ pb-mapper server --port 7666 ### 管理员密钥与临时凭据 中继首次启动时会在认证状态目录生成随机管理员密钥（`admin.key`）。Linux 系统服务 -默认写到 `/var/lib/pb-mapper/auth/admin.key`；macOS 与 Windows 桌面构建则使用当前 -用户可写的应用目录。系统不再提供内置默认 key。管理员密钥留在中继机器上，用它为 -业务签发临时凭据： +默认写到 `/var/lib/pb-mapper/auth/admin.key`；无特权 Linux、macOS 与 Windows +桌面构建则使用当前用户可写的应用目录。系统不再提供内置默认 key。管理员密钥留在 +中继机器上，用它为业务签发临时凭据： ```bash export MSG_HEADER_KEY="$(sudo cat /var/lib/pb-mapper/auth/admin.key)" @@ -228,7 +228,7 @@ flutter run - `PB_MAPPER_SERVER`：CLI 默认服务器地址 - `MSG_HEADER_KEY`：32 字符管理员密钥或 `pbmt1_` 临时凭据 -- `PB_MAPPER_AUTH_STATE_DIR`：中继认证状态目录（Linux 系统服务默认 `/var/lib/pb-mapper/auth`；macOS 与 Windows 使用当前用户可写的应用目录） +- `PB_MAPPER_AUTH_STATE_DIR`：中继认证状态目录（Linux 系统服务默认 `/var/lib/pb-mapper/auth`；无特权 Linux、macOS 与 Windows 使用当前用户可写的应用目录） - `PB_MAPPER_AUTH_MAX_TEMP_KEYS`：临时 key 固定容量，默认 `65536` - `PB_MAPPER_AUTH_MAX_TEMP_TTL_SECS`：临时 key 最大 TTL，默认 30 天 - `PB_MAPPER_LEGACY_PROTOCOL`：`allow` 或 `deny`，默认 `allow` diff --git a/src/bin/pb-mapper.rs b/src/bin/pb-mapper.rs index 40325d9..e1ec083 100644 --- a/src/bin/pb-mapper.rs +++ b/src/bin/pb-mapper.rs @@ -90,8 +90,8 @@ struct ServerArgs { #[arg(long, default_value_t = false)] use_machine_msg_header_key: bool, /// Directory containing encrypted authentication state and the administrator key file. - /// Defaults to /var/lib/pb-mapper/auth on Linux, or a user-writable application - /// directory on macOS and Windows. + /// Defaults to /var/lib/pb-mapper/auth for Linux services or a writable system + /// directory; otherwise a user-writable application directory. #[arg(long)] auth_state_dir: Option, /// Create a random administrator key before starting the relay. diff --git a/src/common/auth.rs b/src/common/auth.rs index 9bd69de..402e37b 100644 --- a/src/common/auth.rs +++ b/src/common/auth.rs @@ -90,9 +90,9 @@ pub fn default_auth_state_dir() -> PathBuf { .unwrap_or_else(platform_default_auth_state_dir) } -/// Linux systemd/Docker keep `/var/lib/pb-mapper/auth`. Desktop macOS and -/// Windows binaries run as a normal user, so they need an application data -/// directory instead of a root-owned system path. +/// Linux systemd/Docker keep `/var/lib/pb-mapper/auth` when that path is usable +/// (root, or an already-writable service directory). Unprivileged Linux, +/// macOS, and Windows binaries need an application data directory instead. pub(crate) fn platform_default_auth_state_dir() -> PathBuf { #[cfg(windows)] { @@ -115,8 +115,66 @@ pub(crate) fn platform_default_auth_state_dir() -> PathBuf { } #[cfg(not(any(windows, target_os = "macos")))] { - PathBuf::from(DEFAULT_AUTH_STATE_DIR) + linux_default_auth_state_dir( + unix_effective_uid(), + linux_system_auth_dir_usable(), + std::env::var_os("XDG_DATA_HOME").as_deref(), + std::env::var_os("HOME").as_deref(), + ) + } +} + +pub(crate) fn linux_default_auth_state_dir( + euid: u32, + system_dir_usable: bool, + xdg_data_home: Option<&std::ffi::OsStr>, + home: Option<&std::ffi::OsStr>, +) -> PathBuf { + if euid == 0 || system_dir_usable { + return PathBuf::from(DEFAULT_AUTH_STATE_DIR); + } + if let Some(xdg) = xdg_data_home { + if !xdg.is_empty() { + return PathBuf::from(xdg).join("pb-mapper").join("auth"); + } + } + if let Some(home) = home { + if !home.is_empty() { + return PathBuf::from(home) + .join(".local") + .join("share") + .join("pb-mapper") + .join("auth"); + } + } + PathBuf::from(DEFAULT_AUTH_STATE_DIR) +} + +#[cfg(not(any(windows, target_os = "macos")))] +fn unix_effective_uid() -> u32 { + extern "C" { + fn geteuid() -> u32; + } + unsafe { geteuid() } +} + +#[cfg(not(any(windows, target_os = "macos")))] +fn linux_system_auth_dir_usable() -> bool { + let path = Path::new(DEFAULT_AUTH_STATE_DIR); + path.is_dir() && unix_path_is_writable(path) +} + +#[cfg(not(any(windows, target_os = "macos")))] +fn unix_path_is_writable(path: &Path) -> bool { + use std::os::unix::ffi::OsStrExt; + let Ok(c_path) = std::ffi::CString::new(path.as_os_str().as_bytes()) else { + return false; + }; + extern "C" { + fn access(pathname: *const std::os::raw::c_char, mode: i32) -> i32; } + const W_OK: i32 = 2; + unsafe { access(c_path.as_ptr(), W_OK) == 0 } } impl Default for AuthConfig { diff --git a/src/common/auth/actor.rs b/src/common/auth/actor.rs index 225e288..fc57bdc 100644 --- a/src/common/auth/actor.rs +++ b/src/common/auth/actor.rs @@ -697,21 +697,23 @@ fn actor_reset( let mut snapshot = empty_snapshot(inner, new_instance_id, admin_replays); push_persisted_audit(&mut snapshot.audit_records, reset_audit.clone()); let admin_key = inner.admin_key(); - if let Err(error) = write_snapshot_and_truncate_wal(config, &admin_key, &snapshot) { + let next_instance_path = config.state_dir.join("server-instance-id.next"); + if let Err(error) = atomic_write(&next_instance_path, &new_instance_id, 0o600) + .and_then(|()| write_snapshot_and_truncate_wal(config, &admin_key, &snapshot)) + .and_then(|()| { + atomic_write( + &config.state_dir.join("server-instance-id"), + &new_instance_id, + 0o600, + ) + }) + { inner.safe_mode.store(true, Ordering::Release); cancel_all_temporary_leases(inner); return Err(error); } + let _ = std::fs::remove_file(&next_instance_path); push_audit_record(inner, reset_audit); - if let Err(error) = atomic_write( - &config.state_dir.join("server-instance-id"), - &new_instance_id, - 0o600, - ) { - inner.safe_mode.store(true, Ordering::Release); - cancel_all_temporary_leases(inner); - return Err(error); - } cancel_all_temporary_leases(inner); { diff --git a/src/common/auth/persistence.rs b/src/common/auth/persistence.rs index 61ee713..efcdebd 100644 --- a/src/common/auth/persistence.rs +++ b/src/common/auth/persistence.rs @@ -1,7 +1,7 @@ //! Durable, encrypted authentication state and audit/replay retention. //! //! ```text -//! startup: admin.key -> decrypt snapshot -> replay WAL -> normalize -> in-memory state +//! startup: lock -> admin.key -> recover instance id -> decrypt snapshot -> replay WAL //! mutation: command -> fsync encrypted WAL -> publish hot-state change //! compact: hot state + audit + replay set -> snapshot -> truncate WAL //! ``` @@ -27,6 +27,13 @@ pub fn encrypted_auth_state_exists(state_dir: &Path) -> bool { auth_snapshot_path(state_dir).exists() || auth_wal_path(state_dir).exists() } +/// Create the state directory and take `auth.lock` before any credential or +/// snapshot file is read or written. +pub(super) fn prepare_state_dir_and_lock(state_dir: &Path) -> Result, AuthFailure> { + prepare_state_dir(state_dir)?; + Ok(Arc::new(acquire_state_dir_lock(state_dir)?)) +} + pub fn acquire_state_dir_lock(state_dir: &Path) -> Result { let path = state_dir.join("auth.lock"); let file = OpenOptions::new() @@ -807,27 +814,80 @@ pub(super) fn load_or_create_instance_id( path: &Path, ) -> Result<[u8; INSTANCE_ID_LEN], AuthFailure> { let instance_path = path.join("server-instance-id"); - if instance_path.exists() { - let bytes = std::fs::read(&instance_path).map_err(|error| { - AuthFailure::new( - "auth_state_unavailable", - format!("failed to read `{}`: {error}", instance_path.display()), - false, - ) - })?; - return bytes.try_into().map_err(|_| { - AuthFailure::new( - "auth_state_unavailable", - "server instance id must be exactly 16 bytes", - false, - ) - }); + if let Some(instance_id) = read_instance_id_file(&instance_path)? { + return Ok(instance_id); } let instance_id = random_instance_id(); atomic_write(&instance_path, &instance_id, 0o600)?; Ok(instance_id) } +pub(super) fn read_instance_id_file( + path: &Path, +) -> Result, AuthFailure> { + if !path.exists() { + return Ok(None); + } + let bytes = std::fs::read(path).map_err(|error| { + AuthFailure::new( + "auth_state_unavailable", + format!("failed to read `{}`: {error}", path.display()), + false, + ) + })?; + bytes.try_into().map(Some).map_err(|_| { + AuthFailure::new( + "auth_state_unavailable", + "server instance id must be exactly 16 bytes", + false, + ) + }) +} + +/// Promote `server-instance-id.next` when the snapshot already belongs to it. +/// +/// Reset writes that staged file, then the empty snapshot, then the live +/// instance-id file. A crash after the snapshot lands would otherwise fail +/// closed on the next start because the live file still has the old id. +pub(super) fn recover_instance_id_after_reset( + state_dir: &Path, + admin_key: &AesKeyType, + current: [u8; INSTANCE_ID_LEN], +) -> Result<[u8; INSTANCE_ID_LEN], AuthFailure> { + let next_path = state_dir.join("server-instance-id.next"); + let Some(next) = read_instance_id_file(&next_path)? else { + return Ok(current); + }; + let snapshot_path = auth_snapshot_path(state_dir); + if !snapshot_path.exists() { + let _ = std::fs::remove_file(&next_path); + return Ok(current); + } + let bytes = std::fs::read(&snapshot_path).map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to read `{}`: {error}", snapshot_path.display()), + false, + ) + })?; + let Ok(plain) = open_blob(admin_key, &bytes) else { + return Ok(current); + }; + let Ok(snapshot) = serde_json::from_slice::(&plain) else { + return Ok(current); + }; + if snapshot.instance_id == current { + let _ = std::fs::remove_file(&next_path); + return Ok(current); + } + if snapshot.instance_id != next { + return Ok(current); + } + atomic_write(&state_dir.join("server-instance-id"), &next, 0o600)?; + let _ = std::fs::remove_file(&next_path); + Ok(next) +} + pub(super) fn random_instance_id() -> [u8; INSTANCE_ID_LEN] { let mut instance_id = [0_u8; INSTANCE_ID_LEN]; let mut rng = rand::rng(); diff --git a/src/common/auth/runtime.rs b/src/common/auth/runtime.rs index bd5a0e0..ef8fde3 100644 --- a/src/common/auth/runtime.rs +++ b/src/common/auth/runtime.rs @@ -17,7 +17,7 @@ use super::*; impl AuthRuntime { pub async fn from_process(config: AuthConfig) -> Result { - prepare_state_dir(&config.state_dir)?; + let state_lock = prepare_state_dir_and_lock(&config.state_dir)?; let credential = load_server_admin_credential(&config.state_dir)?; let Credential::Admin(admin_key) = credential else { return Err(AuthFailure::new( @@ -26,7 +26,7 @@ impl AuthRuntime { false, )); }; - Self::start_with_process_sync(admin_key, config, true).await + Self::start_locked(admin_key, config, true, state_lock).await } /// Start an embedded relay with an administrator key owned only by its state directory. @@ -34,7 +34,7 @@ impl AuthRuntime { /// This deliberately leaves the process credential untouched because the containing UI uses /// that credential for its outbound register, connect, status, and stream connections. pub async fn from_isolated_state(config: AuthConfig) -> Result { - prepare_state_dir(&config.state_dir)?; + let state_lock = prepare_state_dir_and_lock(&config.state_dir)?; let credential = load_isolated_server_admin_credential(&config.state_dir)?; let Credential::Admin(admin_key) = credential else { return Err(AuthFailure::new( @@ -43,21 +43,23 @@ impl AuthRuntime { false, )); }; - Self::start_with_process_sync(admin_key, config, false).await + Self::start_locked(admin_key, config, false, state_lock).await } pub async fn start(admin_key: AesKeyType, config: AuthConfig) -> Result { - Self::start_with_process_sync(admin_key, config, true).await + let state_lock = prepare_state_dir_and_lock(&config.state_dir)?; + Self::start_locked(admin_key, config, true, state_lock).await } - async fn start_with_process_sync( + async fn start_locked( admin_key: AesKeyType, config: AuthConfig, sync_process_credential: bool, + state_lock: Arc, ) -> Result { - prepare_state_dir(&config.state_dir)?; - let state_lock = Arc::new(acquire_state_dir_lock(&config.state_dir)?); let instance_id = load_or_create_instance_id(&config.state_dir)?; + let instance_id = + recover_instance_id_after_reset(&config.state_dir, &admin_key, instance_id)?; let (mut loaded, safe_mode) = load_persisted_state(&config, &admin_key, instance_id); let now = unix_seconds(); if let Some(state) = loaded.as_mut() { diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index 2339c9b..a4656db 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -168,6 +168,31 @@ async fn overlapping_runtimes_cannot_share_an_auth_state_directory() { let _ = std::fs::remove_dir_all(state_dir); } +#[tokio::test] +async fn from_isolated_state_takes_the_state_lock_before_creating_admin_key() { + let state_dir = temp_state_dir("lock-before-key"); + prepare_state_dir(&state_dir).unwrap(); + let _lock = acquire_state_dir_lock(&state_dir).unwrap(); + let config = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 4, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + let error = match AuthRuntime::from_isolated_state(config).await { + Ok(_) => panic!("a locked start should not create a second runtime"), + Err(error) => error, + }; + assert_eq!(error.code, "auth_state_locked"); + assert!( + !state_dir.join("admin.key").exists(), + "a locked start must not create a competing administrator key" + ); + drop(_lock); + tokio::time::sleep(Duration::from_millis(20)).await; + let _ = std::fs::remove_dir_all(state_dir); +} + #[test] fn safe_mode_startup_does_not_allow_compaction() { assert!(!compaction_is_allowed(true)); @@ -264,10 +289,53 @@ fn platform_default_auth_state_dir_is_writable_outside_linux_system_paths() { } #[cfg(not(any(windows, target_os = "macos")))] { - assert_eq!(dir, PathBuf::from(DEFAULT_AUTH_STATE_DIR)); + let expected = linux_default_auth_state_dir( + unix_effective_uid(), + linux_system_auth_dir_usable(), + std::env::var_os("XDG_DATA_HOME").as_deref(), + std::env::var_os("HOME").as_deref(), + ); + assert_eq!(dir, expected); + if unix_effective_uid() != 0 && !linux_system_auth_dir_usable() { + assert_ne!( + dir, + PathBuf::from(DEFAULT_AUTH_STATE_DIR), + "unprivileged Linux should not default to the system auth directory: {}", + dir.display() + ); + } } } +#[test] +fn linux_default_auth_state_dir_prefers_user_data_when_system_dir_is_unusable() { + assert_eq!( + linux_default_auth_state_dir(0, false, None, Some(std::ffi::OsStr::new("/home/op"))), + PathBuf::from(DEFAULT_AUTH_STATE_DIR) + ); + assert_eq!( + linux_default_auth_state_dir(1000, true, None, Some(std::ffi::OsStr::new("/home/op"))), + PathBuf::from(DEFAULT_AUTH_STATE_DIR) + ); + assert_eq!( + linux_default_auth_state_dir( + 1000, + false, + Some(std::ffi::OsStr::new("/xdg")), + Some(std::ffi::OsStr::new("/home/op")) + ), + PathBuf::from("/xdg/pb-mapper/auth") + ); + assert_eq!( + linux_default_auth_state_dir(1000, false, None, Some(std::ffi::OsStr::new("/home/op"))), + PathBuf::from("/home/op/.local/share/pb-mapper/auth") + ); + assert_eq!( + linux_default_auth_state_dir(1000, false, None, None), + PathBuf::from(DEFAULT_AUTH_STATE_DIR) + ); +} + #[test] fn legacy_protocol_policy_trims_valid_values_and_rejects_unknown_values() { assert_eq!( @@ -467,6 +535,129 @@ async fn reset_rotates_instance_and_prevents_old_key_id_reuse() { let _ = std::fs::remove_dir_all(state_dir); } +#[test] +fn recover_instance_id_promotes_next_when_snapshot_matches() { + let state_dir = temp_state_dir("instance-next-promote"); + prepare_state_dir(&state_dir).unwrap(); + let admin_key = *b"0123456789abcdefghijklmnopqrstuv"; + let current = [1_u8; INSTANCE_ID_LEN]; + let next = [2_u8; INSTANCE_ID_LEN]; + atomic_write(&state_dir.join("server-instance-id"), ¤t, 0o600).unwrap(); + atomic_write(&state_dir.join("server-instance-id.next"), &next, 0o600).unwrap(); + let config = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 1, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + let snapshot = PersistedSnapshot { + schema_version: SNAPSHOT_SCHEMA_VERSION, + instance_id: next, + generations: vec![0; 1], + entries: Vec::new(), + legacy_protocol: LegacyProtocolPolicy::Allow, + admin_replays: Vec::new(), + audit_records: VecDeque::new(), + }; + write_snapshot_and_truncate_wal(&config, &admin_key, &snapshot).unwrap(); + + let recovered = recover_instance_id_after_reset(&state_dir, &admin_key, current).unwrap(); + assert_eq!(recovered, next); + assert_eq!( + read_instance_id_file(&state_dir.join("server-instance-id")).unwrap(), + Some(next) + ); + assert!(!state_dir.join("server-instance-id.next").exists()); + let _ = std::fs::remove_dir_all(state_dir); +} + +#[test] +fn recover_instance_id_discards_stale_next_when_snapshot_still_matches_current() { + let state_dir = temp_state_dir("instance-next-stale"); + prepare_state_dir(&state_dir).unwrap(); + let admin_key = *b"0123456789abcdefghijklmnopqrstuv"; + let current = [3_u8; INSTANCE_ID_LEN]; + let next = [4_u8; INSTANCE_ID_LEN]; + atomic_write(&state_dir.join("server-instance-id"), ¤t, 0o600).unwrap(); + atomic_write(&state_dir.join("server-instance-id.next"), &next, 0o600).unwrap(); + let config = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 1, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + let snapshot = PersistedSnapshot { + schema_version: SNAPSHOT_SCHEMA_VERSION, + instance_id: current, + generations: vec![0; 1], + entries: Vec::new(), + legacy_protocol: LegacyProtocolPolicy::Allow, + admin_replays: Vec::new(), + audit_records: VecDeque::new(), + }; + write_snapshot_and_truncate_wal(&config, &admin_key, &snapshot).unwrap(); + + let recovered = recover_instance_id_after_reset(&state_dir, &admin_key, current).unwrap(); + assert_eq!(recovered, current); + assert!(!state_dir.join("server-instance-id.next").exists()); + let _ = std::fs::remove_dir_all(state_dir); +} + +#[tokio::test] +async fn interrupted_reset_recovers_the_staged_instance_id_on_restart() { + let state_dir = temp_state_dir("reset-recover"); + let admin_key = *b"0123456789abcdefghijklmnopqrstuv"; + let config = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 4, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + let runtime = AuthRuntime::start(admin_key, config.clone()).await.unwrap(); + let admin = authenticate_for_test(&runtime, 0).unwrap(); + let issued = runtime + .issue( + &admin, + Duration::from_secs(60), + Some("before-interrupted-reset".to_string()), + ) + .await + .unwrap(); + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + + let old_instance_id = load_or_create_instance_id(&state_dir).unwrap(); + let next = random_instance_id(); + atomic_write(&state_dir.join("server-instance-id.next"), &next, 0o600).unwrap(); + let snapshot = PersistedSnapshot { + schema_version: SNAPSHOT_SCHEMA_VERSION, + instance_id: next, + generations: vec![0; 4], + entries: Vec::new(), + legacy_protocol: LegacyProtocolPolicy::Allow, + admin_replays: Vec::new(), + audit_records: VecDeque::new(), + }; + write_snapshot_and_truncate_wal(&config, &admin_key, &snapshot).unwrap(); + atomic_write( + &state_dir.join("server-instance-id"), + &old_instance_id, + 0o600, + ) + .unwrap(); + + let restored = AuthRuntime::start(admin_key, config).await.unwrap(); + let restored_admin = authenticate_for_test(&restored, 0).unwrap(); + let status = restored.status(&restored_admin).await.unwrap(); + assert!(!status.safe_mode); + assert_eq!(status.server_instance_id, hex(&next)); + assert!(authenticate_for_test(&restored, issued.metadata.key_id).is_err()); + assert!(!state_dir.join("server-instance-id.next").exists()); + drop(restored); + tokio::time::sleep(Duration::from_millis(20)).await; + let _ = std::fs::remove_dir_all(state_dir); +} + #[tokio::test] async fn corrupt_wal_fails_temporary_keys_closed_until_admin_reset() { let state_dir = temp_state_dir("auth-safe-mode"); From f85bbb01c8bf7c7277c3d1989747cce9cb8ca985 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 04:15:40 +0800 Subject: [PATCH 29/74] Discard leftover WAL when recovering an interrupted reset A crash after the reset snapshot lands can leave the previous instance's WAL beside the new snapshot. Promoting server-instance-id.next now truncates that WAL first so startup cannot replay old mutations onto the new derivation identity. --- CHANGELOG.md | 2 +- src/common/auth/persistence.rs | 9 ++++++++- src/common/auth/tests.rs | 3 +++ 3 files changed, 12 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3ac17d1..1a8cf66 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,7 +19,7 @@ All notable changes to this project will be documented in this file. - Compacted the durable first-flight replay log while the relay is running, rolled back torn replay-log appends, rewrote that log atomically, sized first-flight admission from `PB_MAPPER_NEW_STREAMS_PER_SECOND`, and took an exclusive lock on the authentication state directory. - Fsynced the replay-log directory on first creation, took the state lock before `--init-admin-key`, aborted accepted connection tasks on relay shutdown, and staged `admin.key.next` so an interrupted root rotation can recover a matching key and snapshot. - Stopped returning the embedded relay administrator key from routine config fetches; revealing it now requires an explicit FFI/UI action. -- Took `auth.lock` before loading or creating `admin.key`, staged `server-instance-id.next` so an interrupted reset can recover a matching snapshot, and used a user-writable Linux auth directory when `/var/lib/pb-mapper/auth` is not usable. +- Took `auth.lock` before loading or creating `admin.key`, staged `server-instance-id.next` so an interrupted reset can recover a matching snapshot, discarded leftover WAL from the previous instance during that recovery, and used a user-writable Linux auth directory when `/var/lib/pb-mapper/auth` is not usable. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/src/common/auth/persistence.rs b/src/common/auth/persistence.rs index efcdebd..8fac356 100644 --- a/src/common/auth/persistence.rs +++ b/src/common/auth/persistence.rs @@ -693,7 +693,11 @@ pub(super) fn write_snapshot_and_truncate_wal( let sealed = seal_blob(admin_key, &plain)?; let snapshot_path = auth_snapshot_path(&config.state_dir); atomic_write(&snapshot_path, &sealed, 0o600)?; - let wal_path = auth_wal_path(&config.state_dir); + truncate_auth_wal(&config.state_dir) +} + +pub(super) fn truncate_auth_wal(state_dir: &Path) -> Result<(), AuthFailure> { + let wal_path = auth_wal_path(state_dir); let created = !wal_path.exists(); let wal = OpenOptions::new() .create(true) @@ -883,6 +887,9 @@ pub(super) fn recover_instance_id_after_reset( if snapshot.instance_id != next { return Ok(current); } + // The reset snapshot is complete. Any leftover WAL still belongs to the + // previous instance and must not be replayed onto the new derivation id. + truncate_auth_wal(state_dir)?; atomic_write(&state_dir.join("server-instance-id"), &next, 0o600)?; let _ = std::fs::remove_file(&next_path); Ok(next) diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index a4656db..3c1dc91 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -560,6 +560,7 @@ fn recover_instance_id_promotes_next_when_snapshot_matches() { audit_records: VecDeque::new(), }; write_snapshot_and_truncate_wal(&config, &admin_key, &snapshot).unwrap(); + std::fs::write(state_dir.join("auth.wal"), b"old-instance-wal").unwrap(); let recovered = recover_instance_id_after_reset(&state_dir, &admin_key, current).unwrap(); assert_eq!(recovered, next); @@ -568,6 +569,7 @@ fn recover_instance_id_promotes_next_when_snapshot_matches() { Some(next) ); assert!(!state_dir.join("server-instance-id.next").exists()); + assert_eq!(std::fs::read(state_dir.join("auth.wal")).unwrap(), b""); let _ = std::fs::remove_dir_all(state_dir); } @@ -639,6 +641,7 @@ async fn interrupted_reset_recovers_the_staged_instance_id_on_restart() { audit_records: VecDeque::new(), }; write_snapshot_and_truncate_wal(&config, &admin_key, &snapshot).unwrap(); + std::fs::write(state_dir.join("auth.wal"), b"old-instance-wal").unwrap(); atomic_write( &state_dir.join("server-instance-id"), &old_instance_id, From 47b84ea9d4dc0fa420c435b9f2997a9a0e0c37ba Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 04:18:32 +0800 Subject: [PATCH 30/74] Flush the parent directory after Windows auth-state replacements Open the parent with FILE_FLAG_BACKUP_SEMANTICS and FlushFileBuffers so a crash after renaming auth.snapshot or admin.key cannot drop the directory entry. Unix still fsyncs the directory after rename. --- CHANGELOG.md | 1 + src/common/auth/persistence.rs | 38 ++++++++++++++++++++++------------ src/common/auth/tests.rs | 10 +++++++++ 3 files changed, 36 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1a8cf66..db14732 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,7 @@ All notable changes to this project will be documented in this file. - Fsynced the replay-log directory on first creation, took the state lock before `--init-admin-key`, aborted accepted connection tasks on relay shutdown, and staged `admin.key.next` so an interrupted root rotation can recover a matching key and snapshot. - Stopped returning the embedded relay administrator key from routine config fetches; revealing it now requires an explicit FFI/UI action. - Took `auth.lock` before loading or creating `admin.key`, staged `server-instance-id.next` so an interrupted reset can recover a matching snapshot, discarded leftover WAL from the previous instance during that recovery, and used a user-writable Linux auth directory when `/var/lib/pb-mapper/auth` is not usable. +- Flushed the parent directory after Windows auth-state replacements, matching the Unix `fsync` after rename. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/src/common/auth/persistence.rs b/src/common/auth/persistence.rs index 8fac356..b99e52f 100644 --- a/src/common/auth/persistence.rs +++ b/src/common/auth/persistence.rs @@ -181,20 +181,32 @@ pub(crate) fn replace_file(from: &Path, to: &Path) -> std::io::Result<()> { } pub(crate) fn sync_parent_directory(path: &Path) -> Result<(), AuthFailure> { - #[cfg(unix)] - if let Some(parent) = path.parent() { - File::open(parent) - .and_then(|directory| directory.sync_all()) - .map_err(|error| { - AuthFailure::new( - "temporary_key_store_unavailable", - format!("failed to sync `{}`: {error}", parent.display()), - false, - ) - })?; + let Some(parent) = path.parent() else { + return Ok(()); + }; + open_directory_for_sync(parent) + .and_then(|directory| directory.sync_all()) + .map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to sync `{}`: {error}", parent.display()), + false, + ) + }) +} + +fn open_directory_for_sync(path: &Path) -> std::io::Result { + #[cfg(windows)] + { + use std::os::windows::fs::OpenOptionsExt; + const FILE_FLAG_BACKUP_SEMANTICS: u32 = 0x0200_0000; + OpenOptions::new() + .read(true) + .custom_flags(FILE_FLAG_BACKUP_SEMANTICS) + .open(path) } - let _ = path; - Ok(()) + #[cfg(not(windows))] + File::open(path) } pub(super) fn push_audit_record(inner: &AuthStateInner, record: AuditRecord) { diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index 3c1dc91..b297ae4 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -307,6 +307,16 @@ fn platform_default_auth_state_dir_is_writable_outside_linux_system_paths() { } } +#[test] +fn sync_parent_directory_succeeds_for_a_local_file() { + let state_dir = temp_state_dir("dirsync"); + prepare_state_dir(&state_dir).unwrap(); + let path = state_dir.join("probe"); + std::fs::write(&path, b"x").unwrap(); + sync_parent_directory(&path).unwrap(); + let _ = std::fs::remove_dir_all(state_dir); +} + #[test] fn linux_default_auth_state_dir_prefers_user_data_when_system_dir_is_unusable() { assert_eq!( From 0042cf4c47abdc0196230fd1dcd6b69601e6ffb2 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 04:20:16 +0800 Subject: [PATCH 31/74] Discard leftover WAL when recovering an interrupted root rotation Promoting admin.key.next now truncates WAL first. Those records are still encrypted with the previous root key, so replaying them after recovery would fail closed or mix instances. --- CHANGELOG.md | 1 + src/common/auth.rs | 3 +++ src/common/auth/tests.rs | 35 +++++++++++++++++++++++++++++++++++ 3 files changed, 39 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index db14732..ca059f2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,6 +21,7 @@ All notable changes to this project will be documented in this file. - Stopped returning the embedded relay administrator key from routine config fetches; revealing it now requires an explicit FFI/UI action. - Took `auth.lock` before loading or creating `admin.key`, staged `server-instance-id.next` so an interrupted reset can recover a matching snapshot, discarded leftover WAL from the previous instance during that recovery, and used a user-writable Linux auth directory when `/var/lib/pb-mapper/auth` is not usable. - Flushed the parent directory after Windows auth-state replacements, matching the Unix `fsync` after rename. +- Discarded leftover WAL encrypted under the previous administrator key when promoting `admin.key.next`. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/src/common/auth.rs b/src/common/auth.rs index 402e37b..26c1555 100644 --- a/src/common/auth.rs +++ b/src/common/auth.rs @@ -706,6 +706,9 @@ fn recover_admin_key_after_rotation( if open_blob(&next_key, &bytes).is_err() { return Ok(current.to_string()); } + // The rotation snapshot is complete under the staged key. Leftover WAL + // records are still encrypted with the previous key. + truncate_auth_wal(state_dir)?; write_admin_key(state_dir, next.trim())?; let _ = std::fs::remove_file(state_dir.join("admin.key.next")); Ok(next) diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index b297ae4..37487c8 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -615,6 +615,41 @@ fn recover_instance_id_discards_stale_next_when_snapshot_still_matches_current() let _ = std::fs::remove_dir_all(state_dir); } +#[test] +fn recover_admin_key_discards_leftover_wal_from_the_old_key() { + let state_dir = temp_state_dir("admin-next-wal"); + prepare_state_dir(&state_dir).unwrap(); + let old_key = *b"0123456789abcdefghijklmnopqrstuv"; + let new_key = *b"abcdefghijklmnopqrstuvwxyz012345"; + let old_key_str = std::str::from_utf8(&old_key).unwrap(); + let new_key_str = std::str::from_utf8(&new_key).unwrap(); + write_admin_key(&state_dir, old_key_str).unwrap(); + write_admin_key_file(&state_dir.join("admin.key.next"), new_key_str, true).unwrap(); + let config = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 1, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + let snapshot = PersistedSnapshot { + schema_version: SNAPSHOT_SCHEMA_VERSION, + instance_id: [9_u8; INSTANCE_ID_LEN], + generations: vec![0; 1], + entries: Vec::new(), + legacy_protocol: LegacyProtocolPolicy::Allow, + admin_replays: Vec::new(), + audit_records: VecDeque::new(), + }; + write_snapshot_and_truncate_wal(&config, &new_key, &snapshot).unwrap(); + std::fs::write(state_dir.join("auth.wal"), b"old-key-wal").unwrap(); + + let recovered = recover_admin_key_after_rotation(&state_dir, old_key_str).unwrap(); + assert_eq!(recovered.trim(), new_key_str); + assert_eq!(std::fs::read(state_dir.join("auth.wal")).unwrap(), b""); + assert!(!state_dir.join("admin.key.next").exists()); + let _ = std::fs::remove_dir_all(state_dir); +} + #[tokio::test] async fn interrupted_reset_recovers_the_staged_instance_id_on_restart() { let state_dir = temp_state_dir("reset-recover"); From d54a0535f108abd3ac267f381fc0a8878b7a3515 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 04:29:49 +0800 Subject: [PATCH 32/74] Wait for aborted connections and flush Windows dirs with write access Shutdown now awaits aborted connection tasks so auth.lock is released before return. Windows directory sync opens the parent with GENERIC_WRITE because FlushFileBuffers rejects a read-only handle. --- CHANGELOG.md | 1 + src/common/auth/persistence.rs | 4 +- src/pb_server/runtime.rs | 90 +++++++++++++++++++++++++++++----- 3 files changed, 83 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ca059f2..20fd71e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,7 @@ All notable changes to this project will be documented in this file. - Took `auth.lock` before loading or creating `admin.key`, staged `server-instance-id.next` so an interrupted reset can recover a matching snapshot, discarded leftover WAL from the previous instance during that recovery, and used a user-writable Linux auth directory when `/var/lib/pb-mapper/auth` is not usable. - Flushed the parent directory after Windows auth-state replacements, matching the Unix `fsync` after rename. - Discarded leftover WAL encrypted under the previous administrator key when promoting `admin.key.next`. +- Opened Windows parent directories with write access before `FlushFileBuffers`, and waited for aborted connection tasks to finish so `auth.lock` is released before shutdown returns. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/src/common/auth/persistence.rs b/src/common/auth/persistence.rs index b99e52f..b1c0ea2 100644 --- a/src/common/auth/persistence.rs +++ b/src/common/auth/persistence.rs @@ -199,9 +199,11 @@ fn open_directory_for_sync(path: &Path) -> std::io::Result { #[cfg(windows)] { use std::os::windows::fs::OpenOptionsExt; + const GENERIC_READ: u32 = 0x8000_0000; + const GENERIC_WRITE: u32 = 0x4000_0000; const FILE_FLAG_BACKUP_SEMANTICS: u32 = 0x0200_0000; OpenOptions::new() - .read(true) + .access_mode(GENERIC_READ | GENERIC_WRITE) .custom_flags(FILE_FLAG_BACKUP_SEMANTICS) .open(path) } diff --git a/src/pb_server/runtime.rs b/src/pb_server/runtime.rs index ec29ce5..7086dfa 100644 --- a/src/pb_server/runtime.rs +++ b/src/pb_server/runtime.rs @@ -925,23 +925,91 @@ pub async fn run_server_on_listener( } ManagerTask::Shutdown => { tracing::info!("Server shutdown requested, stopping main loop"); - for handle in connection_tasks.drain(..) { - handle.abort(); - } break; } } } - // Gracefully shutdown the listener - listener_handle.abort(); - shutdown_handle.abort(); - if let Some(handle) = status_forward_handle { + // Abort first, then wait. Dropping a JoinHandle after abort() does not + // wait for the task to drop its AuthRuntime clone, so a UI restart can + // still see auth.lock held. + abort_and_wait( + connection_tasks + .into_iter() + .chain(std::iter::once(listener_handle)) + .chain(std::iter::once(shutdown_handle)) + .chain(status_forward_handle), + ) + .await; + tracing::info!("Server shutdown completed"); + Ok(()) +} + +async fn abort_and_wait(handles: impl IntoIterator>) { + let handles: Vec<_> = handles.into_iter().collect(); + for handle in &handles { handle.abort(); } - for handle in connection_tasks { - handle.abort(); + for handle in handles { + let _ = handle.await; + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::common::auth::{ + AuthConfig, AuthRuntime, LegacyProtocolPolicy, PROCESS_CREDENTIAL_TEST_LOCK, + }; + use rand::RngExt; + use std::path::PathBuf; + use std::time::Duration; + + fn temp_state_dir(name: &str) -> PathBuf { + let mut suffix = [0_u8; 8]; + let mut rng = rand::rng(); + for byte in &mut suffix { + *byte = rng.random(); + } + std::env::temp_dir().join(format!( + "pb-mapper-{name}-{}", + suffix + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::() + )) + } + + #[tokio::test] + async fn shutdown_releases_auth_lock_while_a_connection_is_open() { + let _process_credential_guard = PROCESS_CREDENTIAL_TEST_LOCK.lock().await; + let state_dir = temp_state_dir("shutdown-lock"); + let admin_key = *b"0123456789abcdefghijklmnopqrstuv"; + let config = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 4, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + let auth = AuthRuntime::start(admin_key, config.clone()).await.unwrap(); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let shutdown_token = CancellationToken::new(); + let server = tokio::spawn({ + let shutdown_token = shutdown_token.clone(); + async move { run_server_on_listener(listener, shutdown_token, None, false, auth).await } + }); + let _client = TcpStream::connect(addr).await.unwrap(); + tokio::time::sleep(Duration::from_millis(80)).await; + shutdown_token.cancel(); + tokio::time::timeout(Duration::from_secs(2), server) + .await + .expect("server shutdown should finish after aborting connections") + .unwrap() + .unwrap(); + let restarted = AuthRuntime::start(admin_key, config).await.unwrap(); + drop(restarted); + tokio::time::sleep(Duration::from_millis(20)).await; + let _ = std::fs::remove_dir_all(state_dir); } - tracing::info!("Server shutdown completed"); - Ok(()) } From 90c2b35e3cdcd476e0ce270f2abcf7e6957f7d88 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 04:37:24 +0800 Subject: [PATCH 33/74] Return temporary_key_rotated after root rotation or reset A previously issued temporary credential that no longer matches current key material now fails with a distinct code instead of looking like a typo. Live-key mismatches stay temporary_key_invalid. Docs and CLI output state that rotation is a global invalidate. --- CHANGELOG.md | 1 + docs/authentication-v2.md | 15 ++++++----- docs/authentication-v2.zh-CN.md | 8 +++--- src/bin/pb-mapper/admin.rs | 1 + src/common/auth/runtime.rs | 44 +++++++++++++++++++++++++++++---- src/common/auth/tests.rs | 34 ++++++++++++++++++++++++- 6 files changed, 88 insertions(+), 15 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 20fd71e..c781db5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,7 @@ All notable changes to this project will be documented in this file. - Flushed the parent directory after Windows auth-state replacements, matching the Unix `fsync` after rename. - Discarded leftover WAL encrypted under the previous administrator key when promoting `admin.key.next`. - Opened Windows parent directories with write access before `FlushFileBuffers`, and waited for aborted connection tasks to finish so `auth.lock` is released before shutdown returns. +- Distinguished a post-rotation or post-reset temporary credential as `temporary_key_rotated` instead of the generic `temporary_key_invalid` used for a mistyped live key. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/docs/authentication-v2.md b/docs/authentication-v2.md index fcf0dbb..150a40b 100644 --- a/docs/authentication-v2.md +++ b/docs/authentication-v2.md @@ -162,12 +162,15 @@ pb-mapper admin --server relay.example.com:7666 key gc Root rotation writes an empty snapshot encrypted with the new key, preserves the bounded audit history, persists `admin.key`, and then switches the key and -administrator lease as one state transition. It invalidates all temporary -credentials and closes connections authenticated with the old administrator or -temporary keys. The CLI stages the candidate key before the request and verifies -the new key with an authenticated status call. When `--key-file` is omitted, -the recovery copy is written below `$XDG_CONFIG_HOME/pb-mapper` (or -`$HOME/.config/pb-mapper`) rather than requiring local `/var/lib` access. +administrator lease as one state transition. It is a global invalidate: every +temporary credential stops authenticating, including unexpired keys in other +namespaces, and live connections using those keys are cancelled. A later +first flight with one of those credentials returns `temporary_key_rotated`, +not the generic `temporary_key_invalid` used for a mistyped live key. The CLI +stages the candidate key before the request and verifies the new key with an +authenticated status call. When `--key-file` is omitted, the recovery copy is +written below `$XDG_CONFIG_HOME/pb-mapper` (or `$HOME/.config/pb-mapper`) +rather than requiring local `/var/lib` access. An explicit auth-state reset also invalidates all temporary credentials. It rotates the server instance ID so credentials from a corrupted or lost slot diff --git a/docs/authentication-v2.zh-CN.md b/docs/authentication-v2.zh-CN.md index eb7b7ff..c744664 100644 --- a/docs/authentication-v2.zh-CN.md +++ b/docs/authentication-v2.zh-CN.md @@ -121,9 +121,11 @@ tombstone 以给出稳定错误后，槽位可以复用。显式 `key gc` 可立 ### 根密钥轮换与状态重置 根密钥轮换先用新密钥写空 snapshot，同时保留有上限的审计历史，持久化 -`admin.key`，再把密钥与管理员 lease 作为一次状态变更切换。它会使全部临时凭据失效， -并关闭旧管理员或临时凭据建立的连接。CLI 在发请求前保存候选 key，完成后再用新 key -执行一次 `admin status` 验证。未指定 `--key-file` 时，恢复副本默认写到 +`admin.key`，再把密钥与管理员 lease 作为一次状态变更切换。这是全局作废：所有临时 +凭据都会立刻失效，包括其他命名空间里尚未过期的 key，并用这些 key 建立的活动连接 +会被取消。之后再用这些凭据做 first flight 会得到 `temporary_key_rotated`，而不是 +活 key 输错时的 `temporary_key_invalid`。CLI 在发请求前保存候选 key，完成后再用 +新 key 执行一次 `admin status` 验证。未指定 `--key-file` 时，恢复副本默认写到 `$XDG_CONFIG_HOME/pb-mapper`（或 `$HOME/.config/pb-mapper`），不要求本机能写 `/var/lib`。 diff --git a/src/bin/pb-mapper/admin.rs b/src/bin/pb-mapper/admin.rs index bf026eb..a58f661 100644 --- a/src/bin/pb-mapper/admin.rs +++ b/src/bin/pb-mapper/admin.rs @@ -289,6 +289,7 @@ pub(super) async fn run_admin(args: AdminArgs) -> Result<(), Box> { } if args.output == OutputFormat::Human { println!("administrator key rotated and verified"); + println!("all temporary credentials are now invalid (temporary_key_rotated)"); println!("key file: {}", key_file.display()); } else { print_admin_response(args.output, &response)?; diff --git a/src/common/auth/runtime.rs b/src/common/auth/runtime.rs index ef8fde3..9ac74e4 100644 --- a/src/common/auth/runtime.rs +++ b/src/common/auth/runtime.rs @@ -259,11 +259,7 @@ impl AuthRuntime { let expected_key = derive_temporary_key(&inner.admin_key(), &inner.instance_id(), key_id)?; if !bool::from(presented_key.ct_eq(&expected_key)) { inner.auth_failures.fetch_add(1, Ordering::Relaxed); - return Err(AuthFailure::new( - "temporary_key_invalid", - "temporary credential does not match the active relay key material", - false, - )); + return Err(temporary_key_material_mismatch(&inner, key_id)); } let index = key_slot(key_id) as usize; @@ -541,3 +537,41 @@ impl AuthRuntime { .await } } + +fn temporary_key_material_mismatch(inner: &AuthStateInner, key_id: u64) -> AuthFailure { + let index = key_slot(key_id) as usize; + let generation = key_generation(key_id); + let slots = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let generation_was_issued = match slots.get(index) { + Some(slot) => slot.generation == generation && slot.generation > 0, + None => { + let high = inner + .high_slot_generations + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + index + .checked_sub(slots.len()) + .and_then(|offset| high.get(offset).copied()) + == Some(generation) + && generation > 0 + } + }; + let slot_is_active = slots + .get(index) + .is_some_and(|slot| slot.state == SlotState::Active && slot.generation == generation); + if generation_was_issued && !slot_is_active { + return AuthFailure::new( + "temporary_key_rotated", + "temporary credential was invalidated by administrator root rotation or auth-state reset", + false, + ); + } + AuthFailure::new( + "temporary_key_invalid", + "temporary credential does not match the active relay key material", + false, + ) +} diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index 37487c8..c706151 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -522,13 +522,20 @@ async fn reset_rotates_instance_and_prevents_old_key_id_reuse() { .unwrap(); let old_context = authenticate_for_test(&runtime, old.metadata.key_id).unwrap(); let old_cancellation = old_context.cancellation_token().unwrap(); + let old_presented = runtime.derive_key(old.metadata.key_id).unwrap(); runtime.reset(&admin).await.unwrap(); let after = runtime.status(&admin).await.unwrap().server_instance_id; assert_ne!(after, before); assert!(old_cancellation.is_cancelled()); - assert!(authenticate_for_test(&runtime, old.metadata.key_id).is_err()); + assert_eq!( + runtime + .authenticate_presented(old.metadata.key_id, &old_presented) + .unwrap_err() + .code, + "temporary_key_rotated" + ); let replacement = runtime .issue( &admin, @@ -761,6 +768,24 @@ async fn root_rotation_rejects_old_key_and_in_flight_admin_context() { }; let runtime = AuthRuntime::start(old_key, config).await.unwrap(); let old_admin = runtime.authenticate_presented(0, &old_key).unwrap(); + let issued = runtime + .issue( + &old_admin, + Duration::from_secs(60), + Some("before-rotate".to_string()), + ) + .await + .unwrap(); + let old_temporary = runtime.derive_key(issued.metadata.key_id).unwrap(); + let mut mistyped_temporary = old_temporary; + mistyped_temporary[0] ^= 0x01; + assert_eq!( + runtime + .authenticate_presented(issued.metadata.key_id, &mistyped_temporary) + .unwrap_err() + .code, + "temporary_key_invalid" + ); let mistyped_key = *b"1123456789abcdefghijklmnopqrstuv"; assert_eq!( runtime @@ -774,6 +799,13 @@ async fn root_rotation_rejects_old_key_and_in_flight_admin_context() { .rotate_root(&old_admin, new_key) .await .expect("root rotation should succeed"); + assert_eq!( + runtime + .authenticate_presented(issued.metadata.key_id, &old_temporary) + .unwrap_err() + .code, + "temporary_key_rotated" + ); assert_eq!( runtime From 20268016b1df426accf2eb8acc96abb7a3adc13b Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 04:52:50 +0800 Subject: [PATCH 34/74] Fail closed on replay rollback and keep auth.lock in the actor A torn replay log now stays unavailable instead of accepting later appends. The actor holds auth.lock until it exits. temporary_key_rotated requires a root-epoch change. --use-machine-msg-header-key refuses an existing admin.key, and the installer persists MSG_HEADER_KEY. --- CHANGELOG.md | 1 + scripts/install-server-gitee.sh | 16 +++++++++++++++- scripts/install-server-github.sh | 16 +++++++++++++++- src/bin/pb-mapper.rs | 8 ++++++++ src/common/auth.rs | 5 +++++ src/common/auth/actor.rs | 4 ++++ src/common/auth/persistence.rs | 3 +++ src/common/auth/runtime.rs | 8 ++++++-- src/common/auth/tests.rs | 15 +++++++++++++++ src/common/message/secure/replay.rs | 18 +++++++++++++++--- 10 files changed, 87 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c781db5..8e90650 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,7 @@ All notable changes to this project will be documented in this file. - Discarded leftover WAL encrypted under the previous administrator key when promoting `admin.key.next`. - Opened Windows parent directories with write access before `FlushFileBuffers`, and waited for aborted connection tasks to finish so `auth.lock` is released before shutdown returns. - Distinguished a post-rotation or post-reset temporary credential as `temporary_key_rotated` instead of the generic `temporary_key_invalid` used for a mistyped live key. +- Failed closed when first-flight replay-log rollback cannot restore the previous length, kept `auth.lock` in the actor until it exits, reserved `temporary_key_rotated` for a real root-epoch change, refused `--use-machine-msg-header-key` when `admin.key` already exists, and persisted installer `MSG_HEADER_KEY` into `admin.key`. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/scripts/install-server-gitee.sh b/scripts/install-server-gitee.sh index d4cd398..c33c7e2 100755 --- a/scripts/install-server-gitee.sh +++ b/scripts/install-server-gitee.sh @@ -95,7 +95,21 @@ install -m 0755 "$BIN_PATH" "${INSTALL_DIR}/pb-mapper" # key in the environment or /etc/pb-mapper/server.env must win; otherwise the # runtime would prefer the newly copied admin.key and lock operators out. install -d -m 0700 "$AUTH_DIR" -if [ -z "$(configured_msg_header_key)" ] && [ ! -s "$ADMIN_KEY_PATH" ] && [ -s "$LEGACY_KEY_PATH" ]; then +if [ -n "${MSG_HEADER_KEY:-}" ] && [ ! -s "$ADMIN_KEY_PATH" ]; then + case "$MSG_HEADER_KEY" in + pbmt1_*) + echo "MSG_HEADER_KEY is a temporary credential; write a 32-character administrator key to $ADMIN_KEY_PATH" >&2 + exit 1 + ;; + esac + if [ "${#MSG_HEADER_KEY}" -ne 32 ]; then + echo "MSG_HEADER_KEY must be a 32-character administrator key" >&2 + exit 1 + fi + printf '%s\n' "$MSG_HEADER_KEY" > "$ADMIN_KEY_PATH" + chmod 0600 "$ADMIN_KEY_PATH" + echo "Persisted installer MSG_HEADER_KEY to $ADMIN_KEY_PATH" +elif [ -z "$(configured_msg_header_key)" ] && [ ! -s "$ADMIN_KEY_PATH" ] && [ -s "$LEGACY_KEY_PATH" ]; then install -m 0600 "$LEGACY_KEY_PATH" "$ADMIN_KEY_PATH" echo "Migrated the legacy machine-derived key into $ADMIN_KEY_PATH" fi diff --git a/scripts/install-server-github.sh b/scripts/install-server-github.sh index 038d491..07f5ef7 100755 --- a/scripts/install-server-github.sh +++ b/scripts/install-server-github.sh @@ -95,7 +95,21 @@ install -m 0755 "$BIN_PATH" "${INSTALL_DIR}/pb-mapper" # key in the environment or /etc/pb-mapper/server.env must win; otherwise the # runtime would prefer the newly copied admin.key and lock operators out. install -d -m 0700 "$AUTH_DIR" -if [ -z "$(configured_msg_header_key)" ] && [ ! -s "$ADMIN_KEY_PATH" ] && [ -s "$LEGACY_KEY_PATH" ]; then +if [ -n "${MSG_HEADER_KEY:-}" ] && [ ! -s "$ADMIN_KEY_PATH" ]; then + case "$MSG_HEADER_KEY" in + pbmt1_*) + echo "MSG_HEADER_KEY is a temporary credential; write a 32-character administrator key to $ADMIN_KEY_PATH" >&2 + exit 1 + ;; + esac + if [ "${#MSG_HEADER_KEY}" -ne 32 ]; then + echo "MSG_HEADER_KEY must be a 32-character administrator key" >&2 + exit 1 + fi + printf '%s\n' "$MSG_HEADER_KEY" > "$ADMIN_KEY_PATH" + chmod 0600 "$ADMIN_KEY_PATH" + echo "Persisted installer MSG_HEADER_KEY to $ADMIN_KEY_PATH" +elif [ -z "$(configured_msg_header_key)" ] && [ ! -s "$ADMIN_KEY_PATH" ] && [ -s "$LEGACY_KEY_PATH" ]; then install -m 0600 "$LEGACY_KEY_PATH" "$ADMIN_KEY_PATH" echo "Migrated the legacy machine-derived key into $ADMIN_KEY_PATH" fi diff --git a/src/bin/pb-mapper.rs b/src/bin/pb-mapper.rs index e1ec083..cd9cac6 100644 --- a/src/bin/pb-mapper.rs +++ b/src/bin/pb-mapper.rs @@ -281,6 +281,14 @@ async fn run_server(args: ServerArgs) -> Result<(), Box> { set_process_msg_header_key(Some(&key))?; eprintln!("administrator key initialized at {}", key_path.display()); } else if args.use_machine_msg_header_key { + let admin_key_path = auth_config.state_dir.join("admin.key"); + if admin_key_path.exists() { + return Err(format!( + "--use-machine-msg-header-key cannot replace `{}`; use `pb-mapper admin root-key rotate` to change the root key", + admin_key_path.display() + ) + .into()); + } tracing::warn!( "--use-machine-msg-header-key is a legacy compatibility option; prefer a random administrator key" ); diff --git a/src/common/auth.rs b/src/common/auth.rs index 26c1555..95f80f3 100644 --- a/src/common/auth.rs +++ b/src/common/auth.rs @@ -424,6 +424,7 @@ struct SlotHot { generation: u32, state: SlotState, expires_at: u64, + issued_epoch: u64, lease: Weak, } @@ -433,6 +434,7 @@ impl Default for SlotHot { generation: 0, state: SlotState::Free, expires_at: 0, + issued_epoch: 0, lease: Weak::new(), } } @@ -460,6 +462,7 @@ struct AuthStateInner { last_legacy_connection_at: AtomicU64, auth_successes: AtomicU64, auth_failures: AtomicU64, + root_epoch: AtomicU64, audit_records: RwLock>, } @@ -854,6 +857,8 @@ struct PersistedSnapshot { admin_replays: Vec, #[serde(default)] audit_records: VecDeque, + #[serde(default)] + root_epoch: u64, } #[derive(Clone, Debug, Serialize, Deserialize)] diff --git a/src/common/auth/actor.rs b/src/common/auth/actor.rs index fc57bdc..d65b46c 100644 --- a/src/common/auth/actor.rs +++ b/src/common/auth/actor.rs @@ -47,6 +47,7 @@ pub(super) async fn run_auth_actor( mut command_rx: mpsc::Receiver, config: AuthConfig, state: AuthActorState, + _state_lock: Arc, ) { let AuthActorState { mut cold, @@ -435,6 +436,7 @@ fn actor_issue( slot.generation = generation; slot.state = SlotState::Active; slot.expires_at = expires_at; + slot.issued_epoch = inner.root_epoch.load(Ordering::Acquire); slot.lease = Arc::downgrade(&lease); cold.insert( key_id, @@ -693,6 +695,7 @@ fn actor_reset( action: &str, ) -> Result<(), AuthFailure> { let new_instance_id = random_instance_id(); + inner.root_epoch.fetch_add(1, Ordering::AcqRel); let reset_audit = audit(action, None, None); let mut snapshot = empty_snapshot(inner, new_instance_id, admin_replays); push_persisted_audit(&mut snapshot.audit_records, reset_audit.clone()); @@ -763,6 +766,7 @@ fn actor_rotate_root( let new_key_string = String::from_utf8(new_key.to_vec()).expect("printable ASCII is valid UTF-8"); + inner.root_epoch.fetch_add(1, Ordering::AcqRel); let rotate_audit = audit("administrator_key_rotate", None, None); let mut snapshot = empty_snapshot(inner, inner.instance_id(), &VecDeque::new()); push_persisted_audit(&mut snapshot.audit_records, rotate_audit.clone()); diff --git a/src/common/auth/persistence.rs b/src/common/auth/persistence.rs index b1c0ea2..fa7f6df 100644 --- a/src/common/auth/persistence.rs +++ b/src/common/auth/persistence.rs @@ -313,6 +313,7 @@ pub(super) fn build_snapshot( .read() .unwrap_or_else(|poisoned| poisoned.into_inner()) .clone(), + root_epoch: inner.root_epoch.load(Ordering::Acquire), } } @@ -364,6 +365,7 @@ pub(super) fn empty_snapshot( .read() .unwrap_or_else(|poisoned| poisoned.into_inner()) .clone(), + root_epoch: inner.root_epoch.load(Ordering::Acquire), } } @@ -418,6 +420,7 @@ pub(super) fn try_load_persisted_state( legacy_protocol: config.legacy_protocol, admin_replays: Vec::new(), audit_records: VecDeque::new(), + root_epoch: 0, } }; if snapshot.schema_version != SNAPSHOT_SCHEMA_VERSION || snapshot.instance_id != instance_id { diff --git a/src/common/auth/runtime.rs b/src/common/auth/runtime.rs index 9ac74e4..7ba342b 100644 --- a/src/common/auth/runtime.rs +++ b/src/common/auth/runtime.rs @@ -172,6 +172,7 @@ impl AuthRuntime { last_legacy_connection_at: AtomicU64::new(0), auth_successes: AtomicU64::new(0), auth_failures: AtomicU64::new(0), + root_epoch: AtomicU64::new(loaded.as_ref().map(|state| state.root_epoch).unwrap_or(0)), audit_records: RwLock::new(audit_records), }); let (command_tx, command_rx) = mpsc::channel(256); @@ -179,7 +180,7 @@ impl AuthRuntime { inner: Arc::downgrade(&inner), command_tx, config: config.clone(), - _state_lock: state_lock, + _state_lock: state_lock.clone(), }; tokio::spawn(run_auth_actor( @@ -188,6 +189,7 @@ impl AuthRuntime { command_rx, config, AuthActorState::new(cold, wheel, admin_replays, admin_replay_order), + state_lock, )); Ok(runtime) } @@ -562,7 +564,9 @@ fn temporary_key_material_mismatch(inner: &AuthStateInner, key_id: u64) -> AuthF let slot_is_active = slots .get(index) .is_some_and(|slot| slot.state == SlotState::Active && slot.generation == generation); - if generation_was_issued && !slot_is_active { + let issued_epoch = slots.get(index).map(|slot| slot.issued_epoch).unwrap_or(0); + let current_epoch = inner.root_epoch.load(Ordering::Acquire); + if generation_was_issued && !slot_is_active && issued_epoch < current_epoch { return AuthFailure::new( "temporary_key_rotated", "temporary credential was invalidated by administrator root rotation or auth-state reset", diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index c706151..6b01150 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -462,10 +462,20 @@ async fn issue_renew_revoke_and_persist() { assert_eq!(renewed.metadata.key_id, issued.metadata.key_id); assert_eq!(renewed.credential, issued.credential); assert!(renewed.metadata.expires_at > issued.metadata.expires_at); + let presented = runtime.derive_key(issued.metadata.key_id).unwrap(); runtime .revoke(&admin, issued.metadata.key_id) .await .unwrap(); + let mut mistyped = presented; + mistyped[0] ^= 0x01; + assert_eq!( + runtime + .authenticate_presented(issued.metadata.key_id, &mistyped) + .unwrap_err() + .code, + "temporary_key_invalid" + ); assert!(cancellation.is_cancelled()); assert_eq!( context.ensure_active().unwrap_err().code, @@ -575,6 +585,7 @@ fn recover_instance_id_promotes_next_when_snapshot_matches() { legacy_protocol: LegacyProtocolPolicy::Allow, admin_replays: Vec::new(), audit_records: VecDeque::new(), + root_epoch: 0, }; write_snapshot_and_truncate_wal(&config, &admin_key, &snapshot).unwrap(); std::fs::write(state_dir.join("auth.wal"), b"old-instance-wal").unwrap(); @@ -613,6 +624,7 @@ fn recover_instance_id_discards_stale_next_when_snapshot_still_matches_current() legacy_protocol: LegacyProtocolPolicy::Allow, admin_replays: Vec::new(), audit_records: VecDeque::new(), + root_epoch: 0, }; write_snapshot_and_truncate_wal(&config, &admin_key, &snapshot).unwrap(); @@ -646,6 +658,7 @@ fn recover_admin_key_discards_leftover_wal_from_the_old_key() { legacy_protocol: LegacyProtocolPolicy::Allow, admin_replays: Vec::new(), audit_records: VecDeque::new(), + root_epoch: 0, }; write_snapshot_and_truncate_wal(&config, &new_key, &snapshot).unwrap(); std::fs::write(state_dir.join("auth.wal"), b"old-key-wal").unwrap(); @@ -691,6 +704,7 @@ async fn interrupted_reset_recovers_the_staged_instance_id_on_restart() { legacy_protocol: LegacyProtocolPolicy::Allow, admin_replays: Vec::new(), audit_records: VecDeque::new(), + root_epoch: 0, }; write_snapshot_and_truncate_wal(&config, &admin_key, &snapshot).unwrap(); std::fs::write(state_dir.join("auth.wal"), b"old-instance-wal").unwrap(); @@ -1021,6 +1035,7 @@ fn tombstone_migration_prefers_audit_time_and_persists_fail_closed_fallback() { key_id: Some(make_key_id(1, 0)), label: None, }]), + root_epoch: 0, }; assert!(normalize_tombstone_times(&mut snapshot, now)); diff --git a/src/common/message/secure/replay.rs b/src/common/message/secure/replay.rs index e46dbb4..cddbb08 100644 --- a/src/common/message/secure/replay.rs +++ b/src/common/message/secure/replay.rs @@ -109,6 +109,7 @@ pub(super) struct ReplayGuard { max_per_key: u32, log_path: Option, last_compact_at: u64, + log_failed: bool, } impl ReplayGuard { @@ -122,6 +123,7 @@ impl ReplayGuard { max_per_key: first_flight_budget(window_seconds), log_path, last_compact_at: now, + log_failed: false, }; guard.load_persisted(); guard @@ -165,7 +167,12 @@ impl ReplayGuard { self.counts_started_at = now; } - fn persist(&self, fingerprint: &[u8; 32], now: u64) -> std::io::Result<()> { + fn persist(&mut self, fingerprint: &[u8; 32], now: u64) -> std::io::Result<()> { + if self.log_failed { + return Err(std::io::Error::other( + "durable first-flight replay log is unavailable after a failed rollback", + )); + } let Some(path) = &self.log_path else { return Ok(()); }; @@ -179,8 +186,13 @@ impl ReplayGuard { record[..32].copy_from_slice(fingerprint); record[32..].copy_from_slice(&now.to_be_bytes()); if let Err(error) = file.write_all(&record).and_then(|()| file.sync_data()) { - let _ = file.set_len(start_len); - let _ = file.sync_data(); + if file + .set_len(start_len) + .and_then(|()| file.sync_data()) + .is_err() + { + self.log_failed = true; + } return Err(error); } if created { From a81fa13295659dc80f2d4d3941c406cf796eb5fe Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 04:57:47 +0800 Subject: [PATCH 35/74] Refuse replacing live admin.key while encrypted state exists write_admin_key_file now shares the initialize_admin_key guard for the live admin.key path so a forced write cannot leave snapshot/WAL encrypted under the previous root. Staging admin.key.next stays allowed. --- CHANGELOG.md | 1 + src/common/auth/persistence.rs | 41 +++++++++++++++++++++------------- src/common/auth/tests.rs | 9 ++++++++ 3 files changed, 35 insertions(+), 16 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8e90650..e6255a5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,6 +25,7 @@ All notable changes to this project will be documented in this file. - Opened Windows parent directories with write access before `FlushFileBuffers`, and waited for aborted connection tasks to finish so `auth.lock` is released before shutdown returns. - Distinguished a post-rotation or post-reset temporary credential as `temporary_key_rotated` instead of the generic `temporary_key_invalid` used for a mistyped live key. - Failed closed when first-flight replay-log rollback cannot restore the previous length, kept `auth.lock` in the actor until it exits, reserved `temporary_key_rotated` for a real root-epoch change, refused `--use-machine-msg-header-key` when `admin.key` already exists, and persisted installer `MSG_HEADER_KEY` into `admin.key`. +- Refused `write_admin_key_file` of the live `admin.key` while encrypted auth state exists, matching `initialize_admin_key`; `admin.key.next` remains allowed. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/src/common/auth/persistence.rs b/src/common/auth/persistence.rs index fa7f6df..b9a6ff1 100644 --- a/src/common/auth/persistence.rs +++ b/src/common/auth/persistence.rs @@ -945,22 +945,7 @@ pub fn initialize_admin_key(path: &Path, force: bool) -> Result Result<(), A false, )); } + if path.file_name() == Some(std::ffi::OsStr::new("admin.key")) { + refuse_write_if_encrypted_state(path, force)?; + } atomic_write(path, format!("{key}\n").as_bytes(), 0o600) } +fn refuse_write_if_encrypted_state(path: &Path, force: bool) -> Result<(), AuthFailure> { + // Creating or replacing the live root while snapshot/WAL remain leaves + // those files encrypted under the previous key. Staging `admin.key.next` + // is the rotate path and must stay allowed. + let Some(state_dir) = path.parent() else { + return Ok(()); + }; + if !encrypted_auth_state_exists(state_dir) { + return Ok(()); + } + Err(AuthFailure::new( + "administrator_key_state_exists", + format!( + "refusing to {} `{}` while encrypted auth state exists; use `pb-mapper admin root-key rotate` or `pb-mapper admin auth-state reset --confirm`", + if force { "replace" } else { "create" }, + path.display() + ), + false, + )) +} + pub(super) fn atomic_write(path: &Path, data: &[u8], mode: u32) -> Result<(), AuthFailure> { if let Some(parent) = path.parent() { std::fs::create_dir_all(parent).map_err(|error| { diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index 6b01150..bd3e83c 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -38,6 +38,15 @@ fn initialize_admin_key_refuses_to_replace_a_key_when_encrypted_state_exists() { let missing = state_dir.join("missing-admin.key"); let error = initialize_admin_key(&missing, false).unwrap_err(); assert_eq!(error.code, "administrator_key_state_exists"); + let error = + write_admin_key_file(&key_path, "abcdefghijklmnopqrstuvwxyz012345", true).unwrap_err(); + assert_eq!(error.code, "administrator_key_state_exists"); + write_admin_key_file( + &state_dir.join("admin.key.next"), + "abcdefghijklmnopqrstuvwxyz012345", + true, + ) + .unwrap(); let _ = std::fs::remove_dir_all(state_dir); } From a142d0f6242776a627f00cd45eb97f55bc9d781d Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 05:11:19 +0800 Subject: [PATCH 36/74] Keep rotated errors after slot reuse and abort timed-out UI shutdown A pre-rotation temporary credential still returns temporary_key_rotated after a later issue in the same slot. Embedded relay stop now aborts and awaits the task on timeout. Replay logs fail closed if directory sync fails, a post-rotation write of the live key is allowed when it decrypts the snapshot, and a zero stream-rate env falls back to 100. --- CHANGELOG.md | 1 + src/common/auth/persistence.rs | 21 +++++++++++++++++++- src/common/auth/runtime.rs | 19 ++++++++++++------ src/common/auth/tests.rs | 17 +++++++++++++++- src/common/message/secure/replay.rs | 9 ++++++--- ui/native/pb_mapper_ffi/src/state/runtime.rs | 7 +++++-- 6 files changed, 61 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e6255a5..7d4e3af 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,6 +26,7 @@ All notable changes to this project will be documented in this file. - Distinguished a post-rotation or post-reset temporary credential as `temporary_key_rotated` instead of the generic `temporary_key_invalid` used for a mistyped live key. - Failed closed when first-flight replay-log rollback cannot restore the previous length, kept `auth.lock` in the actor until it exits, reserved `temporary_key_rotated` for a real root-epoch change, refused `--use-machine-msg-header-key` when `admin.key` already exists, and persisted installer `MSG_HEADER_KEY` into `admin.key`. - Refused `write_admin_key_file` of the live `admin.key` while encrypted auth state exists, matching `initialize_admin_key`; `admin.key.next` remains allowed. +- Kept `temporary_key_rotated` after a later issue in the same slot, aborted a timed-out embedded-relay shutdown, fail-closed replay logs whose directory sync failed, allowed a post-rotation write of the live key that already decrypts the snapshot, and treated `PB_MAPPER_NEW_STREAMS_PER_SECOND=0` as the default 100. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/src/common/auth/persistence.rs b/src/common/auth/persistence.rs index b9a6ff1..76d0ea4 100644 --- a/src/common/auth/persistence.rs +++ b/src/common/auth/persistence.rs @@ -971,12 +971,31 @@ pub fn write_admin_key_file(path: &Path, key: &str, force: bool) -> Result<(), A false, )); } - if path.file_name() == Some(std::ffi::OsStr::new("admin.key")) { + if path.file_name() == Some(std::ffi::OsStr::new("admin.key")) + && !key_matches_existing_snapshot(path.parent(), key) + { refuse_write_if_encrypted_state(path, force)?; } atomic_write(path, format!("{key}\n").as_bytes(), 0o600) } +fn key_matches_existing_snapshot(state_dir: Option<&Path>, key: &str) -> bool { + let Some(state_dir) = state_dir else { + return false; + }; + let snapshot_path = auth_snapshot_path(state_dir); + if !snapshot_path.exists() { + return false; + } + let Ok(Credential::Admin(admin_key)) = parse_credential(key) else { + return false; + }; + let Ok(bytes) = std::fs::read(&snapshot_path) else { + return false; + }; + open_blob(&admin_key, &bytes).is_ok() +} + fn refuse_write_if_encrypted_state(path: &Path, force: bool) -> Result<(), AuthFailure> { // Creating or replacing the live root while snapshot/WAL remain leaves // those files encrypted under the previous key. Staging `admin.key.next` diff --git a/src/common/auth/runtime.rs b/src/common/auth/runtime.rs index 7ba342b..260abee 100644 --- a/src/common/auth/runtime.rs +++ b/src/common/auth/runtime.rs @@ -547,8 +547,8 @@ fn temporary_key_material_mismatch(inner: &AuthStateInner, key_id: u64) -> AuthF .slots .read() .unwrap_or_else(|poisoned| poisoned.into_inner()); - let generation_was_issued = match slots.get(index) { - Some(slot) => slot.generation == generation && slot.generation > 0, + let current_generation = match slots.get(index) { + Some(slot) => Some(slot.generation), None => { let high = inner .high_slot_generations @@ -557,16 +557,23 @@ fn temporary_key_material_mismatch(inner: &AuthStateInner, key_id: u64) -> AuthF index .checked_sub(slots.len()) .and_then(|offset| high.get(offset).copied()) - == Some(generation) - && generation > 0 } }; let slot_is_active = slots .get(index) .is_some_and(|slot| slot.state == SlotState::Active && slot.generation == generation); - let issued_epoch = slots.get(index).map(|slot| slot.issued_epoch).unwrap_or(0); + if slot_is_active { + return AuthFailure::new( + "temporary_key_invalid", + "temporary credential does not match the active relay key material", + false, + ); + } let current_epoch = inner.root_epoch.load(Ordering::Acquire); - if generation_was_issued && !slot_is_active && issued_epoch < current_epoch { + if current_epoch > 0 + && generation > 0 + && current_generation.is_some_and(|issued| generation <= issued) + { return AuthFailure::new( "temporary_key_rotated", "temporary credential was invalidated by administrator root rotation or auth-state reset", diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index bd3e83c..6b8e69d 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -829,6 +829,22 @@ async fn root_rotation_rejects_old_key_and_in_flight_admin_context() { .code, "temporary_key_rotated" ); + let new_admin = runtime.authenticate_presented(0, &new_key).unwrap(); + let _replacement = runtime + .issue( + &new_admin, + Duration::from_secs(60), + Some("after-rotate".to_string()), + ) + .await + .unwrap(); + assert_eq!( + runtime + .authenticate_presented(issued.metadata.key_id, &old_temporary) + .unwrap_err() + .code, + "temporary_key_rotated" + ); assert_eq!( runtime @@ -845,7 +861,6 @@ async fn root_rotation_rejects_old_key_and_in_flight_admin_context() { .code, "administrator_key_rotated" ); - let new_admin = runtime.authenticate_presented(0, &new_key).unwrap(); assert!(runtime.status(&new_admin).await.is_ok()); drop(runtime); diff --git a/src/common/message/secure/replay.rs b/src/common/message/secure/replay.rs index cddbb08..a4114ca 100644 --- a/src/common/message/secure/replay.rs +++ b/src/common/message/secure/replay.rs @@ -30,8 +30,9 @@ fn first_flight_budget(window_seconds: u64) -> u32 { let streams_per_sec = std::env::var("PB_MAPPER_NEW_STREAMS_PER_SECOND") .ok() .and_then(|value| value.parse().ok()) + .filter(|value| *value > 0) .unwrap_or(DEFAULT_NEW_STREAMS_PER_SECOND) - .clamp(1, 1_000_000); + .min(1_000_000); let window = u32::try_from(window_seconds).unwrap_or(u32::MAX); streams_per_sec .saturating_mul(2) @@ -196,8 +197,10 @@ impl ReplayGuard { return Err(error); } if created { - crate::common::auth::sync_parent_directory(path) - .map_err(|error| std::io::Error::other(error.to_string()))?; + if let Err(error) = crate::common::auth::sync_parent_directory(path) { + self.log_failed = true; + return Err(std::io::Error::other(error.to_string())); + } } Ok(()) } diff --git a/ui/native/pb_mapper_ffi/src/state/runtime.rs b/ui/native/pb_mapper_ffi/src/state/runtime.rs index b7d6a83..5decac0 100644 --- a/ui/native/pb_mapper_ffi/src/state/runtime.rs +++ b/ui/native/pb_mapper_ffi/src/state/runtime.rs @@ -94,12 +94,15 @@ impl PbMapperState { let shutdown_timeout = tokio::time::Duration::from_secs(5); - match tokio::time::timeout(shutdown_timeout, handle).await { + let mut handle = handle; + match tokio::time::timeout(shutdown_timeout, &mut handle).await { Ok(_) => { tracing::info!("Server shutdown gracefully"); } Err(_) => { - tracing::warn!("Server shutdown timed out, may not have closed gracefully"); + handle.abort(); + let _ = handle.await; + tracing::warn!("Server shutdown timed out; aborted the relay task"); } } From 82d9b0b832923d94a92ed06238216027cbdcff12 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 05:25:10 +0800 Subject: [PATCH 37/74] Harden replay recovery and stop accumulating wheel owners Replay compaction now fsyncs the parent directory after replace, and an existing unreadable connection.replay file fails closed. Timing-wheel buckets hold Weak leases with one current owner per key so renewals do not keep extra strong references for the previous TTL. --- CHANGELOG.md | 1 + src/common/auth/actor.rs | 1 + src/common/auth/timing_wheel.rs | 70 ++++++++++++++++++----------- src/common/message/secure/replay.rs | 13 +++++- 4 files changed, 58 insertions(+), 27 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7d4e3af..b498562 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,6 +27,7 @@ All notable changes to this project will be documented in this file. - Failed closed when first-flight replay-log rollback cannot restore the previous length, kept `auth.lock` in the actor until it exits, reserved `temporary_key_rotated` for a real root-epoch change, refused `--use-machine-msg-header-key` when `admin.key` already exists, and persisted installer `MSG_HEADER_KEY` into `admin.key`. - Refused `write_admin_key_file` of the live `admin.key` while encrypted auth state exists, matching `initialize_admin_key`; `admin.key.next` remains allowed. - Kept `temporary_key_rotated` after a later issue in the same slot, aborted a timed-out embedded-relay shutdown, fail-closed replay logs whose directory sync failed, allowed a post-rotation write of the live key that already decrypts the snapshot, and treated `PB_MAPPER_NEW_STREAMS_PER_SECOND=0` as the default 100. +- Fsynced the replay-log directory after compacting replacements, treated an unreadable existing replay log as unavailable, and made timing-wheel buckets hold `Weak` leases so renewals no longer accumulate day-long strong references. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/src/common/auth/actor.rs b/src/common/auth/actor.rs index d65b46c..5189b9c 100644 --- a/src/common/auth/actor.rs +++ b/src/common/auth/actor.rs @@ -573,6 +573,7 @@ fn actor_renew( slot.expires_at = expires_at; lease.expires_at.store(expires_at, Ordering::Release); lease.wheel_version.fetch_add(1, Ordering::AcqRel); + wheel.release(key_id); wheel.insert(lease); drop(slots); metadata_with_credential(inner, cold, key_id, true) diff --git a/src/common/auth/timing_wheel.rs b/src/common/auth/timing_wheel.rs index 32208a7..bb53554 100644 --- a/src/common/auth/timing_wheel.rs +++ b/src/common/auth/timing_wheel.rs @@ -9,21 +9,23 @@ //! reset/rotation -> cancel every wheel-owned lease -> clear all buckets //! ``` //! -//! The wheel owns strong `Arc` references. Request-facing structures retain -//! only `Weak` references, so expiry, revoke, reset, and root rotation have one clear -//! cancellation owner without keeping dead credentials alive indefinitely. +//! The wheel's current-owner map holds the strong `Arc` for each key. +//! Bucket entries are `Weak`, so a renew replaces the previous owner instead of +//! accumulating day-long stale strong references. Request-facing structures also +//! retain only `Weak` references. use super::*; const MAX_INCREMENTAL_ADVANCE_SECONDS: u64 = 256; struct WheelEntry { - lease: Arc, + lease: Weak, version: u64, } pub(super) struct TimingWheel { now: u64, + owners: HashMap>, immediate_due: Vec, level0: Vec>, level1: Vec>, @@ -35,6 +37,7 @@ impl TimingWheel { pub(super) fn new(now: u64) -> Self { Self { now, + owners: HashMap::new(), immediate_due: Vec::new(), level0: empty_buckets(256), level1: empty_buckets(64), @@ -48,10 +51,18 @@ impl TimingWheel { self.insert_with_version(lease, version); } + pub(super) fn release(&mut self, key_id: u64) { + self.owners.remove(&key_id); + } + pub(super) fn insert_with_version(&mut self, lease: Arc, version: u64) { + self.owners.insert(lease.key_id(), lease.clone()); let expires_at = lease.expires_at(); let delta = expires_at.saturating_sub(self.now); - let entry = WheelEntry { lease, version }; + let entry = WheelEntry { + lease: Arc::downgrade(&lease), + version, + }; if expires_at <= self.now { self.immediate_due.push(entry); } else if delta < 1 << 8 { @@ -85,12 +96,14 @@ impl TimingWheel { due.extend(self.take_immediate_due(self.now)); let index = (self.now & 0xff) as usize; for entry in std::mem::take(&mut self.level0[index]) { - if entry.version == entry.lease.wheel_version.load(Ordering::Acquire) { - if entry.lease.expires_at() <= self.now { - due.push(entry.lease); - } else { - self.insert(entry.lease); - } + let Some(lease) = live_lease(&entry) else { + continue; + }; + if lease.expires_at() <= self.now { + self.owners.remove(&lease.key_id()); + due.push(lease); + } else { + self.insert_with_version(lease, entry.version); } } } @@ -107,13 +120,14 @@ impl TimingWheel { let mut due = Vec::new(); for entry in entries { - if entry.version != entry.lease.wheel_version.load(Ordering::Acquire) { + let Some(lease) = live_lease(&entry) else { continue; - } - if entry.lease.expires_at() <= target { - due.push(entry.lease); + }; + if lease.expires_at() <= target { + self.owners.remove(&lease.key_id()); + due.push(lease); } else { - self.insert_with_version(entry.lease, entry.version); + self.insert_with_version(lease, entry.version); } } due @@ -122,13 +136,14 @@ impl TimingWheel { fn take_immediate_due(&mut self, target: u64) -> Vec> { let mut due = Vec::new(); for entry in std::mem::take(&mut self.immediate_due) { - if entry.version != entry.lease.wheel_version.load(Ordering::Acquire) { + let Some(lease) = live_lease(&entry) else { continue; - } - if entry.lease.expires_at() <= target { - due.push(entry.lease); + }; + if lease.expires_at() <= target { + self.owners.remove(&lease.key_id()); + due.push(lease); } else { - self.insert_with_version(entry.lease, entry.version); + self.insert_with_version(lease, entry.version); } } due @@ -151,8 +166,8 @@ impl TimingWheel { _ => Vec::new(), }; for entry in entries { - if entry.version == entry.lease.wheel_version.load(Ordering::Acquire) { - self.insert_with_version(entry.lease, entry.version); + if let Some(lease) = live_lease(&entry) { + self.insert_with_version(lease, entry.version); } } } @@ -163,13 +178,18 @@ impl TimingWheel { take_all_entries(&mut self.level1, &mut entries); take_all_entries(&mut self.level2, &mut entries); take_all_entries(&mut self.level3, &mut entries); - for entry in entries { - entry.lease.cancellation.cancel(); + for lease in self.owners.values() { + lease.cancellation.cancel(); } *self = Self::new(now); } } +fn live_lease(entry: &WheelEntry) -> Option> { + let lease = entry.lease.upgrade()?; + (entry.version == lease.wheel_version.load(Ordering::Acquire)).then_some(lease) +} + fn empty_buckets(count: usize) -> Vec> { std::iter::repeat_with(Vec::new).take(count).collect() } diff --git a/src/common/message/secure/replay.rs b/src/common/message/secure/replay.rs index a4114ca..60914d2 100644 --- a/src/common/message/secure/replay.rs +++ b/src/common/message/secure/replay.rs @@ -209,8 +209,15 @@ impl ReplayGuard { let Some(path) = self.log_path.clone() else { return; }; - let Ok(mut file) = File::open(&path) else { + if !path.exists() { return; + } + let mut file = match File::open(&path) { + Ok(file) => file, + Err(_) => { + self.log_failed = true; + return; + } }; let now = unix_seconds(); let mut live = Vec::new(); @@ -280,7 +287,9 @@ impl ReplayGuard { file.write_all(&live.concat())?; file.sync_all()?; drop(file); - crate::common::auth::replace_file(&temporary, path) + crate::common::auth::replace_file(&temporary, path)?; + crate::common::auth::sync_parent_directory(path) + .map_err(|error| std::io::Error::other(error.to_string())) })(); if result.is_err() { let _ = std::fs::remove_file(&temporary); From 45458e76b34946841ac26b67eafdf14ca6705193 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 05:35:00 +0800 Subject: [PATCH 38/74] Retain admin replay claims from server acceptance time Prune and reload administrator mutation fingerprints using accepted_at instead of the client-supplied timestamp, so a backdated first flight cannot shrink the ten-minute replay window. Older snapshots without the field still fall back to client_timestamp. --- CHANGELOG.md | 1 + docs/authentication-v2.md | 6 ++-- docs/authentication-v2.zh-CN.md | 5 ++-- src/common/auth.rs | 15 ++++++++++ src/common/auth/actor.rs | 3 +- src/common/auth/runtime.rs | 5 +--- src/common/auth/tests.rs | 50 +++++++++++++++++++++++++++++++++ 7 files changed, 76 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b498562..378d814 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -28,6 +28,7 @@ All notable changes to this project will be documented in this file. - Refused `write_admin_key_file` of the live `admin.key` while encrypted auth state exists, matching `initialize_admin_key`; `admin.key.next` remains allowed. - Kept `temporary_key_rotated` after a later issue in the same slot, aborted a timed-out embedded-relay shutdown, fail-closed replay logs whose directory sync failed, allowed a post-rotation write of the live key that already decrypts the snapshot, and treated `PB_MAPPER_NEW_STREAMS_PER_SECOND=0` as the default 100. - Fsynced the replay-log directory after compacting replacements, treated an unreadable existing replay log as unavailable, and made timing-wheel buckets hold `Weak` leases so renewals no longer accumulate day-long strong references. +- Retained administrator mutation replay claims from the server acceptance time, so a backdated client timestamp cannot shrink the ten-minute replay window. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/docs/authentication-v2.md b/docs/authentication-v2.md index 150a40b..2df5ee8 100644 --- a/docs/authentication-v2.md +++ b/docs/authentication-v2.md @@ -119,8 +119,10 @@ retention. A probable duplicate returns the stable retryable error `connection_salt_replayed`; one-shot administrator CLI operations retry once with a fresh salt. Mutating administrator requests additionally claim their exact fingerprint in the encrypted WAL before dispatch. Those claims survive -restart and compaction for ten minutes, so an old captured mutation cannot be -replayed after the Bloom window or a process restart. +restart and compaction for ten minutes after the server accepted them, so an +old captured mutation cannot be replayed after the Bloom window or a process +restart. The client-supplied first-flight timestamp is still checked for +freshness, but it does not control how long the claim is retained. ## Credential lifecycle diff --git a/docs/authentication-v2.zh-CN.md b/docs/authentication-v2.zh-CN.md index c744664..26bd94e 100644 --- a/docs/authentication-v2.zh-CN.md +++ b/docs/authentication-v2.zh-CN.md @@ -91,8 +91,9 @@ HKDF-SHA256 使用 connection salt 作为 salt，凭据的 32 字节 secret 作 1 MiB Bloom filter 的检查与写入，覆盖当前与上一个 600 秒窗口，使首帧允许的 最大未来时间戳无法在过滤器遗忘后继续重放。疑似重复会返回 可重试错误 `connection_salt_replayed`；一次性 admin CLI 会自动换 salt 重试一次。 -会修改状态的管理员请求还会在分发前把精确指纹写入加密 WAL；该记录在十分钟内跨 -重启、跨 compact 保留，不能通过等待 Bloom 窗口结束或重启进程来重放旧操作。 +会修改状态的管理员请求还会在分发前把精确指纹写入加密 WAL；该记录自服务端接受起 +在十分钟内跨重启、跨 compact 保留，不能通过等待 Bloom 窗口结束或重启进程来重放 +旧操作。客户端首帧时间戳仍用于新鲜度检查，但不决定这条记录保留多久。 ## 临时凭据生命周期 diff --git a/src/common/auth.rs b/src/common/auth.rs index 95f80f3..65c1f1b 100644 --- a/src/common/auth.rs +++ b/src/common/auth.rs @@ -865,6 +865,21 @@ struct PersistedSnapshot { struct AdminReplayRecord { fingerprint: [u8; 32], client_timestamp: u64, + /// Server receipt time used for retention. Older snapshots omit this field + /// (`0` after serde default) and fall back to `client_timestamp`. + #[serde(default)] + accepted_at: u64, +} + +impl AdminReplayRecord { + fn within_retention(&self, now: u64) -> bool { + let anchor = if self.accepted_at == 0 { + self.client_timestamp + } else { + self.accepted_at + }; + now.saturating_sub(anchor) <= ADMIN_REPLAY_RETENTION.as_secs() + } } #[derive(Clone, Debug, Serialize, Deserialize)] diff --git a/src/common/auth/actor.rs b/src/common/auth/actor.rs index 5189b9c..ae80864 100644 --- a/src/common/auth/actor.rs +++ b/src/common/auth/actor.rs @@ -307,6 +307,7 @@ fn actor_claim_admin_mutation( let record = AdminReplayRecord { fingerprint, client_timestamp, + accepted_at: now, }; fail_closed_on_uncertain_wal( inner, @@ -327,7 +328,7 @@ pub(super) fn prune_expired_admin_replays( admin_replay_order: &mut VecDeque, ) { admin_replay_order.retain(|record| { - let keep = now.saturating_sub(record.client_timestamp) <= ADMIN_REPLAY_RETENTION.as_secs(); + let keep = record.within_retention(now); if !keep { admin_replays.remove(&record.fingerprint); } diff --git a/src/common/auth/runtime.rs b/src/common/auth/runtime.rs index 260abee..97d5cf8 100644 --- a/src/common/auth/runtime.rs +++ b/src/common/auth/runtime.rs @@ -130,10 +130,7 @@ impl AuthRuntime { state .admin_replays .iter() - .filter(|record| { - now.saturating_sub(record.client_timestamp) - <= ADMIN_REPLAY_RETENTION.as_secs() - }) + .filter(|record| record.within_retention(now)) .cloned() .collect::>() }) diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index 6b8e69d..d471aa3 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -1011,10 +1011,12 @@ fn replay_pruning_removes_only_records_outside_the_retention_window() { let expired = AdminReplayRecord { fingerprint: [1; 32], client_timestamp: now - ADMIN_REPLAY_RETENTION.as_secs() - 1, + accepted_at: now - ADMIN_REPLAY_RETENTION.as_secs() - 1, }; let current = AdminReplayRecord { fingerprint: [2; 32], client_timestamp: now, + accepted_at: now, }; let mut replay_set = HashSet::from([expired.fingerprint, current.fingerprint]); let mut replay_order = VecDeque::from([expired, current.clone()]); @@ -1026,6 +1028,54 @@ fn replay_pruning_removes_only_records_outside_the_retention_window() { assert_eq!(replay_order[0].fingerprint, current.fingerprint); } +#[test] +fn replay_pruning_uses_server_acceptance_not_client_timestamp() { + let now = 10_000; + let retention = ADMIN_REPLAY_RETENTION.as_secs(); + let backdated = AdminReplayRecord { + fingerprint: [3; 32], + client_timestamp: now - retention - 1, + accepted_at: now - 1, + }; + let future_dated_but_expired = AdminReplayRecord { + fingerprint: [4; 32], + client_timestamp: now + retention / 2, + accepted_at: now - retention - 1, + }; + let mut replay_set = + HashSet::from([backdated.fingerprint, future_dated_but_expired.fingerprint]); + let mut replay_order = VecDeque::from([backdated.clone(), future_dated_but_expired]); + + super::actor::prune_expired_admin_replays(now, &mut replay_set, &mut replay_order); + + assert_eq!(replay_set, HashSet::from([backdated.fingerprint])); + assert_eq!(replay_order.len(), 1); + assert_eq!(replay_order[0].fingerprint, backdated.fingerprint); +} + +#[test] +fn replay_pruning_falls_back_to_client_timestamp_for_legacy_records() { + let now = 10_000; + let legacy_expired = AdminReplayRecord { + fingerprint: [5; 32], + client_timestamp: now - ADMIN_REPLAY_RETENTION.as_secs() - 1, + accepted_at: 0, + }; + let legacy_current = AdminReplayRecord { + fingerprint: [6; 32], + client_timestamp: now, + accepted_at: 0, + }; + let mut replay_set = HashSet::from([legacy_expired.fingerprint, legacy_current.fingerprint]); + let mut replay_order = VecDeque::from([legacy_expired, legacy_current.clone()]); + + super::actor::prune_expired_admin_replays(now, &mut replay_set, &mut replay_order); + + assert_eq!(replay_set, HashSet::from([legacy_current.fingerprint])); + assert_eq!(replay_order.len(), 1); + assert_eq!(replay_order[0].fingerprint, legacy_current.fingerprint); +} + #[test] fn tombstone_migration_prefers_audit_time_and_persists_fail_closed_fallback() { let now = 10_000; From 3873d690dd270810dd11e873fad1c114153f8a91 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 05:58:33 +0800 Subject: [PATCH 39/74] Harden replay I/O and keep previous root for rotated errors Replay logs now read exact 40-byte records and fail closed after an incomplete read or a rewrite that may already have replaced the file. Compaction temporary names include a random suffix so a leftover PID-1 file cannot block later rewrites. Revoked wheel owners are released when the tombstone is recycled or GC frees the slot, not at revoke time, so authenticate still reports temporary_key_revoked. The previous root key and instance id stay in memory so a stale first flight can decrypt and return temporary_key_rotated. Installers reject MSG_HEADER_KEY values that are not 32 printable ASCII bytes. --- CHANGELOG.md | 4 + docs/authentication-v2.md | 6 ++ docs/authentication-v2.zh-CN.md | 4 + scripts/install-server-gitee.sh | 12 ++- scripts/install-server-github.sh | 12 ++- src/common/auth.rs | 7 ++ src/common/auth/actor.rs | 16 +++ src/common/auth/runtime.rs | 15 +++ src/common/auth/tests.rs | 13 +++ src/common/auth/timing_wheel.rs | 5 + src/common/message/secure.rs | 135 +++++++++++++++++++----- src/common/message/secure/frame.rs | 20 ++++ src/common/message/secure/replay.rs | 88 ++++++++++------ src/common/message/secure/tests.rs | 153 ++++++++++++++++++++++++++++ 14 files changed, 432 insertions(+), 58 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 378d814..407196e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -29,6 +29,10 @@ All notable changes to this project will be documented in this file. - Kept `temporary_key_rotated` after a later issue in the same slot, aborted a timed-out embedded-relay shutdown, fail-closed replay logs whose directory sync failed, allowed a post-rotation write of the live key that already decrypts the snapshot, and treated `PB_MAPPER_NEW_STREAMS_PER_SECOND=0` as the default 100. - Fsynced the replay-log directory after compacting replacements, treated an unreadable existing replay log as unavailable, and made timing-wheel buckets hold `Weak` leases so renewals no longer accumulate day-long strong references. - Retained administrator mutation replay claims from the server acceptance time, so a backdated client timestamp cannot shrink the ten-minute replay window. +- Read first-flight replay logs with exact-record I/O and fail closed after an incomplete read or a rewrite that may have already replaced the log; compaction temporary files now use a random suffix. +- Released revoked timing-wheel owners when the tombstone is recycled or GC frees the slot, instead of keeping them until the original TTL. +- Kept the previous root key and instance id in memory so a stale first flight after rotation or reset can decrypt and return `temporary_key_rotated`. +- Validated installer `MSG_HEADER_KEY` values as 32 printable ASCII bytes before writing `admin.key`. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/docs/authentication-v2.md b/docs/authentication-v2.md index 2df5ee8..6d0f942 100644 --- a/docs/authentication-v2.md +++ b/docs/authentication-v2.md @@ -124,6 +124,12 @@ old captured mutation cannot be replayed after the Bloom window or a process restart. The client-supplied first-flight timestamp is still checked for freshness, but it does not control how long the claim is retained. +After a live root rotation or auth-state reset, the relay keeps the immediately +previous root key and instance id in memory. A first flight still encrypted with +the old temporary credential is decrypted with that previous material so the +client can read the stable `temporary_key_rotated` error instead of a decrypt +failure. + ## Credential lifecycle ### Issuance and renewal diff --git a/docs/authentication-v2.zh-CN.md b/docs/authentication-v2.zh-CN.md index 26bd94e..a4cc4f9 100644 --- a/docs/authentication-v2.zh-CN.md +++ b/docs/authentication-v2.zh-CN.md @@ -95,6 +95,10 @@ HKDF-SHA256 使用 connection salt 作为 salt，凭据的 32 字节 secret 作 在十分钟内跨重启、跨 compact 保留，不能通过等待 Bloom 窗口结束或重启进程来重放 旧操作。客户端首帧时间戳仍用于新鲜度检查，但不决定这条记录保留多久。 +根密钥轮换或认证状态重置之后，进程会在内存里保留上一份根密钥和 instance id。 +仍用旧临时凭据加密的首帧会用这份材料解密，从而把稳定的 +`temporary_key_rotated` 错误返回给客户端，而不是解密失败。 + ## 临时凭据生命周期 ### 签发与续期 diff --git a/scripts/install-server-gitee.sh b/scripts/install-server-gitee.sh index c33c7e2..b776776 100755 --- a/scripts/install-server-gitee.sh +++ b/scripts/install-server-gitee.sh @@ -15,6 +15,14 @@ ADMIN_KEY_PATH="${AUTH_DIR}/admin.key" LEGACY_KEY_PATH="/var/lib/pb-mapper-server/msg_header_key" SERVER_ENV_FILE="/etc/pb-mapper/server.env" +admin_key_is_env_safe() { + local key="$1" + local bytes + bytes=$(printf '%s' "$key" | wc -c) + [ "$bytes" -eq 32 ] || return 1 + printf '%s' "$key" | LC_ALL=C grep -qx '[[:graph:]]\{32\}' +} + configured_msg_header_key() { if [ -n "${MSG_HEADER_KEY:-}" ]; then printf '%s' "$MSG_HEADER_KEY" @@ -102,8 +110,8 @@ if [ -n "${MSG_HEADER_KEY:-}" ] && [ ! -s "$ADMIN_KEY_PATH" ]; then exit 1 ;; esac - if [ "${#MSG_HEADER_KEY}" -ne 32 ]; then - echo "MSG_HEADER_KEY must be a 32-character administrator key" >&2 + if ! admin_key_is_env_safe "$MSG_HEADER_KEY"; then + echo "MSG_HEADER_KEY must be exactly 32 printable ASCII bytes without whitespace or NUL" >&2 exit 1 fi printf '%s\n' "$MSG_HEADER_KEY" > "$ADMIN_KEY_PATH" diff --git a/scripts/install-server-github.sh b/scripts/install-server-github.sh index 07f5ef7..20cc033 100755 --- a/scripts/install-server-github.sh +++ b/scripts/install-server-github.sh @@ -15,6 +15,14 @@ ADMIN_KEY_PATH="${AUTH_DIR}/admin.key" LEGACY_KEY_PATH="/var/lib/pb-mapper-server/msg_header_key" SERVER_ENV_FILE="/etc/pb-mapper/server.env" +admin_key_is_env_safe() { + local key="$1" + local bytes + bytes=$(printf '%s' "$key" | wc -c) + [ "$bytes" -eq 32 ] || return 1 + printf '%s' "$key" | LC_ALL=C grep -qx '[[:graph:]]\{32\}' +} + configured_msg_header_key() { if [ -n "${MSG_HEADER_KEY:-}" ]; then printf '%s' "$MSG_HEADER_KEY" @@ -102,8 +110,8 @@ if [ -n "${MSG_HEADER_KEY:-}" ] && [ ! -s "$ADMIN_KEY_PATH" ]; then exit 1 ;; esac - if [ "${#MSG_HEADER_KEY}" -ne 32 ]; then - echo "MSG_HEADER_KEY must be a 32-character administrator key" >&2 + if ! admin_key_is_env_safe "$MSG_HEADER_KEY"; then + echo "MSG_HEADER_KEY must be exactly 32 printable ASCII bytes without whitespace or NUL" >&2 exit 1 fi printf '%s\n' "$MSG_HEADER_KEY" > "$ADMIN_KEY_PATH" diff --git a/src/common/auth.rs b/src/common/auth.rs index 65c1f1b..8a2beff 100644 --- a/src/common/auth.rs +++ b/src/common/auth.rs @@ -446,6 +446,12 @@ struct AdminState { lease: Weak, } +#[derive(Clone, Debug)] +struct PreviousRoot { + admin_key: AesKeyType, + instance_id: [u8; INSTANCE_ID_LEN], +} + #[derive(Debug)] struct AuthStateInner { admin: RwLock, @@ -463,6 +469,7 @@ struct AuthStateInner { auth_successes: AtomicU64, auth_failures: AtomicU64, root_epoch: AtomicU64, + previous_root: RwLock>, audit_records: RwLock>, } diff --git a/src/common/auth/actor.rs b/src/common/auth/actor.rs index ae80864..c747045 100644 --- a/src/common/auth/actor.rs +++ b/src/common/auth/actor.rs @@ -129,6 +129,7 @@ pub(super) async fn run_auth_actor( slot.expires_at = 0; slot.lease = Weak::new(); cold.remove(&key_id); + wheel.release(key_id); } } } @@ -221,6 +222,7 @@ pub(super) async fn run_auth_actor( &inner, &config, &mut cold, + &mut wheel, &mut tombstones, &admin_replay_order, )); @@ -638,6 +640,16 @@ fn actor_revoke( }) } +fn remember_previous_root(inner: &AuthStateInner) { + *inner + .previous_root + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(PreviousRoot { + admin_key: inner.admin_key(), + instance_id: inner.instance_id(), + }); +} + fn push_tombstone(tombstones: &mut VecDeque<(u64, u64)>, tombstoned_at: u64, key_id: u64) { let cleanup_at = tombstoned_at.saturating_add(TOMBSTONE_RETENTION.as_secs()); let index = tombstones.partition_point(|(current, _)| *current <= cleanup_at); @@ -648,6 +660,7 @@ fn actor_gc( inner: &Arc, config: &AuthConfig, cold: &mut HashMap, + wheel: &mut TimingWheel, tombstones: &mut VecDeque<(u64, u64)>, admin_replays: &VecDeque, ) -> Result { @@ -670,6 +683,7 @@ fn actor_gc( slot.expires_at = 0; slot.lease = Weak::new(); cold.remove(&key_id); + wheel.release(key_id); removed = removed.saturating_add(1); } } @@ -719,6 +733,7 @@ fn actor_reset( } let _ = std::fs::remove_file(&next_instance_path); push_audit_record(inner, reset_audit); + remember_previous_root(inner); cancel_all_temporary_leases(inner); { @@ -783,6 +798,7 @@ fn actor_rotate_root( } let _ = std::fs::remove_file(&next_key_path); push_audit_record(inner, rotate_audit); + remember_previous_root(inner); cancel_all_temporary_leases(inner); let old_admin_lease = admin_lease.clone(); diff --git a/src/common/auth/runtime.rs b/src/common/auth/runtime.rs index 97d5cf8..d6ef3fb 100644 --- a/src/common/auth/runtime.rs +++ b/src/common/auth/runtime.rs @@ -170,6 +170,7 @@ impl AuthRuntime { auth_successes: AtomicU64::new(0), auth_failures: AtomicU64::new(0), root_epoch: AtomicU64::new(loaded.as_ref().map(|state| state.root_epoch).unwrap_or(0)), + previous_root: RwLock::new(None), audit_records: RwLock::new(audit_records), }); let (command_tx, command_rx) = mpsc::channel(256); @@ -217,6 +218,20 @@ impl AuthRuntime { derive_temporary_key(&inner.admin_key(), &inner.instance_id(), key_id) } + pub(crate) fn derive_previous_key(&self, key_id: u64) -> Option { + let inner = self.inner().ok()?; + let previous = inner + .previous_root + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clone()?; + if key_id == 0 { + Some(previous.admin_key) + } else { + derive_temporary_key(&previous.admin_key, &previous.instance_id, key_id).ok() + } + } + pub fn authenticate_presented( &self, key_id: u64, diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index d471aa3..3f2b9c0 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -993,6 +993,19 @@ fn timing_wheel_expires_cascaded_boundary_entry_without_an_extra_tick() { assert_eq!(due[0].key_id(), lease.key_id()); } +#[test] +fn timing_wheel_release_drops_the_current_owner() { + let now = 1_000; + let lease = Arc::new(AuthLease::new(make_key_id(1, 0), now + 60)); + let mut wheel = TimingWheel::new(now); + wheel.insert(lease.clone()); + assert!(wheel.owns(lease.key_id())); + + wheel.release(lease.key_id()); + assert!(!wheel.owns(lease.key_id())); + assert!(!lease.cancellation_token().is_cancelled()); +} + #[test] fn timing_wheel_clear_cancels_owned_leases() { let now = 1_000; diff --git a/src/common/auth/timing_wheel.rs b/src/common/auth/timing_wheel.rs index bb53554..a410593 100644 --- a/src/common/auth/timing_wheel.rs +++ b/src/common/auth/timing_wheel.rs @@ -55,6 +55,11 @@ impl TimingWheel { self.owners.remove(&key_id); } + #[cfg(test)] + pub(super) fn owns(&self, key_id: u64) -> bool { + self.owners.contains_key(&key_id) + } + pub(super) fn insert_with_version(&mut self, lease: Arc, version: u64) { self.owners.insert(lease.key_id(), lease.clone()); let expires_at = lease.expires_at(); diff --git a/src/common/message/secure.rs b/src/common/message/secure.rs index 6cb7da9..a86183b 100644 --- a/src/common/message/secure.rs +++ b/src/common/message/secure.rs @@ -511,31 +511,42 @@ impl ServerSecurity { ), response_session: None, })?; - let mut session = ServerHeaderSession { - protocol: HeaderProtocol::V2, - legacy_key: key, - v2: Some(material.clone()), - context: None, - _legacy_guard: None, - }; - let mut message_reader = V2MessageReader::new( - reader, - material, + let mut session = v2_session(key, material.clone()); + let (counter, ciphertext) = + read_initial_v2_ciphertext(reader) + .await + .map_err(|error| ServerInitialError { + failure: AuthFailure::new( + "protocol_v2_decrypt_failed", + error.to_string(), + false, + ), + response_session: Some(session_without_context(&session)), + })?; + let mut current_ciphertext = ciphertext.clone(); + let payload = match open_v2_payload( + &material, DIRECTION_CLIENT_TO_SERVER, - 0, - ) - .map_err(|error| ServerInitialError { - failure: AuthFailure::new("protocol_v2_decrypt_failed", error.to_string(), false), - response_session: Some(session_without_context(&session)), - })?; - let payload = message_reader - .read_msg_with_limit(MAX_INITIAL_PLAINTEXT_LEN) - .await - .map_err(|error| ServerInitialError { - failure: AuthFailure::new("protocol_v2_decrypt_failed", error.to_string(), false), - response_session: Some(session_without_context(&session)), - })? - .to_vec(); + counter, + &mut current_ciphertext, + ) { + Ok(payload) => payload, + Err(error) => { + if let Some(stale) = + stale_root_first_flight(&self.auth, key_id, salt, counter, &ciphertext) + { + return Err(stale); + } + return Err(ServerInitialError { + failure: AuthFailure::new( + "protocol_v2_decrypt_failed", + error.to_string(), + false, + ), + response_session: Some(session_without_context(&session)), + }); + } + }; let context = self .auth @@ -596,6 +607,80 @@ impl ServerSecurity { mod limiter; pub use limiter::FailureLogDecision; use limiter::FailureLogLimiter; +async fn read_initial_v2_ciphertext( + reader: &mut T, +) -> Result<(u64, Vec)> { + let counter = reader + .read_u64() + .await + .map_err(|error| protocol_error(format!("failed to read v2 counter: {error}")))?; + if counter != 0 { + return Err(protocol_error(format!( + "protocol-v2 counter mismatch: expected 0, got {counter}" + ))); + } + let datalen = reader + .read_u32() + .await + .map_err(|error| protocol_error(format!("failed to read v2 length: {error}")))?; + let max_encrypted_len = MAX_INITIAL_PLAINTEXT_LEN.saturating_add(AES_256_GCM.tag_len() as u32); + if datalen < AES_256_GCM.tag_len() as u32 || datalen > max_encrypted_len { + return Err(protocol_error(format!( + "protocol-v2 payload length {datalen} exceeds the {MAX_INITIAL_PLAINTEXT_LEN}-byte limit" + ))); + } + let mut ciphertext = vec![0_u8; datalen as usize]; + reader + .read_exact(&mut ciphertext) + .await + .map_err(|error| protocol_error(format!("failed to read v2 payload: {error}")))?; + Ok((counter, ciphertext)) +} + +fn stale_root_first_flight( + auth: &AuthRuntime, + key_id: u64, + salt: [u8; CONNECTION_SALT_LEN], + counter: u64, + ciphertext: &[u8], +) -> Option { + let previous_key = auth.derive_previous_key(key_id)?; + let previous_material = derive_material(key_id, &previous_key, salt).ok()?; + let mut previous_ciphertext = ciphertext.to_vec(); + open_v2_payload( + &previous_material, + DIRECTION_CLIENT_TO_SERVER, + counter, + &mut previous_ciphertext, + ) + .ok()?; + let (code, message) = if key_id == 0 { + ( + "administrator_key_invalid", + "administrator credential does not match the active root key", + ) + } else { + ( + "temporary_key_rotated", + "temporary credential was invalidated by administrator root rotation or auth-state reset", + ) + }; + Some(ServerInitialError { + failure: AuthFailure::new(code, message, false), + response_session: Some(v2_session(previous_key, previous_material)), + }) +} + +fn v2_session(key: AesKeyType, material: V2Material) -> ServerHeaderSession { + ServerHeaderSession { + protocol: HeaderProtocol::V2, + legacy_key: key, + v2: Some(material), + context: None, + _legacy_guard: None, + } +} + fn session_without_context(session: &ServerHeaderSession) -> ServerHeaderSession { ServerHeaderSession { protocol: session.protocol, @@ -635,7 +720,7 @@ impl MessageWriter for HeaderMessageWriter<'_, T> { } mod frame; -use frame::{derive_material, first_prefix, V2Material}; +use frame::{derive_material, first_prefix, open_v2_payload, V2Material}; pub use frame::{V2MessageReader, V2MessageWriter}; mod replay; #[cfg(test)] diff --git a/src/common/message/secure/frame.rs b/src/common/message/secure/frame.rs index 659a7cd..da0afd0 100644 --- a/src/common/message/secure/frame.rs +++ b/src/common/message/secure/frame.rs @@ -170,6 +170,26 @@ impl MessageWriter for V2MessageWriter<'_, T> { } } +pub(super) fn open_v2_payload( + material: &V2Material, + direction: u8, + counter: u64, + ciphertext: &mut [u8], +) -> Result> { + let key_bytes = direction_key(material, direction); + let key = LessSafeKey::new( + UnboundKey::new(&AES_256_GCM, key_bytes) + .map_err(|_| protocol_error("invalid protocol-v2 read key"))?, + ); + let datalen = u32::try_from(ciphertext.len()) + .map_err(|_| protocol_error("protocol-v2 payload length is invalid"))?; + let aad = frame_aad(material, direction, counter, datalen); + let plain = key + .open_in_place(nonce(counter), Aad::from(aad.as_slice()), ciphertext) + .map_err(|_| protocol_error("protocol-v2 payload authentication failed"))?; + Ok(plain.to_vec()) +} + pub(super) fn derive_material( key_id: u64, credential_key: &AesKeyType, diff --git a/src/common/message/secure/replay.rs b/src/common/message/secure/replay.rs index 60914d2..83f31ac 100644 --- a/src/common/message/secure/replay.rs +++ b/src/common/message/secure/replay.rs @@ -17,9 +17,11 @@ use std::collections::HashMap; use std::fs::{File, OpenOptions}; -use std::io::{Read, Write}; +use std::io::{ErrorKind, Read, Write}; use std::path::PathBuf; +use rand::RngExt; + use super::*; const DEFAULT_NEW_STREAMS_PER_SECOND: u32 = 100; @@ -220,14 +222,15 @@ impl ReplayGuard { } }; let now = unix_seconds(); - let mut live = Vec::new(); - let mut record = [0_u8; REPLAY_RECORD_LEN]; - loop { - match file.read(&mut record) { - Ok(0) => break, - Ok(n) if n == REPLAY_RECORD_LEN => {} - _ => break, + let records = match read_complete_replay_records(&mut file) { + Ok(records) => records, + Err(_) => { + self.log_failed = true; + return; } + }; + let mut live = Vec::new(); + for record in records { let timestamp = u64::from_be_bytes(record[32..].try_into().expect("timestamp width")); if now.saturating_sub(timestamp) > self.window_seconds { continue; @@ -236,9 +239,11 @@ impl ReplayGuard { self.bloom.insert(&fingerprint, timestamp); live.push(record); } - if self.rewrite_live(&live).is_ok() { - self.last_compact_at = now; + if self.rewrite_live(&live).is_err() { + self.log_failed = true; + return; } + self.last_compact_at = now; } fn compact(&mut self, now: u64) { @@ -246,38 +251,51 @@ impl ReplayGuard { self.last_compact_at = now; return; }; - let Ok(mut file) = File::open(path) else { - self.last_compact_at = now; - return; - }; - let mut live = Vec::new(); - let mut record = [0_u8; REPLAY_RECORD_LEN]; - loop { - match file.read(&mut record) { - Ok(0) => break, - Ok(n) if n == REPLAY_RECORD_LEN => {} - _ => break, + let mut file = match File::open(path) { + Ok(file) => file, + Err(_) => { + self.log_failed = true; + return; } - let timestamp = u64::from_be_bytes(record[32..].try_into().expect("timestamp width")); - if now.saturating_sub(timestamp) <= self.window_seconds { - live.push(record); + }; + let records = match read_complete_replay_records(&mut file) { + Ok(records) => records, + Err(_) => { + self.log_failed = true; + return; } + }; + let live = records + .into_iter() + .filter(|record| { + let timestamp = + u64::from_be_bytes(record[32..].try_into().expect("timestamp width")); + now.saturating_sub(timestamp) <= self.window_seconds + }) + .collect::>(); + if self.rewrite_live(&live).is_err() { + self.log_failed = true; + return; } - if self.rewrite_live(&live).is_ok() { - self.last_compact_at = now; - } + self.last_compact_at = now; } fn rewrite_live(&self, live: &[[u8; REPLAY_RECORD_LEN]]) -> std::io::Result<()> { let Some(path) = &self.log_path else { return Ok(()); }; + let mut random_suffix = [0_u8; 8]; + let mut rng = rand::rng(); + for byte in &mut random_suffix { + *byte = rng.random(); + } let temporary = path.with_file_name(format!( - ".{}.tmp-{}", + ".{}.tmp-{}-{:016x}", path.file_name() .and_then(|name| name.to_str()) .unwrap_or("connection.replay"), - std::process::id() + std::process::id(), + u64::from_be_bytes(random_suffix) )); let result = (|| { let mut file = OpenOptions::new() @@ -298,6 +316,18 @@ impl ReplayGuard { } } +fn read_complete_replay_records(file: &mut File) -> std::io::Result> { + let mut records = Vec::new(); + loop { + let mut record = [0_u8; REPLAY_RECORD_LEN]; + match file.read_exact(&mut record) { + Ok(()) => records.push(record), + Err(error) if error.kind() == ErrorKind::UnexpectedEof => return Ok(records), + Err(error) => return Err(error), + } + } +} + fn bloom_positions(filter_len: usize, fingerprint: &[u8; 32]) -> [usize; 4] { let bits = filter_len * 8; std::array::from_fn(|index| { diff --git a/src/common/message/secure/tests.rs b/src/common/message/secure/tests.rs index 0c88c3c..a4b2c95 100644 --- a/src/common/message/secure/tests.rs +++ b/src/common/message/secure/tests.rs @@ -175,6 +175,99 @@ async fn revoked_first_flights_do_not_consume_the_replay_filter() { let _ = std::fs::remove_dir_all(config.state_dir); } +#[tokio::test] +async fn rotated_temporary_first_flight_returns_a_readable_rotated_error() { + let _process_credential_guard = PROCESS_CREDENTIAL_TEST_LOCK.lock().await; + let admin = *b"0123456789abcdefghijklmnopqrstuv"; + let new_admin = *b"abcdefghijklmnopqrstuvwxyz012345"; + let config = temp_config(); + let auth = AuthRuntime::start(admin, config.clone()).await.unwrap(); + let admin_context = auth.authenticate_presented(0, &admin).unwrap(); + let issued = auth + .issue(&admin_context, std::time::Duration::from_secs(60), None) + .await + .unwrap(); + let Credential::Temporary { key_id, key } = parse_credential(&issued.credential).unwrap() + else { + panic!("expected temporary credential"); + }; + let client = ClientHeaderSession::new_v2(&Credential::Temporary { key_id, key }).unwrap(); + auth.rotate_root(&admin_context, new_admin).await.unwrap(); + + let security = ServerSecurity::new(auth); + let (mut client_io, mut server_io) = tokio::io::duplex(4096); + let client_task = async { + client + .write_initial(&mut client_io, b"stale-after-rotate") + .await + .unwrap(); + let mut reader = client.response_reader(&mut client_io).unwrap(); + reader.read_msg().await.unwrap().to_vec() + }; + let server_task = async { + let error = match security.read_initial(&mut server_io).await { + Ok(_) => panic!("rotated credential should fail"), + Err(error) => error, + }; + assert_eq!(error.failure.code, "temporary_key_rotated"); + let session = error.response_session.expect("readable error session"); + let mut writer = session.response_writer(&mut server_io).unwrap(); + writer.write_msg(b"temporary_key_rotated").await.unwrap(); + error.failure.code + }; + let (plaintext, code) = tokio::join!(client_task, server_task); + assert_eq!(code, "temporary_key_rotated"); + assert_eq!(plaintext, b"temporary_key_rotated"); + + let _ = std::fs::remove_dir_all(config.state_dir); +} + +#[tokio::test] +async fn reset_temporary_first_flight_returns_a_readable_rotated_error() { + let _process_credential_guard = PROCESS_CREDENTIAL_TEST_LOCK.lock().await; + let admin = *b"0123456789abcdefghijklmnopqrstuv"; + let config = temp_config(); + let auth = AuthRuntime::start(admin, config.clone()).await.unwrap(); + let admin_context = auth.authenticate_presented(0, &admin).unwrap(); + let issued = auth + .issue(&admin_context, std::time::Duration::from_secs(60), None) + .await + .unwrap(); + let Credential::Temporary { key_id, key } = parse_credential(&issued.credential).unwrap() + else { + panic!("expected temporary credential"); + }; + let client = ClientHeaderSession::new_v2(&Credential::Temporary { key_id, key }).unwrap(); + auth.reset(&admin_context).await.unwrap(); + + let security = ServerSecurity::new(auth); + let (mut client_io, mut server_io) = tokio::io::duplex(4096); + let client_task = async { + client + .write_initial(&mut client_io, b"stale-after-reset") + .await + .unwrap(); + let mut reader = client.response_reader(&mut client_io).unwrap(); + reader.read_msg().await.unwrap().to_vec() + }; + let server_task = async { + let error = match security.read_initial(&mut server_io).await { + Ok(_) => panic!("reset credential should fail"), + Err(error) => error, + }; + assert_eq!(error.failure.code, "temporary_key_rotated"); + let session = error.response_session.expect("readable error session"); + let mut writer = session.response_writer(&mut server_io).unwrap(); + writer.write_msg(b"temporary_key_rotated").await.unwrap(); + error.failure.code + }; + let (plaintext, code) = tokio::join!(client_task, server_task); + assert_eq!(code, "temporary_key_rotated"); + assert_eq!(plaintext, b"temporary_key_rotated"); + + let _ = std::fs::remove_dir_all(config.state_dir); +} + #[tokio::test] async fn oversized_initial_frame_is_rejected_before_reading_its_body() { let credential = Credential::Admin(*b"0123456789abcdefghijklmnopqrstuv"); @@ -273,6 +366,66 @@ fn per_credential_admission_limit_does_not_consume_other_keys() { assert_eq!(guard.admit(1, &[1_u8; 32], now), FirstFlightAdmit::Replayed); } +#[test] +fn persisted_first_flights_survive_a_torn_trailing_record() { + let mut random = [0_u8; 8]; + let mut rng = rand::rng(); + for byte in &mut random { + *byte = rng.random(); + } + let path = std::env::temp_dir().join(format!( + "pb-mapper-replay-torn-{}", + u64::from_be_bytes(random) + )); + let now = unix_seconds(); + let fingerprint = [17_u8; 32]; + { + let mut guard = ReplayGuard::open(Some(path.clone()), 1024, DEFAULT_REPLAY_WINDOW_SECONDS); + assert_eq!(guard.admit(7, &fingerprint, now), FirstFlightAdmit::Fresh); + } + let mut torn = std::fs::read(&path).unwrap(); + torn.extend_from_slice(&[0_u8; 10]); + std::fs::write(&path, torn).unwrap(); + let mut restored = ReplayGuard::open(Some(path.clone()), 1024, DEFAULT_REPLAY_WINDOW_SECONDS); + assert_eq!( + restored.admit(7, &fingerprint, now), + FirstFlightAdmit::Replayed + ); + let _ = std::fs::remove_file(path); +} + +#[test] +fn replay_rewrite_succeeds_when_a_pid_temporary_file_already_exists() { + let mut random = [0_u8; 8]; + let mut rng = rand::rng(); + for byte in &mut random { + *byte = rng.random(); + } + let path = std::env::temp_dir().join(format!( + "pb-mapper-replay-tmp-{}", + u64::from_be_bytes(random) + )); + let now = unix_seconds(); + let fingerprint = [19_u8; 32]; + { + let mut guard = ReplayGuard::open(Some(path.clone()), 1024, DEFAULT_REPLAY_WINDOW_SECONDS); + assert_eq!(guard.admit(9, &fingerprint, now), FirstFlightAdmit::Fresh); + } + let leftover = path.with_file_name(format!( + ".{}.tmp-{}", + path.file_name().unwrap().to_str().unwrap(), + std::process::id() + )); + std::fs::write(&leftover, b"stale").unwrap(); + let mut restored = ReplayGuard::open(Some(path.clone()), 1024, DEFAULT_REPLAY_WINDOW_SECONDS); + assert_eq!( + restored.admit(9, &fingerprint, now), + FirstFlightAdmit::Replayed + ); + let _ = std::fs::remove_file(leftover); + let _ = std::fs::remove_file(path); +} + #[test] fn persisted_first_flights_survive_replay_guard_restart() { let mut random = [0_u8; 8]; From e0be4d27f2723cab478d2b27ffb21c7383a05e5d Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 06:14:44 +0800 Subject: [PATCH 40/74] Skip machine-key reinit and pin running client credentials Container entrypoints now pass --use-machine-msg-header-key only when admin.key is missing, so a persistent volume restart does not hit the existing-key refusal. Local client listeners capture the process credential at start and reuse it for probes and accepted streams. key gc removes expired and revoked high-slot entries while keeping their generations so shrinking capacity can reclaim them. --- CHANGELOG.md | 3 + scripts/release/entrypoint/pb-mapper.sh | 8 ++- src/common/auth/actor.rs | 22 +++++++- src/common/auth/runtime.rs | 13 +++++ src/common/auth/tests.rs | 74 +++++++++++++++++++++++++ src/local/client/mod.rs | 39 ++++++++++--- src/local/client/status.rs | 26 ++++++++- src/local/client/stream.rs | 4 +- 8 files changed, 174 insertions(+), 15 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 407196e..ed57049 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -33,6 +33,9 @@ All notable changes to this project will be documented in this file. - Released revoked timing-wheel owners when the tombstone is recycled or GC frees the slot, instead of keeping them until the original TTL. - Kept the previous root key and instance id in memory so a stale first flight after rotation or reset can decrypt and return `temporary_key_rotated`. - Validated installer `MSG_HEADER_KEY` values as 32 printable ASCII bytes before writing `admin.key`. +- Passed `--use-machine-msg-header-key` only when `admin.key` is missing, so a container restart with a persistent auth volume does not fail after first boot. +- Pinned each local client listener to the credential captured at start, so a later Flutter config change cannot switch an existing port onto another tenant. +- Garbage-collected expired and revoked high-slot entries while keeping their generations so shrinking capacity and running `key gc` can reclaim them. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/scripts/release/entrypoint/pb-mapper.sh b/scripts/release/entrypoint/pb-mapper.sh index 15124f7..5b706c4 100644 --- a/scripts/release/entrypoint/pb-mapper.sh +++ b/scripts/release/entrypoint/pb-mapper.sh @@ -28,8 +28,12 @@ else fi if [ "$USE_MACHINE_MSG_HEADER_KEY" = "true" ]; then - echo "WARNING: USE_MACHINE_MSG_HEADER_KEY is a legacy compatibility mode" - ARGS+=(--use-machine-msg-header-key) + if [ -s "$ADMIN_KEY_PATH" ]; then + echo "admin.key already exists; skipping --use-machine-msg-header-key" + else + echo "WARNING: USE_MACHINE_MSG_HEADER_KEY is a legacy compatibility mode" + ARGS+=(--use-machine-msg-header-key) + fi else echo "USE_MACHINE_MSG_HEADER_KEY is set to false" fi diff --git a/src/common/auth/actor.rs b/src/common/auth/actor.rs index c747045..5bf8920 100644 --- a/src/common/auth/actor.rs +++ b/src/common/auth/actor.rs @@ -687,8 +687,28 @@ fn actor_gc( removed = removed.saturating_add(1); } } - tombstones.clear(); drop(slots); + { + let mut high = inner + .high_slot_entries + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + high.retain(|entry| { + let keep = match entry.state { + SlotState::Active if entry.expires_at > now => true, + SlotState::Active | SlotState::Expired | SlotState::Revoked | SlotState::Free => { + false + } + }; + if !keep { + cold.remove(&entry.key_id); + wheel.release(entry.key_id); + removed = removed.saturating_add(1); + } + keep + }); + } + tombstones.clear(); let gc_audit = audit("temporary_key_gc", None, Some(format!("removed={removed}"))); let mut snapshot = build_snapshot(inner, cold, admin_replays); push_persisted_audit(&mut snapshot.audit_records, gc_audit.clone()); diff --git a/src/common/auth/runtime.rs b/src/common/auth/runtime.rs index d6ef3fb..4cb7b9d 100644 --- a/src/common/auth/runtime.rs +++ b/src/common/auth/runtime.rs @@ -218,6 +218,19 @@ impl AuthRuntime { derive_temporary_key(&inner.admin_key(), &inner.instance_id(), key_id) } + #[cfg(test)] + pub(crate) fn high_slot_entry_count(&self) -> usize { + self.inner() + .map(|inner| { + inner + .high_slot_entries + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .len() + }) + .unwrap_or(0) + } + pub(crate) fn derive_previous_key(&self, key_id: u64) -> Option { let inner = self.inner().ok()?; let previous = inner diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index 3f2b9c0..b2d34b7 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -123,6 +123,80 @@ async fn shrinking_then_expanding_capacity_does_not_reuse_old_key_ids() { let _ = std::fs::remove_dir_all(state_dir); } +#[tokio::test] +async fn gc_removes_inactive_high_slot_entries_and_keeps_their_generations() { + let state_dir = temp_state_dir("gc-high-slots"); + let admin_key = *b"0123456789abcdefghijklmnopqrstuv"; + let config_two = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 2, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + let runtime = AuthRuntime::start(admin_key, config_two.clone()) + .await + .unwrap(); + let admin = authenticate_for_test(&runtime, 0).unwrap(); + let first = runtime + .issue(&admin, Duration::from_secs(60), Some("first".to_string())) + .await + .unwrap(); + let second = runtime + .issue(&admin, Duration::from_secs(60), Some("second".to_string())) + .await + .unwrap(); + let Credential::Temporary { + key_id: first_id, .. + } = parse_credential(&first.credential).unwrap() + else { + panic!("expected temporary credential"); + }; + let Credential::Temporary { + key_id: second_id, .. + } = parse_credential(&second.credential).unwrap() + else { + panic!("expected temporary credential"); + }; + runtime.revoke(&admin, first_id).await.unwrap(); + runtime.revoke(&admin, second_id).await.unwrap(); + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + + let config_one = AuthConfig { + max_temporary_keys: 1, + ..config_two.clone() + }; + let runtime = AuthRuntime::start(admin_key, config_one).await.unwrap(); + let admin = authenticate_for_test(&runtime, 0).unwrap(); + assert_eq!(runtime.high_slot_entry_count(), 1); + let removed = runtime.gc(&admin).await.unwrap(); + assert!(removed >= 1); + assert_eq!(runtime.high_slot_entry_count(), 0); + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + + let runtime = AuthRuntime::start(admin_key, config_two).await.unwrap(); + let admin = authenticate_for_test(&runtime, 0).unwrap(); + let _low = runtime + .issue(&admin, Duration::from_secs(60), Some("low".to_string())) + .await + .unwrap(); + let high = runtime + .issue(&admin, Duration::from_secs(60), Some("high".to_string())) + .await + .unwrap(); + let Credential::Temporary { + key_id: reused_id, .. + } = parse_credential(&high.credential).unwrap() + else { + panic!("expected temporary credential"); + }; + assert_ne!(reused_id, second_id); + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + let _ = std::fs::remove_dir_all(state_dir); +} + #[tokio::test] async fn safe_mode_denies_legacy_protocol_instead_of_restoring_the_default() { let state_dir = temp_state_dir("safe-mode-legacy"); diff --git a/src/local/client/mod.rs b/src/local/client/mod.rs index b46b4c5..a43e1f9 100644 --- a/src/local/client/mod.rs +++ b/src/local/client/mod.rs @@ -13,8 +13,9 @@ use tokio::time::MissedTickBehavior; use uni_stream::udp::set_custom_timeout; use self::error::{AcceptLocalStreamSnafu, BindLocalListenerSnafu}; -use self::status::{get_status, get_status_scoped}; +use self::status::{get_status, get_status_scoped, get_status_with_credential}; use self::stream::handle_local_stream; +use crate::common::checksum::{get_process_credential, Credential}; use crate::common::config::{ client_health_check_interval, client_health_check_timeout, client_health_failure_threshold, StatusOp, @@ -123,6 +124,16 @@ pub async fn run_client_side_cli_with_callback_scoped< return; } }; + let credential = match get_process_credential() { + Ok(credential) => credential, + Err(e) => { + tracing::error!("load client credential failed: {e}"); + if let Some(ref callback) = status_callback { + callback("failed"); + } + return; + } + }; let mut retry_backoff = RetryBackoff::default(); @@ -136,7 +147,9 @@ pub async fn run_client_side_cli_with_callback_scoped< "client probing remote server" ); - if let Err(reason) = probe_remote_key(remote_addr, key.as_ref(), namespace).await { + if let Err(reason) = + probe_remote_key(remote_addr, key.as_ref(), namespace, credential).await + { let retry_delay = retry_backoff.next_delay(); tracing::warn!( event = "client_remote_probe_failed", @@ -222,7 +235,7 @@ pub async fn run_client_side_cli_with_callback_scoped< let failure_tx = stream_failure_tx.clone(); tokio::spawn(async move { if let Err(e) = - handle_local_stream(stream, key, remote_addr, keep_alive, namespace).await + handle_local_stream(stream, key, remote_addr, keep_alive, namespace, credential).await { let reason = snafu::Report::from_error(e).to_string(); tracing::warn!( @@ -236,7 +249,7 @@ pub async fn run_client_side_cli_with_callback_scoped< }); } _ = health_interval.tick() => { - if let Err(reason) = probe_remote_key(remote_addr, key.as_ref(), namespace).await { + if let Err(reason) = probe_remote_key(remote_addr, key.as_ref(), namespace, credential).await { consecutive_health_failures = consecutive_health_failures.saturating_add(1); if consecutive_health_failures < health_failure_threshold { tracing::warn!( @@ -278,7 +291,7 @@ pub async fn run_client_side_cli_with_callback_scoped< stream_failure = %stream_failure, "local stream failure reported; probing remote key" ); - if let Err(reason) = probe_remote_key(remote_addr, key.as_ref(), namespace).await { + if let Err(reason) = probe_remote_key(remote_addr, key.as_ref(), namespace, credential).await { tracing::warn!( event = "client_remote_probe_failed_after_stream_error", key = %key, @@ -316,9 +329,15 @@ async fn probe_remote_key( remote_addr: SocketAddr, key: &str, namespace: Option, + credential: Credential, ) -> std::result::Result<(), String> { let timeout = client_health_check_timeout(); - match tokio::time::timeout(timeout, probe_remote_key_once(remote_addr, key, namespace)).await { + match tokio::time::timeout( + timeout, + probe_remote_key_once(remote_addr, key, namespace, credential), + ) + .await + { Ok(result) => result, Err(_) => Err(format!("remote key probe timed out after {timeout:?}")), } @@ -328,6 +347,7 @@ async fn probe_remote_key_once( remote_addr: SocketAddr, key: &str, namespace: Option, + credential: Credential, ) -> std::result::Result<(), String> { match fetch_remote_status( remote_addr, @@ -335,6 +355,7 @@ async fn probe_remote_key_once( key: key.to_string(), }, namespace, + credential, ) .await { @@ -362,7 +383,8 @@ async fn probe_remote_key_once( } } - let status_resp = fetch_remote_status(remote_addr, PbConnStatusReq::Keys, namespace).await?; + let status_resp = + fetch_remote_status(remote_addr, PbConnStatusReq::Keys, namespace, credential).await?; let PbConnStatusResp::Keys(keys) = status_resp else { return Err(format!( "expected keys status response, got {status_resp:?}" @@ -381,11 +403,12 @@ async fn fetch_remote_status( remote_addr: SocketAddr, req: PbConnStatusReq, namespace: Option, + credential: Credential, ) -> std::result::Result { let mut stream = each_addr(remote_addr, TcpStream::connect) .await .map_err(|e| format!("connect remote stream failed: {e}"))?; - get_status_scoped(&mut stream, req, namespace) + get_status_with_credential(&mut stream, req, namespace, &credential) .await .map_err(|e| format!("get status failed: {}", snafu::Report::from_error(e))) } diff --git a/src/local/client/status.rs b/src/local/client/status.rs index 1902889..a479ab3 100644 --- a/src/local/client/status.rs +++ b/src/local/client/status.rs @@ -5,6 +5,7 @@ use super::error::{ ControlIoTimeoutSnafu, CreateHeaderToolSnafu, DecodeStatusRespSnafu, EncodeStatusReqSnafu, ReadStatusRespSnafu, StatusRespNotMatchSnafu, WriteStatusReqSnafu, }; +use crate::common::checksum::Credential; use crate::common::config::control_io_timeout; use crate::common::message::command::{ MessageSerializer, PbConnRequest, PbConnResponse, PbConnStatusReq, PbConnStatusResp, @@ -23,6 +24,28 @@ pub async fn get_status_scoped( remote_stream: &mut S, req: PbConnStatusReq, namespace: Option, +) -> super::error::Result { + let session = + ClientHeaderSession::from_process().context(CreateHeaderToolSnafu { action: "session" })?; + get_status_with_session(remote_stream, req, namespace, session).await +} + +pub async fn get_status_with_credential( + remote_stream: &mut S, + req: PbConnStatusReq, + namespace: Option, + credential: &Credential, +) -> super::error::Result { + let session = ClientHeaderSession::new_v2(credential) + .context(CreateHeaderToolSnafu { action: "session" })?; + get_status_with_session(remote_stream, req, namespace, session).await +} + +async fn get_status_with_session( + remote_stream: &mut S, + req: PbConnStatusReq, + namespace: Option, + session: ClientHeaderSession, ) -> super::error::Result { let timeout = control_io_timeout(); let request = match namespace { @@ -33,9 +56,6 @@ pub async fn get_status_scoped( None => PbConnRequest::Status(req), }; let msg = request.encode().context(EncodeStatusReqSnafu)?; - - let session = - ClientHeaderSession::from_process().context(CreateHeaderToolSnafu { action: "session" })?; match tokio::time::timeout(timeout, session.write_initial(remote_stream, &msg)).await { Ok(result) => result.context(WriteStatusReqSnafu)?, Err(_) => ControlIoTimeoutSnafu { diff --git a/src/local/client/stream.rs b/src/local/client/stream.rs index cf9781c..9815f75 100644 --- a/src/local/client/stream.rs +++ b/src/local/client/stream.rs @@ -10,6 +10,7 @@ use super::error::{ EncodeSubcribeReqSnafu, ReadSubcribeRespSnafu, Result, SubcribeRespNotMatchSnafu, WriteSubcribeReqSnafu, }; +use crate::common::checksum::Credential; use crate::common::config::control_io_timeout; use crate::common::message::command::{MessageSerializer, PbConnRequest, PbConnResponse}; use crate::common::message::forward::StreamForward; @@ -30,6 +31,7 @@ pub async fn handle_local_stream< remote_addr: A, keep_alive: bool, namespace: Option, + credential: Credential, ) -> Result<()> { let mut remote_stream = each_addr(remote_addr, TcpStream::connect) .await @@ -57,7 +59,7 @@ pub async fn handle_local_stream< }, }; let msg = request.encode().context(EncodeSubcribeReqSnafu)?; - let session = ClientHeaderSession::from_process() + let session = ClientHeaderSession::new_v2(&credential) .context(CreateHeaderToolSnafu { action: "session" })?; match tokio::time::timeout(timeout, session.write_initial(&mut remote_stream, &msg)).await { Ok(result) => result.context(WriteSubcribeReqSnafu)?, From 0d78941ce1be6bdeeaadd460b986e4ef56077b7e Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 06:30:19 +0800 Subject: [PATCH 41/74] Pin registration workers and replace canceled renew leases Local server workers now capture the process credential once and reuse it on reconnect, matching the client listener. A renewal that finds its lease already canceled after WAL sync installs a fresh lease so the extended credential stays usable. Worker debug logs print only the credential key id. --- CHANGELOG.md | 2 ++ src/common/auth/actor.rs | 21 ++++++++++++--------- src/common/auth/tests.rs | 35 +++++++++++++++++++++++++++++++++++ src/local/server/mod.rs | 40 ++++++++++++++++++++++++++++++---------- 4 files changed, 79 insertions(+), 19 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ed57049..b1dbcb4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -36,6 +36,8 @@ All notable changes to this project will be documented in this file. - Passed `--use-machine-msg-header-key` only when `admin.key` is missing, so a container restart with a persistent auth volume does not fail after first boot. - Pinned each local client listener to the credential captured at start, so a later Flutter config change cannot switch an existing port onto another tenant. - Garbage-collected expired and revoked high-slot entries while keeping their generations so shrinking capacity and running `key gc` can reclaim them. +- Pinned local registration workers to the credential captured at start, matching the client listener. +- Replaced a lease that was canceled while a renewal WAL record was syncing, so a successful renew does not keep a dead cancellation token. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/src/common/auth/actor.rs b/src/common/auth/actor.rs index 5bf8920..25228bf 100644 --- a/src/common/auth/actor.rs +++ b/src/common/auth/actor.rs @@ -566,16 +566,19 @@ fn actor_renew( true, )); } - let lease = slot.lease.upgrade().ok_or_else(|| { - AuthFailure::new( - "temporary_key_inactive", - "temporary key lease is no longer active", - true, - ) - })?; slot.expires_at = expires_at; - lease.expires_at.store(expires_at, Ordering::Release); - lease.wheel_version.fetch_add(1, Ordering::AcqRel); + let lease = match slot.lease.upgrade() { + Some(lease) if !lease.cancellation_token().is_cancelled() => { + lease.expires_at.store(expires_at, Ordering::Release); + lease.wheel_version.fetch_add(1, Ordering::AcqRel); + lease + } + _ => { + let lease = Arc::new(AuthLease::new(key_id, expires_at)); + slot.lease = Arc::downgrade(&lease); + lease + } + }; wheel.release(key_id); wheel.insert(lease); drop(slots); diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index b2d34b7..2306923 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -592,6 +592,41 @@ async fn issue_renew_revoke_and_persist() { let _ = std::fs::remove_dir_all(state_dir); } +#[tokio::test] +async fn renew_replaces_a_lease_canceled_during_persistence() { + let state_dir = temp_state_dir("renew-canceled-lease"); + let admin_key = *b"0123456789abcdefghijklmnopqrstuv"; + let config = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 4, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + let runtime = AuthRuntime::start(admin_key, config).await.unwrap(); + let admin = authenticate_for_test(&runtime, 0).unwrap(); + let issued = runtime + .issue(&admin, Duration::from_secs(60), Some("renew".to_string())) + .await + .unwrap(); + let context = authenticate_for_test(&runtime, issued.metadata.key_id).unwrap(); + let canceled = context.cancellation_token().unwrap(); + canceled.cancel(); + assert!(canceled.is_cancelled()); + + let renewed = runtime + .renew(&admin, issued.metadata.key_id, Duration::from_secs(120)) + .await + .unwrap(); + assert_eq!(renewed.metadata.key_id, issued.metadata.key_id); + let restored = authenticate_for_test(&runtime, issued.metadata.key_id).unwrap(); + assert!(!restored.cancellation_token().unwrap().is_cancelled()); + assert!(canceled.is_cancelled()); + + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + let _ = std::fs::remove_dir_all(state_dir); +} + #[tokio::test] async fn reset_rotates_instance_and_prevents_old_key_id_reuse() { let state_dir = temp_state_dir("auth-reset"); diff --git a/src/local/server/mod.rs b/src/local/server/mod.rs index 1273de2..3b06d91 100644 --- a/src/local/server/mod.rs +++ b/src/local/server/mod.rs @@ -135,13 +135,28 @@ pub struct ServerTunnelOptions { pub force_namespace: bool, } -#[derive(Clone, Debug)] +#[derive(Clone)] struct ServerCliRunConfig { local_addr: A, remote_addr: A, key: Arc, options: ServerTunnelOptions, worker_index: usize, + credential: Credential, +} + +impl Debug for ServerCliRunConfig { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("ServerCliRunConfig") + .field("local_addr", &self.local_addr) + .field("remote_addr", &self.remote_addr) + .field("key", &self.key) + .field("options", &self.options) + .field("worker_index", &self.worker_index) + .field("credential_key_id", &self.credential.key_id()) + .finish() + } } /// Where a stream request should connect, and how. @@ -325,14 +340,24 @@ async fn run_server_side_cli_worker( LocalStream::Item: StreamForward, A: ToSocketAddrs + Debug + Copy + Send + 'static, { + let mut retry_backoff = RetryBackoff::default(); + let credential = loop { + match get_process_credential() { + Ok(credential) => break credential, + Err(error) => { + tracing::error!("load registration credential failed: {error}"); + tokio::time::sleep(retry_backoff.next_delay()).await; + } + } + }; let run_config = ServerCliRunConfig { local_addr, remote_addr, key: key.clone(), options, worker_index, + credential, }; - let mut retry_backoff = RetryBackoff::default(); loop { let status = if let Err(status) = run_server_side_cli_inner::( &mut retry_backoff, @@ -402,6 +427,7 @@ where force_namespace, }, worker_index, + credential, } = config; let local_addr = match got_one_socket_addr(local_addr).await { Ok(addr) => addr, @@ -446,14 +472,8 @@ where ); // Start registration with a protocol-v2 first frame. The session is reused for all - // subsequent control messages on this TCP connection. - let credential = match get_process_credential() { - Ok(credential) => credential, - Err(error) => { - tracing::error!("load registration credential failed: {error}"); - return Err(Status::ConnectRemote); - } - }; + // subsequent control messages on this TCP connection. The credential is pinned + // when the worker starts so a later process-key change cannot retarget reconnects. let session = match ClientHeaderSession::new_v2(&credential) { Ok(session) => session, Err(error) => { From 8bd52240c70799bd8145c8b20541adb3261793c5 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 06:51:49 +0800 Subject: [PATCH 42/74] Bind tunneled checksums and finish actor/UI credential shutdown Relay forwarding now checksums each hop with that session's authenticated key. Shutdown waits for the auth actor after connection tasks drop. UI workers capture the process credential before spawn and reuse it for status probes. Registration lease probes use the pinned credential, and root rotation still completes when admin.key already matches the new snapshot. --- CHANGELOG.md | 2 + src/common/auth.rs | 4 + src/common/auth/actor.rs | 19 ++++- src/common/auth/persistence.rs | 2 +- src/common/auth/runtime.rs | 35 ++++++--- src/common/message/forward.rs | 28 +++++++ src/common/message/mod.rs | 5 ++ src/common/message/secure.rs | 4 + src/local/client/mod.rs | 47 +++++++++-- src/local/server/mod.rs | 77 ++++++++++++++++--- src/pb_server/client.rs | 75 +++++++++++++----- src/pb_server/runtime.rs | 1 + ui/native/pb_mapper_ffi/src/state.rs | 21 +++-- .../pb_mapper_ffi/src/state/configuration.rs | 4 + ui/native/pb_mapper_ffi/src/state/runtime.rs | 22 +++++- ui/native/pb_mapper_ffi/src/state/status.rs | 6 +- 16 files changed, 289 insertions(+), 63 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b1dbcb4..d4bd7e9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -38,6 +38,8 @@ All notable changes to this project will be documented in this file. - Garbage-collected expired and revoked high-slot entries while keeping their generations so shrinking capacity and running `key gc` can reclaim them. - Pinned local registration workers to the credential captured at start, matching the client listener. - Replaced a lease that was canceled while a renewal WAL record was syncing, so a successful renew does not keep a dead cancellation token. +- Awaited the authentication actor on relay shutdown, used pinned credentials for registration probes and UI tunnel workers/status checks, and finished in-memory root rotation when `admin.key` already matched the new snapshot. +- Bound relay tunneled-frame checksums to each hop's authenticated session key instead of the process administrator key. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/src/common/auth.rs b/src/common/auth.rs index 8a2beff..d393537 100644 --- a/src/common/auth.rs +++ b/src/common/auth.rs @@ -495,6 +495,7 @@ pub struct AuthRuntime { command_tx: mpsc::Sender, config: AuthConfig, _state_lock: Arc, + actor: Arc>>>, } #[derive(Clone, Debug, Serialize, Deserialize)] @@ -608,6 +609,9 @@ enum AuthCommand { detail: Option, response: oneshot::Sender>, }, + Shutdown { + response: oneshot::Sender<()>, + }, } mod runtime; diff --git a/src/common/auth/actor.rs b/src/common/auth/actor.rs index 25228bf..87cc9b3 100644 --- a/src/common/auth/actor.rs +++ b/src/common/auth/actor.rs @@ -269,6 +269,12 @@ pub(super) async fn run_auth_actor( }); let _ = response.send(result); } + AuthCommand::Shutdown { response } => { + admin_lease.cancellation.cancel(); + cancel_all_temporary_leases(&inner); + let _ = response.send(()); + break; + } } } } @@ -815,9 +821,16 @@ fn actor_rotate_root( .and_then(|()| write_snapshot_and_truncate_wal(config, &new_key, &snapshot)) .and_then(|()| write_admin_key(&config.state_dir, &new_key_string)) { - inner.safe_mode.store(true, Ordering::Release); - cancel_all_temporary_leases(inner); - return Err(error); + if !key_matches_existing_snapshot(Some(&config.state_dir), &new_key_string) { + inner.safe_mode.store(true, Ordering::Release); + cancel_all_temporary_leases(inner); + return Err(error); + } + tracing::warn!( + event = "administrator_key_rotate_finalized_after_sync_error", + error = %error, + "admin.key replacement reported an error, but the new snapshot already decrypts with the new key; finishing in-memory rotation" + ); } let _ = std::fs::remove_file(&next_key_path); push_audit_record(inner, rotate_audit); diff --git a/src/common/auth/persistence.rs b/src/common/auth/persistence.rs index 76d0ea4..31044e3 100644 --- a/src/common/auth/persistence.rs +++ b/src/common/auth/persistence.rs @@ -979,7 +979,7 @@ pub fn write_admin_key_file(path: &Path, key: &str, force: bool) -> Result<(), A atomic_write(path, format!("{key}\n").as_bytes(), 0o600) } -fn key_matches_existing_snapshot(state_dir: Option<&Path>, key: &str) -> bool { +pub(super) fn key_matches_existing_snapshot(state_dir: Option<&Path>, key: &str) -> bool { let Some(state_dir) = state_dir else { return false; }; diff --git a/src/common/auth/runtime.rs b/src/common/auth/runtime.rs index 4cb7b9d..0804fe4 100644 --- a/src/common/auth/runtime.rs +++ b/src/common/auth/runtime.rs @@ -174,24 +174,41 @@ impl AuthRuntime { audit_records: RwLock::new(audit_records), }); let (command_tx, command_rx) = mpsc::channel(256); + let actor = tokio::spawn(run_auth_actor( + inner.clone(), + admin_lease, + command_rx, + config.clone(), + AuthActorState::new(cold, wheel, admin_replays, admin_replay_order), + state_lock.clone(), + )); let runtime = Self { inner: Arc::downgrade(&inner), command_tx, config: config.clone(), _state_lock: state_lock.clone(), + actor: Arc::new(std::sync::Mutex::new(Some(actor))), }; - - tokio::spawn(run_auth_actor( - inner, - admin_lease, - command_rx, - config, - AuthActorState::new(cold, wheel, admin_replays, admin_replay_order), - state_lock, - )); Ok(runtime) } + pub async fn shutdown_actor(&self) { + let (response, receiver) = oneshot::channel(); + let _ = self + .command_tx + .send(AuthCommand::Shutdown { response }) + .await; + let _ = receiver.await; + let handle = self + .actor + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .take(); + if let Some(handle) = handle { + let _ = handle.await; + } + } + pub fn config(&self) -> &AuthConfig { &self.config } diff --git a/src/common/message/forward.rs b/src/common/message/forward.rs index d25cd82..4701fa5 100644 --- a/src/common/message/forward.rs +++ b/src/common/message/forward.rs @@ -96,6 +96,12 @@ impl<'a, T: AsyncReadExt + Unpin + Send> NormalDatagramReader<'a, T> { reader: NormalMessageReader::new(reader), } } + + pub fn with_checksum_key(self, key: AesKeyType) -> Self { + Self { + reader: self.reader.with_checksum_key(key), + } + } } impl<'a, T: AsyncReadExt + Unpin + Send> DatagramReader for NormalDatagramReader<'a, T> { @@ -142,6 +148,12 @@ impl<'a, T: AsyncWriteExt + Unpin + Send> NormalDatagramWriter<'a, T> { writer: NormalMessageWriter::new(writer), } } + + pub fn with_checksum_key(self, key: AesKeyType) -> Self { + Self { + writer: self.writer.with_checksum_key(key), + } + } } impl<'a, T: AsyncWriteExt + Unpin + Send> DatagramWriter for NormalDatagramWriter<'a, T> { @@ -158,6 +170,10 @@ impl<'a, T: AsyncReadExt + Send + Unpin, D: Decryptor> CodecForwardReader<'a, T, pub fn new(reader: &'a mut T, decryptor: D) -> Self { Self(CodecMessageReader::new(reader, decryptor)) } + + pub fn with_checksum_key(self, key: AesKeyType) -> Self { + Self(self.0.with_checksum_key(key)) + } } impl<'a, T: AsyncReadExt + Send + Unpin, D: Decryptor> ForwardReader @@ -176,6 +192,10 @@ impl<'a, T: AsyncReadExt + Send + Unpin, D: Decryptor> CodecDatagramReader<'a, T pub fn new(reader: &'a mut T, decryptor: D) -> Self { Self(CodecMessageReader::new(reader, decryptor)) } + + pub fn with_checksum_key(self, key: AesKeyType) -> Self { + Self(self.0.with_checksum_key(key)) + } } impl<'a, T: AsyncReadExt + Send + Unpin, D: Decryptor> DatagramReader @@ -195,6 +215,10 @@ impl<'a, T: AsyncWriteExt + Send + Unpin, E: Encryptor> CodecForwardWriter<'a, T pub fn new(writer: &'a mut T, encryptor: E) -> Self { Self(CodecMessageWriter::new(writer, encryptor)) } + + pub fn with_checksum_key(self, key: AesKeyType) -> Self { + Self(self.0.with_checksum_key(key)) + } } impl<'a, T: AsyncWriteExt + Send + Unpin, E: Encryptor> ForwardWriter @@ -218,6 +242,10 @@ impl<'a, T: AsyncWriteExt + Send + Unpin, E: Encryptor> CodecDatagramWriter<'a, pub fn new(writer: &'a mut T, encryptor: E) -> Self { Self(CodecMessageWriter::new(writer, encryptor)) } + + pub fn with_checksum_key(self, key: AesKeyType) -> Self { + Self(self.0.with_checksum_key(key)) + } } impl<'a, T: AsyncWriteExt + Send + Unpin, E: Encryptor> DatagramWriter diff --git a/src/common/message/mod.rs b/src/common/message/mod.rs index f41867d..87ae650 100644 --- a/src/common/message/mod.rs +++ b/src/common/message/mod.rs @@ -241,6 +241,11 @@ impl<'a, T: AsyncWriteExt + Unpin> NormalMessageWriter<'a, T> { } } + pub fn with_checksum_key(mut self, key: AesKeyType) -> Self { + self.checksum_key = Some(key); + self + } + async fn write_msg_inner(&mut self, msg: &[u8]) -> Result<()> { set_msg_len( &mut self.writer, diff --git a/src/common/message/secure.rs b/src/common/message/secure.rs index a86183b..0c53510 100644 --- a/src/common/message/secure.rs +++ b/src/common/message/secure.rs @@ -188,6 +188,10 @@ impl ServerHeaderSession { self.protocol } + pub fn framing_key(&self) -> AesKeyType { + self.legacy_key + } + pub fn key_id(&self) -> u64 { self.context .as_ref() diff --git a/src/local/client/mod.rs b/src/local/client/mod.rs index a43e1f9..db7a451 100644 --- a/src/local/client/mod.rs +++ b/src/local/client/mod.rs @@ -65,6 +65,7 @@ pub async fn run_client_side_cli_scoped( + local_addr: A, + remote_addr: A, + key: Arc, + keep_alive: bool, + status_callback: Option, + credential: crate::common::checksum::Credential, +) where + ::Item: StreamForward, +{ + run_client_side_cli_with_callback_scoped::( + local_addr, + remote_addr, + key, + keep_alive, + None, + status_callback, + Some(credential), ) .await } @@ -99,6 +126,7 @@ pub async fn run_client_side_cli_with_callback_scoped< keep_alive: bool, namespace: Option, status_callback: Option, + pinned_credential: Option, ) where ::Item: StreamForward, { @@ -124,15 +152,18 @@ pub async fn run_client_side_cli_with_callback_scoped< return; } }; - let credential = match get_process_credential() { - Ok(credential) => credential, - Err(e) => { - tracing::error!("load client credential failed: {e}"); - if let Some(ref callback) = status_callback { - callback("failed"); + let credential = match pinned_credential { + Some(credential) => credential, + None => match get_process_credential() { + Ok(credential) => credential, + Err(e) => { + tracing::error!("load client credential failed: {e}"); + if let Some(ref callback) = status_callback { + callback("failed"); + } + return; } - return; - } + }, }; let mut retry_backoff = RetryBackoff::default(); diff --git a/src/local/server/mod.rs b/src/local/server/mod.rs index 3b06d91..b1cb959 100644 --- a/src/local/server/mod.rs +++ b/src/local/server/mod.rs @@ -186,18 +186,20 @@ async fn probe_remote_registration( key: Arc, registration: ControlRegistration, namespace: Option, + credential: Credential, ) -> RegistrationProbeResult { let timeout = registration_probe_timeout(); let result = tokio::time::timeout(timeout, async { let mut stream = each_addr(remote_addr, TcpStream::connect) .await .map_err(|e| format!("connect remote status stream failed: {e}"))?; - crate::local::client::status::get_status_scoped( + crate::local::client::status::get_status_with_credential( &mut stream, PbConnStatusReq::Service { key: key.to_string(), }, namespace, + &credential, ) .await .map_err(|e| { @@ -259,6 +261,29 @@ pub async fn run_server_side_cli( .await } +pub async fn run_server_side_cli_with_pinned_credential( + local_addr: A, + remote_addr: A, + key: Arc, + options: ServerTunnelOptions, + status_callback: Option, + credential: Credential, +) where + LocalStream: StreamProvider + Send + 'static, + LocalStream::Item: StreamForward, + A: ToSocketAddrs + Debug + Copy, +{ + run_server_side_cli_pool::( + local_addr, + remote_addr, + key, + options, + status_callback, + Some(credential), + ) + .await; +} + pub async fn run_server_side_cli_with_callback( local_addr: A, remote_addr: A, @@ -269,6 +294,29 @@ pub async fn run_server_side_cli_with_callback( LocalStream: StreamProvider + Send + 'static, LocalStream::Item: StreamForward, A: ToSocketAddrs + Debug + Copy, +{ + run_server_side_cli_pool::( + local_addr, + remote_addr, + key, + options, + status_callback, + None, + ) + .await; +} + +async fn run_server_side_cli_pool( + local_addr: A, + remote_addr: A, + key: Arc, + options: ServerTunnelOptions, + status_callback: Option, + pinned_credential: Option, +) where + LocalStream: StreamProvider + Send + 'static, + LocalStream::Item: StreamForward, + A: ToSocketAddrs + Debug + Copy, { let local_addr = match got_one_socket_addr(local_addr).await { Ok(addr) => addr, @@ -296,25 +344,27 @@ pub async fn run_server_side_cli_with_callback( for worker_index in 1..pool_size { let worker_key = key.clone(); worker_handles.push(tokio::spawn(async move { - run_server_side_cli_worker::( + run_server_side_cli_worker_with_credential::( local_addr, remote_addr, worker_key, options, None, worker_index, + pinned_credential, ) .await; })); } } - run_server_side_cli_worker::( + run_server_side_cli_worker_with_credential::( local_addr, remote_addr, key, options, status_callback, 0, + pinned_credential, ) .await; for handle in worker_handles { @@ -328,27 +378,31 @@ pub async fn run_server_side_cli_with_callback( } } -async fn run_server_side_cli_worker( +async fn run_server_side_cli_worker_with_credential( local_addr: A, remote_addr: A, key: Arc, options: ServerTunnelOptions, status_callback: Option, worker_index: usize, + pinned_credential: Option, ) where LocalStream: StreamProvider + Send + 'static, LocalStream::Item: StreamForward, A: ToSocketAddrs + Debug + Copy + Send + 'static, { let mut retry_backoff = RetryBackoff::default(); - let credential = loop { - match get_process_credential() { - Ok(credential) => break credential, - Err(error) => { - tracing::error!("load registration credential failed: {error}"); - tokio::time::sleep(retry_backoff.next_delay()).await; + let credential = match pinned_credential { + Some(credential) => credential, + None => loop { + match get_process_credential() { + Ok(credential) => break credential, + Err(error) => { + tracing::error!("load registration credential failed: {error}"); + tokio::time::sleep(retry_backoff.next_delay()).await; + } } - } + }, }; let run_config = ServerCliRunConfig { local_addr, @@ -734,6 +788,7 @@ where probe_key, registration, namespace, + credential, ) .await; let _ = probe_tx.send(result); diff --git a/src/pb_server/client.rs b/src/pb_server/client.rs index faf8e70..3b5dc39 100644 --- a/src/pb_server/client.rs +++ b/src/pb_server/client.rs @@ -28,7 +28,7 @@ use crate::pb_server::error::{ ClientConnCreateHeaderToolSnafu, ClientConnEncodeStreamRespSnafu, ClientConnWriteStreamRespSnafu, }; -use crate::{create_component, snafu_error_get_or_return_ok, start_forward_with_codec_key}; +use crate::snafu_error_get_or_return_ok; /// Ensure that client-side connections are properly deregistered before a normal connection is /// disconnected or an exception occurs @@ -220,6 +220,8 @@ pub async fn handle_client_conn( })?; } + let client_framing = session.framing_key(); + let server_framing = server_session.framing_key(); if is_datagram { match codec_key { Some(key) => { @@ -227,44 +229,77 @@ pub async fn handle_client_conn( CodecDatagramReader::new( &mut client_reader, snafu_error_get_or_return_ok!(get_decodec(&key)), - ), + ) + .with_checksum_key(client_framing), CodecDatagramWriter::new( &mut client_writer, snafu_error_get_or_return_ok!(get_encodec(&key)), - ), + ) + .with_checksum_key(client_framing), CodecDatagramReader::new( &mut server_reader, snafu_error_get_or_return_ok!(get_decodec(&key)), - ), + ) + .with_checksum_key(server_framing), CodecDatagramWriter::new( &mut server_writer, snafu_error_get_or_return_ok!(get_encodec(&key)), - ), + ) + .with_checksum_key(server_framing), ) .await; } None => { start_datagram_forward( - NormalDatagramReader::new(&mut client_reader), - NormalDatagramWriter::new(&mut client_writer), - NormalDatagramReader::new(&mut server_reader), - NormalDatagramWriter::new(&mut server_writer), + NormalDatagramReader::new(&mut client_reader) + .with_checksum_key(client_framing), + NormalDatagramWriter::new(&mut client_writer) + .with_checksum_key(client_framing), + NormalDatagramReader::new(&mut server_reader) + .with_checksum_key(server_framing), + NormalDatagramWriter::new(&mut server_writer) + .with_checksum_key(server_framing), ) .await; } } } else { - start_forward_with_codec_key!( - codec_key, - &mut client_reader, - &mut client_writer, - &mut server_reader, - &mut server_writer, - true, - true, - true, - true - ); + match codec_key { + Some(key) => { + start_forward( + CodecForwardReader::new( + &mut client_reader, + snafu_error_get_or_return_ok!(get_decodec(&key)), + ) + .with_checksum_key(client_framing), + CodecForwardWriter::new( + &mut client_writer, + snafu_error_get_or_return_ok!(get_encodec(&key)), + ) + .with_checksum_key(client_framing), + CodecForwardReader::new( + &mut server_reader, + snafu_error_get_or_return_ok!(get_decodec(&key)), + ) + .with_checksum_key(server_framing), + CodecForwardWriter::new( + &mut server_writer, + snafu_error_get_or_return_ok!(get_encodec(&key)), + ) + .with_checksum_key(server_framing), + ) + .await; + } + None => { + start_forward( + NormalForwardReader::new(&mut client_reader), + NormalForwardWriter::new(&mut client_writer), + NormalForwardReader::new(&mut server_reader), + NormalForwardWriter::new(&mut server_writer), + ) + .await; + } + } } Ok(()) diff --git a/src/pb_server/runtime.rs b/src/pb_server/runtime.rs index 7086dfa..eb192ab 100644 --- a/src/pb_server/runtime.rs +++ b/src/pb_server/runtime.rs @@ -941,6 +941,7 @@ pub async fn run_server_on_listener( .chain(status_forward_handle), ) .await; + security.auth().shutdown_actor().await; tracing::info!("Server shutdown completed"); Ok(()) } diff --git a/ui/native/pb_mapper_ffi/src/state.rs b/ui/native/pb_mapper_ffi/src/state.rs index a3dc1fd..bb2174b 100644 --- a/ui/native/pb_mapper_ffi/src/state.rs +++ b/ui/native/pb_mapper_ffi/src/state.rs @@ -24,13 +24,15 @@ use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; use pb_mapper::common::auth::{AuthConfig, AuthRuntime}; -use pb_mapper::common::checksum::{parse_credential, set_process_msg_header_key}; +use pb_mapper::common::checksum::{ + get_process_credential, parse_credential, set_process_msg_header_key, Credential, +}; use pb_mapper::common::config::{get_pb_mapper_server_async, get_sockaddr_async}; use pb_mapper::common::message::command::{PbConnStatusReq, PbConnStatusResp}; -use pb_mapper::local::client::status::get_status; -use pb_mapper::local::client::{run_client_side_cli_with_callback, ClientStatusCallback}; +use pb_mapper::local::client::status::{get_status, get_status_with_credential}; +use pb_mapper::local::client::{run_client_side_cli_with_pinned_credential, ClientStatusCallback}; use pb_mapper::local::server::{ - run_server_side_cli_with_callback, ServerTunnelOptions, StatusCallback, + run_server_side_cli_with_pinned_credential, ServerTunnelOptions, StatusCallback, }; use pb_mapper::pb_server::{run_server_on_listener, ServerStatusInfo}; use pb_mapper::utils::addr::each_addr; @@ -58,6 +60,7 @@ struct StatusCacheEntry { async fn check_service_with_get_status( server_addr: &str, service_key: &str, + credential: Option, ) -> Result { let addr = get_sockaddr_async(server_addr) .await @@ -66,7 +69,13 @@ async fn check_service_with_get_status( match TcpStreamProvider::from_addr(addr).await { Ok(mut stream) => { let status_req = PbConnStatusReq::Keys; - match get_status(&mut stream, status_req).await { + let status = match credential { + Some(credential) => { + get_status_with_credential(&mut stream, status_req, None, &credential).await + } + None => get_status(&mut stream, status_req).await, + }; + match status { Ok(status_resp) => match status_resp { PbConnStatusResp::Keys(keys) => { if keys.contains(&service_key.to_string()) { @@ -445,6 +454,8 @@ pub struct PbMapperState { active_connections: Arc>>, service_handles: HashMap>, client_handles: HashMap>, + service_credentials: HashMap, + client_credentials: HashMap, config: AppConfig, config_dir: PathBuf, app_directory_path: Option, diff --git a/ui/native/pb_mapper_ffi/src/state/configuration.rs b/ui/native/pb_mapper_ffi/src/state/configuration.rs index 2ad22b8..4571517 100644 --- a/ui/native/pb_mapper_ffi/src/state/configuration.rs +++ b/ui/native/pb_mapper_ffi/src/state/configuration.rs @@ -54,6 +54,8 @@ impl PbMapperState { active_connections: Arc::new(RwLock::new(HashMap::new())), service_handles: HashMap::new(), client_handles: HashMap::new(), + service_credentials: HashMap::new(), + client_credentials: HashMap::new(), config: AppConfig::default(), config_dir: config_dir.clone(), app_directory_path: app_directory_path.clone(), @@ -89,6 +91,8 @@ impl PbMapperState { active_connections: Arc::new(RwLock::new(HashMap::new())), service_handles: HashMap::new(), client_handles: HashMap::new(), + service_credentials: HashMap::new(), + client_credentials: HashMap::new(), config, config_dir, app_directory_path, diff --git a/ui/native/pb_mapper_ffi/src/state/runtime.rs b/ui/native/pb_mapper_ffi/src/state/runtime.rs index 5decac0..dd762dd 100644 --- a/ui/native/pb_mapper_ffi/src/state/runtime.rs +++ b/ui/native/pb_mapper_ffi/src/state/runtime.rs @@ -125,10 +125,12 @@ impl PbMapperState { for (_, handle) in self.service_handles.drain() { handle.abort(); } + self.service_credentials.clear(); for (_, handle) in self.client_handles.drain() { handle.abort(); } + self.client_credentials.clear(); self.registered_services.write().await.clear(); self.active_connections.write().await.clear(); @@ -151,6 +153,7 @@ impl PbMapperState { remote_sock_addr, } = commit; + self.service_credentials.remove(&service_key); if let Some(previous) = self.service_handles.remove(&service_key) { tracing::warn!( "Service '{service_key}' is already registered, replacing existing handle" @@ -189,9 +192,12 @@ impl PbMapperState { ); }); + let credential = get_process_credential().map_err(CtlError::invalid_argument)?; + self.service_credentials + .insert(service_key.clone(), credential); let handle = if protocol.to_uppercase() == "TCP" { tokio::spawn(async move { - let _ = run_server_side_cli_with_callback::( + let _ = run_server_side_cli_with_pinned_credential::( local_sock_addr, remote_sock_addr, key_clone.into(), @@ -203,12 +209,13 @@ impl PbMapperState { force_namespace: false, }, Some(callback), + credential, ) .await; }) } else { tokio::spawn(async move { - let _ = run_server_side_cli_with_callback::( + let _ = run_server_side_cli_with_pinned_credential::( local_sock_addr, remote_sock_addr, key_clone.into(), @@ -220,6 +227,7 @@ impl PbMapperState { force_namespace: false, }, Some(callback), + credential, ) .await; }) @@ -300,6 +308,7 @@ impl PbMapperState { remote_sock_addr, } = commit; + self.client_credentials.remove(&service_key); if let Some(previous) = self.client_handles.remove(&service_key) { tracing::warn!( "Client for service '{service_key}' is already connected, replacing handle" @@ -328,25 +337,30 @@ impl PbMapperState { }) }; + let credential = get_process_credential().map_err(CtlError::invalid_argument)?; + self.client_credentials + .insert(service_key.clone(), credential); let handle = if protocol_upper == "TCP" { tokio::spawn(async move { - run_client_side_cli_with_callback::( + run_client_side_cli_with_pinned_credential::( local_sock_addr, remote_sock_addr, key_clone.into(), enable_keep_alive, Some(status_callback), + credential, ) .await; }) } else { tokio::spawn(async move { - run_client_side_cli_with_callback::( + run_client_side_cli_with_pinned_credential::( local_sock_addr, remote_sock_addr, key_clone.into(), enable_keep_alive, Some(status_callback), + credential, ) .await; }) diff --git a/ui/native/pb_mapper_ffi/src/state/status.rs b/ui/native/pb_mapper_ffi/src/state/status.rs index 489e567..9cf93be 100644 --- a/ui/native/pb_mapper_ffi/src/state/status.rs +++ b/ui/native/pb_mapper_ffi/src/state/status.rs @@ -364,11 +364,12 @@ impl PbMapperState { let cache = self.service_status_cache.clone(); let refreshing = self.service_status_refreshing.clone(); let key = service_key.to_string(); + let credential = self.service_credentials.get(service_key).copied(); tokio::spawn(async move { let result = tokio::time::timeout( STATUS_REFRESH_TIMEOUT, - check_service_with_get_status(&server_addr, &key), + check_service_with_get_status(&server_addr, &key, credential), ) .await; @@ -427,11 +428,12 @@ impl PbMapperState { let cache = self.client_status_cache.clone(); let refreshing = self.client_status_refreshing.clone(); let key = service_key.to_string(); + let credential = self.client_credentials.get(service_key).copied(); tokio::spawn(async move { let result = tokio::time::timeout( STATUS_REFRESH_TIMEOUT, - check_service_with_get_status(&server_addr, &key), + check_service_with_get_status(&server_addr, &key, credential), ) .await; From 33c88116f7aae01ea988f779f5e5e3767a64cbd4 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 07:02:41 +0800 Subject: [PATCH 43/74] Bind local tunneled frames to the pinned credential key UDP and codec forwarding now checksum remote frames with the credential captured when the client or registration worker started, so a later process-key change cannot desynchronize payloads from the relay. --- CHANGELOG.md | 1 + src/common/message/forward.rs | 97 ++++++++++++++++++++++++++--------- src/local/client/stream.rs | 1 + src/local/server/stream.rs | 1 + 4 files changed, 77 insertions(+), 23 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d4bd7e9..823b2de 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -40,6 +40,7 @@ All notable changes to this project will be documented in this file. - Replaced a lease that was canceled while a renewal WAL record was syncing, so a successful renew does not keep a dead cancellation token. - Awaited the authentication actor on relay shutdown, used pinned credentials for registration probes and UI tunnel workers/status checks, and finished in-memory root rotation when `admin.key` already matched the new snapshot. - Bound relay tunneled-frame checksums to each hop's authenticated session key instead of the process administrator key. +- Bound local UDP and codec tunnels to the pinned credential's checksum key so a later process-key change cannot desynchronize framed payloads. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/src/common/message/forward.rs b/src/common/message/forward.rs index 4701fa5..3ae3084 100644 --- a/src/common/message/forward.rs +++ b/src/common/message/forward.rs @@ -14,12 +14,8 @@ use super::{ }; use crate::common::checksum::AesKeyType; use crate::common::config::duration_from_env; -use crate::common::message::{get_decodec, get_encodec}; +use crate::snafu_error_get_or_return_ok; use crate::utils::codec::{Decryptor, Encryptor}; -use crate::{ - create_component, snafu_error_get_or_return_ok, start_datagram_forward_with_codec_key, - start_forward_with_codec_key, -}; use uni_stream::stream::{StreamSplit, TcpStreamImpl, UdpStreamImpl}; use uni_stream::udp::{UdpStreamReadHalf, UdpStreamWriteHalf}; @@ -558,6 +554,7 @@ impl DatagramWriter for UdpStreamWriteHalf<'_> { pub trait StreamForward: StreamSplit + Sized { fn forward_local_to_remote<'a, R, W>( codec_key: Option, + framing_key: AesKeyType, local_reader: Self::ReaderRef<'a>, local_writer: Self::WriterRef<'a>, remote_reader: R, @@ -571,6 +568,7 @@ pub trait StreamForward: StreamSplit + Sized { impl StreamForward for TcpStreamImpl { fn forward_local_to_remote<'a, R, W>( codec_key: Option, + framing_key: AesKeyType, local_reader: Self::ReaderRef<'a>, local_writer: Self::WriterRef<'a>, remote_reader: R, @@ -585,17 +583,40 @@ impl StreamForward for TcpStreamImpl { let mut local_writer = local_writer; let mut remote_reader = remote_reader; let mut remote_writer = remote_writer; - start_forward_with_codec_key!( - codec_key, - &mut local_reader, - &mut local_writer, - &mut remote_reader, - &mut remote_writer, - false, - false, - true, - true - ); + match codec_key { + Some(key) => { + start_forward( + NormalForwardReader::new(&mut local_reader), + NormalForwardWriter::new(&mut local_writer), + CodecForwardReader::new( + &mut remote_reader, + snafu_error_get_or_return_ok!( + super::get_decodec(&key), + "failed to create decoder when remote forward" + ), + ) + .with_checksum_key(framing_key), + CodecForwardWriter::new( + &mut remote_writer, + snafu_error_get_or_return_ok!( + super::get_encodec(&key), + "failed to create encoder when remote forward" + ), + ) + .with_checksum_key(framing_key), + ) + .await; + } + None => { + start_forward( + NormalForwardReader::new(&mut local_reader), + NormalForwardWriter::new(&mut local_writer), + NormalForwardReader::new(&mut remote_reader), + NormalForwardWriter::new(&mut remote_writer), + ) + .await; + } + } Ok(()) }) } @@ -604,6 +625,7 @@ impl StreamForward for TcpStreamImpl { impl StreamForward for UdpStreamImpl { fn forward_local_to_remote<'a, R, W>( codec_key: Option, + framing_key: AesKeyType, local_reader: Self::ReaderRef<'a>, local_writer: Self::WriterRef<'a>, remote_reader: R, @@ -616,13 +638,42 @@ impl StreamForward for UdpStreamImpl { Box::pin(async move { let mut remote_reader = remote_reader; let mut remote_writer = remote_writer; - start_datagram_forward_with_codec_key!( - codec_key, - local_reader, - local_writer, - &mut remote_reader, - &mut remote_writer - ); + match codec_key { + Some(key) => { + start_datagram_forward( + local_reader, + local_writer, + CodecDatagramReader::new( + &mut remote_reader, + snafu_error_get_or_return_ok!( + super::get_decodec(&key), + "failed to create decoder when datagram forward" + ), + ) + .with_checksum_key(framing_key), + CodecDatagramWriter::new( + &mut remote_writer, + snafu_error_get_or_return_ok!( + super::get_encodec(&key), + "failed to create encoder when datagram forward" + ), + ) + .with_checksum_key(framing_key), + ) + .await; + } + None => { + start_datagram_forward( + local_reader, + local_writer, + NormalDatagramReader::new(&mut remote_reader) + .with_checksum_key(framing_key), + NormalDatagramWriter::new(&mut remote_writer) + .with_checksum_key(framing_key), + ) + .await; + } + } Ok(()) }) } diff --git a/src/local/client/stream.rs b/src/local/client/stream.rs index 9815f75..03062b0 100644 --- a/src/local/client/stream.rs +++ b/src/local/client/stream.rs @@ -107,6 +107,7 @@ pub async fn handle_local_stream< snafu_error_handle!( ::forward_local_to_remote( codec_key, + *credential.key(), client_reader, client_writer, server_reader, diff --git a/src/local/server/stream.rs b/src/local/server/stream.rs index 5c9186c..5224310 100644 --- a/src/local/server/stream.rs +++ b/src/local/server/stream.rs @@ -137,6 +137,7 @@ where snafu_error_handle!( ::forward_local_to_remote( codec_key, + *credential.key(), server_reader, server_writer, client_reader, From 566b21397d6ee05d3ce9a4824c9bde6c8b2b3a2e Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 07:31:31 +0800 Subject: [PATCH 44/74] Cover high-slot keys and unread first-flight errors Admin list/show/renew/revoke/status now include keys retained above the live table. Decrypt failures omit an unreadable error frame, and UI stop paths drop pinned credentials. --- CHANGELOG.md | 1 + docs/authentication-v2.md | 9 +- docs/authentication-v2.zh-CN.md | 7 +- src/common/auth/actor.rs | 352 +++++++++++++++---- src/common/auth/runtime.rs | 8 +- src/common/auth/tests.rs | 68 ++++ src/common/message/secure.rs | 4 +- src/common/message/secure/tests.rs | 35 ++ ui/native/pb_mapper_ffi/src/state/runtime.rs | 4 + 9 files changed, 408 insertions(+), 80 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 823b2de..6604db4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -41,6 +41,7 @@ All notable changes to this project will be documented in this file. - Awaited the authentication actor on relay shutdown, used pinned credentials for registration probes and UI tunnel workers/status checks, and finished in-memory root rotation when `admin.key` already matched the new snapshot. - Bound relay tunneled-frame checksums to each hop's authenticated session key instead of the process administrator key. - Bound local UDP and codec tunnels to the pinned credential's checksum key so a later process-key change cannot desynchronize framed payloads. +- Looked up high-slot credentials for admin list/show/renew/revoke, counted them in status, scheduled their tombstones across restart, and expired due high-slot entries on the actor tick. UI tunnel stop now drops the pinned credential. First-flight decrypt failures no longer send an error frame the presenter cannot read. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/docs/authentication-v2.md b/docs/authentication-v2.md index 6d0f942..9b757dd 100644 --- a/docs/authentication-v2.md +++ b/docs/authentication-v2.md @@ -173,8 +173,13 @@ the bounded audit history, persists `admin.key`, and then switches the key and administrator lease as one state transition. It is a global invalidate: every temporary credential stops authenticating, including unexpired keys in other namespaces, and live connections using those keys are cancelled. A later -first flight with one of those credentials returns `temporary_key_rotated`, -not the generic `temporary_key_invalid` used for a mistyped live key. The CLI +first flight that still decrypts under the previous root returns +`temporary_key_rotated` using that previous session, so the client can read +the structured error. A first flight that cannot decrypt — a mistyped, +foreign, or corrupted credential — fails as `protocol_v2_decrypt_failed` +without an encrypted error frame, because the relay cannot derive a session +the presenter can open. `temporary_key_invalid` is the in-process result when +presented material does not match after derivation. The CLI stages the candidate key before the request and verifies the new key with an authenticated status call. When `--key-file` is omitted, the recovery copy is written below `$XDG_CONFIG_HOME/pb-mapper` (or `$HOME/.config/pb-mapper`) diff --git a/docs/authentication-v2.zh-CN.md b/docs/authentication-v2.zh-CN.md index a4cc4f9..5c2298d 100644 --- a/docs/authentication-v2.zh-CN.md +++ b/docs/authentication-v2.zh-CN.md @@ -128,8 +128,11 @@ tombstone 以给出稳定错误后，槽位可以复用。显式 `key gc` 可立 根密钥轮换先用新密钥写空 snapshot，同时保留有上限的审计历史，持久化 `admin.key`，再把密钥与管理员 lease 作为一次状态变更切换。这是全局作废：所有临时 凭据都会立刻失效，包括其他命名空间里尚未过期的 key，并用这些 key 建立的活动连接 -会被取消。之后再用这些凭据做 first flight 会得到 `temporary_key_rotated`，而不是 -活 key 输错时的 `temporary_key_invalid`。CLI 在发请求前保存候选 key，完成后再用 +会被取消。之后如果 first flight 仍能用上一轮根密钥解密，会得到客户端可读的 +`temporary_key_rotated`。无法解密的 first flight（输错、外站或损坏的凭据）则是 +`protocol_v2_decrypt_failed`，且不会附带加密错误帧，因为中继推导不出对端能打开 +的 session。`temporary_key_invalid` 只用于派生之后材料仍不匹配的进程内校验。 +CLI 在发请求前保存候选 key，完成后再用 新 key 执行一次 `admin status` 验证。未指定 `--key-file` 时，恢复副本默认写到 `$XDG_CONFIG_HOME/pb-mapper`（或 `$HOME/.config/pb-mapper`），不要求本机能写 `/var/lib`。 diff --git a/src/common/auth/actor.rs b/src/common/auth/actor.rs index 87cc9b3..f90bfad 100644 --- a/src/common/auth/actor.rs +++ b/src/common/auth/actor.rs @@ -77,6 +77,32 @@ pub(super) async fn run_auth_actor( )) }) .collect::>(); + tombstones.extend( + inner + .high_slot_entries + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .iter() + .filter_map(|entry| { + let expired_active = entry.state == SlotState::Active && entry.expires_at <= now; + if !matches!(entry.state, SlotState::Expired | SlotState::Revoked) + && !expired_active + { + return None; + } + let tombstoned_at = entry + .tombstoned_at + .or_else(|| { + cold.get(&entry.key_id) + .map(|metadata| metadata.tombstoned_at) + }) + .unwrap_or(entry.expires_at.max(now)); + Some(( + tombstoned_at.saturating_add(TOMBSTONE_RETENTION.as_secs()), + entry.key_id, + )) + }), + ); tombstones.sort_unstable_by_key(|(cleanup_at, _)| *cleanup_at); let mut tombstones = VecDeque::from(tombstones); let mut last_snapshot_at = unix_seconds(); @@ -115,8 +141,11 @@ pub(super) async fn run_auth_actor( "temporary key expired and active work was cancelled" ); } + } else { + lease.cancellation.cancel(); } } + expire_due_high_slots(&inner, &mut cold, &mut tombstones, now); while let Some((cleanup_at, key_id)) = tombstones.front().copied() { if cleanup_at > now { break; @@ -131,6 +160,14 @@ pub(super) async fn run_auth_actor( cold.remove(&key_id); wheel.release(key_id); } + } else { + let mut high = inner + .high_slot_entries + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + high.retain(|entry| entry.key_id != key_id); + cold.remove(&key_id); + wheel.release(key_id); } } prune_expired_admin_replays( @@ -490,6 +527,15 @@ fn actor_list( }) }) .collect::>(); + all.extend( + inner + .high_slot_entries + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .iter() + .filter(|entry| entry.state != SlotState::Free) + .map(high_slot_metadata), + ); all.sort_by_key(|item| std::cmp::Reverse(item.issued_at)); let items = all.iter().skip(start).take(page_size).cloned().collect(); let next_page = (start.saturating_add(page_size) < all.len()).then_some(page.saturating_add(1)); @@ -540,19 +586,34 @@ fn actor_renew( .slots .read() .unwrap_or_else(|poisoned| poisoned.into_inner()); - let slot = slots.get(index).ok_or_else(|| key_not_found(key_id))?; - validate_slot_identity(slot, key_id)?; - if slot.state != SlotState::Active || slot.expires_at <= unix_seconds() { - return Err(AuthFailure::new( - "temporary_key_not_renewable", - "only an active, unexpired temporary key can be renewed", - false, - )); + if let Some(slot) = slots.get(index) { + validate_slot_identity(slot, key_id)?; + if slot.state != SlotState::Active || slot.expires_at <= unix_seconds() { + return Err(key_not_renewable()); + } + } else { + let high = inner + .high_slot_entries + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let entry = high_slot_entry(&high, key_id)?; + if entry.state != SlotState::Active || entry.expires_at <= unix_seconds() { + return Err(key_not_renewable()); + } } } let label = cold .get(&key_id) - .and_then(|metadata| metadata.label.clone()); + .and_then(|metadata| metadata.label.clone()) + .or_else(|| { + inner + .high_slot_entries + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .iter() + .find(|entry| entry.key_id == key_id) + .and_then(|entry| entry.label.clone()) + }); append_mutation( config, inner, @@ -563,31 +624,50 @@ fn actor_renew( .slots .write() .unwrap_or_else(|poisoned| poisoned.into_inner()); - let slot = slots.get_mut(index).ok_or_else(|| key_not_found(key_id))?; - validate_slot_identity(slot, key_id)?; - if slot.state != SlotState::Active { - return Err(AuthFailure::new( - "temporary_key_inactive", - "temporary key lease is no longer active", - true, - )); - } - slot.expires_at = expires_at; - let lease = match slot.lease.upgrade() { - Some(lease) if !lease.cancellation_token().is_cancelled() => { - lease.expires_at.store(expires_at, Ordering::Release); - lease.wheel_version.fetch_add(1, Ordering::AcqRel); - lease - } - _ => { - let lease = Arc::new(AuthLease::new(key_id, expires_at)); - slot.lease = Arc::downgrade(&lease); - lease + if let Some(slot) = slots.get_mut(index) { + validate_slot_identity(slot, key_id)?; + if slot.state != SlotState::Active { + return Err(AuthFailure::new( + "temporary_key_inactive", + "temporary key lease is no longer active", + true, + )); } - }; - wheel.release(key_id); - wheel.insert(lease); + slot.expires_at = expires_at; + let lease = match slot.lease.upgrade() { + Some(lease) if !lease.cancellation_token().is_cancelled() => { + lease.expires_at.store(expires_at, Ordering::Release); + lease.wheel_version.fetch_add(1, Ordering::AcqRel); + lease + } + _ => { + let lease = Arc::new(AuthLease::new(key_id, expires_at)); + slot.lease = Arc::downgrade(&lease); + lease + } + }; + wheel.release(key_id); + wheel.insert(lease); + drop(slots); + return metadata_with_credential(inner, cold, key_id, true); + } drop(slots); + { + let mut high = inner + .high_slot_entries + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let entry = high_slot_entry_mut(&mut high, key_id)?; + if entry.state != SlotState::Active { + return Err(AuthFailure::new( + "temporary_key_inactive", + "temporary key lease is no longer active", + true, + )); + } + entry.expires_at = expires_at; + entry.tombstoned_at = None; + } metadata_with_credential(inner, cold, key_id, true) } @@ -601,26 +681,30 @@ fn actor_revoke( ensure_store_available(inner)?; let now = unix_seconds(); let index = key_slot(key_id) as usize; - { + let (label, issued_at, expires_at) = { let slots = inner .slots .read() .unwrap_or_else(|poisoned| poisoned.into_inner()); - let slot = slots.get(index).ok_or_else(|| key_not_found(key_id))?; - validate_slot_identity(slot, key_id)?; - if slot.state != SlotState::Active { - return Err(AuthFailure::new( - "temporary_key_not_active", - "temporary key is not active", - false, - )); + if let Some(slot) = slots.get(index) { + validate_slot_identity(slot, key_id)?; + if slot.state != SlotState::Active { + return Err(key_not_active()); + } + let metadata = cold.get(&key_id).ok_or_else(|| key_not_found(key_id))?; + (metadata.label.clone(), metadata.issued_at, slot.expires_at) + } else { + let high = inner + .high_slot_entries + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let entry = high_slot_entry(&high, key_id)?; + if entry.state != SlotState::Active { + return Err(key_not_active()); + } + (entry.label.clone(), entry.issued_at, entry.expires_at) } - } - let label = cold - .get(&key_id) - .ok_or_else(|| key_not_found(key_id))? - .label - .clone(); + }; append_mutation( config, inner, @@ -631,21 +715,44 @@ fn actor_revoke( .slots .write() .unwrap_or_else(|poisoned| poisoned.into_inner()); - let slot = slots.get_mut(index).ok_or_else(|| key_not_found(key_id))?; - validate_slot_identity(slot, key_id)?; - slot.state = SlotState::Revoked; - if let Some(lease) = slot.lease.upgrade() { - lease.cancellation.cancel(); + if let Some(slot) = slots.get_mut(index) { + validate_slot_identity(slot, key_id)?; + slot.state = SlotState::Revoked; + if let Some(lease) = slot.lease.upgrade() { + lease.cancellation.cancel(); + } + let cold_metadata = cold.get_mut(&key_id).ok_or_else(|| key_not_found(key_id))?; + cold_metadata.tombstoned_at = now; + push_tombstone(tombstones, now, key_id); + return Ok(TemporaryKeyMetadata { + key_id, + state: slot_state_name(slot.state).to_string(), + issued_at: cold_metadata.issued_at, + expires_at: slot.expires_at, + label: cold_metadata.label.clone(), + }); + } + drop(slots); + let mut high = inner + .high_slot_entries + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let entry = high_slot_entry_mut(&mut high, key_id)?; + if entry.state != SlotState::Active { + return Err(key_not_active()); + } + entry.state = SlotState::Revoked; + entry.tombstoned_at = Some(now); + if let Some(metadata) = cold.get_mut(&key_id) { + metadata.tombstoned_at = now; } - let cold_metadata = cold.get_mut(&key_id).ok_or_else(|| key_not_found(key_id))?; - cold_metadata.tombstoned_at = now; push_tombstone(tombstones, now, key_id); Ok(TemporaryKeyMetadata { key_id, - state: slot_state_name(slot.state).to_string(), - issued_at: cold_metadata.issued_at, - expires_at: slot.expires_at, - label: cold_metadata.label.clone(), + state: slot_state_name(entry.state).to_string(), + issued_at, + expires_at, + label, }) } @@ -892,18 +999,34 @@ fn actor_status(inner: &Arc) -> AuthStatus { .slots .read() .unwrap_or_else(|poisoned| poisoned.into_inner()); + let high = inner + .high_slot_entries + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); let active_keys = slots .iter() .filter(|slot| slot.state == SlotState::Active) - .count(); + .count() + + high + .iter() + .filter(|entry| entry.state == SlotState::Active) + .count(); let expired_keys = slots .iter() .filter(|slot| slot.state == SlotState::Expired) - .count(); + .count() + + high + .iter() + .filter(|entry| entry.state == SlotState::Expired) + .count(); let revoked_keys = slots .iter() .filter(|slot| slot.state == SlotState::Revoked) - .count(); + .count() + + high + .iter() + .filter(|entry| entry.state == SlotState::Revoked) + .count(); let last_legacy_connection_at = inner.last_legacy_connection_at.load(Ordering::Acquire); AuthStatus { schema_version: 1, @@ -995,6 +1118,82 @@ fn key_not_found(key_id: u64) -> AuthFailure { ) } +fn key_not_renewable() -> AuthFailure { + AuthFailure::new( + "temporary_key_not_renewable", + "only an active, unexpired temporary key can be renewed", + false, + ) +} + +fn key_not_active() -> AuthFailure { + AuthFailure::new( + "temporary_key_not_active", + "temporary key is not active", + false, + ) +} + +fn high_slot_entry(high: &[PersistedEntry], key_id: u64) -> Result<&PersistedEntry, AuthFailure> { + high.iter() + .find(|entry| entry.key_id == key_id) + .ok_or_else(|| key_not_found(key_id)) +} + +fn high_slot_entry_mut( + high: &mut [PersistedEntry], + key_id: u64, +) -> Result<&mut PersistedEntry, AuthFailure> { + high.iter_mut() + .find(|entry| entry.key_id == key_id) + .ok_or_else(|| key_not_found(key_id)) +} + +fn high_slot_metadata(entry: &PersistedEntry) -> TemporaryKeyMetadata { + TemporaryKeyMetadata { + key_id: entry.key_id, + state: slot_state_name(entry.state).to_string(), + issued_at: entry.issued_at, + expires_at: entry.expires_at, + label: entry.label.clone(), + } +} + +fn expire_due_high_slots( + inner: &Arc, + cold: &mut HashMap, + tombstones: &mut VecDeque<(u64, u64)>, + now: u64, +) { + let mut high = inner + .high_slot_entries + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let mut expired = Vec::new(); + for entry in high.iter_mut() { + if entry.state == SlotState::Active && entry.expires_at <= now { + entry.state = SlotState::Expired; + let tombstoned_at = entry.expires_at; + entry.tombstoned_at = Some(tombstoned_at); + if let Some(metadata) = cold.get_mut(&entry.key_id) { + metadata.tombstoned_at = tombstoned_at; + } + expired.push((tombstoned_at, entry.key_id, entry.expires_at)); + } + } + drop(high); + for (tombstoned_at, key_id, expires_at) in expired { + push_tombstone(tombstones, tombstoned_at, key_id); + tracing::info!( + event = "temporary_key_expired", + auth_stage = "expiry", + key_id, + expires_at, + "high-slot temporary key expired" + ); + } +} + fn metadata_with_credential( inner: &Arc, cold: &HashMap, @@ -1005,25 +1204,32 @@ fn metadata_with_credential( .slots .read() .unwrap_or_else(|poisoned| poisoned.into_inner()); - let slot = slots - .get(key_slot(key_id) as usize) - .ok_or_else(|| key_not_found(key_id))?; - validate_slot_identity(slot, key_id)?; - let cold = cold.get(&key_id).ok_or_else(|| key_not_found(key_id))?; let credential = if reveal { let key = derive_temporary_key(&inner.admin_key(), &inner.instance_id(), key_id)?; encode_temporary_credential(key_id, &key) } else { String::new() }; + if let Some(slot) = slots.get(key_slot(key_id) as usize) { + validate_slot_identity(slot, key_id)?; + let cold = cold.get(&key_id).ok_or_else(|| key_not_found(key_id))?; + return Ok(IssuedTemporaryKey { + metadata: TemporaryKeyMetadata { + key_id, + state: slot_state_name(slot.state).to_string(), + issued_at: cold.issued_at, + expires_at: slot.expires_at, + label: cold.label.clone(), + }, + credential, + }); + } + let high = inner + .high_slot_entries + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); Ok(IssuedTemporaryKey { - metadata: TemporaryKeyMetadata { - key_id, - state: slot_state_name(slot.state).to_string(), - issued_at: cold.issued_at, - expires_at: slot.expires_at, - label: cold.label.clone(), - }, + metadata: high_slot_metadata(high_slot_entry(&high, key_id)?), credential, }) } diff --git a/src/common/auth/runtime.rs b/src/common/auth/runtime.rs index 0804fe4..91f40a1 100644 --- a/src/common/auth/runtime.rs +++ b/src/common/auth/runtime.rs @@ -149,10 +149,16 @@ impl AuthRuntime { while audit_records.len() > AUDIT_RECORD_CAPACITY { audit_records.pop_front(); } - let (high_slot_generations, high_slot_entries) = loaded + let (high_slot_generations, mut high_slot_entries) = loaded .as_ref() .map(|state| split_high_slot_state(state, config.max_temporary_keys)) .unwrap_or_default(); + for entry in &mut high_slot_entries { + if entry.state == SlotState::Active && entry.expires_at <= now { + entry.state = SlotState::Expired; + entry.tombstoned_at = Some(entry.tombstoned_at.unwrap_or(entry.expires_at)); + } + } let inner = Arc::new(AuthStateInner { admin: RwLock::new(AdminState { key: admin_key, diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index 2306923..ac67f71 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -197,6 +197,74 @@ async fn gc_removes_inactive_high_slot_entries_and_keeps_their_generations() { let _ = std::fs::remove_dir_all(state_dir); } +#[tokio::test] +async fn admin_lifecycle_covers_high_slot_keys_after_capacity_shrink() { + let state_dir = temp_state_dir("high-slot-admin"); + let admin_key = *b"0123456789abcdefghijklmnopqrstuv"; + let config_two = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 2, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + let runtime = AuthRuntime::start(admin_key, config_two.clone()) + .await + .unwrap(); + let admin = authenticate_for_test(&runtime, 0).unwrap(); + let first = runtime + .issue(&admin, Duration::from_secs(60), Some("first".to_string())) + .await + .unwrap(); + let second = runtime + .issue(&admin, Duration::from_secs(60), Some("second".to_string())) + .await + .unwrap(); + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + + let config_one = AuthConfig { + max_temporary_keys: 1, + ..config_two.clone() + }; + let runtime = AuthRuntime::start(admin_key, config_one).await.unwrap(); + let admin = authenticate_for_test(&runtime, 0).unwrap(); + assert_eq!(runtime.high_slot_entry_count(), 1); + let high_id = [first.metadata.key_id, second.metadata.key_id] + .into_iter() + .find(|key_id| key_slot(*key_id) as usize >= 1) + .expect("one issued key should land above the shrunken table"); + let page = runtime.list(&admin, 0, 100).await.unwrap(); + assert_eq!(page.items.len(), 2); + assert!(page.items.iter().any(|item| item.key_id == high_id)); + let shown = runtime.show(&admin, high_id, false).await.unwrap(); + assert_eq!(shown.metadata.key_id, high_id); + assert_eq!(shown.metadata.state, "active"); + assert_eq!( + authenticate_for_test(&runtime, high_id).unwrap_err().code, + "temporary_key_not_found" + ); + let status = runtime.status(&admin).await.unwrap(); + assert_eq!(status.active_keys, 2); + let renewed = runtime + .renew(&admin, high_id, Duration::from_secs(120)) + .await + .unwrap(); + assert!(renewed.metadata.expires_at > shown.metadata.expires_at); + runtime.revoke(&admin, high_id).await.unwrap(); + let revoked = runtime.show(&admin, high_id, false).await.unwrap(); + assert_eq!(revoked.metadata.state, "revoked"); + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + + let runtime = AuthRuntime::start(admin_key, config_two).await.unwrap(); + let admin = authenticate_for_test(&runtime, 0).unwrap(); + let restored = runtime.show(&admin, high_id, false).await.unwrap(); + assert_eq!(restored.metadata.state, "revoked"); + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + let _ = std::fs::remove_dir_all(state_dir); +} + #[tokio::test] async fn safe_mode_denies_legacy_protocol_instead_of_restoring_the_default() { let state_dir = temp_state_dir("safe-mode-legacy"); diff --git a/src/common/message/secure.rs b/src/common/message/secure.rs index 0c53510..7808bf2 100644 --- a/src/common/message/secure.rs +++ b/src/common/message/secure.rs @@ -525,7 +525,7 @@ impl ServerSecurity { error.to_string(), false, ), - response_session: Some(session_without_context(&session)), + response_session: None, })?; let mut current_ciphertext = ciphertext.clone(); let payload = match open_v2_payload( @@ -547,7 +547,7 @@ impl ServerSecurity { error.to_string(), false, ), - response_session: Some(session_without_context(&session)), + response_session: None, }); } }; diff --git a/src/common/message/secure/tests.rs b/src/common/message/secure/tests.rs index a4b2c95..355a77b 100644 --- a/src/common/message/secure/tests.rs +++ b/src/common/message/secure/tests.rs @@ -268,6 +268,41 @@ async fn reset_temporary_first_flight_returns_a_readable_rotated_error() { let _ = std::fs::remove_dir_all(config.state_dir); } +#[tokio::test] +async fn mistyped_temporary_first_flight_does_not_send_an_unreadable_error() { + let admin = *b"0123456789abcdefghijklmnopqrstuv"; + let config = temp_config(); + let auth = AuthRuntime::start(admin, config.clone()).await.unwrap(); + let admin_context = auth.authenticate_presented(0, &admin).unwrap(); + let issued = auth + .issue(&admin_context, std::time::Duration::from_secs(60), None) + .await + .unwrap(); + let Credential::Temporary { key_id, mut key } = parse_credential(&issued.credential).unwrap() + else { + panic!("expected temporary credential"); + }; + key[0] ^= 0x01; + let client = ClientHeaderSession::new_v2(&Credential::Temporary { key_id, key }).unwrap(); + let security = ServerSecurity::new(auth); + let (mut client_io, mut server_io) = tokio::io::duplex(4096); + let client_task = client.write_initial(&mut client_io, b"mistyped"); + let server_task = security.read_initial(&mut server_io); + let (client_result, server_result) = tokio::join!(client_task, server_task); + client_result.unwrap(); + let error = match server_result { + Ok(_) => panic!("mistyped temporary credential should fail decryption"), + Err(error) => error, + }; + assert_eq!(error.failure.code, "protocol_v2_decrypt_failed"); + assert!( + error.response_session.is_none(), + "the presenter cannot open a session derived from the live key" + ); + + let _ = std::fs::remove_dir_all(config.state_dir); +} + #[tokio::test] async fn oversized_initial_frame_is_rejected_before_reading_its_body() { let credential = Credential::Admin(*b"0123456789abcdefghijklmnopqrstuv"); diff --git a/ui/native/pb_mapper_ffi/src/state/runtime.rs b/ui/native/pb_mapper_ffi/src/state/runtime.rs index dd762dd..11dc49f 100644 --- a/ui/native/pb_mapper_ffi/src/state/runtime.rs +++ b/ui/native/pb_mapper_ffi/src/state/runtime.rs @@ -265,6 +265,7 @@ impl PbMapperState { } pub async fn unregister_service(&mut self, service_key: String) -> Result<(), CtlError> { + self.service_credentials.remove(&service_key); if let Some(handle) = self.service_handles.remove(&service_key) { handle.abort(); } @@ -289,6 +290,7 @@ impl PbMapperState { &mut self, service_key: String, ) -> Result<(), CtlError> { + self.service_credentials.remove(&service_key); if let Some(handle) = self.service_handles.remove(&service_key) { handle.abort(); } @@ -421,6 +423,7 @@ impl PbMapperState { pub async fn disconnect_service(&mut self, service_key: String) -> Result<(), CtlError> { // Aborting the task is the part that matters: it is what stops the // retry loop still dialling in the background. + self.client_credentials.remove(&service_key); let aborted = match self.client_handles.remove(&service_key) { Some(handle) => { handle.abort(); @@ -454,6 +457,7 @@ impl PbMapperState { &mut self, service_key: String, ) -> Result<(), CtlError> { + self.client_credentials.remove(&service_key); if let Some(handle) = self.client_handles.remove(&service_key) { handle.abort(); } From b71bf55effe43771c894725088bce2fdccf82a00 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 08:14:41 +0800 Subject: [PATCH 45/74] Keep lease cancel causes and bound first-flight load Record why a lease was cancelled so expiry is not reported as revocation, abort the auth actor on UI shutdown timeout, and cap aggregate Bloom inserts. Unreadable first flights keep the parsed key id, and reset finalizes when the new instance id is already on disk. --- CHANGELOG.md | 1 + src/common/auth.rs | 81 ++++++++++++++++--- src/common/auth/actor.rs | 31 ++++--- src/common/auth/persistence.rs | 25 +++++- src/common/auth/runtime.rs | 23 +++++- src/common/auth/tests.rs | 65 +++++++++++++++ src/common/auth/timing_wheel.rs | 2 +- src/common/message/secure.rs | 51 +++++++++--- src/common/message/secure/replay.rs | 28 ++++++- src/common/message/secure/tests.rs | 12 +++ src/pb_server/connection.rs | 1 + ui/native/pb_mapper_ffi/src/state.rs | 1 + .../pb_mapper_ffi/src/state/configuration.rs | 2 + ui/native/pb_mapper_ffi/src/state/runtime.rs | 11 ++- 14 files changed, 296 insertions(+), 38 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6604db4..b020cd4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -42,6 +42,7 @@ All notable changes to this project will be documented in this file. - Bound relay tunneled-frame checksums to each hop's authenticated session key instead of the process administrator key. - Bound local UDP and codec tunnels to the pinned credential's checksum key so a later process-key change cannot desynchronize framed payloads. - Looked up high-slot credentials for admin list/show/renew/revoke, counted them in status, scheduled their tombstones across restart, and expired due high-slot entries on the actor tick. UI tunnel stop now drops the pinned credential. First-flight decrypt failures no longer send an error frame the presenter cannot read. +- Reported cancelled leases by recorded cause, aborted the auth actor on UI shutdown timeout, validated replacement credentials before stopping a live tunnel, preserved presented key IDs on unreadable first flights, finalized reset when the new instance id was already installed, and bounded aggregate first-flight Bloom inserts. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/src/common/auth.rs b/src/common/auth.rs index d393537..e4f6413 100644 --- a/src/common/auth.rs +++ b/src/common/auth.rs @@ -24,7 +24,7 @@ use std::io::{Read, Write}; #[cfg(unix)] use std::os::unix::fs::PermissionsExt; use std::path::{Path, PathBuf}; -use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, AtomicU8, Ordering}; use std::sync::{Arc, RwLock, Weak}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; @@ -299,12 +299,18 @@ impl fmt::Display for AuthFailure { impl std::error::Error for AuthFailure {} +const LEASE_CANCEL_NONE: u8 = 0; +const LEASE_CANCEL_EXPIRED: u8 = 1; +const LEASE_CANCEL_REVOKED: u8 = 2; +const LEASE_CANCEL_ROTATED: u8 = 3; + #[derive(Debug)] pub struct AuthLease { key_id: u64, expires_at: AtomicU64, wheel_version: AtomicU64, cancellation: CancellationToken, + cancel_reason: AtomicU8, } impl AuthLease { @@ -314,6 +320,7 @@ impl AuthLease { expires_at: AtomicU64::new(expires_at), wheel_version: AtomicU64::new(1), cancellation: CancellationToken::new(), + cancel_reason: AtomicU8::new(LEASE_CANCEL_NONE), } } @@ -328,6 +335,33 @@ impl AuthLease { pub fn cancellation_token(&self) -> CancellationToken { self.cancellation.clone() } + + fn record_cancel(&self, reason: u8) { + let _ = self.cancel_reason.compare_exchange( + LEASE_CANCEL_NONE, + reason, + Ordering::AcqRel, + Ordering::Acquire, + ); + self.cancellation.cancel(); + } + + pub(crate) fn cancel_expired(&self) { + self.record_cancel(LEASE_CANCEL_EXPIRED); + } + + pub(crate) fn cancel_revoked(&self) { + self.record_cancel(LEASE_CANCEL_REVOKED); + } + + pub(crate) fn cancel_rotated(&self) { + self.record_cancel(LEASE_CANCEL_ROTATED); + } + + #[cfg(test)] + pub(crate) fn expire_now(&self) { + self.expires_at.store(0, Ordering::Release); + } } #[derive(Clone, Debug)] @@ -361,18 +395,10 @@ impl AuthContext { ) })?; if lease.cancellation.is_cancelled() { - return Err(AuthFailure::new( - if self.is_admin { - "administrator_key_rotated" - } else { - "temporary_key_revoked" - }, - "credential lease has been cancelled", - false, - )); + return Err(cancelled_lease_failure(self.is_admin, &lease)); } if !self.is_admin && lease.expires_at() <= unix_seconds() { - lease.cancellation.cancel(); + lease.cancel_expired(); return Err(AuthFailure::new( "temporary_key_expired", "temporary key has expired", @@ -410,6 +436,38 @@ impl AuthContext { } } +fn cancelled_lease_failure(is_admin: bool, lease: &AuthLease) -> AuthFailure { + if is_admin { + return AuthFailure::new( + "administrator_key_rotated", + "credential lease has been cancelled", + false, + ); + } + match lease.cancel_reason.load(Ordering::Acquire) { + LEASE_CANCEL_EXPIRED => AuthFailure::new( + "temporary_key_expired", + "temporary key has expired", + false, + ), + LEASE_CANCEL_ROTATED => AuthFailure::new( + "temporary_key_rotated", + "temporary credential was invalidated by administrator root rotation or auth-state reset", + false, + ), + LEASE_CANCEL_REVOKED => AuthFailure::new( + "temporary_key_revoked", + "temporary key was revoked", + false, + ), + _ => AuthFailure::new( + "temporary_key_inactive", + "credential lease has been cancelled", + false, + ), + } +} + #[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] enum SlotState { @@ -496,6 +554,7 @@ pub struct AuthRuntime { config: AuthConfig, _state_lock: Arc, actor: Arc>>>, + actor_abort: tokio::task::AbortHandle, } #[derive(Clone, Debug, Serialize, Deserialize)] diff --git a/src/common/auth/actor.rs b/src/common/auth/actor.rs index f90bfad..cf350c7 100644 --- a/src/common/auth/actor.rs +++ b/src/common/auth/actor.rs @@ -123,7 +123,7 @@ pub(super) async fn run_auth_actor( if let Some(slot) = slots.get_mut(key_slot(key_id) as usize) { if slot.generation == key_generation(key_id) && slot.state == SlotState::Active { slot.state = SlotState::Expired; - lease.cancellation.cancel(); + lease.cancel_expired(); let tombstoned_at = slot.expires_at; if let Some(metadata) = cold.get_mut(&key_id) { metadata.tombstoned_at = tombstoned_at; @@ -142,7 +142,7 @@ pub(super) async fn run_auth_actor( ); } } else { - lease.cancellation.cancel(); + lease.cancel_expired(); } } expire_due_high_slots(&inner, &mut cold, &mut tombstones, now); @@ -205,7 +205,7 @@ pub(super) async fn run_auth_actor( } command = command_rx.recv() => { let Some(command) = command else { - admin_lease.cancellation.cancel(); + admin_lease.cancel_rotated(); cancel_all_temporary_leases(&inner); break; }; @@ -307,7 +307,7 @@ pub(super) async fn run_auth_actor( let _ = response.send(result); } AuthCommand::Shutdown { response } => { - admin_lease.cancellation.cancel(); + admin_lease.cancel_rotated(); cancel_all_temporary_leases(&inner); let _ = response.send(()); break; @@ -719,7 +719,7 @@ fn actor_revoke( validate_slot_identity(slot, key_id)?; slot.state = SlotState::Revoked; if let Some(lease) = slot.lease.upgrade() { - lease.cancellation.cancel(); + lease.cancel_revoked(); } let cold_metadata = cold.get_mut(&key_id).ok_or_else(|| key_not_found(key_id))?; cold_metadata.tombstoned_at = now; @@ -793,7 +793,11 @@ fn actor_gc( { let key_id = make_key_id(slot.generation, index as u32); if let Some(lease) = slot.lease.upgrade() { - lease.cancellation.cancel(); + if slot.state == SlotState::Revoked { + lease.cancel_revoked(); + } else { + lease.cancel_expired(); + } } slot.state = SlotState::Free; slot.expires_at = 0; @@ -863,9 +867,16 @@ fn actor_reset( ) }) { - inner.safe_mode.store(true, Ordering::Release); - cancel_all_temporary_leases(inner); - return Err(error); + if !reset_already_installed(&config.state_dir, &admin_key, &new_instance_id) { + inner.safe_mode.store(true, Ordering::Release); + cancel_all_temporary_leases(inner); + return Err(error); + } + tracing::warn!( + event = "auth_state_reset_finalized_after_sync_error", + error = %error, + "server-instance-id replacement reported an error, but the live id and snapshot already match the new instance; finishing in-memory reset" + ); } let _ = std::fs::remove_file(&next_instance_path); push_audit_record(inner, reset_audit); @@ -971,7 +982,7 @@ fn actor_rotate_root( set_process_msg_header_key(Some(&new_key_string)).map_err(AuthFailure::internal)?; } inner.safe_mode.store(false, Ordering::Release); - old_admin_lease.cancellation.cancel(); + old_admin_lease.cancel_rotated(); *admin_lease = new_admin_lease; Ok(()) } diff --git a/src/common/auth/persistence.rs b/src/common/auth/persistence.rs index 31044e3..8426146 100644 --- a/src/common/auth/persistence.rs +++ b/src/common/auth/persistence.rs @@ -228,7 +228,7 @@ pub(super) fn cancel_all_temporary_leases(inner: &AuthStateInner) { .read() .unwrap_or_else(|poisoned| poisoned.into_inner()); for lease in slots.iter().filter_map(|slot| slot.lease.upgrade()) { - lease.cancellation.cancel(); + lease.cancel_rotated(); } } @@ -979,6 +979,29 @@ pub fn write_admin_key_file(path: &Path, key: &str, force: bool) -> Result<(), A atomic_write(path, format!("{key}\n").as_bytes(), 0o600) } +pub(super) fn reset_already_installed( + state_dir: &Path, + admin_key: &AesKeyType, + new_instance_id: &[u8; INSTANCE_ID_LEN], +) -> bool { + let Ok(Some(live)) = read_instance_id_file(&state_dir.join("server-instance-id")) else { + return false; + }; + if live != *new_instance_id { + return false; + } + let Ok(bytes) = std::fs::read(auth_snapshot_path(state_dir)) else { + return false; + }; + let Ok(plain) = open_blob(admin_key, &bytes) else { + return false; + }; + let Ok(snapshot) = serde_json::from_slice::(&plain) else { + return false; + }; + snapshot.instance_id == *new_instance_id +} + pub(super) fn key_matches_existing_snapshot(state_dir: Option<&Path>, key: &str) -> bool { let Some(state_dir) = state_dir else { return false; diff --git a/src/common/auth/runtime.rs b/src/common/auth/runtime.rs index 91f40a1..7fce2d9 100644 --- a/src/common/auth/runtime.rs +++ b/src/common/auth/runtime.rs @@ -188,12 +188,14 @@ impl AuthRuntime { AuthActorState::new(cold, wheel, admin_replays, admin_replay_order), state_lock.clone(), )); + let actor_abort = actor.abort_handle(); let runtime = Self { inner: Arc::downgrade(&inner), command_tx, config: config.clone(), _state_lock: state_lock.clone(), actor: Arc::new(std::sync::Mutex::new(Some(actor))), + actor_abort, }; Ok(runtime) } @@ -215,6 +217,25 @@ impl AuthRuntime { } } + pub async fn abort_actor(&self) { + self.actor_abort.abort(); + let handle = self + .actor + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .take(); + if let Some(handle) = handle { + let _ = handle.await; + return; + } + for _ in 0..100 { + if self.inner.upgrade().is_none() { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + } + pub fn config(&self) -> &AuthConfig { &self.config } @@ -352,7 +373,7 @@ impl AuthRuntime { )), SlotState::Active if slot.expires_at <= unix_seconds() => { if let Some(lease) = slot.lease.upgrade() { - lease.cancellation.cancel(); + lease.cancel_expired(); } Some(AuthFailure::new( "temporary_key_expired", diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index ac67f71..13dd22a 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -660,6 +660,38 @@ async fn issue_renew_revoke_and_persist() { let _ = std::fs::remove_dir_all(state_dir); } +#[tokio::test] +async fn ensure_active_keeps_expiry_after_the_lease_is_cancelled() { + let state_dir = temp_state_dir("lease-expiry-reason"); + let admin_key = *b"0123456789abcdefghijklmnopqrstuv"; + let config = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 4, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + let runtime = AuthRuntime::start(admin_key, config).await.unwrap(); + let admin = authenticate_for_test(&runtime, 0).unwrap(); + let issued = runtime + .issue(&admin, Duration::from_secs(60), Some("exp".to_string())) + .await + .unwrap(); + let context = authenticate_for_test(&runtime, issued.metadata.key_id).unwrap(); + let lease = context.ensure_active().unwrap(); + lease.expire_now(); + assert_eq!( + context.ensure_active().unwrap_err().code, + "temporary_key_expired" + ); + assert_eq!( + context.ensure_active().unwrap_err().code, + "temporary_key_expired" + ); + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + let _ = std::fs::remove_dir_all(state_dir); +} + #[tokio::test] async fn renew_replaces_a_lease_canceled_during_persistence() { let state_dir = temp_state_dir("renew-canceled-lease"); @@ -787,6 +819,39 @@ fn recover_instance_id_promotes_next_when_snapshot_matches() { let _ = std::fs::remove_dir_all(state_dir); } +#[test] +fn reset_already_installed_accepts_matching_live_id_and_snapshot() { + let state_dir = temp_state_dir("reset-already-installed"); + prepare_state_dir(&state_dir).unwrap(); + let admin_key = *b"0123456789abcdefghijklmnopqrstuv"; + let new_id = [9_u8; INSTANCE_ID_LEN]; + atomic_write(&state_dir.join("server-instance-id"), &new_id, 0o600).unwrap(); + let config = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 1, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + let snapshot = PersistedSnapshot { + schema_version: SNAPSHOT_SCHEMA_VERSION, + instance_id: new_id, + generations: vec![0; 1], + entries: Vec::new(), + legacy_protocol: LegacyProtocolPolicy::Allow, + admin_replays: Vec::new(), + audit_records: VecDeque::new(), + root_epoch: 1, + }; + write_snapshot_and_truncate_wal(&config, &admin_key, &snapshot).unwrap(); + assert!(reset_already_installed(&state_dir, &admin_key, &new_id)); + assert!(!reset_already_installed( + &state_dir, + &admin_key, + &[8_u8; INSTANCE_ID_LEN] + )); + let _ = std::fs::remove_dir_all(state_dir); +} + #[test] fn recover_instance_id_discards_stale_next_when_snapshot_still_matches_current() { let state_dir = temp_state_dir("instance-next-stale"); diff --git a/src/common/auth/timing_wheel.rs b/src/common/auth/timing_wheel.rs index a410593..950e1f4 100644 --- a/src/common/auth/timing_wheel.rs +++ b/src/common/auth/timing_wheel.rs @@ -184,7 +184,7 @@ impl TimingWheel { take_all_entries(&mut self.level2, &mut entries); take_all_entries(&mut self.level3, &mut entries); for lease in self.owners.values() { - lease.cancellation.cancel(); + lease.cancel_rotated(); } *self = Self::new(now); } diff --git a/src/common/message/secure.rs b/src/common/message/secure.rs index 7808bf2..521e4b6 100644 --- a/src/common/message/secure.rs +++ b/src/common/message/secure.rs @@ -265,6 +265,17 @@ pub struct ServerInitialMessage { pub struct ServerInitialError { pub failure: AuthFailure, pub response_session: Option, + pub presented_key_id: Option, +} + +impl ServerInitialError { + fn new(failure: AuthFailure) -> Self { + Self { + failure, + response_session: None, + presented_key_id: None, + } + } } impl fmt::Debug for ServerInitialError { @@ -273,6 +284,7 @@ impl fmt::Debug for ServerInitialError { .debug_struct("ServerInitialError") .field("failure", &self.failure) .field("has_response_session", &self.response_session.is_some()) + .field("presented_key_id", &self.presented_key_id) .finish() } } @@ -329,17 +341,13 @@ impl ServerSecurity { reader: &mut T, ) -> std::result::Result { let mut first = [0_u8; 4]; - reader - .read_exact(&mut first) - .await - .map_err(|error| ServerInitialError { - failure: AuthFailure::new( - "protocol_header_read_failed", - format!("failed to read initial protocol header: {error}"), - true, - ), - response_session: None, - })?; + reader.read_exact(&mut first).await.map_err(|error| { + ServerInitialError::new(AuthFailure::new( + "protocol_header_read_failed", + format!("failed to read initial protocol header: {error}"), + true, + )) + })?; if first == PROTOCOL_V2_MAGIC { self.read_v2_initial(reader).await } else { @@ -360,6 +368,7 @@ impl ServerSecurity { false, ), response_session: None, + presented_key_id: None, }); } let key = self @@ -368,6 +377,7 @@ impl ServerSecurity { .map_err(|failure| ServerInitialError { failure, response_session: None, + presented_key_id: None, })?; let checksum = u32::from_be_bytes(checksum_bytes); let datalen = reader @@ -380,6 +390,7 @@ impl ServerSecurity { true, ), response_session: None, + presented_key_id: None, })?; if !valid_checksum_for_key(datalen, checksum, &key) || datalen > MAX_INITIAL_CIPHERTEXT_LEN { @@ -390,6 +401,7 @@ impl ServerSecurity { false, ), response_session: None, + presented_key_id: None, }); } let mut encrypted = vec![0_u8; datalen as usize]; @@ -403,6 +415,7 @@ impl ServerSecurity { true, ), response_session: None, + presented_key_id: None, })?; let mut codec = Aes256GcmDeCodec::try_new(&key).map_err(|_| ServerInitialError { failure: AuthFailure::new( @@ -411,6 +424,7 @@ impl ServerSecurity { false, ), response_session: None, + presented_key_id: None, })?; let plain = codec .decrypt(&mut encrypted) @@ -421,6 +435,7 @@ impl ServerSecurity { false, ), response_session: None, + presented_key_id: None, })?; let context = self .auth @@ -428,6 +443,7 @@ impl ServerSecurity { .map_err(|failure| ServerInitialError { failure, response_session: None, + presented_key_id: None, })?; let legacy_guard = self.auth @@ -435,6 +451,7 @@ impl ServerSecurity { .map_err(|failure| ServerInitialError { failure, response_session: None, + presented_key_id: None, })?; Ok(ServerInitialMessage { payload: plain.to_vec(), @@ -465,6 +482,7 @@ impl ServerSecurity { true, ), response_session: None, + presented_key_id: None, })?; let version = remainder[0]; let flags = remainder[1]; @@ -483,6 +501,7 @@ impl ServerSecurity { false, ), response_session: None, + presented_key_id: None, }); } let key_id = u64::from_be_bytes(remainder[4..12].try_into().expect("fixed key id")); @@ -498,6 +517,7 @@ impl ServerSecurity { false, ), response_session: None, + presented_key_id: Some(key_id), }); } let key = self @@ -506,6 +526,7 @@ impl ServerSecurity { .map_err(|failure| ServerInitialError { failure, response_session: None, + presented_key_id: Some(key_id), })?; let material = derive_material(key_id, &key, salt).map_err(|error| ServerInitialError { failure: AuthFailure::new( @@ -514,6 +535,7 @@ impl ServerSecurity { false, ), response_session: None, + presented_key_id: Some(key_id), })?; let mut session = v2_session(key, material.clone()); let (counter, ciphertext) = @@ -526,6 +548,7 @@ impl ServerSecurity { false, ), response_session: None, + presented_key_id: Some(key_id), })?; let mut current_ciphertext = ciphertext.clone(); let payload = match open_v2_payload( @@ -548,6 +571,7 @@ impl ServerSecurity { false, ), response_session: None, + presented_key_id: Some(key_id), }); } }; @@ -558,6 +582,7 @@ impl ServerSecurity { .map_err(|failure| ServerInitialError { failure, response_session: Some(session_without_context(&session)), + presented_key_id: Some(key_id), })?; let fingerprint = replay_fingerprint(key_id, &salt); match self @@ -574,6 +599,7 @@ impl ServerSecurity { true, ), response_session: Some(session), + presented_key_id: Some(key_id), }); } FirstFlightAdmit::Limited => { @@ -584,6 +610,7 @@ impl ServerSecurity { true, ), response_session: Some(session), + presented_key_id: Some(key_id), }); } FirstFlightAdmit::Unavailable => { @@ -594,6 +621,7 @@ impl ServerSecurity { true, ), response_session: Some(session), + presented_key_id: Some(key_id), }); } FirstFlightAdmit::Fresh => {} @@ -672,6 +700,7 @@ fn stale_root_first_flight( Some(ServerInitialError { failure: AuthFailure::new(code, message, false), response_session: Some(v2_session(previous_key, previous_material)), + presented_key_id: Some(key_id), }) } diff --git a/src/common/message/secure/replay.rs b/src/common/message/secure/replay.rs index 83f31ac..b3d7de0 100644 --- a/src/common/message/secure/replay.rs +++ b/src/common/message/secure/replay.rs @@ -43,6 +43,11 @@ fn first_flight_budget(window_seconds: u64) -> u32 { .max(8_192) } +fn bloom_insert_capacity(bytes: usize) -> u32 { + let bits = (bytes as u64).saturating_mul(8); + u32::try_from(bits / 16).unwrap_or(u32::MAX).max(1) +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(super) enum FirstFlightAdmit { Fresh, @@ -110,6 +115,9 @@ pub(super) struct ReplayGuard { counts_started_at: u64, window_seconds: u64, max_per_key: u32, + total: u32, + total_started_at: u64, + max_total: u32, log_path: Option, last_compact_at: u64, log_failed: bool, @@ -118,12 +126,16 @@ pub(super) struct ReplayGuard { impl ReplayGuard { pub(super) fn open(log_path: Option, bytes: usize, window_seconds: u64) -> Self { let now = unix_seconds(); + let configured = first_flight_budget(window_seconds); let mut guard = Self { bloom: RotatingBloom::new(bytes, window_seconds), counts: HashMap::new(), counts_started_at: now, window_seconds, - max_per_key: first_flight_budget(window_seconds), + max_per_key: configured, + total: 0, + total_started_at: now, + max_total: configured.min(bloom_insert_capacity(bytes)), log_path, last_compact_at: now, log_failed: false, @@ -138,6 +150,12 @@ impl ReplayGuard { self } + #[cfg(test)] + pub(super) fn with_max_total(mut self, max_total: u32) -> Self { + self.max_total = max_total; + self + } + pub(super) fn admit( &mut self, key_id: u64, @@ -148,14 +166,22 @@ impl ReplayGuard { if self.bloom.contains(fingerprint, now) { return FirstFlightAdmit::Replayed; } + if self.bloom.current_started_at != self.total_started_at { + self.total = 0; + self.total_started_at = self.bloom.current_started_at; + } if self.counts.get(&key_id).copied().unwrap_or(0) >= self.max_per_key { return FirstFlightAdmit::Limited; } + if self.total >= self.max_total { + return FirstFlightAdmit::Limited; + } if self.persist(fingerprint, now).is_err() { return FirstFlightAdmit::Unavailable; } self.bloom.insert(fingerprint, now); *self.counts.entry(key_id).or_insert(0) += 1; + self.total = self.total.saturating_add(1); if now.saturating_sub(self.last_compact_at) >= REPLAY_COMPACT_INTERVAL_SECONDS { self.compact(now); } diff --git a/src/common/message/secure/tests.rs b/src/common/message/secure/tests.rs index 355a77b..c5c7783 100644 --- a/src/common/message/secure/tests.rs +++ b/src/common/message/secure/tests.rs @@ -299,6 +299,7 @@ async fn mistyped_temporary_first_flight_does_not_send_an_unreadable_error() { error.response_session.is_none(), "the presenter cannot open a session derived from the live key" ); + assert_eq!(error.presented_key_id, Some(key_id)); let _ = std::fs::remove_dir_all(config.state_dir); } @@ -401,6 +402,17 @@ fn per_credential_admission_limit_does_not_consume_other_keys() { assert_eq!(guard.admit(1, &[1_u8; 32], now), FirstFlightAdmit::Replayed); } +#[test] +fn aggregate_admission_limit_covers_all_keys() { + let now = unix_seconds(); + let mut guard = ReplayGuard::open(None, 1024, DEFAULT_REPLAY_WINDOW_SECONDS) + .with_max_per_key(100) + .with_max_total(2); + assert_eq!(guard.admit(1, &[1_u8; 32], now), FirstFlightAdmit::Fresh); + assert_eq!(guard.admit(2, &[2_u8; 32], now), FirstFlightAdmit::Fresh); + assert_eq!(guard.admit(3, &[3_u8; 32], now), FirstFlightAdmit::Limited); +} + #[test] fn persisted_first_flights_survive_a_torn_trailing_record() { let mut random = [0_u8; 8]; diff --git a/src/pb_server/connection.rs b/src/pb_server/connection.rs index 7502ebe..4ea3ee0 100644 --- a/src/pb_server/connection.rs +++ b/src/pb_server/connection.rs @@ -60,6 +60,7 @@ pub(super) async fn handle_conn( .response_session .as_ref() .map(|session| session.key_id()) + .or(error.presented_key_id) .unwrap_or_default(); let decision = security.record_failure_log(peer_addr.ip(), key_id, &error.failure.code); if decision.suppressed > 0 { diff --git a/ui/native/pb_mapper_ffi/src/state.rs b/ui/native/pb_mapper_ffi/src/state.rs index bb2174b..7c969d6 100644 --- a/ui/native/pb_mapper_ffi/src/state.rs +++ b/ui/native/pb_mapper_ffi/src/state.rs @@ -446,6 +446,7 @@ struct ConnectCommit { pub struct PbMapperState { server_handle: Option>, + server_auth: Option, server_shutdown_token: Option, server_status_sender: Option>>, diff --git a/ui/native/pb_mapper_ffi/src/state/configuration.rs b/ui/native/pb_mapper_ffi/src/state/configuration.rs index 4571517..fc65331 100644 --- a/ui/native/pb_mapper_ffi/src/state/configuration.rs +++ b/ui/native/pb_mapper_ffi/src/state/configuration.rs @@ -47,6 +47,7 @@ impl PbMapperState { let temp_state = Self { server_handle: None, + server_auth: None, server_shutdown_token: None, server_status_sender: None, server_start_time: None, @@ -84,6 +85,7 @@ impl PbMapperState { let state = Self { server_handle: None, + server_auth: None, server_shutdown_token: None, server_status_sender: None, server_start_time: None, diff --git a/ui/native/pb_mapper_ffi/src/state/runtime.rs b/ui/native/pb_mapper_ffi/src/state/runtime.rs index 11dc49f..9ad35a8 100644 --- a/ui/native/pb_mapper_ffi/src/state/runtime.rs +++ b/ui/native/pb_mapper_ffi/src/state/runtime.rs @@ -47,6 +47,7 @@ impl PbMapperState { let (status_sender, status_receiver) = tokio::sync::mpsc::unbounded_channel(); + self.server_auth = Some(auth.clone()); let handle = tokio::spawn(async move { if let Err(e) = run_server_on_listener( listener, @@ -100,11 +101,17 @@ impl PbMapperState { tracing::info!("Server shutdown gracefully"); } Err(_) => { + if let Some(auth) = self.server_auth.as_ref() { + auth.abort_actor().await; + } handle.abort(); let _ = handle.await; tracing::warn!("Server shutdown timed out; aborted the relay task"); } } + if let Some(auth) = self.server_auth.take() { + auth.shutdown_actor().await; + } self.server_start_time = None; @@ -153,6 +160,7 @@ impl PbMapperState { remote_sock_addr, } = commit; + let credential = get_process_credential().map_err(CtlError::invalid_argument)?; self.service_credentials.remove(&service_key); if let Some(previous) = self.service_handles.remove(&service_key) { tracing::warn!( @@ -192,7 +200,6 @@ impl PbMapperState { ); }); - let credential = get_process_credential().map_err(CtlError::invalid_argument)?; self.service_credentials .insert(service_key.clone(), credential); let handle = if protocol.to_uppercase() == "TCP" { @@ -310,6 +317,7 @@ impl PbMapperState { remote_sock_addr, } = commit; + let credential = get_process_credential().map_err(CtlError::invalid_argument)?; self.client_credentials.remove(&service_key); if let Some(previous) = self.client_handles.remove(&service_key) { tracing::warn!( @@ -339,7 +347,6 @@ impl PbMapperState { }) }; - let credential = get_process_credential().map_err(CtlError::invalid_argument)?; self.client_credentials .insert(service_key.clone(), credential); let handle = if protocol_upper == "TCP" { From 80730a0668aa2ef9982d6d57c9786d0d1670e265 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 08:46:53 +0800 Subject: [PATCH 46/74] Wait for actor abort and avoid replay GCM reuse abort_actor now waits until the actor drops. Replay errors no longer reuse the accepted first-response nonce. Restarts restore the Bloom insert count, and root rotation finalizes only when the live key already matches the new snapshot. --- CHANGELOG.md | 1 + docs/authentication-v2.md | 7 +++--- docs/authentication-v2.zh-CN.md | 6 +++-- src/bin/pb-mapper/admin.rs | 10 ++++++-- src/common/auth/actor.rs | 2 +- src/common/auth/persistence.rs | 15 ++++++++++++ src/common/auth/runtime.rs | 6 +---- src/common/auth/tests.rs | 32 +++++++++++++++++++++++++ src/common/message/secure.rs | 2 +- src/common/message/secure/replay.rs | 4 ++++ src/common/message/secure/tests.rs | 36 +++++++++++++++++++++++++++++ 11 files changed, 107 insertions(+), 14 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b020cd4..1b8e0f9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -43,6 +43,7 @@ All notable changes to this project will be documented in this file. - Bound local UDP and codec tunnels to the pinned credential's checksum key so a later process-key change cannot desynchronize framed payloads. - Looked up high-slot credentials for admin list/show/renew/revoke, counted them in status, scheduled their tombstones across restart, and expired due high-slot entries on the actor tick. UI tunnel stop now drops the pinned credential. First-flight decrypt failures no longer send an error frame the presenter cannot read. - Reported cancelled leases by recorded cause, aborted the auth actor on UI shutdown timeout, validated replacement credentials before stopping a live tunnel, preserved presented key IDs on unreadable first flights, finalized reset when the new instance id was already installed, and bounded aggregate first-flight Bloom inserts. +- Waited for the auth actor to drop before `abort_actor` returns, omitted a reused GCM nonce on salt-replay errors, restored the aggregate replay count from the durable log, and finalized root rotation only when the live `admin.key` already matches the new snapshot. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/docs/authentication-v2.md b/docs/authentication-v2.md index 9b757dd..be5a8da 100644 --- a/docs/authentication-v2.md +++ b/docs/authentication-v2.md @@ -115,9 +115,10 @@ uses server-to-client counter `0`. Later control frames continue from counter The relay fingerprints `(key_id, connection_salt)` and atomically checks and inserts it in two rotating 1 MiB Bloom filters covering the current and previous 600-second windows, so a max-future first-flight timestamp cannot outlive replay -retention. A probable duplicate returns the stable retryable error -`connection_salt_replayed`; one-shot administrator CLI operations retry once -with a fresh salt. Mutating administrator requests additionally claim their +retention. A probable duplicate is `connection_salt_replayed`. The relay does not encrypt +that error with the already-used first-response nonce, so the presenter may +see a decrypt/EOF instead of a readable frame. One-shot administrator CLI +operations retry once with a fresh salt in either case. Mutating administrator requests additionally claim their exact fingerprint in the encrypted WAL before dispatch. Those claims survive restart and compaction for ten minutes after the server accepted them, so an old captured mutation cannot be replayed after the Bloom window or a process diff --git a/docs/authentication-v2.zh-CN.md b/docs/authentication-v2.zh-CN.md index 5c2298d..118f904 100644 --- a/docs/authentication-v2.zh-CN.md +++ b/docs/authentication-v2.zh-CN.md @@ -89,8 +89,10 @@ HKDF-SHA256 使用 connection salt 作为 salt，凭据的 32 字节 secret 作 服务端对 `(key_id, connection_salt)` 做指纹，并在同一个临界区内完成两个轮换的 1 MiB Bloom filter 的检查与写入，覆盖当前与上一个 600 秒窗口，使首帧允许的 -最大未来时间戳无法在过滤器遗忘后继续重放。疑似重复会返回 -可重试错误 `connection_salt_replayed`；一次性 admin CLI 会自动换 salt 重试一次。 +最大未来时间戳无法在过滤器遗忘后继续重放。疑似重复是 +`connection_salt_replayed`。中继不会用已经用过的首响 nonce 加密该错误，所以 +对端可能看到解密失败/EOF 而不是可读错误帧。一次性 admin CLI 在这两种情况下 +都会自动换 salt 重试一次。 会修改状态的管理员请求还会在分发前把精确指纹写入加密 WAL；该记录自服务端接受起 在十分钟内跨重启、跨 compact 保留，不能通过等待 Bloom 窗口结束或重启进程来重放 旧操作。客户端首帧时间戳仍用于新鲜度检查，但不决定这条记录保留多久。 diff --git a/src/bin/pb-mapper/admin.rs b/src/bin/pb-mapper/admin.rs index a58f661..35e5946 100644 --- a/src/bin/pb-mapper/admin.rs +++ b/src/bin/pb-mapper/admin.rs @@ -325,7 +325,7 @@ async fn send_admin_request_with_timeout( ) -> Result> { let encoded = PbConnRequest::Admin(request).encode()?; for attempt in 0..2 { - let response = tokio::time::timeout(io_timeout, async { + let attempt_result = tokio::time::timeout(io_timeout, async { let mut stream = TcpStream::connect(remote_addr) .await .map_err(|error| -> Box { Box::new(error) })?; @@ -344,7 +344,13 @@ async fn send_admin_request_with_timeout( io_timeout.as_millis() ), ) - })??; + }); + let response = match attempt_result { + Ok(Ok(response)) => response, + Ok(Err(_)) if attempt == 0 => continue, + Ok(Err(error)) => return Err(error), + Err(error) => return Err(error.into()), + }; match response { PbConnResponse::Admin(response) => return Ok(response), PbConnResponse::Error(error) diff --git a/src/common/auth/actor.rs b/src/common/auth/actor.rs index cf350c7..4f4989d 100644 --- a/src/common/auth/actor.rs +++ b/src/common/auth/actor.rs @@ -939,7 +939,7 @@ fn actor_rotate_root( .and_then(|()| write_snapshot_and_truncate_wal(config, &new_key, &snapshot)) .and_then(|()| write_admin_key(&config.state_dir, &new_key_string)) { - if !key_matches_existing_snapshot(Some(&config.state_dir), &new_key_string) { + if !rotation_already_installed(&config.state_dir, &new_key_string) { inner.safe_mode.store(true, Ordering::Release); cancel_all_temporary_leases(inner); return Err(error); diff --git a/src/common/auth/persistence.rs b/src/common/auth/persistence.rs index 8426146..3787718 100644 --- a/src/common/auth/persistence.rs +++ b/src/common/auth/persistence.rs @@ -1002,6 +1002,21 @@ pub(super) fn reset_already_installed( snapshot.instance_id == *new_instance_id } +pub(super) fn rotation_already_installed(state_dir: &Path, new_key: &str) -> bool { + key_matches_existing_snapshot(Some(state_dir), new_key) + && live_admin_key_matches(state_dir, new_key) +} + +fn live_admin_key_matches(state_dir: &Path, new_key: &str) -> bool { + let Ok(raw) = std::fs::read(state_dir.join("admin.key")) else { + return false; + }; + let Ok(text) = std::str::from_utf8(&raw) else { + return false; + }; + text.trim().as_bytes() == new_key.trim().as_bytes() +} + pub(super) fn key_matches_existing_snapshot(state_dir: Option<&Path>, key: &str) -> bool { let Some(state_dir) = state_dir else { return false; diff --git a/src/common/auth/runtime.rs b/src/common/auth/runtime.rs index 7fce2d9..76d4e03 100644 --- a/src/common/auth/runtime.rs +++ b/src/common/auth/runtime.rs @@ -226,12 +226,8 @@ impl AuthRuntime { .take(); if let Some(handle) = handle { let _ = handle.await; - return; } - for _ in 0..100 { - if self.inner.upgrade().is_none() { - break; - } + while self.inner.upgrade().is_some() { tokio::time::sleep(Duration::from_millis(10)).await; } } diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index 13dd22a..6559c34 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -921,6 +921,38 @@ fn recover_admin_key_discards_leftover_wal_from_the_old_key() { let _ = std::fs::remove_dir_all(state_dir); } +#[test] +fn rotation_finalize_requires_the_live_admin_key() { + let state_dir = temp_state_dir("rotate-requires-live-key"); + prepare_state_dir(&state_dir).unwrap(); + let old_key = *b"0123456789abcdefghijklmnopqrstuv"; + let new_key = *b"abcdefghijklmnopqrstuvwxyz012345"; + let old_key_str = std::str::from_utf8(&old_key).unwrap(); + let new_key_str = std::str::from_utf8(&new_key).unwrap(); + write_admin_key(&state_dir, old_key_str).unwrap(); + let config = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 1, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + let snapshot = PersistedSnapshot { + schema_version: SNAPSHOT_SCHEMA_VERSION, + instance_id: [3_u8; INSTANCE_ID_LEN], + generations: vec![0; 1], + entries: Vec::new(), + legacy_protocol: LegacyProtocolPolicy::Allow, + admin_replays: Vec::new(), + audit_records: VecDeque::new(), + root_epoch: 1, + }; + write_snapshot_and_truncate_wal(&config, &new_key, &snapshot).unwrap(); + assert!(!rotation_already_installed(&state_dir, new_key_str)); + write_admin_key(&state_dir, new_key_str).unwrap(); + assert!(rotation_already_installed(&state_dir, new_key_str)); + let _ = std::fs::remove_dir_all(state_dir); +} + #[tokio::test] async fn interrupted_reset_recovers_the_staged_instance_id_on_restart() { let state_dir = temp_state_dir("reset-recover"); diff --git a/src/common/message/secure.rs b/src/common/message/secure.rs index 521e4b6..fdf43e9 100644 --- a/src/common/message/secure.rs +++ b/src/common/message/secure.rs @@ -598,7 +598,7 @@ impl ServerSecurity { "protocol-v2 connection salt was already accepted", true, ), - response_session: Some(session), + response_session: None, presented_key_id: Some(key_id), }); } diff --git a/src/common/message/secure/replay.rs b/src/common/message/secure/replay.rs index b3d7de0..3fef5f1 100644 --- a/src/common/message/secure/replay.rs +++ b/src/common/message/secure/replay.rs @@ -265,6 +265,10 @@ impl ReplayGuard { self.bloom.insert(&fingerprint, timestamp); live.push(record); } + self.total = u32::try_from(live.len()) + .unwrap_or(u32::MAX) + .min(self.max_total); + self.total_started_at = self.bloom.current_started_at; if self.rewrite_live(&live).is_err() { self.log_failed = true; return; diff --git a/src/common/message/secure/tests.rs b/src/common/message/secure/tests.rs index c5c7783..00ecf40 100644 --- a/src/common/message/secure/tests.rs +++ b/src/common/message/secure/tests.rs @@ -132,6 +132,13 @@ async fn identical_initial_frames_are_admitted_only_once() { .code, "connection_salt_replayed" ); + assert!(results + .iter() + .filter_map(|result| result.as_ref().err()) + .next() + .unwrap() + .response_session + .is_none()); let _ = std::fs::remove_dir_all(config.state_dir); } @@ -496,6 +503,35 @@ fn persisted_first_flights_survive_replay_guard_restart() { let _ = std::fs::remove_file(path); } +#[test] +fn restored_replay_log_consumes_the_aggregate_budget() { + let mut random = [0_u8; 8]; + let mut rng = rand::rng(); + for byte in &mut random { + *byte = rng.random(); + } + let path = std::env::temp_dir().join(format!( + "pb-mapper-replay-total-{}", + u64::from_be_bytes(random) + )); + let now = unix_seconds(); + { + let mut guard = ReplayGuard::open(Some(path.clone()), 1024, DEFAULT_REPLAY_WINDOW_SECONDS) + .with_max_per_key(100) + .with_max_total(2); + assert_eq!(guard.admit(1, &[1_u8; 32], now), FirstFlightAdmit::Fresh); + assert_eq!(guard.admit(2, &[2_u8; 32], now), FirstFlightAdmit::Fresh); + } + let mut restored = ReplayGuard::open(Some(path.clone()), 1024, DEFAULT_REPLAY_WINDOW_SECONDS) + .with_max_per_key(100) + .with_max_total(2); + assert_eq!( + restored.admit(3, &[3_u8; 32], now), + FirstFlightAdmit::Limited + ); + let _ = std::fs::remove_file(path); +} + #[tokio::test] async fn legacy_initial_frame_validates_against_isolated_relay_key() { let _process_credential_guard = PROCESS_CREDENTIAL_TEST_LOCK.lock().await; From 59144737993ed699c805552360051c6104fed2c4 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 08:54:24 +0800 Subject: [PATCH 47/74] Retry admin commands only before they are sent A dropped response must not resend key issue or other mutating requests. The CLI still retries a readable salt-replay error and failures that happen before the first flight is written. --- CHANGELOG.md | 1 + docs/authentication-v2.md | 5 ++++- docs/authentication-v2.zh-CN.md | 5 +++-- src/bin/pb-mapper/admin.rs | 6 +++++- 4 files changed, 13 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1b8e0f9..cb6f44e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -44,6 +44,7 @@ All notable changes to this project will be documented in this file. - Looked up high-slot credentials for admin list/show/renew/revoke, counted them in status, scheduled their tombstones across restart, and expired due high-slot entries on the actor tick. UI tunnel stop now drops the pinned credential. First-flight decrypt failures no longer send an error frame the presenter cannot read. - Reported cancelled leases by recorded cause, aborted the auth actor on UI shutdown timeout, validated replacement credentials before stopping a live tunnel, preserved presented key IDs on unreadable first flights, finalized reset when the new instance id was already installed, and bounded aggregate first-flight Bloom inserts. - Waited for the auth actor to drop before `abort_actor` returns, omitted a reused GCM nonce on salt-replay errors, restored the aggregate replay count from the durable log, and finalized root rotation only when the live `admin.key` already matches the new snapshot. +- Restricted administrator CLI retries to pre-send failures and a readable `connection_salt_replayed`, so a dropped response cannot issue a second credential. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/docs/authentication-v2.md b/docs/authentication-v2.md index be5a8da..179f1d2 100644 --- a/docs/authentication-v2.md +++ b/docs/authentication-v2.md @@ -118,7 +118,10 @@ inserts it in two rotating 1 MiB Bloom filters covering the current and previous retention. A probable duplicate is `connection_salt_replayed`. The relay does not encrypt that error with the already-used first-response nonce, so the presenter may see a decrypt/EOF instead of a readable frame. One-shot administrator CLI -operations retry once with a fresh salt in either case. Mutating administrator requests additionally claim their +operations retry once with a fresh salt only for a readable +`connection_salt_replayed` result or a failure before the request is written. +They do not resend after a dropped response, which would duplicate +non-idempotent commands such as `key issue`. Mutating administrator requests additionally claim their exact fingerprint in the encrypted WAL before dispatch. Those claims survive restart and compaction for ten minutes after the server accepted them, so an old captured mutation cannot be replayed after the Bloom window or a process diff --git a/docs/authentication-v2.zh-CN.md b/docs/authentication-v2.zh-CN.md index 118f904..a23502e 100644 --- a/docs/authentication-v2.zh-CN.md +++ b/docs/authentication-v2.zh-CN.md @@ -91,8 +91,9 @@ HKDF-SHA256 使用 connection salt 作为 salt，凭据的 32 字节 secret 作 1 MiB Bloom filter 的检查与写入，覆盖当前与上一个 600 秒窗口，使首帧允许的 最大未来时间戳无法在过滤器遗忘后继续重放。疑似重复是 `connection_salt_replayed`。中继不会用已经用过的首响 nonce 加密该错误，所以 -对端可能看到解密失败/EOF 而不是可读错误帧。一次性 admin CLI 在这两种情况下 -都会自动换 salt 重试一次。 +对端可能看到解密失败/EOF 而不是可读错误帧。一次性 admin CLI 只在读到该错误、 +或请求尚未写出时换 salt 重试一次，不会在响应丢失后再发，以免 `key issue` +这类非幂等命令被执行两次。 会修改状态的管理员请求还会在分发前把精确指纹写入加密 WAL；该记录自服务端接受起 在十分钟内跨重启、跨 compact 保留，不能通过等待 Bloom 窗口结束或重启进程来重放 旧操作。客户端首帧时间戳仍用于新鲜度检查，但不决定这条记录保留多久。 diff --git a/src/bin/pb-mapper/admin.rs b/src/bin/pb-mapper/admin.rs index 35e5946..a950bd8 100644 --- a/src/bin/pb-mapper/admin.rs +++ b/src/bin/pb-mapper/admin.rs @@ -325,11 +325,13 @@ async fn send_admin_request_with_timeout( ) -> Result> { let encoded = PbConnRequest::Admin(request).encode()?; for attempt in 0..2 { + let sent = std::sync::atomic::AtomicBool::new(false); let attempt_result = tokio::time::timeout(io_timeout, async { let mut stream = TcpStream::connect(remote_addr) .await .map_err(|error| -> Box { Box::new(error) })?; let session = ClientHeaderSession::from_process()?; + sent.store(true, std::sync::atomic::Ordering::Release); session.write_initial(&mut stream, &encoded).await?; let mut reader = session.response_reader(&mut stream)?; let message = reader.read_msg().await?; @@ -345,10 +347,12 @@ async fn send_admin_request_with_timeout( ), ) }); + let pre_send = !sent.load(std::sync::atomic::Ordering::Acquire); let response = match attempt_result { Ok(Ok(response)) => response, - Ok(Err(_)) if attempt == 0 => continue, + Ok(Err(_)) if attempt == 0 && pre_send => continue, Ok(Err(error)) => return Err(error), + Err(_) if attempt == 0 && pre_send => continue, Err(error) => return Err(error.into()), }; match response { From 717edae9c3099d7943e177e1b050008d54c01210 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 09:15:16 +0800 Subject: [PATCH 48/74] Keep replay nonces unique and stream quotas live Auth errors after an already-accepted salt no longer reuse the first response nonce. Namespace stream counts stay until the client stream deregisters, even if the registration control socket drops. --- CHANGELOG.md | 1 + src/common/message/secure.rs | 18 ++++++++++--- src/common/message/secure/replay.rs | 4 +++ src/common/message/secure/tests.rs | 42 +++++++++++++++++++++++++++++ src/pb_server/connection.rs | 17 ------------ src/pb_server/mod.rs | 3 +-- src/pb_server/runtime.rs | 12 --------- 7 files changed, 63 insertions(+), 34 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cb6f44e..ed422d5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -45,6 +45,7 @@ All notable changes to this project will be documented in this file. - Reported cancelled leases by recorded cause, aborted the auth actor on UI shutdown timeout, validated replacement credentials before stopping a live tunnel, preserved presented key IDs on unreadable first flights, finalized reset when the new instance id was already installed, and bounded aggregate first-flight Bloom inserts. - Waited for the auth actor to drop before `abort_actor` returns, omitted a reused GCM nonce on salt-replay errors, restored the aggregate replay count from the durable log, and finalized root rotation only when the live `admin.key` already matches the new snapshot. - Restricted administrator CLI retries to pre-send failures and a readable `connection_salt_replayed`, so a dropped response cannot issue a second credential. +- Omitted a response session for already-admitted first flights that later fail authentication, and kept namespace stream accounting until the client stream deregisters even if the registration control socket drops. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/src/common/message/secure.rs b/src/common/message/secure.rs index fdf43e9..ad3d3c9 100644 --- a/src/common/message/secure.rs +++ b/src/common/message/secure.rs @@ -551,6 +551,12 @@ impl ServerSecurity { presented_key_id: Some(key_id), })?; let mut current_ciphertext = ciphertext.clone(); + let fingerprint = replay_fingerprint(key_id, &salt); + let already_admitted = self + .replay + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .already_admitted(&fingerprint, unix_seconds()); let payload = match open_v2_payload( &material, DIRECTION_CLIENT_TO_SERVER, @@ -559,9 +565,12 @@ impl ServerSecurity { ) { Ok(payload) => payload, Err(error) => { - if let Some(stale) = + if let Some(mut stale) = stale_root_first_flight(&self.auth, key_id, salt, counter, &ciphertext) { + if already_admitted { + stale.response_session = None; + } return Err(stale); } return Err(ServerInitialError { @@ -581,10 +590,13 @@ impl ServerSecurity { .authenticate_presented(key_id, &key) .map_err(|failure| ServerInitialError { failure, - response_session: Some(session_without_context(&session)), + response_session: if already_admitted { + None + } else { + Some(session_without_context(&session)) + }, presented_key_id: Some(key_id), })?; - let fingerprint = replay_fingerprint(key_id, &salt); match self .replay .lock() diff --git a/src/common/message/secure/replay.rs b/src/common/message/secure/replay.rs index 3fef5f1..a555e08 100644 --- a/src/common/message/secure/replay.rs +++ b/src/common/message/secure/replay.rs @@ -156,6 +156,10 @@ impl ReplayGuard { self } + pub(super) fn already_admitted(&mut self, fingerprint: &[u8; 32], now: u64) -> bool { + self.bloom.contains(fingerprint, now) + } + pub(super) fn admit( &mut self, key_id: u64, diff --git a/src/common/message/secure/tests.rs b/src/common/message/secure/tests.rs index 00ecf40..ee82ad1 100644 --- a/src/common/message/secure/tests.rs +++ b/src/common/message/secure/tests.rs @@ -182,6 +182,48 @@ async fn revoked_first_flights_do_not_consume_the_replay_filter() { let _ = std::fs::remove_dir_all(config.state_dir); } +#[tokio::test] +async fn accepted_then_revoked_replay_does_not_reuse_the_session_nonce() { + let admin = *b"0123456789abcdefghijklmnopqrstuv"; + let config = temp_config(); + let auth = AuthRuntime::start(admin, config.clone()).await.unwrap(); + let admin_context = auth.authenticate_presented(0, &admin).unwrap(); + let issued = auth + .issue(&admin_context, std::time::Duration::from_secs(60), None) + .await + .unwrap(); + let Credential::Temporary { key_id, key } = parse_credential(&issued.credential).unwrap() + else { + panic!("expected temporary credential"); + }; + let client = ClientHeaderSession::new_v2(&Credential::Temporary { key_id, key }).unwrap(); + let bytes = encode_initial(&client, b"accepted-then-revoked").await; + let security = ServerSecurity::new(auth); + security + .read_initial(&mut std::io::Cursor::new(bytes.clone())) + .await + .expect("first flight should be accepted"); + security + .auth() + .revoke(&admin_context, key_id) + .await + .unwrap(); + let replayed = match security + .read_initial(&mut std::io::Cursor::new(bytes)) + .await + { + Ok(_) => panic!("replay after revoke should fail"), + Err(error) => error, + }; + assert_eq!(replayed.failure.code, "temporary_key_revoked"); + assert!( + replayed.response_session.is_none(), + "a previously accepted salt must not reuse nonce 0 for the revoke error" + ); + + let _ = std::fs::remove_dir_all(config.state_dir); +} + #[tokio::test] async fn rotated_temporary_first_flight_returns_a_readable_rotated_error() { let _process_credential_guard = PROCESS_CREDENTIAL_TEST_LOCK.lock().await; diff --git a/src/pb_server/connection.rs b/src/pb_server/connection.rs index 4ea3ee0..064a1bb 100644 --- a/src/pb_server/connection.rs +++ b/src/pb_server/connection.rs @@ -586,20 +586,3 @@ pub(super) fn release_namespace_rate_limit_if_idle( namespace_rate_limits.remove(&namespace); } } - -pub(super) fn remove_pending_streams_for_server( - pending_streams: &mut hashbrown::HashMap, - namespace_stream_counts: &mut hashbrown::HashMap, - server_id_to_remove: RemoteConnId, -) -> usize { - let mut removed = 0; - pending_streams.retain(|_, (server_id, _, key)| { - if *server_id != server_id_to_remove { - return true; - } - decrement_namespace_stream_count(namespace_stream_counts, split_scoped_service_key(key).0); - removed += 1; - false - }); - removed -} diff --git a/src/pb_server/mod.rs b/src/pb_server/mod.rs index 5bcbf01..c80004a 100644 --- a/src/pb_server/mod.rs +++ b/src/pb_server/mod.rs @@ -381,8 +381,7 @@ pub use runtime::{ mod connection; use connection::{ decrement_namespace_stream_count, handle_conn, handle_listener, - release_namespace_rate_limit_if_idle, remove_pending_streams_for_server, - split_scoped_service_key, + release_namespace_rate_limit_if_idle, split_scoped_service_key, }; pub async fn get_init_request( conn: &mut TcpStream, diff --git a/src/pb_server/runtime.rs b/src/pb_server/runtime.rs index eb192ab..5957af5 100644 --- a/src/pb_server/runtime.rs +++ b/src/pb_server/runtime.rs @@ -372,11 +372,6 @@ pub async fn run_server_on_listener( let removed_from_service_map = remove_server_conn(&mut server_conn_map, &key, conn_id); let removed_from_active_map = manager.deregister_conn(conn_id); - let removed_pending_streams = remove_pending_streams_for_server( - &mut pending_streams, - &mut namespace_stream_counts, - conn_id, - ); release_namespace_rate_limit_if_idle( split_scoped_service_key(&key).0, &server_conn_map, @@ -389,7 +384,6 @@ pub async fn run_server_on_listener( conn_id = %conn_id, removed_from_service_map, removed_from_active_map, - removed_pending_streams, registered_services = server_conn_map.len(), server_connections = registered_server_conn_count(&server_conn_map), active_connections = manager.active_conn_count(), @@ -416,11 +410,6 @@ pub async fn run_server_on_listener( let removed_from_service_map = remove_server_conn(&mut server_conn_map, &key, conn_id); let removed_from_active_map = manager.deregister_conn(conn_id); - let removed_pending_streams = remove_pending_streams_for_server( - &mut pending_streams, - &mut namespace_stream_counts, - conn_id, - ); release_namespace_rate_limit_if_idle( split_scoped_service_key(&key).0, &server_conn_map, @@ -444,7 +433,6 @@ pub async fn run_server_on_listener( reason = %reason, removed_from_service_map, removed_from_active_map, - removed_pending_streams, retire_notified, registered_services = server_conn_map.len(), server_connections = registered_server_conn_count(&server_conn_map), From ec32d78a8ee82f2b03763c25af293ebdd885c5c5 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 09:49:26 +0800 Subject: [PATCH 49/74] Admit first flights under one blocking lock Authentication and replay admission share one mutex off the Tokio worker. Unavailable admissions omit a session, restored Bloom generations age from loaded records, and a recovery MSG_HEADER_KEY is not written unless it decrypts existing state. --- CHANGELOG.md | 1 + scripts/install-server-gitee.sh | 11 +- scripts/install-server-github.sh | 11 +- src/common/auth.rs | 9 ++ src/common/auth/actor.rs | 17 ++- src/common/auth/tests.rs | 40 ++++++ src/common/message/secure.rs | 139 ++++++++++++------- src/common/message/secure/replay.rs | 14 +- src/common/message/secure/tests.rs | 33 +++++ ui/native/pb_mapper_ffi/src/state.rs | 25 +++- ui/native/pb_mapper_ffi/src/state/runtime.rs | 4 +- 11 files changed, 229 insertions(+), 75 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ed422d5..cbe3c60 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -46,6 +46,7 @@ All notable changes to this project will be documented in this file. - Waited for the auth actor to drop before `abort_actor` returns, omitted a reused GCM nonce on salt-replay errors, restored the aggregate replay count from the durable log, and finalized root rotation only when the live `admin.key` already matches the new snapshot. - Restricted administrator CLI retries to pre-send failures and a readable `connection_salt_replayed`, so a dropped response cannot issue a second credential. - Omitted a response session for already-admitted first flights that later fail authentication, and kept namespace stream accounting until the client stream deregisters even if the registration control socket drops. +- Evaluated first-flight admission on a blocking thread under one replay lock, omitted a session when durable admission is unavailable, aged restored Bloom generations from loaded record timestamps, validated installer keys from `server.env`, refused to persist a recovery `MSG_HEADER_KEY` that cannot decrypt existing state, batched high-slot tombstone cleanup, and captured the UI credential together with the relay address before DNS. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/scripts/install-server-gitee.sh b/scripts/install-server-gitee.sh index b776776..ee92f7c 100755 --- a/scripts/install-server-gitee.sh +++ b/scripts/install-server-gitee.sh @@ -103,21 +103,22 @@ install -m 0755 "$BIN_PATH" "${INSTALL_DIR}/pb-mapper" # key in the environment or /etc/pb-mapper/server.env must win; otherwise the # runtime would prefer the newly copied admin.key and lock operators out. install -d -m 0700 "$AUTH_DIR" -if [ -n "${MSG_HEADER_KEY:-}" ] && [ ! -s "$ADMIN_KEY_PATH" ]; then - case "$MSG_HEADER_KEY" in +INSTALLER_KEY="$(configured_msg_header_key)" +if [ -n "$INSTALLER_KEY" ] && [ ! -s "$ADMIN_KEY_PATH" ]; then + case "$INSTALLER_KEY" in pbmt1_*) echo "MSG_HEADER_KEY is a temporary credential; write a 32-character administrator key to $ADMIN_KEY_PATH" >&2 exit 1 ;; esac - if ! admin_key_is_env_safe "$MSG_HEADER_KEY"; then + if ! admin_key_is_env_safe "$INSTALLER_KEY"; then echo "MSG_HEADER_KEY must be exactly 32 printable ASCII bytes without whitespace or NUL" >&2 exit 1 fi - printf '%s\n' "$MSG_HEADER_KEY" > "$ADMIN_KEY_PATH" + printf '%s\n' "$INSTALLER_KEY" > "$ADMIN_KEY_PATH" chmod 0600 "$ADMIN_KEY_PATH" echo "Persisted installer MSG_HEADER_KEY to $ADMIN_KEY_PATH" -elif [ -z "$(configured_msg_header_key)" ] && [ ! -s "$ADMIN_KEY_PATH" ] && [ -s "$LEGACY_KEY_PATH" ]; then +elif [ -z "$INSTALLER_KEY" ] && [ ! -s "$ADMIN_KEY_PATH" ] && [ -s "$LEGACY_KEY_PATH" ]; then install -m 0600 "$LEGACY_KEY_PATH" "$ADMIN_KEY_PATH" echo "Migrated the legacy machine-derived key into $ADMIN_KEY_PATH" fi diff --git a/scripts/install-server-github.sh b/scripts/install-server-github.sh index 20cc033..13cd1b1 100755 --- a/scripts/install-server-github.sh +++ b/scripts/install-server-github.sh @@ -103,21 +103,22 @@ install -m 0755 "$BIN_PATH" "${INSTALL_DIR}/pb-mapper" # key in the environment or /etc/pb-mapper/server.env must win; otherwise the # runtime would prefer the newly copied admin.key and lock operators out. install -d -m 0700 "$AUTH_DIR" -if [ -n "${MSG_HEADER_KEY:-}" ] && [ ! -s "$ADMIN_KEY_PATH" ]; then - case "$MSG_HEADER_KEY" in +INSTALLER_KEY="$(configured_msg_header_key)" +if [ -n "$INSTALLER_KEY" ] && [ ! -s "$ADMIN_KEY_PATH" ]; then + case "$INSTALLER_KEY" in pbmt1_*) echo "MSG_HEADER_KEY is a temporary credential; write a 32-character administrator key to $ADMIN_KEY_PATH" >&2 exit 1 ;; esac - if ! admin_key_is_env_safe "$MSG_HEADER_KEY"; then + if ! admin_key_is_env_safe "$INSTALLER_KEY"; then echo "MSG_HEADER_KEY must be exactly 32 printable ASCII bytes without whitespace or NUL" >&2 exit 1 fi - printf '%s\n' "$MSG_HEADER_KEY" > "$ADMIN_KEY_PATH" + printf '%s\n' "$INSTALLER_KEY" > "$ADMIN_KEY_PATH" chmod 0600 "$ADMIN_KEY_PATH" echo "Persisted installer MSG_HEADER_KEY to $ADMIN_KEY_PATH" -elif [ -z "$(configured_msg_header_key)" ] && [ ! -s "$ADMIN_KEY_PATH" ] && [ -s "$LEGACY_KEY_PATH" ]; then +elif [ -z "$INSTALLER_KEY" ] && [ ! -s "$ADMIN_KEY_PATH" ] && [ -s "$LEGACY_KEY_PATH" ]; then install -m 0600 "$LEGACY_KEY_PATH" "$ADMIN_KEY_PATH" echo "Migrated the legacy machine-derived key into $ADMIN_KEY_PATH" fi diff --git a/src/common/auth.rs b/src/common/auth.rs index e4f6413..dadcd74 100644 --- a/src/common/auth.rs +++ b/src/common/auth.rs @@ -808,6 +808,15 @@ fn load_server_admin_credential(state_dir: &Path) -> Result now { break; @@ -161,11 +162,17 @@ pub(super) async fn run_auth_actor( wheel.release(key_id); } } else { - let mut high = inner - .high_slot_entries - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - high.retain(|entry| entry.key_id != key_id); + due_high.push(key_id); + } + } + if !due_high.is_empty() { + let due = due_high.iter().copied().collect::>(); + inner + .high_slot_entries + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .retain(|entry| !due.contains(&entry.key_id)); + for key_id in due_high { cold.remove(&key_id); wheel.release(key_id); } diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index 6559c34..11f80d6 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -319,6 +319,46 @@ async fn overlapping_runtimes_cannot_share_an_auth_state_directory() { let _ = std::fs::remove_dir_all(state_dir); } +#[tokio::test] +async fn env_recovery_key_is_not_written_when_it_cannot_decrypt_existing_state() { + let _process_credential_guard = PROCESS_CREDENTIAL_TEST_LOCK.lock().await; + let state_dir = temp_state_dir("env-key-must-match-snapshot"); + prepare_state_dir(&state_dir).unwrap(); + let good = *b"0123456789abcdefghijklmnopqrstuv"; + let bad = *b"abcdefghijklmnopqrstuvwxyz012345"; + let config = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 1, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + let snapshot = PersistedSnapshot { + schema_version: SNAPSHOT_SCHEMA_VERSION, + instance_id: [4_u8; INSTANCE_ID_LEN], + generations: vec![0; 1], + entries: Vec::new(), + legacy_protocol: LegacyProtocolPolicy::Allow, + admin_replays: Vec::new(), + audit_records: VecDeque::new(), + root_epoch: 0, + }; + write_snapshot_and_truncate_wal(&config, &good, &snapshot).unwrap(); + set_process_msg_header_key(Some(std::str::from_utf8(&bad).unwrap())).unwrap(); + std::env::set_var(ENV_MSG_HEADER_KEY, std::str::from_utf8(&bad).unwrap()); + let error = match AuthRuntime::from_process(config).await { + Ok(_) => panic!("a mismatched recovery key must not start the runtime"), + Err(error) => error, + }; + std::env::remove_var(ENV_MSG_HEADER_KEY); + set_process_msg_header_key(None).unwrap(); + assert_eq!(error.code, "administrator_key_invalid"); + assert!( + !state_dir.join("admin.key").exists(), + "a mismatched MSG_HEADER_KEY must not become the live administrator key" + ); + let _ = std::fs::remove_dir_all(state_dir); +} + #[tokio::test] async fn from_isolated_state_takes_the_state_lock_before_creating_admin_key() { let state_dir = temp_state_dir("lock-before-key"); diff --git a/src/common/message/secure.rs b/src/common/message/secure.rs index ad3d3c9..2a248ee 100644 --- a/src/common/message/secure.rs +++ b/src/common/message/secure.rs @@ -585,59 +585,31 @@ impl ServerSecurity { } }; - let context = self - .auth - .authenticate_presented(key_id, &key) - .map_err(|failure| ServerInitialError { - failure, - response_session: if already_admitted { - None - } else { - Some(session_without_context(&session)) - }, - presented_key_id: Some(key_id), - })?; - match self - .replay - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .admit(key_id, &fingerprint, unix_seconds()) - { - FirstFlightAdmit::Replayed => { - return Err(ServerInitialError { - failure: AuthFailure::new( - "connection_salt_replayed", - "protocol-v2 connection salt was already accepted", - true, - ), - response_session: None, - presented_key_id: Some(key_id), - }); - } - FirstFlightAdmit::Limited => { - return Err(ServerInitialError { - failure: AuthFailure::new( - "connection_admission_limited", - "this credential has opened too many new connections in the current window", - true, - ), - response_session: Some(session), - presented_key_id: Some(key_id), - }); - } - FirstFlightAdmit::Unavailable => { - return Err(ServerInitialError { - failure: AuthFailure::new( - "connection_replay_store_unavailable", - "failed to persist first-flight replay admission", - true, - ), - response_session: Some(session), - presented_key_id: Some(key_id), - }); - } - FirstFlightAdmit::Fresh => {} - } + let replay = self.replay.clone(); + let auth = self.auth.clone(); + let error_session = session_without_context(&session); + let limited_session = session_without_context(&session); + let context = tokio::task::spawn_blocking(move || { + authenticate_and_admit( + &auth, + &replay, + key_id, + key, + fingerprint, + error_session, + limited_session, + ) + }) + .await + .map_err(|_| ServerInitialError { + failure: AuthFailure::new( + "connection_replay_store_unavailable", + "failed to evaluate first-flight admission", + true, + ), + response_session: None, + presented_key_id: Some(key_id), + })??; session.context = Some(context); Ok(ServerInitialMessage { payload, @@ -681,6 +653,67 @@ async fn read_initial_v2_ciphertext( Ok((counter, ciphertext)) } +#[allow(clippy::result_large_err)] +fn authenticate_and_admit( + auth: &AuthRuntime, + replay: &std::sync::Mutex, + key_id: u64, + key: AesKeyType, + fingerprint: [u8; 32], + error_session: ServerHeaderSession, + limited_session: ServerHeaderSession, +) -> std::result::Result { + let mut replay = replay + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let already_admitted = replay.already_admitted(&fingerprint, unix_seconds()); + let context = auth + .authenticate_presented(key_id, &key) + .map_err(|failure| ServerInitialError { + failure, + response_session: if already_admitted { + None + } else { + Some(error_session) + }, + presented_key_id: Some(key_id), + })?; + match replay.admit(key_id, &fingerprint, unix_seconds()) { + FirstFlightAdmit::Fresh => Ok(context), + FirstFlightAdmit::Replayed => Err(ServerInitialError { + failure: AuthFailure::new( + "connection_salt_replayed", + "protocol-v2 connection salt was already accepted", + true, + ), + response_session: None, + presented_key_id: Some(key_id), + }), + FirstFlightAdmit::Limited => Err(ServerInitialError { + failure: AuthFailure::new( + "connection_admission_limited", + "this credential has opened too many new connections in the current window", + true, + ), + response_session: if already_admitted { + None + } else { + Some(limited_session) + }, + presented_key_id: Some(key_id), + }), + FirstFlightAdmit::Unavailable => Err(ServerInitialError { + failure: AuthFailure::new( + "connection_replay_store_unavailable", + "failed to persist first-flight replay admission", + true, + ), + response_session: None, + presented_key_id: Some(key_id), + }), + } +} + fn stale_root_first_flight( auth: &AuthRuntime, key_id: u64, diff --git a/src/common/message/secure/replay.rs b/src/common/message/secure/replay.rs index a555e08..04cf067 100644 --- a/src/common/message/secure/replay.rs +++ b/src/common/message/secure/replay.rs @@ -157,7 +157,12 @@ impl ReplayGuard { } pub(super) fn already_admitted(&mut self, fingerprint: &[u8; 32], now: u64) -> bool { - self.bloom.contains(fingerprint, now) + let present = self.bloom.contains(fingerprint, now); + if self.bloom.current_started_at != self.total_started_at { + self.total = 0; + self.total_started_at = self.bloom.current_started_at; + } + present } pub(super) fn admit( @@ -269,6 +274,13 @@ impl ReplayGuard { self.bloom.insert(&fingerprint, timestamp); live.push(record); } + if let Some(oldest) = live + .iter() + .map(|record| u64::from_be_bytes(record[32..].try_into().expect("timestamp width"))) + .min() + { + self.bloom.current_started_at = oldest; + } self.total = u32::try_from(live.len()) .unwrap_or(u32::MAX) .min(self.max_total); diff --git a/src/common/message/secure/tests.rs b/src/common/message/secure/tests.rs index ee82ad1..5f61de3 100644 --- a/src/common/message/secure/tests.rs +++ b/src/common/message/secure/tests.rs @@ -574,6 +574,39 @@ fn restored_replay_log_consumes_the_aggregate_budget() { let _ = std::fs::remove_file(path); } +#[test] +fn restored_replay_generation_ages_with_loaded_records() { + let mut random = [0_u8; 8]; + let mut rng = rand::rng(); + for byte in &mut random { + *byte = rng.random(); + } + let path = std::env::temp_dir().join(format!( + "pb-mapper-replay-age-{}", + u64::from_be_bytes(random) + )); + let start = unix_seconds().saturating_sub(DEFAULT_REPLAY_WINDOW_SECONDS / 2); + { + let mut guard = ReplayGuard::open(Some(path.clone()), 1024, DEFAULT_REPLAY_WINDOW_SECONDS) + .with_max_per_key(100) + .with_max_total(2); + assert_eq!(guard.admit(1, &[1_u8; 32], start), FirstFlightAdmit::Fresh); + assert_eq!(guard.admit(2, &[2_u8; 32], start), FirstFlightAdmit::Fresh); + } + let mut restored = ReplayGuard::open(Some(path.clone()), 1024, DEFAULT_REPLAY_WINDOW_SECONDS) + .with_max_per_key(100) + .with_max_total(2); + assert_eq!( + restored.admit( + 3, + &[3_u8; 32], + start.saturating_add(DEFAULT_REPLAY_WINDOW_SECONDS) + ), + FirstFlightAdmit::Fresh + ); + let _ = std::fs::remove_file(path); +} + #[tokio::test] async fn legacy_initial_frame_validates_against_isolated_relay_key() { let _process_credential_guard = PROCESS_CREDENTIAL_TEST_LOCK.lock().await; diff --git a/ui/native/pb_mapper_ffi/src/state.rs b/ui/native/pb_mapper_ffi/src/state.rs index 7c969d6..703635e 100644 --- a/ui/native/pb_mapper_ffi/src/state.rs +++ b/ui/native/pb_mapper_ffi/src/state.rs @@ -432,6 +432,7 @@ struct RegisterCommit { enable_keep_alive: bool, local_sock_addr: SocketAddr, remote_sock_addr: SocketAddr, + credential: Credential, } /// Everything [`PbMapperState::finish_connect`] needs once the slow work is done. @@ -442,6 +443,7 @@ struct ConnectCommit { enable_keep_alive: bool, local_sock_addr: SocketAddr, remote_sock_addr: SocketAddr, + credential: Credential, } pub struct PbMapperState { @@ -494,10 +496,13 @@ pub async fn register_service( // 1. Claim the key and take what the slow work needs. Microseconds. // `_claim` is held to the end of the function on purpose: dropping it // early would release the key while the setup is still running. - let (_claim, server_address) = { + // The credential is captured with the relay address so a later config + // change cannot pair a new key with the already-resolved socket. + let (_claim, server_address, credential) = { let state = state.lock().await; let claim = state.claim_registering(&service_key)?; - (claim, state.config.server_address.clone()) + let credential = get_process_credential().map_err(CtlError::invalid_argument)?; + (claim, state.config.server_address.clone(), credential) }; // 2. The slow parts, with the lock released. @@ -524,6 +529,7 @@ pub async fn register_service( enable_keep_alive, local_sock_addr, remote_sock_addr, + credential, }) .await } @@ -537,10 +543,11 @@ pub async fn connect_service( protocol: String, enable_keep_alive: bool, ) -> Result<(), CtlError> { - let (_claim, server_address) = { + let (_claim, server_address, credential) = { let state = state.lock().await; let claim = state.claim_connecting(&service_key)?; - (claim, state.config.server_address.clone()) + let credential = get_process_credential().map_err(CtlError::invalid_argument)?; + (claim, state.config.server_address.clone(), credential) }; let local_sock_addr = get_sockaddr_async(&local_address) @@ -581,6 +588,7 @@ pub async fn connect_service( enable_keep_alive, local_sock_addr, remote_sock_addr, + credential, }) .await } @@ -688,6 +696,15 @@ mod tests { #[tokio::test] async fn a_failed_registration_releases_its_claim() { let (state, root) = temp_state("release"); + struct RestoreProcessKey; + impl Drop for RestoreProcessKey { + fn drop(&mut self) { + let _ = set_process_msg_header_key(None); + } + } + set_process_msg_header_key(Some("0123456789abcdefghijklmnopqrstuv")) + .expect("test credential"); + let _restore_process_key = RestoreProcessKey; // Fails in phase 2, while the claim is held. let first = register_service( diff --git a/ui/native/pb_mapper_ffi/src/state/runtime.rs b/ui/native/pb_mapper_ffi/src/state/runtime.rs index 9ad35a8..edb48f4 100644 --- a/ui/native/pb_mapper_ffi/src/state/runtime.rs +++ b/ui/native/pb_mapper_ffi/src/state/runtime.rs @@ -158,9 +158,9 @@ impl PbMapperState { enable_keep_alive, local_sock_addr, remote_sock_addr, + credential, } = commit; - let credential = get_process_credential().map_err(CtlError::invalid_argument)?; self.service_credentials.remove(&service_key); if let Some(previous) = self.service_handles.remove(&service_key) { tracing::warn!( @@ -315,9 +315,9 @@ impl PbMapperState { enable_keep_alive, local_sock_addr, remote_sock_addr, + credential, } = commit; - let credential = get_process_credential().map_err(CtlError::invalid_argument)?; self.client_credentials.remove(&service_key); if let Some(previous) = self.client_handles.remove(&service_key) { tracing::warn!( From baddd546c188303492dbe82bb2eede6ffdbc05ca Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 10:42:06 +0800 Subject: [PATCH 50/74] Claim limited and stale-root first flights Replay waits stay off Tokio workers. Limited and stale-root salts are reserved before a nonce-0 error. Recovery keys must decrypt WAL-only state, and UI config rolls back if persistence fails. --- CHANGELOG.md | 1 + scripts/install-server-gitee.sh | 18 +++-- scripts/install-server-github.sh | 18 +++-- src/common/auth.rs | 11 ++- src/common/auth/persistence.rs | 47 +++++++++++- src/common/auth/tests.rs | 76 +++++++++++++++++++ src/common/message/secure.rs | 70 ++++++++++++++--- src/common/message/secure/replay.rs | 23 ++++-- src/common/message/secure/tests.rs | 49 +++++++++++- ui/native/pb_mapper_ffi/src/state.rs | 2 + .../pb_mapper_ffi/src/state/configuration.rs | 8 +- ui/native/pb_mapper_ffi/src/state/runtime.rs | 10 +++ ui/native/pb_mapper_ffi/src/state/status.rs | 31 ++++++-- 13 files changed, 326 insertions(+), 38 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cbe3c60..de19f79 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -47,6 +47,7 @@ All notable changes to this project will be documented in this file. - Restricted administrator CLI retries to pre-send failures and a readable `connection_salt_replayed`, so a dropped response cannot issue a second credential. - Omitted a response session for already-admitted first flights that later fail authentication, and kept namespace stream accounting until the client stream deregisters even if the registration control socket drops. - Evaluated first-flight admission on a blocking thread under one replay lock, omitted a session when durable admission is unavailable, aged restored Bloom generations from loaded record timestamps, validated installer keys from `server.env`, refused to persist a recovery `MSG_HEADER_KEY` that cannot decrypt existing state, batched high-slot tombstone cleanup, and captured the UI credential together with the relay address before DNS. +- Kept replay-lock waits off Tokio workers, claimed limited and stale-root first flights before sending a nonce-0 error, accepted a recovery key that decrypts WAL-only state, verified legacy and installer keys against existing state, rolled back UI config when persistence failed, and probed tunnels at their pinned relay endpoints. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/scripts/install-server-gitee.sh b/scripts/install-server-gitee.sh index ee92f7c..d8acf0a 100755 --- a/scripts/install-server-gitee.sh +++ b/scripts/install-server-gitee.sh @@ -115,12 +115,20 @@ if [ -n "$INSTALLER_KEY" ] && [ ! -s "$ADMIN_KEY_PATH" ]; then echo "MSG_HEADER_KEY must be exactly 32 printable ASCII bytes without whitespace or NUL" >&2 exit 1 fi - printf '%s\n' "$INSTALLER_KEY" > "$ADMIN_KEY_PATH" - chmod 0600 "$ADMIN_KEY_PATH" - echo "Persisted installer MSG_HEADER_KEY to $ADMIN_KEY_PATH" + if [ -s "${AUTH_DIR}/auth.snapshot" ] || [ -s "${AUTH_DIR}/auth.wal" ]; then + echo "Leaving $ADMIN_KEY_PATH unset so the service can verify MSG_HEADER_KEY against existing authentication state" + else + printf '%s\n' "$INSTALLER_KEY" > "$ADMIN_KEY_PATH" + chmod 0600 "$ADMIN_KEY_PATH" + echo "Persisted installer MSG_HEADER_KEY to $ADMIN_KEY_PATH" + fi elif [ -z "$INSTALLER_KEY" ] && [ ! -s "$ADMIN_KEY_PATH" ] && [ -s "$LEGACY_KEY_PATH" ]; then - install -m 0600 "$LEGACY_KEY_PATH" "$ADMIN_KEY_PATH" - echo "Migrated the legacy machine-derived key into $ADMIN_KEY_PATH" + if [ -s "${AUTH_DIR}/auth.snapshot" ] || [ -s "${AUTH_DIR}/auth.wal" ]; then + echo "Leaving $ADMIN_KEY_PATH unset so the service can verify the legacy key against existing authentication state" + else + install -m 0600 "$LEGACY_KEY_PATH" "$ADMIN_KEY_PATH" + echo "Migrated the legacy machine-derived key into $ADMIN_KEY_PATH" + fi fi # Stop and remove existing service if present diff --git a/scripts/install-server-github.sh b/scripts/install-server-github.sh index 13cd1b1..13acdc4 100755 --- a/scripts/install-server-github.sh +++ b/scripts/install-server-github.sh @@ -115,12 +115,20 @@ if [ -n "$INSTALLER_KEY" ] && [ ! -s "$ADMIN_KEY_PATH" ]; then echo "MSG_HEADER_KEY must be exactly 32 printable ASCII bytes without whitespace or NUL" >&2 exit 1 fi - printf '%s\n' "$INSTALLER_KEY" > "$ADMIN_KEY_PATH" - chmod 0600 "$ADMIN_KEY_PATH" - echo "Persisted installer MSG_HEADER_KEY to $ADMIN_KEY_PATH" + if [ -s "${AUTH_DIR}/auth.snapshot" ] || [ -s "${AUTH_DIR}/auth.wal" ]; then + echo "Leaving $ADMIN_KEY_PATH unset so the service can verify MSG_HEADER_KEY against existing authentication state" + else + printf '%s\n' "$INSTALLER_KEY" > "$ADMIN_KEY_PATH" + chmod 0600 "$ADMIN_KEY_PATH" + echo "Persisted installer MSG_HEADER_KEY to $ADMIN_KEY_PATH" + fi elif [ -z "$INSTALLER_KEY" ] && [ ! -s "$ADMIN_KEY_PATH" ] && [ -s "$LEGACY_KEY_PATH" ]; then - install -m 0600 "$LEGACY_KEY_PATH" "$ADMIN_KEY_PATH" - echo "Migrated the legacy machine-derived key into $ADMIN_KEY_PATH" + if [ -s "${AUTH_DIR}/auth.snapshot" ] || [ -s "${AUTH_DIR}/auth.wal" ]; then + echo "Leaving $ADMIN_KEY_PATH unset so the service can verify the legacy key against existing authentication state" + else + install -m 0600 "$LEGACY_KEY_PATH" "$ADMIN_KEY_PATH" + echo "Migrated the legacy machine-derived key into $ADMIN_KEY_PATH" + fi fi # Stop and remove existing service if present diff --git a/src/common/auth.rs b/src/common/auth.rs index dadcd74..512a14c 100644 --- a/src/common/auth.rs +++ b/src/common/auth.rs @@ -809,7 +809,7 @@ fn load_server_admin_credential(state_dir: &Path) -> Result Result Result<(), A )); } if path.file_name() == Some(std::ffi::OsStr::new("admin.key")) - && !key_matches_existing_snapshot(path.parent(), key) + && !key_matches_existing_state(path.parent(), key) { refuse_write_if_encrypted_state(path, force)?; } @@ -1034,6 +1034,51 @@ pub(super) fn key_matches_existing_snapshot(state_dir: Option<&Path>, key: &str) open_blob(&admin_key, &bytes).is_ok() } +pub(super) fn key_matches_existing_state(state_dir: Option<&Path>, key: &str) -> bool { + if key_matches_existing_snapshot(state_dir, key) { + return true; + } + let Some(state_dir) = state_dir else { + return false; + }; + if auth_snapshot_path(state_dir).exists() { + return false; + } + let wal_path = auth_wal_path(state_dir); + if !wal_path.exists() { + return false; + } + let Ok(Credential::Admin(admin_key)) = parse_credential(key) else { + return false; + }; + wal_decrypts_with_key(&wal_path, &admin_key) +} + +fn wal_decrypts_with_key(path: &Path, admin_key: &AesKeyType) -> bool { + let Ok(mut file) = File::open(path) else { + return false; + }; + let Ok(metadata) = file.metadata() else { + return false; + }; + if metadata.len() == 0 { + return true; + } + let mut length = [0_u8; 4]; + if file.read_exact(&mut length).is_err() { + return false; + } + let length = u32::from_be_bytes(length) as usize; + if length == 0 || length > 1024 * 1024 { + return false; + } + let mut sealed = vec![0_u8; length]; + if file.read_exact(&mut sealed).is_err() { + return false; + } + open_blob(admin_key, &sealed).is_ok() +} + fn refuse_write_if_encrypted_state(path: &Path, force: bool) -> Result<(), AuthFailure> { // Creating or replacing the live root while snapshot/WAL remain leaves // those files encrypted under the previous key. Staging `admin.key.next` diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index 11f80d6..77510ba 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -359,6 +359,82 @@ async fn env_recovery_key_is_not_written_when_it_cannot_decrypt_existing_state() let _ = std::fs::remove_dir_all(state_dir); } +#[tokio::test] +async fn env_recovery_key_is_accepted_for_wal_only_state() { + let _process_credential_guard = PROCESS_CREDENTIAL_TEST_LOCK.lock().await; + let state_dir = temp_state_dir("env-key-matches-wal"); + prepare_state_dir(&state_dir).unwrap(); + let good = *b"0123456789abcdefghijklmnopqrstuv"; + let config = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 1, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + append_wal( + &config, + &good, + &WalRecord::Audit(AuditRecord { + at: 1, + action: "issue".to_string(), + key_id: None, + label: None, + }), + ) + .unwrap(); + set_process_msg_header_key(Some(std::str::from_utf8(&good).unwrap())).unwrap(); + std::env::set_var(ENV_MSG_HEADER_KEY, std::str::from_utf8(&good).unwrap()); + let started = AuthRuntime::from_process(config).await; + std::env::remove_var(ENV_MSG_HEADER_KEY); + set_process_msg_header_key(None).unwrap(); + started.expect("a matching recovery key must start from WAL-only state"); + assert!( + state_dir.join("admin.key").exists(), + "a matching MSG_HEADER_KEY should become the live administrator key" + ); + let _ = std::fs::remove_dir_all(state_dir); +} + +#[tokio::test] +async fn env_recovery_key_is_not_written_when_wal_only_state_does_not_match() { + let _process_credential_guard = PROCESS_CREDENTIAL_TEST_LOCK.lock().await; + let state_dir = temp_state_dir("env-key-must-match-wal"); + prepare_state_dir(&state_dir).unwrap(); + let good = *b"0123456789abcdefghijklmnopqrstuv"; + let bad = *b"abcdefghijklmnopqrstuvwxyz012345"; + let config = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 1, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + append_wal( + &config, + &good, + &WalRecord::Audit(AuditRecord { + at: 1, + action: "issue".to_string(), + key_id: None, + label: None, + }), + ) + .unwrap(); + set_process_msg_header_key(Some(std::str::from_utf8(&bad).unwrap())).unwrap(); + std::env::set_var(ENV_MSG_HEADER_KEY, std::str::from_utf8(&bad).unwrap()); + let error = match AuthRuntime::from_process(config).await { + Ok(_) => panic!("a mismatched recovery key must not start from WAL-only state"), + Err(error) => error, + }; + std::env::remove_var(ENV_MSG_HEADER_KEY); + set_process_msg_header_key(None).unwrap(); + assert_eq!(error.code, "administrator_key_invalid"); + assert!( + !state_dir.join("admin.key").exists(), + "a mismatched MSG_HEADER_KEY must not become the live administrator key" + ); + let _ = std::fs::remove_dir_all(state_dir); +} + #[tokio::test] async fn from_isolated_state_takes_the_state_lock_before_creating_admin_key() { let state_dir = temp_state_dir("lock-before-key"); diff --git a/src/common/message/secure.rs b/src/common/message/secure.rs index 2a248ee..2c2d192 100644 --- a/src/common/message/secure.rs +++ b/src/common/message/secure.rs @@ -552,11 +552,6 @@ impl ServerSecurity { })?; let mut current_ciphertext = ciphertext.clone(); let fingerprint = replay_fingerprint(key_id, &salt); - let already_admitted = self - .replay - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .already_admitted(&fingerprint, unix_seconds()); let payload = match open_v2_payload( &material, DIRECTION_CLIENT_TO_SERVER, @@ -565,13 +560,31 @@ impl ServerSecurity { ) { Ok(payload) => payload, Err(error) => { - if let Some(mut stale) = - stale_root_first_flight(&self.auth, key_id, salt, counter, &ciphertext) + if stale_root_first_flight(&self.auth, key_id, salt, counter, &ciphertext).is_some() { - if already_admitted { - stale.response_session = None; - } - return Err(stale); + let replay = self.replay.clone(); + let auth = self.auth.clone(); + return Err(tokio::task::spawn_blocking(move || { + claim_stale_root_first_flight( + &auth, + &replay, + key_id, + salt, + counter, + ciphertext, + fingerprint, + ) + }) + .await + .unwrap_or_else(|_| ServerInitialError { + failure: AuthFailure::new( + "connection_replay_store_unavailable", + "failed to evaluate first-flight admission", + true, + ), + response_session: None, + presented_key_id: Some(key_id), + })); } return Err(ServerInitialError { failure: AuthFailure::new( @@ -714,6 +727,41 @@ fn authenticate_and_admit( } } +#[allow(clippy::result_large_err)] +fn claim_stale_root_first_flight( + auth: &AuthRuntime, + replay: &std::sync::Mutex, + key_id: u64, + salt: [u8; CONNECTION_SALT_LEN], + counter: u64, + ciphertext: Vec, + fingerprint: [u8; 32], +) -> ServerInitialError { + let mut error = stale_root_first_flight(auth, key_id, salt, counter, &ciphertext) + .unwrap_or_else(|| ServerInitialError { + failure: AuthFailure::new( + "protocol_v2_decrypt_failed", + "stale-root first flight could not be classified", + false, + ), + response_session: None, + presented_key_id: Some(key_id), + }); + let mut replay = replay + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let already_admitted = replay.already_admitted(&fingerprint, unix_seconds()); + if already_admitted + || !matches!( + replay.claim(&fingerprint, unix_seconds()), + FirstFlightAdmit::Fresh + ) + { + error.response_session = None; + } + error +} + fn stale_root_first_flight( auth: &AuthRuntime, key_id: u64, diff --git a/src/common/message/secure/replay.rs b/src/common/message/secure/replay.rs index 04cf067..4deb5ec 100644 --- a/src/common/message/secure/replay.rs +++ b/src/common/message/secure/replay.rs @@ -179,11 +179,13 @@ impl ReplayGuard { self.total = 0; self.total_started_at = self.bloom.current_started_at; } - if self.counts.get(&key_id).copied().unwrap_or(0) >= self.max_per_key { - return FirstFlightAdmit::Limited; - } - if self.total >= self.max_total { - return FirstFlightAdmit::Limited; + if self.counts.get(&key_id).copied().unwrap_or(0) >= self.max_per_key + || self.total >= self.max_total + { + return match self.claim(fingerprint, now) { + FirstFlightAdmit::Unavailable => FirstFlightAdmit::Unavailable, + _ => FirstFlightAdmit::Limited, + }; } if self.persist(fingerprint, now).is_err() { return FirstFlightAdmit::Unavailable; @@ -197,6 +199,17 @@ impl ReplayGuard { FirstFlightAdmit::Fresh } + pub(super) fn claim(&mut self, fingerprint: &[u8; 32], now: u64) -> FirstFlightAdmit { + if self.bloom.contains(fingerprint, now) { + return FirstFlightAdmit::Replayed; + } + if self.persist(fingerprint, now).is_err() { + return FirstFlightAdmit::Unavailable; + } + self.bloom.insert(fingerprint, now); + FirstFlightAdmit::Fresh + } + fn rotate_counts(&mut self, now: u64) { if now.saturating_sub(self.counts_started_at) < self.window_seconds { return; diff --git a/src/common/message/secure/tests.rs b/src/common/message/secure/tests.rs index 5f61de3..3da7b64 100644 --- a/src/common/message/secure/tests.rs +++ b/src/common/message/secure/tests.rs @@ -271,6 +271,51 @@ async fn rotated_temporary_first_flight_returns_a_readable_rotated_error() { let _ = std::fs::remove_dir_all(config.state_dir); } +#[tokio::test] +async fn stale_root_replay_omits_the_rotated_error_session() { + let _process_credential_guard = PROCESS_CREDENTIAL_TEST_LOCK.lock().await; + let admin = *b"0123456789abcdefghijklmnopqrstuv"; + let new_admin = *b"abcdefghijklmnopqrstuvwxyz012345"; + let config = temp_config(); + let auth = AuthRuntime::start(admin, config.clone()).await.unwrap(); + let admin_context = auth.authenticate_presented(0, &admin).unwrap(); + let issued = auth + .issue(&admin_context, std::time::Duration::from_secs(60), None) + .await + .unwrap(); + let Credential::Temporary { key_id, key } = parse_credential(&issued.credential).unwrap() + else { + panic!("expected temporary credential"); + }; + let client = ClientHeaderSession::new_v2(&Credential::Temporary { key_id, key }).unwrap(); + let bytes = encode_initial(&client, b"stale-replay").await; + auth.rotate_root(&admin_context, new_admin).await.unwrap(); + let security = ServerSecurity::new(auth); + let first = match security + .read_initial(&mut std::io::Cursor::new(bytes.clone())) + .await + { + Ok(_) => panic!("rotated credential should fail"), + Err(error) => error, + }; + let second = match security + .read_initial(&mut std::io::Cursor::new(bytes)) + .await + { + Ok(_) => panic!("rotated credential replay should fail"), + Err(error) => error, + }; + assert_eq!(first.failure.code, "temporary_key_rotated"); + assert!(first.response_session.is_some()); + assert_eq!(second.failure.code, "temporary_key_rotated"); + assert!( + second.response_session.is_none(), + "a claimed stale-root salt must not reuse nonce 0" + ); + + let _ = std::fs::remove_dir_all(config.state_dir); +} + #[tokio::test] async fn reset_temporary_first_flight_returns_a_readable_rotated_error() { let _process_credential_guard = PROCESS_CREDENTIAL_TEST_LOCK.lock().await; @@ -447,7 +492,8 @@ fn per_credential_admission_limit_does_not_consume_other_keys() { assert_eq!(guard.admit(1, &[1_u8; 32], now), FirstFlightAdmit::Fresh); assert_eq!(guard.admit(1, &[2_u8; 32], now), FirstFlightAdmit::Fresh); assert_eq!(guard.admit(1, &[3_u8; 32], now), FirstFlightAdmit::Limited); - assert_eq!(guard.admit(2, &[3_u8; 32], now), FirstFlightAdmit::Fresh); + assert_eq!(guard.admit(1, &[3_u8; 32], now), FirstFlightAdmit::Replayed); + assert_eq!(guard.admit(2, &[4_u8; 32], now), FirstFlightAdmit::Fresh); assert_eq!(guard.admit(1, &[1_u8; 32], now), FirstFlightAdmit::Replayed); } @@ -460,6 +506,7 @@ fn aggregate_admission_limit_covers_all_keys() { assert_eq!(guard.admit(1, &[1_u8; 32], now), FirstFlightAdmit::Fresh); assert_eq!(guard.admit(2, &[2_u8; 32], now), FirstFlightAdmit::Fresh); assert_eq!(guard.admit(3, &[3_u8; 32], now), FirstFlightAdmit::Limited); + assert_eq!(guard.admit(3, &[3_u8; 32], now), FirstFlightAdmit::Replayed); } #[test] diff --git a/ui/native/pb_mapper_ffi/src/state.rs b/ui/native/pb_mapper_ffi/src/state.rs index 703635e..01c88b4 100644 --- a/ui/native/pb_mapper_ffi/src/state.rs +++ b/ui/native/pb_mapper_ffi/src/state.rs @@ -459,6 +459,8 @@ pub struct PbMapperState { client_handles: HashMap>, service_credentials: HashMap, client_credentials: HashMap, + service_endpoints: HashMap, + client_endpoints: HashMap, config: AppConfig, config_dir: PathBuf, app_directory_path: Option, diff --git a/ui/native/pb_mapper_ffi/src/state/configuration.rs b/ui/native/pb_mapper_ffi/src/state/configuration.rs index fc65331..601856c 100644 --- a/ui/native/pb_mapper_ffi/src/state/configuration.rs +++ b/ui/native/pb_mapper_ffi/src/state/configuration.rs @@ -57,6 +57,8 @@ impl PbMapperState { client_handles: HashMap::new(), service_credentials: HashMap::new(), client_credentials: HashMap::new(), + service_endpoints: HashMap::new(), + client_endpoints: HashMap::new(), config: AppConfig::default(), config_dir: config_dir.clone(), app_directory_path: app_directory_path.clone(), @@ -95,6 +97,8 @@ impl PbMapperState { client_handles: HashMap::new(), service_credentials: HashMap::new(), client_credentials: HashMap::new(), + service_endpoints: HashMap::new(), + client_endpoints: HashMap::new(), config, config_dir, app_directory_path, @@ -199,10 +203,10 @@ impl PbMapperState { } } - pub fn save_config(&self) -> Result<(), CtlError> { + pub(super) fn write_config_file(&self, config: &AppConfig) -> Result<(), CtlError> { let config_path = self.get_config_file_path(); let contents = - serde_json::to_string_pretty(&self.config).map_err(|e| CtlError::io(e.to_string()))?; + serde_json::to_string_pretty(config).map_err(|e| CtlError::io(e.to_string()))?; fs::write(config_path, contents).map_err(|e| CtlError::io(e.to_string()))?; Ok(()) } diff --git a/ui/native/pb_mapper_ffi/src/state/runtime.rs b/ui/native/pb_mapper_ffi/src/state/runtime.rs index edb48f4..1e17703 100644 --- a/ui/native/pb_mapper_ffi/src/state/runtime.rs +++ b/ui/native/pb_mapper_ffi/src/state/runtime.rs @@ -133,11 +133,13 @@ impl PbMapperState { handle.abort(); } self.service_credentials.clear(); + self.service_endpoints.clear(); for (_, handle) in self.client_handles.drain() { handle.abort(); } self.client_credentials.clear(); + self.client_endpoints.clear(); self.registered_services.write().await.clear(); self.active_connections.write().await.clear(); @@ -202,6 +204,8 @@ impl PbMapperState { self.service_credentials .insert(service_key.clone(), credential); + self.service_endpoints + .insert(service_key.clone(), remote_sock_addr); let handle = if protocol.to_uppercase() == "TCP" { tokio::spawn(async move { let _ = run_server_side_cli_with_pinned_credential::( @@ -273,6 +277,7 @@ impl PbMapperState { pub async fn unregister_service(&mut self, service_key: String) -> Result<(), CtlError> { self.service_credentials.remove(&service_key); + self.service_endpoints.remove(&service_key); if let Some(handle) = self.service_handles.remove(&service_key) { handle.abort(); } @@ -298,6 +303,7 @@ impl PbMapperState { service_key: String, ) -> Result<(), CtlError> { self.service_credentials.remove(&service_key); + self.service_endpoints.remove(&service_key); if let Some(handle) = self.service_handles.remove(&service_key) { handle.abort(); } @@ -349,6 +355,8 @@ impl PbMapperState { self.client_credentials .insert(service_key.clone(), credential); + self.client_endpoints + .insert(service_key.clone(), remote_sock_addr); let handle = if protocol_upper == "TCP" { tokio::spawn(async move { run_client_side_cli_with_pinned_credential::( @@ -431,6 +439,7 @@ impl PbMapperState { // Aborting the task is the part that matters: it is what stops the // retry loop still dialling in the background. self.client_credentials.remove(&service_key); + self.client_endpoints.remove(&service_key); let aborted = match self.client_handles.remove(&service_key) { Some(handle) => { handle.abort(); @@ -465,6 +474,7 @@ impl PbMapperState { service_key: String, ) -> Result<(), CtlError> { self.client_credentials.remove(&service_key); + self.client_endpoints.remove(&service_key); if let Some(handle) = self.client_handles.remove(&service_key) { handle.abort(); } diff --git a/ui/native/pb_mapper_ffi/src/state/status.rs b/ui/native/pb_mapper_ffi/src/state/status.rs index 9cf93be..e6fd8e2 100644 --- a/ui/native/pb_mapper_ffi/src/state/status.rs +++ b/ui/native/pb_mapper_ffi/src/state/status.rs @@ -33,11 +33,20 @@ impl PbMapperState { msg_header_key: String, ) -> Result<(), CtlError> { let msg_header_key = normalize_msg_header_key(msg_header_key)?; - self.config.server_address = server_address; - self.config.keep_alive_enabled = keep_alive; - self.config.msg_header_key = msg_header_key; - self.apply_msg_header_key_env()?; - self.save_config()?; + let previous = self.config.clone(); + let candidate = AppConfig { + server_address, + keep_alive_enabled: keep_alive, + msg_header_key, + }; + self.write_config_file(&candidate)?; + self.config = candidate; + if let Err(error) = self.apply_msg_header_key_env() { + self.config = previous.clone(); + let _ = self.write_config_file(&previous); + let _ = self.apply_msg_header_key_env(); + return Err(error); + } self.reset_status_caches().await; Ok(()) } @@ -360,7 +369,11 @@ impl PbMapperState { refreshing.insert(service_key.to_string()); } - let server_addr = self.config.server_address.clone(); + let server_addr = self + .service_endpoints + .get(service_key) + .map(ToString::to_string) + .unwrap_or_else(|| self.config.server_address.clone()); let cache = self.service_status_cache.clone(); let refreshing = self.service_status_refreshing.clone(); let key = service_key.to_string(); @@ -424,7 +437,11 @@ impl PbMapperState { refreshing.insert(service_key.to_string()); } - let server_addr = self.config.server_address.clone(); + let server_addr = self + .client_endpoints + .get(service_key) + .map(ToString::to_string) + .unwrap_or_else(|| self.config.server_address.clone()); let cache = self.client_status_cache.clone(); let refreshing = self.client_status_refreshing.clone(); let key = service_key.to_string(); From 9ef2ddeba329b7a3dbd407b153a023ca1417f181 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 10:52:47 +0800 Subject: [PATCH 51/74] Bound denial claims and verify container keys Limited first flights no longer persist unique salts after the Bloom budget is full, and they omit a nonce-0 error session. The container entrypoint leaves admin.key unset when encrypted state remains so the runtime can verify a legacy recovery key. --- CHANGELOG.md | 1 + scripts/release/entrypoint/pb-mapper.sh | 8 ++++++-- src/common/message/secure.rs | 18 ++---------------- src/common/message/secure/replay.rs | 13 +++++++++---- src/common/message/secure/tests.rs | 5 +++-- 5 files changed, 21 insertions(+), 24 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index de19f79..b8dfcfd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -48,6 +48,7 @@ All notable changes to this project will be documented in this file. - Omitted a response session for already-admitted first flights that later fail authentication, and kept namespace stream accounting until the client stream deregisters even if the registration control socket drops. - Evaluated first-flight admission on a blocking thread under one replay lock, omitted a session when durable admission is unavailable, aged restored Bloom generations from loaded record timestamps, validated installer keys from `server.env`, refused to persist a recovery `MSG_HEADER_KEY` that cannot decrypt existing state, batched high-slot tombstone cleanup, and captured the UI credential together with the relay address before DNS. - Kept replay-lock waits off Tokio workers, claimed limited and stale-root first flights before sending a nonce-0 error, accepted a recovery key that decrypts WAL-only state, verified legacy and installer keys against existing state, rolled back UI config when persistence failed, and probed tunnels at their pinned relay endpoints. +- Stopped persisting limited first flights once the Bloom budget is full, omitted their nonce-0 error sessions, and left container `admin.key` unset when encrypted state remains so the runtime can verify a legacy recovery key. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/scripts/release/entrypoint/pb-mapper.sh b/scripts/release/entrypoint/pb-mapper.sh index 5b706c4..2c4258e 100644 --- a/scripts/release/entrypoint/pb-mapper.sh +++ b/scripts/release/entrypoint/pb-mapper.sh @@ -16,8 +16,12 @@ LEGACY_KEY_PATH="/var/lib/pb-mapper-server/msg_header_key" install -d -m 0700 "$AUTH_DIR" if [ -z "${MSG_HEADER_KEY:-}" ] && [ ! -s "$ADMIN_KEY_PATH" ] && [ -s "$LEGACY_KEY_PATH" ]; then - install -m 0600 "$LEGACY_KEY_PATH" "$ADMIN_KEY_PATH" - echo "Migrated the legacy machine-derived key into $ADMIN_KEY_PATH" + if [ -s "$AUTH_DIR/auth.snapshot" ] || [ -s "$AUTH_DIR/auth.wal" ]; then + echo "Leaving $ADMIN_KEY_PATH unset so the service can verify the legacy key against existing authentication state" + else + install -m 0600 "$LEGACY_KEY_PATH" "$ADMIN_KEY_PATH" + echo "Migrated the legacy machine-derived key into $ADMIN_KEY_PATH" + fi fi if [ "${USE_IPV6:-false}" = "true" ]; then diff --git a/src/common/message/secure.rs b/src/common/message/secure.rs index 2c2d192..217b7e8 100644 --- a/src/common/message/secure.rs +++ b/src/common/message/secure.rs @@ -601,17 +601,8 @@ impl ServerSecurity { let replay = self.replay.clone(); let auth = self.auth.clone(); let error_session = session_without_context(&session); - let limited_session = session_without_context(&session); let context = tokio::task::spawn_blocking(move || { - authenticate_and_admit( - &auth, - &replay, - key_id, - key, - fingerprint, - error_session, - limited_session, - ) + authenticate_and_admit(&auth, &replay, key_id, key, fingerprint, error_session) }) .await .map_err(|_| ServerInitialError { @@ -674,7 +665,6 @@ fn authenticate_and_admit( key: AesKeyType, fingerprint: [u8; 32], error_session: ServerHeaderSession, - limited_session: ServerHeaderSession, ) -> std::result::Result { let mut replay = replay .lock() @@ -708,11 +698,7 @@ fn authenticate_and_admit( "this credential has opened too many new connections in the current window", true, ), - response_session: if already_admitted { - None - } else { - Some(limited_session) - }, + response_session: None, presented_key_id: Some(key_id), }), FirstFlightAdmit::Unavailable => Err(ServerInitialError { diff --git a/src/common/message/secure/replay.rs b/src/common/message/secure/replay.rs index 4deb5ec..a27fe79 100644 --- a/src/common/message/secure/replay.rs +++ b/src/common/message/secure/replay.rs @@ -182,10 +182,7 @@ impl ReplayGuard { if self.counts.get(&key_id).copied().unwrap_or(0) >= self.max_per_key || self.total >= self.max_total { - return match self.claim(fingerprint, now) { - FirstFlightAdmit::Unavailable => FirstFlightAdmit::Unavailable, - _ => FirstFlightAdmit::Limited, - }; + return FirstFlightAdmit::Limited; } if self.persist(fingerprint, now).is_err() { return FirstFlightAdmit::Unavailable; @@ -203,10 +200,18 @@ impl ReplayGuard { if self.bloom.contains(fingerprint, now) { return FirstFlightAdmit::Replayed; } + if self.bloom.current_started_at != self.total_started_at { + self.total = 0; + self.total_started_at = self.bloom.current_started_at; + } + if self.total >= self.max_total { + return FirstFlightAdmit::Unavailable; + } if self.persist(fingerprint, now).is_err() { return FirstFlightAdmit::Unavailable; } self.bloom.insert(fingerprint, now); + self.total = self.total.saturating_add(1); FirstFlightAdmit::Fresh } diff --git a/src/common/message/secure/tests.rs b/src/common/message/secure/tests.rs index 3da7b64..c499e3d 100644 --- a/src/common/message/secure/tests.rs +++ b/src/common/message/secure/tests.rs @@ -492,7 +492,7 @@ fn per_credential_admission_limit_does_not_consume_other_keys() { assert_eq!(guard.admit(1, &[1_u8; 32], now), FirstFlightAdmit::Fresh); assert_eq!(guard.admit(1, &[2_u8; 32], now), FirstFlightAdmit::Fresh); assert_eq!(guard.admit(1, &[3_u8; 32], now), FirstFlightAdmit::Limited); - assert_eq!(guard.admit(1, &[3_u8; 32], now), FirstFlightAdmit::Replayed); + assert_eq!(guard.admit(1, &[3_u8; 32], now), FirstFlightAdmit::Limited); assert_eq!(guard.admit(2, &[4_u8; 32], now), FirstFlightAdmit::Fresh); assert_eq!(guard.admit(1, &[1_u8; 32], now), FirstFlightAdmit::Replayed); } @@ -506,7 +506,8 @@ fn aggregate_admission_limit_covers_all_keys() { assert_eq!(guard.admit(1, &[1_u8; 32], now), FirstFlightAdmit::Fresh); assert_eq!(guard.admit(2, &[2_u8; 32], now), FirstFlightAdmit::Fresh); assert_eq!(guard.admit(3, &[3_u8; 32], now), FirstFlightAdmit::Limited); - assert_eq!(guard.admit(3, &[3_u8; 32], now), FirstFlightAdmit::Replayed); + assert_eq!(guard.admit(3, &[3_u8; 32], now), FirstFlightAdmit::Limited); + assert_eq!(guard.admit(4, &[4_u8; 32], now), FirstFlightAdmit::Limited); } #[test] From ec52ab4515d0279baeb307764729901d99f8022c Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 11:09:46 +0800 Subject: [PATCH 52/74] Claim auth errors and schedule high-slot expiry Revoked or expired first flights reserve their salt before a nonce-0 error. High-slot expiries are scheduled instead of scanned every tick. Admin retries stay pre-send, connect fails without a credential, and truncated replay records fail closed. --- CHANGELOG.md | 1 + src/bin/pb-mapper.rs | 3 + src/bin/pb-mapper/admin.rs | 2 +- src/common/auth/actor.rs | 83 ++++++++++++++++++++++-- src/common/message/secure.rs | 20 ++++-- src/common/message/secure/replay.rs | 14 ++-- src/common/message/secure/tests.rs | 7 +- ui/lib/src/views/configuration_view.dart | 8 ++- ui/lib/src/views/setup_wizard_view.dart | 2 +- 9 files changed, 120 insertions(+), 20 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b8dfcfd..fc83c4a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -49,6 +49,7 @@ All notable changes to this project will be documented in this file. - Evaluated first-flight admission on a blocking thread under one replay lock, omitted a session when durable admission is unavailable, aged restored Bloom generations from loaded record timestamps, validated installer keys from `server.env`, refused to persist a recovery `MSG_HEADER_KEY` that cannot decrypt existing state, batched high-slot tombstone cleanup, and captured the UI credential together with the relay address before DNS. - Kept replay-lock waits off Tokio workers, claimed limited and stale-root first flights before sending a nonce-0 error, accepted a recovery key that decrypts WAL-only state, verified legacy and installer keys against existing state, rolled back UI config when persistence failed, and probed tunnels at their pinned relay endpoints. - Stopped persisting limited first flights once the Bloom budget is full, omitted their nonce-0 error sessions, and left container `admin.key` unset when encrypted state remains so the runtime can verify a legacy recovery key. +- Claimed revoked or expired first flights before returning a nonce-0 error, scheduled high-slot expiries instead of scanning every tick, marked administrator requests sent only after the first flight is written, allowed clearing the UI credential, failed `pb-mapper connect` without a credential, and fail-closed truncated replay records. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/src/bin/pb-mapper.rs b/src/bin/pb-mapper.rs index cd9cac6..117c4d3 100644 --- a/src/bin/pb-mapper.rs +++ b/src/bin/pb-mapper.rs @@ -348,6 +348,9 @@ async fn register( } async fn run_connect(args: ConnectArgs) -> Result<(), Box> { + pb_mapper::common::checksum::get_process_credential().map_err(|error| { + std::io::Error::other(format!("client credential is required: {error}")) + })?; let local_addr = get_sockaddr_async(&args.addr).await?; let remote_addr = get_pb_mapper_server_async(args.relay.server.as_deref()).await?; let key = args.key.into(); diff --git a/src/bin/pb-mapper/admin.rs b/src/bin/pb-mapper/admin.rs index a950bd8..1f9d620 100644 --- a/src/bin/pb-mapper/admin.rs +++ b/src/bin/pb-mapper/admin.rs @@ -331,8 +331,8 @@ async fn send_admin_request_with_timeout( .await .map_err(|error| -> Box { Box::new(error) })?; let session = ClientHeaderSession::from_process()?; - sent.store(true, std::sync::atomic::Ordering::Release); session.write_initial(&mut stream, &encoded).await?; + sent.store(true, std::sync::atomic::Ordering::Release); let mut reader = session.response_reader(&mut stream)?; let message = reader.read_msg().await?; Ok::<_, Box>(PbConnResponse::decode(message)?) diff --git a/src/common/auth/actor.rs b/src/common/auth/actor.rs index bec0ebd..f76c7b8 100644 --- a/src/common/auth/actor.rs +++ b/src/common/auth/actor.rs @@ -16,6 +16,8 @@ //! authenticated before root rotation from executing against the new administrator //! state. The actor is also the sole strong owner of temporary-key leases. +use std::collections::BTreeMap; + use super::*; pub(super) struct AuthActorState { @@ -105,6 +107,20 @@ pub(super) async fn run_auth_actor( ); tombstones.sort_unstable_by_key(|(cleanup_at, _)| *cleanup_at); let mut tombstones = VecDeque::from(tombstones); + let mut high_expiries = BTreeMap::>::new(); + for entry in inner + .high_slot_entries + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .iter() + { + if entry.state == SlotState::Active && entry.expires_at > now { + high_expiries + .entry(entry.expires_at) + .or_default() + .insert(entry.key_id); + } + } let mut last_snapshot_at = unix_seconds(); let mut tick = tokio::time::interval(Duration::from_secs(1)); tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); @@ -145,7 +161,7 @@ pub(super) async fn run_auth_actor( lease.cancel_expired(); } } - expire_due_high_slots(&inner, &mut cold, &mut tombstones, now); + expire_due_high_slots(&inner, &mut cold, &mut tombstones, &mut high_expiries, now); let mut due_high = Vec::new(); while let Some((cleanup_at, key_id)) = tombstones.front().copied() { if cleanup_at > now { @@ -252,12 +268,12 @@ pub(super) async fn run_auth_actor( } AuthCommand::Renew { authority, key_id, ttl, response } => { let result = validate_admin_authority(&inner, &authority) - .and_then(|()| actor_renew(&inner, &config, &cold, &mut wheel, key_id, ttl)); + .and_then(|()| actor_renew(&inner, &config, &cold, &mut wheel, &mut high_expiries, key_id, ttl)); let _ = response.send(result); } AuthCommand::Revoke { authority, key_id, response } => { let result = validate_admin_authority(&inner, &authority) - .and_then(|()| actor_revoke(&inner, &config, &mut cold, &mut tombstones, key_id)); + .and_then(|()| actor_revoke(&inner, &config, &mut cold, &mut tombstones, &mut high_expiries, key_id)); let _ = response.send(result); } AuthCommand::Gc { authority, response } => { @@ -268,6 +284,7 @@ pub(super) async fn run_auth_actor( &mut cold, &mut wheel, &mut tombstones, + &mut high_expiries, &admin_replay_order, )); let _ = response.send(result); @@ -280,13 +297,14 @@ pub(super) async fn run_auth_actor( &mut cold, &mut wheel, &admin_replay_order, + &mut high_expiries, "auth_state_reset", )); let _ = response.send(result); } AuthCommand::RotateRoot { authority, new_key, response } => { let result = validate_admin_authority(&inner, &authority) - .and_then(|()| actor_rotate_root(&inner, &config, &mut cold, &mut wheel, &mut admin_lease, new_key)); + .and_then(|()| actor_rotate_root(&inner, &config, &mut cold, &mut wheel, &mut admin_lease, &mut high_expiries, new_key)); if result.is_ok() { admin_replays.clear(); admin_replay_order.clear(); @@ -582,6 +600,7 @@ fn actor_renew( config: &AuthConfig, cold: &HashMap, wheel: &mut TimingWheel, + high_expiries: &mut BTreeMap>, key_id: u64, ttl: Duration, ) -> Result { @@ -672,8 +691,10 @@ fn actor_renew( true, )); } + unschedule_high_expiry(high_expiries, entry.expires_at, key_id); entry.expires_at = expires_at; entry.tombstoned_at = None; + schedule_high_expiry(high_expiries, expires_at, key_id); } metadata_with_credential(inner, cold, key_id, true) } @@ -683,6 +704,7 @@ fn actor_revoke( config: &AuthConfig, cold: &mut HashMap, tombstones: &mut VecDeque<(u64, u64)>, + high_expiries: &mut BTreeMap>, key_id: u64, ) -> Result { ensure_store_available(inner)?; @@ -748,6 +770,7 @@ fn actor_revoke( if entry.state != SlotState::Active { return Err(key_not_active()); } + unschedule_high_expiry(high_expiries, entry.expires_at, key_id); entry.state = SlotState::Revoked; entry.tombstoned_at = Some(now); if let Some(metadata) = cold.get_mut(&key_id) { @@ -785,6 +808,7 @@ fn actor_gc( cold: &mut HashMap, wheel: &mut TimingWheel, tombstones: &mut VecDeque<(u64, u64)>, + high_expiries: &mut BTreeMap>, admin_replays: &VecDeque, ) -> Result { ensure_store_available(inner)?; @@ -834,6 +858,12 @@ fn actor_gc( } keep }); + high_expiries.clear(); + for entry in high.iter() { + if entry.state == SlotState::Active && entry.expires_at > now { + schedule_high_expiry(high_expiries, entry.expires_at, entry.key_id); + } + } } tombstones.clear(); let gc_audit = audit("temporary_key_gc", None, Some(format!("removed={removed}"))); @@ -855,6 +885,7 @@ fn actor_reset( cold: &mut HashMap, wheel: &mut TimingWheel, admin_replays: &VecDeque, + high_expiries: &mut BTreeMap>, action: &str, ) -> Result<(), AuthFailure> { let new_instance_id = random_instance_id(); @@ -907,6 +938,7 @@ fn actor_reset( .unwrap_or_else(|poisoned| poisoned.into_inner()) = new_instance_id; cold.clear(); wheel.clear(unix_seconds()); + high_expiries.clear(); clear_retained_high_slot_entries(inner); inner.safe_mode.store(false, Ordering::Release); Ok(()) @@ -918,6 +950,7 @@ fn actor_rotate_root( cold: &mut HashMap, wheel: &mut TimingWheel, admin_lease: &mut Arc, + high_expiries: &mut BTreeMap>, new_key: AesKeyType, ) -> Result<(), AuthFailure> { if new_key == inner.admin_key() { @@ -976,6 +1009,7 @@ fn actor_rotate_root( } cold.clear(); wheel.clear(unix_seconds()); + high_expiries.clear(); clear_retained_high_slot_entries(inner); let new_admin_lease = Arc::new(AuthLease::new(0, u64::MAX)); *inner @@ -1177,18 +1211,55 @@ fn high_slot_metadata(entry: &PersistedEntry) -> TemporaryKeyMetadata { } } +fn schedule_high_expiry( + high_expiries: &mut BTreeMap>, + expires_at: u64, + key_id: u64, +) { + high_expiries.entry(expires_at).or_default().insert(key_id); +} + +fn unschedule_high_expiry( + high_expiries: &mut BTreeMap>, + expires_at: u64, + key_id: u64, +) { + if let Some(ids) = high_expiries.get_mut(&expires_at) { + ids.remove(&key_id); + if ids.is_empty() { + high_expiries.remove(&expires_at); + } + } +} + fn expire_due_high_slots( inner: &Arc, cold: &mut HashMap, tombstones: &mut VecDeque<(u64, u64)>, + high_expiries: &mut BTreeMap>, now: u64, ) { + let mut due = Vec::new(); + while let Some((&expires_at, _)) = high_expiries.first_key_value() { + if expires_at > now { + break; + } + if let Some(ids) = high_expiries.remove(&expires_at) { + due.extend(ids.into_iter().map(|key_id| (expires_at, key_id))); + } + } + if due.is_empty() { + return; + } let mut high = inner .high_slot_entries .write() .unwrap_or_else(|poisoned| poisoned.into_inner()); let mut expired = Vec::new(); - for entry in high.iter_mut() { + for (_scheduled_at, key_id) in due { + let Some(entry) = high.iter_mut().find(|entry| entry.key_id == key_id) else { + continue; + }; if entry.state == SlotState::Active && entry.expires_at <= now { entry.state = SlotState::Expired; let tombstoned_at = entry.expires_at; @@ -1197,6 +1268,8 @@ fn expire_due_high_slots( metadata.tombstoned_at = tombstoned_at; } expired.push((tombstoned_at, entry.key_id, entry.expires_at)); + } else if entry.state == SlotState::Active && entry.expires_at > now { + schedule_high_expiry(high_expiries, entry.expires_at, key_id); } } drop(high); diff --git a/src/common/message/secure.rs b/src/common/message/secure.rs index 217b7e8..368d412 100644 --- a/src/common/message/secure.rs +++ b/src/common/message/secure.rs @@ -672,14 +672,22 @@ fn authenticate_and_admit( let already_admitted = replay.already_admitted(&fingerprint, unix_seconds()); let context = auth .authenticate_presented(key_id, &key) - .map_err(|failure| ServerInitialError { - failure, - response_session: if already_admitted { + .map_err(|failure| { + let response_session = if already_admitted { None - } else { + } else if matches!( + replay.claim(&fingerprint, unix_seconds()), + FirstFlightAdmit::Fresh + ) { Some(error_session) - }, - presented_key_id: Some(key_id), + } else { + None + }; + ServerInitialError { + failure, + response_session, + presented_key_id: Some(key_id), + } })?; match replay.admit(key_id, &fingerprint, unix_seconds()) { FirstFlightAdmit::Fresh => Ok(context), diff --git a/src/common/message/secure/replay.rs b/src/common/message/secure/replay.rs index a27fe79..6fb8bbc 100644 --- a/src/common/message/secure/replay.rs +++ b/src/common/message/secure/replay.rs @@ -384,11 +384,17 @@ fn read_complete_replay_records(file: &mut File) -> std::io::Result records.push(record), - Err(error) if error.kind() == ErrorKind::UnexpectedEof => return Ok(records), - Err(error) => return Err(error), + let read = file.read(&mut record)?; + if read == 0 { + return Ok(records); } + if read < REPLAY_RECORD_LEN { + return Err(std::io::Error::new( + ErrorKind::UnexpectedEof, + "truncated first-flight replay record", + )); + } + records.push(record); } } diff --git a/src/common/message/secure/tests.rs b/src/common/message/secure/tests.rs index c499e3d..860e59c 100644 --- a/src/common/message/secure/tests.rs +++ b/src/common/message/secure/tests.rs @@ -177,7 +177,12 @@ async fn revoked_first_flights_do_not_consume_the_replay_filter() { Err(error) => error, }; assert_eq!(first.failure.code, "temporary_key_revoked"); + assert!(first.response_session.is_some()); assert_eq!(second.failure.code, "temporary_key_revoked"); + assert!( + second.response_session.is_none(), + "a claimed revoked salt must not reuse nonce 0" + ); let _ = std::fs::remove_dir_all(config.state_dir); } @@ -533,7 +538,7 @@ fn persisted_first_flights_survive_a_torn_trailing_record() { let mut restored = ReplayGuard::open(Some(path.clone()), 1024, DEFAULT_REPLAY_WINDOW_SECONDS); assert_eq!( restored.admit(7, &fingerprint, now), - FirstFlightAdmit::Replayed + FirstFlightAdmit::Unavailable ); let _ = std::fs::remove_file(path); } diff --git a/ui/lib/src/views/configuration_view.dart b/ui/lib/src/views/configuration_view.dart index e74fa3c..a639f83 100644 --- a/ui/lib/src/views/configuration_view.dart +++ b/ui/lib/src/views/configuration_view.dart @@ -87,7 +87,9 @@ class _ConfigurationViewState extends State { Future _saveConfiguration() async { if (_isSaving) return; // Prevent multiple simultaneous saves final msgHeaderKey = _msgHeaderKeyController.text.trim(); - if (msgHeaderKey.length != 32 && !msgHeaderKey.startsWith('pbmt1_')) { + if (msgHeaderKey.isNotEmpty && + msgHeaderKey.length != 32 && + !msgHeaderKey.startsWith('pbmt1_')) { showToast(context, context.l10n.keyLengthInvalid, kind: ToastKind.error); return; } @@ -313,7 +315,9 @@ class _ConfigurationViewState extends State { if (serverAddress.isEmpty) { throw const FormatException('serverAddress is required'); } - if (msgHeaderKey.length != 32 && !msgHeaderKey.startsWith('pbmt1_')) { + if (msgHeaderKey.isNotEmpty && + msgHeaderKey.length != 32 && + !msgHeaderKey.startsWith('pbmt1_')) { throw const FormatException( 'MSG_HEADER_KEY must be a 32-character administrator key or a pbmt1_ temporary credential', ); diff --git a/ui/lib/src/views/setup_wizard_view.dart b/ui/lib/src/views/setup_wizard_view.dart index 538d5be..51aa3ed 100644 --- a/ui/lib/src/views/setup_wizard_view.dart +++ b/ui/lib/src/views/setup_wizard_view.dart @@ -158,7 +158,7 @@ class _SetupWizardViewState extends State { setState(() => _error = l10n.setupServerInvalid); return; } - if (key.length != 32 && !key.startsWith('pbmt1_')) { + if (key.isNotEmpty && key.length != 32 && !key.startsWith('pbmt1_')) { setState(() => _error = l10n.setupKeyInvalid); return; } From 22ae9d76e931e2f8f7c61b2d2cfcbdbfb0137381 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 11:23:24 +0800 Subject: [PATCH 53/74] Bound actor abort and fail register/status abort_actor now times out instead of spinning if AuthStateInner does not drop. pb-mapper register and status return a nonzero exit status when the process credential is missing or rejected. --- CHANGELOG.md | 1 + src/bin/pb-mapper.rs | 6 ++++-- src/common/auth/runtime.rs | 11 ++++++++++- src/local/client/mod.rs | 20 +++++++++++--------- ui/native/pb_mapper_ffi/src/state/runtime.rs | 6 +++++- 5 files changed, 31 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index fc83c4a..4e96c79 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -50,6 +50,7 @@ All notable changes to this project will be documented in this file. - Kept replay-lock waits off Tokio workers, claimed limited and stale-root first flights before sending a nonce-0 error, accepted a recovery key that decrypts WAL-only state, verified legacy and installer keys against existing state, rolled back UI config when persistence failed, and probed tunnels at their pinned relay endpoints. - Stopped persisting limited first flights once the Bloom budget is full, omitted their nonce-0 error sessions, and left container `admin.key` unset when encrypted state remains so the runtime can verify a legacy recovery key. - Claimed revoked or expired first flights before returning a nonce-0 error, scheduled high-slot expiries instead of scanning every tick, marked administrator requests sent only after the first flight is written, allowed clearing the UI credential, failed `pb-mapper connect` without a credential, and fail-closed truncated replay records. +- Bounded `abort_actor` so a leaked `AuthStateInner` cannot hang shutdown, and made `register`/`status` fail with a nonzero exit status when the process credential is missing or rejected. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/src/bin/pb-mapper.rs b/src/bin/pb-mapper.rs index 117c4d3..572538a 100644 --- a/src/bin/pb-mapper.rs +++ b/src/bin/pb-mapper.rs @@ -315,6 +315,9 @@ async fn run_server(args: ServerArgs) -> Result<(), Box> { } async fn run_register(args: RegisterArgs) -> Result<(), Box> { + pb_mapper::common::checksum::get_process_credential().map_err(|error| { + std::io::Error::other(format!("registration credential is required: {error}")) + })?; let local_addr = get_sockaddr_async(&args.addr).await?; let remote_addr = get_pb_mapper_server_async(args.relay.server.as_deref()).await?; let options = ServerTunnelOptions { @@ -383,8 +386,7 @@ async fn run_connect(args: ConnectArgs) -> Result<(), Box> { async fn run_status(args: StatusArgs) -> Result<(), Box> { let remote_addr = get_pb_mapper_server_async(args.server.as_deref()).await?; - handle_status_cli_scoped(args.op, remote_addr, args.namespace).await; - Ok(()) + handle_status_cli_scoped(args.op, remote_addr, args.namespace).await } fn parse_duration(raw: &str) -> Result { diff --git a/src/common/auth/runtime.rs b/src/common/auth/runtime.rs index 76d4e03..e7d8c88 100644 --- a/src/common/auth/runtime.rs +++ b/src/common/auth/runtime.rs @@ -217,7 +217,7 @@ impl AuthRuntime { } } - pub async fn abort_actor(&self) { + pub async fn abort_actor(&self) -> Result<(), AuthFailure> { self.actor_abort.abort(); let handle = self .actor @@ -227,9 +227,18 @@ impl AuthRuntime { if let Some(handle) = handle { let _ = handle.await; } + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); while self.inner.upgrade().is_some() { + if tokio::time::Instant::now() >= deadline { + return Err(AuthFailure::new( + "auth_state_unavailable", + "authentication actor did not drop after abort", + true, + )); + } tokio::time::sleep(Duration::from_millis(10)).await; } + Ok(()) } pub fn config(&self) -> &AuthConfig { diff --git a/src/local/client/mod.rs b/src/local/client/mod.rs index db7a451..0c0fff7 100644 --- a/src/local/client/mod.rs +++ b/src/local/client/mod.rs @@ -462,14 +462,16 @@ pub async fn handle_status_cli op: StatusOp, addr: A, ) { - handle_status_cli_scoped(op, addr, None).await + if let Err(error) = handle_status_cli_scoped(op, addr, None).await { + tracing::error!("{error}"); + } } pub async fn handle_status_cli_scoped( op: StatusOp, addr: A, namespace: Option, -) { +) -> Result<(), Box> { match op { StatusOp::RemoteId => show_status_scoped(addr, PbConnStatusReq::RemoteId, namespace).await, StatusOp::Keys => show_status_scoped(addr, PbConnStatusReq::Keys, namespace).await, @@ -480,12 +482,12 @@ pub async fn show_status_scoped, -) { - let mut stream = snafu_error_get_or_return!( - each_addr(remote_addr, TcpStream::connect).await, - "get status stream" - ); - let status = snafu_error_get_or_return!(get_status_scoped(&mut stream, req, namespace).await); - let status = snafu_error_get_or_return!(serde_json::to_string_pretty(&status)); +) -> Result<(), Box> { + let mut stream = each_addr(remote_addr, TcpStream::connect) + .await + .map_err(|error| format!("get status stream: {error}"))?; + let status = get_status_scoped(&mut stream, req, namespace).await?; + let status = serde_json::to_string_pretty(&status)?; println!("Status:{status}"); + Ok(()) } diff --git a/ui/native/pb_mapper_ffi/src/state/runtime.rs b/ui/native/pb_mapper_ffi/src/state/runtime.rs index 1e17703..3e228dd 100644 --- a/ui/native/pb_mapper_ffi/src/state/runtime.rs +++ b/ui/native/pb_mapper_ffi/src/state/runtime.rs @@ -102,7 +102,11 @@ impl PbMapperState { } Err(_) => { if let Some(auth) = self.server_auth.as_ref() { - auth.abort_actor().await; + if let Err(error) = auth.abort_actor().await { + tracing::warn!( + "timed out waiting for the authentication actor to drop: {error}" + ); + } } handle.abort(); let _ = handle.await; From f113b3d230535b33f908024de580af1126f689bd Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 11:36:00 +0800 Subject: [PATCH 54/74] Abort pooled workers and keep Compose keys Dropping a registration task now aborts every control-pool worker. Compose upgrades keep the machine-derived key and persist the legacy /var/lib/pb-mapper-server path. --- CHANGELOG.md | 1 + docker/docker-compose.yml | 6 +++++- src/local/server/mod.rs | 26 +++++++++++++++++++++++--- 3 files changed, 29 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4e96c79..face267 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -51,6 +51,7 @@ All notable changes to this project will be documented in this file. - Stopped persisting limited first flights once the Bloom budget is full, omitted their nonce-0 error sessions, and left container `admin.key` unset when encrypted state remains so the runtime can verify a legacy recovery key. - Claimed revoked or expired first flights before returning a nonce-0 error, scheduled high-slot expiries instead of scanning every tick, marked administrator requests sent only after the first flight is written, allowed clearing the UI credential, failed `pb-mapper connect` without a credential, and fail-closed truncated replay records. - Bounded `abort_actor` so a leaked `AuthStateInner` cannot hang shutdown, and made `register`/`status` fail with a nonzero exit status when the process credential is missing or rejected. +- Aborted pooled registration workers when the outer service task is cancelled, and kept Compose upgrades on the machine-derived key plus a persisted `/var/lib/pb-mapper-server` volume. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml index d25b769..699519f 100644 --- a/docker/docker-compose.yml +++ b/docker/docker-compose.yml @@ -5,13 +5,17 @@ services: image: ackingliu/pb-mapper:x86_64_musl environment: PB_MAPPER_PORT: 7666 - USE_MACHINE_MSG_HEADER_KEY: false + # Keep the machine-derived key on first boot of an empty auth volume so + # recreating this service does not mint a new administrator key. + USE_MACHINE_MSG_HEADER_KEY: true RUST_LOG: error volumes: - pb-mapper-auth:/var/lib/pb-mapper/auth + - pb-mapper-legacy:/var/lib/pb-mapper-server ports: - "7666:7666" restart: unless-stopped volumes: pb-mapper-auth: + pb-mapper-legacy: diff --git a/src/local/server/mod.rs b/src/local/server/mod.rs index b1cb959..338adf3 100644 --- a/src/local/server/mod.rs +++ b/src/local/server/mod.rs @@ -8,6 +8,7 @@ use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use snafu::ResultExt; use tokio::net::TcpStream; +use tokio::task::JoinHandle; use tokio::time::MissedTickBehavior; use tracing::instrument; @@ -343,7 +344,7 @@ async fn run_server_side_cli_pool( if pool_size > 1 { for worker_index in 1..pool_size { let worker_key = key.clone(); - worker_handles.push(tokio::spawn(async move { + worker_handles.push(AbortOnDropHandle(Some(tokio::spawn(async move { run_server_side_cli_worker_with_credential::( local_addr, remote_addr, @@ -354,7 +355,7 @@ async fn run_server_side_cli_pool( pinned_credential, ) .await; - })); + })))); } } run_server_side_cli_worker_with_credential::( @@ -368,7 +369,7 @@ async fn run_server_side_cli_pool( ) .await; for handle in worker_handles { - if let Err(e) = handle.await { + if let Err(e) = handle.join().await { tracing::warn!( event = "local_server_control_worker_join_failed", error = %e, @@ -378,6 +379,25 @@ async fn run_server_side_cli_pool( } } +struct AbortOnDropHandle(Option>); + +impl AbortOnDropHandle { + async fn join(mut self) -> Result<(), tokio::task::JoinError> { + match self.0.take() { + Some(handle) => handle.await, + None => Ok(()), + } + } +} + +impl Drop for AbortOnDropHandle { + fn drop(&mut self) { + if let Some(handle) = self.0.take() { + handle.abort(); + } + } +} + async fn run_server_side_cli_worker_with_credential( local_addr: A, remote_addr: A, From 86e4c2b92ad4be32239158a60331fc2217d84594 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 11:50:03 +0800 Subject: [PATCH 55/74] Reap connections and batch tombstones Accept no longer walks every live connection handle. Timing-wheel fast-forwards collect due tombstones and merge them in one sort. --- CHANGELOG.md | 1 + src/common/auth/actor.rs | 59 ++++++++++++++++++++++++++++++++++++---- src/pb_server/runtime.rs | 15 +++++----- 3 files changed, 62 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index face267..d6db113 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -52,6 +52,7 @@ All notable changes to this project will be documented in this file. - Claimed revoked or expired first flights before returning a nonce-0 error, scheduled high-slot expiries instead of scanning every tick, marked administrator requests sent only after the first flight is written, allowed clearing the UI credential, failed `pb-mapper connect` without a credential, and fail-closed truncated replay records. - Bounded `abort_actor` so a leaked `AuthStateInner` cannot hang shutdown, and made `register`/`status` fail with a nonzero exit status when the process credential is missing or rejected. - Aborted pooled registration workers when the outer service task is cancelled, and kept Compose upgrades on the machine-derived key plus a persisted `/var/lib/pb-mapper-server` volume. +- Reaped finished connection tasks with a `JoinSet` instead of scanning every live handle on accept, and batch-sorted timing-wheel tombstones after a large clock jump. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/src/common/auth/actor.rs b/src/common/auth/actor.rs index f76c7b8..db22616 100644 --- a/src/common/auth/actor.rs +++ b/src/common/auth/actor.rs @@ -128,6 +128,7 @@ pub(super) async fn run_auth_actor( tokio::select! { _ = tick.tick() => { let now = unix_seconds(); + let mut due_tombstones = Vec::new(); for lease in wheel.advance(now) { let key_id = lease.key_id(); let version = lease.wheel_version.load(Ordering::Acquire); @@ -144,11 +145,7 @@ pub(super) async fn run_auth_actor( if let Some(metadata) = cold.get_mut(&key_id) { metadata.tombstoned_at = tombstoned_at; } - push_tombstone( - &mut tombstones, - tombstoned_at, - key_id, - ); + due_tombstones.push((tombstoned_at, key_id)); tracing::info!( event = "temporary_key_expired", auth_stage = "expiry", @@ -161,6 +158,7 @@ pub(super) async fn run_auth_actor( lease.cancel_expired(); } } + extend_tombstones(&mut tombstones, due_tombstones); expire_due_high_slots(&inner, &mut cold, &mut tombstones, &mut high_expiries, now); let mut due_high = Vec::new(); while let Some((cleanup_at, key_id)) = tombstones.front().copied() { @@ -802,6 +800,57 @@ fn push_tombstone(tombstones: &mut VecDeque<(u64, u64)>, tombstoned_at: u64, key tombstones.insert(index, (cleanup_at, key_id)); } +fn extend_tombstones(tombstones: &mut VecDeque<(u64, u64)>, due: Vec<(u64, u64)>) { + if due.is_empty() { + return; + } + if due.len() == 1 { + let (tombstoned_at, key_id) = due[0]; + push_tombstone(tombstones, tombstoned_at, key_id); + return; + } + let mut extra = due + .into_iter() + .map(|(tombstoned_at, key_id)| { + ( + tombstoned_at.saturating_add(TOMBSTONE_RETENTION.as_secs()), + key_id, + ) + }) + .collect::>(); + extra.sort_unstable_by_key(|(cleanup_at, _)| *cleanup_at); + if tombstones.is_empty() { + *tombstones = extra.into(); + return; + } + let existing = tombstones.drain(..).collect::>(); + *tombstones = merge_sorted_tombstones(existing, extra).into(); +} + +fn merge_sorted_tombstones(left: Vec<(u64, u64)>, right: Vec<(u64, u64)>) -> Vec<(u64, u64)> { + let mut merged = Vec::with_capacity(left.len().saturating_add(right.len())); + let mut left = left.into_iter().peekable(); + let mut right = right.into_iter().peekable(); + loop { + match (left.peek(), right.peek()) { + (Some((left_at, _)), Some((right_at, _))) if left_at <= right_at => { + merged.push(left.next().expect("peeked left tombstone")); + } + (Some(_), Some(_)) => merged.push(right.next().expect("peeked right tombstone")), + (Some(_), None) => { + merged.extend(left); + break; + } + (None, Some(_)) => { + merged.extend(right); + break; + } + (None, None) => break, + } + } + merged +} + fn actor_gc( inner: &Arc, config: &AuthConfig, diff --git a/src/pb_server/runtime.rs b/src/pb_server/runtime.rs index 5957af5..8345365 100644 --- a/src/pb_server/runtime.rs +++ b/src/pb_server/runtime.rs @@ -104,7 +104,7 @@ pub async fn run_server_on_listener( let new_streams_per_second = env_limit("PB_MAPPER_NEW_STREAMS_PER_SECOND", 100); let new_streams_burst = env_limit("PB_MAPPER_NEW_STREAMS_BURST", 200); let mut next_server_generation = 1_u64; - let mut connection_tasks = Vec::new(); + let mut connection_tasks = tokio::task::JoinSet::new(); let listen_addr = listener.local_addr()?; tracing::info!( @@ -359,14 +359,13 @@ pub async fn run_server_on_listener( ); let manager_task_sender = manager.get_task_sender(); let security = security.clone(); - connection_tasks - .retain(|handle: &tokio::task::JoinHandle<()>| !handle.is_finished()); - connection_tasks.push(tokio::spawn(async move { + while connection_tasks.try_join_next().is_some() {} + connection_tasks.spawn(async move { snafu_error_handle!( handle_conn(conn_id, peer_addr, manager_task_sender, stream, security) .await ); - })); + }); } ManagerTask::DeRegisterServerConn { key, conn_id } => { let removed_from_service_map = @@ -921,10 +920,10 @@ pub async fn run_server_on_listener( // Abort first, then wait. Dropping a JoinHandle after abort() does not // wait for the task to drop its AuthRuntime clone, so a UI restart can // still see auth.lock held. + connection_tasks.abort_all(); + while connection_tasks.join_next().await.is_some() {} abort_and_wait( - connection_tasks - .into_iter() - .chain(std::iter::once(listener_handle)) + std::iter::once(listener_handle) .chain(std::iter::once(shutdown_handle)) .chain(status_forward_handle), ) From 71b99cd33bce99257161bb8d512889d8a6de6098 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 12:21:57 +0800 Subject: [PATCH 56/74] Simplify auth expiry, first-flight, and UI pinning Drop scheduled high-slot maps, merge-sorted tombstones, and split credential/endpoint maps. Keep nonce-0 uniqueness, WAL-first recovery, and pinned tunnels. --- CHANGELOG.md | 53 +--- src/common/auth/actor.rs | 162 ++----------- src/common/auth/runtime.rs | 23 +- src/common/message/secure.rs | 229 ++++++++---------- src/common/message/secure/replay.rs | 44 ++-- src/local/server/mod.rs | 31 +-- ui/native/pb_mapper_ffi/src/state.rs | 12 +- .../pb_mapper_ffi/src/state/configuration.rs | 12 +- ui/native/pb_mapper_ffi/src/state/runtime.rs | 44 ++-- ui/native/pb_mapper_ffi/src/state/status.rs | 18 +- 10 files changed, 201 insertions(+), 427 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d6db113..9d6e546 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,55 +4,14 @@ All notable changes to this project will be documented in this file. ## [0.4.0] - 2026-08-18 - Added a sole administrator credential plus renewable, expiring, and immediately revocable `pbmt1_` temporary credentials with fixed-slot O(1) lookup and isolated per-key service namespaces. -- Added single-flight protocol-v2 authentication with directional AES-256-GCM keys, monotonic frame counters, authenticated routing metadata, replay detection, stable structured errors, and optional legacy framing during migration. -- Added encrypted snapshot/WAL authentication state, lifecycle audit records, hierarchical timing-wheel expiry, hard closure of revoked live connections, safe-mode recovery, root-key rotation, and explicit auth-state reset. +- Added single-flight protocol-v2 authentication with directional AES-256-GCM keys, monotonic frame counters, authenticated routing metadata, durable first-flight replay protection, and optional legacy framing during migration. +- Added encrypted snapshot/WAL authentication state, exclusive `auth.lock`, lifecycle audit records, hierarchical timing-wheel expiry, hard closure of revoked live connections, recoverable root-key rotation, and explicit auth-state reset. - Extended the unified CLI with temporary-key lifecycle, service/connection inventory, auth status, protocol policy, root rotation, namespace targeting, and human/JSON/NDJSON output. - Replaced insecure default-key fallback with first-start random administrator-key generation, retained machine-derived keys only for explicit compatibility, and updated Flutter, installers, systemd, Docker, release metadata, and bilingual documentation. -- Fixed remaining review findings: installer migration now honors `MSG_HEADER_KEY` from `/etc/pb-mapper/server.env`, isolated relays validate legacy frames with their own administrator key, first-flight replay retention covers the full clock-skew window, and desktop macOS/Windows servers use a user-writable auth directory. -- Rejected NUL/non-printable rotated administrator keys, cancelled in-flight status reads on credential revocation, refused `--force-init-admin-key` when encrypted auth state already exists, bound isolated-relay legacy continuation checksums to the relay key, and doubled first-flight Bloom retention so a max-future timestamp cannot outlive the filter. -- Centralized env-safe administrator-key checks, isolated-relay legacy codec construction, credential-cancellation races, and auth snapshot/WAL paths so later protocol and lifecycle changes reuse one implementation. -- Refused administrator-key initialization whenever encrypted auth state is present, preserved discarded slot generations across capacity changes, rolled back or fail-closed uncertain WAL appends, authenticated first flights before consuming the replay filter, and reused the registration credential for provider streams. -- Capped legacy first-flight allocations and kept legacy framing denied when authentication state enters safe mode. -- Skipped snapshot compaction while startup is in safe mode, fsynced the auth directory when creating `auth.wal`, cleared retained high-slot entries on reset/rotate, dropped the slot write lock before WAL fail-closed cancellation, fail-closed process checksums after the credential is cleared, and rate-limited first flights per credential before consuming the shared replay filter. -- Persisted first-flight replay admissions across restarts, replaced existing auth files atomically on Windows, rejected explicit out-of-range server auth flags, and exposed the embedded relay's isolated administrator key through FFI/UI. -- Made a second administrator first-flight salt replay surface the dedicated retry-exhausted error instead of leaving that path unreachable. -- Compacted the durable first-flight replay log while the relay is running, rolled back torn replay-log appends, rewrote that log atomically, sized first-flight admission from `PB_MAPPER_NEW_STREAMS_PER_SECOND`, and took an exclusive lock on the authentication state directory. -- Fsynced the replay-log directory on first creation, took the state lock before `--init-admin-key`, aborted accepted connection tasks on relay shutdown, and staged `admin.key.next` so an interrupted root rotation can recover a matching key and snapshot. -- Stopped returning the embedded relay administrator key from routine config fetches; revealing it now requires an explicit FFI/UI action. -- Took `auth.lock` before loading or creating `admin.key`, staged `server-instance-id.next` so an interrupted reset can recover a matching snapshot, discarded leftover WAL from the previous instance during that recovery, and used a user-writable Linux auth directory when `/var/lib/pb-mapper/auth` is not usable. -- Flushed the parent directory after Windows auth-state replacements, matching the Unix `fsync` after rename. -- Discarded leftover WAL encrypted under the previous administrator key when promoting `admin.key.next`. -- Opened Windows parent directories with write access before `FlushFileBuffers`, and waited for aborted connection tasks to finish so `auth.lock` is released before shutdown returns. -- Distinguished a post-rotation or post-reset temporary credential as `temporary_key_rotated` instead of the generic `temporary_key_invalid` used for a mistyped live key. -- Failed closed when first-flight replay-log rollback cannot restore the previous length, kept `auth.lock` in the actor until it exits, reserved `temporary_key_rotated` for a real root-epoch change, refused `--use-machine-msg-header-key` when `admin.key` already exists, and persisted installer `MSG_HEADER_KEY` into `admin.key`. -- Refused `write_admin_key_file` of the live `admin.key` while encrypted auth state exists, matching `initialize_admin_key`; `admin.key.next` remains allowed. -- Kept `temporary_key_rotated` after a later issue in the same slot, aborted a timed-out embedded-relay shutdown, fail-closed replay logs whose directory sync failed, allowed a post-rotation write of the live key that already decrypts the snapshot, and treated `PB_MAPPER_NEW_STREAMS_PER_SECOND=0` as the default 100. -- Fsynced the replay-log directory after compacting replacements, treated an unreadable existing replay log as unavailable, and made timing-wheel buckets hold `Weak` leases so renewals no longer accumulate day-long strong references. -- Retained administrator mutation replay claims from the server acceptance time, so a backdated client timestamp cannot shrink the ten-minute replay window. -- Read first-flight replay logs with exact-record I/O and fail closed after an incomplete read or a rewrite that may have already replaced the log; compaction temporary files now use a random suffix. -- Released revoked timing-wheel owners when the tombstone is recycled or GC frees the slot, instead of keeping them until the original TTL. -- Kept the previous root key and instance id in memory so a stale first flight after rotation or reset can decrypt and return `temporary_key_rotated`. -- Validated installer `MSG_HEADER_KEY` values as 32 printable ASCII bytes before writing `admin.key`. -- Passed `--use-machine-msg-header-key` only when `admin.key` is missing, so a container restart with a persistent auth volume does not fail after first boot. -- Pinned each local client listener to the credential captured at start, so a later Flutter config change cannot switch an existing port onto another tenant. -- Garbage-collected expired and revoked high-slot entries while keeping their generations so shrinking capacity and running `key gc` can reclaim them. -- Pinned local registration workers to the credential captured at start, matching the client listener. -- Replaced a lease that was canceled while a renewal WAL record was syncing, so a successful renew does not keep a dead cancellation token. -- Awaited the authentication actor on relay shutdown, used pinned credentials for registration probes and UI tunnel workers/status checks, and finished in-memory root rotation when `admin.key` already matched the new snapshot. -- Bound relay tunneled-frame checksums to each hop's authenticated session key instead of the process administrator key. -- Bound local UDP and codec tunnels to the pinned credential's checksum key so a later process-key change cannot desynchronize framed payloads. -- Looked up high-slot credentials for admin list/show/renew/revoke, counted them in status, scheduled their tombstones across restart, and expired due high-slot entries on the actor tick. UI tunnel stop now drops the pinned credential. First-flight decrypt failures no longer send an error frame the presenter cannot read. -- Reported cancelled leases by recorded cause, aborted the auth actor on UI shutdown timeout, validated replacement credentials before stopping a live tunnel, preserved presented key IDs on unreadable first flights, finalized reset when the new instance id was already installed, and bounded aggregate first-flight Bloom inserts. -- Waited for the auth actor to drop before `abort_actor` returns, omitted a reused GCM nonce on salt-replay errors, restored the aggregate replay count from the durable log, and finalized root rotation only when the live `admin.key` already matches the new snapshot. -- Restricted administrator CLI retries to pre-send failures and a readable `connection_salt_replayed`, so a dropped response cannot issue a second credential. -- Omitted a response session for already-admitted first flights that later fail authentication, and kept namespace stream accounting until the client stream deregisters even if the registration control socket drops. -- Evaluated first-flight admission on a blocking thread under one replay lock, omitted a session when durable admission is unavailable, aged restored Bloom generations from loaded record timestamps, validated installer keys from `server.env`, refused to persist a recovery `MSG_HEADER_KEY` that cannot decrypt existing state, batched high-slot tombstone cleanup, and captured the UI credential together with the relay address before DNS. -- Kept replay-lock waits off Tokio workers, claimed limited and stale-root first flights before sending a nonce-0 error, accepted a recovery key that decrypts WAL-only state, verified legacy and installer keys against existing state, rolled back UI config when persistence failed, and probed tunnels at their pinned relay endpoints. -- Stopped persisting limited first flights once the Bloom budget is full, omitted their nonce-0 error sessions, and left container `admin.key` unset when encrypted state remains so the runtime can verify a legacy recovery key. -- Claimed revoked or expired first flights before returning a nonce-0 error, scheduled high-slot expiries instead of scanning every tick, marked administrator requests sent only after the first flight is written, allowed clearing the UI credential, failed `pb-mapper connect` without a credential, and fail-closed truncated replay records. -- Bounded `abort_actor` so a leaked `AuthStateInner` cannot hang shutdown, and made `register`/`status` fail with a nonzero exit status when the process credential is missing or rejected. -- Aborted pooled registration workers when the outer service task is cancelled, and kept Compose upgrades on the machine-derived key plus a persisted `/var/lib/pb-mapper-server` volume. -- Reaped finished connection tasks with a `JoinSet` instead of scanning every live handle on accept, and batch-sorted timing-wheel tombstones after a large clock jump. +- Recovery keys must decrypt existing snapshot or WAL state before they are persisted. Interrupted rotation and reset recover from staged `admin.key.next` and `server-instance-id.next`. +- First-flight salts are unique for nonce 0: admission is atomic under one lock, torn replay records fail closed, and a nonce-0 error frame is sent only after that salt is reserved. +- Pinned UI and local tunnels to the credential and relay address captured at start, and bound tunneled-frame checksums to each hop's authenticated session key. +- Aborted pooled registration workers and accepted connection tasks on shutdown; the relay reaps connection tasks with a `JoinSet`. ## [0.3.0] - 2026-08-18 - Replaced the three role-specific executables with one `pb-mapper` CLI and explicit `server`, `register`, `connect`, and `status` commands. diff --git a/src/common/auth/actor.rs b/src/common/auth/actor.rs index db22616..ce8047d 100644 --- a/src/common/auth/actor.rs +++ b/src/common/auth/actor.rs @@ -16,8 +16,6 @@ //! authenticated before root rotation from executing against the new administrator //! state. The actor is also the sole strong owner of temporary-key leases. -use std::collections::BTreeMap; - use super::*; pub(super) struct AuthActorState { @@ -107,20 +105,6 @@ pub(super) async fn run_auth_actor( ); tombstones.sort_unstable_by_key(|(cleanup_at, _)| *cleanup_at); let mut tombstones = VecDeque::from(tombstones); - let mut high_expiries = BTreeMap::>::new(); - for entry in inner - .high_slot_entries - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .iter() - { - if entry.state == SlotState::Active && entry.expires_at > now { - high_expiries - .entry(entry.expires_at) - .or_default() - .insert(entry.key_id); - } - } let mut last_snapshot_at = unix_seconds(); let mut tick = tokio::time::interval(Duration::from_secs(1)); tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); @@ -128,7 +112,6 @@ pub(super) async fn run_auth_actor( tokio::select! { _ = tick.tick() => { let now = unix_seconds(); - let mut due_tombstones = Vec::new(); for lease in wheel.advance(now) { let key_id = lease.key_id(); let version = lease.wheel_version.load(Ordering::Acquire); @@ -145,7 +128,7 @@ pub(super) async fn run_auth_actor( if let Some(metadata) = cold.get_mut(&key_id) { metadata.tombstoned_at = tombstoned_at; } - due_tombstones.push((tombstoned_at, key_id)); + push_tombstone(&mut tombstones, tombstoned_at, key_id); tracing::info!( event = "temporary_key_expired", auth_stage = "expiry", @@ -158,8 +141,7 @@ pub(super) async fn run_auth_actor( lease.cancel_expired(); } } - extend_tombstones(&mut tombstones, due_tombstones); - expire_due_high_slots(&inner, &mut cold, &mut tombstones, &mut high_expiries, now); + expire_due_high_slots(&inner, &mut cold, &mut tombstones, now); let mut due_high = Vec::new(); while let Some((cleanup_at, key_id)) = tombstones.front().copied() { if cleanup_at > now { @@ -266,12 +248,12 @@ pub(super) async fn run_auth_actor( } AuthCommand::Renew { authority, key_id, ttl, response } => { let result = validate_admin_authority(&inner, &authority) - .and_then(|()| actor_renew(&inner, &config, &cold, &mut wheel, &mut high_expiries, key_id, ttl)); + .and_then(|()| actor_renew(&inner, &config, &cold, &mut wheel, key_id, ttl)); let _ = response.send(result); } AuthCommand::Revoke { authority, key_id, response } => { let result = validate_admin_authority(&inner, &authority) - .and_then(|()| actor_revoke(&inner, &config, &mut cold, &mut tombstones, &mut high_expiries, key_id)); + .and_then(|()| actor_revoke(&inner, &config, &mut cold, &mut tombstones, key_id)); let _ = response.send(result); } AuthCommand::Gc { authority, response } => { @@ -282,7 +264,6 @@ pub(super) async fn run_auth_actor( &mut cold, &mut wheel, &mut tombstones, - &mut high_expiries, &admin_replay_order, )); let _ = response.send(result); @@ -295,14 +276,13 @@ pub(super) async fn run_auth_actor( &mut cold, &mut wheel, &admin_replay_order, - &mut high_expiries, "auth_state_reset", )); let _ = response.send(result); } AuthCommand::RotateRoot { authority, new_key, response } => { let result = validate_admin_authority(&inner, &authority) - .and_then(|()| actor_rotate_root(&inner, &config, &mut cold, &mut wheel, &mut admin_lease, &mut high_expiries, new_key)); + .and_then(|()| actor_rotate_root(&inner, &config, &mut cold, &mut wheel, &mut admin_lease, new_key)); if result.is_ok() { admin_replays.clear(); admin_replay_order.clear(); @@ -598,7 +578,6 @@ fn actor_renew( config: &AuthConfig, cold: &HashMap, wheel: &mut TimingWheel, - high_expiries: &mut BTreeMap>, key_id: u64, ttl: Duration, ) -> Result { @@ -689,10 +668,8 @@ fn actor_renew( true, )); } - unschedule_high_expiry(high_expiries, entry.expires_at, key_id); entry.expires_at = expires_at; entry.tombstoned_at = None; - schedule_high_expiry(high_expiries, expires_at, key_id); } metadata_with_credential(inner, cold, key_id, true) } @@ -702,7 +679,6 @@ fn actor_revoke( config: &AuthConfig, cold: &mut HashMap, tombstones: &mut VecDeque<(u64, u64)>, - high_expiries: &mut BTreeMap>, key_id: u64, ) -> Result { ensure_store_available(inner)?; @@ -768,7 +744,6 @@ fn actor_revoke( if entry.state != SlotState::Active { return Err(key_not_active()); } - unschedule_high_expiry(high_expiries, entry.expires_at, key_id); entry.state = SlotState::Revoked; entry.tombstoned_at = Some(now); if let Some(metadata) = cold.get_mut(&key_id) { @@ -800,64 +775,12 @@ fn push_tombstone(tombstones: &mut VecDeque<(u64, u64)>, tombstoned_at: u64, key tombstones.insert(index, (cleanup_at, key_id)); } -fn extend_tombstones(tombstones: &mut VecDeque<(u64, u64)>, due: Vec<(u64, u64)>) { - if due.is_empty() { - return; - } - if due.len() == 1 { - let (tombstoned_at, key_id) = due[0]; - push_tombstone(tombstones, tombstoned_at, key_id); - return; - } - let mut extra = due - .into_iter() - .map(|(tombstoned_at, key_id)| { - ( - tombstoned_at.saturating_add(TOMBSTONE_RETENTION.as_secs()), - key_id, - ) - }) - .collect::>(); - extra.sort_unstable_by_key(|(cleanup_at, _)| *cleanup_at); - if tombstones.is_empty() { - *tombstones = extra.into(); - return; - } - let existing = tombstones.drain(..).collect::>(); - *tombstones = merge_sorted_tombstones(existing, extra).into(); -} - -fn merge_sorted_tombstones(left: Vec<(u64, u64)>, right: Vec<(u64, u64)>) -> Vec<(u64, u64)> { - let mut merged = Vec::with_capacity(left.len().saturating_add(right.len())); - let mut left = left.into_iter().peekable(); - let mut right = right.into_iter().peekable(); - loop { - match (left.peek(), right.peek()) { - (Some((left_at, _)), Some((right_at, _))) if left_at <= right_at => { - merged.push(left.next().expect("peeked left tombstone")); - } - (Some(_), Some(_)) => merged.push(right.next().expect("peeked right tombstone")), - (Some(_), None) => { - merged.extend(left); - break; - } - (None, Some(_)) => { - merged.extend(right); - break; - } - (None, None) => break, - } - } - merged -} - fn actor_gc( inner: &Arc, config: &AuthConfig, cold: &mut HashMap, wheel: &mut TimingWheel, tombstones: &mut VecDeque<(u64, u64)>, - high_expiries: &mut BTreeMap>, admin_replays: &VecDeque, ) -> Result { ensure_store_available(inner)?; @@ -907,12 +830,6 @@ fn actor_gc( } keep }); - high_expiries.clear(); - for entry in high.iter() { - if entry.state == SlotState::Active && entry.expires_at > now { - schedule_high_expiry(high_expiries, entry.expires_at, entry.key_id); - } - } } tombstones.clear(); let gc_audit = audit("temporary_key_gc", None, Some(format!("removed={removed}"))); @@ -934,7 +851,6 @@ fn actor_reset( cold: &mut HashMap, wheel: &mut TimingWheel, admin_replays: &VecDeque, - high_expiries: &mut BTreeMap>, action: &str, ) -> Result<(), AuthFailure> { let new_instance_id = random_instance_id(); @@ -987,7 +903,6 @@ fn actor_reset( .unwrap_or_else(|poisoned| poisoned.into_inner()) = new_instance_id; cold.clear(); wheel.clear(unix_seconds()); - high_expiries.clear(); clear_retained_high_slot_entries(inner); inner.safe_mode.store(false, Ordering::Release); Ok(()) @@ -999,7 +914,6 @@ fn actor_rotate_root( cold: &mut HashMap, wheel: &mut TimingWheel, admin_lease: &mut Arc, - high_expiries: &mut BTreeMap>, new_key: AesKeyType, ) -> Result<(), AuthFailure> { if new_key == inner.admin_key() { @@ -1058,7 +972,6 @@ fn actor_rotate_root( } cold.clear(); wheel.clear(unix_seconds()); - high_expiries.clear(); clear_retained_high_slot_entries(inner); let new_admin_lease = Arc::new(AuthLease::new(0, u64::MAX)); *inner @@ -1260,75 +1173,32 @@ fn high_slot_metadata(entry: &PersistedEntry) -> TemporaryKeyMetadata { } } -fn schedule_high_expiry( - high_expiries: &mut BTreeMap>, - expires_at: u64, - key_id: u64, -) { - high_expiries.entry(expires_at).or_default().insert(key_id); -} - -fn unschedule_high_expiry( - high_expiries: &mut BTreeMap>, - expires_at: u64, - key_id: u64, -) { - if let Some(ids) = high_expiries.get_mut(&expires_at) { - ids.remove(&key_id); - if ids.is_empty() { - high_expiries.remove(&expires_at); - } - } -} - fn expire_due_high_slots( inner: &Arc, cold: &mut HashMap, tombstones: &mut VecDeque<(u64, u64)>, - high_expiries: &mut BTreeMap>, now: u64, ) { - let mut due = Vec::new(); - while let Some((&expires_at, _)) = high_expiries.first_key_value() { - if expires_at > now { - break; - } - if let Some(ids) = high_expiries.remove(&expires_at) { - due.extend(ids.into_iter().map(|key_id| (expires_at, key_id))); - } - } - if due.is_empty() { - return; - } let mut high = inner .high_slot_entries .write() .unwrap_or_else(|poisoned| poisoned.into_inner()); - let mut expired = Vec::new(); - for (_scheduled_at, key_id) in due { - let Some(entry) = high.iter_mut().find(|entry| entry.key_id == key_id) else { + for entry in high.iter_mut() { + if entry.state != SlotState::Active || entry.expires_at > now { continue; - }; - if entry.state == SlotState::Active && entry.expires_at <= now { - entry.state = SlotState::Expired; - let tombstoned_at = entry.expires_at; - entry.tombstoned_at = Some(tombstoned_at); - if let Some(metadata) = cold.get_mut(&entry.key_id) { - metadata.tombstoned_at = tombstoned_at; - } - expired.push((tombstoned_at, entry.key_id, entry.expires_at)); - } else if entry.state == SlotState::Active && entry.expires_at > now { - schedule_high_expiry(high_expiries, entry.expires_at, key_id); } - } - drop(high); - for (tombstoned_at, key_id, expires_at) in expired { - push_tombstone(tombstones, tombstoned_at, key_id); + entry.state = SlotState::Expired; + let tombstoned_at = entry.expires_at; + entry.tombstoned_at = Some(tombstoned_at); + if let Some(metadata) = cold.get_mut(&entry.key_id) { + metadata.tombstoned_at = tombstoned_at; + } + push_tombstone(tombstones, tombstoned_at, entry.key_id); tracing::info!( event = "temporary_key_expired", auth_stage = "expiry", - key_id, - expires_at, + key_id = entry.key_id, + expires_at = entry.expires_at, "high-slot temporary key expired" ); } diff --git a/src/common/auth/runtime.rs b/src/common/auth/runtime.rs index e7d8c88..28e72ab 100644 --- a/src/common/auth/runtime.rs +++ b/src/common/auth/runtime.rs @@ -227,18 +227,19 @@ impl AuthRuntime { if let Some(handle) = handle { let _ = handle.await; } - let deadline = tokio::time::Instant::now() + Duration::from_secs(5); - while self.inner.upgrade().is_some() { - if tokio::time::Instant::now() >= deadline { - return Err(AuthFailure::new( - "auth_state_unavailable", - "authentication actor did not drop after abort", - true, - )); + tokio::time::timeout(Duration::from_secs(5), async { + while self.inner.upgrade().is_some() { + tokio::time::sleep(Duration::from_millis(10)).await; } - tokio::time::sleep(Duration::from_millis(10)).await; - } - Ok(()) + }) + .await + .map_err(|_| { + AuthFailure::new( + "auth_state_unavailable", + "authentication actor did not drop after abort", + true, + ) + }) } pub fn config(&self) -> &AuthConfig { diff --git a/src/common/message/secure.rs b/src/common/message/secure.rs index 368d412..0d0ff37 100644 --- a/src/common/message/secure.rs +++ b/src/common/message/secure.rs @@ -552,68 +552,45 @@ impl ServerSecurity { })?; let mut current_ciphertext = ciphertext.clone(); let fingerprint = replay_fingerprint(key_id, &salt); - let payload = match open_v2_payload( + let work = match open_v2_payload( &material, DIRECTION_CLIENT_TO_SERVER, counter, &mut current_ciphertext, ) { - Ok(payload) => payload, + Ok(payload) => FirstFlightWork::Live { + key, + payload, + error_session: session_without_context(&session), + }, Err(error) => { - if stale_root_first_flight(&self.auth, key_id, salt, counter, &ciphertext).is_some() - { - let replay = self.replay.clone(); - let auth = self.auth.clone(); - return Err(tokio::task::spawn_blocking(move || { - claim_stale_root_first_flight( - &auth, - &replay, + match stale_root_first_flight(&self.auth, key_id, salt, counter, &ciphertext) { + Some(stale) => FirstFlightWork::Stale(stale), + None => { + return Err(first_flight_error( + "protocol_v2_decrypt_failed", + error.to_string(), + false, key_id, - salt, - counter, - ciphertext, - fingerprint, - ) - }) - .await - .unwrap_or_else(|_| ServerInitialError { - failure: AuthFailure::new( - "connection_replay_store_unavailable", - "failed to evaluate first-flight admission", - true, - ), - response_session: None, - presented_key_id: Some(key_id), - })); + )) + } } - return Err(ServerInitialError { - failure: AuthFailure::new( - "protocol_v2_decrypt_failed", - error.to_string(), - false, - ), - response_session: None, - presented_key_id: Some(key_id), - }); } }; - let replay = self.replay.clone(); let auth = self.auth.clone(); - let error_session = session_without_context(&session); - let context = tokio::task::spawn_blocking(move || { - authenticate_and_admit(&auth, &replay, key_id, key, fingerprint, error_session) + let (payload, context) = tokio::task::spawn_blocking(move || { + evaluate_first_flight(&auth, &replay, key_id, fingerprint, work) }) .await - .map_err(|_| ServerInitialError { - failure: AuthFailure::new( + .unwrap_or_else(|_| { + Err(first_flight_error( "connection_replay_store_unavailable", "failed to evaluate first-flight admission", true, - ), - response_session: None, - presented_key_id: Some(key_id), - })??; + key_id, + )) + })?; session.context = Some(context); Ok(ServerInitialMessage { payload, @@ -657,103 +634,97 @@ async fn read_initial_v2_ciphertext( Ok((counter, ciphertext)) } -#[allow(clippy::result_large_err)] -fn authenticate_and_admit( - auth: &AuthRuntime, - replay: &std::sync::Mutex, +enum FirstFlightWork { + Live { + key: AesKeyType, + payload: Vec, + error_session: ServerHeaderSession, + }, + Stale(ServerInitialError), +} + +fn first_flight_error( + code: &'static str, + message: impl Into, + retryable: bool, key_id: u64, - key: AesKeyType, - fingerprint: [u8; 32], - error_session: ServerHeaderSession, -) -> std::result::Result { - let mut replay = replay - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let already_admitted = replay.already_admitted(&fingerprint, unix_seconds()); - let context = auth - .authenticate_presented(key_id, &key) - .map_err(|failure| { - let response_session = if already_admitted { - None - } else if matches!( - replay.claim(&fingerprint, unix_seconds()), - FirstFlightAdmit::Fresh - ) { - Some(error_session) - } else { - None - }; - ServerInitialError { - failure, - response_session, - presented_key_id: Some(key_id), - } - })?; - match replay.admit(key_id, &fingerprint, unix_seconds()) { - FirstFlightAdmit::Fresh => Ok(context), - FirstFlightAdmit::Replayed => Err(ServerInitialError { - failure: AuthFailure::new( - "connection_salt_replayed", - "protocol-v2 connection salt was already accepted", - true, - ), - response_session: None, - presented_key_id: Some(key_id), - }), - FirstFlightAdmit::Limited => Err(ServerInitialError { - failure: AuthFailure::new( - "connection_admission_limited", - "this credential has opened too many new connections in the current window", - true, - ), - response_session: None, - presented_key_id: Some(key_id), - }), - FirstFlightAdmit::Unavailable => Err(ServerInitialError { - failure: AuthFailure::new( - "connection_replay_store_unavailable", - "failed to persist first-flight replay admission", - true, - ), - response_session: None, - presented_key_id: Some(key_id), - }), +) -> ServerInitialError { + ServerInitialError { + failure: AuthFailure::new(code, message, retryable), + response_session: None, + presented_key_id: Some(key_id), } } +fn reserved_error_session( + replay: &mut ReplayGuard, + fingerprint: &[u8; 32], + session: ServerHeaderSession, +) -> Option { + matches!( + replay.claim(fingerprint, unix_seconds()), + FirstFlightAdmit::Fresh + ) + .then_some(session) +} + #[allow(clippy::result_large_err)] -fn claim_stale_root_first_flight( +fn evaluate_first_flight( auth: &AuthRuntime, replay: &std::sync::Mutex, key_id: u64, - salt: [u8; CONNECTION_SALT_LEN], - counter: u64, - ciphertext: Vec, fingerprint: [u8; 32], -) -> ServerInitialError { - let mut error = stale_root_first_flight(auth, key_id, salt, counter, &ciphertext) - .unwrap_or_else(|| ServerInitialError { - failure: AuthFailure::new( - "protocol_v2_decrypt_failed", - "stale-root first flight could not be classified", - false, - ), - response_session: None, - presented_key_id: Some(key_id), - }); + work: FirstFlightWork, +) -> std::result::Result<(Vec, AuthContext), ServerInitialError> { let mut replay = replay .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()); - let already_admitted = replay.already_admitted(&fingerprint, unix_seconds()); - if already_admitted - || !matches!( - replay.claim(&fingerprint, unix_seconds()), - FirstFlightAdmit::Fresh - ) - { - error.response_session = None; + match work { + FirstFlightWork::Live { + key, + payload, + error_session, + } => { + let context = auth + .authenticate_presented(key_id, &key) + .map_err(|failure| ServerInitialError { + failure, + response_session: reserved_error_session( + &mut replay, + &fingerprint, + error_session, + ), + presented_key_id: Some(key_id), + })?; + match replay.admit(key_id, &fingerprint, unix_seconds()) { + FirstFlightAdmit::Fresh => Ok((payload, context)), + FirstFlightAdmit::Replayed => Err(first_flight_error( + "connection_salt_replayed", + "protocol-v2 connection salt was already accepted", + true, + key_id, + )), + FirstFlightAdmit::Limited => Err(first_flight_error( + "connection_admission_limited", + "this credential has opened too many new connections in the current window", + true, + key_id, + )), + FirstFlightAdmit::Unavailable => Err(first_flight_error( + "connection_replay_store_unavailable", + "failed to persist first-flight replay admission", + true, + key_id, + )), + } + } + FirstFlightWork::Stale(mut error) => { + if let Some(session) = error.response_session.take() { + error.response_session = reserved_error_session(&mut replay, &fingerprint, session); + } + Err(error) + } } - error } fn stale_root_first_flight( diff --git a/src/common/message/secure/replay.rs b/src/common/message/secure/replay.rs index 6fb8bbc..3214b56 100644 --- a/src/common/message/secure/replay.rs +++ b/src/common/message/secure/replay.rs @@ -156,15 +156,6 @@ impl ReplayGuard { self } - pub(super) fn already_admitted(&mut self, fingerprint: &[u8; 32], now: u64) -> bool { - let present = self.bloom.contains(fingerprint, now); - if self.bloom.current_started_at != self.total_started_at { - self.total = 0; - self.total_started_at = self.bloom.current_started_at; - } - present - } - pub(super) fn admit( &mut self, key_id: u64, @@ -175,38 +166,41 @@ impl ReplayGuard { if self.bloom.contains(fingerprint, now) { return FirstFlightAdmit::Replayed; } - if self.bloom.current_started_at != self.total_started_at { - self.total = 0; - self.total_started_at = self.bloom.current_started_at; - } + self.reset_total_if_rotated(); if self.counts.get(&key_id).copied().unwrap_or(0) >= self.max_per_key || self.total >= self.max_total { return FirstFlightAdmit::Limited; } - if self.persist(fingerprint, now).is_err() { - return FirstFlightAdmit::Unavailable; - } - self.bloom.insert(fingerprint, now); - *self.counts.entry(key_id).or_insert(0) += 1; - self.total = self.total.saturating_add(1); - if now.saturating_sub(self.last_compact_at) >= REPLAY_COMPACT_INTERVAL_SECONDS { - self.compact(now); + let result = self.reserve(fingerprint, now); + if result == FirstFlightAdmit::Fresh { + *self.counts.entry(key_id).or_insert(0) += 1; + if now.saturating_sub(self.last_compact_at) >= REPLAY_COMPACT_INTERVAL_SECONDS { + self.compact(now); + } } - FirstFlightAdmit::Fresh + result } pub(super) fn claim(&mut self, fingerprint: &[u8; 32], now: u64) -> FirstFlightAdmit { if self.bloom.contains(fingerprint, now) { return FirstFlightAdmit::Replayed; } + self.reset_total_if_rotated(); + if self.total >= self.max_total { + return FirstFlightAdmit::Unavailable; + } + self.reserve(fingerprint, now) + } + + fn reset_total_if_rotated(&mut self) { if self.bloom.current_started_at != self.total_started_at { self.total = 0; self.total_started_at = self.bloom.current_started_at; } - if self.total >= self.max_total { - return FirstFlightAdmit::Unavailable; - } + } + + fn reserve(&mut self, fingerprint: &[u8; 32], now: u64) -> FirstFlightAdmit { if self.persist(fingerprint, now).is_err() { return FirstFlightAdmit::Unavailable; } diff --git a/src/local/server/mod.rs b/src/local/server/mod.rs index 338adf3..9e4a306 100644 --- a/src/local/server/mod.rs +++ b/src/local/server/mod.rs @@ -8,7 +8,7 @@ use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use snafu::ResultExt; use tokio::net::TcpStream; -use tokio::task::JoinHandle; +use tokio::task::JoinSet; use tokio::time::MissedTickBehavior; use tracing::instrument; @@ -340,11 +340,11 @@ async fn run_server_side_cli_pool( pool_size, "starting local server control connection pool" ); - let mut worker_handles = Vec::new(); + let mut workers = JoinSet::new(); if pool_size > 1 { for worker_index in 1..pool_size { let worker_key = key.clone(); - worker_handles.push(AbortOnDropHandle(Some(tokio::spawn(async move { + workers.spawn(async move { run_server_side_cli_worker_with_credential::( local_addr, remote_addr, @@ -355,7 +355,7 @@ async fn run_server_side_cli_pool( pinned_credential, ) .await; - })))); + }); } } run_server_side_cli_worker_with_credential::( @@ -368,8 +368,8 @@ async fn run_server_side_cli_pool( pinned_credential, ) .await; - for handle in worker_handles { - if let Err(e) = handle.join().await { + while let Some(result) = workers.join_next().await { + if let Err(e) = result { tracing::warn!( event = "local_server_control_worker_join_failed", error = %e, @@ -379,25 +379,6 @@ async fn run_server_side_cli_pool( } } -struct AbortOnDropHandle(Option>); - -impl AbortOnDropHandle { - async fn join(mut self) -> Result<(), tokio::task::JoinError> { - match self.0.take() { - Some(handle) => handle.await, - None => Ok(()), - } - } -} - -impl Drop for AbortOnDropHandle { - fn drop(&mut self) { - if let Some(handle) = self.0.take() { - handle.abort(); - } - } -} - async fn run_server_side_cli_worker_with_credential( local_addr: A, remote_addr: A, diff --git a/ui/native/pb_mapper_ffi/src/state.rs b/ui/native/pb_mapper_ffi/src/state.rs index 01c88b4..bae25e1 100644 --- a/ui/native/pb_mapper_ffi/src/state.rs +++ b/ui/native/pb_mapper_ffi/src/state.rs @@ -423,6 +423,12 @@ fn claim_key( }) } +#[derive(Clone, Copy)] +struct PinnedTunnel { + credential: Credential, + endpoint: SocketAddr, +} + /// Everything [`PbMapperState::finish_register`] needs once the slow work is done. struct RegisterCommit { service_key: String, @@ -457,10 +463,8 @@ pub struct PbMapperState { active_connections: Arc>>, service_handles: HashMap>, client_handles: HashMap>, - service_credentials: HashMap, - client_credentials: HashMap, - service_endpoints: HashMap, - client_endpoints: HashMap, + service_tunnels: HashMap, + client_tunnels: HashMap, config: AppConfig, config_dir: PathBuf, app_directory_path: Option, diff --git a/ui/native/pb_mapper_ffi/src/state/configuration.rs b/ui/native/pb_mapper_ffi/src/state/configuration.rs index 601856c..d78fee0 100644 --- a/ui/native/pb_mapper_ffi/src/state/configuration.rs +++ b/ui/native/pb_mapper_ffi/src/state/configuration.rs @@ -55,10 +55,8 @@ impl PbMapperState { active_connections: Arc::new(RwLock::new(HashMap::new())), service_handles: HashMap::new(), client_handles: HashMap::new(), - service_credentials: HashMap::new(), - client_credentials: HashMap::new(), - service_endpoints: HashMap::new(), - client_endpoints: HashMap::new(), + service_tunnels: HashMap::new(), + client_tunnels: HashMap::new(), config: AppConfig::default(), config_dir: config_dir.clone(), app_directory_path: app_directory_path.clone(), @@ -95,10 +93,8 @@ impl PbMapperState { active_connections: Arc::new(RwLock::new(HashMap::new())), service_handles: HashMap::new(), client_handles: HashMap::new(), - service_credentials: HashMap::new(), - client_credentials: HashMap::new(), - service_endpoints: HashMap::new(), - client_endpoints: HashMap::new(), + service_tunnels: HashMap::new(), + client_tunnels: HashMap::new(), config, config_dir, app_directory_path, diff --git a/ui/native/pb_mapper_ffi/src/state/runtime.rs b/ui/native/pb_mapper_ffi/src/state/runtime.rs index 3e228dd..b803fa9 100644 --- a/ui/native/pb_mapper_ffi/src/state/runtime.rs +++ b/ui/native/pb_mapper_ffi/src/state/runtime.rs @@ -136,14 +136,12 @@ impl PbMapperState { for (_, handle) in self.service_handles.drain() { handle.abort(); } - self.service_credentials.clear(); - self.service_endpoints.clear(); + self.service_tunnels.clear(); for (_, handle) in self.client_handles.drain() { handle.abort(); } - self.client_credentials.clear(); - self.client_endpoints.clear(); + self.client_tunnels.clear(); self.registered_services.write().await.clear(); self.active_connections.write().await.clear(); @@ -167,7 +165,7 @@ impl PbMapperState { credential, } = commit; - self.service_credentials.remove(&service_key); + self.service_tunnels.remove(&service_key); if let Some(previous) = self.service_handles.remove(&service_key) { tracing::warn!( "Service '{service_key}' is already registered, replacing existing handle" @@ -206,10 +204,13 @@ impl PbMapperState { ); }); - self.service_credentials - .insert(service_key.clone(), credential); - self.service_endpoints - .insert(service_key.clone(), remote_sock_addr); + self.service_tunnels.insert( + service_key.clone(), + PinnedTunnel { + credential, + endpoint: remote_sock_addr, + }, + ); let handle = if protocol.to_uppercase() == "TCP" { tokio::spawn(async move { let _ = run_server_side_cli_with_pinned_credential::( @@ -280,8 +281,7 @@ impl PbMapperState { } pub async fn unregister_service(&mut self, service_key: String) -> Result<(), CtlError> { - self.service_credentials.remove(&service_key); - self.service_endpoints.remove(&service_key); + self.service_tunnels.remove(&service_key); if let Some(handle) = self.service_handles.remove(&service_key) { handle.abort(); } @@ -306,8 +306,7 @@ impl PbMapperState { &mut self, service_key: String, ) -> Result<(), CtlError> { - self.service_credentials.remove(&service_key); - self.service_endpoints.remove(&service_key); + self.service_tunnels.remove(&service_key); if let Some(handle) = self.service_handles.remove(&service_key) { handle.abort(); } @@ -328,7 +327,7 @@ impl PbMapperState { credential, } = commit; - self.client_credentials.remove(&service_key); + self.client_tunnels.remove(&service_key); if let Some(previous) = self.client_handles.remove(&service_key) { tracing::warn!( "Client for service '{service_key}' is already connected, replacing handle" @@ -357,10 +356,13 @@ impl PbMapperState { }) }; - self.client_credentials - .insert(service_key.clone(), credential); - self.client_endpoints - .insert(service_key.clone(), remote_sock_addr); + self.client_tunnels.insert( + service_key.clone(), + PinnedTunnel { + credential, + endpoint: remote_sock_addr, + }, + ); let handle = if protocol_upper == "TCP" { tokio::spawn(async move { run_client_side_cli_with_pinned_credential::( @@ -442,8 +444,7 @@ impl PbMapperState { pub async fn disconnect_service(&mut self, service_key: String) -> Result<(), CtlError> { // Aborting the task is the part that matters: it is what stops the // retry loop still dialling in the background. - self.client_credentials.remove(&service_key); - self.client_endpoints.remove(&service_key); + self.client_tunnels.remove(&service_key); let aborted = match self.client_handles.remove(&service_key) { Some(handle) => { handle.abort(); @@ -477,8 +478,7 @@ impl PbMapperState { &mut self, service_key: String, ) -> Result<(), CtlError> { - self.client_credentials.remove(&service_key); - self.client_endpoints.remove(&service_key); + self.client_tunnels.remove(&service_key); if let Some(handle) = self.client_handles.remove(&service_key) { handle.abort(); } diff --git a/ui/native/pb_mapper_ffi/src/state/status.rs b/ui/native/pb_mapper_ffi/src/state/status.rs index e6fd8e2..9c4f38a 100644 --- a/ui/native/pb_mapper_ffi/src/state/status.rs +++ b/ui/native/pb_mapper_ffi/src/state/status.rs @@ -369,15 +369,14 @@ impl PbMapperState { refreshing.insert(service_key.to_string()); } - let server_addr = self - .service_endpoints - .get(service_key) - .map(ToString::to_string) + let tunnel = self.service_tunnels.get(service_key); + let server_addr = tunnel + .map(|tunnel| tunnel.endpoint.to_string()) .unwrap_or_else(|| self.config.server_address.clone()); let cache = self.service_status_cache.clone(); let refreshing = self.service_status_refreshing.clone(); let key = service_key.to_string(); - let credential = self.service_credentials.get(service_key).copied(); + let credential = tunnel.map(|tunnel| tunnel.credential); tokio::spawn(async move { let result = tokio::time::timeout( @@ -437,15 +436,14 @@ impl PbMapperState { refreshing.insert(service_key.to_string()); } - let server_addr = self - .client_endpoints - .get(service_key) - .map(ToString::to_string) + let tunnel = self.client_tunnels.get(service_key); + let server_addr = tunnel + .map(|tunnel| tunnel.endpoint.to_string()) .unwrap_or_else(|| self.config.server_address.clone()); let cache = self.client_status_cache.clone(); let refreshing = self.client_status_refreshing.clone(); let key = service_key.to_string(); - let credential = self.client_credentials.get(service_key).copied(); + let credential = tunnel.map(|tunnel| tunnel.credential); tokio::spawn(async move { let result = tokio::time::timeout( From 82817803f80cf7f237a2725099c868dac8e4dd8c Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 12:44:41 +0800 Subject: [PATCH 57/74] Split auth persistence and first-flight modules Keep types in the auth root. Move config, key loading, snapshot/WAL/fs, and actor lifecycle/epoch into focused files. Share reset/rotate wipe, snapshot construction, and recovery-key persistence. --- src/common/auth.rs | 487 +------- src/common/auth/actor.rs | 1263 --------------------- src/common/auth/actor/epoch.rs | 166 +++ src/common/auth/actor/lifecycle.rs | 516 +++++++++ src/common/auth/actor/mod.rs | 585 ++++++++++ src/common/auth/config.rs | 188 +++ src/common/auth/keys.rs | 243 ++++ src/common/auth/persistence.rs | 1187 ------------------- src/common/auth/persistence/admin_key.rs | 284 +++++ src/common/auth/persistence/blob.rs | 73 ++ src/common/auth/persistence/fs.rs | 275 +++++ src/common/auth/persistence/mod.rs | 77 ++ src/common/auth/persistence/snapshot.rs | 322 ++++++ src/common/auth/persistence/wal.rs | 229 ++++ src/common/message/secure.rs | 150 +-- src/common/message/secure/first_flight.rs | 150 +++ 16 files changed, 3157 insertions(+), 3038 deletions(-) delete mode 100644 src/common/auth/actor.rs create mode 100644 src/common/auth/actor/epoch.rs create mode 100644 src/common/auth/actor/lifecycle.rs create mode 100644 src/common/auth/actor/mod.rs create mode 100644 src/common/auth/config.rs create mode 100644 src/common/auth/keys.rs delete mode 100644 src/common/auth/persistence.rs create mode 100644 src/common/auth/persistence/admin_key.rs create mode 100644 src/common/auth/persistence/blob.rs create mode 100644 src/common/auth/persistence/fs.rs create mode 100644 src/common/auth/persistence/mod.rs create mode 100644 src/common/auth/persistence/snapshot.rs create mode 100644 src/common/auth/persistence/wal.rs create mode 100644 src/common/message/secure/first_flight.rs diff --git a/src/common/auth.rs b/src/common/auth.rs index 512a14c..0545153 100644 --- a/src/common/auth.rs +++ b/src/common/auth.rs @@ -84,192 +84,6 @@ pub struct AuthConfig { pub legacy_protocol: LegacyProtocolPolicy, } -pub fn default_auth_state_dir() -> PathBuf { - std::env::var_os("PB_MAPPER_AUTH_STATE_DIR") - .map(PathBuf::from) - .unwrap_or_else(platform_default_auth_state_dir) -} - -/// Linux systemd/Docker keep `/var/lib/pb-mapper/auth` when that path is usable -/// (root, or an already-writable service directory). Unprivileged Linux, -/// macOS, and Windows binaries need an application data directory instead. -pub(crate) fn platform_default_auth_state_dir() -> PathBuf { - #[cfg(windows)] - { - let base = std::env::var_os("LOCALAPPDATA") - .or_else(|| std::env::var_os("APPDATA")) - .map(PathBuf::from) - .unwrap_or_else(|| PathBuf::from(r"C:\ProgramData")); - base.join("pb-mapper").join("auth") - } - #[cfg(target_os = "macos")] - { - match std::env::var_os("HOME") { - Some(home) => PathBuf::from(home) - .join("Library") - .join("Application Support") - .join("pb-mapper") - .join("auth"), - None => PathBuf::from("/Library/Application Support/pb-mapper/auth"), - } - } - #[cfg(not(any(windows, target_os = "macos")))] - { - linux_default_auth_state_dir( - unix_effective_uid(), - linux_system_auth_dir_usable(), - std::env::var_os("XDG_DATA_HOME").as_deref(), - std::env::var_os("HOME").as_deref(), - ) - } -} - -pub(crate) fn linux_default_auth_state_dir( - euid: u32, - system_dir_usable: bool, - xdg_data_home: Option<&std::ffi::OsStr>, - home: Option<&std::ffi::OsStr>, -) -> PathBuf { - if euid == 0 || system_dir_usable { - return PathBuf::from(DEFAULT_AUTH_STATE_DIR); - } - if let Some(xdg) = xdg_data_home { - if !xdg.is_empty() { - return PathBuf::from(xdg).join("pb-mapper").join("auth"); - } - } - if let Some(home) = home { - if !home.is_empty() { - return PathBuf::from(home) - .join(".local") - .join("share") - .join("pb-mapper") - .join("auth"); - } - } - PathBuf::from(DEFAULT_AUTH_STATE_DIR) -} - -#[cfg(not(any(windows, target_os = "macos")))] -fn unix_effective_uid() -> u32 { - extern "C" { - fn geteuid() -> u32; - } - unsafe { geteuid() } -} - -#[cfg(not(any(windows, target_os = "macos")))] -fn linux_system_auth_dir_usable() -> bool { - let path = Path::new(DEFAULT_AUTH_STATE_DIR); - path.is_dir() && unix_path_is_writable(path) -} - -#[cfg(not(any(windows, target_os = "macos")))] -fn unix_path_is_writable(path: &Path) -> bool { - use std::os::unix::ffi::OsStrExt; - let Ok(c_path) = std::ffi::CString::new(path.as_os_str().as_bytes()) else { - return false; - }; - extern "C" { - fn access(pathname: *const std::os::raw::c_char, mode: i32) -> i32; - } - const W_OK: i32 = 2; - unsafe { access(c_path.as_ptr(), W_OK) == 0 } -} - -impl Default for AuthConfig { - fn default() -> Self { - Self { - state_dir: default_auth_state_dir(), - max_temporary_keys: env_usize( - "PB_MAPPER_AUTH_MAX_TEMP_KEYS", - DEFAULT_TEMP_KEY_CAPACITY, - 1, - MAX_TEMP_KEY_CAPACITY, - ), - max_temporary_key_ttl: Duration::from_secs(env_u64( - "PB_MAPPER_AUTH_MAX_TEMP_TTL_SECS", - DEFAULT_MAX_TEMP_KEY_TTL.as_secs(), - MIN_TEMP_KEY_TTL.as_secs(), - MAX_TEMP_KEY_TTL.as_secs(), - )), - legacy_protocol: legacy_protocol_from_env(), - } - } -} - -fn legacy_protocol_from_env() -> LegacyProtocolPolicy { - match std::env::var("PB_MAPPER_LEGACY_PROTOCOL") { - Err(std::env::VarError::NotPresent) => LegacyProtocolPolicy::Allow, - Err(std::env::VarError::NotUnicode(_)) => { - tracing::error!( - event = "legacy_protocol_config_invalid", - "PB_MAPPER_LEGACY_PROTOCOL is not UTF-8; denying legacy framing" - ); - LegacyProtocolPolicy::Deny - } - Ok(value) => parse_legacy_protocol_policy(&value).unwrap_or_else(|| { - tracing::error!( - event = "legacy_protocol_config_invalid", - value, - "PB_MAPPER_LEGACY_PROTOCOL must be `allow` or `deny`; denying legacy framing" - ); - LegacyProtocolPolicy::Deny - }), - } -} - -fn parse_legacy_protocol_policy(value: &str) -> Option { - match value.trim().to_ascii_lowercase().as_str() { - "allow" => Some(LegacyProtocolPolicy::Allow), - "deny" => Some(LegacyProtocolPolicy::Deny), - _ => None, - } -} - -fn env_usize(name: &str, default: usize, min: usize, max: usize) -> usize { - env_bounded(name, default, min, max) -} - -fn env_u64(name: &str, default: u64, min: u64, max: u64) -> u64 { - env_bounded(name, default, min, max) -} - -fn env_bounded(name: &str, default: T, min: T, max: T) -> T -where - T: std::str::FromStr + PartialOrd + Copy + fmt::Display, -{ - match std::env::var(name) { - Err(std::env::VarError::NotPresent) => default, - Ok(raw) => match raw.parse::() { - Ok(value) if value >= min && value <= max => value, - _ => { - tracing::warn!( - event = "auth_config_value_invalid", - variable = name, - value = raw, - min = %min, - max = %max, - fallback = %default, - "invalid authentication configuration value; using the default" - ); - default - } - }, - Err(std::env::VarError::NotUnicode(_)) => { - tracing::warn!( - event = "auth_config_value_invalid", - variable = name, - min = %min, - max = %max, - fallback = %default, - "authentication configuration value is not UTF-8; using the default" - ); - default - } - } -} - #[derive(Clone, Debug, Serialize, Deserialize)] pub struct AuthFailure { pub code: String, @@ -413,26 +227,26 @@ impl AuthContext { } pub(crate) fn admin_cancellation_token(&self) -> Result { - if !self.is_admin { - return Err(AuthFailure::new( - "admin_permission_required", - "administrator credential is required for this operation", - false, - )); - } + self.require_admin()?; self.cancellation_token() } fn admin_authority(&self) -> Result, AuthFailure> { - if !self.is_admin { - return Err(AuthFailure::new( + self.require_admin()?; + self.ensure_active()?; + Ok(self.lease.clone()) + } + + fn require_admin(&self) -> Result<(), AuthFailure> { + if self.is_admin { + Ok(()) + } else { + Err(AuthFailure::new( "admin_permission_required", "administrator credential is required for this operation", false, - )); + )) } - self.ensure_active()?; - Ok(self.lease.clone()) } } @@ -531,19 +345,17 @@ struct AuthStateInner { audit_records: RwLock>, } +fn recover_lock(result: std::sync::LockResult) -> T { + result.unwrap_or_else(|poisoned| poisoned.into_inner()) +} + impl AuthStateInner { fn admin_key(&self) -> AesKeyType { - self.admin - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .key + recover_lock(self.admin.read()).key } fn instance_id(&self) -> [u8; INSTANCE_ID_LEN] { - *self - .instance_id - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()) + *recover_lock(self.instance_id.read()) } } @@ -673,6 +485,21 @@ enum AuthCommand { }, } +mod config; +pub use config::default_auth_state_dir; +#[cfg(test)] +pub(in crate::common::auth) use config::parse_legacy_protocol_policy; +#[cfg(test)] +pub(crate) use config::{linux_default_auth_state_dir, platform_default_auth_state_dir}; +#[cfg(all(test, not(any(windows, target_os = "macos"))))] +pub(in crate::common::auth) use config::{linux_system_auth_dir_usable, unix_effective_uid}; +mod keys; +#[cfg(test)] +pub(in crate::common::auth) use keys::recover_admin_key_after_rotation; +pub use keys::{derive_temporary_key, key_generation, key_slot, make_key_id}; +pub(in crate::common::auth) use keys::{ + load_isolated_server_admin_credential, load_server_admin_credential, +}; mod runtime; pub struct LegacyConnectionGuard { @@ -689,240 +516,6 @@ impl Drop for LegacyConnectionGuard { } } -fn read_admin_key(path: &Path) -> Result, AuthFailure> { - if !path.exists() { - return Ok(None); - } - #[cfg(unix)] - { - let metadata = std::fs::metadata(path).map_err(|error| { - AuthFailure::new( - "administrator_key_required", - format!( - "administrator key file `{}` metadata could not be read: {error}", - path.display() - ), - false, - ) - })?; - if metadata.permissions().mode() & 0o077 != 0 { - std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600)).map_err( - |error| { - AuthFailure::new( - "administrator_key_required", - format!( - "administrator key file `{}` permissions could not be secured: {error}", - path.display() - ), - false, - ) - }, - )?; - tracing::warn!( - event = "administrator_key_permissions_repaired", - path = %path.display(), - "restricted administrator key file permissions to 0600" - ); - } - } - std::fs::read_to_string(path).map(Some).map_err(|error| { - AuthFailure::new( - "administrator_key_required", - format!( - "administrator key file `{}` could not be read: {error}", - path.display() - ), - false, - ) - }) -} - -fn validate_admin_credential(raw: &str) -> Result { - let credential = parse_credential(raw.trim()) - .map_err(|error| AuthFailure::new("administrator_key_invalid", error, false))?; - if !credential.is_admin() { - return Err(AuthFailure::new( - "administrator_key_required", - "the server key file contains a temporary credential", - false, - )); - } - Ok(credential) -} - -fn recover_admin_key_after_rotation( - state_dir: &Path, - current: &str, -) -> Result { - let snapshot_path = auth_snapshot_path(state_dir); - if !snapshot_path.exists() { - return Ok(current.to_string()); - } - let bytes = std::fs::read(&snapshot_path).map_err(|error| { - AuthFailure::new( - "temporary_key_store_unavailable", - format!("failed to read `{}`: {error}", snapshot_path.display()), - false, - ) - })?; - if let Ok(Credential::Admin(current_key)) = parse_credential(current.trim()) { - if open_blob(¤t_key, &bytes).is_ok() { - return Ok(current.to_string()); - } - } - let Some(next) = read_admin_key(&state_dir.join("admin.key.next"))? else { - return Ok(current.to_string()); - }; - let Ok(Credential::Admin(next_key)) = parse_credential(next.trim()) else { - return Ok(current.to_string()); - }; - if open_blob(&next_key, &bytes).is_err() { - return Ok(current.to_string()); - } - // The rotation snapshot is complete under the staged key. Leftover WAL - // records are still encrypted with the previous key. - truncate_auth_wal(state_dir)?; - write_admin_key(state_dir, next.trim())?; - let _ = std::fs::remove_file(state_dir.join("admin.key.next")); - Ok(next) -} - -fn load_server_admin_credential(state_dir: &Path) -> Result { - let path = state_dir.join("admin.key"); - let raw = if let Some(raw) = read_admin_key(&path)? { - raw - } else if std::env::var_os(ENV_MSG_HEADER_KEY).is_some() { - let credential = get_process_credential() - .map_err(|error| AuthFailure::new("administrator_key_invalid", error, false))?; - let Credential::Admin(key) = credential else { - return Err(AuthFailure::new( - "administrator_key_required", - "the relay server cannot start with a temporary credential", - false, - )); - }; - let key = String::from_utf8(key.to_vec()).map_err(|_| { - AuthFailure::new( - "administrator_key_invalid", - "the relay administrator key must be printable UTF-8 so it can be persisted", - false, - ) - })?; - if encrypted_auth_state_exists(state_dir) - && !key_matches_existing_state(Some(state_dir), &key) - { - return Err(AuthFailure::new( - "administrator_key_invalid", - "MSG_HEADER_KEY does not decrypt the existing authentication state; refusing to write admin.key", - false, - )); - } - write_admin_key(state_dir, &key)?; - key - } else if Path::new(MACHINE_MSG_HEADER_KEY_PATH).is_file() { - let key = std::fs::read_to_string(MACHINE_MSG_HEADER_KEY_PATH).map_err(|error| { - AuthFailure::new( - "administrator_key_required", - format!( - "legacy administrator key file `{MACHINE_MSG_HEADER_KEY_PATH}` could not be read: {error}" - ), - false, - ) - })?; - validate_admin_credential(&key)?; - if encrypted_auth_state_exists(state_dir) - && !key_matches_existing_state(Some(state_dir), key.trim()) - { - return Err(AuthFailure::new( - "administrator_key_invalid", - "legacy administrator key does not decrypt the existing authentication state; refusing to write admin.key", - false, - )); - } - write_admin_key(state_dir, key.trim())?; - tracing::warn!( - event = "administrator_key_migrated", - source = MACHINE_MSG_HEADER_KEY_PATH, - destination = %path.display(), - "migrated the legacy administrator key into the v0.4 authentication state directory" - ); - key - } else { - let key = initialize_admin_key(&path, false)?; - tracing::warn!( - event = "administrator_key_initialized", - path = %path.display(), - "no administrator credential was configured; generated a random key file" - ); - key - }; - let raw = recover_admin_key_after_rotation(state_dir, &raw)?; - let credential = validate_admin_credential(&raw)?; - set_process_msg_header_key(Some(raw.trim())).map_err(AuthFailure::internal)?; - Ok(credential) -} - -/// Load or create an app-local relay root without reading or mutating the process credential. -/// -/// The Flutter process uses its configured process credential for the remote relay, while its -/// optional embedded relay owns an independent administrator key under the app data directory. -fn load_isolated_server_admin_credential(state_dir: &Path) -> Result { - let path = state_dir.join("admin.key"); - let raw = match read_admin_key(&path)? { - Some(raw) => raw, - None => { - let key = initialize_admin_key(&path, false)?; - tracing::warn!( - event = "isolated_administrator_key_initialized", - path = %path.display(), - "generated an administrator key for an embedded relay" - ); - key - } - }; - let raw = recover_admin_key_after_rotation(state_dir, &raw)?; - validate_admin_credential(&raw) -} - -pub fn make_key_id(generation: u32, slot: u32) -> u64 { - (u64::from(generation) << 32) | u64::from(slot) -} - -pub fn key_generation(key_id: u64) -> u32 { - (key_id >> 32) as u32 -} - -pub fn key_slot(key_id: u64) -> u32 { - key_id as u32 -} - -pub fn derive_temporary_key( - admin_key: &AesKeyType, - instance_id: &[u8; INSTANCE_ID_LEN], - key_id: u64, -) -> Result { - let salt = Salt::new(HKDF_SHA256, instance_id); - let pseudo_random_key = salt.extract(admin_key); - let key_id_bytes = key_id.to_be_bytes(); - let info = [b"pb-mapper-temp-key-v1".as_slice(), key_id_bytes.as_slice()]; - let output = pseudo_random_key - .expand(&info, HkdfLen(32)) - .map_err(|_| AuthFailure::internal("failed to expand temporary key"))?; - let mut key = [0_u8; 32]; - output - .fill(&mut key) - .map_err(|_| AuthFailure::internal("failed to fill temporary key"))?; - Ok(key) -} - -struct HkdfLen(usize); - -impl ring::hkdf::KeyType for HkdfLen { - fn len(&self) -> usize { - self.0 - } -} - #[derive(Clone, Debug, Serialize, Deserialize)] struct PersistedEntry { key_id: u64, @@ -1000,6 +593,20 @@ mod actor; use actor::{run_auth_actor, AuthActorState}; mod persistence; pub use persistence::*; +pub(in crate::common::auth) use persistence::{ + append_audit, append_mutation, append_wal, atomic_write, auth_snapshot_path, build_snapshot, + cancel_all_temporary_leases, clear_retained_high_slot_entries, compaction_is_allowed, + empty_snapshot, fail_closed_on_uncertain_wal, hex, key_matches_existing_state, + load_or_create_instance_id, load_persisted_state, normalize_tombstone_times, open_blob, + prepare_state_dir_and_lock, push_audit_record, push_persisted_audit, random_instance_id, + recover_instance_id_after_reset, reset_already_installed, rotation_already_installed, + split_high_slot_state, truncate_auth_wal, unix_seconds, write_admin_key, + write_snapshot_and_truncate_wal, +}; +#[cfg(test)] +pub(in crate::common::auth) use persistence::{ + prepare_state_dir, read_instance_id_file, try_load_persisted_state, +}; mod timing_wheel; use timing_wheel::TimingWheel; #[cfg(test)] diff --git a/src/common/auth/actor.rs b/src/common/auth/actor.rs deleted file mode 100644 index ce8047d..0000000 --- a/src/common/auth/actor.rs +++ /dev/null @@ -1,1263 +0,0 @@ -//! Serialized owner of mutable authentication lifecycle state. -//! -//! ```text -//! authenticated admin command -//! | -//! v -//! validate current admin lease -//! | -//! v -//! append encrypted WAL -> mutate slots / leases / timing wheel -//! | -//! +-> periodic snapshot + bounded replay/audit retention -//! ``` -//! -//! Keeping authorization revalidation and mutations in one actor prevents a request -//! authenticated before root rotation from executing against the new administrator -//! state. The actor is also the sole strong owner of temporary-key leases. - -use super::*; - -pub(super) struct AuthActorState { - cold: HashMap, - wheel: TimingWheel, - admin_replays: HashSet<[u8; 32]>, - admin_replay_order: VecDeque, -} - -impl AuthActorState { - pub(super) fn new( - cold: HashMap, - wheel: TimingWheel, - admin_replays: HashSet<[u8; 32]>, - admin_replay_order: VecDeque, - ) -> Self { - Self { - cold, - wheel, - admin_replays, - admin_replay_order, - } - } -} - -pub(super) async fn run_auth_actor( - inner: Arc, - mut admin_lease: Arc, - mut command_rx: mpsc::Receiver, - config: AuthConfig, - state: AuthActorState, - _state_lock: Arc, -) { - let AuthActorState { - mut cold, - mut wheel, - mut admin_replays, - mut admin_replay_order, - } = state; - let now = unix_seconds(); - let mut tombstones = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .iter() - .enumerate() - .filter_map(|(index, slot)| { - if !matches!(slot.state, SlotState::Expired | SlotState::Revoked) { - return None; - } - let key_id = make_key_id(slot.generation, index as u32); - let tombstoned_at = cold - .get(&key_id) - .map(|metadata| metadata.tombstoned_at) - .unwrap_or(now); - Some(( - tombstoned_at.saturating_add(TOMBSTONE_RETENTION.as_secs()), - key_id, - )) - }) - .collect::>(); - tombstones.extend( - inner - .high_slot_entries - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .iter() - .filter_map(|entry| { - let expired_active = entry.state == SlotState::Active && entry.expires_at <= now; - if !matches!(entry.state, SlotState::Expired | SlotState::Revoked) - && !expired_active - { - return None; - } - let tombstoned_at = entry - .tombstoned_at - .or_else(|| { - cold.get(&entry.key_id) - .map(|metadata| metadata.tombstoned_at) - }) - .unwrap_or(entry.expires_at.max(now)); - Some(( - tombstoned_at.saturating_add(TOMBSTONE_RETENTION.as_secs()), - entry.key_id, - )) - }), - ); - tombstones.sort_unstable_by_key(|(cleanup_at, _)| *cleanup_at); - let mut tombstones = VecDeque::from(tombstones); - let mut last_snapshot_at = unix_seconds(); - let mut tick = tokio::time::interval(Duration::from_secs(1)); - tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - loop { - tokio::select! { - _ = tick.tick() => { - let now = unix_seconds(); - for lease in wheel.advance(now) { - let key_id = lease.key_id(); - let version = lease.wheel_version.load(Ordering::Acquire); - if lease.expires_at() > now { - wheel.insert_with_version(lease, version); - continue; - } - let mut slots = inner.slots.write().unwrap_or_else(|poisoned| poisoned.into_inner()); - if let Some(slot) = slots.get_mut(key_slot(key_id) as usize) { - if slot.generation == key_generation(key_id) && slot.state == SlotState::Active { - slot.state = SlotState::Expired; - lease.cancel_expired(); - let tombstoned_at = slot.expires_at; - if let Some(metadata) = cold.get_mut(&key_id) { - metadata.tombstoned_at = tombstoned_at; - } - push_tombstone(&mut tombstones, tombstoned_at, key_id); - tracing::info!( - event = "temporary_key_expired", - auth_stage = "expiry", - key_id, - expires_at = lease.expires_at(), - "temporary key expired and active work was cancelled" - ); - } - } else { - lease.cancel_expired(); - } - } - expire_due_high_slots(&inner, &mut cold, &mut tombstones, now); - let mut due_high = Vec::new(); - while let Some((cleanup_at, key_id)) = tombstones.front().copied() { - if cleanup_at > now { - break; - } - tombstones.pop_front(); - let mut slots = inner.slots.write().unwrap_or_else(|poisoned| poisoned.into_inner()); - if let Some(slot) = slots.get_mut(key_slot(key_id) as usize) { - if slot.generation == key_generation(key_id) && matches!(slot.state, SlotState::Expired | SlotState::Revoked) { - slot.state = SlotState::Free; - slot.expires_at = 0; - slot.lease = Weak::new(); - cold.remove(&key_id); - wheel.release(key_id); - } - } else { - due_high.push(key_id); - } - } - if !due_high.is_empty() { - let due = due_high.iter().copied().collect::>(); - inner - .high_slot_entries - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .retain(|entry| !due.contains(&entry.key_id)); - for key_id in due_high { - cold.remove(&key_id); - wheel.release(key_id); - } - } - prune_expired_admin_replays( - now, - &mut admin_replays, - &mut admin_replay_order, - ); - // WHY: A failed load starts safe mode with empty in-memory - // generations. Compacting that reconstruction would replace the - // damaged snapshot, truncate the WAL, and let the next start - // exit safe mode without rotating the instance id. - if compaction_is_allowed(inner.safe_mode.load(Ordering::Acquire)) - && now.saturating_sub(last_snapshot_at) - >= SNAPSHOT_COMPACTION_INTERVAL.as_secs() - { - let snapshot = build_snapshot(&inner, &cold, &admin_replay_order); - if let Err(error) = write_snapshot_and_truncate_wal( - &config, - &inner.admin_key(), - &snapshot, - ) { - inner.safe_mode.store(true, Ordering::Release); - cancel_all_temporary_leases(&inner); - tracing::error!( - event = "auth_state_safe_mode", - auth_stage = "snapshot_compaction", - reason = %error.code, - error = %error, - "authentication state compaction failed closed" - ); - } else { - last_snapshot_at = now; - } - } - } - command = command_rx.recv() => { - let Some(command) = command else { - admin_lease.cancel_rotated(); - cancel_all_temporary_leases(&inner); - break; - }; - match command { - AuthCommand::ClaimAdminMutation { - authority, - fingerprint, - client_timestamp, - response, - } => { - let result = validate_admin_authority(&inner, &authority).and_then(|()| { - actor_claim_admin_mutation( - &inner, - &config, - &mut admin_replays, - &mut admin_replay_order, - fingerprint, - client_timestamp, - ) - }); - let _ = response.send(result); - } - AuthCommand::Issue { authority, ttl, label, response } => { - let result = validate_admin_authority(&inner, &authority) - .and_then(|()| actor_issue(&inner, &config, &mut cold, &mut wheel, ttl, label)); - let _ = response.send(result); - } - AuthCommand::List { authority, page, page_size, response } => { - let result = validate_admin_authority(&inner, &authority) - .and_then(|()| actor_list(&inner, &cold, page, page_size)); - let _ = response.send(result); - } - AuthCommand::Show { authority, key_id, reveal, response } => { - let result = validate_admin_authority(&inner, &authority) - .and_then(|()| actor_show(&inner, &config, &cold, key_id, reveal)); - let _ = response.send(result); - } - AuthCommand::Renew { authority, key_id, ttl, response } => { - let result = validate_admin_authority(&inner, &authority) - .and_then(|()| actor_renew(&inner, &config, &cold, &mut wheel, key_id, ttl)); - let _ = response.send(result); - } - AuthCommand::Revoke { authority, key_id, response } => { - let result = validate_admin_authority(&inner, &authority) - .and_then(|()| actor_revoke(&inner, &config, &mut cold, &mut tombstones, key_id)); - let _ = response.send(result); - } - AuthCommand::Gc { authority, response } => { - let result = validate_admin_authority(&inner, &authority) - .and_then(|()| actor_gc( - &inner, - &config, - &mut cold, - &mut wheel, - &mut tombstones, - &admin_replay_order, - )); - let _ = response.send(result); - } - AuthCommand::Reset { authority, response } => { - let result = validate_admin_authority(&inner, &authority) - .and_then(|()| actor_reset( - &inner, - &config, - &mut cold, - &mut wheel, - &admin_replay_order, - "auth_state_reset", - )); - let _ = response.send(result); - } - AuthCommand::RotateRoot { authority, new_key, response } => { - let result = validate_admin_authority(&inner, &authority) - .and_then(|()| actor_rotate_root(&inner, &config, &mut cold, &mut wheel, &mut admin_lease, new_key)); - if result.is_ok() { - admin_replays.clear(); - admin_replay_order.clear(); - } - let _ = response.send(result); - } - AuthCommand::SetLegacyProtocol { authority, policy, response } => { - let result = validate_admin_authority(&inner, &authority) - .and_then(|()| actor_set_legacy_protocol(&inner, &config, policy)); - let _ = response.send(result); - } - AuthCommand::Status { authority, response } => { - let result = validate_admin_authority(&inner, &authority) - .map(|()| actor_status(&inner)); - let _ = response.send(result); - } - AuthCommand::Audit { authority, action, key_id, detail, response } => { - let result = validate_admin_authority(&inner, &authority).and_then(|()| { - append_audit( - &config, - &inner, - audit(&action, key_id, detail), - ) - }); - let _ = response.send(result); - } - AuthCommand::Shutdown { response } => { - admin_lease.cancel_rotated(); - cancel_all_temporary_leases(&inner); - let _ = response.send(()); - break; - } - } - } - } - } -} - -fn actor_claim_admin_mutation( - inner: &AuthStateInner, - config: &AuthConfig, - admin_replays: &mut HashSet<[u8; 32]>, - admin_replay_order: &mut VecDeque, - fingerprint: [u8; 32], - client_timestamp: u64, -) -> Result<(), AuthFailure> { - let now = unix_seconds(); - prune_expired_admin_replays(now, admin_replays, admin_replay_order); - if admin_replays.contains(&fingerprint) { - return Err(AuthFailure::new( - "admin_request_replayed", - "administrator mutation was already admitted", - false, - )); - } - if admin_replays.len() >= ADMIN_REPLAY_CAPACITY { - return Err(AuthFailure::new( - "admin_replay_capacity_exhausted", - "administrator mutation replay window is full; retry after older claims expire", - true, - )); - } - if now.abs_diff(client_timestamp) > ADMIN_REPLAY_RETENTION.as_secs() / 2 { - return Err(AuthFailure::new( - "admin_request_timestamp_invalid", - "administrator mutation timestamp is outside the accepted window", - false, - )); - } - let record = AdminReplayRecord { - fingerprint, - client_timestamp, - accepted_at: now, - }; - fail_closed_on_uncertain_wal( - inner, - append_wal( - config, - &inner.admin_key(), - &WalRecord::AdminReplay(record.clone()), - ), - )?; - admin_replays.insert(fingerprint); - admin_replay_order.push_back(record); - Ok(()) -} - -pub(super) fn prune_expired_admin_replays( - now: u64, - admin_replays: &mut HashSet<[u8; 32]>, - admin_replay_order: &mut VecDeque, -) { - admin_replay_order.retain(|record| { - let keep = record.within_retention(now); - if !keep { - admin_replays.remove(&record.fingerprint); - } - keep - }); -} - -fn validate_ttl(config: &AuthConfig, ttl: Duration) -> Result { - if ttl < MIN_TEMP_KEY_TTL { - return Err(AuthFailure::new( - "temporary_key_ttl_too_short", - format!( - "temporary key TTL must be at least {} seconds", - MIN_TEMP_KEY_TTL.as_secs() - ), - false, - )); - } - if ttl > config.max_temporary_key_ttl { - return Err(AuthFailure::new( - "temporary_key_ttl_too_long", - format!( - "temporary key TTL exceeds the configured maximum of {} seconds", - config.max_temporary_key_ttl.as_secs() - ), - false, - )); - } - Ok(unix_seconds().saturating_add(ttl.as_secs())) -} - -fn validate_label(label: Option) -> Result, AuthFailure> { - let label = label - .map(|label| label.trim().to_string()) - .filter(|label| !label.is_empty()); - if label.as_ref().is_some_and(|label| label.len() > 64) { - return Err(AuthFailure::new( - "temporary_key_label_too_long", - "temporary key label must not exceed 64 UTF-8 bytes", - false, - )); - } - Ok(label) -} - -fn actor_issue( - inner: &Arc, - config: &AuthConfig, - cold: &mut HashMap, - wheel: &mut TimingWheel, - ttl: Duration, - label: Option, -) -> Result { - ensure_store_available(inner)?; - let expires_at = validate_ttl(config, ttl)?; - let label = validate_label(label)?; - let issued_at = unix_seconds(); - let (index, generation, key_id, entry) = { - let slots = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let Some((index, slot)) = slots - .iter() - .enumerate() - .find(|(_, slot)| slot.state == SlotState::Free && slot.generation < u32::MAX) - else { - return Err(AuthFailure::new( - "temporary_key_capacity_exhausted", - "temporary key slot table is full", - true, - )); - }; - let generation = slot.generation + 1; - let key_id = make_key_id(generation, index as u32); - ( - index, - generation, - key_id, - PersistedEntry { - key_id, - state: SlotState::Active, - issued_at, - expires_at, - label: label.clone(), - tombstoned_at: None, - }, - ) - }; - // Persist before taking the slot write lock. A fail-closed WAL error - // cancels leases via slots.read() and must not nest under slots.write(). - append_mutation( - config, - inner, - StateMutation::Issue(entry), - audit("temporary_key_issue", Some(key_id), label.clone()), - )?; - let mut slots = inner - .slots - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let slot = slots - .get_mut(index) - .ok_or_else(|| AuthFailure::internal("issued slot disappeared"))?; - let lease = Arc::new(AuthLease::new(key_id, expires_at)); - slot.generation = generation; - slot.state = SlotState::Active; - slot.expires_at = expires_at; - slot.issued_epoch = inner.root_epoch.load(Ordering::Acquire); - slot.lease = Arc::downgrade(&lease); - cold.insert( - key_id, - ColdMetadata { - issued_at, - label, - tombstoned_at: 0, - }, - ); - wheel.insert(lease); - drop(slots); - metadata_with_credential(inner, cold, key_id, true) -} - -fn actor_list( - inner: &Arc, - cold: &HashMap, - page: u32, - page_size: u16, -) -> Result { - let page_size = page_size.clamp(1, 1000) as usize; - let start = (page as usize).saturating_mul(page_size); - let slots = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let mut all = slots - .iter() - .enumerate() - .filter_map(|(index, slot)| { - if slot.state == SlotState::Free { - return None; - } - let key_id = make_key_id(slot.generation, index as u32); - let cold = cold.get(&key_id)?; - Some(TemporaryKeyMetadata { - key_id, - state: slot_state_name(slot.state).to_string(), - issued_at: cold.issued_at, - expires_at: slot.expires_at, - label: cold.label.clone(), - }) - }) - .collect::>(); - all.extend( - inner - .high_slot_entries - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .iter() - .filter(|entry| entry.state != SlotState::Free) - .map(high_slot_metadata), - ); - all.sort_by_key(|item| std::cmp::Reverse(item.issued_at)); - let items = all.iter().skip(start).take(page_size).cloned().collect(); - let next_page = (start.saturating_add(page_size) < all.len()).then_some(page.saturating_add(1)); - Ok(KeyPage { - schema_version: 1, - items, - next_page, - }) -} - -fn actor_show( - inner: &Arc, - config: &AuthConfig, - cold: &HashMap, - key_id: u64, - reveal: bool, -) -> Result { - let result = metadata_with_credential(inner, cold, key_id, reveal)?; - append_audit( - config, - inner, - audit( - if reveal { - "temporary_key_reveal" - } else { - "temporary_key_show" - }, - Some(key_id), - result.metadata.label.clone(), - ), - )?; - Ok(result) -} - -fn actor_renew( - inner: &Arc, - config: &AuthConfig, - cold: &HashMap, - wheel: &mut TimingWheel, - key_id: u64, - ttl: Duration, -) -> Result { - ensure_store_available(inner)?; - let expires_at = validate_ttl(config, ttl)?; - let index = key_slot(key_id) as usize; - { - let slots = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - if let Some(slot) = slots.get(index) { - validate_slot_identity(slot, key_id)?; - if slot.state != SlotState::Active || slot.expires_at <= unix_seconds() { - return Err(key_not_renewable()); - } - } else { - let high = inner - .high_slot_entries - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let entry = high_slot_entry(&high, key_id)?; - if entry.state != SlotState::Active || entry.expires_at <= unix_seconds() { - return Err(key_not_renewable()); - } - } - } - let label = cold - .get(&key_id) - .and_then(|metadata| metadata.label.clone()) - .or_else(|| { - inner - .high_slot_entries - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .iter() - .find(|entry| entry.key_id == key_id) - .and_then(|entry| entry.label.clone()) - }); - append_mutation( - config, - inner, - StateMutation::Renew { key_id, expires_at }, - audit("temporary_key_renew", Some(key_id), label), - )?; - let mut slots = inner - .slots - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - if let Some(slot) = slots.get_mut(index) { - validate_slot_identity(slot, key_id)?; - if slot.state != SlotState::Active { - return Err(AuthFailure::new( - "temporary_key_inactive", - "temporary key lease is no longer active", - true, - )); - } - slot.expires_at = expires_at; - let lease = match slot.lease.upgrade() { - Some(lease) if !lease.cancellation_token().is_cancelled() => { - lease.expires_at.store(expires_at, Ordering::Release); - lease.wheel_version.fetch_add(1, Ordering::AcqRel); - lease - } - _ => { - let lease = Arc::new(AuthLease::new(key_id, expires_at)); - slot.lease = Arc::downgrade(&lease); - lease - } - }; - wheel.release(key_id); - wheel.insert(lease); - drop(slots); - return metadata_with_credential(inner, cold, key_id, true); - } - drop(slots); - { - let mut high = inner - .high_slot_entries - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let entry = high_slot_entry_mut(&mut high, key_id)?; - if entry.state != SlotState::Active { - return Err(AuthFailure::new( - "temporary_key_inactive", - "temporary key lease is no longer active", - true, - )); - } - entry.expires_at = expires_at; - entry.tombstoned_at = None; - } - metadata_with_credential(inner, cold, key_id, true) -} - -fn actor_revoke( - inner: &Arc, - config: &AuthConfig, - cold: &mut HashMap, - tombstones: &mut VecDeque<(u64, u64)>, - key_id: u64, -) -> Result { - ensure_store_available(inner)?; - let now = unix_seconds(); - let index = key_slot(key_id) as usize; - let (label, issued_at, expires_at) = { - let slots = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - if let Some(slot) = slots.get(index) { - validate_slot_identity(slot, key_id)?; - if slot.state != SlotState::Active { - return Err(key_not_active()); - } - let metadata = cold.get(&key_id).ok_or_else(|| key_not_found(key_id))?; - (metadata.label.clone(), metadata.issued_at, slot.expires_at) - } else { - let high = inner - .high_slot_entries - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let entry = high_slot_entry(&high, key_id)?; - if entry.state != SlotState::Active { - return Err(key_not_active()); - } - (entry.label.clone(), entry.issued_at, entry.expires_at) - } - }; - append_mutation( - config, - inner, - StateMutation::Revoke { key_id, at: now }, - audit("temporary_key_revoke", Some(key_id), label.clone()), - )?; - let mut slots = inner - .slots - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - if let Some(slot) = slots.get_mut(index) { - validate_slot_identity(slot, key_id)?; - slot.state = SlotState::Revoked; - if let Some(lease) = slot.lease.upgrade() { - lease.cancel_revoked(); - } - let cold_metadata = cold.get_mut(&key_id).ok_or_else(|| key_not_found(key_id))?; - cold_metadata.tombstoned_at = now; - push_tombstone(tombstones, now, key_id); - return Ok(TemporaryKeyMetadata { - key_id, - state: slot_state_name(slot.state).to_string(), - issued_at: cold_metadata.issued_at, - expires_at: slot.expires_at, - label: cold_metadata.label.clone(), - }); - } - drop(slots); - let mut high = inner - .high_slot_entries - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let entry = high_slot_entry_mut(&mut high, key_id)?; - if entry.state != SlotState::Active { - return Err(key_not_active()); - } - entry.state = SlotState::Revoked; - entry.tombstoned_at = Some(now); - if let Some(metadata) = cold.get_mut(&key_id) { - metadata.tombstoned_at = now; - } - push_tombstone(tombstones, now, key_id); - Ok(TemporaryKeyMetadata { - key_id, - state: slot_state_name(entry.state).to_string(), - issued_at, - expires_at, - label, - }) -} - -fn remember_previous_root(inner: &AuthStateInner) { - *inner - .previous_root - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(PreviousRoot { - admin_key: inner.admin_key(), - instance_id: inner.instance_id(), - }); -} - -fn push_tombstone(tombstones: &mut VecDeque<(u64, u64)>, tombstoned_at: u64, key_id: u64) { - let cleanup_at = tombstoned_at.saturating_add(TOMBSTONE_RETENTION.as_secs()); - let index = tombstones.partition_point(|(current, _)| *current <= cleanup_at); - tombstones.insert(index, (cleanup_at, key_id)); -} - -fn actor_gc( - inner: &Arc, - config: &AuthConfig, - cold: &mut HashMap, - wheel: &mut TimingWheel, - tombstones: &mut VecDeque<(u64, u64)>, - admin_replays: &VecDeque, -) -> Result { - ensure_store_available(inner)?; - let now = unix_seconds(); - let mut removed = 0_u64; - let mut slots = inner - .slots - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - for (index, slot) in slots.iter_mut().enumerate() { - if matches!(slot.state, SlotState::Expired | SlotState::Revoked) - || (slot.state == SlotState::Active && slot.expires_at <= now) - { - let key_id = make_key_id(slot.generation, index as u32); - if let Some(lease) = slot.lease.upgrade() { - if slot.state == SlotState::Revoked { - lease.cancel_revoked(); - } else { - lease.cancel_expired(); - } - } - slot.state = SlotState::Free; - slot.expires_at = 0; - slot.lease = Weak::new(); - cold.remove(&key_id); - wheel.release(key_id); - removed = removed.saturating_add(1); - } - } - drop(slots); - { - let mut high = inner - .high_slot_entries - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - high.retain(|entry| { - let keep = match entry.state { - SlotState::Active if entry.expires_at > now => true, - SlotState::Active | SlotState::Expired | SlotState::Revoked | SlotState::Free => { - false - } - }; - if !keep { - cold.remove(&entry.key_id); - wheel.release(entry.key_id); - removed = removed.saturating_add(1); - } - keep - }); - } - tombstones.clear(); - let gc_audit = audit("temporary_key_gc", None, Some(format!("removed={removed}"))); - let mut snapshot = build_snapshot(inner, cold, admin_replays); - push_persisted_audit(&mut snapshot.audit_records, gc_audit.clone()); - let admin_key = inner.admin_key(); - if let Err(error) = write_snapshot_and_truncate_wal(config, &admin_key, &snapshot) { - inner.safe_mode.store(true, Ordering::Release); - cancel_all_temporary_leases(inner); - return Err(error); - } - push_audit_record(inner, gc_audit); - Ok(removed) -} - -fn actor_reset( - inner: &Arc, - config: &AuthConfig, - cold: &mut HashMap, - wheel: &mut TimingWheel, - admin_replays: &VecDeque, - action: &str, -) -> Result<(), AuthFailure> { - let new_instance_id = random_instance_id(); - inner.root_epoch.fetch_add(1, Ordering::AcqRel); - let reset_audit = audit(action, None, None); - let mut snapshot = empty_snapshot(inner, new_instance_id, admin_replays); - push_persisted_audit(&mut snapshot.audit_records, reset_audit.clone()); - let admin_key = inner.admin_key(); - let next_instance_path = config.state_dir.join("server-instance-id.next"); - if let Err(error) = atomic_write(&next_instance_path, &new_instance_id, 0o600) - .and_then(|()| write_snapshot_and_truncate_wal(config, &admin_key, &snapshot)) - .and_then(|()| { - atomic_write( - &config.state_dir.join("server-instance-id"), - &new_instance_id, - 0o600, - ) - }) - { - if !reset_already_installed(&config.state_dir, &admin_key, &new_instance_id) { - inner.safe_mode.store(true, Ordering::Release); - cancel_all_temporary_leases(inner); - return Err(error); - } - tracing::warn!( - event = "auth_state_reset_finalized_after_sync_error", - error = %error, - "server-instance-id replacement reported an error, but the live id and snapshot already match the new instance; finishing in-memory reset" - ); - } - let _ = std::fs::remove_file(&next_instance_path); - push_audit_record(inner, reset_audit); - remember_previous_root(inner); - - cancel_all_temporary_leases(inner); - { - let mut slots = inner - .slots - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - for slot in slots.iter_mut() { - slot.state = SlotState::Free; - slot.expires_at = 0; - slot.lease = Weak::new(); - } - } - *inner - .instance_id - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()) = new_instance_id; - cold.clear(); - wheel.clear(unix_seconds()); - clear_retained_high_slot_entries(inner); - inner.safe_mode.store(false, Ordering::Release); - Ok(()) -} - -fn actor_rotate_root( - inner: &Arc, - config: &AuthConfig, - cold: &mut HashMap, - wheel: &mut TimingWheel, - admin_lease: &mut Arc, - new_key: AesKeyType, -) -> Result<(), AuthFailure> { - if new_key == inner.admin_key() { - return Err(AuthFailure::new( - "administrator_key_unchanged", - "new administrator key must differ from the current key", - false, - )); - } - if !is_env_safe_admin_key(&new_key) { - return Err(AuthFailure::new( - "administrator_key_invalid", - env_safe_admin_key_error(), - false, - )); - } - let new_key_string = - String::from_utf8(new_key.to_vec()).expect("printable ASCII is valid UTF-8"); - - inner.root_epoch.fetch_add(1, Ordering::AcqRel); - let rotate_audit = audit("administrator_key_rotate", None, None); - let mut snapshot = empty_snapshot(inner, inner.instance_id(), &VecDeque::new()); - push_persisted_audit(&mut snapshot.audit_records, rotate_audit.clone()); - let next_key_path = config.state_dir.join("admin.key.next"); - if let Err(error) = write_admin_key_file(&next_key_path, &new_key_string, true) - .and_then(|()| write_snapshot_and_truncate_wal(config, &new_key, &snapshot)) - .and_then(|()| write_admin_key(&config.state_dir, &new_key_string)) - { - if !rotation_already_installed(&config.state_dir, &new_key_string) { - inner.safe_mode.store(true, Ordering::Release); - cancel_all_temporary_leases(inner); - return Err(error); - } - tracing::warn!( - event = "administrator_key_rotate_finalized_after_sync_error", - error = %error, - "admin.key replacement reported an error, but the new snapshot already decrypts with the new key; finishing in-memory rotation" - ); - } - let _ = std::fs::remove_file(&next_key_path); - push_audit_record(inner, rotate_audit); - remember_previous_root(inner); - - cancel_all_temporary_leases(inner); - let old_admin_lease = admin_lease.clone(); - { - let mut slots = inner - .slots - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - for slot in slots.iter_mut() { - slot.state = SlotState::Free; - slot.expires_at = 0; - slot.lease = Weak::new(); - } - } - cold.clear(); - wheel.clear(unix_seconds()); - clear_retained_high_slot_entries(inner); - let new_admin_lease = Arc::new(AuthLease::new(0, u64::MAX)); - *inner - .admin - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()) = AdminState { - key: new_key, - lease: Arc::downgrade(&new_admin_lease), - }; - if inner.sync_process_credential { - set_process_msg_header_key(Some(&new_key_string)).map_err(AuthFailure::internal)?; - } - inner.safe_mode.store(false, Ordering::Release); - old_admin_lease.cancel_rotated(); - *admin_lease = new_admin_lease; - Ok(()) -} - -fn actor_set_legacy_protocol( - inner: &Arc, - config: &AuthConfig, - policy: LegacyProtocolPolicy, -) -> Result<(), AuthFailure> { - ensure_store_available(inner)?; - append_mutation( - config, - inner, - StateMutation::LegacyProtocol(policy), - audit("legacy_protocol_update", None, Some(format!("{policy:?}"))), - )?; - inner - .legacy_protocol_allowed - .store(policy.is_allowed(), Ordering::Release); - Ok(()) -} - -fn actor_status(inner: &Arc) -> AuthStatus { - let slots = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let high = inner - .high_slot_entries - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let active_keys = slots - .iter() - .filter(|slot| slot.state == SlotState::Active) - .count() - + high - .iter() - .filter(|entry| entry.state == SlotState::Active) - .count(); - let expired_keys = slots - .iter() - .filter(|slot| slot.state == SlotState::Expired) - .count() - + high - .iter() - .filter(|entry| entry.state == SlotState::Expired) - .count(); - let revoked_keys = slots - .iter() - .filter(|slot| slot.state == SlotState::Revoked) - .count() - + high - .iter() - .filter(|entry| entry.state == SlotState::Revoked) - .count(); - let last_legacy_connection_at = inner.last_legacy_connection_at.load(Ordering::Acquire); - AuthStatus { - schema_version: 1, - safe_mode: inner.safe_mode.load(Ordering::Acquire), - capacity: slots.len(), - active_keys, - expired_keys, - revoked_keys, - legacy_protocol: if inner.legacy_protocol_allowed.load(Ordering::Acquire) { - LegacyProtocolPolicy::Allow - } else { - LegacyProtocolPolicy::Deny - }, - active_legacy_connections: inner.active_legacy_connections.load(Ordering::Acquire), - last_legacy_connection_at: (last_legacy_connection_at != 0) - .then_some(last_legacy_connection_at), - auth_successes: inner.auth_successes.load(Ordering::Relaxed), - auth_failures: inner.auth_failures.load(Ordering::Relaxed), - server_instance_id: hex(&inner.instance_id()), - } -} - -fn validate_admin_authority( - inner: &AuthStateInner, - authority: &Weak, -) -> Result<(), AuthFailure> { - let presented = authority.upgrade().ok_or_else(|| { - AuthFailure::new( - "administrator_key_rotated", - "administrator credential lease is no longer active", - false, - ) - })?; - if presented.cancellation.is_cancelled() { - return Err(AuthFailure::new( - "administrator_key_rotated", - "administrator credential lease has been cancelled", - false, - )); - } - let current = inner - .admin - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .lease - .upgrade() - .ok_or_else(|| { - AuthFailure::new( - "administrator_key_rotated", - "active administrator credential lease is unavailable", - false, - ) - })?; - if !Arc::ptr_eq(&presented, ¤t) { - return Err(AuthFailure::new( - "administrator_key_rotated", - "administrator request was authenticated before the latest root-key rotation", - false, - )); - } - Ok(()) -} - -fn ensure_store_available(inner: &AuthStateInner) -> Result<(), AuthFailure> { - if inner.safe_mode.load(Ordering::Acquire) { - Err(AuthFailure::new( - "temporary_key_store_unavailable", - "temporary key store is in administrator safe mode", - false, - )) - } else { - Ok(()) - } -} - -fn validate_slot_identity(slot: &SlotHot, key_id: u64) -> Result<(), AuthFailure> { - if slot.generation != key_generation(key_id) || slot.state == SlotState::Free { - Err(key_not_found(key_id)) - } else { - Ok(()) - } -} - -fn key_not_found(key_id: u64) -> AuthFailure { - AuthFailure::new( - "temporary_key_not_found", - format!("temporary key {key_id} does not exist"), - false, - ) -} - -fn key_not_renewable() -> AuthFailure { - AuthFailure::new( - "temporary_key_not_renewable", - "only an active, unexpired temporary key can be renewed", - false, - ) -} - -fn key_not_active() -> AuthFailure { - AuthFailure::new( - "temporary_key_not_active", - "temporary key is not active", - false, - ) -} - -fn high_slot_entry(high: &[PersistedEntry], key_id: u64) -> Result<&PersistedEntry, AuthFailure> { - high.iter() - .find(|entry| entry.key_id == key_id) - .ok_or_else(|| key_not_found(key_id)) -} - -fn high_slot_entry_mut( - high: &mut [PersistedEntry], - key_id: u64, -) -> Result<&mut PersistedEntry, AuthFailure> { - high.iter_mut() - .find(|entry| entry.key_id == key_id) - .ok_or_else(|| key_not_found(key_id)) -} - -fn high_slot_metadata(entry: &PersistedEntry) -> TemporaryKeyMetadata { - TemporaryKeyMetadata { - key_id: entry.key_id, - state: slot_state_name(entry.state).to_string(), - issued_at: entry.issued_at, - expires_at: entry.expires_at, - label: entry.label.clone(), - } -} - -fn expire_due_high_slots( - inner: &Arc, - cold: &mut HashMap, - tombstones: &mut VecDeque<(u64, u64)>, - now: u64, -) { - let mut high = inner - .high_slot_entries - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - for entry in high.iter_mut() { - if entry.state != SlotState::Active || entry.expires_at > now { - continue; - } - entry.state = SlotState::Expired; - let tombstoned_at = entry.expires_at; - entry.tombstoned_at = Some(tombstoned_at); - if let Some(metadata) = cold.get_mut(&entry.key_id) { - metadata.tombstoned_at = tombstoned_at; - } - push_tombstone(tombstones, tombstoned_at, entry.key_id); - tracing::info!( - event = "temporary_key_expired", - auth_stage = "expiry", - key_id = entry.key_id, - expires_at = entry.expires_at, - "high-slot temporary key expired" - ); - } -} - -fn metadata_with_credential( - inner: &Arc, - cold: &HashMap, - key_id: u64, - reveal: bool, -) -> Result { - let slots = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let credential = if reveal { - let key = derive_temporary_key(&inner.admin_key(), &inner.instance_id(), key_id)?; - encode_temporary_credential(key_id, &key) - } else { - String::new() - }; - if let Some(slot) = slots.get(key_slot(key_id) as usize) { - validate_slot_identity(slot, key_id)?; - let cold = cold.get(&key_id).ok_or_else(|| key_not_found(key_id))?; - return Ok(IssuedTemporaryKey { - metadata: TemporaryKeyMetadata { - key_id, - state: slot_state_name(slot.state).to_string(), - issued_at: cold.issued_at, - expires_at: slot.expires_at, - label: cold.label.clone(), - }, - credential, - }); - } - let high = inner - .high_slot_entries - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - Ok(IssuedTemporaryKey { - metadata: high_slot_metadata(high_slot_entry(&high, key_id)?), - credential, - }) -} - -fn slot_state_name(state: SlotState) -> &'static str { - match state { - SlotState::Free => "free", - SlotState::Active => "active", - SlotState::Expired => "expired", - SlotState::Revoked => "revoked", - } -} - -fn audit(action: &str, key_id: Option, label: Option) -> AuditRecord { - AuditRecord { - at: unix_seconds(), - action: action.to_string(), - key_id, - label, - } -} diff --git a/src/common/auth/actor/epoch.rs b/src/common/auth/actor/epoch.rs new file mode 100644 index 0000000..0c61529 --- /dev/null +++ b/src/common/auth/actor/epoch.rs @@ -0,0 +1,166 @@ +//! Root rotation, auth-state reset, and live temporary-key wipe. +use super::super::*; +use super::{audit, ensure_store_available}; + +fn wipe_temporary_keys( + inner: &AuthStateInner, + cold: &mut HashMap, + wheel: &mut TimingWheel, +) { + cancel_all_temporary_leases(inner); + let mut slots = inner + .slots + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + for slot in slots.iter_mut() { + slot.state = SlotState::Free; + slot.expires_at = 0; + slot.lease = Weak::new(); + } + cold.clear(); + wheel.clear(unix_seconds()); + clear_retained_high_slot_entries(inner); +} + +fn remember_previous_root(inner: &AuthStateInner) { + *inner + .previous_root + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(PreviousRoot { + admin_key: inner.admin_key(), + instance_id: inner.instance_id(), + }); +} + +pub(super) fn actor_reset( + inner: &Arc, + config: &AuthConfig, + cold: &mut HashMap, + wheel: &mut TimingWheel, + admin_replays: &VecDeque, + action: &str, +) -> Result<(), AuthFailure> { + let new_instance_id = random_instance_id(); + inner.root_epoch.fetch_add(1, Ordering::AcqRel); + let reset_audit = audit(action, None, None); + let mut snapshot = empty_snapshot(inner, new_instance_id, admin_replays); + push_persisted_audit(&mut snapshot.audit_records, reset_audit.clone()); + let admin_key = inner.admin_key(); + let next_instance_path = config.state_dir.join("server-instance-id.next"); + if let Err(error) = atomic_write(&next_instance_path, &new_instance_id, 0o600) + .and_then(|()| write_snapshot_and_truncate_wal(config, &admin_key, &snapshot)) + .and_then(|()| { + atomic_write( + &config.state_dir.join("server-instance-id"), + &new_instance_id, + 0o600, + ) + }) + { + if !reset_already_installed(&config.state_dir, &admin_key, &new_instance_id) { + inner.safe_mode.store(true, Ordering::Release); + cancel_all_temporary_leases(inner); + return Err(error); + } + tracing::warn!( + event = "auth_state_reset_finalized_after_sync_error", + error = %error, + "server-instance-id replacement reported an error, but the live id and snapshot already match the new instance; finishing in-memory reset" + ); + } + let _ = std::fs::remove_file(&next_instance_path); + push_audit_record(inner, reset_audit); + remember_previous_root(inner); + wipe_temporary_keys(inner, cold, wheel); + *inner + .instance_id + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = new_instance_id; + inner.safe_mode.store(false, Ordering::Release); + Ok(()) +} + +pub(super) fn actor_rotate_root( + inner: &Arc, + config: &AuthConfig, + cold: &mut HashMap, + wheel: &mut TimingWheel, + admin_lease: &mut Arc, + new_key: AesKeyType, +) -> Result<(), AuthFailure> { + if new_key == inner.admin_key() { + return Err(AuthFailure::new( + "administrator_key_unchanged", + "new administrator key must differ from the current key", + false, + )); + } + if !is_env_safe_admin_key(&new_key) { + return Err(AuthFailure::new( + "administrator_key_invalid", + env_safe_admin_key_error(), + false, + )); + } + let new_key_string = + String::from_utf8(new_key.to_vec()).expect("printable ASCII is valid UTF-8"); + + inner.root_epoch.fetch_add(1, Ordering::AcqRel); + let rotate_audit = audit("administrator_key_rotate", None, None); + let mut snapshot = empty_snapshot(inner, inner.instance_id(), &VecDeque::new()); + push_persisted_audit(&mut snapshot.audit_records, rotate_audit.clone()); + let next_key_path = config.state_dir.join("admin.key.next"); + if let Err(error) = write_admin_key_file(&next_key_path, &new_key_string, true) + .and_then(|()| write_snapshot_and_truncate_wal(config, &new_key, &snapshot)) + .and_then(|()| write_admin_key(&config.state_dir, &new_key_string)) + { + if !rotation_already_installed(&config.state_dir, &new_key_string) { + inner.safe_mode.store(true, Ordering::Release); + cancel_all_temporary_leases(inner); + return Err(error); + } + tracing::warn!( + event = "administrator_key_rotate_finalized_after_sync_error", + error = %error, + "admin.key replacement reported an error, but the new snapshot already decrypts with the new key; finishing in-memory rotation" + ); + } + let _ = std::fs::remove_file(&next_key_path); + push_audit_record(inner, rotate_audit); + remember_previous_root(inner); + wipe_temporary_keys(inner, cold, wheel); + let old_admin_lease = admin_lease.clone(); + let new_admin_lease = Arc::new(AuthLease::new(0, u64::MAX)); + *inner + .admin + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = AdminState { + key: new_key, + lease: Arc::downgrade(&new_admin_lease), + }; + if inner.sync_process_credential { + set_process_msg_header_key(Some(&new_key_string)).map_err(AuthFailure::internal)?; + } + inner.safe_mode.store(false, Ordering::Release); + old_admin_lease.cancel_rotated(); + *admin_lease = new_admin_lease; + Ok(()) +} + +pub(super) fn actor_set_legacy_protocol( + inner: &Arc, + config: &AuthConfig, + policy: LegacyProtocolPolicy, +) -> Result<(), AuthFailure> { + ensure_store_available(inner)?; + append_mutation( + config, + inner, + StateMutation::LegacyProtocol(policy), + audit("legacy_protocol_update", None, Some(format!("{policy:?}"))), + )?; + inner + .legacy_protocol_allowed + .store(policy.is_allowed(), Ordering::Release); + Ok(()) +} diff --git a/src/common/auth/actor/lifecycle.rs b/src/common/auth/actor/lifecycle.rs new file mode 100644 index 0000000..2606e6b --- /dev/null +++ b/src/common/auth/actor/lifecycle.rs @@ -0,0 +1,516 @@ +//! Issue, inspect, renew, revoke, and collect temporary keys. +use super::super::*; +use super::{ + audit, ensure_store_available, key_not_active, key_not_found, key_not_renewable, + push_tombstone, slot_state_name, validate_slot_identity, +}; + +fn validate_ttl(config: &AuthConfig, ttl: Duration) -> Result { + if ttl < MIN_TEMP_KEY_TTL { + return Err(AuthFailure::new( + "temporary_key_ttl_too_short", + format!( + "temporary key TTL must be at least {} seconds", + MIN_TEMP_KEY_TTL.as_secs() + ), + false, + )); + } + if ttl > config.max_temporary_key_ttl { + return Err(AuthFailure::new( + "temporary_key_ttl_too_long", + format!( + "temporary key TTL exceeds the configured maximum of {} seconds", + config.max_temporary_key_ttl.as_secs() + ), + false, + )); + } + Ok(unix_seconds().saturating_add(ttl.as_secs())) +} + +fn validate_label(label: Option) -> Result, AuthFailure> { + let label = label + .map(|label| label.trim().to_string()) + .filter(|label| !label.is_empty()); + if label.as_ref().is_some_and(|label| label.len() > 64) { + return Err(AuthFailure::new( + "temporary_key_label_too_long", + "temporary key label must not exceed 64 UTF-8 bytes", + false, + )); + } + Ok(label) +} + +pub(super) fn actor_issue( + inner: &Arc, + config: &AuthConfig, + cold: &mut HashMap, + wheel: &mut TimingWheel, + ttl: Duration, + label: Option, +) -> Result { + ensure_store_available(inner)?; + let expires_at = validate_ttl(config, ttl)?; + let label = validate_label(label)?; + let issued_at = unix_seconds(); + let (index, generation, key_id, entry) = { + let slots = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let Some((index, slot)) = slots + .iter() + .enumerate() + .find(|(_, slot)| slot.state == SlotState::Free && slot.generation < u32::MAX) + else { + return Err(AuthFailure::new( + "temporary_key_capacity_exhausted", + "temporary key slot table is full", + true, + )); + }; + let generation = slot.generation + 1; + let key_id = make_key_id(generation, index as u32); + ( + index, + generation, + key_id, + PersistedEntry { + key_id, + state: SlotState::Active, + issued_at, + expires_at, + label: label.clone(), + tombstoned_at: None, + }, + ) + }; + // Persist before taking the slot write lock. A fail-closed WAL error + // cancels leases via slots.read() and must not nest under slots.write(). + append_mutation( + config, + inner, + StateMutation::Issue(entry), + audit("temporary_key_issue", Some(key_id), label.clone()), + )?; + let mut slots = inner + .slots + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let slot = slots + .get_mut(index) + .ok_or_else(|| AuthFailure::internal("issued slot disappeared"))?; + let lease = Arc::new(AuthLease::new(key_id, expires_at)); + slot.generation = generation; + slot.state = SlotState::Active; + slot.expires_at = expires_at; + slot.issued_epoch = inner.root_epoch.load(Ordering::Acquire); + slot.lease = Arc::downgrade(&lease); + cold.insert( + key_id, + ColdMetadata { + issued_at, + label, + tombstoned_at: 0, + }, + ); + wheel.insert(lease); + drop(slots); + metadata_with_credential(inner, cold, key_id, true) +} + +pub(super) fn actor_list( + inner: &Arc, + cold: &HashMap, + page: u32, + page_size: u16, +) -> Result { + let page_size = page_size.clamp(1, 1000) as usize; + let start = (page as usize).saturating_mul(page_size); + let slots = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let mut all = slots + .iter() + .enumerate() + .filter_map(|(index, slot)| { + if slot.state == SlotState::Free { + return None; + } + let key_id = make_key_id(slot.generation, index as u32); + let cold = cold.get(&key_id)?; + Some(TemporaryKeyMetadata { + key_id, + state: slot_state_name(slot.state).to_string(), + issued_at: cold.issued_at, + expires_at: slot.expires_at, + label: cold.label.clone(), + }) + }) + .collect::>(); + all.extend( + inner + .high_slot_entries + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .iter() + .filter(|entry| entry.state != SlotState::Free) + .map(high_slot_metadata), + ); + all.sort_by_key(|item| std::cmp::Reverse(item.issued_at)); + let items = all.iter().skip(start).take(page_size).cloned().collect(); + let next_page = (start.saturating_add(page_size) < all.len()).then_some(page.saturating_add(1)); + Ok(KeyPage { + schema_version: 1, + items, + next_page, + }) +} + +pub(super) fn actor_show( + inner: &Arc, + config: &AuthConfig, + cold: &HashMap, + key_id: u64, + reveal: bool, +) -> Result { + let result = metadata_with_credential(inner, cold, key_id, reveal)?; + append_audit( + config, + inner, + audit( + if reveal { + "temporary_key_reveal" + } else { + "temporary_key_show" + }, + Some(key_id), + result.metadata.label.clone(), + ), + )?; + Ok(result) +} + +pub(super) fn actor_renew( + inner: &Arc, + config: &AuthConfig, + cold: &HashMap, + wheel: &mut TimingWheel, + key_id: u64, + ttl: Duration, +) -> Result { + ensure_store_available(inner)?; + let expires_at = validate_ttl(config, ttl)?; + let index = key_slot(key_id) as usize; + { + let slots = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if let Some(slot) = slots.get(index) { + validate_slot_identity(slot, key_id)?; + if slot.state != SlotState::Active || slot.expires_at <= unix_seconds() { + return Err(key_not_renewable()); + } + } else { + let high = inner + .high_slot_entries + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let entry = high_slot_entry(&high, key_id)?; + if entry.state != SlotState::Active || entry.expires_at <= unix_seconds() { + return Err(key_not_renewable()); + } + } + } + let label = cold + .get(&key_id) + .and_then(|metadata| metadata.label.clone()) + .or_else(|| { + inner + .high_slot_entries + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .iter() + .find(|entry| entry.key_id == key_id) + .and_then(|entry| entry.label.clone()) + }); + append_mutation( + config, + inner, + StateMutation::Renew { key_id, expires_at }, + audit("temporary_key_renew", Some(key_id), label), + )?; + let mut slots = inner + .slots + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if let Some(slot) = slots.get_mut(index) { + validate_slot_identity(slot, key_id)?; + if slot.state != SlotState::Active { + return Err(AuthFailure::new( + "temporary_key_inactive", + "temporary key lease is no longer active", + true, + )); + } + slot.expires_at = expires_at; + let lease = match slot.lease.upgrade() { + Some(lease) if !lease.cancellation_token().is_cancelled() => { + lease.expires_at.store(expires_at, Ordering::Release); + lease.wheel_version.fetch_add(1, Ordering::AcqRel); + lease + } + _ => { + let lease = Arc::new(AuthLease::new(key_id, expires_at)); + slot.lease = Arc::downgrade(&lease); + lease + } + }; + wheel.release(key_id); + wheel.insert(lease); + drop(slots); + return metadata_with_credential(inner, cold, key_id, true); + } + drop(slots); + { + let mut high = inner + .high_slot_entries + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let entry = high_slot_entry_mut(&mut high, key_id)?; + if entry.state != SlotState::Active { + return Err(AuthFailure::new( + "temporary_key_inactive", + "temporary key lease is no longer active", + true, + )); + } + entry.expires_at = expires_at; + entry.tombstoned_at = None; + } + metadata_with_credential(inner, cold, key_id, true) +} + +pub(super) fn actor_revoke( + inner: &Arc, + config: &AuthConfig, + cold: &mut HashMap, + tombstones: &mut VecDeque<(u64, u64)>, + key_id: u64, +) -> Result { + ensure_store_available(inner)?; + let now = unix_seconds(); + let index = key_slot(key_id) as usize; + let (label, issued_at, expires_at) = { + let slots = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if let Some(slot) = slots.get(index) { + validate_slot_identity(slot, key_id)?; + if slot.state != SlotState::Active { + return Err(key_not_active()); + } + let metadata = cold.get(&key_id).ok_or_else(|| key_not_found(key_id))?; + (metadata.label.clone(), metadata.issued_at, slot.expires_at) + } else { + let high = inner + .high_slot_entries + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let entry = high_slot_entry(&high, key_id)?; + if entry.state != SlotState::Active { + return Err(key_not_active()); + } + (entry.label.clone(), entry.issued_at, entry.expires_at) + } + }; + append_mutation( + config, + inner, + StateMutation::Revoke { key_id, at: now }, + audit("temporary_key_revoke", Some(key_id), label.clone()), + )?; + let mut slots = inner + .slots + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if let Some(slot) = slots.get_mut(index) { + validate_slot_identity(slot, key_id)?; + slot.state = SlotState::Revoked; + if let Some(lease) = slot.lease.upgrade() { + lease.cancel_revoked(); + } + let cold_metadata = cold.get_mut(&key_id).ok_or_else(|| key_not_found(key_id))?; + cold_metadata.tombstoned_at = now; + push_tombstone(tombstones, now, key_id); + return Ok(TemporaryKeyMetadata { + key_id, + state: slot_state_name(slot.state).to_string(), + issued_at: cold_metadata.issued_at, + expires_at: slot.expires_at, + label: cold_metadata.label.clone(), + }); + } + drop(slots); + let mut high = inner + .high_slot_entries + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let entry = high_slot_entry_mut(&mut high, key_id)?; + if entry.state != SlotState::Active { + return Err(key_not_active()); + } + entry.state = SlotState::Revoked; + entry.tombstoned_at = Some(now); + if let Some(metadata) = cold.get_mut(&key_id) { + metadata.tombstoned_at = now; + } + push_tombstone(tombstones, now, key_id); + Ok(TemporaryKeyMetadata { + key_id, + state: slot_state_name(entry.state).to_string(), + issued_at, + expires_at, + label, + }) +} + +pub(super) fn actor_gc( + inner: &Arc, + config: &AuthConfig, + cold: &mut HashMap, + wheel: &mut TimingWheel, + tombstones: &mut VecDeque<(u64, u64)>, + admin_replays: &VecDeque, +) -> Result { + ensure_store_available(inner)?; + let now = unix_seconds(); + let mut removed = 0_u64; + let mut slots = inner + .slots + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + for (index, slot) in slots.iter_mut().enumerate() { + if matches!(slot.state, SlotState::Expired | SlotState::Revoked) + || (slot.state == SlotState::Active && slot.expires_at <= now) + { + let key_id = make_key_id(slot.generation, index as u32); + if let Some(lease) = slot.lease.upgrade() { + if slot.state == SlotState::Revoked { + lease.cancel_revoked(); + } else { + lease.cancel_expired(); + } + } + slot.state = SlotState::Free; + slot.expires_at = 0; + slot.lease = Weak::new(); + cold.remove(&key_id); + wheel.release(key_id); + removed = removed.saturating_add(1); + } + } + drop(slots); + { + let mut high = inner + .high_slot_entries + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + high.retain(|entry| { + let keep = match entry.state { + SlotState::Active if entry.expires_at > now => true, + SlotState::Active | SlotState::Expired | SlotState::Revoked | SlotState::Free => { + false + } + }; + if !keep { + cold.remove(&entry.key_id); + wheel.release(entry.key_id); + removed = removed.saturating_add(1); + } + keep + }); + } + tombstones.clear(); + let gc_audit = audit("temporary_key_gc", None, Some(format!("removed={removed}"))); + let mut snapshot = build_snapshot(inner, cold, admin_replays); + push_persisted_audit(&mut snapshot.audit_records, gc_audit.clone()); + let admin_key = inner.admin_key(); + if let Err(error) = write_snapshot_and_truncate_wal(config, &admin_key, &snapshot) { + inner.safe_mode.store(true, Ordering::Release); + cancel_all_temporary_leases(inner); + return Err(error); + } + push_audit_record(inner, gc_audit); + Ok(removed) +} + +fn high_slot_entry(high: &[PersistedEntry], key_id: u64) -> Result<&PersistedEntry, AuthFailure> { + high.iter() + .find(|entry| entry.key_id == key_id) + .ok_or_else(|| key_not_found(key_id)) +} + +fn high_slot_entry_mut( + high: &mut [PersistedEntry], + key_id: u64, +) -> Result<&mut PersistedEntry, AuthFailure> { + high.iter_mut() + .find(|entry| entry.key_id == key_id) + .ok_or_else(|| key_not_found(key_id)) +} + +fn high_slot_metadata(entry: &PersistedEntry) -> TemporaryKeyMetadata { + TemporaryKeyMetadata { + key_id: entry.key_id, + state: slot_state_name(entry.state).to_string(), + issued_at: entry.issued_at, + expires_at: entry.expires_at, + label: entry.label.clone(), + } +} + +fn metadata_with_credential( + inner: &Arc, + cold: &HashMap, + key_id: u64, + reveal: bool, +) -> Result { + let slots = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let credential = if reveal { + let key = derive_temporary_key(&inner.admin_key(), &inner.instance_id(), key_id)?; + encode_temporary_credential(key_id, &key) + } else { + String::new() + }; + if let Some(slot) = slots.get(key_slot(key_id) as usize) { + validate_slot_identity(slot, key_id)?; + let cold = cold.get(&key_id).ok_or_else(|| key_not_found(key_id))?; + return Ok(IssuedTemporaryKey { + metadata: TemporaryKeyMetadata { + key_id, + state: slot_state_name(slot.state).to_string(), + issued_at: cold.issued_at, + expires_at: slot.expires_at, + label: cold.label.clone(), + }, + credential, + }); + } + let high = inner + .high_slot_entries + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + Ok(IssuedTemporaryKey { + metadata: high_slot_metadata(high_slot_entry(&high, key_id)?), + credential, + }) +} diff --git a/src/common/auth/actor/mod.rs b/src/common/auth/actor/mod.rs new file mode 100644 index 0000000..ff01083 --- /dev/null +++ b/src/common/auth/actor/mod.rs @@ -0,0 +1,585 @@ +//! Serialized owner of mutable authentication lifecycle state. +//! +//! ```text +//! authenticated admin command +//! | +//! v +//! validate current admin lease +//! | +//! v +//! append encrypted WAL -> mutate slots / leases / timing wheel +//! | +//! +-> periodic snapshot + bounded replay/audit retention +//! ``` +//! +//! Keeping authorization revalidation and mutations in one actor prevents a request +//! authenticated before root rotation from executing against the new administrator +//! state. The actor is also the sole strong owner of temporary-key leases. + +use super::*; + +mod epoch; +mod lifecycle; +use epoch::*; +use lifecycle::*; + +pub(super) struct AuthActorState { + cold: HashMap, + wheel: TimingWheel, + admin_replays: HashSet<[u8; 32]>, + admin_replay_order: VecDeque, +} + +impl AuthActorState { + pub(super) fn new( + cold: HashMap, + wheel: TimingWheel, + admin_replays: HashSet<[u8; 32]>, + admin_replay_order: VecDeque, + ) -> Self { + Self { + cold, + wheel, + admin_replays, + admin_replay_order, + } + } +} + +pub(super) async fn run_auth_actor( + inner: Arc, + mut admin_lease: Arc, + mut command_rx: mpsc::Receiver, + config: AuthConfig, + state: AuthActorState, + _state_lock: Arc, +) { + let AuthActorState { + mut cold, + mut wheel, + mut admin_replays, + mut admin_replay_order, + } = state; + let now = unix_seconds(); + let mut tombstones = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .iter() + .enumerate() + .filter_map(|(index, slot)| { + if !matches!(slot.state, SlotState::Expired | SlotState::Revoked) { + return None; + } + let key_id = make_key_id(slot.generation, index as u32); + let tombstoned_at = cold + .get(&key_id) + .map(|metadata| metadata.tombstoned_at) + .unwrap_or(now); + Some(( + tombstoned_at.saturating_add(TOMBSTONE_RETENTION.as_secs()), + key_id, + )) + }) + .collect::>(); + tombstones.extend( + inner + .high_slot_entries + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .iter() + .filter_map(|entry| { + let expired_active = entry.state == SlotState::Active && entry.expires_at <= now; + if !matches!(entry.state, SlotState::Expired | SlotState::Revoked) + && !expired_active + { + return None; + } + let tombstoned_at = entry + .tombstoned_at + .or_else(|| { + cold.get(&entry.key_id) + .map(|metadata| metadata.tombstoned_at) + }) + .unwrap_or(entry.expires_at.max(now)); + Some(( + tombstoned_at.saturating_add(TOMBSTONE_RETENTION.as_secs()), + entry.key_id, + )) + }), + ); + tombstones.sort_unstable_by_key(|(cleanup_at, _)| *cleanup_at); + let mut tombstones = VecDeque::from(tombstones); + let mut last_snapshot_at = unix_seconds(); + let mut tick = tokio::time::interval(Duration::from_secs(1)); + tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + tokio::select! { + _ = tick.tick() => { + let now = unix_seconds(); + for lease in wheel.advance(now) { + let key_id = lease.key_id(); + let version = lease.wheel_version.load(Ordering::Acquire); + if lease.expires_at() > now { + wheel.insert_with_version(lease, version); + continue; + } + let mut slots = inner.slots.write().unwrap_or_else(|poisoned| poisoned.into_inner()); + if let Some(slot) = slots.get_mut(key_slot(key_id) as usize) { + if slot.generation == key_generation(key_id) && slot.state == SlotState::Active { + slot.state = SlotState::Expired; + lease.cancel_expired(); + let tombstoned_at = slot.expires_at; + if let Some(metadata) = cold.get_mut(&key_id) { + metadata.tombstoned_at = tombstoned_at; + } + push_tombstone(&mut tombstones, tombstoned_at, key_id); + tracing::info!( + event = "temporary_key_expired", + auth_stage = "expiry", + key_id, + expires_at = lease.expires_at(), + "temporary key expired and active work was cancelled" + ); + } + } else { + lease.cancel_expired(); + } + } + expire_due_high_slots(&inner, &mut cold, &mut tombstones, now); + let mut due_high = Vec::new(); + while let Some((cleanup_at, key_id)) = tombstones.front().copied() { + if cleanup_at > now { + break; + } + tombstones.pop_front(); + let mut slots = inner.slots.write().unwrap_or_else(|poisoned| poisoned.into_inner()); + if let Some(slot) = slots.get_mut(key_slot(key_id) as usize) { + if slot.generation == key_generation(key_id) && matches!(slot.state, SlotState::Expired | SlotState::Revoked) { + slot.state = SlotState::Free; + slot.expires_at = 0; + slot.lease = Weak::new(); + cold.remove(&key_id); + wheel.release(key_id); + } + } else { + due_high.push(key_id); + } + } + if !due_high.is_empty() { + let due = due_high.iter().copied().collect::>(); + inner + .high_slot_entries + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .retain(|entry| !due.contains(&entry.key_id)); + for key_id in due_high { + cold.remove(&key_id); + wheel.release(key_id); + } + } + prune_expired_admin_replays( + now, + &mut admin_replays, + &mut admin_replay_order, + ); + // WHY: A failed load starts safe mode with empty in-memory + // generations. Compacting that reconstruction would replace the + // damaged snapshot, truncate the WAL, and let the next start + // exit safe mode without rotating the instance id. + if compaction_is_allowed(inner.safe_mode.load(Ordering::Acquire)) + && now.saturating_sub(last_snapshot_at) + >= SNAPSHOT_COMPACTION_INTERVAL.as_secs() + { + let snapshot = build_snapshot(&inner, &cold, &admin_replay_order); + if let Err(error) = write_snapshot_and_truncate_wal( + &config, + &inner.admin_key(), + &snapshot, + ) { + inner.safe_mode.store(true, Ordering::Release); + cancel_all_temporary_leases(&inner); + tracing::error!( + event = "auth_state_safe_mode", + auth_stage = "snapshot_compaction", + reason = %error.code, + error = %error, + "authentication state compaction failed closed" + ); + } else { + last_snapshot_at = now; + } + } + } + command = command_rx.recv() => { + let Some(command) = command else { + admin_lease.cancel_rotated(); + cancel_all_temporary_leases(&inner); + break; + }; + match command { + AuthCommand::ClaimAdminMutation { + authority, + fingerprint, + client_timestamp, + response, + } => { + let result = validate_admin_authority(&inner, &authority).and_then(|()| { + actor_claim_admin_mutation( + &inner, + &config, + &mut admin_replays, + &mut admin_replay_order, + fingerprint, + client_timestamp, + ) + }); + let _ = response.send(result); + } + AuthCommand::Issue { authority, ttl, label, response } => { + let result = validate_admin_authority(&inner, &authority) + .and_then(|()| actor_issue(&inner, &config, &mut cold, &mut wheel, ttl, label)); + let _ = response.send(result); + } + AuthCommand::List { authority, page, page_size, response } => { + let result = validate_admin_authority(&inner, &authority) + .and_then(|()| actor_list(&inner, &cold, page, page_size)); + let _ = response.send(result); + } + AuthCommand::Show { authority, key_id, reveal, response } => { + let result = validate_admin_authority(&inner, &authority) + .and_then(|()| actor_show(&inner, &config, &cold, key_id, reveal)); + let _ = response.send(result); + } + AuthCommand::Renew { authority, key_id, ttl, response } => { + let result = validate_admin_authority(&inner, &authority) + .and_then(|()| actor_renew(&inner, &config, &cold, &mut wheel, key_id, ttl)); + let _ = response.send(result); + } + AuthCommand::Revoke { authority, key_id, response } => { + let result = validate_admin_authority(&inner, &authority) + .and_then(|()| actor_revoke(&inner, &config, &mut cold, &mut tombstones, key_id)); + let _ = response.send(result); + } + AuthCommand::Gc { authority, response } => { + let result = validate_admin_authority(&inner, &authority) + .and_then(|()| actor_gc( + &inner, + &config, + &mut cold, + &mut wheel, + &mut tombstones, + &admin_replay_order, + )); + let _ = response.send(result); + } + AuthCommand::Reset { authority, response } => { + let result = validate_admin_authority(&inner, &authority) + .and_then(|()| actor_reset( + &inner, + &config, + &mut cold, + &mut wheel, + &admin_replay_order, + "auth_state_reset", + )); + let _ = response.send(result); + } + AuthCommand::RotateRoot { authority, new_key, response } => { + let result = validate_admin_authority(&inner, &authority) + .and_then(|()| actor_rotate_root(&inner, &config, &mut cold, &mut wheel, &mut admin_lease, new_key)); + if result.is_ok() { + admin_replays.clear(); + admin_replay_order.clear(); + } + let _ = response.send(result); + } + AuthCommand::SetLegacyProtocol { authority, policy, response } => { + let result = validate_admin_authority(&inner, &authority) + .and_then(|()| actor_set_legacy_protocol(&inner, &config, policy)); + let _ = response.send(result); + } + AuthCommand::Status { authority, response } => { + let result = validate_admin_authority(&inner, &authority) + .map(|()| actor_status(&inner)); + let _ = response.send(result); + } + AuthCommand::Audit { authority, action, key_id, detail, response } => { + let result = validate_admin_authority(&inner, &authority).and_then(|()| { + append_audit( + &config, + &inner, + audit(&action, key_id, detail), + ) + }); + let _ = response.send(result); + } + AuthCommand::Shutdown { response } => { + admin_lease.cancel_rotated(); + cancel_all_temporary_leases(&inner); + let _ = response.send(()); + break; + } + } + } + } + } +} + +fn actor_claim_admin_mutation( + inner: &AuthStateInner, + config: &AuthConfig, + admin_replays: &mut HashSet<[u8; 32]>, + admin_replay_order: &mut VecDeque, + fingerprint: [u8; 32], + client_timestamp: u64, +) -> Result<(), AuthFailure> { + let now = unix_seconds(); + prune_expired_admin_replays(now, admin_replays, admin_replay_order); + if admin_replays.contains(&fingerprint) { + return Err(AuthFailure::new( + "admin_request_replayed", + "administrator mutation was already admitted", + false, + )); + } + if admin_replays.len() >= ADMIN_REPLAY_CAPACITY { + return Err(AuthFailure::new( + "admin_replay_capacity_exhausted", + "administrator mutation replay window is full; retry after older claims expire", + true, + )); + } + if now.abs_diff(client_timestamp) > ADMIN_REPLAY_RETENTION.as_secs() / 2 { + return Err(AuthFailure::new( + "admin_request_timestamp_invalid", + "administrator mutation timestamp is outside the accepted window", + false, + )); + } + let record = AdminReplayRecord { + fingerprint, + client_timestamp, + accepted_at: now, + }; + fail_closed_on_uncertain_wal( + inner, + append_wal( + config, + &inner.admin_key(), + &WalRecord::AdminReplay(record.clone()), + ), + )?; + admin_replays.insert(fingerprint); + admin_replay_order.push_back(record); + Ok(()) +} + +pub(super) fn prune_expired_admin_replays( + now: u64, + admin_replays: &mut HashSet<[u8; 32]>, + admin_replay_order: &mut VecDeque, +) { + admin_replay_order.retain(|record| { + let keep = record.within_retention(now); + if !keep { + admin_replays.remove(&record.fingerprint); + } + keep + }); +} + +fn push_tombstone(tombstones: &mut VecDeque<(u64, u64)>, tombstoned_at: u64, key_id: u64) { + let cleanup_at = tombstoned_at.saturating_add(TOMBSTONE_RETENTION.as_secs()); + let index = tombstones.partition_point(|(current, _)| *current <= cleanup_at); + tombstones.insert(index, (cleanup_at, key_id)); +} + +fn actor_status(inner: &Arc) -> AuthStatus { + let slots = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let high = inner + .high_slot_entries + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let active_keys = slots + .iter() + .filter(|slot| slot.state == SlotState::Active) + .count() + + high + .iter() + .filter(|entry| entry.state == SlotState::Active) + .count(); + let expired_keys = slots + .iter() + .filter(|slot| slot.state == SlotState::Expired) + .count() + + high + .iter() + .filter(|entry| entry.state == SlotState::Expired) + .count(); + let revoked_keys = slots + .iter() + .filter(|slot| slot.state == SlotState::Revoked) + .count() + + high + .iter() + .filter(|entry| entry.state == SlotState::Revoked) + .count(); + let last_legacy_connection_at = inner.last_legacy_connection_at.load(Ordering::Acquire); + AuthStatus { + schema_version: 1, + safe_mode: inner.safe_mode.load(Ordering::Acquire), + capacity: slots.len(), + active_keys, + expired_keys, + revoked_keys, + legacy_protocol: if inner.legacy_protocol_allowed.load(Ordering::Acquire) { + LegacyProtocolPolicy::Allow + } else { + LegacyProtocolPolicy::Deny + }, + active_legacy_connections: inner.active_legacy_connections.load(Ordering::Acquire), + last_legacy_connection_at: (last_legacy_connection_at != 0) + .then_some(last_legacy_connection_at), + auth_successes: inner.auth_successes.load(Ordering::Relaxed), + auth_failures: inner.auth_failures.load(Ordering::Relaxed), + server_instance_id: hex(&inner.instance_id()), + } +} + +fn validate_admin_authority( + inner: &AuthStateInner, + authority: &Weak, +) -> Result<(), AuthFailure> { + let presented = authority.upgrade().ok_or_else(|| { + AuthFailure::new( + "administrator_key_rotated", + "administrator credential lease is no longer active", + false, + ) + })?; + if presented.cancellation.is_cancelled() { + return Err(AuthFailure::new( + "administrator_key_rotated", + "administrator credential lease has been cancelled", + false, + )); + } + let current = inner + .admin + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .lease + .upgrade() + .ok_or_else(|| { + AuthFailure::new( + "administrator_key_rotated", + "active administrator credential lease is unavailable", + false, + ) + })?; + if !Arc::ptr_eq(&presented, ¤t) { + return Err(AuthFailure::new( + "administrator_key_rotated", + "administrator request was authenticated before the latest root-key rotation", + false, + )); + } + Ok(()) +} + +fn ensure_store_available(inner: &AuthStateInner) -> Result<(), AuthFailure> { + if inner.safe_mode.load(Ordering::Acquire) { + Err(AuthFailure::new( + "temporary_key_store_unavailable", + "temporary key store is in administrator safe mode", + false, + )) + } else { + Ok(()) + } +} + +fn validate_slot_identity(slot: &SlotHot, key_id: u64) -> Result<(), AuthFailure> { + if slot.generation != key_generation(key_id) || slot.state == SlotState::Free { + Err(key_not_found(key_id)) + } else { + Ok(()) + } +} + +fn key_not_found(key_id: u64) -> AuthFailure { + AuthFailure::new( + "temporary_key_not_found", + format!("temporary key {key_id} does not exist"), + false, + ) +} + +fn key_not_renewable() -> AuthFailure { + AuthFailure::new( + "temporary_key_not_renewable", + "only an active, unexpired temporary key can be renewed", + false, + ) +} + +fn key_not_active() -> AuthFailure { + AuthFailure::new( + "temporary_key_not_active", + "temporary key is not active", + false, + ) +} + +fn expire_due_high_slots( + inner: &Arc, + cold: &mut HashMap, + tombstones: &mut VecDeque<(u64, u64)>, + now: u64, +) { + let mut high = inner + .high_slot_entries + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + for entry in high.iter_mut() { + if entry.state != SlotState::Active || entry.expires_at > now { + continue; + } + entry.state = SlotState::Expired; + let tombstoned_at = entry.expires_at; + entry.tombstoned_at = Some(tombstoned_at); + if let Some(metadata) = cold.get_mut(&entry.key_id) { + metadata.tombstoned_at = tombstoned_at; + } + push_tombstone(tombstones, tombstoned_at, entry.key_id); + tracing::info!( + event = "temporary_key_expired", + auth_stage = "expiry", + key_id = entry.key_id, + expires_at = entry.expires_at, + "high-slot temporary key expired" + ); + } +} + +fn slot_state_name(state: SlotState) -> &'static str { + match state { + SlotState::Free => "free", + SlotState::Active => "active", + SlotState::Expired => "expired", + SlotState::Revoked => "revoked", + } +} + +fn audit(action: &str, key_id: Option, label: Option) -> AuditRecord { + AuditRecord { + at: unix_seconds(), + action: action.to_string(), + key_id, + label, + } +} diff --git a/src/common/auth/config.rs b/src/common/auth/config.rs new file mode 100644 index 0000000..32e0bba --- /dev/null +++ b/src/common/auth/config.rs @@ -0,0 +1,188 @@ +//! Authentication configuration and platform state-directory defaults. +use super::*; + +pub fn default_auth_state_dir() -> PathBuf { + std::env::var_os("PB_MAPPER_AUTH_STATE_DIR") + .map(PathBuf::from) + .unwrap_or_else(platform_default_auth_state_dir) +} + +/// Linux systemd/Docker keep `/var/lib/pb-mapper/auth` when that path is usable +/// (root, or an already-writable service directory). Unprivileged Linux, +/// macOS, and Windows binaries need an application data directory instead. +pub(crate) fn platform_default_auth_state_dir() -> PathBuf { + #[cfg(windows)] + { + let base = std::env::var_os("LOCALAPPDATA") + .or_else(|| std::env::var_os("APPDATA")) + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from(r"C:\ProgramData")); + base.join("pb-mapper").join("auth") + } + #[cfg(target_os = "macos")] + { + match std::env::var_os("HOME") { + Some(home) => PathBuf::from(home) + .join("Library") + .join("Application Support") + .join("pb-mapper") + .join("auth"), + None => PathBuf::from("/Library/Application Support/pb-mapper/auth"), + } + } + #[cfg(not(any(windows, target_os = "macos")))] + { + linux_default_auth_state_dir( + unix_effective_uid(), + linux_system_auth_dir_usable(), + std::env::var_os("XDG_DATA_HOME").as_deref(), + std::env::var_os("HOME").as_deref(), + ) + } +} + +pub(crate) fn linux_default_auth_state_dir( + euid: u32, + system_dir_usable: bool, + xdg_data_home: Option<&std::ffi::OsStr>, + home: Option<&std::ffi::OsStr>, +) -> PathBuf { + if euid == 0 || system_dir_usable { + return PathBuf::from(DEFAULT_AUTH_STATE_DIR); + } + if let Some(xdg) = xdg_data_home { + if !xdg.is_empty() { + return PathBuf::from(xdg).join("pb-mapper").join("auth"); + } + } + if let Some(home) = home { + if !home.is_empty() { + return PathBuf::from(home) + .join(".local") + .join("share") + .join("pb-mapper") + .join("auth"); + } + } + PathBuf::from(DEFAULT_AUTH_STATE_DIR) +} + +#[cfg(not(any(windows, target_os = "macos")))] +pub(super) fn unix_effective_uid() -> u32 { + extern "C" { + fn geteuid() -> u32; + } + unsafe { geteuid() } +} + +#[cfg(not(any(windows, target_os = "macos")))] +pub(super) fn linux_system_auth_dir_usable() -> bool { + let path = Path::new(DEFAULT_AUTH_STATE_DIR); + path.is_dir() && unix_path_is_writable(path) +} + +#[cfg(not(any(windows, target_os = "macos")))] +fn unix_path_is_writable(path: &Path) -> bool { + use std::os::unix::ffi::OsStrExt; + let Ok(c_path) = std::ffi::CString::new(path.as_os_str().as_bytes()) else { + return false; + }; + extern "C" { + fn access(pathname: *const std::os::raw::c_char, mode: i32) -> i32; + } + const W_OK: i32 = 2; + unsafe { access(c_path.as_ptr(), W_OK) == 0 } +} + +impl Default for AuthConfig { + fn default() -> Self { + Self { + state_dir: default_auth_state_dir(), + max_temporary_keys: env_usize( + "PB_MAPPER_AUTH_MAX_TEMP_KEYS", + DEFAULT_TEMP_KEY_CAPACITY, + 1, + MAX_TEMP_KEY_CAPACITY, + ), + max_temporary_key_ttl: Duration::from_secs(env_u64( + "PB_MAPPER_AUTH_MAX_TEMP_TTL_SECS", + DEFAULT_MAX_TEMP_KEY_TTL.as_secs(), + MIN_TEMP_KEY_TTL.as_secs(), + MAX_TEMP_KEY_TTL.as_secs(), + )), + legacy_protocol: legacy_protocol_from_env(), + } + } +} + +fn legacy_protocol_from_env() -> LegacyProtocolPolicy { + match std::env::var("PB_MAPPER_LEGACY_PROTOCOL") { + Err(std::env::VarError::NotPresent) => LegacyProtocolPolicy::Allow, + Err(std::env::VarError::NotUnicode(_)) => { + tracing::error!( + event = "legacy_protocol_config_invalid", + "PB_MAPPER_LEGACY_PROTOCOL is not UTF-8; denying legacy framing" + ); + LegacyProtocolPolicy::Deny + } + Ok(value) => parse_legacy_protocol_policy(&value).unwrap_or_else(|| { + tracing::error!( + event = "legacy_protocol_config_invalid", + value, + "PB_MAPPER_LEGACY_PROTOCOL must be `allow` or `deny`; denying legacy framing" + ); + LegacyProtocolPolicy::Deny + }), + } +} + +pub(super) fn parse_legacy_protocol_policy(value: &str) -> Option { + match value.trim().to_ascii_lowercase().as_str() { + "allow" => Some(LegacyProtocolPolicy::Allow), + "deny" => Some(LegacyProtocolPolicy::Deny), + _ => None, + } +} + +fn env_usize(name: &str, default: usize, min: usize, max: usize) -> usize { + env_bounded(name, default, min, max) +} + +fn env_u64(name: &str, default: u64, min: u64, max: u64) -> u64 { + env_bounded(name, default, min, max) +} + +fn env_bounded(name: &str, default: T, min: T, max: T) -> T +where + T: std::str::FromStr + PartialOrd + Copy + fmt::Display, +{ + match std::env::var(name) { + Err(std::env::VarError::NotPresent) => default, + Ok(raw) => match raw.parse::() { + Ok(value) if value >= min && value <= max => value, + _ => { + tracing::warn!( + event = "auth_config_value_invalid", + variable = name, + value = raw, + min = %min, + max = %max, + fallback = %default, + "invalid authentication configuration value; using the default" + ); + default + } + }, + Err(std::env::VarError::NotUnicode(_)) => { + tracing::warn!( + event = "auth_config_value_invalid", + variable = name, + min = %min, + max = %max, + fallback = %default, + "authentication configuration value is not UTF-8; using the default" + ); + default + } + } +} diff --git a/src/common/auth/keys.rs b/src/common/auth/keys.rs new file mode 100644 index 0000000..8cf718f --- /dev/null +++ b/src/common/auth/keys.rs @@ -0,0 +1,243 @@ +//! Administrator credential load, recovery, and temporary-key derivation. +use super::*; + +fn read_admin_key(path: &Path) -> Result, AuthFailure> { + if !path.exists() { + return Ok(None); + } + #[cfg(unix)] + { + let metadata = std::fs::metadata(path).map_err(|error| { + AuthFailure::new( + "administrator_key_required", + format!( + "administrator key file `{}` metadata could not be read: {error}", + path.display() + ), + false, + ) + })?; + if metadata.permissions().mode() & 0o077 != 0 { + std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600)).map_err( + |error| { + AuthFailure::new( + "administrator_key_required", + format!( + "administrator key file `{}` permissions could not be secured: {error}", + path.display() + ), + false, + ) + }, + )?; + tracing::warn!( + event = "administrator_key_permissions_repaired", + path = %path.display(), + "restricted administrator key file permissions to 0600" + ); + } + } + std::fs::read_to_string(path).map(Some).map_err(|error| { + AuthFailure::new( + "administrator_key_required", + format!( + "administrator key file `{}` could not be read: {error}", + path.display() + ), + false, + ) + }) +} + +fn persist_recovery_admin_key( + state_dir: &Path, + key: &str, + mismatch_message: &'static str, +) -> Result<(), AuthFailure> { + if encrypted_auth_state_exists(state_dir) && !key_matches_existing_state(Some(state_dir), key) { + return Err(AuthFailure::new( + "administrator_key_invalid", + mismatch_message, + false, + )); + } + write_admin_key(state_dir, key) +} + +fn validate_admin_credential(raw: &str) -> Result { + let credential = parse_credential(raw.trim()) + .map_err(|error| AuthFailure::new("administrator_key_invalid", error, false))?; + if !credential.is_admin() { + return Err(AuthFailure::new( + "administrator_key_required", + "the server key file contains a temporary credential", + false, + )); + } + Ok(credential) +} + +pub(super) fn recover_admin_key_after_rotation( + state_dir: &Path, + current: &str, +) -> Result { + let snapshot_path = auth_snapshot_path(state_dir); + if !snapshot_path.exists() { + return Ok(current.to_string()); + } + let bytes = std::fs::read(&snapshot_path).map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to read `{}`: {error}", snapshot_path.display()), + false, + ) + })?; + if let Ok(Credential::Admin(current_key)) = parse_credential(current.trim()) { + if open_blob(¤t_key, &bytes).is_ok() { + return Ok(current.to_string()); + } + } + let Some(next) = read_admin_key(&state_dir.join("admin.key.next"))? else { + return Ok(current.to_string()); + }; + let Ok(Credential::Admin(next_key)) = parse_credential(next.trim()) else { + return Ok(current.to_string()); + }; + if open_blob(&next_key, &bytes).is_err() { + return Ok(current.to_string()); + } + // The rotation snapshot is complete under the staged key. Leftover WAL + // records are still encrypted with the previous key. + truncate_auth_wal(state_dir)?; + write_admin_key(state_dir, next.trim())?; + let _ = std::fs::remove_file(state_dir.join("admin.key.next")); + Ok(next) +} + +pub(super) fn load_server_admin_credential(state_dir: &Path) -> Result { + let path = state_dir.join("admin.key"); + let raw = if let Some(raw) = read_admin_key(&path)? { + raw + } else if std::env::var_os(ENV_MSG_HEADER_KEY).is_some() { + let credential = get_process_credential() + .map_err(|error| AuthFailure::new("administrator_key_invalid", error, false))?; + let Credential::Admin(key) = credential else { + return Err(AuthFailure::new( + "administrator_key_required", + "the relay server cannot start with a temporary credential", + false, + )); + }; + let key = String::from_utf8(key.to_vec()).map_err(|_| { + AuthFailure::new( + "administrator_key_invalid", + "the relay administrator key must be printable UTF-8 so it can be persisted", + false, + ) + })?; + persist_recovery_admin_key( + state_dir, + &key, + "MSG_HEADER_KEY does not decrypt the existing authentication state; refusing to write admin.key", + )?; + key + } else if Path::new(MACHINE_MSG_HEADER_KEY_PATH).is_file() { + let key = std::fs::read_to_string(MACHINE_MSG_HEADER_KEY_PATH).map_err(|error| { + AuthFailure::new( + "administrator_key_required", + format!( + "legacy administrator key file `{MACHINE_MSG_HEADER_KEY_PATH}` could not be read: {error}" + ), + false, + ) + })?; + validate_admin_credential(&key)?; + persist_recovery_admin_key( + state_dir, + key.trim(), + "legacy administrator key does not decrypt the existing authentication state; refusing to write admin.key", + )?; + tracing::warn!( + event = "administrator_key_migrated", + source = MACHINE_MSG_HEADER_KEY_PATH, + destination = %path.display(), + "migrated the legacy administrator key into the v0.4 authentication state directory" + ); + key + } else { + let key = initialize_admin_key(&path, false)?; + tracing::warn!( + event = "administrator_key_initialized", + path = %path.display(), + "no administrator credential was configured; generated a random key file" + ); + key + }; + let raw = recover_admin_key_after_rotation(state_dir, &raw)?; + let credential = validate_admin_credential(&raw)?; + set_process_msg_header_key(Some(raw.trim())).map_err(AuthFailure::internal)?; + Ok(credential) +} + +/// Load or create an app-local relay root without reading or mutating the process credential. +/// +/// The Flutter process uses its configured process credential for the remote relay, while its +/// optional embedded relay owns an independent administrator key under the app data directory. +pub(super) fn load_isolated_server_admin_credential( + state_dir: &Path, +) -> Result { + let path = state_dir.join("admin.key"); + let raw = match read_admin_key(&path)? { + Some(raw) => raw, + None => { + let key = initialize_admin_key(&path, false)?; + tracing::warn!( + event = "isolated_administrator_key_initialized", + path = %path.display(), + "generated an administrator key for an embedded relay" + ); + key + } + }; + let raw = recover_admin_key_after_rotation(state_dir, &raw)?; + validate_admin_credential(&raw) +} + +pub fn make_key_id(generation: u32, slot: u32) -> u64 { + (u64::from(generation) << 32) | u64::from(slot) +} + +pub fn key_generation(key_id: u64) -> u32 { + (key_id >> 32) as u32 +} + +pub fn key_slot(key_id: u64) -> u32 { + key_id as u32 +} + +pub fn derive_temporary_key( + admin_key: &AesKeyType, + instance_id: &[u8; INSTANCE_ID_LEN], + key_id: u64, +) -> Result { + let salt = Salt::new(HKDF_SHA256, instance_id); + let pseudo_random_key = salt.extract(admin_key); + let key_id_bytes = key_id.to_be_bytes(); + let info = [b"pb-mapper-temp-key-v1".as_slice(), key_id_bytes.as_slice()]; + let output = pseudo_random_key + .expand(&info, HkdfLen(32)) + .map_err(|_| AuthFailure::internal("failed to expand temporary key"))?; + let mut key = [0_u8; 32]; + output + .fill(&mut key) + .map_err(|_| AuthFailure::internal("failed to fill temporary key"))?; + Ok(key) +} + +struct HkdfLen(usize); + +impl ring::hkdf::KeyType for HkdfLen { + fn len(&self) -> usize { + self.0 + } +} diff --git a/src/common/auth/persistence.rs b/src/common/auth/persistence.rs deleted file mode 100644 index 67bf606..0000000 --- a/src/common/auth/persistence.rs +++ /dev/null @@ -1,1187 +0,0 @@ -//! Durable, encrypted authentication state and audit/replay retention. -//! -//! ```text -//! startup: lock -> admin.key -> recover instance id -> decrypt snapshot -> replay WAL -//! mutation: command -> fsync encrypted WAL -> publish hot-state change -//! compact: hot state + audit + replay set -> snapshot -> truncate WAL -//! ``` -//! -//! Snapshot replacement and administrator-key files use atomic rename. Bounded audit -//! and replay collections are carried through compaction so security history does not -//! disappear when the WAL is truncated. - -use super::*; - -pub(super) const AUTH_SNAPSHOT_FILE: &str = "auth.snapshot"; -pub(super) const AUTH_WAL_FILE: &str = "auth.wal"; - -pub(super) fn auth_snapshot_path(state_dir: &Path) -> PathBuf { - state_dir.join(AUTH_SNAPSHOT_FILE) -} - -pub(super) fn auth_wal_path(state_dir: &Path) -> PathBuf { - state_dir.join(AUTH_WAL_FILE) -} - -pub fn encrypted_auth_state_exists(state_dir: &Path) -> bool { - auth_snapshot_path(state_dir).exists() || auth_wal_path(state_dir).exists() -} - -/// Create the state directory and take `auth.lock` before any credential or -/// snapshot file is read or written. -pub(super) fn prepare_state_dir_and_lock(state_dir: &Path) -> Result, AuthFailure> { - prepare_state_dir(state_dir)?; - Ok(Arc::new(acquire_state_dir_lock(state_dir)?)) -} - -pub fn acquire_state_dir_lock(state_dir: &Path) -> Result { - let path = state_dir.join("auth.lock"); - let file = OpenOptions::new() - .create(true) - .read(true) - .write(true) - .truncate(false) - .open(&path) - .map_err(|error| { - AuthFailure::new( - "auth_state_unavailable", - format!("failed to open `{}`: {error}", path.display()), - false, - ) - })?; - lock_exclusive_nonblock(&file).map_err(|error| { - AuthFailure::new( - "auth_state_locked", - format!( - "authentication state directory `{}` is already in use: {error}", - state_dir.display() - ), - false, - ) - })?; - Ok(file) -} - -fn lock_exclusive_nonblock(file: &File) -> std::io::Result<()> { - #[cfg(unix)] - { - extern "C" { - fn flock(fd: i32, operation: i32) -> i32; - } - const LOCK_EX: i32 = 2; - const LOCK_NB: i32 = 4; - use std::os::unix::io::AsRawFd; - if unsafe { flock(file.as_raw_fd(), LOCK_EX | LOCK_NB) } != 0 { - return Err(std::io::Error::last_os_error()); - } - Ok(()) - } - #[cfg(windows)] - { - use std::os::windows::io::AsRawHandle; - const LOCKFILE_FAIL_IMMEDIATELY: u32 = 0x1; - const LOCKFILE_EXCLUSIVE_LOCK: u32 = 0x2; - #[repr(C)] - struct Overlapped { - internal: usize, - internal_high: usize, - offset: u32, - offset_high: u32, - event: *mut core::ffi::c_void, - } - extern "system" { - fn LockFileEx( - file: *mut core::ffi::c_void, - flags: u32, - reserved: u32, - bytes_low: u32, - bytes_high: u32, - overlapped: *mut Overlapped, - ) -> i32; - } - let mut overlapped = Overlapped { - internal: 0, - internal_high: 0, - offset: 0, - offset_high: 0, - event: core::ptr::null_mut(), - }; - let ok = unsafe { - LockFileEx( - file.as_raw_handle(), - LOCKFILE_FAIL_IMMEDIATELY | LOCKFILE_EXCLUSIVE_LOCK, - 0, - 1, - 0, - &mut overlapped, - ) - }; - if ok == 0 { - Err(std::io::Error::last_os_error()) - } else { - Ok(()) - } - } - #[cfg(not(any(unix, windows)))] - { - let _ = file; - Ok(()) - } -} - -pub(super) fn compaction_is_allowed(safe_mode: bool) -> bool { - !safe_mode -} - -pub(super) fn clear_retained_high_slot_entries(inner: &AuthStateInner) { - inner - .high_slot_entries - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .clear(); -} - -pub(crate) fn replace_file(from: &Path, to: &Path) -> std::io::Result<()> { - #[cfg(windows)] - { - use std::os::windows::ffi::OsStrExt; - - const MOVEFILE_REPLACE_EXISTING: u32 = 0x1; - const MOVEFILE_WRITE_THROUGH: u32 = 0x8; - extern "system" { - fn MoveFileExW( - lp_existing_file_name: *const u16, - lp_new_file_name: *const u16, - dw_flags: u32, - ) -> i32; - } - fn wide(path: &Path) -> Vec { - path.as_os_str() - .encode_wide() - .chain(std::iter::once(0)) - .collect() - } - let from_w = wide(from); - let to_w = wide(to); - let ok = unsafe { - MoveFileExW( - from_w.as_ptr(), - to_w.as_ptr(), - MOVEFILE_REPLACE_EXISTING | MOVEFILE_WRITE_THROUGH, - ) - }; - if ok == 0 { - Err(std::io::Error::last_os_error()) - } else { - Ok(()) - } - } - #[cfg(not(windows))] - std::fs::rename(from, to) -} - -pub(crate) fn sync_parent_directory(path: &Path) -> Result<(), AuthFailure> { - let Some(parent) = path.parent() else { - return Ok(()); - }; - open_directory_for_sync(parent) - .and_then(|directory| directory.sync_all()) - .map_err(|error| { - AuthFailure::new( - "temporary_key_store_unavailable", - format!("failed to sync `{}`: {error}", parent.display()), - false, - ) - }) -} - -fn open_directory_for_sync(path: &Path) -> std::io::Result { - #[cfg(windows)] - { - use std::os::windows::fs::OpenOptionsExt; - const GENERIC_READ: u32 = 0x8000_0000; - const GENERIC_WRITE: u32 = 0x4000_0000; - const FILE_FLAG_BACKUP_SEMANTICS: u32 = 0x0200_0000; - OpenOptions::new() - .access_mode(GENERIC_READ | GENERIC_WRITE) - .custom_flags(FILE_FLAG_BACKUP_SEMANTICS) - .open(path) - } - #[cfg(not(windows))] - File::open(path) -} - -pub(super) fn push_audit_record(inner: &AuthStateInner, record: AuditRecord) { - let mut records = inner - .audit_records - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - while records.len() >= AUDIT_RECORD_CAPACITY { - records.pop_front(); - } - records.push_back(record); -} - -pub(super) fn cancel_all_temporary_leases(inner: &AuthStateInner) { - let slots = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - for lease in slots.iter().filter_map(|slot| slot.lease.upgrade()) { - lease.cancel_rotated(); - } -} - -fn snapshot_generations(inner: &AuthStateInner) -> Vec { - let slots = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let extra = inner - .high_slot_generations - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let mut generations = slots.iter().map(|slot| slot.generation).collect::>(); - generations.extend_from_slice(&extra); - generations -} - -pub(super) fn split_high_slot_state( - snapshot: &PersistedSnapshot, - capacity: usize, -) -> (Vec, Vec) { - let high_generations = snapshot.generations.get(capacity..).unwrap_or(&[]).to_vec(); - let high_entries = snapshot - .entries - .iter() - .filter(|entry| key_slot(entry.key_id) as usize >= capacity) - .cloned() - .collect(); - (high_generations, high_entries) -} - -pub(super) fn build_snapshot( - inner: &AuthStateInner, - cold: &HashMap, - admin_replays: &VecDeque, -) -> PersistedSnapshot { - let slots = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let generations = snapshot_generations(inner); - let mut entries = slots - .iter() - .enumerate() - .filter_map(|(index, slot)| { - if slot.state == SlotState::Free { - return None; - } - let key_id = make_key_id(slot.generation, index as u32); - let cold = cold.get(&key_id)?; - Some(PersistedEntry { - key_id, - state: slot.state, - issued_at: cold.issued_at, - expires_at: slot.expires_at, - label: cold.label.clone(), - tombstoned_at: (cold.tombstoned_at != 0).then_some(cold.tombstoned_at), - }) - }) - .collect::>(); - entries.extend( - inner - .high_slot_entries - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .iter() - .cloned(), - ); - PersistedSnapshot { - schema_version: SNAPSHOT_SCHEMA_VERSION, - instance_id: inner.instance_id(), - generations, - entries, - legacy_protocol: if inner.legacy_protocol_allowed.load(Ordering::Acquire) { - LegacyProtocolPolicy::Allow - } else { - LegacyProtocolPolicy::Deny - }, - admin_replays: admin_replays.iter().cloned().collect(), - audit_records: inner - .audit_records - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .clone(), - root_epoch: inner.root_epoch.load(Ordering::Acquire), - } -} - -pub(super) fn normalize_tombstone_times(snapshot: &mut PersistedSnapshot, now: u64) -> bool { - let mut changed = false; - for entry in &mut snapshot.entries { - if entry.tombstoned_at.is_some() { - continue; - } - let tombstoned_at = match entry.state { - SlotState::Expired => Some(entry.expires_at), - SlotState::Revoked => snapshot - .audit_records - .iter() - .rev() - .find(|record| { - record.action == "temporary_key_revoke" && record.key_id == Some(entry.key_id) - }) - .map(|record| record.at) - .or(Some(now)), - SlotState::Free | SlotState::Active => None, - }; - if tombstoned_at.is_some() { - entry.tombstoned_at = tombstoned_at; - changed = true; - } - } - changed -} - -pub(super) fn empty_snapshot( - inner: &AuthStateInner, - instance_id: [u8; INSTANCE_ID_LEN], - admin_replays: &VecDeque, -) -> PersistedSnapshot { - PersistedSnapshot { - schema_version: SNAPSHOT_SCHEMA_VERSION, - instance_id, - generations: snapshot_generations(inner), - entries: Vec::new(), - legacy_protocol: if inner.legacy_protocol_allowed.load(Ordering::Acquire) { - LegacyProtocolPolicy::Allow - } else { - LegacyProtocolPolicy::Deny - }, - admin_replays: admin_replays.iter().cloned().collect(), - audit_records: inner - .audit_records - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .clone(), - root_epoch: inner.root_epoch.load(Ordering::Acquire), - } -} - -pub(super) fn load_persisted_state( - config: &AuthConfig, - admin_key: &AesKeyType, - instance_id: [u8; INSTANCE_ID_LEN], -) -> (Option, bool) { - match try_load_persisted_state(config, admin_key, instance_id) { - Ok(state) => (Some(state), false), - Err(error) => { - tracing::error!( - event = "auth_state_safe_mode", - auth_stage = "state_load", - reason = %error.code, - error = %error, - "temporary key store failed closed in administrator safe mode" - ); - (None, true) - } - } -} - -pub(super) fn try_load_persisted_state( - config: &AuthConfig, - admin_key: &AesKeyType, - instance_id: [u8; INSTANCE_ID_LEN], -) -> Result { - let snapshot_path = auth_snapshot_path(&config.state_dir); - let mut snapshot = if snapshot_path.exists() { - let bytes = std::fs::read(&snapshot_path).map_err(|error| { - AuthFailure::new( - "temporary_key_store_unavailable", - format!("failed to read `{}`: {error}", snapshot_path.display()), - false, - ) - })?; - let plain = open_blob(admin_key, &bytes)?; - serde_json::from_slice::(&plain).map_err(|error| { - AuthFailure::new( - "temporary_key_store_unavailable", - format!("failed to decode auth snapshot: {error}"), - false, - ) - })? - } else { - PersistedSnapshot { - schema_version: SNAPSHOT_SCHEMA_VERSION, - instance_id, - generations: vec![0; config.max_temporary_keys], - entries: Vec::new(), - legacy_protocol: config.legacy_protocol, - admin_replays: Vec::new(), - audit_records: VecDeque::new(), - root_epoch: 0, - } - }; - if snapshot.schema_version != SNAPSHOT_SCHEMA_VERSION || snapshot.instance_id != instance_id { - return Err(AuthFailure::new( - "temporary_key_store_unavailable", - "auth snapshot schema or server instance id does not match", - false, - )); - } - if snapshot.generations.len() < config.max_temporary_keys { - snapshot.generations.resize(config.max_temporary_keys, 0); - } - - let wal_path = auth_wal_path(&config.state_dir); - if wal_path.exists() { - for record in read_wal(&wal_path, admin_key)? { - match record { - WalRecord::Mutation { mutation, audit } => { - apply_persisted_mutation(&mut snapshot, mutation, config.max_temporary_keys)?; - push_persisted_audit(&mut snapshot.audit_records, audit); - } - WalRecord::AdminReplay(record) => snapshot.admin_replays.push(record), - WalRecord::Audit(audit) => push_persisted_audit(&mut snapshot.audit_records, audit), - } - } - } - Ok(snapshot) -} - -pub(super) fn apply_persisted_mutation( - snapshot: &mut PersistedSnapshot, - mutation: StateMutation, - capacity: usize, -) -> Result<(), AuthFailure> { - match mutation { - StateMutation::Issue(entry) => { - let index = key_slot(entry.key_id) as usize; - if snapshot.generations.len() <= index { - snapshot.generations.resize(index + 1, 0); - } - snapshot.generations[index] = key_generation(entry.key_id); - if index >= capacity { - snapshot - .entries - .retain(|current| key_slot(current.key_id) as usize != index); - snapshot.entries.push(entry); - return Ok(()); - } - snapshot - .entries - .retain(|current| key_slot(current.key_id) as usize != index); - snapshot.entries.push(entry); - } - StateMutation::Renew { key_id, expires_at } => { - let entry = snapshot - .entries - .iter_mut() - .find(|entry| entry.key_id == key_id) - .ok_or_else(|| { - AuthFailure::new( - "temporary_key_store_unavailable", - "WAL renew record references an unknown key", - false, - ) - })?; - entry.expires_at = expires_at; - entry.state = SlotState::Active; - entry.tombstoned_at = None; - } - StateMutation::Revoke { key_id, at } => { - let entry = snapshot - .entries - .iter_mut() - .find(|entry| entry.key_id == key_id) - .ok_or_else(|| { - AuthFailure::new( - "temporary_key_store_unavailable", - "WAL revoke record references an unknown key", - false, - ) - })?; - entry.state = SlotState::Revoked; - entry.tombstoned_at = Some(at); - } - StateMutation::LegacyProtocol(policy) => snapshot.legacy_protocol = policy, - } - Ok(()) -} - -// WHY: `append_wal` sets `retryable` only after it restores the file to the -// pre-append length. Callers persist first and publish hot state only after -// that returns, so a retryable error is a clean no-op. Fail closed only when -// that rollback itself failed and the WAL length is unknown. -pub(super) fn fail_closed_on_uncertain_wal( - inner: &AuthStateInner, - result: Result<(), AuthFailure>, -) -> Result<(), AuthFailure> { - if let Err(error) = &result { - if !error.retryable { - inner.safe_mode.store(true, Ordering::Release); - cancel_all_temporary_leases(inner); - } - } - result -} - -pub(super) fn append_mutation( - config: &AuthConfig, - inner: &AuthStateInner, - mutation: StateMutation, - audit: AuditRecord, -) -> Result<(), AuthFailure> { - fail_closed_on_uncertain_wal( - inner, - append_wal( - config, - &inner.admin_key(), - &WalRecord::Mutation { - mutation, - audit: audit.clone(), - }, - ), - )?; - push_audit_record(inner, audit); - Ok(()) -} - -pub(super) fn append_audit( - config: &AuthConfig, - inner: &AuthStateInner, - audit: AuditRecord, -) -> Result<(), AuthFailure> { - fail_closed_on_uncertain_wal( - inner, - append_wal(config, &inner.admin_key(), &WalRecord::Audit(audit.clone())), - )?; - push_audit_record(inner, audit); - Ok(()) -} - -pub(super) fn push_persisted_audit(records: &mut VecDeque, record: AuditRecord) { - while records.len() >= AUDIT_RECORD_CAPACITY { - records.pop_front(); - } - records.push_back(record); -} - -pub(super) fn append_wal( - config: &AuthConfig, - admin_key: &AesKeyType, - record: &WalRecord, -) -> Result<(), AuthFailure> { - let plain = serde_json::to_vec(record).map_err(|error| { - AuthFailure::internal(format!("failed to encode auth WAL record: {error}")) - })?; - let sealed = seal_blob(admin_key, &plain)?; - let length = u32::try_from(sealed.len()) - .map_err(|_| AuthFailure::internal("auth WAL record is too large"))?; - let path = auth_wal_path(&config.state_dir); - let created = !path.exists(); - let mut file = OpenOptions::new() - .create(true) - .append(true) - .open(&path) - .map_err(|error| { - AuthFailure::new( - "temporary_key_store_unavailable", - format!("failed to open `{}`: {error}", path.display()), - true, - ) - })?; - #[cfg(unix)] - file.set_permissions(std::fs::Permissions::from_mode(0o600)) - .map_err(|error| { - AuthFailure::new( - "temporary_key_store_unavailable", - format!("failed to secure `{}`: {error}", path.display()), - false, - ) - })?; - let start_len = file - .metadata() - .map(|metadata| metadata.len()) - .map_err(|error| { - AuthFailure::new( - "temporary_key_store_unavailable", - format!("failed to inspect `{}`: {error}", path.display()), - true, - ) - })?; - if let Err(error) = file - .write_all(&length.to_be_bytes()) - .and_then(|()| file.write_all(&sealed)) - .and_then(|()| file.sync_data()) - { - // retryable == rolled_back. A later append can then start at a known - // good offset. If truncation fails, the next record would be unreadable. - let rolled_back = file - .set_len(start_len) - .and_then(|()| file.sync_data()) - .is_ok(); - return Err(AuthFailure::new( - "temporary_key_store_unavailable", - if rolled_back { - format!("failed to durably append `{}`: {error}", path.display()) - } else { - format!( - "failed to durably append `{}` and could not restore the previous WAL length: {error}", - path.display() - ) - }, - rolled_back, - )); - } - if created { - sync_parent_directory(&path)?; - } - Ok(()) -} - -pub(super) fn read_wal(path: &Path, admin_key: &AesKeyType) -> Result, AuthFailure> { - let mut file = File::open(path).map_err(|error| { - AuthFailure::new( - "temporary_key_store_unavailable", - format!("failed to open `{}`: {error}", path.display()), - false, - ) - })?; - let mut records = Vec::new(); - loop { - let mut length = [0_u8; 4]; - match file.read(&mut length[..1]) { - Ok(0) => break, - Ok(1) => {} - Ok(_) => unreachable!("single-byte WAL prefix read"), - Err(error) => { - return Err(AuthFailure::new( - "temporary_key_store_unavailable", - format!("failed to read auth WAL length: {error}"), - false, - )); - } - } - file.read_exact(&mut length[1..]).map_err(|error| { - AuthFailure::new( - "temporary_key_store_unavailable", - format!("truncated auth WAL length: {error}"), - false, - ) - })?; - let length = u32::from_be_bytes(length) as usize; - if length > 1024 * 1024 { - return Err(AuthFailure::new( - "temporary_key_store_unavailable", - "auth WAL record exceeds 1 MiB", - false, - )); - } - let mut sealed = vec![0_u8; length]; - file.read_exact(&mut sealed).map_err(|error| { - AuthFailure::new( - "temporary_key_store_unavailable", - format!("truncated auth WAL record: {error}"), - false, - ) - })?; - let plain = open_blob(admin_key, &sealed)?; - records.push(serde_json::from_slice(&plain).map_err(|error| { - AuthFailure::new( - "temporary_key_store_unavailable", - format!("failed to decode auth WAL record: {error}"), - false, - ) - })?); - } - Ok(records) -} - -pub(super) fn write_snapshot_and_truncate_wal( - config: &AuthConfig, - admin_key: &AesKeyType, - snapshot: &PersistedSnapshot, -) -> Result<(), AuthFailure> { - let plain = serde_json::to_vec(snapshot).map_err(|error| { - AuthFailure::internal(format!("failed to encode auth snapshot: {error}")) - })?; - let sealed = seal_blob(admin_key, &plain)?; - let snapshot_path = auth_snapshot_path(&config.state_dir); - atomic_write(&snapshot_path, &sealed, 0o600)?; - truncate_auth_wal(&config.state_dir) -} - -pub(super) fn truncate_auth_wal(state_dir: &Path) -> Result<(), AuthFailure> { - let wal_path = auth_wal_path(state_dir); - let created = !wal_path.exists(); - let wal = OpenOptions::new() - .create(true) - .write(true) - .truncate(true) - .open(&wal_path) - .map_err(|error| { - AuthFailure::new( - "temporary_key_store_unavailable", - format!("failed to truncate `{}`: {error}", wal_path.display()), - true, - ) - })?; - wal.sync_all().map_err(|error| { - AuthFailure::new( - "temporary_key_store_unavailable", - format!("failed to sync `{}`: {error}", wal_path.display()), - true, - ) - })?; - if created { - sync_parent_directory(&wal_path)?; - } - Ok(()) -} - -pub(super) fn seal_blob(admin_key: &AesKeyType, plain: &[u8]) -> Result, AuthFailure> { - let key = LessSafeKey::new( - UnboundKey::new(&AES_256_GCM, admin_key) - .map_err(|_| AuthFailure::internal("failed to initialize state encryption key"))?, - ); - let mut nonce_bytes = [0_u8; 12]; - let mut rng = rand::rng(); - for byte in &mut nonce_bytes { - *byte = rng.random(); - } - let mut output = plain.to_vec(); - key.seal_in_place_append_tag( - Nonce::assume_unique_for_key(nonce_bytes), - Aad::from(STATE_AAD), - &mut output, - ) - .map_err(|_| AuthFailure::internal("failed to encrypt authentication state"))?; - let mut sealed = Vec::with_capacity(STATE_BLOB_MAGIC.len() + nonce_bytes.len() + output.len()); - sealed.extend_from_slice(STATE_BLOB_MAGIC); - sealed.extend_from_slice(&nonce_bytes); - sealed.extend_from_slice(&output); - Ok(sealed) -} - -pub(super) fn open_blob(admin_key: &AesKeyType, sealed: &[u8]) -> Result, AuthFailure> { - if sealed.len() < STATE_BLOB_MAGIC.len() + 12 + AES_256_GCM.tag_len() - || &sealed[..STATE_BLOB_MAGIC.len()] != STATE_BLOB_MAGIC - { - return Err(AuthFailure::new( - "temporary_key_store_unavailable", - "authentication state blob has an invalid header", - false, - )); - } - let nonce_start = STATE_BLOB_MAGIC.len(); - let nonce_end = nonce_start + 12; - let nonce_bytes: [u8; 12] = sealed[nonce_start..nonce_end] - .try_into() - .expect("validated nonce width"); - let mut plain = sealed[nonce_end..].to_vec(); - let key = LessSafeKey::new(UnboundKey::new(&AES_256_GCM, admin_key).map_err(|_| { - AuthFailure::new( - "temporary_key_store_unavailable", - "failed to initialize state decryption key", - false, - ) - })?); - let opened = key - .open_in_place( - Nonce::assume_unique_for_key(nonce_bytes), - Aad::from(STATE_AAD), - &mut plain, - ) - .map_err(|_| { - AuthFailure::new( - "temporary_key_store_unavailable", - "authentication state integrity check failed", - false, - ) - })?; - let len = opened.len(); - plain.truncate(len); - Ok(plain) -} - -pub(super) fn prepare_state_dir(path: &Path) -> Result<(), AuthFailure> { - std::fs::create_dir_all(path).map_err(|error| { - AuthFailure::new( - "auth_state_unavailable", - format!( - "failed to create auth state directory `{}`: {error}", - path.display() - ), - false, - ) - })?; - #[cfg(unix)] - std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o700)).map_err(|error| { - AuthFailure::new( - "auth_state_unavailable", - format!( - "failed to secure auth state directory `{}`: {error}", - path.display() - ), - false, - ) - })?; - Ok(()) -} - -pub(super) fn load_or_create_instance_id( - path: &Path, -) -> Result<[u8; INSTANCE_ID_LEN], AuthFailure> { - let instance_path = path.join("server-instance-id"); - if let Some(instance_id) = read_instance_id_file(&instance_path)? { - return Ok(instance_id); - } - let instance_id = random_instance_id(); - atomic_write(&instance_path, &instance_id, 0o600)?; - Ok(instance_id) -} - -pub(super) fn read_instance_id_file( - path: &Path, -) -> Result, AuthFailure> { - if !path.exists() { - return Ok(None); - } - let bytes = std::fs::read(path).map_err(|error| { - AuthFailure::new( - "auth_state_unavailable", - format!("failed to read `{}`: {error}", path.display()), - false, - ) - })?; - bytes.try_into().map(Some).map_err(|_| { - AuthFailure::new( - "auth_state_unavailable", - "server instance id must be exactly 16 bytes", - false, - ) - }) -} - -/// Promote `server-instance-id.next` when the snapshot already belongs to it. -/// -/// Reset writes that staged file, then the empty snapshot, then the live -/// instance-id file. A crash after the snapshot lands would otherwise fail -/// closed on the next start because the live file still has the old id. -pub(super) fn recover_instance_id_after_reset( - state_dir: &Path, - admin_key: &AesKeyType, - current: [u8; INSTANCE_ID_LEN], -) -> Result<[u8; INSTANCE_ID_LEN], AuthFailure> { - let next_path = state_dir.join("server-instance-id.next"); - let Some(next) = read_instance_id_file(&next_path)? else { - return Ok(current); - }; - let snapshot_path = auth_snapshot_path(state_dir); - if !snapshot_path.exists() { - let _ = std::fs::remove_file(&next_path); - return Ok(current); - } - let bytes = std::fs::read(&snapshot_path).map_err(|error| { - AuthFailure::new( - "temporary_key_store_unavailable", - format!("failed to read `{}`: {error}", snapshot_path.display()), - false, - ) - })?; - let Ok(plain) = open_blob(admin_key, &bytes) else { - return Ok(current); - }; - let Ok(snapshot) = serde_json::from_slice::(&plain) else { - return Ok(current); - }; - if snapshot.instance_id == current { - let _ = std::fs::remove_file(&next_path); - return Ok(current); - } - if snapshot.instance_id != next { - return Ok(current); - } - // The reset snapshot is complete. Any leftover WAL still belongs to the - // previous instance and must not be replayed onto the new derivation id. - truncate_auth_wal(state_dir)?; - atomic_write(&state_dir.join("server-instance-id"), &next, 0o600)?; - let _ = std::fs::remove_file(&next_path); - Ok(next) -} - -pub(super) fn random_instance_id() -> [u8; INSTANCE_ID_LEN] { - let mut instance_id = [0_u8; INSTANCE_ID_LEN]; - let mut rng = rand::rng(); - for byte in &mut instance_id { - *byte = rng.random(); - } - instance_id -} - -pub(super) fn write_admin_key(state_dir: &Path, key: &str) -> Result<(), AuthFailure> { - atomic_write( - &state_dir.join("admin.key"), - format!("{key}\n").as_bytes(), - 0o600, - ) -} - -pub fn generate_admin_key() -> String { - const CHARSET: &[u8] = b"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"; - let mut rng = rand::rng(); - (0..32) - .map(|_| CHARSET[rng.random_range(0..CHARSET.len())] as char) - .collect() -} - -pub fn initialize_admin_key(path: &Path, force: bool) -> Result { - if path.exists() && !force { - return Err(AuthFailure::new( - "administrator_key_exists", - format!("administrator key file `{}` already exists", path.display()), - false, - )); - } - refuse_write_if_encrypted_state(path, force)?; - let key = generate_admin_key(); - atomic_write(path, format!("{key}\n").as_bytes(), 0o600)?; - Ok(key) -} - -pub fn write_admin_key_file(path: &Path, key: &str, force: bool) -> Result<(), AuthFailure> { - let Credential::Admin(_) = parse_credential(key) - .map_err(|error| AuthFailure::new("administrator_key_invalid", error, false))? - else { - return Err(AuthFailure::new( - "administrator_key_invalid", - "administrator key file requires a 32-byte administrator key", - false, - )); - }; - if path.exists() && !force { - return Err(AuthFailure::new( - "administrator_key_exists", - format!( - "administrator key file `{}` already exists; pass --force to replace it", - path.display() - ), - false, - )); - } - if path.file_name() == Some(std::ffi::OsStr::new("admin.key")) - && !key_matches_existing_state(path.parent(), key) - { - refuse_write_if_encrypted_state(path, force)?; - } - atomic_write(path, format!("{key}\n").as_bytes(), 0o600) -} - -pub(super) fn reset_already_installed( - state_dir: &Path, - admin_key: &AesKeyType, - new_instance_id: &[u8; INSTANCE_ID_LEN], -) -> bool { - let Ok(Some(live)) = read_instance_id_file(&state_dir.join("server-instance-id")) else { - return false; - }; - if live != *new_instance_id { - return false; - } - let Ok(bytes) = std::fs::read(auth_snapshot_path(state_dir)) else { - return false; - }; - let Ok(plain) = open_blob(admin_key, &bytes) else { - return false; - }; - let Ok(snapshot) = serde_json::from_slice::(&plain) else { - return false; - }; - snapshot.instance_id == *new_instance_id -} - -pub(super) fn rotation_already_installed(state_dir: &Path, new_key: &str) -> bool { - key_matches_existing_snapshot(Some(state_dir), new_key) - && live_admin_key_matches(state_dir, new_key) -} - -fn live_admin_key_matches(state_dir: &Path, new_key: &str) -> bool { - let Ok(raw) = std::fs::read(state_dir.join("admin.key")) else { - return false; - }; - let Ok(text) = std::str::from_utf8(&raw) else { - return false; - }; - text.trim().as_bytes() == new_key.trim().as_bytes() -} - -pub(super) fn key_matches_existing_snapshot(state_dir: Option<&Path>, key: &str) -> bool { - let Some(state_dir) = state_dir else { - return false; - }; - let snapshot_path = auth_snapshot_path(state_dir); - if !snapshot_path.exists() { - return false; - } - let Ok(Credential::Admin(admin_key)) = parse_credential(key) else { - return false; - }; - let Ok(bytes) = std::fs::read(&snapshot_path) else { - return false; - }; - open_blob(&admin_key, &bytes).is_ok() -} - -pub(super) fn key_matches_existing_state(state_dir: Option<&Path>, key: &str) -> bool { - if key_matches_existing_snapshot(state_dir, key) { - return true; - } - let Some(state_dir) = state_dir else { - return false; - }; - if auth_snapshot_path(state_dir).exists() { - return false; - } - let wal_path = auth_wal_path(state_dir); - if !wal_path.exists() { - return false; - } - let Ok(Credential::Admin(admin_key)) = parse_credential(key) else { - return false; - }; - wal_decrypts_with_key(&wal_path, &admin_key) -} - -fn wal_decrypts_with_key(path: &Path, admin_key: &AesKeyType) -> bool { - let Ok(mut file) = File::open(path) else { - return false; - }; - let Ok(metadata) = file.metadata() else { - return false; - }; - if metadata.len() == 0 { - return true; - } - let mut length = [0_u8; 4]; - if file.read_exact(&mut length).is_err() { - return false; - } - let length = u32::from_be_bytes(length) as usize; - if length == 0 || length > 1024 * 1024 { - return false; - } - let mut sealed = vec![0_u8; length]; - if file.read_exact(&mut sealed).is_err() { - return false; - } - open_blob(admin_key, &sealed).is_ok() -} - -fn refuse_write_if_encrypted_state(path: &Path, force: bool) -> Result<(), AuthFailure> { - // Creating or replacing the live root while snapshot/WAL remain leaves - // those files encrypted under the previous key. Staging `admin.key.next` - // is the rotate path and must stay allowed. - let Some(state_dir) = path.parent() else { - return Ok(()); - }; - if !encrypted_auth_state_exists(state_dir) { - return Ok(()); - } - Err(AuthFailure::new( - "administrator_key_state_exists", - format!( - "refusing to {} `{}` while encrypted auth state exists; use `pb-mapper admin root-key rotate` or `pb-mapper admin auth-state reset --confirm`", - if force { "replace" } else { "create" }, - path.display() - ), - false, - )) -} - -pub(super) fn atomic_write(path: &Path, data: &[u8], mode: u32) -> Result<(), AuthFailure> { - if let Some(parent) = path.parent() { - std::fs::create_dir_all(parent).map_err(|error| { - AuthFailure::new( - "auth_state_unavailable", - format!("failed to create `{}`: {error}", parent.display()), - false, - ) - })?; - } - let file_name = path - .file_name() - .and_then(|name| name.to_str()) - .unwrap_or("auth-state"); - let mut random_suffix = [0_u8; 8]; - let mut rng = rand::rng(); - for byte in &mut random_suffix { - *byte = rng.random(); - } - let temporary = path.with_file_name(format!( - ".{file_name}.tmp-{}-{}", - std::process::id(), - hex(&random_suffix) - )); - let mut file = OpenOptions::new() - .create_new(true) - .write(true) - .open(&temporary) - .map_err(|error| { - AuthFailure::new( - "auth_state_unavailable", - format!("failed to open `{}`: {error}", temporary.display()), - false, - ) - })?; - let result = (|| { - #[cfg(unix)] - file.set_permissions(std::fs::Permissions::from_mode(mode)) - .map_err(|error| { - AuthFailure::internal(format!("failed to set key permissions: {error}")) - })?; - #[cfg(not(unix))] - let _ = mode; - file.write_all(data) - .and_then(|()| file.sync_all()) - .map_err(|error| { - AuthFailure::new( - "auth_state_unavailable", - format!("failed to write `{}`: {error}", temporary.display()), - false, - ) - })?; - drop(file); - replace_file(&temporary, path).map_err(|error| { - AuthFailure::new( - "auth_state_unavailable", - format!("failed to replace `{}`: {error}", path.display()), - false, - ) - })?; - sync_parent_directory(path)?; - Ok(()) - })(); - if result.is_err() { - let _ = std::fs::remove_file(&temporary); - } - result -} - -pub(super) fn unix_seconds() -> u64 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_secs() -} - -pub(super) fn hex(bytes: &[u8]) -> String { - const DIGITS: &[u8; 16] = b"0123456789abcdef"; - let mut output = String::with_capacity(bytes.len() * 2); - for byte in bytes { - output.push(DIGITS[(byte >> 4) as usize] as char); - output.push(DIGITS[(byte & 0x0f) as usize] as char); - } - output -} diff --git a/src/common/auth/persistence/admin_key.rs b/src/common/auth/persistence/admin_key.rs new file mode 100644 index 0000000..33cb7fa --- /dev/null +++ b/src/common/auth/persistence/admin_key.rs @@ -0,0 +1,284 @@ +//! Administrator key files, instance id, and recovery-key identity checks. +use super::super::*; +use super::{ + atomic_write, auth_snapshot_path, auth_wal_path, encrypted_auth_state_exists, open_blob, + truncate_auth_wal, +}; + +pub(in crate::common::auth) fn load_or_create_instance_id( + path: &Path, +) -> Result<[u8; INSTANCE_ID_LEN], AuthFailure> { + let instance_path = path.join("server-instance-id"); + if let Some(instance_id) = read_instance_id_file(&instance_path)? { + return Ok(instance_id); + } + let instance_id = random_instance_id(); + atomic_write(&instance_path, &instance_id, 0o600)?; + Ok(instance_id) +} + +pub(in crate::common::auth) fn read_instance_id_file( + path: &Path, +) -> Result, AuthFailure> { + if !path.exists() { + return Ok(None); + } + let bytes = std::fs::read(path).map_err(|error| { + AuthFailure::new( + "auth_state_unavailable", + format!("failed to read `{}`: {error}", path.display()), + false, + ) + })?; + bytes.try_into().map(Some).map_err(|_| { + AuthFailure::new( + "auth_state_unavailable", + "server instance id must be exactly 16 bytes", + false, + ) + }) +} + +/// Promote `server-instance-id.next` when the snapshot already belongs to it. +/// +/// Reset writes that staged file, then the empty snapshot, then the live +/// instance-id file. A crash after the snapshot lands would otherwise fail +/// closed on the next start because the live file still has the old id. +pub(in crate::common::auth) fn recover_instance_id_after_reset( + state_dir: &Path, + admin_key: &AesKeyType, + current: [u8; INSTANCE_ID_LEN], +) -> Result<[u8; INSTANCE_ID_LEN], AuthFailure> { + let next_path = state_dir.join("server-instance-id.next"); + let Some(next) = read_instance_id_file(&next_path)? else { + return Ok(current); + }; + let snapshot_path = auth_snapshot_path(state_dir); + if !snapshot_path.exists() { + let _ = std::fs::remove_file(&next_path); + return Ok(current); + } + let bytes = std::fs::read(&snapshot_path).map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to read `{}`: {error}", snapshot_path.display()), + false, + ) + })?; + let Ok(plain) = open_blob(admin_key, &bytes) else { + return Ok(current); + }; + let Ok(snapshot) = serde_json::from_slice::(&plain) else { + return Ok(current); + }; + if snapshot.instance_id == current { + let _ = std::fs::remove_file(&next_path); + return Ok(current); + } + if snapshot.instance_id != next { + return Ok(current); + } + // The reset snapshot is complete. Any leftover WAL still belongs to the + // previous instance and must not be replayed onto the new derivation id. + truncate_auth_wal(state_dir)?; + atomic_write(&state_dir.join("server-instance-id"), &next, 0o600)?; + let _ = std::fs::remove_file(&next_path); + Ok(next) +} + +pub(in crate::common::auth) fn random_instance_id() -> [u8; INSTANCE_ID_LEN] { + let mut instance_id = [0_u8; INSTANCE_ID_LEN]; + let mut rng = rand::rng(); + for byte in &mut instance_id { + *byte = rng.random(); + } + instance_id +} + +pub(in crate::common::auth) fn write_admin_key( + state_dir: &Path, + key: &str, +) -> Result<(), AuthFailure> { + atomic_write( + &state_dir.join("admin.key"), + format!("{key}\n").as_bytes(), + 0o600, + ) +} + +pub fn generate_admin_key() -> String { + const CHARSET: &[u8] = b"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"; + let mut rng = rand::rng(); + (0..32) + .map(|_| CHARSET[rng.random_range(0..CHARSET.len())] as char) + .collect() +} + +pub fn initialize_admin_key(path: &Path, force: bool) -> Result { + if path.exists() && !force { + return Err(AuthFailure::new( + "administrator_key_exists", + format!("administrator key file `{}` already exists", path.display()), + false, + )); + } + refuse_write_if_encrypted_state(path, force)?; + let key = generate_admin_key(); + atomic_write(path, format!("{key}\n").as_bytes(), 0o600)?; + Ok(key) +} + +pub fn write_admin_key_file(path: &Path, key: &str, force: bool) -> Result<(), AuthFailure> { + let Credential::Admin(_) = parse_credential(key) + .map_err(|error| AuthFailure::new("administrator_key_invalid", error, false))? + else { + return Err(AuthFailure::new( + "administrator_key_invalid", + "administrator key file requires a 32-byte administrator key", + false, + )); + }; + if path.exists() && !force { + return Err(AuthFailure::new( + "administrator_key_exists", + format!( + "administrator key file `{}` already exists; pass --force to replace it", + path.display() + ), + false, + )); + } + if path.file_name() == Some(std::ffi::OsStr::new("admin.key")) + && !key_matches_existing_state(path.parent(), key) + { + refuse_write_if_encrypted_state(path, force)?; + } + atomic_write(path, format!("{key}\n").as_bytes(), 0o600) +} + +pub(in crate::common::auth) fn reset_already_installed( + state_dir: &Path, + admin_key: &AesKeyType, + new_instance_id: &[u8; INSTANCE_ID_LEN], +) -> bool { + let Ok(Some(live)) = read_instance_id_file(&state_dir.join("server-instance-id")) else { + return false; + }; + if live != *new_instance_id { + return false; + } + let Ok(bytes) = std::fs::read(auth_snapshot_path(state_dir)) else { + return false; + }; + let Ok(plain) = open_blob(admin_key, &bytes) else { + return false; + }; + let Ok(snapshot) = serde_json::from_slice::(&plain) else { + return false; + }; + snapshot.instance_id == *new_instance_id +} + +pub(in crate::common::auth) fn rotation_already_installed(state_dir: &Path, new_key: &str) -> bool { + key_matches_existing_snapshot(Some(state_dir), new_key) + && live_admin_key_matches(state_dir, new_key) +} + +fn live_admin_key_matches(state_dir: &Path, new_key: &str) -> bool { + let Ok(raw) = std::fs::read(state_dir.join("admin.key")) else { + return false; + }; + let Ok(text) = std::str::from_utf8(&raw) else { + return false; + }; + text.trim().as_bytes() == new_key.trim().as_bytes() +} + +pub(in crate::common::auth) fn key_matches_existing_snapshot( + state_dir: Option<&Path>, + key: &str, +) -> bool { + let Some(state_dir) = state_dir else { + return false; + }; + let snapshot_path = auth_snapshot_path(state_dir); + if !snapshot_path.exists() { + return false; + } + let Ok(Credential::Admin(admin_key)) = parse_credential(key) else { + return false; + }; + let Ok(bytes) = std::fs::read(&snapshot_path) else { + return false; + }; + open_blob(&admin_key, &bytes).is_ok() +} + +pub(in crate::common::auth) fn key_matches_existing_state( + state_dir: Option<&Path>, + key: &str, +) -> bool { + if key_matches_existing_snapshot(state_dir, key) { + return true; + } + let Some(state_dir) = state_dir else { + return false; + }; + if auth_snapshot_path(state_dir).exists() { + return false; + } + let wal_path = auth_wal_path(state_dir); + if !wal_path.exists() { + return false; + } + let Ok(Credential::Admin(admin_key)) = parse_credential(key) else { + return false; + }; + wal_decrypts_with_key(&wal_path, &admin_key) +} + +fn wal_decrypts_with_key(path: &Path, admin_key: &AesKeyType) -> bool { + let Ok(mut file) = File::open(path) else { + return false; + }; + let Ok(metadata) = file.metadata() else { + return false; + }; + if metadata.len() == 0 { + return true; + } + let mut length = [0_u8; 4]; + if file.read_exact(&mut length).is_err() { + return false; + } + let length = u32::from_be_bytes(length) as usize; + if length == 0 || length > 1024 * 1024 { + return false; + } + let mut sealed = vec![0_u8; length]; + if file.read_exact(&mut sealed).is_err() { + return false; + } + open_blob(admin_key, &sealed).is_ok() +} + +fn refuse_write_if_encrypted_state(path: &Path, force: bool) -> Result<(), AuthFailure> { + // Creating or replacing the live root while snapshot/WAL remain leaves + // those files encrypted under the previous key. Staging `admin.key.next` + // is the rotate path and must stay allowed. + let Some(state_dir) = path.parent() else { + return Ok(()); + }; + if !encrypted_auth_state_exists(state_dir) { + return Ok(()); + } + Err(AuthFailure::new( + "administrator_key_state_exists", + format!( + "refusing to {} `{}` while encrypted auth state exists; use `pb-mapper admin root-key rotate` or `pb-mapper admin auth-state reset --confirm`", + if force { "replace" } else { "create" }, + path.display() + ), + false, + )) +} diff --git a/src/common/auth/persistence/blob.rs b/src/common/auth/persistence/blob.rs new file mode 100644 index 0000000..288e261 --- /dev/null +++ b/src/common/auth/persistence/blob.rs @@ -0,0 +1,73 @@ +//! AEAD wrap/unwrap for snapshot and WAL payloads. +use super::super::*; + +pub(in crate::common::auth) fn seal_blob( + admin_key: &AesKeyType, + plain: &[u8], +) -> Result, AuthFailure> { + let key = LessSafeKey::new( + UnboundKey::new(&AES_256_GCM, admin_key) + .map_err(|_| AuthFailure::internal("failed to initialize state encryption key"))?, + ); + let mut nonce_bytes = [0_u8; 12]; + let mut rng = rand::rng(); + for byte in &mut nonce_bytes { + *byte = rng.random(); + } + let mut output = plain.to_vec(); + key.seal_in_place_append_tag( + Nonce::assume_unique_for_key(nonce_bytes), + Aad::from(STATE_AAD), + &mut output, + ) + .map_err(|_| AuthFailure::internal("failed to encrypt authentication state"))?; + let mut sealed = Vec::with_capacity(STATE_BLOB_MAGIC.len() + nonce_bytes.len() + output.len()); + sealed.extend_from_slice(STATE_BLOB_MAGIC); + sealed.extend_from_slice(&nonce_bytes); + sealed.extend_from_slice(&output); + Ok(sealed) +} + +pub(in crate::common::auth) fn open_blob( + admin_key: &AesKeyType, + sealed: &[u8], +) -> Result, AuthFailure> { + if sealed.len() < STATE_BLOB_MAGIC.len() + 12 + AES_256_GCM.tag_len() + || &sealed[..STATE_BLOB_MAGIC.len()] != STATE_BLOB_MAGIC + { + return Err(AuthFailure::new( + "temporary_key_store_unavailable", + "authentication state blob has an invalid header", + false, + )); + } + let nonce_start = STATE_BLOB_MAGIC.len(); + let nonce_end = nonce_start + 12; + let nonce_bytes: [u8; 12] = sealed[nonce_start..nonce_end] + .try_into() + .expect("validated nonce width"); + let mut plain = sealed[nonce_end..].to_vec(); + let key = LessSafeKey::new(UnboundKey::new(&AES_256_GCM, admin_key).map_err(|_| { + AuthFailure::new( + "temporary_key_store_unavailable", + "failed to initialize state decryption key", + false, + ) + })?); + let opened = key + .open_in_place( + Nonce::assume_unique_for_key(nonce_bytes), + Aad::from(STATE_AAD), + &mut plain, + ) + .map_err(|_| { + AuthFailure::new( + "temporary_key_store_unavailable", + "authentication state integrity check failed", + false, + ) + })?; + let len = opened.len(); + plain.truncate(len); + Ok(plain) +} diff --git a/src/common/auth/persistence/fs.rs b/src/common/auth/persistence/fs.rs new file mode 100644 index 0000000..a0283d4 --- /dev/null +++ b/src/common/auth/persistence/fs.rs @@ -0,0 +1,275 @@ +//! Directory lock, atomic replace, and parent-directory durability. +use super::super::*; +use super::hex; + +/// Create the state directory and take `auth.lock` before any credential or +/// snapshot file is read or written. +pub(in crate::common::auth) fn prepare_state_dir_and_lock( + state_dir: &Path, +) -> Result, AuthFailure> { + prepare_state_dir(state_dir)?; + Ok(Arc::new(acquire_state_dir_lock(state_dir)?)) +} + +pub fn acquire_state_dir_lock(state_dir: &Path) -> Result { + let path = state_dir.join("auth.lock"); + let file = OpenOptions::new() + .create(true) + .read(true) + .write(true) + .truncate(false) + .open(&path) + .map_err(|error| { + AuthFailure::new( + "auth_state_unavailable", + format!("failed to open `{}`: {error}", path.display()), + false, + ) + })?; + lock_exclusive_nonblock(&file).map_err(|error| { + AuthFailure::new( + "auth_state_locked", + format!( + "authentication state directory `{}` is already in use: {error}", + state_dir.display() + ), + false, + ) + })?; + Ok(file) +} + +fn lock_exclusive_nonblock(file: &File) -> std::io::Result<()> { + #[cfg(unix)] + { + extern "C" { + fn flock(fd: i32, operation: i32) -> i32; + } + const LOCK_EX: i32 = 2; + const LOCK_NB: i32 = 4; + use std::os::unix::io::AsRawFd; + if unsafe { flock(file.as_raw_fd(), LOCK_EX | LOCK_NB) } != 0 { + return Err(std::io::Error::last_os_error()); + } + Ok(()) + } + #[cfg(windows)] + { + use std::os::windows::io::AsRawHandle; + const LOCKFILE_FAIL_IMMEDIATELY: u32 = 0x1; + const LOCKFILE_EXCLUSIVE_LOCK: u32 = 0x2; + #[repr(C)] + struct Overlapped { + internal: usize, + internal_high: usize, + offset: u32, + offset_high: u32, + event: *mut core::ffi::c_void, + } + extern "system" { + fn LockFileEx( + file: *mut core::ffi::c_void, + flags: u32, + reserved: u32, + bytes_low: u32, + bytes_high: u32, + overlapped: *mut Overlapped, + ) -> i32; + } + let mut overlapped = Overlapped { + internal: 0, + internal_high: 0, + offset: 0, + offset_high: 0, + event: core::ptr::null_mut(), + }; + let ok = unsafe { + LockFileEx( + file.as_raw_handle(), + LOCKFILE_FAIL_IMMEDIATELY | LOCKFILE_EXCLUSIVE_LOCK, + 0, + 1, + 0, + &mut overlapped, + ) + }; + if ok == 0 { + Err(std::io::Error::last_os_error()) + } else { + Ok(()) + } + } + #[cfg(not(any(unix, windows)))] + { + let _ = file; + Ok(()) + } +} + +pub(crate) fn replace_file(from: &Path, to: &Path) -> std::io::Result<()> { + #[cfg(windows)] + { + use std::os::windows::ffi::OsStrExt; + + const MOVEFILE_REPLACE_EXISTING: u32 = 0x1; + const MOVEFILE_WRITE_THROUGH: u32 = 0x8; + extern "system" { + fn MoveFileExW( + lp_existing_file_name: *const u16, + lp_new_file_name: *const u16, + dw_flags: u32, + ) -> i32; + } + fn wide(path: &Path) -> Vec { + path.as_os_str() + .encode_wide() + .chain(std::iter::once(0)) + .collect() + } + let from_w = wide(from); + let to_w = wide(to); + let ok = unsafe { + MoveFileExW( + from_w.as_ptr(), + to_w.as_ptr(), + MOVEFILE_REPLACE_EXISTING | MOVEFILE_WRITE_THROUGH, + ) + }; + if ok == 0 { + Err(std::io::Error::last_os_error()) + } else { + Ok(()) + } + } + #[cfg(not(windows))] + std::fs::rename(from, to) +} + +pub(crate) fn sync_parent_directory(path: &Path) -> Result<(), AuthFailure> { + let Some(parent) = path.parent() else { + return Ok(()); + }; + open_directory_for_sync(parent) + .and_then(|directory| directory.sync_all()) + .map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to sync `{}`: {error}", parent.display()), + false, + ) + }) +} + +fn open_directory_for_sync(path: &Path) -> std::io::Result { + #[cfg(windows)] + { + use std::os::windows::fs::OpenOptionsExt; + const GENERIC_READ: u32 = 0x8000_0000; + const GENERIC_WRITE: u32 = 0x4000_0000; + const FILE_FLAG_BACKUP_SEMANTICS: u32 = 0x0200_0000; + OpenOptions::new() + .access_mode(GENERIC_READ | GENERIC_WRITE) + .custom_flags(FILE_FLAG_BACKUP_SEMANTICS) + .open(path) + } + #[cfg(not(windows))] + File::open(path) +} + +pub(in crate::common::auth) fn prepare_state_dir(path: &Path) -> Result<(), AuthFailure> { + std::fs::create_dir_all(path).map_err(|error| { + AuthFailure::new( + "auth_state_unavailable", + format!( + "failed to create auth state directory `{}`: {error}", + path.display() + ), + false, + ) + })?; + #[cfg(unix)] + std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o700)).map_err(|error| { + AuthFailure::new( + "auth_state_unavailable", + format!( + "failed to secure auth state directory `{}`: {error}", + path.display() + ), + false, + ) + })?; + Ok(()) +} + +pub(in crate::common::auth) fn atomic_write( + path: &Path, + data: &[u8], + mode: u32, +) -> Result<(), AuthFailure> { + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent).map_err(|error| { + AuthFailure::new( + "auth_state_unavailable", + format!("failed to create `{}`: {error}", parent.display()), + false, + ) + })?; + } + let file_name = path + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or("auth-state"); + let mut random_suffix = [0_u8; 8]; + let mut rng = rand::rng(); + for byte in &mut random_suffix { + *byte = rng.random(); + } + let temporary = path.with_file_name(format!( + ".{file_name}.tmp-{}-{}", + std::process::id(), + hex(&random_suffix) + )); + let mut file = OpenOptions::new() + .create_new(true) + .write(true) + .open(&temporary) + .map_err(|error| { + AuthFailure::new( + "auth_state_unavailable", + format!("failed to open `{}`: {error}", temporary.display()), + false, + ) + })?; + let result = (|| { + #[cfg(unix)] + file.set_permissions(std::fs::Permissions::from_mode(mode)) + .map_err(|error| { + AuthFailure::internal(format!("failed to set key permissions: {error}")) + })?; + #[cfg(not(unix))] + let _ = mode; + file.write_all(data) + .and_then(|()| file.sync_all()) + .map_err(|error| { + AuthFailure::new( + "auth_state_unavailable", + format!("failed to write `{}`: {error}", temporary.display()), + false, + ) + })?; + drop(file); + replace_file(&temporary, path).map_err(|error| { + AuthFailure::new( + "auth_state_unavailable", + format!("failed to replace `{}`: {error}", path.display()), + false, + ) + })?; + sync_parent_directory(path)?; + Ok(()) + })(); + if result.is_err() { + let _ = std::fs::remove_file(&temporary); + } + result +} diff --git a/src/common/auth/persistence/mod.rs b/src/common/auth/persistence/mod.rs new file mode 100644 index 0000000..c3f130e --- /dev/null +++ b/src/common/auth/persistence/mod.rs @@ -0,0 +1,77 @@ +//! Durable, encrypted authentication state and audit/replay retention. +//! +//! ```text +//! startup: lock -> admin.key -> recover instance id -> decrypt snapshot -> replay WAL +//! mutation: command -> fsync encrypted WAL -> publish hot-state change +//! compact: hot state + audit + replay set -> snapshot -> truncate WAL +//! ``` +//! +//! Snapshot replacement and administrator-key files use atomic rename. Bounded audit +//! and replay collections are carried through compaction so security history does not +//! disappear when the WAL is truncated. + +use super::*; + +mod admin_key; +mod blob; +mod fs; +mod snapshot; +mod wal; + +#[cfg(test)] +pub(in crate::common::auth) use admin_key::read_instance_id_file; +pub use admin_key::{generate_admin_key, initialize_admin_key, write_admin_key_file}; +pub(in crate::common::auth) use admin_key::{ + key_matches_existing_state, load_or_create_instance_id, random_instance_id, + recover_instance_id_after_reset, reset_already_installed, rotation_already_installed, + write_admin_key, +}; +pub(in crate::common::auth) use blob::{open_blob, seal_blob}; +pub use fs::acquire_state_dir_lock; +#[cfg(test)] +pub(in crate::common::auth) use fs::prepare_state_dir; +pub(in crate::common::auth) use fs::{atomic_write, prepare_state_dir_and_lock}; +pub(crate) use fs::{replace_file, sync_parent_directory}; +#[cfg(test)] +pub(in crate::common::auth) use snapshot::try_load_persisted_state; +pub(in crate::common::auth) use snapshot::{ + build_snapshot, cancel_all_temporary_leases, clear_retained_high_slot_entries, + compaction_is_allowed, empty_snapshot, load_persisted_state, normalize_tombstone_times, + push_audit_record, push_persisted_audit, split_high_slot_state, +}; +pub(in crate::common::auth) use wal::{ + append_audit, append_mutation, append_wal, fail_closed_on_uncertain_wal, read_wal, + truncate_auth_wal, write_snapshot_and_truncate_wal, +}; + +pub(in crate::common::auth) const AUTH_SNAPSHOT_FILE: &str = "auth.snapshot"; +pub(in crate::common::auth) const AUTH_WAL_FILE: &str = "auth.wal"; + +pub(in crate::common::auth) fn auth_snapshot_path(state_dir: &Path) -> PathBuf { + state_dir.join(AUTH_SNAPSHOT_FILE) +} + +pub(in crate::common::auth) fn auth_wal_path(state_dir: &Path) -> PathBuf { + state_dir.join(AUTH_WAL_FILE) +} + +pub fn encrypted_auth_state_exists(state_dir: &Path) -> bool { + auth_snapshot_path(state_dir).exists() || auth_wal_path(state_dir).exists() +} + +pub(in crate::common::auth) fn unix_seconds() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() +} + +pub(in crate::common::auth) fn hex(bytes: &[u8]) -> String { + const DIGITS: &[u8; 16] = b"0123456789abcdef"; + let mut output = String::with_capacity(bytes.len() * 2); + for byte in bytes { + output.push(DIGITS[(byte >> 4) as usize] as char); + output.push(DIGITS[(byte & 0x0f) as usize] as char); + } + output +} diff --git a/src/common/auth/persistence/snapshot.rs b/src/common/auth/persistence/snapshot.rs new file mode 100644 index 0000000..b1730b0 --- /dev/null +++ b/src/common/auth/persistence/snapshot.rs @@ -0,0 +1,322 @@ +//! Snapshot construction, load, and mutation replay onto persisted entries. +use super::super::*; +use super::{auth_snapshot_path, auth_wal_path, open_blob, read_wal}; + +pub(in crate::common::auth) fn compaction_is_allowed(safe_mode: bool) -> bool { + !safe_mode +} + +pub(in crate::common::auth) fn clear_retained_high_slot_entries(inner: &AuthStateInner) { + inner + .high_slot_entries + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clear(); +} + +pub(in crate::common::auth) fn push_audit_record(inner: &AuthStateInner, record: AuditRecord) { + let mut records = inner + .audit_records + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + while records.len() >= AUDIT_RECORD_CAPACITY { + records.pop_front(); + } + records.push_back(record); +} + +pub(in crate::common::auth) fn cancel_all_temporary_leases(inner: &AuthStateInner) { + let slots = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + for lease in slots.iter().filter_map(|slot| slot.lease.upgrade()) { + lease.cancel_rotated(); + } +} + +fn snapshot_generations(inner: &AuthStateInner) -> Vec { + let slots = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let extra = inner + .high_slot_generations + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let mut generations = slots.iter().map(|slot| slot.generation).collect::>(); + generations.extend_from_slice(&extra); + generations +} + +pub(in crate::common::auth) fn split_high_slot_state( + snapshot: &PersistedSnapshot, + capacity: usize, +) -> (Vec, Vec) { + let high_generations = snapshot.generations.get(capacity..).unwrap_or(&[]).to_vec(); + let high_entries = snapshot + .entries + .iter() + .filter(|entry| key_slot(entry.key_id) as usize >= capacity) + .cloned() + .collect(); + (high_generations, high_entries) +} + +pub(in crate::common::auth) fn build_snapshot( + inner: &AuthStateInner, + cold: &HashMap, + admin_replays: &VecDeque, +) -> PersistedSnapshot { + let slots = inner + .slots + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let generations = snapshot_generations(inner); + let mut entries = slots + .iter() + .enumerate() + .filter_map(|(index, slot)| { + if slot.state == SlotState::Free { + return None; + } + let key_id = make_key_id(slot.generation, index as u32); + let cold = cold.get(&key_id)?; + Some(PersistedEntry { + key_id, + state: slot.state, + issued_at: cold.issued_at, + expires_at: slot.expires_at, + label: cold.label.clone(), + tombstoned_at: (cold.tombstoned_at != 0).then_some(cold.tombstoned_at), + }) + }) + .collect::>(); + entries.extend( + inner + .high_slot_entries + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .iter() + .cloned(), + ); + snapshot_with( + inner, + inner.instance_id(), + generations, + entries, + admin_replays, + ) +} + +pub(in crate::common::auth) fn normalize_tombstone_times( + snapshot: &mut PersistedSnapshot, + now: u64, +) -> bool { + let mut changed = false; + for entry in &mut snapshot.entries { + if entry.tombstoned_at.is_some() { + continue; + } + let tombstoned_at = match entry.state { + SlotState::Expired => Some(entry.expires_at), + SlotState::Revoked => snapshot + .audit_records + .iter() + .rev() + .find(|record| { + record.action == "temporary_key_revoke" && record.key_id == Some(entry.key_id) + }) + .map(|record| record.at) + .or(Some(now)), + SlotState::Free | SlotState::Active => None, + }; + if tombstoned_at.is_some() { + entry.tombstoned_at = tombstoned_at; + changed = true; + } + } + changed +} + +pub(in crate::common::auth) fn empty_snapshot( + inner: &AuthStateInner, + instance_id: [u8; INSTANCE_ID_LEN], + admin_replays: &VecDeque, +) -> PersistedSnapshot { + snapshot_with( + inner, + instance_id, + snapshot_generations(inner), + Vec::new(), + admin_replays, + ) +} + +fn snapshot_with( + inner: &AuthStateInner, + instance_id: [u8; INSTANCE_ID_LEN], + generations: Vec, + entries: Vec, + admin_replays: &VecDeque, +) -> PersistedSnapshot { + PersistedSnapshot { + schema_version: SNAPSHOT_SCHEMA_VERSION, + instance_id, + generations, + entries, + legacy_protocol: if inner.legacy_protocol_allowed.load(Ordering::Acquire) { + LegacyProtocolPolicy::Allow + } else { + LegacyProtocolPolicy::Deny + }, + admin_replays: admin_replays.iter().cloned().collect(), + audit_records: inner + .audit_records + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clone(), + root_epoch: inner.root_epoch.load(Ordering::Acquire), + } +} + +pub(in crate::common::auth) fn load_persisted_state( + config: &AuthConfig, + admin_key: &AesKeyType, + instance_id: [u8; INSTANCE_ID_LEN], +) -> (Option, bool) { + match try_load_persisted_state(config, admin_key, instance_id) { + Ok(state) => (Some(state), false), + Err(error) => { + tracing::error!( + event = "auth_state_safe_mode", + auth_stage = "state_load", + reason = %error.code, + error = %error, + "temporary key store failed closed in administrator safe mode" + ); + (None, true) + } + } +} + +pub(in crate::common::auth) fn try_load_persisted_state( + config: &AuthConfig, + admin_key: &AesKeyType, + instance_id: [u8; INSTANCE_ID_LEN], +) -> Result { + let snapshot_path = auth_snapshot_path(&config.state_dir); + let mut snapshot = if snapshot_path.exists() { + let bytes = std::fs::read(&snapshot_path).map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to read `{}`: {error}", snapshot_path.display()), + false, + ) + })?; + let plain = open_blob(admin_key, &bytes)?; + serde_json::from_slice::(&plain).map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to decode auth snapshot: {error}"), + false, + ) + })? + } else { + PersistedSnapshot { + schema_version: SNAPSHOT_SCHEMA_VERSION, + instance_id, + generations: vec![0; config.max_temporary_keys], + entries: Vec::new(), + legacy_protocol: config.legacy_protocol, + admin_replays: Vec::new(), + audit_records: VecDeque::new(), + root_epoch: 0, + } + }; + if snapshot.schema_version != SNAPSHOT_SCHEMA_VERSION || snapshot.instance_id != instance_id { + return Err(AuthFailure::new( + "temporary_key_store_unavailable", + "auth snapshot schema or server instance id does not match", + false, + )); + } + if snapshot.generations.len() < config.max_temporary_keys { + snapshot.generations.resize(config.max_temporary_keys, 0); + } + + let wal_path = auth_wal_path(&config.state_dir); + if wal_path.exists() { + for record in read_wal(&wal_path, admin_key)? { + match record { + WalRecord::Mutation { mutation, audit } => { + apply_persisted_mutation(&mut snapshot, mutation)?; + push_persisted_audit(&mut snapshot.audit_records, audit); + } + WalRecord::AdminReplay(record) => snapshot.admin_replays.push(record), + WalRecord::Audit(audit) => push_persisted_audit(&mut snapshot.audit_records, audit), + } + } + } + Ok(snapshot) +} + +pub(in crate::common::auth) fn apply_persisted_mutation( + snapshot: &mut PersistedSnapshot, + mutation: StateMutation, +) -> Result<(), AuthFailure> { + match mutation { + StateMutation::Issue(entry) => { + let index = key_slot(entry.key_id) as usize; + if snapshot.generations.len() <= index { + snapshot.generations.resize(index + 1, 0); + } + snapshot.generations[index] = key_generation(entry.key_id); + snapshot + .entries + .retain(|current| key_slot(current.key_id) as usize != index); + snapshot.entries.push(entry); + } + StateMutation::Renew { key_id, expires_at } => { + let entry = snapshot_entry_mut(snapshot, key_id, "renew")?; + entry.expires_at = expires_at; + entry.state = SlotState::Active; + entry.tombstoned_at = None; + } + StateMutation::Revoke { key_id, at } => { + let entry = snapshot_entry_mut(snapshot, key_id, "revoke")?; + entry.state = SlotState::Revoked; + entry.tombstoned_at = Some(at); + } + StateMutation::LegacyProtocol(policy) => snapshot.legacy_protocol = policy, + } + Ok(()) +} + +fn snapshot_entry_mut<'a>( + snapshot: &'a mut PersistedSnapshot, + key_id: u64, + operation: &str, +) -> Result<&'a mut PersistedEntry, AuthFailure> { + snapshot + .entries + .iter_mut() + .find(|entry| entry.key_id == key_id) + .ok_or_else(|| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("WAL {operation} record references an unknown key"), + false, + ) + }) +} + +pub(in crate::common::auth) fn push_persisted_audit( + records: &mut VecDeque, + record: AuditRecord, +) { + while records.len() >= AUDIT_RECORD_CAPACITY { + records.pop_front(); + } + records.push_back(record); +} diff --git a/src/common/auth/persistence/wal.rs b/src/common/auth/persistence/wal.rs new file mode 100644 index 0000000..e5690eb --- /dev/null +++ b/src/common/auth/persistence/wal.rs @@ -0,0 +1,229 @@ +//! Encrypted WAL append, replay, and snapshot compaction. +use super::super::*; +use super::{ + atomic_write, auth_snapshot_path, auth_wal_path, cancel_all_temporary_leases, open_blob, + push_audit_record, seal_blob, sync_parent_directory, +}; + +pub(in crate::common::auth) fn fail_closed_on_uncertain_wal( + inner: &AuthStateInner, + result: Result<(), AuthFailure>, +) -> Result<(), AuthFailure> { + if let Err(error) = &result { + if !error.retryable { + inner.safe_mode.store(true, Ordering::Release); + cancel_all_temporary_leases(inner); + } + } + result +} + +pub(in crate::common::auth) fn append_mutation( + config: &AuthConfig, + inner: &AuthStateInner, + mutation: StateMutation, + audit: AuditRecord, +) -> Result<(), AuthFailure> { + fail_closed_on_uncertain_wal( + inner, + append_wal( + config, + &inner.admin_key(), + &WalRecord::Mutation { + mutation, + audit: audit.clone(), + }, + ), + )?; + push_audit_record(inner, audit); + Ok(()) +} + +pub(in crate::common::auth) fn append_audit( + config: &AuthConfig, + inner: &AuthStateInner, + audit: AuditRecord, +) -> Result<(), AuthFailure> { + fail_closed_on_uncertain_wal( + inner, + append_wal(config, &inner.admin_key(), &WalRecord::Audit(audit.clone())), + )?; + push_audit_record(inner, audit); + Ok(()) +} + +pub(in crate::common::auth) fn append_wal( + config: &AuthConfig, + admin_key: &AesKeyType, + record: &WalRecord, +) -> Result<(), AuthFailure> { + let plain = serde_json::to_vec(record).map_err(|error| { + AuthFailure::internal(format!("failed to encode auth WAL record: {error}")) + })?; + let sealed = seal_blob(admin_key, &plain)?; + let length = u32::try_from(sealed.len()) + .map_err(|_| AuthFailure::internal("auth WAL record is too large"))?; + let path = auth_wal_path(&config.state_dir); + let created = !path.exists(); + let mut file = OpenOptions::new() + .create(true) + .append(true) + .open(&path) + .map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to open `{}`: {error}", path.display()), + true, + ) + })?; + #[cfg(unix)] + file.set_permissions(std::fs::Permissions::from_mode(0o600)) + .map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to secure `{}`: {error}", path.display()), + false, + ) + })?; + let start_len = file + .metadata() + .map(|metadata| metadata.len()) + .map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to inspect `{}`: {error}", path.display()), + true, + ) + })?; + if let Err(error) = file + .write_all(&length.to_be_bytes()) + .and_then(|()| file.write_all(&sealed)) + .and_then(|()| file.sync_data()) + { + // retryable == rolled_back. A later append can then start at a known + // good offset. If truncation fails, the next record would be unreadable. + let rolled_back = file + .set_len(start_len) + .and_then(|()| file.sync_data()) + .is_ok(); + return Err(AuthFailure::new( + "temporary_key_store_unavailable", + if rolled_back { + format!("failed to durably append `{}`: {error}", path.display()) + } else { + format!( + "failed to durably append `{}` and could not restore the previous WAL length: {error}", + path.display() + ) + }, + rolled_back, + )); + } + if created { + sync_parent_directory(&path)?; + } + Ok(()) +} + +pub(in crate::common::auth) fn read_wal( + path: &Path, + admin_key: &AesKeyType, +) -> Result, AuthFailure> { + let mut file = File::open(path).map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to open `{}`: {error}", path.display()), + false, + ) + })?; + let mut records = Vec::new(); + loop { + let mut length = [0_u8; 4]; + match file.read(&mut length[..1]) { + Ok(0) => break, + Ok(1) => {} + Ok(_) => unreachable!("single-byte WAL prefix read"), + Err(error) => { + return Err(AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to read auth WAL length: {error}"), + false, + )); + } + } + file.read_exact(&mut length[1..]).map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("truncated auth WAL length: {error}"), + false, + ) + })?; + let length = u32::from_be_bytes(length) as usize; + if length > 1024 * 1024 { + return Err(AuthFailure::new( + "temporary_key_store_unavailable", + "auth WAL record exceeds 1 MiB", + false, + )); + } + let mut sealed = vec![0_u8; length]; + file.read_exact(&mut sealed).map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("truncated auth WAL record: {error}"), + false, + ) + })?; + let plain = open_blob(admin_key, &sealed)?; + records.push(serde_json::from_slice(&plain).map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to decode auth WAL record: {error}"), + false, + ) + })?); + } + Ok(records) +} + +pub(in crate::common::auth) fn write_snapshot_and_truncate_wal( + config: &AuthConfig, + admin_key: &AesKeyType, + snapshot: &PersistedSnapshot, +) -> Result<(), AuthFailure> { + let plain = serde_json::to_vec(snapshot).map_err(|error| { + AuthFailure::internal(format!("failed to encode auth snapshot: {error}")) + })?; + let sealed = seal_blob(admin_key, &plain)?; + let snapshot_path = auth_snapshot_path(&config.state_dir); + atomic_write(&snapshot_path, &sealed, 0o600)?; + truncate_auth_wal(&config.state_dir) +} + +pub(in crate::common::auth) fn truncate_auth_wal(state_dir: &Path) -> Result<(), AuthFailure> { + let wal_path = auth_wal_path(state_dir); + let created = !wal_path.exists(); + let wal = OpenOptions::new() + .create(true) + .write(true) + .truncate(true) + .open(&wal_path) + .map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to truncate `{}`: {error}", wal_path.display()), + true, + ) + })?; + wal.sync_all().map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to sync `{}`: {error}", wal_path.display()), + true, + ) + })?; + if created { + sync_parent_directory(&wal_path)?; + } + Ok(()) +} diff --git a/src/common/message/secure.rs b/src/common/message/secure.rs index 0d0ff37..354a95b 100644 --- a/src/common/message/secure.rs +++ b/src/common/message/secure.rs @@ -634,154 +634,6 @@ async fn read_initial_v2_ciphertext( Ok((counter, ciphertext)) } -enum FirstFlightWork { - Live { - key: AesKeyType, - payload: Vec, - error_session: ServerHeaderSession, - }, - Stale(ServerInitialError), -} - -fn first_flight_error( - code: &'static str, - message: impl Into, - retryable: bool, - key_id: u64, -) -> ServerInitialError { - ServerInitialError { - failure: AuthFailure::new(code, message, retryable), - response_session: None, - presented_key_id: Some(key_id), - } -} - -fn reserved_error_session( - replay: &mut ReplayGuard, - fingerprint: &[u8; 32], - session: ServerHeaderSession, -) -> Option { - matches!( - replay.claim(fingerprint, unix_seconds()), - FirstFlightAdmit::Fresh - ) - .then_some(session) -} - -#[allow(clippy::result_large_err)] -fn evaluate_first_flight( - auth: &AuthRuntime, - replay: &std::sync::Mutex, - key_id: u64, - fingerprint: [u8; 32], - work: FirstFlightWork, -) -> std::result::Result<(Vec, AuthContext), ServerInitialError> { - let mut replay = replay - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - match work { - FirstFlightWork::Live { - key, - payload, - error_session, - } => { - let context = auth - .authenticate_presented(key_id, &key) - .map_err(|failure| ServerInitialError { - failure, - response_session: reserved_error_session( - &mut replay, - &fingerprint, - error_session, - ), - presented_key_id: Some(key_id), - })?; - match replay.admit(key_id, &fingerprint, unix_seconds()) { - FirstFlightAdmit::Fresh => Ok((payload, context)), - FirstFlightAdmit::Replayed => Err(first_flight_error( - "connection_salt_replayed", - "protocol-v2 connection salt was already accepted", - true, - key_id, - )), - FirstFlightAdmit::Limited => Err(first_flight_error( - "connection_admission_limited", - "this credential has opened too many new connections in the current window", - true, - key_id, - )), - FirstFlightAdmit::Unavailable => Err(first_flight_error( - "connection_replay_store_unavailable", - "failed to persist first-flight replay admission", - true, - key_id, - )), - } - } - FirstFlightWork::Stale(mut error) => { - if let Some(session) = error.response_session.take() { - error.response_session = reserved_error_session(&mut replay, &fingerprint, session); - } - Err(error) - } - } -} - -fn stale_root_first_flight( - auth: &AuthRuntime, - key_id: u64, - salt: [u8; CONNECTION_SALT_LEN], - counter: u64, - ciphertext: &[u8], -) -> Option { - let previous_key = auth.derive_previous_key(key_id)?; - let previous_material = derive_material(key_id, &previous_key, salt).ok()?; - let mut previous_ciphertext = ciphertext.to_vec(); - open_v2_payload( - &previous_material, - DIRECTION_CLIENT_TO_SERVER, - counter, - &mut previous_ciphertext, - ) - .ok()?; - let (code, message) = if key_id == 0 { - ( - "administrator_key_invalid", - "administrator credential does not match the active root key", - ) - } else { - ( - "temporary_key_rotated", - "temporary credential was invalidated by administrator root rotation or auth-state reset", - ) - }; - Some(ServerInitialError { - failure: AuthFailure::new(code, message, false), - response_session: Some(v2_session(previous_key, previous_material)), - presented_key_id: Some(key_id), - }) -} - -fn v2_session(key: AesKeyType, material: V2Material) -> ServerHeaderSession { - ServerHeaderSession { - protocol: HeaderProtocol::V2, - legacy_key: key, - v2: Some(material), - context: None, - _legacy_guard: None, - } -} - -fn session_without_context(session: &ServerHeaderSession) -> ServerHeaderSession { - ServerHeaderSession { - protocol: session.protocol, - legacy_key: session.legacy_key, - v2: session.v2.clone(), - context: None, - _legacy_guard: None, - } -} - pub enum HeaderMessageReader<'a, T: AsyncReadExt + Unpin> { Legacy(CodecMessageReader<'a, T, Aes256GcmDeCodec>), V2(V2MessageReader<'a, T>), @@ -817,6 +669,8 @@ mod replay; #[cfg(test)] use replay::RotatingBloom; use replay::{replay_fingerprint, FirstFlightAdmit, ReplayGuard}; +mod first_flight; +use first_flight::*; fn legacy_message_reader<'a, T: AsyncReadExt + Unpin>( reader: &'a mut T, key: &AesKeyType, diff --git a/src/common/message/secure/first_flight.rs b/src/common/message/secure/first_flight.rs new file mode 100644 index 0000000..45a017b --- /dev/null +++ b/src/common/message/secure/first_flight.rs @@ -0,0 +1,150 @@ +//! First-flight authentication, replay admission, and stale-root classification. +use super::*; + +pub(super) enum FirstFlightWork { + Live { + key: AesKeyType, + payload: Vec, + error_session: ServerHeaderSession, + }, + Stale(ServerInitialError), +} + +pub(super) fn first_flight_error( + code: &'static str, + message: impl Into, + retryable: bool, + key_id: u64, +) -> ServerInitialError { + ServerInitialError { + failure: AuthFailure::new(code, message, retryable), + response_session: None, + presented_key_id: Some(key_id), + } +} + +fn reserved_error_session( + replay: &mut ReplayGuard, + fingerprint: &[u8; 32], + session: ServerHeaderSession, +) -> Option { + matches!( + replay.claim(fingerprint, unix_seconds()), + FirstFlightAdmit::Fresh + ) + .then_some(session) +} + +#[allow(clippy::result_large_err)] +pub(super) fn evaluate_first_flight( + auth: &AuthRuntime, + replay: &std::sync::Mutex, + key_id: u64, + fingerprint: [u8; 32], + work: FirstFlightWork, +) -> std::result::Result<(Vec, AuthContext), ServerInitialError> { + let mut replay = replay + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + match work { + FirstFlightWork::Live { + key, + payload, + error_session, + } => { + let context = auth + .authenticate_presented(key_id, &key) + .map_err(|failure| ServerInitialError { + failure, + response_session: reserved_error_session( + &mut replay, + &fingerprint, + error_session, + ), + presented_key_id: Some(key_id), + })?; + match replay.admit(key_id, &fingerprint, unix_seconds()) { + FirstFlightAdmit::Fresh => Ok((payload, context)), + FirstFlightAdmit::Replayed => Err(first_flight_error( + "connection_salt_replayed", + "protocol-v2 connection salt was already accepted", + true, + key_id, + )), + FirstFlightAdmit::Limited => Err(first_flight_error( + "connection_admission_limited", + "this credential has opened too many new connections in the current window", + true, + key_id, + )), + FirstFlightAdmit::Unavailable => Err(first_flight_error( + "connection_replay_store_unavailable", + "failed to persist first-flight replay admission", + true, + key_id, + )), + } + } + FirstFlightWork::Stale(mut error) => { + if let Some(session) = error.response_session.take() { + error.response_session = reserved_error_session(&mut replay, &fingerprint, session); + } + Err(error) + } + } +} + +pub(super) fn stale_root_first_flight( + auth: &AuthRuntime, + key_id: u64, + salt: [u8; CONNECTION_SALT_LEN], + counter: u64, + ciphertext: &[u8], +) -> Option { + let previous_key = auth.derive_previous_key(key_id)?; + let previous_material = derive_material(key_id, &previous_key, salt).ok()?; + let mut previous_ciphertext = ciphertext.to_vec(); + open_v2_payload( + &previous_material, + DIRECTION_CLIENT_TO_SERVER, + counter, + &mut previous_ciphertext, + ) + .ok()?; + let (code, message) = if key_id == 0 { + ( + "administrator_key_invalid", + "administrator credential does not match the active root key", + ) + } else { + ( + "temporary_key_rotated", + "temporary credential was invalidated by administrator root rotation or auth-state reset", + ) + }; + Some(ServerInitialError { + failure: AuthFailure::new(code, message, false), + response_session: Some(v2_session(previous_key, previous_material)), + presented_key_id: Some(key_id), + }) +} + +pub(super) fn v2_session(key: AesKeyType, material: V2Material) -> ServerHeaderSession { + ServerHeaderSession { + protocol: HeaderProtocol::V2, + legacy_key: key, + v2: Some(material), + context: None, + _legacy_guard: None, + } +} + +pub(super) fn session_without_context(session: &ServerHeaderSession) -> ServerHeaderSession { + ServerHeaderSession { + protocol: session.protocol, + legacy_key: session.legacy_key, + v2: session.v2.clone(), + context: None, + _legacy_guard: None, + } +} From d25c448384f4628501dab50f45599c7753063607 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Wed, 19 Aug 2026 13:16:33 +0800 Subject: [PATCH 58/74] Pin tunnel credentials once and share first-flight IO Resolve the register-pool credential before spawn, put every control worker in one JoinSet, and reuse ClientHeaderSession::exchange plus a shared v2 frame reader. FFI now keeps handle and pin together per role. --- src/bin/pb-mapper.rs | 35 +- src/common/message/secure.rs | 298 ++++++++---------- src/common/message/secure/first_flight.rs | 6 +- src/common/message/secure/frame.rs | 64 ++-- src/local/client/status.rs | 32 +- src/local/client/stream.rs | 32 +- src/local/server/mod.rs | 110 +++---- src/local/server/stream.rs | 31 +- src/pb_server/error.rs | 38 --- src/pb_server/mod.rs | 18 +- ui/native/pb_mapper_ffi/src/state.rs | 27 +- .../pb_mapper_ffi/src/state/configuration.rs | 50 +-- ui/native/pb_mapper_ffi/src/state/runtime.rs | 236 +++++++------- ui/native/pb_mapper_ffi/src/state/status.rs | 143 ++++----- 14 files changed, 484 insertions(+), 636 deletions(-) diff --git a/src/bin/pb-mapper.rs b/src/bin/pb-mapper.rs index 572538a..1f6b4fa 100644 --- a/src/bin/pb-mapper.rs +++ b/src/bin/pb-mapper.rs @@ -36,8 +36,10 @@ use pb_mapper::common::message::command::{ use pb_mapper::common::message::forward::StreamForward; use pb_mapper::common::message::secure::ClientHeaderSession; use pb_mapper::common::message::MessageReader; -use pb_mapper::local::client::{handle_status_cli_scoped, run_client_side_cli_scoped}; -use pb_mapper::local::server::{run_server_side_cli, ServerTunnelOptions}; +use pb_mapper::local::client::{ + handle_status_cli_scoped, run_client_side_cli_with_callback_scoped, +}; +use pb_mapper::local::server::{run_server_side_cli_with_pinned_credential, ServerTunnelOptions}; use pb_mapper::pb_server::run_server_with_shutdown; use tokio::net::TcpStream; use tokio_util::sync::CancellationToken; @@ -315,7 +317,7 @@ async fn run_server(args: ServerArgs) -> Result<(), Box> { } async fn run_register(args: RegisterArgs) -> Result<(), Box> { - pb_mapper::common::checksum::get_process_credential().map_err(|error| { + let credential = pb_mapper::common::checksum::get_process_credential().map_err(|error| { std::io::Error::other(format!("registration credential is required: {error}")) })?; let local_addr = get_sockaddr_async(&args.addr).await?; @@ -330,10 +332,12 @@ async fn run_register(args: RegisterArgs) -> Result<(), Box> { match args.transport { Transport::Tcp => { - register::(local_addr, remote_addr, args.key, options).await + register::(local_addr, remote_addr, args.key, options, credential) + .await } Transport::Udp => { - register::(local_addr, remote_addr, args.key, options).await + register::(local_addr, remote_addr, args.key, options, credential) + .await } } Ok(()) @@ -344,14 +348,23 @@ async fn register( remote_addr: std::net::SocketAddr, key: String, options: ServerTunnelOptions, + credential: pb_mapper::common::checksum::Credential, ) where LocalStream::Item: StreamForward, { - run_server_side_cli::(local_addr, remote_addr, key.into(), options).await; + run_server_side_cli_with_pinned_credential::( + local_addr, + remote_addr, + key.into(), + options, + None, + credential, + ) + .await; } async fn run_connect(args: ConnectArgs) -> Result<(), Box> { - pb_mapper::common::checksum::get_process_credential().map_err(|error| { + let credential = pb_mapper::common::checksum::get_process_credential().map_err(|error| { std::io::Error::other(format!("client credential is required: {error}")) })?; let local_addr = get_sockaddr_async(&args.addr).await?; @@ -361,22 +374,26 @@ async fn run_connect(args: ConnectArgs) -> Result<(), Box> { match args.transport { Transport::Tcp => { - run_client_side_cli_scoped::( + run_client_side_cli_with_callback_scoped::( local_addr, remote_addr, key, keep_alive, args.namespace, + None, + Some(credential), ) .await; } Transport::Udp => { - run_client_side_cli_scoped::( + run_client_side_cli_with_callback_scoped::( local_addr, remote_addr, key, keep_alive, args.namespace, + None, + Some(credential), ) .await; } diff --git a/src/common/message/secure.rs b/src/common/message/secure.rs index 354a95b..57aac69 100644 --- a/src/common/message/secure.rs +++ b/src/common/message/secure.rs @@ -18,7 +18,7 @@ //! log suppression, and protocol tests are isolated in focused child modules. use std::sync::{Arc, Mutex}; -use std::time::{SystemTime, UNIX_EPOCH}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; use rand::RngExt; use ring::aead::{Aad, LessSafeKey, Nonce, UnboundKey, AES_256_GCM}; @@ -144,6 +144,32 @@ impl ClientHeaderSession { } } + pub async fn exchange( + &self, + stream: &mut T, + payload: &[u8], + timeout: Duration, + ) -> Result> { + match tokio::time::timeout(timeout, self.write_initial(stream, payload)).await { + Ok(result) => result?, + Err(_) => { + return Err(protocol_error(format!( + "timed out writing first-flight request after {timeout:?}" + ))) + } + } + let mut reader = self.response_reader(stream)?; + let message = match tokio::time::timeout(timeout, reader.read_msg()).await { + Ok(result) => result?, + Err(_) => { + return Err(protocol_error(format!( + "timed out reading first-flight response after {timeout:?}" + ))) + } + }; + Ok(message.to_vec()) + } + pub fn continuation_writer<'a, T: AsyncWriteExt + Unpin>( &self, writer: &'a mut T, @@ -276,6 +302,31 @@ impl ServerInitialError { presented_key_id: None, } } + + fn fail(code: &'static str, message: impl Into, retryable: bool) -> Self { + Self::new(AuthFailure::new(code, message, retryable)) + } + + fn fail_key( + code: &'static str, + message: impl Into, + retryable: bool, + key_id: u64, + ) -> Self { + Self { + failure: AuthFailure::new(code, message, retryable), + response_session: None, + presented_key_id: Some(key_id), + } + } + + fn from_failure_key(failure: AuthFailure, key_id: u64) -> Self { + Self { + failure, + response_session: None, + presented_key_id: Some(key_id), + } + } } impl fmt::Debug for ServerInitialError { @@ -361,98 +412,59 @@ impl ServerSecurity { checksum_bytes: [u8; 4], ) -> std::result::Result { if !self.auth.legacy_protocol_allowed().unwrap_or(false) { - return Err(ServerInitialError { - failure: AuthFailure::new( - "legacy_protocol_disabled", - "legacy protocol is disabled by the administrator", - false, - ), - response_session: None, - presented_key_id: None, - }); + return Err(ServerInitialError::fail( + "legacy_protocol_disabled", + "legacy protocol is disabled by the administrator", + false, + )); } - let key = self - .auth - .admin_key() - .map_err(|failure| ServerInitialError { - failure, - response_session: None, - presented_key_id: None, - })?; + let key = self.auth.admin_key().map_err(ServerInitialError::new)?; let checksum = u32::from_be_bytes(checksum_bytes); - let datalen = reader - .read_u32() - .await - .map_err(|error| ServerInitialError { - failure: AuthFailure::new( - "legacy_frame_invalid", - format!("failed to read legacy frame length: {error}"), - true, - ), - response_session: None, - presented_key_id: None, - })?; + let datalen = reader.read_u32().await.map_err(|error| { + ServerInitialError::fail( + "legacy_frame_invalid", + format!("failed to read legacy frame length: {error}"), + true, + ) + })?; if !valid_checksum_for_key(datalen, checksum, &key) || datalen > MAX_INITIAL_CIPHERTEXT_LEN { - return Err(ServerInitialError { - failure: AuthFailure::new( - "legacy_frame_invalid", - "legacy frame checksum or length is invalid", - false, - ), - response_session: None, - presented_key_id: None, - }); + return Err(ServerInitialError::fail( + "legacy_frame_invalid", + "legacy frame checksum or length is invalid", + false, + )); } let mut encrypted = vec![0_u8; datalen as usize]; - reader - .read_exact(&mut encrypted) - .await - .map_err(|error| ServerInitialError { - failure: AuthFailure::new( - "legacy_frame_invalid", - format!("failed to read legacy frame body: {error}"), - true, - ), - response_session: None, - presented_key_id: None, - })?; - let mut codec = Aes256GcmDeCodec::try_new(&key).map_err(|_| ServerInitialError { - failure: AuthFailure::new( + reader.read_exact(&mut encrypted).await.map_err(|error| { + ServerInitialError::fail( + "legacy_frame_invalid", + format!("failed to read legacy frame body: {error}"), + true, + ) + })?; + let mut codec = Aes256GcmDeCodec::try_new(&key).map_err(|_| { + ServerInitialError::fail( "legacy_decrypt_failed", "failed to initialize legacy decryption", false, - ), - response_session: None, - presented_key_id: None, + ) + })?; + let plain = codec.decrypt(&mut encrypted).map_err(|_| { + ServerInitialError::fail( + "legacy_decrypt_failed", + "legacy credential or encrypted frame is invalid", + false, + ) })?; - let plain = codec - .decrypt(&mut encrypted) - .map_err(|_| ServerInitialError { - failure: AuthFailure::new( - "legacy_decrypt_failed", - "legacy credential or encrypted frame is invalid", - false, - ), - response_session: None, - presented_key_id: None, - })?; let context = self .auth .authenticate_presented(0, &key) - .map_err(|failure| ServerInitialError { - failure, - response_session: None, - presented_key_id: None, - })?; - let legacy_guard = - self.auth - .record_legacy_connection() - .map_err(|failure| ServerInitialError { - failure, - response_session: None, - presented_key_id: None, - })?; + .map_err(ServerInitialError::new)?; + let legacy_guard = self + .auth + .record_legacy_connection() + .map_err(ServerInitialError::new)?; Ok(ServerInitialMessage { payload: plain.to_vec(), session: ServerHeaderSession { @@ -472,37 +484,28 @@ impl ServerSecurity { reader: &mut T, ) -> std::result::Result { let mut remainder = [0_u8; FIRST_PREFIX_REMAINDER_LEN]; - reader - .read_exact(&mut remainder) - .await - .map_err(|error| ServerInitialError { - failure: AuthFailure::new( - "protocol_v2_header_invalid", - format!("failed to read protocol-v2 header: {error}"), - true, - ), - response_session: None, - presented_key_id: None, - })?; + reader.read_exact(&mut remainder).await.map_err(|error| { + ServerInitialError::fail( + "protocol_v2_header_invalid", + format!("failed to read protocol-v2 header: {error}"), + true, + ) + })?; let version = remainder[0]; let flags = remainder[1]; let reserved = u16::from_be_bytes([remainder[2], remainder[3]]); if version != PROTOCOL_V2_VERSION || flags != 0 || reserved != 0 { - return Err(ServerInitialError { - failure: AuthFailure::new( - if version != PROTOCOL_V2_VERSION { - "protocol_version_unsupported" - } else { - "protocol_v2_header_invalid" - }, - format!( - "unsupported protocol header version={version} flags={flags} reserved={reserved}" - ), - false, + return Err(ServerInitialError::fail( + if version != PROTOCOL_V2_VERSION { + "protocol_version_unsupported" + } else { + "protocol_v2_header_invalid" + }, + format!( + "unsupported protocol header version={version} flags={flags} reserved={reserved}" ), - response_session: None, - presented_key_id: None, - }); + false, + )); } let key_id = u64::from_be_bytes(remainder[4..12].try_into().expect("fixed key id")); let salt: [u8; CONNECTION_SALT_LEN] = @@ -510,46 +513,36 @@ impl ServerSecurity { let client_timestamp = u64::from_be_bytes(salt[..8].try_into().expect("fixed timestamp")); let now = unix_seconds(); if now.abs_diff(client_timestamp) > MAX_CONNECTION_CLOCK_SKEW_SECONDS { - return Err(ServerInitialError { - failure: AuthFailure::new( - "connection_timestamp_invalid", - "protocol-v2 connection timestamp is outside the accepted clock-skew window", - false, - ), - response_session: None, - presented_key_id: Some(key_id), - }); + return Err(ServerInitialError::fail_key( + "connection_timestamp_invalid", + "protocol-v2 connection timestamp is outside the accepted clock-skew window", + false, + key_id, + )); } let key = self .auth .derive_key(key_id) - .map_err(|failure| ServerInitialError { - failure, - response_session: None, - presented_key_id: Some(key_id), - })?; - let material = derive_material(key_id, &key, salt).map_err(|error| ServerInitialError { - failure: AuthFailure::new( + .map_err(|failure| ServerInitialError::from_failure_key(failure, key_id))?; + let material = derive_material(key_id, &key, salt).map_err(|error| { + ServerInitialError::fail_key( "protocol_v2_key_derivation_failed", error.to_string(), false, - ), - response_session: None, - presented_key_id: Some(key_id), + key_id, + ) })?; let mut session = v2_session(key, material.clone()); - let (counter, ciphertext) = - read_initial_v2_ciphertext(reader) - .await - .map_err(|error| ServerInitialError { - failure: AuthFailure::new( - "protocol_v2_decrypt_failed", - error.to_string(), - false, - ), - response_session: None, - presented_key_id: Some(key_id), - })?; + let (counter, ciphertext) = read_v2_frame(reader, 0, MAX_INITIAL_PLAINTEXT_LEN) + .await + .map_err(|error| { + ServerInitialError::fail_key( + "protocol_v2_decrypt_failed", + error.to_string(), + false, + key_id, + ) + })?; let mut current_ciphertext = ciphertext.clone(); let fingerprint = replay_fingerprint(key_id, &salt); let work = match open_v2_payload( @@ -604,35 +597,6 @@ impl ServerSecurity { mod limiter; pub use limiter::FailureLogDecision; use limiter::FailureLogLimiter; -async fn read_initial_v2_ciphertext( - reader: &mut T, -) -> Result<(u64, Vec)> { - let counter = reader - .read_u64() - .await - .map_err(|error| protocol_error(format!("failed to read v2 counter: {error}")))?; - if counter != 0 { - return Err(protocol_error(format!( - "protocol-v2 counter mismatch: expected 0, got {counter}" - ))); - } - let datalen = reader - .read_u32() - .await - .map_err(|error| protocol_error(format!("failed to read v2 length: {error}")))?; - let max_encrypted_len = MAX_INITIAL_PLAINTEXT_LEN.saturating_add(AES_256_GCM.tag_len() as u32); - if datalen < AES_256_GCM.tag_len() as u32 || datalen > max_encrypted_len { - return Err(protocol_error(format!( - "protocol-v2 payload length {datalen} exceeds the {MAX_INITIAL_PLAINTEXT_LEN}-byte limit" - ))); - } - let mut ciphertext = vec![0_u8; datalen as usize]; - reader - .read_exact(&mut ciphertext) - .await - .map_err(|error| protocol_error(format!("failed to read v2 payload: {error}")))?; - Ok((counter, ciphertext)) -} pub enum HeaderMessageReader<'a, T: AsyncReadExt + Unpin> { Legacy(CodecMessageReader<'a, T, Aes256GcmDeCodec>), @@ -663,7 +627,7 @@ impl MessageWriter for HeaderMessageWriter<'_, T> { } mod frame; -use frame::{derive_material, first_prefix, open_v2_payload, V2Material}; +use frame::{derive_material, first_prefix, open_v2_payload, read_v2_frame, V2Material}; pub use frame::{V2MessageReader, V2MessageWriter}; mod replay; #[cfg(test)] diff --git a/src/common/message/secure/first_flight.rs b/src/common/message/secure/first_flight.rs index 45a017b..9b8468a 100644 --- a/src/common/message/secure/first_flight.rs +++ b/src/common/message/secure/first_flight.rs @@ -16,11 +16,7 @@ pub(super) fn first_flight_error( retryable: bool, key_id: u64, ) -> ServerInitialError { - ServerInitialError { - failure: AuthFailure::new(code, message, retryable), - response_session: None, - presented_key_id: Some(key_id), - } + ServerInitialError::fail_key(code, message, retryable, key_id) } fn reserved_error_session( diff --git a/src/common/message/secure/frame.rs b/src/common/message/secure/frame.rs index da0afd0..d76f75b 100644 --- a/src/common/message/secure/frame.rs +++ b/src/common/message/secure/frame.rs @@ -53,33 +53,11 @@ impl<'a, T: AsyncReadExt + Unpin> V2MessageReader<'a, T> { } pub(super) async fn read_msg_with_limit(&mut self, max_plaintext_len: u32) -> Result<&'_ [u8]> { - let counter = self - .reader - .read_u64() - .await - .map_err(|error| protocol_error(format!("failed to read v2 counter: {error}")))?; - if counter != self.expected_counter { - return Err(protocol_error(format!( - "protocol-v2 counter mismatch: expected {}, got {counter}", - self.expected_counter - ))); - } - let datalen = self - .reader - .read_u32() - .await - .map_err(|error| protocol_error(format!("failed to read v2 length: {error}")))?; - let max_encrypted_len = max_plaintext_len.saturating_add(AES_256_GCM.tag_len() as u32); - if datalen < AES_256_GCM.tag_len() as u32 || datalen > max_encrypted_len { - return Err(protocol_error(format!( - "protocol-v2 payload length {datalen} exceeds the {max_plaintext_len}-byte limit" - ))); - } - self.buffer.resize(datalen as usize, 0); - self.reader - .read_exact(&mut self.buffer) - .await - .map_err(|error| protocol_error(format!("failed to read v2 payload: {error}")))?; + let (counter, ciphertext) = + read_v2_frame(self.reader, self.expected_counter, max_plaintext_len).await?; + self.buffer = ciphertext; + let datalen = u32::try_from(self.buffer.len()) + .map_err(|_| protocol_error("protocol-v2 payload exceeds u32 length"))?; let aad = frame_aad(&self.material, self.direction, counter, datalen); let plain = self .key @@ -95,6 +73,38 @@ impl<'a, T: AsyncReadExt + Unpin> V2MessageReader<'a, T> { } } +pub(super) async fn read_v2_frame( + reader: &mut T, + expected_counter: u64, + max_plaintext_len: u32, +) -> Result<(u64, Vec)> { + let counter = reader + .read_u64() + .await + .map_err(|error| protocol_error(format!("failed to read v2 counter: {error}")))?; + if counter != expected_counter { + return Err(protocol_error(format!( + "protocol-v2 counter mismatch: expected {expected_counter}, got {counter}" + ))); + } + let datalen = reader + .read_u32() + .await + .map_err(|error| protocol_error(format!("failed to read v2 length: {error}")))?; + let max_encrypted_len = max_plaintext_len.saturating_add(AES_256_GCM.tag_len() as u32); + if datalen < AES_256_GCM.tag_len() as u32 || datalen > max_encrypted_len { + return Err(protocol_error(format!( + "protocol-v2 payload length {datalen} exceeds the {max_plaintext_len}-byte limit" + ))); + } + let mut ciphertext = vec![0_u8; datalen as usize]; + reader + .read_exact(&mut ciphertext) + .await + .map_err(|error| protocol_error(format!("failed to read v2 payload: {error}")))?; + Ok((counter, ciphertext)) +} + impl MessageReader for V2MessageReader<'_, T> { async fn read_msg(&mut self) -> Result<&'_ [u8]> { self.read_msg_with_limit(MAX_MSG_LEN - AES_256_GCM.tag_len() as u32) diff --git a/src/local/client/status.rs b/src/local/client/status.rs index a479ab3..1760c35 100644 --- a/src/local/client/status.rs +++ b/src/local/client/status.rs @@ -2,8 +2,8 @@ use snafu::ResultExt; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use super::error::{ - ControlIoTimeoutSnafu, CreateHeaderToolSnafu, DecodeStatusRespSnafu, EncodeStatusReqSnafu, - ReadStatusRespSnafu, StatusRespNotMatchSnafu, WriteStatusReqSnafu, + CreateHeaderToolSnafu, DecodeStatusRespSnafu, EncodeStatusReqSnafu, StatusRespNotMatchSnafu, + WriteStatusReqSnafu, }; use crate::common::checksum::Credential; use crate::common::config::control_io_timeout; @@ -11,7 +11,6 @@ use crate::common::message::command::{ MessageSerializer, PbConnRequest, PbConnResponse, PbConnStatusReq, PbConnStatusResp, }; use crate::common::message::secure::ClientHeaderSession; -use crate::common::message::MessageReader; pub async fn get_status( remote_stream: &mut S, @@ -56,28 +55,11 @@ async fn get_status_with_session None => PbConnRequest::Status(req), }; let msg = request.encode().context(EncodeStatusReqSnafu)?; - match tokio::time::timeout(timeout, session.write_initial(remote_stream, &msg)).await { - Ok(result) => result.context(WriteStatusReqSnafu)?, - Err(_) => ControlIoTimeoutSnafu { - action: "write status request", - timeout, - } - .fail()?, - } - - // get status - let mut msg_reader = session - .response_reader(remote_stream) - .context(CreateHeaderToolSnafu { action: "reader" })?; - let msg = match tokio::time::timeout(timeout, msg_reader.read_msg()).await { - Ok(result) => result.context(ReadStatusRespSnafu)?, - Err(_) => ControlIoTimeoutSnafu { - action: "read status response", - timeout, - } - .fail()?, - }; - let resp = PbConnResponse::decode(msg).context(DecodeStatusRespSnafu)?; + let response = session + .exchange(remote_stream, &msg, timeout) + .await + .context(WriteStatusReqSnafu)?; + let resp = PbConnResponse::decode(&response).context(DecodeStatusRespSnafu)?; match resp { PbConnResponse::Status(status) => Ok(status), PbConnResponse::Error(error) => StatusRespNotMatchSnafu { diff --git a/src/local/client/stream.rs b/src/local/client/stream.rs index 03062b0..539aff4 100644 --- a/src/local/client/stream.rs +++ b/src/local/client/stream.rs @@ -6,16 +6,14 @@ use tokio::net::TcpStream; use tracing::{info_span, instrument}; use super::error::{ - ConnectRemoteStreamSnafu, ControlIoTimeoutSnafu, DecodeSubcribeRespSnafu, - EncodeSubcribeReqSnafu, ReadSubcribeRespSnafu, Result, SubcribeRespNotMatchSnafu, - WriteSubcribeReqSnafu, + ConnectRemoteStreamSnafu, DecodeSubcribeRespSnafu, EncodeSubcribeReqSnafu, Result, + SubcribeRespNotMatchSnafu, WriteSubcribeReqSnafu, }; use crate::common::checksum::Credential; use crate::common::config::control_io_timeout; use crate::common::message::command::{MessageSerializer, PbConnRequest, PbConnResponse}; use crate::common::message::forward::StreamForward; use crate::common::message::secure::ClientHeaderSession; -use crate::common::message::MessageReader; use crate::local::client::error::CreateHeaderToolSnafu; use crate::snafu_error_handle; use uni_stream::addr::{each_addr, ToSocketAddrs}; @@ -61,27 +59,11 @@ pub async fn handle_local_stream< let msg = request.encode().context(EncodeSubcribeReqSnafu)?; let session = ClientHeaderSession::new_v2(&credential) .context(CreateHeaderToolSnafu { action: "session" })?; - match tokio::time::timeout(timeout, session.write_initial(&mut remote_stream, &msg)).await { - Ok(result) => result.context(WriteSubcribeReqSnafu)?, - Err(_) => ControlIoTimeoutSnafu { - action: "write subcribe request", - timeout, - } - .fail()?, - } - // handle response - let mut msg_reader = session - .response_reader(&mut remote_stream) - .context(CreateHeaderToolSnafu { action: "reader" })?; - let msg = match tokio::time::timeout(timeout, msg_reader.read_msg()).await { - Ok(result) => result.context(ReadSubcribeRespSnafu)?, - Err(_) => ControlIoTimeoutSnafu { - action: "read subcribe response", - timeout, - } - .fail()?, - }; - let resp = PbConnResponse::decode(msg).context(DecodeSubcribeRespSnafu)?; + let response = session + .exchange(&mut remote_stream, &msg, timeout) + .await + .context(WriteSubcribeReqSnafu)?; + let resp = PbConnResponse::decode(&response).context(DecodeSubcribeRespSnafu)?; match resp { PbConnResponse::Subcribe { codec_key, diff --git a/src/local/server/mod.rs b/src/local/server/mod.rs index 9e4a306..13ca2c3 100644 --- a/src/local/server/mod.rs +++ b/src/local/server/mod.rs @@ -160,16 +160,6 @@ impl Debug for ServerCliRunConfig { } } -/// Where a stream request should connect, and how. -#[derive(Clone, Copy, Debug)] -struct StreamTarget { - local_addr: A, - remote_addr: A, - keep_alive: bool, - namespace: Option, - credential: Credential, -} - fn duration_to_millis(duration: Duration) -> u64 { duration.as_millis().min(u128::from(u64::MAX)) as u64 } @@ -307,6 +297,22 @@ pub async fn run_server_side_cli_with_callback( .await; } +async fn resolve_registration_credential(pinned: Option) -> Credential { + if let Some(credential) = pinned { + return credential; + } + let mut retry_backoff = RetryBackoff::default(); + loop { + match get_process_credential() { + Ok(credential) => return credential, + Err(error) => { + tracing::error!("load registration credential failed: {error}"); + tokio::time::sleep(retry_backoff.next_delay()).await; + } + } + } +} + async fn run_server_side_cli_pool( local_addr: A, remote_addr: A, @@ -333,7 +339,8 @@ async fn run_server_side_cli_pool( return; } }; - let pool_size = control_conn_pool_size(); + let credential = resolve_registration_credential(pinned_credential).await; + let pool_size = control_conn_pool_size().max(1); tracing::info!( event = "local_server_control_pool_starting", key = %key, @@ -341,33 +348,27 @@ async fn run_server_side_cli_pool( "starting local server control connection pool" ); let mut workers = JoinSet::new(); - if pool_size > 1 { - for worker_index in 1..pool_size { - let worker_key = key.clone(); - workers.spawn(async move { - run_server_side_cli_worker_with_credential::( - local_addr, - remote_addr, - worker_key, - options, - None, - worker_index, - pinned_credential, - ) - .await; - }); - } + let mut status_callback = status_callback; + for worker_index in 0..pool_size { + let worker_key = key.clone(); + let callback = if worker_index == 0 { + status_callback.take() + } else { + None + }; + workers.spawn(async move { + run_server_side_cli_worker::( + local_addr, + remote_addr, + worker_key, + options, + callback, + worker_index, + credential, + ) + .await; + }); } - run_server_side_cli_worker_with_credential::( - local_addr, - remote_addr, - key, - options, - status_callback, - 0, - pinned_credential, - ) - .await; while let Some(result) = workers.join_next().await { if let Err(e) = result { tracing::warn!( @@ -379,32 +380,20 @@ async fn run_server_side_cli_pool( } } -async fn run_server_side_cli_worker_with_credential( +async fn run_server_side_cli_worker( local_addr: A, remote_addr: A, key: Arc, options: ServerTunnelOptions, status_callback: Option, worker_index: usize, - pinned_credential: Option, + credential: Credential, ) where LocalStream: StreamProvider + Send + 'static, LocalStream::Item: StreamForward, A: ToSocketAddrs + Debug + Copy + Send + 'static, { let mut retry_backoff = RetryBackoff::default(); - let credential = match pinned_credential { - Some(credential) => credential, - None => loop { - match get_process_credential() { - Ok(credential) => break credential, - Err(error) => { - tracing::error!("load registration credential failed: {error}"); - tokio::time::sleep(retry_backoff.next_delay()).await; - } - } - }, - }; let run_config = ServerCliRunConfig { local_addr, remote_addr, @@ -728,7 +717,7 @@ where snafu_error_get_or_continue!( handle_request::( msg, - StreamTarget { + StreamConnect { local_addr, remote_addr, keep_alive, @@ -896,7 +885,7 @@ async fn handle_request< A: ToSocketAddrs + Debug + Copy + Clone + Send + 'static, >( msg: &[u8], - target: StreamTarget, + target: StreamConnect, key: Arc, conn_id: u32, write_tx: &tokio::sync::mpsc::UnboundedSender, @@ -931,19 +920,8 @@ where let key = key.clone(); tokio::spawn(async move { snafu_error_handle!( - handle_stream::( - key, - client_id, - server_generation, - StreamConnect { - local_addr: target.local_addr, - remote_addr: target.remote_addr, - keep_alive: target.keep_alive, - namespace: target.namespace, - credential: target.credential, - }, - ) - .await + handle_stream::(key, client_id, server_generation, target,) + .await ) }); } diff --git a/src/local/server/stream.rs b/src/local/server/stream.rs index 5224310..e843af5 100644 --- a/src/local/server/stream.rs +++ b/src/local/server/stream.rs @@ -7,20 +7,20 @@ use tracing::info_span; use super::error::{ ConnectLocalStreamSnafu, ConnectRemoteStreamSnafu, ControlIoTimeoutSnafu, - DecodePbConnStreamRespSnafu, EncodePbConnStreamReqSnafu, PbConnStreamRespNotMatchSnafu, - ReadPbConnStreamRespSnafu, Result, WritePbConnStreamReqSnafu, + DecodePbConnStreamRespSnafu, EncodePbConnStreamReqSnafu, PbConnStreamRespNotMatchSnafu, Result, + WritePbConnStreamReqSnafu, }; use crate::common::checksum::Credential; use crate::common::config::control_io_timeout; use crate::common::message::command::{MessageSerializer, PbConnRequest, PbConnResponse}; use crate::common::message::forward::StreamForward; use crate::common::message::secure::ClientHeaderSession; -use crate::common::message::MessageReader; use crate::local::server::error::CreateHeaderToolSnafu; use crate::snafu_error_handle; use uni_stream::addr::{each_addr, ToSocketAddrs}; use uni_stream::stream::{set_tcp_keep_alive, set_tcp_nodelay, StreamProvider, StreamSplit}; +#[derive(Clone, Copy, Debug)] pub struct StreamConnect { pub local_addr: A, pub remote_addr: A, @@ -93,26 +93,11 @@ where let codec_key = { let session = ClientHeaderSession::new_v2(&credential) .context(CreateHeaderToolSnafu { action: "session" })?; - match tokio::time::timeout(timeout, session.write_initial(&mut remote_stream, &msg)).await { - Ok(result) => result.context(WritePbConnStreamReqSnafu)?, - Err(_) => ControlIoTimeoutSnafu { - action: "write pb conn stream request", - timeout, - } - .fail()?, - } - let mut msg_reader = session - .response_reader(&mut remote_stream) - .context(CreateHeaderToolSnafu { action: "reader" })?; - let msg = match tokio::time::timeout(timeout, msg_reader.read_msg()).await { - Ok(result) => result.context(ReadPbConnStreamRespSnafu)?, - Err(_) => ControlIoTimeoutSnafu { - action: "read pb conn stream response", - timeout, - } - .fail()?, - }; - let resp = PbConnResponse::decode(msg).context(DecodePbConnStreamRespSnafu)?; + let response = session + .exchange(&mut remote_stream, &msg, timeout) + .await + .context(WritePbConnStreamReqSnafu)?; + let resp = PbConnResponse::decode(&response).context(DecodePbConnStreamRespSnafu)?; match resp { PbConnResponse::Stream { codec_key } => codec_key, PbConnResponse::Error(error) => PbConnStreamRespNotMatchSnafu { diff --git a/src/pb_server/error.rs b/src/pb_server/error.rs index 552c68e..07ac55f 100644 --- a/src/pb_server/error.rs +++ b/src/pb_server/error.rs @@ -10,12 +10,6 @@ use crate::common::{self}; pub enum Error { #[snafu(display("administrator operation failed: {detail}"))] AdminOperation { detail: String }, - /// server task center error - #[snafu(display("read pb conn init request with `conn_id:{conn_id}`"))] - TaskCenterReadInitRequest { - conn_id: RemoteConnId, - source: common::error::Error, - }, #[snafu(display( "timed out reading pb conn init request with `conn_id:{conn_id}` after {timeout:?}" ))] @@ -23,11 +17,6 @@ pub enum Error { conn_id: RemoteConnId, timeout: Duration, }, - #[snafu(display("decode pb conn init request with `conn_id:{conn_id}`"))] - TaskCenterDecodeInitRequest { - conn_id: RemoteConnId, - source: common::error::Error, - }, #[snafu(display("send listener task error, type:{source:?} detail:{source}"))] TaskCenterSendListener { source: kanal::SendError<()> }, @@ -178,15 +167,6 @@ pub enum Error { conn_id: RemoteConnId, source: common::error::Error, }, - #[snafu(display( - "send deregister server task error with `key:{key}` `conn_id:{conn_id}`, \ - type:{source:?} detail:{source}" - ))] - ServerConnSendDeregisterServer { - key: Arc, - conn_id: RemoteConnId, - source: kanal::SendTimeoutError<()>, - }, #[snafu(display( "send register task error with `key:{key}` `conn_id:{conn_id}`, \ type:{source:?} detail:{source}" @@ -205,16 +185,6 @@ pub enum Error { }, #[snafu(display("client data stream credential is inactive: {detail}"))] ClientConnAuthInactive { detail: String }, - #[snafu(display( - "send deregister client task error with `key:{key}` `server:{server_id:?}` <-> \ - `client:{client_id}`, type:{source:?} detail:{source}" - ))] - ClientConnSendDeregisterClient { - key: Arc, - server_id: Option, - client_id: RemoteConnId, - source: kanal::SendTimeoutError<()>, - }, #[snafu(display( "send subcribe task error with `key:{key}` `conn_id:{conn_id}`, type:{source:?} \ detail:{source}" @@ -325,14 +295,6 @@ pub enum Error { StatusEncodeResp { source: common::error::Error }, #[snafu(display("write status response error"))] StatusWriteResp { source: common::error::Error }, - #[snafu(display( - "send deregister request error with `conn_id:{conn_id}`, type:{source:?} \ - detail:{source}" - ))] - StatusSendDeregister { - conn_id: RemoteConnId, - source: kanal::SendTimeoutError<()>, - }, #[snafu(display("Server listen error"))] ServerListen { source: std::io::Error }, } diff --git a/src/pb_server/mod.rs b/src/pb_server/mod.rs index c80004a..14a0098 100644 --- a/src/pb_server/mod.rs +++ b/src/pb_server/mod.rs @@ -29,8 +29,7 @@ use tracing::instrument; use self::admin::handle_admin_request; use self::client::handle_client_conn; use self::error::{ - TaskCenterDecodeInitRequestSnafu, TaskCenterInitRequestTimeoutSnafu, - TaskCenterReadInitRequestSnafu, TaskCenterSendListenerSnafu, TaskCenterSendStatusRespSnafu, + TaskCenterInitRequestTimeoutSnafu, TaskCenterSendListenerSnafu, TaskCenterSendStatusRespSnafu, TaskCenterSendStreamRespToManagerSnafu, TaskCenterSetKeepAliveSnafu, }; use self::server::{handle_server_conn, ServerRegistration}; @@ -45,7 +44,7 @@ use crate::common::message::command::{ PbServiceConnStatus, }; use crate::common::message::secure::{HeaderProtocol, ServerHeaderSession, ServerSecurity}; -use crate::common::message::{get_header_msg_reader, MessageReader, MessageWriter}; +use crate::common::message::MessageWriter; use crate::pb_server::error::{ ServerListenSnafu, TaskCenterClientSendStreamSnafu, TaskCenterSendRegisterRespSnafu, TaskCenterSendStreamRespToClientSnafu, TaskCenterSendSubcribeRespSnafu, @@ -383,16 +382,3 @@ use connection::{ decrement_namespace_stream_count, handle_conn, handle_listener, release_namespace_rate_limit_if_idle, split_scoped_service_key, }; -pub async fn get_init_request( - conn: &mut TcpStream, - conn_id: RemoteConnId, -) -> Result { - let mut reader = - get_header_msg_reader(conn).context(TaskCenterReadInitRequestSnafu { conn_id })?; - let timeout = control_io_timeout(); - let msg = match tokio::time::timeout(timeout, reader.read_msg()).await { - Ok(result) => result.context(TaskCenterReadInitRequestSnafu { conn_id })?, - Err(_) => TaskCenterInitRequestTimeoutSnafu { conn_id, timeout }.fail()?, - }; - PbConnRequest::decode(msg).context(TaskCenterDecodeInitRequestSnafu { conn_id }) -} diff --git a/ui/native/pb_mapper_ffi/src/state.rs b/ui/native/pb_mapper_ffi/src/state.rs index bae25e1..dd1bed2 100644 --- a/ui/native/pb_mapper_ffi/src/state.rs +++ b/ui/native/pb_mapper_ffi/src/state.rs @@ -429,6 +429,27 @@ struct PinnedTunnel { endpoint: SocketAddr, } +struct TunnelRuntime { + handle: JoinHandle<()>, + pin: PinnedTunnel, +} + +fn abort_runtime(map: &mut HashMap, key: &str) -> bool { + match map.remove(key) { + Some(runtime) => { + runtime.handle.abort(); + true + } + None => false, + } +} + +fn replace_runtime(map: &mut HashMap, key: &str, runtime: TunnelRuntime) { + if let Some(previous) = map.insert(key.to_string(), runtime) { + previous.handle.abort(); + } +} + /// Everything [`PbMapperState::finish_register`] needs once the slow work is done. struct RegisterCommit { service_key: String, @@ -461,10 +482,8 @@ pub struct PbMapperState { server_start_time: Option, registered_services: Arc>>, active_connections: Arc>>, - service_handles: HashMap>, - client_handles: HashMap>, - service_tunnels: HashMap, - client_tunnels: HashMap, + service_runtime: HashMap, + client_runtime: HashMap, config: AppConfig, config_dir: PathBuf, app_directory_path: Option, diff --git a/ui/native/pb_mapper_ffi/src/state/configuration.rs b/ui/native/pb_mapper_ffi/src/state/configuration.rs index d78fee0..ced14e7 100644 --- a/ui/native/pb_mapper_ffi/src/state/configuration.rs +++ b/ui/native/pb_mapper_ffi/src/state/configuration.rs @@ -45,7 +45,7 @@ impl PbMapperState { uptime_seconds: 0, })); - let temp_state = Self { + let mut state = Self { server_handle: None, server_auth: None, server_shutdown_token: None, @@ -53,14 +53,12 @@ impl PbMapperState { server_start_time: None, registered_services: Arc::new(RwLock::new(HashMap::new())), active_connections: Arc::new(RwLock::new(HashMap::new())), - service_handles: HashMap::new(), - client_handles: HashMap::new(), - service_tunnels: HashMap::new(), - client_tunnels: HashMap::new(), + service_runtime: HashMap::new(), + client_runtime: HashMap::new(), config: AppConfig::default(), - config_dir: config_dir.clone(), - app_directory_path: app_directory_path.clone(), - local_server_status_cache: local_server_status_cache.clone(), + config_dir, + app_directory_path, + local_server_status_cache, local_server_status_last_update: Arc::new(RwLock::new(None)), local_server_status_refreshing: Arc::new(AtomicBool::new(false)), service_status_cache: Arc::new(RwLock::new(HashMap::new())), @@ -70,44 +68,16 @@ impl PbMapperState { registering: Arc::new(StdMutex::new(HashSet::new())), connecting: Arc::new(StdMutex::new(HashSet::new())), }; - - let config = temp_state.load_config().unwrap_or_else(|e| { + state.config = state.load_config().unwrap_or_else(|e| { tracing::warn!("Could not load config: {}, using defaults", e); AppConfig::default() }); - tracing::info!( "Loaded configuration: server_address={}, keep_alive={}, msg_header_key_set={}", - config.server_address, - config.keep_alive_enabled, - !config.msg_header_key.is_empty() + state.config.server_address, + state.config.keep_alive_enabled, + !state.config.msg_header_key.is_empty() ); - - let state = Self { - server_handle: None, - server_auth: None, - server_shutdown_token: None, - server_status_sender: None, - server_start_time: None, - registered_services: Arc::new(RwLock::new(HashMap::new())), - active_connections: Arc::new(RwLock::new(HashMap::new())), - service_handles: HashMap::new(), - client_handles: HashMap::new(), - service_tunnels: HashMap::new(), - client_tunnels: HashMap::new(), - config, - config_dir, - app_directory_path, - local_server_status_cache, - local_server_status_last_update: Arc::new(RwLock::new(None)), - local_server_status_refreshing: Arc::new(AtomicBool::new(false)), - service_status_cache: Arc::new(RwLock::new(HashMap::new())), - client_status_cache: Arc::new(RwLock::new(HashMap::new())), - service_status_refreshing: Arc::new(RwLock::new(HashSet::new())), - client_status_refreshing: Arc::new(RwLock::new(HashSet::new())), - registering: Arc::new(StdMutex::new(HashSet::new())), - connecting: Arc::new(StdMutex::new(HashSet::new())), - }; if let Err(e) = state.apply_msg_header_key_env() { tracing::error!("Failed to apply MSG_HEADER_KEY during init: {}", e); } diff --git a/ui/native/pb_mapper_ffi/src/state/runtime.rs b/ui/native/pb_mapper_ffi/src/state/runtime.rs index b803fa9..ba326f2 100644 --- a/ui/native/pb_mapper_ffi/src/state/runtime.rs +++ b/ui/native/pb_mapper_ffi/src/state/runtime.rs @@ -133,15 +133,12 @@ impl PbMapperState { *last_update = Some(Instant::now()); } - for (_, handle) in self.service_handles.drain() { - handle.abort(); + for (_, runtime) in self.service_runtime.drain() { + runtime.handle.abort(); } - self.service_tunnels.clear(); - - for (_, handle) in self.client_handles.drain() { - handle.abort(); + for (_, runtime) in self.client_runtime.drain() { + runtime.handle.abort(); } - self.client_tunnels.clear(); self.registered_services.write().await.clear(); self.active_connections.write().await.clear(); @@ -165,15 +162,10 @@ impl PbMapperState { credential, } = commit; - self.service_tunnels.remove(&service_key); - if let Some(previous) = self.service_handles.remove(&service_key) { + if abort_runtime(&mut self.service_runtime, &service_key) { tracing::warn!( "Service '{service_key}' is already registered, replacing existing handle" ); - // Dropping a `JoinHandle` does not stop the task. Without this the - // replaced tunnel kept running and retrying, with nothing left - // holding a handle able to abort it. - previous.abort(); } tracing::info!( @@ -204,52 +196,32 @@ impl PbMapperState { ); }); - self.service_tunnels.insert( - service_key.clone(), - PinnedTunnel { - credential, - endpoint: remote_sock_addr, + let handle = spawn_register_tunnel( + &protocol, + local_sock_addr, + remote_sock_addr, + key_clone, + ServerTunnelOptions { + need_codec: enable_encryption, + is_datagram: !protocol.eq_ignore_ascii_case("TCP"), + keep_alive: enable_keep_alive, + namespace: None, + force_namespace: false, }, + callback, + credential, ); - let handle = if protocol.to_uppercase() == "TCP" { - tokio::spawn(async move { - let _ = run_server_side_cli_with_pinned_credential::( - local_sock_addr, - remote_sock_addr, - key_clone.into(), - ServerTunnelOptions { - need_codec: enable_encryption, - is_datagram: false, - keep_alive: enable_keep_alive, - namespace: None, - force_namespace: false, - }, - Some(callback), - credential, - ) - .await; - }) - } else { - tokio::spawn(async move { - let _ = run_server_side_cli_with_pinned_credential::( - local_sock_addr, - remote_sock_addr, - key_clone.into(), - ServerTunnelOptions { - need_codec: enable_encryption, - is_datagram: true, - keep_alive: enable_keep_alive, - namespace: None, - force_namespace: false, - }, - Some(callback), + replace_runtime( + &mut self.service_runtime, + &service_key, + TunnelRuntime { + handle, + pin: PinnedTunnel { credential, - ) - .await; - }) - }; - - self.service_handles.insert(service_key.clone(), handle); + endpoint: remote_sock_addr, + }, + }, + ); { let mut cache = self.service_status_cache.write().await; @@ -281,10 +253,7 @@ impl PbMapperState { } pub async fn unregister_service(&mut self, service_key: String) -> Result<(), CtlError> { - self.service_tunnels.remove(&service_key); - if let Some(handle) = self.service_handles.remove(&service_key) { - handle.abort(); - } + abort_runtime(&mut self.service_runtime, &service_key); if self .registered_services @@ -306,10 +275,7 @@ impl PbMapperState { &mut self, service_key: String, ) -> Result<(), CtlError> { - self.service_tunnels.remove(&service_key); - if let Some(handle) = self.service_handles.remove(&service_key) { - handle.abort(); - } + abort_runtime(&mut self.service_runtime, &service_key); self.registered_services.write().await.remove(&service_key); @@ -327,14 +293,10 @@ impl PbMapperState { credential, } = commit; - self.client_tunnels.remove(&service_key); - if let Some(previous) = self.client_handles.remove(&service_key) { + if abort_runtime(&mut self.client_runtime, &service_key) { tracing::warn!( "Client for service '{service_key}' is already connected, replacing handle" ); - // As in `finish_register`: dropping the handle leaves the old - // client's retry loop running with nothing able to stop it. - previous.abort(); } let protocol_upper = protocol.to_uppercase(); @@ -356,40 +318,26 @@ impl PbMapperState { }) }; - self.client_tunnels.insert( - service_key.clone(), - PinnedTunnel { - credential, - endpoint: remote_sock_addr, - }, + let handle = spawn_connect_tunnel( + &protocol_upper, + local_sock_addr, + remote_sock_addr, + key_clone, + enable_keep_alive, + status_callback, + credential, ); - let handle = if protocol_upper == "TCP" { - tokio::spawn(async move { - run_client_side_cli_with_pinned_credential::( - local_sock_addr, - remote_sock_addr, - key_clone.into(), - enable_keep_alive, - Some(status_callback), - credential, - ) - .await; - }) - } else { - tokio::spawn(async move { - run_client_side_cli_with_pinned_credential::( - local_sock_addr, - remote_sock_addr, - key_clone.into(), - enable_keep_alive, - Some(status_callback), + replace_runtime( + &mut self.client_runtime, + &service_key, + TunnelRuntime { + handle, + pin: PinnedTunnel { credential, - ) - .await; - }) - }; - - self.client_handles.insert(service_key.clone(), handle); + endpoint: remote_sock_addr, + }, + }, + ); { let mut cache = self.client_status_cache.write().await; @@ -444,14 +392,7 @@ impl PbMapperState { pub async fn disconnect_service(&mut self, service_key: String) -> Result<(), CtlError> { // Aborting the task is the part that matters: it is what stops the // retry loop still dialling in the background. - self.client_tunnels.remove(&service_key); - let aborted = match self.client_handles.remove(&service_key) { - Some(handle) => { - handle.abort(); - true - } - None => false, - }; + let aborted = abort_runtime(&mut self.client_runtime, &service_key); let was_listed = self .active_connections @@ -478,13 +419,82 @@ impl PbMapperState { &mut self, service_key: String, ) -> Result<(), CtlError> { - self.client_tunnels.remove(&service_key); - if let Some(handle) = self.client_handles.remove(&service_key) { - handle.abort(); - } + abort_runtime(&mut self.client_runtime, &service_key); self.active_connections.write().await.remove(&service_key); self.delete_client_config(&service_key) } } + +fn spawn_register_tunnel( + protocol: &str, + local_sock_addr: SocketAddr, + remote_sock_addr: SocketAddr, + key: String, + options: ServerTunnelOptions, + callback: StatusCallback, + credential: Credential, +) -> JoinHandle<()> { + if protocol.eq_ignore_ascii_case("TCP") { + tokio::spawn(async move { + let _ = run_server_side_cli_with_pinned_credential::( + local_sock_addr, + remote_sock_addr, + key.into(), + options, + Some(callback), + credential, + ) + .await; + }) + } else { + tokio::spawn(async move { + let _ = run_server_side_cli_with_pinned_credential::( + local_sock_addr, + remote_sock_addr, + key.into(), + options, + Some(callback), + credential, + ) + .await; + }) + } +} + +fn spawn_connect_tunnel( + protocol: &str, + local_sock_addr: SocketAddr, + remote_sock_addr: SocketAddr, + key: String, + enable_keep_alive: bool, + callback: ClientStatusCallback, + credential: Credential, +) -> JoinHandle<()> { + if protocol.eq_ignore_ascii_case("TCP") { + tokio::spawn(async move { + run_client_side_cli_with_pinned_credential::( + local_sock_addr, + remote_sock_addr, + key.into(), + enable_keep_alive, + Some(callback), + credential, + ) + .await; + }) + } else { + tokio::spawn(async move { + run_client_side_cli_with_pinned_credential::( + local_sock_addr, + remote_sock_addr, + key.into(), + enable_keep_alive, + Some(callback), + credential, + ) + .await; + }) + } +} diff --git a/ui/native/pb_mapper_ffi/src/state/status.rs b/ui/native/pb_mapper_ffi/src/state/status.rs index 9c4f38a..5927111 100644 --- a/ui/native/pb_mapper_ffi/src/state/status.rs +++ b/ui/native/pb_mapper_ffi/src/state/status.rs @@ -59,7 +59,7 @@ impl PbMapperState { sorted_configs.sort_by_key(|config| config.created_at); for config in sorted_configs { - let (status, message) = self.calculate_service_status(&config.service_key).await; + let (status, message) = self.get_cached_service_status(&config.service_key).await; services.push(ServiceConfigInfo { service_key: config.service_key.clone(), @@ -85,7 +85,7 @@ impl PbMapperState { } pub async fn get_service_status(&self, service_key: String) -> ServiceStatusResponse { - let (status, message) = self.calculate_service_status(&service_key).await; + let (status, message) = self.get_cached_service_status(&service_key).await; ServiceStatusResponse { service_key, status, @@ -98,7 +98,7 @@ impl PbMapperState { let mut client_infos = Vec::new(); for (service_key, config) in store.clients.iter() { - let (status, status_message) = self.calculate_client_status(service_key).await; + let (status, status_message) = self.get_cached_client_status(service_key).await; client_infos.push(ClientConfigInfo { service_key: config.service_key.clone(), @@ -125,7 +125,7 @@ impl PbMapperState { } pub async fn get_client_status(&self, service_key: String) -> ClientStatusResponse { - let (status, message) = self.calculate_client_status(&service_key).await; + let (status, message) = self.get_cached_client_status(&service_key).await; ClientStatusResponse { service_key, status, @@ -285,79 +285,68 @@ impl PbMapperState { // Cache service status to avoid blocking UI with network checks on every paint. async fn get_cached_service_status(&self, service_key: &str) -> (String, String) { - if let Some(handle) = self.service_handles.get(service_key) { - if handle.is_finished() { - return ( - "failed".to_string(), - "Service connection terminated".to_string(), - ); - } - - let cached = { - let cache = self.service_status_cache.read().await; - cache.get(service_key).cloned() - }; - - let should_refresh = cached - .as_ref() - .map(|entry| entry.updated_at.elapsed() > STATUS_CACHE_TTL) - .unwrap_or(true); - - if should_refresh { - self.schedule_service_status_refresh(service_key).await; - } - - if let Some(entry) = cached { - return (entry.status, entry.message); - } - + let Some(runtime) = self.service_runtime.get(service_key) else { return ( - "retrying".to_string(), - "Checking service status...".to_string(), + "stopped".to_string(), + "Service is not registered".to_string(), + ); + }; + if runtime.handle.is_finished() { + return ( + "failed".to_string(), + "Service connection terminated".to_string(), ); } - - ( - "stopped".to_string(), - "Service is not registered".to_string(), - ) + let cached = { + let cache = self.service_status_cache.read().await; + cache.get(service_key).cloned() + }; + if cached + .as_ref() + .map(|entry| entry.updated_at.elapsed() > STATUS_CACHE_TTL) + .unwrap_or(true) + { + self.schedule_service_status_refresh(service_key).await; + } + cached + .map(|entry| (entry.status, entry.message)) + .unwrap_or_else(|| { + ( + "retrying".to_string(), + "Checking service status...".to_string(), + ) + }) } - // Cache client status to avoid blocking UI with network checks on every paint. async fn get_cached_client_status(&self, service_key: &str) -> (String, String) { - if let Some(handle) = self.client_handles.get(service_key) { - if handle.is_finished() { - return ( - "failed".to_string(), - "Client connection terminated".to_string(), - ); - } - - let cached = { - let cache = self.client_status_cache.read().await; - cache.get(service_key).cloned() - }; - - let should_refresh = cached - .as_ref() - .map(|entry| entry.updated_at.elapsed() > STATUS_CACHE_TTL) - .unwrap_or(true); - - if should_refresh { - self.schedule_client_status_refresh(service_key).await; - } - - if let Some(entry) = cached { - return (entry.status, entry.message); - } - + let Some(runtime) = self.client_runtime.get(service_key) else { + return ("stopped".to_string(), "Client is not connected".to_string()); + }; + if runtime.handle.is_finished() { return ( - "retrying".to_string(), - "Checking client status...".to_string(), + "failed".to_string(), + "Client connection terminated".to_string(), ); } - - ("stopped".to_string(), "Client is not connected".to_string()) + let cached = { + let cache = self.client_status_cache.read().await; + cache.get(service_key).cloned() + }; + if cached + .as_ref() + .map(|entry| entry.updated_at.elapsed() > STATUS_CACHE_TTL) + .unwrap_or(true) + { + self.schedule_client_status_refresh(service_key).await; + } + cached + .map(|entry| (entry.status, entry.message)) + .unwrap_or_else(|| { + ( + "retrying".to_string(), + "Checking client status...".to_string(), + ) + }) } pub(super) async fn schedule_service_status_refresh(&self, service_key: &str) { @@ -369,7 +358,10 @@ impl PbMapperState { refreshing.insert(service_key.to_string()); } - let tunnel = self.service_tunnels.get(service_key); + let tunnel = self + .service_runtime + .get(service_key) + .map(|runtime| runtime.pin); let server_addr = tunnel .map(|tunnel| tunnel.endpoint.to_string()) .unwrap_or_else(|| self.config.server_address.clone()); @@ -436,7 +428,10 @@ impl PbMapperState { refreshing.insert(service_key.to_string()); } - let tunnel = self.client_tunnels.get(service_key); + let tunnel = self + .client_runtime + .get(service_key) + .map(|runtime| runtime.pin); let server_addr = tunnel .map(|tunnel| tunnel.endpoint.to_string()) .unwrap_or_else(|| self.config.server_address.clone()); @@ -493,12 +488,4 @@ impl PbMapperState { refreshing.remove(&key); }); } - - async fn calculate_service_status(&self, service_key: &str) -> (String, String) { - self.get_cached_service_status(service_key).await - } - - async fn calculate_client_status(&self, service_key: &str) -> (String, String) { - self.get_cached_client_status(service_key).await - } } From 6af7d24da7d184355856665e7ec5867193e4487b Mon Sep 17 00:00:00 2001 From: ackingliu Date: Thu, 20 Aug 2026 20:33:11 +0800 Subject: [PATCH 59/74] Gate the Linux auth state dir helper by platform linux_default_auth_state_dir is only compiled off Windows/macOS, but its re-export and its test were gated on cfg(test) alone, so cargo test failed to build on macOS. Co-authored-by: Cursor --- src/common/auth.rs | 4 +++- src/common/auth/config.rs | 1 + src/common/auth/tests.rs | 1 + 3 files changed, 5 insertions(+), 1 deletion(-) diff --git a/src/common/auth.rs b/src/common/auth.rs index 0545153..bd51962 100644 --- a/src/common/auth.rs +++ b/src/common/auth.rs @@ -487,10 +487,12 @@ enum AuthCommand { mod config; pub use config::default_auth_state_dir; +#[cfg(all(test, not(any(windows, target_os = "macos"))))] +pub(crate) use config::linux_default_auth_state_dir; #[cfg(test)] pub(in crate::common::auth) use config::parse_legacy_protocol_policy; #[cfg(test)] -pub(crate) use config::{linux_default_auth_state_dir, platform_default_auth_state_dir}; +pub(crate) use config::platform_default_auth_state_dir; #[cfg(all(test, not(any(windows, target_os = "macos"))))] pub(in crate::common::auth) use config::{linux_system_auth_dir_usable, unix_effective_uid}; mod keys; diff --git a/src/common/auth/config.rs b/src/common/auth/config.rs index 32e0bba..b2b3b47 100644 --- a/src/common/auth/config.rs +++ b/src/common/auth/config.rs @@ -41,6 +41,7 @@ pub(crate) fn platform_default_auth_state_dir() -> PathBuf { } } +#[cfg(not(any(windows, target_os = "macos")))] pub(crate) fn linux_default_auth_state_dir( euid: u32, system_dir_usable: bool, diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index 77510ba..ea1f3a8 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -584,6 +584,7 @@ fn sync_parent_directory_succeeds_for_a_local_file() { let _ = std::fs::remove_dir_all(state_dir); } +#[cfg(not(any(windows, target_os = "macos")))] #[test] fn linux_default_auth_state_dir_prefers_user_data_when_system_dir_is_unusable() { assert_eq!( From bc2d708230fb0e3efdc073b91890a00cb88d79e3 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Fri, 21 Aug 2026 03:52:37 +0800 Subject: [PATCH 60/74] Drive credential cleanup from timing-wheel callbacks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The timing wheel scanned every bucket on a clock jump and walked every second in between, so elapsed time cost work even with nothing to do. Its entries also held `Weak` leases, which forced a separate owner map plus a `wheel_version` counter to tell a renewed entry from a stale one, and left the actual cleanup — cancelling the lease, marking the row, queuing the tombstone, freeing the row, dropping the metadata — spread across the actor tick, the GC sweep, revoke, and the rotation wipe. Entries now own a callback that performs one phase of a teardown and returns when it next wants to run. A key's whole life is one entry: phase one cancels the lease and marks the row dead, phase two frees the row and forgets the key. Dropping an entry runs whatever phases it has left, so expiry, revoke, GC, root rotation, and shutdown all reach the same code, and no call site performs cleanup. That removes the tombstone queue, the owner map, `wheel_version`, and `clear_retained_high_slot_entries`. Advancing is now one bucket per level that turns over, matching the one-second tick, with a single-pass drain reserved for a jump past the longest schedulable delay so a corrected hardware clock cannot spin for hours. Also give the identity types names: `KeyId`, `Generation`, and `SlotIndex` replace interchangeable integers, so `KeyId::new(generation, slot)` cannot take its arguments in either order and a key id cannot be used as an array index. All three are `#[serde(transparent)]`, so snapshots and the admin protocol keep their plain-integer encoding. Level and slot indices narrow to `u8`, `issued_epoch` was written but never read and is gone, and the slot table's layout, generation rule, and tombstone window are documented where the types are defined. Co-Authored-By: Claude Opus 5 (1M context) --- src/common/auth.rs | 274 +++++++++++--- src/common/auth/actor/epoch.rs | 47 +-- src/common/auth/actor/lifecycle.rs | 338 +++++++---------- src/common/auth/actor/mod.rs | 216 ++--------- src/common/auth/ids.rs | 129 +++++++ src/common/auth/keys.rs | 14 +- src/common/auth/leases.rs | 280 ++++++++++++++ src/common/auth/persistence/mod.rs | 6 +- src/common/auth/persistence/snapshot.rs | 76 ++-- src/common/auth/runtime.rs | 101 ++--- src/common/auth/tests.rs | 438 ++++++++++++++++------ src/common/auth/timing_wheel.rs | 406 ++++++++++++-------- src/common/message/secure.rs | 25 +- src/common/message/secure/first_flight.rs | 8 +- src/common/message/secure/frame.rs | 6 +- src/common/message/secure/limiter.rs | 7 +- src/common/message/secure/replay.rs | 8 +- src/common/message/secure/tests.rs | 124 ++++-- src/pb_server/admin.rs | 23 +- src/pb_server/connection.rs | 10 +- src/pb_server/mod.rs | 2 +- tests/regression.rs | 25 +- 22 files changed, 1600 insertions(+), 963 deletions(-) create mode 100644 src/common/auth/ids.rs create mode 100644 src/common/auth/leases.rs diff --git a/src/common/auth.rs b/src/common/auth.rs index bd51962..75372f4 100644 --- a/src/common/auth.rs +++ b/src/common/auth.rs @@ -1,21 +1,51 @@ //! Authentication state for protocol-v2 connections and administrator operations. //! -//! The root administrator key is never copied into a temporary credential. Temporary -//! keys are derived from `(root key, server instance id, key id)` and the hot slot table -//! stores only lifecycle metadata plus a weak lease reference. The background actor owns -//! the strong leases through a hierarchical timing wheel. +//! # How a temporary credential works +//! +//! Nothing secret is stored per key. A temporary credential is *derived* from the +//! root key, the server instance id, and the key id, so the server can verify a +//! credential it holds no copy of, and a key id is all the state a key needs: +//! +//! ```text +//! issue: root key + instance id + key id --HKDF--> credential handed to the client +//! verify: root key + instance id + key id --HKDF--> compare against what was presented +//! ``` +//! +//! Because the material is derived, invalidating every key at once is a matter of +//! changing an input: a root rotation replaces the root key, a state reset replaces +//! the instance id. Neither has to touch individual keys. +//! +//! # Where the state lives //! //! ```text -//! administrator key + instance id + key id -> derived temporary credential -//! | -//! request -> AuthContext -> Weak lease -+-> actor-owned Arc lease -> timing wheel -//! +-> cancel on expiry/revoke/reset/rotation +//! key_id = generation:slot +//! | +//! request ──> derive & compare ──> slots[slot] ── lifecycle: Free/Active/ +//! │ Expired/Revoked, expires_at +//! │ +//! Weak lease ──> Arc lease, owned by the actor's +//! ^ timing wheel — the single place +//! │ a lease's lifetime ends +//! AuthContext (also Weak) ─────────────┘ +//! ``` +//! +//! The slot table is a preallocated array indexed straight off the key id, so +//! verification costs an array index and churn does not grow memory. The +//! `SlotState` docs below cover the table's layout, why generations exist, and +//! why dead rows linger. `Leases` (`leases.rs`) owns the three structures a +//! key's lifetime spans; `timing_wheel.rs` schedules the expiries. //! -//! AuthRuntime facade -> serialized actor -> encrypted snapshot + WAL +//! # Where mutations happen +//! +//! ```text +//! AuthRuntime (facade) ──channel──> one actor ──> encrypted snapshot + WAL //! ``` //! -//! The facade/model types remain in this root module; runtime checks, actor mutations, -//! persistence, expiry scheduling, and focused tests live in their respective children. +//! Every mutation is serialized through a single actor, so a request authorized +//! before a root rotation cannot execute against the state that replaced it. +//! +//! The facade and model types stay in this root module; runtime checks, actor +//! mutations, persistence, expiry scheduling, and tests live in the children. use std::collections::{HashMap, HashSet, VecDeque}; use std::fmt; @@ -42,7 +72,9 @@ use super::checksum::{ ENV_MSG_HEADER_KEY, MACHINE_MSG_HEADER_KEY_PATH, }; -pub const ADMIN_NAMESPACE: u64 = 0; +/// The namespace administrator connections operate in. Tenant namespaces are the +/// key id that owns them, so this mirrors [`ADMIN_KEY_ID`]. +pub const ADMIN_NAMESPACE: u64 = ADMIN_KEY_ID.as_u64(); pub const DEFAULT_AUTH_STATE_DIR: &str = "/var/lib/pb-mapper/auth"; pub const DEFAULT_TEMP_KEY_CAPACITY: usize = 65_536; pub const MAX_TEMP_KEY_CAPACITY: usize = 1_048_576; @@ -50,6 +82,10 @@ pub const DEFAULT_MAX_TEMP_KEY_TTL: Duration = Duration::from_secs(30 * 24 * 60 pub const MIN_TEMP_KEY_TTL: Duration = Duration::from_secs(10); pub const MAX_TEMP_KEY_TTL: Duration = Duration::from_secs(365 * 24 * 60 * 60); const TOMBSTONE_RETENTION: Duration = Duration::from_secs(60); +/// Longest delay any scheduled cleanup can ask for, so the timing wheel can tell +/// a plausible wait from a clock correction. +const MAX_SCHEDULABLE_DELAY: Duration = + Duration::from_secs(MAX_TEMP_KEY_TTL.as_secs() + TOMBSTONE_RETENTION.as_secs()); const SNAPSHOT_COMPACTION_INTERVAL: Duration = Duration::from_secs(5 * 60); const SNAPSHOT_SCHEMA_VERSION: u16 = 1; const STATE_BLOB_MAGIC: &[u8; 5] = b"PBAS1"; @@ -120,25 +156,23 @@ const LEASE_CANCEL_ROTATED: u8 = 3; #[derive(Debug)] pub struct AuthLease { - key_id: u64, + key_id: KeyId, expires_at: AtomicU64, - wheel_version: AtomicU64, cancellation: CancellationToken, cancel_reason: AtomicU8, } impl AuthLease { - fn new(key_id: u64, expires_at: u64) -> Self { + fn new(key_id: KeyId, expires_at: u64) -> Self { Self { key_id, expires_at: AtomicU64::new(expires_at), - wheel_version: AtomicU64::new(1), cancellation: CancellationToken::new(), cancel_reason: AtomicU8::new(LEASE_CANCEL_NONE), } } - pub fn key_id(&self) -> u64 { + pub fn key_id(&self) -> KeyId { self.key_id } @@ -180,17 +214,21 @@ impl AuthLease { #[derive(Clone, Debug)] pub struct AuthContext { - pub key_id: u64, + pub key_id: KeyId, pub namespace: u64, pub is_admin: bool, lease: Weak, } impl AuthContext { - fn from_lease(key_id: u64, is_admin: bool, lease: &Arc) -> Self { + fn from_lease(key_id: KeyId, is_admin: bool, lease: &Arc) -> Self { Self { key_id, - namespace: if is_admin { ADMIN_NAMESPACE } else { key_id }, + namespace: if is_admin { + ADMIN_NAMESPACE + } else { + key_id.as_u64() + }, is_admin, lease: Arc::downgrade(lease), } @@ -282,31 +320,132 @@ fn cancelled_lease_failure(is_admin: bool, lease: &AuthLease) -> AuthFailure { } } +/// # The slot table +/// +/// A temporary key is never stored. It is *derived* on demand from +/// `(root key, instance id, key id)`, so the server can verify a credential it +/// has no copy of. That makes the key id the whole identity of a key, and a +/// key id is a slot index plus a generation counter: +/// +/// ```text +/// key_id: u64 +/// ┌───────────────────────────┬───────────────────────────┐ +/// │ generation (high 32) │ slot index (low 32) │ +/// └───────────────────────────┴───────────────────────────┘ +/// ^ bumped on reuse ^ where the row lives +/// ``` +/// +/// The slot index is a direct offset into `AuthStateInner::slots`, a +/// preallocated `Box<[SlotHot]>`. So verifying a credential is an array index, +/// not a map lookup or a scan, and the table's memory does not grow with churn: +/// +/// ```text +/// slots: [ SlotHot; max_temporary_keys ] +/// idx 0 gen 7 Active expires_at=… lease─┐ +/// idx 1 gen 0 Free │ Weak, so the actor's +/// idx 2 gen 3 Expired (tombstoned) │ timing wheel is the +/// idx 3 gen 9 Active expires_at=… lease─┴─ only strong owner +/// ``` +/// +/// ## Why the generation counter +/// +/// A freed slot is reused, so the index alone would let a *retired* credential +/// authenticate against the *new* tenant of that row. The generation bump makes +/// the old key id refer to a row that no longer exists: +/// +/// ```text +/// issue -> idx 2, gen 3 => key_id 0x0000_0003_0000_0002 +/// expire -> idx 2 retired, generation kept at 3 +/// reissue -> idx 2, gen 4 => key_id 0x0000_0004_0000_0002 +/// the old key id still names gen 3, which nothing matches +/// ``` +/// +/// This is why [`SlotHot::retire`] clears the row but preserves `generation`, +/// and why a generation is never reset — not by expiry, GC, root rotation, or a +/// full state reset. +/// +/// ## The lifecycle +/// +/// ```text +/// issue deadline passes / revoke +/// Free ─────────> Active ──────────────────────────────> Expired +/// ^ │ Revoked +/// │ └── renew: same row, later expires_at │ +/// │ │ +/// └──────────── retire, after TOMBSTONE_RETENTION ─────────┘ +/// ``` +/// +/// `Expired`/`Revoked` are tombstones, not garbage. A row lingers in that state +/// for `TOMBSTONE_RETENTION` so a client that presents a dead credential is told +/// *why* ("expired", "revoked") instead of receiving the indistinguishable +/// "unknown key" it would get from an already-recycled row. `Leases` owns that +/// delay; see `leases.rs`. +/// +/// ## Slots above capacity +/// +/// `max_temporary_keys` is configurable, so a restart can shrink the table below +/// what the persisted state used. Those rows cannot be indexed any more, but +/// their generations still have to be honoured — otherwise growing the table +/// again would reissue a key id that was already handed out. They are retained +/// out-of-line in `high_slot_generations` / `high_slot_entries`, which is why so +/// many operations check the array first and fall back to a scan of that vector. #[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] enum SlotState { + /// Never used, or retired and past its tombstone. Free, + /// A live credential. `expires_at` is authoritative. Active, + /// Dead. Retained for `TOMBSTONE_RETENTION` so the reason survives. Expired, Revoked, } +/// One row of the slot table. `generation` outlives every other field. #[derive(Debug)] struct SlotHot { - generation: u32, + generation: Generation, state: SlotState, expires_at: u64, - issued_epoch: u64, + /// `Weak`, because the actor's timing wheel holds the strong reference and is + /// the single place a lease's lifetime ends. See `timing_wheel.rs`. lease: Weak, } +impl SlotHot { + /// Whether this row still belongs to `key_id`'s generation. A row that has + /// been reissued belongs to a newer tenant and must not be touched on the + /// old one's behalf. + fn holds(&self, key_id: KeyId) -> bool { + self.generation == key_id.generation() && self.state != SlotState::Free + } + + /// Whether a garbage collection should free this row: it is already dead, or + /// it is active but past its deadline. + fn is_collectable(&self, now: u64) -> bool { + match self.state { + SlotState::Expired | SlotState::Revoked => true, + SlotState::Active => self.expires_at <= now, + SlotState::Free => false, + } + } + + /// Frees the slot for reuse while keeping its generation, so a key id that + /// has been handed out is never issued a second time. + fn retire(&mut self) { + *self = Self { + generation: self.generation, + ..Self::default() + }; + } +} + impl Default for SlotHot { fn default() -> Self { Self { - generation: 0, + generation: Generation::FIRST, state: SlotState::Free, expires_at: 0, - issued_epoch: 0, lease: Weak::new(), } } @@ -329,11 +468,27 @@ struct AuthStateInner { admin: RwLock, sync_process_credential: bool, instance_id: RwLock<[u8; INSTANCE_ID_LEN]>, + /// Preallocated, indexed directly by `key_slot(key_id)`. Documented on + /// [`SlotState`]. slots: RwLock>, - /// Generations and entries for slots above the current capacity. Kept so a - /// later capacity increase cannot reuse a discarded slot's key id. - high_slot_generations: RwLock>, + /// Rows the configured capacity no longer covers, because a restart shrank + /// the table below what the persisted state used: + /// + /// ```text + /// slots: [ 0 1 2 3 ] <- indexable + /// high: [ 4 5 ] <- generations still honoured, out-of-line + /// ``` + /// + /// Their generations must be kept so growing the table again cannot reissue + /// a key id that was already handed out, and their entries so a still-live + /// credential in that range keeps working. This is the fallback path that + /// operations take after missing in `slots`. + high_slot_generations: RwLock>, high_slot_entries: RwLock>, + /// Per-key description that no authentication check needs, kept out of the + /// hot slot row. Lives here rather than inside the actor so a key's handle + /// can drop it without the actor being involved; see `leases.rs`. + cold: RwLock>, safe_mode: AtomicBool, legacy_protocol_allowed: AtomicBool, active_legacy_connections: AtomicU64, @@ -350,6 +505,32 @@ fn recover_lock(result: std::sync::LockResult) -> T { } impl AuthStateInner { + /// Rows the configured capacity no longer covers. See the field's docs; the + /// fallback is a scan because the range is small and rarely touched. + fn high(&self) -> std::sync::RwLockReadGuard<'_, Vec> { + recover_lock(self.high_slot_entries.read()) + } + + fn high_mut(&self) -> std::sync::RwLockWriteGuard<'_, Vec> { + recover_lock(self.high_slot_entries.write()) + } + + fn slots(&self) -> std::sync::RwLockReadGuard<'_, Box<[SlotHot]>> { + recover_lock(self.slots.read()) + } + + fn slots_mut(&self) -> std::sync::RwLockWriteGuard<'_, Box<[SlotHot]>> { + recover_lock(self.slots.write()) + } + + fn cold(&self) -> std::sync::RwLockReadGuard<'_, HashMap> { + recover_lock(self.cold.read()) + } + + fn cold_mut(&self) -> std::sync::RwLockWriteGuard<'_, HashMap> { + recover_lock(self.cold.write()) + } + fn admin_key(&self) -> AesKeyType { recover_lock(self.admin.read()).key } @@ -371,7 +552,7 @@ pub struct AuthRuntime { #[derive(Clone, Debug, Serialize, Deserialize)] pub struct TemporaryKeyMetadata { - pub key_id: u64, + pub key_id: KeyId, pub state: String, pub issued_at: u64, pub expires_at: u64, @@ -436,19 +617,19 @@ enum AuthCommand { }, Show { authority: Weak, - key_id: u64, + key_id: KeyId, reveal: bool, response: oneshot::Sender>, }, Renew { authority: Weak, - key_id: u64, + key_id: KeyId, ttl: Duration, response: oneshot::Sender>, }, Revoke { authority: Weak, - key_id: u64, + key_id: KeyId, response: oneshot::Sender>, }, Gc { @@ -476,7 +657,7 @@ enum AuthCommand { Audit { authority: Weak, action: String, - key_id: Option, + key_id: Option, detail: Option, response: oneshot::Sender>, }, @@ -496,9 +677,9 @@ pub(crate) use config::platform_default_auth_state_dir; #[cfg(all(test, not(any(windows, target_os = "macos"))))] pub(in crate::common::auth) use config::{linux_system_auth_dir_usable, unix_effective_uid}; mod keys; +pub use keys::derive_temporary_key; #[cfg(test)] pub(in crate::common::auth) use keys::recover_admin_key_after_rotation; -pub use keys::{derive_temporary_key, key_generation, key_slot, make_key_id}; pub(in crate::common::auth) use keys::{ load_isolated_server_admin_credential, load_server_admin_credential, }; @@ -520,7 +701,7 @@ impl Drop for LegacyConnectionGuard { #[derive(Clone, Debug, Serialize, Deserialize)] struct PersistedEntry { - key_id: u64, + key_id: KeyId, state: SlotState, issued_at: u64, expires_at: u64, @@ -533,7 +714,7 @@ struct PersistedEntry { struct PersistedSnapshot { schema_version: u16, instance_id: [u8; INSTANCE_ID_LEN], - generations: Vec, + generations: Vec, entries: Vec, legacy_protocol: LegacyProtocolPolicy, #[serde(default)] @@ -568,8 +749,8 @@ impl AdminReplayRecord { #[derive(Clone, Debug, Serialize, Deserialize)] enum StateMutation { Issue(PersistedEntry), - Renew { key_id: u64, expires_at: u64 }, - Revoke { key_id: u64, at: u64 }, + Renew { key_id: KeyId, expires_at: u64 }, + Revoke { key_id: KeyId, at: u64 }, LegacyProtocol(LegacyProtocolPolicy), } @@ -577,7 +758,7 @@ enum StateMutation { struct AuditRecord { at: u64, action: String, - key_id: Option, + key_id: Option, label: Option, } @@ -597,18 +778,21 @@ mod persistence; pub use persistence::*; pub(in crate::common::auth) use persistence::{ append_audit, append_mutation, append_wal, atomic_write, auth_snapshot_path, build_snapshot, - cancel_all_temporary_leases, clear_retained_high_slot_entries, compaction_is_allowed, - empty_snapshot, fail_closed_on_uncertain_wal, hex, key_matches_existing_state, - load_or_create_instance_id, load_persisted_state, normalize_tombstone_times, open_blob, - prepare_state_dir_and_lock, push_audit_record, push_persisted_audit, random_instance_id, - recover_instance_id_after_reset, reset_already_installed, rotation_already_installed, - split_high_slot_state, truncate_auth_wal, unix_seconds, write_admin_key, - write_snapshot_and_truncate_wal, + cancel_all_temporary_leases, compaction_is_allowed, empty_snapshot, + fail_closed_on_uncertain_wal, hex, key_matches_existing_state, load_or_create_instance_id, + load_persisted_state, normalize_tombstone_times, open_blob, prepare_state_dir_and_lock, + push_audit_record, push_persisted_audit, random_instance_id, recover_instance_id_after_reset, + reset_already_installed, rotation_already_installed, split_high_slot_state, truncate_auth_wal, + unix_seconds, write_admin_key, write_snapshot_and_truncate_wal, }; #[cfg(test)] pub(in crate::common::auth) use persistence::{ prepare_state_dir, read_instance_id_file, try_load_persisted_state, }; +mod ids; +pub use ids::{Generation, KeyId, SlotIndex, ADMIN_KEY_ID}; +mod leases; +use leases::Leases; mod timing_wheel; use timing_wheel::TimingWheel; #[cfg(test)] diff --git a/src/common/auth/actor/epoch.rs b/src/common/auth/actor/epoch.rs index 0c61529..84497e5 100644 --- a/src/common/auth/actor/epoch.rs +++ b/src/common/auth/actor/epoch.rs @@ -2,31 +2,8 @@ use super::super::*; use super::{audit, ensure_store_available}; -fn wipe_temporary_keys( - inner: &AuthStateInner, - cold: &mut HashMap, - wheel: &mut TimingWheel, -) { - cancel_all_temporary_leases(inner); - let mut slots = inner - .slots - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - for slot in slots.iter_mut() { - slot.state = SlotState::Free; - slot.expires_at = 0; - slot.lease = Weak::new(); - } - cold.clear(); - wheel.clear(unix_seconds()); - clear_retained_high_slot_entries(inner); -} - fn remember_previous_root(inner: &AuthStateInner) { - *inner - .previous_root - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(PreviousRoot { + *recover_lock(inner.previous_root.write()) = Some(PreviousRoot { admin_key: inner.admin_key(), instance_id: inner.instance_id(), }); @@ -35,8 +12,7 @@ fn remember_previous_root(inner: &AuthStateInner) { pub(super) fn actor_reset( inner: &Arc, config: &AuthConfig, - cold: &mut HashMap, - wheel: &mut TimingWheel, + leases: &mut Leases, admin_replays: &VecDeque, action: &str, ) -> Result<(), AuthFailure> { @@ -71,11 +47,8 @@ pub(super) fn actor_reset( let _ = std::fs::remove_file(&next_instance_path); push_audit_record(inner, reset_audit); remember_previous_root(inner); - wipe_temporary_keys(inner, cold, wheel); - *inner - .instance_id - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()) = new_instance_id; + leases.wipe(unix_seconds()); + *recover_lock(inner.instance_id.write()) = new_instance_id; inner.safe_mode.store(false, Ordering::Release); Ok(()) } @@ -83,8 +56,7 @@ pub(super) fn actor_reset( pub(super) fn actor_rotate_root( inner: &Arc, config: &AuthConfig, - cold: &mut HashMap, - wheel: &mut TimingWheel, + leases: &mut Leases, admin_lease: &mut Arc, new_key: AesKeyType, ) -> Result<(), AuthFailure> { @@ -128,13 +100,10 @@ pub(super) fn actor_rotate_root( let _ = std::fs::remove_file(&next_key_path); push_audit_record(inner, rotate_audit); remember_previous_root(inner); - wipe_temporary_keys(inner, cold, wheel); + leases.wipe(unix_seconds()); let old_admin_lease = admin_lease.clone(); - let new_admin_lease = Arc::new(AuthLease::new(0, u64::MAX)); - *inner - .admin - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()) = AdminState { + let new_admin_lease = Arc::new(AuthLease::new(ADMIN_KEY_ID, u64::MAX)); + *recover_lock(inner.admin.write()) = AdminState { key: new_key, lease: Arc::downgrade(&new_admin_lease), }; diff --git a/src/common/auth/actor/lifecycle.rs b/src/common/auth/actor/lifecycle.rs index 2606e6b..ee3fa89 100644 --- a/src/common/auth/actor/lifecycle.rs +++ b/src/common/auth/actor/lifecycle.rs @@ -2,7 +2,7 @@ use super::super::*; use super::{ audit, ensure_store_available, key_not_active, key_not_found, key_not_renewable, - push_tombstone, slot_state_name, validate_slot_identity, + slot_state_name, validate_slot_identity, }; fn validate_ttl(config: &AuthConfig, ttl: Duration) -> Result { @@ -43,11 +43,49 @@ fn validate_label(label: Option) -> Result, AuthFailure> Ok(label) } +/// One key's lifecycle, read from wherever that key lives. +struct KeyState { + state: SlotState, + expires_at: u64, + issued_at: u64, + label: Option, +} + +/// Reads a key's lifecycle from the slot table, falling back to the entries +/// retained for slots the configured capacity no longer covers. Every operation +/// that accepts any live key id needs both paths; see +/// `AuthStateInner::high_slot_generations`. +fn key_state(inner: &AuthStateInner, key_id: KeyId) -> Result { + let slots = inner.slots(); + if let Some(slot) = slots.get(key_id.slot().as_index()) { + validate_slot_identity(slot, key_id)?; + let metadata = inner + .cold() + .get(&key_id) + .cloned() + .ok_or_else(|| key_not_found(key_id))?; + return Ok(KeyState { + state: slot.state, + expires_at: slot.expires_at, + issued_at: metadata.issued_at, + label: metadata.label.clone(), + }); + } + drop(slots); + let high = inner.high(); + let entry = high_slot_entry(&high, key_id)?; + Ok(KeyState { + state: entry.state, + expires_at: entry.expires_at, + issued_at: entry.issued_at, + label: entry.label.clone(), + }) +} + pub(super) fn actor_issue( inner: &Arc, config: &AuthConfig, - cold: &mut HashMap, - wheel: &mut TimingWheel, + leases: &mut Leases, ttl: Duration, label: Option, ) -> Result { @@ -56,25 +94,24 @@ pub(super) fn actor_issue( let label = validate_label(label)?; let issued_at = unix_seconds(); let (index, generation, key_id, entry) = { - let slots = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let Some((index, slot)) = slots - .iter() - .enumerate() - .find(|(_, slot)| slot.state == SlotState::Free && slot.generation < u32::MAX) - else { + let slots = inner.slots(); + // A row whose generation cannot advance is skipped rather than reused: it + // has no unused identity left to hand out. + let Some((slot_index, generation)) = slots.iter().enumerate().find_map(|(index, slot)| { + (slot.state == SlotState::Free) + .then(|| slot.generation.next()) + .flatten() + .map(|generation| (SlotIndex::from_index(index), generation)) + }) else { return Err(AuthFailure::new( "temporary_key_capacity_exhausted", "temporary key slot table is full", true, )); }; - let generation = slot.generation + 1; - let key_id = make_key_id(generation, index as u32); + let key_id = KeyId::new(generation, slot_index); ( - index, + slot_index.as_index(), generation, key_id, PersistedEntry { @@ -95,10 +132,7 @@ pub(super) fn actor_issue( StateMutation::Issue(entry), audit("temporary_key_issue", Some(key_id), label.clone()), )?; - let mut slots = inner - .slots - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); + let mut slots = inner.slots_mut(); let slot = slots .get_mut(index) .ok_or_else(|| AuthFailure::internal("issued slot disappeared"))?; @@ -106,33 +140,21 @@ pub(super) fn actor_issue( slot.generation = generation; slot.state = SlotState::Active; slot.expires_at = expires_at; - slot.issued_epoch = inner.root_epoch.load(Ordering::Acquire); slot.lease = Arc::downgrade(&lease); - cold.insert( - key_id, - ColdMetadata { - issued_at, - label, - tombstoned_at: 0, - }, - ); - wheel.insert(lease); drop(slots); - metadata_with_credential(inner, cold, key_id, true) + leases.issue(&lease, issued_at, label); + metadata_with_credential(inner, key_id, true) } pub(super) fn actor_list( inner: &Arc, - cold: &HashMap, page: u32, page_size: u16, ) -> Result { let page_size = page_size.clamp(1, 1000) as usize; let start = (page as usize).saturating_mul(page_size); - let slots = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); + let slots = inner.slots(); + let cold = inner.cold(); let mut all = slots .iter() .enumerate() @@ -140,7 +162,7 @@ pub(super) fn actor_list( if slot.state == SlotState::Free { return None; } - let key_id = make_key_id(slot.generation, index as u32); + let key_id = KeyId::new(slot.generation, SlotIndex::from_index(index)); let cold = cold.get(&key_id)?; Some(TemporaryKeyMetadata { key_id, @@ -153,9 +175,7 @@ pub(super) fn actor_list( .collect::>(); all.extend( inner - .high_slot_entries - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()) + .high() .iter() .filter(|entry| entry.state != SlotState::Free) .map(high_slot_metadata), @@ -173,11 +193,10 @@ pub(super) fn actor_list( pub(super) fn actor_show( inner: &Arc, config: &AuthConfig, - cold: &HashMap, - key_id: u64, + key_id: KeyId, reveal: bool, ) -> Result { - let result = metadata_with_credential(inner, cold, key_id, reveal)?; + let result = metadata_with_credential(inner, key_id, reveal)?; append_audit( config, inner, @@ -197,57 +216,25 @@ pub(super) fn actor_show( pub(super) fn actor_renew( inner: &Arc, config: &AuthConfig, - cold: &HashMap, - wheel: &mut TimingWheel, - key_id: u64, + leases: &mut Leases, + key_id: KeyId, ttl: Duration, ) -> Result { ensure_store_available(inner)?; let expires_at = validate_ttl(config, ttl)?; - let index = key_slot(key_id) as usize; - { - let slots = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - if let Some(slot) = slots.get(index) { - validate_slot_identity(slot, key_id)?; - if slot.state != SlotState::Active || slot.expires_at <= unix_seconds() { - return Err(key_not_renewable()); - } - } else { - let high = inner - .high_slot_entries - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let entry = high_slot_entry(&high, key_id)?; - if entry.state != SlotState::Active || entry.expires_at <= unix_seconds() { - return Err(key_not_renewable()); - } - } + let index = key_id.slot().as_index(); + let current = key_state(inner, key_id)?; + if current.state != SlotState::Active || current.expires_at <= unix_seconds() { + return Err(key_not_renewable()); } - let label = cold - .get(&key_id) - .and_then(|metadata| metadata.label.clone()) - .or_else(|| { - inner - .high_slot_entries - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .iter() - .find(|entry| entry.key_id == key_id) - .and_then(|entry| entry.label.clone()) - }); + let label = current.label; append_mutation( config, inner, StateMutation::Renew { key_id, expires_at }, - audit("temporary_key_renew", Some(key_id), label), + audit("temporary_key_renew", Some(key_id), label.clone()), )?; - let mut slots = inner - .slots - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); + let mut slots = inner.slots_mut(); if let Some(slot) = slots.get_mut(index) { validate_slot_identity(slot, key_id)?; if slot.state != SlotState::Active { @@ -258,29 +245,26 @@ pub(super) fn actor_renew( )); } slot.expires_at = expires_at; - let lease = match slot.lease.upgrade() { + match slot.lease.upgrade() { Some(lease) if !lease.cancellation_token().is_cancelled() => { lease.expires_at.store(expires_at, Ordering::Release); - lease.wheel_version.fetch_add(1, Ordering::AcqRel); - lease + drop(slots); + leases.renew(key_id, expires_at); } + // A cancelled lease cannot be revived, so the renewal installs a + // replacement; that drops the handle on the lease it succeeds. _ => { let lease = Arc::new(AuthLease::new(key_id, expires_at)); slot.lease = Arc::downgrade(&lease); - lease + drop(slots); + leases.adopt(&lease); } - }; - wheel.release(key_id); - wheel.insert(lease); - drop(slots); - return metadata_with_credential(inner, cold, key_id, true); + } + return metadata_with_credential(inner, key_id, true); } drop(slots); { - let mut high = inner - .high_slot_entries - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); + let mut high = inner.high_mut(); let entry = high_slot_entry_mut(&mut high, key_id)?; if entry.state != SlotState::Active { return Err(AuthFailure::new( @@ -292,88 +276,75 @@ pub(super) fn actor_renew( entry.expires_at = expires_at; entry.tombstoned_at = None; } - metadata_with_credential(inner, cold, key_id, true) + metadata_with_credential(inner, key_id, true) } pub(super) fn actor_revoke( inner: &Arc, config: &AuthConfig, - cold: &mut HashMap, - tombstones: &mut VecDeque<(u64, u64)>, - key_id: u64, + leases: &mut Leases, + key_id: KeyId, ) -> Result { ensure_store_available(inner)?; let now = unix_seconds(); - let index = key_slot(key_id) as usize; - let (label, issued_at, expires_at) = { - let slots = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - if let Some(slot) = slots.get(index) { - validate_slot_identity(slot, key_id)?; - if slot.state != SlotState::Active { - return Err(key_not_active()); - } - let metadata = cold.get(&key_id).ok_or_else(|| key_not_found(key_id))?; - (metadata.label.clone(), metadata.issued_at, slot.expires_at) - } else { - let high = inner - .high_slot_entries - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let entry = high_slot_entry(&high, key_id)?; - if entry.state != SlotState::Active { - return Err(key_not_active()); - } - (entry.label.clone(), entry.issued_at, entry.expires_at) - } - }; + let index = key_id.slot().as_index(); + let current = key_state(inner, key_id)?; + if current.state != SlotState::Active { + return Err(key_not_active()); + } + let KeyState { + label, + issued_at, + expires_at, + .. + } = current; append_mutation( config, inner, StateMutation::Revoke { key_id, at: now }, audit("temporary_key_revoke", Some(key_id), label.clone()), )?; - let mut slots = inner - .slots - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); + let mut slots = inner.slots_mut(); if let Some(slot) = slots.get_mut(index) { validate_slot_identity(slot, key_id)?; slot.state = SlotState::Revoked; if let Some(lease) = slot.lease.upgrade() { lease.cancel_revoked(); } - let cold_metadata = cold.get_mut(&key_id).ok_or_else(|| key_not_found(key_id))?; - cold_metadata.tombstoned_at = now; - push_tombstone(tombstones, now, key_id); + let state = slot_state_name(slot.state).to_string(); + let expires_at = slot.expires_at; + drop(slots); + // Retire only: the row stays until its retention elapses, because the + // slot table holds a `Weak` and a later request has to be able to read + // the revoked reason rather than find a recycled row. + leases.retire_now(key_id); + let metadata = inner + .cold() + .get(&key_id) + .cloned() + .ok_or_else(|| key_not_found(key_id))?; return Ok(TemporaryKeyMetadata { key_id, - state: slot_state_name(slot.state).to_string(), - issued_at: cold_metadata.issued_at, - expires_at: slot.expires_at, - label: cold_metadata.label.clone(), + state, + issued_at: metadata.issued_at, + expires_at, + label: metadata.label.clone(), }); } drop(slots); - let mut high = inner - .high_slot_entries - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); + let mut high = inner.high_mut(); let entry = high_slot_entry_mut(&mut high, key_id)?; if entry.state != SlotState::Active { return Err(key_not_active()); } entry.state = SlotState::Revoked; entry.tombstoned_at = Some(now); - if let Some(metadata) = cold.get_mut(&key_id) { - metadata.tombstoned_at = now; - } - push_tombstone(tombstones, now, key_id); + let state = slot_state_name(entry.state).to_string(); + drop(high); + leases.retire_now(key_id); Ok(TemporaryKeyMetadata { key_id, - state: slot_state_name(entry.state).to_string(), + state, issued_at, expires_at, label, @@ -383,62 +354,13 @@ pub(super) fn actor_revoke( pub(super) fn actor_gc( inner: &Arc, config: &AuthConfig, - cold: &mut HashMap, - wheel: &mut TimingWheel, - tombstones: &mut VecDeque<(u64, u64)>, + leases: &mut Leases, admin_replays: &VecDeque, ) -> Result { ensure_store_available(inner)?; - let now = unix_seconds(); - let mut removed = 0_u64; - let mut slots = inner - .slots - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - for (index, slot) in slots.iter_mut().enumerate() { - if matches!(slot.state, SlotState::Expired | SlotState::Revoked) - || (slot.state == SlotState::Active && slot.expires_at <= now) - { - let key_id = make_key_id(slot.generation, index as u32); - if let Some(lease) = slot.lease.upgrade() { - if slot.state == SlotState::Revoked { - lease.cancel_revoked(); - } else { - lease.cancel_expired(); - } - } - slot.state = SlotState::Free; - slot.expires_at = 0; - slot.lease = Weak::new(); - cold.remove(&key_id); - wheel.release(key_id); - removed = removed.saturating_add(1); - } - } - drop(slots); - { - let mut high = inner - .high_slot_entries - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - high.retain(|entry| { - let keep = match entry.state { - SlotState::Active if entry.expires_at > now => true, - SlotState::Active | SlotState::Expired | SlotState::Revoked | SlotState::Free => { - false - } - }; - if !keep { - cold.remove(&entry.key_id); - wheel.release(entry.key_id); - removed = removed.saturating_add(1); - } - keep - }); - } - tombstones.clear(); + let removed = leases.collect_garbage(unix_seconds()); let gc_audit = audit("temporary_key_gc", None, Some(format!("removed={removed}"))); - let mut snapshot = build_snapshot(inner, cold, admin_replays); + let mut snapshot = build_snapshot(inner, admin_replays); push_persisted_audit(&mut snapshot.audit_records, gc_audit.clone()); let admin_key = inner.admin_key(); if let Err(error) = write_snapshot_and_truncate_wal(config, &admin_key, &snapshot) { @@ -450,7 +372,7 @@ pub(super) fn actor_gc( Ok(removed) } -fn high_slot_entry(high: &[PersistedEntry], key_id: u64) -> Result<&PersistedEntry, AuthFailure> { +fn high_slot_entry(high: &[PersistedEntry], key_id: KeyId) -> Result<&PersistedEntry, AuthFailure> { high.iter() .find(|entry| entry.key_id == key_id) .ok_or_else(|| key_not_found(key_id)) @@ -458,7 +380,7 @@ fn high_slot_entry(high: &[PersistedEntry], key_id: u64) -> Result<&PersistedEnt fn high_slot_entry_mut( high: &mut [PersistedEntry], - key_id: u64, + key_id: KeyId, ) -> Result<&mut PersistedEntry, AuthFailure> { high.iter_mut() .find(|entry| entry.key_id == key_id) @@ -477,22 +399,19 @@ fn high_slot_metadata(entry: &PersistedEntry) -> TemporaryKeyMetadata { fn metadata_with_credential( inner: &Arc, - cold: &HashMap, - key_id: u64, + key_id: KeyId, reveal: bool, ) -> Result { - let slots = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); + let slots = inner.slots(); let credential = if reveal { let key = derive_temporary_key(&inner.admin_key(), &inner.instance_id(), key_id)?; - encode_temporary_credential(key_id, &key) + encode_temporary_credential(key_id.as_u64(), &key) } else { String::new() }; - if let Some(slot) = slots.get(key_slot(key_id) as usize) { + if let Some(slot) = slots.get(key_id.slot().as_index()) { validate_slot_identity(slot, key_id)?; + let cold = inner.cold(); let cold = cold.get(&key_id).ok_or_else(|| key_not_found(key_id))?; return Ok(IssuedTemporaryKey { metadata: TemporaryKeyMetadata { @@ -505,10 +424,7 @@ fn metadata_with_credential( credential, }); } - let high = inner - .high_slot_entries - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); + let high = inner.high(); Ok(IssuedTemporaryKey { metadata: high_slot_metadata(high_slot_entry(&high, key_id)?), credential, diff --git a/src/common/auth/actor/mod.rs b/src/common/auth/actor/mod.rs index ff01083..20e2c75 100644 --- a/src/common/auth/actor/mod.rs +++ b/src/common/auth/actor/mod.rs @@ -24,22 +24,19 @@ use epoch::*; use lifecycle::*; pub(super) struct AuthActorState { - cold: HashMap, - wheel: TimingWheel, + leases: Leases, admin_replays: HashSet<[u8; 32]>, admin_replay_order: VecDeque, } impl AuthActorState { pub(super) fn new( - cold: HashMap, - wheel: TimingWheel, + leases: Leases, admin_replays: HashSet<[u8; 32]>, admin_replay_order: VecDeque, ) -> Self { Self { - cold, - wheel, + leases, admin_replays, admin_replay_order, } @@ -55,61 +52,10 @@ pub(super) async fn run_auth_actor( _state_lock: Arc, ) { let AuthActorState { - mut cold, - mut wheel, + mut leases, mut admin_replays, mut admin_replay_order, } = state; - let now = unix_seconds(); - let mut tombstones = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .iter() - .enumerate() - .filter_map(|(index, slot)| { - if !matches!(slot.state, SlotState::Expired | SlotState::Revoked) { - return None; - } - let key_id = make_key_id(slot.generation, index as u32); - let tombstoned_at = cold - .get(&key_id) - .map(|metadata| metadata.tombstoned_at) - .unwrap_or(now); - Some(( - tombstoned_at.saturating_add(TOMBSTONE_RETENTION.as_secs()), - key_id, - )) - }) - .collect::>(); - tombstones.extend( - inner - .high_slot_entries - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .iter() - .filter_map(|entry| { - let expired_active = entry.state == SlotState::Active && entry.expires_at <= now; - if !matches!(entry.state, SlotState::Expired | SlotState::Revoked) - && !expired_active - { - return None; - } - let tombstoned_at = entry - .tombstoned_at - .or_else(|| { - cold.get(&entry.key_id) - .map(|metadata| metadata.tombstoned_at) - }) - .unwrap_or(entry.expires_at.max(now)); - Some(( - tombstoned_at.saturating_add(TOMBSTONE_RETENTION.as_secs()), - entry.key_id, - )) - }), - ); - tombstones.sort_unstable_by_key(|(cleanup_at, _)| *cleanup_at); - let mut tombstones = VecDeque::from(tombstones); let mut last_snapshot_at = unix_seconds(); let mut tick = tokio::time::interval(Duration::from_secs(1)); tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); @@ -117,67 +63,7 @@ pub(super) async fn run_auth_actor( tokio::select! { _ = tick.tick() => { let now = unix_seconds(); - for lease in wheel.advance(now) { - let key_id = lease.key_id(); - let version = lease.wheel_version.load(Ordering::Acquire); - if lease.expires_at() > now { - wheel.insert_with_version(lease, version); - continue; - } - let mut slots = inner.slots.write().unwrap_or_else(|poisoned| poisoned.into_inner()); - if let Some(slot) = slots.get_mut(key_slot(key_id) as usize) { - if slot.generation == key_generation(key_id) && slot.state == SlotState::Active { - slot.state = SlotState::Expired; - lease.cancel_expired(); - let tombstoned_at = slot.expires_at; - if let Some(metadata) = cold.get_mut(&key_id) { - metadata.tombstoned_at = tombstoned_at; - } - push_tombstone(&mut tombstones, tombstoned_at, key_id); - tracing::info!( - event = "temporary_key_expired", - auth_stage = "expiry", - key_id, - expires_at = lease.expires_at(), - "temporary key expired and active work was cancelled" - ); - } - } else { - lease.cancel_expired(); - } - } - expire_due_high_slots(&inner, &mut cold, &mut tombstones, now); - let mut due_high = Vec::new(); - while let Some((cleanup_at, key_id)) = tombstones.front().copied() { - if cleanup_at > now { - break; - } - tombstones.pop_front(); - let mut slots = inner.slots.write().unwrap_or_else(|poisoned| poisoned.into_inner()); - if let Some(slot) = slots.get_mut(key_slot(key_id) as usize) { - if slot.generation == key_generation(key_id) && matches!(slot.state, SlotState::Expired | SlotState::Revoked) { - slot.state = SlotState::Free; - slot.expires_at = 0; - slot.lease = Weak::new(); - cold.remove(&key_id); - wheel.release(key_id); - } - } else { - due_high.push(key_id); - } - } - if !due_high.is_empty() { - let due = due_high.iter().copied().collect::>(); - inner - .high_slot_entries - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .retain(|entry| !due.contains(&entry.key_id)); - for key_id in due_high { - cold.remove(&key_id); - wheel.release(key_id); - } - } + leases.tick(now); prune_expired_admin_replays( now, &mut admin_replays, @@ -191,7 +77,7 @@ pub(super) async fn run_auth_actor( && now.saturating_sub(last_snapshot_at) >= SNAPSHOT_COMPACTION_INTERVAL.as_secs() { - let snapshot = build_snapshot(&inner, &cold, &admin_replay_order); + let snapshot = build_snapshot(&inner, &admin_replay_order); if let Err(error) = write_snapshot_and_truncate_wal( &config, &inner.admin_key(), @@ -238,56 +124,42 @@ pub(super) async fn run_auth_actor( } AuthCommand::Issue { authority, ttl, label, response } => { let result = validate_admin_authority(&inner, &authority) - .and_then(|()| actor_issue(&inner, &config, &mut cold, &mut wheel, ttl, label)); + .and_then(|()| actor_issue(&inner, &config, &mut leases, ttl, label)); let _ = response.send(result); } AuthCommand::List { authority, page, page_size, response } => { let result = validate_admin_authority(&inner, &authority) - .and_then(|()| actor_list(&inner, &cold, page, page_size)); + .and_then(|()| actor_list(&inner, page, page_size)); let _ = response.send(result); } AuthCommand::Show { authority, key_id, reveal, response } => { let result = validate_admin_authority(&inner, &authority) - .and_then(|()| actor_show(&inner, &config, &cold, key_id, reveal)); + .and_then(|()| actor_show(&inner, &config, key_id, reveal)); let _ = response.send(result); } AuthCommand::Renew { authority, key_id, ttl, response } => { let result = validate_admin_authority(&inner, &authority) - .and_then(|()| actor_renew(&inner, &config, &cold, &mut wheel, key_id, ttl)); + .and_then(|()| actor_renew(&inner, &config, &mut leases, key_id, ttl)); let _ = response.send(result); } AuthCommand::Revoke { authority, key_id, response } => { let result = validate_admin_authority(&inner, &authority) - .and_then(|()| actor_revoke(&inner, &config, &mut cold, &mut tombstones, key_id)); + .and_then(|()| actor_revoke(&inner, &config, &mut leases, key_id)); let _ = response.send(result); } AuthCommand::Gc { authority, response } => { let result = validate_admin_authority(&inner, &authority) - .and_then(|()| actor_gc( - &inner, - &config, - &mut cold, - &mut wheel, - &mut tombstones, - &admin_replay_order, - )); + .and_then(|()| actor_gc(&inner, &config, &mut leases, &admin_replay_order)); let _ = response.send(result); } AuthCommand::Reset { authority, response } => { let result = validate_admin_authority(&inner, &authority) - .and_then(|()| actor_reset( - &inner, - &config, - &mut cold, - &mut wheel, - &admin_replay_order, - "auth_state_reset", - )); + .and_then(|()| actor_reset(&inner, &config, &mut leases, &admin_replay_order, "auth_state_reset")); let _ = response.send(result); } AuthCommand::RotateRoot { authority, new_key, response } => { let result = validate_admin_authority(&inner, &authority) - .and_then(|()| actor_rotate_root(&inner, &config, &mut cold, &mut wheel, &mut admin_lease, new_key)); + .and_then(|()| actor_rotate_root(&inner, &config, &mut leases, &mut admin_lease, new_key)); if result.is_ok() { admin_replays.clear(); admin_replay_order.clear(); @@ -389,21 +261,9 @@ pub(super) fn prune_expired_admin_replays( }); } -fn push_tombstone(tombstones: &mut VecDeque<(u64, u64)>, tombstoned_at: u64, key_id: u64) { - let cleanup_at = tombstoned_at.saturating_add(TOMBSTONE_RETENTION.as_secs()); - let index = tombstones.partition_point(|(current, _)| *current <= cleanup_at); - tombstones.insert(index, (cleanup_at, key_id)); -} - fn actor_status(inner: &Arc) -> AuthStatus { - let slots = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let high = inner - .high_slot_entries - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); + let slots = inner.slots(); + let high = inner.high(); let active_keys = slots .iter() .filter(|slot| slot.state == SlotState::Active) @@ -468,10 +328,7 @@ fn validate_admin_authority( false, )); } - let current = inner - .admin - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()) + let current = recover_lock(inner.admin.read()) .lease .upgrade() .ok_or_else(|| { @@ -503,15 +360,15 @@ fn ensure_store_available(inner: &AuthStateInner) -> Result<(), AuthFailure> { } } -fn validate_slot_identity(slot: &SlotHot, key_id: u64) -> Result<(), AuthFailure> { - if slot.generation != key_generation(key_id) || slot.state == SlotState::Free { +fn validate_slot_identity(slot: &SlotHot, key_id: KeyId) -> Result<(), AuthFailure> { + if slot.generation != key_id.generation() || slot.state == SlotState::Free { Err(key_not_found(key_id)) } else { Ok(()) } } -fn key_not_found(key_id: u64) -> AuthFailure { +fn key_not_found(key_id: KeyId) -> AuthFailure { AuthFailure::new( "temporary_key_not_found", format!("temporary key {key_id} does not exist"), @@ -535,37 +392,6 @@ fn key_not_active() -> AuthFailure { ) } -fn expire_due_high_slots( - inner: &Arc, - cold: &mut HashMap, - tombstones: &mut VecDeque<(u64, u64)>, - now: u64, -) { - let mut high = inner - .high_slot_entries - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - for entry in high.iter_mut() { - if entry.state != SlotState::Active || entry.expires_at > now { - continue; - } - entry.state = SlotState::Expired; - let tombstoned_at = entry.expires_at; - entry.tombstoned_at = Some(tombstoned_at); - if let Some(metadata) = cold.get_mut(&entry.key_id) { - metadata.tombstoned_at = tombstoned_at; - } - push_tombstone(tombstones, tombstoned_at, entry.key_id); - tracing::info!( - event = "temporary_key_expired", - auth_stage = "expiry", - key_id = entry.key_id, - expires_at = entry.expires_at, - "high-slot temporary key expired" - ); - } -} - fn slot_state_name(state: SlotState) -> &'static str { match state { SlotState::Free => "free", @@ -575,7 +401,7 @@ fn slot_state_name(state: SlotState) -> &'static str { } } -fn audit(action: &str, key_id: Option, label: Option) -> AuditRecord { +fn audit(action: &str, key_id: Option, label: Option) -> AuditRecord { AuditRecord { at: unix_seconds(), action: action.to_string(), diff --git a/src/common/auth/ids.rs b/src/common/auth/ids.rs new file mode 100644 index 0000000..fdd5a0e --- /dev/null +++ b/src/common/auth/ids.rs @@ -0,0 +1,129 @@ +//! Identity types for temporary credentials. +//! +//! ```text +//! KeyId (u64) — what a client presents +//! ┌──────────────────────────┬──────────────────────────┐ +//! │ Generation (high 32) │ SlotIndex (low 32) │ +//! └──────────────────────────┴──────────────────────────┘ +//! which tenant of the row which row of the table +//! ``` +//! +//! These were all bare integers, which made `make_key_id(generation, slot)` +//! accept its arguments in either order and let a slot index be compared against +//! a generation without complaint. Separate types make both a compile error, and +//! keep a `KeyId` from being used as an array index by mistake — the only way to +//! get one is [`KeyId::slot`], which is also the only place the truncation to a +//! row number is expressed. +//! +//! All three are `#[serde(transparent)]`, so persisted snapshots and the admin +//! wire protocol keep the plain-integer encoding they already had. + +use std::fmt; + +use serde::{Deserialize, Serialize}; + +/// The identity a client presents: a [`SlotIndex`] paired with the +/// [`Generation`] of the row it was issued from. +#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash, Serialize, Deserialize)] +#[serde(transparent)] +pub struct KeyId(u64); + +/// Which tenant of a slot a credential belongs to. Bumped every time the row is +/// reissued, and never reset, so a retired credential can never match the row +/// that replaced it. +#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash, Serialize, Deserialize)] +#[serde(transparent)] +pub struct Generation(u32); + +/// Which row of the slot table a credential lives in. +#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash, Serialize, Deserialize)] +#[serde(transparent)] +pub struct SlotIndex(u32); + +/// The administrator, which owns no slot and never expires. +pub const ADMIN_KEY_ID: KeyId = KeyId(0); + +impl KeyId { + pub const fn new(generation: Generation, slot: SlotIndex) -> Self { + Self(((generation.0 as u64) << 32) | slot.0 as u64) + } + + pub const fn generation(self) -> Generation { + Generation((self.0 >> 32) as u32) + } + + pub const fn slot(self) -> SlotIndex { + SlotIndex(self.0 as u32) + } + + pub const fn is_admin(self) -> bool { + self.0 == ADMIN_KEY_ID.0 + } + + /// The bytes mixed into the credential's key derivation. + pub const fn to_be_bytes(self) -> [u8; 8] { + self.0.to_be_bytes() + } + + pub const fn as_u64(self) -> u64 { + self.0 + } + + pub const fn from_u64(raw: u64) -> Self { + Self(raw) + } +} + +impl Generation { + pub const FIRST: Self = Self(0); + + /// The generation for a reissue of this row, or `None` once the row has been + /// cycled `u32::MAX` times and can no longer produce a fresh identity. + pub fn next(self) -> Option { + self.0.checked_add(1).map(Self) + } + + pub const fn as_u32(self) -> u32 { + self.0 + } + + pub const fn from_u32(raw: u32) -> Self { + Self(raw) + } +} + +impl SlotIndex { + pub const fn as_index(self) -> usize { + self.0 as usize + } + + /// # Panics + /// + /// If `index` exceeds `u32::MAX`. `MAX_TEMP_KEY_CAPACITY` caps the table far + /// below that, so a real index cannot reach it; panicking keeps a future + /// capacity change from silently wrapping into another row's identity. + pub fn from_index(index: usize) -> Self { + match u32::try_from(index) { + Ok(index) => Self(index), + Err(_) => panic!("slot index exceeds the addressable slot table"), + } + } +} + +impl fmt::Display for KeyId { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt(formatter) + } +} + +impl fmt::Display for Generation { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt(formatter) + } +} + +impl fmt::Display for SlotIndex { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt(formatter) + } +} diff --git a/src/common/auth/keys.rs b/src/common/auth/keys.rs index 8cf718f..a0592d4 100644 --- a/src/common/auth/keys.rs +++ b/src/common/auth/keys.rs @@ -203,22 +203,10 @@ pub(super) fn load_isolated_server_admin_credential( validate_admin_credential(&raw) } -pub fn make_key_id(generation: u32, slot: u32) -> u64 { - (u64::from(generation) << 32) | u64::from(slot) -} - -pub fn key_generation(key_id: u64) -> u32 { - (key_id >> 32) as u32 -} - -pub fn key_slot(key_id: u64) -> u32 { - key_id as u32 -} - pub fn derive_temporary_key( admin_key: &AesKeyType, instance_id: &[u8; INSTANCE_ID_LEN], - key_id: u64, + key_id: KeyId, ) -> Result { let salt = Salt::new(HKDF_SHA256, instance_id); let pseudo_random_key = salt.extract(admin_key); diff --git a/src/common/auth/leases.rs b/src/common/auth/leases.rs new file mode 100644 index 0000000..64c83ca --- /dev/null +++ b/src/common/auth/leases.rs @@ -0,0 +1,280 @@ +//! Temporary-key lifetimes, expressed as one self-advancing cleanup callback per +//! key. +//! +//! ```text +//! issue schedule(expires_at) ── slots[i] Active, lease live +//! | +//! v phase 1: deadline reached, or the entry is cancelled +//! retire lease cancelled, slots[i] Expired, tombstoned_at recorded +//! | +//! v returns expires_at + TOMBSTONE_RETENTION +//! (waiting) <- a client presenting the dead credential is told +//! | "expired", not the "unknown key" it would get from +//! v an already-recycled row +//! reap phase 2: slots[i] Free (generation kept), cold metadata and any +//! high-slot row removed. Nothing left; the entry is done. +//! ``` +//! +//! One entry covers a key's whole life, so there is no tombstone queue and no +//! sweep to keep in step with the wheel. Every way a key can end runs the same +//! two phases: reaching a deadline runs them on schedule, and cancelling the +//! entry — for a revoke, a GC, a root rotation, or the wheel being dropped — runs +//! whichever phases remain immediately. No call site performs cleanup, which is +//! what keeps a forgotten call from stranding a lease past its row's reuse or +//! leaking a metadata entry per issued key. +//! +//! The callbacks hold a `Weak`, so they neither keep the state +//! alive nor touch it after a runtime has shut down. + +use super::*; + +/// Whether a new entry tears down the one it replaces, or takes over its work. +enum Schedule { + /// Any entry already held is torn down first. + Fresh, + /// The previous entry is discarded without running its phases, because this + /// one now owes them. + Supersede, +} + +pub(super) struct Leases { + inner: Weak, + wheel: TimingWheel, +} + +impl Leases { + /// Rebuilds a loaded state's schedule: live keys wait for their expiry, and + /// keys that were already dead wait out the rest of their retention. + pub(super) fn restored(inner: &Arc, now: u64) -> Self { + let mut leases = Self { + inner: Arc::downgrade(inner), + wheel: TimingWheel::new(now), + }; + let mut live = Vec::new(); + let mut dead = Vec::new(); + for (index, slot) in inner.slots().iter().enumerate() { + let key_id = KeyId::new(slot.generation, SlotIndex::from_index(index)); + match slot.state { + SlotState::Active => live.extend(slot.lease.upgrade().map(|l| (key_id, l))), + SlotState::Expired | SlotState::Revoked => dead.push(key_id), + SlotState::Free => {} + } + } + dead.extend( + inner + .high() + .iter() + .filter(|entry| entry.state != SlotState::Active) + .map(|entry| entry.key_id), + ); + for (key_id, lease) in live { + leases.watch(key_id, lease, Schedule::Fresh); + } + for key_id in dead { + let tombstoned_at = inner + .cold() + .get(&key_id) + .map(|cold| cold.tombstoned_at) + .filter(|at| *at != 0) + .unwrap_or(now); + leases.entomb(key_id, tombstoned_at); + } + leases + } + + /// Takes over a newly issued key: records its description and schedules the + /// retirement its expiry, or any earlier cancellation, will run. + pub(super) fn issue(&mut self, lease: &Arc, issued_at: u64, label: Option) { + let Some(inner) = self.inner.upgrade() else { + return; + }; + inner.cold_mut().insert( + lease.key_id(), + ColdMetadata { + issued_at, + label, + tombstoned_at: 0, + }, + ); + self.watch(lease.key_id(), lease.clone(), Schedule::Fresh); + } + + /// Hands a key's remaining life to a replacement lease, for a renewal whose + /// original lease had already been cancelled. The row stays alive, so the + /// entry being replaced must not run its teardown. + pub(super) fn adopt(&mut self, lease: &Arc) { + self.watch(lease.key_id(), lease.clone(), Schedule::Supersede); + } + + /// Moves a renewed key to its new expiry. Returns `false` for a key the + /// wheel is not watching, as for a high slot. + pub(super) fn renew(&mut self, key_id: KeyId, expires_at: u64) -> bool { + self.wheel.reschedule(key_id, expires_at) + } + + /// Retires a key now rather than at its expiry, leaving its tombstone to run + /// on schedule. This is what a revoke needs: the credential stops working + /// immediately, but the row is still held long enough to report *why*. + pub(super) fn retire_now(&mut self, key_id: KeyId) { + self.wheel.advance_one_phase(key_id); + } + + /// Ends a key outright, running whichever of its phases remain: an active + /// key is retired and reaped, and a tombstoned one is reaped. Skips the + /// retention wait, so it is for a caller that wants the row back now. + pub(super) fn end(&mut self, key_id: KeyId) { + self.wheel.cancel(key_id); + } + + /// Runs every phase whose deadline has passed. + pub(super) fn tick(&mut self, now: u64) { + self.wheel.advance(now); + } + + /// Ends every key at once, for a root rotation or state reset. Dropping the + /// wheel runs all remaining phases, so no row, lease, or metadata entry + /// survives it. + pub(super) fn wipe(&mut self, now: u64) { + // Rotation is the one reason a phase cannot infer, so it is recorded + // before the drop; `record_cancel` keeps the first reason. + if let Some(inner) = self.inner.upgrade() { + for lease in inner.slots().iter().filter_map(|slot| slot.lease.upgrade()) { + lease.cancel_rotated(); + } + } + self.wheel = TimingWheel::new(now); + } + + /// Ends every key that is dead or past its deadline, skipping the tombstone + /// wait. Returns how many keys were ended. + pub(super) fn collect_garbage(&mut self, now: u64) -> u64 { + let Some(inner) = self.inner.upgrade() else { + return 0; + }; + let mut due = inner + .slots() + .iter() + .enumerate() + .filter(|(_, slot)| slot.is_collectable(now)) + .map(|(index, slot)| KeyId::new(slot.generation, SlotIndex::from_index(index))) + .collect::>(); + due.extend( + inner + .high() + .iter() + .filter(|entry| entry.state != SlotState::Active || entry.expires_at <= now) + .map(|entry| entry.key_id), + ); + for key_id in &due { + self.end(*key_id); + } + due.len() as u64 + } + + /// Schedules a live key's two phases, starting at its lease's expiry. + /// + /// The entry owns the strong `Arc`, so the lease lives exactly as + /// long as the wheel is watching it: request-facing structures hold only + /// `Weak` references, and dropping the entry is what ends the lease. + fn watch(&mut self, key_id: KeyId, lease: Arc, schedule: Schedule) { + let inner = self.inner.clone(); + let deadline = lease.expires_at(); + // WHY the closure keeps the lease across both phases: the slot table + // holds only a `Weak`, so this is the reference that lets a request + // during the tombstone read *why* the key died instead of finding a + // vanished lease. It is released when the entry itself is dropped, after + // the reap. + let mut retired = false; + let phase = move || { + let inner = inner.upgrade()?; + if std::mem::replace(&mut retired, true) { + reap(&inner, key_id); + return None; + } + Some(retire(&inner, key_id, Some(&lease))) + }; + match schedule { + Schedule::Fresh => self.wheel.schedule(key_id, deadline, phase), + Schedule::Supersede => self.wheel.supersede(key_id, deadline, phase), + } + } + + /// Schedules only the reap phase, for a key that is already dead. + fn entomb(&mut self, key_id: KeyId, tombstoned_at: u64) { + let inner = self.inner.clone(); + self.wheel + .schedule(key_id, retention_ends(tombstoned_at), move || { + reap(&inner.upgrade()?, key_id); + None + }); + } +} + +fn retention_ends(tombstoned_at: u64) -> u64 { + tombstoned_at.saturating_add(TOMBSTONE_RETENTION.as_secs()) +} + +/// Phase 1: cancels the lease, marks the row dead, and records when its +/// retention starts. Returns the deadline of the reap that follows. +fn retire(inner: &Arc, key_id: KeyId, lease: Option<&Arc>) -> u64 { + if let Some(lease) = lease { + // WHY expiry is the fallback reason: a key ended for any other reason was + // already cancelled by the code that knew that reason, and `record_cancel` + // keeps the first one, so this cannot mislabel it. + lease.cancel_expired(); + } + let mut slots = inner.slots_mut(); + let tombstoned_at = match slots.get_mut(key_id.slot().as_index()) { + Some(slot) if slot.holds(key_id) && slot.state == SlotState::Active => { + slot.state = SlotState::Expired; + slot.expires_at + } + // Already marked dead by a revoke, or the row moved on. Its retention + // still has to be honoured, timed from whenever it was marked. + _ => { + drop(slots); + let mut cold = inner.cold_mut(); + let tombstoned_at = match cold.get_mut(&key_id) { + Some(cold) if cold.tombstoned_at != 0 => cold.tombstoned_at, + Some(cold) => { + cold.tombstoned_at = unix_seconds(); + cold.tombstoned_at + } + None => unix_seconds(), + }; + return retention_ends(tombstoned_at); + } + }; + drop(slots); + inner + .cold_mut() + .entry(key_id) + .and_modify(|cold| cold.tombstoned_at = tombstoned_at); + tracing::info!( + event = "temporary_key_expired", + auth_stage = "expiry", + key_id = key_id.as_u64(), + expires_at = tombstoned_at, + "temporary key expired and active work was cancelled" + ); + retention_ends(tombstoned_at) +} + +/// Phase 2: frees the row and forgets the key. +fn reap(inner: &Arc, key_id: KeyId) { + let mut slots = inner.slots_mut(); + match slots.get_mut(key_id.slot().as_index()) { + Some(slot) if slot.holds(key_id) => { + slot.retire(); + drop(slots); + } + Some(_) => return, + // Above the addressable table: the retained row is dropped outright, + // since only its generation has to survive. + None => { + drop(slots); + inner.high_mut().retain(|entry| entry.key_id != key_id); + } + } + inner.cold_mut().remove(&key_id); +} diff --git a/src/common/auth/persistence/mod.rs b/src/common/auth/persistence/mod.rs index c3f130e..f367efa 100644 --- a/src/common/auth/persistence/mod.rs +++ b/src/common/auth/persistence/mod.rs @@ -35,9 +35,9 @@ pub(crate) use fs::{replace_file, sync_parent_directory}; #[cfg(test)] pub(in crate::common::auth) use snapshot::try_load_persisted_state; pub(in crate::common::auth) use snapshot::{ - build_snapshot, cancel_all_temporary_leases, clear_retained_high_slot_entries, - compaction_is_allowed, empty_snapshot, load_persisted_state, normalize_tombstone_times, - push_audit_record, push_persisted_audit, split_high_slot_state, + build_snapshot, cancel_all_temporary_leases, compaction_is_allowed, empty_snapshot, + load_persisted_state, normalize_tombstone_times, push_audit_record, push_persisted_audit, + split_high_slot_state, }; pub(in crate::common::auth) use wal::{ append_audit, append_mutation, append_wal, fail_closed_on_uncertain_wal, read_wal, diff --git a/src/common/auth/persistence/snapshot.rs b/src/common/auth/persistence/snapshot.rs index b1730b0..a8bead0 100644 --- a/src/common/auth/persistence/snapshot.rs +++ b/src/common/auth/persistence/snapshot.rs @@ -6,19 +6,8 @@ pub(in crate::common::auth) fn compaction_is_allowed(safe_mode: bool) -> bool { !safe_mode } -pub(in crate::common::auth) fn clear_retained_high_slot_entries(inner: &AuthStateInner) { - inner - .high_slot_entries - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .clear(); -} - pub(in crate::common::auth) fn push_audit_record(inner: &AuthStateInner, record: AuditRecord) { - let mut records = inner - .audit_records - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); + let mut records = recover_lock(inner.audit_records.write()); while records.len() >= AUDIT_RECORD_CAPACITY { records.pop_front(); } @@ -26,24 +15,15 @@ pub(in crate::common::auth) fn push_audit_record(inner: &AuthStateInner, record: } pub(in crate::common::auth) fn cancel_all_temporary_leases(inner: &AuthStateInner) { - let slots = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); + let slots = inner.slots(); for lease in slots.iter().filter_map(|slot| slot.lease.upgrade()) { lease.cancel_rotated(); } } -fn snapshot_generations(inner: &AuthStateInner) -> Vec { - let slots = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - let extra = inner - .high_slot_generations - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); +fn snapshot_generations(inner: &AuthStateInner) -> Vec { + let slots = inner.slots(); + let extra = recover_lock(inner.high_slot_generations.read()); let mut generations = slots.iter().map(|slot| slot.generation).collect::>(); generations.extend_from_slice(&extra); generations @@ -52,12 +32,12 @@ fn snapshot_generations(inner: &AuthStateInner) -> Vec { pub(in crate::common::auth) fn split_high_slot_state( snapshot: &PersistedSnapshot, capacity: usize, -) -> (Vec, Vec) { +) -> (Vec, Vec) { let high_generations = snapshot.generations.get(capacity..).unwrap_or(&[]).to_vec(); let high_entries = snapshot .entries .iter() - .filter(|entry| key_slot(entry.key_id) as usize >= capacity) + .filter(|entry| entry.key_id.slot().as_index() >= capacity) .cloned() .collect(); (high_generations, high_entries) @@ -65,13 +45,10 @@ pub(in crate::common::auth) fn split_high_slot_state( pub(in crate::common::auth) fn build_snapshot( inner: &AuthStateInner, - cold: &HashMap, admin_replays: &VecDeque, ) -> PersistedSnapshot { - let slots = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); + let slots = inner.slots(); + let cold = inner.cold(); let generations = snapshot_generations(inner); let mut entries = slots .iter() @@ -80,7 +57,7 @@ pub(in crate::common::auth) fn build_snapshot( if slot.state == SlotState::Free { return None; } - let key_id = make_key_id(slot.generation, index as u32); + let key_id = KeyId::new(slot.generation, SlotIndex::from_index(index)); let cold = cold.get(&key_id)?; Some(PersistedEntry { key_id, @@ -92,14 +69,7 @@ pub(in crate::common::auth) fn build_snapshot( }) }) .collect::>(); - entries.extend( - inner - .high_slot_entries - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .iter() - .cloned(), - ); + entries.extend(inner.high().iter().cloned()); snapshot_with( inner, inner.instance_id(), @@ -156,7 +126,7 @@ pub(in crate::common::auth) fn empty_snapshot( fn snapshot_with( inner: &AuthStateInner, instance_id: [u8; INSTANCE_ID_LEN], - generations: Vec, + generations: Vec, entries: Vec, admin_replays: &VecDeque, ) -> PersistedSnapshot { @@ -171,11 +141,7 @@ fn snapshot_with( LegacyProtocolPolicy::Deny }, admin_replays: admin_replays.iter().cloned().collect(), - audit_records: inner - .audit_records - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .clone(), + audit_records: recover_lock(inner.audit_records.read()).clone(), root_epoch: inner.root_epoch.load(Ordering::Acquire), } } @@ -226,7 +192,7 @@ pub(in crate::common::auth) fn try_load_persisted_state( PersistedSnapshot { schema_version: SNAPSHOT_SCHEMA_VERSION, instance_id, - generations: vec![0; config.max_temporary_keys], + generations: vec![Generation::FIRST; config.max_temporary_keys], entries: Vec::new(), legacy_protocol: config.legacy_protocol, admin_replays: Vec::new(), @@ -242,7 +208,9 @@ pub(in crate::common::auth) fn try_load_persisted_state( )); } if snapshot.generations.len() < config.max_temporary_keys { - snapshot.generations.resize(config.max_temporary_keys, 0); + snapshot + .generations + .resize(config.max_temporary_keys, Generation::FIRST); } let wal_path = auth_wal_path(&config.state_dir); @@ -267,14 +235,14 @@ pub(in crate::common::auth) fn apply_persisted_mutation( ) -> Result<(), AuthFailure> { match mutation { StateMutation::Issue(entry) => { - let index = key_slot(entry.key_id) as usize; + let index = entry.key_id.slot().as_index(); if snapshot.generations.len() <= index { - snapshot.generations.resize(index + 1, 0); + snapshot.generations.resize(index + 1, Generation::FIRST); } - snapshot.generations[index] = key_generation(entry.key_id); + snapshot.generations[index] = entry.key_id.generation(); snapshot .entries - .retain(|current| key_slot(current.key_id) as usize != index); + .retain(|current| current.key_id.slot().as_index() != index); snapshot.entries.push(entry); } StateMutation::Renew { key_id, expires_at } => { @@ -295,7 +263,7 @@ pub(in crate::common::auth) fn apply_persisted_mutation( fn snapshot_entry_mut<'a>( snapshot: &'a mut PersistedSnapshot, - key_id: u64, + key_id: KeyId, operation: &str, ) -> Result<&'a mut PersistedEntry, AuthFailure> { snapshot diff --git a/src/common/auth/runtime.rs b/src/common/auth/runtime.rs index 28e72ab..c04b4d4 100644 --- a/src/common/auth/runtime.rs +++ b/src/common/auth/runtime.rs @@ -72,9 +72,9 @@ impl AuthRuntime { .collect::>() .into_boxed_slice(); let mut cold = HashMap::new(); - let mut wheel = TimingWheel::new(now); + let mut restored_leases = Vec::new(); - let admin_lease = Arc::new(AuthLease::new(0, u64::MAX)); + let admin_lease = Arc::new(AuthLease::new(ADMIN_KEY_ID, u64::MAX)); if let Some(state) = loaded.as_ref() { for (index, generation) in state.generations.iter().copied().enumerate() { if let Some(slot) = slots.get_mut(index) { @@ -82,11 +82,11 @@ impl AuthRuntime { } } for entry in &state.entries { - let index = key_slot(entry.key_id) as usize; + let index = entry.key_id.slot().as_index(); let Some(slot) = slots.get_mut(index) else { continue; }; - if slot.generation != key_generation(entry.key_id) { + if slot.generation != entry.key_id.generation() { continue; } let state = if entry.state == SlotState::Active && entry.expires_at <= now { @@ -111,7 +111,9 @@ impl AuthRuntime { if state == SlotState::Active { let lease = Arc::new(AuthLease::new(entry.key_id, entry.expires_at)); slot.lease = Arc::downgrade(&lease); - wheel.insert(lease); + // Held only until the schedule below adopts them; the wheel + // is the lasting owner. + restored_leases.push(lease); } } } @@ -178,6 +180,7 @@ impl AuthRuntime { root_epoch: AtomicU64::new(loaded.as_ref().map(|state| state.root_epoch).unwrap_or(0)), previous_root: RwLock::new(None), audit_records: RwLock::new(audit_records), + cold: RwLock::new(cold), }); let (command_tx, command_rx) = mpsc::channel(256); let actor = tokio::spawn(run_auth_actor( @@ -185,7 +188,11 @@ impl AuthRuntime { admin_lease, command_rx, config.clone(), - AuthActorState::new(cold, wheel, admin_replays, admin_replay_order), + AuthActorState::new( + Leases::restored(&inner, now), + admin_replays, + admin_replay_order, + ), state_lock.clone(), )); let actor_abort = actor.abort_handle(); @@ -207,11 +214,7 @@ impl AuthRuntime { .send(AuthCommand::Shutdown { response }) .await; let _ = receiver.await; - let handle = self - .actor - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .take(); + let handle = recover_lock(self.actor.lock()).take(); if let Some(handle) = handle { let _ = handle.await; } @@ -219,11 +222,7 @@ impl AuthRuntime { pub async fn abort_actor(&self) -> Result<(), AuthFailure> { self.actor_abort.abort(); - let handle = self - .actor - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .take(); + let handle = recover_lock(self.actor.lock()).take(); if let Some(handle) = handle { let _ = handle.await; } @@ -260,9 +259,9 @@ impl AuthRuntime { Ok(self.inner()?.admin_key()) } - pub(crate) fn derive_key(&self, key_id: u64) -> Result { + pub(crate) fn derive_key(&self, key_id: KeyId) -> Result { let inner = self.inner()?; - if key_id == 0 { + if key_id.is_admin() { return Ok(inner.admin_key()); } derive_temporary_key(&inner.admin_key(), &inner.instance_id(), key_id) @@ -270,25 +269,13 @@ impl AuthRuntime { #[cfg(test)] pub(crate) fn high_slot_entry_count(&self) -> usize { - self.inner() - .map(|inner| { - inner - .high_slot_entries - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .len() - }) - .unwrap_or(0) + self.inner().map(|inner| inner.high().len()).unwrap_or(0) } - pub(crate) fn derive_previous_key(&self, key_id: u64) -> Option { + pub(crate) fn derive_previous_key(&self, key_id: KeyId) -> Option { let inner = self.inner().ok()?; - let previous = inner - .previous_root - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .clone()?; - if key_id == 0 { + let previous = recover_lock(inner.previous_root.read()).clone()?; + if key_id.is_admin() { Some(previous.admin_key) } else { derive_temporary_key(&previous.admin_key, &previous.instance_id, key_id).ok() @@ -297,15 +284,12 @@ impl AuthRuntime { pub fn authenticate_presented( &self, - key_id: u64, + key_id: KeyId, presented_key: &AesKeyType, ) -> Result { let inner = self.inner()?; - if key_id == 0 { - let admin = inner - .admin - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); + if key_id.is_admin() { + let admin = recover_lock(inner.admin.read()); if !bool::from(presented_key.ct_eq(&admin.key)) { inner.auth_failures.fetch_add(1, Ordering::Relaxed); return Err(AuthFailure::new( @@ -322,7 +306,7 @@ impl AuthRuntime { ) })?; inner.auth_successes.fetch_add(1, Ordering::Relaxed); - return Ok(AuthContext::from_lease(0, true, &lease)); + return Ok(AuthContext::from_lease(ADMIN_KEY_ID, true, &lease)); } if inner.safe_mode.load(Ordering::Acquire) { inner.auth_failures.fetch_add(1, Ordering::Relaxed); @@ -339,12 +323,9 @@ impl AuthRuntime { return Err(temporary_key_material_mismatch(&inner, key_id)); } - let index = key_slot(key_id) as usize; - let generation = key_generation(key_id); - let slots = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); + let index = key_id.slot().as_index(); + let generation = key_id.generation(); + let slots = inner.slots(); let Some(slot) = slots.get(index) else { inner.auth_failures.fetch_add(1, Ordering::Relaxed); return Err(AuthFailure::new( @@ -497,7 +478,7 @@ impl AuthRuntime { pub async fn show( &self, authorization: &AuthContext, - key_id: u64, + key_id: KeyId, reveal: bool, ) -> Result { let authority = authorization.admin_authority()?; @@ -513,7 +494,7 @@ impl AuthRuntime { pub async fn renew( &self, authorization: &AuthContext, - key_id: u64, + key_id: KeyId, ttl: Duration, ) -> Result { let authority = authorization.admin_authority()?; @@ -529,7 +510,7 @@ impl AuthRuntime { pub async fn revoke( &self, authorization: &AuthContext, - key_id: u64, + key_id: KeyId, ) -> Result { let authority = authorization.admin_authority()?; self.request(|response| AuthCommand::Revoke { @@ -599,7 +580,7 @@ impl AuthRuntime { &self, authorization: &AuthContext, action: impl Into, - key_id: Option, + key_id: Option, detail: Option, ) -> Result<(), AuthFailure> { let authority = authorization.admin_authority()?; @@ -615,20 +596,14 @@ impl AuthRuntime { } } -fn temporary_key_material_mismatch(inner: &AuthStateInner, key_id: u64) -> AuthFailure { - let index = key_slot(key_id) as usize; - let generation = key_generation(key_id); - let slots = inner - .slots - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); +fn temporary_key_material_mismatch(inner: &AuthStateInner, key_id: KeyId) -> AuthFailure { + let index = key_id.slot().as_index(); + let generation = key_id.generation(); + let slots = inner.slots(); let current_generation = match slots.get(index) { Some(slot) => Some(slot.generation), None => { - let high = inner - .high_slot_generations - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()); + let high = recover_lock(inner.high_slot_generations.read()); index .checked_sub(slots.len()) .and_then(|offset| high.get(offset).copied()) @@ -646,7 +621,7 @@ fn temporary_key_material_mismatch(inner: &AuthStateInner, key_id: u64) -> AuthF } let current_epoch = inner.root_epoch.load(Ordering::Acquire); if current_epoch > 0 - && generation > 0 + && generation > Generation::FIRST && current_generation.is_some_and(|issued| generation <= issued) { return AuthFailure::new( diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index ea1f3a8..04b4f48 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -21,7 +21,7 @@ fn temp_state_dir(name: &str) -> PathBuf { std::env::temp_dir().join(format!("pb-mapper-{name}-{}", hex(&suffix))) } -fn authenticate_for_test(runtime: &AuthRuntime, key_id: u64) -> Result { +fn authenticate_for_test(runtime: &AuthRuntime, key_id: KeyId) -> Result { let key = runtime.derive_key(key_id)?; runtime.authenticate_presented(key_id, &key) } @@ -63,7 +63,7 @@ async fn shrinking_then_expanding_capacity_does_not_reuse_old_key_ids() { let runtime = AuthRuntime::start(admin_key, config_two.clone()) .await .unwrap(); - let admin = authenticate_for_test(&runtime, 0).unwrap(); + let admin = authenticate_for_test(&runtime, ADMIN_KEY_ID).unwrap(); let first = runtime .issue(&admin, Duration::from_secs(60), Some("first".to_string())) .await @@ -84,8 +84,14 @@ async fn shrinking_then_expanding_capacity_does_not_reuse_old_key_ids() { else { panic!("expected temporary credential"); }; - runtime.revoke(&admin, first_id).await.unwrap(); - runtime.revoke(&admin, second_id).await.unwrap(); + runtime + .revoke(&admin, KeyId::from_u64(first_id)) + .await + .unwrap(); + runtime + .revoke(&admin, KeyId::from_u64(second_id)) + .await + .unwrap(); runtime.gc(&admin).await.unwrap(); drop(runtime); tokio::time::sleep(Duration::from_millis(20)).await; @@ -95,7 +101,7 @@ async fn shrinking_then_expanding_capacity_does_not_reuse_old_key_ids() { ..config_two.clone() }; let runtime = AuthRuntime::start(admin_key, config_one).await.unwrap(); - let admin = authenticate_for_test(&runtime, 0).unwrap(); + let admin = authenticate_for_test(&runtime, ADMIN_KEY_ID).unwrap(); assert!(!runtime.status(&admin).await.unwrap().safe_mode); let _third = runtime .issue(&admin, Duration::from_secs(60), Some("third".to_string())) @@ -105,7 +111,7 @@ async fn shrinking_then_expanding_capacity_does_not_reuse_old_key_ids() { tokio::time::sleep(Duration::from_millis(20)).await; let runtime = AuthRuntime::start(admin_key, config_two).await.unwrap(); - let admin = authenticate_for_test(&runtime, 0).unwrap(); + let admin = authenticate_for_test(&runtime, ADMIN_KEY_ID).unwrap(); assert!(!runtime.status(&admin).await.unwrap().safe_mode); let fourth = runtime .issue(&admin, Duration::from_secs(60), Some("fourth".to_string())) @@ -136,7 +142,7 @@ async fn gc_removes_inactive_high_slot_entries_and_keeps_their_generations() { let runtime = AuthRuntime::start(admin_key, config_two.clone()) .await .unwrap(); - let admin = authenticate_for_test(&runtime, 0).unwrap(); + let admin = authenticate_for_test(&runtime, ADMIN_KEY_ID).unwrap(); let first = runtime .issue(&admin, Duration::from_secs(60), Some("first".to_string())) .await @@ -157,8 +163,14 @@ async fn gc_removes_inactive_high_slot_entries_and_keeps_their_generations() { else { panic!("expected temporary credential"); }; - runtime.revoke(&admin, first_id).await.unwrap(); - runtime.revoke(&admin, second_id).await.unwrap(); + runtime + .revoke(&admin, KeyId::from_u64(first_id)) + .await + .unwrap(); + runtime + .revoke(&admin, KeyId::from_u64(second_id)) + .await + .unwrap(); drop(runtime); tokio::time::sleep(Duration::from_millis(20)).await; @@ -167,7 +179,7 @@ async fn gc_removes_inactive_high_slot_entries_and_keeps_their_generations() { ..config_two.clone() }; let runtime = AuthRuntime::start(admin_key, config_one).await.unwrap(); - let admin = authenticate_for_test(&runtime, 0).unwrap(); + let admin = authenticate_for_test(&runtime, ADMIN_KEY_ID).unwrap(); assert_eq!(runtime.high_slot_entry_count(), 1); let removed = runtime.gc(&admin).await.unwrap(); assert!(removed >= 1); @@ -176,7 +188,7 @@ async fn gc_removes_inactive_high_slot_entries_and_keeps_their_generations() { tokio::time::sleep(Duration::from_millis(20)).await; let runtime = AuthRuntime::start(admin_key, config_two).await.unwrap(); - let admin = authenticate_for_test(&runtime, 0).unwrap(); + let admin = authenticate_for_test(&runtime, ADMIN_KEY_ID).unwrap(); let _low = runtime .issue(&admin, Duration::from_secs(60), Some("low".to_string())) .await @@ -210,7 +222,7 @@ async fn admin_lifecycle_covers_high_slot_keys_after_capacity_shrink() { let runtime = AuthRuntime::start(admin_key, config_two.clone()) .await .unwrap(); - let admin = authenticate_for_test(&runtime, 0).unwrap(); + let admin = authenticate_for_test(&runtime, ADMIN_KEY_ID).unwrap(); let first = runtime .issue(&admin, Duration::from_secs(60), Some("first".to_string())) .await @@ -227,11 +239,11 @@ async fn admin_lifecycle_covers_high_slot_keys_after_capacity_shrink() { ..config_two.clone() }; let runtime = AuthRuntime::start(admin_key, config_one).await.unwrap(); - let admin = authenticate_for_test(&runtime, 0).unwrap(); + let admin = authenticate_for_test(&runtime, ADMIN_KEY_ID).unwrap(); assert_eq!(runtime.high_slot_entry_count(), 1); let high_id = [first.metadata.key_id, second.metadata.key_id] .into_iter() - .find(|key_id| key_slot(*key_id) as usize >= 1) + .find(|key_id| key_id.slot().as_index() >= 1) .expect("one issued key should land above the shrunken table"); let page = runtime.list(&admin, 0, 100).await.unwrap(); assert_eq!(page.items.len(), 2); @@ -257,7 +269,7 @@ async fn admin_lifecycle_covers_high_slot_keys_after_capacity_shrink() { tokio::time::sleep(Duration::from_millis(20)).await; let runtime = AuthRuntime::start(admin_key, config_two).await.unwrap(); - let admin = authenticate_for_test(&runtime, 0).unwrap(); + let admin = authenticate_for_test(&runtime, ADMIN_KEY_ID).unwrap(); let restored = runtime.show(&admin, high_id, false).await.unwrap(); assert_eq!(restored.metadata.state, "revoked"); drop(runtime); @@ -276,7 +288,7 @@ async fn safe_mode_denies_legacy_protocol_instead_of_restoring_the_default() { legacy_protocol: LegacyProtocolPolicy::Allow, }; let runtime = AuthRuntime::start(admin_key, config.clone()).await.unwrap(); - let admin = authenticate_for_test(&runtime, 0).unwrap(); + let admin = authenticate_for_test(&runtime, ADMIN_KEY_ID).unwrap(); runtime .set_legacy_protocol(&admin, LegacyProtocolPolicy::Deny) .await @@ -286,7 +298,7 @@ async fn safe_mode_denies_legacy_protocol_instead_of_restoring_the_default() { std::fs::write(state_dir.join("auth.wal"), b"broken-wal").unwrap(); let runtime = AuthRuntime::start(admin_key, config).await.unwrap(); - let admin = authenticate_for_test(&runtime, 0).unwrap(); + let admin = authenticate_for_test(&runtime, ADMIN_KEY_ID).unwrap(); let status = runtime.status(&admin).await.unwrap(); assert!(status.safe_mode); assert_eq!(status.legacy_protocol, LegacyProtocolPolicy::Deny); @@ -335,7 +347,7 @@ async fn env_recovery_key_is_not_written_when_it_cannot_decrypt_existing_state() let snapshot = PersistedSnapshot { schema_version: SNAPSHOT_SCHEMA_VERSION, instance_id: [4_u8; INSTANCE_ID_LEN], - generations: vec![0; 1], + generations: vec![Generation::FIRST; 1], entries: Vec::new(), legacy_protocol: LegacyProtocolPolicy::Allow, admin_replays: Vec::new(), @@ -479,7 +491,7 @@ async fn reset_clears_retained_high_slot_entries() { let runtime = AuthRuntime::start(admin_key, config_two.clone()) .await .unwrap(); - let admin = authenticate_for_test(&runtime, 0).unwrap(); + let admin = authenticate_for_test(&runtime, ADMIN_KEY_ID).unwrap(); let first = runtime .issue(&admin, Duration::from_secs(60), Some("first".to_string())) .await @@ -496,13 +508,13 @@ async fn reset_clears_retained_high_slot_entries() { ..config_two.clone() }; let runtime = AuthRuntime::start(admin_key, config_one).await.unwrap(); - let admin = authenticate_for_test(&runtime, 0).unwrap(); + let admin = authenticate_for_test(&runtime, ADMIN_KEY_ID).unwrap(); runtime.reset(&admin).await.unwrap(); drop(runtime); tokio::time::sleep(Duration::from_millis(20)).await; let runtime = AuthRuntime::start(admin_key, config_two).await.unwrap(); - let admin = authenticate_for_test(&runtime, 0).unwrap(); + let admin = authenticate_for_test(&runtime, ADMIN_KEY_ID).unwrap(); let page = runtime.list(&admin, 0, 100).await.unwrap(); assert!(page.items.is_empty()); assert!(authenticate_for_test(&runtime, first.metadata.key_id).is_err()); @@ -523,7 +535,7 @@ async fn rotate_root_rejects_a_nul_containing_key() { legacy_protocol: LegacyProtocolPolicy::Allow, }; let runtime = AuthRuntime::start(admin_key, config).await.unwrap(); - let admin = authenticate_for_test(&runtime, 0).unwrap(); + let admin = authenticate_for_test(&runtime, ADMIN_KEY_ID).unwrap(); let mut bad = *b"0123456789abcdefghijklmnopqrstuv"; bad[4] = 0; let error = runtime.rotate_root(&admin, bad).await.unwrap_err(); @@ -628,11 +640,25 @@ fn legacy_protocol_policy_trims_valid_values_and_rejects_unknown_values() { assert_eq!(parse_legacy_protocol_policy(""), None); } +#[test] +fn key_id_serializes_as_a_plain_integer() { + let key_id = KeyId::new(Generation::from_u32(3), SlotIndex::from_index(2)); + assert_eq!(serde_json::to_string(&key_id).unwrap(), "12884901890"); + assert_eq!( + serde_json::from_str::("12884901890").unwrap(), + key_id + ); + assert_eq!( + serde_json::to_string(&Generation::from_u32(7)).unwrap(), + "7" + ); +} + #[test] fn key_id_round_trip() { - let key_id = make_key_id(42, 65_535); - assert_eq!(key_generation(key_id), 42); - assert_eq!(key_slot(key_id), 65_535); + let key_id = KeyId::new(Generation::from_u32(42), SlotIndex::from_index(65_535)); + assert_eq!(key_id.generation(), Generation::from_u32(42)); + assert_eq!(key_id.slot(), SlotIndex::from_index(65_535)); } #[test] @@ -640,18 +666,38 @@ fn derived_key_is_bound_to_instance_and_key_id() { let admin_key = *b"0123456789abcdefghijklmnopqrstuv"; let instance_a = [1_u8; INSTANCE_ID_LEN]; let instance_b = [2_u8; INSTANCE_ID_LEN]; - let key = derive_temporary_key(&admin_key, &instance_a, make_key_id(1, 7)).unwrap(); + let key = derive_temporary_key( + &admin_key, + &instance_a, + KeyId::new(Generation::from_u32(1), SlotIndex::from_index(7)), + ) + .unwrap(); assert_eq!( key, - derive_temporary_key(&admin_key, &instance_a, make_key_id(1, 7)).unwrap() + derive_temporary_key( + &admin_key, + &instance_a, + KeyId::new(Generation::from_u32(1), SlotIndex::from_index(7)) + ) + .unwrap() ); assert_ne!( key, - derive_temporary_key(&admin_key, &instance_b, make_key_id(1, 7)).unwrap() + derive_temporary_key( + &admin_key, + &instance_b, + KeyId::new(Generation::from_u32(1), SlotIndex::from_index(7)) + ) + .unwrap() ); assert_ne!( key, - derive_temporary_key(&admin_key, &instance_a, make_key_id(2, 7)).unwrap() + derive_temporary_key( + &admin_key, + &instance_a, + KeyId::new(Generation::from_u32(2), SlotIndex::from_index(7)) + ) + .unwrap() ); } @@ -659,9 +705,10 @@ fn derived_key_is_bound_to_instance_and_key_id() { async fn isolated_runtime_preserves_remote_temporary_process_credential() { let _process_credential_guard = PROCESS_CREDENTIAL_TEST_LOCK.lock().await; let state_dir = temp_state_dir("isolated-relay"); - let temporary_key_id = make_key_id(1, 0); + let temporary_key_id = KeyId::new(Generation::from_u32(1), SlotIndex::from_index(0)); let temporary_key = *b"temporary-remote-key-0123456789a"; - let temporary_credential = encode_temporary_credential(temporary_key_id, &temporary_key); + let temporary_credential = + encode_temporary_credential(temporary_key_id.as_u64(), &temporary_key); set_process_msg_header_key(Some(&temporary_credential)).unwrap(); let config = AuthConfig { state_dir: state_dir.clone(), @@ -674,7 +721,7 @@ async fn isolated_runtime_preserves_remote_temporary_process_credential() { assert_eq!( get_process_credential().unwrap(), Credential::Temporary { - key_id: temporary_key_id, + key_id: temporary_key_id.as_u64(), key: temporary_key, } ); @@ -684,7 +731,9 @@ async fn isolated_runtime_preserves_remote_temporary_process_credential() { else { panic!("isolated relay key should be an administrator credential"); }; - let local_admin = runtime.authenticate_presented(0, &local_admin_key).unwrap(); + let local_admin = runtime + .authenticate_presented(ADMIN_KEY_ID, &local_admin_key) + .unwrap(); runtime .rotate_root(&local_admin, *b"isolated-new-admin-key-012345678") .await @@ -692,7 +741,7 @@ async fn isolated_runtime_preserves_remote_temporary_process_credential() { assert_eq!( get_process_credential().unwrap(), Credential::Temporary { - key_id: temporary_key_id, + key_id: temporary_key_id.as_u64(), key: temporary_key, } ); @@ -714,7 +763,7 @@ async fn issue_renew_revoke_and_persist() { legacy_protocol: LegacyProtocolPolicy::Allow, }; let runtime = AuthRuntime::start(admin_key, config.clone()).await.unwrap(); - let admin = authenticate_for_test(&runtime, 0).unwrap(); + let admin = authenticate_for_test(&runtime, ADMIN_KEY_ID).unwrap(); let issued = runtime .issue(&admin, Duration::from_secs(60), Some("demo".to_string())) .await @@ -788,7 +837,7 @@ async fn ensure_active_keeps_expiry_after_the_lease_is_cancelled() { legacy_protocol: LegacyProtocolPolicy::Allow, }; let runtime = AuthRuntime::start(admin_key, config).await.unwrap(); - let admin = authenticate_for_test(&runtime, 0).unwrap(); + let admin = authenticate_for_test(&runtime, ADMIN_KEY_ID).unwrap(); let issued = runtime .issue(&admin, Duration::from_secs(60), Some("exp".to_string())) .await @@ -820,7 +869,7 @@ async fn renew_replaces_a_lease_canceled_during_persistence() { legacy_protocol: LegacyProtocolPolicy::Allow, }; let runtime = AuthRuntime::start(admin_key, config).await.unwrap(); - let admin = authenticate_for_test(&runtime, 0).unwrap(); + let admin = authenticate_for_test(&runtime, ADMIN_KEY_ID).unwrap(); let issued = runtime .issue(&admin, Duration::from_secs(60), Some("renew".to_string())) .await @@ -855,7 +904,7 @@ async fn reset_rotates_instance_and_prevents_old_key_id_reuse() { legacy_protocol: LegacyProtocolPolicy::Allow, }; let runtime = AuthRuntime::start(admin_key, config).await.unwrap(); - let admin = authenticate_for_test(&runtime, 0).unwrap(); + let admin = authenticate_for_test(&runtime, ADMIN_KEY_ID).unwrap(); let before = runtime.status(&admin).await.unwrap().server_instance_id; let old = runtime .issue( @@ -915,7 +964,7 @@ fn recover_instance_id_promotes_next_when_snapshot_matches() { let snapshot = PersistedSnapshot { schema_version: SNAPSHOT_SCHEMA_VERSION, instance_id: next, - generations: vec![0; 1], + generations: vec![Generation::FIRST; 1], entries: Vec::new(), legacy_protocol: LegacyProtocolPolicy::Allow, admin_replays: Vec::new(), @@ -952,7 +1001,7 @@ fn reset_already_installed_accepts_matching_live_id_and_snapshot() { let snapshot = PersistedSnapshot { schema_version: SNAPSHOT_SCHEMA_VERSION, instance_id: new_id, - generations: vec![0; 1], + generations: vec![Generation::FIRST; 1], entries: Vec::new(), legacy_protocol: LegacyProtocolPolicy::Allow, admin_replays: Vec::new(), @@ -987,7 +1036,7 @@ fn recover_instance_id_discards_stale_next_when_snapshot_still_matches_current() let snapshot = PersistedSnapshot { schema_version: SNAPSHOT_SCHEMA_VERSION, instance_id: current, - generations: vec![0; 1], + generations: vec![Generation::FIRST; 1], entries: Vec::new(), legacy_protocol: LegacyProtocolPolicy::Allow, admin_replays: Vec::new(), @@ -1021,7 +1070,7 @@ fn recover_admin_key_discards_leftover_wal_from_the_old_key() { let snapshot = PersistedSnapshot { schema_version: SNAPSHOT_SCHEMA_VERSION, instance_id: [9_u8; INSTANCE_ID_LEN], - generations: vec![0; 1], + generations: vec![Generation::FIRST; 1], entries: Vec::new(), legacy_protocol: LegacyProtocolPolicy::Allow, admin_replays: Vec::new(), @@ -1056,7 +1105,7 @@ fn rotation_finalize_requires_the_live_admin_key() { let snapshot = PersistedSnapshot { schema_version: SNAPSHOT_SCHEMA_VERSION, instance_id: [3_u8; INSTANCE_ID_LEN], - generations: vec![0; 1], + generations: vec![Generation::FIRST; 1], entries: Vec::new(), legacy_protocol: LegacyProtocolPolicy::Allow, admin_replays: Vec::new(), @@ -1081,7 +1130,7 @@ async fn interrupted_reset_recovers_the_staged_instance_id_on_restart() { legacy_protocol: LegacyProtocolPolicy::Allow, }; let runtime = AuthRuntime::start(admin_key, config.clone()).await.unwrap(); - let admin = authenticate_for_test(&runtime, 0).unwrap(); + let admin = authenticate_for_test(&runtime, ADMIN_KEY_ID).unwrap(); let issued = runtime .issue( &admin, @@ -1099,7 +1148,7 @@ async fn interrupted_reset_recovers_the_staged_instance_id_on_restart() { let snapshot = PersistedSnapshot { schema_version: SNAPSHOT_SCHEMA_VERSION, instance_id: next, - generations: vec![0; 4], + generations: vec![Generation::FIRST; 4], entries: Vec::new(), legacy_protocol: LegacyProtocolPolicy::Allow, admin_replays: Vec::new(), @@ -1116,7 +1165,7 @@ async fn interrupted_reset_recovers_the_staged_instance_id_on_restart() { .unwrap(); let restored = AuthRuntime::start(admin_key, config).await.unwrap(); - let restored_admin = authenticate_for_test(&restored, 0).unwrap(); + let restored_admin = authenticate_for_test(&restored, ADMIN_KEY_ID).unwrap(); let status = restored.status(&restored_admin).await.unwrap(); assert!(!status.safe_mode); assert_eq!(status.server_instance_id, hex(&next)); @@ -1138,7 +1187,7 @@ async fn corrupt_wal_fails_temporary_keys_closed_until_admin_reset() { legacy_protocol: LegacyProtocolPolicy::Allow, }; let runtime = AuthRuntime::start(admin_key, config.clone()).await.unwrap(); - let admin = authenticate_for_test(&runtime, 0).unwrap(); + let admin = authenticate_for_test(&runtime, ADMIN_KEY_ID).unwrap(); let issued = runtime .issue( &admin, @@ -1152,7 +1201,7 @@ async fn corrupt_wal_fails_temporary_keys_closed_until_admin_reset() { std::fs::write(state_dir.join("auth.wal"), b"broken-wal").unwrap(); let recovered = AuthRuntime::start(admin_key, config).await.unwrap(); - let recovered_admin = authenticate_for_test(&recovered, 0).unwrap(); + let recovered_admin = authenticate_for_test(&recovered, ADMIN_KEY_ID).unwrap(); assert!(recovered.status(&recovered_admin).await.unwrap().safe_mode); assert_eq!( authenticate_for_test(&recovered, issued.metadata.key_id) @@ -1181,7 +1230,9 @@ async fn root_rotation_rejects_old_key_and_in_flight_admin_context() { legacy_protocol: LegacyProtocolPolicy::Allow, }; let runtime = AuthRuntime::start(old_key, config).await.unwrap(); - let old_admin = runtime.authenticate_presented(0, &old_key).unwrap(); + let old_admin = runtime + .authenticate_presented(ADMIN_KEY_ID, &old_key) + .unwrap(); let issued = runtime .issue( &old_admin, @@ -1203,7 +1254,7 @@ async fn root_rotation_rejects_old_key_and_in_flight_admin_context() { let mistyped_key = *b"1123456789abcdefghijklmnopqrstuv"; assert_eq!( runtime - .authenticate_presented(0, &mistyped_key) + .authenticate_presented(ADMIN_KEY_ID, &mistyped_key) .unwrap_err() .code, "administrator_key_invalid" @@ -1220,7 +1271,9 @@ async fn root_rotation_rejects_old_key_and_in_flight_admin_context() { .code, "temporary_key_rotated" ); - let new_admin = runtime.authenticate_presented(0, &new_key).unwrap(); + let new_admin = runtime + .authenticate_presented(ADMIN_KEY_ID, &new_key) + .unwrap(); let _replacement = runtime .issue( &new_admin, @@ -1239,7 +1292,7 @@ async fn root_rotation_rejects_old_key_and_in_flight_admin_context() { assert_eq!( runtime - .authenticate_presented(0, &old_key) + .authenticate_presented(ADMIN_KEY_ID, &old_key) .unwrap_err() .code, "administrator_key_invalid" @@ -1272,7 +1325,7 @@ async fn admitted_admin_mutation_replay_survives_restart() { let fingerprint = [0x5a; 32]; let timestamp = unix_seconds(); let runtime = AuthRuntime::start(admin_key, config.clone()).await.unwrap(); - let admin = authenticate_for_test(&runtime, 0).unwrap(); + let admin = authenticate_for_test(&runtime, ADMIN_KEY_ID).unwrap(); runtime .claim_admin_mutation(&admin, fingerprint, timestamp) .await @@ -1281,7 +1334,7 @@ async fn admitted_admin_mutation_replay_survives_restart() { tokio::time::sleep(Duration::from_millis(20)).await; let restored = AuthRuntime::start(admin_key, config).await.unwrap(); - let restored_admin = authenticate_for_test(&restored, 0).unwrap(); + let restored_admin = authenticate_for_test(&restored, ADMIN_KEY_ID).unwrap(); assert_eq!( restored .claim_admin_mutation(&restored_admin, fingerprint, timestamp) @@ -1307,7 +1360,7 @@ async fn snapshot_compaction_preserves_audit_records() { legacy_protocol: LegacyProtocolPolicy::Allow, }; let runtime = AuthRuntime::start(admin_key, config.clone()).await.unwrap(); - let admin = authenticate_for_test(&runtime, 0).unwrap(); + let admin = authenticate_for_test(&runtime, ADMIN_KEY_ID).unwrap(); runtime .issue(&admin, Duration::from_secs(60), Some("audited".to_string())) .await @@ -1329,84 +1382,240 @@ async fn snapshot_compaction_preserves_audit_records() { let _ = std::fs::remove_dir_all(state_dir); } +#[tokio::test] +async fn revoking_keeps_the_row_until_its_retention_elapses() { + let state_dir = temp_state_dir("revoke-retention"); + let admin_key = *b"0123456789abcdefghijklmnopqrstuv"; + let config = AuthConfig { + state_dir: state_dir.clone(), + max_temporary_keys: 4, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + }; + let runtime = AuthRuntime::start(admin_key, config).await.unwrap(); + let admin = authenticate_for_test(&runtime, ADMIN_KEY_ID).unwrap(); + let issued = runtime + .issue(&admin, Duration::from_secs(60), Some("revoked".to_string())) + .await + .unwrap(); + let key_id = issued.metadata.key_id; + let presented = runtime.derive_key(key_id).unwrap(); + + runtime.revoke(&admin, key_id).await.unwrap(); + + // The credential stops working at once, but the row survives so the reason + // is still reportable rather than degrading to "unknown key". + assert_eq!( + runtime + .authenticate_presented(key_id, &presented) + .unwrap_err() + .code, + "temporary_key_revoked" + ); + assert!(runtime + .list(&admin, 0, 100) + .await + .unwrap() + .items + .iter() + .any(|item| item.key_id == key_id && item.state == "revoked")); + + drop(runtime); + tokio::time::sleep(Duration::from_millis(20)).await; + let _ = std::fs::remove_dir_all(state_dir); +} + +/// Records which phases ran, so a test can assert on the callback's effects +/// rather than on a return value the wheel no longer produces. +#[derive(Clone, Default)] +struct PhaseLog(Arc>>); + +impl PhaseLog { + fn push(&self, phase: &'static str) { + recover_lock(self.0.lock()).push(phase); + } + + fn phases(&self) -> Vec<&'static str> { + recover_lock(self.0.lock()).clone() + } +} + +/// Schedules the two-phase shape `Leases` uses: a first phase that asks for a +/// second one `gap` seconds later, then a final phase. +fn schedule_two_phases( + wheel: &mut TimingWheel, + key_id: KeyId, + deadline: u64, + gap: u64, +) -> PhaseLog { + let log = PhaseLog::default(); + let recorder = log.clone(); + let mut first = true; + wheel.schedule(key_id, deadline, move || { + if std::mem::take(&mut first) { + recorder.push("retire"); + return Some(deadline + gap); + } + recorder.push("reap"); + None + }); + log +} + +fn schedule_once(wheel: &mut TimingWheel, key_id: KeyId, deadline: u64) -> PhaseLog { + let log = PhaseLog::default(); + let recorder = log.clone(); + wheel.schedule(key_id, deadline, move || { + recorder.push("fired"); + None + }); + log +} + #[test] -fn timing_wheel_ignores_stale_renewal_entry() { - let now = 1_000; - let lease = Arc::new(AuthLease::new(make_key_id(1, 0), now + 5)); - let mut wheel = TimingWheel::new(now); - wheel.insert(lease.clone()); - lease.expires_at.store(now + 20, Ordering::Release); - lease.wheel_version.fetch_add(1, Ordering::AcqRel); - wheel.insert(lease.clone()); - assert!(wheel.advance(now + 6).is_empty()); - assert_eq!(wheel.advance(now + 20).len(), 1); +fn timing_wheel_runs_the_next_phase_at_the_deadline_it_asked_for() { + let key_id = KeyId::new(Generation::from_u32(1), SlotIndex::from_index(0)); + let mut wheel = TimingWheel::new(1_000); + let log = schedule_two_phases(&mut wheel, key_id, 1_005, 60); + + wheel.advance(1_004); + assert!(log.phases().is_empty()); + wheel.advance(1_005); + assert_eq!(log.phases(), ["retire"]); + // The second phase waits for the deadline the first one returned. + wheel.advance(1_064); + assert_eq!(log.phases(), ["retire"]); + wheel.advance(1_065); + assert_eq!(log.phases(), ["retire", "reap"]); + assert!(!wheel.holds(key_id)); } #[test] -fn timing_wheel_fast_forwards_large_clock_jumps() { - let now = 1_000; - let target = now + 7 * 24 * 60 * 60; - let expired = Arc::new(AuthLease::new(make_key_id(1, 0), now + 5)); - let future = Arc::new(AuthLease::new(make_key_id(1, 1), target + 20)); - let mut wheel = TimingWheel::new(now); - wheel.insert(expired.clone()); - wheel.insert(future.clone()); - - let due = wheel.advance(target); - assert_eq!(due.len(), 1); - assert_eq!(due[0].key_id(), expired.key_id()); - assert!(wheel.advance(target + 19).is_empty()); - assert_eq!(wheel.advance(target + 20).len(), 1); +fn timing_wheel_cancel_runs_every_remaining_phase_at_once() { + let key_id = KeyId::new(Generation::from_u32(1), SlotIndex::from_index(0)); + let mut wheel = TimingWheel::new(1_000); + let log = schedule_two_phases(&mut wheel, key_id, 1_005, 60); + + wheel.cancel(key_id); + assert_eq!(log.phases(), ["retire", "reap"]); + assert!(!wheel.holds(key_id)); } #[test] -fn timing_wheel_returns_already_expired_insert_without_wrapping() { - let now = 1_000; - let expired = Arc::new(AuthLease::new(make_key_id(1, 0), now - 1)); - let mut wheel = TimingWheel::new(now); - wheel.insert(expired.clone()); +fn timing_wheel_cancel_after_the_first_phase_runs_only_the_rest() { + let key_id = KeyId::new(Generation::from_u32(1), SlotIndex::from_index(0)); + let mut wheel = TimingWheel::new(1_000); + let log = schedule_two_phases(&mut wheel, key_id, 1_005, 60); + + wheel.advance(1_005); + wheel.cancel(key_id); + assert_eq!(log.phases(), ["retire", "reap"]); +} + +#[test] +fn timing_wheel_drop_runs_every_remaining_phase() { + let key_id = KeyId::new(Generation::from_u32(1), SlotIndex::from_index(0)); + let mut wheel = TimingWheel::new(1_000); + let log = schedule_two_phases(&mut wheel, key_id, 1_005, 60); - let due = wheel.advance(now); - assert_eq!(due.len(), 1); - assert_eq!(due[0].key_id(), expired.key_id()); + drop(wheel); + assert_eq!(log.phases(), ["retire", "reap"]); } #[test] -fn timing_wheel_expires_cascaded_boundary_entry_without_an_extra_tick() { - let now = 700; - let expires_at = 1_024; - let lease = Arc::new(AuthLease::new(make_key_id(1, 0), expires_at)); - let mut wheel = TimingWheel::new(now); - wheel.insert(lease.clone()); +fn timing_wheel_reschedule_moves_an_entry_without_running_it() { + let key_id = KeyId::new(Generation::from_u32(1), SlotIndex::from_index(0)); + let mut wheel = TimingWheel::new(1_000); + let log = schedule_once(&mut wheel, key_id, 1_005); + + assert!(wheel.reschedule(key_id, 1_020)); + wheel.advance(1_019); + assert!(log.phases().is_empty()); + wheel.advance(1_020); + assert_eq!(log.phases(), ["fired"]); +} - let due = wheel.advance(expires_at); - assert_eq!(due.len(), 1); - assert_eq!(due[0].key_id(), lease.key_id()); +#[test] +fn timing_wheel_reschedule_reports_a_key_it_does_not_hold() { + let mut wheel = TimingWheel::new(1_000); + assert!(!wheel.reschedule( + KeyId::new(Generation::from_u32(1), SlotIndex::from_index(0)), + 2_000 + )); +} + +#[test] +fn timing_wheel_scheduling_over_an_entry_finishes_the_one_it_replaces() { + let key_id = KeyId::new(Generation::from_u32(1), SlotIndex::from_index(0)); + let mut wheel = TimingWheel::new(1_000); + let stale = schedule_two_phases(&mut wheel, key_id, 1_005, 60); + let fresh = schedule_once(&mut wheel, key_id, 1_020); + + assert_eq!(stale.phases(), ["retire", "reap"]); + wheel.advance(1_020); + assert_eq!(fresh.phases(), ["fired"]); } #[test] -fn timing_wheel_release_drops_the_current_owner() { +fn timing_wheel_fires_an_entry_scheduled_in_the_past_without_wrapping() { + let key_id = KeyId::new(Generation::from_u32(1), SlotIndex::from_index(0)); + let mut wheel = TimingWheel::new(1_000); + let log = schedule_once(&mut wheel, key_id, 999); + + wheel.advance(1_000); + assert_eq!(log.phases(), ["fired"]); +} + +#[test] +fn timing_wheel_cascades_an_entry_down_two_levels() { + // A deadline two levels up has to reach level 0 before it can be drained. let now = 1_000; - let lease = Arc::new(AuthLease::new(make_key_id(1, 0), now + 60)); + let deadline = now + (1 << (6 * 2)); + let key_id = KeyId::new(Generation::from_u32(1), SlotIndex::from_index(0)); let mut wheel = TimingWheel::new(now); - wheel.insert(lease.clone()); - assert!(wheel.owns(lease.key_id())); + let log = schedule_once(&mut wheel, key_id, deadline); - wheel.release(lease.key_id()); - assert!(!wheel.owns(lease.key_id())); - assert!(!lease.cancellation_token().is_cancelled()); + wheel.advance(deadline - 1); + assert!(log.phases().is_empty()); + wheel.advance(deadline); + assert_eq!(log.phases(), ["fired"]); } #[test] -fn timing_wheel_clear_cancels_owned_leases() { +fn timing_wheel_fires_a_boundary_deadline_without_an_extra_tick() { + let key_id = KeyId::new(Generation::from_u32(1), SlotIndex::from_index(0)); + let mut wheel = TimingWheel::new(700); + let log = schedule_once(&mut wheel, key_id, 1_024); + + wheel.advance(1_024); + assert_eq!(log.phases(), ["fired"]); +} + +#[test] +fn timing_wheel_drains_in_one_pass_when_the_clock_jumps_past_every_deadline() { + // A correction longer than the longest schedulable delay leaves nothing + // pending, so the wheel empties without ticking through the elapsed years. let now = 1_000; - let lease = Arc::new(AuthLease::new(make_key_id(1, 0), now + 60)); - let cancellation = lease.cancellation_token(); + let key_id = KeyId::new(Generation::from_u32(1), SlotIndex::from_index(0)); let mut wheel = TimingWheel::new(now); - wheel.insert(lease); + let log = schedule_two_phases(&mut wheel, key_id, now + 60, 60); - wheel.clear(now + 1); - assert!(cancellation.is_cancelled()); + wheel.advance(now + 4 * MAX_TEMP_KEY_TTL.as_secs()); + assert_eq!(log.phases(), ["retire", "reap"]); + assert!(!wheel.holds(key_id)); +} + +#[test] +fn timing_wheel_keeps_its_position_when_the_clock_steps_backwards() { + let key_id = KeyId::new(Generation::from_u32(1), SlotIndex::from_index(0)); + let mut wheel = TimingWheel::new(1_000); + let log = schedule_once(&mut wheel, key_id, 1_030); + + wheel.advance(1_020); + wheel.advance(1_005); + assert!(log.phases().is_empty()); + wheel.advance(1_030); + assert_eq!(log.phases(), ["fired"]); } #[test] @@ -1484,7 +1693,7 @@ fn replay_pruning_falls_back_to_client_timestamp_for_legacy_records() { fn tombstone_migration_prefers_audit_time_and_persists_fail_closed_fallback() { let now = 10_000; let revoked_with_audit = PersistedEntry { - key_id: make_key_id(1, 0), + key_id: KeyId::new(Generation::from_u32(1), SlotIndex::from_index(0)), state: SlotState::Revoked, issued_at: 100, expires_at: 20_000, @@ -1492,7 +1701,7 @@ fn tombstone_migration_prefers_audit_time_and_persists_fail_closed_fallback() { tombstoned_at: None, }; let revoked_without_audit = PersistedEntry { - key_id: make_key_id(1, 1), + key_id: KeyId::new(Generation::from_u32(1), SlotIndex::from_index(1)), state: SlotState::Revoked, issued_at: 100, expires_at: 20_000, @@ -1503,14 +1712,17 @@ fn tombstone_migration_prefers_audit_time_and_persists_fail_closed_fallback() { let mut snapshot = PersistedSnapshot { schema_version: SNAPSHOT_SCHEMA_VERSION, instance_id: [1; INSTANCE_ID_LEN], - generations: vec![1, 1], + generations: vec![Generation::from_u32(1), Generation::from_u32(1)], entries: vec![revoked_with_audit, revoked_without_audit], legacy_protocol: LegacyProtocolPolicy::Deny, admin_replays: Vec::new(), audit_records: VecDeque::from([AuditRecord { at: audit_at, action: "temporary_key_revoke".to_string(), - key_id: Some(make_key_id(1, 0)), + key_id: Some(KeyId::new( + Generation::from_u32(1), + SlotIndex::from_index(0), + )), label: None, }]), root_epoch: 0, diff --git a/src/common/auth/timing_wheel.rs b/src/common/auth/timing_wheel.rs index 950e1f4..10ca9ae 100644 --- a/src/common/auth/timing_wheel.rs +++ b/src/common/auth/timing_wheel.rs @@ -1,206 +1,314 @@ -//! Hierarchical expiry scheduler for temporary credential leases. +//! Hierarchical timer wheel whose entries are self-advancing cleanup callbacks. //! //! ```text -//! lease(expires_at) -> level/slot bucket -> one-second actor tick -> expired leases -//! renew ----> version bump ------^ stale bucket entries are ignored -//! | -//! large clock jump -> bounded bucket scan + rebuild (never second-by-second catch-up) -//! overdue insert --> immediate-due queue --> next advance, without a wheel revolution -//! reset/rotation -> cancel every wheel-owned lease -> clear all buckets +//! schedule(deadline, callback) -> one level/slot bucket +//! deadline reached -> callback runs -> reschedules itself, or is done +//! dropped early -> callback runs to completion right there +//! +//! one-second tick -> drain level 0's current slot +//! -> every 64th tick, cascade level 1 into finer levels, and so on //! ``` //! -//! The wheel's current-owner map holds the strong `Arc` for each key. -//! Bucket entries are `Weak`, so a renew replaces the previous owner instead of -//! accumulating day-long stale strong references. Request-facing structures also -//! retain only `Weak` references. +//! The wheel knows nothing about what it schedules. An entry owns a callback that +//! performs one phase of some cleanup and returns when it next wants to run, or +//! `None` when finished, so a multi-stage teardown is written once at the point +//! the entry is created. Dropping an entry early runs every phase it has left, in +//! order, immediately — which is what lets cancelling one entry stand in for a +//! whole cleanup routine, and lets dropping the wheel tear down everything it +//! owns without a caller walking any of it. +//! +//! Each entry lives in exactly one bucket, and `positions` records which, so +//! rescheduling or cancelling one is a map lookup rather than a search. A tick +//! touches one slot per level that turns over, never the whole wheel, so cost +//! tracks the entries that actually cascade or fire. use super::*; -const MAX_INCREMENTAL_ADVANCE_SECONDS: u64 = 256; +/// Bits of a deadline that one level's slot field covers, so a level holds +/// `1 << 6` slots and each is 64 times coarser than the level below it. +const SLOT_BITS: u32 = 6; +const SLOTS: usize = 1 << SLOT_BITS; +const SLOT_MASK: u64 = SLOTS as u64 - 1; +/// Six 64-slot levels span `64^6` seconds, which keeps every deadline a +/// configured TTL can produce inside a level whose range really contains it. +const NUM_LEVELS: usize = 6; +const TOP_LEVEL: Level = NUM_LEVELS as Level - 1; + +/// Which level of the hierarchy a bucket belongs to: `0..NUM_LEVELS`. +type Level = u8; + +/// Which bucket within one level: `0..SLOTS`. +type Slot = u8; + +/// One phase of a scheduled teardown: does its work and returns the deadline of +/// the phase after it, or `None` once nothing remains. +type Phase = Box Option + Send>; + +struct Entry { + deadline: u64, + phase: Phase, + /// Set once a phase has returned `None`, so a completed entry's drop does + /// not call into its callback again. + finished: bool, +} + +impl Entry { + /// Runs the next phase and reports the deadline it wants, if any. + fn fire(&mut self) -> Option { + if self.finished { + return None; + } + let next = (self.phase)(); + self.finished = next.is_none(); + next + } +} + +impl Drop for Entry { + /// An entry let go of before its deadline still owes every phase it has + /// left, so they all run here. This is why cancelling an entry and letting + /// it expire have the same effect, only sooner. + fn drop(&mut self) { + while self.fire().is_some() {} + } +} -struct WheelEntry { - lease: Weak, - version: u64, +/// Where a key's entry currently sits, so a reschedule or cancel does not have +/// to search the wheel. +#[derive(Clone, Copy)] +enum Position { + /// Scheduled for a deadline the wheel had already passed. + Overdue, + Wheel { + level: Level, + slot: Slot, + }, } +type Bucket = HashMap; + pub(super) struct TimingWheel { now: u64, - owners: HashMap>, - immediate_due: Vec, - level0: Vec>, - level1: Vec>, - level2: Vec>, - level3: Vec>, + positions: HashMap, + /// Entries whose deadline was already past when they were filed. Level 0's + /// slot for `now` was drained this tick, so filing them there would delay + /// them by a full revolution. + overdue: Bucket, + levels: [Vec; NUM_LEVELS], } impl TimingWheel { pub(super) fn new(now: u64) -> Self { Self { now, - owners: HashMap::new(), - immediate_due: Vec::new(), - level0: empty_buckets(256), - level1: empty_buckets(64), - level2: empty_buckets(64), - level3: empty_buckets(64), + positions: HashMap::new(), + overdue: Bucket::new(), + levels: std::array::from_fn(|_| { + std::iter::repeat_with(Bucket::new).take(SLOTS).collect() + }), } } - pub(super) fn insert(&mut self, lease: Arc) { - let version = lease.wheel_version.load(Ordering::Acquire); - self.insert_with_version(lease, version); + /// Schedules `phase` to run once `deadline` has passed. Any entry already + /// held for `key_id` is dropped, which runs the phases it had left. + pub(super) fn schedule( + &mut self, + key_id: KeyId, + deadline: u64, + phase: impl FnMut() -> Option + Send + 'static, + ) { + self.cancel(key_id); + self.place( + key_id, + Entry { + deadline, + phase: Box::new(phase), + finished: false, + }, + ); } - pub(super) fn release(&mut self, key_id: u64) { - self.owners.remove(&key_id); + /// Replaces the entry for `key_id`, discarding the previous one *without* + /// running its remaining phases. Use this only when the new entry takes over + /// the same cleanup, so nothing the old one owed is lost; otherwise + /// [`Self::schedule`] is what you want. + pub(super) fn supersede( + &mut self, + key_id: KeyId, + deadline: u64, + phase: impl FnMut() -> Option + Send + 'static, + ) { + if let Some(mut previous) = self.detach(key_id) { + previous.finished = true; + } + self.place( + key_id, + Entry { + deadline, + phase: Box::new(phase), + finished: false, + }, + ); } - #[cfg(test)] - pub(super) fn owns(&self, key_id: u64) -> bool { - self.owners.contains_key(&key_id) + /// Moves an entry to a new deadline without running anything. Returns + /// `false` when the wheel holds no entry for `key_id`. + pub(super) fn reschedule(&mut self, key_id: KeyId, deadline: u64) -> bool { + let Some(mut entry) = self.detach(key_id) else { + return false; + }; + entry.deadline = deadline; + self.place(key_id, entry); + true } - pub(super) fn insert_with_version(&mut self, lease: Arc, version: u64) { - self.owners.insert(lease.key_id(), lease.clone()); - let expires_at = lease.expires_at(); - let delta = expires_at.saturating_sub(self.now); - let entry = WheelEntry { - lease: Arc::downgrade(&lease), - version, - }; - if expires_at <= self.now { - self.immediate_due.push(entry); - } else if delta < 1 << 8 { - self.level0[(expires_at & 0xff) as usize].push(entry); - } else if delta < 1 << 14 { - self.level1[((expires_at >> 8) & 0x3f) as usize].push(entry); - } else if delta < 1 << 20 { - self.level2[((expires_at >> 14) & 0x3f) as usize].push(entry); - } else { - self.level3[((expires_at >> 20) & 0x3f) as usize].push(entry); - } + /// Runs everything the entry for `key_id` still owes, now rather than at its + /// deadline. A no-op when the wheel holds no entry for it. + pub(super) fn cancel(&mut self, key_id: KeyId) { + drop(self.detach(key_id)); } - pub(super) fn advance(&mut self, target: u64) -> Vec> { - if target.saturating_sub(self.now) > MAX_INCREMENTAL_ADVANCE_SECONDS { - return self.fast_forward(target); + /// Runs only the entry's next phase, then waits for the deadline that phase + /// asked for. Use this where a stage has arrived early but the stages after + /// it must still keep their own timing — a revoke ends a key's active phase + /// without skipping the retention that follows it. + pub(super) fn advance_one_phase(&mut self, key_id: KeyId) -> bool { + let Some(mut entry) = self.detach(key_id) else { + return false; + }; + match entry.fire() { + Some(next) => { + entry.deadline = next; + self.place(key_id, entry); + } + // Finished, so the drop below has nothing left to run. + None => drop(entry), } + true + } - let mut due = self.take_immediate_due(target); - while self.now < target { - self.now = self.now.saturating_add(1); - if self.now & 0xff == 0 { - self.cascade(1); - if (self.now >> 8) & 0x3f == 0 { - self.cascade(2); - if (self.now >> 14) & 0x3f == 0 { - self.cascade(3); - } + /// Runs the wheel up to `target`, firing every entry whose deadline has + /// passed and re-filing the phases they schedule next. + pub(super) fn advance(&mut self, target: u64) { + let overdue = std::mem::take(&mut self.overdue); + self.settle(overdue, target); + // A jump longer than the longest lifetime the config can produce means + // every entry is already past its deadline, so the whole wheel can be + // drained in one pass. Ticking through it instead would spin for hours + // when a bad hardware clock is corrected forward by years. + if target.saturating_sub(self.now) > MAX_SCHEDULABLE_DELAY.as_secs() { + self.now = target; + // A phase can schedule a successor that is also already overdue, so + // keep draining until a pass leaves nothing due. + loop { + let due = self + .levels + .iter_mut() + .flat_map(|level| level.iter_mut()) + .fold(Bucket::new(), |mut due, bucket| { + due.extend(std::mem::take(bucket)); + due + }); + let overdue = std::mem::take(&mut self.overdue); + if due.is_empty() && overdue.is_empty() { + return; } + self.settle(due, target); + self.settle(overdue, target); } - due.extend(self.take_immediate_due(self.now)); - let index = (self.now & 0xff) as usize; - for entry in std::mem::take(&mut self.level0[index]) { - let Some(lease) = live_lease(&entry) else { + } + while self.now < target { + self.now += 1; + // Coarse to fine, so an entry cascading several levels down still + // reaches level 0 in time to be drained by this same tick. + for level in (1..=TOP_LEVEL).rev() { + // A level turns over once every `slot_range` seconds, exactly + // when `now` has no bits left below that level's slot field. + if self.now & (slot_range(level) - 1) != 0 { continue; - }; - if lease.expires_at() <= self.now { - self.owners.remove(&lease.key_id()); - due.push(lease); - } else { - self.insert_with_version(lease, entry.version); } + let entries = self.take_bucket(level, self.now); + self.settle(entries, self.now); } + let entries = self.take_bucket(0, self.now); + self.settle(entries, self.now); } - due + // A clock stepping backwards must not rewind the wheel: buckets are + // indexed relative to `now`, so re-indexing against an earlier `now` + // would file entries into slots the wheel has already drained. + self.now = self.now.max(target); } - fn fast_forward(&mut self, target: u64) -> Vec> { - self.now = target; - let mut entries = std::mem::take(&mut self.immediate_due); - take_all_entries(&mut self.level0, &mut entries); - take_all_entries(&mut self.level1, &mut entries); - take_all_entries(&mut self.level2, &mut entries); - take_all_entries(&mut self.level3, &mut entries); - - let mut due = Vec::new(); - for entry in entries { - let Some(lease) = live_lease(&entry) else { + /// Fires the entries due at `deadline` and re-files both the phases they + /// schedule next and the entries that are not due yet. + fn settle(&mut self, entries: Bucket, deadline: u64) { + for (key_id, mut entry) in entries { + if entry.deadline > deadline { + self.place(key_id, entry); continue; - }; - if lease.expires_at() <= target { - self.owners.remove(&lease.key_id()); - due.push(lease); - } else { - self.insert_with_version(lease, entry.version); } - } - due - } - - fn take_immediate_due(&mut self, target: u64) -> Vec> { - let mut due = Vec::new(); - for entry in std::mem::take(&mut self.immediate_due) { - let Some(lease) = live_lease(&entry) else { - continue; - }; - if lease.expires_at() <= target { - self.owners.remove(&lease.key_id()); - due.push(lease); - } else { - self.insert_with_version(lease, entry.version); + match entry.fire() { + Some(next) => { + entry.deadline = next; + self.place(key_id, entry); + } + None => { + self.positions.remove(&key_id); + } } } - due } - fn cascade(&mut self, level: u8) { - let entries = match level { - 1 => { - let index = ((self.now >> 8) & 0x3f) as usize; - std::mem::take(&mut self.level1[index]) - } - 2 => { - let index = ((self.now >> 14) & 0x3f) as usize; - std::mem::take(&mut self.level2[index]) - } - 3 => { - let index = ((self.now >> 20) & 0x3f) as usize; - std::mem::take(&mut self.level3[index]) - } - _ => Vec::new(), + fn place(&mut self, key_id: KeyId, entry: Entry) { + let position = if entry.deadline <= self.now { + self.overdue.insert(key_id, entry); + Position::Overdue + } else { + let level = level_for(self.now, entry.deadline); + let slot = slot_for(level, entry.deadline); + self.bucket(level, slot).insert(key_id, entry); + Position::Wheel { level, slot } }; - for entry in entries { - if let Some(lease) = live_lease(&entry) { - self.insert_with_version(lease, entry.version); - } - } + self.positions.insert(key_id, position); } - pub(super) fn clear(&mut self, now: u64) { - let mut entries = std::mem::take(&mut self.immediate_due); - take_all_entries(&mut self.level0, &mut entries); - take_all_entries(&mut self.level1, &mut entries); - take_all_entries(&mut self.level2, &mut entries); - take_all_entries(&mut self.level3, &mut entries); - for lease in self.owners.values() { - lease.cancel_rotated(); + fn detach(&mut self, key_id: KeyId) -> Option { + match self.positions.remove(&key_id)? { + Position::Overdue => self.overdue.remove(&key_id), + Position::Wheel { level, slot } => self.bucket(level, slot).remove(&key_id), } - *self = Self::new(now); + } + + fn bucket(&mut self, level: Level, slot: Slot) -> &mut Bucket { + &mut self.levels[level as usize][slot as usize] + } + + /// Empties the bucket that `when` falls in at `level`. + fn take_bucket(&mut self, level: Level, when: u64) -> Bucket { + let slot = slot_for(level, when); + std::mem::take(self.bucket(level, slot)) + } + + #[cfg(test)] + pub(super) fn holds(&self, key_id: KeyId) -> bool { + self.positions.contains_key(&key_id) } } -fn live_lease(entry: &WheelEntry) -> Option> { - let lease = entry.lease.upgrade()?; - (entry.version == lease.wheel_version.load(Ordering::Acquire)).then_some(lease) +/// Seconds covered by one of `level`'s slots. +fn slot_range(level: Level) -> u64 { + 1 << (SLOT_BITS * level as u32) } -fn empty_buckets(count: usize) -> Vec> { - std::iter::repeat_with(Vec::new).take(count).collect() +fn slot_for(level: Level, when: u64) -> Slot { + ((when >> (SLOT_BITS * level as u32)) & SLOT_MASK) as Slot } -fn take_all_entries(buckets: &mut [Vec], entries: &mut Vec) { - for bucket in buckets { - entries.append(bucket); - } +/// Finest level able to hold `deadline`: the one whose slot field covers the +/// highest bit in which `now` and `deadline` differ. A deadline past the top +/// level is clamped into it and cannot fire early, because a drained entry is +/// only fired once its own deadline has passed. +fn level_for(now: u64, deadline: u64) -> Level { + let significant = 63 - ((now ^ deadline) | SLOT_MASK).leading_zeros(); + ((significant / SLOT_BITS) as Level).min(TOP_LEVEL) } diff --git a/src/common/message/secure.rs b/src/common/message/secure.rs index 57aac69..fcef449 100644 --- a/src/common/message/secure.rs +++ b/src/common/message/secure.rs @@ -29,7 +29,9 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt}; use super::{ CodecMessageReader, CodecMessageWriter, DataLenType, MessageReader, MessageWriter, MAX_MSG_LEN, }; -use crate::common::auth::{AuthContext, AuthFailure, AuthRuntime, LegacyConnectionGuard}; +use crate::common::auth::{ + AuthContext, AuthFailure, AuthRuntime, KeyId, LegacyConnectionGuard, ADMIN_KEY_ID, +}; use crate::common::checksum::{ get_process_credential, valid_checksum_for_key, AesKeyType, Credential, }; @@ -78,7 +80,8 @@ impl ClientHeaderSession { for byte in &mut salt[8..] { *byte = rng.random(); } - let material = derive_material(credential.key_id(), credential.key(), salt)?; + let material = + derive_material(KeyId::from_u64(credential.key_id()), credential.key(), salt)?; Ok(Self { protocol: HeaderProtocol::V2, legacy_key: *credential.key(), @@ -218,7 +221,7 @@ impl ServerHeaderSession { self.legacy_key } - pub fn key_id(&self) -> u64 { + pub fn key_id(&self) -> KeyId { self.context .as_ref() .map(|context| context.key_id) @@ -226,7 +229,7 @@ impl ServerHeaderSession { self.v2 .as_ref() .map(|material| material.key_id) - .unwrap_or_default() + .unwrap_or(ADMIN_KEY_ID) }) } @@ -291,7 +294,7 @@ pub struct ServerInitialMessage { pub struct ServerInitialError { pub failure: AuthFailure, pub response_session: Option, - pub presented_key_id: Option, + pub presented_key_id: Option, } impl ServerInitialError { @@ -311,7 +314,7 @@ impl ServerInitialError { code: &'static str, message: impl Into, retryable: bool, - key_id: u64, + key_id: KeyId, ) -> Self { Self { failure: AuthFailure::new(code, message, retryable), @@ -320,7 +323,7 @@ impl ServerInitialError { } } - fn from_failure_key(failure: AuthFailure, key_id: u64) -> Self { + fn from_failure_key(failure: AuthFailure, key_id: KeyId) -> Self { Self { failure, response_session: None, @@ -378,7 +381,7 @@ impl ServerSecurity { pub fn record_failure_log( &self, peer_ip: std::net::IpAddr, - key_id: u64, + key_id: KeyId, reason: &str, ) -> FailureLogDecision { self.failure_logs @@ -459,7 +462,7 @@ impl ServerSecurity { })?; let context = self .auth - .authenticate_presented(0, &key) + .authenticate_presented(ADMIN_KEY_ID, &key) .map_err(ServerInitialError::new)?; let legacy_guard = self .auth @@ -507,7 +510,9 @@ impl ServerSecurity { false, )); } - let key_id = u64::from_be_bytes(remainder[4..12].try_into().expect("fixed key id")); + let key_id = KeyId::from_u64(u64::from_be_bytes( + remainder[4..12].try_into().expect("fixed key id"), + )); let salt: [u8; CONNECTION_SALT_LEN] = remainder[12..28].try_into().expect("fixed connection salt"); let client_timestamp = u64::from_be_bytes(salt[..8].try_into().expect("fixed timestamp")); diff --git a/src/common/message/secure/first_flight.rs b/src/common/message/secure/first_flight.rs index 9b8468a..2de259c 100644 --- a/src/common/message/secure/first_flight.rs +++ b/src/common/message/secure/first_flight.rs @@ -14,7 +14,7 @@ pub(super) fn first_flight_error( code: &'static str, message: impl Into, retryable: bool, - key_id: u64, + key_id: KeyId, ) -> ServerInitialError { ServerInitialError::fail_key(code, message, retryable, key_id) } @@ -35,7 +35,7 @@ fn reserved_error_session( pub(super) fn evaluate_first_flight( auth: &AuthRuntime, replay: &std::sync::Mutex, - key_id: u64, + key_id: KeyId, fingerprint: [u8; 32], work: FirstFlightWork, ) -> std::result::Result<(Vec, AuthContext), ServerInitialError> { @@ -92,7 +92,7 @@ pub(super) fn evaluate_first_flight( pub(super) fn stale_root_first_flight( auth: &AuthRuntime, - key_id: u64, + key_id: KeyId, salt: [u8; CONNECTION_SALT_LEN], counter: u64, ciphertext: &[u8], @@ -107,7 +107,7 @@ pub(super) fn stale_root_first_flight( &mut previous_ciphertext, ) .ok()?; - let (code, message) = if key_id == 0 { + let (code, message) = if key_id.is_admin() { ( "administrator_key_invalid", "administrator credential does not match the active root key", diff --git a/src/common/message/secure/frame.rs b/src/common/message/secure/frame.rs index d76f75b..14536f8 100644 --- a/src/common/message/secure/frame.rs +++ b/src/common/message/secure/frame.rs @@ -10,11 +10,13 @@ //! The initial reader can impose a smaller pre-authentication limit before allocating //! a body; continuation frames retain the normal protocol maximum. +use crate::common::auth::KeyId; + use super::*; #[derive(Clone)] pub(super) struct V2Material { - pub(super) key_id: u64, + pub(super) key_id: KeyId, pub(super) flags: u8, pub(super) salt: [u8; CONNECTION_SALT_LEN], pub(super) client_to_server: AesKeyType, @@ -201,7 +203,7 @@ pub(super) fn open_v2_payload( } pub(super) fn derive_material( - key_id: u64, + key_id: KeyId, credential_key: &AesKeyType, salt_bytes: [u8; CONNECTION_SALT_LEN], ) -> Result { diff --git a/src/common/message/secure/limiter.rs b/src/common/message/secure/limiter.rs index a28912a..46e7dc6 100644 --- a/src/common/message/secure/limiter.rs +++ b/src/common/message/secure/limiter.rs @@ -9,6 +9,8 @@ //! decisions: every authentication failure is still rejected, only duplicate logging //! is coalesced. +use crate::common::auth::KeyId; + #[derive(Clone, Copy, Debug)] pub struct FailureLogDecision { pub emit: bool, @@ -23,7 +25,8 @@ pub(super) struct FailureLogEntry { #[derive(Default)] pub(super) struct FailureLogLimiter { - pub(super) entries: std::collections::HashMap<(std::net::IpAddr, u64, String), FailureLogEntry>, + pub(super) entries: + std::collections::HashMap<(std::net::IpAddr, KeyId, String), FailureLogEntry>, pub(super) overflow: Option, } @@ -31,7 +34,7 @@ impl FailureLogLimiter { pub(super) fn record( &mut self, peer_ip: std::net::IpAddr, - key_id: u64, + key_id: KeyId, reason: &str, now: u64, ) -> FailureLogDecision { diff --git a/src/common/message/secure/replay.rs b/src/common/message/secure/replay.rs index 3214b56..54727b4 100644 --- a/src/common/message/secure/replay.rs +++ b/src/common/message/secure/replay.rs @@ -15,6 +15,8 @@ //! Per-credential counts stop one tenant from filling the shared filter with //! unique salts before the request payload is decoded. +use crate::common::auth::KeyId; + use std::collections::HashMap; use std::fs::{File, OpenOptions}; use std::io::{ErrorKind, Read, Write}; @@ -56,7 +58,7 @@ pub(super) enum FirstFlightAdmit { Unavailable, } -pub(super) fn replay_fingerprint(key_id: u64, salt: &[u8; CONNECTION_SALT_LEN]) -> [u8; 32] { +pub(super) fn replay_fingerprint(key_id: KeyId, salt: &[u8; CONNECTION_SALT_LEN]) -> [u8; 32] { let mut input = [0_u8; 8 + CONNECTION_SALT_LEN]; input[..8].copy_from_slice(&key_id.to_be_bytes()); input[8..].copy_from_slice(salt); @@ -111,7 +113,7 @@ impl RotatingBloom { pub(super) struct ReplayGuard { bloom: RotatingBloom, - counts: HashMap, + counts: HashMap, counts_started_at: u64, window_seconds: u64, max_per_key: u32, @@ -158,7 +160,7 @@ impl ReplayGuard { pub(super) fn admit( &mut self, - key_id: u64, + key_id: KeyId, fingerprint: &[u8; 32], now: u64, ) -> FirstFlightAdmit { diff --git a/src/common/message/secure/tests.rs b/src/common/message/secure/tests.rs index 860e59c..a1b3762 100644 --- a/src/common/message/secure/tests.rs +++ b/src/common/message/secure/tests.rs @@ -64,7 +64,7 @@ async fn temporary_credential_authenticates_without_storing_secret() { let admin = *b"0123456789abcdefghijklmnopqrstuv"; let config = temp_config(); let auth = AuthRuntime::start(admin, config.clone()).await.unwrap(); - let admin_context = auth.authenticate_presented(0, &admin).unwrap(); + let admin_context = auth.authenticate_presented(ADMIN_KEY_ID, &admin).unwrap(); let issued = auth .issue(&admin_context, std::time::Duration::from_secs(60), None) .await @@ -148,7 +148,7 @@ async fn revoked_first_flights_do_not_consume_the_replay_filter() { let admin = *b"0123456789abcdefghijklmnopqrstuv"; let config = temp_config(); let auth = AuthRuntime::start(admin, config.clone()).await.unwrap(); - let admin_context = auth.authenticate_presented(0, &admin).unwrap(); + let admin_context = auth.authenticate_presented(ADMIN_KEY_ID, &admin).unwrap(); let issued = auth .issue(&admin_context, std::time::Duration::from_secs(60), None) .await @@ -159,7 +159,9 @@ async fn revoked_first_flights_do_not_consume_the_replay_filter() { }; let client = ClientHeaderSession::new_v2(&Credential::Temporary { key_id, key }).unwrap(); let bytes = encode_initial(&client, b"revoked").await; - auth.revoke(&admin_context, key_id).await.unwrap(); + auth.revoke(&admin_context, KeyId::from_u64(key_id)) + .await + .unwrap(); let security = ServerSecurity::new(auth); let first = match security @@ -192,7 +194,7 @@ async fn accepted_then_revoked_replay_does_not_reuse_the_session_nonce() { let admin = *b"0123456789abcdefghijklmnopqrstuv"; let config = temp_config(); let auth = AuthRuntime::start(admin, config.clone()).await.unwrap(); - let admin_context = auth.authenticate_presented(0, &admin).unwrap(); + let admin_context = auth.authenticate_presented(ADMIN_KEY_ID, &admin).unwrap(); let issued = auth .issue(&admin_context, std::time::Duration::from_secs(60), None) .await @@ -210,7 +212,7 @@ async fn accepted_then_revoked_replay_does_not_reuse_the_session_nonce() { .expect("first flight should be accepted"); security .auth() - .revoke(&admin_context, key_id) + .revoke(&admin_context, KeyId::from_u64(key_id)) .await .unwrap(); let replayed = match security @@ -236,7 +238,7 @@ async fn rotated_temporary_first_flight_returns_a_readable_rotated_error() { let new_admin = *b"abcdefghijklmnopqrstuvwxyz012345"; let config = temp_config(); let auth = AuthRuntime::start(admin, config.clone()).await.unwrap(); - let admin_context = auth.authenticate_presented(0, &admin).unwrap(); + let admin_context = auth.authenticate_presented(ADMIN_KEY_ID, &admin).unwrap(); let issued = auth .issue(&admin_context, std::time::Duration::from_secs(60), None) .await @@ -283,7 +285,7 @@ async fn stale_root_replay_omits_the_rotated_error_session() { let new_admin = *b"abcdefghijklmnopqrstuvwxyz012345"; let config = temp_config(); let auth = AuthRuntime::start(admin, config.clone()).await.unwrap(); - let admin_context = auth.authenticate_presented(0, &admin).unwrap(); + let admin_context = auth.authenticate_presented(ADMIN_KEY_ID, &admin).unwrap(); let issued = auth .issue(&admin_context, std::time::Duration::from_secs(60), None) .await @@ -327,7 +329,7 @@ async fn reset_temporary_first_flight_returns_a_readable_rotated_error() { let admin = *b"0123456789abcdefghijklmnopqrstuv"; let config = temp_config(); let auth = AuthRuntime::start(admin, config.clone()).await.unwrap(); - let admin_context = auth.authenticate_presented(0, &admin).unwrap(); + let admin_context = auth.authenticate_presented(ADMIN_KEY_ID, &admin).unwrap(); let issued = auth .issue(&admin_context, std::time::Duration::from_secs(60), None) .await @@ -372,7 +374,7 @@ async fn mistyped_temporary_first_flight_does_not_send_an_unreadable_error() { let admin = *b"0123456789abcdefghijklmnopqrstuv"; let config = temp_config(); let auth = AuthRuntime::start(admin, config.clone()).await.unwrap(); - let admin_context = auth.authenticate_presented(0, &admin).unwrap(); + let admin_context = auth.authenticate_presented(ADMIN_KEY_ID, &admin).unwrap(); let issued = auth .issue(&admin_context, std::time::Duration::from_secs(60), None) .await @@ -398,7 +400,7 @@ async fn mistyped_temporary_first_flight_does_not_send_an_unreadable_error() { error.response_session.is_none(), "the presenter cannot open a session derived from the live key" ); - assert_eq!(error.presented_key_id, Some(key_id)); + assert_eq!(error.presented_key_id, Some(KeyId::from_u64(key_id))); let _ = std::fs::remove_dir_all(config.state_dir); } @@ -494,12 +496,30 @@ fn per_credential_admission_limit_does_not_consume_other_keys() { let now = unix_seconds(); let mut guard = ReplayGuard::open(None, 1024, DEFAULT_REPLAY_WINDOW_SECONDS).with_max_per_key(2); - assert_eq!(guard.admit(1, &[1_u8; 32], now), FirstFlightAdmit::Fresh); - assert_eq!(guard.admit(1, &[2_u8; 32], now), FirstFlightAdmit::Fresh); - assert_eq!(guard.admit(1, &[3_u8; 32], now), FirstFlightAdmit::Limited); - assert_eq!(guard.admit(1, &[3_u8; 32], now), FirstFlightAdmit::Limited); - assert_eq!(guard.admit(2, &[4_u8; 32], now), FirstFlightAdmit::Fresh); - assert_eq!(guard.admit(1, &[1_u8; 32], now), FirstFlightAdmit::Replayed); + assert_eq!( + guard.admit(KeyId::from_u64(1), &[1_u8; 32], now), + FirstFlightAdmit::Fresh + ); + assert_eq!( + guard.admit(KeyId::from_u64(1), &[2_u8; 32], now), + FirstFlightAdmit::Fresh + ); + assert_eq!( + guard.admit(KeyId::from_u64(1), &[3_u8; 32], now), + FirstFlightAdmit::Limited + ); + assert_eq!( + guard.admit(KeyId::from_u64(1), &[3_u8; 32], now), + FirstFlightAdmit::Limited + ); + assert_eq!( + guard.admit(KeyId::from_u64(2), &[4_u8; 32], now), + FirstFlightAdmit::Fresh + ); + assert_eq!( + guard.admit(KeyId::from_u64(1), &[1_u8; 32], now), + FirstFlightAdmit::Replayed + ); } #[test] @@ -508,11 +528,26 @@ fn aggregate_admission_limit_covers_all_keys() { let mut guard = ReplayGuard::open(None, 1024, DEFAULT_REPLAY_WINDOW_SECONDS) .with_max_per_key(100) .with_max_total(2); - assert_eq!(guard.admit(1, &[1_u8; 32], now), FirstFlightAdmit::Fresh); - assert_eq!(guard.admit(2, &[2_u8; 32], now), FirstFlightAdmit::Fresh); - assert_eq!(guard.admit(3, &[3_u8; 32], now), FirstFlightAdmit::Limited); - assert_eq!(guard.admit(3, &[3_u8; 32], now), FirstFlightAdmit::Limited); - assert_eq!(guard.admit(4, &[4_u8; 32], now), FirstFlightAdmit::Limited); + assert_eq!( + guard.admit(KeyId::from_u64(1), &[1_u8; 32], now), + FirstFlightAdmit::Fresh + ); + assert_eq!( + guard.admit(KeyId::from_u64(2), &[2_u8; 32], now), + FirstFlightAdmit::Fresh + ); + assert_eq!( + guard.admit(KeyId::from_u64(3), &[3_u8; 32], now), + FirstFlightAdmit::Limited + ); + assert_eq!( + guard.admit(KeyId::from_u64(3), &[3_u8; 32], now), + FirstFlightAdmit::Limited + ); + assert_eq!( + guard.admit(KeyId::from_u64(4), &[4_u8; 32], now), + FirstFlightAdmit::Limited + ); } #[test] @@ -530,14 +565,17 @@ fn persisted_first_flights_survive_a_torn_trailing_record() { let fingerprint = [17_u8; 32]; { let mut guard = ReplayGuard::open(Some(path.clone()), 1024, DEFAULT_REPLAY_WINDOW_SECONDS); - assert_eq!(guard.admit(7, &fingerprint, now), FirstFlightAdmit::Fresh); + assert_eq!( + guard.admit(KeyId::from_u64(7), &fingerprint, now), + FirstFlightAdmit::Fresh + ); } let mut torn = std::fs::read(&path).unwrap(); torn.extend_from_slice(&[0_u8; 10]); std::fs::write(&path, torn).unwrap(); let mut restored = ReplayGuard::open(Some(path.clone()), 1024, DEFAULT_REPLAY_WINDOW_SECONDS); assert_eq!( - restored.admit(7, &fingerprint, now), + restored.admit(KeyId::from_u64(7), &fingerprint, now), FirstFlightAdmit::Unavailable ); let _ = std::fs::remove_file(path); @@ -558,7 +596,10 @@ fn replay_rewrite_succeeds_when_a_pid_temporary_file_already_exists() { let fingerprint = [19_u8; 32]; { let mut guard = ReplayGuard::open(Some(path.clone()), 1024, DEFAULT_REPLAY_WINDOW_SECONDS); - assert_eq!(guard.admit(9, &fingerprint, now), FirstFlightAdmit::Fresh); + assert_eq!( + guard.admit(KeyId::from_u64(9), &fingerprint, now), + FirstFlightAdmit::Fresh + ); } let leftover = path.with_file_name(format!( ".{}.tmp-{}", @@ -568,7 +609,7 @@ fn replay_rewrite_succeeds_when_a_pid_temporary_file_already_exists() { std::fs::write(&leftover, b"stale").unwrap(); let mut restored = ReplayGuard::open(Some(path.clone()), 1024, DEFAULT_REPLAY_WINDOW_SECONDS); assert_eq!( - restored.admit(9, &fingerprint, now), + restored.admit(KeyId::from_u64(9), &fingerprint, now), FirstFlightAdmit::Replayed ); let _ = std::fs::remove_file(leftover); @@ -588,11 +629,14 @@ fn persisted_first_flights_survive_replay_guard_restart() { let fingerprint = [13_u8; 32]; { let mut guard = ReplayGuard::open(Some(path.clone()), 1024, DEFAULT_REPLAY_WINDOW_SECONDS); - assert_eq!(guard.admit(7, &fingerprint, now), FirstFlightAdmit::Fresh); + assert_eq!( + guard.admit(KeyId::from_u64(7), &fingerprint, now), + FirstFlightAdmit::Fresh + ); } let mut restored = ReplayGuard::open(Some(path.clone()), 1024, DEFAULT_REPLAY_WINDOW_SECONDS); assert_eq!( - restored.admit(7, &fingerprint, now), + restored.admit(KeyId::from_u64(7), &fingerprint, now), FirstFlightAdmit::Replayed ); let _ = std::fs::remove_file(path); @@ -614,14 +658,20 @@ fn restored_replay_log_consumes_the_aggregate_budget() { let mut guard = ReplayGuard::open(Some(path.clone()), 1024, DEFAULT_REPLAY_WINDOW_SECONDS) .with_max_per_key(100) .with_max_total(2); - assert_eq!(guard.admit(1, &[1_u8; 32], now), FirstFlightAdmit::Fresh); - assert_eq!(guard.admit(2, &[2_u8; 32], now), FirstFlightAdmit::Fresh); + assert_eq!( + guard.admit(KeyId::from_u64(1), &[1_u8; 32], now), + FirstFlightAdmit::Fresh + ); + assert_eq!( + guard.admit(KeyId::from_u64(2), &[2_u8; 32], now), + FirstFlightAdmit::Fresh + ); } let mut restored = ReplayGuard::open(Some(path.clone()), 1024, DEFAULT_REPLAY_WINDOW_SECONDS) .with_max_per_key(100) .with_max_total(2); assert_eq!( - restored.admit(3, &[3_u8; 32], now), + restored.admit(KeyId::from_u64(3), &[3_u8; 32], now), FirstFlightAdmit::Limited ); let _ = std::fs::remove_file(path); @@ -643,15 +693,21 @@ fn restored_replay_generation_ages_with_loaded_records() { let mut guard = ReplayGuard::open(Some(path.clone()), 1024, DEFAULT_REPLAY_WINDOW_SECONDS) .with_max_per_key(100) .with_max_total(2); - assert_eq!(guard.admit(1, &[1_u8; 32], start), FirstFlightAdmit::Fresh); - assert_eq!(guard.admit(2, &[2_u8; 32], start), FirstFlightAdmit::Fresh); + assert_eq!( + guard.admit(KeyId::from_u64(1), &[1_u8; 32], start), + FirstFlightAdmit::Fresh + ); + assert_eq!( + guard.admit(KeyId::from_u64(2), &[2_u8; 32], start), + FirstFlightAdmit::Fresh + ); } let mut restored = ReplayGuard::open(Some(path.clone()), 1024, DEFAULT_REPLAY_WINDOW_SECONDS) .with_max_per_key(100) .with_max_total(2); assert_eq!( restored.admit( - 3, + KeyId::from_u64(3), &[3_u8; 32], start.saturating_add(DEFAULT_REPLAY_WINDOW_SECONDS) ), @@ -709,7 +765,7 @@ fn failure_log_limiter_has_a_hard_cardinality_bound() { let mut limiter = FailureLogLimiter::default(); let peer = "127.0.0.1".parse().unwrap(); for key_id in 0..10_000 { - limiter.record(peer, key_id, "invalid", 1_000); + limiter.record(peer, KeyId::from_u64(key_id), "invalid", 1_000); } assert_eq!(limiter.entries.len(), 4096); assert!(limiter.overflow.is_some()); diff --git a/src/pb_server/admin.rs b/src/pb_server/admin.rs index fa057dd..9773540 100644 --- a/src/pb_server/admin.rs +++ b/src/pb_server/admin.rs @@ -15,7 +15,7 @@ use tokio::net::TcpStream; use super::error::Error; use super::{ManagerTask, ManagerTaskSender, Result}; -use crate::common::auth::{AuthContext, AuthFailure, AuthRuntime}; +use crate::common::auth::{AuthContext, AuthFailure, AuthRuntime, KeyId}; use crate::common::checksum::{parse_credential, Credential}; use crate::common::conn_id::RemoteConnId; use crate::common::message::command::{ @@ -90,22 +90,26 @@ async fn execute( .map(AdminResponse::KeyList) } AdminRequest::KeyShow { key_id } => auth - .show(authorization, key_id, false) + .show(authorization, KeyId::from_u64(key_id), false) .await .map(AdminResponse::KeyShown), AdminRequest::KeyReveal { key_id } => auth - .show(authorization, key_id, true) + .show(authorization, KeyId::from_u64(key_id), true) .await .map(AdminResponse::KeyShown), AdminRequest::KeyRenew { key_id, ttl_seconds, } => auth - .renew(authorization, key_id, Duration::from_secs(ttl_seconds)) + .renew( + authorization, + KeyId::from_u64(key_id), + Duration::from_secs(ttl_seconds), + ) .await .map(AdminResponse::KeyRenewed), AdminRequest::KeyRevoke { key_id } => auth - .revoke(authorization, key_id) + .revoke(authorization, KeyId::from_u64(key_id)) .await .map(AdminResponse::KeyRevoked), AdminRequest::KeyGc => auth @@ -161,7 +165,7 @@ async fn execute( &auth, authorization, "service_list", - key_id, + key_id.map(KeyId::from_u64), Some(format!("page={page},page_size={page_size}")), ) .await; @@ -190,7 +194,7 @@ async fn execute( &auth, authorization, "connection_list", - key_id, + key_id.map(KeyId::from_u64), Some(format!("page={page},page_size={page_size}")), ) .await; @@ -268,7 +272,7 @@ async fn audit_read( auth: &AuthRuntime, authorization: &AuthContext, action: &str, - key_id: Option, + key_id: Option, detail: Option, ) { if let Err(error) = auth @@ -289,6 +293,7 @@ async fn audit_read( #[cfg(test)] mod tests { use super::*; + use crate::common::auth::ADMIN_KEY_ID; use crate::common::auth::{AuthConfig, LegacyProtocolPolicy}; fn temp_state_dir(name: &str) -> std::path::PathBuf { @@ -324,7 +329,7 @@ mod tests { .await .expect("authentication runtime should start"); let admin = runtime - .authenticate_presented(0, &old_key) + .authenticate_presented(ADMIN_KEY_ID, &old_key) .expect("old administrator key should authenticate"); let request = if connection_query { AdminRequest::ConnectionList { diff --git a/src/pb_server/connection.rs b/src/pb_server/connection.rs index 064a1bb..6cb0545 100644 --- a/src/pb_server/connection.rs +++ b/src/pb_server/connection.rs @@ -61,13 +61,13 @@ pub(super) async fn handle_conn( .as_ref() .map(|session| session.key_id()) .or(error.presented_key_id) - .unwrap_or_default(); + .unwrap_or(ADMIN_KEY_ID); let decision = security.record_failure_log(peer_addr.ip(), key_id, &error.failure.code); if decision.suppressed > 0 { tracing::warn!( event = "auth_failures_suppressed", peer_ip = %peer_addr.ip(), - key_id, + key_id = key_id.as_u64(), reason = %error.failure.code, suppressed = decision.suppressed, "suppressed repeated authentication failures in the previous window" @@ -79,7 +79,7 @@ pub(super) async fn handle_conn( auth_stage = "initial_frame", conn_id = %conn_id, peer_addr = %peer_addr, - key_id, + key_id = key_id.as_u64(), reason = %error.failure.code, retryable = error.failure.retryable, error = %error.failure.message, @@ -181,7 +181,7 @@ pub(super) async fn handle_conn( auth_stage = "session", conn_id = %conn_id, peer_addr = %peer_addr, - key_id = auth_context.key_id, + key_id = auth_context.key_id.as_u64(), namespace = auth_context.namespace, protocol = ?session.protocol(), is_admin = auth_context.is_admin, @@ -459,7 +459,7 @@ where _ = cancellation.cancelled() => { tracing::info!( event = "connection_auth_expired", - key_id = auth_context.key_id, + key_id = auth_context.key_id.as_u64(), conn_id = %conn_id, "closing {closed_what}" ); diff --git a/src/pb_server/mod.rs b/src/pb_server/mod.rs index 14a0098..40835fc 100644 --- a/src/pb_server/mod.rs +++ b/src/pb_server/mod.rs @@ -34,7 +34,7 @@ use self::error::{ }; use self::server::{handle_server_conn, ServerRegistration}; use self::status::handle_show_status; -use crate::common::auth::{AuthConfig, AuthContext, AuthRuntime}; +use crate::common::auth::{AuthConfig, AuthContext, AuthRuntime, ADMIN_KEY_ID}; use crate::common::config::{control_io_timeout, keep_alive_from_env, server_lease_timeout}; use crate::common::conn_id::{ConnIdProvider, RemoteConnId}; use crate::common::manager::{ForwardMessage, SenderChan, TaskManager}; diff --git a/tests/regression.rs b/tests/regression.rs index 4e3e62f..6c78140 100644 --- a/tests/regression.rs +++ b/tests/regression.rs @@ -4,7 +4,7 @@ use std::sync::{Arc, Mutex}; use std::time::Duration; use pb_mapper::common::auth::{ - write_admin_key_file, AuthConfig, AuthRuntime, LegacyProtocolPolicy, + write_admin_key_file, AuthConfig, AuthRuntime, LegacyProtocolPolicy, ADMIN_KEY_ID, }; use pb_mapper::common::checksum::{parse_credential, set_process_msg_header_key, Credential}; use pb_mapper::common::message::command::{ @@ -108,7 +108,10 @@ async fn admin_all_preserves_json_output_mode() { .await .unwrap(); let admin = runtime - .authenticate_presented(0, TEST_ADMIN_KEY.as_bytes().first_chunk::<32>().unwrap()) + .authenticate_presented( + ADMIN_KEY_ID, + TEST_ADMIN_KEY.as_bytes().first_chunk::<32>().unwrap(), + ) .unwrap(); runtime .issue( @@ -316,7 +319,10 @@ async fn temporary_credentials_are_isolated_denied_admin_and_revoked_live() { .await .unwrap(); let admin = runtime - .authenticate_presented(0, TEST_ADMIN_KEY.as_bytes().first_chunk::<32>().unwrap()) + .authenticate_presented( + ADMIN_KEY_ID, + TEST_ADMIN_KEY.as_bytes().first_chunk::<32>().unwrap(), + ) .unwrap(); let first = runtime .issue(&admin, Duration::from_secs(120), Some("first".to_string())) @@ -418,7 +424,7 @@ async fn temporary_credentials_are_isolated_denied_admin_and_revoked_live() { server_addr, &admin_credential, PbConnRequest::Admin(AdminRequest::KeyRevoke { - key_id: first.metadata.key_id, + key_id: first.metadata.key_id.as_u64(), }), ) .await; @@ -461,7 +467,10 @@ async fn revoking_subscriber_credential_closes_cross_credential_data_stream() { .await .unwrap(); let admin = runtime - .authenticate_presented(0, TEST_ADMIN_KEY.as_bytes().first_chunk::<32>().unwrap()) + .authenticate_presented( + ADMIN_KEY_ID, + TEST_ADMIN_KEY.as_bytes().first_chunk::<32>().unwrap(), + ) .unwrap(); let issued = runtime .issue( @@ -493,7 +502,7 @@ async fn revoking_subscriber_credential_closes_cross_credential_data_stream() { need_codec: false, is_datagram: false, key: service.to_string(), - namespace: issued.metadata.key_id, + namespace: issued.metadata.key_id.as_u64(), force_namespace: true, protocol_version: Some(2), client_instance_id: Some("active-stream-test".to_string()), @@ -562,7 +571,7 @@ async fn revoking_subscriber_credential_closes_cross_credential_data_stream() { &mut provider, &PbConnRequest::StreamScoped { key: service.to_string(), - namespace: issued.metadata.key_id, + namespace: issued.metadata.key_id.as_u64(), dst_id: client_id, server_generation, } @@ -606,7 +615,7 @@ async fn revoking_subscriber_credential_closes_cross_credential_data_stream() { server_addr, &admin_credential, PbConnRequest::Admin(AdminRequest::KeyRevoke { - key_id: issued.metadata.key_id, + key_id: issued.metadata.key_id.as_u64(), }), ) .await; From b8a61d4f122c06a942f7ebba1db4f34f24c426a8 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Fri, 21 Aug 2026 04:21:58 +0800 Subject: [PATCH 61/74] Separate the timing wheel from what it schedules MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The wheel had a `HashMap` beside its buckets so it could find a key's entry, which put credential identity inside a scheduler and made every insert maintain two structures. `advance` was worse: a clock jump past the longest schedulable delay folded every bucket in the wheel into one pass, exactly the whole-structure walk a timing wheel exists to avoid. The wheel now schedules opaque `Arc`s and ticks one second at a time, draining one bucket per level that turns over. It holds the only strong references, so a timer runs when the last entry referring to it is dropped, and it never looks anything up: no key map, no positions, no identity comparison. 314 lines to 168. `Leases` keeps the `KeyId -> Weak` map instead, which is what the business side needed all along. Renewing schedules the same timer again at the later deadline — the earlier placement still drains, but it is no longer the last reference, so it fires nothing. Revoking fires the retire timer early and leaves its scheduled placement inert, so the retention window that follows still runs on its own schedule. A clock correction past any schedulable delay drops the schedule wholesale rather than making the wheel walk itself. Each key now owns two independent timers, retire and reap, both scheduled at issue. The reap timer holds the lease `Arc`, because firing a timer consumes its callback: an `Arc` held by the retire stage would be released the moment that stage ran, and the slot table holds only a `Weak`, so a request during the retention window would find a vanished lease instead of learning the key was revoked. Co-Authored-By: Claude Opus 5 (1M context) --- src/common/auth.rs | 2 +- src/common/auth/leases.rs | 284 +++++++++++++++++----------- src/common/auth/tests.rs | 213 +++++++-------------- src/common/auth/timing_wheel.rs | 322 +++++++++----------------------- 4 files changed, 325 insertions(+), 496 deletions(-) diff --git a/src/common/auth.rs b/src/common/auth.rs index 75372f4..a85789a 100644 --- a/src/common/auth.rs +++ b/src/common/auth.rs @@ -794,6 +794,6 @@ pub use ids::{Generation, KeyId, SlotIndex, ADMIN_KEY_ID}; mod leases; use leases::Leases; mod timing_wheel; -use timing_wheel::TimingWheel; +use timing_wheel::{Timer, TimingWheel}; #[cfg(test)] mod tests; diff --git a/src/common/auth/leases.rs b/src/common/auth/leases.rs index 64c83ca..b663e2a 100644 --- a/src/common/auth/leases.rs +++ b/src/common/auth/leases.rs @@ -1,45 +1,49 @@ -//! Temporary-key lifetimes, expressed as one self-advancing cleanup callback per -//! key. +//! Temporary-key lifetimes, scheduled on the timing wheel. //! //! ```text -//! issue schedule(expires_at) ── slots[i] Active, lease live -//! | -//! v phase 1: deadline reached, or the entry is cancelled -//! retire lease cancelled, slots[i] Expired, tombstoned_at recorded -//! | -//! v returns expires_at + TOMBSTONE_RETENTION -//! (waiting) <- a client presenting the dead credential is told -//! | "expired", not the "unknown key" it would get from -//! v an already-recycled row -//! reap phase 2: slots[i] Free (generation kept), cold metadata and any -//! high-slot row removed. Nothing left; the entry is done. +//! issue schedule(expires_at) ── slots[i] Active, lease live +//! | +//! v deadline arrives, or a revoke fires the timer early +//! retire lease cancelled, slots[i] Expired, tombstoned_at recorded, +//! a reap timer scheduled for +TOMBSTONE_RETENTION +//! | <- a client presenting the dead credential is told +//! | "expired", not the "unknown key" it would get from +//! v an already-recycled row +//! reap slots[i] Free (generation kept), cold metadata and any high-slot +//! row removed //! ``` //! -//! One entry covers a key's whole life, so there is no tombstone queue and no -//! sweep to keep in step with the wheel. Every way a key can end runs the same -//! two phases: reaching a deadline runs them on schedule, and cancelling the -//! entry — for a revoke, a GC, a root rotation, or the wheel being dropped — runs -//! whichever phases remain immediately. No call site performs cleanup, which is -//! what keeps a forgotten call from stranding a lease past its row's reuse or -//! leaking a metadata entry per issued key. +//! Both stages are timers, so nothing sweeps and no queue has to stay in step +//! with the wheel. Every way a key can end runs the same callback: a deadline +//! arriving runs it on schedule, [`Timer::fire`] runs it early for a revoke or a +//! GC, and dropping the wheel runs it for a rotation or shutdown. +//! +//! `timers` maps each key to a `Weak` handle on its current timer, which is what +//! keeps key identity out of the wheel. Renewing upgrades the handle and +//! schedules the same timer at the later deadline: the earlier placement still +//! drains, but it is no longer the last reference, so nothing fires. Because the +//! map holds only `Weak` references, an entry whose timer has fired costs nothing +//! but a stale key, cleared by the callback itself. //! //! The callbacks hold a `Weak`, so they neither keep the state //! alive nor touch it after a runtime has shut down. use super::*; -/// Whether a new entry tears down the one it replaces, or takes over its work. -enum Schedule { - /// Any entry already held is torn down first. - Fresh, - /// The previous entry is discarded without running its phases, because this - /// one now owes them. - Supersede, +/// A key's two scheduled stages. Both are `Weak`, so a stage that has already +/// run costs nothing but a stale map key. +#[derive(Default)] +struct Stages { + retire: Weak, + reap: Weak, } pub(super) struct Leases { inner: Weak, wheel: TimingWheel, + /// Each key's stages, so a renew or an early end can reach them without the + /// wheel knowing what a key is. + stages: HashMap, } impl Leases { @@ -49,6 +53,7 @@ impl Leases { let mut leases = Self { inner: Arc::downgrade(inner), wheel: TimingWheel::new(now), + stages: HashMap::new(), }; let mut live = Vec::new(); let mut dead = Vec::new(); @@ -68,7 +73,7 @@ impl Leases { .map(|entry| entry.key_id), ); for (key_id, lease) in live { - leases.watch(key_id, lease, Schedule::Fresh); + leases.watch(key_id, lease); } for key_id in dead { let tombstoned_at = inner @@ -77,13 +82,13 @@ impl Leases { .map(|cold| cold.tombstoned_at) .filter(|at| *at != 0) .unwrap_or(now); - leases.entomb(key_id, tombstoned_at); + leases.schedule_reap(key_id, retention_ends(tombstoned_at)); } leases } - /// Takes over a newly issued key: records its description and schedules the - /// retirement its expiry, or any earlier cancellation, will run. + /// Takes over a newly issued key: records its description and schedules both + /// stages of its teardown. pub(super) fn issue(&mut self, lease: &Arc, issued_at: u64, label: Option) { let Some(inner) = self.inner.upgrade() else { return; @@ -96,56 +101,85 @@ impl Leases { tombstoned_at: 0, }, ); - self.watch(lease.key_id(), lease.clone(), Schedule::Fresh); + self.watch(lease.key_id(), lease.clone()); } /// Hands a key's remaining life to a replacement lease, for a renewal whose - /// original lease had already been cancelled. The row stays alive, so the - /// entry being replaced must not run its teardown. + /// original lease had already been cancelled. pub(super) fn adopt(&mut self, lease: &Arc) { - self.watch(lease.key_id(), lease.clone(), Schedule::Supersede); + self.watch(lease.key_id(), lease.clone()); } - /// Moves a renewed key to its new expiry. Returns `false` for a key the - /// wheel is not watching, as for a high slot. + /// Moves a renewed key to its new expiry, and its reap along with it. Returns + /// `false` for a key with no live stages, as for a high slot. + /// + /// Each timer is scheduled a second time rather than moved: its earlier + /// placement drains on the old deadline but is no longer the last reference, + /// so it fires nothing. pub(super) fn renew(&mut self, key_id: KeyId, expires_at: u64) -> bool { - self.wheel.reschedule(key_id, expires_at) + let Some(stages) = self.stages.get(&key_id) else { + return false; + }; + let (Some(retire), Some(reap)) = (stages.retire.upgrade(), stages.reap.upgrade()) else { + return false; + }; + self.wheel.schedule(expires_at, retire); + self.wheel.schedule(retention_ends(expires_at), reap); + true } - /// Retires a key now rather than at its expiry, leaving its tombstone to run - /// on schedule. This is what a revoke needs: the credential stops working - /// immediately, but the row is still held long enough to report *why*. + /// Retires a key now rather than at its expiry, leaving its reap on schedule. + /// This is what a revoke needs: the credential stops working immediately, but + /// the row is held long enough to report *why*. pub(super) fn retire_now(&mut self, key_id: KeyId) { - self.wheel.advance_one_phase(key_id); + if let Some(retire) = self.stage(key_id, |stages| &stages.retire) { + retire.fire(); + } } - /// Ends a key outright, running whichever of its phases remain: an active - /// key is retired and reaped, and a tombstoned one is reaped. Skips the - /// retention wait, so it is for a caller that wants the row back now. + /// Ends a key outright, running both stages. Skips the retention wait, so it + /// is for a caller that wants the row back now. pub(super) fn end(&mut self, key_id: KeyId) { - self.wheel.cancel(key_id); + self.retire_now(key_id); + if let Some(reap) = self.stage(key_id, |stages| &stages.reap) { + reap.fire(); + } + self.stages.remove(&key_id); } - /// Runs every phase whose deadline has passed. + /// Runs every callback whose deadline has passed. pub(super) fn tick(&mut self, now: u64) { - self.wheel.advance(now); + // A jump longer than anything a key can be scheduled for means every + // timer is due, so the schedule is dropped wholesale instead of ticked up + // to. That keeps a corrected hardware clock from spinning for hours. + // + // A clock stepping backwards is ignored: buckets are indexed relative to + // the wheel's `now`, so re-filing against an earlier one would place + // entries in slots it has already drained. + if now.saturating_sub(self.wheel.now()) > MAX_SCHEDULABLE_DELAY.as_secs() { + self.drop_schedule(now); + return; + } + while self.wheel.now() < now { + self.wheel.tick(); + } } /// Ends every key at once, for a root rotation or state reset. Dropping the - /// wheel runs all remaining phases, so no row, lease, or metadata entry - /// survives it. + /// wheel releases the last reference to every timer, so no row, lease, or + /// metadata entry survives it. pub(super) fn wipe(&mut self, now: u64) { - // Rotation is the one reason a phase cannot infer, so it is recorded + // Rotation is the one reason a callback cannot infer, so it is recorded // before the drop; `record_cancel` keeps the first reason. if let Some(inner) = self.inner.upgrade() { for lease in inner.slots().iter().filter_map(|slot| slot.lease.upgrade()) { lease.cancel_rotated(); } } - self.wheel = TimingWheel::new(now); + self.drop_schedule(now); } - /// Ends every key that is dead or past its deadline, skipping the tombstone + /// Ends every key that is dead or past its deadline, skipping the retention /// wait. Returns how many keys were ended. pub(super) fn collect_garbage(&mut self, now: u64) -> u64 { let Some(inner) = self.inner.upgrade() else { @@ -171,42 +205,78 @@ impl Leases { due.len() as u64 } - /// Schedules a live key's two phases, starting at its lease's expiry. + /// Replaces the whole schedule, running every callback the old one held. + fn drop_schedule(&mut self, now: u64) { + self.stages.clear(); + self.wheel = TimingWheel::new(now); + } + + /// Upgrades one of a key's stages, forgetting the key once both have run. + fn stage( + &mut self, + key_id: KeyId, + which: impl Fn(&Stages) -> &Weak, + ) -> Option> { + let stages = self.stages.get(&key_id)?; + let timer = which(stages).upgrade(); + if stages.retire.strong_count() == 0 && stages.reap.strong_count() == 0 { + self.stages.remove(&key_id); + } + timer + } + + /// Schedules both stages of a live key: retirement at its lease's expiry, and + /// the reap a retention window later. /// - /// The entry owns the strong `Arc`, so the lease lives exactly as - /// long as the wheel is watching it: request-facing structures hold only - /// `Weak` references, and dropping the entry is what ends the lease. - fn watch(&mut self, key_id: KeyId, lease: Arc, schedule: Schedule) { + /// WHY the reap timer owns the lease rather than the retire timer: the slot + /// table holds only a `Weak`, so this is the reference that lets a request + /// during the retention window read *why* the key died instead of finding a + /// vanished lease. Firing a timer consumes its callback, so an `Arc` held by + /// the retire stage would be released the moment that stage ran. + fn watch(&mut self, key_id: KeyId, lease: Arc) { let inner = self.inner.clone(); - let deadline = lease.expires_at(); - // WHY the closure keeps the lease across both phases: the slot table - // holds only a `Weak`, so this is the reference that lets a request - // during the tombstone read *why* the key died instead of finding a - // vanished lease. It is released when the entry itself is dropped, after - // the reap. - let mut retired = false; - let phase = move || { - let inner = inner.upgrade()?; - if std::mem::replace(&mut retired, true) { - reap(&inner, key_id); - return None; + let expires_at = lease.expires_at(); + let retire_lease = lease.clone(); + let retire = Timer::new(move || { + if let Some(inner) = inner.upgrade() { + retire(&inner, key_id, &retire_lease); } - Some(retire(&inner, key_id, Some(&lease))) - }; - match schedule { - Schedule::Fresh => self.wheel.schedule(key_id, deadline, phase), - Schedule::Supersede => self.wheel.supersede(key_id, deadline, phase), - } + }); + let reap = self.reap_timer(key_id, Some(lease)); + self.stages.insert( + key_id, + Stages { + retire: Arc::downgrade(&retire), + reap: Arc::downgrade(&reap), + }, + ); + self.wheel.schedule(expires_at, retire); + self.wheel.schedule(retention_ends(expires_at), reap); + } + + /// Schedules only the reap, for a key that is already dead. + fn schedule_reap(&mut self, key_id: KeyId, deadline: u64) { + let reap = self.reap_timer(key_id, None); + self.stages.insert( + key_id, + Stages { + reap: Arc::downgrade(&reap), + ..Stages::default() + }, + ); + self.wheel.schedule(deadline, reap); } - /// Schedules only the reap phase, for a key that is already dead. - fn entomb(&mut self, key_id: KeyId, tombstoned_at: u64) { + /// Builds the reap stage. `lease` is the key's live lease when there is one, + /// kept alive by this timer until the row is recycled. + fn reap_timer(&self, key_id: KeyId, lease: Option>) -> Arc { let inner = self.inner.clone(); - self.wheel - .schedule(key_id, retention_ends(tombstoned_at), move || { - reap(&inner.upgrade()?, key_id); - None - }); + Timer::new(move || { + drop(lease); + if let Some(inner) = inner.upgrade() { + reap(&inner, key_id); + } + }) } } @@ -214,37 +284,24 @@ fn retention_ends(tombstoned_at: u64) -> u64 { tombstoned_at.saturating_add(TOMBSTONE_RETENTION.as_secs()) } -/// Phase 1: cancels the lease, marks the row dead, and records when its -/// retention starts. Returns the deadline of the reap that follows. -fn retire(inner: &Arc, key_id: KeyId, lease: Option<&Arc>) -> u64 { - if let Some(lease) = lease { - // WHY expiry is the fallback reason: a key ended for any other reason was - // already cancelled by the code that knew that reason, and `record_cancel` - // keeps the first one, so this cannot mislabel it. - lease.cancel_expired(); - } +/// Ends a key's active stage: cancels the lease, marks the row dead, records when +/// its retention starts, and schedules the reap that frees the row. +fn retire(inner: &Arc, key_id: KeyId, lease: &Arc) { + // WHY expiry is the fallback reason: a key ended for any other reason was + // already cancelled by the code that knew that reason, and `record_cancel` + // keeps the first one, so this cannot mislabel it. + lease.cancel_expired(); let mut slots = inner.slots_mut(); - let tombstoned_at = match slots.get_mut(key_id.slot().as_index()) { - Some(slot) if slot.holds(key_id) && slot.state == SlotState::Active => { - slot.state = SlotState::Expired; - slot.expires_at - } - // Already marked dead by a revoke, or the row moved on. Its retention - // still has to be honoured, timed from whenever it was marked. - _ => { - drop(slots); - let mut cold = inner.cold_mut(); - let tombstoned_at = match cold.get_mut(&key_id) { - Some(cold) if cold.tombstoned_at != 0 => cold.tombstoned_at, - Some(cold) => { - cold.tombstoned_at = unix_seconds(); - cold.tombstoned_at - } - None => unix_seconds(), - }; - return retention_ends(tombstoned_at); - } + let Some(slot) = slots.get_mut(key_id.slot().as_index()) else { + return; }; + // Already dead, or the row moved on: a revoke marked it and recorded its + // tombstone time, and the reap is already scheduled either way. + if !slot.holds(key_id) || slot.state != SlotState::Active { + return; + } + slot.state = SlotState::Expired; + let tombstoned_at = slot.expires_at; drop(slots); inner .cold_mut() @@ -257,10 +314,9 @@ fn retire(inner: &Arc, key_id: KeyId, lease: Option<&Arc, key_id: KeyId) { let mut slots = inner.slots_mut(); match slots.get_mut(key_id.slot().as_index()) { diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index 04b4f48..16a0fb5 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -1425,197 +1425,116 @@ async fn revoking_keeps_the_row_until_its_retention_elapses() { let _ = std::fs::remove_dir_all(state_dir); } -/// Records which phases ran, so a test can assert on the callback's effects -/// rather than on a return value the wheel no longer produces. +/// Records the timers that fired, so a test can assert on callback effects +/// rather than on a return value the wheel does not produce. #[derive(Clone, Default)] -struct PhaseLog(Arc>>); +struct FireLog(Arc>>); -impl PhaseLog { - fn push(&self, phase: &'static str) { - recover_lock(self.0.lock()).push(phase); +impl FireLog { + fn timer(&self, name: &'static str) -> Arc { + let log = self.0.clone(); + Timer::new(move || recover_lock(log.lock()).push(name)) } - fn phases(&self) -> Vec<&'static str> { + fn fired(&self) -> Vec<&'static str> { recover_lock(self.0.lock()).clone() } } -/// Schedules the two-phase shape `Leases` uses: a first phase that asks for a -/// second one `gap` seconds later, then a final phase. -fn schedule_two_phases( - wheel: &mut TimingWheel, - key_id: KeyId, - deadline: u64, - gap: u64, -) -> PhaseLog { - let log = PhaseLog::default(); - let recorder = log.clone(); - let mut first = true; - wheel.schedule(key_id, deadline, move || { - if std::mem::take(&mut first) { - recorder.push("retire"); - return Some(deadline + gap); - } - recorder.push("reap"); - None - }); - log -} - -fn schedule_once(wheel: &mut TimingWheel, key_id: KeyId, deadline: u64) -> PhaseLog { - let log = PhaseLog::default(); - let recorder = log.clone(); - wheel.schedule(key_id, deadline, move || { - recorder.push("fired"); - None - }); - log -} - -#[test] -fn timing_wheel_runs_the_next_phase_at_the_deadline_it_asked_for() { - let key_id = KeyId::new(Generation::from_u32(1), SlotIndex::from_index(0)); - let mut wheel = TimingWheel::new(1_000); - let log = schedule_two_phases(&mut wheel, key_id, 1_005, 60); - - wheel.advance(1_004); - assert!(log.phases().is_empty()); - wheel.advance(1_005); - assert_eq!(log.phases(), ["retire"]); - // The second phase waits for the deadline the first one returned. - wheel.advance(1_064); - assert_eq!(log.phases(), ["retire"]); - wheel.advance(1_065); - assert_eq!(log.phases(), ["retire", "reap"]); - assert!(!wheel.holds(key_id)); -} - -#[test] -fn timing_wheel_cancel_runs_every_remaining_phase_at_once() { - let key_id = KeyId::new(Generation::from_u32(1), SlotIndex::from_index(0)); - let mut wheel = TimingWheel::new(1_000); - let log = schedule_two_phases(&mut wheel, key_id, 1_005, 60); - - wheel.cancel(key_id); - assert_eq!(log.phases(), ["retire", "reap"]); - assert!(!wheel.holds(key_id)); -} - -#[test] -fn timing_wheel_cancel_after_the_first_phase_runs_only_the_rest() { - let key_id = KeyId::new(Generation::from_u32(1), SlotIndex::from_index(0)); - let mut wheel = TimingWheel::new(1_000); - let log = schedule_two_phases(&mut wheel, key_id, 1_005, 60); - - wheel.advance(1_005); - wheel.cancel(key_id); - assert_eq!(log.phases(), ["retire", "reap"]); -} - -#[test] -fn timing_wheel_drop_runs_every_remaining_phase() { - let key_id = KeyId::new(Generation::from_u32(1), SlotIndex::from_index(0)); - let mut wheel = TimingWheel::new(1_000); - let log = schedule_two_phases(&mut wheel, key_id, 1_005, 60); - - drop(wheel); - assert_eq!(log.phases(), ["retire", "reap"]); +/// Runs the wheel forward to `target`, one tick at a time. +fn run_to(wheel: &mut TimingWheel, target: u64) { + while wheel.now() < target { + wheel.tick(); + } } #[test] -fn timing_wheel_reschedule_moves_an_entry_without_running_it() { - let key_id = KeyId::new(Generation::from_u32(1), SlotIndex::from_index(0)); +fn timing_wheel_fires_a_timer_at_its_deadline() { + let log = FireLog::default(); let mut wheel = TimingWheel::new(1_000); - let log = schedule_once(&mut wheel, key_id, 1_005); + wheel.schedule(1_005, log.timer("timer")); - assert!(wheel.reschedule(key_id, 1_020)); - wheel.advance(1_019); - assert!(log.phases().is_empty()); - wheel.advance(1_020); - assert_eq!(log.phases(), ["fired"]); + run_to(&mut wheel, 1_004); + assert!(log.fired().is_empty()); + run_to(&mut wheel, 1_005); + assert_eq!(log.fired(), ["timer"]); } #[test] -fn timing_wheel_reschedule_reports_a_key_it_does_not_hold() { +fn timing_wheel_fires_each_timer_once() { + let log = FireLog::default(); let mut wheel = TimingWheel::new(1_000); - assert!(!wheel.reschedule( - KeyId::new(Generation::from_u32(1), SlotIndex::from_index(0)), - 2_000 - )); + let timer = log.timer("timer"); + // The same timer placed twice: the first placement to drain is not the last + // reference, so only the later one fires it. + wheel.schedule(1_005, timer.clone()); + wheel.schedule(1_020, timer); + + run_to(&mut wheel, 1_005); + assert!(log.fired().is_empty()); + run_to(&mut wheel, 1_020); + assert_eq!(log.fired(), ["timer"]); } #[test] -fn timing_wheel_scheduling_over_an_entry_finishes_the_one_it_replaces() { - let key_id = KeyId::new(Generation::from_u32(1), SlotIndex::from_index(0)); +fn timing_wheel_firing_early_makes_the_scheduled_placement_inert() { + let log = FireLog::default(); let mut wheel = TimingWheel::new(1_000); - let stale = schedule_two_phases(&mut wheel, key_id, 1_005, 60); - let fresh = schedule_once(&mut wheel, key_id, 1_020); - - assert_eq!(stale.phases(), ["retire", "reap"]); - wheel.advance(1_020); - assert_eq!(fresh.phases(), ["fired"]); + let timer = log.timer("timer"); + wheel.schedule(1_005, timer.clone()); + + timer.fire(); + assert_eq!(log.fired(), ["timer"]); + // The deadline arriving drops the placement, which must not fire it again. + run_to(&mut wheel, 1_005); + assert_eq!(log.fired(), ["timer"]); } #[test] -fn timing_wheel_fires_an_entry_scheduled_in_the_past_without_wrapping() { - let key_id = KeyId::new(Generation::from_u32(1), SlotIndex::from_index(0)); +fn timing_wheel_drop_fires_everything_it_holds() { + let log = FireLog::default(); let mut wheel = TimingWheel::new(1_000); - let log = schedule_once(&mut wheel, key_id, 999); + wheel.schedule(1_005, log.timer("early")); + wheel.schedule(9_999_999, log.timer("late")); - wheel.advance(1_000); - assert_eq!(log.phases(), ["fired"]); + drop(wheel); + let mut fired = log.fired(); + fired.sort_unstable(); + assert_eq!(fired, ["early", "late"]); } #[test] -fn timing_wheel_cascades_an_entry_down_two_levels() { +fn timing_wheel_cascades_a_timer_down_two_levels() { // A deadline two levels up has to reach level 0 before it can be drained. let now = 1_000; let deadline = now + (1 << (6 * 2)); - let key_id = KeyId::new(Generation::from_u32(1), SlotIndex::from_index(0)); + let log = FireLog::default(); let mut wheel = TimingWheel::new(now); - let log = schedule_once(&mut wheel, key_id, deadline); + wheel.schedule(deadline, log.timer("timer")); - wheel.advance(deadline - 1); - assert!(log.phases().is_empty()); - wheel.advance(deadline); - assert_eq!(log.phases(), ["fired"]); + run_to(&mut wheel, deadline - 1); + assert!(log.fired().is_empty()); + run_to(&mut wheel, deadline); + assert_eq!(log.fired(), ["timer"]); } #[test] fn timing_wheel_fires_a_boundary_deadline_without_an_extra_tick() { - let key_id = KeyId::new(Generation::from_u32(1), SlotIndex::from_index(0)); + let log = FireLog::default(); let mut wheel = TimingWheel::new(700); - let log = schedule_once(&mut wheel, key_id, 1_024); - - wheel.advance(1_024); - assert_eq!(log.phases(), ["fired"]); -} - -#[test] -fn timing_wheel_drains_in_one_pass_when_the_clock_jumps_past_every_deadline() { - // A correction longer than the longest schedulable delay leaves nothing - // pending, so the wheel empties without ticking through the elapsed years. - let now = 1_000; - let key_id = KeyId::new(Generation::from_u32(1), SlotIndex::from_index(0)); - let mut wheel = TimingWheel::new(now); - let log = schedule_two_phases(&mut wheel, key_id, now + 60, 60); + wheel.schedule(1_024, log.timer("timer")); - wheel.advance(now + 4 * MAX_TEMP_KEY_TTL.as_secs()); - assert_eq!(log.phases(), ["retire", "reap"]); - assert!(!wheel.holds(key_id)); + run_to(&mut wheel, 1_024); + assert_eq!(log.fired(), ["timer"]); } #[test] -fn timing_wheel_keeps_its_position_when_the_clock_steps_backwards() { - let key_id = KeyId::new(Generation::from_u32(1), SlotIndex::from_index(0)); +fn timing_wheel_fires_a_deadline_already_in_the_past() { + let log = FireLog::default(); let mut wheel = TimingWheel::new(1_000); - let log = schedule_once(&mut wheel, key_id, 1_030); - - wheel.advance(1_020); - wheel.advance(1_005); - assert!(log.phases().is_empty()); - wheel.advance(1_030); - assert_eq!(log.phases(), ["fired"]); + // Scheduling into the past drops the placement immediately. + wheel.schedule(999, log.timer("timer")); + assert_eq!(log.fired(), ["timer"]); } #[test] diff --git a/src/common/auth/timing_wheel.rs b/src/common/auth/timing_wheel.rs index 10ca9ae..8c24c17 100644 --- a/src/common/auth/timing_wheel.rs +++ b/src/common/auth/timing_wheel.rs @@ -1,26 +1,25 @@ -//! Hierarchical timer wheel whose entries are self-advancing cleanup callbacks. +//! Hierarchical timer wheel. It schedules opaque timers and knows nothing about +//! what they mean. //! //! ```text -//! schedule(deadline, callback) -> one level/slot bucket -//! deadline reached -> callback runs -> reschedules itself, or is done -//! dropped early -> callback runs to completion right there +//! schedule(deadline, timer) -> the bucket covering that deadline //! -//! one-second tick -> drain level 0's current slot -//! -> every 64th tick, cascade level 1 into finer levels, and so on +//! tick() -> now += 1 +//! -> drain level 0's slot for `now`; every 64th tick level 1's, and so on +//! -> an entry whose deadline has arrived is dropped; the rest are re-filed +//! into the finer level that now covers them //! ``` //! -//! The wheel knows nothing about what it schedules. An entry owns a callback that -//! performs one phase of some cleanup and returns when it next wants to run, or -//! `None` when finished, so a multi-stage teardown is written once at the point -//! the entry is created. Dropping an entry early runs every phase it has left, in -//! order, immediately — which is what lets cancelling one entry stand in for a -//! whole cleanup routine, and lets dropping the wheel tear down everything it -//! owns without a caller walking any of it. +//! The wheel holds the only strong references to its timers, so a timer runs when +//! the last entry referring to it is dropped. That is what keeps the wheel +//! indifferent to its users: it never looks a timer up, compares identities, or +//! has to be told that one was superseded. //! -//! Each entry lives in exactly one bucket, and `positions` records which, so -//! rescheduling or cancelling one is a map lookup rather than a search. A tick -//! touches one slot per level that turns over, never the whole wheel, so cost -//! tracks the entries that actually cascade or fire. +//! Rescheduling exploits that directly. Holding a `Weak`, a caller inserts +//! the same timer again at a later deadline; the earlier entry still drains on its +//! own schedule, but dropping it no longer brings the count to zero, so the timer +//! waits for the last entry to go. Cancelling is the mirror image: [`Timer::fire`] +//! runs the callback early and leaves the remaining entries inert. use super::*; @@ -29,8 +28,8 @@ use super::*; const SLOT_BITS: u32 = 6; const SLOTS: usize = 1 << SLOT_BITS; const SLOT_MASK: u64 = SLOTS as u64 - 1; -/// Six 64-slot levels span `64^6` seconds, which keeps every deadline a -/// configured TTL can produce inside a level whose range really contains it. +/// Six 64-slot levels span `64^6` seconds, so any deadline a caller can ask for +/// lands in a level whose range really contains it. const NUM_LEVELS: usize = 6; const TOP_LEVEL: Level = NUM_LEVELS as Level - 1; @@ -40,258 +39,113 @@ type Level = u8; /// Which bucket within one level: `0..SLOTS`. type Slot = u8; -/// One phase of a scheduled teardown: does its work and returns the deadline of -/// the phase after it, or `None` once nothing remains. -type Phase = Box Option + Send>; - -struct Entry { - deadline: u64, - phase: Phase, - /// Set once a phase has returned `None`, so a completed entry's drop does - /// not call into its callback again. - finished: bool, +/// A callback that runs once: when its deadline arrives, or when it is cancelled, +/// whichever comes first. +pub(super) struct Timer { + /// `None` once the callback has run, so any remaining wheel entries for this + /// timer are inert and a cancelled timer cannot fire twice. + callback: std::sync::Mutex>>, } -impl Entry { - /// Runs the next phase and reports the deadline it wants, if any. - fn fire(&mut self) -> Option { - if self.finished { - return None; +impl Timer { + pub(super) fn new(callback: impl FnOnce() + Send + 'static) -> Arc { + Arc::new(Self { + callback: std::sync::Mutex::new(Some(Box::new(callback))), + }) + } + + /// Runs the callback unless it has run already. Called by the wheel when a + /// deadline arrives, and by a caller cancelling ahead of that. + pub(super) fn fire(&self) { + let callback = recover_lock(self.callback.lock()).take(); + if let Some(callback) = callback { + callback(); } - let next = (self.phase)(); - self.finished = next.is_none(); - next } } -impl Drop for Entry { - /// An entry let go of before its deadline still owes every phase it has - /// left, so they all run here. This is why cancelling an entry and letting - /// it expire have the same effect, only sooner. +impl Drop for Timer { + /// Releasing the last reference is what fires a timer, so dropping the wheel + /// tears down everything it was holding. fn drop(&mut self) { - while self.fire().is_some() {} + self.fire(); } } -/// Where a key's entry currently sits, so a reschedule or cancel does not have -/// to search the wheel. -#[derive(Clone, Copy)] -enum Position { - /// Scheduled for a deadline the wheel had already passed. - Overdue, - Wheel { - level: Level, - slot: Slot, - }, +/// One placement of a timer. The deadline lives here rather than in the `Timer`, +/// so a timer rescheduled later leaves its earlier placements draining harmlessly +/// instead of dragging them forward. +struct Entry { + deadline: u64, + /// Never read: holding the reference *is* the entry's job, and releasing it + /// is what can fire the timer. + #[allow(dead_code)] + timer: Arc, } -type Bucket = HashMap; - pub(super) struct TimingWheel { now: u64, - positions: HashMap, - /// Entries whose deadline was already past when they were filed. Level 0's - /// slot for `now` was drained this tick, so filing them there would delay - /// them by a full revolution. - overdue: Bucket, - levels: [Vec; NUM_LEVELS], + levels: [Vec>; NUM_LEVELS], } impl TimingWheel { pub(super) fn new(now: u64) -> Self { Self { now, - positions: HashMap::new(), - overdue: Bucket::new(), - levels: std::array::from_fn(|_| { - std::iter::repeat_with(Bucket::new).take(SLOTS).collect() - }), + levels: std::array::from_fn(|_| std::iter::repeat_with(Vec::new).take(SLOTS).collect()), } } - /// Schedules `phase` to run once `deadline` has passed. Any entry already - /// held for `key_id` is dropped, which runs the phases it had left. - pub(super) fn schedule( - &mut self, - key_id: KeyId, - deadline: u64, - phase: impl FnMut() -> Option + Send + 'static, - ) { - self.cancel(key_id); - self.place( - key_id, - Entry { - deadline, - phase: Box::new(phase), - finished: false, - }, - ); + pub(super) fn now(&self) -> u64 { + self.now } - /// Replaces the entry for `key_id`, discarding the previous one *without* - /// running its remaining phases. Use this only when the new entry takes over - /// the same cleanup, so nothing the old one owed is lost; otherwise - /// [`Self::schedule`] is what you want. - pub(super) fn supersede( - &mut self, - key_id: KeyId, - deadline: u64, - phase: impl FnMut() -> Option + Send + 'static, - ) { - if let Some(mut previous) = self.detach(key_id) { - previous.finished = true; - } - self.place( - key_id, - Entry { - deadline, - phase: Box::new(phase), - finished: false, - }, - ); - } - - /// Moves an entry to a new deadline without running anything. Returns - /// `false` when the wheel holds no entry for `key_id`. - pub(super) fn reschedule(&mut self, key_id: KeyId, deadline: u64) -> bool { - let Some(mut entry) = self.detach(key_id) else { - return false; - }; - entry.deadline = deadline; - self.place(key_id, entry); - true + /// Holds `timer` until `deadline`. Scheduling a timer the wheel already holds + /// adds a placement rather than replacing one, which is how a caller moves a + /// deadline outward without the wheel having to find the old entry. + pub(super) fn schedule(&mut self, deadline: u64, timer: Arc) { + self.place(Entry { deadline, timer }); } - /// Runs everything the entry for `key_id` still owes, now rather than at its - /// deadline. A no-op when the wheel holds no entry for it. - pub(super) fn cancel(&mut self, key_id: KeyId) { - drop(self.detach(key_id)); - } - - /// Runs only the entry's next phase, then waits for the deadline that phase - /// asked for. Use this where a stage has arrived early but the stages after - /// it must still keep their own timing — a revoke ends a key's active phase - /// without skipping the retention that follows it. - pub(super) fn advance_one_phase(&mut self, key_id: KeyId) -> bool { - let Some(mut entry) = self.detach(key_id) else { - return false; - }; - match entry.fire() { - Some(next) => { - entry.deadline = next; - self.place(key_id, entry); - } - // Finished, so the drop below has nothing left to run. - None => drop(entry), - } - true - } - - /// Runs the wheel up to `target`, firing every entry whose deadline has - /// passed and re-filing the phases they schedule next. - pub(super) fn advance(&mut self, target: u64) { - let overdue = std::mem::take(&mut self.overdue); - self.settle(overdue, target); - // A jump longer than the longest lifetime the config can produce means - // every entry is already past its deadline, so the whole wheel can be - // drained in one pass. Ticking through it instead would spin for hours - // when a bad hardware clock is corrected forward by years. - if target.saturating_sub(self.now) > MAX_SCHEDULABLE_DELAY.as_secs() { - self.now = target; - // A phase can schedule a successor that is also already overdue, so - // keep draining until a pass leaves nothing due. - loop { - let due = self - .levels - .iter_mut() - .flat_map(|level| level.iter_mut()) - .fold(Bucket::new(), |mut due, bucket| { - due.extend(std::mem::take(bucket)); - due - }); - let overdue = std::mem::take(&mut self.overdue); - if due.is_empty() && overdue.is_empty() { - return; - } - self.settle(due, target); - self.settle(overdue, target); - } - } - while self.now < target { - self.now += 1; - // Coarse to fine, so an entry cascading several levels down still - // reaches level 0 in time to be drained by this same tick. - for level in (1..=TOP_LEVEL).rev() { - // A level turns over once every `slot_range` seconds, exactly - // when `now` has no bits left below that level's slot field. - if self.now & (slot_range(level) - 1) != 0 { - continue; - } - let entries = self.take_bucket(level, self.now); - self.settle(entries, self.now); - } - let entries = self.take_bucket(0, self.now); - self.settle(entries, self.now); - } - // A clock stepping backwards must not rewind the wheel: buckets are - // indexed relative to `now`, so re-indexing against an earlier `now` - // would file entries into slots the wheel has already drained. - self.now = self.now.max(target); - } - - /// Fires the entries due at `deadline` and re-files both the phases they - /// schedule next and the entries that are not due yet. - fn settle(&mut self, entries: Bucket, deadline: u64) { - for (key_id, mut entry) in entries { - if entry.deadline > deadline { - self.place(key_id, entry); + /// Advances one second and drains whatever that turnover exposes. + pub(super) fn tick(&mut self) { + self.now += 1; + // Coarse to fine, so a timer cascading several levels down still reaches + // level 0 in time to be drained by this same tick. + for level in (1..=TOP_LEVEL).rev() { + // A level turns over once every `slot_range` seconds, exactly when + // `now` has no bits left below that level's slot field. + if self.now & (slot_range(level) - 1) != 0 { continue; } - match entry.fire() { - Some(next) => { - entry.deadline = next; - self.place(key_id, entry); - } - None => { - self.positions.remove(&key_id); - } - } + let entries = self.take_bucket(level, self.now); + self.refile(entries); } + let entries = self.take_bucket(0, self.now); + self.refile(entries); } - fn place(&mut self, key_id: KeyId, entry: Entry) { - let position = if entry.deadline <= self.now { - self.overdue.insert(key_id, entry); - Position::Overdue - } else { - let level = level_for(self.now, entry.deadline); - let slot = slot_for(level, entry.deadline); - self.bucket(level, slot).insert(key_id, entry); - Position::Wheel { level, slot } - }; - self.positions.insert(key_id, position); - } - - fn detach(&mut self, key_id: KeyId) -> Option { - match self.positions.remove(&key_id)? { - Position::Overdue => self.overdue.remove(&key_id), - Position::Wheel { level, slot } => self.bucket(level, slot).remove(&key_id), + fn refile(&mut self, entries: Vec) { + for entry in entries { + self.place(entry); } } - fn bucket(&mut self, level: Level, slot: Slot) -> &mut Bucket { - &mut self.levels[level as usize][slot as usize] + fn place(&mut self, entry: Entry) { + // An arrived deadline means this placement is done: returning drops the + // entry, which fires the timer if this was its last reference. + if entry.deadline <= self.now { + return; + } + let level = level_for(self.now, entry.deadline); + let slot = slot_for(level, entry.deadline); + self.levels[level as usize][slot as usize].push(entry); } /// Empties the bucket that `when` falls in at `level`. - fn take_bucket(&mut self, level: Level, when: u64) -> Bucket { + fn take_bucket(&mut self, level: Level, when: u64) -> Vec { let slot = slot_for(level, when); - std::mem::take(self.bucket(level, slot)) - } - - #[cfg(test)] - pub(super) fn holds(&self, key_id: KeyId) -> bool { - self.positions.contains_key(&key_id) + std::mem::take(&mut self.levels[level as usize][slot as usize]) } } @@ -306,8 +160,8 @@ fn slot_for(level: Level, when: u64) -> Slot { /// Finest level able to hold `deadline`: the one whose slot field covers the /// highest bit in which `now` and `deadline` differ. A deadline past the top -/// level is clamped into it and cannot fire early, because a drained entry is -/// only fired once its own deadline has passed. +/// level is clamped into it and cannot fire early, because an entry is dropped +/// only once its own deadline has arrived. fn level_for(now: u64, deadline: u64) -> Level { let significant = 63 - ((now ^ deadline) | SLOT_MASK).leading_zeros(); ((significant / SLOT_BITS) as Level).min(TOP_LEVEL) From d3d6ac44da2e8eb1259aa0a6746583bf894620d0 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Fri, 21 Aug 2026 06:29:21 +0800 Subject: [PATCH 62/74] Rebuild the timer wheel around nested routes; move to parking_lot MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The wheel kept a `HashMap` beside its buckets and, on a large clock jump, folded every bucket into a single pass — a whole-structure walk, which is the one thing a timing wheel exists to avoid. Entries were also indexed by absolute deadline, so draining a coarse bucket had to recompute where each entry belonged and refile it. Buckets now hold routes instead. Scheduling decomposes the delay into base-`radix` digits and builds one nested `Link` per digit, coarsest outermost; each `Link::Relay` waits in one bucket and, when that bucket comes off the front, hands the leg nested inside it to the next, finer bucket. The chain *is* the route, so a tick is `pop_front`, `push_back`, and handing legs on — no arithmetic per entry, no key map, no positions. Level count is derived from the longest delay the wheel must support rather than hardcoded, and radix is a parameter. Placement is verified exhaustively rather than sampled: every delay from every starting offset must fire on exactly the tick it asked for. A radix decomposition is easy to get wrong by one bucket, and the first draft was. Locks: the wheel's per-level mutexes are gone, because they were not a requirement. They existed only so `Link::drop` could file its own successor, which needed shared access to the queues. Dropping now yields data and the wheel files it with `&mut self`, so the sharing — and the locks — disappear. `Timer::callback` keeps a lock: two routes share one timer, and `Arc: Send` (which `tokio::spawn` demands) implies `T: Sync`, so `Cell` will not do. The common path skips it anyway, since `Drop` has `&mut self`. Replace the remaining std locks with parking_lot. Nothing used poisoning — every call site was `unwrap_or_else(|poisoned| poisoned.into_inner())`, plus a `recover_lock` helper to hide the `LockResult`. All of that goes, along with an unreachable "state is poisoned" error path in the FFI crate. Document the reasoning in docs/rust-shared-mutability-and-locks.zh-CN.md: `Arc` governs lifetime, `Send`/`Sync` govern cross-thread safety, interior mutability governs writing through `&T` — three orthogonal mechanisms — and the judgement order that follows (can it be exclusive? must it cross threads? integer or critical section?), worked through the locks above. Co-Authored-By: Claude Opus 5 (1M context) --- Cargo.lock | 2 + Cargo.toml | 2 + .../rust-shared-mutability-and-locks.zh-CN.md | 370 +++++++++++++ src/common/auth.rs | 37 +- src/common/auth/actor/epoch.rs | 6 +- src/common/auth/actor/mod.rs | 17 +- src/common/auth/leases.rs | 38 +- src/common/auth/persistence/snapshot.rs | 6 +- src/common/auth/runtime.rs | 12 +- src/common/auth/tests.rs | 112 ---- src/common/auth/timing_wheel.rs | 487 ++++++++++++++---- src/common/checksum.rs | 33 +- src/common/message/forward.rs | 18 +- src/common/message/secure.rs | 5 +- src/common/message/secure/first_flight.rs | 6 +- ui/native/pb_mapper_ffi/Cargo.toml | 1 + ui/native/pb_mapper_ffi/src/state.rs | 21 +- .../pb_mapper_ffi/src/state/configuration.rs | 4 +- 18 files changed, 854 insertions(+), 323 deletions(-) create mode 100644 docs/rust-shared-mutability-and-locks.zh-CN.md diff --git a/Cargo.lock b/Cargo.lock index bc6876c..4c3b7f0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -788,6 +788,7 @@ dependencies = [ "hashbrown 0.16.1", "kanal", "once_cell", + "parking_lot", "rand 0.10.0", "ring", "serde", @@ -810,6 +811,7 @@ dependencies = [ "better_mimalloc_rs", "clap", "dirs", + "parking_lot", "pb-mapper", "serde", "serde_json", diff --git a/Cargo.toml b/Cargo.toml index 47d0b77..79bce79 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -28,6 +28,7 @@ uni-stream.workspace = true kanal.workspace = true base64.workspace = true subtle.workspace = true +parking_lot.workspace = true [dev-dependencies] dotenvy = "0.15.7" @@ -68,6 +69,7 @@ ring = "0.17.14" once_cell = "1.20.2" base64 = "0.22.1" subtle = "2.6.1" +parking_lot = "0.12" uni-stream = { git = "https://github.com/acking-you/uni-stream.git", branch = "master" } kanal = { git = "https://github.com/acking-you/kanal.git", branch = "dev/pb-mapper" } diff --git a/docs/rust-shared-mutability-and-locks.zh-CN.md b/docs/rust-shared-mutability-and-locks.zh-CN.md new file mode 100644 index 0000000..08a4a5f --- /dev/null +++ b/docs/rust-shared-mutability-and-locks.zh-CN.md @@ -0,0 +1,370 @@ +# 共享可变性、Arc 与锁：从时间轮重构中得到的判断顺序 + +面向场景：你在设计一个数据结构时，发现「好像得加个锁」，但不确定这个锁到底是需求还是自己造出来的。 + +> **Code Version**: pb-mapper 工作区，`feat/temporary-credential-auth`（2026-08-21） +> +> 相关文档：[Send/Sync/Pin 与 async 状态机深度解析](./rust-async-send-sync-pin-deep-dive.md) +> 覆盖 `async fn` 如何编译成状态机、`Pin` 为什么必要。本文讲的是它的另一面： +> **数据结构层面**该不该共享、该不该加锁。 + +## 1. 问题从哪里来 + +重构 `src/common/auth/timing_wheel.rs` 时，中间某一版长成这样： + +```rust +struct Queues { + levels: Vec>>>, // 每一级一把锁 +} +``` + +时间轮的每一级都上了一把锁。而这个时间轮是被 auth actor **独占**的——`Leases` +从头到尾以 `&mut Leases` 传递，从未进过 `Arc`。既然没有任何并发，这些锁是从哪来的？ + +答案是我自己造出来的：我让 `Link::drop` 自己把下一跳投递进桶里。`Drop::drop` +只有 `&mut self`（指向 Link 自己），拿不到轮子的 `&mut`，所以只能让 Link 持有一个 +`Weak` 共享轮子——**一旦共享，就必须加锁**。 + +去掉共享，锁就自己消失了：让 `Drop` 只保留数据，由 `tick` 拿着 `&mut self` 投递。 + +这件事暴露出一个常见的思维捷径： + +> ❌ 「共享可变 → 加锁」 + +它跳过了两个更该先问的问题。本文把正确的判断顺序拆出来。 + +## 2. 三个正交的机制 + +先把三件经常被混为一谈的事分开。它们各管一件事，互不替代。 + +| 机制 | 管什么 | 典型工具 | +|---|---|---| +| **所有权 / 借用** | 谁能改、谁能读 | `&mut T` / `&T` | +| **`Arc`** | 生命周期：owner 何时死 | `Arc` / `Rc` | +| **`Send` / `Sync`** | 跨线程访问的安全性 | marker trait，编译器自动推导 | +| **内部可变性** | 通过 `&T` 修改 | `Cell` / `RefCell` / `Atomic` / `Mutex` | + +``` + 「我能改它吗」 「它还活着吗」 「换线程安全吗」 + │ │ │ + 借用规则 Arc/Rc Send/Sync + │ │ │ + └──── 三者独立，缺一不可，且不能互相代替 ────────┘ +``` + +一个具体的反直觉例子：`Vec` 本身就是 `Send + Sync`，但这**不代表**你可以 +不用 `Arc` 就把它共享给 `tokio::spawn` 的任务。`Send + Sync` 解决的是安全性， +`'static` 生命周期要求得靠 `Arc` 解决。反过来，`Arc>` 生命周期没问题， +但因为 `Cell` 不是 `Sync`，一样过不了 `spawn`。 + +## 3. Send 与 Sync 到底是什么 + +两个 marker trait，不含任何方法，编译器**自动推导**（结构体所有字段都满足则它满足）： + +- **`Send`**：这个值可以**移动**到另一个线程。 +- **`Sync`**：这个值可以被**多个线程同时引用**。 + +第二条有个更精确、也更好用的等价定义： + +``` +T: Sync ⟺ &T: Send +``` + +即「把它的引用发给别的线程是否安全」。这比「线程安全」这种模糊说法准确得多。 + +### 3.1 为什么必须区分这两件事 + +看 `Cell`： + +```rust +let c = Cell::new(0); +thread::spawn(move || c.set(1)); // 独占移过去 —— 安全，故 Cell: Send +// 但： +thread::scope(|s| { + s.spawn(|| c.set(1)); // 两个线程同时 set + s.spawn(|| c.set(2)); // 数据竞争，故 Cell: !Sync +}); +``` + +`Cell::set` 只是一条普通写入，没有任何同步。**移动**给一个线程完全安全（原线程 +再也碰不到它）；**同时借给两个线程**就是 UB。这恰好是 `Send` 与 `Sync` 的分界。 + +### 3.2 Sync 只赋予「共享读」，不赋予「共享写」 + +这是最容易混淆的一点。大多数类型（`u64`、`String`、`Vec`、`HashMap`）都是 +`Send + Sync`——但你拿到 `&Vec` 依然改不了它： + +```rust +let shared = Arc::new(vec![1_u8, 2, 3]); +shared.push(4); +// error[E0596]: cannot borrow data in an `Arc` as mutable +``` + +`Vec` 是 `Sync` 的**原因**正是「通过 `&T` 改不了它」。它 `Sync` 是因为它老实, +不是因为它做了同步。于是 `Sync` 的类型分成两类： + +| 类别 | 例子 | 为什么 Sync | +|---|---|---| +| **老实类型** | `u64`, `Vec`, `HashMap` | 通过 `&T` 根本改不了，无从竞争 | +| **内部可变 + 自带同步** | `Mutex`, `AtomicU64`, `RwLock` | 通过 `&T` 能改，但访问被串行化 | + +`Cell` 是第三类：通过 `&T` 能改，**且不带同步**——所以它被排除在 `Sync` 之外。 +三类划清，`Sync` 的定义就自洽了。 + +> 💡 **Key Point**：需要锁的条件不是「T 不是 Sync」，而是**「我要通过 `&T` 修改它」**。 + +### 3.3 Mutex 的类型学意义 + +``` +T: Send ──[ 包一层 Mutex ]──> Mutex: Sync +``` + +**锁的作用就是把「可移动」升级成「可共享」。** 所以 `Mutex>` 是 `Sync` +的——一个类型不是 `Sync` 从来不是死路。 + +## 4. 为什么 Rc 既不 Send 也不 Sync + +根因只有一个：**`Rc` 的引用计数是普通 `usize` 加减，非原子**。那正是它比 `Arc` 快的 +全部原因。但这一个根因导致两个独立后果： + +**`!Sync`**（直觉的那一半）：`Rc::clone` 只要 `&self`，而它会 `count += 1`。两个线程 +各持 `&Rc` 同时 clone，两次非原子递增丢一次计数 → 提前释放 → use-after-free。 + +**`!Send`**（更微妙，也更关键）：你可能想「整个移过去不就独占了吗」。但 +**`Rc` 从来不是唯一的那一份**： + +```rust +let here = Rc::new(0); +let there = here.clone(); // 两个 handle，一个共享的非原子计数 +thread::spawn(move || drop(there)); // 那边 count -= 1 +drop(here); // 这边 count -= 1，同时进行 +``` + +`there` 移走了，`here` 还在原线程。两边同时递减同一个非原子计数 → 泄漏或双重释放。 + +> 💡 **Key Point**：`Rc: Send` 不安全，不是因为被移动的那一份，而是因为 +> **留在原地的那些**。类型系统无法表达「仅当这是最后一份 handle 时才允许移动」, +> 所以只能整个禁掉。 + +对比 `Cell` 就完整了： + +| | 计数 | Send | Sync | +|---|---|---|---| +| `Rc` | 非原子 | ✗ 别的 handle 会同时改计数 | ✗ `&Rc` 就能 clone | +| `Arc` | 原子 | ✓（需 `T: Send + Sync`）| ✓（同）| +| `Cell` | — | ✓ 移走后原线程什么都不剩 | ✗ `&Cell` 就能 set | + +`Cell: Send` 而 `Rc: !Send`，差别正在于「移走后原线程手里还有没有东西」。 + +## 5. Arc 管的是生命周期，不是安全性 + +既然 `Vec` 本身就 `Send + Sync`，为什么还需要 `Arc`？直接 `&Vec` 不行吗？ + +**在能证明作用域的场景里，确实不需要**： + +```rust +let data = vec![1_u8, 2, 3]; +thread::scope(|s| { + s.spawn(|| println!("{:?}", &data)); // 零 Arc + s.spawn(|| println!("{:?}", &data)); +}); +println!("still owned: {:?}", data); // 依然是 owner +``` + +`thread::scope` 保证所有子线程在它返回前 join，所以编译器**能证明** `data` 活得更久。 + +`tokio::spawn` 和 `thread::spawn` 则不然——任务是 detached 的，何时结束由运行时决定。 +编译器无法证明任何栈上的东西活得比它久，于是有了 `'static` 约束。满足它只有两条路： + +1. 把所有权 `move` 进去 → 只有一个任务能拿到，没法共享； +2. 用 `Arc` → 所有权归引用计数集体所有，**没有任何栈帧是它的 owner**， + 于是每个持有者天然满足 `'static`。 + +``` + 能证明作用域 不能（detached / 'static） + 只读 &T（零成本） Arc + 要改 &mut T（零成本） Arc> +``` + +> 💡 **Key Point**：`Arc` 把生命周期问题转成运行时引用计数，代价是一次原子加减。 +> 这跟 `T` 是否 `Sync` 无关——`Sync` 只决定 `Arc` 能不能 `Send`。 + +## 6. tokio::spawn 的签名从哪来 + +```rust +pub fn spawn(future: F) -> JoinHandle +where F: Future + Send + 'static, F::Output: Send + 'static +``` + +三个约束各有其因： + +- **`Send`**：tokio 多线程调度器有 work-stealing——空闲 worker 会从别的 worker + 队列里偷任务。你的 future 可能在线程 A 上 poll 一次、挂起、然后在线程 B 上 poll + 下一次。它是被**移动**过去的，故需 `Send`。 +- **`'static`**：future 存活时间由运行时决定，不受调用处作用域约束，不能借用栈上的东西。 +- **`F::Output: Send`**：结果要从 worker 线程送回 `JoinHandle` 的等待方。 + +对 `async fn`，编译器把它编译成状态机，**所有跨 `.await` 存活的局部变量都成为该状态机 +的字段**。于是「future 是否 `Send`」= 「这些字段是否全部 `Send`」。这就是为什么一个 +持有 `Rc` 的 async 函数无法 `spawn`——哪怕只在两个 `.await` 之间用了一下。 + +（状态机的展开细节见 +[Send/Sync/Pin 与 async 状态机深度解析 §3](./rust-async-send-sync-pin-deep-dive.md)。） + +逃逸口：`spawn_local` 没有 `Send` 约束，代价是任务被钉在单线程 `LocalSet` 上， +拿不到 work-stealing。 + +## 7. 判断顺序 + +把上面几节合起来，得到一个可执行的检查表。**顺序很重要**——跳过前两问就会得到 +第 1 节那种每级一把锁的东西。 + +``` +① 这里为什么是共享的？能不能改成独占？ + │ + ├─ 能 ──> 用 &mut T，到此结束（零成本，无锁） + │ + ↓ 不能 +② 共享是否真的要跨线程？ + │ + ├─ 不必 ──> Cell / RefCell（零成本 / 一个计数器） + │ + ↓ 必须（Send/Sync 约束逼上来了） +③ 要改的是什么粒度？ + │ + ├─ 单个整数/指针 ──> Atomic（无锁） + │ + └─ 一段临界区 ────> Mutex / RwLock +``` + +第 ① 步最容易被跳过，而它恰恰是收益最大的一步：**共享是可以被设计掉的**。 + +## 8. 落到 pb-mapper 的真实代码 + +### 8.1 被设计掉的锁：时间轮的 Queues + +第 1 节那版每级一把 `Mutex`，走的是「`Drop` 里投递 → 需要共享 `Queues` → 加锁」。 +现在 `Link::Relay` 只是纯数据，`tick` 拿 `&mut self` 自己投递 +（`src/common/auth/timing_wheel.rs`）： + +```rust +Link::Relay { level, slot, next } => self.file(level as usize, slot as usize, *next), +``` + +停在第 ① 步。`Queues` 类型、`Weak`、以及那一排锁全部消失。 + +### 8.2 无锁的共享可变：AuthLease.expires_at + +```rust +pub struct AuthLease { + expires_at: AtomicU64, // src/common/auth.rs + ... +} +``` + +lease 通过 `Arc` 共享（请求侧持 `Weak`，时间轮持强引用），续期要改 `expires_at`, +所以是货真价实的「跨线程共享可变」。但它只读写一个 `u64`，停在第 ③ 步的 `Atomic` 分支, +不需要锁。 + +### 8.3 无法避免的锁：Timer.callback + +```rust +pub(super) struct Timer { + callback: Mutex>>, +} +``` + +这把锁走完了全部三步，每一步都无路可退： + +``` +① 能独占吗？ 不能 —— 续期时 retire/reap 两条路径共享同一个 Timer， + 且「最后一个引用被丢弃时触发」这个语义本身就是引用计数 +② 能只单线程吗？ 不能 —— 推导链如下 +③ 能用 Atomic 吗？ 不能 —— FnOnce 只能按值调用，必须把整个 Box 移出来 +``` + +第 ② 步的推导链值得完整写出来，它是本文所有概念的汇合点： + +``` +tokio::spawn(run_auth_actor(...)) 要求 future: Send + → actor future 跨 .await 持有 Leases ⇒ Leases: Send + → Leases 持有 Arc ⇒ Arc: Send + → Arc: Send 需要 T: Send + Sync ⇒ Timer: Sync + → Timer: Sync 需要 callback 字段: Sync + → Cell 不是 Sync，Mutex 是（当 T: Send） +``` + +（`Arc: Send` 为何要 `T: Sync`：clone 出去后另一个线程通过它拿到 `&T`， +那正是 `Sync` 管的事；同时也要 `T: Send`，因为那个线程可能持有最后一份引用 +并在自己那里析构 `T`。） + +不过实际开销接近零：**绝大多数 timer 从不加锁触发**。`Drop for Timer` 有 +`&mut self`，走 `Mutex::get_mut()`： + +```rust +impl Drop for Timer { + fn drop(&mut self) { + let callback = self.callback.get_mut().take(); // 无锁 + run(callback); + } +} +``` + +只有显式提前 `fire()`（revoke、GC）才真正 lock，而那时也无人竞争。 + +### 8.4 停在第一步：Leases.stages + +```rust +pub(super) struct Leases { + stages: HashMap, // 没有 Arc，没有锁 +} +``` + +`HashMap` 是 `Send + Sync`，但这在这里毫不相关——actor 独占 `Leases`，所有方法都是 +`&mut self`。**`HashMap` 是不是 `Sync` 根本不影响这个决定。** + +### 8.5 必要的锁：AuthStateInner + +```rust +struct AuthStateInner { + slots: RwLock>, + cold: RwLock>, + ... +} +``` + +它被 `Arc` 共享给请求处理路径和 actor 两边，双方都要改，且要保护的是「查槽位 → +校验 generation → 改状态」这样的临界区而非单个整数。三步走完，`RwLock` 是对的。 + +## 9. 为什么用 parking_lot + +标准库的 `Mutex`/`RwLock` 带 **poisoning**：持锁线程 panic 后，锁被标记为「有毒」, +后续 `lock()` 返回 `Err`。本项目从不利用这个信号——迁移前每个调用点都是同一句样板： + +```rust +lock.read().unwrap_or_else(|poisoned| poisoned.into_inner()) +``` + +即「无论如何都取出内部值」，等于把 poisoning 显式关掉。`src/common/auth.rs` 里还为此 +养了一个 `recover_lock` 辅助函数专门抹掉 `LockResult`。 + +`parking_lot` 不做 poisoning，于是： + +- `lock()` 直接返回 guard，所有样板和 `recover_lock` 一起消失； +- FFI 侧 `claim_key` 里那条「state is poisoned」错误分支变成不可达，直接删掉—— + 少一个永远不会发生的错误码； +- 未竞争时是纯自旋 + 无系统调用，比 std 的 futex 路径更快； +- 锁本身不必为 poisoning 保留状态，`Mutex` 只占一个字节加 `T`。 + +代价：panic 后锁会正常释放，其他线程可能看到中间状态的值。本项目原先就用 +`into_inner()` 接受了这个行为，所以迁移是纯简化，语义不变。 + +## 10. 小结 + +- **`Arc` 管生命周期，`Send`/`Sync` 管跨线程安全性，内部可变性管「通过 `&T` 修改」。** + 三件事正交，不能互相代替。 +- `Sync` 只赋予共享**读**。需要锁的条件是「我要通过 `&T` 改它」，而非「T 不是 Sync」。 +- `Mutex` 的类型学意义：把 `Send` 升级成 `Sync`。 +- 判断顺序：**能不能不共享 → 是否真要跨线程 → 粒度是整数还是临界区**。 + 第一步收益最大，也最常被跳过。 +- 一把锁如果无法说清它走完了这三步，它大概是被设计出来的，而不是需求。 diff --git a/src/common/auth.rs b/src/common/auth.rs index a85789a..25098eb 100644 --- a/src/common/auth.rs +++ b/src/common/auth.rs @@ -55,9 +55,10 @@ use std::io::{Read, Write}; use std::os::unix::fs::PermissionsExt; use std::path::{Path, PathBuf}; use std::sync::atomic::{AtomicBool, AtomicU64, AtomicU8, Ordering}; -use std::sync::{Arc, RwLock, Weak}; +use std::sync::{Arc, Weak}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use parking_lot::{Mutex, RwLock}; use rand::RngExt; use ring::aead::{Aad, LessSafeKey, Nonce, UnboundKey, AES_256_GCM}; use ring::hkdf::{Salt, HKDF_SHA256}; @@ -500,43 +501,39 @@ struct AuthStateInner { audit_records: RwLock>, } -fn recover_lock(result: std::sync::LockResult) -> T { - result.unwrap_or_else(|poisoned| poisoned.into_inner()) -} - impl AuthStateInner { /// Rows the configured capacity no longer covers. See the field's docs; the /// fallback is a scan because the range is small and rarely touched. - fn high(&self) -> std::sync::RwLockReadGuard<'_, Vec> { - recover_lock(self.high_slot_entries.read()) + fn high(&self) -> parking_lot::RwLockReadGuard<'_, Vec> { + self.high_slot_entries.read() } - fn high_mut(&self) -> std::sync::RwLockWriteGuard<'_, Vec> { - recover_lock(self.high_slot_entries.write()) + fn high_mut(&self) -> parking_lot::RwLockWriteGuard<'_, Vec> { + self.high_slot_entries.write() } - fn slots(&self) -> std::sync::RwLockReadGuard<'_, Box<[SlotHot]>> { - recover_lock(self.slots.read()) + fn slots(&self) -> parking_lot::RwLockReadGuard<'_, Box<[SlotHot]>> { + self.slots.read() } - fn slots_mut(&self) -> std::sync::RwLockWriteGuard<'_, Box<[SlotHot]>> { - recover_lock(self.slots.write()) + fn slots_mut(&self) -> parking_lot::RwLockWriteGuard<'_, Box<[SlotHot]>> { + self.slots.write() } - fn cold(&self) -> std::sync::RwLockReadGuard<'_, HashMap> { - recover_lock(self.cold.read()) + fn cold(&self) -> parking_lot::RwLockReadGuard<'_, HashMap> { + self.cold.read() } - fn cold_mut(&self) -> std::sync::RwLockWriteGuard<'_, HashMap> { - recover_lock(self.cold.write()) + fn cold_mut(&self) -> parking_lot::RwLockWriteGuard<'_, HashMap> { + self.cold.write() } fn admin_key(&self) -> AesKeyType { - recover_lock(self.admin.read()).key + self.admin.read().key } fn instance_id(&self) -> [u8; INSTANCE_ID_LEN] { - *recover_lock(self.instance_id.read()) + *self.instance_id.read() } } @@ -546,7 +543,7 @@ pub struct AuthRuntime { command_tx: mpsc::Sender, config: AuthConfig, _state_lock: Arc, - actor: Arc>>>, + actor: Arc>>>, actor_abort: tokio::task::AbortHandle, } diff --git a/src/common/auth/actor/epoch.rs b/src/common/auth/actor/epoch.rs index 84497e5..d72cbc6 100644 --- a/src/common/auth/actor/epoch.rs +++ b/src/common/auth/actor/epoch.rs @@ -3,7 +3,7 @@ use super::super::*; use super::{audit, ensure_store_available}; fn remember_previous_root(inner: &AuthStateInner) { - *recover_lock(inner.previous_root.write()) = Some(PreviousRoot { + *inner.previous_root.write() = Some(PreviousRoot { admin_key: inner.admin_key(), instance_id: inner.instance_id(), }); @@ -48,7 +48,7 @@ pub(super) fn actor_reset( push_audit_record(inner, reset_audit); remember_previous_root(inner); leases.wipe(unix_seconds()); - *recover_lock(inner.instance_id.write()) = new_instance_id; + *inner.instance_id.write() = new_instance_id; inner.safe_mode.store(false, Ordering::Release); Ok(()) } @@ -103,7 +103,7 @@ pub(super) fn actor_rotate_root( leases.wipe(unix_seconds()); let old_admin_lease = admin_lease.clone(); let new_admin_lease = Arc::new(AuthLease::new(ADMIN_KEY_ID, u64::MAX)); - *recover_lock(inner.admin.write()) = AdminState { + *inner.admin.write() = AdminState { key: new_key, lease: Arc::downgrade(&new_admin_lease), }; diff --git a/src/common/auth/actor/mod.rs b/src/common/auth/actor/mod.rs index 20e2c75..2945614 100644 --- a/src/common/auth/actor/mod.rs +++ b/src/common/auth/actor/mod.rs @@ -328,16 +328,13 @@ fn validate_admin_authority( false, )); } - let current = recover_lock(inner.admin.read()) - .lease - .upgrade() - .ok_or_else(|| { - AuthFailure::new( - "administrator_key_rotated", - "active administrator credential lease is unavailable", - false, - ) - })?; + let current = inner.admin.read().lease.upgrade().ok_or_else(|| { + AuthFailure::new( + "administrator_key_rotated", + "active administrator credential lease is unavailable", + false, + ) + })?; if !Arc::ptr_eq(&presented, ¤t) { return Err(AuthFailure::new( "administrator_key_rotated", diff --git a/src/common/auth/leases.rs b/src/common/auth/leases.rs index b663e2a..8fb741c 100644 --- a/src/common/auth/leases.rs +++ b/src/common/auth/leases.rs @@ -41,6 +41,10 @@ struct Stages { pub(super) struct Leases { inner: Weak, wheel: TimingWheel, + /// The wall-clock second the wheel's current tick corresponds to. The wheel + /// itself only counts ticks, so this is where absolute deadlines are turned + /// into the relative delays it takes. + now: u64, /// Each key's stages, so a renew or an early end can reach them without the /// wheel knowing what a key is. stages: HashMap, @@ -52,7 +56,8 @@ impl Leases { pub(super) fn restored(inner: &Arc, now: u64) -> Self { let mut leases = Self { inner: Arc::downgrade(inner), - wheel: TimingWheel::new(now), + wheel: new_wheel(), + now, stages: HashMap::new(), }; let mut live = Vec::new(); @@ -123,8 +128,8 @@ impl Leases { let (Some(retire), Some(reap)) = (stages.retire.upgrade(), stages.reap.upgrade()) else { return false; }; - self.wheel.schedule(expires_at, retire); - self.wheel.schedule(retention_ends(expires_at), reap); + self.schedule_at(expires_at, retire); + self.schedule_at(retention_ends(expires_at), reap); true } @@ -156,11 +161,12 @@ impl Leases { // A clock stepping backwards is ignored: buckets are indexed relative to // the wheel's `now`, so re-filing against an earlier one would place // entries in slots it has already drained. - if now.saturating_sub(self.wheel.now()) > MAX_SCHEDULABLE_DELAY.as_secs() { + if now.saturating_sub(self.now) > self.wheel.max_delay() { self.drop_schedule(now); return; } - while self.wheel.now() < now { + while self.now < now { + self.now += 1; self.wheel.tick(); } } @@ -208,7 +214,15 @@ impl Leases { /// Replaces the whole schedule, running every callback the old one held. fn drop_schedule(&mut self, now: u64) { self.stages.clear(); - self.wheel = TimingWheel::new(now); + self.now = now; + self.wheel = new_wheel(); + } + + /// Schedules `timer` for an absolute second, as the delay from now the wheel + /// works in. A deadline already past releases the timer at once. + fn schedule_at(&mut self, deadline: u64, timer: Arc) { + self.wheel + .schedule(deadline.saturating_sub(self.now), timer); } /// Upgrades one of a key's stages, forgetting the key once both have run. @@ -250,8 +264,8 @@ impl Leases { reap: Arc::downgrade(&reap), }, ); - self.wheel.schedule(expires_at, retire); - self.wheel.schedule(retention_ends(expires_at), reap); + self.schedule_at(expires_at, retire); + self.schedule_at(retention_ends(expires_at), reap); } /// Schedules only the reap, for a key that is already dead. @@ -264,7 +278,7 @@ impl Leases { ..Stages::default() }, ); - self.wheel.schedule(deadline, reap); + self.schedule_at(deadline, reap); } /// Builds the reap stage. `lease` is the key's live lease when there is one, @@ -334,3 +348,9 @@ fn reap(inner: &Arc, key_id: KeyId) { } inner.cold_mut().remove(&key_id); } + +/// The wheel every schedule uses: wide enough for the longest lifetime a key can +/// have, at 64 buckets per level. +fn new_wheel() -> TimingWheel { + TimingWheel::new(MAX_SCHEDULABLE_DELAY.as_secs(), 64) +} diff --git a/src/common/auth/persistence/snapshot.rs b/src/common/auth/persistence/snapshot.rs index a8bead0..c3c5dff 100644 --- a/src/common/auth/persistence/snapshot.rs +++ b/src/common/auth/persistence/snapshot.rs @@ -7,7 +7,7 @@ pub(in crate::common::auth) fn compaction_is_allowed(safe_mode: bool) -> bool { } pub(in crate::common::auth) fn push_audit_record(inner: &AuthStateInner, record: AuditRecord) { - let mut records = recover_lock(inner.audit_records.write()); + let mut records = inner.audit_records.write(); while records.len() >= AUDIT_RECORD_CAPACITY { records.pop_front(); } @@ -23,7 +23,7 @@ pub(in crate::common::auth) fn cancel_all_temporary_leases(inner: &AuthStateInne fn snapshot_generations(inner: &AuthStateInner) -> Vec { let slots = inner.slots(); - let extra = recover_lock(inner.high_slot_generations.read()); + let extra = inner.high_slot_generations.read(); let mut generations = slots.iter().map(|slot| slot.generation).collect::>(); generations.extend_from_slice(&extra); generations @@ -141,7 +141,7 @@ fn snapshot_with( LegacyProtocolPolicy::Deny }, admin_replays: admin_replays.iter().cloned().collect(), - audit_records: recover_lock(inner.audit_records.read()).clone(), + audit_records: inner.audit_records.read().clone(), root_epoch: inner.root_epoch.load(Ordering::Acquire), } } diff --git a/src/common/auth/runtime.rs b/src/common/auth/runtime.rs index c04b4d4..c33f011 100644 --- a/src/common/auth/runtime.rs +++ b/src/common/auth/runtime.rs @@ -201,7 +201,7 @@ impl AuthRuntime { command_tx, config: config.clone(), _state_lock: state_lock.clone(), - actor: Arc::new(std::sync::Mutex::new(Some(actor))), + actor: Arc::new(Mutex::new(Some(actor))), actor_abort, }; Ok(runtime) @@ -214,7 +214,7 @@ impl AuthRuntime { .send(AuthCommand::Shutdown { response }) .await; let _ = receiver.await; - let handle = recover_lock(self.actor.lock()).take(); + let handle = self.actor.lock().take(); if let Some(handle) = handle { let _ = handle.await; } @@ -222,7 +222,7 @@ impl AuthRuntime { pub async fn abort_actor(&self) -> Result<(), AuthFailure> { self.actor_abort.abort(); - let handle = recover_lock(self.actor.lock()).take(); + let handle = self.actor.lock().take(); if let Some(handle) = handle { let _ = handle.await; } @@ -274,7 +274,7 @@ impl AuthRuntime { pub(crate) fn derive_previous_key(&self, key_id: KeyId) -> Option { let inner = self.inner().ok()?; - let previous = recover_lock(inner.previous_root.read()).clone()?; + let previous = inner.previous_root.read().clone()?; if key_id.is_admin() { Some(previous.admin_key) } else { @@ -289,7 +289,7 @@ impl AuthRuntime { ) -> Result { let inner = self.inner()?; if key_id.is_admin() { - let admin = recover_lock(inner.admin.read()); + let admin = inner.admin.read(); if !bool::from(presented_key.ct_eq(&admin.key)) { inner.auth_failures.fetch_add(1, Ordering::Relaxed); return Err(AuthFailure::new( @@ -603,7 +603,7 @@ fn temporary_key_material_mismatch(inner: &AuthStateInner, key_id: KeyId) -> Aut let current_generation = match slots.get(index) { Some(slot) => Some(slot.generation), None => { - let high = recover_lock(inner.high_slot_generations.read()); + let high = inner.high_slot_generations.read(); index .checked_sub(slots.len()) .and_then(|offset| high.get(offset).copied()) diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index 16a0fb5..6d82cc6 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -1425,118 +1425,6 @@ async fn revoking_keeps_the_row_until_its_retention_elapses() { let _ = std::fs::remove_dir_all(state_dir); } -/// Records the timers that fired, so a test can assert on callback effects -/// rather than on a return value the wheel does not produce. -#[derive(Clone, Default)] -struct FireLog(Arc>>); - -impl FireLog { - fn timer(&self, name: &'static str) -> Arc { - let log = self.0.clone(); - Timer::new(move || recover_lock(log.lock()).push(name)) - } - - fn fired(&self) -> Vec<&'static str> { - recover_lock(self.0.lock()).clone() - } -} - -/// Runs the wheel forward to `target`, one tick at a time. -fn run_to(wheel: &mut TimingWheel, target: u64) { - while wheel.now() < target { - wheel.tick(); - } -} - -#[test] -fn timing_wheel_fires_a_timer_at_its_deadline() { - let log = FireLog::default(); - let mut wheel = TimingWheel::new(1_000); - wheel.schedule(1_005, log.timer("timer")); - - run_to(&mut wheel, 1_004); - assert!(log.fired().is_empty()); - run_to(&mut wheel, 1_005); - assert_eq!(log.fired(), ["timer"]); -} - -#[test] -fn timing_wheel_fires_each_timer_once() { - let log = FireLog::default(); - let mut wheel = TimingWheel::new(1_000); - let timer = log.timer("timer"); - // The same timer placed twice: the first placement to drain is not the last - // reference, so only the later one fires it. - wheel.schedule(1_005, timer.clone()); - wheel.schedule(1_020, timer); - - run_to(&mut wheel, 1_005); - assert!(log.fired().is_empty()); - run_to(&mut wheel, 1_020); - assert_eq!(log.fired(), ["timer"]); -} - -#[test] -fn timing_wheel_firing_early_makes_the_scheduled_placement_inert() { - let log = FireLog::default(); - let mut wheel = TimingWheel::new(1_000); - let timer = log.timer("timer"); - wheel.schedule(1_005, timer.clone()); - - timer.fire(); - assert_eq!(log.fired(), ["timer"]); - // The deadline arriving drops the placement, which must not fire it again. - run_to(&mut wheel, 1_005); - assert_eq!(log.fired(), ["timer"]); -} - -#[test] -fn timing_wheel_drop_fires_everything_it_holds() { - let log = FireLog::default(); - let mut wheel = TimingWheel::new(1_000); - wheel.schedule(1_005, log.timer("early")); - wheel.schedule(9_999_999, log.timer("late")); - - drop(wheel); - let mut fired = log.fired(); - fired.sort_unstable(); - assert_eq!(fired, ["early", "late"]); -} - -#[test] -fn timing_wheel_cascades_a_timer_down_two_levels() { - // A deadline two levels up has to reach level 0 before it can be drained. - let now = 1_000; - let deadline = now + (1 << (6 * 2)); - let log = FireLog::default(); - let mut wheel = TimingWheel::new(now); - wheel.schedule(deadline, log.timer("timer")); - - run_to(&mut wheel, deadline - 1); - assert!(log.fired().is_empty()); - run_to(&mut wheel, deadline); - assert_eq!(log.fired(), ["timer"]); -} - -#[test] -fn timing_wheel_fires_a_boundary_deadline_without_an_extra_tick() { - let log = FireLog::default(); - let mut wheel = TimingWheel::new(700); - wheel.schedule(1_024, log.timer("timer")); - - run_to(&mut wheel, 1_024); - assert_eq!(log.fired(), ["timer"]); -} - -#[test] -fn timing_wheel_fires_a_deadline_already_in_the_past() { - let log = FireLog::default(); - let mut wheel = TimingWheel::new(1_000); - // Scheduling into the past drops the placement immediately. - wheel.schedule(999, log.timer("timer")); - assert_eq!(log.fired(), ["timer"]); -} - #[test] fn replay_pruning_removes_only_records_outside_the_retention_window() { let now = 10_000; diff --git a/src/common/auth/timing_wheel.rs b/src/common/auth/timing_wheel.rs index 8c24c17..d219645 100644 --- a/src/common/auth/timing_wheel.rs +++ b/src/common/auth/timing_wheel.rs @@ -1,168 +1,439 @@ -//! Hierarchical timer wheel. It schedules opaque timers and knows nothing about -//! what they mean. +//! Hierarchical timer wheel: rotating bucket queues indexed by relative delay. +//! +//! Levels are digit positions in base `radix`, and how many exist is derived from +//! the longest delay the wheel must support. Scheduling decomposes the delay into +//! those digits and builds one nested link per digit, coarsest outermost: //! //! ```text -//! schedule(deadline, timer) -> the bucket covering that deadline +//! radix = 64, schedule(delay = 1*64² + 5*64 + 3) +//! +//! level 2 [ ][A][ ]… A pops after 1 rotation of 64² ticks; dropping it +//! level 1 [ ]…[B][ ]… files B, which pops 5 rotations of 64 ticks later +//! level 0 [ ][ ][C]… and files C, which pops 3 ticks later and fires +//! ``` //! -//! tick() -> now += 1 -//! -> drain level 0's slot for `now`; every 64th tick level 1's, and so on -//! -> an entry whose deadline has arrived is dropped; the rest are re-filed -//! into the finer level that now covers them +//! `A` holds `B` holds `C` holds the timer, so the chain *is* the route: no per +//! timer list of future placements, and nothing to look up or recompute. A bucket +//! is just `Vec`, and dropping a link is what files the next one. +//! +//! ```text +//! tick() -> ticks += 1 +//! -> level 0 always rotates; level i rotates when ticks % radix^i == 0 +//! -> rotate = pop_front, push_back an empty bucket; dropping the popped +//! bucket files each link's successor, or fires the timer if the link +//! was the innermost //! ``` //! -//! The wheel holds the only strong references to its timers, so a timer runs when -//! the last entry referring to it is dropped. That is what keeps the wheel -//! indifferent to its users: it never looks a timer up, compares identities, or -//! has to be told that one was superseded. +//! So a tick moves one bucket per level that turns over and performs no +//! arithmetic per entry: the queues rotate, which keeps a bucket's index equal to +//! its distance from now. //! -//! Rescheduling exploits that directly. Holding a `Weak`, a caller inserts -//! the same timer again at a later deadline; the earlier entry still drains on its -//! own schedule, but dropping it no longer brings the count to zero, so the timer -//! waits for the last entry to go. Cancelling is the mirror image: [`Timer::fire`] -//! runs the callback early and leaves the remaining entries inert. +//! Only the wheel holds strong references to a timer, so it runs when the last +//! chain holding it is dropped. The wheel never looks a timer up, compares +//! identities, or has to be told one was superseded: to move a deadline, schedule +//! the same timer again — the earlier chain still drains, but it is no longer the +//! last reference, so it fires nothing. use super::*; -/// Bits of a deadline that one level's slot field covers, so a level holds -/// `1 << 6` slots and each is 64 times coarser than the level below it. -const SLOT_BITS: u32 = 6; -const SLOTS: usize = 1 << SLOT_BITS; -const SLOT_MASK: u64 = SLOTS as u64 - 1; -/// Six 64-slot levels span `64^6` seconds, so any deadline a caller can ask for -/// lands in a level whose range really contains it. -const NUM_LEVELS: usize = 6; -const TOP_LEVEL: Level = NUM_LEVELS as Level - 1; - -/// Which level of the hierarchy a bucket belongs to: `0..NUM_LEVELS`. -type Level = u8; - -/// Which bucket within one level: `0..SLOTS`. -type Slot = u8; - -/// A callback that runs once: when its deadline arrives, or when it is cancelled, +/// A callback that runs once: when its delay elapses, or when it is cancelled, /// whichever comes first. pub(super) struct Timer { - /// `None` once the callback has run, so any remaining wheel entries for this - /// timer are inert and a cancelled timer cannot fire twice. - callback: std::sync::Mutex>>, + /// `None` once the callback has run, so any route still holding this timer is + /// inert and a cancelled timer cannot fire twice. + /// + /// WHY a lock for state a single task owns: running a `FnOnce` moves it out, + /// which needs `&mut`, but a timer is reached through a shared handle so that + /// two routes can hold one. `Arc: Send` — which `tokio::spawn` requires of + /// the actor this runs in — implies `T: Sync`, and shared mutability that is + /// `Sync` needs a lock; `Cell` would be cheaper but is not `Sync`. It is never + /// contended, and the path that fires almost every timer skips it: `Drop` has + /// `&mut self`, so it reaches the callback directly. + callback: Mutex>>, } impl Timer { pub(super) fn new(callback: impl FnOnce() + Send + 'static) -> Arc { Arc::new(Self { - callback: std::sync::Mutex::new(Some(Box::new(callback))), + callback: Mutex::new(Some(Box::new(callback))), }) } - /// Runs the callback unless it has run already. Called by the wheel when a - /// deadline arrives, and by a caller cancelling ahead of that. + /// Runs the callback unless it has run already, for a caller cancelling ahead + /// of the deadline. pub(super) fn fire(&self) { - let callback = recover_lock(self.callback.lock()).take(); - if let Some(callback) = callback { - callback(); - } + let callback = self.callback.lock().take(); + run(callback); } } impl Drop for Timer { /// Releasing the last reference is what fires a timer, so dropping the wheel - /// tears down everything it was holding. + /// runs everything it was holding. Owning `&mut self` here is what lets the + /// usual path take the callback without locking. fn drop(&mut self) { - self.fire(); + let callback = self.callback.get_mut().take(); + run(callback); } } -/// One placement of a timer. The deadline lives here rather than in the `Timer`, -/// so a timer rescheduled later leaves its earlier placements draining harmlessly -/// instead of dragging them forward. -struct Entry { - deadline: u64, - /// Never read: holding the reference *is* the entry's job, and releasing it - /// is what can fire the timer. - #[allow(dead_code)] - timer: Arc, +fn run(callback: Option>) { + if let Some(callback) = callback { + callback(); + } +} + +/// One leg of a timer's route through the levels. +/// +/// A delay spanning several digits cannot be filed in one bucket, so the route is +/// a chain: each [`Link::Relay`] waits in one bucket and, once that bucket comes +/// off the front, hands the leg nested inside it to the wheel, which files it in +/// the next, finer bucket. Only the outermost leg is ever in a bucket, and only +/// [`Link::Deliver`] holds the timer, so the chain unwinding one bucket at a time +/// *is* the timer descending the levels. That is what leaves a tick with nothing +/// to compute. +/// +/// ```text +/// delay = 1*64² + 5*64 + 3 +/// Relay{L2,slot 1} -> Relay{L1,slot 5} -> Relay{L0,slot 3} -> Deliver(timer) +/// ^ filed now ^ filed when the ^ …and so on ^ dropping this +/// one before it fires the timer +/// comes off +/// ``` +enum Link { + /// The end of a route. Never read: holding the reference *is* this leg's job, + /// and releasing it is what fires the timer. + Deliver(#[allow(dead_code)] Arc), + /// Files `next` into `level`'s `slot` when this leg comes off the front. + /// + /// `Box`, not `Arc`: exactly one bucket owns a route at a time, so a leg needs + /// no reference count of its own — only the timer at the end is shared. + Relay { + level: u8, + slot: u16, + next: Box, + }, } +/// A bucket's worth of routes. Dropping one without draining it releases the +/// timers at the end of every route it holds, which is how dropping the wheel +/// fires everything. +type Bucket = Vec; + pub(super) struct TimingWheel { - now: u64, - levels: [Vec>; NUM_LEVELS], + /// Ticks elapsed since construction. Buckets are indexed relative to it, so + /// advancing re-indexes nothing. + ticks: u64, + radix: u64, + levels: Vec>, } impl TimingWheel { - pub(super) fn new(now: u64) -> Self { + /// Builds the smallest wheel that can place `max_delay` ticks, adding a level + /// at a time until the levels together span it. + pub(super) fn new(max_delay: u64, radix: u64) -> Self { + assert!(radix > 1, "a level needs at least two buckets"); + let mut levels = 1_usize; + let mut span = radix; + while span < max_delay { + levels += 1; + span = span.saturating_mul(radix); + } Self { - now, - levels: std::array::from_fn(|_| std::iter::repeat_with(Vec::new).take(SLOTS).collect()), + ticks: 0, + radix, + levels: (0..levels) + .map(|_| { + std::iter::repeat_with(Bucket::new) + .take(radix as usize) + .collect() + }) + .collect(), } } - pub(super) fn now(&self) -> u64 { - self.now + /// Longest delay this wheel can place exactly. + pub(super) fn max_delay(&self) -> u64 { + self.period(self.levels.len()) } - /// Holds `timer` until `deadline`. Scheduling a timer the wheel already holds - /// adds a placement rather than replacing one, which is how a caller moves a - /// deadline outward without the wheel having to find the old entry. - pub(super) fn schedule(&mut self, deadline: u64, timer: Arc) { - self.place(Entry { deadline, timer }); + /// Holds `timer` for `delay` ticks. A delay of zero, or one past + /// [`Self::max_delay`], releases the timer at once rather than misplacing it. + /// + /// Scheduling a timer the wheel already holds builds a second route rather + /// than replacing the first, which is how a caller moves a deadline without + /// the wheel having to find the old one. + pub(super) fn schedule(&mut self, delay: u64, timer: Arc) { + if delay == 0 || delay > self.max_delay() { + // Dropping `timer` here fires it if this was the last reference. + return; + } + let deliver = Link::Deliver(timer); + // The coarsest reachable level absorbs however far the current tick sits + // into its rotation, so its bucket comes off on a rotation boundary. Every + // finer level is at zero offset there, which makes the delay still + // remaining a plain base-`radix` decomposition from that point down. + let (level, slot, remaining) = (0..self.levels.len()) + .rev() + .find_map(|level| self.entry_leg(level, self.ticks + delay)) + .expect("a delay within max_delay reaches some level"); + let route = match remaining { + 0 => deliver, + remaining => self.route(remaining, deliver), + }; + self.file(level, slot, route); } - /// Advances one second and drains whatever that turnover exposes. + /// Advances one tick, rotating every level that turns over. pub(super) fn tick(&mut self) { - self.now += 1; - // Coarse to fine, so a timer cascading several levels down still reaches - // level 0 in time to be drained by this same tick. - for level in (1..=TOP_LEVEL).rev() { - // A level turns over once every `slot_range` seconds, exactly when - // `now` has no bits left below that level's slot field. - if self.now & (slot_range(level) - 1) != 0 { - continue; + self.ticks += 1; + for level in 0..self.levels.len() { + // A level turns over every `radix^level` ticks. Once one does not, no + // coarser one can either, since its period divides theirs. + if self.ticks % self.period(level) != 0 { + break; + } + let bucket = self.rotate(level); + for link in bucket { + match link { + // Out of legs: dropping it releases the timer, firing it if + // this was the last route holding it. + Link::Deliver(timer) => drop(timer), + Link::Relay { level, slot, next } => { + self.file(level as usize, slot as usize, *next) + } + } } - let entries = self.take_bucket(level, self.now); - self.refile(entries); } - let entries = self.take_bucket(0, self.now); - self.refile(entries); } - fn refile(&mut self, entries: Vec) { - for entry in entries { - self.place(entry); + /// Takes `level`'s front bucket off and puts an empty one on the back, so + /// bucket indices stay relative to the current tick. + fn rotate(&mut self, level: usize) -> Bucket { + let queue = &mut self.levels[level]; + let bucket = queue.pop_front().unwrap_or_default(); + queue.push_back(Bucket::new()); + bucket + } + + fn file(&mut self, level: usize, slot: usize, link: Link) { + if let Some(bucket) = self.levels[level].get_mut(slot) { + bucket.push(link); } } - fn place(&mut self, entry: Entry) { - // An arrived deadline means this placement is done: returning drops the - // entry, which fires the timer if this was its last reference. - if entry.deadline <= self.now { - return; + /// The route to file for a timer due `remaining` ticks after a rotation + /// boundary, built by recursing into the finer levels so each leg owns the + /// part of the route it hands on. `remaining` must be non-zero. + fn route(&self, remaining: u64, inner: Link) -> Link { + let (level, slot, rest) = self.next_leg(remaining); + let next = match rest { + 0 => inner, + rest => self.route(rest, inner), + }; + Link::Relay { + level: level as u8, + slot: slot as u16, + next: Box::new(next), } - let level = level_for(self.now, entry.deadline); - let slot = slot_for(level, entry.deadline); - self.levels[level as usize][slot as usize].push(entry); } - /// Empties the bucket that `when` falls in at `level`. - fn take_bucket(&mut self, level: Level, when: u64) -> Vec { - let slot = slot_for(level, when); - std::mem::take(&mut self.levels[level as usize][slot as usize]) + /// The route's first leg if it starts at `level`: which bucket holds a timer + /// due at tick `target`, and how much delay that leaves for the legs after it. + /// `None` when this level's next rotation already overshoots `target`, or when + /// `target` is more than one revolution away. + fn entry_leg(&self, level: usize, target: u64) -> Option<(usize, usize, u64)> { + let period = self.period(level); + let next_rotation = self.ticks - self.ticks % period + period; + let ahead = target.checked_sub(next_rotation)?; + let slot = ahead / period; + (slot < self.radix).then_some((level, slot as usize, ahead % period)) } -} -/// Seconds covered by one of `level`'s slots. -fn slot_range(level: Level) -> u64 { - 1 << (SLOT_BITS * level as u32) -} + /// The next leg for a timer due `remaining` ticks after a rotation boundary: + /// the coarsest level whose rotation still fits. From a boundary that level's + /// front bucket comes off one period out, so bucket `j` comes off after + /// `j + 1` of them. + fn next_leg(&self, remaining: u64) -> (usize, usize, u64) { + let level = (0..self.levels.len()) + .rev() + .find(|level| self.period(*level) <= remaining) + .unwrap_or(0); + let period = self.period(level); + (level, (remaining / period - 1) as usize, remaining % period) + } -fn slot_for(level: Level, when: u64) -> Slot { - ((when >> (SLOT_BITS * level as u32)) & SLOT_MASK) as Slot + /// Ticks spanned by `level` and every level below it: `radix^level`. + fn period(&self, level: usize) -> u64 { + self.radix.saturating_pow(level as u32) + } } -/// Finest level able to hold `deadline`: the one whose slot field covers the -/// highest bit in which `now` and `deadline` differ. A deadline past the top -/// level is clamped into it and cannot fire early, because an entry is dropped -/// only once its own deadline has arrived. -fn level_for(now: u64, deadline: u64) -> Level { - let significant = 63 - ((now ^ deadline) | SLOT_MASK).leading_zeros(); - ((significant / SLOT_BITS) as Level).min(TOP_LEVEL) +#[cfg(test)] +mod tests { + use super::*; + + /// A wheel plus the tick each scheduled timer actually fired on. + struct Harness { + wheel: TimingWheel, + fired: Arc>>, + ticks: Arc, + } + + impl Harness { + fn new(max_delay: u64, radix: u64) -> Self { + Self { + wheel: TimingWheel::new(max_delay, radix), + fired: Arc::new(Mutex::new(Vec::new())), + ticks: Arc::new(AtomicU64::new(0)), + } + } + + fn timer(&self, id: u32) -> Arc { + let fired = self.fired.clone(); + let ticks = self.ticks.clone(); + Timer::new(move || { + fired.lock().push((id, ticks.load(Ordering::Acquire))); + }) + } + + fn schedule(&mut self, id: u32, delay: u64) { + let timer = self.timer(id); + self.wheel.schedule(delay, timer); + } + + fn tick(&mut self) { + self.ticks.fetch_add(1, Ordering::AcqRel); + self.wheel.tick(); + } + + fn fired_at(&self, id: u32) -> Option { + self.fired + .lock() + .iter() + .find(|(fired, _)| *fired == id) + .map(|(_, at)| *at) + } + } + + /// Every delay, from every starting offset, must fire on exactly the tick it + /// asked for. This is the wheel's whole contract, and a radix decomposition is + /// easy to get wrong by one bucket, so it is checked exhaustively rather than + /// sampled. + #[test] + fn every_delay_fires_on_its_exact_tick() { + let radix = 4; + let max_delay = radix * radix * radix; + for offset in 0..2 * radix * radix { + let mut harness = Harness::new(max_delay, radix); + for _ in 0..offset { + harness.tick(); + } + for delay in 1..=max_delay { + harness.schedule(delay as u32, delay); + } + for _ in 0..max_delay { + harness.tick(); + } + for delay in 1..=max_delay { + assert_eq!( + harness.fired_at(delay as u32), + Some(offset + delay), + "radix {radix}, offset {offset}, delay {delay}" + ); + } + } + } + + /// The same contract at the shape the wheel actually runs with. + #[test] + fn every_short_delay_fires_on_its_exact_tick_at_radix_64() { + let radix = 64; + let mut harness = Harness::new(radix * radix * radix * radix, radix); + for _ in 0..100 { + harness.tick(); + } + let delays = (1..=200).chain([radix - 1, radix, radix + 1, radix * radix, 4095, 4096]); + for delay in delays.clone() { + harness.schedule(delay as u32, delay); + } + for _ in 0..5000 { + harness.tick(); + } + for delay in delays { + assert_eq!( + harness.fired_at(delay as u32), + Some(100 + delay), + "delay {delay}" + ); + } + } + + #[test] + fn level_count_covers_the_requested_delay() { + assert_eq!(TimingWheel::new(64, 64).max_delay(), 64); + assert_eq!(TimingWheel::new(65, 64).max_delay(), 4096); + assert_eq!(TimingWheel::new(4096, 64).max_delay(), 4096); + assert_eq!(TimingWheel::new(4097, 64).max_delay(), 262_144); + } + + #[test] + fn a_delay_the_wheel_cannot_place_fires_at_once() { + let mut harness = Harness::new(64, 64); + harness.schedule(1, 65); + assert_eq!(harness.fired_at(1), Some(0)); + harness.schedule(2, 0); + assert_eq!(harness.fired_at(2), Some(0)); + } + + #[test] + fn rescheduling_the_same_timer_defers_it_to_the_later_route() { + let mut harness = Harness::new(4096, 64); + let timer = harness.timer(1); + harness.wheel.schedule(5, timer.clone()); + harness.wheel.schedule(20, timer); + + for _ in 0..5 { + harness.tick(); + } + assert_eq!( + harness.fired_at(1), + None, + "the earlier route must not fire the timer" + ); + for _ in 5..20 { + harness.tick(); + } + assert_eq!(harness.fired_at(1), Some(20)); + } + + #[test] + fn firing_early_makes_the_scheduled_route_inert() { + let mut harness = Harness::new(4096, 64); + let timer = harness.timer(1); + harness.wheel.schedule(10, timer.clone()); + + timer.fire(); + assert_eq!(harness.fired_at(1), Some(0)); + for _ in 0..10 { + harness.tick(); + } + assert_eq!(harness.fired.lock().len(), 1); + } + + #[test] + fn dropping_the_wheel_fires_everything_it_holds() { + let mut harness = Harness::new(262_144, 64); + harness.schedule(1, 5); + harness.schedule(2, 200_000); + + let fired = harness.fired.clone(); + drop(harness); + let ids = fired + .lock() + .iter() + .map(|(id, _)| *id) + .collect::>(); + assert_eq!(ids, HashSet::from([1, 2])); + } } diff --git a/src/common/checksum.rs b/src/common/checksum.rs index 91576d2..6efe4b2 100644 --- a/src/common/checksum.rs +++ b/src/common/checksum.rs @@ -4,7 +4,9 @@ use std::os::unix::fs::PermissionsExt; use std::path::Path; use std::process::Command; use std::sync::atomic::{AtomicU32, Ordering}; -use std::sync::{LazyLock, RwLock}; +use std::sync::LazyLock; + +use parking_lot::RwLock; use base64::engine::general_purpose::URL_SAFE_NO_PAD; use base64::Engine; @@ -92,15 +94,9 @@ fn update_runtime_credential(credential: Option) { .as_ref() .map(|credential| gen_checksum_by_key(credential.key())) .unwrap_or_default(); - let mut guard = MSG_HEADER_KEY_STATE - .credential - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()); + let mut guard = MSG_HEADER_KEY_STATE.credential.write(); *guard = credential; - *MSG_HEADER_KEY_STATE - .load_error - .write() - .unwrap_or_else(|poisoned| poisoned.into_inner()) = None; + *MSG_HEADER_KEY_STATE.load_error.write() = None; MSG_HEADER_KEY_STATE.hash.store(hash, Ordering::Release); } @@ -129,23 +125,12 @@ static MSG_HEADER_KEY_STATE: LazyLock = LazyLock::new(|| { /// Return the configured process credential, failing closed when none exists. pub fn get_process_credential() -> Result { - if let Some(error) = MSG_HEADER_KEY_STATE - .load_error - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .clone() - { + if let Some(error) = MSG_HEADER_KEY_STATE.load_error.read().clone() { return Err(error); } - MSG_HEADER_KEY_STATE - .credential - .read() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .ok_or_else(|| { - format!( - "`{ENV_MSG_HEADER_KEY}` is required; no insecure default credential is available" - ) - }) + MSG_HEADER_KEY_STATE.credential.read().ok_or_else(|| { + format!("`{ENV_MSG_HEADER_KEY}` is required; no insecure default credential is available") + }) } /// Get current message header key bytes. diff --git a/src/common/message/forward.rs b/src/common/message/forward.rs index 3ae3084..e45699d 100644 --- a/src/common/message/forward.rs +++ b/src/common/message/forward.rs @@ -829,7 +829,9 @@ macro_rules! start_datagram_forward_with_codec_key { mod tests { use std::collections::VecDeque; use std::io; - use std::sync::{Arc, Mutex}; + use std::sync::Arc; + + use parking_lot::Mutex; use std::time::Duration; use super::*; @@ -891,22 +893,22 @@ mod tests { impl ScriptedWriter { fn chunks(&self) -> Vec> { - self.state.lock().unwrap().chunks.clone() + self.state.lock().chunks.clone() } fn shutdowns(&self) -> usize { - self.state.lock().unwrap().shutdowns + self.state.lock().shutdowns } } impl ForwardWriter for ScriptedWriter { async fn write(&mut self, src: &[u8]) -> Result<()> { - self.state.lock().unwrap().chunks.push(src.to_vec()); + self.state.lock().chunks.push(src.to_vec()); Ok(()) } async fn shutdown(&mut self) { - self.state.lock().unwrap().shutdowns += 1; + self.state.lock().shutdowns += 1; } } @@ -954,7 +956,7 @@ mod tests { } fn chunks(&self) -> Vec> { - self.state.lock().unwrap().chunks.clone() + self.state.lock().chunks.clone() } } @@ -962,12 +964,12 @@ mod tests { async fn write(&mut self, src: &[u8]) -> Result<()> { self.write_started.notify_one(); tokio::time::sleep(self.delay).await; - self.state.lock().unwrap().chunks.push(src.to_vec()); + self.state.lock().chunks.push(src.to_vec()); Ok(()) } async fn shutdown(&mut self) { - self.state.lock().unwrap().shutdowns += 1; + self.state.lock().shutdowns += 1; } } diff --git a/src/common/message/secure.rs b/src/common/message/secure.rs index fcef449..6038a88 100644 --- a/src/common/message/secure.rs +++ b/src/common/message/secure.rs @@ -17,7 +17,9 @@ //! This root module coordinates client/server sessions. Frame mechanics, replay admission, //! log suppression, and protocol tests are isolated in focused child modules. -use std::sync::{Arc, Mutex}; +use std::sync::Arc; + +use parking_lot::Mutex; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use rand::RngExt; @@ -386,7 +388,6 @@ impl ServerSecurity { ) -> FailureLogDecision { self.failure_logs .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) .record(peer_ip, key_id, reason, unix_seconds()) } diff --git a/src/common/message/secure/first_flight.rs b/src/common/message/secure/first_flight.rs index 2de259c..fb9183c 100644 --- a/src/common/message/secure/first_flight.rs +++ b/src/common/message/secure/first_flight.rs @@ -34,14 +34,12 @@ fn reserved_error_session( #[allow(clippy::result_large_err)] pub(super) fn evaluate_first_flight( auth: &AuthRuntime, - replay: &std::sync::Mutex, + replay: &parking_lot::Mutex, key_id: KeyId, fingerprint: [u8; 32], work: FirstFlightWork, ) -> std::result::Result<(Vec, AuthContext), ServerInitialError> { - let mut replay = replay - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); + let mut replay = replay.lock(); match work { FirstFlightWork::Live { key, diff --git a/ui/native/pb_mapper_ffi/Cargo.toml b/ui/native/pb_mapper_ffi/Cargo.toml index 2b30fdc..5cc391b 100644 --- a/ui/native/pb_mapper_ffi/Cargo.toml +++ b/ui/native/pb_mapper_ffi/Cargo.toml @@ -23,3 +23,4 @@ dirs = "5.0" pb-mapper = { path = "../../../" } uni-stream.workspace = true +parking_lot.workspace = true diff --git a/ui/native/pb_mapper_ffi/src/state.rs b/ui/native/pb_mapper_ffi/src/state.rs index dd1bed2..08bab5a 100644 --- a/ui/native/pb_mapper_ffi/src/state.rs +++ b/ui/native/pb_mapper_ffi/src/state.rs @@ -14,9 +14,10 @@ use std::fs; use std::net::{IpAddr, Ipv4Addr, SocketAddr}; use std::path::PathBuf; use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::{Arc, Mutex as StdMutex}; +use std::sync::Arc; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; +use parking_lot::Mutex as SyncMutex; use serde::{Deserialize, Serialize}; use tokio::net::{TcpListener, TcpStream}; use tokio::sync::{Mutex, RwLock}; @@ -388,30 +389,26 @@ struct ConnectionInfo { /// the first's [`JoinHandle`], leaving a tunnel running that nothing could /// abort. Holding the key across the gap is what closes it again. /// -/// The set is behind a `std::sync::Mutex` rather than tokio's so that `Drop` can +/// The set is behind a blocking mutex rather than tokio's so that `Drop` can /// release it; it is only ever held for a set insert or remove. struct KeyClaim { key: String, - claims: Arc>>, + claims: Arc>>, } impl Drop for KeyClaim { fn drop(&mut self) { - if let Ok(mut claims) = self.claims.lock() { - claims.remove(&self.key); - } + self.claims.lock().remove(&self.key); } } /// Claims `key`, or reports that someone else is already setting it up. fn claim_key( - claims: &Arc>>, + claims: &Arc>>, key: &str, what: &str, ) -> Result { - let mut guard = claims - .lock() - .map_err(|_| CtlError::internal(format!("{what} state for '{key}' is poisoned")))?; + let mut guard = claims.lock(); if !guard.insert(key.to_string()) { return Err(CtlError::already_in_progress(format!( "'{key}' is already {what}" @@ -496,8 +493,8 @@ pub struct PbMapperState { client_status_refreshing: Arc>>, /// Keys currently being set up. See [`KeyClaim`]. Separate sets because a /// key can legitimately be registered and connected to at the same time. - registering: Arc>>, - connecting: Arc>>, + registering: Arc>>, + connecting: Arc>>, } mod configuration; diff --git a/ui/native/pb_mapper_ffi/src/state/configuration.rs b/ui/native/pb_mapper_ffi/src/state/configuration.rs index ced14e7..feb4bd5 100644 --- a/ui/native/pb_mapper_ffi/src/state/configuration.rs +++ b/ui/native/pb_mapper_ffi/src/state/configuration.rs @@ -65,8 +65,8 @@ impl PbMapperState { client_status_cache: Arc::new(RwLock::new(HashMap::new())), service_status_refreshing: Arc::new(RwLock::new(HashSet::new())), client_status_refreshing: Arc::new(RwLock::new(HashSet::new())), - registering: Arc::new(StdMutex::new(HashSet::new())), - connecting: Arc::new(StdMutex::new(HashSet::new())), + registering: Arc::new(SyncMutex::new(HashSet::new())), + connecting: Arc::new(SyncMutex::new(HashSet::new())), }; state.config = state.load_config().unwrap_or_else(|e| { tracing::warn!("Could not load config: {}, using defaults", e); From 06542ba3224f93310ed5f677fa491a3d5e464e1e Mon Sep 17 00:00:00 2001 From: LB7666 Date: Fri, 21 Aug 2026 12:56:53 +0800 Subject: [PATCH 63/74] Teach the CI path filter about crates/ MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The `case` that decides whether to run the Rust job matched `src/*` and a root-level `Cargo.toml`, and nothing else. Once sources move under `crates/`, no path would match: every subsequent commit would report green without having been compiled or tested. Add `crates/*`, plus a `*/Cargo.toml` arm for per-crate manifests — the existing unanchored entry only ever catches the root one. `ui/native/*` still precedes it, so the FFI crate keeps landing on the Rust side. Co-Authored-By: Claude Fable 5 --- .github/workflows/syntax-check.yml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.github/workflows/syntax-check.yml b/.github/workflows/syntax-check.yml index 03e5e8d..6be4f0d 100644 --- a/.github/workflows/syntax-check.yml +++ b/.github/workflows/syntax-check.yml @@ -62,8 +62,11 @@ jobs: # belongs to the Rust side and must be matched before ui/*. ui/native/*) rust=true ;; ui/*) flutter=true ;; - src/*|tests/*|examples/*) rust=true ;; + crates/*|src/*|tests/*|examples/*) rust=true ;; Cargo.toml|Cargo.lock|rust-toolchain.toml|rustfmt.toml) rust=true ;; + # Per-crate manifests. `case` patterns match the whole path, so + # the unanchored entry above only ever catches the root manifest. + */Cargo.toml|*/Cargo.lock) rust=true ;; # A change to this workflow has to prove itself on both. .github/workflows/syntax-check.yml) rust=true; flutter=true ;; esac From c31eb3b428ed195426bab4c5a686d822f16919c7 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Fri, 21 Aug 2026 12:59:30 +0800 Subject: [PATCH 64/74] Drop dead code left over from earlier refactors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit None of this has a consumer, and carrying it into the new crates would just relocate the confusion: - Five `common::error::Error` variants (`Stm*`, `Lsn*`) that nothing has constructed since streams and listeners moved out to `uni-stream`. Verified by searching for both the variant and its generated snafu context selector. - `once_cell`, declared twice and imported nowhere; the code already uses `std::sync::LazyLock`. - Three `#[macro_export]` macros in `message/forward.rs` (`create_component` and the two `start_*_with_codec_key`) with no callers. This also retires the two `$crate::common::message::` paths inside them, which the split would otherwise have had to rewrite. - `Aes256GcmCodec::try_new_with_default_key`, uncalled, and the only reason `utils::codec` reached back into `common::checksum` — so `utils` no longer depends on `common` at all. - Profiles `wasm-dev`, `server-dev`, `android-dev` and the empty `[profile]` table. Nothing references them, the latter two are no-op `inherits = "dev"`, and there is no wasm target here. Co-Authored-By: Claude Fable 5 --- Cargo.lock | 1 - Cargo.toml | 14 ---- src/common/error.rs | 24 ------ src/common/message/forward.rs | 146 ---------------------------------- src/utils/codec.rs | 7 -- 5 files changed, 192 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 4c3b7f0..d0e664c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -787,7 +787,6 @@ dependencies = [ "futures", "hashbrown 0.16.1", "kanal", - "once_cell", "parking_lot", "rand 0.10.0", "ring", diff --git a/Cargo.toml b/Cargo.toml index 79bce79..fbfc215 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,7 +23,6 @@ better_mimalloc_rs.workspace = true bytes.workspace = true trust-dns-resolver.workspace = true ring.workspace = true -once_cell.workspace = true uni-stream.workspace = true kanal.workspace = true base64.workspace = true @@ -66,21 +65,8 @@ better_mimalloc_rs = { version = "0.1.2", features = ["config"] } bytes = "1.11" trust-dns-resolver = { version = "0.23.2" } ring = "0.17.14" -once_cell = "1.20.2" base64 = "0.22.1" subtle = "2.6.1" parking_lot = "0.12" uni-stream = { git = "https://github.com/acking-you/uni-stream.git", branch = "master" } kanal = { git = "https://github.com/acking-you/kanal.git", branch = "dev/pb-mapper" } - -[profile] - -[profile.wasm-dev] -inherits = "dev" -opt-level = 1 - -[profile.server-dev] -inherits = "dev" - -[profile.android-dev] -inherits = "dev" diff --git a/src/common/error.rs b/src/common/error.rs index 9717721..cd5609e 100644 --- a/src/common/error.rs +++ b/src/common/error.rs @@ -73,30 +73,6 @@ pub enum Error { /// Error for forward #[snafu(display("failed to forward message to write in normal text"))] FwdNetworkWriteWithNormal { source: std::io::Error }, - /// Error for stream - #[snafu(display("failed to connect stream, type:`{stream_type}`"))] - StmConnectStream { - // must be "UDP" or "TCP" - stream_type: &'static str, - source: std::io::Error, - }, - #[snafu(display("failed to got one addr from iter"))] - StmGotOneAddrFromIter, - #[snafu(display("failed to got one addr when parsing address"))] - StmGotOneAddr { source: std::io::Error }, - /// Error for listener - #[snafu(display("listener failed to bind addr, type:`{listener_type}`"))] - LsnListenerBind { - // must be "UDP" or "TCP" - listener_type: &'static str, - source: std::io::Error, - }, - #[snafu(display("listener failed to accept stream, type:`{listener_type}`"))] - LsnListenerAccept { - // must be "UDP" or "TCP" - listener_type: &'static str, - source: std::io::Error, - }, /// Error for config #[snafu(display("parse socket address from string:`{string}` error"))] CfgParseSockAddr { diff --git a/src/common/message/forward.rs b/src/common/message/forward.rs index e45699d..7207ee1 100644 --- a/src/common/message/forward.rs +++ b/src/common/message/forward.rs @@ -679,152 +679,6 @@ impl StreamForward for UdpStreamImpl { } } -#[macro_export] -macro_rules! create_component { - (Reader, $stream:expr,true, $key:expr, $get_codec:ident, $name:expr) => { - CodecForwardReader::new( - $stream, - snafu_error_get_or_return_ok!( - $get_codec(&$key), - concat!("failed to create decoder when `", $name, "` forward msg") - ), - ) - }; - (Reader, $stream:expr,false, $key:expr, $get_codec:ident, $name:expr) => { - NormalForwardReader::new($stream) - }; - (Writer, $stream:expr,true, $key:expr, $get_codec:ident, $name:expr) => { - CodecForwardWriter::new( - $stream, - snafu_error_get_or_return_ok!( - $get_codec(&$key), - concat!("failed to create encoder when `", $name, "` forward msg") - ), - ) - }; - (Writer, $stream:expr,false, $key:expr, $get_codec:ident, $name:expr) => { - NormalForwardWriter::new($stream) - }; -} - -/// When using it, please remember to manually import the following symbols: -/// - [`start_forward`] -/// - [`crate::create_component`] -/// - [`ForwardReader`] -/// - [`ForwardWriter`] -/// - [`CodecForwardReader`] -/// - [`CodecForwardWriter`] -/// - [`crate::snafu_error_get_or_return_ok`] -/// - [`super::get_decodec`] -/// - [`super::get_encodec`] -#[macro_export] -macro_rules! start_forward_with_codec_key { - ( - $codec_key:expr, - $client_reader:expr, - $client_writer:expr, - $server_reader:expr, - $server_writer:expr, - $client_reader_codec:tt, - $client_writer_codec:tt, - $server_reader_codec:tt, - $server_writer_codec:tt - ) => { - match $codec_key { - Some(key) => { - (start_forward( - create_component!( - Reader, - $client_reader, - $client_reader_codec, - key, - get_decodec, - "client_reader" - ), - create_component!( - Writer, - $client_writer, - $client_writer_codec, - key, - get_encodec, - "client_writer" - ), - create_component!( - Reader, - $server_reader, - $server_reader_codec, - key, - get_decodec, - "server_reader" - ), - create_component!( - Writer, - $server_writer, - $server_writer_codec, - key, - get_encodec, - "server_writer" - ), - ) - .await) - } - None => { - (start_forward( - NormalForwardReader::new($client_reader), - NormalForwardWriter::new($client_writer), - NormalForwardReader::new($server_reader), - NormalForwardWriter::new($server_writer), - ) - .await) - } - } - }; -} - -#[macro_export] -macro_rules! start_datagram_forward_with_codec_key { - ( - $codec_key:expr, - $udp_reader:expr, - $udp_writer:expr, - $tcp_reader:expr, - $tcp_writer:expr - ) => { - match $codec_key { - Some(key) => { - (start_datagram_forward( - $udp_reader, - $udp_writer, - CodecDatagramReader::new( - $tcp_reader, - snafu_error_get_or_return_ok!( - $crate::common::message::get_decodec(&key), - "failed to create decoder when datagram forward" - ), - ), - CodecDatagramWriter::new( - $tcp_writer, - snafu_error_get_or_return_ok!( - $crate::common::message::get_encodec(&key), - "failed to create encoder when datagram forward" - ), - ), - ) - .await) - } - None => { - (start_datagram_forward( - $udp_reader, - $udp_writer, - NormalDatagramReader::new($tcp_reader), - NormalDatagramWriter::new($tcp_writer), - ) - .await) - } - } - }; -} - #[cfg(test)] mod tests { use std::collections::VecDeque; diff --git a/src/utils/codec.rs b/src/utils/codec.rs index 59be873..438f3d4 100644 --- a/src/utils/codec.rs +++ b/src/utils/codec.rs @@ -5,8 +5,6 @@ use ring::aead::{ NONCE_LEN, }; -use crate::common::checksum::get_msg_header_key; - #[derive(Clone, Copy, Default)] struct Counter(u32); @@ -55,11 +53,6 @@ impl Aes256GcmCodec { }) } - pub fn try_new_with_default_key() -> RingResult { - let key = get_msg_header_key().map_err(|_| ring::error::Unspecified)?; - Aes256GcmCodec::try_new(key.as_ref()) - } - pub fn encrypt(&mut self, data: &mut [u8]) -> RingResult { self.seal.encrypt(data) } From cadcb5e88b338d5e9ce6af4767096a16fd03be2f Mon Sep 17 00:00:00 2001 From: LB7666 Date: Fri, 21 Aug 2026 13:03:31 +0800 Subject: [PATCH 65/74] Move to Rust 1.98.0, still on edition 2021 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `rust-toolchain.toml` plus the eight `toolchain:` pins across the four workflows. Deliberately no edition change yet: the edition migration is its own commit, so a failure there cannot be confused with a toolchain regression. Two clippy lints new in this release had to be settled first. `cargo check` would not have caught either, only `clippy -D warnings` does: - `manual_is_multiple_of` in the timing wheel — taken as written, since `is_multiple_of` says what the modulo was checking. - `result_large_err` on five signatures in `impl ServerSecurity`. Allowed rather than boxed, with the reason recorded at the impl: I measured both halves, and at 264 bytes `ServerInitialMessage` is already wider than the 256-byte `ServerInitialError`, so the `Result` is sized by its `Ok` variant. Boxing the error would add an allocation and change a public type while leaving `Result` at 264 bytes. Co-Authored-By: Claude Fable 5 --- .github/workflows/docker-publish.yml | 2 +- .github/workflows/release-ui.yml | 10 +++++----- .github/workflows/release.yml | 2 +- .github/workflows/syntax-check.yml | 2 +- rust-toolchain.toml | 2 +- src/common/auth/timing_wheel.rs | 2 +- src/common/message/secure.rs | 4 ++++ 7 files changed, 14 insertions(+), 10 deletions(-) diff --git a/.github/workflows/docker-publish.yml b/.github/workflows/docker-publish.yml index cf43090..3b22465 100644 --- a/.github/workflows/docker-publish.yml +++ b/.github/workflows/docker-publish.yml @@ -61,7 +61,7 @@ jobs: uses: houseabsolute/actions-rust-cross@v1 with: target: ${{ matrix.target.arch }}-unknown-linux-musl - toolchain: 1.88.0 + toolchain: 1.98.0 args: "--locked --release --bin pb-mapper" strip: true diff --git a/.github/workflows/release-ui.yml b/.github/workflows/release-ui.yml index a108439..6cefaf3 100644 --- a/.github/workflows/release-ui.yml +++ b/.github/workflows/release-ui.yml @@ -80,7 +80,7 @@ jobs: - name: Install Rust uses: dtolnay/rust-toolchain@stable with: - toolchain: 1.88.0 + toolchain: 1.98.0 - name: Build latest Windows FFI run: | make build-pb-mapper-ffi-windows @@ -128,7 +128,7 @@ jobs: - name: Install Rust uses: dtolnay/rust-toolchain@stable with: - toolchain: 1.88.0 + toolchain: 1.98.0 - name: Install dependencies run: | sudo apt-get update -y @@ -236,7 +236,7 @@ jobs: - name: Install Rust uses: dtolnay/rust-toolchain@stable with: - toolchain: 1.88.0 + toolchain: 1.98.0 - name: Set up Android NDK uses: nttld/setup-ndk@v1 with: @@ -399,7 +399,7 @@ jobs: - name: Install Rust uses: dtolnay/rust-toolchain@stable with: - toolchain: 1.88.0 + toolchain: 1.98.0 - name: Install appdmg run: | npm install -g appdmg @@ -480,7 +480,7 @@ jobs: - name: Install Rust uses: dtolnay/rust-toolchain@stable with: - toolchain: 1.88.0 + toolchain: 1.98.0 - name: Build latest iOS FFI run: | make build-pb-mapper-ffi-ios diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 7c12da9..dabf432 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -61,7 +61,7 @@ jobs: uses: houseabsolute/actions-rust-cross@v1 with: target: ${{ matrix.platform.target }} - toolchain: 1.88.0 + toolchain: 1.98.0 args: "--locked --release --bin pb-mapper" strip: true diff --git a/.github/workflows/syntax-check.yml b/.github/workflows/syntax-check.yml index 6be4f0d..e2804bc 100644 --- a/.github/workflows/syntax-check.yml +++ b/.github/workflows/syntax-check.yml @@ -91,7 +91,7 @@ jobs: - name: Install Rust uses: dtolnay/rust-toolchain@stable with: - toolchain: 1.88.0 + toolchain: 1.98.0 components: clippy, rustfmt - name: Cache Rust dependencies diff --git a/rust-toolchain.toml b/rust-toolchain.toml index e88baf1..b73c15e 100644 --- a/rust-toolchain.toml +++ b/rust-toolchain.toml @@ -1,2 +1,2 @@ [toolchain] -channel = "1.88.0" +channel = "1.98.0" diff --git a/src/common/auth/timing_wheel.rs b/src/common/auth/timing_wheel.rs index d219645..989c340 100644 --- a/src/common/auth/timing_wheel.rs +++ b/src/common/auth/timing_wheel.rs @@ -190,7 +190,7 @@ impl TimingWheel { for level in 0..self.levels.len() { // A level turns over every `radix^level` ticks. Once one does not, no // coarser one can either, since its period divides theirs. - if self.ticks % self.period(level) != 0 { + if !self.ticks.is_multiple_of(self.period(level)) { break; } let bucket = self.rotate(level); diff --git a/src/common/message/secure.rs b/src/common/message/secure.rs index 6038a88..b96595d 100644 --- a/src/common/message/secure.rs +++ b/src/common/message/secure.rs @@ -362,6 +362,10 @@ pub struct ServerSecurity { failure_logs: Arc>, } +// `ServerInitialError` is 256 bytes, but the success type it is paired with, +// `ServerInitialMessage`, is 264 — so the `Result` is already sized by its `Ok` +// variant and boxing the error would buy an allocation for no size reduction. +#[allow(clippy::result_large_err)] impl ServerSecurity { pub fn new(auth: AuthRuntime) -> Self { let replay_path = auth.config().state_dir.join("connection.replay"); From 6c763ba56c85c11ddcbb4c83ea4aff49839b7318 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Fri, 21 Aug 2026 13:09:24 +0800 Subject: [PATCH 66/74] Move DNS resolution to hickory-resolver MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `trust-dns-resolver` 0.23.2 has been frozen since 2023; hickory is its continuation. Kept separate from the other dependency bumps because it is the one change here with runtime behavior attached. hickory removed the blocking `Resolver` outright, so the sync path can no longer reach the custom DNS servers and now goes straight to `std`. That is the behavior it already had: `get_custom_resolver` returned `None` inside a Tokio runtime, and outside one the caller fell back to `std` whenever the resolver was missing. The async path is unaffected and still prefers the custom servers. `build()` is fallible, so the resolver is a `LazyLock>` built once — retrying per lookup would just repeat the same failure — and `get_ip_addrs_async` reports a missing resolver through the existing fallback. The unused `get_ip_addrs` and `DNS_QUERY_PORT` go away with it; `NameServerConfig::udp_and_tcp` already implies port 53 and trusting negative responses, which is what the old call passed explicitly. Verified both paths at runtime against `localhost` and a public domain. `trust-dns-resolver` remains in the tree only via `uni-stream`. Co-Authored-By: Claude Fable 5 --- Cargo.lock | 470 +++++++++++++++++++++++++++++++++++++++++++--- Cargo.toml | 4 +- src/utils/addr.rs | 80 ++++---- 3 files changed, 484 insertions(+), 70 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d0e664c..f8065c7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -81,9 +81,15 @@ checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.114", ] +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + [[package]] name = "base64" version = "0.22.1" @@ -125,6 +131,12 @@ dependencies = [ "rustc_version", ] +[[package]] +name = "bumpalo" +version = "3.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" + [[package]] name = "bytes" version = "1.11.0" @@ -195,7 +207,7 @@ dependencies = [ "heck", "proc-macro2", "quote", - "syn", + "syn 2.0.114", ] [[package]] @@ -210,6 +222,32 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b05b61dc5112cbb17e4b6cd61790d9845d13888356391624cbe7e41efeac1e75" +[[package]] +name = "combine" +version = "4.6.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba5a308b75df32fe02788e748662718f03fde005016435c444eea572398219fd" +dependencies = [ + "bytes", + "memchr", +] + +[[package]] +name = "core-foundation" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e195e091a93c46f7102ec7818a2aa394e1e1771c3ab4825963fa03e45afb8f" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "core-foundation-sys" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" + [[package]] name = "cpufeatures" version = "0.3.0" @@ -219,6 +257,36 @@ dependencies = [ "libc", ] +[[package]] +name = "critical-section" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "790eea4361631c5e7d22598ecd5723ff611904e3344ce8720784c93e3d83d40b" + +[[package]] +name = "crossbeam-channel" +version = "0.5.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d85363c37faeca707aef026efa9f3b34d077bce547e48f770770625c6013679e" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17" + [[package]] name = "cty" version = "0.2.2" @@ -260,7 +328,7 @@ checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.114", ] [[package]] @@ -269,6 +337,12 @@ version = "0.15.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1aaf95b3e5c8f23aa320147307562d361db0ae0d51242340f558153b4eb2439b" +[[package]] +name = "either" +version = "1.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "252afb9ae5eaa683babdc6a068b3f5726eb19e05070c731f9b2a23a7c3e8ed34" + [[package]] name = "enum-as-inner" version = "0.6.1" @@ -278,7 +352,7 @@ dependencies = [ "heck", "proc-macro2", "quote", - "syn", + "syn 2.0.114", ] [[package]] @@ -380,7 +454,7 @@ checksum = "162ee34ebcb7c64a8abebc059ce0fee27c2262618d7b60ed8faf72fef13c3650" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.114", ] [[package]] @@ -464,6 +538,76 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "hickory-net" +version = "0.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2295ed2f9c31e471e1428a8f88a3f0e1f4b27c15049592138d1eebe9c35b183" +dependencies = [ + "async-trait", + "cfg-if", + "data-encoding", + "futures-channel", + "futures-io", + "futures-util", + "hickory-proto", + "idna 1.1.0", + "ipnet", + "jni", + "rand 0.10.0", + "thiserror 2.0.20", + "tinyvec", + "tokio", + "tracing", + "url", +] + +[[package]] +name = "hickory-proto" +version = "0.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bab31817bfb44672a252e97fe81cd0c18d1b2cf892108922f6818820df8c643" +dependencies = [ + "data-encoding", + "idna 1.1.0", + "ipnet", + "jni", + "once_cell", + "prefix-trie", + "rand 0.10.0", + "ring", + "thiserror 2.0.20", + "tinyvec", + "tracing", + "url", +] + +[[package]] +name = "hickory-resolver" +version = "0.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0d58d28879ceecde6607729660c2667a081ccdc082e082675042793960f178c" +dependencies = [ + "cfg-if", + "futures-util", + "hickory-net", + "hickory-proto", + "ipconfig", + "ipnet", + "jni", + "moka", + "ndk-context", + "once_cell", + "parking_lot", + "rand 0.10.0", + "resolv-conf", + "smallvec", + "system-configuration", + "thiserror 2.0.20", + "tokio", + "tracing", +] + [[package]] name = "icu_collections" version = "2.1.1" @@ -611,6 +755,9 @@ name = "ipnet" version = "2.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "469fb0b9cefa57e3ef31275ee7cacb78f2fdca44e4765491884a2b119d4eb130" +dependencies = [ + "serde", +] [[package]] name = "is_terminal_polyfill" @@ -624,6 +771,66 @@ version = "1.0.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2" +[[package]] +name = "jni" +version = "0.22.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5efd9a482cf3a427f00d6b35f14332adc7902ce91efb778580e180ff90fa3498" +dependencies = [ + "cfg-if", + "combine", + "jni-macros", + "jni-sys", + "log", + "simd_cesu8", + "thiserror 2.0.20", + "walkdir", + "windows-link", +] + +[[package]] +name = "jni-macros" +version = "0.22.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a00109accc170f0bdb141fed3e393c565b6f5e072365c3bd58f5b062591560a3" +dependencies = [ + "proc-macro2", + "quote", + "rustc_version", + "simd_cesu8", + "syn 2.0.114", +] + +[[package]] +name = "jni-sys" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6377a88cb3910bee9b0fa88d4f42e1d2da8e79915598f65fb0c7ee14c878af2" +dependencies = [ + "jni-sys-macros", +] + +[[package]] +name = "jni-sys-macros" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38c0b942f458fe50cdac086d2f946512305e5631e720728f2a61aabcd47a6264" +dependencies = [ + "quote", + "syn 2.0.114", +] + +[[package]] +name = "js-sys" +version = "0.3.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e0c1080212aad755ea003d18543e8768dd432c48819efd73a7bf1e39b7a5a3a" +dependencies = [ + "cfg-if", + "futures-util", + "wasm-bindgen", +] + [[package]] name = "kanal" version = "0.2.0-beta2" @@ -725,6 +932,29 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "moka" +version = "0.12.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4293f18e7567a1caf3c584855554377025c65e0aa445344d04171f5ad63d19b9" +dependencies = [ + "crossbeam-channel", + "crossbeam-epoch", + "crossbeam-utils", + "equivalent", + "parking_lot", + "portable-atomic", + "smallvec", + "tagptr", + "uuid", +] + +[[package]] +name = "ndk-context" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "27b02d87554356db9e9a873add8782d4ea6e3e58ea071a9adb9a2e8ddb884a8b" + [[package]] name = "nu-ansi-term" version = "0.50.3" @@ -734,11 +964,24 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + [[package]] name = "once_cell" version = "1.21.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" +dependencies = [ + "critical-section", + "portable-atomic", +] [[package]] name = "once_cell_polyfill" @@ -786,6 +1029,7 @@ dependencies = [ "dotenvy", "futures", "hashbrown 0.16.1", + "hickory-resolver", "kanal", "parking_lot", "rand 0.10.0", @@ -799,7 +1043,6 @@ dependencies = [ "tokio-util", "tracing", "tracing-subscriber", - "trust-dns-resolver", "uni-stream", ] @@ -839,6 +1082,12 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" +[[package]] +name = "portable-atomic" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05c8b63e8d9609db387f0324918f81d68fe27748f084ef092fb35954d0539a85" + [[package]] name = "potential_utf" version = "0.1.4" @@ -857,6 +1106,17 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "prefix-trie" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cf6e3177f0684016a5c209b00882e15f8bdd3f3bb48f0491df10cd102d0c6e7" +dependencies = [ + "either", + "ipnet", + "num-traits", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -864,7 +1124,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" dependencies = [ "proc-macro2", - "syn", + "syn 2.0.114", ] [[package]] @@ -955,7 +1215,7 @@ checksum = "ba009ff324d1fc1b900bd1fdb31564febe58a8ccc8a6fdbb93b543d33b13ca43" dependencies = [ "getrandom 0.2.17", "libredox", - "thiserror", + "thiserror 1.0.69", ] [[package]] @@ -1004,6 +1264,21 @@ dependencies = [ "semver", ] +[[package]] +name = "rustversion" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" + +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + [[package]] name = "scopeguard" version = "1.2.0" @@ -1043,7 +1318,7 @@ checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.114", ] [[package]] @@ -1084,6 +1359,22 @@ dependencies = [ "libc", ] +[[package]] +name = "simd_cesu8" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11031e251abf8611c80f460e19dbdeb54a66db918e49c65a7065b46ac7aec520" +dependencies = [ + "rustc_version", + "simdutf8", +] + +[[package]] +name = "simdutf8" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" + [[package]] name = "slab" version = "0.4.11" @@ -1114,7 +1405,7 @@ dependencies = [ "heck", "proc-macro2", "quote", - "syn", + "syn 2.0.114", ] [[package]] @@ -1166,6 +1457,17 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "syn" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + [[package]] name = "synstructure" version = "0.13.2" @@ -1174,16 +1476,52 @@ checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.114", +] + +[[package]] +name = "system-configuration" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a13f3d0daba03132c0aa9767f98351b3488edc2c100cda2d2ec2b04f3d8d3c8b" +dependencies = [ + "bitflags", + "core-foundation", + "system-configuration-sys", ] +[[package]] +name = "system-configuration-sys" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e1d1b10ced5ca923a1fcb8d03e96b8d3268065d724548c0211415ff6ac6bac4" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "tagptr" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417" + [[package]] name = "thiserror" version = "1.0.69" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" dependencies = [ - "thiserror-impl", + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec86235f5fcc2a73650310756d2ac5b138a5780bbbdfae3eeccec992c435ba4f" +dependencies = [ + "thiserror-impl 2.0.20", ] [[package]] @@ -1194,7 +1532,18 @@ checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.114", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc04cd3e1236dd4a98afca4569f2deb3f120e5422a4023be2cb683f8486292af" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", ] [[package]] @@ -1256,7 +1605,7 @@ checksum = "af407857209536a95c8e56f8231ef2c2e2aff839b22e07a1ffcbc617e9db9fa5" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.114", ] [[package]] @@ -1291,7 +1640,7 @@ checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.114", ] [[package]] @@ -1364,7 +1713,7 @@ dependencies = [ "once_cell", "rand 0.8.5", "smallvec", - "thiserror", + "thiserror 1.0.69", "tinyvec", "tokio", "tracing", @@ -1386,7 +1735,7 @@ dependencies = [ "rand 0.8.5", "resolv-conf", "smallvec", - "thiserror", + "thiserror 1.0.69", "tokio", "tracing", "trust-dns-proto", @@ -1465,12 +1814,33 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" +[[package]] +name = "uuid" +version = "1.24.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2cefc03fd367c0c6d4305de1b312cf00248c4114f4a0418ce6a6af769e3b0bd9" +dependencies = [ + "getrandom 0.4.1", + "js-sys", + "wasm-bindgen", +] + [[package]] name = "valuable" version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + [[package]] name = "wasi" version = "0.11.1+wasi-snapshot-preview1" @@ -1495,6 +1865,51 @@ dependencies = [ "wit-bindgen", ] +[[package]] +name = "wasm-bindgen" +version = "0.2.127" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b70935747edd64d89de3efa29d73789b806c15798f8e7dca4d8ac356b50ce70" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.127" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77775f8f3f7217702089053b94958f8f54061a3f663417df76e19cbdcca29bc1" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.127" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e11d33f857dc2fb11b8bc75aee111aa9cbeb12cd9f25efd3d4c2a3dd4e235284" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn 2.0.114", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.127" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ef64dbcc55df09c7e5a46182d181c2cfa3e925f3da937ea764728b4bbb9dcbf" +dependencies = [ + "unicode-ident", +] + [[package]] name = "wasm-encoder" version = "0.244.0" @@ -1535,6 +1950,15 @@ version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72069c3113ab32ab29e5584db3c6ec55d416895e60715417b5b883a357c3e471" +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "windows-link" version = "0.2.1" @@ -1803,7 +2227,7 @@ dependencies = [ "heck", "indexmap", "prettyplease", - "syn", + "syn 2.0.114", "wasm-metadata", "wit-bindgen-core", "wit-component", @@ -1819,7 +2243,7 @@ dependencies = [ "prettyplease", "proc-macro2", "quote", - "syn", + "syn 2.0.114", "wit-bindgen-core", "wit-bindgen-rust", ] @@ -1886,7 +2310,7 @@ checksum = "b659052874eb698efe5b9e8cf382204678a0086ebf46982b79d6ca3182927e5d" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.114", "synstructure", ] @@ -1907,7 +2331,7 @@ checksum = "2c7962b26b0a8685668b671ee4b54d007a67d4eaf05fda79ac0ecf41e32270f1" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.114", ] [[package]] @@ -1927,7 +2351,7 @@ checksum = "d71e5d6e06ab090c67b5e44993ec16b72dcbaabc526db883a360057678b48502" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.114", "synstructure", ] @@ -1961,7 +2385,7 @@ checksum = "eadce39539ca5cb3985590102671f2567e659fca9666581ad3411d59207951f3" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.114", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index fbfc215..f4e4ed9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,7 +21,7 @@ clap.workspace = true futures.workspace = true better_mimalloc_rs.workspace = true bytes.workspace = true -trust-dns-resolver.workspace = true +hickory-resolver.workspace = true ring.workspace = true uni-stream.workspace = true kanal.workspace = true @@ -63,7 +63,7 @@ clap = { version = "4.5", features = ["derive"] } futures = "0.3.31" better_mimalloc_rs = { version = "0.1.2", features = ["config"] } bytes = "1.11" -trust-dns-resolver = { version = "0.23.2" } +hickory-resolver = { version = "0.26.1" } ring = "0.17.14" base64 = "0.22.1" subtle = "2.6.1" diff --git a/src/utils/addr.rs b/src/utils/addr.rs index a8485ca..8e5ce8a 100644 --- a/src/utils/addr.rs +++ b/src/utils/addr.rs @@ -4,9 +4,10 @@ use std::pin::Pin; use std::sync::LazyLock; use std::task::{ready, Context, Poll}; +use hickory_resolver::config::{NameServerConfig, ResolverConfig, ResolverOpts}; +use hickory_resolver::net::runtime::TokioRuntimeProvider; +use hickory_resolver::{Resolver, TokioResolver}; use tokio::task::JoinHandle; -use trust_dns_resolver::config::{NameServerConfigGroup, ResolverConfig, ResolverOpts}; -use trust_dns_resolver::{Resolver, TokioAsyncResolver}; type Result = std::result::Result; type ReadyFuture = future::Ready>; @@ -243,42 +244,40 @@ const DEFAULT_DNS_SERVER_GROUP: &[IpAddr] = &[ IpAddr::V6(Ipv6Addr::new(0x2001, 0x4860, 0x4860, 0, 0, 0, 0, 0x8888)), // google ]; -const DNS_QUERY_PORT: u16 = 53; - #[inline] fn custom_resolver_config() -> ResolverConfig { + // `udp_and_tcp` uses the standard DNS port and trusts negative responses, + // matching what this passed explicitly before. ResolverConfig::from_parts( None, vec![], - NameServerConfigGroup::from_ips_clear(DEFAULT_DNS_SERVER_GROUP, DNS_QUERY_PORT, true), + DEFAULT_DNS_SERVER_GROUP + .iter() + .copied() + .map(NameServerConfig::udp_and_tcp) + .collect::>(), ) } +/// The custom resolver, or `None` if it could not be built. +/// +/// Built once: `build()` can fail, and retrying per lookup would repeat the +/// same failure. Callers fall back to the system resolver. #[inline] -pub fn get_custom_resolver() -> Option { - // The sync resolver uses `block_on` internally and will panic if called from a Tokio runtime - // thread. Keep a guard here so callers can fall back to system DNS, and use the async helpers - // below when running inside async code. - if tokio::runtime::Handle::try_current().is_ok() { - tracing::debug!("Skipping sync custom DNS resolver inside Tokio runtime thread"); - return None; - } - - match Resolver::new(custom_resolver_config(), ResolverOpts::default()) { - Ok(r) => Some(r), - Err(e) => { - tracing::error!( - "Create custom dns resolver error:{e},we will use default dns resolver" - ); - None +fn get_custom_async_resolver() -> Option { + static RESOLVER: LazyLock> = LazyLock::new(|| { + let mut builder = Resolver::builder_with_config( + custom_resolver_config(), + TokioRuntimeProvider::default(), + ); + *builder.options_mut() = ResolverOpts::default(); + match builder.build() { + Ok(resolver) => Some(resolver), + Err(e) => { + tracing::error!("Create custom dns resolver error:{e},falling back to system dns"); + None + } } - } -} - -#[inline] -fn get_custom_async_resolver() -> TokioAsyncResolver { - static RESOLVER: LazyLock = LazyLock::new(|| { - TokioAsyncResolver::tokio(custom_resolver_config(), ResolverOpts::default()) }); RESOLVER.clone() } @@ -298,24 +297,15 @@ macro_rules! try_opt { }; } -fn get_ip_addrs(s: &str) -> Result> { - thread_local! { - static RESOLVER:Option = get_custom_resolver(); - } - let result = RESOLVER.with(|r| r.as_ref().map(|r| r.lookup_ip(s))); - try_opt!(result, "custom resolver not exist") - .map(|v| v.into_iter().collect()) - .map_err(|e| invalid_input!(e)) -} - -/// Blocking DNS lookup. Avoid calling this from inside a Tokio runtime thread. +/// Blocking DNS lookup, via the system resolver. +/// +/// The custom DNS servers are only reachable from the async helpers below. +/// hickory has no blocking resolver, and the sync path never reached the +/// custom one in practice anyway: it was skipped inside a Tokio runtime, and +/// outside one this fell back to `std` whenever it was unavailable. #[inline] pub fn get_socket_addrs_from_host_port(host: &str, port: u16) -> Result> { - match get_ip_addrs(host) { - Ok(r) => Ok(r.into_iter().map(|ip| SocketAddr::new(ip, port)).collect()), - // Resolve dns properly with the standard library - Err(_) => std::net::ToSocketAddrs::to_socket_addrs(&(host, port)).map(|v| v.collect()), - } + std::net::ToSocketAddrs::to_socket_addrs(&(host, port)).map(|v| v.collect()) } /// Blocking DNS lookup. Avoid calling this from inside a Tokio runtime thread. @@ -328,7 +318,7 @@ pub fn get_socket_addrs(s: &str) -> Result> { /// Async DNS lookup using the custom resolver, safe to call inside Tokio runtimes. pub async fn get_ip_addrs_async(s: &str) -> Result> { - let resolver = get_custom_async_resolver(); + let resolver = try_opt!(get_custom_async_resolver(), "custom resolver not exist"); resolver .lookup_ip(s) .await From 9d9f65fe67b2083826ba4317ae362080c118a1d2 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Fri, 21 Aug 2026 13:11:20 +0800 Subject: [PATCH 67/74] Upgrade the remaining dependencies MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The four bumps that cross a compatibility boundary: snafu 0.8.7 → 0.9.2, hashbrown 0.16 → 0.17.1, base64 0.22.1 → 0.23.1, and dirs 5.0 → 6.0.0 in the FFI manifest, which declares it directly rather than through the workspace. None needed a source change. snafu 0.9 was the one worth checking, since every error type here goes through it, and all of `display`, `visibility(pub(super))`, `context`, `ensure!` and `.fail()` compile as written. hashbrown 0.15 and 0.16 stay in the lockfile via `uni-stream`; our own crates resolve to 0.17.1. Co-Authored-By: Claude Fable 5 --- Cargo.lock | 41 +++++++++++++++++++----------- Cargo.toml | 6 ++--- ui/native/pb_mapper_ffi/Cargo.toml | 2 +- 3 files changed, 30 insertions(+), 19 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f8065c7..2e6535e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -92,9 +92,9 @@ checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" [[package]] name = "base64" -version = "0.22.1" +version = "0.23.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5" [[package]] name = "better_mimalloc_rs" @@ -301,23 +301,23 @@ checksum = "d7a1e2f27636f116493b8b860f5546edb47c8d8f8ea73e1d2a20be88e28d1fea" [[package]] name = "dirs" -version = "5.0.1" +version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "44c45a9d03d6676652bcb5e724c7e988de1acad23a711b5217ab9cbecbec2225" +checksum = "c3e8aa94d75141228480295a7d0e7feb620b1a5ad9f12bc40be62411e38cce4e" dependencies = [ "dirs-sys", ] [[package]] name = "dirs-sys" -version = "0.4.1" +version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "520f05a5cbd335fae5a99ff7a6ab8627577660ee5cfd6a94a6a929b52ff0321c" +checksum = "e01a3366d27ee9890022452ee61b2b63a67e6f13f58900b651ff5665f0bb1fab" dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] @@ -532,6 +532,17 @@ dependencies = [ "foldhash 0.2.0", ] +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash 0.2.0", +] + [[package]] name = "heck" version = "0.5.0" @@ -1028,7 +1039,7 @@ dependencies = [ "clap", "dotenvy", "futures", - "hashbrown 0.16.1", + "hashbrown 0.17.1", "hickory-resolver", "kanal", "parking_lot", @@ -1209,13 +1220,13 @@ dependencies = [ [[package]] name = "redox_users" -version = "0.4.6" +version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ba009ff324d1fc1b900bd1fdb31564febe58a8ccc8a6fdbb93b543d33b13ca43" +checksum = "a4e608c6638b9c18977b00b475ac1f28d14e84b27d8d42f70e0bf1e3dec127ac" dependencies = [ "getrandom 0.2.17", "libredox", - "thiserror 1.0.69", + "thiserror 2.0.20", ] [[package]] @@ -1389,18 +1400,18 @@ checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" [[package]] name = "snafu" -version = "0.8.9" +version = "0.9.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6e84b3f4eacbf3a1ce05eac6763b4d629d60cbc94d632e4092c54ade71f1e1a2" +checksum = "e45cb604038abb7b926b679887b3226d8d0f23874b66623625a0454be425a4b7" dependencies = [ "snafu-derive", ] [[package]] name = "snafu-derive" -version = "0.8.9" +version = "0.9.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c1c97747dbf44bb1ca44a561ece23508e99cb592e862f22222dcf42f51d1e451" +checksum = "287f59010008f0d7cf5e3b03196d666c1acc46c8d3e9cf34c28a1a7157601e72" dependencies = [ "heck", "proc-macro2", diff --git a/Cargo.toml b/Cargo.toml index f4e4ed9..cf2ed2c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -49,7 +49,7 @@ rand = "0.10" socket2 = "0.6" tokio = { version = "1", features = ["full"] } tokio-util = "0.7" -snafu = "0.8.7" +snafu = "0.9.2" serde = { version = "1.0", features = ["derive"] } serde_json = { version = "1.0", default-features = false, features = ["alloc"] } tracing = "0.1.40" @@ -58,14 +58,14 @@ tracing-subscriber = { version = "0.3.18", features = [ "fmt", "json", ], default-features = true } -hashbrown = { version = "0.16" } +hashbrown = { version = "0.17.1" } clap = { version = "4.5", features = ["derive"] } futures = "0.3.31" better_mimalloc_rs = { version = "0.1.2", features = ["config"] } bytes = "1.11" hickory-resolver = { version = "0.26.1" } ring = "0.17.14" -base64 = "0.22.1" +base64 = "0.23.1" subtle = "2.6.1" parking_lot = "0.12" uni-stream = { git = "https://github.com/acking-you/uni-stream.git", branch = "master" } diff --git a/ui/native/pb_mapper_ffi/Cargo.toml b/ui/native/pb_mapper_ffi/Cargo.toml index 5cc391b..ae239d5 100644 --- a/ui/native/pb_mapper_ffi/Cargo.toml +++ b/ui/native/pb_mapper_ffi/Cargo.toml @@ -19,7 +19,7 @@ tokio.workspace = true clap.workspace = true better_mimalloc_rs.workspace = true tokio-util = "0.7" -dirs = "5.0" +dirs = "6.0.0" pb-mapper = { path = "../../../" } uni-stream.workspace = true From 601bb9ee6348e61b5960292b11c81bfd038deac2 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Fri, 21 Aug 2026 13:22:34 +0800 Subject: [PATCH 68/74] Migrate to edition 2024 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The FFI crate hardcodes its edition rather than inheriting the workspace one, so it needs the same change and turned out to carry most of the work. Mechanical, in four groups: - Four `extern "C"` blocks become `unsafe extern "C"`. - 28 `#[no_mangle]` become `#[unsafe(no_mangle)]`, all in the FFI crate. These name the symbols Dart dlopens, so I diffed `nm` output before and after: the same 28 `pb_mapper_*` symbols, unchanged. - 17 `env::set_var`/`remove_var` calls now need `unsafe`. Each got a safety note saying why it is sound rather than a blanket wrap, because the reasons differ: `set_process_msg_header_key` keeps the environment only as a mirror while the authoritative credential lives behind an `RwLock`; the CLI overrides run on the main thread before any task spawns; the tests either hold `PROCESS_CREDENTIAL_TEST_LOCK` or own the variable outright and restore it. - 17 nested `if let` blocks become let-chains, which edition 2024 stabilised — `clippy::collapsible_if` now flags them. Not in the plan's count: the `#[no_mangle]` attributes and the let-chains, both of which only appear once the earlier group compiles. Co-Authored-By: Claude Fable 5 --- Cargo.toml | 2 +- examples/pb_local_server.rs | 2 +- src/bin/pb-mapper.rs | 89 ++++++++++++-------- src/common/auth.rs | 30 +++---- src/common/auth/config.rs | 28 +++--- src/common/auth/keys.rs | 8 +- src/common/auth/persistence/fs.rs | 2 +- src/common/auth/persistence/wal.rs | 10 +-- src/common/auth/runtime.rs | 8 +- src/common/auth/tests.rs | 52 +++++++++--- src/common/checksum.rs | 61 ++++++++------ src/common/config.rs | 34 ++++++-- src/common/manager.rs | 10 +-- src/common/message/mod.rs | 6 +- src/common/message/secure.rs | 22 ++--- src/common/message/secure/replay.rs | 8 +- src/common/message/secure/tests.rs | 16 ++-- src/local/client/mod.rs | 8 +- src/local/client/status.rs | 6 +- src/local/client/stream.rs | 4 +- src/local/server/mod.rs | 38 +++++---- src/local/server/stream.rs | 4 +- src/pb_server/admin.rs | 4 +- src/pb_server/client.rs | 12 +-- src/pb_server/error.rs | 8 +- src/pb_server/mod.rs | 20 ++--- src/pb_server/runtime.rs | 88 ++++++++++--------- src/pb_server/server.rs | 4 +- src/pb_server/status.rs | 2 +- src/utils/addr.rs | 2 +- src/utils/codec.rs | 8 +- tests/regression.rs | 34 +++++--- tests/test_delay.rs | 14 +-- ui/native/pb_mapper_ffi/Cargo.toml | 2 +- ui/native/pb_mapper_ffi/src/cli.rs | 6 +- ui/native/pb_mapper_ffi/src/client.rs | 10 +-- ui/native/pb_mapper_ffi/src/config.rs | 6 +- ui/native/pb_mapper_ffi/src/ctl/endpoint.rs | 18 ++-- ui/native/pb_mapper_ffi/src/ctl/server.rs | 4 +- ui/native/pb_mapper_ffi/src/events.rs | 4 +- ui/native/pb_mapper_ffi/src/handle.rs | 10 +-- ui/native/pb_mapper_ffi/src/lib.rs | 6 +- ui/native/pb_mapper_ffi/src/logging.rs | 10 +-- ui/native/pb_mapper_ffi/src/response.rs | 2 +- ui/native/pb_mapper_ffi/src/server.rs | 12 +-- ui/native/pb_mapper_ffi/src/service.rs | 10 +-- ui/native/pb_mapper_ffi/src/state.rs | 10 +-- ui/native/pb_mapper_ffi/src/state/runtime.rs | 12 +-- ui/native/pb_mapper_ffi/src/state/status.rs | 13 ++- 49 files changed, 433 insertions(+), 346 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index cf2ed2c..d6582e6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -42,7 +42,7 @@ exclude = ["deps/uni-stream", "deps/kanal"] [workspace.package] version = "0.4.0" authors = ["L_B__"] -edition = "2021" +edition = "2024" [workspace.dependencies] rand = "0.10" diff --git a/examples/pb_local_server.rs b/examples/pb_local_server.rs index a282436..b01c15d 100644 --- a/examples/pb_local_server.rs +++ b/examples/pb_local_server.rs @@ -1,5 +1,5 @@ use pb_mapper::common::config::init_tracing; -use pb_mapper::local::server::{run_server_side_cli, ServerTunnelOptions}; +use pb_mapper::local::server::{ServerTunnelOptions, run_server_side_cli}; use uni_stream::stream::TcpStreamProvider; #[tokio::main] diff --git a/src/bin/pb-mapper.rs b/src/bin/pb-mapper.rs index 1f6b4fa..fc10703 100644 --- a/src/bin/pb-mapper.rs +++ b/src/bin/pb-mapper.rs @@ -19,27 +19,27 @@ use std::time::Duration; use better_mimalloc_rs::MiMalloc; use clap::{Args, Parser, Subcommand, ValueEnum}; use pb_mapper::common::auth::{ - acquire_state_dir_lock, generate_admin_key, initialize_admin_key, write_admin_key_file, AuthConfig, KeyPage, LegacyProtocolPolicy, MAX_TEMP_KEY_CAPACITY, MAX_TEMP_KEY_TTL, - MIN_TEMP_KEY_TTL, + MIN_TEMP_KEY_TTL, acquire_state_dir_lock, generate_admin_key, initialize_admin_key, + write_admin_key_file, }; use pb_mapper::common::checksum::set_process_msg_header_key; -use pb_mapper::common::checksum::{setup_machine_msg_header_key, MACHINE_MSG_HEADER_KEY_PATH}; +use pb_mapper::common::checksum::{MACHINE_MSG_HEADER_KEY_PATH, setup_machine_msg_header_key}; use pb_mapper::common::config::{ - control_io_timeout, get_pb_mapper_server_async, get_sockaddr_async, init_tracing, - keep_alive_from_env, StatusOp, + StatusOp, control_io_timeout, get_pb_mapper_server_async, get_sockaddr_async, init_tracing, + keep_alive_from_env, }; +use pb_mapper::common::message::MessageReader; use pb_mapper::common::message::command::{ AdminConnectionPage, AdminRequest, AdminResponse, AdminServicePage, MessageSerializer, PbConnRequest, PbConnResponse, }; use pb_mapper::common::message::forward::StreamForward; use pb_mapper::common::message::secure::ClientHeaderSession; -use pb_mapper::common::message::MessageReader; use pb_mapper::local::client::{ handle_status_cli_scoped, run_client_side_cli_with_callback_scoped, }; -use pb_mapper::local::server::{run_server_side_cli_with_pinned_credential, ServerTunnelOptions}; +use pb_mapper::local::server::{ServerTunnelOptions, run_server_side_cli_with_pinned_credential}; use pb_mapper::pb_server::run_server_with_shutdown; use tokio::net::TcpStream; use tokio_util::sync::CancellationToken; @@ -229,9 +229,18 @@ async fn run(cli: Cli) -> Result<(), Box> { Ok(()) } +/// Publishes CLI flags as environment variables, which is how the auth +/// subsystem reads its configuration. +/// +/// # Safety note +/// +/// Mutating the environment is unsafe in edition 2024 because it races +/// concurrent readers. Every call here happens during argument handling on the +/// main thread, before any runtime task or thread is spawned, so there is no +/// concurrent reader to race. fn apply_server_auth_overrides(args: &ServerArgs) -> Result<(), Box> { if let Some(auth_state_dir) = &args.auth_state_dir { - std::env::set_var("PB_MAPPER_AUTH_STATE_DIR", auth_state_dir); + unsafe { std::env::set_var("PB_MAPPER_AUTH_STATE_DIR", auth_state_dir) }; } if let Some(max_temporary_keys) = args.max_temporary_keys { if !(1..=MAX_TEMP_KEY_CAPACITY).contains(&max_temporary_keys) { @@ -240,10 +249,12 @@ fn apply_server_auth_overrides(args: &ServerArgs) -> Result<(), Box> ) .into()); } - std::env::set_var( - "PB_MAPPER_AUTH_MAX_TEMP_KEYS", - max_temporary_keys.to_string(), - ); + unsafe { + std::env::set_var( + "PB_MAPPER_AUTH_MAX_TEMP_KEYS", + max_temporary_keys.to_string(), + ) + }; } if let Some(max_temporary_key_ttl) = args.max_temporary_key_ttl { if max_temporary_key_ttl < MIN_TEMP_KEY_TTL || max_temporary_key_ttl > MAX_TEMP_KEY_TTL { @@ -254,10 +265,12 @@ fn apply_server_auth_overrides(args: &ServerArgs) -> Result<(), Box> ) .into()); } - std::env::set_var( - "PB_MAPPER_AUTH_MAX_TEMP_TTL_SECS", - max_temporary_key_ttl.as_secs().to_string(), - ); + unsafe { + std::env::set_var( + "PB_MAPPER_AUTH_MAX_TEMP_TTL_SECS", + max_temporary_key_ttl.as_secs().to_string(), + ) + }; } Ok(()) } @@ -265,13 +278,17 @@ fn apply_server_auth_overrides(args: &ServerArgs) -> Result<(), Box> async fn run_server(args: ServerArgs) -> Result<(), Box> { apply_server_auth_overrides(&args)?; if let Some(legacy_protocol) = args.legacy_protocol { - std::env::set_var( - "PB_MAPPER_LEGACY_PROTOCOL", - match legacy_protocol { - LegacyProtocolArg::Allow => "allow", - LegacyProtocolArg::Deny => "deny", - }, - ); + // SAFETY: as in `apply_server_auth_overrides` — this runs before the + // server spawns anything that reads the environment. + unsafe { + std::env::set_var( + "PB_MAPPER_LEGACY_PROTOCOL", + match legacy_protocol { + LegacyProtocolArg::Allow => "allow", + LegacyProtocolArg::Deny => "deny", + }, + ) + }; } let auth_config = AuthConfig::default(); if args.init_admin_key { @@ -426,7 +443,7 @@ fn parse_duration(raw: &str) -> Result { _ => { return Err(format!( "unsupported duration unit `{unit}`; use s, m, h, or d" - )) + )); } }; value @@ -548,17 +565,19 @@ mod tests { Cli::try_parse_from(["pb-mapper", "admin", "key", "list", "--page-size", "1001",]) .is_err() ); - assert!(Cli::try_parse_from( - ["pb-mapper", "admin", "key", "issue", "--ttl", "1fortnight",] - ) - .is_err()); - assert!(Cli::try_parse_from([ - "pb-mapper", - "server", - "--init-admin-key", - "--use-machine-msg-header-key", - ]) - .is_err()); + assert!( + Cli::try_parse_from(["pb-mapper", "admin", "key", "issue", "--ttl", "1fortnight",]) + .is_err() + ); + assert!( + Cli::try_parse_from([ + "pb-mapper", + "server", + "--init-admin-key", + "--use-machine-msg-header-key", + ]) + .is_err() + ); } #[test] diff --git a/src/common/auth.rs b/src/common/auth.rs index 25098eb..eeaece0 100644 --- a/src/common/auth.rs +++ b/src/common/auth.rs @@ -54,23 +54,23 @@ use std::io::{Read, Write}; #[cfg(unix)] use std::os::unix::fs::PermissionsExt; use std::path::{Path, PathBuf}; -use std::sync::atomic::{AtomicBool, AtomicU64, AtomicU8, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU64, Ordering}; use std::sync::{Arc, Weak}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use parking_lot::{Mutex, RwLock}; use rand::RngExt; -use ring::aead::{Aad, LessSafeKey, Nonce, UnboundKey, AES_256_GCM}; -use ring::hkdf::{Salt, HKDF_SHA256}; +use ring::aead::{AES_256_GCM, Aad, LessSafeKey, Nonce, UnboundKey}; +use ring::hkdf::{HKDF_SHA256, Salt}; use serde::{Deserialize, Serialize}; use subtle::ConstantTimeEq; use tokio::sync::{mpsc, oneshot}; use tokio_util::sync::CancellationToken; use super::checksum::{ + AesKeyType, Credential, ENV_MSG_HEADER_KEY, MACHINE_MSG_HEADER_KEY_PATH, encode_temporary_credential, env_safe_admin_key_error, get_process_credential, - is_env_safe_admin_key, parse_credential, set_process_msg_header_key, AesKeyType, Credential, - ENV_MSG_HEADER_KEY, MACHINE_MSG_HEADER_KEY_PATH, + is_env_safe_admin_key, parse_credential, set_process_msg_header_key, }; /// The namespace administrator connections operate in. Tenant namespaces are the @@ -298,21 +298,17 @@ fn cancelled_lease_failure(is_admin: bool, lease: &AuthLease) -> AuthFailure { ); } match lease.cancel_reason.load(Ordering::Acquire) { - LEASE_CANCEL_EXPIRED => AuthFailure::new( - "temporary_key_expired", - "temporary key has expired", - false, - ), + LEASE_CANCEL_EXPIRED => { + AuthFailure::new("temporary_key_expired", "temporary key has expired", false) + } LEASE_CANCEL_ROTATED => AuthFailure::new( "temporary_key_rotated", "temporary credential was invalidated by administrator root rotation or auth-state reset", false, ), - LEASE_CANCEL_REVOKED => AuthFailure::new( - "temporary_key_revoked", - "temporary key was revoked", - false, - ), + LEASE_CANCEL_REVOKED => { + AuthFailure::new("temporary_key_revoked", "temporary key was revoked", false) + } _ => AuthFailure::new( "temporary_key_inactive", "credential lease has been cancelled", @@ -770,7 +766,7 @@ enum WalRecord { } mod actor; -use actor::{run_auth_actor, AuthActorState}; +use actor::{AuthActorState, run_auth_actor}; mod persistence; pub use persistence::*; pub(in crate::common::auth) use persistence::{ @@ -787,7 +783,7 @@ pub(in crate::common::auth) use persistence::{ prepare_state_dir, read_instance_id_file, try_load_persisted_state, }; mod ids; -pub use ids::{Generation, KeyId, SlotIndex, ADMIN_KEY_ID}; +pub use ids::{ADMIN_KEY_ID, Generation, KeyId, SlotIndex}; mod leases; use leases::Leases; mod timing_wheel; diff --git a/src/common/auth/config.rs b/src/common/auth/config.rs index b2b3b47..e320419 100644 --- a/src/common/auth/config.rs +++ b/src/common/auth/config.rs @@ -51,26 +51,26 @@ pub(crate) fn linux_default_auth_state_dir( if euid == 0 || system_dir_usable { return PathBuf::from(DEFAULT_AUTH_STATE_DIR); } - if let Some(xdg) = xdg_data_home { - if !xdg.is_empty() { - return PathBuf::from(xdg).join("pb-mapper").join("auth"); - } + if let Some(xdg) = xdg_data_home + && !xdg.is_empty() + { + return PathBuf::from(xdg).join("pb-mapper").join("auth"); } - if let Some(home) = home { - if !home.is_empty() { - return PathBuf::from(home) - .join(".local") - .join("share") - .join("pb-mapper") - .join("auth"); - } + if let Some(home) = home + && !home.is_empty() + { + return PathBuf::from(home) + .join(".local") + .join("share") + .join("pb-mapper") + .join("auth"); } PathBuf::from(DEFAULT_AUTH_STATE_DIR) } #[cfg(not(any(windows, target_os = "macos")))] pub(super) fn unix_effective_uid() -> u32 { - extern "C" { + unsafe extern "C" { fn geteuid() -> u32; } unsafe { geteuid() } @@ -88,7 +88,7 @@ fn unix_path_is_writable(path: &Path) -> bool { let Ok(c_path) = std::ffi::CString::new(path.as_os_str().as_bytes()) else { return false; }; - extern "C" { + unsafe extern "C" { fn access(pathname: *const std::os::raw::c_char, mode: i32) -> i32; } const W_OK: i32 = 2; diff --git a/src/common/auth/keys.rs b/src/common/auth/keys.rs index a0592d4..5df9e99 100644 --- a/src/common/auth/keys.rs +++ b/src/common/auth/keys.rs @@ -92,10 +92,10 @@ pub(super) fn recover_admin_key_after_rotation( false, ) })?; - if let Ok(Credential::Admin(current_key)) = parse_credential(current.trim()) { - if open_blob(¤t_key, &bytes).is_ok() { - return Ok(current.to_string()); - } + if let Ok(Credential::Admin(current_key)) = parse_credential(current.trim()) + && open_blob(¤t_key, &bytes).is_ok() + { + return Ok(current.to_string()); } let Some(next) = read_admin_key(&state_dir.join("admin.key.next"))? else { return Ok(current.to_string()); diff --git a/src/common/auth/persistence/fs.rs b/src/common/auth/persistence/fs.rs index a0283d4..4bca35c 100644 --- a/src/common/auth/persistence/fs.rs +++ b/src/common/auth/persistence/fs.rs @@ -42,7 +42,7 @@ pub fn acquire_state_dir_lock(state_dir: &Path) -> Result { fn lock_exclusive_nonblock(file: &File) -> std::io::Result<()> { #[cfg(unix)] { - extern "C" { + unsafe extern "C" { fn flock(fd: i32, operation: i32) -> i32; } const LOCK_EX: i32 = 2; diff --git a/src/common/auth/persistence/wal.rs b/src/common/auth/persistence/wal.rs index e5690eb..d0797e3 100644 --- a/src/common/auth/persistence/wal.rs +++ b/src/common/auth/persistence/wal.rs @@ -9,11 +9,11 @@ pub(in crate::common::auth) fn fail_closed_on_uncertain_wal( inner: &AuthStateInner, result: Result<(), AuthFailure>, ) -> Result<(), AuthFailure> { - if let Err(error) = &result { - if !error.retryable { - inner.safe_mode.store(true, Ordering::Release); - cancel_all_temporary_leases(inner); - } + if let Err(error) = &result + && !error.retryable + { + inner.safe_mode.store(true, Ordering::Release); + cancel_all_temporary_leases(inner); } result } diff --git a/src/common/auth/runtime.rs b/src/common/auth/runtime.rs index c33f011..b5b74d8 100644 --- a/src/common/auth/runtime.rs +++ b/src/common/auth/runtime.rs @@ -62,10 +62,10 @@ impl AuthRuntime { recover_instance_id_after_reset(&config.state_dir, &admin_key, instance_id)?; let (mut loaded, safe_mode) = load_persisted_state(&config, &admin_key, instance_id); let now = unix_seconds(); - if let Some(state) = loaded.as_mut() { - if normalize_tombstone_times(state, now) { - write_snapshot_and_truncate_wal(&config, &admin_key, state)?; - } + if let Some(state) = loaded.as_mut() + && normalize_tombstone_times(state, now) + { + write_snapshot_and_truncate_wal(&config, &admin_key, state)?; } let mut slots = (0..config.max_temporary_keys) .map(|_| SlotHot::default()) diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index 6d82cc6..e5518f1 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -356,12 +356,20 @@ async fn env_recovery_key_is_not_written_when_it_cannot_decrypt_existing_state() }; write_snapshot_and_truncate_wal(&config, &good, &snapshot).unwrap(); set_process_msg_header_key(Some(std::str::from_utf8(&bad).unwrap())).unwrap(); - std::env::set_var(ENV_MSG_HEADER_KEY, std::str::from_utf8(&bad).unwrap()); + // SAFETY: this test holds `PROCESS_CREDENTIAL_TEST_LOCK`, which + // serialises every test that touches the process credential. + unsafe { + std::env::set_var(ENV_MSG_HEADER_KEY, std::str::from_utf8(&bad).unwrap()); + }; let error = match AuthRuntime::from_process(config).await { Ok(_) => panic!("a mismatched recovery key must not start the runtime"), Err(error) => error, }; - std::env::remove_var(ENV_MSG_HEADER_KEY); + // SAFETY: this test holds `PROCESS_CREDENTIAL_TEST_LOCK`, which + // serialises every test that touches the process credential. + unsafe { + std::env::remove_var(ENV_MSG_HEADER_KEY); + }; set_process_msg_header_key(None).unwrap(); assert_eq!(error.code, "administrator_key_invalid"); assert!( @@ -395,9 +403,17 @@ async fn env_recovery_key_is_accepted_for_wal_only_state() { ) .unwrap(); set_process_msg_header_key(Some(std::str::from_utf8(&good).unwrap())).unwrap(); - std::env::set_var(ENV_MSG_HEADER_KEY, std::str::from_utf8(&good).unwrap()); + // SAFETY: this test holds `PROCESS_CREDENTIAL_TEST_LOCK`, which + // serialises every test that touches the process credential. + unsafe { + std::env::set_var(ENV_MSG_HEADER_KEY, std::str::from_utf8(&good).unwrap()); + }; let started = AuthRuntime::from_process(config).await; - std::env::remove_var(ENV_MSG_HEADER_KEY); + // SAFETY: this test holds `PROCESS_CREDENTIAL_TEST_LOCK`, which + // serialises every test that touches the process credential. + unsafe { + std::env::remove_var(ENV_MSG_HEADER_KEY); + }; set_process_msg_header_key(None).unwrap(); started.expect("a matching recovery key must start from WAL-only state"); assert!( @@ -432,12 +448,20 @@ async fn env_recovery_key_is_not_written_when_wal_only_state_does_not_match() { ) .unwrap(); set_process_msg_header_key(Some(std::str::from_utf8(&bad).unwrap())).unwrap(); - std::env::set_var(ENV_MSG_HEADER_KEY, std::str::from_utf8(&bad).unwrap()); + // SAFETY: this test holds `PROCESS_CREDENTIAL_TEST_LOCK`, which + // serialises every test that touches the process credential. + unsafe { + std::env::set_var(ENV_MSG_HEADER_KEY, std::str::from_utf8(&bad).unwrap()); + }; let error = match AuthRuntime::from_process(config).await { Ok(_) => panic!("a mismatched recovery key must not start from WAL-only state"), Err(error) => error, }; - std::env::remove_var(ENV_MSG_HEADER_KEY); + // SAFETY: this test holds `PROCESS_CREDENTIAL_TEST_LOCK`, which + // serialises every test that touches the process credential. + unsafe { + std::env::remove_var(ENV_MSG_HEADER_KEY); + }; set_process_msg_header_key(None).unwrap(); assert_eq!(error.code, "administrator_key_invalid"); assert!( @@ -1412,13 +1436,15 @@ async fn revoking_keeps_the_row_until_its_retention_elapses() { .code, "temporary_key_revoked" ); - assert!(runtime - .list(&admin, 0, 100) - .await - .unwrap() - .items - .iter() - .any(|item| item.key_id == key_id && item.state == "revoked")); + assert!( + runtime + .list(&admin, 0, 100) + .await + .unwrap() + .items + .iter() + .any(|item| item.key_id == key_id && item.state == "revoked") + ); drop(runtime); tokio::time::sleep(Duration::from_millis(20)).await; diff --git a/src/common/checksum.rs b/src/common/checksum.rs index 6efe4b2..bf57883 100644 --- a/src/common/checksum.rs +++ b/src/common/checksum.rs @@ -3,15 +3,15 @@ use std::io; use std::os::unix::fs::PermissionsExt; use std::path::Path; use std::process::Command; -use std::sync::atomic::{AtomicU32, Ordering}; use std::sync::LazyLock; +use std::sync::atomic::{AtomicU32, Ordering}; use parking_lot::RwLock; -use base64::engine::general_purpose::URL_SAFE_NO_PAD; use base64::Engine; +use base64::engine::general_purpose::URL_SAFE_NO_PAD; use rand::RngExt; -use ring::digest::{digest, SHA256}; +use ring::digest::{SHA256, digest}; use super::message::DataLenType; @@ -145,14 +145,21 @@ pub fn get_msg_header_key() -> Result, String> { pub fn set_process_msg_header_key(msg_header_key: Option<&str>) -> Result<(), String> { let normalized = msg_header_key.map(str::trim).unwrap_or(""); if normalized.is_empty() { - std::env::remove_var(ENV_MSG_HEADER_KEY); + // SAFETY: edition 2024 makes these unsafe because the environment is + // process-global and another thread reading it concurrently is a data + // race. Here the environment is only a mirror for child processes and + // for the initial read at startup; the credential every operation + // actually consults is `MSG_HEADER_KEY_STATE`, behind an `RwLock`, and + // it is updated immediately below. + unsafe { std::env::remove_var(ENV_MSG_HEADER_KEY) }; update_runtime_credential(None); return Ok(()); } let credential = parse_credential(normalized)?; - std::env::set_var(ENV_MSG_HEADER_KEY, normalized); + // SAFETY: as above. + unsafe { std::env::set_var(ENV_MSG_HEADER_KEY, normalized) }; update_runtime_credential(Some(credential)); Ok(()) } @@ -240,18 +247,18 @@ fn get_machine_hostname() -> io::Result { return Ok(hostname); } - if let Ok(content) = std::fs::read_to_string("/etc/hostname") { - if let Some(hostname) = normalize_non_empty(Some(content.as_str())) { - return Ok(hostname); - } + if let Ok(content) = std::fs::read_to_string("/etc/hostname") + && let Some(hostname) = normalize_non_empty(Some(content.as_str())) + { + return Ok(hostname); } - if let Ok(output) = Command::new("hostname").output() { - if output.status.success() { - let stdout = String::from_utf8_lossy(&output.stdout); - if let Some(hostname) = normalize_non_empty(Some(stdout.as_ref())) { - return Ok(hostname); - } + if let Ok(output) = Command::new("hostname").output() + && output.status.success() + { + let stdout = String::from_utf8_lossy(&output.stdout); + if let Some(hostname) = normalize_non_empty(Some(stdout.as_ref())) { + return Ok(hostname); } } @@ -272,22 +279,22 @@ fn normalize_non_empty(input: Option<&str>) -> Option { } fn get_machine_mac_addresses() -> io::Result> { - if let Ok(mac_addresses) = get_machine_mac_addresses_from_sysfs() { - if !mac_addresses.is_empty() { - return Ok(mac_addresses); - } + if let Ok(mac_addresses) = get_machine_mac_addresses_from_sysfs() + && !mac_addresses.is_empty() + { + return Ok(mac_addresses); } - if let Ok(mac_addresses) = get_machine_mac_addresses_from_ip_link() { - if !mac_addresses.is_empty() { - return Ok(mac_addresses); - } + if let Ok(mac_addresses) = get_machine_mac_addresses_from_ip_link() + && !mac_addresses.is_empty() + { + return Ok(mac_addresses); } - if let Ok(mac_addresses) = get_machine_mac_addresses_from_ifconfig() { - if !mac_addresses.is_empty() { - return Ok(mac_addresses); - } + if let Ok(mac_addresses) = get_machine_mac_addresses_from_ifconfig() + && !mac_addresses.is_empty() + { + return Ok(mac_addresses); } Err(io::Error::new( diff --git a/src/common/config.rs b/src/common/config.rs index 0c0fcce..d4e3aaf 100644 --- a/src/common/config.rs +++ b/src/common/config.rs @@ -5,7 +5,7 @@ use std::time::Duration; use clap::ValueEnum; use snafu::ResultExt; use tracing_subscriber::layer::SubscriberExt; -use tracing_subscriber::{fmt, EnvFilter, Layer}; +use tracing_subscriber::{EnvFilter, Layer, fmt}; use super::error::{CfgPbServerEnvNotExistSnafu, Result}; @@ -410,17 +410,27 @@ mod tests { fn keep_alive_reads_the_environment_every_time() { let restore = std::env::var(PB_MAPPER_KEEP_ALIVE).ok(); - std::env::remove_var(PB_MAPPER_KEEP_ALIVE); + // SAFETY: mutating the environment is unsafe in edition 2024 because + // it is process-global. This is the only test that touches + // `PB_MAPPER_KEEP_ALIVE` — which is why it is one test and not several + // — and it restores the original value before returning. + unsafe { + std::env::remove_var(PB_MAPPER_KEEP_ALIVE); + } assert!(!keep_alive_from_env(), "absent means off"); - std::env::set_var(PB_MAPPER_KEEP_ALIVE, "ON"); + unsafe { + std::env::set_var(PB_MAPPER_KEEP_ALIVE, "ON"); + } assert!(keep_alive_from_env(), "the documented spelling"); // The regression. This used to be a `LazyLock`, so the answer was // whatever the first caller in the process saw and could never change — // which is why the UI's per-service toggle did nothing after the first // tunnel started. - std::env::set_var(PB_MAPPER_KEEP_ALIVE, "OFF"); + unsafe { + std::env::set_var(PB_MAPPER_KEEP_ALIVE, "OFF"); + } assert!( !keep_alive_from_env(), "OFF must mean off; the old check was `is_ok()`, so any value at \ @@ -428,17 +438,23 @@ mod tests { ); for truthy in ["on", "1", "true", "yes", " ON "] { - std::env::set_var(PB_MAPPER_KEEP_ALIVE, truthy); + unsafe { + std::env::set_var(PB_MAPPER_KEEP_ALIVE, truthy); + } assert!(keep_alive_from_env(), "{truthy:?} should enable"); } for falsy in ["", "off", "0", "false", "no"] { - std::env::set_var(PB_MAPPER_KEEP_ALIVE, falsy); + unsafe { + std::env::set_var(PB_MAPPER_KEEP_ALIVE, falsy); + } assert!(!keep_alive_from_env(), "{falsy:?} should not enable"); } - match restore { - Some(value) => std::env::set_var(PB_MAPPER_KEEP_ALIVE, value), - None => std::env::remove_var(PB_MAPPER_KEEP_ALIVE), + unsafe { + match restore { + Some(value) => std::env::set_var(PB_MAPPER_KEEP_ALIVE, value), + None => std::env::remove_var(PB_MAPPER_KEEP_ALIVE), + } } } } diff --git a/src/common/manager.rs b/src/common/manager.rs index c6348d6..0c325cf 100644 --- a/src/common/manager.rs +++ b/src/common/manager.rs @@ -31,11 +31,11 @@ pub struct TaskManager, - > TaskManager + MangerChanType, + ConnChanType, + ConnIdType: ConnIdTrait, + ConnIdProviderType: ConnIdProvider, +> TaskManager { pub fn new( conn_id_provider: ConnIdProviderType, diff --git a/src/common/message/mod.rs b/src/common/message/mod.rs index 87ae650..6c20a0c 100644 --- a/src/common/message/mod.rs +++ b/src/common/message/mod.rs @@ -3,13 +3,13 @@ pub mod command; pub mod forward; pub mod secure; -use snafu::{ensure, ResultExt}; +use snafu::{ResultExt, ensure}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use super::buffer::{BufferGetter, CommonBuffer, FixedSizeBuffer}; use super::checksum::{ - get_checksum, get_checksum_for_key, get_msg_header_key, process_checksum_is_ready, - valid_checksum, valid_checksum_for_key, AesKeyType, + AesKeyType, get_checksum, get_checksum_for_key, get_msg_header_key, process_checksum_is_ready, + valid_checksum, valid_checksum_for_key, }; use super::error::{ self, MsgDatalenValidateSnafu, MsgNetworkReadBodySnafu, MsgNetworkReadCheckSumSnafu, diff --git a/src/common/message/secure.rs b/src/common/message/secure.rs index b96595d..de73081 100644 --- a/src/common/message/secure.rs +++ b/src/common/message/secure.rs @@ -23,19 +23,19 @@ use parking_lot::Mutex; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use rand::RngExt; -use ring::aead::{Aad, LessSafeKey, Nonce, UnboundKey, AES_256_GCM}; -use ring::digest::{digest, SHA256}; -use ring::hkdf::{Salt, HKDF_SHA256}; +use ring::aead::{AES_256_GCM, Aad, LessSafeKey, Nonce, UnboundKey}; +use ring::digest::{SHA256, digest}; +use ring::hkdf::{HKDF_SHA256, Salt}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use super::{ - CodecMessageReader, CodecMessageWriter, DataLenType, MessageReader, MessageWriter, MAX_MSG_LEN, + CodecMessageReader, CodecMessageWriter, DataLenType, MAX_MSG_LEN, MessageReader, MessageWriter, }; use crate::common::auth::{ - AuthContext, AuthFailure, AuthRuntime, KeyId, LegacyConnectionGuard, ADMIN_KEY_ID, + ADMIN_KEY_ID, AuthContext, AuthFailure, AuthRuntime, KeyId, LegacyConnectionGuard, }; use crate::common::checksum::{ - get_process_credential, valid_checksum_for_key, AesKeyType, Credential, + AesKeyType, Credential, get_process_credential, valid_checksum_for_key, }; use crate::common::error::{Error, Result}; use crate::utils::codec::{Aes256GcmDeCodec, Aes256GcmEnCodec, Decryptor}; @@ -160,7 +160,7 @@ impl ClientHeaderSession { Err(_) => { return Err(protocol_error(format!( "timed out writing first-flight request after {timeout:?}" - ))) + ))); } } let mut reader = self.response_reader(stream)?; @@ -169,7 +169,7 @@ impl ClientHeaderSession { Err(_) => { return Err(protocol_error(format!( "timed out reading first-flight response after {timeout:?}" - ))) + ))); } }; Ok(message.to_vec()) @@ -575,7 +575,7 @@ impl ServerSecurity { error.to_string(), false, key_id, - )) + )); } } } @@ -637,12 +637,12 @@ impl MessageWriter for HeaderMessageWriter<'_, T> { } mod frame; -use frame::{derive_material, first_prefix, open_v2_payload, read_v2_frame, V2Material}; +use frame::{V2Material, derive_material, first_prefix, open_v2_payload, read_v2_frame}; pub use frame::{V2MessageReader, V2MessageWriter}; mod replay; #[cfg(test)] use replay::RotatingBloom; -use replay::{replay_fingerprint, FirstFlightAdmit, ReplayGuard}; +use replay::{FirstFlightAdmit, ReplayGuard, replay_fingerprint}; mod first_flight; use first_flight::*; fn legacy_message_reader<'a, T: AsyncReadExt + Unpin>( diff --git a/src/common/message/secure/replay.rs b/src/common/message/secure/replay.rs index 54727b4..1596966 100644 --- a/src/common/message/secure/replay.rs +++ b/src/common/message/secure/replay.rs @@ -247,11 +247,9 @@ impl ReplayGuard { } return Err(error); } - if created { - if let Err(error) = crate::common::auth::sync_parent_directory(path) { - self.log_failed = true; - return Err(std::io::Error::other(error.to_string())); - } + if created && let Err(error) = crate::common::auth::sync_parent_directory(path) { + self.log_failed = true; + return Err(std::io::Error::other(error.to_string())); } Ok(()) } diff --git a/src/common/message/secure/tests.rs b/src/common/message/secure/tests.rs index a1b3762..e957918 100644 --- a/src/common/message/secure/tests.rs +++ b/src/common/message/secure/tests.rs @@ -132,13 +132,15 @@ async fn identical_initial_frames_are_admitted_only_once() { .code, "connection_salt_replayed" ); - assert!(results - .iter() - .filter_map(|result| result.as_ref().err()) - .next() - .unwrap() - .response_session - .is_none()); + assert!( + results + .iter() + .filter_map(|result| result.as_ref().err()) + .next() + .unwrap() + .response_session + .is_none() + ); let _ = std::fs::remove_dir_all(config.state_dir); } diff --git a/src/local/client/mod.rs b/src/local/client/mod.rs index 0c0fff7..d10c764 100644 --- a/src/local/client/mod.rs +++ b/src/local/client/mod.rs @@ -15,16 +15,16 @@ use uni_stream::udp::set_custom_timeout; use self::error::{AcceptLocalStreamSnafu, BindLocalListenerSnafu}; use self::status::{get_status, get_status_scoped, get_status_with_credential}; use self::stream::handle_local_stream; -use crate::common::checksum::{get_process_credential, Credential}; +use crate::common::checksum::{Credential, get_process_credential}; use crate::common::config::{ - client_health_check_interval, client_health_check_timeout, client_health_failure_threshold, - StatusOp, + StatusOp, client_health_check_interval, client_health_check_timeout, + client_health_failure_threshold, }; use crate::common::message::command::{PbConnStatusReq, PbConnStatusResp}; use crate::common::message::forward::StreamForward; use crate::snafu_error_get_or_return; use crate::utils::timeout::RetryBackoff; -use uni_stream::addr::{each_addr, ToSocketAddrs}; +use uni_stream::addr::{ToSocketAddrs, each_addr}; use uni_stream::stream::got_one_socket_addr; use uni_stream::stream::{ListenerProvider, StreamAccept}; diff --git a/src/local/client/status.rs b/src/local/client/status.rs index 1760c35..7faee86 100644 --- a/src/local/client/status.rs +++ b/src/local/client/status.rs @@ -81,7 +81,8 @@ mod tests { #[tokio::test] async fn get_status_times_out_when_peer_stalls_after_request() { - std::env::set_var("PB_MAPPER_CONTROL_IO_TIMEOUT", "20ms"); + // SAFETY: no other thread in this test reads the environment. + unsafe { std::env::set_var("PB_MAPPER_CONTROL_IO_TIMEOUT", "20ms") }; let (mut client, _server) = tokio::io::duplex(1024); let result = tokio::time::timeout( @@ -91,7 +92,8 @@ mod tests { .await .expect("get_status ignored PB_MAPPER_CONTROL_IO_TIMEOUT"); - std::env::remove_var("PB_MAPPER_CONTROL_IO_TIMEOUT"); + // SAFETY: as above. + unsafe { std::env::remove_var("PB_MAPPER_CONTROL_IO_TIMEOUT") }; assert!(result.is_err()); } } diff --git a/src/local/client/stream.rs b/src/local/client/stream.rs index 539aff4..63ab7f7 100644 --- a/src/local/client/stream.rs +++ b/src/local/client/stream.rs @@ -16,8 +16,8 @@ use crate::common::message::forward::StreamForward; use crate::common::message::secure::ClientHeaderSession; use crate::local::client::error::CreateHeaderToolSnafu; use crate::snafu_error_handle; -use uni_stream::addr::{each_addr, ToSocketAddrs}; -use uni_stream::stream::{set_tcp_keep_alive, set_tcp_nodelay, NetworkStream}; +use uni_stream::addr::{ToSocketAddrs, each_addr}; +use uni_stream::stream::{NetworkStream, set_tcp_keep_alive, set_tcp_nodelay}; #[instrument(skip(local_stream))] pub async fn handle_local_stream< diff --git a/src/local/server/mod.rs b/src/local/server/mod.rs index 13ca2c3..68fcd6c 100644 --- a/src/local/server/mod.rs +++ b/src/local/server/mod.rs @@ -17,15 +17,15 @@ use self::error::{ EncodeRegisterReqSnafu, EncodeStreamAckMsgSnafu, ReadRegisterRespSnafu, ReadStreamReqSnafu, RegisterRespNotMatchSnafu, SendRegisterReqSnafu, WritePingMsgSnafu, WriteStreamAckMsgSnafu, }; -use self::stream::{handle_stream, StreamConnect}; -use crate::common::checksum::{get_process_credential, Credential}; +use self::stream::{StreamConnect, handle_stream}; +use crate::common::checksum::{Credential, get_process_credential}; use crate::common::config::{ control_conn_pool_size, control_heartbeat_interval, control_heartbeat_tolerance, control_io_timeout, control_suspect_grace, registration_probe_timeout, }; use crate::common::message::command::{ - LocalServer, MessageSerializer, PbConnRequest, PbConnResponse, PbConnStatusReq, - PbConnStatusResp, PbServerRequest, CONTROL_PROTOCOL_V2, + CONTROL_PROTOCOL_V2, LocalServer, MessageSerializer, PbConnRequest, PbConnResponse, + PbConnStatusReq, PbConnStatusResp, PbServerRequest, }; use crate::common::message::forward::StreamForward; use crate::common::message::secure::ClientHeaderSession; @@ -35,9 +35,9 @@ use crate::{ snafu_error_get_or_continue, snafu_error_get_or_return, snafu_error_get_or_return_ok, snafu_error_handle, }; -use uni_stream::addr::{each_addr, ToSocketAddrs}; +use uni_stream::addr::{ToSocketAddrs, each_addr}; use uni_stream::stream::{ - got_one_socket_addr, set_tcp_keep_alive, set_tcp_nodelay, StreamProvider, + StreamProvider, got_one_socket_addr, set_tcp_keep_alive, set_tcp_nodelay, }; fn get_ping_message(protocol_version: u16, seq: u64) -> error::Result> { @@ -208,7 +208,7 @@ async fn probe_remote_registration( Err(_) => { return RegistrationProbeResult::Failed(format!( "status probe timed out after {timeout:?}" - )) + )); } }; @@ -553,11 +553,13 @@ where let msg = snafu_error_get_or_return_ok!(request.encode().context(EncodeRegisterReqSnafu)); match tokio::time::timeout(timeout, session.write_initial(&mut manager_stream, &msg)).await { Ok(result) => snafu_error_get_or_return_ok!(result.context(SendRegisterReqSnafu)), - Err(_) => snafu_error_get_or_return_ok!(ControlIoTimeoutSnafu { - action: "send register request", - timeout, - } - .fail()), + Err(_) => snafu_error_get_or_return_ok!( + ControlIoTimeoutSnafu { + action: "send register request", + timeout, + } + .fail() + ), } let (mut reader, mut writer) = manager_stream.into_split(); let mut msg_reader = match session.response_reader(&mut reader) { @@ -572,11 +574,13 @@ where let timeout = control_io_timeout(); let msg = match tokio::time::timeout(timeout, msg_reader.read_msg()).await { Ok(result) => snafu_error_get_or_return_ok!(result.context(ReadRegisterRespSnafu)), - Err(_) => snafu_error_get_or_return_ok!(ControlIoTimeoutSnafu { - action: "read register response", - timeout, - } - .fail()), + Err(_) => snafu_error_get_or_return_ok!( + ControlIoTimeoutSnafu { + action: "read register response", + timeout, + } + .fail() + ), }; let resp = snafu_error_get_or_return_ok!( PbConnResponse::decode(msg).context(DecodeRegisterRespSnafu) diff --git a/src/local/server/stream.rs b/src/local/server/stream.rs index e843af5..52ade70 100644 --- a/src/local/server/stream.rs +++ b/src/local/server/stream.rs @@ -17,8 +17,8 @@ use crate::common::message::forward::StreamForward; use crate::common::message::secure::ClientHeaderSession; use crate::local::server::error::CreateHeaderToolSnafu; use crate::snafu_error_handle; -use uni_stream::addr::{each_addr, ToSocketAddrs}; -use uni_stream::stream::{set_tcp_keep_alive, set_tcp_nodelay, StreamProvider, StreamSplit}; +use uni_stream::addr::{ToSocketAddrs, each_addr}; +use uni_stream::stream::{StreamProvider, StreamSplit, set_tcp_keep_alive, set_tcp_nodelay}; #[derive(Clone, Copy, Debug)] pub struct StreamConnect { diff --git a/src/pb_server/admin.rs b/src/pb_server/admin.rs index 9773540..e135792 100644 --- a/src/pb_server/admin.rs +++ b/src/pb_server/admin.rs @@ -16,13 +16,13 @@ use tokio::net::TcpStream; use super::error::Error; use super::{ManagerTask, ManagerTaskSender, Result}; use crate::common::auth::{AuthContext, AuthFailure, AuthRuntime, KeyId}; -use crate::common::checksum::{parse_credential, Credential}; +use crate::common::checksum::{Credential, parse_credential}; use crate::common::conn_id::RemoteConnId; +use crate::common::message::MessageWriter; use crate::common::message::command::{ AdminRequest, AdminResponse, MessageSerializer, PbConnResponse, }; use crate::common::message::secure::ServerHeaderSession; -use crate::common::message::MessageWriter; pub async fn handle_admin_request( request: AdminRequest, diff --git a/src/pb_server/client.rs b/src/pb_server/client.rs index 3b5dc39..df3e762 100644 --- a/src/pb_server/client.rs +++ b/src/pb_server/client.rs @@ -3,7 +3,7 @@ use std::time::Duration; use snafu::ResultExt; use tokio::net::TcpStream; -use tokio::time::{timeout, Instant}; +use tokio::time::{Instant, timeout}; use tracing::instrument; use super::error::{ @@ -13,17 +13,17 @@ use super::error::{ ClientConnSubcribeRespNotMatchSnafu, ClientConnWriteSubcribeRespSnafu, }; use super::{ConnTask, ImutableKey, ManagerTask, ManagerTaskSender, Result}; -use crate::common::checksum::{gen_random_key, AesKeyType}; +use crate::common::checksum::{AesKeyType, gen_random_key}; use crate::common::config::{stream_ack_timeout, stream_ready_timeout, stream_recovery_timeout}; use crate::common::conn_id::RemoteConnId; use crate::common::message::command::{MessageSerializer, PbConnResponse}; use crate::common::message::forward::{ - start_datagram_forward, start_forward, CodecDatagramReader, CodecDatagramWriter, - CodecForwardReader, CodecForwardWriter, NormalDatagramReader, NormalDatagramWriter, - NormalForwardReader, NormalForwardWriter, + CodecDatagramReader, CodecDatagramWriter, CodecForwardReader, CodecForwardWriter, + NormalDatagramReader, NormalDatagramWriter, NormalForwardReader, NormalForwardWriter, + start_datagram_forward, start_forward, }; use crate::common::message::secure::ServerHeaderSession; -use crate::common::message::{get_decodec, get_encodec, MessageWriter}; +use crate::common::message::{MessageWriter, get_decodec, get_encodec}; use crate::pb_server::error::{ ClientConnCreateHeaderToolSnafu, ClientConnEncodeStreamRespSnafu, ClientConnWriteStreamRespSnafu, diff --git a/src/pb_server/error.rs b/src/pb_server/error.rs index 07ac55f..e131712 100644 --- a/src/pb_server/error.rs +++ b/src/pb_server/error.rs @@ -119,9 +119,7 @@ pub enum Error { conn_id: RemoteConnId, source: common::error::Error, }, - #[snafu(display( - "server conn write register resp error with `key:{key}` `conn_id:{conn_id}`" - ))] + #[snafu(display("server conn write register resp error with `key:{key}` `conn_id:{conn_id}`"))] ServerConnWriteRegisteredOk { key: Arc, conn_id: RemoteConnId, @@ -262,9 +260,7 @@ pub enum Error { conn_id: RemoteConnId, source: common::error::Error, }, - #[snafu(display( - "client conn write subcribe resp error with `key:{key}` `conn_id:{conn_id}`" - ))] + #[snafu(display("client conn write subcribe resp error with `key:{key}` `conn_id:{conn_id}`"))] ClientConnWriteSubcribeResp { key: Arc, conn_id: RemoteConnId, diff --git a/src/pb_server/mod.rs b/src/pb_server/mod.rs index 40835fc..fbaf9cf 100644 --- a/src/pb_server/mod.rs +++ b/src/pb_server/mod.rs @@ -32,19 +32,19 @@ use self::error::{ TaskCenterInitRequestTimeoutSnafu, TaskCenterSendListenerSnafu, TaskCenterSendStatusRespSnafu, TaskCenterSendStreamRespToManagerSnafu, TaskCenterSetKeepAliveSnafu, }; -use self::server::{handle_server_conn, ServerRegistration}; +use self::server::{ServerRegistration, handle_server_conn}; use self::status::handle_show_status; -use crate::common::auth::{AuthConfig, AuthContext, AuthRuntime, ADMIN_KEY_ID}; +use crate::common::auth::{ADMIN_KEY_ID, AuthConfig, AuthContext, AuthRuntime}; use crate::common::config::{control_io_timeout, keep_alive_from_env, server_lease_timeout}; use crate::common::conn_id::{ConnIdProvider, RemoteConnId}; use crate::common::manager::{ForwardMessage, SenderChan, TaskManager}; +use crate::common::message::MessageWriter; use crate::common::message::command::{ AdminConnectionInfo, AdminConnectionPage, AdminServiceInfo, AdminServicePage, MessageSerializer, PbConnRequest, PbConnResponse, PbConnStatusReq, PbConnStatusResp, PbServiceConnStatus, }; use crate::common::message::secure::{HeaderProtocol, ServerHeaderSession, ServerSecurity}; -use crate::common::message::MessageWriter; use crate::pb_server::error::{ ServerListenSnafu, TaskCenterClientSendStreamSnafu, TaskCenterSendRegisterRespSnafu, TaskCenterSendStreamRespToClientSnafu, TaskCenterSendSubcribeRespSnafu, @@ -249,14 +249,14 @@ fn remove_server_conn( key: &ImutableKey, conn_id: RemoteConnId, ) -> bool { - if let Some(ids) = server_conn_map.get_mut(key) { - if let Some(idx) = ids.iter().position(|info| info.conn_id == conn_id) { - ids.remove(idx); - if ids.is_empty() { - server_conn_map.remove(key); - } - return true; + if let Some(ids) = server_conn_map.get_mut(key) + && let Some(idx) = ids.iter().position(|info| info.conn_id == conn_id) + { + ids.remove(idx); + if ids.is_empty() { + server_conn_map.remove(key); } + return true; } false } diff --git a/src/pb_server/runtime.rs b/src/pb_server/runtime.rs index 8345365..c5d4a99 100644 --- a/src/pb_server/runtime.rs +++ b/src/pb_server/runtime.rs @@ -335,11 +335,13 @@ pub async fn run_server_on_listener( }) } }; - snafu_error_get_or_continue!(conn_sender - .send(ConnTask::StatusResp(resp)) - .await - .map_err(|_| kanal::SendError(())) - .context(TaskCenterSendStatusRespSnafu { conn_id })); + snafu_error_get_or_continue!( + conn_sender + .send(ConnTask::StatusResp(resp)) + .await + .map_err(|_| kanal::SendError(())) + .context(TaskCenterSendStatusRespSnafu { conn_id }) + ); } ManagerTask::Accept { stream, peer_addr } => { let conn_id = manager.get_conn_id( @@ -588,15 +590,17 @@ pub async fn run_server_on_listener( idle_connections = manager.idle_conn_count(), "server connection registered" ); - snafu_error_get_or_continue!(conn_sender - .send(ConnTask::RegisterResp { - generation, - protocol_version, - lease_ttl_ms: server_lease_timeout().as_millis() as u64, - }) - .await - .map_err(|_| kanal::SendError(())) - .context(TaskCenterSendRegisterRespSnafu { key, conn_id })); + snafu_error_get_or_continue!( + conn_sender + .send(ConnTask::RegisterResp { + generation, + protocol_version, + lease_ttl_ms: server_lease_timeout().as_millis() as u64, + }) + .await + .map_err(|_| kanal::SendError(())) + .context(TaskCenterSendRegisterRespSnafu { key, conn_id }) + ); } ManagerTask::Stream { key, @@ -660,19 +664,23 @@ pub async fn run_server_on_listener( active_connections = manager.active_conn_count(), "server stream ready for client" ); - let client_sender = snafu_error_get_or_continue!(manager - .get_conn_sender_chan(&client_id) - .context(TaskCenterStreamConnIdNotExistSnafu { conn_id: client_id })); - snafu_error_handle!(client_sender - .send(ConnTask::StreamResp { - server_id, - server_generation: expected_generation, - stream, - session, - }) - .await - .map_err(|_| kanal::SendError(())) - .context(TaskCenterSendStreamRespToClientSnafu { conn_id: client_id })); + let client_sender = snafu_error_get_or_continue!( + manager + .get_conn_sender_chan(&client_id) + .context(TaskCenterStreamConnIdNotExistSnafu { conn_id: client_id }) + ); + snafu_error_handle!( + client_sender + .send(ConnTask::StreamResp { + server_id, + server_generation: expected_generation, + stream, + session, + }) + .await + .map_err(|_| kanal::SendError(())) + .context(TaskCenterSendStreamRespToClientSnafu { conn_id: client_id }) + ); } ManagerTask::StreamAck { server_id, @@ -714,17 +722,21 @@ pub async fn run_server_on_listener( { info.health = ServerConnHealth::Healthy; } - let client_sender = snafu_error_get_or_continue!(manager - .get_conn_sender_chan(&client_id) - .context(TaskCenterStreamConnIdNotExistSnafu { conn_id: client_id })); - snafu_error_handle!(client_sender - .send(ConnTask::StreamAck { - server_id, - server_generation, - }) - .await - .map_err(|_| kanal::SendError(())) - .context(TaskCenterSendStreamRespToClientSnafu { conn_id: client_id })); + let client_sender = snafu_error_get_or_continue!( + manager + .get_conn_sender_chan(&client_id) + .context(TaskCenterStreamConnIdNotExistSnafu { conn_id: client_id }) + ); + snafu_error_handle!( + client_sender + .send(ConnTask::StreamAck { + server_id, + server_generation, + }) + .await + .map_err(|_| kanal::SendError(())) + .context(TaskCenterSendStreamRespToClientSnafu { conn_id: client_id }) + ); } ManagerTask::Subcribe { key, diff --git a/src/pb_server/server.rs b/src/pb_server/server.rs index b379057..e5631ee 100644 --- a/src/pb_server/server.rs +++ b/src/pb_server/server.rs @@ -15,7 +15,7 @@ use super::{ConnTask, ImutableKey, ManagerTask, ManagerTaskSender, Result}; use crate::common::config::server_lease_timeout; use crate::common::conn_id::RemoteConnId; use crate::common::message::command::{ - LocalServer, MessageSerializer, PbConnResponse, PbServerRequest, CONTROL_PROTOCOL_V2, + CONTROL_PROTOCOL_V2, LocalServer, MessageSerializer, PbConnResponse, PbServerRequest, }; use crate::common::message::secure::ServerHeaderSession; use crate::common::message::{MessageReader, MessageWriter}; @@ -582,7 +582,7 @@ mod tests { use crate::common::conn_id::RemoteConnId; use crate::pb_server::ManagerTask; - use super::{ServerConnGuard, SERVER_TIMEOUT}; + use super::{SERVER_TIMEOUT, ServerConnGuard}; #[test] fn server_timeout_has_slack_over_local_server_ping_interval() { diff --git a/src/pb_server/status.rs b/src/pb_server/status.rs index 77c48ef..7c3cca1 100644 --- a/src/pb_server/status.rs +++ b/src/pb_server/status.rs @@ -8,9 +8,9 @@ use super::error::{ }; use super::{ConnTask, ManagerTask, ManagerTaskSender}; use crate::common::conn_id::RemoteConnId; +use crate::common::message::MessageWriter; use crate::common::message::command::{MessageSerializer, PbConnStatusReq}; use crate::common::message::secure::ServerHeaderSession; -use crate::common::message::MessageWriter; struct StatusConnGuard { conn_id: RemoteConnId, diff --git a/src/utils/addr.rs b/src/utils/addr.rs index 8e5ce8a..618f151 100644 --- a/src/utils/addr.rs +++ b/src/utils/addr.rs @@ -2,7 +2,7 @@ use std::future::{self, Future}; use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6}; use std::pin::Pin; use std::sync::LazyLock; -use std::task::{ready, Context, Poll}; +use std::task::{Context, Poll, ready}; use hickory_resolver::config::{NameServerConfig, ResolverConfig, ResolverOpts}; use hickory_resolver::net::runtime::TokioRuntimeProvider; diff --git a/src/utils/codec.rs b/src/utils/codec.rs index 438f3d4..c909e3f 100644 --- a/src/utils/codec.rs +++ b/src/utils/codec.rs @@ -1,8 +1,8 @@ use std::mem::size_of; use ring::aead::{ - Aad, BoundKey, Nonce, NonceSequence, OpeningKey, SealingKey, Tag, UnboundKey, AES_256_GCM, - NONCE_LEN, + AES_256_GCM, Aad, BoundKey, NONCE_LEN, Nonce, NonceSequence, OpeningKey, SealingKey, Tag, + UnboundKey, }; #[derive(Clone, Copy, Default)] @@ -191,7 +191,9 @@ mod tests { #[test] fn test_codec() { const TEST_KEY: [u8; 32] = [0x42; 32]; - let data = String::from("fdafas反对fdasfasfasfsdafdasfsdfasd范德萨发顺🤣❤️😁😍👍👍丰十大大师傅士大夫大撒发射点发士大夫大师傅大师傅士大夫士大夫阿斯蒂芬大师傅阿斯顿法大师傅看叫阿三的发就可是大家发开始打客服开始大幅喀什的开发点卡收费就开始打客服就是的咖啡肯定撒法开始打客服就是的咖啡就开始大幅扣税的急啊看发叫阿三的发生的开发就是大家可是大家发看大数据开发大数据开发大家ask发就是的咖啡的萨芬就卡死的房价开始打家开发商的JFK上的飞机卡上的纠纷开始打飞机宽带技术开发就开始大家开发建设的卡JFK大数据风控静安寺的看法角度看萨芬卡上的纠纷看静安寺的看法角度思考积分可是大家发卡是大家看法就大肆砍伐尽快打算减肥肯定是积分开始大幅技术大咖积分开始打飞机扣税的急啊看发的技术开发就是JFK十大福克斯大家开发大撒发射点幅度萨芬撒旦发发收范德萨发顺丰士大夫十大阿斯蒂芬大师傅阿斯顿附件是的客服对接撒巨大石块积分的课时费阿斯蒂芬法大师傅大师傅十大法大师傅阿斯蒂芬阿斯顿法大师傅阿斯蒂芬大师傅阿斯顿法大师傅大师傅阿斯蒂芬阿斯蒂芬士大夫阿斯蒂芬大师傅的萨芬打算减肥上岛咖啡加快速度大数据开发就是打客服看大数据开发就开始减肥卡萨丁JFK是大家看法加快速度JFK技术大咖积分喀什的开发独守空房技术大咖积分空手道解放扣税的开发商的开发接口是大家看法角度看是否扣税的急啊看发生的开发的快速减肥开始大幅就是打客服卡上的纠纷啊撒旦解放扣税的急啊看发加快速度点卡JFK啥的但是法大师傅技术大咖积分卡萨丁就反馈是大家看法啊是大家看法卡上的纠纷可是大家发喀什的开发大卡司喀什的开发就是打客服法大师傅士大夫的式咖啡机上岛咖啡就是的咖啡艰苦大师傅看上雕刻技法喀什的开发上岛咖啡就喀什的开发就是打客服卡上的纠纷技术的咖啡机肯定撒开发啊十大科技开发速度加啊反馈就是的咖啡开始大幅大师傅似的十大放假啊上岛咖啡就可是大家发空间的是否撒旦士大夫的撒娇开发是大家看法大肆砍伐就喀什的开发氨基酸的考虑非军事对抗疗法金克拉撒旦发艰苦拉萨的飞机喀什打开发就可是大家发可是大家看附件卡上的纠纷卡刷点卡技术的咖啡机可是大家发卡是大家看法静安寺的看法就可是大家发卡萨丁就开发商的急啊看飞机迪斯科发技术的咖啡机可是大家发看电视剧开发商大开始打到发大水发大水"); + let data = String::from( + "fdafas反对fdasfasfasfsdafdasfsdfasd范德萨发顺🤣❤️😁😍👍👍丰十大大师傅士大夫大撒发射点发士大夫大师傅大师傅士大夫士大夫阿斯蒂芬大师傅阿斯顿法大师傅看叫阿三的发就可是大家发开始打客服开始大幅喀什的开发点卡收费就开始打客服就是的咖啡肯定撒法开始打客服就是的咖啡就开始大幅扣税的急啊看发叫阿三的发生的开发就是大家可是大家发看大数据开发大数据开发大家ask发就是的咖啡的萨芬就卡死的房价开始打家开发商的JFK上的飞机卡上的纠纷开始打飞机宽带技术开发就开始大家开发建设的卡JFK大数据风控静安寺的看法角度看萨芬卡上的纠纷看静安寺的看法角度思考积分可是大家发卡是大家看法就大肆砍伐尽快打算减肥肯定是积分开始大幅技术大咖积分开始打飞机扣税的急啊看发的技术开发就是JFK十大福克斯大家开发大撒发射点幅度萨芬撒旦发发收范德萨发顺丰士大夫十大阿斯蒂芬大师傅阿斯顿附件是的客服对接撒巨大石块积分的课时费阿斯蒂芬法大师傅大师傅十大法大师傅阿斯蒂芬阿斯顿法大师傅阿斯蒂芬大师傅阿斯顿法大师傅大师傅阿斯蒂芬阿斯蒂芬士大夫阿斯蒂芬大师傅的萨芬打算减肥上岛咖啡加快速度大数据开发就是打客服看大数据开发就开始减肥卡萨丁JFK是大家看法加快速度JFK技术大咖积分喀什的开发独守空房技术大咖积分空手道解放扣税的开发商的开发接口是大家看法角度看是否扣税的急啊看发生的开发的快速减肥开始大幅就是打客服卡上的纠纷啊撒旦解放扣税的急啊看发加快速度点卡JFK啥的但是法大师傅技术大咖积分卡萨丁就反馈是大家看法啊是大家看法卡上的纠纷可是大家发喀什的开发大卡司喀什的开发就是打客服法大师傅士大夫的式咖啡机上岛咖啡就是的咖啡艰苦大师傅看上雕刻技法喀什的开发上岛咖啡就喀什的开发就是打客服卡上的纠纷技术的咖啡机肯定撒开发啊十大科技开发速度加啊反馈就是的咖啡开始大幅大师傅似的十大放假啊上岛咖啡就可是大家发空间的是否撒旦士大夫的撒娇开发是大家看法大肆砍伐就喀什的开发氨基酸的考虑非军事对抗疗法金克拉撒旦发艰苦拉萨的飞机喀什打开发就可是大家发可是大家看附件卡上的纠纷卡刷点卡技术的咖啡机可是大家发卡是大家看法静安寺的看法就可是大家发卡萨丁就开发商的急啊看飞机迪斯科发技术的咖啡机可是大家发看电视剧开发商大开始打到发大水发大水", + ); let mut cryption = Aes256GcmCodec::try_new(&TEST_KEY).unwrap(); let mut out_buf = data.as_bytes().to_vec(); let tag = { diff --git a/tests/regression.rs b/tests/regression.rs index 6c78140..793c943 100644 --- a/tests/regression.rs +++ b/tests/regression.rs @@ -4,9 +4,9 @@ use std::sync::{Arc, Mutex}; use std::time::Duration; use pb_mapper::common::auth::{ - write_admin_key_file, AuthConfig, AuthRuntime, LegacyProtocolPolicy, ADMIN_KEY_ID, + ADMIN_KEY_ID, AuthConfig, AuthRuntime, LegacyProtocolPolicy, write_admin_key_file, }; -use pb_mapper::common::checksum::{parse_credential, set_process_msg_header_key, Credential}; +use pb_mapper::common::checksum::{Credential, parse_credential, set_process_msg_header_key}; use pb_mapper::common::message::command::{ AdminRequest, AdminResponse, LocalServer, MessageSerializer, PbConnRequest, PbConnResponse, PbConnStatusReq, PbConnStatusResp, PbServerRequest, PbServiceConnStatus, @@ -15,10 +15,10 @@ use pb_mapper::common::message::secure::{ ClientHeaderSession, ServerHeaderSession, ServerSecurity, }; use pb_mapper::common::message::{ - get_header_msg_reader, get_header_msg_writer, MessageReader, MessageWriter, + MessageReader, MessageWriter, get_header_msg_reader, get_header_msg_writer, }; use pb_mapper::local::client::run_client_side_cli_with_callback; -use pb_mapper::local::server::{run_server_side_cli_with_callback, ServerTunnelOptions}; +use pb_mapper::local::server::{ServerTunnelOptions, run_server_side_cli_with_callback}; use pb_mapper::pb_server::run_server_with_auth_config; use tokio::io::AsyncReadExt; use tokio::io::AsyncWriteExt; @@ -33,19 +33,27 @@ struct EnvVarGuard { } impl EnvVarGuard { + /// # Safety note + /// + /// Mutating the environment is unsafe in edition 2024 because it races + /// concurrent readers. The tests that use this guard set a variable no + /// other test reads, and the guard restores the previous value on drop. fn set(key: &'static str, value: &'static str) -> Self { let old_value = std::env::var(key).ok(); - std::env::set_var(key, value); + unsafe { std::env::set_var(key, value) }; Self { key, old_value } } } impl Drop for EnvVarGuard { fn drop(&mut self) { - if let Some(value) = self.old_value.take() { - std::env::set_var(self.key, value); - } else { - std::env::remove_var(self.key); + // SAFETY: as in `set`. + unsafe { + if let Some(value) = self.old_value.take() { + std::env::set_var(self.key, value); + } else { + std::env::remove_var(self.key); + } } } } @@ -743,10 +751,10 @@ async fn local_server_reconnects_when_registered_conn_is_missing_from_remote_sta .unwrap(); let mut writer = session.response_writer(&mut stream).unwrap(); writer.write_msg(&response).await.unwrap(); - if count == 2 { - if let Some(tx) = second_register_tx.lock().await.take() { - tx.send(()).unwrap(); - } + if count == 2 + && let Some(tx) = second_register_tx.lock().await.take() + { + tx.send(()).unwrap(); } tracing::debug!(key, count, "fake server accepted register"); std::future::pending::<()>().await; diff --git a/tests/test_delay.rs b/tests/test_delay.rs index b80b87d..09da3a1 100644 --- a/tests/test_delay.rs +++ b/tests/test_delay.rs @@ -8,12 +8,12 @@ use pb_mapper::common::message::{ MessageReader, MessageWriter, NormalMessageReader, NormalMessageWriter, }; use pb_mapper::local::client::run_client_side_cli; -use pb_mapper::local::server::{run_server_side_cli, ServerTunnelOptions}; +use pb_mapper::local::server::{ServerTunnelOptions, run_server_side_cli}; use pb_mapper::pb_server::run_server_with_auth_config; use rand::RngExt; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::UdpSocket; -use tokio::time::{timeout, Instant}; +use tokio::time::{Instant, timeout}; use tokio_util::sync::CancellationToken; use uni_stream::addr::ToSocketAddrs; use uni_stream::stream::{ListenerProvider, TcpListenerProvider, UdpListenerProvider}; @@ -219,11 +219,11 @@ async fn run_udp_datagram_echo(addr: &str, rounds: usize, burst: usize) { let mut ready = false; for _ in 0..10 { socket.send(probe).await.unwrap(); - if let Ok(Ok(len)) = timeout(Duration::from_millis(300), socket.recv(&mut buf)).await { - if &buf[..len] == probe { - ready = true; - break; - } + if let Ok(Ok(len)) = timeout(Duration::from_millis(300), socket.recv(&mut buf)).await + && &buf[..len] == probe + { + ready = true; + break; } tokio::time::sleep(Duration::from_millis(50)).await; } diff --git a/ui/native/pb_mapper_ffi/Cargo.toml b/ui/native/pb_mapper_ffi/Cargo.toml index ae239d5..4c814bf 100644 --- a/ui/native/pb_mapper_ffi/Cargo.toml +++ b/ui/native/pb_mapper_ffi/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "pb-mapper-ffi" version = "0.1.0" -edition = "2021" +edition = "2024" [lib] crate-type = ["cdylib", "staticlib"] diff --git a/ui/native/pb_mapper_ffi/src/cli.rs b/ui/native/pb_mapper_ffi/src/cli.rs index e270994..d749c4e 100644 --- a/ui/native/pb_mapper_ffi/src/cli.rs +++ b/ui/native/pb_mapper_ffi/src/cli.rs @@ -11,12 +11,12 @@ //! Dart's whole part is: hand argv over, print nothing, exit with what comes //! back. -use std::ffi::{c_char, c_int, CStr}; +use std::ffi::{CStr, c_char, c_int}; use clap::{CommandFactory, Parser}; use crate::ctl::proto::Response; -use crate::ctl::{endpoint, server, Command}; +use crate::ctl::{Command, endpoint, server}; /// Returned when argv is not a command at all, so Dart knows to run the GUI. /// Chosen so it cannot collide with a real exit code. @@ -143,7 +143,7 @@ fn run(args: Vec) -> c_int { /// /// # Safety /// `argv` must point to `argc` valid, NUL-terminated C strings. -#[no_mangle] +#[unsafe(no_mangle)] pub unsafe extern "C" fn pb_mapper_cli_main(argc: c_int, argv: *const *const c_char) -> c_int { if argv.is_null() || argc <= 0 { return NOT_A_COMMAND; diff --git a/ui/native/pb_mapper_ffi/src/client.rs b/ui/native/pb_mapper_ffi/src/client.rs index 7db2abd..3caec10 100644 --- a/ui/native/pb_mapper_ffi/src/client.rs +++ b/ui/native/pb_mapper_ffi/src/client.rs @@ -12,7 +12,7 @@ use crate::response::{err_ctl, err_null_handle, ok_data, ok_message, parse_c_str use crate::state::{ClientConfigInfo, ClientStatusResponse}; /// Connect client to service. -#[no_mangle] +#[unsafe(no_mangle)] pub unsafe extern "C" fn pb_mapper_connect_service( handle: *mut PbMapperHandle, service_key: *const c_char, @@ -65,7 +65,7 @@ pub unsafe extern "C" fn pb_mapper_connect_service( } /// Disconnect client from service. -#[no_mangle] +#[unsafe(no_mangle)] pub unsafe extern "C" fn pb_mapper_disconnect_service( handle: *mut PbMapperHandle, service_key: *const c_char, @@ -96,7 +96,7 @@ pub unsafe extern "C" fn pb_mapper_disconnect_service( } /// Delete client config (also stops client if running). -#[no_mangle] +#[unsafe(no_mangle)] pub unsafe extern "C" fn pb_mapper_delete_client_config( handle: *mut PbMapperHandle, service_key: *const c_char, @@ -127,7 +127,7 @@ pub unsafe extern "C" fn pb_mapper_delete_client_config( } /// Get client configs list. -#[no_mangle] +#[unsafe(no_mangle)] pub unsafe extern "C" fn pb_mapper_get_client_configs_json( handle: *mut PbMapperHandle, ) -> *mut c_char { @@ -146,7 +146,7 @@ pub unsafe extern "C" fn pb_mapper_get_client_configs_json( } /// Get client status for a specific key. -#[no_mangle] +#[unsafe(no_mangle)] pub unsafe extern "C" fn pb_mapper_get_client_status_json( handle: *mut PbMapperHandle, service_key: *const c_char, diff --git a/ui/native/pb_mapper_ffi/src/config.rs b/ui/native/pb_mapper_ffi/src/config.rs index 3ddda8a..b6f542a 100644 --- a/ui/native/pb_mapper_ffi/src/config.rs +++ b/ui/native/pb_mapper_ffi/src/config.rs @@ -11,7 +11,7 @@ use crate::handle::PbMapperHandle; use crate::response::{err_ctl, err_null_handle, ok_data, ok_message, parse_c_string}; /// Get current app config. -#[no_mangle] +#[unsafe(no_mangle)] pub unsafe extern "C" fn pb_mapper_get_config_json(handle: *mut PbMapperHandle) -> *mut c_char { if handle.is_null() { return err_null_handle(); @@ -37,7 +37,7 @@ pub unsafe extern "C" fn pb_mapper_get_config_json(handle: *mut PbMapperHandle) /// Reveal the embedded relay administrator key. This is a separate call so /// routine config fetches cannot leak the root secret. -#[no_mangle] +#[unsafe(no_mangle)] pub unsafe extern "C" fn pb_mapper_reveal_isolated_admin_key( handle: *mut PbMapperHandle, ) -> *mut c_char { @@ -58,7 +58,7 @@ pub unsafe extern "C" fn pb_mapper_reveal_isolated_admin_key( } /// Update app config. -#[no_mangle] +#[unsafe(no_mangle)] pub unsafe extern "C" fn pb_mapper_update_config( handle: *mut PbMapperHandle, server_address: *const c_char, diff --git a/ui/native/pb_mapper_ffi/src/ctl/endpoint.rs b/ui/native/pb_mapper_ffi/src/ctl/endpoint.rs index 27e1e23..b808c2c 100644 --- a/ui/native/pb_mapper_ffi/src/ctl/endpoint.rs +++ b/ui/native/pb_mapper_ffi/src/ctl/endpoint.rs @@ -105,10 +105,10 @@ mod imp { } pub fn endpoint() -> String { - if let Ok(custom) = std::env::var(super::ENDPOINT_ENV) { - if !custom.is_empty() { - return custom; - } + if let Ok(custom) = std::env::var(super::ENDPOINT_ENV) + && !custom.is_empty() + { + return custom; } // Per user, so two accounts on one machine do not collide. let user = std::env::var("USERNAME").unwrap_or_else(|_| "default".into()); @@ -160,10 +160,10 @@ mod imp { const SUN_PATH_MAX: usize = 100; pub fn endpoint() -> String { - if let Ok(custom) = std::env::var(super::ENDPOINT_ENV) { - if !custom.is_empty() { - return custom; - } + if let Ok(custom) = std::env::var(super::ENDPOINT_ENV) + && !custom.is_empty() + { + return custom; } // XDG_RUNTIME_DIR is per-user and cleaned on logout, which is what a // socket wants. TMPDIR is the macOS equivalent and matters more there @@ -195,7 +195,7 @@ mod imp { format!("/tmp/pb-mapper-ui-{uid}.sock") } - extern "C" { + unsafe extern "C" { #[link_name = "getuid"] fn libc_getuid() -> u32; } diff --git a/ui/native/pb_mapper_ffi/src/ctl/server.rs b/ui/native/pb_mapper_ffi/src/ctl/server.rs index a4570e7..a4ac90a 100644 --- a/ui/native/pb_mapper_ffi/src/ctl/server.rs +++ b/ui/native/pb_mapper_ffi/src/ctl/server.rs @@ -11,8 +11,8 @@ use tokio::sync::Mutex; use tokio_util::sync::CancellationToken; use crate::ctl::endpoint; -use crate::ctl::proto::{self, Request, Response, PROTOCOL_VERSION}; -use crate::ctl::{dispatch, Origin}; +use crate::ctl::proto::{self, PROTOCOL_VERSION, Request, Response}; +use crate::ctl::{Origin, dispatch}; use crate::error::CtlError; use crate::state::PbMapperState; diff --git a/ui/native/pb_mapper_ffi/src/events.rs b/ui/native/pb_mapper_ffi/src/events.rs index d05abfa..a5521e4 100644 --- a/ui/native/pb_mapper_ffi/src/events.rs +++ b/ui/native/pb_mapper_ffi/src/events.rs @@ -7,7 +7,7 @@ //! mean taking the state lock on the emit path and guessing which projection //! the receiver wants. -use std::ffi::{c_char, CString}; +use std::ffi::{CString, c_char}; use std::sync::atomic::{AtomicU64, Ordering}; use serde::Serialize; @@ -71,7 +71,7 @@ pub fn emit(kind: ChangeKind, key: Option<&str>, origin: Origin) { /// /// # Safety /// `callback` must stay valid until it is replaced or cleared with null. -#[no_mangle] +#[unsafe(no_mangle)] pub extern "C" fn pb_mapper_set_change_callback(callback: Option) { CHANGE_CALLBACK.store(callback); } diff --git a/ui/native/pb_mapper_ffi/src/handle.rs b/ui/native/pb_mapper_ffi/src/handle.rs index b97a07e..a6ada13 100644 --- a/ui/native/pb_mapper_ffi/src/handle.rs +++ b/ui/native/pb_mapper_ffi/src/handle.rs @@ -1,7 +1,7 @@ //! FFI handle lifecycle and app directory configuration. #![allow(clippy::missing_safety_doc)] -use std::ffi::{c_char, CStr}; +use std::ffi::{CStr, c_char}; use std::ptr; use std::sync::Arc; @@ -31,7 +31,7 @@ impl Drop for PbMapperHandle { /// /// # Safety /// Returns a pointer that must be freed with `pb_mapper_destroy`. -#[no_mangle] +#[unsafe(no_mangle)] pub extern "C" fn pb_mapper_create() -> *mut PbMapperHandle { let runtime = match Runtime::new() { Ok(rt) => rt, @@ -57,7 +57,7 @@ pub extern "C" fn pb_mapper_create() -> *mut PbMapperHandle { /// /// # Safety /// `handle` must be a valid pointer returned by `pb_mapper_create`. -#[no_mangle] +#[unsafe(no_mangle)] pub unsafe extern "C" fn pb_mapper_start_control_server( handle: *mut PbMapperHandle, ) -> *mut c_char { @@ -82,7 +82,7 @@ pub unsafe extern "C" fn pb_mapper_start_control_server( /// /// # Safety /// `handle` must be a valid pointer returned by `pb_mapper_create`. -#[no_mangle] +#[unsafe(no_mangle)] pub unsafe extern "C" fn pb_mapper_destroy(handle: *mut PbMapperHandle) { if !handle.is_null() { unsafe { drop(Box::from_raw(handle)) }; @@ -93,7 +93,7 @@ pub unsafe extern "C" fn pb_mapper_destroy(handle: *mut PbMapperHandle) { /// /// # Safety /// `handle` must be valid. `path` must be valid C string or null. -#[no_mangle] +#[unsafe(no_mangle)] pub unsafe extern "C" fn pb_mapper_set_app_dir( handle: *mut PbMapperHandle, path: *const c_char, diff --git a/ui/native/pb_mapper_ffi/src/lib.rs b/ui/native/pb_mapper_ffi/src/lib.rs index 4635488..5853746 100644 --- a/ui/native/pb_mapper_ffi/src/lib.rs +++ b/ui/native/pb_mapper_ffi/src/lib.rs @@ -17,7 +17,7 @@ mod state; // Re-export public FFI functions and handle type. use better_mimalloc_rs::MiMalloc; -pub use cli::{pb_mapper_cli_main, NOT_A_COMMAND}; +pub use cli::{NOT_A_COMMAND, pb_mapper_cli_main}; pub use client::{ pb_mapper_connect_service, pb_mapper_delete_client_config, pb_mapper_disconnect_service, pb_mapper_get_client_configs_json, pb_mapper_get_client_status_json, @@ -27,8 +27,8 @@ pub use config::{ }; pub use events::pb_mapper_set_change_callback; pub use handle::{ - pb_mapper_create, pb_mapper_destroy, pb_mapper_set_app_dir, pb_mapper_start_control_server, - PbMapperHandle, + PbMapperHandle, pb_mapper_create, pb_mapper_destroy, pb_mapper_set_app_dir, + pb_mapper_start_control_server, }; pub use logging::{pb_mapper_free_string, pb_mapper_init_logging, pb_mapper_set_log_callback}; pub use server::{ diff --git a/ui/native/pb_mapper_ffi/src/logging.rs b/ui/native/pb_mapper_ffi/src/logging.rs index 2419029..2768e21 100644 --- a/ui/native/pb_mapper_ffi/src/logging.rs +++ b/ui/native/pb_mapper_ffi/src/logging.rs @@ -1,6 +1,6 @@ //! Logging system for FFI interface. -use std::ffi::{c_char, c_int, CString}; +use std::ffi::{CString, c_char, c_int}; use std::time::{SystemTime, UNIX_EPOCH}; use crate::callback::CallbackSlot; @@ -34,7 +34,7 @@ pub(crate) fn send_log(level: c_int, message: &str) { /// /// # Safety /// `callback` must be a valid function pointer or null to disable logging. -#[no_mangle] +#[unsafe(no_mangle)] pub extern "C" fn pb_mapper_set_log_callback(callback: Option) { LOG_CALLBACK.store(callback); } @@ -43,7 +43,7 @@ pub extern "C" fn pb_mapper_set_log_callback(callback: Option) { /// /// # Safety /// `s` must be a valid pointer returned from this library, or null. -#[no_mangle] +#[unsafe(no_mangle)] pub unsafe extern "C" fn pb_mapper_free_string(s: *mut c_char) { if !s.is_null() { unsafe { drop(CString::from_raw(s)) }; @@ -112,11 +112,11 @@ impl tracing::field::Visit for MessageVisitor { /// /// # Safety /// Can be called multiple times safely. -#[no_mangle] +#[unsafe(no_mangle)] pub extern "C" fn pb_mapper_init_logging() { + use tracing_subscriber::Layer; use tracing_subscriber::layer::SubscriberExt; use tracing_subscriber::util::SubscriberInitExt; - use tracing_subscriber::Layer; let _ = tracing_subscriber::registry() .with(FfiLogLayer) diff --git a/ui/native/pb_mapper_ffi/src/response.rs b/ui/native/pb_mapper_ffi/src/response.rs index aa92880..180ae4a 100644 --- a/ui/native/pb_mapper_ffi/src/response.rs +++ b/ui/native/pb_mapper_ffi/src/response.rs @@ -1,6 +1,6 @@ //! Shared helpers for FFI response formatting and argument parsing. -use std::ffi::{c_char, CStr, CString}; +use std::ffi::{CStr, CString, c_char}; use std::ptr; use serde_json::json; diff --git a/ui/native/pb_mapper_ffi/src/server.rs b/ui/native/pb_mapper_ffi/src/server.rs index ef0df8f..aac58c2 100644 --- a/ui/native/pb_mapper_ffi/src/server.rs +++ b/ui/native/pb_mapper_ffi/src/server.rs @@ -11,7 +11,7 @@ use crate::handle::PbMapperHandle; use crate::response::{err_ctl, err_null_handle, ok_data, ok_message, parse_c_string}; /// Start pb-mapper server. -#[no_mangle] +#[unsafe(no_mangle)] pub unsafe extern "C" fn pb_mapper_start_server( handle: *mut PbMapperHandle, port: u16, @@ -38,7 +38,7 @@ pub unsafe extern "C" fn pb_mapper_start_server( } /// Stop pb-mapper server. -#[no_mangle] +#[unsafe(no_mangle)] pub unsafe extern "C" fn pb_mapper_stop_server(handle: *mut PbMapperHandle) -> *mut c_char { if handle.is_null() { return err_null_handle(); @@ -61,7 +61,7 @@ pub unsafe extern "C" fn pb_mapper_stop_server(handle: *mut PbMapperHandle) -> * } /// Get local server status (running/uptime). -#[no_mangle] +#[unsafe(no_mangle)] pub unsafe extern "C" fn pb_mapper_get_local_server_status_json( handle: *mut PbMapperHandle, ) -> *mut c_char { @@ -84,7 +84,7 @@ pub unsafe extern "C" fn pb_mapper_get_local_server_status_json( /// The status detail's `serverMap` is a Debug dump of the whole map and is not /// something a UI should be parsing. This answers the same question with the /// protocol's own structured query, one key at a time. -#[no_mangle] +#[unsafe(no_mangle)] pub unsafe extern "C" fn pb_mapper_get_service_conns_json( handle: *mut PbMapperHandle, service_key: *const c_char, @@ -112,7 +112,7 @@ pub unsafe extern "C" fn pb_mapper_get_service_conns_json( } /// Get server status detail (remote server). -#[no_mangle] +#[unsafe(no_mangle)] pub unsafe extern "C" fn pb_mapper_get_server_status_detail_json( handle: *mut PbMapperHandle, ) -> *mut c_char { @@ -134,7 +134,7 @@ pub unsafe extern "C" fn pb_mapper_get_server_status_detail_json( } /// Force-refresh server status (blocks until network result is available). -#[no_mangle] +#[unsafe(no_mangle)] pub unsafe extern "C" fn pb_mapper_force_refresh_server_status_json( handle: *mut PbMapperHandle, ) -> *mut c_char { diff --git a/ui/native/pb_mapper_ffi/src/service.rs b/ui/native/pb_mapper_ffi/src/service.rs index 5b479d6..9dc4d53 100644 --- a/ui/native/pb_mapper_ffi/src/service.rs +++ b/ui/native/pb_mapper_ffi/src/service.rs @@ -12,7 +12,7 @@ use crate::response::{err_ctl, err_null_handle, ok_data, ok_message, parse_c_str use crate::state::{ServiceConfigInfo, ServiceStatusResponse}; /// Register a service. -#[no_mangle] +#[unsafe(no_mangle)] pub unsafe extern "C" fn pb_mapper_register_service( handle: *mut PbMapperHandle, service_key: *const c_char, @@ -65,7 +65,7 @@ pub unsafe extern "C" fn pb_mapper_register_service( } /// Unregister a service (stop running but keep config). -#[no_mangle] +#[unsafe(no_mangle)] pub unsafe extern "C" fn pb_mapper_unregister_service( handle: *mut PbMapperHandle, service_key: *const c_char, @@ -96,7 +96,7 @@ pub unsafe extern "C" fn pb_mapper_unregister_service( } /// Delete service config (also stops service if running). -#[no_mangle] +#[unsafe(no_mangle)] pub unsafe extern "C" fn pb_mapper_delete_service_config( handle: *mut PbMapperHandle, service_key: *const c_char, @@ -127,7 +127,7 @@ pub unsafe extern "C" fn pb_mapper_delete_service_config( } /// Get service configs list. -#[no_mangle] +#[unsafe(no_mangle)] pub unsafe extern "C" fn pb_mapper_get_service_configs_json( handle: *mut PbMapperHandle, ) -> *mut c_char { @@ -146,7 +146,7 @@ pub unsafe extern "C" fn pb_mapper_get_service_configs_json( } /// Get service status for a specific key. -#[no_mangle] +#[unsafe(no_mangle)] pub unsafe extern "C" fn pb_mapper_get_service_status_json( handle: *mut PbMapperHandle, service_key: *const c_char, diff --git a/ui/native/pb_mapper_ffi/src/state.rs b/ui/native/pb_mapper_ffi/src/state.rs index 08bab5a..8242cc0 100644 --- a/ui/native/pb_mapper_ffi/src/state.rs +++ b/ui/native/pb_mapper_ffi/src/state.rs @@ -13,8 +13,8 @@ use std::collections::{HashMap, HashSet}; use std::fs; use std::net::{IpAddr, Ipv4Addr, SocketAddr}; use std::path::PathBuf; -use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use parking_lot::Mutex as SyncMutex; @@ -26,16 +26,16 @@ use tokio_util::sync::CancellationToken; use pb_mapper::common::auth::{AuthConfig, AuthRuntime}; use pb_mapper::common::checksum::{ - get_process_credential, parse_credential, set_process_msg_header_key, Credential, + Credential, get_process_credential, parse_credential, set_process_msg_header_key, }; use pb_mapper::common::config::{get_pb_mapper_server_async, get_sockaddr_async}; use pb_mapper::common::message::command::{PbConnStatusReq, PbConnStatusResp}; use pb_mapper::local::client::status::{get_status, get_status_with_credential}; -use pb_mapper::local::client::{run_client_side_cli_with_pinned_credential, ClientStatusCallback}; +use pb_mapper::local::client::{ClientStatusCallback, run_client_side_cli_with_pinned_credential}; use pb_mapper::local::server::{ - run_server_side_cli_with_pinned_credential, ServerTunnelOptions, StatusCallback, + ServerTunnelOptions, StatusCallback, run_server_side_cli_with_pinned_credential, }; -use pb_mapper::pb_server::{run_server_on_listener, ServerStatusInfo}; +use pb_mapper::pb_server::{ServerStatusInfo, run_server_on_listener}; use pb_mapper::utils::addr::each_addr; use uni_stream::stream::got_one_socket_addr; use uni_stream::stream::{ diff --git a/ui/native/pb_mapper_ffi/src/state/runtime.rs b/ui/native/pb_mapper_ffi/src/state/runtime.rs index ba326f2..9786fd3 100644 --- a/ui/native/pb_mapper_ffi/src/state/runtime.rs +++ b/ui/native/pb_mapper_ffi/src/state/runtime.rs @@ -101,12 +101,12 @@ impl PbMapperState { tracing::info!("Server shutdown gracefully"); } Err(_) => { - if let Some(auth) = self.server_auth.as_ref() { - if let Err(error) = auth.abort_actor().await { - tracing::warn!( - "timed out waiting for the authentication actor to drop: {error}" - ); - } + if let Some(auth) = self.server_auth.as_ref() + && let Err(error) = auth.abort_actor().await + { + tracing::warn!( + "timed out waiting for the authentication actor to drop: {error}" + ); } handle.abort(); let _ = handle.await; diff --git a/ui/native/pb_mapper_ffi/src/state/status.rs b/ui/native/pb_mapper_ffi/src/state/status.rs index 5927111..b4219cf 100644 --- a/ui/native/pb_mapper_ffi/src/state/status.rs +++ b/ui/native/pb_mapper_ffi/src/state/status.rs @@ -193,14 +193,13 @@ impl PbMapperState { if let Some(sender) = sender { let (response_sender, response_receiver) = tokio::sync::oneshot::channel(); - if sender.send(response_sender).is_ok() { - if let Ok(Ok(info)) = + if sender.send(response_sender).is_ok() + && let Ok(Ok(info)) = tokio::time::timeout(Duration::from_millis(200), response_receiver).await - { - status.active_connections = info.active_connections; - status.registered_services = info.registered_services; - status.uptime_seconds = info.uptime_seconds; - } + { + status.active_connections = info.active_connections; + status.registered_services = info.registered_services; + status.uptime_seconds = info.uptime_seconds; } } From 9351bfa910986360f4c4d8b00205055a6377e00b Mon Sep 17 00:00:00 2001 From: LB7666 Date: Fri, 21 Aug 2026 13:36:21 +0800 Subject: [PATCH 69/74] Extract pb-mapper-core MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The bottom layer: checksum, config, conn_id, error, addr, codec, timeout. All seven land as git renames, so blame follows them. The old paths keep working through re-exports in `common/mod.rs` and a new `utils.rs`, so the rest of the tree is untouched and every later crate can move on its own. Three things had to move rather than be copied, because each was a cycle or would not survive the boundary: - `DataLenType` now lives in `core`, since `checksum` and `error` both name it while `message` depends on them. `message` re-exports it rather than redeclaring, or the two would be distinct names for the same width. - `PROCESS_CREDENTIAL_TEST_LOCK` moves to `core::test_support`, next to the credential it guards. It was `#[cfg(test)] pub(crate)` in auth, and a test-only item is invisible to another crate's tests — four future crates read it, so it is now unconditionally `pub`. - `replace_file` and `sync_parent_directory` become `core::durable_file`, reporting `io::Result`. They are general file primitives that happened to be written where the first caller was. `MngWaitForTask` did not come along: only `manager` constructs it, so it is now `manager`'s own error and `kanal` stays out of the bottom layer. The `Error` selectors widen from `pub(super)` to `pub` — 13 of them are built from other crates. Also dropped `socket2` and `futures`, declared but referenced nowhere, and gated `lib.rs`'s `mod tests` behind `#[cfg(test)]` — it was compiling into release builds. 136 tests still pass, now 87 in the old crate plus 12 that moved to core. Co-Authored-By: Claude Fable 5 --- Cargo.lock | 21 ++++- Cargo.toml | 26 +++++-- crates/pb-mapper-core/Cargo.toml | 18 +++++ .../pb-mapper-core/src}/addr.rs | 0 .../pb-mapper-core/src}/checksum.rs | 4 +- .../pb-mapper-core/src}/codec.rs | 2 +- .../pb-mapper-core/src}/config.rs | 26 +++---- .../pb-mapper-core/src}/conn_id.rs | 0 crates/pb-mapper-core/src/durable_file.rs | 76 +++++++++++++++++++ .../pb-mapper-core/src}/error.rs | 11 ++- crates/pb-mapper-core/src/lib.rs | 20 +++++ crates/pb-mapper-core/src/test_support.rs | 14 ++++ .../pb-mapper-core/src}/timeout.rs | 0 src/common/auth.rs | 4 - src/common/auth/tests.rs | 2 + src/common/manager.rs | 13 +++- src/common/message/mod.rs | 6 +- src/common/message/secure/tests.rs | 3 +- src/common/mod.rs | 8 +- src/lib.rs | 10 +++ src/pb_server/admin.rs | 2 +- src/pb_server/runtime.rs | 5 +- src/utils.rs | 3 + src/utils/mod.rs | 3 - 24 files changed, 226 insertions(+), 51 deletions(-) create mode 100644 crates/pb-mapper-core/Cargo.toml rename {src/utils => crates/pb-mapper-core/src}/addr.rs (100%) rename {src/common => crates/pb-mapper-core/src}/checksum.rs (99%) rename {src/utils => crates/pb-mapper-core/src}/codec.rs (99%) rename {src/common => crates/pb-mapper-core/src}/config.rs (95%) rename {src/common => crates/pb-mapper-core/src}/conn_id.rs (100%) create mode 100644 crates/pb-mapper-core/src/durable_file.rs rename {src/common => crates/pb-mapper-core/src}/error.rs (96%) create mode 100644 crates/pb-mapper-core/src/lib.rs create mode 100644 crates/pb-mapper-core/src/test_support.rs rename {src/utils => crates/pb-mapper-core/src}/timeout.rs (100%) create mode 100644 src/utils.rs delete mode 100644 src/utils/mod.rs diff --git a/Cargo.lock b/Cargo.lock index 2e6535e..4bf3432 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1038,17 +1038,15 @@ dependencies = [ "bytes", "clap", "dotenvy", - "futures", "hashbrown 0.17.1", - "hickory-resolver", "kanal", "parking_lot", + "pb-mapper-core", "rand 0.10.0", "ring", "serde", "serde_json", "snafu", - "socket2 0.6.1", "subtle", "tokio", "tokio-util", @@ -1057,6 +1055,23 @@ dependencies = [ "uni-stream", ] +[[package]] +name = "pb-mapper-core" +version = "0.4.0" +dependencies = [ + "base64", + "clap", + "hickory-resolver", + "parking_lot", + "rand 0.10.0", + "ring", + "serde_json", + "snafu", + "tokio", + "tracing", + "tracing-subscriber", +] + [[package]] name = "pb-mapper-ffi" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index d6582e6..214f0d5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,8 +7,8 @@ authors.workspace = true # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] +pb-mapper-core.workspace = true rand.workspace = true -socket2.workspace = true tokio.workspace = true tokio-util.workspace = true snafu.workspace = true @@ -18,10 +18,8 @@ tracing.workspace = true tracing-subscriber.workspace = true hashbrown.workspace = true clap.workspace = true -futures.workspace = true better_mimalloc_rs.workspace = true bytes.workspace = true -hickory-resolver.workspace = true ring.workspace = true uni-stream.workspace = true kanal.workspace = true @@ -30,23 +28,36 @@ subtle.workspace = true parking_lot.workspace = true [dev-dependencies] -dotenvy = "0.15.7" +dotenvy.workspace = true [features] udp-timeout = ["uni-stream/udp-timeout"] [workspace] -members = ["ui/native/pb_mapper_ffi"] +members = ["crates/*", "ui/native/pb_mapper_ffi"] exclude = ["deps/uni-stream", "deps/kanal"] +# Spelled out because a virtual manifest does not infer the resolver from the +# edition, and this root becomes virtual once `src/` is empty. +resolver = "3" [workspace.package] +# `version` must stay the first key here: `release.yml` and `docker-publish.yml` +# both parse it positionally with awk to check the tag against it. version = "0.4.0" authors = ["L_B__"] edition = "2024" +[workspace.lints.clippy] +# The FFI crate has denied these since it was written — a panic crossing the C +# ABI is undefined behaviour rather than a stack trace. Everything here is +# reachable from that boundary, so the whole workspace inherits it. +unwrap_used = "deny" +expect_used = "deny" + [workspace.dependencies] +pb-mapper-core = { path = "crates/pb-mapper-core" } + rand = "0.10" -socket2 = "0.6" tokio = { version = "1", features = ["full"] } tokio-util = "0.7" snafu = "0.9.2" @@ -60,7 +71,6 @@ tracing-subscriber = { version = "0.3.18", features = [ ], default-features = true } hashbrown = { version = "0.17.1" } clap = { version = "4.5", features = ["derive"] } -futures = "0.3.31" better_mimalloc_rs = { version = "0.1.2", features = ["config"] } bytes = "1.11" hickory-resolver = { version = "0.26.1" } @@ -68,5 +78,7 @@ ring = "0.17.14" base64 = "0.23.1" subtle = "2.6.1" parking_lot = "0.12" +dirs = "6.0.0" +dotenvy = "0.15.7" uni-stream = { git = "https://github.com/acking-you/uni-stream.git", branch = "master" } kanal = { git = "https://github.com/acking-you/kanal.git", branch = "dev/pb-mapper" } diff --git a/crates/pb-mapper-core/Cargo.toml b/crates/pb-mapper-core/Cargo.toml new file mode 100644 index 0000000..f054a61 --- /dev/null +++ b/crates/pb-mapper-core/Cargo.toml @@ -0,0 +1,18 @@ +[package] +name = "pb-mapper-core" +version.workspace = true +edition.workspace = true +authors.workspace = true + +[dependencies] +base64.workspace = true +clap.workspace = true +hickory-resolver.workspace = true +parking_lot.workspace = true +rand.workspace = true +ring.workspace = true +serde_json.workspace = true +snafu.workspace = true +tokio.workspace = true +tracing.workspace = true +tracing-subscriber.workspace = true diff --git a/src/utils/addr.rs b/crates/pb-mapper-core/src/addr.rs similarity index 100% rename from src/utils/addr.rs rename to crates/pb-mapper-core/src/addr.rs diff --git a/src/common/checksum.rs b/crates/pb-mapper-core/src/checksum.rs similarity index 99% rename from src/common/checksum.rs rename to crates/pb-mapper-core/src/checksum.rs index bf57883..ce4da1f 100644 --- a/src/common/checksum.rs +++ b/crates/pb-mapper-core/src/checksum.rs @@ -13,7 +13,7 @@ use base64::engine::general_purpose::URL_SAFE_NO_PAD; use rand::RngExt; use ring::digest::{SHA256, digest}; -use super::message::DataLenType; +use crate::DataLenType; pub type ChecksumType = u32; @@ -576,7 +576,7 @@ mod tests { #[tokio::test] async fn clearing_the_process_credential_fails_closed_for_checksums() { use super::*; - use crate::common::auth::PROCESS_CREDENTIAL_TEST_LOCK; + use crate::test_support::PROCESS_CREDENTIAL_TEST_LOCK; let _guard = PROCESS_CREDENTIAL_TEST_LOCK.lock().await; set_process_msg_header_key(Some("0123456789abcdefghijklmnopqrstuv")).unwrap(); diff --git a/src/utils/codec.rs b/crates/pb-mapper-core/src/codec.rs similarity index 99% rename from src/utils/codec.rs rename to crates/pb-mapper-core/src/codec.rs index c909e3f..ae6c186 100644 --- a/src/utils/codec.rs +++ b/crates/pb-mapper-core/src/codec.rs @@ -152,7 +152,7 @@ mod tests { use std::slice::from_raw_parts_mut; use std::time::Instant; - use crate::utils::codec::Aes256GcmCodec; + use crate::codec::Aes256GcmCodec; struct Timer { ins: Instant, diff --git a/src/common/config.rs b/crates/pb-mapper-core/src/config.rs similarity index 95% rename from src/common/config.rs rename to crates/pb-mapper-core/src/config.rs index d4e3aaf..bef77f3 100644 --- a/src/common/config.rs +++ b/crates/pb-mapper-core/src/config.rs @@ -7,7 +7,7 @@ use snafu::ResultExt; use tracing_subscriber::layer::SubscriberExt; use tracing_subscriber::{EnvFilter, Layer, fmt}; -use super::error::{CfgPbServerEnvNotExistSnafu, Result}; +use crate::error::{CfgPbServerEnvNotExistSnafu, Result}; #[derive(ValueEnum, Debug, Clone, Copy)] pub enum StatusOp { @@ -30,24 +30,24 @@ pub fn get_sockaddr(addr: &str) -> Result { Ok(mut socket_addrs) => { socket_addrs .next() - .ok_or_else(|| super::error::Error::CfgParseSockAddr { + .ok_or_else(|| crate::error::Error::CfgParseSockAddr { string: addr.to_string(), source: original_parse_error, }) } - Err(_) => Err(super::error::Error::CfgParseSockAddr { + Err(_) => Err(crate::error::Error::CfgParseSockAddr { string: addr.to_string(), source: original_parse_error, }), } } else { // For other hostnames, use the custom DNS resolution - use crate::utils::addr::get_socket_addrs; + use crate::addr::get_socket_addrs; match get_socket_addrs(addr) { Ok(socket_addrs) => { // Return the first resolved address socket_addrs.into_iter().next().ok_or_else(|| { - super::error::Error::CfgParseSockAddr { + crate::error::Error::CfgParseSockAddr { string: addr.to_string(), source: original_parse_error, } @@ -57,12 +57,12 @@ pub fn get_sockaddr(addr: &str) -> Result { // If custom DNS resolution fails, fallback to system resolver match std::net::ToSocketAddrs::to_socket_addrs(addr) { Ok(mut socket_addrs) => socket_addrs.next().ok_or_else(|| { - super::error::Error::CfgParseSockAddr { + crate::error::Error::CfgParseSockAddr { string: addr.to_string(), source: original_parse_error, } }), - Err(_) => Err(super::error::Error::CfgParseSockAddr { + Err(_) => Err(crate::error::Error::CfgParseSockAddr { string: addr.to_string(), source: original_parse_error, }), @@ -87,22 +87,22 @@ pub async fn get_sockaddr_async(addr: &str) -> Result { Ok(mut socket_addrs) => { socket_addrs .next() - .ok_or_else(|| super::error::Error::CfgParseSockAddr { + .ok_or_else(|| crate::error::Error::CfgParseSockAddr { string: addr.to_string(), source: original_parse_error, }) } - Err(_) => Err(super::error::Error::CfgParseSockAddr { + Err(_) => Err(crate::error::Error::CfgParseSockAddr { string: addr.to_string(), source: original_parse_error, }), } } else { // For other hostnames, use the custom DNS resolution - use crate::utils::addr::get_socket_addrs_async; + use crate::addr::get_socket_addrs_async; match get_socket_addrs_async(addr).await { Ok(socket_addrs) => socket_addrs.into_iter().next().ok_or_else(|| { - super::error::Error::CfgParseSockAddr { + crate::error::Error::CfgParseSockAddr { string: addr.to_string(), source: original_parse_error, } @@ -111,12 +111,12 @@ pub async fn get_sockaddr_async(addr: &str) -> Result { // If custom DNS resolution fails, fallback to system resolver match tokio::net::lookup_host(addr).await { Ok(mut socket_addrs) => socket_addrs.next().ok_or_else(|| { - super::error::Error::CfgParseSockAddr { + crate::error::Error::CfgParseSockAddr { string: addr.to_string(), source: original_parse_error, } }), - Err(_) => Err(super::error::Error::CfgParseSockAddr { + Err(_) => Err(crate::error::Error::CfgParseSockAddr { string: addr.to_string(), source: original_parse_error, }), diff --git a/src/common/conn_id.rs b/crates/pb-mapper-core/src/conn_id.rs similarity index 100% rename from src/common/conn_id.rs rename to crates/pb-mapper-core/src/conn_id.rs diff --git a/crates/pb-mapper-core/src/durable_file.rs b/crates/pb-mapper-core/src/durable_file.rs new file mode 100644 index 0000000..95b730d --- /dev/null +++ b/crates/pb-mapper-core/src/durable_file.rs @@ -0,0 +1,76 @@ +//! Atomic replace and parent-directory durability. +//! +//! These are general file primitives, not credential ones — they lived in the +//! auth persistence layer only because that is where the first caller was. Both +//! report `io::Result` and leave it to the caller to map into its own error +//! type. + +use std::fs::File; +use std::path::Path; + +/// Replaces `to` with `from`, atomically where the platform allows it. +pub fn replace_file(from: &Path, to: &Path) -> std::io::Result<()> { + #[cfg(windows)] + { + use std::os::windows::ffi::OsStrExt; + + const MOVEFILE_REPLACE_EXISTING: u32 = 0x1; + const MOVEFILE_WRITE_THROUGH: u32 = 0x8; + unsafe extern "system" { + fn MoveFileExW( + lp_existing_file_name: *const u16, + lp_new_file_name: *const u16, + dw_flags: u32, + ) -> i32; + } + fn wide(path: &Path) -> Vec { + path.as_os_str() + .encode_wide() + .chain(std::iter::once(0)) + .collect() + } + let from_w = wide(from); + let to_w = wide(to); + let ok = unsafe { + MoveFileExW( + from_w.as_ptr(), + to_w.as_ptr(), + MOVEFILE_REPLACE_EXISTING | MOVEFILE_WRITE_THROUGH, + ) + }; + if ok == 0 { + Err(std::io::Error::last_os_error()) + } else { + Ok(()) + } + } + #[cfg(not(windows))] + std::fs::rename(from, to) +} + +/// Fsyncs the directory holding `path`, so a rename into it survives a crash. +/// +/// A path with no parent is a no-op rather than an error. +pub fn sync_parent_directory(path: &Path) -> std::io::Result<()> { + let Some(parent) = path.parent() else { + return Ok(()); + }; + open_directory_for_sync(parent).and_then(|directory| directory.sync_all()) +} + +fn open_directory_for_sync(path: &Path) -> std::io::Result { + #[cfg(windows)] + { + use std::fs::OpenOptions; + use std::os::windows::fs::OpenOptionsExt; + const GENERIC_READ: u32 = 0x8000_0000; + const GENERIC_WRITE: u32 = 0x4000_0000; + const FILE_FLAG_BACKUP_SEMANTICS: u32 = 0x0200_0000; + OpenOptions::new() + .access_mode(GENERIC_READ | GENERIC_WRITE) + .custom_flags(FILE_FLAG_BACKUP_SEMANTICS) + .open(path) + } + #[cfg(not(windows))] + File::open(path) +} diff --git a/src/common/error.rs b/crates/pb-mapper-core/src/error.rs similarity index 96% rename from src/common/error.rs rename to crates/pb-mapper-core/src/error.rs index cd5609e..b4fc24f 100644 --- a/src/common/error.rs +++ b/crates/pb-mapper-core/src/error.rs @@ -3,11 +3,13 @@ use std::net::AddrParseError; use snafu::Snafu; -use super::checksum::ChecksumType; -use super::message::DataLenType; +use crate::DataLenType; +use crate::checksum::ChecksumType; #[derive(Debug, Snafu)] -#[snafu(visibility(pub(super)))] +// The generated context selectors are constructed by the protocol, client, and +// server crates, so they have to be reachable from outside this one. +#[snafu(visibility(pub))] pub enum Error { /// Error handling for message #[snafu(display("read `checksum` from network error"))] @@ -67,9 +69,6 @@ pub enum Error { action: &'static str, source: std::io::Error, }, - /// Error for manager - #[snafu(display("`TaskManager` fails while waiting for a task"))] - MngWaitForTask { source: kanal::ReceiveError }, /// Error for forward #[snafu(display("failed to forward message to write in normal text"))] FwdNetworkWriteWithNormal { source: std::io::Error }, diff --git a/crates/pb-mapper-core/src/lib.rs b/crates/pb-mapper-core/src/lib.rs new file mode 100644 index 0000000..0d13345 --- /dev/null +++ b/crates/pb-mapper-core/src/lib.rs @@ -0,0 +1,20 @@ +//! The bottom layer: credential primitives, framing checksums, configuration, +//! address resolution, and the file primitives the durable stores build on. +//! +//! Nothing here depends on another `pb-mapper` crate, which is what makes it +//! the bottom. `DataLenType` lives here rather than with the message framing +//! that names it, so that `checksum` and `error` can use it without depending +//! on the protocol layer. + +pub mod addr; +pub mod checksum; +pub mod codec; +pub mod config; +pub mod conn_id; +pub mod durable_file; +pub mod error; +pub mod test_support; +pub mod timeout; + +/// The width of the length prefix on a framed message. +pub type DataLenType = u32; diff --git a/crates/pb-mapper-core/src/test_support.rs b/crates/pb-mapper-core/src/test_support.rs new file mode 100644 index 0000000..12f5cee --- /dev/null +++ b/crates/pb-mapper-core/src/test_support.rs @@ -0,0 +1,14 @@ +//! Shared serialisation for tests that mutate process-global credential state. + +/// Serialises tests that set or clear the process credential. +/// +/// The credential is process-global, so two such tests running on different +/// runner threads would see each other's writes. Every test that calls +/// `set_process_msg_header_key` — in this crate and in the auth, protocol, and +/// server crates — takes this first. +/// +/// It lives here, next to the state it guards, and is unconditionally `pub` +/// rather than `#[cfg(test)]`: a test-only item is not visible to another +/// crate's tests, because each crate compiles its own test configuration. +pub static PROCESS_CREDENTIAL_TEST_LOCK: std::sync::LazyLock> = + std::sync::LazyLock::new(|| tokio::sync::Mutex::new(())); diff --git a/src/utils/timeout.rs b/crates/pb-mapper-core/src/timeout.rs similarity index 100% rename from src/utils/timeout.rs rename to crates/pb-mapper-core/src/timeout.rs diff --git a/src/common/auth.rs b/src/common/auth.rs index eeaece0..ae4f4e6 100644 --- a/src/common/auth.rs +++ b/src/common/auth.rs @@ -96,10 +96,6 @@ const ADMIN_REPLAY_RETENTION: Duration = Duration::from_secs(10 * 60); const ADMIN_REPLAY_CAPACITY: usize = 65_536; const AUDIT_RECORD_CAPACITY: usize = 4096; -#[cfg(test)] -pub(crate) static PROCESS_CREDENTIAL_TEST_LOCK: std::sync::LazyLock> = - std::sync::LazyLock::new(|| tokio::sync::Mutex::new(())); - #[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum LegacyProtocolPolicy { diff --git a/src/common/auth/tests.rs b/src/common/auth/tests.rs index e5518f1..be2755b 100644 --- a/src/common/auth/tests.rs +++ b/src/common/auth/tests.rs @@ -10,6 +10,8 @@ //! Protocol framing has its own tests under `common::message::secure::tests`; this //! module focuses on lifecycle, persistence, audit, replay, and timing-wheel behavior. +use pb_mapper_core::test_support::PROCESS_CREDENTIAL_TEST_LOCK; + use super::*; fn temp_state_dir(name: &str) -> PathBuf { diff --git a/src/common/manager.rs b/src/common/manager.rs index 0c325cf..8644927 100644 --- a/src/common/manager.rs +++ b/src/common/manager.rs @@ -1,8 +1,17 @@ -use snafu::ResultExt; +use snafu::{ResultExt, Snafu}; use tracing::instrument; use super::conn_id::{ConnId, ConnIdProvider, ConnIdTrait}; -use super::error::{MngWaitForTaskSnafu, Result}; + +/// The manager owns this rather than the core error enum: it is the only thing +/// that waits on a task channel, and it keeps `kanal` out of the bottom layer. +#[derive(Debug, Snafu)] +#[snafu(display("`TaskManager` fails while waiting for a task"))] +pub struct MngWaitForTaskError { + source: kanal::ReceiveError, +} + +type Result = std::result::Result; /// The [`ConnId::local_id`] of the server is the same as the client. and it is only generated /// by the client. The [`ConnId::remote_id`] and [`ConnId::local_id`] of the client can be used to diff --git a/src/common/message/mod.rs b/src/common/message/mod.rs index 6c20a0c..607658b 100644 --- a/src/common/message/mod.rs +++ b/src/common/message/mod.rs @@ -69,7 +69,11 @@ const CODEC_TAG_LEN: DataLenType = 16; /// For encrypted frames, the tag is appended to the payload. const MAX_MSG_LEN: DataLenType = MAX_PLAINTEXT_LEN + CODEC_TAG_LEN; -pub type DataLenType = u32; +// Defined in `pb-mapper-core` so that the checksum and error types can name it +// without depending on this module. Re-exported rather than redeclared: a second +// `pub type` would be a distinct name for the same width, and the two would read +// as unrelated at the crate boundary. +pub use pb_mapper_core::DataLenType; macro_rules! gen_read_network_with_error { ($func_name:ident, $read_method:ident, $error:expr, $return_ty:ty) => { diff --git a/src/common/message/secure/tests.rs b/src/common/message/secure/tests.rs index e957918..432820a 100644 --- a/src/common/message/secure/tests.rs +++ b/src/common/message/secure/tests.rs @@ -10,10 +10,11 @@ //! credentials, while lifecycle persistence remains covered by `common::auth::tests`. use super::*; -use crate::common::auth::{AuthConfig, LegacyProtocolPolicy, PROCESS_CREDENTIAL_TEST_LOCK}; +use crate::common::auth::{AuthConfig, LegacyProtocolPolicy}; use crate::common::checksum::{ encode_temporary_credential, parse_credential, set_process_msg_header_key, }; +use pb_mapper_core::test_support::PROCESS_CREDENTIAL_TEST_LOCK; fn temp_config() -> AuthConfig { let mut random = [0_u8; 8]; diff --git a/src/common/mod.rs b/src/common/mod.rs index 4d7513e..81e632d 100644 --- a/src/common/mod.rs +++ b/src/common/mod.rs @@ -1,8 +1,8 @@ pub mod auth; pub mod buffer; -pub mod checksum; -pub mod config; -pub mod conn_id; -pub mod error; pub mod manager; pub mod message; + +// Moved to `pb-mapper-core`. Re-exported while the split is in progress so the +// modules below keep their existing paths. +pub use pb_mapper_core::{checksum, config, conn_id, error}; diff --git a/src/lib.rs b/src/lib.rs index 7bd3c9f..5b17647 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,6 +4,16 @@ pub mod local; pub mod pb_server; pub mod utils; +// The `snafu_error_*` macros moved to `pb-mapper-core` with the error type. +// `#[macro_export]` puts them at that crate's root, so re-export them here to +// keep `crate::snafu_error_handle!` working while the split is in progress. +pub use pb_mapper_core::{ + snafu_error_get_or_continue, snafu_error_get_or_return, snafu_error_get_or_return_ok, + snafu_error_handle, +}; + +// This was missing `#[cfg(test)]`, so it compiled into every release build. +#[cfg(test)] mod tests { #[test] diff --git a/src/pb_server/admin.rs b/src/pb_server/admin.rs index e135792..f8c85dd 100644 --- a/src/pb_server/admin.rs +++ b/src/pb_server/admin.rs @@ -307,7 +307,7 @@ mod tests { } async fn inventory_query_rejects_rotation(connection_query: bool) { - let _process_credential_guard = crate::common::auth::PROCESS_CREDENTIAL_TEST_LOCK + let _process_credential_guard = pb_mapper_core::test_support::PROCESS_CREDENTIAL_TEST_LOCK .lock() .await; let state_dir = temp_state_dir(if connection_query { diff --git a/src/pb_server/runtime.rs b/src/pb_server/runtime.rs index c5d4a99..9340865 100644 --- a/src/pb_server/runtime.rs +++ b/src/pb_server/runtime.rs @@ -958,9 +958,8 @@ async fn abort_and_wait(handles: impl IntoIterator Date: Fri, 21 Aug 2026 13:40:18 +0800 Subject: [PATCH 70/74] Extract pb-mapper-auth MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit All 17 auth files, as git renames. It is the largest subsystem and the most self-contained — its only outward reference was `checksum`, now reached through `pb-mapper-core` — so it gains the most from standing alone. The 11 platform `cfg` blocks are all in here too. `pub(in crate::common::auth)` becomes `pub(crate)` throughout: the same scope, now that the module is the crate. Four items the protocol layer calls had to widen to `pub` — `admin_key`, `derive_key`, `derive_previous_key`, and `admin_cancellation_token` — which the compiler found by reporting them as dead code from inside the crate before reporting them as private from outside it. `replace_file` is gone from auth; `sync_parent_directory` stays as a four-line wrapper mapping `core`'s `io::Result` onto `AuthFailure`, so its callers here are unchanged. The replay log in the protocol layer calls `core` directly now, which drops two `io::Error::other(error.to_string())` round trips that existed only to get back the error it started with. Co-Authored-By: Claude Fable 5 --- Cargo.lock | 17 ++++ Cargo.toml | 2 + crates/pb-mapper-auth/Cargo.toml | 18 ++++ .../pb-mapper-auth/src}/actor/epoch.rs | 0 .../pb-mapper-auth/src}/actor/lifecycle.rs | 0 .../pb-mapper-auth/src}/actor/mod.rs | 0 .../pb-mapper-auth/src}/config.rs | 0 .../auth => crates/pb-mapper-auth/src}/ids.rs | 0 .../pb-mapper-auth/src}/keys.rs | 0 .../pb-mapper-auth/src}/leases.rs | 0 .../pb-mapper-auth/src/lib.rs | 20 ++--- .../src}/persistence/admin_key.rs | 27 ++---- .../pb-mapper-auth/src}/persistence/blob.rs | 10 +-- .../pb-mapper-auth/src}/persistence/fs.rs | 89 +++---------------- .../pb-mapper-auth/src}/persistence/mod.rs | 30 +++---- .../src}/persistence/snapshot.rs | 28 +++--- .../pb-mapper-auth/src}/persistence/wal.rs | 17 ++-- .../pb-mapper-auth/src}/runtime.rs | 6 +- .../pb-mapper-auth/src}/tests.rs | 0 .../pb-mapper-auth/src}/timing_wheel.rs | 0 src/common/message/secure/replay.rs | 9 +- src/common/mod.rs | 6 +- 22 files changed, 111 insertions(+), 168 deletions(-) create mode 100644 crates/pb-mapper-auth/Cargo.toml rename {src/common/auth => crates/pb-mapper-auth/src}/actor/epoch.rs (100%) rename {src/common/auth => crates/pb-mapper-auth/src}/actor/lifecycle.rs (100%) rename {src/common/auth => crates/pb-mapper-auth/src}/actor/mod.rs (100%) rename {src/common/auth => crates/pb-mapper-auth/src}/config.rs (100%) rename {src/common/auth => crates/pb-mapper-auth/src}/ids.rs (100%) rename {src/common/auth => crates/pb-mapper-auth/src}/keys.rs (100%) rename {src/common/auth => crates/pb-mapper-auth/src}/leases.rs (100%) rename src/common/auth.rs => crates/pb-mapper-auth/src/lib.rs (97%) rename {src/common/auth => crates/pb-mapper-auth/src}/persistence/admin_key.rs (91%) rename {src/common/auth => crates/pb-mapper-auth/src}/persistence/blob.rs (89%) rename {src/common/auth => crates/pb-mapper-auth/src}/persistence/fs.rs (70%) rename {src/common/auth => crates/pb-mapper-auth/src}/persistence/mod.rs (67%) rename {src/common/auth => crates/pb-mapper-auth/src}/persistence/snapshot.rs (91%) rename {src/common/auth => crates/pb-mapper-auth/src}/persistence/wal.rs (93%) rename {src/common/auth => crates/pb-mapper-auth/src}/runtime.rs (98%) rename {src/common/auth => crates/pb-mapper-auth/src}/tests.rs (100%) rename {src/common/auth => crates/pb-mapper-auth/src}/timing_wheel.rs (100%) diff --git a/Cargo.lock b/Cargo.lock index 4bf3432..114882c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1041,6 +1041,7 @@ dependencies = [ "hashbrown 0.17.1", "kanal", "parking_lot", + "pb-mapper-auth", "pb-mapper-core", "rand 0.10.0", "ring", @@ -1055,6 +1056,22 @@ dependencies = [ "uni-stream", ] +[[package]] +name = "pb-mapper-auth" +version = "0.4.0" +dependencies = [ + "parking_lot", + "pb-mapper-core", + "rand 0.10.0", + "ring", + "serde", + "serde_json", + "subtle", + "tokio", + "tokio-util", + "tracing", +] + [[package]] name = "pb-mapper-core" version = "0.4.0" diff --git a/Cargo.toml b/Cargo.toml index 214f0d5..8b9bd0e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,6 +7,7 @@ authors.workspace = true # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] +pb-mapper-auth.workspace = true pb-mapper-core.workspace = true rand.workspace = true tokio.workspace = true @@ -55,6 +56,7 @@ unwrap_used = "deny" expect_used = "deny" [workspace.dependencies] +pb-mapper-auth = { path = "crates/pb-mapper-auth" } pb-mapper-core = { path = "crates/pb-mapper-core" } rand = "0.10" diff --git a/crates/pb-mapper-auth/Cargo.toml b/crates/pb-mapper-auth/Cargo.toml new file mode 100644 index 0000000..8f053ea --- /dev/null +++ b/crates/pb-mapper-auth/Cargo.toml @@ -0,0 +1,18 @@ +[package] +name = "pb-mapper-auth" +version.workspace = true +edition.workspace = true +authors.workspace = true + +[dependencies] +pb-mapper-core.workspace = true + +parking_lot.workspace = true +rand.workspace = true +ring.workspace = true +serde.workspace = true +serde_json.workspace = true +subtle.workspace = true +tokio.workspace = true +tokio-util.workspace = true +tracing.workspace = true diff --git a/src/common/auth/actor/epoch.rs b/crates/pb-mapper-auth/src/actor/epoch.rs similarity index 100% rename from src/common/auth/actor/epoch.rs rename to crates/pb-mapper-auth/src/actor/epoch.rs diff --git a/src/common/auth/actor/lifecycle.rs b/crates/pb-mapper-auth/src/actor/lifecycle.rs similarity index 100% rename from src/common/auth/actor/lifecycle.rs rename to crates/pb-mapper-auth/src/actor/lifecycle.rs diff --git a/src/common/auth/actor/mod.rs b/crates/pb-mapper-auth/src/actor/mod.rs similarity index 100% rename from src/common/auth/actor/mod.rs rename to crates/pb-mapper-auth/src/actor/mod.rs diff --git a/src/common/auth/config.rs b/crates/pb-mapper-auth/src/config.rs similarity index 100% rename from src/common/auth/config.rs rename to crates/pb-mapper-auth/src/config.rs diff --git a/src/common/auth/ids.rs b/crates/pb-mapper-auth/src/ids.rs similarity index 100% rename from src/common/auth/ids.rs rename to crates/pb-mapper-auth/src/ids.rs diff --git a/src/common/auth/keys.rs b/crates/pb-mapper-auth/src/keys.rs similarity index 100% rename from src/common/auth/keys.rs rename to crates/pb-mapper-auth/src/keys.rs diff --git a/src/common/auth/leases.rs b/crates/pb-mapper-auth/src/leases.rs similarity index 100% rename from src/common/auth/leases.rs rename to crates/pb-mapper-auth/src/leases.rs diff --git a/src/common/auth.rs b/crates/pb-mapper-auth/src/lib.rs similarity index 97% rename from src/common/auth.rs rename to crates/pb-mapper-auth/src/lib.rs index ae4f4e6..ec1772b 100644 --- a/src/common/auth.rs +++ b/crates/pb-mapper-auth/src/lib.rs @@ -67,7 +67,7 @@ use subtle::ConstantTimeEq; use tokio::sync::{mpsc, oneshot}; use tokio_util::sync::CancellationToken; -use super::checksum::{ +use pb_mapper_core::checksum::{ AesKeyType, Credential, ENV_MSG_HEADER_KEY, MACHINE_MSG_HEADER_KEY_PATH, encode_temporary_credential, env_safe_admin_key_error, get_process_credential, is_env_safe_admin_key, parse_credential, set_process_msg_header_key, @@ -261,7 +261,7 @@ impl AuthContext { Ok(self.ensure_active()?.cancellation_token()) } - pub(crate) fn admin_cancellation_token(&self) -> Result { + pub fn admin_cancellation_token(&self) -> Result { self.require_admin()?; self.cancellation_token() } @@ -660,18 +660,16 @@ pub use config::default_auth_state_dir; #[cfg(all(test, not(any(windows, target_os = "macos"))))] pub(crate) use config::linux_default_auth_state_dir; #[cfg(test)] -pub(in crate::common::auth) use config::parse_legacy_protocol_policy; +pub(crate) use config::parse_legacy_protocol_policy; #[cfg(test)] pub(crate) use config::platform_default_auth_state_dir; #[cfg(all(test, not(any(windows, target_os = "macos"))))] -pub(in crate::common::auth) use config::{linux_system_auth_dir_usable, unix_effective_uid}; +pub(crate) use config::{linux_system_auth_dir_usable, unix_effective_uid}; mod keys; pub use keys::derive_temporary_key; #[cfg(test)] -pub(in crate::common::auth) use keys::recover_admin_key_after_rotation; -pub(in crate::common::auth) use keys::{ - load_isolated_server_admin_credential, load_server_admin_credential, -}; +pub(crate) use keys::recover_admin_key_after_rotation; +pub(crate) use keys::{load_isolated_server_admin_credential, load_server_admin_credential}; mod runtime; pub struct LegacyConnectionGuard { @@ -765,7 +763,7 @@ mod actor; use actor::{AuthActorState, run_auth_actor}; mod persistence; pub use persistence::*; -pub(in crate::common::auth) use persistence::{ +pub(crate) use persistence::{ append_audit, append_mutation, append_wal, atomic_write, auth_snapshot_path, build_snapshot, cancel_all_temporary_leases, compaction_is_allowed, empty_snapshot, fail_closed_on_uncertain_wal, hex, key_matches_existing_state, load_or_create_instance_id, @@ -775,9 +773,7 @@ pub(in crate::common::auth) use persistence::{ unix_seconds, write_admin_key, write_snapshot_and_truncate_wal, }; #[cfg(test)] -pub(in crate::common::auth) use persistence::{ - prepare_state_dir, read_instance_id_file, try_load_persisted_state, -}; +pub(crate) use persistence::{prepare_state_dir, read_instance_id_file, try_load_persisted_state}; mod ids; pub use ids::{ADMIN_KEY_ID, Generation, KeyId, SlotIndex}; mod leases; diff --git a/src/common/auth/persistence/admin_key.rs b/crates/pb-mapper-auth/src/persistence/admin_key.rs similarity index 91% rename from src/common/auth/persistence/admin_key.rs rename to crates/pb-mapper-auth/src/persistence/admin_key.rs index 33cb7fa..f523e42 100644 --- a/src/common/auth/persistence/admin_key.rs +++ b/crates/pb-mapper-auth/src/persistence/admin_key.rs @@ -5,7 +5,7 @@ use super::{ truncate_auth_wal, }; -pub(in crate::common::auth) fn load_or_create_instance_id( +pub(crate) fn load_or_create_instance_id( path: &Path, ) -> Result<[u8; INSTANCE_ID_LEN], AuthFailure> { let instance_path = path.join("server-instance-id"); @@ -17,7 +17,7 @@ pub(in crate::common::auth) fn load_or_create_instance_id( Ok(instance_id) } -pub(in crate::common::auth) fn read_instance_id_file( +pub(crate) fn read_instance_id_file( path: &Path, ) -> Result, AuthFailure> { if !path.exists() { @@ -44,7 +44,7 @@ pub(in crate::common::auth) fn read_instance_id_file( /// Reset writes that staged file, then the empty snapshot, then the live /// instance-id file. A crash after the snapshot lands would otherwise fail /// closed on the next start because the live file still has the old id. -pub(in crate::common::auth) fn recover_instance_id_after_reset( +pub(crate) fn recover_instance_id_after_reset( state_dir: &Path, admin_key: &AesKeyType, current: [u8; INSTANCE_ID_LEN], @@ -86,7 +86,7 @@ pub(in crate::common::auth) fn recover_instance_id_after_reset( Ok(next) } -pub(in crate::common::auth) fn random_instance_id() -> [u8; INSTANCE_ID_LEN] { +pub(crate) fn random_instance_id() -> [u8; INSTANCE_ID_LEN] { let mut instance_id = [0_u8; INSTANCE_ID_LEN]; let mut rng = rand::rng(); for byte in &mut instance_id { @@ -95,10 +95,7 @@ pub(in crate::common::auth) fn random_instance_id() -> [u8; INSTANCE_ID_LEN] { instance_id } -pub(in crate::common::auth) fn write_admin_key( - state_dir: &Path, - key: &str, -) -> Result<(), AuthFailure> { +pub(crate) fn write_admin_key(state_dir: &Path, key: &str) -> Result<(), AuthFailure> { atomic_write( &state_dir.join("admin.key"), format!("{key}\n").as_bytes(), @@ -156,7 +153,7 @@ pub fn write_admin_key_file(path: &Path, key: &str, force: bool) -> Result<(), A atomic_write(path, format!("{key}\n").as_bytes(), 0o600) } -pub(in crate::common::auth) fn reset_already_installed( +pub(crate) fn reset_already_installed( state_dir: &Path, admin_key: &AesKeyType, new_instance_id: &[u8; INSTANCE_ID_LEN], @@ -179,7 +176,7 @@ pub(in crate::common::auth) fn reset_already_installed( snapshot.instance_id == *new_instance_id } -pub(in crate::common::auth) fn rotation_already_installed(state_dir: &Path, new_key: &str) -> bool { +pub(crate) fn rotation_already_installed(state_dir: &Path, new_key: &str) -> bool { key_matches_existing_snapshot(Some(state_dir), new_key) && live_admin_key_matches(state_dir, new_key) } @@ -194,10 +191,7 @@ fn live_admin_key_matches(state_dir: &Path, new_key: &str) -> bool { text.trim().as_bytes() == new_key.trim().as_bytes() } -pub(in crate::common::auth) fn key_matches_existing_snapshot( - state_dir: Option<&Path>, - key: &str, -) -> bool { +pub(crate) fn key_matches_existing_snapshot(state_dir: Option<&Path>, key: &str) -> bool { let Some(state_dir) = state_dir else { return false; }; @@ -214,10 +208,7 @@ pub(in crate::common::auth) fn key_matches_existing_snapshot( open_blob(&admin_key, &bytes).is_ok() } -pub(in crate::common::auth) fn key_matches_existing_state( - state_dir: Option<&Path>, - key: &str, -) -> bool { +pub(crate) fn key_matches_existing_state(state_dir: Option<&Path>, key: &str) -> bool { if key_matches_existing_snapshot(state_dir, key) { return true; } diff --git a/src/common/auth/persistence/blob.rs b/crates/pb-mapper-auth/src/persistence/blob.rs similarity index 89% rename from src/common/auth/persistence/blob.rs rename to crates/pb-mapper-auth/src/persistence/blob.rs index 288e261..3697072 100644 --- a/src/common/auth/persistence/blob.rs +++ b/crates/pb-mapper-auth/src/persistence/blob.rs @@ -1,10 +1,7 @@ //! AEAD wrap/unwrap for snapshot and WAL payloads. use super::super::*; -pub(in crate::common::auth) fn seal_blob( - admin_key: &AesKeyType, - plain: &[u8], -) -> Result, AuthFailure> { +pub(crate) fn seal_blob(admin_key: &AesKeyType, plain: &[u8]) -> Result, AuthFailure> { let key = LessSafeKey::new( UnboundKey::new(&AES_256_GCM, admin_key) .map_err(|_| AuthFailure::internal("failed to initialize state encryption key"))?, @@ -28,10 +25,7 @@ pub(in crate::common::auth) fn seal_blob( Ok(sealed) } -pub(in crate::common::auth) fn open_blob( - admin_key: &AesKeyType, - sealed: &[u8], -) -> Result, AuthFailure> { +pub(crate) fn open_blob(admin_key: &AesKeyType, sealed: &[u8]) -> Result, AuthFailure> { if sealed.len() < STATE_BLOB_MAGIC.len() + 12 + AES_256_GCM.tag_len() || &sealed[..STATE_BLOB_MAGIC.len()] != STATE_BLOB_MAGIC { diff --git a/src/common/auth/persistence/fs.rs b/crates/pb-mapper-auth/src/persistence/fs.rs similarity index 70% rename from src/common/auth/persistence/fs.rs rename to crates/pb-mapper-auth/src/persistence/fs.rs index 4bca35c..c83100a 100644 --- a/src/common/auth/persistence/fs.rs +++ b/crates/pb-mapper-auth/src/persistence/fs.rs @@ -4,9 +4,7 @@ use super::hex; /// Create the state directory and take `auth.lock` before any credential or /// snapshot file is read or written. -pub(in crate::common::auth) fn prepare_state_dir_and_lock( - state_dir: &Path, -) -> Result, AuthFailure> { +pub(crate) fn prepare_state_dir_and_lock(state_dir: &Path) -> Result, AuthFailure> { prepare_state_dir(state_dir)?; Ok(Arc::new(acquire_state_dir_lock(state_dir)?)) } @@ -106,77 +104,18 @@ fn lock_exclusive_nonblock(file: &File) -> std::io::Result<()> { } } -pub(crate) fn replace_file(from: &Path, to: &Path) -> std::io::Result<()> { - #[cfg(windows)] - { - use std::os::windows::ffi::OsStrExt; - - const MOVEFILE_REPLACE_EXISTING: u32 = 0x1; - const MOVEFILE_WRITE_THROUGH: u32 = 0x8; - extern "system" { - fn MoveFileExW( - lp_existing_file_name: *const u16, - lp_new_file_name: *const u16, - dw_flags: u32, - ) -> i32; - } - fn wide(path: &Path) -> Vec { - path.as_os_str() - .encode_wide() - .chain(std::iter::once(0)) - .collect() - } - let from_w = wide(from); - let to_w = wide(to); - let ok = unsafe { - MoveFileExW( - from_w.as_ptr(), - to_w.as_ptr(), - MOVEFILE_REPLACE_EXISTING | MOVEFILE_WRITE_THROUGH, - ) - }; - if ok == 0 { - Err(std::io::Error::last_os_error()) - } else { - Ok(()) - } - } - #[cfg(not(windows))] - std::fs::rename(from, to) -} - +/// `core`'s durability primitive, reported as an `AuthFailure`. pub(crate) fn sync_parent_directory(path: &Path) -> Result<(), AuthFailure> { - let Some(parent) = path.parent() else { - return Ok(()); - }; - open_directory_for_sync(parent) - .and_then(|directory| directory.sync_all()) - .map_err(|error| { - AuthFailure::new( - "temporary_key_store_unavailable", - format!("failed to sync `{}`: {error}", parent.display()), - false, - ) - }) -} - -fn open_directory_for_sync(path: &Path) -> std::io::Result { - #[cfg(windows)] - { - use std::os::windows::fs::OpenOptionsExt; - const GENERIC_READ: u32 = 0x8000_0000; - const GENERIC_WRITE: u32 = 0x4000_0000; - const FILE_FLAG_BACKUP_SEMANTICS: u32 = 0x0200_0000; - OpenOptions::new() - .access_mode(GENERIC_READ | GENERIC_WRITE) - .custom_flags(FILE_FLAG_BACKUP_SEMANTICS) - .open(path) - } - #[cfg(not(windows))] - File::open(path) + pb_mapper_core::durable_file::sync_parent_directory(path).map_err(|error| { + AuthFailure::new( + "temporary_key_store_unavailable", + format!("failed to sync `{}`: {error}", path.display()), + false, + ) + }) } -pub(in crate::common::auth) fn prepare_state_dir(path: &Path) -> Result<(), AuthFailure> { +pub(crate) fn prepare_state_dir(path: &Path) -> Result<(), AuthFailure> { std::fs::create_dir_all(path).map_err(|error| { AuthFailure::new( "auth_state_unavailable", @@ -201,11 +140,7 @@ pub(in crate::common::auth) fn prepare_state_dir(path: &Path) -> Result<(), Auth Ok(()) } -pub(in crate::common::auth) fn atomic_write( - path: &Path, - data: &[u8], - mode: u32, -) -> Result<(), AuthFailure> { +pub(crate) fn atomic_write(path: &Path, data: &[u8], mode: u32) -> Result<(), AuthFailure> { if let Some(parent) = path.parent() { std::fs::create_dir_all(parent).map_err(|error| { AuthFailure::new( @@ -258,7 +193,7 @@ pub(in crate::common::auth) fn atomic_write( ) })?; drop(file); - replace_file(&temporary, path).map_err(|error| { + pb_mapper_core::durable_file::replace_file(&temporary, path).map_err(|error| { AuthFailure::new( "auth_state_unavailable", format!("failed to replace `{}`: {error}", path.display()), diff --git a/src/common/auth/persistence/mod.rs b/crates/pb-mapper-auth/src/persistence/mod.rs similarity index 67% rename from src/common/auth/persistence/mod.rs rename to crates/pb-mapper-auth/src/persistence/mod.rs index f367efa..144190d 100644 --- a/src/common/auth/persistence/mod.rs +++ b/crates/pb-mapper-auth/src/persistence/mod.rs @@ -19,39 +19,39 @@ mod snapshot; mod wal; #[cfg(test)] -pub(in crate::common::auth) use admin_key::read_instance_id_file; +pub(crate) use admin_key::read_instance_id_file; pub use admin_key::{generate_admin_key, initialize_admin_key, write_admin_key_file}; -pub(in crate::common::auth) use admin_key::{ +pub(crate) use admin_key::{ key_matches_existing_state, load_or_create_instance_id, random_instance_id, recover_instance_id_after_reset, reset_already_installed, rotation_already_installed, write_admin_key, }; -pub(in crate::common::auth) use blob::{open_blob, seal_blob}; +pub(crate) use blob::{open_blob, seal_blob}; pub use fs::acquire_state_dir_lock; #[cfg(test)] -pub(in crate::common::auth) use fs::prepare_state_dir; -pub(in crate::common::auth) use fs::{atomic_write, prepare_state_dir_and_lock}; -pub(crate) use fs::{replace_file, sync_parent_directory}; +pub(crate) use fs::prepare_state_dir; +pub(crate) use fs::sync_parent_directory; +pub(crate) use fs::{atomic_write, prepare_state_dir_and_lock}; #[cfg(test)] -pub(in crate::common::auth) use snapshot::try_load_persisted_state; -pub(in crate::common::auth) use snapshot::{ +pub(crate) use snapshot::try_load_persisted_state; +pub(crate) use snapshot::{ build_snapshot, cancel_all_temporary_leases, compaction_is_allowed, empty_snapshot, load_persisted_state, normalize_tombstone_times, push_audit_record, push_persisted_audit, split_high_slot_state, }; -pub(in crate::common::auth) use wal::{ +pub(crate) use wal::{ append_audit, append_mutation, append_wal, fail_closed_on_uncertain_wal, read_wal, truncate_auth_wal, write_snapshot_and_truncate_wal, }; -pub(in crate::common::auth) const AUTH_SNAPSHOT_FILE: &str = "auth.snapshot"; -pub(in crate::common::auth) const AUTH_WAL_FILE: &str = "auth.wal"; +pub(crate) const AUTH_SNAPSHOT_FILE: &str = "auth.snapshot"; +pub(crate) const AUTH_WAL_FILE: &str = "auth.wal"; -pub(in crate::common::auth) fn auth_snapshot_path(state_dir: &Path) -> PathBuf { +pub(crate) fn auth_snapshot_path(state_dir: &Path) -> PathBuf { state_dir.join(AUTH_SNAPSHOT_FILE) } -pub(in crate::common::auth) fn auth_wal_path(state_dir: &Path) -> PathBuf { +pub(crate) fn auth_wal_path(state_dir: &Path) -> PathBuf { state_dir.join(AUTH_WAL_FILE) } @@ -59,14 +59,14 @@ pub fn encrypted_auth_state_exists(state_dir: &Path) -> bool { auth_snapshot_path(state_dir).exists() || auth_wal_path(state_dir).exists() } -pub(in crate::common::auth) fn unix_seconds() -> u64 { +pub(crate) fn unix_seconds() -> u64 { SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap_or_default() .as_secs() } -pub(in crate::common::auth) fn hex(bytes: &[u8]) -> String { +pub(crate) fn hex(bytes: &[u8]) -> String { const DIGITS: &[u8; 16] = b"0123456789abcdef"; let mut output = String::with_capacity(bytes.len() * 2); for byte in bytes { diff --git a/src/common/auth/persistence/snapshot.rs b/crates/pb-mapper-auth/src/persistence/snapshot.rs similarity index 91% rename from src/common/auth/persistence/snapshot.rs rename to crates/pb-mapper-auth/src/persistence/snapshot.rs index c3c5dff..18d72f0 100644 --- a/src/common/auth/persistence/snapshot.rs +++ b/crates/pb-mapper-auth/src/persistence/snapshot.rs @@ -2,11 +2,11 @@ use super::super::*; use super::{auth_snapshot_path, auth_wal_path, open_blob, read_wal}; -pub(in crate::common::auth) fn compaction_is_allowed(safe_mode: bool) -> bool { +pub(crate) fn compaction_is_allowed(safe_mode: bool) -> bool { !safe_mode } -pub(in crate::common::auth) fn push_audit_record(inner: &AuthStateInner, record: AuditRecord) { +pub(crate) fn push_audit_record(inner: &AuthStateInner, record: AuditRecord) { let mut records = inner.audit_records.write(); while records.len() >= AUDIT_RECORD_CAPACITY { records.pop_front(); @@ -14,7 +14,7 @@ pub(in crate::common::auth) fn push_audit_record(inner: &AuthStateInner, record: records.push_back(record); } -pub(in crate::common::auth) fn cancel_all_temporary_leases(inner: &AuthStateInner) { +pub(crate) fn cancel_all_temporary_leases(inner: &AuthStateInner) { let slots = inner.slots(); for lease in slots.iter().filter_map(|slot| slot.lease.upgrade()) { lease.cancel_rotated(); @@ -29,7 +29,7 @@ fn snapshot_generations(inner: &AuthStateInner) -> Vec { generations } -pub(in crate::common::auth) fn split_high_slot_state( +pub(crate) fn split_high_slot_state( snapshot: &PersistedSnapshot, capacity: usize, ) -> (Vec, Vec) { @@ -43,7 +43,7 @@ pub(in crate::common::auth) fn split_high_slot_state( (high_generations, high_entries) } -pub(in crate::common::auth) fn build_snapshot( +pub(crate) fn build_snapshot( inner: &AuthStateInner, admin_replays: &VecDeque, ) -> PersistedSnapshot { @@ -79,10 +79,7 @@ pub(in crate::common::auth) fn build_snapshot( ) } -pub(in crate::common::auth) fn normalize_tombstone_times( - snapshot: &mut PersistedSnapshot, - now: u64, -) -> bool { +pub(crate) fn normalize_tombstone_times(snapshot: &mut PersistedSnapshot, now: u64) -> bool { let mut changed = false; for entry in &mut snapshot.entries { if entry.tombstoned_at.is_some() { @@ -109,7 +106,7 @@ pub(in crate::common::auth) fn normalize_tombstone_times( changed } -pub(in crate::common::auth) fn empty_snapshot( +pub(crate) fn empty_snapshot( inner: &AuthStateInner, instance_id: [u8; INSTANCE_ID_LEN], admin_replays: &VecDeque, @@ -146,7 +143,7 @@ fn snapshot_with( } } -pub(in crate::common::auth) fn load_persisted_state( +pub(crate) fn load_persisted_state( config: &AuthConfig, admin_key: &AesKeyType, instance_id: [u8; INSTANCE_ID_LEN], @@ -166,7 +163,7 @@ pub(in crate::common::auth) fn load_persisted_state( } } -pub(in crate::common::auth) fn try_load_persisted_state( +pub(crate) fn try_load_persisted_state( config: &AuthConfig, admin_key: &AesKeyType, instance_id: [u8; INSTANCE_ID_LEN], @@ -229,7 +226,7 @@ pub(in crate::common::auth) fn try_load_persisted_state( Ok(snapshot) } -pub(in crate::common::auth) fn apply_persisted_mutation( +pub(crate) fn apply_persisted_mutation( snapshot: &mut PersistedSnapshot, mutation: StateMutation, ) -> Result<(), AuthFailure> { @@ -279,10 +276,7 @@ fn snapshot_entry_mut<'a>( }) } -pub(in crate::common::auth) fn push_persisted_audit( - records: &mut VecDeque, - record: AuditRecord, -) { +pub(crate) fn push_persisted_audit(records: &mut VecDeque, record: AuditRecord) { while records.len() >= AUDIT_RECORD_CAPACITY { records.pop_front(); } diff --git a/src/common/auth/persistence/wal.rs b/crates/pb-mapper-auth/src/persistence/wal.rs similarity index 93% rename from src/common/auth/persistence/wal.rs rename to crates/pb-mapper-auth/src/persistence/wal.rs index d0797e3..a049ab8 100644 --- a/src/common/auth/persistence/wal.rs +++ b/crates/pb-mapper-auth/src/persistence/wal.rs @@ -5,7 +5,7 @@ use super::{ push_audit_record, seal_blob, sync_parent_directory, }; -pub(in crate::common::auth) fn fail_closed_on_uncertain_wal( +pub(crate) fn fail_closed_on_uncertain_wal( inner: &AuthStateInner, result: Result<(), AuthFailure>, ) -> Result<(), AuthFailure> { @@ -18,7 +18,7 @@ pub(in crate::common::auth) fn fail_closed_on_uncertain_wal( result } -pub(in crate::common::auth) fn append_mutation( +pub(crate) fn append_mutation( config: &AuthConfig, inner: &AuthStateInner, mutation: StateMutation, @@ -39,7 +39,7 @@ pub(in crate::common::auth) fn append_mutation( Ok(()) } -pub(in crate::common::auth) fn append_audit( +pub(crate) fn append_audit( config: &AuthConfig, inner: &AuthStateInner, audit: AuditRecord, @@ -52,7 +52,7 @@ pub(in crate::common::auth) fn append_audit( Ok(()) } -pub(in crate::common::auth) fn append_wal( +pub(crate) fn append_wal( config: &AuthConfig, admin_key: &AesKeyType, record: &WalRecord, @@ -125,10 +125,7 @@ pub(in crate::common::auth) fn append_wal( Ok(()) } -pub(in crate::common::auth) fn read_wal( - path: &Path, - admin_key: &AesKeyType, -) -> Result, AuthFailure> { +pub(crate) fn read_wal(path: &Path, admin_key: &AesKeyType) -> Result, AuthFailure> { let mut file = File::open(path).map_err(|error| { AuthFailure::new( "temporary_key_store_unavailable", @@ -186,7 +183,7 @@ pub(in crate::common::auth) fn read_wal( Ok(records) } -pub(in crate::common::auth) fn write_snapshot_and_truncate_wal( +pub(crate) fn write_snapshot_and_truncate_wal( config: &AuthConfig, admin_key: &AesKeyType, snapshot: &PersistedSnapshot, @@ -200,7 +197,7 @@ pub(in crate::common::auth) fn write_snapshot_and_truncate_wal( truncate_auth_wal(&config.state_dir) } -pub(in crate::common::auth) fn truncate_auth_wal(state_dir: &Path) -> Result<(), AuthFailure> { +pub(crate) fn truncate_auth_wal(state_dir: &Path) -> Result<(), AuthFailure> { let wal_path = auth_wal_path(state_dir); let created = !wal_path.exists(); let wal = OpenOptions::new() diff --git a/src/common/auth/runtime.rs b/crates/pb-mapper-auth/src/runtime.rs similarity index 98% rename from src/common/auth/runtime.rs rename to crates/pb-mapper-auth/src/runtime.rs index b5b74d8..0fb3266 100644 --- a/src/common/auth/runtime.rs +++ b/crates/pb-mapper-auth/src/runtime.rs @@ -255,11 +255,11 @@ impl AuthRuntime { }) } - pub(crate) fn admin_key(&self) -> Result { + pub fn admin_key(&self) -> Result { Ok(self.inner()?.admin_key()) } - pub(crate) fn derive_key(&self, key_id: KeyId) -> Result { + pub fn derive_key(&self, key_id: KeyId) -> Result { let inner = self.inner()?; if key_id.is_admin() { return Ok(inner.admin_key()); @@ -272,7 +272,7 @@ impl AuthRuntime { self.inner().map(|inner| inner.high().len()).unwrap_or(0) } - pub(crate) fn derive_previous_key(&self, key_id: KeyId) -> Option { + pub fn derive_previous_key(&self, key_id: KeyId) -> Option { let inner = self.inner().ok()?; let previous = inner.previous_root.read().clone()?; if key_id.is_admin() { diff --git a/src/common/auth/tests.rs b/crates/pb-mapper-auth/src/tests.rs similarity index 100% rename from src/common/auth/tests.rs rename to crates/pb-mapper-auth/src/tests.rs diff --git a/src/common/auth/timing_wheel.rs b/crates/pb-mapper-auth/src/timing_wheel.rs similarity index 100% rename from src/common/auth/timing_wheel.rs rename to crates/pb-mapper-auth/src/timing_wheel.rs diff --git a/src/common/message/secure/replay.rs b/src/common/message/secure/replay.rs index 1596966..e6334bc 100644 --- a/src/common/message/secure/replay.rs +++ b/src/common/message/secure/replay.rs @@ -247,9 +247,9 @@ impl ReplayGuard { } return Err(error); } - if created && let Err(error) = crate::common::auth::sync_parent_directory(path) { + if created && let Err(error) = pb_mapper_core::durable_file::sync_parent_directory(path) { self.log_failed = true; - return Err(std::io::Error::other(error.to_string())); + return Err(error); } Ok(()) } @@ -363,9 +363,8 @@ impl ReplayGuard { file.write_all(&live.concat())?; file.sync_all()?; drop(file); - crate::common::auth::replace_file(&temporary, path)?; - crate::common::auth::sync_parent_directory(path) - .map_err(|error| std::io::Error::other(error.to_string())) + pb_mapper_core::durable_file::replace_file(&temporary, path)?; + pb_mapper_core::durable_file::sync_parent_directory(path) })(); if result.is_err() { let _ = std::fs::remove_file(&temporary); diff --git a/src/common/mod.rs b/src/common/mod.rs index 81e632d..195b530 100644 --- a/src/common/mod.rs +++ b/src/common/mod.rs @@ -1,8 +1,8 @@ -pub mod auth; pub mod buffer; pub mod manager; pub mod message; -// Moved to `pb-mapper-core`. Re-exported while the split is in progress so the -// modules below keep their existing paths. +// Moved to `pb-mapper-core` and `pb-mapper-auth`. Re-exported while the split is +// in progress so the modules below keep their existing paths. +pub use pb_mapper_auth as auth; pub use pb_mapper_core::{checksum, config, conn_id, error}; From c064b46ec862c2cfbbd7cbf3e7cf6405900fbba8 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Fri, 21 Aug 2026 13:42:36 +0800 Subject: [PATCH 71/74] Extract pb-mapper-protocol Message framing, the v2 secure session, forwarding, and `buffer`, which moves with them because the framing is its only consumer. It sits above auth rather than beside it: `message::secure` reads `AuthRuntime` and `KeyId`, and auth never looks back the other way. `DataLenType` is re-exported from `core` instead of redeclared, so the protocol and the checksum that validates it name the same type. The `#![allow(async_fn_in_trait)]` that used to sit on `mod common` in the old lib.rs is now at this crate's root, where the traits that need it live, with a note on why sending them across tasks is not a case that arises. Co-Authored-By: Claude Fable 5 --- Cargo.lock | 19 +++++++++++++++ Cargo.toml | 4 +++- crates/pb-mapper-protocol/Cargo.toml | 23 +++++++++++++++++++ .../pb-mapper-protocol/src}/buffer.rs | 8 +++---- .../pb-mapper-protocol/src}/command.rs | 6 ++--- .../pb-mapper-protocol/src}/forward.rs | 20 ++++++++-------- .../pb-mapper-protocol/src/lib.rs | 17 ++++++++++---- .../pb-mapper-protocol/src}/secure.rs | 8 +++---- .../src}/secure/first_flight.rs | 0 .../pb-mapper-protocol/src}/secure/frame.rs | 2 +- .../pb-mapper-protocol/src}/secure/limiter.rs | 2 +- .../pb-mapper-protocol/src}/secure/replay.rs | 2 +- .../pb-mapper-protocol/src}/secure/tests.rs | 8 +++---- src/common/mod.rs | 8 +++---- 14 files changed, 89 insertions(+), 38 deletions(-) create mode 100644 crates/pb-mapper-protocol/Cargo.toml rename {src/common => crates/pb-mapper-protocol/src}/buffer.rs (92%) rename {src/common/message => crates/pb-mapper-protocol/src}/command.rs (98%) rename {src/common/message => crates/pb-mapper-protocol/src}/forward.rs (98%) rename src/common/message/mod.rs => crates/pb-mapper-protocol/src/lib.rs (96%) rename {src/common/message => crates/pb-mapper-protocol/src}/secure.rs (99%) rename {src/common/message => crates/pb-mapper-protocol/src}/secure/first_flight.rs (100%) rename {src/common/message => crates/pb-mapper-protocol/src}/secure/frame.rs (99%) rename {src/common/message => crates/pb-mapper-protocol/src}/secure/limiter.rs (98%) rename {src/common/message => crates/pb-mapper-protocol/src}/secure/replay.rs (99%) rename {src/common/message => crates/pb-mapper-protocol/src}/secure/tests.rs (99%) diff --git a/Cargo.lock b/Cargo.lock index 114882c..64182ad 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1043,6 +1043,7 @@ dependencies = [ "parking_lot", "pb-mapper-auth", "pb-mapper-core", + "pb-mapper-protocol", "rand 0.10.0", "ring", "serde", @@ -1107,6 +1108,24 @@ dependencies = [ "uni-stream", ] +[[package]] +name = "pb-mapper-protocol" +version = "0.4.0" +dependencies = [ + "bytes", + "parking_lot", + "pb-mapper-auth", + "pb-mapper-core", + "rand 0.10.0", + "ring", + "serde", + "serde_json", + "snafu", + "tokio", + "tracing", + "uni-stream", +] + [[package]] name = "percent-encoding" version = "2.3.2" diff --git a/Cargo.toml b/Cargo.toml index 8b9bd0e..34c4a96 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,6 +9,7 @@ authors.workspace = true [dependencies] pb-mapper-auth.workspace = true pb-mapper-core.workspace = true +pb-mapper-protocol.workspace = true rand.workspace = true tokio.workspace = true tokio-util.workspace = true @@ -32,7 +33,7 @@ parking_lot.workspace = true dotenvy.workspace = true [features] -udp-timeout = ["uni-stream/udp-timeout"] +udp-timeout = ["uni-stream/udp-timeout", "pb-mapper-protocol/udp-timeout"] [workspace] members = ["crates/*", "ui/native/pb_mapper_ffi"] @@ -58,6 +59,7 @@ expect_used = "deny" [workspace.dependencies] pb-mapper-auth = { path = "crates/pb-mapper-auth" } pb-mapper-core = { path = "crates/pb-mapper-core" } +pb-mapper-protocol = { path = "crates/pb-mapper-protocol" } rand = "0.10" tokio = { version = "1", features = ["full"] } diff --git a/crates/pb-mapper-protocol/Cargo.toml b/crates/pb-mapper-protocol/Cargo.toml new file mode 100644 index 0000000..eae6617 --- /dev/null +++ b/crates/pb-mapper-protocol/Cargo.toml @@ -0,0 +1,23 @@ +[package] +name = "pb-mapper-protocol" +version.workspace = true +edition.workspace = true +authors.workspace = true + +[dependencies] +pb-mapper-auth.workspace = true +pb-mapper-core.workspace = true + +bytes.workspace = true +parking_lot.workspace = true +rand.workspace = true +ring.workspace = true +serde.workspace = true +serde_json.workspace = true +snafu.workspace = true +tokio.workspace = true +tracing.workspace = true +uni-stream.workspace = true + +[features] +udp-timeout = ["uni-stream/udp-timeout"] diff --git a/src/common/buffer.rs b/crates/pb-mapper-protocol/src/buffer.rs similarity index 92% rename from src/common/buffer.rs rename to crates/pb-mapper-protocol/src/buffer.rs index 57e3326..b53f572 100644 --- a/src/common/buffer.rs +++ b/crates/pb-mapper-protocol/src/buffer.rs @@ -3,7 +3,7 @@ use snafu::ResultExt; use tokio::io::AsyncReadExt; -use super::error::MsgNetworkReadBufferdRawDataSnafu; +use pb_mapper_core::error::MsgNetworkReadBufferdRawDataSnafu; const INIT_BUF_SIZE: usize = 8 * 1024; const MAX_BUF_SIZE: usize = 8 * 1024 * 1024; @@ -104,7 +104,7 @@ impl BufferGetter for CommonBuffer { /// This trait is used for buffered reads where the packet length is not known pub trait BufferedReader { - async fn read(&mut self) -> super::error::Result<&'_ [u8]>; + async fn read(&mut self) -> pb_mapper_core::error::Result<&'_ [u8]>; } pub struct BufferReader<'a, T> { @@ -119,7 +119,7 @@ impl<'reader, T: AsyncReadExt + Unpin> BufferReader<'reader, T> { } } - async fn read_inner(&mut self) -> super::error::Result<&[u8]> { + async fn read_inner(&mut self) -> pb_mapper_core::error::Result<&[u8]> { if self.buffer.need_resize() { self.buffer.dyn_resize() } @@ -134,7 +134,7 @@ impl<'reader, T: AsyncReadExt + Unpin> BufferReader<'reader, T> { } impl<'reader, T: AsyncReadExt + Unpin> BufferedReader for BufferReader<'reader, T> { - async fn read(&mut self) -> super::error::Result<&'_ [u8]> { + async fn read(&mut self) -> pb_mapper_core::error::Result<&'_ [u8]> { self.read_inner().await } } diff --git a/src/common/message/command.rs b/crates/pb-mapper-protocol/src/command.rs similarity index 98% rename from src/common/message/command.rs rename to crates/pb-mapper-protocol/src/command.rs index f660c6e..f52f087 100644 --- a/src/common/message/command.rs +++ b/crates/pb-mapper-protocol/src/command.rs @@ -1,11 +1,11 @@ use serde::{Deserialize, Serialize}; use snafu::ResultExt; -use super::super::error::{MsgSerializeSnafu, Result}; -use crate::common::auth::{ +use pb_mapper_auth::{ AuthStatus, IssuedTemporaryKey, KeyPage, LegacyProtocolPolicy, TemporaryKeyMetadata, }; -use crate::common::checksum::AesKeyType; +use pb_mapper_core::checksum::AesKeyType; +use pb_mapper_core::error::{MsgSerializeSnafu, Result}; pub const CONTROL_PROTOCOL_V2: u16 = 2; diff --git a/src/common/message/forward.rs b/crates/pb-mapper-protocol/src/forward.rs similarity index 98% rename from src/common/message/forward.rs rename to crates/pb-mapper-protocol/src/forward.rs index 7207ee1..d8af9a7 100644 --- a/src/common/message/forward.rs +++ b/crates/pb-mapper-protocol/src/forward.rs @@ -6,16 +6,16 @@ use std::time::Duration; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::time::Instant; -use super::super::buffer::{BufferReader, BufferedReader}; -use super::error::{FwdNetworkWriteWithNormalSnafu, Result}; use super::{ CodecMessageReader, CodecMessageWriter, MessageReader, MessageWriter, NormalMessageReader, NormalMessageWriter, }; -use crate::common::checksum::AesKeyType; -use crate::common::config::duration_from_env; -use crate::snafu_error_get_or_return_ok; -use crate::utils::codec::{Decryptor, Encryptor}; +use crate::buffer::{BufferReader, BufferedReader}; +use pb_mapper_core::checksum::AesKeyType; +use pb_mapper_core::codec::{Decryptor, Encryptor}; +use pb_mapper_core::config::duration_from_env; +use pb_mapper_core::error::{FwdNetworkWriteWithNormalSnafu, Result}; +use pb_mapper_core::snafu_error_get_or_return_ok; use uni_stream::stream::{StreamSplit, TcpStreamImpl, UdpStreamImpl}; use uni_stream::udp::{UdpStreamReadHalf, UdpStreamWriteHalf}; @@ -533,7 +533,7 @@ impl DatagramReader for UdpStreamReadHalf { async fn recv(&mut self) -> Result { self.recv_datagram() .await - .map_err(|e| super::error::Error::MsgForward { + .map_err(|e| pb_mapper_core::error::Error::MsgForward { action: "read", source: e, }) @@ -544,7 +544,7 @@ impl DatagramWriter for UdpStreamWriteHalf<'_> { async fn send(&mut self, src: &[u8]) -> Result<()> { self.send_datagram(src) .await - .map_err(|e| super::error::Error::MsgForward { + .map_err(|e| pb_mapper_core::error::Error::MsgForward { action: "write", source: e, }) @@ -689,8 +689,8 @@ mod tests { use std::time::Duration; use super::*; - use crate::common::config::parse_duration; - use crate::common::error::Error; + use pb_mapper_core::config::parse_duration; + use pb_mapper_core::error::Error; use tokio::sync::Notify; enum ReadAction { diff --git a/src/common/message/mod.rs b/crates/pb-mapper-protocol/src/lib.rs similarity index 96% rename from src/common/message/mod.rs rename to crates/pb-mapper-protocol/src/lib.rs index 607658b..29ae36c 100644 --- a/src/common/message/mod.rs +++ b/crates/pb-mapper-protocol/src/lib.rs @@ -1,24 +1,31 @@ //! Define message protocols and tools for reading and writing //! messages +//! +//! The reader and writer traits are `async fn` in a public trait, which cannot +//! state its auto-trait bounds. That is deliberate: these are only ever awaited +//! on the connection task that owns the stream, never sent across one. +#![allow(async_fn_in_trait)] + +pub mod buffer; pub mod command; pub mod forward; pub mod secure; use snafu::{ResultExt, ensure}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; -use super::buffer::{BufferGetter, CommonBuffer, FixedSizeBuffer}; -use super::checksum::{ +use crate::buffer::{BufferGetter, CommonBuffer, FixedSizeBuffer}; +use pb_mapper_core::checksum::{ AesKeyType, get_checksum, get_checksum_for_key, get_msg_header_key, process_checksum_is_ready, valid_checksum, valid_checksum_for_key, }; -use super::error::{ +use pb_mapper_core::codec::{Aes256GcmDeCodec, Aes256GcmEnCodec, Decryptor, Encryptor}; +use pb_mapper_core::error::MsgDatalenExceededSnafu; +use pb_mapper_core::error::{ self, MsgDatalenValidateSnafu, MsgNetworkReadBodySnafu, MsgNetworkReadCheckSumSnafu, MsgNetworkReadDatalenSnafu, MsgNetworkWriteBodySnafu, MsgNetworkWriteCheckSumSnafu, MsgNetworkWriteCodecMsgSnafu, MsgNetworkWriteCodecTagSnafu, MsgNetworkWriteDatalenSnafu, Result, }; -use crate::common::error::MsgDatalenExceededSnafu; -use crate::utils::codec::{Aes256GcmDeCodec, Aes256GcmEnCodec, Decryptor, Encryptor}; /// This message protocol contains header and body, and the header /// includes checksum, datalen,respectively, u32, u32, where datalen diff --git a/src/common/message/secure.rs b/crates/pb-mapper-protocol/src/secure.rs similarity index 99% rename from src/common/message/secure.rs rename to crates/pb-mapper-protocol/src/secure.rs index de73081..9468c50 100644 --- a/src/common/message/secure.rs +++ b/crates/pb-mapper-protocol/src/secure.rs @@ -31,14 +31,14 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt}; use super::{ CodecMessageReader, CodecMessageWriter, DataLenType, MAX_MSG_LEN, MessageReader, MessageWriter, }; -use crate::common::auth::{ +use pb_mapper_auth::{ ADMIN_KEY_ID, AuthContext, AuthFailure, AuthRuntime, KeyId, LegacyConnectionGuard, }; -use crate::common::checksum::{ +use pb_mapper_core::checksum::{ AesKeyType, Credential, get_process_credential, valid_checksum_for_key, }; -use crate::common::error::{Error, Result}; -use crate::utils::codec::{Aes256GcmDeCodec, Aes256GcmEnCodec, Decryptor}; +use pb_mapper_core::codec::{Aes256GcmDeCodec, Aes256GcmEnCodec, Decryptor}; +use pb_mapper_core::error::{Error, Result}; pub const PROTOCOL_V2_MAGIC: [u8; 4] = *b"PBM2"; pub const PROTOCOL_V2_VERSION: u8 = 2; diff --git a/src/common/message/secure/first_flight.rs b/crates/pb-mapper-protocol/src/secure/first_flight.rs similarity index 100% rename from src/common/message/secure/first_flight.rs rename to crates/pb-mapper-protocol/src/secure/first_flight.rs diff --git a/src/common/message/secure/frame.rs b/crates/pb-mapper-protocol/src/secure/frame.rs similarity index 99% rename from src/common/message/secure/frame.rs rename to crates/pb-mapper-protocol/src/secure/frame.rs index 14536f8..09a3c8f 100644 --- a/src/common/message/secure/frame.rs +++ b/crates/pb-mapper-protocol/src/secure/frame.rs @@ -10,7 +10,7 @@ //! The initial reader can impose a smaller pre-authentication limit before allocating //! a body; continuation frames retain the normal protocol maximum. -use crate::common::auth::KeyId; +use pb_mapper_auth::KeyId; use super::*; diff --git a/src/common/message/secure/limiter.rs b/crates/pb-mapper-protocol/src/secure/limiter.rs similarity index 98% rename from src/common/message/secure/limiter.rs rename to crates/pb-mapper-protocol/src/secure/limiter.rs index 46e7dc6..e67a7a6 100644 --- a/src/common/message/secure/limiter.rs +++ b/crates/pb-mapper-protocol/src/secure/limiter.rs @@ -9,7 +9,7 @@ //! decisions: every authentication failure is still rejected, only duplicate logging //! is coalesced. -use crate::common::auth::KeyId; +use pb_mapper_auth::KeyId; #[derive(Clone, Copy, Debug)] pub struct FailureLogDecision { diff --git a/src/common/message/secure/replay.rs b/crates/pb-mapper-protocol/src/secure/replay.rs similarity index 99% rename from src/common/message/secure/replay.rs rename to crates/pb-mapper-protocol/src/secure/replay.rs index e6334bc..bb25e2d 100644 --- a/src/common/message/secure/replay.rs +++ b/crates/pb-mapper-protocol/src/secure/replay.rs @@ -15,7 +15,7 @@ //! Per-credential counts stop one tenant from filling the shared filter with //! unique salts before the request payload is decoded. -use crate::common::auth::KeyId; +use pb_mapper_auth::KeyId; use std::collections::HashMap; use std::fs::{File, OpenOptions}; diff --git a/src/common/message/secure/tests.rs b/crates/pb-mapper-protocol/src/secure/tests.rs similarity index 99% rename from src/common/message/secure/tests.rs rename to crates/pb-mapper-protocol/src/secure/tests.rs index 432820a..a1f982c 100644 --- a/src/common/message/secure/tests.rs +++ b/crates/pb-mapper-protocol/src/secure/tests.rs @@ -10,8 +10,8 @@ //! credentials, while lifecycle persistence remains covered by `common::auth::tests`. use super::*; -use crate::common::auth::{AuthConfig, LegacyProtocolPolicy}; -use crate::common::checksum::{ +use pb_mapper_auth::{AuthConfig, LegacyProtocolPolicy}; +use pb_mapper_core::checksum::{ encode_temporary_credential, parse_credential, set_process_msg_header_key, }; use pb_mapper_core::test_support::PROCESS_CREDENTIAL_TEST_LOCK; @@ -71,7 +71,7 @@ async fn temporary_credential_authenticates_without_storing_secret() { .await .unwrap(); let Credential::Temporary { key_id, key } = - crate::common::checksum::parse_credential(&issued.credential).unwrap() + pb_mapper_core::checksum::parse_credential(&issued.credential).unwrap() else { panic!("expected temporary credential") }; @@ -435,7 +435,7 @@ async fn oversized_initial_frame_is_rejected_before_reading_its_body() { #[tokio::test] async fn oversized_legacy_initial_frame_is_rejected_before_reading_its_body() { - use crate::common::checksum::get_checksum_for_key; + use pb_mapper_core::checksum::get_checksum_for_key; let admin = *b"0123456789abcdefghijklmnopqrstuv"; let config = temp_config(); diff --git a/src/common/mod.rs b/src/common/mod.rs index 195b530..4cf23d9 100644 --- a/src/common/mod.rs +++ b/src/common/mod.rs @@ -1,8 +1,8 @@ -pub mod buffer; pub mod manager; -pub mod message; -// Moved to `pb-mapper-core` and `pb-mapper-auth`. Re-exported while the split is -// in progress so the modules below keep their existing paths. +// Moved out to their own crates. Re-exported while the split is in progress so +// the modules below keep their existing paths. pub use pb_mapper_auth as auth; pub use pb_mapper_core::{checksum, config, conn_id, error}; +pub use pb_mapper_protocol as message; +pub use pb_mapper_protocol::buffer; From 96a627e5e277ace349952a0c266f5b9ea300659e Mon Sep 17 00:00:00 2001 From: LB7666 Date: Fri, 21 Aug 2026 13:52:14 +0800 Subject: [PATCH 72/74] Split out server, client, and the CLI; root becomes virtual MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The last three crates, and the end of `src/`. `pb-mapper-server` and `pb-mapper-client` are peers, not layers — they have never referenced each other, in either direction. `manager` moves into the server crate, whose routing runtime is its only caller. `pb-mapper-cli` owns the binary, plus the integration tests and examples, because it is the crate that depends on everything they exercise. The binary is still called `pb-mapper`: the crate is named differently but the file is still `src/bin/pb-mapper.rs`, and `cargo build --bin pb-mapper` resolves it from the workspace root, so the release workflows, both Dockerfiles, and the install scripts need no change. Verified by building it and checking the path. The 26 downstream files now name the crate they actually want instead of going through a facade. In the three error modules that alias the old module and then write `common::error::Error` on nearly every variant, the alias is retargeted rather than the variants rewritten. The FFI crate replaces its one `path = "../../../"` dependency with the five crates it uses. `libpb_mapper_ffi.so` and all 28 exported symbols are byte-identical, which is what the Dart loader and the release-ui hash checks depend on. `test_serde_mapper_header` moved to `protocol::command`, next to the type it pins; it was in the old crate root and the test count caught it going missing. 136 tests pass, matching the pre-split baseline exactly. Co-Authored-By: Claude Fable 5 --- Cargo.lock | 64 ++++++++++----- Cargo.toml | 82 +++++-------------- crates/pb-mapper-cli/Cargo.toml | 37 +++++++++ .../examples}/echo_tcp_client.rs | 0 .../examples}/echo_tcp_server.rs | 0 .../examples}/echo_udp_client.rs | 2 +- .../examples}/echo_udp_server.rs | 2 +- .../examples}/pb_local_client.rs | 4 +- .../examples}/pb_local_server.rs | 4 +- .../pb-mapper-cli/examples}/pb_server.rs | 4 +- .../pb-mapper-cli/src}/bin/pb-mapper.rs | 32 ++++---- .../pb-mapper-cli/src}/bin/pb-mapper/admin.rs | 0 {tests => crates/pb-mapper-cli/tests}/.env | 0 .../pb-mapper-cli/tests}/regression.rs | 18 ++-- .../pb-mapper-cli/tests}/test_delay.rs | 14 ++-- crates/pb-mapper-client/Cargo.toml | 18 ++++ .../pb-mapper-client/src}/client/error.rs | 4 +- .../pb-mapper-client/src}/client/mod.rs | 16 ++-- .../pb-mapper-client/src}/client/status.rs | 8 +- .../pb-mapper-client/src}/client/stream.rs | 14 ++-- .../pb-mapper-client/src/lib.rs | 0 .../pb-mapper-client/src}/server/error.rs | 4 +- .../pb-mapper-client/src}/server/mod.rs | 24 +++--- .../pb-mapper-client/src}/server/stream.rs | 14 ++-- crates/pb-mapper-protocol/src/command.rs | 28 +++++++ crates/pb-mapper-server/Cargo.toml | 24 ++++++ .../pb-mapper-server/src}/admin.rs | 31 ++++--- .../pb-mapper-server/src}/client.rs | 24 +++--- .../pb-mapper-server/src}/connection.rs | 20 ++--- .../pb-mapper-server/src}/error.rs | 6 +- .../pb-mapper-server/src/lib.rs | 28 ++++--- .../pb-mapper-server/src}/manager.rs | 2 +- .../pb-mapper-server/src}/runtime.rs | 2 +- .../pb-mapper-server/src}/server.rs | 14 ++-- .../pb-mapper-server/src}/status.rs | 8 +- src/common/mod.rs | 8 -- src/lib.rs | 39 --------- src/utils.rs | 3 - ui/native/pb_mapper_ffi/Cargo.toml | 10 ++- ui/native/pb_mapper_ffi/src/state.rs | 22 ++--- 40 files changed, 341 insertions(+), 293 deletions(-) create mode 100644 crates/pb-mapper-cli/Cargo.toml rename {examples => crates/pb-mapper-cli/examples}/echo_tcp_client.rs (100%) rename {examples => crates/pb-mapper-cli/examples}/echo_tcp_server.rs (100%) rename {examples => crates/pb-mapper-cli/examples}/echo_udp_client.rs (93%) rename {examples => crates/pb-mapper-cli/examples}/echo_udp_server.rs (96%) rename {examples => crates/pb-mapper-cli/examples}/pb_local_client.rs (72%) rename {examples => crates/pb-mapper-cli/examples}/pb_local_server.rs (80%) rename {examples => crates/pb-mapper-cli/examples}/pb_server.rs (57%) rename {src => crates/pb-mapper-cli/src}/bin/pb-mapper.rs (96%) rename {src => crates/pb-mapper-cli/src}/bin/pb-mapper/admin.rs (100%) rename {tests => crates/pb-mapper-cli/tests}/.env (100%) rename {tests => crates/pb-mapper-cli/tests}/regression.rs (99%) rename {tests => crates/pb-mapper-cli/tests}/test_delay.rs (97%) create mode 100644 crates/pb-mapper-client/Cargo.toml rename {src/local => crates/pb-mapper-client/src}/client/error.rs (93%) rename {src/local => crates/pb-mapper-client/src}/client/mod.rs (97%) rename {src/local => crates/pb-mapper-client/src}/client/status.rs (94%) rename {src/local => crates/pb-mapper-client/src}/client/stream.rs (89%) rename src/local/mod.rs => crates/pb-mapper-client/src/lib.rs (100%) rename {src/local => crates/pb-mapper-client/src}/server/error.rs (94%) rename {src/local => crates/pb-mapper-client/src}/server/mod.rs (98%) rename {src/local => crates/pb-mapper-client/src}/server/stream.rs (91%) create mode 100644 crates/pb-mapper-server/Cargo.toml rename {src/pb_server => crates/pb-mapper-server/src}/admin.rs (94%) rename {src/pb_server => crates/pb-mapper-server/src}/client.rs (97%) rename {src/pb_server => crates/pb-mapper-server/src}/connection.rs (96%) rename {src/pb_server => crates/pb-mapper-server/src}/error.rs (98%) rename src/pb_server/mod.rs => crates/pb-mapper-server/src/lib.rs (94%) rename {src/common => crates/pb-mapper-server/src}/manager.rs (99%) rename {src/pb_server => crates/pb-mapper-server/src}/runtime.rs (99%) rename {src/pb_server => crates/pb-mapper-server/src}/server.rs (98%) rename {src/pb_server => crates/pb-mapper-server/src}/status.rs (94%) delete mode 100644 src/common/mod.rs delete mode 100644 src/lib.rs delete mode 100644 src/utils.rs diff --git a/Cargo.lock b/Cargo.lock index 64182ad..3e44e6f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1030,47 +1030,52 @@ dependencies = [ ] [[package]] -name = "pb-mapper" +name = "pb-mapper-auth" version = "0.4.0" dependencies = [ - "base64", - "better_mimalloc_rs", - "bytes", - "clap", - "dotenvy", - "hashbrown 0.17.1", - "kanal", "parking_lot", - "pb-mapper-auth", "pb-mapper-core", - "pb-mapper-protocol", "rand 0.10.0", "ring", "serde", "serde_json", - "snafu", "subtle", "tokio", "tokio-util", "tracing", - "tracing-subscriber", - "uni-stream", ] [[package]] -name = "pb-mapper-auth" +name = "pb-mapper-cli" version = "0.4.0" dependencies = [ - "parking_lot", + "better_mimalloc_rs", + "clap", + "dotenvy", + "pb-mapper-auth", + "pb-mapper-client", "pb-mapper-core", + "pb-mapper-protocol", + "pb-mapper-server", "rand 0.10.0", - "ring", - "serde", "serde_json", - "subtle", "tokio", "tokio-util", "tracing", + "uni-stream", +] + +[[package]] +name = "pb-mapper-client" +version = "0.4.0" +dependencies = [ + "pb-mapper-core", + "pb-mapper-protocol", + "serde_json", + "snafu", + "tokio", + "tracing", + "uni-stream", ] [[package]] @@ -1098,7 +1103,11 @@ dependencies = [ "clap", "dirs", "parking_lot", - "pb-mapper", + "pb-mapper-auth", + "pb-mapper-client", + "pb-mapper-core", + "pb-mapper-protocol", + "pb-mapper-server", "serde", "serde_json", "tokio", @@ -1126,6 +1135,23 @@ dependencies = [ "uni-stream", ] +[[package]] +name = "pb-mapper-server" +version = "0.4.0" +dependencies = [ + "hashbrown 0.17.1", + "kanal", + "pb-mapper-auth", + "pb-mapper-core", + "pb-mapper-protocol", + "rand 0.10.0", + "snafu", + "tokio", + "tokio-util", + "tracing", + "uni-stream", +] + [[package]] name = "percent-encoding" version = "2.3.2" diff --git a/Cargo.toml b/Cargo.toml index 34c4a96..9ddacda 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,45 +1,10 @@ -[package] -name = "pb-mapper" -version.workspace = true -edition.workspace = true -authors.workspace = true - -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html - -[dependencies] -pb-mapper-auth.workspace = true -pb-mapper-core.workspace = true -pb-mapper-protocol.workspace = true -rand.workspace = true -tokio.workspace = true -tokio-util.workspace = true -snafu.workspace = true -serde.workspace = true -serde_json.workspace = true -tracing.workspace = true -tracing-subscriber.workspace = true -hashbrown.workspace = true -clap.workspace = true -better_mimalloc_rs.workspace = true -bytes.workspace = true -ring.workspace = true -uni-stream.workspace = true -kanal.workspace = true -base64.workspace = true -subtle.workspace = true -parking_lot.workspace = true - -[dev-dependencies] -dotenvy.workspace = true - -[features] -udp-timeout = ["uni-stream/udp-timeout", "pb-mapper-protocol/udp-timeout"] - [workspace] +# Every library and the CLI live under `crates/`. The FFI cdylib stays next to +# the Flutter code that loads it. members = ["crates/*", "ui/native/pb_mapper_ffi"] exclude = ["deps/uni-stream", "deps/kanal"] -# Spelled out because a virtual manifest does not infer the resolver from the -# edition, and this root becomes virtual once `src/` is empty. +# Spelled out: a virtual manifest does not infer the resolver from the edition, +# and without this it silently falls back to resolver 1. resolver = "3" [workspace.package] @@ -49,40 +14,35 @@ version = "0.4.0" authors = ["L_B__"] edition = "2024" -[workspace.lints.clippy] -# The FFI crate has denied these since it was written — a panic crossing the C -# ABI is undefined behaviour rather than a stack trace. Everything here is -# reachable from that boundary, so the whole workspace inherits it. -unwrap_used = "deny" -expect_used = "deny" - [workspace.dependencies] pb-mapper-auth = { path = "crates/pb-mapper-auth" } +pb-mapper-client = { path = "crates/pb-mapper-client" } pb-mapper-core = { path = "crates/pb-mapper-core" } pb-mapper-protocol = { path = "crates/pb-mapper-protocol" } +pb-mapper-server = { path = "crates/pb-mapper-server" } +base64 = "0.23.1" +better_mimalloc_rs = { version = "0.1.2", features = ["config"] } +bytes = "1.11" +clap = { version = "4.5", features = ["derive"] } +dirs = "6.0.0" +dotenvy = "0.15.7" +hashbrown = { version = "0.17.1" } +hickory-resolver = { version = "0.26.1" } +kanal = { git = "https://github.com/acking-you/kanal.git", branch = "dev/pb-mapper" } +parking_lot = "0.12" rand = "0.10" -tokio = { version = "1", features = ["full"] } -tokio-util = "0.7" -snafu = "0.9.2" +ring = "0.17.14" serde = { version = "1.0", features = ["derive"] } serde_json = { version = "1.0", default-features = false, features = ["alloc"] } +snafu = "0.9.2" +subtle = "2.6.1" +tokio = { version = "1", features = ["full"] } +tokio-util = "0.7" tracing = "0.1.40" tracing-subscriber = { version = "0.3.18", features = [ "env-filter", "fmt", "json", ], default-features = true } -hashbrown = { version = "0.17.1" } -clap = { version = "4.5", features = ["derive"] } -better_mimalloc_rs = { version = "0.1.2", features = ["config"] } -bytes = "1.11" -hickory-resolver = { version = "0.26.1" } -ring = "0.17.14" -base64 = "0.23.1" -subtle = "2.6.1" -parking_lot = "0.12" -dirs = "6.0.0" -dotenvy = "0.15.7" uni-stream = { git = "https://github.com/acking-you/uni-stream.git", branch = "master" } -kanal = { git = "https://github.com/acking-you/kanal.git", branch = "dev/pb-mapper" } diff --git a/crates/pb-mapper-cli/Cargo.toml b/crates/pb-mapper-cli/Cargo.toml new file mode 100644 index 0000000..a560734 --- /dev/null +++ b/crates/pb-mapper-cli/Cargo.toml @@ -0,0 +1,37 @@ +[package] +name = "pb-mapper-cli" +version.workspace = true +edition.workspace = true +authors.workspace = true + +# The binary keeps the name `pb-mapper`, discovered from +# `src/bin/pb-mapper.rs`. Release workflows, both Dockerfiles, and the install +# scripts all hardcode it, and `cargo build --bin pb-mapper` resolves it from +# the workspace root regardless of which crate holds it. + +[dependencies] +pb-mapper-auth.workspace = true +pb-mapper-client.workspace = true +pb-mapper-core.workspace = true +pb-mapper-protocol.workspace = true +pb-mapper-server.workspace = true + +better_mimalloc_rs.workspace = true +clap.workspace = true +serde_json.workspace = true +tokio.workspace = true +tokio-util.workspace = true +tracing.workspace = true +uni-stream.workspace = true + +[dev-dependencies] +dotenvy.workspace = true +rand.workspace = true + +[features] +udp-timeout = [ + "uni-stream/udp-timeout", + "pb-mapper-protocol/udp-timeout", + "pb-mapper-client/udp-timeout", + "pb-mapper-server/udp-timeout", +] diff --git a/examples/echo_tcp_client.rs b/crates/pb-mapper-cli/examples/echo_tcp_client.rs similarity index 100% rename from examples/echo_tcp_client.rs rename to crates/pb-mapper-cli/examples/echo_tcp_client.rs diff --git a/examples/echo_tcp_server.rs b/crates/pb-mapper-cli/examples/echo_tcp_server.rs similarity index 100% rename from examples/echo_tcp_server.rs rename to crates/pb-mapper-cli/examples/echo_tcp_server.rs diff --git a/examples/echo_udp_client.rs b/crates/pb-mapper-cli/examples/echo_udp_client.rs similarity index 93% rename from examples/echo_udp_client.rs rename to crates/pb-mapper-cli/examples/echo_udp_client.rs index 1a69b4e..9db1a3a 100644 --- a/examples/echo_udp_client.rs +++ b/crates/pb-mapper-cli/examples/echo_udp_client.rs @@ -1,6 +1,6 @@ use std::error::Error; -use pb_mapper::common::config::init_tracing; +use pb_mapper_core::config::init_tracing; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use uni_stream::udp::UdpStream; diff --git a/examples/echo_udp_server.rs b/crates/pb-mapper-cli/examples/echo_udp_server.rs similarity index 96% rename from examples/echo_udp_server.rs rename to crates/pb-mapper-cli/examples/echo_udp_server.rs index 135aae1..23a8574 100644 --- a/examples/echo_udp_server.rs +++ b/crates/pb-mapper-cli/examples/echo_udp_server.rs @@ -2,7 +2,7 @@ use std::error::Error; use std::net::SocketAddr; use std::str::FromStr; -use pb_mapper::common::config::init_tracing; +use pb_mapper_core::config::init_tracing; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use uni_stream::udp::UdpListener; diff --git a/examples/pb_local_client.rs b/crates/pb-mapper-cli/examples/pb_local_client.rs similarity index 72% rename from examples/pb_local_client.rs rename to crates/pb-mapper-cli/examples/pb_local_client.rs index 3e7d1bc..e42731b 100644 --- a/examples/pb_local_client.rs +++ b/crates/pb-mapper-cli/examples/pb_local_client.rs @@ -1,5 +1,5 @@ -use pb_mapper::common::config::init_tracing; -use pb_mapper::local::client::run_client_side_cli; +use pb_mapper_client::client::run_client_side_cli; +use pb_mapper_core::config::init_tracing; use uni_stream::stream::TcpListenerProvider; #[tokio::main] diff --git a/examples/pb_local_server.rs b/crates/pb-mapper-cli/examples/pb_local_server.rs similarity index 80% rename from examples/pb_local_server.rs rename to crates/pb-mapper-cli/examples/pb_local_server.rs index b01c15d..584824d 100644 --- a/examples/pb_local_server.rs +++ b/crates/pb-mapper-cli/examples/pb_local_server.rs @@ -1,5 +1,5 @@ -use pb_mapper::common::config::init_tracing; -use pb_mapper::local::server::{ServerTunnelOptions, run_server_side_cli}; +use pb_mapper_client::server::{ServerTunnelOptions, run_server_side_cli}; +use pb_mapper_core::config::init_tracing; use uni_stream::stream::TcpStreamProvider; #[tokio::main] diff --git a/examples/pb_server.rs b/crates/pb-mapper-cli/examples/pb_server.rs similarity index 57% rename from examples/pb_server.rs rename to crates/pb-mapper-cli/examples/pb_server.rs index 1c09ff0..b8e3ba5 100644 --- a/examples/pb_server.rs +++ b/crates/pb-mapper-cli/examples/pb_server.rs @@ -1,5 +1,5 @@ -use pb_mapper::common::config::init_tracing; -use pb_mapper::pb_server::run_server; +use pb_mapper_core::config::init_tracing; +use pb_mapper_server::run_server; #[tokio::main] async fn main() -> std::io::Result<()> { diff --git a/src/bin/pb-mapper.rs b/crates/pb-mapper-cli/src/bin/pb-mapper.rs similarity index 96% rename from src/bin/pb-mapper.rs rename to crates/pb-mapper-cli/src/bin/pb-mapper.rs index fc10703..14d6be5 100644 --- a/src/bin/pb-mapper.rs +++ b/crates/pb-mapper-cli/src/bin/pb-mapper.rs @@ -18,29 +18,29 @@ use std::time::Duration; use better_mimalloc_rs::MiMalloc; use clap::{Args, Parser, Subcommand, ValueEnum}; -use pb_mapper::common::auth::{ +use pb_mapper_auth::{ AuthConfig, KeyPage, LegacyProtocolPolicy, MAX_TEMP_KEY_CAPACITY, MAX_TEMP_KEY_TTL, MIN_TEMP_KEY_TTL, acquire_state_dir_lock, generate_admin_key, initialize_admin_key, write_admin_key_file, }; -use pb_mapper::common::checksum::set_process_msg_header_key; -use pb_mapper::common::checksum::{MACHINE_MSG_HEADER_KEY_PATH, setup_machine_msg_header_key}; -use pb_mapper::common::config::{ +use pb_mapper_client::client::{ + handle_status_cli_scoped, run_client_side_cli_with_callback_scoped, +}; +use pb_mapper_client::server::{ServerTunnelOptions, run_server_side_cli_with_pinned_credential}; +use pb_mapper_core::checksum::set_process_msg_header_key; +use pb_mapper_core::checksum::{MACHINE_MSG_HEADER_KEY_PATH, setup_machine_msg_header_key}; +use pb_mapper_core::config::{ StatusOp, control_io_timeout, get_pb_mapper_server_async, get_sockaddr_async, init_tracing, keep_alive_from_env, }; -use pb_mapper::common::message::MessageReader; -use pb_mapper::common::message::command::{ +use pb_mapper_protocol::MessageReader; +use pb_mapper_protocol::command::{ AdminConnectionPage, AdminRequest, AdminResponse, AdminServicePage, MessageSerializer, PbConnRequest, PbConnResponse, }; -use pb_mapper::common::message::forward::StreamForward; -use pb_mapper::common::message::secure::ClientHeaderSession; -use pb_mapper::local::client::{ - handle_status_cli_scoped, run_client_side_cli_with_callback_scoped, -}; -use pb_mapper::local::server::{ServerTunnelOptions, run_server_side_cli_with_pinned_credential}; -use pb_mapper::pb_server::run_server_with_shutdown; +use pb_mapper_protocol::forward::StreamForward; +use pb_mapper_protocol::secure::ClientHeaderSession; +use pb_mapper_server::run_server_with_shutdown; use tokio::net::TcpStream; use tokio_util::sync::CancellationToken; use uni_stream::stream::{ @@ -334,7 +334,7 @@ async fn run_server(args: ServerArgs) -> Result<(), Box> { } async fn run_register(args: RegisterArgs) -> Result<(), Box> { - let credential = pb_mapper::common::checksum::get_process_credential().map_err(|error| { + let credential = pb_mapper_core::checksum::get_process_credential().map_err(|error| { std::io::Error::other(format!("registration credential is required: {error}")) })?; let local_addr = get_sockaddr_async(&args.addr).await?; @@ -365,7 +365,7 @@ async fn register( remote_addr: std::net::SocketAddr, key: String, options: ServerTunnelOptions, - credential: pb_mapper::common::checksum::Credential, + credential: pb_mapper_core::checksum::Credential, ) where LocalStream::Item: StreamForward, { @@ -381,7 +381,7 @@ async fn register( } async fn run_connect(args: ConnectArgs) -> Result<(), Box> { - let credential = pb_mapper::common::checksum::get_process_credential().map_err(|error| { + let credential = pb_mapper_core::checksum::get_process_credential().map_err(|error| { std::io::Error::other(format!("client credential is required: {error}")) })?; let local_addr = get_sockaddr_async(&args.addr).await?; diff --git a/src/bin/pb-mapper/admin.rs b/crates/pb-mapper-cli/src/bin/pb-mapper/admin.rs similarity index 100% rename from src/bin/pb-mapper/admin.rs rename to crates/pb-mapper-cli/src/bin/pb-mapper/admin.rs diff --git a/tests/.env b/crates/pb-mapper-cli/tests/.env similarity index 100% rename from tests/.env rename to crates/pb-mapper-cli/tests/.env diff --git a/tests/regression.rs b/crates/pb-mapper-cli/tests/regression.rs similarity index 99% rename from tests/regression.rs rename to crates/pb-mapper-cli/tests/regression.rs index 793c943..a28b8cb 100644 --- a/tests/regression.rs +++ b/crates/pb-mapper-cli/tests/regression.rs @@ -3,23 +3,21 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use std::time::Duration; -use pb_mapper::common::auth::{ +use pb_mapper_auth::{ ADMIN_KEY_ID, AuthConfig, AuthRuntime, LegacyProtocolPolicy, write_admin_key_file, }; -use pb_mapper::common::checksum::{Credential, parse_credential, set_process_msg_header_key}; -use pb_mapper::common::message::command::{ +use pb_mapper_client::client::run_client_side_cli_with_callback; +use pb_mapper_client::server::{ServerTunnelOptions, run_server_side_cli_with_callback}; +use pb_mapper_core::checksum::{Credential, parse_credential, set_process_msg_header_key}; +use pb_mapper_protocol::command::{ AdminRequest, AdminResponse, LocalServer, MessageSerializer, PbConnRequest, PbConnResponse, PbConnStatusReq, PbConnStatusResp, PbServerRequest, PbServiceConnStatus, }; -use pb_mapper::common::message::secure::{ - ClientHeaderSession, ServerHeaderSession, ServerSecurity, -}; -use pb_mapper::common::message::{ +use pb_mapper_protocol::secure::{ClientHeaderSession, ServerHeaderSession, ServerSecurity}; +use pb_mapper_protocol::{ MessageReader, MessageWriter, get_header_msg_reader, get_header_msg_writer, }; -use pb_mapper::local::client::run_client_side_cli_with_callback; -use pb_mapper::local::server::{ServerTunnelOptions, run_server_side_cli_with_callback}; -use pb_mapper::pb_server::run_server_with_auth_config; +use pb_mapper_server::run_server_with_auth_config; use tokio::io::AsyncReadExt; use tokio::io::AsyncWriteExt; use tokio::net::{TcpListener, TcpStream}; diff --git a/tests/test_delay.rs b/crates/pb-mapper-cli/tests/test_delay.rs similarity index 97% rename from tests/test_delay.rs rename to crates/pb-mapper-cli/tests/test_delay.rs index 09da3a1..4efb96a 100644 --- a/tests/test_delay.rs +++ b/crates/pb-mapper-cli/tests/test_delay.rs @@ -2,14 +2,12 @@ use std::env; use std::sync::LazyLock; use std::time::Duration; -use pb_mapper::common::auth::{AuthConfig, LegacyProtocolPolicy}; -use pb_mapper::common::config::init_tracing; -use pb_mapper::common::message::{ - MessageReader, MessageWriter, NormalMessageReader, NormalMessageWriter, -}; -use pb_mapper::local::client::run_client_side_cli; -use pb_mapper::local::server::{ServerTunnelOptions, run_server_side_cli}; -use pb_mapper::pb_server::run_server_with_auth_config; +use pb_mapper_auth::{AuthConfig, LegacyProtocolPolicy}; +use pb_mapper_client::client::run_client_side_cli; +use pb_mapper_client::server::{ServerTunnelOptions, run_server_side_cli}; +use pb_mapper_core::config::init_tracing; +use pb_mapper_protocol::{MessageReader, MessageWriter, NormalMessageReader, NormalMessageWriter}; +use pb_mapper_server::run_server_with_auth_config; use rand::RngExt; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::UdpSocket; diff --git a/crates/pb-mapper-client/Cargo.toml b/crates/pb-mapper-client/Cargo.toml new file mode 100644 index 0000000..5890441 --- /dev/null +++ b/crates/pb-mapper-client/Cargo.toml @@ -0,0 +1,18 @@ +[package] +name = "pb-mapper-client" +version.workspace = true +edition.workspace = true +authors.workspace = true + +[dependencies] +pb-mapper-core.workspace = true +pb-mapper-protocol.workspace = true + +serde_json.workspace = true +snafu.workspace = true +tokio.workspace = true +tracing.workspace = true +uni-stream.workspace = true + +[features] +udp-timeout = ["uni-stream/udp-timeout", "pb-mapper-protocol/udp-timeout"] diff --git a/src/local/client/error.rs b/crates/pb-mapper-client/src/client/error.rs similarity index 93% rename from src/local/client/error.rs rename to crates/pb-mapper-client/src/client/error.rs index 7a83d16..bdae3c3 100644 --- a/src/local/client/error.rs +++ b/crates/pb-mapper-client/src/client/error.rs @@ -2,7 +2,9 @@ use std::time::Duration; use snafu::Snafu; -use crate::common; +// The `common::error::Error` spellings below are the source type on nearly every +// variant; aliasing keeps them as they were. +use pb_mapper_core as common; #[derive(Debug, Snafu)] #[snafu(visibility(pub(super)))] diff --git a/src/local/client/mod.rs b/crates/pb-mapper-client/src/client/mod.rs similarity index 97% rename from src/local/client/mod.rs rename to crates/pb-mapper-client/src/client/mod.rs index d10c764..768f5d5 100644 --- a/src/local/client/mod.rs +++ b/crates/pb-mapper-client/src/client/mod.rs @@ -15,15 +15,15 @@ use uni_stream::udp::set_custom_timeout; use self::error::{AcceptLocalStreamSnafu, BindLocalListenerSnafu}; use self::status::{get_status, get_status_scoped, get_status_with_credential}; use self::stream::handle_local_stream; -use crate::common::checksum::{Credential, get_process_credential}; -use crate::common::config::{ +use pb_mapper_core::checksum::{Credential, get_process_credential}; +use pb_mapper_core::config::{ StatusOp, client_health_check_interval, client_health_check_timeout, client_health_failure_threshold, }; -use crate::common::message::command::{PbConnStatusReq, PbConnStatusResp}; -use crate::common::message::forward::StreamForward; -use crate::snafu_error_get_or_return; -use crate::utils::timeout::RetryBackoff; +use pb_mapper_core::snafu_error_get_or_return; +use pb_mapper_core::timeout::RetryBackoff; +use pb_mapper_protocol::command::{PbConnStatusReq, PbConnStatusResp}; +use pb_mapper_protocol::forward::StreamForward; use uni_stream::addr::{ToSocketAddrs, each_addr}; use uni_stream::stream::got_one_socket_addr; use uni_stream::stream::{ListenerProvider, StreamAccept}; @@ -100,7 +100,7 @@ pub async fn run_client_side_cli_with_pinned_credential< key: Arc, keep_alive: bool, status_callback: Option, - credential: crate::common::checksum::Credential, + credential: pb_mapper_core::checksum::Credential, ) where ::Item: StreamForward, { @@ -126,7 +126,7 @@ pub async fn run_client_side_cli_with_callback_scoped< keep_alive: bool, namespace: Option, status_callback: Option, - pinned_credential: Option, + pinned_credential: Option, ) where ::Item: StreamForward, { diff --git a/src/local/client/status.rs b/crates/pb-mapper-client/src/client/status.rs similarity index 94% rename from src/local/client/status.rs rename to crates/pb-mapper-client/src/client/status.rs index 7faee86..6f686d7 100644 --- a/src/local/client/status.rs +++ b/crates/pb-mapper-client/src/client/status.rs @@ -5,12 +5,12 @@ use super::error::{ CreateHeaderToolSnafu, DecodeStatusRespSnafu, EncodeStatusReqSnafu, StatusRespNotMatchSnafu, WriteStatusReqSnafu, }; -use crate::common::checksum::Credential; -use crate::common::config::control_io_timeout; -use crate::common::message::command::{ +use pb_mapper_core::checksum::Credential; +use pb_mapper_core::config::control_io_timeout; +use pb_mapper_protocol::command::{ MessageSerializer, PbConnRequest, PbConnResponse, PbConnStatusReq, PbConnStatusResp, }; -use crate::common::message::secure::ClientHeaderSession; +use pb_mapper_protocol::secure::ClientHeaderSession; pub async fn get_status( remote_stream: &mut S, diff --git a/src/local/client/stream.rs b/crates/pb-mapper-client/src/client/stream.rs similarity index 89% rename from src/local/client/stream.rs rename to crates/pb-mapper-client/src/client/stream.rs index 63ab7f7..ef8bf6b 100644 --- a/src/local/client/stream.rs +++ b/crates/pb-mapper-client/src/client/stream.rs @@ -9,13 +9,13 @@ use super::error::{ ConnectRemoteStreamSnafu, DecodeSubcribeRespSnafu, EncodeSubcribeReqSnafu, Result, SubcribeRespNotMatchSnafu, WriteSubcribeReqSnafu, }; -use crate::common::checksum::Credential; -use crate::common::config::control_io_timeout; -use crate::common::message::command::{MessageSerializer, PbConnRequest, PbConnResponse}; -use crate::common::message::forward::StreamForward; -use crate::common::message::secure::ClientHeaderSession; -use crate::local::client::error::CreateHeaderToolSnafu; -use crate::snafu_error_handle; +use crate::client::error::CreateHeaderToolSnafu; +use pb_mapper_core::checksum::Credential; +use pb_mapper_core::config::control_io_timeout; +use pb_mapper_core::snafu_error_handle; +use pb_mapper_protocol::command::{MessageSerializer, PbConnRequest, PbConnResponse}; +use pb_mapper_protocol::forward::StreamForward; +use pb_mapper_protocol::secure::ClientHeaderSession; use uni_stream::addr::{ToSocketAddrs, each_addr}; use uni_stream::stream::{NetworkStream, set_tcp_keep_alive, set_tcp_nodelay}; diff --git a/src/local/mod.rs b/crates/pb-mapper-client/src/lib.rs similarity index 100% rename from src/local/mod.rs rename to crates/pb-mapper-client/src/lib.rs diff --git a/src/local/server/error.rs b/crates/pb-mapper-client/src/server/error.rs similarity index 94% rename from src/local/server/error.rs rename to crates/pb-mapper-client/src/server/error.rs index 1029071..19bdd4c 100644 --- a/src/local/server/error.rs +++ b/crates/pb-mapper-client/src/server/error.rs @@ -2,7 +2,9 @@ use std::time::Duration; use snafu::Snafu; -use crate::common::{self}; +// The `common::error::Error` spellings below are the source type on nearly every +// variant; aliasing keeps them as they were. +use pb_mapper_core as common; #[derive(Debug, Snafu)] #[snafu(visibility(pub(super)))] diff --git a/src/local/server/mod.rs b/crates/pb-mapper-client/src/server/mod.rs similarity index 98% rename from src/local/server/mod.rs rename to crates/pb-mapper-client/src/server/mod.rs index 68fcd6c..b47ed2c 100644 --- a/src/local/server/mod.rs +++ b/crates/pb-mapper-client/src/server/mod.rs @@ -18,23 +18,23 @@ use self::error::{ RegisterRespNotMatchSnafu, SendRegisterReqSnafu, WritePingMsgSnafu, WriteStreamAckMsgSnafu, }; use self::stream::{StreamConnect, handle_stream}; -use crate::common::checksum::{Credential, get_process_credential}; -use crate::common::config::{ +use pb_mapper_core::checksum::{Credential, get_process_credential}; +use pb_mapper_core::config::{ control_conn_pool_size, control_heartbeat_interval, control_heartbeat_tolerance, control_io_timeout, control_suspect_grace, registration_probe_timeout, }; -use crate::common::message::command::{ - CONTROL_PROTOCOL_V2, LocalServer, MessageSerializer, PbConnRequest, PbConnResponse, - PbConnStatusReq, PbConnStatusResp, PbServerRequest, -}; -use crate::common::message::forward::StreamForward; -use crate::common::message::secure::ClientHeaderSession; -use crate::common::message::{MessageReader, MessageWriter}; -use crate::utils::timeout::RetryBackoff; -use crate::{ +use pb_mapper_core::timeout::RetryBackoff; +use pb_mapper_core::{ snafu_error_get_or_continue, snafu_error_get_or_return, snafu_error_get_or_return_ok, snafu_error_handle, }; +use pb_mapper_protocol::command::{ + CONTROL_PROTOCOL_V2, LocalServer, MessageSerializer, PbConnRequest, PbConnResponse, + PbConnStatusReq, PbConnStatusResp, PbServerRequest, +}; +use pb_mapper_protocol::forward::StreamForward; +use pb_mapper_protocol::secure::ClientHeaderSession; +use pb_mapper_protocol::{MessageReader, MessageWriter}; use uni_stream::addr::{ToSocketAddrs, each_addr}; use uni_stream::stream::{ StreamProvider, got_one_socket_addr, set_tcp_keep_alive, set_tcp_nodelay, @@ -184,7 +184,7 @@ async fn probe_remote_registration( let mut stream = each_addr(remote_addr, TcpStream::connect) .await .map_err(|e| format!("connect remote status stream failed: {e}"))?; - crate::local::client::status::get_status_with_credential( + crate::client::status::get_status_with_credential( &mut stream, PbConnStatusReq::Service { key: key.to_string(), diff --git a/src/local/server/stream.rs b/crates/pb-mapper-client/src/server/stream.rs similarity index 91% rename from src/local/server/stream.rs rename to crates/pb-mapper-client/src/server/stream.rs index 52ade70..c8b5201 100644 --- a/src/local/server/stream.rs +++ b/crates/pb-mapper-client/src/server/stream.rs @@ -10,13 +10,13 @@ use super::error::{ DecodePbConnStreamRespSnafu, EncodePbConnStreamReqSnafu, PbConnStreamRespNotMatchSnafu, Result, WritePbConnStreamReqSnafu, }; -use crate::common::checksum::Credential; -use crate::common::config::control_io_timeout; -use crate::common::message::command::{MessageSerializer, PbConnRequest, PbConnResponse}; -use crate::common::message::forward::StreamForward; -use crate::common::message::secure::ClientHeaderSession; -use crate::local::server::error::CreateHeaderToolSnafu; -use crate::snafu_error_handle; +use crate::server::error::CreateHeaderToolSnafu; +use pb_mapper_core::checksum::Credential; +use pb_mapper_core::config::control_io_timeout; +use pb_mapper_core::snafu_error_handle; +use pb_mapper_protocol::command::{MessageSerializer, PbConnRequest, PbConnResponse}; +use pb_mapper_protocol::forward::StreamForward; +use pb_mapper_protocol::secure::ClientHeaderSession; use uni_stream::addr::{ToSocketAddrs, each_addr}; use uni_stream::stream::{StreamProvider, StreamSplit, set_tcp_keep_alive, set_tcp_nodelay}; diff --git a/crates/pb-mapper-protocol/src/command.rs b/crates/pb-mapper-protocol/src/command.rs index f52f087..0da74fb 100644 --- a/crates/pb-mapper-protocol/src/command.rs +++ b/crates/pb-mapper-protocol/src/command.rs @@ -333,3 +333,31 @@ gen_impl_msg_serializer!(PbConnRequest); gen_impl_msg_serializer!(PbConnResponse); gen_impl_msg_serializer!(PbServerRequest); gen_impl_msg_serializer!(LocalServer); + +#[cfg(test)] +mod tests { + use super::PbConnRequest; + + /// The wire form of `Register` is load-bearing: a running peer on the other + /// side of an upgrade has to keep parsing it. The `None` fields must stay + /// absent from the JSON rather than serialise as null. + #[test] + fn test_serde_mapper_header() { + let mapper = PbConnRequest::Register { + key: "test".into(), + need_codec: false, + is_datagram: false, + protocol_version: None, + client_instance_id: None, + heartbeat_interval_ms: None, + heartbeat_tolerance_ms: None, + }; + let json_value = serde_json::to_string(&mapper).unwrap(); + let raw_json_str = + r##"{"Register":{"need_codec":false,"is_datagram":false,"key":"test"}}"##; + assert_eq!(raw_json_str, json_value); + + let value: PbConnRequest = serde_json::from_str(raw_json_str).unwrap(); + assert_eq!(mapper, value) + } +} diff --git a/crates/pb-mapper-server/Cargo.toml b/crates/pb-mapper-server/Cargo.toml new file mode 100644 index 0000000..21d3d3d --- /dev/null +++ b/crates/pb-mapper-server/Cargo.toml @@ -0,0 +1,24 @@ +[package] +name = "pb-mapper-server" +version.workspace = true +edition.workspace = true +authors.workspace = true + +[dependencies] +pb-mapper-auth.workspace = true +pb-mapper-core.workspace = true +pb-mapper-protocol.workspace = true + +hashbrown.workspace = true +kanal.workspace = true +snafu.workspace = true +tokio.workspace = true +tokio-util.workspace = true +tracing.workspace = true +uni-stream.workspace = true + +[dev-dependencies] +rand.workspace = true + +[features] +udp-timeout = ["uni-stream/udp-timeout", "pb-mapper-protocol/udp-timeout"] diff --git a/src/pb_server/admin.rs b/crates/pb-mapper-server/src/admin.rs similarity index 94% rename from src/pb_server/admin.rs rename to crates/pb-mapper-server/src/admin.rs index f8c85dd..b119d6d 100644 --- a/src/pb_server/admin.rs +++ b/crates/pb-mapper-server/src/admin.rs @@ -15,14 +15,12 @@ use tokio::net::TcpStream; use super::error::Error; use super::{ManagerTask, ManagerTaskSender, Result}; -use crate::common::auth::{AuthContext, AuthFailure, AuthRuntime, KeyId}; -use crate::common::checksum::{Credential, parse_credential}; -use crate::common::conn_id::RemoteConnId; -use crate::common::message::MessageWriter; -use crate::common::message::command::{ - AdminRequest, AdminResponse, MessageSerializer, PbConnResponse, -}; -use crate::common::message::secure::ServerHeaderSession; +use pb_mapper_auth::{AuthContext, AuthFailure, AuthRuntime, KeyId}; +use pb_mapper_core::checksum::{Credential, parse_credential}; +use pb_mapper_core::conn_id::RemoteConnId; +use pb_mapper_protocol::MessageWriter; +use pb_mapper_protocol::command::{AdminRequest, AdminResponse, MessageSerializer, PbConnResponse}; +use pb_mapper_protocol::secure::ServerHeaderSession; pub async fn handle_admin_request( request: AdminRequest, @@ -293,8 +291,8 @@ async fn audit_read( #[cfg(test)] mod tests { use super::*; - use crate::common::auth::ADMIN_KEY_ID; - use crate::common::auth::{AuthConfig, LegacyProtocolPolicy}; + use pb_mapper_auth::ADMIN_KEY_ID; + use pb_mapper_auth::{AuthConfig, LegacyProtocolPolicy}; fn temp_state_dir(name: &str) -> std::path::PathBuf { std::env::temp_dir().join(format!( @@ -363,7 +361,7 @@ mod tests { ManagerTask::AdminServiceList { response_sender, .. } => { - let _ = response_sender.send(crate::common::message::command::AdminServicePage { + let _ = response_sender.send(pb_mapper_protocol::command::AdminServicePage { schema_version: 1, items: Vec::new(), next_page: None, @@ -372,12 +370,11 @@ mod tests { ManagerTask::AdminConnectionList { response_sender, .. } => { - let _ = - response_sender.send(crate::common::message::command::AdminConnectionPage { - schema_version: 1, - items: Vec::new(), - next_page: None, - }); + let _ = response_sender.send(pb_mapper_protocol::command::AdminConnectionPage { + schema_version: 1, + items: Vec::new(), + next_page: None, + }); } _ => panic!("expected an administrator inventory manager task"), } diff --git a/src/pb_server/client.rs b/crates/pb-mapper-server/src/client.rs similarity index 97% rename from src/pb_server/client.rs rename to crates/pb-mapper-server/src/client.rs index df3e762..960cad8 100644 --- a/src/pb_server/client.rs +++ b/crates/pb-mapper-server/src/client.rs @@ -13,22 +13,22 @@ use super::error::{ ClientConnSubcribeRespNotMatchSnafu, ClientConnWriteSubcribeRespSnafu, }; use super::{ConnTask, ImutableKey, ManagerTask, ManagerTaskSender, Result}; -use crate::common::checksum::{AesKeyType, gen_random_key}; -use crate::common::config::{stream_ack_timeout, stream_ready_timeout, stream_recovery_timeout}; -use crate::common::conn_id::RemoteConnId; -use crate::common::message::command::{MessageSerializer, PbConnResponse}; -use crate::common::message::forward::{ +use crate::error::{ + ClientConnCreateHeaderToolSnafu, ClientConnEncodeStreamRespSnafu, + ClientConnWriteStreamRespSnafu, +}; +use pb_mapper_core::checksum::{AesKeyType, gen_random_key}; +use pb_mapper_core::config::{stream_ack_timeout, stream_ready_timeout, stream_recovery_timeout}; +use pb_mapper_core::conn_id::RemoteConnId; +use pb_mapper_core::snafu_error_get_or_return_ok; +use pb_mapper_protocol::command::{MessageSerializer, PbConnResponse}; +use pb_mapper_protocol::forward::{ CodecDatagramReader, CodecDatagramWriter, CodecForwardReader, CodecForwardWriter, NormalDatagramReader, NormalDatagramWriter, NormalForwardReader, NormalForwardWriter, start_datagram_forward, start_forward, }; -use crate::common::message::secure::ServerHeaderSession; -use crate::common::message::{MessageWriter, get_decodec, get_encodec}; -use crate::pb_server::error::{ - ClientConnCreateHeaderToolSnafu, ClientConnEncodeStreamRespSnafu, - ClientConnWriteStreamRespSnafu, -}; -use crate::snafu_error_get_or_return_ok; +use pb_mapper_protocol::secure::ServerHeaderSession; +use pb_mapper_protocol::{MessageWriter, get_decodec, get_encodec}; /// Ensure that client-side connections are properly deregistered before a normal connection is /// disconnected or an exception occurs diff --git a/src/pb_server/connection.rs b/crates/pb-mapper-server/src/connection.rs similarity index 96% rename from src/pb_server/connection.rs rename to crates/pb-mapper-server/src/connection.rs index 6cb0545..1134200 100644 --- a/src/pb_server/connection.rs +++ b/crates/pb-mapper-server/src/connection.rs @@ -109,7 +109,7 @@ pub(super) async fn handle_conn( write_protocol_error( &mut conn, &initial.session, - &crate::common::auth::AuthFailure::new( + &pb_mapper_auth::AuthFailure::new( "request_decode_failed", "authenticated request payload is malformed", false, @@ -362,7 +362,7 @@ pub(super) async fn handle_conn( write_protocol_error( &mut conn, &session, - &crate::common::auth::AuthFailure::new( + &pb_mapper_auth::AuthFailure::new( "admin_permission_required", "administrator credential is required for this operation", false, @@ -375,7 +375,7 @@ pub(super) async fn handle_conn( write_protocol_error( &mut conn, &session, - &crate::common::auth::AuthFailure::new( + &pb_mapper_auth::AuthFailure::new( "admin_protocol_v2_required", "administrator operations require protocol v2", false, @@ -471,7 +471,7 @@ where async fn write_protocol_error( conn: &mut TcpStream, session: &ServerHeaderSession, - failure: &crate::common::auth::AuthFailure, + failure: &pb_mapper_auth::AuthFailure, ) { let response = PbConnResponse::error( failure.code.clone(), @@ -494,17 +494,17 @@ fn resolve_namespace( requested: Option, force_register_namespace: bool, is_register: bool, -) -> std::result::Result { +) -> std::result::Result { let namespace = requested.unwrap_or(context.namespace); if !context.is_admin && namespace != context.namespace { - return Err(crate::common::auth::AuthFailure::new( + return Err(pb_mapper_auth::AuthFailure::new( "namespace_access_denied", "temporary credentials can only access their own namespace", false, )); } if context.is_admin && is_register && namespace != 0 && !force_register_namespace { - return Err(crate::common::auth::AuthFailure::new( + return Err(pb_mapper_auth::AuthFailure::new( "namespace_force_required", "administrator registration in a temporary namespace requires --force", false, @@ -517,9 +517,9 @@ fn scoped_service_key( context: &AuthContext, namespace: u64, service_name: &str, -) -> std::result::Result { +) -> std::result::Result { if service_name.is_empty() || service_name.len() > 1024 || service_name.contains('\0') { - return Err(crate::common::auth::AuthFailure::new( + return Err(pb_mapper_auth::AuthFailure::new( "service_name_invalid", "service names must be 1-1024 bytes and must not contain NUL", false, @@ -531,7 +531,7 @@ fn scoped_service_key( .bytes() .all(|byte| byte.is_ascii_alphanumeric() || b"._:-".contains(&byte))) { - return Err(crate::common::auth::AuthFailure::new( + return Err(pb_mapper_auth::AuthFailure::new( "service_name_invalid", "temporary-key service names must be 1-128 ASCII bytes from [A-Za-z0-9._:-]", false, diff --git a/src/pb_server/error.rs b/crates/pb-mapper-server/src/error.rs similarity index 98% rename from src/pb_server/error.rs rename to crates/pb-mapper-server/src/error.rs index e131712..6149158 100644 --- a/src/pb_server/error.rs +++ b/crates/pb-mapper-server/src/error.rs @@ -2,8 +2,10 @@ use std::{sync::Arc, time::Duration}; use snafu::Snafu; -use crate::common::conn_id::RemoteConnId; -use crate::common::{self}; +use pb_mapper_core::conn_id::RemoteConnId; +// The `common::error::Error` spellings below are the source type on nearly every +// variant; aliasing keeps them as they were. +use pb_mapper_core as common; #[derive(Debug, Snafu)] #[snafu(visibility(pub(super)))] diff --git a/src/pb_server/mod.rs b/crates/pb-mapper-server/src/lib.rs similarity index 94% rename from src/pb_server/mod.rs rename to crates/pb-mapper-server/src/lib.rs index fbaf9cf..9d4c9a5 100644 --- a/src/pb_server/mod.rs +++ b/crates/pb-mapper-server/src/lib.rs @@ -13,6 +13,8 @@ mod admin; mod client; mod error; +// Moved here from `common`: the routing runtime is its only caller. +pub mod manager; mod server; mod status; @@ -34,23 +36,23 @@ use self::error::{ }; use self::server::{ServerRegistration, handle_server_conn}; use self::status::handle_show_status; -use crate::common::auth::{ADMIN_KEY_ID, AuthConfig, AuthContext, AuthRuntime}; -use crate::common::config::{control_io_timeout, keep_alive_from_env, server_lease_timeout}; -use crate::common::conn_id::{ConnIdProvider, RemoteConnId}; -use crate::common::manager::{ForwardMessage, SenderChan, TaskManager}; -use crate::common::message::MessageWriter; -use crate::common::message::command::{ - AdminConnectionInfo, AdminConnectionPage, AdminServiceInfo, AdminServicePage, - MessageSerializer, PbConnRequest, PbConnResponse, PbConnStatusReq, PbConnStatusResp, - PbServiceConnStatus, -}; -use crate::common::message::secure::{HeaderProtocol, ServerHeaderSession, ServerSecurity}; -use crate::pb_server::error::{ +use crate::error::{ ServerListenSnafu, TaskCenterClientSendStreamSnafu, TaskCenterSendRegisterRespSnafu, TaskCenterSendStreamRespToClientSnafu, TaskCenterSendSubcribeRespSnafu, TaskCenterStreamConnIdNotExistSnafu, }; -use crate::{snafu_error_get_or_continue, snafu_error_handle}; +use crate::manager::{ForwardMessage, SenderChan, TaskManager}; +use pb_mapper_auth::{ADMIN_KEY_ID, AuthConfig, AuthContext, AuthRuntime}; +use pb_mapper_core::config::{control_io_timeout, keep_alive_from_env, server_lease_timeout}; +use pb_mapper_core::conn_id::{ConnIdProvider, RemoteConnId}; +use pb_mapper_core::{snafu_error_get_or_continue, snafu_error_handle}; +use pb_mapper_protocol::MessageWriter; +use pb_mapper_protocol::command::{ + AdminConnectionInfo, AdminConnectionPage, AdminServiceInfo, AdminServicePage, + MessageSerializer, PbConnRequest, PbConnResponse, PbConnStatusReq, PbConnStatusResp, + PbServiceConnStatus, +}; +use pb_mapper_protocol::secure::{HeaderProtocol, ServerHeaderSession, ServerSecurity}; use uni_stream::stream::{set_tcp_keep_alive, set_tcp_nodelay}; pub enum ManagerTask { diff --git a/src/common/manager.rs b/crates/pb-mapper-server/src/manager.rs similarity index 99% rename from src/common/manager.rs rename to crates/pb-mapper-server/src/manager.rs index 8644927..d94a34b 100644 --- a/src/common/manager.rs +++ b/crates/pb-mapper-server/src/manager.rs @@ -1,7 +1,7 @@ use snafu::{ResultExt, Snafu}; use tracing::instrument; -use super::conn_id::{ConnId, ConnIdProvider, ConnIdTrait}; +use pb_mapper_core::conn_id::{ConnId, ConnIdProvider, ConnIdTrait}; /// The manager owns this rather than the core error enum: it is the only thing /// that waits on a task channel, and it keeps `kanal` out of the bottom layer. diff --git a/src/pb_server/runtime.rs b/crates/pb-mapper-server/src/runtime.rs similarity index 99% rename from src/pb_server/runtime.rs rename to crates/pb-mapper-server/src/runtime.rs index 9340865..4efaee5 100644 --- a/src/pb_server/runtime.rs +++ b/crates/pb-mapper-server/src/runtime.rs @@ -958,7 +958,7 @@ async fn abort_and_wait(handles: impl IntoIterator Date: Fri, 21 Aug 2026 14:02:21 +0800 Subject: [PATCH 73/74] Deny unwrap and expect across the workspace MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The FFI crate has denied these since it was written, because a panic crossing the C ABI is undefined behaviour rather than a stack trace. Every crate is now reachable from that boundary, so the deny moves to `[workspace.lints]` and the FFI crate inherits it instead of keeping its own copy. A `clippy.toml` exempts test code — a panicking assertion is a failing test, which is the point. That covers `#[cfg(test)]` modules but not `tests/` or `examples/` targets, whose entire body is test code, so those carry a file-level allow with the same reasoning. That left nine real sites in production paths. Six are now errors rather than panics, and each was already inside a function returning `Result`: credential and v2-prefix width conversions that parse bytes from an unauthenticated peer, the administrator-key UTF-8 conversion, and the state blob's nonce, which reads a file that may have been truncated. Two accessors replace `expect("v2 session material")`, where `Some` and `HeaderProtocol::V2` are the same condition without the type saying so. The timing wheel's `schedule` returns `()`, so its unreachable case now takes the same path as an out-of-range delay: drop the timer, which fires it — the safe direction for a credential deadline. `replay.rs` keeps a module-level allow, with the reason recorded there: its conversions slice a `[u8; 40]` at constant offsets or take SHA-256's 32-byte output, widths the array types already fix, in functions returning `()`. Two more `mod tests` were missing `#[cfg(test)]` and compiled into release builds; the lint is what surfaced them. Co-Authored-By: Claude Fable 5 --- Cargo.toml | 4 ++ clippy.toml | 5 +++ crates/pb-mapper-auth/Cargo.toml | 3 ++ crates/pb-mapper-auth/src/actor/epoch.rs | 11 +++++- crates/pb-mapper-auth/src/persistence/blob.rs | 13 +++++-- crates/pb-mapper-auth/src/timing_wheel.rs | 9 ++++- crates/pb-mapper-cli/Cargo.toml | 3 ++ .../pb-mapper-cli/examples/echo_udp_server.rs | 3 ++ crates/pb-mapper-cli/tests/regression.rs | 5 +++ crates/pb-mapper-cli/tests/test_delay.rs | 3 ++ crates/pb-mapper-client/Cargo.toml | 3 ++ crates/pb-mapper-core/Cargo.toml | 3 ++ crates/pb-mapper-core/src/checksum.rs | 16 ++++++-- crates/pb-mapper-protocol/Cargo.toml | 3 ++ crates/pb-mapper-protocol/src/secure.rs | 39 +++++++++++++++---- .../pb-mapper-protocol/src/secure/replay.rs | 6 +++ crates/pb-mapper-server/Cargo.toml | 3 ++ ui/native/pb_mapper_ffi/Cargo.toml | 5 +-- 18 files changed, 116 insertions(+), 21 deletions(-) create mode 100644 clippy.toml diff --git a/Cargo.toml b/Cargo.toml index 9ddacda..4848d5f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,6 +14,10 @@ version = "0.4.0" authors = ["L_B__"] edition = "2024" +[workspace.lints.clippy] +unwrap_used = "deny" +expect_used = "deny" + [workspace.dependencies] pb-mapper-auth = { path = "crates/pb-mapper-auth" } pb-mapper-client = { path = "crates/pb-mapper-client" } diff --git a/clippy.toml b/clippy.toml new file mode 100644 index 0000000..96ba758 --- /dev/null +++ b/clippy.toml @@ -0,0 +1,5 @@ +# `unwrap_used` and `expect_used` are denied workspace-wide (see the root +# `[workspace.lints.clippy]`). A panic in a test is a failing test, which is the +# point, so exempt test code rather than annotating every assertion. +allow-unwrap-in-tests = true +allow-expect-in-tests = true diff --git a/crates/pb-mapper-auth/Cargo.toml b/crates/pb-mapper-auth/Cargo.toml index 8f053ea..a76d7a3 100644 --- a/crates/pb-mapper-auth/Cargo.toml +++ b/crates/pb-mapper-auth/Cargo.toml @@ -16,3 +16,6 @@ subtle.workspace = true tokio.workspace = true tokio-util.workspace = true tracing.workspace = true + +[lints] +workspace = true diff --git a/crates/pb-mapper-auth/src/actor/epoch.rs b/crates/pb-mapper-auth/src/actor/epoch.rs index d72cbc6..c69ba8a 100644 --- a/crates/pb-mapper-auth/src/actor/epoch.rs +++ b/crates/pb-mapper-auth/src/actor/epoch.rs @@ -74,8 +74,15 @@ pub(super) fn actor_rotate_root( false, )); } - let new_key_string = - String::from_utf8(new_key.to_vec()).expect("printable ASCII is valid UTF-8"); + // Unreachable: `is_env_safe_admin_key` above accepts only printable ASCII. + // Reported rather than panicked, since this already returns `Result`. + let new_key_string = String::from_utf8(new_key.to_vec()).map_err(|_| { + AuthFailure::new( + "administrator_key_invalid", + env_safe_admin_key_error(), + false, + ) + })?; inner.root_epoch.fetch_add(1, Ordering::AcqRel); let rotate_audit = audit("administrator_key_rotate", None, None); diff --git a/crates/pb-mapper-auth/src/persistence/blob.rs b/crates/pb-mapper-auth/src/persistence/blob.rs index 3697072..dbebd7f 100644 --- a/crates/pb-mapper-auth/src/persistence/blob.rs +++ b/crates/pb-mapper-auth/src/persistence/blob.rs @@ -37,9 +37,16 @@ pub(crate) fn open_blob(admin_key: &AesKeyType, sealed: &[u8]) -> Result } let nonce_start = STATE_BLOB_MAGIC.len(); let nonce_end = nonce_start + 12; - let nonce_bytes: [u8; 12] = sealed[nonce_start..nonce_end] - .try_into() - .expect("validated nonce width"); + // Unreachable: the length check above guarantees these 12 bytes exist. This + // parses a file that may have been truncated or corrupted, so it reports + // rather than panics. + let nonce_bytes: [u8; 12] = sealed[nonce_start..nonce_end].try_into().map_err(|_| { + AuthFailure::new( + "temporary_key_store_unavailable", + "authentication state blob has an invalid nonce", + false, + ) + })?; let mut plain = sealed[nonce_end..].to_vec(); let key = LessSafeKey::new(UnboundKey::new(&AES_256_GCM, admin_key).map_err(|_| { AuthFailure::new( diff --git a/crates/pb-mapper-auth/src/timing_wheel.rs b/crates/pb-mapper-auth/src/timing_wheel.rs index 989c340..aefcfb4 100644 --- a/crates/pb-mapper-auth/src/timing_wheel.rs +++ b/crates/pb-mapper-auth/src/timing_wheel.rs @@ -173,10 +173,15 @@ impl TimingWheel { // into its rotation, so its bucket comes off on a rotation boundary. Every // finer level is at zero offset there, which makes the delay still // remaining a plain base-`radix` decomposition from that point down. - let (level, slot, remaining) = (0..self.levels.len()) + // The guard above bounds `delay`, so some level always takes it. + // Unreachable, and treated like the out-of-range case: dropping the + // timer fires it, which is the safe direction for a credential deadline. + let Some((level, slot, remaining)) = (0..self.levels.len()) .rev() .find_map(|level| self.entry_leg(level, self.ticks + delay)) - .expect("a delay within max_delay reaches some level"); + else { + return; + }; let route = match remaining { 0 => deliver, remaining => self.route(remaining, deliver), diff --git a/crates/pb-mapper-cli/Cargo.toml b/crates/pb-mapper-cli/Cargo.toml index a560734..f45b297 100644 --- a/crates/pb-mapper-cli/Cargo.toml +++ b/crates/pb-mapper-cli/Cargo.toml @@ -35,3 +35,6 @@ udp-timeout = [ "pb-mapper-client/udp-timeout", "pb-mapper-server/udp-timeout", ] + +[lints] +workspace = true diff --git a/crates/pb-mapper-cli/examples/echo_udp_server.rs b/crates/pb-mapper-cli/examples/echo_udp_server.rs index 23a8574..6dc17c5 100644 --- a/crates/pb-mapper-cli/examples/echo_udp_server.rs +++ b/crates/pb-mapper-cli/examples/echo_udp_server.rs @@ -1,3 +1,6 @@ +// An example: panicking on a failed bind is the clearest thing it can do. +#![allow(clippy::unwrap_used, clippy::expect_used)] + use std::error::Error; use std::net::SocketAddr; use std::str::FromStr; diff --git a/crates/pb-mapper-cli/tests/regression.rs b/crates/pb-mapper-cli/tests/regression.rs index a28b8cb..5b3e7b2 100644 --- a/crates/pb-mapper-cli/tests/regression.rs +++ b/crates/pb-mapper-cli/tests/regression.rs @@ -1,3 +1,8 @@ +// An integration test: a failed `unwrap` is a failed test, which is the report +// this file exists to produce. `allow-unwrap-in-tests` covers `#[cfg(test)]` +// modules but not a `tests/` target, whose whole body is test code. +#![allow(clippy::unwrap_used, clippy::expect_used)] + use std::net::SocketAddr; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; diff --git a/crates/pb-mapper-cli/tests/test_delay.rs b/crates/pb-mapper-cli/tests/test_delay.rs index 4efb96a..78c747b 100644 --- a/crates/pb-mapper-cli/tests/test_delay.rs +++ b/crates/pb-mapper-cli/tests/test_delay.rs @@ -1,3 +1,6 @@ +// See the note in `regression.rs`: the whole file is test code. +#![allow(clippy::unwrap_used, clippy::expect_used)] + use std::env; use std::sync::LazyLock; use std::time::Duration; diff --git a/crates/pb-mapper-client/Cargo.toml b/crates/pb-mapper-client/Cargo.toml index 5890441..771b2e8 100644 --- a/crates/pb-mapper-client/Cargo.toml +++ b/crates/pb-mapper-client/Cargo.toml @@ -16,3 +16,6 @@ uni-stream.workspace = true [features] udp-timeout = ["uni-stream/udp-timeout", "pb-mapper-protocol/udp-timeout"] + +[lints] +workspace = true diff --git a/crates/pb-mapper-core/Cargo.toml b/crates/pb-mapper-core/Cargo.toml index f054a61..1ae3b9c 100644 --- a/crates/pb-mapper-core/Cargo.toml +++ b/crates/pb-mapper-core/Cargo.toml @@ -16,3 +16,6 @@ snafu.workspace = true tokio.workspace = true tracing.workspace = true tracing-subscriber.workspace = true + +[lints] +workspace = true diff --git a/crates/pb-mapper-core/src/checksum.rs b/crates/pb-mapper-core/src/checksum.rs index ce4da1f..190ced2 100644 --- a/crates/pb-mapper-core/src/checksum.rs +++ b/crates/pb-mapper-core/src/checksum.rs @@ -185,13 +185,20 @@ pub fn parse_credential(raw: &str) -> Result { if expected.as_ref()[..4] != payload[41..45] { return Err("temporary credential checksum mismatch".to_string()); } - let key_id = u64::from_be_bytes(payload[1..9].try_into().expect("fixed key id width")); + // The 45-byte check above already guarantees both widths, so neither arm + // is reachable — but this returns `Result` anyway, so saying so costs a + // line and removes a panic from a path that parses network input. + let key_id = u64::from_be_bytes( + payload[1..9] + .try_into() + .map_err(|_| "temporary credential key id is malformed".to_string())?, + ); if key_id == 0 { return Err("temporary credential key id must not be zero".to_string()); } let key = payload[9..41] .try_into() - .expect("fixed temporary key width"); + .map_err(|_| "temporary credential key is malformed".to_string())?; return Ok(Credential::Temporary { key_id, key }); } @@ -202,8 +209,9 @@ pub fn parse_credential(raw: &str) -> Result { if !is_env_safe_admin_key(bytes) { return Err(env_safe_admin_key_error()); } + // Unreachable after the `ADMIN_KEY_LEN` check above, for the same reason. Ok(Credential::Admin( - bytes.try_into().expect("validated admin key width"), + bytes.try_into().map_err(|_| key_len_error(raw))?, )) } @@ -539,6 +547,8 @@ pub fn gen_random_key() -> [u8; 32] { random_key } +// This was missing its `#[cfg(test)]`, so it compiled into release builds. +#[cfg(test)] mod tests { #[test] fn test_random_checksum() { diff --git a/crates/pb-mapper-protocol/Cargo.toml b/crates/pb-mapper-protocol/Cargo.toml index eae6617..b4a948a 100644 --- a/crates/pb-mapper-protocol/Cargo.toml +++ b/crates/pb-mapper-protocol/Cargo.toml @@ -21,3 +21,6 @@ uni-stream.workspace = true [features] udp-timeout = ["uni-stream/udp-timeout"] + +[lints] +workspace = true diff --git a/crates/pb-mapper-protocol/src/secure.rs b/crates/pb-mapper-protocol/src/secure.rs index 9468c50..98a6c16 100644 --- a/crates/pb-mapper-protocol/src/secure.rs +++ b/crates/pb-mapper-protocol/src/secure.rs @@ -69,6 +69,16 @@ pub struct ClientHeaderSession { } impl ClientHeaderSession { + /// The v2 material, which is `Some` exactly when `protocol` is `V2`. + /// + /// The type does not tie the two together, so this reports instead of + /// panicking on a state that construction never produces. + fn v2_material(&self) -> Result<&V2Material> { + self.v2 + .as_ref() + .ok_or_else(|| protocol_error("v2 session is missing its key material")) + } + /// New clients always use protocol v2, for both administrator and temporary credentials. pub fn from_process() -> Result { let credential = get_process_credential().map_err(protocol_error)?; @@ -116,7 +126,7 @@ impl ClientHeaderSession { .await } HeaderProtocol::V2 => { - let material = self.v2.as_ref().expect("v2 session material"); + let material = self.v2_material()?; writer .write_all(&first_prefix(material)) .await @@ -142,7 +152,7 @@ impl ClientHeaderSession { )?)), HeaderProtocol::V2 => Ok(HeaderMessageReader::V2(V2MessageReader::new( reader, - self.v2.as_ref().expect("v2 session material").clone(), + self.v2_material()?.clone(), DIRECTION_SERVER_TO_CLIENT, 0, )?)), @@ -187,7 +197,7 @@ impl ClientHeaderSession { )?)), HeaderProtocol::V2 => Ok(HeaderMessageWriter::V2(V2MessageWriter::new( writer, - self.v2.as_ref().expect("v2 session material").clone(), + self.v2_material()?.clone(), DIRECTION_CLIENT_TO_SERVER, 1, )?)), @@ -215,6 +225,14 @@ impl fmt::Debug for ServerHeaderSession { } impl ServerHeaderSession { + /// The v2 material, `Some` exactly when `protocol` is `V2` — see + /// [`ClientHeaderSession::v2_material`]. + fn v2_material(&self) -> Result<&V2Material> { + self.v2 + .as_ref() + .ok_or_else(|| protocol_error("v2 session is missing its key material")) + } + pub fn protocol(&self) -> HeaderProtocol { self.protocol } @@ -259,7 +277,7 @@ impl ServerHeaderSession { )?)), HeaderProtocol::V2 => Ok(HeaderMessageWriter::V2(V2MessageWriter::new( writer, - self.v2.as_ref().expect("v2 session material").clone(), + self.v2_material()?.clone(), DIRECTION_SERVER_TO_CLIENT, 0, )?)), @@ -278,7 +296,7 @@ impl ServerHeaderSession { )?)), HeaderProtocol::V2 => Ok(HeaderMessageReader::V2(V2MessageReader::new( reader, - self.v2.as_ref().expect("v2 session material").clone(), + self.v2_material()?.clone(), DIRECTION_CLIENT_TO_SERVER, 1, )?)), @@ -515,12 +533,17 @@ impl ServerSecurity { false, )); } + // The prefix-length check above fixes all three widths, so none of these + // can fail. Reported rather than asserted: this parses the first bytes an + // unauthenticated peer sends, and a panic there is a remote abort. + let malformed = + || ServerInitialError::fail("protocol_error", "v2 prefix is malformed", false); let key_id = KeyId::from_u64(u64::from_be_bytes( - remainder[4..12].try_into().expect("fixed key id"), + remainder[4..12].try_into().map_err(|_| malformed())?, )); let salt: [u8; CONNECTION_SALT_LEN] = - remainder[12..28].try_into().expect("fixed connection salt"); - let client_timestamp = u64::from_be_bytes(salt[..8].try_into().expect("fixed timestamp")); + remainder[12..28].try_into().map_err(|_| malformed())?; + let client_timestamp = u64::from_be_bytes(salt[..8].try_into().map_err(|_| malformed())?); let now = unix_seconds(); if now.abs_diff(client_timestamp) > MAX_CONNECTION_CLOCK_SKEW_SECONDS { return Err(ServerInitialError::fail_key( diff --git a/crates/pb-mapper-protocol/src/secure/replay.rs b/crates/pb-mapper-protocol/src/secure/replay.rs index bb25e2d..8d240c0 100644 --- a/crates/pb-mapper-protocol/src/secure/replay.rs +++ b/crates/pb-mapper-protocol/src/secure/replay.rs @@ -14,6 +14,12 @@ //! end of a window with a max-future timestamp cannot be replayed after rotation. //! Per-credential counts stop one tenant from filling the shared filter with //! unique salts before the request payload is decoded. +//! +//! The `expect`s here are all slices of a `[u8; REPLAY_RECORD_LEN]` at constant +//! offsets, or SHA-256's 32-byte output — widths the array types already fix, so +//! the conversions cannot fail. Unlike the parsing paths, these sit in functions +//! that return `()`, so there is nothing to report a width error to. +#![allow(clippy::expect_used)] use pb_mapper_auth::KeyId; diff --git a/crates/pb-mapper-server/Cargo.toml b/crates/pb-mapper-server/Cargo.toml index 21d3d3d..6efb20a 100644 --- a/crates/pb-mapper-server/Cargo.toml +++ b/crates/pb-mapper-server/Cargo.toml @@ -22,3 +22,6 @@ rand.workspace = true [features] udp-timeout = ["uni-stream/udp-timeout", "pb-mapper-protocol/udp-timeout"] + +[lints] +workspace = true diff --git a/ui/native/pb_mapper_ffi/Cargo.toml b/ui/native/pb_mapper_ffi/Cargo.toml index 03948b4..be8a7aa 100644 --- a/ui/native/pb_mapper_ffi/Cargo.toml +++ b/ui/native/pb_mapper_ffi/Cargo.toml @@ -6,9 +6,8 @@ edition = "2024" [lib] crate-type = ["cdylib", "staticlib"] -[lints.clippy] -unwrap_used = "deny" -expect_used = "deny" +[lints] +workspace = true [dependencies] serde = { version = "1.0.219", features = ["derive"] } From 7b7cebe279f3ff2b37095677e7a33cbbecde503d Mon Sep 17 00:00:00 2001 From: LB7666 Date: Fri, 21 Aug 2026 14:14:29 +0800 Subject: [PATCH 74/74] Correct the documentation for the split, and the drift it uncovered MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `claude-code-review.yml` feeds `AGENTS.md` and `CLAUDE.md` into every PR review, so a stale path there actively misleads. Both now describe the crate layout, and the paths in them exist. Most of what was wrong predates this refactor. `CLAUDE.md` documented `src/common/stream.rs`, `src/common/listener.rs`, and `src/utils/udp.rs`, all deleted when the stream abstractions moved out to `uni-stream`; five Flutter files that do not exist, including a "Server Management" section for a view that was never there; `LocalService` for a type named `LocalServer`; and four role commands where there are five — `admin` was missing from both files. It listed three environment variables against roughly twenty read by the code, so it now names the common ones and points at the `PB_MAPPER_*` constants for the rest. Three places claimed web/wasm support; there is no wasm target and the UI loads a native library over `dart:ffi`, which the web cannot do. The dead build-profiles section went with the profiles. `AGENTS.md` was the only file naming both the edition and the toolchain version, and now points at `[workspace.package]` and `rust-toolchain.toml` instead, so the next upgrade does not have to touch prose. The auth docs used brace expansion — `auth/{actor,persistence,...}.rs` — which expanded to `actor.rs` and `persistence.rs`; both are directories. The README badges and the intro doc said Rust 2021. Both architecture diagrams were regenerated, since their six boxes were labelled `src/...`. `ui-cli-mode-spec.md` and `rust-async-send-sync-pin-deep-dive.md` are narrative articles whose line references had already drifted before this work — one cites `pb_server/mod.rs:932` in a 384-line file. They carry a historical-archive note rather than line-by-line fixes, which would drift again on the next change. Co-Authored-By: Claude Fable 5 --- AGENTS.md | 21 +- CLAUDE.md | 215 ++++++++++++------ README.md | 8 +- README.zh-CN.md | 8 +- docs/authentication-v2.md | 23 +- docs/authentication-v2.zh-CN.md | 21 +- docs/pb-mapper-intro.zh-CN.md | 2 +- docs/rust-async-send-sync-pin-deep-dive.md | 4 + .../rust-shared-mutability-and-locks.zh-CN.md | 8 +- docs/ui-cli-mode-spec.md | 5 + 10 files changed, 207 insertions(+), 108 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 5c96352..e21de91 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,19 +1,24 @@ # Repository Guidelines ## Architecture Overview -- One `pb-mapper` binary in `src/bin/` with four role commands: +- One `pb-mapper` binary in `crates/pb-mapper-cli/src/bin/` with five role commands: - `server`: central router (default port 7666) - `register`: registers local TCP/UDP services with the router - `connect`: connects to a registered service and exposes a local port - `status`: queries router IDs and registered keys -- Core crates: `src/pb_server`, `src/local/{server,client}`, `src/common` (protocol, streams, listeners), `src/utils`. + - `admin`: issues, lists, and revokes credentials; rotates the administrator key +- Crates, bottom-up: `pb-mapper-core` (credentials, checksum, config, addressing) + → `pb-mapper-auth` (credential lifecycle and persistence) → `pb-mapper-protocol` + (framing and secure sessions) → `pb-mapper-server` and `pb-mapper-client`, which + are peers → `pb-mapper-cli`. `ui/native/pb_mapper_ffi` is the C ABI cdylib. ## Project Structure & Modules -- `src/`: Rust backend and CLI - - `src/bin/pb-mapper.rs`: unified CLI entry point - - `src/pb_server`, `src/local`, `src/common`, `src/utils` +- `crates/`: the Rust workspace; the root `Cargo.toml` is a virtual manifest + - `crates/pb-mapper-cli/src/bin/pb-mapper.rs`: unified CLI entry point + - `crates/pb-mapper-{core,auth,protocol,server,client,cli}` + - `crates/pb-mapper-cli/tests/`: integration tests; loads env from `tests/.env` + - `crates/pb-mapper-cli/examples/`: runnable examples - `ui/`: Flutter UI; Rust bridge under `ui/native/*` -- `tests/`: integration tests; loads env from `tests/.env` - `docker/`, `services/`, `scripts/`: container, systemd, build/release ## Build, Test, and Development Commands @@ -28,7 +33,9 @@ Notes: CI builds release artifacts on tags `vX.Y.Z` (see `.github/workflows/release.yml`). ## Coding Style & Naming Conventions -- Rust 2021; toolchain pinned via `rust-toolchain.toml` (CI uses 1.88.0) +- Edition is set once in `[workspace.package]`; the toolchain is pinned in + `rust-toolchain.toml`, which CI installs. Both are deliberately not repeated + here — a version in prose goes stale on the next upgrade. - Format: `cargo fmt --all` (4 spaces; import grouping per `rustfmt.toml`) - Lint: `cargo clippy --all-targets -- -D warnings` - Naming: modules/functions `snake_case`, types/traits `PascalCase`, consts `SCREAMING_SNAKE_CASE` diff --git a/CLAUDE.md b/CLAUDE.md index 37d6151..8cb35e4 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -6,7 +6,7 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co This is a Rust-based network tunneling/proxy system called `pb-mapper` that allows exposing local services to clients over a public network. The project enables users to access their home services (like file transfer servers) from anywhere by creating secure tunnels through a public server. -The system uses one **pb-mapper** binary (`src/bin/pb-mapper.rs`) with explicit role commands: +The system uses one **pb-mapper** binary (`crates/pb-mapper-cli/src/bin/pb-mapper.rs`) with explicit role commands: 1. **`pb-mapper server`**: Central server that manages connections between local services and clients - Runs on port 7666 by default @@ -25,7 +25,11 @@ The system uses one **pb-mapper** binary (`src/bin/pb-mapper.rs`) with explicit 4. **`pb-mapper status`**: Queries remote IDs and registered service keys -5. **UI Module** (`ui/`): Flutter graphical interface +5. **`pb-mapper admin`**: Administrator operations against a running server — + issuing, listing, and revoking temporary credentials, rotating the + administrator key, and listing services and connections + +6. **UI Module** (`ui/`): Flutter graphical interface - Replaces all CLI functionality with a user-friendly GUI - Calls into Rust through raw `dart:ffi` against the `pb-mapper-ffi` crate - Provides comprehensive service management interface @@ -36,62 +40,101 @@ The system works by creating a bridge between local services and remote clients ### Project Structure +The root `Cargo.toml` is a virtual manifest; every crate lives under `crates/`, +except the FFI cdylib, which sits next to the Flutter code that loads it. + ``` pb-mapper/ -├── src/ # Main Rust codebase -│ ├── bin/ # Unified pb-mapper CLI entry point -│ ├── pb_server/ # Central server implementation -│ ├── local/ # Local service handlers (server/client) -│ ├── common/ # Shared utilities and protocols -│ └── utils/ # Helper functions +├── crates/ +│ ├── pb-mapper-core/ # Bottom layer: checksum, config, conn_id, error, +│ │ # addr, codec, timeout, durable_file, DataLenType +│ ├── pb-mapper-auth/ # Credential lifecycle, persistence, timing wheel +│ ├── pb-mapper-protocol/ # Message framing, v2 secure sessions, forwarding +│ ├── pb-mapper-server/ # Central relay server, plus the task manager +│ ├── pb-mapper-client/ # Both tunnel ends: `register` and `connect` +│ └── pb-mapper-cli/ # The `pb-mapper` binary, integration tests, examples ├── ui/ # Flutter UI, talking to Rust over dart:ffi │ ├── lib/ # Flutter application code │ │ ├── l10n/ # ARB sources and generated AppLocalizations │ │ └── src/ffi/ # The Dart side of the FFI boundary │ ├── native/pb_mapper_ffi/ # C ABI crate (a workspace member) │ └── test/ # Widget tests -├── examples/ # Example implementations -├── tests/ # Integration tests ├── docker/ # Docker deployment configuration └── services/ # Systemd service files ``` +The dependency graph is a DAG, and the layering is what the crate split +encodes: + +``` +pb-mapper-cli pb-mapper-ffi + │ │ + └────┬─────────────────┤ + ▼ ▼ + pb-mapper-server pb-mapper-client (peers: no reference either way) + └──────┬──────────┘ + ▼ + pb-mapper-protocol + ▼ + pb-mapper-auth + ▼ + pb-mapper-core +``` + +Note that the binary is still named `pb-mapper`, discovered from +`src/bin/pb-mapper.rs` inside `pb-mapper-cli`. The release workflows, both +Dockerfiles, and the install scripts hardcode that name, and `cargo build --bin +pb-mapper` resolves it from the workspace root regardless of the crate name. +Likewise `pb-mapper-ffi` keeps its package name, because it determines the +`libpb_mapper_ffi.{so,dylib,a}` / `pb_mapper_ffi.dll` filenames that the Dart +loader, two CMakeLists, four xcconfigs, and the release-ui hash checks expect. + ### Core Modules -#### Rust Backend (`src/`) -- **`src/pb_server/`**: Central server implementation - - `server.rs`: Main server logic with connection management - - `client.rs`: Client connection handling - - `status.rs`: Server status reporting - - `mod.rs`: Server manager with ManagerTask and ConnTask enums - -- **`src/local/server/`**: Local service registration (`register` functionality) - - `stream.rs`: Stream handling for service registration - - `mod.rs`: Registration logic and server-side CLI implementation - - `error.rs`: Server-specific error handling - -- **`src/local/client/`**: Client connection handling (`connect` functionality) - - `stream.rs`: Stream management for client connections - - `status.rs`: Status checking and reporting - - `mod.rs`: Client-side CLI implementation - - `error.rs`: Client-specific error handling - -- **`src/common/`**: Shared utilities and protocols - - `message/`: Protocol definitions (command.rs, forward.rs) - - `config.rs`: Configuration management and environment variables - - `stream.rs`: Stream abstractions (TcpStreamProvider, UdpStreamProvider) - - `listener.rs`: Listener abstractions (TcpListenerProvider, UdpListenerProvider) - - `manager.rs`: Connection management utilities - - `buffer.rs`: Buffer management for data streaming - - `checksum.rs`: Data integrity verification - - `conn_id.rs`: Connection ID management - - `error.rs`: Common error definitions - -- **`src/utils/`**: Helper functions - - `addr.rs`: Address resolution with OneOrMore enum for multiple addresses - - `codec.rs`: Encryption/decryption utilities - - `timeout.rs`: Timeout handling mechanisms - - `udp.rs`: UDP-specific utilities +#### Rust Backend (`crates/`) +- **`pb-mapper-core/`**: The bottom layer; depends on no other crate here + - `checksum.rs`: The process credential, and the framing checksum over `datalen` + - `config.rs`: Environment configuration and address resolution entry points + - `conn_id.rs`: Connection ID types + - `error.rs`: The shared error type, plus the `snafu_error_*` macros + - `addr.rs`: Address resolution; custom DNS servers on the async path + - `codec.rs`: AES-256-GCM encrypt/decrypt + - `timeout.rs`: `RetryBackoff` + - `durable_file.rs`: Atomic replace and parent-directory fsync + - `test_support.rs`: `PROCESS_CREDENTIAL_TEST_LOCK`, shared across crates' tests + - `lib.rs`: `DataLenType`, which lives here so `checksum` and `error` can name it + +- **`pb-mapper-auth/`**: The credential subsystem, and the largest one + - `lib.rs`: `AuthRuntime`, `AuthContext`, `AuthFailure`, `KeyId` + - `runtime.rs`: Key derivation and authentication of a presented key + - `actor/`: The lifecycle actor — `epoch.rs` for root rotation + - `persistence/`: `snapshot.rs`, `wal.rs`, `blob.rs`, `admin_key.rs`, `fs.rs` + - `timing_wheel.rs`: Hierarchical wheel driving credential expiry + - `leases.rs`, `keys.rs`, `ids.rs`, `config.rs`: Leases, key material, platform dirs + +- **`pb-mapper-protocol/`**: Framing and the authenticated session + - `lib.rs`: The checksum + length framing, and the reader/writer traits + - `command.rs`: Request/response types (`PbConnRequest`, `LocalServer`, `AdminRequest`, …) + - `secure.rs`: Protocol-v2 single-flight sessions, client and server + - `secure/`: `frame.rs`, `first_flight.rs`, `replay.rs`, `limiter.rs` + - `forward.rs`: Stream and datagram forwarding + - `buffer.rs`: Read buffers for the framing + +- **`pb-mapper-server/`**: The central relay + - `lib.rs`: `ManagerTask` / `ConnTask`, and the routing domain model + - `runtime.rs`: Serialises the global routing maps and quotas (the largest file) + - `connection.rs`: Per-socket authentication and dispatch + - `server.rs`, `client.rs`: The service-side and subscriber-side loops + - `admin.rs`: Administrator request handling + - `status.rs`, `error.rs`, `manager.rs`: Status replies, errors, the task manager + +- **`pb-mapper-client/`**: Both ends of a tunnel + - `server/`: `register` — publishes a local service (`mod.rs`, `stream.rs`, `error.rs`) + - `client/`: `connect` — subscribes and listens locally, plus `status.rs` + +- **`pb-mapper-cli/`**: The binary, integration tests, and examples + - `src/bin/pb-mapper.rs`: Argument parsing and the role commands + - `src/bin/pb-mapper/admin.rs`: The `admin` subcommand #### Flutter UI (`ui/`) - **`lib/src/views/`**: One file per zone the shell can show @@ -126,12 +169,17 @@ pb-mapper/ ### Key Components -1. **Message Protocol** (`src/common/message/`): +1. **Message Protocol** (`crates/pb-mapper-protocol/`): - **Command Protocol** (`command.rs`): Defines request/response types: - `PbConnStatusReq`/`PbConnStatusResp`: Status checking - `PbConnRequest`/`PbConnResponse`: Connection management - `PbServerRequest`: Server operation requests - - `LocalService`: Service type definitions (TCP/UDP) + - `LocalServer`: Service type definitions (TCP/UDP) + - `AdminRequest`/`AdminResponse`: Administrator operations + - **Secure sessions** (`secure.rs`): Protocol-v2 first flight — the initial + frame carries a clear-text routing prefix plus an authenticated encrypted + request, adding no extra round trip, and later frames on the connection use + directional keys with monotonic counters - **Forward Protocol** (`forward.rs`): Data forwarding mechanisms - Uses JSON serialization with custom framing (checksum + length header) - Supports encryption/decryption for secure communication via ring crate @@ -142,15 +190,17 @@ pb-mapper/ - Implements keep-alive and timeout mechanisms - Uses actor model for concurrent connection handling -3. **Stream Abstractions**: - - `StreamProvider` trait for TCP/UDP stream handling - - `ListenerProvider` trait for TCP/UDP listener management - - Unified interface for different transport protocols +3. **Stream Abstractions**: `StreamProvider` and `ListenerProvider` give TCP and + UDP one interface. These live in the external `uni-stream` crate, not in this + repository. + +4. **Authentication** (`crates/pb-mapper-auth/`): An administrator key plus + derived temporary credentials, persisted through a write-ahead log and + snapshots, with expiry driven by a hierarchical timing wheel. See + `docs/authentication-v2.md`. -4. **Configuration System**: - - Environment variable support: - - `PB_MAPPER_SERVER`: Remote server address - - `PB_MAPPER_KEEP_ALIVE`: TCP keep-alive setting +5. **Configuration System**: + - Environment variables (see Environment Variables below) - Command-line argument parsing with clap - Workspace-based dependency management @@ -184,16 +234,20 @@ and it is what lets a widget test substitute `FakePbMapperApi` ### Current UI Implementation Status -The UI is fully implemented with the following structure: +Every view under `ui/lib/src/views/`: - **Main App** (`ui/lib/main.dart`): Entry point with navigation and theme management - **Landing Page** (`main_landing_view.dart`): Central navigation hub -- **Server Management** (`server_management_page.dart`, `server_management_view.dart`): Complete server control -- **Service Registration** (`service_registration_page.dart`, `service_registration_view.dart`): Service registration interface -- **Client Connection** (`client_connection_page.dart`, `client_connection_view.dart`): Client connection management +- **Setup Wizard** (`setup_wizard_view.dart`): First-run guided setup +- **Service Registration** (`service_registration_view.dart`): The register workspace +- **Registered Services** (`registered_services_view.dart`): What this process has registered +- **Client Connection** (`client_connection_view.dart`): The connect workspace - **Status Monitoring** (`status_monitoring_view.dart`): Real-time status dashboard - **Configuration** (`configuration_view.dart`): Environment and settings management -- **Logging** (`log_display_widget.dart`, `log_manager.dart`): Comprehensive log viewing +- **Logging** (`log_view_page.dart`, `src/common/log_manager.dart` (under `ui/lib/`)): The log stream + +There is no separate server-management view: starting and stopping the relay is +part of the landing page and the setup wizard. ### UI Features Implemented @@ -265,29 +319,34 @@ The UI is fully implemented with the following structure: - **FFI Integration**: Direct `dart:ffi` calls into the `pb-mapper-ffi` crate - **Real-time Updates**: Live status monitoring and log streaming - **Configuration Management**: Persistent settings and environment variable management -- **Multi-platform**: Desktop, mobile, and web support +- **Multi-platform**: Desktop and mobile. There is no web/wasm target — the UI + loads a native library over `dart:ffi`, which the web cannot do. ## Development Notes ### Project Structure & Dependencies -- **Workspace Configuration**: Multi-crate workspace with shared dependencies in root `Cargo.toml` +- **Workspace Configuration**: Virtual manifest at the root; versions are pinned + once in `[workspace.dependencies]` and crates take them with `.workspace = true` - **Memory Optimization**: Uses mimalloc-rust for improved memory allocation performance -- **Error Handling**: Comprehensive error handling with snafu crate across all modules +- **Error Handling**: snafu, with each crate owning its own error type and wrapping + the layer below as a `source` rather than sharing one workspace-wide enum - **Async Runtime**: Built on Tokio with full async/await support - **Serialization**: serde and serde_json for message serialization -- **Networking**: socket2 for low-level socket operations, trust-dns-resolver for DNS +- **Networking**: uni-stream for the stream/listener abstractions, hickory-resolver + for DNS (custom resolvers on the async path only — the sync path uses `std`, + since hickory has no blocking resolver) - **Cryptography**: ring crate for encryption/decryption functionality ### Code Quality & Standards -- **Linting**: Strict clippy rules in UI native hub (deny unwrap_used, expect_used, wildcard_imports) +- **Linting**: `unwrap_used` and `expect_used` are denied for the whole workspace + via `[workspace.lints]`; `clippy.toml` exempts test code, and `tests/` and + `examples/` targets carry a file-level allow. A production `unwrap` needs a + reason recorded at the site. - **Formatting**: rustfmt.toml configuration for consistent code style - **Toolchain**: rust-toolchain.toml for reproducible builds -- **Testing**: Comprehensive test suite in `tests/` directory - -### Build Profiles -- **wasm-dev**: Optimized for WebAssembly builds -- **server-dev**: Development profile for server components -- **android-dev**: Android-specific build optimizations +- **Testing**: Unit tests live beside the code; integration tests are in + `crates/pb-mapper-cli/tests/`, which is the crate that depends on every layer + they exercise ### UI Development Guidelines - **Framework**: Flutter 3.44.9, Material 3. CI pins the same version. @@ -299,9 +358,23 @@ The UI is fully implemented with the following structure: - **Responsive Design**: Adaptive layouts for different screen sizes ### Environment Variables + +The commonly used ones: + - **`PB_MAPPER_SERVER`**: Default remote server address for CLI tools -- **`PB_MAPPER_KEEP_ALIVE`**: Global TCP keep-alive setting ("ON" to enable) +- **`PB_MAPPER_KEEP_ALIVE`**: TCP keep-alive ("ON", "1", "true", "yes" to enable). + Read on every call, not cached — the UI's per-service toggle depends on that. +- **`MSG_HEADER_KEY`**: The process credential, administrator or temporary. + Required; there is no insecure default. - **`RUST_LOG`**: Tracing level configuration (supports env-filter) +- **`PB_MAPPER_LOG_FORMAT`**: Log output format + +Timeouts, intervals, and pool sizes are also configurable, and there are more +than a dozen: the authoritative list is the `pub const PB_MAPPER_*` declarations +at the top of `crates/pb-mapper-core/src/config.rs`, each read by the accessor +named after it. Beyond those, `PB_MAPPER_AUTH_STATE_DIR`, +`PB_MAPPER_LEGACY_PROTOCOL`, and `PB_MAPPER_NEW_STREAMS_PER_SECOND` are read by +name where they are used; the first two are also settable as `server` flags. ## Development Workflow diff --git a/README.md b/README.md index 24b76fb..e94280c 100644 --- a/README.md +++ b/README.md @@ -3,7 +3,7 @@ pb-mapper

- Rust 2021 + Rust 2024 Tokio Flutter License: MIT @@ -110,7 +110,7 @@ Open `http://localhost:3000` in the coffee-shop browser — traffic flows throug ## Developer view -- **Rust core** — the unified entry point is `src/bin/pb-mapper.rs`; shared protocol and networking live in `src/common` and `src/utils`; server / register / connect internals live in `src/pb_server`, `src/local/server`, and `src/local/client`. +- **Rust core** — a workspace under `crates/`, layered bottom-up: `pb-mapper-core` (credentials, checksum, config, addressing), `pb-mapper-auth` (credential lifecycle and persistence), `pb-mapper-protocol` (framing and secure sessions), then `pb-mapper-server` and `pb-mapper-client` as peers, with the `pb-mapper` binary in `pb-mapper-cli`. - **Flutter UI** — views in `ui/lib/src/views`, FFI layers in `ui/lib/src/ffi`, Rust bridge in `ui/native/pb_mapper_ffi`. FFI calls run on a background isolate, and Rust returns JSON (`{success, message, data}`) to keep the C ABI stable. ## Documentation @@ -122,10 +122,10 @@ Open `http://localhost:3000` in the coffee-shop browser — traffic flows throug ## Repository layout -- `src/` — Rust backend +- `crates/` — the Rust workspace (six crates; the root manifest is virtual) - `ui/` — Flutter UI + native bridge - `docs/` — documentation and assets -- `docker/`, `services/`, `scripts/`, `tests/` — deployment and tooling +- `docker/`, `services/`, `scripts/` — deployment and tooling - `skills/` — AI coding agent deployment skills (server and connect tunnel) ## License diff --git a/README.zh-CN.md b/README.zh-CN.md index e4d37dd..562351e 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -3,7 +3,7 @@ pb-mapper

- Rust 2021 + Rust 2024 Tokio Flutter License: MIT @@ -110,7 +110,7 @@ pb-mapper connect tcp --server :7666 --key web --addr 127.0.0.1:3000 ## 开发者视角 -- **Rust 核心**：统一入口为 `src/bin/pb-mapper.rs`；协议与网络通用逻辑在 `src/common`、`src/utils`；server/register/connect 实现在 `src/pb_server`、`src/local/server`、`src/local/client`。 +- **Rust 核心**：`crates/` 下的 workspace，自底向上分层：`pb-mapper-core`（凭据、校验和、配置、地址解析）→ `pb-mapper-auth`（凭据生命周期与持久化）→ `pb-mapper-protocol`（帧格式与安全会话）→ `pb-mapper-server` 与 `pb-mapper-client`（二者平级，互不引用）→ `pb-mapper-cli`（`pb-mapper` 二进制所在）。 - **Flutter UI**：界面在 `ui/lib/src/views`，FFI 各层在 `ui/lib/src/ffi`，Rust 桥接在 `ui/native/pb_mapper_ffi`。FFI 调用跑在后台 isolate，Rust 统一返回 JSON（`{success, message, data}`）以保持 C ABI 稳定。 ## 文档 @@ -122,10 +122,10 @@ pb-mapper connect tcp --server :7666 --key web --addr 127.0.0.1:3000 ## 仓库结构 -- `src/` — Rust 后端 +- `crates/` — Rust workspace（六个 crate，根清单为虚拟清单） - `ui/` — Flutter UI + 原生桥接 - `docs/` — 文档与素材 -- `docker/`、`services/`、`scripts/`、`tests/` — 部署与工具 +- `docker/`、`services/`、`scripts/` — 部署与工具 - `skills/` — AI 编程助手部署 skill（服务端、客户端隧道） ## 许可证 diff --git a/docs/authentication-v2.md b/docs/authentication-v2.md index 179f1d2..900a0b2 100644 --- a/docs/authentication-v2.md +++ b/docs/authentication-v2.md @@ -321,17 +321,22 @@ state. ## Code index -- Credential format and process configuration: `src/common/checksum.rs` -- Authentication facade and shared model: `src/common/auth.rs` -- Lifecycle actor, persistence, runtime, and timing wheel: - `src/common/auth/{actor,persistence,runtime,timing_wheel}.rs` +- Credential format and process configuration: + `crates/pb-mapper-core/src/checksum.rs` +- Authentication facade and shared model: `crates/pb-mapper-auth/src/lib.rs` +- Lifecycle actor: `crates/pb-mapper-auth/src/actor/` +- Persistence: `crates/pb-mapper-auth/src/persistence/` +- Runtime and timing wheel: `crates/pb-mapper-auth/src/runtime.rs` and + `crates/pb-mapper-auth/src/timing_wheel.rs` - V2 session facade plus frame, limiter, and replay modules: - `src/common/message/secure.rs` and `src/common/message/secure/` + `crates/pb-mapper-protocol/src/secure.rs` and + `crates/pb-mapper-protocol/src/secure/` - Relay state, runtime loop, and connection dispatch: - `src/pb_server/{mod,runtime,connection}.rs` -- Administrator request execution: `src/pb_server/admin.rs` -- Unified CLI and administrator command module: `src/bin/pb-mapper.rs` and - `src/bin/pb-mapper/admin.rs` + `crates/pb-mapper-server/src/lib.rs`, `runtime.rs`, and `connection.rs` +- Administrator request execution: `crates/pb-mapper-server/src/admin.rs` +- Unified CLI and administrator command module: + `crates/pb-mapper-cli/src/bin/pb-mapper.rs` and + `crates/pb-mapper-cli/src/bin/pb-mapper/admin.rs` ## Summary diff --git a/docs/authentication-v2.zh-CN.md b/docs/authentication-v2.zh-CN.md index a23502e..af9b2e2 100644 --- a/docs/authentication-v2.zh-CN.md +++ b/docs/authentication-v2.zh-CN.md @@ -250,15 +250,20 @@ Docker 必须持久化 `/var/lib/pb-mapper/auth`；否则重建容器会产生 ## 代码索引 -- 凭据格式与进程配置：`src/common/checksum.rs` -- 认证 facade 与共享模型：`src/common/auth.rs` -- 生命周期 actor、持久化、runtime 与时间轮： - `src/common/auth/{actor,persistence,runtime,timing_wheel}.rs` +- 凭据格式与进程配置：`crates/pb-mapper-core/src/checksum.rs` +- 认证 facade 与共享模型：`crates/pb-mapper-auth/src/lib.rs` +- 生命周期 actor：`crates/pb-mapper-auth/src/actor/` +- 持久化：`crates/pb-mapper-auth/src/persistence/` +- runtime 与时间轮：`crates/pb-mapper-auth/src/runtime.rs`、 + `crates/pb-mapper-auth/src/timing_wheel.rs` - V2 session facade、frame、限流与 replay 模块： - `src/common/message/secure.rs` 与 `src/common/message/secure/` -- 中继状态、runtime loop 与连接分发：`src/pb_server/{mod,runtime,connection}.rs` -- 管理请求执行：`src/pb_server/admin.rs` -- 统一 CLI 与管理员命令模块：`src/bin/pb-mapper.rs`、`src/bin/pb-mapper/admin.rs` + `crates/pb-mapper-protocol/src/secure.rs` 与 + `crates/pb-mapper-protocol/src/secure/` +- 中继状态、runtime loop 与连接分发：`crates/pb-mapper-server/src/lib.rs`、 + `runtime.rs`、`connection.rs` +- 管理请求执行：`crates/pb-mapper-server/src/admin.rs` +- 统一 CLI 与管理员命令模块：`crates/pb-mapper-cli/src/bin/pb-mapper.rs`、 + `crates/pb-mapper-cli/src/bin/pb-mapper/admin.rs` ## 总结 diff --git a/docs/pb-mapper-intro.zh-CN.md b/docs/pb-mapper-intro.zh-CN.md index 0497b0f..0f3dc6d 100644 --- a/docs/pb-mapper-intro.zh-CN.md +++ b/docs/pb-mapper-intro.zh-CN.md @@ -45,7 +45,7 @@ flowchart LR ## 技术栈 -- **语言和运行时**：Rust 2021 + Tokio 异步运行时 +- **语言和运行时**：Rust 2024 edition + Tokio 异步运行时 - **内存分配器**：自己 fork 的 [`better_mimalloc_rs`](https://github.com/acking-you/better_mimalloc_rs)，后面会细说为什么 - **网络抽象**：自研 [`uni-stream`](https://github.com/acking-you/uni-stream)，把 TCP 和 UDP 统一成一套流接口；底层用 `socket2` 控制 socket 选项，`trust-dns-resolver` 做 DNS - **协议**：serde_json 序列化，自定义帧格式（checksum + 长度头），可选 `ring` 做 AES-256-GCM 端到端加密 diff --git a/docs/rust-async-send-sync-pin-deep-dive.md b/docs/rust-async-send-sync-pin-deep-dive.md index 3af9ca6..6a4aa5e 100644 --- a/docs/rust-async-send-sync-pin-deep-dive.md +++ b/docs/rust-async-send-sync-pin-deep-dive.md @@ -1,5 +1,9 @@ # Rust 并发安全（Send/Sync/Pin）与 async 状态机深度解析 +> 历史设计文档：反映 2025 年某时期的实现，其中的行号与代码引用可能已漂移 +> （拆分为多 crate 后所有路径都已移到 `crates/` 下），仅供设计意图参考。 +> 当前实现请以代码为准。 + 面向场景：你在实现网络转发（如 pb-mapper）时需要理解 **为什么某些 future 必须 `Send + 'static`**、为什么 `async fn` 能跨 `.await` 持有借用、以及 `Pin` 如何保证自引用安全。 > **Code Version**: pb-mapper 本地工作区（2026-01-17） diff --git a/docs/rust-shared-mutability-and-locks.zh-CN.md b/docs/rust-shared-mutability-and-locks.zh-CN.md index 08a4a5f..2b83971 100644 --- a/docs/rust-shared-mutability-and-locks.zh-CN.md +++ b/docs/rust-shared-mutability-and-locks.zh-CN.md @@ -10,7 +10,7 @@ ## 1. 问题从哪里来 -重构 `src/common/auth/timing_wheel.rs` 时，中间某一版长成这样： +重构 `crates/pb-mapper-auth/src/timing_wheel.rs` 时，中间某一版长成这样： ```rust struct Queues { @@ -245,7 +245,7 @@ where F: Future + Send + 'static, F::Output: Send + 'static 第 1 节那版每级一把 `Mutex`，走的是「`Drop` 里投递 → 需要共享 `Queues` → 加锁」。 现在 `Link::Relay` 只是纯数据，`tick` 拿 `&mut self` 自己投递 -（`src/common/auth/timing_wheel.rs`）： +（`crates/pb-mapper-auth/src/timing_wheel.rs`）： ```rust Link::Relay { level, slot, next } => self.file(level as usize, slot as usize, *next), @@ -257,7 +257,7 @@ Link::Relay { level, slot, next } => self.file(level as usize, slot as usize, *n ```rust pub struct AuthLease { - expires_at: AtomicU64, // src/common/auth.rs + expires_at: AtomicU64, // crates/pb-mapper-auth/src/lib.rs ... } ``` @@ -345,7 +345,7 @@ struct AuthStateInner { lock.read().unwrap_or_else(|poisoned| poisoned.into_inner()) ``` -即「无论如何都取出内部值」，等于把 poisoning 显式关掉。`src/common/auth.rs` 里还为此 +即「无论如何都取出内部值」，等于把 poisoning 显式关掉。`crates/pb-mapper-auth/src/lib.rs` 里还为此 养了一个 `recover_lock` 辅助函数专门抹掉 `LockResult`。 `parking_lot` 不做 poisoning，于是： diff --git a/docs/ui-cli-mode-spec.md b/docs/ui-cli-mode-spec.md index c94d47c..6c5bca3 100644 --- a/docs/ui-cli-mode-spec.md +++ b/docs/ui-cli-mode-spec.md @@ -1,5 +1,10 @@ # pb-mapper UI as a CLI — design spec +> Historical design document: this reflects the implementation at some point in +> 2025. Its line numbers and code references may have drifted — the crate split +> moved every path under `crates/` — and it is kept for design intent only. For +> current behaviour, read the code. + > Historical design note: the standalone CLI was consolidated in v0.3.0. In > current commands, `pb-mapper-server` maps to `pb-mapper server`, > `pb-mapper-server-cli` maps to `pb-mapper register`, and