diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d24466c..656d4a2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -41,7 +41,6 @@ jobs: - { package: ./internal/report, name: FuzzText } - { package: ./internal/report, name: FuzzClues } - { package: ./internal/repo, name: FuzzReadFileStaysInTrackedFiles } - - { package: ./internal/modelproxy, name: FuzzProviderRequest } steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: diff --git a/docker/agent/Dockerfile b/docker/agent/Dockerfile index 55cac83..99f7ec1 100644 --- a/docker/agent/Dockerfile +++ b/docker/agent/Dockerfile @@ -9,7 +9,7 @@ ARG CLAUDE_AMD64_LOCK=v2.1.306@sha256:9dc640d16d8228eff9620bde9e0f70d38385a2124f ARG CLAUDE_ARM64_LOCK=v2.1.306@sha256:871fd7f46210c13a95a38506e03ac90abfc5125fc25a1d4f840fbab551b5cd2e ARG CODEX_AMD64_LOCK=rust-v0.163.0@sha256:582646c6bf0561eb0de540f7fc85b9121b0091a826bc6e96830b7f1722122182 ARG CODEX_ARM64_LOCK=rust-v0.163.0@sha256:9d8ecf2d19d3e184877983902d6c85de48b3888dcb72c382a8a0c975845927fb -ARG HARNESS_VERSION=v0.1.16 +ARG HARNESS_VERSION=v0.1.18 FROM golang:1.27.3-trixie@sha256:1896654cfc5c2b5ad74c07f29cb36f5d2d1cbed42e3054070fbfe444e521e83c AS proxy ARG HARNESS_VERSION diff --git a/docs/development.md b/docs/development.md index 511c2c3..964d302 100644 --- a/docs/development.md +++ b/docs/development.md @@ -9,20 +9,21 @@ go build ./cmd/tamis go test -race ./... ``` -Four fuzz targets cover the comment escaper, the duplicate-search clues, file reading and the provider protocol filter: +Three fuzz targets cover the comment escaper, the duplicate-search clues and file reading: | Target | Package | |---|---| | `FuzzText` | `./internal/report` | | `FuzzClues` | `./internal/report` | | `FuzzReadFileStaysInTrackedFiles` | `./internal/repo` | -| `FuzzProviderRequest` | `./internal/modelproxy` | ```bash go test ./internal/report -run '^$' -fuzz '^FuzzText$' -fuzztime 60s ``` -On each push, CI checks formatting and `go mod tidy`, then runs `go vet`, the tests, the four fuzz targets and the Docker isolation tests of both agent images. Make both isolation jobs required release checks. +The provider protocol filter is fuzzed in harness, with its model proxy. + +On each push, CI checks formatting and `go mod tidy`, then runs `go vet`, the tests, the three fuzz targets and the Docker isolation tests of both agent images. Make both isolation jobs required release checks. ## Isolation tests @@ -79,8 +80,7 @@ tamis posts the comment for real. It skips an advisory that already has a comple |---|---| | `cmd/tamis` | Configuration and entry point | | `internal/triage` | Advisory selection and the triage pipeline | -| `internal/sandbox` | The agent container, through harness | -| `internal/modelproxy` | The proxy that keeps the model key on the host | +| `internal/sandbox` | The agent container and the model proxy that keeps the key on the host, through harness | | `internal/budget` | Request and byte budgets shared by every agent and comparator of a run | | `internal/usage` | Requests and tokens spent per advisory, for the log | | `internal/engine` | The in-process model loop of the exact duplicate confirmation | diff --git a/docs/security.md b/docs/security.md index 13aa6d4..c37cccf 100644 --- a/docs/security.md +++ b/docs/security.md @@ -37,7 +37,7 @@ The agent runs through [harness](https://github.com/alpha-omega-security/harness - a throwaway container per advisory, with a read-only root filesystem, all capabilities dropped and no new privileges, - one single-branch clone of the analysed branch per execution, mounted read-only into each agent run, with no untracked files and an empty `.git/config`, while only tamis reads the other branches, checked out as worktrees, - an internal bridge network with both gateway modes set to `isolated` and IPv6 disabled, whose only peer is an authenticated egress proxy that allows the model proxy and `api.github.com`, -- a read-only workspace except for `report.json`, with memory, CPU, process, file-size, output-capture and Docker log limits, +- a read-only workspace except for `report.json`, with memory, CPU, process and file-size limits, - an image with only the agent CLI, git, curl, ripgrep and `harness-proxy`, each download checked against a pinned checksum, - that image pinned by digest in each release, and checked by the action against its build attestation before use, like the tamis binary of a release. @@ -68,6 +68,7 @@ tamis turns off the Codex ChatGPT features that call other routes: plugins, apps - Request and byte limits bound one run, not the cost of an endless series of runs. Enforce a provider spending ceiling and a workflow timeout as the [setup guide](setup.md#2-create-a-private-repository) describes. - tamis sends the report to the model provider. Check its data retention terms against your disclosure policy. - An internal network without a host gateway needs Docker Engine 28 or newer. +- The agent CLI state is a writable host directory, removed after each advisory. A run can rewrite its own configuration for the correction round and fill the runner disk, though not leave the network or the model proxy. - The advisory comments API is in public preview and may change. CI builds both agent images and runs the Docker isolation tests on `ubuntu-latest`. They test the model path, direct and proxied host access, DNS, and IPv4 and IPv6 egress, without real provider credentials. Treat both `isolation` jobs as required checks before releasing or updating the deployment SHA. diff --git a/go.mod b/go.mod index a641f7c..a295f57 100644 --- a/go.mod +++ b/go.mod @@ -3,7 +3,7 @@ module github.com/php/tamis go 1.27.3 require ( - github.com/alpha-omega-security/harness v0.1.16 + github.com/alpha-omega-security/harness v0.1.18 github.com/anthropics/anthropic-sdk-go v1.84.0 github.com/git-pkgs/cwe v0.1.1 github.com/openai/openai-go/v3 v3.85.1 @@ -20,6 +20,7 @@ require ( github.com/dnephin/pflag v1.0.7 // indirect github.com/fatih/color v1.19.0 // indirect github.com/fsnotify/fsnotify v1.10.1 // indirect + github.com/git-pkgs/skills v0.2.0 // indirect github.com/go-logr/logr v1.4.4 // indirect github.com/go-logr/stdr v1.2.2 // indirect github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 // indirect diff --git a/go.sum b/go.sum index ad0256a..38b050f 100644 --- a/go.sum +++ b/go.sum @@ -8,8 +8,8 @@ github.com/Azure/azure-sdk-for-go/sdk/azcore v1.23.2 h1:utpeoEeZjd+A8J41zvoLsOOr github.com/Azure/azure-sdk-for-go/sdk/azcore v1.23.2/go.mod h1:iptorS+VYKFL2N6PnebpS91dubG35eAOEERnT4PJbQU= github.com/Azure/azure-sdk-for-go/sdk/internal v1.12.0 h1:fhqpLE3UEXi9lPaBRpQ6XuRW0nU7hgg4zlmZZa+a9q4= github.com/Azure/azure-sdk-for-go/sdk/internal v1.12.0/go.mod h1:7dCRMLwisfRH3dBupKeNCioWYUZ4SS09Z14H+7i8ZoY= -github.com/alpha-omega-security/harness v0.1.16 h1:FxYZO4aBRUOJT4tBBDeflbTUB3atu5ucmniy+c2R9t4= -github.com/alpha-omega-security/harness v0.1.16/go.mod h1:wpUi6g2ebGi0OMSHaDTGhj/zMV1Lse1P8GrAp4J3LMI= +github.com/alpha-omega-security/harness v0.1.18 h1:99coqkQjuQ13jAX0ivZo7mu5xQyJW6n/QWj19USGMCI= +github.com/alpha-omega-security/harness v0.1.18/go.mod h1:4JI0CxVRwiWQ9zAtSTnM5p8l0A79bwREtps9BA9oLXY= github.com/anthropics/anthropic-sdk-go v1.84.0 h1:mnAhggc5H5A98JYmhQELT+j4J89PYBFKnnSvUTNwNd8= github.com/anthropics/anthropic-sdk-go v1.84.0/go.mod h1:GThfYqPJoaQ/6pmibCI98Cr4y5su2FXMUHn3NrSSnIc= github.com/aws/aws-sdk-go-v2 v1.47.1 h1:uOIZnp4PK3ZhKI0dNrJrhTEsLxbpXHTAJlwoS1pvAtw= @@ -64,6 +64,8 @@ github.com/fsnotify/fsnotify v1.10.1 h1:b0/UzAf9yR5rhf3RPm9gf3ehBPpf0oZKIjtpKrx5 github.com/fsnotify/fsnotify v1.10.1/go.mod h1:TLheqan6HD6GBK6PrDWyDPBaEV8LspOxvPSjC+bVfgo= github.com/git-pkgs/cwe v0.1.1 h1:FTdif2t/pQCd8IFxJUt3fy4khoQ93djiw2sukqE8KiQ= github.com/git-pkgs/cwe v0.1.1/go.mod h1:6yUxEoA+3rI4W3ftqb7PZP1Z7zZFmV4/wfr7P46HtHU= +github.com/git-pkgs/skills v0.2.0 h1:WcYbOukbrQLhjrg0qocknf16MrLVF5bKdwOeNHF6we0= +github.com/git-pkgs/skills v0.2.0/go.mod h1:fnhLBuE0+QEhLh6bb3x2IDljwu5lhJlrU+FE2LfHAk4= github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= github.com/go-logr/logr v1.4.4 h1:tG4xh9yMsRCAiodLVTxyrkzSZ9+o0L1Kg/+cPVcbP/8= github.com/go-logr/logr v1.4.4/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= diff --git a/internal/modelproxy/live_test.go b/internal/modelproxy/live_test.go deleted file mode 100644 index 012717e..0000000 --- a/internal/modelproxy/live_test.go +++ /dev/null @@ -1,130 +0,0 @@ -//go:build live - -package modelproxy - -import ( - "bufio" - "bytes" - "cmp" - "encoding/json" - "fmt" - "net/http" - "net/http/httptest" - "os" - "os/exec" - "path/filepath" - "strings" - "testing" - - "github.com/stretchr/testify/require" - - "github.com/php/tamis/internal/usage" -) - -// These tests pin that the usage the proxy counts for a real agent CLI run equals the one the CLI reports: -// TAMIS_LIVE_ANTHROPIC_KEY_FILE=path/to/key go test -tags live -run Live -v ./internal/modelproxy -// TAMIS_LIVE_OPENAI_KEY_FILE=path/to/key go test -tags live -run Live -v ./internal/modelproxy - -const livePrompt = "Read go.mod in the current directory with a tool, then reply with the module path only." - -type statusRecorder struct { - http.ResponseWriter - status int -} - -func (s *statusRecorder) WriteHeader(code int) { - s.status = code - s.ResponseWriter.WriteHeader(code) -} - -func (s *statusRecorder) Flush() { - s.ResponseWriter.(http.Flusher).Flush() -} - -// liveRun serves a proxy for provider with the key in the file named by env, and runs cmd in the repository root -// with the URL and the token issued for it. It returns the CLI output and what the proxy counted. -func liveRun(t *testing.T, provider Provider, env, model string, cmd func(url, token string) *exec.Cmd) ([]byte, usage.Usage) { - t.Helper() - path := os.Getenv(env) - if path == "" { - t.Skip(env + " is not set") - } - raw, err := os.ReadFile(path) - require.NoError(t, err) - mp, err := New(provider, strings.TrimSpace(string(raw)), model) - require.NoError(t, err) - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - rec := &statusRecorder{ResponseWriter: w, status: http.StatusOK} - mp.ServeHTTP(rec, r) - t.Logf("%s %s -> %d", r.Method, r.URL.Path, rec.status) - })) - t.Cleanup(srv.Close) - token, revoke := mp.Issue(t.Context()) - - c := cmd(srv.URL+PathPrefix, token) - c.Dir = "../.." - out, err := c.Output() - spent := revoke() - require.NoError(t, err, "the CLI failed: %s", out) - t.Logf("proxy: %+v", spent) - - return out, spent -} - -func TestLiveClaudeCodeUsage(t *testing.T) { - model := cmp.Or(os.Getenv("TAMIS_LIVE_MODEL"), "claude-haiku-5-5") - out, spent := liveRun(t, Anthropic, "TAMIS_LIVE_ANTHROPIC_KEY_FILE", model, func(url, token string) *exec.Cmd { - c := exec.Command("claude", "-p", livePrompt, "--model", model, "--output-format", "json", "--allowedTools", "Read") - c.Env = append(os.Environ(), "ANTHROPIC_BASE_URL="+url, "ANTHROPIC_API_KEY="+token, "CLAUDE_CONFIG_DIR="+t.TempDir(), "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC=1") - return c - }) - - var result struct { - Usage struct { - InputTokens int64 `json:"input_tokens"` - CacheCreation int64 `json:"cache_creation_input_tokens"` - CacheRead int64 `json:"cache_read_input_tokens"` - OutputTokens int64 `json:"output_tokens"` - } `json:"usage"` - } - require.NoError(t, json.Unmarshal(out, &result), "%s", out) - cli := usage.Usage{InputTokens: result.Usage.InputTokens + result.Usage.CacheCreation + result.Usage.CacheRead, OutputTokens: result.Usage.OutputTokens} - t.Logf("cli: %+v", cli) - require.GreaterOrEqual(t, spent.Requests, int64(2)) - require.Equal(t, cli.InputTokens, spent.InputTokens) - require.Equal(t, cli.OutputTokens, spent.OutputTokens) -} - -func TestLiveCodexUsage(t *testing.T) { - model := cmp.Or(os.Getenv("TAMIS_LIVE_MODEL"), "gpt-5.5") - out, spent := liveRun(t, OpenAI, "TAMIS_LIVE_OPENAI_KEY_FILE", model, func(url, token string) *exec.Cmd { - home := t.TempDir() - // The same features as in sandbox.writeConfig are off, the proxy refuses the requests they make. - config := fmt.Sprintf("model_provider = \"tamis\"\nweb_search = \"disabled\"\n\n[analytics]\nenabled = false\n\n[features]\napps = false\nplugins = false\nremote_plugin = false\nplugin_sharing = false\ngoals = false\nimage_generation = false\nskill_mcp_dependency_install = false\n\n[model_providers.tamis]\nname = \"Tamis\"\nbase_url = %q\nenv_key = \"CODEX_API_KEY\"\nwire_api = \"responses\"\nsupports_websockets = false\nrequest_max_retries = 0\nstream_max_retries = 0\n", url+"/v1") - require.NoError(t, os.WriteFile(filepath.Join(home, "config.toml"), []byte(config), 0o600)) - c := exec.Command("codex", "exec", "--json", "--ephemeral", "--skip-git-repo-check", "-s", "read-only", "-m", model, livePrompt) - c.Env = append(os.Environ(), "CODEX_HOME="+home, "CODEX_API_KEY="+token) - return c - }) - - var cli usage.Usage - scanner := bufio.NewScanner(bytes.NewReader(out)) - scanner.Buffer(nil, 1<<20) - for scanner.Scan() { - var event struct { - Type string `json:"type"` - Usage struct { - InputTokens int64 `json:"input_tokens"` - OutputTokens int64 `json:"output_tokens"` - } `json:"usage"` - } - if json.Unmarshal(scanner.Bytes(), &event) == nil && event.Type == "turn.completed" { - cli.InputTokens += event.Usage.InputTokens - cli.OutputTokens += event.Usage.OutputTokens - } - } - t.Logf("cli: %+v", cli) - require.GreaterOrEqual(t, spent.Requests, int64(2), "%s", out) - require.Equal(t, cli.InputTokens, spent.InputTokens, "%s", out) - require.Equal(t, cli.OutputTokens, spent.OutputTokens, "%s", out) -} diff --git a/internal/modelproxy/modelproxy.go b/internal/modelproxy/modelproxy.go deleted file mode 100644 index 139dde2..0000000 --- a/internal/modelproxy/modelproxy.go +++ /dev/null @@ -1,495 +0,0 @@ -// Package modelproxy keeps the model API key on the host. A sandboxed agent gets a token that is only -// good for this proxy, which swaps it for the real key on a fixed list of routes and forwards the request. -// Adapted from scrutineer's internal/worker/model_proxy.go. -package modelproxy - -import ( - "bytes" - "cmp" - "context" - "crypto/rand" - "crypto/sha256" - "encoding/hex" - "encoding/json" - "errors" - "fmt" - "io" - "log/slog" - "maps" - "net/http" - "net/http/httputil" - "net/url" - "path" - "regexp" - "slices" - "strings" - "sync" - - "github.com/php/tamis/internal/budget" - "github.com/php/tamis/internal/usage" -) - -const ( - // PathPrefix is where the proxy serves, the agent's base URL ends with it. - PathPrefix = "/model-proxy" - // maxBody caps a forwarded request so a hostile agent cannot make the host buffer unbounded data. - maxBody = 1 << 20 - maxRequests = 150 - maxTokens = 16_000 - tokenPrefix = "tamis-" - // maxLoggedShape bounds what one refusal adds to the workflow log. - maxLoggedShape = 8 << 10 -) - -var errUpstreamRedirect = errors.New("model proxy: upstream redirect refused") - -var clientVersion = regexp.MustCompile(`^client_version=[0-9]+(\.[0-9]+)*$`) - -// runKey carries the hash of the token a request came with, so its response is counted for that run. -type runKey struct{} - -// Provider describes one model API: where it lives, which routes an agent may call and how the key is sent. -type Provider struct { - Upstream string - // Upstreams sends the routes under each prefix to their own upstream, the prefix itself is dropped. - Upstreams map[string]string - // TokenPrefix starts the tokens the agent gets, when its CLI recognises a credential by its prefix. - TokenPrefix string - Routes func(method, route string) bool - SetKey func(h http.Header, key string) - NoStore bool - // StoreOff forces store to false only, for an upstream that rejects the other NoStore defaults. - StoreOff bool - UntypedTools bool - // Query permits a query string on a route, none is permitted without it. - Query func(route, query string) bool - // Unlogged lists the "METHOD route" refusals a CLI makes on every run, which would hide unexpected ones. - Unlogged map[string]bool - // PrivateWorkspace redacts ChatGPT identity and managed configuration before they reach Codex. - PrivateWorkspace bool -} - -var Anthropic = Provider{ - Upstream: "https://api.anthropic.com", - Routes: func(method, route string) bool { - switch { - case method == http.MethodPost: - return route == "/v1/messages" || route == "/v1/messages/count_tokens" - default: - return false - } - }, - SetKey: func(h http.Header, key string) { h.Set("X-Api-Key", key) }, - UntypedTools: true, - Query: func(_, query string) bool { return query == "beta=true" }, -} - -var OpenAI = Provider{ - Upstream: "https://api.openai.com", - Routes: func(method, route string) bool { - switch { - case method == http.MethodPost: - return route == "/v1/responses" - default: - return false - } - }, - SetKey: func(h http.Header, key string) { h.Set("Authorization", "Bearer "+key) }, - NoStore: true, -} - -// ChatGPT serves Codex signed in with a Codex access token: it checks the token on auth.openai.com, then -// calls the model through the ChatGPT backend instead of the OpenAI API. -var ChatGPT = Provider{ - Upstreams: map[string]string{ - "/auth": "https://auth.openai.com/api/accounts", - "/backend-api": "https://chatgpt.com/backend-api", - }, - Routes: func(method, route string) bool { - switch method { - case http.MethodGet: - // Codex refuses to start without the workspace's managed configuration, and asks for its model - // catalogue again on every turn while it cannot read it. - return route == "/auth/v1/user-auth-credential/whoami" || route == "/backend-api/wham/config/bundle" || route == "/backend-api/codex/models" - case http.MethodPost: - return route == "/backend-api/codex/responses" - default: - return false - } - }, - Query: func(route, query string) bool { - return route == "/backend-api/codex/models" && clientVersion.MatchString(query) - }, - // Codex checks the account and its settings at start-up and runs without them. - Unlogged: map[string]bool{"GET /backend-api/wham/accounts/check": true, "GET /backend-api/wham/settings/user": true}, - TokenPrefix: "at-" + tokenPrefix, - SetKey: func(h http.Header, key string) { h.Set("Authorization", "Bearer "+key) }, - StoreOff: true, - PrivateWorkspace: true, -} - -type Proxy struct { - provider Provider - key string - proxy *httputil.ReverseProxy - - mu sync.RWMutex - tokens map[[sha256.Size]byte]*tokenBudget - spent map[[sha256.Size]byte]usage.Usage - model string - accountID string -} - -type tokenBudget struct { - remaining int - ctx context.Context -} - -func New(provider Provider, key, model string) (*Proxy, error) { - if key == "" { - return nil, errors.New("model proxy: an API key is required") - } - if model == "" { - return nil, errors.New("model proxy: a model is required") - } - upstreams := map[string]*url.URL{} - for prefix, raw := range provider.Upstreams { - u, err := url.Parse(raw) - if err != nil { - return nil, fmt.Errorf("model proxy: %w", err) - } - upstreams[prefix] = u - } - if provider.Upstreams == nil { - u, err := url.Parse(provider.Upstream) - if err != nil { - return nil, fmt.Errorf("model proxy: %w", err) - } - upstreams[""] = u - } - - p := &Proxy{provider: provider, key: key, model: model, tokens: map[[sha256.Size]byte]*tokenBudget{}, spent: map[[sha256.Size]byte]usage.Usage{}} - p.proxy = &httputil.ReverseProxy{ - Rewrite: func(pr *httputil.ProxyRequest) { - // ServeHTTP only forwards permitted routes, each under exactly one prefix. - var upstream *url.URL - route := pr.In.URL.Path - for prefix, u := range upstreams { - if rest, ok := strings.CutPrefix(pr.In.URL.Path, prefix); ok && (prefix == "" || strings.HasPrefix(rest, "/")) { - upstream, route = u, rest - } - } - pr.Out.URL.Scheme, pr.Out.URL.Host, pr.Out.Host = upstream.Scheme, upstream.Host, upstream.Host - pr.Out.URL.Path, pr.Out.URL.RawPath = upstream.Path+route, "" - // New provider capabilities must not be enabled through a client-supplied header. - pr.Out.Header = http.Header{} - pr.Out.Header.Set("Content-Type", "application/json") - pr.Out.Header.Set("Accept", pr.In.Header.Get("Accept")) - if provider.UntypedTools { - pr.Out.Header.Set("Anthropic-Version", "2023-06-01") - var betas []string - for _, beta := range strings.Split(pr.In.Header.Get("Anthropic-Beta"), ",") { - beta = strings.TrimSpace(beta) - if slices.Contains([]string{"claude-code-20250219", "interleaved-thinking-2025-05-14", "fine-grained-tool-streaming-2025-05-14", "context-management-2025-06-27", "effort-2025-11-24", "adaptive-thinking-2026-01-28", "structured-outputs-2025-11-13", "prompt-caching-2024-07-31", "tool-search-tool-2025-10-19", "per-turn-control-2026-07-01", "mid-conversation-tool-changes-2026-07-01", "thinking-display-updates-2026-08-18", "thinking-token-count-2026-05-13", "prompt-caching-scope-2026-01-05"}, beta) { - betas = append(betas, beta) - } - } - if len(betas) > 0 { - pr.Out.Header.Set("Anthropic-Beta", strings.Join(betas, ",")) - } - } - provider.SetKey(pr.Out.Header, key) - if provider.PrivateWorkspace && strings.HasPrefix(pr.In.URL.Path, "/backend-api/") { - p.mu.RLock() - accountID := p.accountID - p.mu.RUnlock() - if accountID != "" { - pr.Out.Header.Set("ChatGPT-Account-Id", accountID) - } - } - }, - // Stream server-sent events to the agent as they arrive. - FlushInterval: -1, - ModifyResponse: func(resp *http.Response) error { - resp.Header.Del("Set-Cookie") - if resp.StatusCode >= http.StatusMultipleChoices && resp.StatusCode < http.StatusBadRequest { - return errUpstreamRedirect - } - if resp.StatusCode >= http.StatusBadRequest { - slog.Warn("model provider refused a request", "method", resp.Request.Method, "route", resp.Request.URL.Path, "status", resp.StatusCode) - } - if provider.PrivateWorkspace && resp.StatusCode < http.StatusMultipleChoices { - switch { - case strings.HasSuffix(resp.Request.URL.Path, "/v1/user-auth-credential/whoami"): - return p.redactIdentity(resp) - case strings.HasSuffix(resp.Request.URL.Path, "/wham/config/bundle"): - return replaceBody(resp, []byte(`{"config_toml":{"enterprise_managed":[]},"requirements_toml":{"enterprise_managed":[]}}`)) - } - } - if resp.Request.Method == http.MethodPost && resp.StatusCode < http.StatusMultipleChoices && !strings.HasSuffix(resp.Request.URL.Path, "/count_tokens") { - sum, _ := resp.Request.Context().Value(runKey{}).([sha256.Size]byte) - resp.Body = &usageReader{ - ReadCloser: resp.Body, - stream: strings.HasPrefix(resp.Header.Get("Content-Type"), "text/event-stream"), - done: func(u usage.Usage) { p.record(sum, u) }, - } - } - - return nil - }, - // Upstream errors are not logged: they could carry parts of the conversation. - ErrorHandler: func(w http.ResponseWriter, _ *http.Request, err error) { - var tooLarge *http.MaxBytesError - if errors.As(err, &tooLarge) { - http.Error(w, "request too large", http.StatusRequestEntityTooLarge) - return - } - http.Error(w, "model proxy could not reach the upstream API", http.StatusBadGateway) - }, - } - - return p, nil -} - -// Issue mints a token for one agent run, revoke invalidates it and returns what the run spent. Only its hash is kept. -func (p *Proxy) Issue(ctx context.Context) (token string, revoke func() usage.Usage) { - var raw [32]byte - _, _ = rand.Read(raw[:]) - token = cmp.Or(p.provider.TokenPrefix, tokenPrefix) + hex.EncodeToString(raw[:]) - sum := sha256.Sum256([]byte(token)) - p.mu.Lock() - p.tokens[sum] = &tokenBudget{remaining: maxRequests, ctx: ctx} - p.mu.Unlock() - - return token, func() usage.Usage { - p.mu.Lock() - defer p.mu.Unlock() - spent := p.spent[sum] - delete(p.tokens, sum) - delete(p.spent, sum) - - return spent - } -} - -func (p *Proxy) record(sum [sha256.Size]byte, u usage.Usage) { - p.mu.Lock() - defer p.mu.Unlock() - if _, issued := p.tokens[sum]; issued { - p.spent[sum] = p.spent[sum].Plus(u) - } -} - -func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) { - route, ok := strings.CutPrefix(r.URL.Path, PathPrefix) - if !ok { - http.NotFound(w, r) - return - } - sum, token := p.authorized(r) - if token == nil { - http.Error(w, "invalid or revoked token", http.StatusUnauthorized) - return - } - r = r.WithContext(context.WithValue(r.Context(), runKey{}, sum)) - if r.URL.RawPath != "" || (r.URL.RawQuery != "" && (p.provider.Query == nil || !p.provider.Query(route, r.URL.RawQuery))) || path.Clean(route) != route || !p.provider.Routes(r.Method, route) { - if p.provider.Unlogged[r.Method+" "+route] { - // These refusals would otherwise spend the requests the agent needs for the model. - p.mu.Lock() - token.remaining++ - p.mu.Unlock() - } else { - slog.Info("model proxy refused a route", "method", r.Method, "route", truncate(route, maxLoggedShape)) - } - http.Error(w, "route not permitted", http.StatusForbidden) - return - } - if r.Method == http.MethodPost { - raw, err := io.ReadAll(io.LimitReader(r.Body, maxBody+1)) - if err != nil { - http.Error(w, "invalid provider request", http.StatusForbidden) - return - } - body, err := p.safeBody(bytes.NewReader(raw)) - if err != nil { - if errors.Is(err, errRequestTooLarge) { - http.Error(w, err.Error(), http.StatusRequestEntityTooLarge) - return - } - var v any - _ = json.Unmarshal(raw, &v) - slog.Warn("model proxy refused a request", "route", route, "shape", truncate(shape(v), maxLoggedShape)) - http.Error(w, err.Error(), http.StatusForbidden) - return - } - r.Body = io.NopCloser(bytes.NewReader(body)) - r.ContentLength = int64(len(body)) - if !budget.Take(token.ctx, len(body)) { - http.Error(w, "provider budget exhausted", http.StatusTooManyRequests) - return - } - } - if strings.Contains(strings.ToLower(r.Header.Get("Anthropic-Beta")), "mcp") { - http.Error(w, "provider capability not permitted", http.StatusForbidden) - return - } - r.URL.Path = route - p.proxy.ServeHTTP(w, r) -} - -func (p *Proxy) authorized(r *http.Request) ([sha256.Size]byte, *tokenBudget) { - token := r.Header.Get("X-Api-Key") - if bearer, ok := strings.CutPrefix(r.Header.Get("Authorization"), "Bearer "); token == "" && ok { - token = bearer - } - if token == "" { - return [sha256.Size]byte{}, nil - } - sum := sha256.Sum256([]byte(token)) - p.mu.Lock() - defer p.mu.Unlock() - b := p.tokens[sum] - if b == nil || b.remaining <= 0 || b.ctx.Err() != nil { - return sum, nil - } - b.remaining-- - - return sum, b -} - -func (p *Proxy) redactIdentity(resp *http.Response) error { - defer resp.Body.Close() - body, err := io.ReadAll(io.LimitReader(resp.Body, maxBody+1)) - if err != nil || len(body) > maxBody { - return errors.New("model proxy: invalid identity response") - } - var identity struct { - AccountID string `json:"chatgpt_account_id"` - PlanType string `json:"chatgpt_plan_type"` - FedRAMP bool `json:"chatgpt_account_is_fedramp"` - } - if err := json.Unmarshal(body, &identity); err != nil || identity.AccountID == "" || identity.PlanType == "" { - return errors.New("model proxy: invalid identity response") - } - p.mu.Lock() - p.accountID = identity.AccountID - p.mu.Unlock() - redacted, _ := json.Marshal(map[string]any{ - "email": nil, - "chatgpt_user_id": "tamis", - "chatgpt_account_id": "tamis", - "chatgpt_plan_type": identity.PlanType, - "chatgpt_account_is_fedramp": identity.FedRAMP, - }) - - return replaceBody(resp, redacted) -} - -func replaceBody(resp *http.Response, body []byte) error { - if resp.Body != nil { - resp.Body.Close() - } - resp.Body = io.NopCloser(bytes.NewReader(body)) - resp.ContentLength = int64(len(body)) - resp.Header.Set("Content-Type", "application/json") - resp.Header.Del("Content-Encoding") - resp.Header.Del("Content-Length") - - return nil -} - -var errRequestTooLarge = errors.New("request too large") - -func (p *Proxy) safeBody(body io.Reader) ([]byte, error) { - b, err := io.ReadAll(io.LimitReader(body, maxBody+1)) - if err != nil { - return nil, errors.New("invalid provider request") - } - if len(b) > maxBody { - return nil, errRequestTooLarge - } - var v map[string]any - dec := json.NewDecoder(bytes.NewReader(b)) - dec.UseNumber() - if err := dec.Decode(&v); err != nil { - return nil, errors.New("invalid provider request") - } - if _, err := dec.Token(); err != io.EOF { - return nil, errors.New("invalid provider request") - } - model, ok := v["model"].(string) - if !ok || model != p.model || !safeRequest(v, p.provider.UntypedTools) { - return nil, errors.New("provider request not permitted") - } - delete(v, "client_metadata") - for _, key := range []string{"previous_response_id", "conversation", "prompt", "fallbacks", "multi_agent"} { - if _, ok := v[key]; ok { - return nil, errors.New("provider request not permitted") - } - } - for _, key := range []string{"max_tokens", "max_output_tokens"} { - if raw, ok := v[key]; ok { - n, ok := raw.(json.Number) - if !ok { - return nil, errors.New("provider request not permitted") - } - value, err := n.Int64() - if err != nil || value < 1 { - return nil, errors.New("provider request not permitted") - } - if value > maxTokens { - v[key] = maxTokens - } - } - } - if p.provider.StoreOff { - v["store"] = false - } - if p.provider.NoStore { - v["store"] = false - v["background"] = false - if _, ok := v["max_output_tokens"]; !ok { - v["max_output_tokens"] = maxTokens - } - } - - return json.Marshal(v) -} - -// shape describes a refused request by its keys and protocol type names, so a log shows which position the -// filter refused without the conversation text it carries. -func shape(v any) string { - switch v := v.(type) { - case map[string]any: - keys := slices.Sorted(maps.Keys(v)) - parts := make([]string, 0, len(keys)) - for _, k := range keys { - if s, ok := v[k].(string); ok && (k == "type" || k == "role" || k == "execution") { - parts = append(parts, k+"="+s) - continue - } - parts = append(parts, k+":"+shape(v[k])) - } - return "{" + strings.Join(parts, " ") + "}" - case []any: - var parts []string - seen := map[string]bool{} - for _, item := range v { - if s := shape(item); !seen[s] { - seen[s] = true - parts = append(parts, s) - } - } - return "[" + strings.Join(parts, " ") + "]" - default: - return fmt.Sprintf("%T", v) - } -} - -func truncate(s string, n int) string { - if len(s) <= n { - return s - } - - return s[:n] + "..." -} diff --git a/internal/modelproxy/modelproxy_test.go b/internal/modelproxy/modelproxy_test.go deleted file mode 100644 index 58c7a5a..0000000 --- a/internal/modelproxy/modelproxy_test.go +++ /dev/null @@ -1,616 +0,0 @@ -package modelproxy - -import ( - "bytes" - "context" - "encoding/json" - "io" - "log/slog" - "net/http" - "net/http/httptest" - "net/url" - "strings" - "testing" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - - "github.com/php/tamis/internal/budget" -) - -type upstreamCall struct { - method, path, apiKey, auth, cookie, accountID, body string -} - -// setup returns a proxy for p, pointed at a fake upstream that records each call. -func setup(t *testing.T, p Provider, status int) (*httptest.Server, *[]upstreamCall, string) { - t.Helper() - var calls []upstreamCall - up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - body, _ := io.ReadAll(r.Body) - calls = append(calls, upstreamCall{r.Method, r.URL.Path, r.Header.Get("X-Api-Key"), r.Header.Get("Authorization"), r.Header.Get("Cookie"), r.Header.Get("ChatGPT-Account-Id"), string(body)}) - if status == http.StatusFound { - w.Header().Set("Location", "https://evil.example/") - } - w.Header().Set("Set-Cookie", "session=1") - w.WriteHeader(status) - response := `{"ok":true}` - if strings.HasSuffix(r.URL.Path, "/v1/user-auth-credential/whoami") { - response = `{"email":"security@example.com","chatgpt_user_id":"user-secret","chatgpt_account_id":"account-secret","chatgpt_plan_type":"enterprise","chatgpt_account_is_fedramp":false}` - } - if strings.HasSuffix(r.URL.Path, "/wham/config/bundle") { - response = `{"config_toml":{"enterprise_managed":[{"id":"secret","name":"secret","contents":"additional_developer_instructions = \\"internal\\""}]}}` - } - io.WriteString(w, response) - })) - t.Cleanup(up.Close) - p.Upstream = up.URL - upstreams := map[string]string{} - for prefix, raw := range p.Upstreams { - u, _ := url.Parse(raw) - upstreams[prefix] = up.URL + u.Path - } - if p.Upstreams != nil { - p.Upstreams = upstreams - } - - mp, err := New(p, "real-key", "allowed") - require.NoError(t, err) - token, revoke := mp.Issue(t.Context()) - srv := httptest.NewServer(mp) - t.Cleanup(srv.Close) - t.Cleanup(func() { revoke() }) - - return srv, &calls, token -} - -func call(t *testing.T, srv *httptest.Server, method, path string, header http.Header) *http.Response { - t.Helper() - return callBody(t, srv, method, path, `{"model":"allowed"}`, header) -} - -func callBody(t *testing.T, srv *httptest.Server, method, path, body string, header http.Header) *http.Response { - t.Helper() - req, err := http.NewRequest(method, srv.URL+path, strings.NewReader(body)) - require.NoError(t, err) - req.Header = header - resp, err := http.DefaultClient.Do(req) - require.NoError(t, err) - t.Cleanup(func() { resp.Body.Close() }) - - return resp -} - -func withKey(name, value string) http.Header { - h := http.Header{} - h.Set(name, value) - h.Set("Cookie", "secret=1") - - return h -} - -func TestAnthropicSwapsTheToken(t *testing.T) { - srv, calls, token := setup(t, Anthropic, http.StatusOK) - - resp := call(t, srv, http.MethodPost, PathPrefix+"/v1/messages?beta=true", withKey("X-Api-Key", token)) - - require.Equal(t, http.StatusOK, resp.StatusCode) - require.Empty(t, resp.Header.Get("Set-Cookie")) - require.Len(t, *calls, 1) - got := (*calls)[0] - require.Equal(t, "/v1/messages", got.path) - require.Equal(t, "real-key", got.apiKey) - require.Empty(t, got.auth) - require.Empty(t, got.cookie) -} - -func TestOpenAISwapsTheBearerToken(t *testing.T) { - srv, calls, token := setup(t, OpenAI, http.StatusOK) - - body := `{"model":"allowed","input":[{"type":"message","id":"msg_current","role":"assistant","content":[{"type":"output_text","text":"ok"}]}]}` - resp := callBody(t, srv, http.MethodPost, PathPrefix+"/v1/responses", body, withKey("Authorization", "Bearer "+token)) - - require.Equal(t, http.StatusOK, resp.StatusCode) - require.Len(t, *calls, 1) - require.Equal(t, "Bearer real-key", (*calls)[0].auth) - require.Empty(t, (*calls)[0].apiKey) - require.Contains(t, (*calls)[0].body, `"store":false`, "unsafe OpenAI defaults were not overridden") - require.Contains(t, (*calls)[0].body, `"background":false`, "unsafe OpenAI defaults were not overridden") -} - -func TestChatGPTSwapsTheAccessTokenOnBothUpstreams(t *testing.T) { - srv, calls, token := setup(t, ChatGPT, http.StatusOK) - require.True(t, strings.HasPrefix(token, "at-tamis-"), "Codex only reads a token starting with at- as an access token: %q", token) - - whoami := call(t, srv, http.MethodGet, PathPrefix+"/auth/v1/user-auth-credential/whoami", withKey("Authorization", "Bearer "+token)) - identity, _ := io.ReadAll(whoami.Body) - backendKey := withKey("Authorization", "Bearer "+token) - backendKey.Set("ChatGPT-Account-Id", "attacker-controlled") - bundle := call(t, srv, http.MethodGet, PathPrefix+"/backend-api/wham/config/bundle", backendKey) - managed, _ := io.ReadAll(bundle.Body) - responses := call(t, srv, http.MethodPost, PathPrefix+"/backend-api/codex/responses", backendKey) - - require.Equal(t, http.StatusOK, whoami.StatusCode) - require.Equal(t, http.StatusOK, bundle.StatusCode) - require.Equal(t, http.StatusOK, responses.StatusCode) - require.Len(t, *calls, 3) - for i, want := range []string{"/api/accounts/v1/user-auth-credential/whoami", "/backend-api/wham/config/bundle", "/backend-api/codex/responses"} { - got := (*calls)[i] - assert.Equal(t, want, got.path) - assert.Equal(t, "Bearer real-key", got.auth, want) - assert.Empty(t, got.cookie, want) - } - assert.Empty(t, (*calls)[0].accountID, "workspace identity was not applied server-side only") - assert.Equal(t, "account-secret", (*calls)[1].accountID, "workspace identity was not applied server-side only") - assert.Equal(t, "account-secret", (*calls)[2].accountID, "workspace identity was not applied server-side only") - assert.NotContains(t, string(identity), "secret", "private workspace response reached the agent") - assert.NotContains(t, string(identity), "security@example.com", "private workspace response reached the agent") - assert.Contains(t, string(identity), `"chatgpt_account_id":"tamis"`, "private workspace response reached the agent") - assert.NotContains(t, string(managed), "internal", "private workspace response reached the agent") - assert.Contains(t, string(managed), `"enterprise_managed":[]`, "private workspace response reached the agent") - body := (*calls)[2].body - assert.Contains(t, body, `"store":false`, "the ChatGPT backend only accepts store set to false") - assert.NotContains(t, body, "background", "the ChatGPT backend only accepts store set to false") - assert.NotContains(t, body, "max_output_tokens", "the ChatGPT backend only accepts store set to false") -} - -func TestChatGPTRefusesAnIdentityWithoutAccount(t *testing.T) { - up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - io.WriteString(w, `{"email":"security@example.com","chatgpt_plan_type":"enterprise"}`) - })) - t.Cleanup(up.Close) - p := ChatGPT - p.Upstreams = map[string]string{"/auth": up.URL + "/api/accounts", "/backend-api": up.URL + "/backend-api"} - mp, err := New(p, "real-key", "allowed") - require.NoError(t, err) - token, revoke := mp.Issue(t.Context()) - t.Cleanup(func() { revoke() }) - srv := httptest.NewServer(mp) - t.Cleanup(srv.Close) - - resp := call(t, srv, http.MethodGet, PathPrefix+"/auth/v1/user-auth-credential/whoami", withKey("Authorization", "Bearer "+token)) - body, _ := io.ReadAll(resp.Body) - - require.Equal(t, http.StatusBadGateway, resp.StatusCode) - require.NotContains(t, string(body), "security@example.com") -} - -func TestChatGPTRefusesTheOtherBackendRoutes(t *testing.T) { - srv, calls, token := setup(t, ChatGPT, http.StatusOK) - - for _, tc := range []struct{ method, path string }{ - {http.MethodGet, "/backend-api/codex/models?client_version=0.163.0&extra=1"}, - {http.MethodGet, "/backend-api/codex/models?model=other"}, - {http.MethodPost, "/backend-api/codex/responses?client_version=0.163.0"}, - {http.MethodGet, "/backend-api/wham/accounts/check"}, - {http.MethodPost, "/backend-api/ps/mcp"}, - {http.MethodPost, "/auth/v1/user-auth-credential/whoami"}, - {http.MethodGet, "/backend-api/codex/responses"}, - {http.MethodPost, "/v1/responses"}, - } { - assert.Equal(t, http.StatusForbidden, call(t, srv, tc.method, PathPrefix+tc.path, withKey("Authorization", "Bearer "+token)).StatusCode, "%s %s", tc.method, tc.path) - } - require.Empty(t, *calls, "refused routes reached the upstream") -} - -func TestRefusals(t *testing.T) { - srv, calls, token := setup(t, Anthropic, http.StatusOK) - - for name, tc := range map[string]struct { - method, path string - header http.Header - status int - }{ - "no token": {http.MethodPost, PathPrefix + "/v1/messages", http.Header{}, http.StatusUnauthorized}, - "unknown token": {http.MethodPost, PathPrefix + "/v1/messages", withKey("X-Api-Key", "tamis-guess"), http.StatusUnauthorized}, - "real key": {http.MethodPost, PathPrefix + "/v1/messages", withKey("X-Api-Key", "real-key"), http.StatusUnauthorized}, - "other route": {http.MethodPost, PathPrefix + "/v1/files", withKey("X-Api-Key", token), http.StatusForbidden}, - "wrong method": {http.MethodDelete, PathPrefix + "/v1/messages", withKey("X-Api-Key", token), http.StatusForbidden}, - "traversal": {http.MethodPost, PathPrefix + "/v1/x/../messages", withKey("X-Api-Key", token), http.StatusForbidden}, - "encoded path": {http.MethodPost, PathPrefix + "/v1/messages%2Fbatches", withKey("X-Api-Key", token), http.StatusForbidden}, - "nested model": {http.MethodGet, PathPrefix + "/v1/models/a/b", withKey("X-Api-Key", token), http.StatusForbidden}, - "query": {http.MethodPost, PathPrefix + "/v1/messages?feature=remote", withKey("X-Api-Key", token), http.StatusForbidden}, - "outside prefix": {http.MethodPost, "/v1/messages", withKey("X-Api-Key", token), http.StatusNotFound}, - "openai route": {http.MethodPost, PathPrefix + "/v1/responses", withKey("X-Api-Key", token), http.StatusForbidden}, - } { - t.Run(name, func(t *testing.T) { - require.Equal(t, tc.status, call(t, srv, tc.method, tc.path, tc.header).StatusCode) - }) - } - require.Empty(t, *calls, "nothing should reach the upstream") -} - -func TestRevokedTokenIsRefused(t *testing.T) { - mp, err := New(Anthropic, "real-key", "allowed") - require.NoError(t, err) - token, revoke := mp.Issue(t.Context()) - srv := httptest.NewServer(mp) - t.Cleanup(srv.Close) - - revoke() - - require.Equal(t, http.StatusUnauthorized, call(t, srv, http.MethodPost, PathPrefix+"/v1/messages", withKey("X-Api-Key", token)).StatusCode) -} - -func TestUpstreamRedirectIsNotRelayed(t *testing.T) { - srv, _, token := setup(t, Anthropic, http.StatusFound) - - resp := call(t, srv, http.MethodPost, PathPrefix+"/v1/messages", withKey("X-Api-Key", token)) - - require.Equal(t, http.StatusBadGateway, resp.StatusCode) - require.Empty(t, resp.Header.Get("Location")) -} - -func TestNewRequiresAKey(t *testing.T) { - _, err := New(Anthropic, "", "allowed") - require.Error(t, err) -} - -func TestNewRequiresAModel(t *testing.T) { - _, err := New(Anthropic, "real-key", "") - require.Error(t, err) -} - -func TestProviderCapabilitiesAndOtherModelsAreRefused(t *testing.T) { - srv, calls, token := setup(t, OpenAI, http.StatusOK) - for name, body := range map[string]string{ - "other model": `{"model":"expensive"}`, - "remote MCP": `{"model":"allowed","tools":[{"type":"mcp","server_url":"https://evil.example"}]}`, - "remote file": `{"model":"allowed","input":[{"type":"input_file","file_url":"https://evil.example"}]}`, - "conversation": `{"model":"allowed","conversation":"conv_other"}`, - "item reference": `{"model":"allowed","input":[{"type":"item_reference","id":"item_other"}]}`, - "implicit item reference": `{"model":"allowed","input":[{"id":"item_other"}]}`, - "null item reference": `{"model":"allowed","input":[{"type":null,"id":"item_other"}]}`, - "multi-agent": `{"model":"allowed","multi_agent":{"enabled":true,"max_concurrent_subagents":100}}`, - "nested multi-agent item": `{"model":"allowed","input":[{"type":"multi_agent_call","name":"spawn_agent"}]}`, - "untyped tool": `{"model":"allowed","tools":[{"name":"remote"}]}`, - } { - t.Run(name, func(t *testing.T) { - require.Equal(t, http.StatusForbidden, callBody(t, srv, http.MethodPost, PathPrefix+"/v1/responses", body, withKey("Authorization", "Bearer "+token)).StatusCode) - }) - } - require.Empty(t, *calls, "nothing should reach the upstream") -} - -func TestAnthropicFallbacksAreRefused(t *testing.T) { - srv, calls, token := setup(t, Anthropic, http.StatusOK) - body := `{"model":"allowed","max_tokens":16000,"fallbacks":[{"model":"expensive","max_tokens":128000}]}` - header := withKey("X-Api-Key", token) - header.Set("Anthropic-Beta", "server-side-fallback-2026-07-01") - - require.Equal(t, http.StatusForbidden, callBody(t, srv, http.MethodPost, PathPrefix+"/v1/messages?beta=true", body, header).StatusCode) - require.Empty(t, *calls, "the fallback reached the upstream") -} - -func TestProviderTokenLimitIsClamped(t *testing.T) { - srv, calls, token := setup(t, OpenAI, http.StatusOK) - resp := callBody(t, srv, http.MethodPost, PathPrefix+"/v1/responses", `{"model":"allowed","max_output_tokens":16001}`, withKey("Authorization", "Bearer "+token)) - require.Equal(t, http.StatusOK, resp.StatusCode) - require.Len(t, *calls, 1) - require.Contains(t, (*calls)[0].body, `"max_output_tokens":16000`) -} - -func TestEveryAccessTokenRequestSpendsTheQuota(t *testing.T) { - srv, calls, token := setup(t, ChatGPT, http.StatusOK) - whoami := PathPrefix + "/auth/v1/user-auth-credential/whoami" - for range maxRequests { - call(t, srv, http.MethodGet, whoami, withKey("Authorization", "Bearer "+token)) - } - require.Equal(t, http.StatusUnauthorized, call(t, srv, http.MethodPost, PathPrefix+"/backend-api/codex/responses", withKey("Authorization", "Bearer "+token)).StatusCode, "request after token checks") - require.Equal(t, http.StatusUnauthorized, call(t, srv, http.MethodGet, whoami, withKey("Authorization", "Bearer at-tamis-guess")).StatusCode, "token check with an unknown token") - require.Len(t, *calls, maxRequests) -} - -func TestProviderRequestQuota(t *testing.T) { - srv, calls, token := setup(t, Anthropic, http.StatusOK) - for i := 0; i < maxRequests; i++ { - require.Equal(t, http.StatusOK, call(t, srv, http.MethodPost, PathPrefix+"/v1/messages", withKey("X-Api-Key", token)).StatusCode, "request %d", i+1) - } - require.Equal(t, http.StatusUnauthorized, call(t, srv, http.MethodPost, PathPrefix+"/v1/messages", withKey("X-Api-Key", token)).StatusCode, "request after quota") - require.Len(t, *calls, maxRequests) -} - -func TestNestedToolsAndUnknownProtocolCapabilitiesFailClosed(t *testing.T) { - for _, provider := range []Provider{OpenAI, ChatGPT} { - t.Run(provider.TokenPrefix, func(t *testing.T) { - for _, body := range []string{ - `{"model":"allowed","input":[{"type":"additional_tools","role":"developer","tools":[{"type":"shell","environment":{"type":"container_auto","network_policy":{"type":"allowlist","allowed_domains":["attacker.invalid"]}}}]}]}`, - `{"model":"allowed","input":[{"type":"tool_search_output","tools":[{"type":"namespace","name":"n","tools":[{"type":"shell"}]}]}]}`, - `{"model":"allowed","tools":[{"type":"namespace","name":"n","tools":[{"type":"mcp","server_url":"https://attacker.invalid"}]}]}`, - `{"model":"allowed","tools":[{"type":"tool_search","execution":"server"}]}`, - `{"model":"allowed","input":[{"type":"configuration_update","tools":[]}]}`, - `{"model":"allowed","input":[{"type":"future_provider_tool"}]}`, - `{"model":"allowed","future_provider_option":true}`, - `{"model":"allowed","background":true}`, - } { - srv, calls, token := setup(t, provider, http.StatusOK) - route := "/v1/responses" - if provider.PrivateWorkspace { - route = "/backend-api/codex/responses" - } - require.Equal(t, http.StatusForbidden, callBody(t, srv, http.MethodPost, PathPrefix+route, body, withKey("Authorization", "Bearer "+token)).StatusCode, "unsafe request reached upstream: %s", body) - require.Empty(t, *calls, "unsafe request reached upstream: %s", body) - } - }) - } -} - -func TestLocalToolConversationPasses(t *testing.T) { - srv, calls, token := setup(t, OpenAI, http.StatusOK) - body := `{"model":"allowed","tools":[{"type":"namespace","name":"functions","tools":[{"type":"function","name":"read","parameters":{"type":"object","properties":{"type":{"type":"string"}}}}]},{"type":"tool_search","execution":"client"}],"input":[{"role":"user","content":[{"type":"input_text","text":"report"}]},{"type":"reasoning","id":"rs_1","encrypted_content":"opaque","summary":[]},{"type":"function_call","call_id":"c","name":"read","arguments":"{}"},{"type":"function_call_output","call_id":"c","output":"source"},{"type":"additional_tools","role":"developer","tools":[{"type":"custom","name":"local","format":{"type":"text"}}]}],"include":["reasoning.encrypted_content"]}` - require.Equal(t, http.StatusOK, callBody(t, srv, http.MethodPost, PathPrefix+"/v1/responses", body, withKey("Authorization", "Bearer "+token)).StatusCode, "local tools rejected") - require.Len(t, *calls, 1, "local tools rejected") -} - -func TestTokensShareTheAdvisoryBudgetAndStopOnCancellation(t *testing.T) { - upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) })) - defer upstream.Close() - provider := OpenAI - provider.Upstream = upstream.URL - mp, err := New(provider, "key", "allowed") - require.NoError(t, err) - srv := httptest.NewServer(mp) - defer srv.Close() - ctx, cancel := context.WithCancel(budget.WithLimit(t.Context(), 1, 1<<20)) - defer cancel() - a, revokeA := mp.Issue(ctx) - defer revokeA() - b, revokeB := mp.Issue(ctx) - defer revokeB() - require.Equal(t, http.StatusOK, call(t, srv, http.MethodPost, PathPrefix+"/v1/responses", withKey("Authorization", "Bearer "+a)).StatusCode) - require.Equal(t, http.StatusTooManyRequests, call(t, srv, http.MethodPost, PathPrefix+"/v1/responses", withKey("Authorization", "Bearer "+b)).StatusCode) - cancel() - require.Equal(t, http.StatusUnauthorized, call(t, srv, http.MethodPost, PathPrefix+"/v1/responses", withKey("Authorization", "Bearer "+a)).StatusCode) -} - -func TestAnthropicLocalToolsAndContextEdits(t *testing.T) { - srv, calls, token := setup(t, Anthropic, http.StatusOK) - body := `{"model":"allowed","messages":[{"role":"user","content":[{"type":"text","text":"report"}]},{"role":"assistant","content":[{"type":"tool_use","id":"t","name":"Read","input":{"path":"a.c"}}]},{"role":"user","content":[{"type":"tool_result","tool_use_id":"t","content":[{"type":"text","text":"source"},{"type":"tool_reference","tool_name":"Read"}]}]}],"tools":[{"name":"Read","description":"read","defer_loading":true,"input_schema":{"type":"object"}}],"context_management":{"edits":[{"type":"clear_thinking_20251015","keep":"all"}]}}` - require.Equal(t, http.StatusOK, callBody(t, srv, http.MethodPost, PathPrefix+"/v1/messages", body, withKey("X-Api-Key", token)).StatusCode, "valid local Anthropic conversation rejected") - require.Len(t, *calls, 1, "valid local Anthropic conversation rejected") - body = strings.ReplaceAll(body, "clear_thinking_20251015", "future_remote_compaction") - require.Equal(t, http.StatusForbidden, callBody(t, srv, http.MethodPost, PathPrefix+"/v1/messages", body, withKey("X-Api-Key", token)).StatusCode, "unknown context capability reached Anthropic") - require.Len(t, *calls, 1, "unknown context capability reached Anthropic") -} - -func TestCapabilityHeadersAreNotForwarded(t *testing.T) { - var got http.Header - up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { got = r.Header.Clone(); w.WriteHeader(http.StatusOK) })) - defer up.Close() - provider := Anthropic - provider.Upstream = up.URL - mp, err := New(provider, "key", "allowed") - require.NoError(t, err) - token, revoke := mp.Issue(t.Context()) - defer revoke() - srv := httptest.NewServer(mp) - defer srv.Close() - h := withKey("X-Api-Key", token) - h.Set("Anthropic-Beta", "context-management-2025-06-27,future-hosted-tools-2026-01-01") - h.Set("OpenAI-Beta", "agents=enabled") - h.Set("X-Unknown-Provider-Capability", "enabled") - require.Equal(t, http.StatusOK, call(t, srv, http.MethodPost, PathPrefix+"/v1/messages", h).StatusCode) - require.Equal(t, "context-management-2025-06-27", got.Get("Anthropic-Beta"), "unreviewed capability header reached upstream") - require.Empty(t, got.Get("OpenAI-Beta"), "unreviewed capability header reached upstream") - require.Empty(t, got.Get("X-Unknown-Provider-Capability"), "unreviewed capability header reached upstream") -} - -// Each reviewed protocol position passes, and one unreviewed field or shape at that position is refused. -func TestRequestShapes(t *testing.T) { - for _, tc := range []struct { - name string - provider Provider - accepted map[string]string - refused map[string]string - }{ - { - name: "anthropic", - provider: Anthropic, - accepted: map[string]string{ - "thinking": `{"model":"allowed","messages":[{"role":"assistant","content":[{"type":"thinking","thinking":"t","signature":"s"},{"type":"redacted_thinking","data":"d"}]}]}`, - "effort": `{"model":"allowed","messages":[{"role":"user","content":"x","output_config":{"effort":"high"}}]}`, - "clear tool uses": `{"model":"allowed","context_management":{"edits":[{"type":"clear_tool_uses_20250919","keep":{"type":"tool_uses","value":3}}]}}`, - "system blocks": `{"model":"allowed","system":[{"type":"text","text":"s","cache_control":{"type":"ephemeral"}}]}`, - }, - refused: map[string]string{ - "tools object": `{"model":"allowed","tools":{}}`, - "tool string": `{"model":"allowed","tools":["Read"]}`, - "hosted tool": `{"model":"allowed","tools":[{"type":"web_search_20250305","name":"web_search"}]}`, - "tool field": `{"model":"allowed","tools":[{"name":"Read","input_schema":{},"mcp_server":"s"}]}`, - "messages object": `{"model":"allowed","messages":{}}`, - "message field": `{"model":"allowed","messages":[{"role":"user","content":"x","container":"c"}]}`, - "output config string": `{"model":"allowed","messages":[{"role":"user","content":"x","output_config":"high"}]}`, - "output config field": `{"model":"allowed","messages":[{"role":"user","content":"x","output_config":{"effort":"high","task_budget":1}}]}`, - "system object": `{"model":"allowed","system":{}}`, - "system string block": `{"model":"allowed","system":["s"]}`, - "text field": `{"model":"allowed","system":[{"type":"text","text":"s","citations":[]}]}`, - "thinking field": `{"model":"allowed","messages":[{"role":"assistant","content":[{"type":"thinking","thinking":"t","signature":"s","source":"x"}]}]}`, - "tool use field": `{"model":"allowed","messages":[{"role":"assistant","content":[{"type":"tool_use","id":"t","name":"Read","input":{},"caller":{}}]}]}`, - "tool result field": `{"model":"allowed","messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"t","container":"c"}]}]}`, - "tool reference field": `{"model":"allowed","messages":[{"role":"user","content":[{"type":"tool_reference","tool_name":"Read","server":"s"}]}]}`, - "server tool block": `{"model":"allowed","messages":[{"role":"assistant","content":[{"type":"server_tool_use","id":"s","name":"web_search","input":{}}]}]}`, - "context management list": `{"model":"allowed","context_management":[]}`, - "context management field": `{"model":"allowed","context_management":{"edits":[],"compaction":true}}`, - "edits object": `{"model":"allowed","context_management":{"edits":{}}}`, - "edit string": `{"model":"allowed","context_management":{"edits":["clear_thinking_20251015"]}}`, - "clear thinking field": `{"model":"allowed","context_management":{"edits":[{"type":"clear_thinking_20251015","keep":"all","trigger":{}}]}}`, - "clear tool uses field": `{"model":"allowed","context_management":{"edits":[{"type":"clear_tool_uses_20250919","server":"s"}]}}`, - "token limit string": `{"model":"allowed","max_tokens":"16000"}`, - "token limit zero": `{"model":"allowed","max_tokens":0}`, - }, - }, - { - name: "openai", - provider: OpenAI, - accepted: map[string]string{ - "string input": `{"model":"allowed","input":"report"}`, - "local tools": `{"model":"allowed","tools":[{"type":"local_shell"},{"type":"apply_patch"},{"type":"custom","name":"c","format":{"type":"text"}}]}`, - "named tool choice": `{"model":"allowed","tool_choice":{"type":"function","name":"read"}}`, - "client search": `{"model":"allowed","input":[{"type":"tool_search_call","call_id":"c","arguments":{},"execution":"client"},{"type":"custom_tool_call_output","call_id":"c","output":[{"type":"input_text","text":"x"}]}]}`, - "assistant output": `{"model":"allowed","input":[{"type":"message","id":"m","role":"assistant","phase":"final_answer","status":"completed","content":[{"type":"output_text","text":"x","annotations":[],"logprobs":[]}]}]}`, - }, - refused: map[string]string{ - "include string": `{"model":"allowed","include":"reasoning.encrypted_content"}`, - "include other": `{"model":"allowed","include":["file_search_call.results"]}`, - "function field": `{"model":"allowed","tools":[{"type":"function","name":"f","server_url":"https://attacker.invalid"}]}`, - "custom field": `{"model":"allowed","tools":[{"type":"custom","name":"c","server_url":"https://attacker.invalid"}]}`, - "local shell field": `{"model":"allowed","tools":[{"type":"local_shell","environment":{}}]}`, - "client search field": `{"model":"allowed","tools":[{"type":"tool_search","execution":"client","server_url":"https://attacker.invalid"}]}`, - "tool choice string": `{"model":"allowed","tool_choice":"web_search"}`, - "tool choice hosted": `{"model":"allowed","tool_choice":{"type":"mcp","server_label":"s"}}`, - "tool choice field": `{"model":"allowed","tool_choice":{"type":"function","name":"f","server":"s"}}`, - "input number": `{"model":"allowed","input":1}`, - "input string item": `{"model":"allowed","input":["report"]}`, - "message field": `{"model":"allowed","input":[{"role":"user","content":"x","file_url":"https://attacker.invalid"}]}`, - "message image": `{"model":"allowed","input":[{"role":"user","content":[{"type":"input_image","image_url":"https://attacker.invalid"}]}]}`, - "output text field": `{"model":"allowed","input":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"x","file_url":"https://attacker.invalid"}]}]}`, - "reasoning field": `{"model":"allowed","input":[{"type":"reasoning","summary":[],"server":"s"}]}`, - "function call field": `{"model":"allowed","input":[{"type":"function_call","call_id":"c","name":"f","arguments":"{}","server":"s"}]}`, - "call output field": `{"model":"allowed","input":[{"type":"function_call_output","call_id":"c","output":"x","server":"s"}]}`, - "additional tools field": `{"model":"allowed","input":[{"type":"additional_tools","tools":[],"server":"s"}]}`, - "server tool search call": `{"model":"allowed","input":[{"type":"tool_search_call","call_id":"c","arguments":{},"execution":"server"}]}`, - "tool search call field": `{"model":"allowed","input":[{"type":"tool_search_call","call_id":"c","execution":"client","server":"s"}]}`, - }, - }, - } { - proxy, err := New(tc.provider, "key", "allowed") - require.NoError(t, err) - for name, body := range tc.accepted { - t.Run(tc.name+"/accepted/"+name, func(t *testing.T) { - _, err := proxy.safeBody(strings.NewReader(body)) - require.NoError(t, err) - }) - } - for name, body := range tc.refused { - t.Run(tc.name+"/refused/"+name, func(t *testing.T) { - _, err := proxy.safeBody(strings.NewReader(body)) - require.EqualError(t, err, "provider request not permitted") - }) - } - } -} - -func TestMalformedRequestsAreRefused(t *testing.T) { - proxy, err := New(OpenAI, "key", "allowed") - require.NoError(t, err) - for name, body := range map[string]string{ - "not JSON": `model=allowed`, - "trailing data": `{"model":"allowed"} {"model":"allowed","tools":[{"type":"mcp"}]}`, - } { - t.Run(name, func(t *testing.T) { - _, err := proxy.safeBody(strings.NewReader(body)) - require.EqualError(t, err, "invalid provider request") - }) - } - _, err = proxy.safeBody(strings.NewReader(`{"model":"allowed","input":"` + strings.Repeat("x", maxBody) + `"}`)) - require.ErrorIs(t, err, errRequestTooLarge) -} - -func TestOversizedRequestIsRefusedBeforeTheUpstream(t *testing.T) { - srv, calls, token := setup(t, OpenAI, http.StatusOK) - body := `{"model":"allowed","input":"` + strings.Repeat("x", maxBody) + `"}` - require.Equal(t, http.StatusRequestEntityTooLarge, callBody(t, srv, http.MethodPost, PathPrefix+"/v1/responses", body, withKey("Authorization", "Bearer "+token)).StatusCode) - require.Empty(t, *calls) -} - -func TestMCPBetaHeaderIsRefused(t *testing.T) { - srv, calls, token := setup(t, Anthropic, http.StatusOK) - header := withKey("X-Api-Key", token) - header.Set("Anthropic-Beta", "mcp-client-2025-04-04") - require.Equal(t, http.StatusForbidden, call(t, srv, http.MethodPost, PathPrefix+"/v1/messages", header).StatusCode) - require.Empty(t, *calls) -} - -func captureLogs(t *testing.T) *bytes.Buffer { - t.Helper() - var logs bytes.Buffer - previous := slog.Default() - slog.SetDefault(slog.New(slog.NewTextHandler(&logs, nil))) - t.Cleanup(func() { slog.SetDefault(previous) }) - - return &logs -} - -func TestRefusalsAreLoggedWithoutTheConversation(t *testing.T) { - logs := captureLogs(t) - srv, _, token := setup(t, OpenAI, http.StatusOK) - header := withKey("Authorization", "Bearer "+token) - - require.Equal(t, http.StatusOK, callBody(t, srv, http.MethodPost, PathPrefix+"/v1/responses", `{"model":"allowed","input":"SECRET-PAYLOAD"}`, header).StatusCode) - require.Empty(t, logs.String(), "a permitted request was logged") - - require.Equal(t, http.StatusForbidden, callBody(t, srv, http.MethodPost, PathPrefix+"/v1/responses", `{"model":"allowed","input":[{"role":"user","content":[{"type":"input_text","text":"SECRET-PAYLOAD","future_field":true}]}]}`, header).StatusCode) - require.Equal(t, http.StatusForbidden, call(t, srv, http.MethodGet, PathPrefix+"/v1/models", header).StatusCode) - got := logs.String() - assert.Contains(t, got, `msg="model proxy refused a request" route=/v1/responses shape="{input:[{content:[{future_field:bool text:string type=input_text}] role=user}] model:string}"`) - assert.Contains(t, got, `msg="model proxy refused a route" method=GET route=/v1/models`) - assert.NotContains(t, got, "SECRET-PAYLOAD") -} - -func TestShapeListsEachItemShapeOnceInOrder(t *testing.T) { - var v any - require.NoError(t, json.Unmarshal([]byte(`[{"type":"b"},{"type":"a"},{"type":"b"},1,{"type":"a"}]`), &v)) - - require.Equal(t, "[{type=b} {type=a} float64]", shape(v)) -} - -func TestProviderRefusalsAreLoggedWithoutTheirBody(t *testing.T) { - logs := captureLogs(t) - srv, _, token := setup(t, OpenAI, http.StatusBadRequest) - - require.Equal(t, http.StatusBadRequest, call(t, srv, http.MethodPost, PathPrefix+"/v1/responses", withKey("Authorization", "Bearer "+token)).StatusCode) - assert.Contains(t, logs.String(), `msg="model provider refused a request" method=POST route=/v1/responses status=400`) - assert.NotContains(t, logs.String(), `"ok"`) - assert.NotContains(t, logs.String(), `\"ok\"`) -} - -func FuzzProviderRequest(f *testing.F) { - for _, seed := range []string{`{"model":"allowed"}`, `{"model":"allowed","tools":[{"type":"function","name":"local","parameters":{}}]}`, `{"model":"allowed","input":[{"type":"additional_tools","tools":[{"type":"shell"}]}]}`, `{"model":"allowed","tools":[{"type":{}}]}`} { - f.Add(seed) - } - f.Fuzz(func(t *testing.T, body string) { - for _, provider := range []Provider{Anthropic, OpenAI, ChatGPT} { - proxy, err := New(provider, "fixture", "allowed") - require.NoError(t, err) - _, _ = proxy.safeBody(strings.NewReader(body)) - } - }) -} - -func TestChatGPTServesTheModelCatalogue(t *testing.T) { - srv, calls, token := setup(t, ChatGPT, http.StatusOK) - - resp := call(t, srv, http.MethodGet, PathPrefix+"/backend-api/codex/models?client_version=0.163.0", withKey("Authorization", "Bearer "+token)) - - require.Equal(t, http.StatusOK, resp.StatusCode) - require.Len(t, *calls, 1) - require.Equal(t, "/backend-api/codex/models", (*calls)[0].path) -} - -func TestExpectedRefusalsAreNotLogged(t *testing.T) { - logs := captureLogs(t) - srv, _, token := setup(t, ChatGPT, http.StatusOK) - header := withKey("Authorization", "Bearer "+token) - - require.Equal(t, http.StatusForbidden, call(t, srv, http.MethodGet, PathPrefix+"/backend-api/wham/accounts/check", header).StatusCode) - require.Equal(t, http.StatusForbidden, call(t, srv, http.MethodGet, PathPrefix+"/backend-api/wham/settings/user", header).StatusCode) - require.Empty(t, logs.String(), "an expected refusal was logged") - - require.Equal(t, http.StatusForbidden, call(t, srv, http.MethodPost, PathPrefix+"/backend-api/wham/accounts/check", header).StatusCode) - require.Contains(t, logs.String(), `msg="model proxy refused a route" method=POST route=/backend-api/wham/accounts/check`) -} - -func TestExpectedRefusalsDoNotSpendTheQuota(t *testing.T) { - srv, calls, token := setup(t, ChatGPT, http.StatusOK) - header := withKey("Authorization", "Bearer "+token) - for range maxRequests + 1 { - require.Equal(t, http.StatusForbidden, call(t, srv, http.MethodGet, PathPrefix+"/backend-api/wham/accounts/check", header).StatusCode) - } - require.Equal(t, http.StatusOK, call(t, srv, http.MethodPost, PathPrefix+"/backend-api/codex/responses", header).StatusCode, "expected refusals spent the quota") - - for range maxRequests { - call(t, srv, http.MethodGet, PathPrefix+"/backend-api/ps/mcp", header) - } - require.Equal(t, http.StatusUnauthorized, call(t, srv, http.MethodPost, PathPrefix+"/backend-api/codex/responses", header).StatusCode, "unexpected refusals did not spend the quota") - require.Len(t, *calls, 1) -} diff --git a/internal/modelproxy/schema.go b/internal/modelproxy/schema.go deleted file mode 100644 index a9f6708..0000000 --- a/internal/modelproxy/schema.go +++ /dev/null @@ -1,248 +0,0 @@ -package modelproxy - -import "slices" - -// Validate protocol positions, not arbitrary JSON in function schemas or tool arguments. -// A new provider capability must be reviewed before it can cross this boundary. -func safeRequest(v map[string]any, anthropic bool) bool { - if anthropic { - return fields(v, "model", "messages", "system", "tools", "tool_choice", "max_tokens", "stream", "temperature", "top_p", "top_k", "stop_sequences", "thinking", "output_config", "metadata", "cache_control", "context_management") && - safeMessages(v["messages"]) && safeContent(v["system"], true) && safeToolList(v["tools"], true) && safeContextEdits(v["context_management"]) - } - if !fields(v, "model", "input", "instructions", "tools", "tool_choice", "parallel_tool_calls", "reasoning", "store", "stream", "stream_options", "include", "text", "max_output_tokens", "max_tokens", "metadata", "client_metadata", "prompt_cache_key", "service_tier", "background") { - return false - } - if background, ok := v["background"]; ok && background != false { - return false - } - if include, ok := v["include"]; ok { - values, ok := include.([]any) - if !ok { - return false - } - for _, value := range values { - if value != "reasoning.encrypted_content" { - return false - } - } - } - return safeItems(v["input"]) && safeToolList(v["tools"], false) && safeToolChoice(v["tool_choice"]) -} - -func fields(v map[string]any, allowed ...string) bool { - for key := range v { - if !slices.Contains(allowed, key) { - return false - } - } - return true -} - -func safeToolList(v any, anthropic bool) bool { - if v == nil { - return true - } - list, ok := v.([]any) - if !ok { - return false - } - for _, item := range list { - t, ok := item.(map[string]any) - if !ok { - return false - } - if anthropic { - if !fields(t, "name", "description", "input_schema", "cache_control", "type", "strict", "defer_loading") || (t["type"] != nil && t["type"] != "custom") { - return false - } - continue - } - switch t["type"] { - case "function": - if !fields(t, "type", "name", "description", "parameters", "strict", "defer_loading") { - return false - } - case "custom": - if !fields(t, "type", "name", "description", "format", "defer_loading") { - return false - } - case "namespace": - if !fields(t, "type", "name", "description", "tools") || !safeToolList(t["tools"], false) { - return false - } - case "local_shell", "apply_patch": - if !fields(t, "type") { - return false - } - case "tool_search": - if t["execution"] != "client" || !fields(t, "type", "execution", "description", "parameters") { - return false - } - default: - return false - } - } - return true -} - -func safeToolChoice(v any) bool { - if v == nil || v == "auto" || v == "none" || v == "required" { - return true - } - t, ok := v.(map[string]any) - return ok && (t["type"] == "function" || t["type"] == "custom") && fields(t, "type", "name") -} - -func safeItems(v any) bool { - if v == nil { - return true - } - if _, ok := v.(string); ok { - return true - } - items, ok := v.([]any) - if !ok { - return false - } - for _, item := range items { - t, ok := item.(map[string]any) - if !ok { - return false - } - switch t["type"] { - case nil, "message": - if (t["type"] == nil && t["id"] != nil) || !fields(t, "type", "id", "role", "content", "status", "phase") || !safeContent(t["content"], false) { - return false - } - case "reasoning": - if !fields(t, "type", "id", "summary", "content", "encrypted_content", "status") { - return false - } - case "function_call", "custom_tool_call": - if !fields(t, "type", "id", "call_id", "name", "arguments", "input", "status", "namespace") { - return false - } - case "function_call_output", "custom_tool_call_output": - if !fields(t, "type", "id", "call_id", "output", "status") || !safeContent(t["output"], false) { - return false - } - case "additional_tools", "tool_search_output": - if !fields(t, "type", "id", "call_id", "role", "tools", "status", "execution") || !safeToolList(t["tools"], false) { - return false - } - case "tool_search_call": - if t["execution"] != "client" || !fields(t, "type", "id", "call_id", "arguments", "status", "execution") { - return false - } - default: - return false - } - } - return true -} - -func safeMessages(v any) bool { - if v == nil { - return true - } - items, ok := v.([]any) - if !ok { - return false - } - for _, item := range items { - m, ok := item.(map[string]any) - if !ok || !fields(m, "role", "content", "output_config") || !safeContent(m["content"], true) { - return false - } - if output, ok := m["output_config"]; ok { - config, ok := output.(map[string]any) - if !ok || !fields(config, "effort") { - return false - } - } - } - return true -} - -func safeContent(v any, anthropic bool) bool { - if v == nil { - return true - } - if _, ok := v.(string); ok { - return true - } - list, ok := v.([]any) - if !ok { - return false - } - for _, item := range list { - t, ok := item.(map[string]any) - if !ok { - return false - } - if !anthropic { - if !slices.Contains([]any{"input_text", "output_text", "refusal"}, t["type"]) || !fields(t, "type", "text", "refusal", "annotations", "logprobs") { - return false - } - continue - } - switch t["type"] { - case "text": - if !fields(t, "type", "text", "cache_control") { - return false - } - case "thinking", "redacted_thinking": - if !fields(t, "type", "thinking", "signature", "data") { - return false - } - case "tool_use": - if !fields(t, "type", "id", "name", "input", "cache_control") { - return false - } - case "tool_result": - if !fields(t, "type", "tool_use_id", "content", "is_error", "cache_control") || !safeContent(t["content"], true) { - return false - } - case "tool_reference": - if !fields(t, "type", "tool_name") { - return false - } - default: - return false - } - } - return true -} - -func safeContextEdits(v any) bool { - if v == nil { - return true - } - m, ok := v.(map[string]any) - if !ok || !fields(m, "edits") { - return false - } - edits, ok := m["edits"].([]any) - if !ok { - return false - } - for _, edit := range edits { - e, ok := edit.(map[string]any) - if !ok { - return false - } - switch e["type"] { - case "clear_thinking_20251015": - if !fields(e, "type", "keep") { - return false - } - case "clear_tool_uses_20250919": - if !fields(e, "type", "trigger", "keep", "clear_at_least", "exclude_tools", "clear_tool_inputs") { - return false - } - default: - return false - } - } - return true -} diff --git a/internal/modelproxy/usage.go b/internal/modelproxy/usage.go deleted file mode 100644 index 8e8664b..0000000 --- a/internal/modelproxy/usage.go +++ /dev/null @@ -1,93 +0,0 @@ -package modelproxy - -import ( - "bytes" - "encoding/json" - "io" - "sync" - - "github.com/php/tamis/internal/usage" -) - -// usageReader reads the tokens a model response reports, as one JSON body or as server-sent events, while the -// response is relayed. Streamed counters are cumulative, so each one keeps its highest value. -type usageReader struct { - io.ReadCloser - stream bool - buf []byte - tokens usage.Usage - done func(usage.Usage) - once sync.Once -} - -func (u *usageReader) Read(b []byte) (int, error) { - n, err := u.ReadCloser.Read(b) - u.buf = append(u.buf, b[:n]...) - for u.stream { - line, rest, ok := bytes.Cut(u.buf, []byte("\n")) - if !ok { - break - } - u.observe(line) - u.buf = rest - } - // The usage of a response larger than any model answer is lost rather than buffered. - if len(u.buf) > maxBody { - u.buf = nil - } - if err != nil { - u.finish() - } - - return n, err -} - -func (u *usageReader) Close() error { - u.finish() - - return u.ReadCloser.Close() -} - -func (u *usageReader) finish() { - u.once.Do(func() { - u.observe(u.buf) - u.tokens.Requests = 1 - u.done(u.tokens) - }) -} - -type reported struct { - InputTokens int64 `json:"input_tokens"` - CacheCreation int64 `json:"cache_creation_input_tokens"` - CacheRead int64 `json:"cache_read_input_tokens"` - OutputTokens int64 `json:"output_tokens"` -} - -// observe looks where Anthropic and the Responses API put the usage: at the top of a message or of a -// message_delta event, under message in message_start and under response in response.completed. -func (u *usageReader) observe(data []byte) { - if !bytes.Contains(data, []byte(`"usage"`)) { - return - } - data = bytes.TrimPrefix(bytes.TrimSpace(data), []byte("data:")) - var event struct { - Usage *reported `json:"usage"` - Message struct { - Usage *reported `json:"usage"` - } `json:"message"` - Response struct { - Usage *reported `json:"usage"` - } `json:"response"` - } - if json.Unmarshal(data, &event) != nil { - return - } - for _, r := range []*reported{event.Usage, event.Message.Usage, event.Response.Usage} { - if r == nil { - continue - } - // Anthropic leaves cached tokens out of input_tokens, OpenAI counts them in. - u.tokens.InputTokens = max(u.tokens.InputTokens, r.InputTokens+r.CacheCreation+r.CacheRead) - u.tokens.OutputTokens = max(u.tokens.OutputTokens, r.OutputTokens) - } -} diff --git a/internal/modelproxy/usage_test.go b/internal/modelproxy/usage_test.go deleted file mode 100644 index 107486b..0000000 --- a/internal/modelproxy/usage_test.go +++ /dev/null @@ -1,129 +0,0 @@ -package modelproxy - -import ( - "io" - "net/http" - "net/http/httptest" - "strings" - "testing" - - "github.com/stretchr/testify/require" - - "github.com/php/tamis/internal/usage" -) - -// spend relays one request to an upstream answering with body, and returns what the run spent once revoked. -func spend(t *testing.T, p Provider, route string, status int, contentType, body string) usage.Usage { - t.Helper() - up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - w.Header().Set("Content-Type", contentType) - w.WriteHeader(status) - io.WriteString(w, body) - })) - t.Cleanup(up.Close) - p.Upstream = up.URL - mp, err := New(p, "real-key", "allowed") - require.NoError(t, err) - token, revoke := mp.Issue(t.Context()) - srv := httptest.NewServer(mp) - t.Cleanup(srv.Close) - - resp := call(t, srv, http.MethodPost, PathPrefix+route, withKey("Authorization", "Bearer "+token)) - io.Copy(io.Discard, resp.Body) - resp.Body.Close() - - return revoke() -} - -func TestUsageOfAResponse(t *testing.T) { - for name, tc := range map[string]struct { - provider Provider - route string - contentType string - body string - want usage.Usage - }{ - "anthropic message": {Anthropic, "/v1/messages", "application/json", - `{"type":"message","usage":{"input_tokens":10,"cache_creation_input_tokens":20,"cache_read_input_tokens":30,"output_tokens":5}}`, - usage.Usage{Requests: 1, InputTokens: 60, OutputTokens: 5}}, - "anthropic stream": {Anthropic, "/v1/messages", "text/event-stream; charset=utf-8", - "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":10,\"cache_read_input_tokens\":30,\"output_tokens\":1}}}\n\n" + - "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"text\":\"usage\"}}\n\n" + - "event: message_delta\ndata: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":42}}\n\n", - usage.Usage{Requests: 1, InputTokens: 40, OutputTokens: 42}}, - "responses stream": {OpenAI, "/v1/responses", "text/event-stream", - "event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"usage\":null}}\n\n" + - "event: response.completed\ndata: {\"type\":\"response.completed\",\"response\":{\"usage\":{\"input_tokens\":100,\"output_tokens\":7}}}\n\n", - usage.Usage{Requests: 1, InputTokens: 100, OutputTokens: 7}}, - "stream without usage": {OpenAI, "/v1/responses", "text/event-stream", - "data: not json\n\ndata: {\"usage\":\"many\"}\n\ndata: [DONE]\n\n", - usage.Usage{Requests: 1}}, - "token count": {Anthropic, "/v1/messages/count_tokens", "application/json", `{"input_tokens":10}`, usage.Usage{}}, - } { - t.Run(name, func(t *testing.T) { - require.Equal(t, tc.want, spend(t, tc.provider, tc.route, http.StatusOK, tc.contentType, tc.body)) - }) - } -} - -func TestUsageReaderKeepsAnUnterminatedLastEvent(t *testing.T) { - var got usage.Usage - r := &usageReader{ - ReadCloser: io.NopCloser(strings.NewReader("data: {\"type\":\"response.completed\",\"response\":{\"usage\":{\"input_tokens\":3,\"output_tokens\":2}}}")), - stream: true, - done: func(u usage.Usage) { got = u }, - } - - io.Copy(io.Discard, r) - - require.Equal(t, usage.Usage{Requests: 1, InputTokens: 3, OutputTokens: 2}, got) -} - -func TestUsageReaderCountsAStreamClosedEarly(t *testing.T) { - var calls []usage.Usage - body := "data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":10,\"output_tokens\":1}}}\n\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"text\":\"partial" - r := &usageReader{ - ReadCloser: io.NopCloser(strings.NewReader(body)), - stream: true, - done: func(u usage.Usage) { calls = append(calls, u) }, - } - - io.ReadFull(r, make([]byte, len(body)-5)) - r.Close() - r.Close() - - require.Len(t, calls, 1) - require.Equal(t, usage.Usage{Requests: 1, InputTokens: 10, OutputTokens: 1}, calls[0]) -} - -func TestUsageSkipsFailedRequests(t *testing.T) { - got := spend(t, Anthropic, "/v1/messages", http.StatusTooManyRequests, "application/json", `{"usage":{"input_tokens":10,"output_tokens":5}}`) - - require.Zero(t, got, "a failed request was counted") -} - -func TestUsageIsCountedPerToken(t *testing.T) { - up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - w.Header().Set("Content-Type", "application/json") - io.WriteString(w, `{"usage":{"input_tokens":10,"output_tokens":5}}`) - })) - t.Cleanup(up.Close) - p := Anthropic - p.Upstream = up.URL - mp, err := New(p, "real-key", "allowed") - require.NoError(t, err) - srv := httptest.NewServer(mp) - t.Cleanup(srv.Close) - first, revokeFirst := mp.Issue(t.Context()) - _, revokeSecond := mp.Issue(t.Context()) - - for range 2 { - resp := call(t, srv, http.MethodPost, PathPrefix+"/v1/messages", withKey("X-Api-Key", first)) - io.Copy(io.Discard, resp.Body) - resp.Body.Close() - } - - require.Zero(t, revokeSecond(), "the second run was charged") - require.Equal(t, usage.Usage{Requests: 2, InputTokens: 20, OutputTokens: 10}, revokeFirst()) - require.Zero(t, revokeFirst(), "a revoked token still reports usage") -} diff --git a/internal/sandbox/runner.go b/internal/sandbox/runner.go index 831af6c..7d76881 100644 --- a/internal/sandbox/runner.go +++ b/internal/sandbox/runner.go @@ -3,183 +3,58 @@ package sandbox import ( "bufio" "bytes" - "context" - "crypto/rand" - "encoding/hex" "encoding/json" "errors" - "fmt" "io" - "net" - "os" - "os/exec" - "path/filepath" - "strconv" - "strings" - "time" "github.com/alpha-omega-security/harness" "github.com/alpha-omega-security/harness/container" - "github.com/alpha-omega-security/harness/egress" ) -const ( - proxyPort = "3128" - workMemory = "2g" - workCPUs = "2" - workPIDs = "256" - workFile = "67108864" - maxAgentOutput = 1 << 20 -) - -type isolatedRun struct { - runtime container.Runtime - image string - network string - proxy string - agent string - proxyIP string - token string - // env is set in the agent container on top of the harness environment. - env []string -} - -// startIsolatedRun starts the egress proxy, letting the agent reach the model proxy and hosts. harness -// refuses tunnels to its API port, so a model proxy served over TLS is reached as a separate host port. -func startIsolatedRun(ctx context.Context, runtime container.Runtime, image, modelPort string, tls bool, hosts []string) (*isolatedRun, error) { - if runtime.Bin != "docker" { - return nil, errors.New("only Docker is supported") - } - run := &isolatedRun{runtime: runtime, image: image} - var err error - if run.network, err = isolatedName(); err != nil { - return nil, err - } - run.proxy = "tamis-proxy-" + strings.TrimPrefix(run.network, "tamis-") - run.agent = "tamis-agent-" + strings.TrimPrefix(run.network, "tamis-") - run.token = egress.NewToken() - - if out, err := exec.CommandContext(ctx, runtime.Bin, "network", "create", "--driver", "bridge", "--internal", "--ipv6=false", "--opt", "com.docker.network.bridge.gateway_mode_ipv4=isolated", "--opt", "com.docker.network.bridge.gateway_mode_ipv6=isolated", "--", run.network).CombinedOutput(); err != nil { - return nil, fmt.Errorf("creating the isolated network: %w: %s", err, strings.TrimSpace(string(out))) - } - failed := true - defer func() { - if failed { - run.Close() - } - }() - - sidecar := container.SidecarConfig{ - Token: run.token, - Allow: append([]string{egress.HostGatewayAlias}, hosts...), - APIPort: modelPort, - GatewayIP: egress.HostGatewayAlias, - } - if tls { - sidecar.APIPort, sidecar.HostPorts = "", []string{modelPort} - } - proxyEnv := container.SidecarEnv(sidecar, egress.ListenFirstIface+":"+proxyPort) - args := []string{"run", "-d", "--name", run.proxy, "--network", run.network, "--cap-drop", "ALL", "--security-opt", "no-new-privileges", "--read-only", "--memory", "256m", "--cpus", "1", "--pids-limit", "128", "--ulimit", "fsize=16777216:16777216", "--log-driver", "json-file", "--log-opt", "max-size=1m", "--log-opt", "max-file=1", "--tmpfs", "/tmp:rw,noexec,nosuid,size=16m", "--add-host", egress.HostGatewayAlias + ":host-gateway", "-e", container.ProxyTokenEnv} - for _, entry := range proxyEnv[1:] { - args = append(args, "-e", entry) - } - args = append(args, "--", image, "harness-proxy", "--require-capability="+egress.CapabilityDenyAPIConnect) - cmd := exec.CommandContext(ctx, runtime.Bin, args...) - cmd.Env = withEnv(os.Environ(), container.ProxyTokenEnv, run.token) - if out, err := cmd.CombinedOutput(); err != nil { - return nil, fmt.Errorf("starting the egress proxy: %w: %s", err, strings.TrimSpace(string(out))) - } - if out, err := exec.CommandContext(ctx, runtime.Bin, "network", "connect", "--", "bridge", run.proxy).CombinedOutput(); err != nil { - return nil, fmt.Errorf("connecting the egress proxy: %w: %s", err, strings.TrimSpace(string(out))) - } - out, err := exec.CommandContext(ctx, runtime.Bin, "inspect", "--format", fmt.Sprintf(`{{(index .NetworkSettings.Networks %q).IPAddress}}`, run.network), "--", run.proxy).Output() - if err != nil { - return nil, errors.New("inspecting the egress proxy") - } - ip := strings.TrimSpace(string(out)) - if net.ParseIP(ip) == nil { - return nil, errors.New("the egress proxy has no private address") - } - run.proxyIP = ip - if err := run.waitReady(ctx); err != nil { - return nil, err - } - failed = false - - return run, nil -} - -func isolatedName() (string, error) { - var raw [8]byte - if _, err := rand.Read(raw[:]); err != nil { - return "", err - } +const maxAgentOutput = 1 << 20 - return "tamis-" + strconv.Itoa(os.Getpid()) + "-" + hex.EncodeToString(raw[:]), nil +// agent keeps the failure events of a run's output, which otherwise holds the report and the agent's +// reasoning about it. +type agent struct { + harness.Harness + failures boundedWriter } -func (r *isolatedRun) waitReady(ctx context.Context) error { - deadline := time.NewTimer(30 * time.Second) - defer deadline.Stop() - for { - out, err := exec.CommandContext(ctx, r.runtime.Bin, "logs", "--", r.proxy).CombinedOutput() - logs := strings.TrimSpace(string(out)) - if err == nil && strings.Contains(logs, "egress proxy listening") { - return nil - } - select { - case <-ctx.Done(): - return ctx.Err() - case <-deadline.C: - return errors.New("the egress proxy did not become ready") - case <-time.After(200 * time.Millisecond): - } +// EgressHosts is empty so the egress proxy allows the model proxy and the request's hosts only, not the +// provider's own API. +func (*agent) EgressHosts() []string { return nil } + +func (a *agent) ParseStream(r io.Reader, _ func(harness.Event)) { keepFailures(r, &a.failures) } + +// runner keeps the agent on a network whose only way out is the egress proxy, allowing the model proxy and +// hosts. harness refuses tunnels to its API port, so a model proxy served over TLS is reached as a separate +// host port. +func (r *Reader) runner(state, source, token string, hosts, env []string) container.Runner { + sidecar := container.SidecarConfig{Allow: hosts, APIPort: r.modelPort} + if r.backend.accessToken { + sidecar.APIPort, sidecar.HostPorts = "", []string{r.modelPort} + } + + return container.Runner{ + Runtime: r.runtime, + Image: r.image, + StateDir: state, + Mounts: []container.Mount{{Host: source, Container: container.WorkMount + "/" + srcDir, ReadOnly: true}}, + // Codex's SQLite state crashes it with SIGBUS on a state directory bind-mounted by Docker Desktop. + Env: append(env, r.backend.keyEnv, "CODEX_SQLITE_HOME=/tmp/codex-sqlite"), + OmitEnv: []string{r.backend.keyEnv, "CLAUDE_CODE_OAUTH_TOKEN", "CODEX_API_KEY"}, + ProcessEnv: []string{r.backend.keyEnv + "=" + token}, + Hardened: true, + IsolateHost: true, + Sidecar: sidecar, + ReadOnlyWorkspace: true, + Memory: "2g", + CPUs: "2", + PIDsLimit: 256, + FileSizeLimit: 64 << 20, } } -func (r *isolatedRun) Run(ctx context.Context, h harness.Harness, job harness.Job, workspace, source, state, keyEnv, key string) error { - if err := harness.WriteSystemPrompt(h, job); err != nil { - return err - } - reportPath := filepath.Join(workspace, outputFile) - if err := os.WriteFile(reportPath, nil, 0o600); err != nil { - return err - } - job.Workspace = container.WorkMount - args := r.agentArgs(h, job, workspace, source, state, reportPath, keyEnv) - - cmd := exec.CommandContext(ctx, r.runtime.Bin, args...) - cmd.Env = withEnv(os.Environ(), keyEnv, key) - for _, env := range []struct{ key, value string }{ - {"HTTPS_PROXY", egress.EndpointURL(r.token, net.JoinHostPort(r.proxyIP, proxyPort))}, - {"HTTP_PROXY", egress.EndpointURL(r.token, net.JoinHostPort(r.proxyIP, proxyPort))}, - {"ALL_PROXY", egress.EndpointURL(r.token, net.JoinHostPort(r.proxyIP, proxyPort))}, - {"NO_PROXY", ""}, - {"https_proxy", egress.EndpointURL(r.token, net.JoinHostPort(r.proxyIP, proxyPort))}, - {"http_proxy", egress.EndpointURL(r.token, net.JoinHostPort(r.proxyIP, proxyPort))}, - {"all_proxy", egress.EndpointURL(r.token, net.JoinHostPort(r.proxyIP, proxyPort))}, - {"no_proxy", ""}, - } { - cmd.Env = withEnv(cmd.Env, env.key, env.value) - } - var stderr, failures boundedWriter - cmd.Stderr = &stderr - stdout, err := cmd.StdoutPipe() - if err != nil { - return err - } - if err := cmd.Start(); err != nil { - return err - } - keepFailures(stdout, &failures) - err = cmd.Wait() - if out := strings.TrimSpace(stderr.String() + "\n" + failures.String()); err != nil && out != "" { - return fmt.Errorf("%w: %s", err, out) - } - - return err -} - // keepFailures copies the error events of a JSONL agent output and discards the rest, lines longer than // maxAgentOutput included. Codex exec --json // reports why it stopped only there, as "error" and "turn.failed" events or "error" items. @@ -207,41 +82,6 @@ func keepFailures(r io.Reader, w io.Writer) { } } -func (r *isolatedRun) agentArgs(h harness.Harness, job harness.Job, workspace, source, state, reportPath, keyEnv string) []string { - args := []string{ - "run", "--rm", "--name", r.agent, "--log-driver", "none", "--cap-drop", "ALL", "--security-opt", "no-new-privileges", "--read-only", - "--memory", workMemory, "--cpus", workCPUs, "--pids-limit", workPIDs, "--ulimit", "fsize=" + workFile + ":" + workFile, - "--user", strconv.Itoa(os.Getuid()) + ":" + strconv.Itoa(os.Getgid()), "-e", "HOME=/tmp", "--tmpfs", "/tmp:rw,exec,nosuid,size=256m", - "--mount", "type=bind,src=" + workspace + ",dst=" + container.WorkMount + ",readonly", - "--mount", "type=bind,src=" + source + ",dst=" + container.WorkMount + "/" + srcDir + ",readonly", - "--mount", "type=bind,src=" + reportPath + ",dst=" + container.WorkMount + "/" + outputFile, - "--tmpfs", container.StateMount + ":rw,noexec,nosuid,size=256m", - "-w", container.WorkMount, "--network", r.network, "--add-host", egress.HostGatewayAlias + ":" + r.proxyIP, - } - for _, name := range []string{"config.toml", caFile} { - if _, err := os.Stat(filepath.Join(state, name)); err == nil { - args = append(args, "--mount", "type=bind,src="+filepath.Join(state, name)+",dst="+container.StateMount+"/"+name+",readonly") - } - } - for _, entry := range r.env { - args = append(args, "-e", entry) - } - for _, entry := range h.StateEnv(container.StateMount) { - args = append(args, "-e", entry) - } - for _, entry := range h.Env(job.BaseURL) { - name, _, _ := strings.Cut(entry, "=") - if name != keyEnv && name != "CLAUDE_CODE_OAUTH_TOKEN" && name != "CODEX_API_KEY" { - args = append(args, "-e", entry) - } - } - args = append(args, "-e", keyEnv, "-e", "HTTPS_PROXY", "-e", "HTTP_PROXY", "-e", "ALL_PROXY", "-e", "NO_PROXY", "-e", "https_proxy", "-e", "http_proxy", "-e", "all_proxy", "-e", "no_proxy") - args = append(args, "--", r.image, h.Binary()) - args = append(args, h.Args(job)...) - - return args -} - type boundedWriter struct { b bytes.Buffer } @@ -255,29 +95,3 @@ func (w *boundedWriter) Write(p []byte) (int, error) { } func (w *boundedWriter) String() string { return w.b.String() } - -func withEnv(env []string, key, value string) []string { - prefix := key + "=" - filtered := env[:0] - for _, entry := range env { - if !strings.HasPrefix(entry, prefix) { - filtered = append(filtered, entry) - } - } - - return append(filtered, prefix+value) -} - -func (r *isolatedRun) Close() { - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) - defer cancel() - if r.agent != "" { - _ = exec.CommandContext(ctx, r.runtime.Bin, "rm", "-f", "--", r.agent).Run() - } - if r.proxy != "" { - _ = exec.CommandContext(ctx, r.runtime.Bin, "rm", "-f", "--", r.proxy).Run() - } - if r.network != "" { - _ = exec.CommandContext(ctx, r.runtime.Bin, "network", "rm", "--", r.network).Run() - } -} diff --git a/internal/sandbox/runner_integration_test.go b/internal/sandbox/runner_integration_test.go index 13dec72..4dda88f 100644 --- a/internal/sandbox/runner_integration_test.go +++ b/internal/sandbox/runner_integration_test.go @@ -14,19 +14,22 @@ import ( "net/http/httptest" "os" "os/exec" + "path/filepath" "strconv" "strings" "sync/atomic" "testing" "time" + "github.com/alpha-omega-security/harness" "github.com/alpha-omega-security/harness/container" "github.com/alpha-omega-security/harness/egress" "github.com/stretchr/testify/require" ) func TestIsolatedNetworkReachesTheModelButNotOtherHostServices(t *testing.T) { - require.NotEmpty(t, os.Getenv("TAMIS_TEST_IMAGE"), "TAMIS_TEST_IMAGE is required") + image := os.Getenv("TAMIS_TEST_IMAGE") + require.NotEmpty(t, image, "TAMIS_TEST_IMAGE is required") runtime, ok := container.DetectRuntime("") require.True(t, ok, "Docker is unavailable") model, modelPort := testServer(t, "model") @@ -34,22 +37,28 @@ func TestIsolatedNetworkReachesTheModelButNotOtherHostServices(t *testing.T) { t.Cleanup(func() { _ = model.Close() }) t.Cleanup(func() { _ = secret.Close() }) - ctx, cancel := context.WithTimeout(t.Context(), time.Minute) + ctx, cancel := context.WithTimeout(t.Context(), 2*time.Minute) defer cancel() - run, err := startIsolatedRun(ctx, runtime, os.Getenv("TAMIS_TEST_IMAGE"), modelPort, false, allow) + hostIP := container.ResolveHostGatewayIPv4(ctx, runtime, image, "") + require.NotEmpty(t, hostIP, "resolving the real host gateway") + h, err := harness.ByName("claude") + require.NoError(t, err) + reader := &Reader{backend: backends["claude"], harness: h, runtime: runtime, image: image, modelPort: modelPort} + env := []string{"MODEL_URL=http://" + egress.HostGatewayAlias + ":" + modelPort, "SECRET_URL=http://" + net.JoinHostPort(hostIP, secretPort), "PROXIED_SECRET_URL=http://" + egress.HostGatewayAlias + ":" + secretPort} + scope, err := reader.runner(t.TempDir(), t.TempDir(), "proxy-token", allow, env).Open(ctx, &agent{Harness: h}) + require.NoError(t, err) + t.Cleanup(func() { scope.Close(nil) }) + + out, err := exec.CommandContext(ctx, runtime.Bin, "network", "ls", "--filter", "name=harness-hardened-"+strconv.Itoa(os.Getpid())+"-", "--format", "{{.Name}}").Output() require.NoError(t, err) - t.Cleanup(run.Close) - out, err := exec.CommandContext(ctx, runtime.Bin, "exec", "--", run.proxy, "getent", "hosts", egress.HostGatewayAlias).Output() - require.NoError(t, err, "resolving the real host gateway") - require.NotEmpty(t, strings.Fields(string(out)), "resolving the real host gateway") - hostIP := strings.Fields(string(out))[0] - out, err = exec.CommandContext(ctx, runtime.Bin, "network", "inspect", "--format", `{{.Internal}} {{.EnableIPv6}} {{index .Options "com.docker.network.bridge.gateway_mode_ipv4"}} {{index .Options "com.docker.network.bridge.gateway_mode_ipv6"}}`, "--", run.network).Output() + networks := strings.Fields(string(out)) + require.Len(t, networks, 1) + out, err = exec.CommandContext(ctx, runtime.Bin, "network", "inspect", "--format", `{{.Internal}} {{.EnableIPv6}} {{index .Options "com.docker.network.bridge.gateway_mode_ipv4"}} {{index .Options "com.docker.network.bridge.gateway_mode_ipv6"}}`, "--", networks[0]).Output() require.NoError(t, err) require.Equal(t, "true false isolated isolated", strings.TrimSpace(string(out))) - proxyURL := egress.EndpointURL(run.token, net.JoinHostPort(run.proxyIP, proxyPort)) - cmd := exec.CommandContext(ctx, runtime.Bin, "run", "--rm", "--network", run.network, "--add-host", egress.HostGatewayAlias+":"+run.proxyIP, "-e", "http_proxy", "-e", "NO_PROXY", "-e", "MODEL_URL", "-e", "SECRET_URL", "-e", "PROXIED_SECRET_URL", "--", os.Getenv("TAMIS_TEST_IMAGE"), "sh", "-c", `test "$(curl -fsS --connect-timeout 5 "$MODEL_URL")" = model && ! curl -fsS --noproxy '*' --connect-timeout 2 "$SECRET_URL" && ! curl -fsS --connect-timeout 2 "$PROXIED_SECRET_URL" && ! curl -fsS --noproxy '*' --connect-timeout 2 'https://1.1.1.1' && ! curl -gfsS --noproxy '*' --connect-timeout 2 'https://[2606:4700:4700::1111]' && ! getent hosts example.com`) - cmd.Env = append(os.Environ(), "http_proxy="+proxyURL, "NO_PROXY=", "MODEL_URL=http://"+egress.HostGatewayAlias+":"+modelPort, "SECRET_URL=http://"+net.JoinHostPort(hostIP, secretPort), "PROXIED_SECRET_URL=http://"+egress.HostGatewayAlias+":"+secretPort) - out, err = cmd.CombinedOutput() + workspace := t.TempDir() + require.NoError(t, os.Mkdir(filepath.Join(workspace, srcDir), 0o700)) + out, err = scope.RunCommand(ctx, harness.Job{Workspace: workspace, OutputFile: outputFile}, container.Command{Args: []string{"sh", "-c", `export http_proxy="$HTTP_PROXY" && test "$(curl -fsS --connect-timeout 5 "$MODEL_URL")" = model && ! curl -fsS --noproxy '*' --connect-timeout 2 "$SECRET_URL" && ! curl -fsS --connect-timeout 2 "$PROXIED_SECRET_URL" && ! curl -fsS --noproxy '*' --connect-timeout 2 'https://1.1.1.1' && ! curl -gfsS --noproxy '*' --connect-timeout 2 'https://[2606:4700:4700::1111]' && ! getent hosts example.com`}}) require.NoError(t, err, "isolation probe: %s", strings.TrimSpace(string(out))) } diff --git a/internal/sandbox/sandbox.go b/internal/sandbox/sandbox.go index b802a07..1ceec96 100644 --- a/internal/sandbox/sandbox.go +++ b/internal/sandbox/sandbox.go @@ -10,6 +10,7 @@ import ( "encoding/json" "errors" "fmt" + "log/slog" "net" "net/http" "os" @@ -22,8 +23,9 @@ import ( "github.com/alpha-omega-security/harness" "github.com/alpha-omega-security/harness/container" "github.com/alpha-omega-security/harness/egress" + "github.com/alpha-omega-security/harness/modelproxy" - "github.com/php/tamis/internal/modelproxy" + "github.com/php/tamis/internal/budget" "github.com/php/tamis/internal/usage" ) @@ -119,7 +121,7 @@ func New(name, model, apiKey, image string) (*Reader, error) { image = Images[b.harness] } - models, err := modelproxy.New(b.provider, apiKey, model) + models, err := modelproxy.New(b.provider, apiKey, model, modelproxy.Options{Charge: budget.Take, Log: slog.Default()}) if err != nil { return nil, err } @@ -131,7 +133,7 @@ func New(name, model, apiKey, image string) (*Reader, error) { var ca []byte if b.accessToken { var cert tls.Certificate - if ca, cert, err = newCertificate(egress.HostGatewayAlias); err != nil { + if ca, cert, err = modelproxy.NewCertificate(egress.HostGatewayAlias); err != nil { ln.Close() return nil, err } @@ -181,21 +183,29 @@ func (r *Reader) Read(ctx context.Context, req Request) (json.RawMessage, error) } token, revoke := r.models.Issue(ctx) - defer func() { usage.Add(ctx, revoke()) }() + defer func() { + spent := revoke() + usage.Add(ctx, usage.Usage{Requests: int64(spent.Requests), InputTokens: int64(spent.InputTokens), OutputTokens: int64(spent.OutputTokens)}) + }() hosts := allow if req.Offline { hosts = nil } - run, err := startIsolatedRun(ctx, r.runtime, r.image, r.modelPort, r.backend.accessToken, hosts) - if err != nil { - return nil, err - } - defer run.Close() origin := "http://" + egress.HostGatewayAlias + ":" + r.modelPort + var env []string if r.backend.accessToken { origin = "https://" + egress.HostGatewayAlias + ":" + r.modelPort - run.env = []string{"CODEX_CA_CERTIFICATE=" + container.StateMount + "/" + caFile, "CODEX_AUTHAPI_BASE_URL=" + origin + modelproxy.PathPrefix + "/auth"} + env = []string{"CODEX_CA_CERTIFICATE=" + container.StateMount + "/" + caFile, "CODEX_AUTHAPI_BASE_URL=" + origin + modelproxy.PathPrefix + "/auth"} } + if err := r.writeConfig(state, origin); err != nil { + return nil, err + } + a := &agent{Harness: r.harness} + scope, err := r.runner(state, source, token, hosts, env).Open(ctx, a) + if err != nil { + return nil, err + } + defer scope.Close(nil) job := harness.Job{ Workspace: workspace, SrcDir: srcDir, @@ -206,14 +216,14 @@ func (r *Reader) Read(ctx context.Context, req Request) (json.RawMessage, error) OutputFile: outputFile, BaseURL: origin + r.backend.basePath, } - if err := r.writeConfig(state, origin); err != nil { - return nil, err - } for attempt := 0; attempt < 2; attempt++ { - // Agent output is discarded: it holds the report and the agent's reasoning about it. - err := run.Run(ctx, r.harness, job, workspace, source, state, r.backend.keyEnv, token) + a.failures = boundedWriter{} + err := scope.Run(ctx, job, nil) if err != nil { + if failures := strings.TrimSpace(a.failures.String()); failures != "" { + err = fmt.Errorf("%w: %s", err, failures) + } return salvage(filepath.Join(workspace, outputFile), req.Validate, err) } @@ -289,7 +299,7 @@ func (r *Reader) writeConfig(state, origin string) error { } } // The ChatGPT features are off: they would call routes the model proxy refuses. - config := fmt.Sprintf("model_provider = \"tamis\"\nmodel_reasoning_effort = \"medium\"\ndaybreak = %t\nweb_search = \"disabled\"\nchatgpt_base_url = %q\n\n[analytics]\nenabled = false\n\n[features]\napps = false\nplugins = false\nremote_plugin = false\nplugin_sharing = false\ngoals = false\nimage_generation = false\nskill_mcp_dependency_install = false\n\n[model_providers.tamis]\nname = \"Tamis\"\nbase_url = %q\n%s\nwire_api = \"responses\"\nsupports_websockets = false\nsupports_standalone_web_search = false\nrequest_max_retries = 0\nstream_max_retries = 0\n", r.Daybreak, origin+modelproxy.PathPrefix+"/backend-api", origin+r.backend.basePath, auth) + config := fmt.Sprintf("model_provider = \"tamis\"\nmodel_reasoning_effort = \"medium\"\ndaybreak = %t\nweb_search = \"disabled\"\nchatgpt_base_url = %q\n\n[features]\nremote_plugin = false\nplugin_sharing = false\ngoals = false\nimage_generation = false\nskill_mcp_dependency_install = false\n\n[model_providers.tamis]\nname = \"Tamis\"\nbase_url = %q\n%s\nwire_api = \"responses\"\nsupports_websockets = false\nsupports_standalone_web_search = false\nrequest_max_retries = 0\nstream_max_retries = 0\n", r.Daybreak, origin+modelproxy.PathPrefix+"/backend-api", origin+r.backend.basePath, auth) return os.WriteFile(filepath.Join(state, "config.toml"), []byte(config), 0o600) } diff --git a/internal/sandbox/sandbox_test.go b/internal/sandbox/sandbox_test.go index 557b833..425a72a 100644 --- a/internal/sandbox/sandbox_test.go +++ b/internal/sandbox/sandbox_test.go @@ -116,36 +116,65 @@ func TestBoundedWriterDiscardsOutputPastTheLimit(t *testing.T) { require.Len(t, w.String(), maxAgentOutput) } -func TestAgentContainerUsesOnlyThePrivateProxyAndLimitsResources(t *testing.T) { +func TestAgentRunsIsolatedWithLimitedResources(t *testing.T) { h, err := harness.ByName("claude") require.NoError(t, err) - run := &isolatedRun{runtime: container.Runtime{Bin: "docker"}, image: "agent", network: "private", proxyIP: "172.20.0.2", agent: "work"} - args := run.agentArgs(h, harness.Job{Workspace: container.WorkMount}, "/host/work", "/host/snapshot", "/host/state", "/host/work/report.json", "ANTHROPIC_API_KEY") - joined := strings.Join(args, " ") - for _, want := range []string{"host.docker.internal:172.20.0.2", "--network private", "--pids-limit 256", "--memory 2g", "--cpus 2", "src=/host/work,dst=/work,readonly", "src=/host/snapshot,dst=/work/src,readonly", "--name work", "--log-driver none"} { - assert.Contains(t, joined, want) - } - require.NotContains(t, joined, "host.docker.internal:host-gateway", "the hostile container can reach the host gateway") + reader := &Reader{backend: backends["claude"], harness: h, image: "agent", modelPort: "1234"} + + runner := reader.runner("/host/state", "/host/snapshot", "proxy-token", allow, nil) + + assert.True(t, runner.Hardened, "the agent shares a network with other containers") + assert.True(t, runner.IsolateHost, "the agent can reach host services") + assert.True(t, runner.ReadOnlyWorkspace, "the agent can change its own instructions") + assert.Equal(t, container.SidecarConfig{Allow: allow, APIPort: "1234"}, runner.Sidecar) + assert.Equal(t, []container.Mount{{Host: "/host/snapshot", Container: "/work/src", ReadOnly: true}}, runner.Mounts) + assert.Equal(t, []any{"2g", "2", 256, int64(64 << 20)}, []any{runner.Memory, runner.CPUs, runner.PIDsLimit, runner.FileSizeLimit}) +} + +func TestOfflineAgentReachesOnlyTheModelProxy(t *testing.T) { + h, err := harness.ByName("claude") + require.NoError(t, err) + reader := &Reader{backend: backends["claude"], harness: h, modelPort: "1234"} + + runner := reader.runner("/host/state", "/host/snapshot", "proxy-token", nil, nil) + + assert.Empty(t, runner.Sidecar.Allow) + assert.Empty(t, (&agent{Harness: h}).EgressHosts(), "the agent can reach the provider around the model proxy") } func TestAccessTokenContainerGetsNoAmbientAPIKey(t *testing.T) { - t.Setenv("CODEX_API_KEY", "ambient-secret") h, err := harness.ByName("codex") require.NoError(t, err) - run := &isolatedRun{image: "agent", network: "private", proxyIP: "172.20.0.2", agent: "work"} - args := run.agentArgs(h, harness.Job{Workspace: container.WorkMount}, "/host/work", "/host/snapshot", "/host/state", "/host/work/report.json", "CODEX_ACCESS_TOKEN") - for i, arg := range args { - require.NotEqual(t, "CODEX_API_KEY", arg, "ambient API key reached argument %d", i) - require.False(t, strings.HasPrefix(arg, "CODEX_API_KEY="), "ambient API key reached argument %d: %s", i, arg) - } - require.Contains(t, args, "CODEX_ACCESS_TOKEN", "access token key was not injected") + reader := &Reader{backend: backends["codex-token"], harness: h, modelPort: "1234"} + + runner := reader.runner("/host/state", "/host/snapshot", "proxy-token", allow, nil) + + assert.Contains(t, runner.OmitEnv, "CODEX_API_KEY", "the ambient API key reaches the container") + assert.Contains(t, runner.Env, "CODEX_ACCESS_TOKEN", "the access token key is not injected") + assert.Equal(t, []string{"CODEX_ACCESS_TOKEN=proxy-token"}, runner.ProcessEnv) + assert.Equal(t, container.SidecarConfig{Allow: allow, HostPorts: []string{"1234"}}, runner.Sidecar) } -func TestOnlyDockerCanStartAnIsolatedRun(t *testing.T) { - _, err := startIsolatedRun(t.Context(), container.Runtime{Bin: "podman"}, "agent", "1234", false, allow) +func TestOnlyDockerCanRunTheAgent(t *testing.T) { + h, err := harness.ByName("claude") + require.NoError(t, err) + reader := &Reader{backend: backends["claude"], harness: h, runtime: container.Runtime{Bin: "podman"}, image: "agent", modelPort: "1234"} + + _, err = reader.runner(t.TempDir(), t.TempDir(), "proxy-token", allow, nil).Open(t.Context(), &agent{Harness: h}) + require.Error(t, err, "an unsupported runtime must fail closed") } +func TestAgentKeepsOnlyTheFailuresOfItsOutput(t *testing.T) { + h, err := harness.ByName("codex") + require.NoError(t, err) + a := &agent{Harness: h} + + a.ParseStream(strings.NewReader(`{"type":"item.completed","item":{"type":"agent_message","text":"report content"}}`+"\n"+`{"type":"turn.failed","error":{"message":"usage limit"}}`), nil) + + require.Equal(t, `{"type":"turn.failed","error":{"message":"usage limit"}}`, a.failures.String()) +} + func TestReadOutputIsBounded(t *testing.T) { path := filepath.Join(t.TempDir(), "report.json") for name, size := range map[string]int{"valid": 1 << 20, "too large": 1<<20 + 1} { @@ -206,7 +235,7 @@ func TestWriteConfig(t *testing.T) { } else { require.Empty(t, got) } - for _, want := range []string{"model_provider = \"tamis\"", "model_reasoning_effort = \"medium\"", "supports_websockets = false", "plugins = false", "[analytics]\nenabled = false", "chatgpt_base_url = \"https://proxy:1/model-proxy/backend-api\"", "base_url = \"https://proxy:1" + backends[tc.backend].basePath + "\""} { + for _, want := range []string{"model_provider = \"tamis\"", "model_reasoning_effort = \"medium\"", "supports_websockets = false", "chatgpt_base_url = \"https://proxy:1/model-proxy/backend-api\"", "base_url = \"https://proxy:1" + backends[tc.backend].basePath + "\""} { if tc.wantConfig { require.Contains(t, string(got), want) } @@ -219,6 +248,17 @@ func TestWriteConfig(t *testing.T) { } } +func TestCodexRunsWithTheChatGPTFeaturesOff(t *testing.T) { + h, err := harness.ByName("codex") + require.NoError(t, err) + + args := strings.Join(h.Args(harness.Job{Workspace: container.WorkMount}), " ") + + for _, want := range []string{"-c analytics.enabled=false", "-c features.apps=false", "-c features.plugins=false"} { + assert.Contains(t, args, want) + } +} + func TestImagesArePinnedByDigest(t *testing.T) { pinned := regexp.MustCompile(`^ghcr\.io/php/tamis-agent@sha256:[0-9a-f]{64}$`) agents := map[string]bool{} diff --git a/internal/sandbox/tls.go b/internal/sandbox/tls.go deleted file mode 100644 index 3d2be4b..0000000 --- a/internal/sandbox/tls.go +++ /dev/null @@ -1,57 +0,0 @@ -package sandbox - -import ( - "crypto/ecdsa" - "crypto/elliptic" - "crypto/rand" - "crypto/tls" - "crypto/x509" - "crypto/x509/pkix" - "encoding/pem" - "math/big" - "time" -) - -// newCertificate returns a throwaway CA, as PEM, and a server certificate it signed for host. Codex only -// sends a ChatGPT credential over HTTPS, so the model proxy needs a certificate the container trusts. The -// CA key is dropped on return, so nothing can sign another certificate with it. -func newCertificate(host string) ([]byte, tls.Certificate, error) { - now := time.Now() - caKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) - if err != nil { - return nil, tls.Certificate{}, err - } - ca := &x509.Certificate{ - SerialNumber: big.NewInt(1), - Subject: pkix.Name{CommonName: "tamis model proxy CA"}, - NotBefore: now.Add(-time.Minute), - NotAfter: now.Add(24 * time.Hour), - KeyUsage: x509.KeyUsageCertSign, - BasicConstraintsValid: true, - IsCA: true, - MaxPathLenZero: true, - } - caDER, err := x509.CreateCertificate(rand.Reader, ca, ca, &caKey.PublicKey, caKey) - if err != nil { - return nil, tls.Certificate{}, err - } - key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) - if err != nil { - return nil, tls.Certificate{}, err - } - leaf := &x509.Certificate{ - SerialNumber: big.NewInt(2), - Subject: pkix.Name{CommonName: host}, - DNSNames: []string{host}, - NotBefore: now.Add(-time.Minute), - NotAfter: now.Add(24 * time.Hour), - KeyUsage: x509.KeyUsageDigitalSignature, - ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth}, - } - der, err := x509.CreateCertificate(rand.Reader, leaf, ca, &key.PublicKey, caKey) - if err != nil { - return nil, tls.Certificate{}, err - } - - return pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: caDER}), tls.Certificate{Certificate: [][]byte{der}, PrivateKey: key}, nil -} diff --git a/internal/sandbox/tls_test.go b/internal/sandbox/tls_test.go deleted file mode 100644 index e4565fe..0000000 --- a/internal/sandbox/tls_test.go +++ /dev/null @@ -1,23 +0,0 @@ -package sandbox - -import ( - "crypto/x509" - "testing" - - "github.com/stretchr/testify/require" -) - -func TestNewCertificateIsOnlyTrustedForTheModelProxyHost(t *testing.T) { - caPEM, cert, err := newCertificate("host.docker.internal") - require.NoError(t, err) - roots := x509.NewCertPool() - require.True(t, roots.AppendCertsFromPEM(caPEM), "the CA is not valid PEM") - leaf, err := x509.ParseCertificate(cert.Certificate[0]) - require.NoError(t, err) - - _, err = leaf.Verify(x509.VerifyOptions{Roots: roots, DNSName: "host.docker.internal"}) - require.NoError(t, err, "the container would reject the model proxy") - _, err = leaf.Verify(x509.VerifyOptions{Roots: roots, DNSName: "chatgpt.com"}) - require.Error(t, err, "the certificate is valid for another host") - require.False(t, leaf.IsCA, "the server certificate can sign certificates") -}