From 6a2a0b9a899ecf93fc2f35f496a4b061f24e1f7c Mon Sep 17 00:00:00 2001 From: Cristian Girlea Date: Wed, 21 Oct 2026 22:19:58 +0300 Subject: [PATCH] feat: scope jobs to the token that submitted them The worker takes its tokens from BRIDGE_TOKEN (named default), BRIDGE_TOKENS or BRIDGE_TOKENS_FILE, at most 64, each checked at start with a message that names the variable and never a token, and both token variables leave the environment before the first child starts. A job belongs to the token that submitted it: another token reads and cancels it as 404, exactly like an unknown job, and idempotency keys are per token. The owner name is kept in jobs.sqlite (an older file gains the column with its rows owned by default) and is not in the job snapshot; the three metrics counters gain an owner label. A two-token compose peer carries the PHP and Python integration checks; the PHP client and the MCP server are unchanged. --- .claude/rules/php-client.md | 2 +- .claude/rules/worker.md | 7 +- .claude/settings.json | 4 +- .claude/skills/check/SKILL.md | 2 +- .github/workflows/tests.yml | 2 +- AGENTS.md | 2 +- CHANGELOG.md | 4 + README.md | 4 +- docker/compose.yaml | 20 +++ docs/integration.md | 2 +- docs/mcp.md | 6 +- docs/protocol.md | 24 ++- docs/security.md | 4 +- docs/testing.md | 8 +- tests/integration.py | 61 ++++++- tests/php/integration.php | 28 +++ tests/python/test_jobs.py | 207 ++++++++++++++++++++-- tests/python/test_server.py | 315 +++++++++++++++++++++++++++++++++- worker/ai_bridge/jobs.py | 86 ++++++---- worker/ai_bridge/metrics.py | 11 +- worker/ai_bridge/server.py | 130 ++++++++++++-- worker/ai_bridge/tasks.py | 20 ++- 22 files changed, 833 insertions(+), 116 deletions(-) diff --git a/.claude/rules/php-client.md b/.claude/rules/php-client.md index 6096b44..312b270 100644 --- a/.claude/rules/php-client.md +++ b/.claude/rules/php-client.md @@ -15,7 +15,7 @@ paths: job's pinned image from `.github/workflows/tests.yml` and rerun the same compose commands. - Tests are `check()` and `rejects()` calls: `tests/php/run.php` runs offline (contract checks and one transport error), and `tests/php/integration.php` loads it and then drives the client against the real - `worker`, `fault-worker` and `hostile` peers over HTTP. The path through FrankenPHP is + `worker`, `fault-worker`, `tokens-worker` (two tokens) and `hostile` peers over HTTP. The path through FrankenPHP is `tests/integration.py`. A new behaviour gets a success case, invalid-input cases and, where the worker is involved, an integration case. Syntax is checked with `php -l` over `src`, `examples` and `tests/php`. - The client never retries a `POST` on its own (a transport failure leaves the submission uncertain), and diff --git a/.claude/rules/worker.md b/.claude/rules/worker.md index f83d613..ee720d3 100644 --- a/.claude/rules/worker.md +++ b/.claude/rules/worker.md @@ -17,9 +17,10 @@ paths: - `MODEL_NAMES` is what every result reports as `model` and what the MCP server's `BACKENDS` table must equal (`tests/mcp/test_models.py`); change both together. - Limits live in three places: concurrency, capacity and retention are `Settings` in `jobs.py` (defaults - from the environment in `server.main`), the body size is `MAX_BODY` in `server.py`, and the field, - count and character limits are the `MAX_*` constants in `tasks.py`; the deadline range is checked in - `JobStore.submit`. `docs/protocol.md` documents them; change the doc with the limit. + from the environment in `server.main`), the body size and the token limits are `MAX_BODY`, `MAX_TOKENS`, + `MAX_TOKENS_FILE` and `TOKEN_NAME` in `server.py`, and the field, count and character limits are the + `MAX_*` constants in `tasks.py`; the deadline range is checked in `JobStore.submit`. `docs/protocol.md` + documents them; change the doc with the limit. - Tests run offline with `python -m unittest discover -s tests/python` inside `python-tests` (`network_mode: none`). A change to the model path belongs on `model-smoke.yml`'s `paths:` list, so the real-model smoke tests run for it; keep that list current when a file moves. diff --git a/.claude/settings.json b/.claude/settings.json index c07235e..9b09ca7 100644 --- a/.claude/settings.json +++ b/.claude/settings.json @@ -25,7 +25,7 @@ "Bash(docker compose -f docker/compose.yaml run --rm --no-deps laravel-tests)", "Bash(docker compose -f docker/compose.yaml run --rm --no-deps symfony-tests)", "Bash(docker compose -f docker/compose.yaml run --rm --no-deps tls-check)", - "Bash(docker compose -f docker/compose.yaml up -d worker frankenphp fault-worker hostile)", + "Bash(docker compose -f docker/compose.yaml up -d worker frankenphp fault-worker hostile tokens-worker)", "Bash(docker compose -f docker/compose.yaml up -d worker frankenphp)", "Bash(docker compose -f docker/compose.yaml up -d worker)", "Bash(docker compose -f docker/compose.yaml up -d tls-proxy)", @@ -83,7 +83,7 @@ "PowerShell(docker compose -f docker/compose.yaml run --rm --no-deps laravel-tests)", "PowerShell(docker compose -f docker/compose.yaml run --rm --no-deps symfony-tests)", "PowerShell(docker compose -f docker/compose.yaml run --rm --no-deps tls-check)", - "PowerShell(docker compose -f docker/compose.yaml up -d worker frankenphp fault-worker hostile)", + "PowerShell(docker compose -f docker/compose.yaml up -d worker frankenphp fault-worker hostile tokens-worker)", "PowerShell(docker compose -f docker/compose.yaml up -d worker frankenphp)", "PowerShell(docker compose -f docker/compose.yaml up -d worker)", "PowerShell(docker compose -f docker/compose.yaml up -d tls-proxy)", diff --git a/.claude/skills/check/SKILL.md b/.claude/skills/check/SKILL.md index 00a2396..d2c5589 100644 --- a/.claude/skills/check/SKILL.md +++ b/.claude/skills/check/SKILL.md @@ -16,7 +16,7 @@ docker compose -f docker/compose.yaml --profile test config docker compose -f docker/compose.yaml run --rm --no-deps python-tests docker compose -f docker/compose.yaml run --rm --no-deps php-tests docker compose -f docker/compose.yaml run --rm --no-deps php-tests sh -ec 'find src examples tests/php -name "*.php" -print0 | xargs -0 -n1 php -l' -docker compose -f docker/compose.yaml up -d worker frankenphp fault-worker hostile +docker compose -f docker/compose.yaml up -d worker frankenphp fault-worker hostile tokens-worker docker compose -f docker/compose.yaml run --rm --no-deps php-integration docker compose -f docker/compose.yaml run --rm --no-deps integration docker compose -f docker/compose.yaml --profile test down diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 2ab21cc..a96dd01 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -46,7 +46,7 @@ jobs: docker compose -f docker/compose.yaml run --rm --no-deps php-tests sh -ec 'find src examples tests/php -name "*.php" -print0 | xargs -0 -n1 php -l' - name: HTTP and FrankenPHP integration run: | - docker compose -f docker/compose.yaml up -d worker frankenphp fault-worker hostile + docker compose -f docker/compose.yaml up -d worker frankenphp fault-worker hostile tokens-worker docker compose -f docker/compose.yaml run --rm --no-deps php-integration docker compose -f docker/compose.yaml run --rm --no-deps integration - name: Service logs on failure diff --git a/AGENTS.md b/AGENTS.md index 4b60664..0263c0d 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -33,7 +33,7 @@ docker compose -f docker/compose.yaml --profile test config docker compose -f docker/compose.yaml run --rm --no-deps python-tests docker compose -f docker/compose.yaml run --rm --no-deps php-tests docker compose -f docker/compose.yaml run --rm --no-deps php-tests sh -ec 'find src examples tests/php -name "*.php" -print0 | xargs -0 -n1 php -l' -docker compose -f docker/compose.yaml up -d worker frankenphp fault-worker hostile +docker compose -f docker/compose.yaml up -d worker frankenphp fault-worker hostile tokens-worker docker compose -f docker/compose.yaml run --rm --no-deps php-integration docker compose -f docker/compose.yaml run --rm --no-deps integration docker compose -f docker/compose.yaml --profile test down diff --git a/CHANGELOG.md b/CHANGELOG.md index b430738..86dfe04 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,10 @@ All notable changes to this project are documented in this file. The format foll ### Added +- Per-client tokens on the worker: `BRIDGE_TOKENS=name=token,...` or `BRIDGE_TOKENS_FILE` (one `name=token` per line, mode 0600 or stricter), at most 64, as an alternative to `BRIDGE_TOKEN`, which is the single token named `default`. A bad token, name, file or combination of the three stops the worker at start with a message naming the variable and never a token. +- A job belongs to the token that submitted it: another token's `GET` or `cancel` of it is 404, exactly as for an unknown job, and idempotency keys are per token. The owner's name is kept in `jobs.sqlite` (a file from an earlier version gains the column, its rows owned by `default`) and is not in the job's snapshot. A worker with only `BRIDGE_TOKEN` behaves as before. +- An `owner` label, the token's name, on `bridge_jobs_total`, `bridge_jobs_submitted_total` and `bridge_submissions_rejected_total`, with every configured name's series present at zero. + - `PhpAiBridge\Laravel\BridgeServiceProvider`, registered by Laravel's package discovery, binds one `Client` built from a publishable `config/ai_bridge.php` (`BRIDGE_URL`, `BRIDGE_TOKEN`, `BRIDGE_REQUEST_TIMEOUT_MS`; `vendor:publish --tag=ai-bridge-config`), so configuration caching works. - `PhpAiBridge\Symfony\PhpAiBridgeBundle` registers `Client` as an autowirable service from a `php_ai_bridge` configuration whose `url`, `token` and `request_timeout_ms` default to the same environment variables. - Both fail on a missing or invalid setting with a message that names it and never quotes the token. `laravel/framework` and `symfony/http-kernel` are suggested, not required: the library's requirements are unchanged. diff --git a/README.md b/README.md index 9929a26..f9cb175 100644 --- a/README.md +++ b/README.md @@ -112,7 +112,7 @@ docker compose -f docker/compose.yaml run --rm --no-deps python-tests docker compose -f docker/compose.yaml run --rm --no-deps php-tests # Internal-only HTTP peers, including intentionally broken peers: -docker compose -f docker/compose.yaml up -d worker frankenphp fault-worker hostile +docker compose -f docker/compose.yaml up -d worker frankenphp fault-worker hostile tokens-worker docker compose -f docker/compose.yaml run --rm --no-deps php-integration docker compose -f docker/compose.yaml run --rm --no-deps integration @@ -141,7 +141,7 @@ CI runs deterministic tests for pushes to main and for pull requests across PHP ## Important limits - **Durable job table only when `BRIDGE_STATE_DIR` is set:** the worker then keeps its job table (never inputs) in SQLite, so finished results and idempotency keys survive a restart until retention ends and unfinished jobs become `failed` with `service_restarted`; unset, a restart loses everything. Completed results expire after 300 seconds by default; 404 can mean expired, lost or unknown. An optional `idempotency_key` makes resubmitting after an uncertain POST safe while the job is retained. There is still no replay of unfinished jobs, durable delivery or automatic retry. -- **Trusted services only:** one bearer token authorizes all jobs. Your application must authenticate users, associate job IDs with their owners, and authorize reads/cancellation. Random IDs are not authorization. +- **Trusted services only:** a bearer token reaches only the jobs it submitted; give each application its own (`BRIDGE_TOKENS`, [the protocol](docs/protocol.md#tokens)), since they still share one worker's capacity. Your application must authenticate users, associate job IDs with their owners, and authorize reads/cancellation. Random IDs are not authorization. - **Private HTTP transport:** deploy behind an authenticated TLS boundary for cross-host use; the [integration guide](docs/integration.md#cross-host-use-behind-tls) has a tested example with Caddy in front of the worker. The sample Python HTTP server is not an internet-facing production server. - **Process isolation is not a sandbox for arbitrary code:** task code is trusted and registered in source. Clients cannot upload Python, import modules or select a shell command. The container limits the entire service, not each job separately. - **One job per warm child, no sharing across jobs:** the worker keeps `BRIDGE_CONCURRENCY` child processes alive and each keeps the models it has loaded, so only a child's first job of a kind pays for imports and model loading; a child is replaced after a cancellation, a deadline or a crash, and after `BRIDGE_WORKER_MAX_JOBS` jobs (default 1000). Within one rerank job the ONNX backend scores documents 32 pairs per forward pass; there is no batching across jobs, no sharing of a loaded model between children, and no claim of GPU support, multi-host scaling or stateful agents. diff --git a/docker/compose.yaml b/docker/compose.yaml index 294a3d4..b4f8531 100644 --- a/docker/compose.yaml +++ b/docker/compose.yaml @@ -250,6 +250,9 @@ services: <<: *python-env FRANKENPHP_URL: http://frankenphp:8080 WORKER_URL: http://worker:8090 + TOKENS_WORKER_URL: http://tokens-worker:8090 + BRIDGE_APP_TOKEN: test-only-app-client-token-never-use-in-production + BRIDGE_OTHER_TOKEN: test-only-other-client-token-never-use-in-production command: [python, tests/integration.py] php-example: @@ -268,6 +271,18 @@ services: BRIDGE_CONCURRENCY: "1" command: [python, -m, ai_bridge.server, --host, 0.0.0.0] + # Two clients of one worker, each with its own test-only literal token, independent of the run's BRIDGE_TOKEN. + # The environment is written out rather than merged, since BRIDGE_TOKEN beside BRIDGE_TOKENS is a start error. + tokens-worker: + <<: *python + profiles: [test] + networks: [isolated] + environment: + PYTHONDONTWRITEBYTECODE: "1" + PYTHONPATH: /work/worker + BRIDGE_TOKENS: app=test-only-app-client-token-never-use-in-production,other=test-only-other-client-token-never-use-in-production + command: [python, -m, ai_bridge.server, --host, 0.0.0.0] + hostile: <<: *python profiles: [test] @@ -278,6 +293,11 @@ services: <<: *php profiles: [test] networks: [isolated] + environment: + <<: *php-env + BRIDGE_TOKENS_URL: http://tokens-worker:8090 + BRIDGE_APP_TOKEN: test-only-app-client-token-never-use-in-production + BRIDGE_OTHER_TOKEN: test-only-other-client-token-never-use-in-production command: [php, tests/php/integration.php] fetcher: diff --git a/docs/integration.md b/docs/integration.md index 7112936..de73f39 100644 --- a/docs/integration.md +++ b/docs/integration.md @@ -255,7 +255,7 @@ Implement the following in your application, for either framework: 4. On a later status request, verify ownership **before** calling `get()`. Handle `queued`/`running`, `succeeded`, `failed`, `timed_out` and `cancelled`. For success, validate using `RerankResult::fromJob()` with the original document count; indexes refer to the original order. 5. Authorize cancellation the same way before calling `cancel()`. A job may complete before cancellation arrives; the returned snapshot is authoritative. -The service token authorizes every job. Do not send it to a browser or rely on random job IDs for authorization. Do not blindly publish the example routes as your production API. Capture required results before retention expires. Without `BRIDGE_STATE_DIR` the worker is in-memory and a restart loses all jobs and results; with it, finished results survive a restart until retention ends, and jobs that were unfinished become `failed` with `service_restarted`. A 404 cannot distinguish expiry, restart loss and an unknown job. +A service token reaches every job submitted with it, whichever of your users it was for; give each application its own ([tokens](protocol.md#tokens)). Do not send it to a browser or rely on random job IDs for authorization. Do not blindly publish the example routes as your production API. Capture required results before retention expires. Without `BRIDGE_STATE_DIR` the worker is in-memory and a restart loses all jobs and results; with it, finished results survive a restart until retention ends, and jobs that were unfinished become `failed` with `service_restarted`. A 404 cannot distinguish expiry, restart loss and an unknown job. When a submission's outcome is uncertain (`transport_error`, or the PHP process died before it stored the job ID), resubmit with the same idempotency key instead of a fresh request: the worker returns the job the first request created, in whatever state it is now, or creates it if the first request never arrived. Pick the key per unit of work (for example a hash of your own record ID and the task) and store it before the first POST: diff --git a/docs/mcp.md b/docs/mcp.md index 00417f4..d510a46 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -84,7 +84,7 @@ On Linux the fixture directory must be writable by container UID 65532 for the b | Variable | Meaning | | --- | --- | -| `BRIDGE_TOKEN` | The worker's service token, required. Read from the environment only; never returned. | +| `BRIDGE_TOKEN` | The worker's service token, required. Read from the environment only; never returned. On a worker shared with an application, this is the MCP server's own entry in the worker's token list (`BRIDGE_TOKENS` or `BRIDGE_TOKENS_FILE`), named for example `mcp`, so an injected agent sees only its own jobs; compose's `mcp-worker` is dedicated and keeps a single token. | | `BRIDGE_URL` | The worker origin, default `http://127.0.0.1:8090`. | | `BRIDGE_MCP_ROOT` | The one directory whose files `documents_path` may name. Unset means paths are refused. | | `BRIDGE_MCP_TIMEOUT_MS` | Per-call wait before cancelling, 100–300000, default 60000. | @@ -127,10 +127,10 @@ docker compose -f docker/compose.yaml run --rm -i -T mcp-model ## Trust model -Nothing about the worker's trust model changes: a private network, one bearer token, never internet-facing. Two things are new and must be understood: +Nothing about the worker's trust model changes: a private network, a bearer token per client, never internet-facing. Two things are new and must be understood: - **Over stdio there is no authentication.** The host process launches the server and inherits its trust; the server holds `BRIDGE_TOKEN` from its environment. Running this server over HTTP or SSE would be a new, unsolved authorisation surface and is out of scope; the bearer token does not cover it. -- **An agent can now submit jobs, and an agent can be prompt-injected.** It cannot read the token, but it can fill the worker's capacity or submit junk. Point the MCP server at a **dedicated worker**, never the one serving a production PHP application; the compose `mcp` service does this with `mcp-worker`. Capacity, concurrency and retention are per worker. +- **An agent can now submit jobs, and an agent can be prompt-injected.** It cannot read the token, but it can fill the worker's capacity or submit junk. Point the MCP server at a **dedicated worker**, never the one serving a production PHP application; the compose `mcp` service does this with `mcp-worker`. Capacity, concurrency and retention are per worker. Where a worker must be shared, give the MCP server its own token (`mcp=...` in the worker's `BRIDGE_TOKENS`): the agent then cannot read or cancel the application's jobs, which answer it 404, though it can still fill the shared capacity. Paths are confined to `BRIDGE_MCP_ROOT`, including an index's vectors file after symlinks are resolved; the server never writes files, and only the operator-run `mcp-index` builder writes, into the directory it is told to. Results contain indexes, scores, snippets of the agent's own files, masked text and spans; the original documents are not echoed by the worker. diff --git a/docs/protocol.md b/docs/protocol.md index 5c6796f..d9f8a22 100644 --- a/docs/protocol.md +++ b/docs/protocol.md @@ -10,6 +10,18 @@ Every route requires `Authorization: Bearer `. JSON requests requ | `GET /v1/jobs/{id}` | 200 with a snapshot, or 404 | | `POST /v1/jobs/{id}/cancel` | Send `{}`; returns the terminal/current snapshot, or 404 | +## Tokens + +The worker takes its tokens from exactly one of three variables, read once at start: + +- `BRIDGE_TOKEN=`: a single token, named `default`. +- `BRIDGE_TOKENS=name1=token1,name2=token2`: one entry per client. A token here cannot contain a comma; use the file for one that does. +- `BRIDGE_TOKENS_FILE=`: one `name=token` per line; lines that are blank or whose first non-space character is `#` are ignored, and entry lines are not trimmed. The file must be a regular file of at most 65536 bytes with mode 0600 or stricter (no access for group or others). + +A name is 1–64 characters of `A-Z`, `a-z`, `0-9`, `_`, `.` and `-`; a token is at least 32 non-whitespace ASCII characters, everything after the first `=` of its entry. At most 64 tokens; no name or token may appear twice. An empty variable counts as unset, so compose's `${X:-}` forms work. Setting `BRIDGE_TOKEN` together with either list, or both lists, is refused. The worker refuses to start on any of these, and on a file it cannot read or that group or others can, with a message naming the variable and never a token. `BRIDGE_TOKEN` and `BRIDGE_TOKENS` leave the worker's environment before it starts any child process; the file's path stays, its content does not enter the environment. + +A request is authenticated by comparing its bearer with every configured token in constant time; the matching token's name is the request's owner. A job belongs to the owner that submitted it: `GET /v1/jobs/{id}` and `POST /v1/jobs/{id}/cancel` with another token answer 404 `not_found`, exactly as for a job that does not exist, never 403. `GET /healthz` and `GET /metrics` answer any configured token. The owner is not part of a job's snapshot. Capacity, concurrency and retention are shared by every owner. + Submit: ```json @@ -18,7 +30,7 @@ Submit: `rerank`, `embed` and `redact` are enabled normally. For `rerank`, query and each document must contain 1–8192 characters with non-whitespace content; 1–512 documents are accepted, and the query plus all documents may total at most 200000 characters. The 262144 byte limit still applies to the encoded request, so multi-byte text and escaped quotes can reach it before the character budget; the PHP client counts the budget in bytes and sends raw UTF-8 to keep the two close. An optional integer `top_k` from 1 to the document count limits how many rankings are returned; omitting it returns every document. The deadline is an integer from 100 to 300000 milliseconds and includes queue time, process startup and execution. Missing `timeout_ms` means 30000. -An optional `idempotency_key` is a string of 1–128 printable ASCII characters (0x21–0x7e); any other value, `null` included, is 400 `invalid_request`. While the job first submitted with a key is retained, submitting the same key again starts nothing and returns 202 with that job's current snapshot, so its status may be `queued`, `running` or terminal; the repeat is answered even when capacity is full, and its `task`, `input` and `timeout_ms` are validated but not compared with the first. Once that job has expired, the key submits a new job. Keys are scoped to one service instance (one token), not to a caller. +An optional `idempotency_key` is a string of 1–128 printable ASCII characters (0x21–0x7e); any other value, `null` included, is 400 `invalid_request`. While the job first submitted with a key is retained, submitting the same key again starts nothing and returns 202 with that job's current snapshot, so its status may be `queued`, `running` or terminal; the repeat is answered even when capacity is full, and its `task`, `input` and `timeout_ms` are validated but not compared with the first. Once that job has expired, the key submits a new job. Keys are scoped to the submitting token within one service instance: the same key from two tokens is two jobs, and a repeat returns only the job its own token submitted. With a single token this is the behaviour it always had. A job contains `id`, `task`, `status`, `timeout_ms`, `result`, `error`, and `progress`. IDs are 32 lowercase hex characters. Result/error/progress are null until applicable. Input documents are not echoed in status responses. @@ -36,7 +48,7 @@ Every 429 carries `Retry-After` in whole delay seconds (an integer of at least 1 The default capacity is 128 total queued, running and retained jobs, concurrency is 2, and terminal retention is 300 seconds. A full result cache rejects submissions until entries expire; it does not silently evict a result early. Configure these with `BRIDGE_CAPACITY`, `BRIDGE_CONCURRENCY`, and `BRIDGE_RETENTION_SECONDS`. The worker keeps as many child processes as its concurrency and replaces a child after `BRIDGE_WORKER_MAX_JOBS` jobs (default 1000, between 1 and 1000000), as well as after a cancellation, a deadline or a crash. A numeric setting that does not parse stops the worker at start with a message naming its variable; one outside its range is refused before any child starts. -With `BRIDGE_STATE_DIR` set, the job table (never a job's input) is written through to SQLite at `/jobs.sqlite` on every status change; unset, it stays in memory and a restart loses every job. When the worker starts on an existing file, a job still `queued` or `running` there becomes `failed` with error `service_restarted` (its input was never stored, so it is not replayed; a shutdown reports such jobs as `cancelled` to whoever polls the stopping process but leaves them unfinished in the file), a finished job whose retention has not run out by the wall clock is served by GET as before until it does, and an expired one is deleted. A 404 for a well-formed ID therefore means the job expired, was never submitted to this instance, or was lost to a restart without a state directory; the worker cannot tell these apart. +With `BRIDGE_STATE_DIR` set, the job table (never a job's input) is written through to SQLite at `/jobs.sqlite` on every status change; unset, it stays in memory and a restart loses every job. When the worker starts on an existing file, a job still `queued` or `running` there becomes `failed` with error `service_restarted` (its input was never stored, so it is not replayed; a shutdown reports such jobs as `cancelled` to whoever polls the stopping process but leaves them unfinished in the file), a finished job whose retention has not run out by the wall clock is served by GET as before until it does, and an expired one is deleted. The file records each job's owner, the token's name and never its value. A file written before owners existed has no `owner` column: the worker adds it when it opens the file, and its rows belong to `default`, the single-token name, so a worker moved to a token list keeps serving them only to a token named `default`. A 404 for a well-formed ID therefore means the job expired, was never submitted to this instance, belongs to another token, or was lost to a restart without a state directory; the worker cannot tell these apart. The PHP client uses a fresh cURL handle per call, refuses redirects, ignores proxy environment variables, caps response bytes, and verifies that GET/cancel responses match the requested ID. Configure its request timeout separately from job execution and local waiting deadlines. There are no transparent submission retries, a 429 included. The MCP server in `mcp_server/` is a second, independent client with the same checks; unlike the PHP client it cancels a job whose local wait expires, because no later request will collect it. @@ -44,13 +56,13 @@ Adding a task requires trusted source changes to `worker/ai_bridge/tasks.py`: an ## Metrics -`GET /metrics` takes the same bearer token as every other route (401 without it; a `POST` is answered as on any unknown route, 404 for `{}`, and another spelling, a query string included, is 404) and answers `Content-Type: text/plain; version=0.0.4; charset=utf-8`, the Prometheus text exposition format, written by the standard library. Every value is an aggregate since the worker process started, held in memory only: a restart resets every counter and histogram to zero, the job table in `BRIDGE_STATE_DIR` does not feed them, and a job read back from that table counts nothing. Each family has one `# HELP` and one `# TYPE` line, and every series exists from the start, at zero. Label values come only from fixed sets: `task` is a task of the allowlist (`rerank`, `embed`, `redact`, plus the test tasks on a worker started with `BRIDGE_TEST_TASKS=1`), `status` one of the four terminal statuses, `reason` is `capacity`, and `backend` the configured backend; no job ID, idempotency key, input, result, error, model path or token appears. +`GET /metrics` takes any of the worker's tokens, as every other route does (401 without one; a `POST` is answered as on any unknown route, 404 for `{}`, and another spelling, a query string included, is 404) and answers `Content-Type: text/plain; version=0.0.4; charset=utf-8`, the Prometheus text exposition format, written by the standard library. Every value is an aggregate since the worker process started, held in memory only: a restart resets every counter and histogram to zero, the job table in `BRIDGE_STATE_DIR` does not feed them, and a job read back from that table counts nothing. Each family has one `# HELP` and one `# TYPE` line, and every series exists from the start, at zero. Label values come only from fixed sets: `owner` is the name of one of the worker's tokens (every configured name has its series from the start, `default` for `BRIDGE_TOKEN`), `task` is a task of the allowlist (`rerank`, `embed`, `redact`, plus the test tasks on a worker started with `BRIDGE_TEST_TASKS=1`), `status` one of the four terminal statuses, `reason` is `capacity`, and `backend` the configured backend; no job ID, idempotency key, input, result, error, model path or token appears. | Family | Type | Labels | What it holds | | --- | --- | --- | --- | -| `bridge_jobs_total` | counter | `task`, `status` | Jobs that reached `succeeded`, `failed`, `timed_out` or `cancelled` | -| `bridge_jobs_submitted_total` | counter | `task` | Accepted submissions that created a job; a repeated `idempotency_key` creates none and is not counted | -| `bridge_submissions_rejected_total` | counter | `reason="capacity"` | Submissions answered 429, for every cause the 429 section names | +| `bridge_jobs_total` | counter | `owner`, `task`, `status` | Jobs that reached `succeeded`, `failed`, `timed_out` or `cancelled` | +| `bridge_jobs_submitted_total` | counter | `owner`, `task` | Accepted submissions that created a job; a repeated `idempotency_key` creates none and is not counted | +| `bridge_submissions_rejected_total` | counter | `owner`, `reason="capacity"` | Submissions answered 429, for every cause the 429 section names | | `bridge_queue_depth` | gauge | | Jobs `queued` now | | `bridge_jobs_running` | gauge | | Jobs `running` now | | `bridge_pool_size` | gauge | | Child processes in the pool, starting or started | diff --git a/docs/security.md b/docs/security.md index 8eec4de..c40fbac 100644 --- a/docs/security.md +++ b/docs/security.md @@ -2,9 +2,9 @@ This experimental service is designed for trusted application-to-service traffic on a private network. Do not expose the example routes or Python server directly to the internet. Across hosts, put a TLS proxy in front of the worker as [Cross-host use behind TLS](integration.md#cross-host-use-behind-tls) shows. -- Generate a service token of at least 32 non-whitespace ASCII characters. Never commit it. A token grants access to every job in that service instance. +- Generate a service token of at least 32 non-whitespace ASCII characters. Never commit it. A token reaches only the jobs it submitted: another token's job answers 404 as if it did not exist, and idempotency keys are per token. Give each application, and the MCP server, its own token in `BRIDGE_TOKENS` or `BRIDGE_TOKENS_FILE` (mode 0600 or stricter), so one can be revoked by removing its entry and restarting the worker without touching the others, and a leaked token exposes only its own client's jobs. Capacity is still shared, so any token can fill the worker; `/healthz` and `/metrics` answer every token, and the metrics carry the token names, never the tokens. The job table stores a job's owner by name only. - Enforce end-user authentication and per-job ownership in the PHP application. The example deliberately omits those application-specific policies and publishes no ports. -- Only trusted task implementations are allowed. The service removes `BRIDGE_TOKEN` from its environment before it starts its child processes, so the token is in neither their environment nor their initial environment block (`/proc//environ`), which a removal inside the child would leave as it was; the children never speak HTTP. Everything else in the service environment stays visible to them. A child serves many jobs in turn and keeps its loaded models; it keeps no reference to a finished job's input or result, but the memory they used is freed, not wiped, so a core dump of a long-lived child can still hold earlier jobs' data where a process that died after each job could not. They are not suitable for arbitrary user-provided scripts. +- Only trusted task implementations are allowed. The service removes `BRIDGE_TOKEN` and `BRIDGE_TOKENS` from its environment before it starts its child processes, and reads `BRIDGE_TOKENS_FILE` without putting its content there, so no token is in either their environment nor their initial environment block (`/proc//environ`), which a removal inside the child would leave as it was; the children never speak HTTP. Everything else in the service environment stays visible to them. A child serves many jobs in turn and keeps its loaded models; it keeps no reference to a finished job's input or result, but the memory they used is freed, not wiped, so a core dump of a long-lived child can still hold earlier jobs' data where a process that died after each job could not. They are not suitable for arbitrary user-provided scripts. - Container configuration uses a non-root user, read-only filesystem, dropped capabilities, no-new-privileges, CPU/memory/PID bounds, and an internal-only network for execution. It mounts only this checkout and the optional dependency/model cache, never a Docker socket or home directory. - Dependency acquisition is separate and internet-enabled. Execution installs hash-checked wheels from the local cache without internet access. The model downloader retrieves fixed files at a pinned revision and checks their hashes; it does not execute model-repository code. - CI runs on GitHub's hosted runners only, because the repository is public. GitHub runs a pull request's workflows as the pull request has them, so a fork's pull request executes its own code, workflow edits included; on a disposable runner of GitHub's that holds nothing of the maintainer's, that is the intended boundary. No workflow names a self-hosted runner, and `tests/python/test_workflows.py` fails if one does again; a fork can edit that test along with the workflows, so while any self-hosted runner is still registered for the repository, keep requiring approval for outside contributors' workflow runs, and unregister it so that there is none a pull request could route its code onto. Every action a workflow uses is pinned by commit hash, and the `actions/checkout` step does not persist credentials; the same test fails if either changes. diff --git a/docs/testing.md b/docs/testing.md index ab70055..d35918d 100644 --- a/docs/testing.md +++ b/docs/testing.md @@ -35,7 +35,7 @@ To select another PHP version, set `BRIDGE_PHP_IMAGE` to its exact pinned image export BRIDGE_PHP_IMAGE='dunglas/frankenphp:1.12.7-php8.2-bookworm@sha256:15646d56650cc5a6c68275ac934eab8f8648a5e776690035f9278a93e0577088' docker compose -f docker/compose.yaml --profile test config docker compose -f docker/compose.yaml run --rm --no-deps php-tests -docker compose -f docker/compose.yaml up -d worker frankenphp fault-worker hostile +docker compose -f docker/compose.yaml up -d worker frankenphp fault-worker hostile tokens-worker docker compose -f docker/compose.yaml run --rm --no-deps php-integration docker compose -f docker/compose.yaml run --rm --no-deps integration ``` @@ -56,11 +56,11 @@ PHP branch support was verified locally on 2026-09-17 on PHP 8.2.33, 8.3.33, 8.4 | Layer | Scope and counting | | --- | --- | -| Python worker | 189 tests: lifecycle, validation and HTTP, the job table in SQLite (a store reopened on its file fails unfinished jobs as `service_restarted`, serves retained results, deletes expired rows from the file and never holds an input; a submission the table refuses is a 429, a later write it refuses leaves the job served, and a key whose row outlived a failed delete runs again) and idempotency keys, process crashes, cancellation, timeouts, retention, the Retry-After a full or closing store gives and its header on every 429, allocation/exit races, the warm pool (that successive jobs share a child in which no live container or frame holds a reference to a finished job's input or result, after a success, a failure and real task payloads, with the detector checked against a planted needle; that a cancelled, timed-out or crashed child is replaced; that a child is recycled after its share of jobs; that reads, submissions and cancellations answer at once while a child is being started, that a child whose start ends after the store closed does not outlive it, and that a job assigned to or waiting for a live child does not fail when no replacement can start; that a job whose child died before taking it is tried once more on another, only once, that a job sent but not yet taken can still be cancelled or time out, and that a job whose child crashed after taking it is not; that a child, idle or running a job, ignores the SIGINT a terminal's Ctrl+C sends the whole process group; that the service token is absent from a child's environment and from its environment block when the service started it; that a message for an earlier job lands nowhere; that a child that never reports ready is replaced and that children dying at start are not restarted in a loop; and that a cancellation does not hold up reads), redaction rules including IPv6, lower-case IBANs and email addresses in any script and their speed on hostile input, NER aggregation, name joining, the second title-cased reading, the uncased model's reading of the whole text and the merge driven by stubs, the fetcher's pins and the model directories an ONNX worker checks for before it starts, batched cross-encoder scoring, the NER measurement's sample and scoring, the workflows' runners, action pins and checkout credentials, both Claude Code hooks (the post-edit rules, and the pre-call push guard through its own cases), progress bounds, the service's reading of its BRIDGE_* settings (defaults, every variable, malformed and out-of-range values, a state directory that cannot hold the table), and the model-name table every result reports from | +| Python worker | 240 tests: lifecycle, validation and HTTP, per-client tokens (each source and every way it is refused at start, a job and its idempotency key scoped to the submitting token, through the job table and a table from before owners, the `owner` metrics series, and no token in an error, the metrics, the table, a child's environment or a start message), the job table in SQLite (a store reopened on its file fails unfinished jobs as `service_restarted`, serves retained results, deletes expired rows from the file and never holds an input; a submission the table refuses is a 429, a later write it refuses leaves the job served, and a key whose row outlived a failed delete runs again) and idempotency keys, process crashes, cancellation, timeouts, retention, the Retry-After a full or closing store gives and its header on every 429, allocation/exit races, the warm pool (that successive jobs share a child in which no live container or frame holds a reference to a finished job's input or result, after a success, a failure and real task payloads, with the detector checked against a planted needle; that a cancelled, timed-out or crashed child is replaced; that a child is recycled after its share of jobs; that reads, submissions and cancellations answer at once while a child is being started, that a child whose start ends after the store closed does not outlive it, and that a job assigned to or waiting for a live child does not fail when no replacement can start; that a job whose child died before taking it is tried once more on another, only once, that a job sent but not yet taken can still be cancelled or time out, and that a job whose child crashed after taking it is not; that a child, idle or running a job, ignores the SIGINT a terminal's Ctrl+C sends the whole process group; that the service token is absent from a child's environment and from its environment block when the service started it; that a message for an earlier job lands nowhere; that a child that never reports ready is replaced and that children dying at start are not restarted in a loop; and that a cancellation does not hold up reads), redaction rules including IPv6, lower-case IBANs and email addresses in any script and their speed on hostile input, NER aggregation, name joining, the second title-cased reading, the uncased model's reading of the whole text and the merge driven by stubs, the fetcher's pins and the model directories an ONNX worker checks for before it starts, batched cross-encoder scoring, the NER measurement's sample and scoring, the workflows' runners, action pins and checkout credentials, both Claude Code hooks (the post-edit rules, and the pre-call push guard through its own cases), progress bounds, the service's reading of its BRIDGE_* settings (defaults, every variable, malformed and out-of-range values, a state directory that cannot hold the table), and the model-name table every result reports from | | MCP server | 125 tests, offline with hash-pinned wheels: the protocol client against an in-process worker, every tool through the SDK's in-memory client with test tasks enabled, the entry point as a subprocess, result validators on synthetic data, the index builder (chunking, the format and each corruption it refuses, building against the worker, caching, its command line) `bridge_search` (the model gate, path confinement, reloads, rerank on and off), the reachability of every answer in the LLM evaluation, that a default search sees only part of its corpus and that its two-step question takes two steps, that the demo worker reaches none of the semantic evaluation's answers, that a tool waits out a full worker for its Retry-After within the call's budget or with its own backoff without one, before each of a search's jobs and within one budget, and says so when it stays full, the evaluation runner's command, isolation checks, scoring and per-backend launcher and index check without starting `claude`, a check that the server's model names match the worker's, that every path parameter says where it resolves and every position in a result says what it counts, and a check that every timing claim in a tool description is a measured median | -| PHP contracts + HTTP | 185 checks per PHP version, a full fault-worker's Retry-After read through the client among them; this **includes** the 148 contract checks (Retry-After from a loopback peer, present, absent, malformed and repeated, and a header block over 16 KiB refused), not 185 + 148; the contract checks include a fake transport's view of each request, the wait deadline across a transport failure, and that no file in `src/` outside `Laravel/` and `Symfony/` names a framework | +| PHP contracts + HTTP | 192 checks per PHP version, a full fault-worker's Retry-After read through the client among them and a two-token `tokens-worker` where another token's read and cancel are 404 and the same idempotency key is two jobs; this **includes** the 148 contract checks (Retry-After from a loopback peer, present, absent, malformed and repeated, and a header block over 16 KiB refused), not 192 + 148; the contract checks include a fake transport's view of each request, the wait deadline across a transport failure, and that no file in `src/` outside `Laravel/` and `Symfony/` names a framework | | PSR-18 transport | The 148 contract checks plus 30 through `Psr18Transport` and in-repo fakes of the PSR-7, PSR-17 and PSR-18 interfaces (`tests/php/psr18_fakes.php`): a round trip with its method, URL, headers and body, a body at the 262144-byte cap and one over it read one byte past the cap and no further, five redirect statuses refused unread, a non-JSON or missing content type, malformed JSON, a 429's Retry-After from one header, a repeated one and none, a client exception and a body read failure as `transport_error` without their text; 178 checks by `php-psr18-tests` with `.cache/php-vendor` on the include path, once, in the MCP server job; `php-tests` skips them with one printed line; no real Symfony HttpClient or Guzzle runs | -| FrankenPHP | 12 jobs submitted by six concurrent test threads through a reused PHP worker, a 512-document `top_k` request, embedding and redaction round trips, and invalid-input, missing-job and multibyte boundary checks | +| FrankenPHP | 12 jobs submitted by six concurrent test threads through a reused PHP worker, a 512-document `top_k` request, embedding and redaction round trips, and invalid-input, missing-job and multibyte boundary checks; then, directly against `tokens-worker`, the same ownership checks as the PHP layer and the `owner` metrics | | PHP syntax | All 25 PHP source/example/test files | | Consumer install | `composer validate --strict` on the library, then a throwaway consumer project in tmpfs, with Packagist disabled and the platform pinned to PHP 8.2.0, installs its `git archive` zip from an `artifact` repository (or from a path repository in a git worktree whose git directory is outside the mount; the archive is built from the committed `HEAD`, so an uncommitted change is not what it tests), checks that `vendor/` holds `src/Client.php` and `docs/` and not `worker/`, `tests/` or `docker/`, that the library brought no other vendor into `vendor/` (so no framework), and autoloads `PhpAiBridge\Client`; runs once, in the MCP server job, with no network | | PHPStan | Level max over `src`, `examples` and `tests/php` with `phpstan.neon`, PHPStan 2.3.1 from `composer.lock` installed into `.cache/php-vendor` by `php-fetcher` and run offline by `php-analysis`, with the frameworks' autoloader from `.cache/php-frameworks-vendor` as its bootstrap file so the provider, the bundle and their tests are analysed against Laravel and Symfony; `phpstan-baseline.neon` holds 9 findings in examples and tests, each with its reason, and none in `src`; runs once, in the MCP server job | diff --git a/tests/integration.py b/tests/integration.py index bac2b66..54f78d1 100644 --- a/tests/integration.py +++ b/tests/integration.py @@ -9,6 +9,7 @@ BASE = os.environ.get("FRANKENPHP_URL", "http://frankenphp:8080") WORKER = os.environ.get("WORKER_URL", "http://worker:8090") +TOKENS_WORKER = os.environ.get("TOKENS_WORKER_URL", "http://tokens-worker:8090") def request(path, body=None): @@ -18,10 +19,10 @@ def request(path, body=None): return response.status, json.load(response), int(response.headers["X-Demo-Worker-Requests"]) -def metrics(token=None): +def metrics(token=None, worker=WORKER): """The worker's counters as {(name, labels): value}; the format itself is checked by the unit tests.""" headers = {"Authorization": "Bearer " + (token if token is not None else os.environ["BRIDGE_TOKEN"])} - with urlopen(Request(WORKER + "/metrics", headers=headers), timeout=5) as response: + with urlopen(Request(worker + "/metrics", headers=headers), timeout=5) as response: assert response.headers["Content-Type"] == "text/plain; version=0.0.4; charset=utf-8" text = response.read().decode("utf-8") samples = {} @@ -32,6 +33,56 @@ def metrics(token=None): return text, samples +def call(token, method, path, body=None): + """Status and decoded body of one request to the two-token worker, an error status included.""" + data = json.dumps(body).encode() if body is not None else None + headers = {"Authorization": "Bearer " + token, "Content-Type": "application/json"} + try: + with urlopen(Request(TOKENS_WORKER + path, data=data, headers=headers, method=method), timeout=5) as response: + return response.status, json.load(response) + except HTTPError as error: + return error.code, json.load(error) + + +def owners(): + """Two tokens on one worker: a job is the submitting token's, and another token reads it as unknown.""" + app, other = os.environ["BRIDGE_APP_TOKEN"], os.environ["BRIDGE_OTHER_TOKEN"] + deadline = time.monotonic() + 30 + while True: + try: + if call(app, "GET", "/healthz")[0] == 200: + break + except (URLError, OSError): + pass + assert time.monotonic() < deadline, "the two-token worker did not become ready" + time.sleep(0.1) + # The PHP integration ran on this worker first, so the counters are compared before and after. + _, counted = metrics(other, TOKENS_WORKER) + rerank = {"task": "rerank", "input": {"query": "needle token", "documents": ["other words", "needle token"]}} + status, job = call(app, "POST", "/v1/jobs", rerank) + assert status == 202, status + unknown = call(other, "GET", "/v1/jobs/" + "a" * 32) + assert unknown[0] == 404, unknown + assert call(other, "GET", "/v1/jobs/" + job["id"]) == unknown + assert call(other, "POST", "/v1/jobs/" + job["id"] + "/cancel", {}) == unknown + assert call(app, "GET", "/v1/jobs/" + job["id"])[1]["id"] == job["id"] + assert call(app, "POST", "/v1/jobs/" + job["id"] + "/cancel", {})[1]["id"] == job["id"] + keyed = dict(rerank, idempotency_key="python-owner-" + job["id"]) + from_app, from_other = call(app, "POST", "/v1/jobs", keyed)[1], call(other, "POST", "/v1/jobs", keyed)[1] + assert from_app["id"] != from_other["id"] + assert call(app, "POST", "/v1/jobs", keyed)[1]["id"] == from_app["id"] + assert call(other, "POST", "/v1/jobs", keyed)[1]["id"] == from_other["id"] + status, body = call("x" * 32, "GET", "/v1/jobs/" + job["id"]) + assert (status, body["error"]["code"]) == (401, "unauthorized"), status + assert call(other, "GET", "/healthz")[0] == 200 + text, moved = metrics(app, TOKENS_WORKER) + for owner, submitted in (("app", 2), ("other", 1)): + name = f'bridge_jobs_submitted_total{{owner="{owner}",task="rerank"}}' + assert moved[name] - counted[name] == submitted, (name, counted[name], moved[name]) + for token in (app, other): + assert token not in text + + def wait_ready(): deadline = time.monotonic() + 30 while time.monotonic() < deadline: @@ -157,12 +208,14 @@ def main(): raise AssertionError("Metrics were served without the token") text, moved = metrics() for task, least in [("rerank", 13), ("embed", 1), ("redact", 1)]: - for name in [f'bridge_jobs_total{{task="{task}",status="succeeded"}}', f'bridge_jobs_submitted_total{{task="{task}"}}', + for name in [f'bridge_jobs_total{{owner="default",task="{task}",status="succeeded"}}', + f'bridge_jobs_submitted_total{{owner="default",task="{task}"}}', f'bridge_job_duration_seconds_count{{task="{task}"}}']: assert moved[name] - counted[name] >= least, (name, counted[name], moved[name]) for document in ["unique0 token0", "needle token", "x@y.io", "4111 1111 1111 1111", "red apple", "café"]: assert document not in text, document - print("PASS: 12 concurrent isolated jobs through a reused FrankenPHP worker; top_k narrowing; embedding; redaction; invalid and missing-job cases; worker metrics") + owners() + print("PASS: 12 concurrent isolated jobs through a reused FrankenPHP worker; top_k narrowing; embedding; redaction; invalid and missing-job cases; worker metrics; jobs scoped to their token") if __name__ == "__main__": diff --git a/tests/php/integration.php b/tests/php/integration.php index 3db39bf..ac4fc44 100644 --- a/tests/php/integration.php +++ b/tests/php/integration.php @@ -108,6 +108,34 @@ function ready(Client $client): void check($repeated->id === $first->id && $repeated->status === 'succeeded' && $repeated->result === $firstDone->result, 'repeated idempotency key returns the same job'); $different = $real->submitRerank('needle token', ['other words', 'needle token'], idempotencyKey: $idempotencyKey . '-2'); check($different->id !== $first->id, 'different idempotency key runs a new job'); +// A worker with two tokens: each client sees only the jobs its own token submitted. +$tokensUrl = getenv('BRIDGE_TOKENS_URL') ?: throw new RuntimeException('Set BRIDGE_TOKENS_URL'); +$app = new Client($tokensUrl, getenv('BRIDGE_APP_TOKEN') ?: throw new RuntimeException('Set BRIDGE_APP_TOKEN')); +$other = new Client($tokensUrl, getenv('BRIDGE_OTHER_TOKEN') ?: throw new RuntimeException('Set BRIDGE_OTHER_TOKEN')); +ready($app); +$owned = $app->submitRerank('needle token', ['other words', 'needle token']); +foreach (['read' => fn () => $other->get($owned->id), 'cancel' => fn () => $other->cancel($owned->id)] as $call => $attempt) { + try { + $attempt(); + throw new RuntimeException("Expected another token's $call to be refused"); + } catch (BridgeException $error) { + check($error->httpStatus === 404, "another token's $call is 404 like an unknown job"); + } +} +check($app->get($owned->id)->id === $owned->id, 'the submitting token reads its job'); +check($app->cancel($owned->id)->id === $owned->id, 'the submitting token cancels its job'); +$sharedKey = 'php-owner-' . bin2hex(random_bytes(8)); +$fromApp = $app->submitRerank('needle token', ['needle token'], idempotencyKey: $sharedKey); +$fromOther = $other->submitRerank('needle token', ['needle token'], idempotencyKey: $sharedKey); +check($fromApp->id !== $fromOther->id, 'the same idempotency key from two tokens is two jobs'); +check($app->submitRerank('needle token', ['needle token'], idempotencyKey: $sharedKey)->id === $fromApp->id, + 'a repeated key returns the job of the token that sent it'); +try { + (new Client($tokensUrl, str_repeat('x', 32)))->get($owned->id); + throw new RuntimeException('Expected authentication failure on the two-token worker'); +} catch (BridgeException $error) { + check($error->httpStatus === 401, 'a token the worker does not hold is 401'); +} // Last, since the fault-worker then stays full for its 300 s retention: its default capacity is 128 jobs, and the // jobs above have finished, so the earliest of them to expire says when a slot returns. $refusal = null; diff --git a/tests/python/test_jobs.py b/tests/python/test_jobs.py index 328dc66..6471af0 100644 --- a/tests/python/test_jobs.py +++ b/tests/python/test_jobs.py @@ -6,6 +6,7 @@ import os import signal import sqlite3 +import sys import tempfile import threading import time @@ -16,6 +17,7 @@ from unittest.mock import Mock, patch from ai_bridge import jobs, tasks +from ai_bridge import server as service from ai_bridge.jobs import CapacityError, JobStore, Settings, TERMINAL from ai_bridge.server import BridgeServer from ai_bridge.tasks import EMBED_DIMENSIONS, MAX_TEXTS, InvalidInput, _entities, _merge, execute, validate @@ -1380,7 +1382,7 @@ def test_the_service_token_is_absent_from_the_child_environment(self): server = BridgeServer(("127.0.0.1", 0), token, Settings(concurrency=1, test_tasks=True)) try: self.assertNotIn("BRIDGE_TOKEN", os.environ) - self.assertEqual(server.token, token) + self.assertEqual([name for name, _ in server.credentials], ["default"]) self.store = server.jobs report = self.inspect() finally: @@ -1395,6 +1397,47 @@ def test_the_service_token_is_absent_from_the_child_environment(self): self.assertIn("BRIDGE_TOKEN", os.environ) self.assertFalse(report["token_in_environment"]) + def test_listed_tokens_are_absent_from_the_child_environment(self): + # main() takes BRIDGE_TOKENS out of the environment, and reads BRIDGE_TOKENS_FILE without putting what it + # holds there, before the store starts a child; only the file's path stays. + self.store.close() + listed = "test-only-listed-token-never-use-in-production" + in_file = "test-only-file-token-never-use-in-production" + with tempfile.TemporaryDirectory() as root: + path = os.path.join(root, "tokens") + with open(path, "w", encoding="ascii") as file: + file.write(f"default={in_file}\n") + os.chmod(path, 0o600) + for environment, value in (({"BRIDGE_TOKENS": f"other=x{in_file},default={listed}"}, listed), + ({"BRIDGE_TOKENS_FILE": path}, in_file)): + with self.subTest(source=next(iter(environment))): + reports = [] + + def serve(bridge, **_): + self.store = bridge.jobs + reports.append(self.inspect(value)) + + with patch.dict(os.environ, {"BRIDGE_TEST_TASKS": "1", "BRIDGE_CONCURRENCY": "1", **environment}), \ + patch.object(sys, "argv", ["ai_bridge.server", "--port", "0"]), \ + patch.object(service.signal, "signal"), \ + patch.object(service.BridgeServer, "serve_forever", autospec=True, side_effect=serve): + os.environ.pop("BRIDGE_TOKEN", None) + service.main() + report, = reports + self.assertFalse(report["token_in_environment"]) + self.assertFalse(report["needle_in_environment"]) + self.assertEqual(report["held"], 0) + if os.path.exists("/proc/self/environ"): + self.assertIs(report["token_in_environment_block"], False) + self.assertIs(report["needle_in_environment_block"], False) + # The same checks see a value that is there. + with patch.dict(os.environ, {"BRIDGE_PROBE": listed}): + self.store = JobStore(Settings(concurrency=1, test_tasks=True)) + report = self.inspect(listed) + self.assertTrue(report["needle_in_environment"]) + if os.path.exists("/proc/self/environ"): + self.assertIs(report["needle_in_environment_block"], True) + @unittest.skipUnless(os.path.exists("/proc/self/environ"), "needs /proc to read a child's environment block") def test_a_child_inherits_one_thread_pins(self): # The pools size themselves when their library loads, so only the block the child starts with counts; @@ -1616,8 +1659,141 @@ def test_a_running_job_is_not_held_up_by_another_being_cancelled(self): self.assertEqual(self.wait(running["id"])["result"]["value"], "kept") +class OwnerTests(unittest.TestCase): + """A job belongs to the token name that submitted it: another owner reads it as unknown, keys are per owner, + and the owner outlives a restart in the job table.""" + + OLD_SCHEMA = """CREATE TABLE jobs ( + id TEXT PRIMARY KEY, task TEXT NOT NULL, status TEXT NOT NULL, result TEXT, error TEXT, progress TEXT, + timeout_ms INTEGER NOT NULL, idempotency_key TEXT, submitted_at REAL NOT NULL, finished_at REAL); + CREATE UNIQUE INDEX jobs_idempotency_key ON jobs (idempotency_key);""" + + def settings(self, **values): + return Settings(concurrency=1, test_tasks=True, owners=("app", "other"), **values) + + def wait(self, store, job_id, owner, timeout=5): + deadline = time.monotonic() + timeout + while store.get(job_id, owner)["status"] not in TERMINAL: + self.assertLess(time.monotonic(), deadline) + time.sleep(0.01) + return store.get(job_id, owner) + + def test_another_owner_reads_and_cancels_nothing(self): + store = JobStore(self.settings()) + try: + job = store.submit("test.delay", {"seconds": 5}, owner="app") + self.assertIsNone(store.get(job["id"], "other")) + self.assertIsNone(store.cancel(job["id"], "other")) + self.assertIsNone(store.get(job["id"], "default")) + self.assertNotEqual(store.get(job["id"], "app")["status"], "cancelled") + self.assertEqual(store.cancel(job["id"], "app")["status"], "cancelled") + self.assertNotIn("owner", job) + finally: + store.close() + + def test_keys_are_per_owner(self): + store = JobStore(self.settings()) + try: + first = store.submit("test.delay", {"seconds": 0}, idempotency_key="k", owner="app") + other = store.submit("test.delay", {"seconds": 0}, idempotency_key="k", owner="other") + self.assertNotEqual(first["id"], other["id"]) + self.assertEqual(store.submit("test.delay", {"seconds": 0}, idempotency_key="k", owner="app")["id"], + first["id"]) + finally: + store.close() + + def test_an_unknown_owner_or_owner_list_is_refused(self): + store = JobStore(self.settings()) + try: + with self.assertRaises(ValueError): + store.submit("test.delay", {"seconds": 0}, owner="default") + finally: + store.close() + for owners in ((), ["app"], ("app", "app"), ("app", 1)): + with self.subTest(owners=owners), self.assertRaises(ValueError): + JobStore(Settings(concurrency=1, owners=owners)) + + def test_a_reopened_store_keeps_each_owner(self): + with tempfile.TemporaryDirectory() as state_dir: + first = JobStore(self.settings(state_dir=state_dir)) + try: + job = first.submit("test.delay", {"seconds": 0}, idempotency_key="kept", owner="app") + self.wait(first, job["id"], "app") + finally: + first.close() + reopened = JobStore(self.settings(state_dir=state_dir)) + try: + self.assertIsNone(reopened.get(job["id"], "other")) + self.assertIsNone(reopened.cancel(job["id"], "other")) + self.assertEqual(reopened.get(job["id"], "app")["status"], "succeeded") + self.assertEqual( + reopened.submit("test.delay", {"seconds": 0}, idempotency_key="kept", owner="app")["id"], job["id"]) + fresh = reopened.submit("test.delay", {"seconds": 0}, idempotency_key="kept", owner="other") + self.assertNotEqual(fresh["id"], job["id"]) + finally: + reopened.close() + with contextlib.closing(sqlite3.connect(os.path.join(state_dir, "jobs.sqlite"))) as db: + self.assertEqual(dict(db.execute("SELECT id, owner FROM jobs")), {job["id"]: "app", fresh["id"]: "other"}) + + def test_a_table_from_before_owners_gets_the_column_and_its_rows_belong_to_default(self): + with tempfile.TemporaryDirectory() as state_dir: + path = os.path.join(state_dir, "jobs.sqlite") + with contextlib.closing(sqlite3.connect(path)) as db: + db.executescript(self.OLD_SCHEMA) + with db: + db.execute("INSERT INTO jobs VALUES ('" + "b" * 32 + "', 'test.delay', 'succeeded', " + "'{\"value\": \"old\"}', 'null', 'null', 30000, 'old-key', ?, ?)", + (time.time(), time.time())) + store = JobStore(Settings(concurrency=1, test_tasks=True, owners=("default", "app"), state_dir=state_dir)) + try: + self.assertEqual(store.get("b" * 32)["result"], {"value": "old"}) + self.assertIsNone(store.get("b" * 32, "app")) + self.assertEqual(store.submit("test.delay", {"seconds": 0}, idempotency_key="old-key")["id"], "b" * 32) + # The old index on the key alone is gone, so another owner's same key is a job of its own. + other = store.submit("test.delay", {"seconds": 0}, idempotency_key="old-key", owner="app") + self.assertNotEqual(other["id"], "b" * 32) + finally: + store.close() + with contextlib.closing(sqlite3.connect(path)) as db: + self.assertIn("owner", [row[1] for row in db.execute("PRAGMA table_info(jobs)")]) + self.assertEqual(db.execute("SELECT owner FROM jobs WHERE id = ?", ("b" * 32,)).fetchone(), ("default",)) + indexes = [row[1] for row in db.execute("PRAGMA index_list(jobs)") if not row[1].startswith("sqlite_")] + self.assertEqual(indexes, ["jobs_owner_idempotency_key"]) + # Opening the migrated file again changes nothing. + store = JobStore(Settings(concurrency=1, test_tasks=True, owners=("default", "app"), state_dir=state_dir)) + try: + self.assertEqual(store.get("b" * 32)["result"], {"value": "old"}) + self.assertIsNone(store.get("b" * 32, "app")) + finally: + store.close() + with contextlib.closing(sqlite3.connect(path)) as db: + self.assertEqual(db.execute("SELECT owner FROM jobs WHERE id = ?", ("b" * 32,)).fetchone(), ("default",)) + indexes = [row[1] for row in db.execute("PRAGMA index_list(jobs)") if not row[1].startswith("sqlite_")] + self.assertEqual(indexes, ["jobs_owner_idempotency_key"]) + + def test_every_owner_has_its_series_from_the_start(self): + store = JobStore(self.settings()) + try: + snapshot = store.snapshot() + self.assertEqual(snapshot["owners"], ("app", "other")) + self.assertEqual(set(snapshot["rejected"]), {("app", "capacity"), ("other", "capacity")}) + self.assertEqual(set(snapshot["submitted"]), {(owner, task) for owner in ("app", "other") + for task in snapshot["tasks"]}) + self.assertEqual(len(snapshot["jobs"]), 2 * len(snapshot["tasks"]) * 4) + self.assertEqual(sum(snapshot["jobs"].values()) + sum(snapshot["submitted"].values()), 0) + job = store.submit("test.delay", {"seconds": 0}, owner="other") + self.wait(store, job["id"], "other") + snapshot = store.snapshot() + self.assertEqual({key: count for key, count in snapshot["jobs"].items() if count}, + {("other", "test.delay", "succeeded"): 1}) + self.assertEqual({key: count for key, count in snapshot["submitted"].items() if count}, + {("other", "test.delay"): 1}) + finally: + store.close() + + class MetricsSnapshotTests(unittest.TestCase): - """What the store counts for GET /metrics: plain numbers keyed only by task, status and reason.""" + """What the store counts for GET /metrics: plain numbers keyed only by owner, task, status and reason.""" def setUp(self): self.store = JobStore(Settings(concurrency=1, test_tasks=True)) @@ -1647,10 +1823,10 @@ def test_a_fresh_store_reports_every_series_at_zero(self): snapshot = self.wait_pool(1, 1) tasks = ("rerank", "embed", "redact", "test.delay", "test.crash", "test.error", "test.inspect") self.assertEqual(snapshot["tasks"], tasks) - self.assertEqual(snapshot["jobs"], {(task, status): 0 for task in tasks + self.assertEqual(snapshot["jobs"], {("default", task, status): 0 for task in tasks for status in ("succeeded", "failed", "timed_out", "cancelled")}) - self.assertEqual(snapshot["submitted"], dict.fromkeys(tasks, 0)) - self.assertEqual(snapshot["rejected"], {"capacity": 0}) + self.assertEqual(snapshot["submitted"], {("default", task): 0 for task in tasks}) + self.assertEqual(snapshot["rejected"], {("default", "capacity"): 0}) self.assertEqual((snapshot["queued"], snapshot["running"]), (0, 0)) self.assertEqual((snapshot["backend"], snapshot["concurrency"], snapshot["capacity"]), ("lexical", 1, 128)) for histogram in list(snapshot["duration"].values()) + list(snapshot["queue"].values()): @@ -1678,10 +1854,11 @@ def test_each_terminal_status_is_counted_once(self): ["succeeded", "failed", "timed_out", "failed"]) snapshot = self.store.snapshot() self.assertEqual({key: count for key, count in snapshot["jobs"].items() if count}, { - ("test.delay", "succeeded"): 1, ("test.error", "failed"): 1, ("test.delay", "cancelled"): 1, - ("test.delay", "timed_out"): 1, ("test.crash", "failed"): 1}) + ("default", "test.delay", "succeeded"): 1, ("default", "test.error", "failed"): 1, + ("default", "test.delay", "cancelled"): 1, ("default", "test.delay", "timed_out"): 1, + ("default", "test.crash", "failed"): 1}) self.assertEqual({task: count for task, count in snapshot["submitted"].items() if count}, - {"test.delay": 3, "test.error": 1, "test.crash": 1}) + {("default", "test.delay"): 3, ("default", "test.error"): 1, ("default", "test.crash"): 1}) def test_histograms_agree_with_the_jobs_run(self): started = time.monotonic() @@ -1742,26 +1919,26 @@ def test_a_capacity_refusal_is_counted_and_a_repeat_or_invalid_input_is_not(self with self.assertRaises(ValueError): self.store.submit("rerank", {"query": "", "documents": ["x"]}) snapshot = self.store.snapshot() - self.assertEqual(snapshot["rejected"], {"capacity": 2}) - self.assertEqual(snapshot["submitted"]["test.delay"], 1) - self.assertEqual(snapshot["submitted"]["rerank"], 0) + self.assertEqual(snapshot["rejected"], {("default", "capacity"): 2}) + self.assertEqual(snapshot["submitted"]["default", "test.delay"], 1) + self.assertEqual(snapshot["submitted"]["default", "rerank"], 0) def test_a_submission_the_table_refuses_counts_as_a_refusal(self): FailingDatabase(self.store, "INSERT") with self.assertRaises(CapacityError): self.store.submit("test.delay", {"seconds": 0}) snapshot = self.store.snapshot() - self.assertEqual(snapshot["rejected"], {"capacity": 1}) - self.assertEqual(snapshot["submitted"]["test.delay"], 0) + self.assertEqual(snapshot["rejected"], {("default", "capacity"): 1}) + self.assertEqual(snapshot["submitted"]["default", "test.delay"], 0) def test_the_snapshot_is_a_copy(self): self.wait(self.store.submit("test.delay", {"seconds": 0})["id"]) snapshot = self.store.snapshot() snapshot["duration"]["test.delay"]["buckets"][0] = 99 - snapshot["jobs"][("test.delay", "succeeded")] = 99 + snapshot["jobs"][("default", "test.delay", "succeeded")] = 99 again = self.store.snapshot() self.assertNotEqual(again["duration"]["test.delay"]["buckets"][0], 99) - self.assertEqual(again["jobs"][("test.delay", "succeeded")], 1) + self.assertEqual(again["jobs"][("default", "test.delay", "succeeded")], 1) def test_jobs_restored_from_the_table_count_nothing(self): with tempfile.TemporaryDirectory() as state_dir: diff --git a/tests/python/test_server.py b/tests/python/test_server.py index dbda6d5..87282ad 100644 --- a/tests/python/test_server.py +++ b/tests/python/test_server.py @@ -109,7 +109,7 @@ def test_a_submission_the_job_table_refuses_is_429(self): status, data = self.request("POST", "/v1/jobs", body, self.auth()) self.assertEqual(status, 429) self.assertEqual(data["error"]["code"], "capacity_exceeded") - self.assertNotIn("http-k3", self.server.jobs._keys) + self.assertNotIn(("default", "http-k3"), self.server.jobs._keys) def test_body_limit(self): self.assertEqual(self.request("POST", "/v1/jobs", " " * 262145, self.auth())[0], 413) @@ -177,7 +177,7 @@ def rejected(): try: connection.request("GET", "/metrics", headers={"Authorization": "Bearer " + TOKEN}) return sample(parse_metrics(connection.getresponse().read().decode()), - "bridge_submissions_rejected_total", reason="capacity") + "bridge_submissions_rejected_total", owner="default", reason="capacity") finally: connection.close() @@ -541,7 +541,7 @@ def test_the_body_parses_and_names_every_family(self): self.assertEqual(sample(families, "bridge_info", protocol="1", backend="lexical"), 1) self.assertEqual(sample(families, "bridge_concurrency"), 1) self.assertEqual(sample(families, "bridge_capacity"), 16) - self.assertIn(sample(families, "bridge_submissions_rejected_total", reason="capacity"), (0, 1, 2)) + self.assertIn(sample(families, "bridge_submissions_rejected_total", owner="default", reason="capacity"), (0, 1, 2)) bounds = sorted({dict(labels)["le"] for name, labels in families["bridge_job_duration_seconds"][1] if name.endswith("_bucket")}, key=float) self.assertEqual(bounds, ["0.05", "0.1", "0.25", "0.5", "1", "2.5", "5", "10", "30", "60", "300", "+Inf"]) @@ -575,12 +575,12 @@ def test_counters_and_histograms_move_with_the_jobs_and_nothing_of_a_job_is_expo def moved(name, **labels): return sample(after, name, **labels) - sample(before, name, **labels) - self.assertEqual(moved("bridge_jobs_total", task="test.delay", status="succeeded"), 1) - self.assertEqual(moved("bridge_jobs_total", task="test.error", status="failed"), 1) - self.assertEqual(moved("bridge_jobs_total", task="test.delay", status="cancelled"), 1) - self.assertEqual(moved("bridge_jobs_total", task="test.delay", status="timed_out"), 1) - self.assertEqual(moved("bridge_jobs_submitted_total", task="test.delay"), 3) - self.assertEqual(moved("bridge_jobs_submitted_total", task="test.error"), 1) + self.assertEqual(moved("bridge_jobs_total", owner="default", task="test.delay", status="succeeded"), 1) + self.assertEqual(moved("bridge_jobs_total", owner="default", task="test.error", status="failed"), 1) + self.assertEqual(moved("bridge_jobs_total", owner="default", task="test.delay", status="cancelled"), 1) + self.assertEqual(moved("bridge_jobs_total", owner="default", task="test.delay", status="timed_out"), 1) + self.assertEqual(moved("bridge_jobs_submitted_total", owner="default", task="test.delay"), 3) + self.assertEqual(moved("bridge_jobs_submitted_total", owner="default", task="test.error"), 1) self.assertEqual(moved("bridge_job_duration_seconds_count", task="test.delay"), 3) self.assertEqual(moved("bridge_job_duration_seconds_count", task="test.error"), 1) self.assertGreaterEqual(moved("bridge_job_duration_seconds_sum", task="test.delay"), 0.15) @@ -623,3 +623,300 @@ def probe(): with patch.object(server, "render", checking): self.scrape() self.assertEqual(held, [False]) + + + +APP = "test-only-app-token-never-use-in-production-1" +OTHER = "test-only-other-token-never-use-in-production-2" + + +class TokenSourceTests(unittest.TestCase): + """main() resolves the tokens from BRIDGE_TOKEN, BRIDGE_TOKENS or BRIDGE_TOKENS_FILE, and refuses to start on a bad + one with a message naming the variable and never a token.""" + + def main(self, environment): + """Run main with exactly these BRIDGE_* variables; returns the tokens BridgeServer was given and the token + variables still in the environment when it was built.""" + seen = {} + + def built(address, tokens, settings): + seen["tokens"] = tokens + seen["left"] = {name for name in ("BRIDGE_TOKEN", "BRIDGE_TOKENS") if name in os.environ} + return MagicMock() + + with patch.dict(os.environ, environment), patch.object(sys, "argv", ["ai_bridge.server", "--port", "0"]), \ + patch.object(server.signal, "signal"), patch.object(server, "BridgeServer", side_effect=built): + for name in [name for name in os.environ if name.startswith("BRIDGE_") and name not in environment]: + del os.environ[name] + server.main() + return seen + + def refused(self, environment, variable): + with self.assertRaises(SystemExit) as raised: + self.main(environment) + message = str(raised.exception.code) + self.assertIn(variable, message) + for value in (APP, OTHER, TOKEN, APP[:20], "short-token"): + self.assertNotIn(value, message) + return message + + def tokens_file(self, root, content, mode=0o600): + path = os.path.join(root, "tokens") + with open(path, "w", encoding="utf-8") as file: + file.write(content) + os.chmod(path, mode) + return path + + def test_the_single_token_is_named_default_and_leaves_the_environment(self): + seen = self.main({"BRIDGE_TOKEN": TOKEN}) + self.assertEqual(seen, {"tokens": {"default": TOKEN}, "left": set()}) + + def test_the_listed_tokens_are_read_and_leave_the_environment(self): + seen = self.main({"BRIDGE_TOKENS": f"app={APP},other={OTHER}"}) + self.assertEqual(seen, {"tokens": {"app": APP, "other": OTHER}, "left": set()}) + # A token may hold "=": only the first one separates the name. + self.assertEqual(self.main({"BRIDGE_TOKENS": f"a.b-c_9={APP}=x"})["tokens"], {"a.b-c_9": APP + "=x"}) + + def test_sixty_four_tokens_are_accepted(self): + listed = ",".join(f"n{index}={APP}{index}" for index in range(64)) + self.assertEqual(len(self.main({"BRIDGE_TOKENS": listed})["tokens"]), 64) + + def test_the_token_file_is_read_with_comments_and_blank_lines(self): + with tempfile.TemporaryDirectory() as root: + for mode in (0o600, 0o400): + with self.subTest(mode=oct(mode)): + path = self.tokens_file(root, f"# the applications\n\napp={APP}\n # indented\nother={OTHER}\n", mode) + seen = self.main({"BRIDGE_TOKENS_FILE": path}) + self.assertEqual(seen["tokens"], {"app": APP, "other": OTHER}) + os.remove(path) + + def test_a_bad_token_is_refused_for_each_source(self): + spaced = APP[:20] + " " + APP[20:] + for value in ("short-token", spaced, APP + "\x7f", APP + "é", ""): + with self.subTest(value=repr(value)[:24]): + self.refused({"BRIDGE_TOKENS": f"app={value}"}, "BRIDGE_TOKENS") + self.refused({"BRIDGE_TOKEN": value or " "}, "BRIDGE_TOKEN") + with tempfile.TemporaryDirectory() as root: + self.refused({"BRIDGE_TOKENS_FILE": self.tokens_file(root, f"app={value}\n")}, "BRIDGE_TOKENS_FILE") + with tempfile.TemporaryDirectory() as root: + # A trailing space in the file is part of the line, so the token is refused, not trimmed. + self.refused({"BRIDGE_TOKENS_FILE": self.tokens_file(root, f"app={APP} \n")}, "BRIDGE_TOKENS_FILE") + + def test_a_bad_name_or_entry_is_refused(self): + for listed in (f"={APP}", f"a b={APP}", "x" * 65 + f"={APP}", APP, f"café={APP}", f"app={APP},", + f"app/1={APP}", f",app={APP}"): + with self.subTest(listed=listed[:24]): + self.refused({"BRIDGE_TOKENS": listed}, "BRIDGE_TOKENS") + self.assertEqual(self.main({"BRIDGE_TOKENS": "x" * 64 + f"={APP}"})["tokens"], {"x" * 64: APP}) + + def test_a_duplicate_name_or_value_is_refused(self): + for listed in (f"app={APP},app={OTHER}", f"app={APP},other={APP}"): + with self.subTest(listed=listed[:12]): + self.refused({"BRIDGE_TOKENS": listed}, "BRIDGE_TOKENS") + with tempfile.TemporaryDirectory() as root: + path = self.tokens_file(root, listed.replace(",", "\n")) + self.refused({"BRIDGE_TOKENS_FILE": path}, "BRIDGE_TOKENS_FILE") + + def test_more_than_sixty_four_tokens_are_refused(self): + listed = ",".join(f"n{index}={APP}{index}" for index in range(65)) + self.assertIn("64", self.refused({"BRIDGE_TOKENS": listed}, "BRIDGE_TOKENS")) + + def test_more_than_sixty_four_tokens_in_the_file_are_refused(self): + with tempfile.TemporaryDirectory() as root: + path = self.tokens_file(root, "".join(f"n{index}={APP}{index}\n" for index in range(65))) + self.assertIn("64", self.refused({"BRIDGE_TOKENS_FILE": path}, "BRIDGE_TOKENS_FILE")) + + def test_a_token_file_over_the_size_limit_is_refused(self): + with tempfile.TemporaryDirectory() as root: + content = f"app={APP}\n#" + "x" * (server.MAX_TOKENS_FILE - len(APP) - 6) + "\n" + self.assertEqual(len(content.encode()), server.MAX_TOKENS_FILE + 1) + message = self.refused({"BRIDGE_TOKENS_FILE": self.tokens_file(root, content)}, "BRIDGE_TOKENS_FILE") + self.assertIn("65536", message) + + def test_a_fifo_at_the_token_file_path_is_refused_at_once(self): + if not hasattr(os, "mkfifo"): + self.skipTest("this platform has no FIFOs") + with tempfile.TemporaryDirectory() as root: + path = os.path.join(root, "tokens") + os.mkfifo(path, 0o600) + started = time.monotonic() + self.assertIn("regular file", self.refused({"BRIDGE_TOKENS_FILE": path}, "BRIDGE_TOKENS_FILE")) + self.assertLess(time.monotonic() - started, 5) + + def test_an_empty_variable_counts_as_unset(self): + self.assertEqual(self.main({"BRIDGE_TOKEN": TOKEN, "BRIDGE_TOKENS": ""})["tokens"], {"default": TOKEN}) + with tempfile.TemporaryDirectory() as root: + path = self.tokens_file(root, f"app={APP}\n") + seen = self.main({"BRIDGE_TOKENS": "", "BRIDGE_TOKENS_FILE": path}) + self.assertEqual(seen["tokens"], {"app": APP}) + + def test_an_empty_list_or_file_is_refused(self): + with tempfile.TemporaryDirectory() as root: + self.refused({"BRIDGE_TOKENS_FILE": self.tokens_file(root, "# nothing\n\n")}, "BRIDGE_TOKENS_FILE") + self.refused({}, "BRIDGE_TOKEN") + + def test_a_token_file_open_to_group_or_others_is_refused(self): + with tempfile.TemporaryDirectory() as root: + for mode in (0o640, 0o604, 0o660, 0o644, 0o602): + with self.subTest(mode=oct(mode)): + path = self.tokens_file(root, f"app={APP}\n", mode) + self.assertIn("0600", self.refused({"BRIDGE_TOKENS_FILE": path}, "BRIDGE_TOKENS_FILE")) + os.remove(path) + + def test_a_token_file_that_cannot_be_read_is_refused(self): + with tempfile.TemporaryDirectory() as root: + paths = [os.path.join(root, "missing"), root] + # Root reads through the mode bits; compose runs these tests as uid 65532. + if os.geteuid() != 0: + paths.append(self.tokens_file(root, f"app={APP}\n", 0o000)) + for path in paths: + with self.subTest(path=path): + self.refused({"BRIDGE_TOKENS_FILE": path}, "BRIDGE_TOKENS_FILE") + binary = os.path.join(root, "binary") + with open(binary, "wb") as file: + file.write(b"app=\xff" + APP.encode()) + os.chmod(binary, 0o600) + self.refused({"BRIDGE_TOKENS_FILE": binary}, "BRIDGE_TOKENS_FILE") + + def test_the_sources_are_exclusive(self): + with tempfile.TemporaryDirectory() as root: + path = self.tokens_file(root, f"app={APP}\n") + for environment, names in ( + ({"BRIDGE_TOKENS": f"app={APP}", "BRIDGE_TOKENS_FILE": path}, ("BRIDGE_TOKENS", "BRIDGE_TOKENS_FILE")), + ({"BRIDGE_TOKEN": TOKEN, "BRIDGE_TOKENS": f"app={APP}"}, ("BRIDGE_TOKEN", "BRIDGE_TOKENS")), + ({"BRIDGE_TOKEN": TOKEN, "BRIDGE_TOKENS_FILE": path}, ("BRIDGE_TOKEN", "BRIDGE_TOKENS_FILE"))): + with self.subTest(names=names): + message = self.refused(environment, names[0]) + self.assertIn(names[1], message) + + def test_the_server_validates_the_tokens_it_is_given(self): + for tokens in ({}, {"app": "short"}, {"a b": APP}, {"app": APP, "other": APP}, + {f"n{index}": APP + str(index) for index in range(65)}): + with self.subTest(tokens=list(tokens)[:2]), self.assertRaises(ValueError) as raised: + BridgeServer(("127.0.0.1", 0), tokens) + self.assertNotIn(APP, str(raised.exception)) + + +class OwnerHttpTests(unittest.TestCase): + """Two tokens on one worker: a job belongs to the token that submitted it.""" + + def setUp(self): + self.state = tempfile.TemporaryDirectory() + self.start(Settings(concurrency=1, capacity=8, test_tasks=True, state_dir=self.state.name)) + + def start(self, settings): + self.server = BridgeServer(("127.0.0.1", 0), {"app": APP, "other": OTHER}, settings) + self.thread = threading.Thread(target=self.server.serve_forever, daemon=True) + self.thread.start() + + def stop(self): + self.server.shutdown() + self.thread.join() + self.server.server_close() + + def tearDown(self): + self.stop() + self.state.cleanup() + + def call(self, method, path, token, body=None): + """Status and the decoded body; no answer to any token may hold a token.""" + connection = http.client.HTTPConnection(*self.server.server_address, timeout=2) + try: + connection.request(method, path, None if body is None else json.dumps(body), + {"Authorization": "Bearer " + token, "Content-Type": "application/json"}) + response = connection.getresponse() + status, data = response.status, response.read() + finally: + connection.close() + for value in (APP, OTHER): + self.assertNotIn(value.encode(), data) + return status, json.loads(data) if data.startswith(b"{") else data.decode() + + def submit(self, token, key=None, seconds=0): + body = {"task": "test.delay", "input": {"seconds": seconds}} + if key is not None: + body["idempotency_key"] = key + return self.call("POST", "/v1/jobs", token, body) + + def finished(self, job_id, owner): + deadline = time.monotonic() + 5 + while self.server.jobs.get(job_id, owner)["status"] != "succeeded": + self.assertLess(time.monotonic(), deadline, "the job never finished") + time.sleep(0.01) + + def test_another_token_cannot_read_or_cancel_a_job(self): + status, job = self.submit(APP, seconds=5) + self.assertEqual(status, 202) + unknown = self.call("GET", "/v1/jobs/" + "a" * 32, OTHER) + self.assertEqual(unknown[0], 404) + # Exactly the answer for an unknown job, so the other token learns nothing of it. + self.assertEqual(self.call("GET", "/v1/jobs/" + job["id"], OTHER), unknown) + self.assertEqual(self.call("POST", "/v1/jobs/" + job["id"] + "/cancel", OTHER, {}), unknown) + status, seen = self.call("GET", "/v1/jobs/" + job["id"], APP) + self.assertEqual(status, 200) + self.assertNotEqual(seen["status"], "cancelled") + self.assertEqual(set(seen), {"id", "task", "status", "result", "error", "progress", "timeout_ms"}) + status, cancelled = self.call("POST", "/v1/jobs/" + job["id"] + "/cancel", APP, {}) + self.assertEqual((status, cancelled["status"]), (200, "cancelled")) + + def test_an_idempotency_key_belongs_to_its_token(self): + first_app = self.submit(APP, "shared-key")[1] + first_other = self.submit(OTHER, "shared-key")[1] + self.assertNotEqual(first_app["id"], first_other["id"]) + self.assertEqual(self.submit(APP, "shared-key")[1]["id"], first_app["id"]) + self.assertEqual(self.submit(OTHER, "shared-key")[1]["id"], first_other["id"]) + + def test_health_and_metrics_answer_every_token(self): + for token in (APP, OTHER): + with self.subTest(token=token[:14]): + self.assertEqual(self.call("GET", "/healthz", token)[0], 200) + self.assertEqual(self.call("GET", "/metrics", token)[0], 200) + for token in ("wrong", APP + "x", APP[:-1], APP + "," + OTHER): + with self.subTest(token=token[:14]): + status, body = self.call("GET", "/healthz", token) + self.assertEqual((status, body["error"]["code"]), (401, "unauthorized")) + + def test_the_owner_series_exist_at_zero_and_move_for_the_right_owner(self): + self.stop() + self.start(Settings(concurrency=1, capacity=1, test_tasks=True)) + families = parse_metrics(self.call("GET", "/metrics", APP)[1]) + for owner in ("app", "other"): + self.assertEqual(sample(families, "bridge_jobs_total", owner=owner, task="test.delay", status="succeeded"), 0) + self.assertEqual(sample(families, "bridge_jobs_submitted_total", owner=owner, task="test.delay"), 0) + self.assertEqual(sample(families, "bridge_submissions_rejected_total", owner=owner, reason="capacity"), 0) + owners = {dict(labels)["owner"] for family in ("bridge_jobs_total", "bridge_jobs_submitted_total", + "bridge_submissions_rejected_total") + for _, labels in families[family][1]} + self.assertEqual(owners, {"app", "other"}) + status, job = self.submit(APP) + self.assertEqual(status, 202) + self.finished(job["id"], "app") + self.assertEqual(self.submit(OTHER)[0], 429) + families = parse_metrics(self.call("GET", "/metrics", OTHER)[1]) + self.assertEqual(sample(families, "bridge_jobs_total", owner="app", task="test.delay", status="succeeded"), 1) + self.assertEqual(sample(families, "bridge_jobs_total", owner="other", task="test.delay", status="succeeded"), 0) + self.assertEqual(sample(families, "bridge_jobs_submitted_total", owner="app", task="test.delay"), 1) + self.assertEqual(sample(families, "bridge_jobs_submitted_total", owner="other", task="test.delay"), 0) + self.assertEqual(sample(families, "bridge_submissions_rejected_total", owner="other", reason="capacity"), 1) + self.assertEqual(sample(families, "bridge_submissions_rejected_total", owner="app", reason="capacity"), 0) + # The families without an owner label are unchanged. + self.assertEqual(sample(families, "bridge_job_duration_seconds_count", task="test.delay"), 1) + + def test_no_token_value_reaches_the_job_table_and_a_restored_job_keeps_its_owner(self): + status, job = self.submit(APP, "table-key") + self.assertEqual(status, 202) + self.submit(OTHER, "table-key") + self.finished(job["id"], "app") + self.stop() + self.start(Settings(concurrency=1, capacity=8, test_tasks=True, state_dir=self.state.name)) + written = b"" + for name in os.listdir(self.state.name): + if name.startswith("jobs.sqlite"): + with open(os.path.join(self.state.name, name), "rb") as file: + written += file.read() + self.assertIn(b"app", written) + for value in (APP, OTHER): + self.assertNotIn(value.encode(), written) + self.assertEqual(self.call("GET", "/v1/jobs/" + job["id"], OTHER)[0], 404) + self.assertEqual(self.call("GET", "/v1/jobs/" + job["id"], APP)[1]["status"], "succeeded") + self.assertEqual(self.submit(APP, "table-key")[1]["id"], job["id"]) diff --git a/worker/ai_bridge/jobs.py b/worker/ai_bridge/jobs.py index 7195c90..4cff25b 100644 --- a/worker/ai_bridge/jobs.py +++ b/worker/ai_bridge/jobs.py @@ -34,10 +34,13 @@ BACKOFF_SECONDS = 1.0 MAX_JOBS_PER_WORKER = 1000 IDEMPOTENCY_KEY = re.compile(r"[\x21-\x7e]{1,128}") +# The owner is a token's name, never its value; _open_db gives a table from before owners the column, rows to "default". _SCHEMA = """CREATE TABLE IF NOT EXISTS jobs ( id TEXT PRIMARY KEY, task TEXT NOT NULL, status TEXT NOT NULL, result TEXT, error TEXT, progress TEXT, - timeout_ms INTEGER NOT NULL, idempotency_key TEXT, submitted_at REAL NOT NULL, finished_at REAL); -CREATE UNIQUE INDEX IF NOT EXISTS jobs_idempotency_key ON jobs (idempotency_key);""" + timeout_ms INTEGER NOT NULL, idempotency_key TEXT, submitted_at REAL NOT NULL, finished_at REAL, owner TEXT);""" +_INDEX = """DROP INDEX IF EXISTS jobs_idempotency_key; +CREATE UNIQUE INDEX IF NOT EXISTS jobs_owner_idempotency_key ON jobs (owner, idempotency_key);""" +DEFAULT_OWNER = "default" class CapacityError(Exception): @@ -58,15 +61,18 @@ class Settings: max_jobs: int = MAX_JOBS_PER_WORKER # Where jobs.sqlite lives; BRIDGE_STATE_DIR in the service. None keeps the table in memory only. state_dir: str | None = None + # The names of the service's tokens: a job belongs to one, and each has its own metrics series. + owners: tuple[str, ...] = (DEFAULT_OWNER,) def _serve(connection, backend, model_dir): """A child's whole life: run the jobs the parent sends until it hangs up.""" # Ctrl+C in a terminal reaches the whole process group; stopping children is the parent's job. signal.signal(signal.SIGINT, signal.SIG_IGN) - # The service removes the token from its environment before it starts any child (server.py), so it is not + # The service removes its tokens from its environment before it starts any child (server.py), so they are not # in the child's environment block; this second pop covers a store built by something that did not. os.environ.pop("BRIDGE_TOKEN", None) + os.environ.pop("BRIDGE_TOKENS", None) try: # Jobs come only once the child is here to read them: the parent never waits on a child still importing. connection.send((None, "ready", None)) @@ -138,17 +144,21 @@ def __init__(self, settings=None): or not math.isfinite(self.settings.retention_seconds) or self.settings.retention_seconds <= 0): raise ValueError("invalid worker settings") + owners = self.settings.owners + if (type(owners) is not tuple or not owners or len(set(owners)) != len(owners) + or any(type(owner) is not str for owner in owners)): + raise ValueError("invalid worker owners") self._context = multiprocessing.get_context("spawn") # Aggregates for GET /metrics, since this process started; every series exists from the start, at zero. tasks = tuple(MODEL_NAMES[self.settings.backend]) + (TEST_TASKS if self.settings.test_tasks else ()) self._tasks = tasks - self._ended = {(task, status): 0 for task in tasks for status in STATUSES} - self._submitted = dict.fromkeys(tasks, 0) - self._rejected = {"capacity": 0} + self._ended = {(owner, task, status): 0 for owner in owners for task in tasks for status in STATUSES} + self._submitted = {(owner, task): 0 for owner in owners for task in tasks} + self._rejected = {(owner, "capacity"): 0 for owner in owners} self._durations = {task: metrics.new_histogram() for task in tasks} self._waits = {task: metrics.new_histogram() for task in tasks} self._jobs = {} - self._keys = {} # idempotency key -> job id, for retained jobs only + self._keys = {} # (owner, idempotency key) -> job id, for retained jobs only self._active = {} # job id -> the worker it was sent to self._pool = [] # children that may take a job, idle or running one self._retiring = [] # children that never serve again, until they are gone @@ -178,6 +188,11 @@ def _open_db(self): # An expired result is overwritten in the file, not left in a free page. self._db.execute("PRAGMA secure_delete=ON") self._db.executescript(_SCHEMA) + if "owner" not in [column[1] for column in self._db.execute("PRAGMA table_info(jobs)")]: + self._db.execute("ALTER TABLE jobs ADD COLUMN owner TEXT") + with self._db: + self._db.execute("UPDATE jobs SET owner = ? WHERE owner IS NULL", (DEFAULT_OWNER,)) + self._db.executescript(_INDEX) with self._db: self._db.execute( "UPDATE jobs SET status = 'failed', result = 'null', error = ?, finished_at = ? " @@ -186,21 +201,21 @@ def _open_db(self): self._db.execute("DELETE FROM jobs WHERE finished_at IS NULL OR finished_at <= ?", (time.time() - self.settings.retention_seconds,)) rows = self._db.execute("SELECT id, task, status, result, error, progress, timeout_ms, " - "idempotency_key, finished_at FROM jobs").fetchall() + "idempotency_key, finished_at, owner FROM jobs").fetchall() except Exception: self._db.close() raise wall, now = time.time(), time.monotonic() - for job_id, task, status, result, error, progress, timeout_ms, key, finished_at in rows: + for job_id, task, status, result, error, progress, timeout_ms, key, finished_at, owner in rows: self._jobs[job_id] = { "id": job_id, "task": task, "status": status, "result": json.loads(result), "error": json.loads(error), "progress": json.loads(progress), "timeout_ms": timeout_ms, # Retention counts from the wall-clock finish, carried over to this process's monotonic clock. "_input": None, "_deadline": now, "_finished": now - max(0.0, wall - finished_at), - "_retried": True, "_key": key, + "_retried": True, "_key": key, "_owner": owner, } if key is not None: - self._keys[key] = job_id + self._keys[owner, key] = job_id def _write(self, job, submitted=False): """Write a job's row through; a no-op once close() has let go of the connection. A submission that cannot @@ -216,9 +231,9 @@ def _write(self, job, submitted=False): if submitted: # REPLACE drops a row an earlier failed delete left holding this key, which memory no longer has. self._db.execute( - "INSERT OR REPLACE INTO jobs (task, timeout_ms, idempotency_key, submitted_at, status, result, " - "error, progress, finished_at, id) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - (job["task"], job["timeout_ms"], job["_key"], time.time()) + values) + "INSERT OR REPLACE INTO jobs (owner, task, timeout_ms, idempotency_key, submitted_at, status, " + "result, error, progress, finished_at, id) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + (job["_owner"], job["task"], job["timeout_ms"], job["_key"], time.time()) + values) else: self._db.execute("UPDATE jobs SET status = ?, result = ?, error = ?, progress = ?, " "finished_at = ? WHERE id = ?", values) @@ -226,7 +241,9 @@ def _write(self, job, submitted=False): if submitted: raise CapacityError("job table cannot be written") from None - def submit(self, task, payload, timeout_ms=30000, idempotency_key=None): + def submit(self, task, payload, timeout_ms=30000, idempotency_key=None, owner=DEFAULT_OWNER): + if owner not in self.settings.owners: + raise ValueError("owner is not one of the service's token names") if not isinstance(task, str): raise ValueError("task must be a string") if type(timeout_ms) is not int or not 100 <= timeout_ms <= 300000: @@ -237,19 +254,19 @@ def submit(self, task, payload, timeout_ms=30000, idempotency_key=None): payload = validate(task, payload, self.settings.test_tasks) with self._lock: try: - return self._accept(task, payload, timeout_ms, idempotency_key) + return self._accept(task, payload, timeout_ms, idempotency_key, owner) except CapacityError: # Every 429 the service sends comes from here, so each is counted exactly once. - self._rejected["capacity"] += 1 + self._rejected[owner, "capacity"] += 1 raise - def _accept(self, task, payload, timeout_ms, idempotency_key): + def _accept(self, task, payload, timeout_ms, idempotency_key, owner): if self._stopping.is_set(): raise CapacityError("worker is shutting down") self._expire() # A repeat adds no job, so it is answered even at capacity. - if idempotency_key in self._keys: - return self._public(self._jobs[self._keys[idempotency_key]]) + if (owner, idempotency_key) in self._keys: + return self._public(self._jobs[self._keys[owner, idempotency_key]]) if len(self._jobs) >= self.settings.capacity: raise CapacityError("job capacity reached; wait for retained jobs to expire", self._retry_after()) now = time.monotonic() @@ -258,24 +275,30 @@ def _accept(self, task, payload, timeout_ms, idempotency_key): "id": job_id, "task": task, "status": "queued", "result": None, "error": None, "progress": None, "timeout_ms": timeout_ms, "_input": copy.deepcopy(payload), "_deadline": now + timeout_ms / 1000, - "_finished": None, "_retried": False, "_key": idempotency_key, "_submitted_at": now, + "_finished": None, "_retried": False, "_key": idempotency_key, "_submitted_at": now, "_owner": owner, } self._write(job, submitted=True) self._jobs[job_id] = job if idempotency_key is not None: - self._keys[idempotency_key] = job_id - self._submitted[task] += 1 + self._keys[owner, idempotency_key] = job_id + self._submitted[owner, task] += 1 return self._public(job) - def get(self, job_id): + def _owned(self, job_id, owner): + """The job if this owner submitted it; another owner's job is as unknown to it as one that never was.""" + job = self._jobs.get(job_id) + return job if job is not None and job["_owner"] == owner else None + + def get(self, job_id, owner=DEFAULT_OWNER): with self._lock: self._expire() - return self._public(self._jobs[job_id]) if job_id in self._jobs else None + job = self._owned(job_id, owner) + return None if job is None else self._public(job) - def cancel(self, job_id): + def cancel(self, job_id, owner=DEFAULT_OWNER): with self._lock: self._expire() - job = self._jobs.get(job_id) + job = self._owned(job_id, owner) if job is None: return None if job["status"] not in TERMINAL: @@ -289,8 +312,8 @@ def snapshot(self): statuses = [job["status"] for job in self._jobs.values()] return { "backend": self.settings.backend, "concurrency": self.settings.concurrency, - "capacity": self.settings.capacity, "tasks": self._tasks, "jobs": dict(self._ended), - "submitted": dict(self._submitted), "rejected": dict(self._rejected), + "capacity": self.settings.capacity, "tasks": self._tasks, "owners": self.settings.owners, + "jobs": dict(self._ended), "submitted": dict(self._submitted), "rejected": dict(self._rejected), "queued": statuses.count("queued"), "running": statuses.count("running"), "pool_size": len(self._pool), "pool_ready": sum(worker.ready for worker in self._pool), "duration": copy.deepcopy(self._durations), "queue": copy.deepcopy(self._waits), @@ -313,7 +336,8 @@ def _expire(self): for job_id in list(self._jobs): finished = self._jobs[job_id]["_finished"] if finished is not None and now - finished >= self.settings.retention_seconds: - self._keys.pop(self._jobs.pop(job_id)["_key"], None) + job = self._jobs.pop(job_id) + self._keys.pop((job["_owner"], job["_key"]), None) expired.append((job_id,)) if expired and self._db is not None: try: @@ -328,7 +352,7 @@ def _finish(self, job, status, result=None, code=None): job["_finished"] = time.monotonic() self._write(job) # A job read back from the table finished in an earlier process and never ends here. - self._ended[job["task"], status] += 1 + self._ended[job["_owner"], job["task"], status] += 1 metrics.observe(self._durations[job["task"]], job["_finished"] - job["_submitted_at"]) def _abort(self, job, status, code=None): diff --git a/worker/ai_bridge/metrics.py b/worker/ai_bridge/metrics.py index 8b5941a..69f6e1c 100644 --- a/worker/ai_bridge/metrics.py +++ b/worker/ai_bridge/metrics.py @@ -1,6 +1,6 @@ """Aggregate counts and timings in the Prometheus text exposition format, version 0.0.4, with the standard library. -Label values come only from fixed sets (the task allowlist, the terminal statuses, the backend name), never from a -job, so nothing a client sent can reach the body and the number of series is bounded.""" +Label values come only from fixed sets (the configured token names, the task allowlist, the terminal statuses, the +backend name), never from a job, so nothing a client sent can reach the body and the number of series is bounded.""" import bisect import math @@ -71,11 +71,12 @@ def render(snapshot): tasks = snapshot["tasks"] return "".join([ counter("bridge_jobs_total", "Jobs that reached a terminal status since the process started.", - [({"task": task, "status": status}, count) for (task, status), count in snapshot["jobs"].items()]), + [({"owner": owner, "task": task, "status": status}, count) + for (owner, task, status), count in snapshot["jobs"].items()]), counter("bridge_jobs_submitted_total", "Jobs accepted since the process started.", - [({"task": task}, snapshot["submitted"][task]) for task in tasks]), + [({"owner": owner, "task": task}, count) for (owner, task), count in snapshot["submitted"].items()]), counter("bridge_submissions_rejected_total", "Submissions refused with 429 since the process started.", - [({"reason": reason}, count) for reason, count in snapshot["rejected"].items()]), + [({"owner": owner, "reason": reason}, count) for (owner, reason), count in snapshot["rejected"].items()]), gauge("bridge_queue_depth", "Jobs waiting for a child process.", [({}, snapshot["queued"])]), gauge("bridge_jobs_running", "Jobs a child process has started.", [({}, snapshot["running"])]), gauge("bridge_pool_size", "Child processes in the pool.", [({}, snapshot["pool_size"])]), diff --git a/worker/ai_bridge/server.py b/worker/ai_bridge/server.py index 666edbf..79d25da 100644 --- a/worker/ai_bridge/server.py +++ b/worker/ai_bridge/server.py @@ -8,29 +8,117 @@ import signal import socket import sqlite3 +import stat import threading +from dataclasses import replace from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from .jobs import CapacityError, JobStore, Settings +from .jobs import DEFAULT_OWNER, CapacityError, JobStore, Settings from .metrics import CONTENT_TYPE, render from .tasks import missing_models MAX_BODY = 262144 +MAX_TOKENS = 64 +MAX_TOKENS_FILE = 65536 +TOKEN_NAME = re.compile(r"[A-Za-z0-9_.-]{1,64}") + + +def check_tokens(tokens, source="the token list"): + """Refuse a token list the service must not start with; the message names the source, never a name or token.""" + if not tokens: + raise ValueError(f"{source} holds no token") + if len(tokens) > MAX_TOKENS: + raise ValueError(f"{source} holds more than {MAX_TOKENS} tokens") + for name, token in tokens.items(): + if type(name) is not str or not TOKEN_NAME.fullmatch(name): + raise ValueError(f"{source}: every name must be 1-64 characters of A-Z, a-z, 0-9, '_', '.' and '-'") + if type(token) is not str or len(token) < 32 or any(not 33 <= ord(char) <= 126 for char in token): + raise ValueError(f"{source}: every token must contain at least 32 non-whitespace ASCII characters") + if len(set(tokens.values())) != len(tokens): + raise ValueError(f"{source} gives two names the same token") + return tokens + + +def _entries(lines, source): + """name=token entries as a mapping; the token is everything after the first "=", since a token may hold one.""" + if len(lines) > MAX_TOKENS: + raise ValueError(f"{source} holds more than {MAX_TOKENS} tokens") + tokens = {} + for line in lines: + name, separator, token = line.partition("=") + if not separator or not TOKEN_NAME.fullmatch(name): + raise ValueError(f"{source}: every entry must be name=token, the name 1-64 characters of A-Z, a-z, 0-9, " + "'_', '.' and '-'") + if name in tokens: + raise ValueError(f"{source} names a token twice") + tokens[name] = token + return check_tokens(tokens, source) + + +def _read_tokens_file(path): + """The file's entries, read once; a file others could read or replace is refused, since it holds every token.""" + try: + # Non-blocking, so a FIFO at the path cannot hold the start up; a regular file reads as usual. + descriptor = os.open(path, os.O_RDONLY | getattr(os, "O_NONBLOCK", 0)) + with os.fdopen(descriptor, "rb") as file: + mode = os.fstat(file.fileno()).st_mode + if not stat.S_ISREG(mode): + raise ValueError("BRIDGE_TOKENS_FILE must be a regular file") + if mode & 0o077: + raise ValueError("BRIDGE_TOKENS_FILE must not be open to group or others: give it mode 0600 or stricter") + raw = file.read(MAX_TOKENS_FILE + 1) + except OSError: + raise ValueError("BRIDGE_TOKENS_FILE cannot be read by the service") from None + if len(raw) > MAX_TOKENS_FILE: + raise ValueError(f"BRIDGE_TOKENS_FILE is larger than {MAX_TOKENS_FILE} bytes") + try: + text = raw.decode("utf-8") + except UnicodeDecodeError: + raise ValueError("BRIDGE_TOKENS_FILE is not UTF-8 text") from None + lines = [line for line in text.splitlines() if line.strip() and not line.lstrip().startswith("#")] + return _entries(lines, "BRIDGE_TOKENS_FILE") + + +def resolve_tokens(): + """The service's tokens by name from BRIDGE_TOKEN (named default), BRIDGE_TOKENS or BRIDGE_TOKENS_FILE, exactly one + of them. Both token variables leave the environment here, before any child exists: the store's children inherit + the environment block, and a variable removed later stays in /proc//environ. The file's path may stay.""" + single = os.environ.pop("BRIDGE_TOKEN", "") + listed = os.environ.pop("BRIDGE_TOKENS", "") + path = os.environ.get("BRIDGE_TOKENS_FILE", "") + try: + if single and (listed or path): + raise ValueError("BRIDGE_TOKEN cannot be set together with BRIDGE_TOKENS or BRIDGE_TOKENS_FILE") + if listed and path: + raise ValueError("BRIDGE_TOKENS and BRIDGE_TOKENS_FILE cannot both be set") + if path: + return _read_tokens_file(path) + if listed: + return _entries(listed.split(","), "BRIDGE_TOKENS") + if single: + return check_tokens({DEFAULT_OWNER: single}, "BRIDGE_TOKEN") + raise ValueError("Set BRIDGE_TOKEN, BRIDGE_TOKENS or BRIDGE_TOKENS_FILE to the service's tokens") + except ValueError as error: + raise SystemExit(str(error)) from None class BridgeServer(ThreadingHTTPServer): daemon_threads = True request_queue_size = 16 - def __init__(self, address, token, settings=None): - if len(token) < 32 or any(not 33 <= ord(char) <= 126 for char in token): - raise ValueError("BRIDGE_TOKEN must contain at least 32 non-whitespace ASCII characters") - self.token = token + def __init__(self, address, tokens, settings=None): + """tokens maps each name to its token; a single token is the name default.""" + if isinstance(tokens, str): + tokens = {DEFAULT_OWNER: tokens} + check_tokens(tokens) + # The header each name is presented with; every request is compared with all of them. + self.credentials = tuple((name, ("Bearer " + token).encode("ascii")) for name, token in tokens.items()) self.slots = threading.BoundedSemaphore(32) # The children the store starts inherit this process's environment block as it is at that moment, and - # removing a variable later leaves it in /proc//environ; so the token leaves before any child exists. + # removing a variable later leaves it in /proc//environ; so the tokens leave before any child exists. os.environ.pop("BRIDGE_TOKEN", None) - self.jobs = JobStore(settings) + os.environ.pop("BRIDGE_TOKENS", None) + self.jobs = JobStore(replace(settings or Settings(), owners=tuple(tokens))) try: super().__init__(address, Handler) except Exception: @@ -93,12 +181,17 @@ def _error(self, status, code, message, headers=()): self._reply(status, {"error": {"code": code, "message": message}}, headers) def _authorized(self): + """The name of the token presented, or None once a 401 is sent. Every token is compared, whichever matches.""" values = self.headers.get_all("Authorization", []) - expected = ("Bearer " + self.server.token).encode("ascii") - if len(values) != 1 or not hmac.compare_digest(values[0].encode("utf-8"), expected): + presented = values[0].encode("utf-8") if len(values) == 1 else b"" + owner = None + for name, expected in self.server.credentials: + if hmac.compare_digest(presented, expected) and owner is None: + owner = name + if len(values) != 1 or owner is None: self._error(401, "unauthorized", "Service authentication required") - return False - return True + return None + return owner def _body(self): lengths = self.headers.get_all("Content-Length", []) @@ -130,7 +223,8 @@ def unique_keys(pairs): return body def do_GET(self): - if not self._authorized(): + owner = self._authorized() + if owner is None: return if self.path == "/healthz": self._reply(200, {"status": "ok", "protocol": 1, "backend": self.server.jobs.settings.backend}) @@ -140,14 +234,15 @@ def do_GET(self): self._text(200, render(self.server.jobs.snapshot())) return match = re.fullmatch(r"/v1/jobs/([a-f0-9]{32})", self.path) - job = self.server.jobs.get(match[1]) if match else None + job = self.server.jobs.get(match[1], owner) if match else None if job is None: self._error(404, "not_found", "Job or route not found") else: self._reply(200, job) def do_POST(self): - if not self._authorized(): + owner = self._authorized() + if owner is None: return try: body = self._body() @@ -159,13 +254,13 @@ def do_POST(self): if "idempotency_key" in body and body["idempotency_key"] is None: raise ValueError("idempotency_key must be a string") job = self.server.jobs.submit(body["task"], body["input"], body.get("timeout_ms", 30000), - body.get("idempotency_key")) + body.get("idempotency_key"), owner) self._reply(202, job) return match = re.fullmatch(r"/v1/jobs/([a-f0-9]{32})/cancel", self.path) if body: raise ValueError("Cancel body must be an empty object") - job = self.server.jobs.cancel(match[1]) if match else None + job = self.server.jobs.cancel(match[1], owner) if match else None if job is None: self._error(404, "not_found", "Job or route not found") else: @@ -194,6 +289,7 @@ def main(): parser.add_argument("--host", default="127.0.0.1") parser.add_argument("--port", type=int, default=8090) args = parser.parse_args() + tokens = resolve_tokens() settings = Settings( concurrency=_number("BRIDGE_CONCURRENCY", "2", int), capacity=_number("BRIDGE_CAPACITY", "128", int), @@ -211,7 +307,7 @@ def main(): raise SystemExit(f"{settings.model_dir} lacks the model files for {', '.join(missing)}; acquire them with " "the fetcher: docker compose -f docker/compose.yaml run --rm --no-deps fetcher") try: - bridge = BridgeServer((args.host, args.port), os.environ.pop("BRIDGE_TOKEN", ""), settings) + bridge = BridgeServer((args.host, args.port), tokens, settings) except sqlite3.Error: raise SystemExit(f"BRIDGE_STATE_DIR {settings.state_dir} must be a directory the service can write " "jobs.sqlite in") from None diff --git a/worker/ai_bridge/tasks.py b/worker/ai_bridge/tasks.py index 9b0c771..2d54613 100644 --- a/worker/ai_bridge/tasks.py +++ b/worker/ai_bridge/tasks.py @@ -184,10 +184,10 @@ def _reachable_holding(needle: str, roots: list) -> int: def _inspect(parts: list) -> dict: - """What a test may ask a worker child about itself: its process id, whether the service token reached its - environment or its initial environment block (which unsetenv leaves in place), and how many strings reachable - from any live object or any frame's locals still hold the text joined from parts, which no earlier job of the - child may have left behind. The parts are joined here so that the question itself does not hold the answer; + """What a test may ask a worker child about itself: its process id, whether a service token variable or the text + joined from parts reached its environment or its initial environment block (which unsetenv leaves in place), and + how many strings reachable from any live object or any frame's locals still hold that text, which no earlier job + of the child may have left behind. The parts are joined here so that the question itself does not hold the answer; the joined copy is always a new object, so it is left out of the count by identity.""" import gc import sys @@ -210,11 +210,15 @@ def _inspect(parts: list) -> dict: held = _reachable_holding(needle, roots) try: with open("/proc/self/environ", "rb") as environ: - in_block = b"BRIDGE_TOKEN=" in environ.read() + block = environ.read().split(b"\0") + in_block = any(entry.startswith((b"BRIDGE_TOKEN=", b"BRIDGE_TOKENS=")) for entry in block) + needle_in_block = bool(needle) and any(needle.encode() in entry for entry in block) except OSError: - in_block = None # no procfs here - return {"pid": os.getpid(), "token_in_environment": "BRIDGE_TOKEN" in os.environ, - "token_in_environment_block": in_block, "held": held} + in_block = needle_in_block = None # no procfs here + return {"pid": os.getpid(), "token_in_environment": "BRIDGE_TOKEN" in os.environ or "BRIDGE_TOKENS" in os.environ, + "token_in_environment_block": in_block, "held": held, + "needle_in_environment": bool(needle) and any(needle in value for value in os.environ.values()), + "needle_in_environment_block": needle_in_block} def _tensors(encodings: list) -> dict: