From 1591f0d3d399855eff654afcf38238bbb5e2a42b Mon Sep 17 00:00:00 2001 From: Eric Wang Date: Thu, 6 Aug 2026 02:21:25 -0700 Subject: [PATCH] fix(gateway): reserve worst-case cost, price challenges at issuance, expire tokens MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Third hardening pass, working the 2026-08-06 review's launch blockers. Four structural holes, each of which made a documented guarantee false: **Escalation priced the wrong event.** Difficulty was computed from *redeemed* mints, so an attacker could collect a day's worth of base-difficulty challenges up front and solve them at leisure — the ladder never engaged. Challenges are now priced and counted at issuance, signed to the /48 bucket they were issued to (challenge wire v2), and the counts persist across restarts so a deploy is not an amnesty. The mint endpoints also throttle by that same bucket instead of the raw IPv6 address, and an off-by-one that gave FreeMints+1 cheap mints is fixed. **The budget was a horizon, not a ceiling.** Requests were admitted unbilled and charged after the model answered, so MAX_INFLIGHT requests could all clear the breaker a dollar before it fired. Admit now reserves the request's worst case — one input token per body byte, plus max_tokens and a chain-of-thought allowance the estimate previously ignored — against both budgets, and Charge/Refund/Release reconcile it. The old 4-bytes-per-token estimate was not an upper bound; this one is. **Accounting failed open.** Journal open/write/fsync errors were ignored, so a full disk meant spend a restart would refund. The ledger now refuses admissions while it cannot write (503, self-healing on reopen), and replay is line-oriented so one torn line no longer eats every entry after it. Client aborts before upstream headers are charged an input-side estimate instead of refunded — send-and-abort was a free drain on the key. The upstream balance is polled every 15 minutes so the gateway stops claiming credit the account no longer has, and NaN budgets are rejected at boot. **Identities were immortal and unbounded.** Tokens now expire (7 days; the CLI re-enrols quietly at day 6, so honest users never notice), each token gets its own rate and concurrency valves so one subject cannot park the whole service, every limiter map has a hard size bound that fails closed instead of OOMing the container, the read-only endpoints are rate limited, /models is answered from a 30-second cache, and server-side tools (web_search) are refused on /responses — their per-search cost never appears in the usage object, so no ceiling could have held it. Also: the upstream HTTP client no longer follows redirects (Go strips Authorization cross-host but not x-api-key — a redirecting upstream would have been handed the key), and the shipped Caddy example now overwrites X-Forwarded-For, which Caddy appends to by default. All gateway tests pass with -race; the interop test confirms the CLI enrols unchanged against the v2 challenge. Co-Authored-By: Claude Fable 5 --- gateway/DESIGN.md | 49 +++++-- gateway/cmd/dsgate/main.go | 48 +++++-- gateway/deploy/README.md | 22 ++- gateway/internal/meter/meter.go | 30 ++-- gateway/internal/mint/mint.go | 138 ++++++++++++++++--- gateway/internal/mint/mint_test.go | 62 ++++++++- gateway/internal/policy/policy.go | 26 ++++ gateway/internal/policy/policy_test.go | 23 ++++ gateway/internal/quota/lifetime_test.go | 2 +- gateway/internal/quota/quota.go | 172 ++++++++++++++++++----- gateway/internal/quota/quota_test.go | 173 +++++++++++++++++++----- gateway/internal/server/anon.go | 27 +++- gateway/internal/server/proxy.go | 119 +++++++++++++--- gateway/internal/server/ratelimit.go | 31 +++-- gateway/internal/server/server.go | 156 +++++++++++++++++++-- gateway/internal/server/server_test.go | 158 ++++++++++++++++++++-- gateway/internal/token/forgery_test.go | 4 +- gateway/internal/token/token.go | 46 +++++-- gateway/internal/token/token_test.go | 14 +- internal/cli/root.go | 27 ++++ 20 files changed, 1138 insertions(+), 189 deletions(-) diff --git a/gateway/DESIGN.md b/gateway/DESIGN.md index bfc9008..07386eb 100644 --- a/gateway/DESIGN.md +++ b/gateway/DESIGN.md @@ -127,11 +127,16 @@ survive a bad day and cannot survive a bad invoice. - **Trap in 2:** PoW difficulty tuned for a laptop is nothing to a server. → Difficulty escalates per address bucket: the fourth mint from one address today costs 4x the first, the eighth costs 1024x. No ASN - database, no blocklist, self-throttling. + database, no blocklist, self-throttling. Escalation prices the + challenge being *issued*, not the redemption — counted at redemption, + a batch of cheap challenges collected up front would bypass the whole + ladder — and each challenge is signed to the bucket it was issued to. - **Trap in the proxy:** metering happens *after* the model runs, so a - single admitted request can overshoot. → Bounded by clamping - `max_tokens` and request size, so the worst single request is under a - cent, and by capping in-flight requests. + single admitted request could overshoot. → The worst case a request + could cost is *reserved* against both budgets at admission and + reconciled to the real charge when it settles, so the budgets are + ceilings, not horizons. Clamped `max_tokens` and request size keep the + worst case small enough that reservations do not choke concurrency. - **Trap in "no signup":** users may not realise their prompts transit a third party. → `deepseek free` prints exactly what leaves the machine and where it goes, before it enrols. Consent is the act of running it. @@ -204,9 +209,16 @@ a streaming request to be able to bill it — it tees the stream, passes every byte through untouched, and parses the tail. That removes the one compromise this design would otherwise have had. -If parsing fails anyway, the request is charged -`max_tokens × output_rate + body_bytes/4 × input_rate` — deliberately -above what it could have cost. Unbillable must never mean free. +If parsing fails anyway, the request is charged its full reservation — +one input token per body byte plus `max_tokens` and a chain-of-thought +allowance of output, priced deliberately above anything it could have +cost. Unbillable must never mean free. The same figure is what Admit +reserves, which is why the two cannot diverge. + +The journal is the budget's memory, so it fails closed: if a debit +cannot be written and fsync'd, the gateway stops admitting requests +until it can (and retries the journal on each refused admit, so a +recovered disk heals without a restart). ### State @@ -236,6 +248,12 @@ subject doubles as the DeepSeek `user_id` (base64url of 16 random bytes is already inside their `[a-zA-Z0-9\-_]+` rule). Quota counters are keyed by subject and expire at UTC midnight. +Tokens themselves expire after `DSGATE_TOKEN_TTL_DAYS` (7 by default), +enforced against the signed `issued` timestamp — otherwise identities +minted cheaply over months could be stockpiled and spent together. The +CLI renews its enrolment quietly at day 6, so an honest user never sees +the expiry. + ### The mint 1. `POST /v1/anon/challenge` → an HMAC-signed challenge carrying its own @@ -244,10 +262,13 @@ by subject and expire at UTC midnight. leading zero bits. 3. `POST /v1/anon/token` → verified, marked single-use, token issued. -Difficulty is chosen server-side from how many tokens that address -bucket has already minted today, and signed into the challenge so the -client cannot argue. The first mint of the day is about a second of one -core; the eighth is a quarter of an hour. +Difficulty is chosen server-side from how many challenges that address +bucket has already been *issued* today — counted at issuance and +persisted across restarts, so neither batching challenges nor a deploy +resets the ladder — and signed into the challenge along with a hash of +the issuing bucket, so it can only be redeemed from where it was asked +for. The first mint of the day is about a second of one core; the eighth +is a quarter of an hour. Minting buckets IPv4 per address and IPv6 per **/48** — the block a site is delegated. Bucketing IPv6 per /64 would let anyone with an ordinary @@ -285,7 +306,11 @@ dollar.** | `DSGATE_TOTAL_BUDGET_USD` | 20.00 | the credit pool; when it empties, the service says so | | `DSGATE_MINT_DAILY_PER_IP` | 3 | beyond this, difficulty escalates rather than refusing | | `DSGATE_POW_BITS` | 20 | ~1s on one core; 22 if we get farmed | -| `DSGATE_MAX_INFLIGHT` | 8 | bounds overshoot and protects a 1 GiB box | +| `DSGATE_MAX_INFLIGHT` | 8 | protects a 1 GiB box; spend is bounded by reservations | +| `DSGATE_SUBJECT_REQUESTS_PER_MINUTE` | 6 | one token cannot dominate the minute | +| `DSGATE_SUBJECT_INFLIGHT` | 2 | one token cannot park the whole service | +| `DSGATE_TOKEN_TTL_DAYS` | 7 | identities age out instead of accumulating | +| `DSGATE_BALANCE_CHECK_MINUTES` | 15 | the ledger's "we have credit" is checked against the real account | **Free tier is flash only.** Pro is 3x the price and the request is *rejected*, not silently downgraded — a user who asked for pro and got diff --git a/gateway/cmd/dsgate/main.go b/gateway/cmd/dsgate/main.go index 3cad865..58af5b8 100644 --- a/gateway/cmd/dsgate/main.go +++ b/gateway/cmd/dsgate/main.go @@ -17,6 +17,7 @@ import ( "errors" "fmt" "log" + "math" "net/http" "os" "os/signal" @@ -75,12 +76,16 @@ Per-user daily limits: DSGATE_ANON_MAX_TOKENS (4096) per-request output cap DSGATE_MAX_BODY_BYTES (131072) per-request body cap DSGATE_REQUESTS_PER_MINUTE (20) per-address burst + DSGATE_SUBJECT_REQUESTS_PER_MINUTE (6) per-token burst + DSGATE_SUBJECT_INFLIGHT (2) per-token concurrency + DSGATE_TOKEN_TTL_DAYS (7) token lifetime; 0 = never expires Service limits — these are what actually bound the spend: DSGATE_DAILY_BUDGET_USD (1.00) circuit breaker, resets 00:00 UTC DSGATE_TOTAL_BUDGET_USD (20.00) the credit pool DSGATE_MAX_INFLIGHT (8) + DSGATE_BALANCE_CHECK_MINUTES (15) poll upstream /user/balance; 0 = off Anti-abuse: @@ -123,6 +128,13 @@ func run() error { DailyBudgetUSD: envFloat("DSGATE_DAILY_BUDGET_USD", 1.00), TotalBudgetUSD: envFloat("DSGATE_TOTAL_BUDGET_USD", 20.00), } + // A NaN budget compares false against everything, which would make + // every admission check pass forever. Money limits have to be numbers. + if math.IsNaN(limits.DailyBudgetUSD) || math.IsInf(limits.DailyBudgetUSD, 0) || + math.IsNaN(limits.TotalBudgetUSD) || math.IsInf(limits.TotalBudgetUSD, 0) || + limits.DailyBudgetUSD < 0 || limits.TotalBudgetUSD < 0 { + return errors.New("DSGATE_DAILY_BUDGET_USD and DSGATE_TOTAL_BUDGET_USD must be finite, non-negative numbers") + } ledger, err := quota.Open(filepath.Join(stateDir, "ledger"), limits) if err != nil { return err @@ -133,25 +145,30 @@ func run() error { BaseBits: uint8(envInt("DSGATE_POW_BITS", 20)), FreeMints: envInt("DSGATE_MINT_DAILY_PER_IP", 3), TTL: 5 * time.Minute, + StatePath: filepath.Join(stateDir, "mint.json"), }) cfg := server.Config{ - UpstreamBaseURL: env("DSGATE_UPSTREAM_BASE_URL", "https://api.deepseek.com"), - UpstreamKey: key, - Model: env("DSGATE_MODEL", "deepseek-v4-flash"), - MaxBodyBytes: int64(envInt("DSGATE_MAX_BODY_BYTES", 131072)), - MaxTokens: envInt("DSGATE_ANON_MAX_TOKENS", 4096), - MaxInflight: envInt("DSGATE_MAX_INFLIGHT", 8), - RequestsPerMinute: envInt("DSGATE_REQUESTS_PER_MINUTE", 20), - TrustProxy: envBool("DSGATE_TRUST_PROXY", false), - Origins: envList("DSGATE_ORIGINS"), - AdminToken: os.Getenv("DSGATE_ADMIN_TOKEN"), - Announce: os.Getenv("DSGATE_ANNOUNCE"), + UpstreamBaseURL: env("DSGATE_UPSTREAM_BASE_URL", "https://api.deepseek.com"), + UpstreamKey: key, + Model: env("DSGATE_MODEL", "deepseek-v4-flash"), + MaxBodyBytes: int64(envInt("DSGATE_MAX_BODY_BYTES", 131072)), + MaxTokens: envInt("DSGATE_ANON_MAX_TOKENS", 4096), + MaxInflight: envInt("DSGATE_MAX_INFLIGHT", 8), + RequestsPerMinute: envInt("DSGATE_REQUESTS_PER_MINUTE", 20), + SubjectRequestsPerMinute: envInt("DSGATE_SUBJECT_REQUESTS_PER_MINUTE", 6), + SubjectInflight: envInt("DSGATE_SUBJECT_INFLIGHT", 2), + TokenTTL: time.Duration(envInt("DSGATE_TOKEN_TTL_DAYS", 7)) * 24 * time.Hour, + TrustProxy: envBool("DSGATE_TRUST_PROXY", false), + Origins: envList("DSGATE_ORIGINS"), + AdminToken: os.Getenv("DSGATE_ADMIN_TOKEN"), + Announce: os.Getenv("DSGATE_ANNOUNCE"), } + gw := server.New(cfg, signer, m, ledger) srv := &http.Server{ Addr: env("DSGATE_ADDR", ":8787"), - Handler: server.New(cfg, signer, m, ledger).Handler(), + Handler: gw.Handler(), // DeepSeek documents holding a request up to ten minutes before // inference starts, so anything shorter here would manufacture // failures out of normal slow starts. The header timeout stays @@ -181,6 +198,13 @@ func run() error { ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) defer stop() + // The local ledger only knows what this gateway spent; the account can + // empty underneath it. Checking the real balance keeps "we have + // credit" honest. + if mins := envInt("DSGATE_BALANCE_CHECK_MINUTES", 15); mins > 0 { + gw.StartBalanceWatch(ctx, time.Duration(mins)*time.Minute) + } + go func() { <-ctx.Done() log.Print("shutting down") diff --git a/gateway/deploy/README.md b/gateway/deploy/README.md index 01f859d..14cdd23 100644 --- a/gateway/deploy/README.md +++ b/gateway/deploy/README.md @@ -52,6 +52,12 @@ So: bind to loopback, terminate TLS in front, and set the flag. Caddy: ```caddyfile free.example.com { reverse_proxy 127.0.0.1:8787 { + # OVERWRITE the forwarded address. Caddy's default is to APPEND + # to whatever X-Forwarded-For the client sent — and with + # TRUST_PROXY=1 the gateway reads the leftmost entry, which + # would be the client's own invention. One line closes it. + header_up X-Forwarded-For {client_ip} + # The API holds a connection for up to ten minutes before # inference starts. A shorter timeout here turns a normal slow # start into a failed request. @@ -64,11 +70,17 @@ free.example.com { } ``` -`flush_interval -1` is not optional. Without it Caddy buffers the -response and every streamed answer arrives in one lump at the end. - -nginx wants `proxy_buffering off;` and `proxy_read_timeout 900s;` for the -same two reasons. +`header_up X-Forwarded-For {client_ip}` is not optional: without it, +anyone who can set a header mints identities from addresses of their +choosing. `flush_interval -1` is not optional either — without it Caddy +buffers the response and every streamed answer arrives in one lump at +the end. + +nginx wants `proxy_set_header X-Forwarded-For $remote_addr;` (overwrite, +never append), `proxy_buffering off;` and `proxy_read_timeout 900s;` for +the same reasons. It also inherits an access log by default — +`access_log off;` in the vhost, or the layer in front of the gateway +breaks the gateway's own "no IPs logged" promise. ## What has to be backed up diff --git a/gateway/internal/meter/meter.go b/gateway/internal/meter/meter.go index b370d2a..02aa17e 100644 --- a/gateway/internal/meter/meter.go +++ b/gateway/internal/meter/meter.go @@ -75,22 +75,34 @@ func Cost(model string, u Usage) float64 { float64(u.OutputTokens)*p.Output/perMillion } -// Estimate is the pessimistic charge for a request whose usage could not -// be read: the most it could possibly have cost, given the caps the -// gateway already enforced on it. +// Estimate is the admission-time ceiling on what a request could cost, +// and the pessimistic charge when its usage could not be read. // -// Four bytes per token under-counts nothing in practice — DeepSeek's -// tokenizer averages closer to 3.3 bytes for English and more for code — -// but this is a ceiling by construction anyway, because a request cannot -// have produced more than maxTokens of output. +// It has to be a true upper bound, because the budget breaker reserves it +// before the request is forwarded — an estimate a request could exceed +// would turn the hard ceiling back into a horizon. So: +// +// - Input is one token per byte. DeepSeek's tokenizer averages 3–4 +// bytes per token, but adversarial text can approach one, and it +// cannot go below: a token never encodes less than a byte. +// - Output is maxTokens plus a reasoning allowance. Thinking is on by +// default upstream, and DeepSeek does not document that reasoning +// tokens respect max_tokens — they are billed as output either way, +// so the bound assumes they do not. func Estimate(model string, requestBytes, maxTokens int) float64 { return Cost(model, Usage{ - InputTokens: requestBytes/4 + 1, - OutputTokens: maxTokens, + InputTokens: requestBytes + 1, + OutputTokens: maxTokens + reasoningAllowance, Found: false, }) } +// reasoningAllowance is the output headroom reserved for chain-of-thought +// tokens on top of the caller's visible max_tokens. 32k covers the +// longest thinking runs measured live; at flash rates it prices at under +// a cent, so over-reserving costs headroom, not money. +const reasoningAllowance = 32 << 10 + // rawUsage is permissive on purpose: it decodes the usage object of every // format at once, using pointers so "absent" and "zero" stay distinct. // Which fields are present is what identifies the format. diff --git a/gateway/internal/mint/mint.go b/gateway/internal/mint/mint.go index 5ea3331..79879af 100644 --- a/gateway/internal/mint/mint.go +++ b/gateway/internal/mint/mint.go @@ -13,9 +13,11 @@ package mint import ( + "encoding/json" "fmt" "net" "net/netip" + "os" "sync" "time" @@ -32,6 +34,10 @@ type Config struct { FreeMints int // TTL is how long a challenge stays solvable. TTL time.Duration + // StatePath, when set, persists the per-address issuance counts so a + // restart does not reset every address's difficulty to base. Without + // it a deploy would hand an attacker a fresh batch of cheap mints. + StatePath string } // DefaultConfig is the shipped policy. @@ -49,6 +55,20 @@ const escalationBits = 2 // busy NAT is throttled rather than permanently locked out. const maxEscalation = 12 +// maxTrackedBuckets bounds the issuance-count map. Past it, an unseen +// bucket is treated as if it had exhausted its free mints: the map cannot +// be grown without bound by an address-distributed attacker, and during +// such an attack maximum difficulty for newcomers is the right answer +// anyway. 64k buckets at ~30 bytes each is about 2 MiB, which a 192 MiB +// container can hold. +const maxTrackedBuckets = 1 << 16 + +// maxOutstanding bounds the single-use set. Every entry cost its solver a +// proof of work, so reaching this bound means someone spent tens of CPU +// hours inside one TTL window — at which point refusing redemptions +// beats being OOM-killed. +const maxOutstanding = 1 << 16 + // Mint issues and redeems challenges. type Mint struct { signer *token.Signer @@ -57,8 +77,11 @@ type Mint struct { mu sync.Mutex // day is the UTC date the per-address counts belong to. day string - // mints counts tokens issued per address bucket today. - mints map[string]int + // issued counts challenges handed out per address bucket today. The + // count moves at issuance, not redemption: difficulty priced off + // completed mints could be bypassed by collecting a batch of cheap + // challenges first and redeeming them later. + issued map[string]int // redeemed is the single-use set for challenges, keyed by the random // component. Bounded by the mint rate times the TTL, and swept. redeemed map[[16]byte]time.Time @@ -70,11 +93,12 @@ func New(signer *token.Signer, cfg Config) *Mint { m := &Mint{ signer: signer, cfg: cfg, - mints: map[string]int{}, + issued: map[string]int{}, redeemed: map[[16]byte]time.Time{}, now: time.Now, } m.day = m.now().UTC().Format("2006-01-02") + m.loadState() return m } @@ -86,23 +110,51 @@ func (m *Mint) SetClock(now func() time.Time) { m.signer.SetClock(now) } -// Challenge issues a puzzle sized to how much this address has already -// minted today. +// Challenge issues a puzzle sized to how many challenges this address has +// already been given today, and bound to the address so it cannot be +// redeemed from anywhere else. +// +// The issuance count is spent here, before the client has solved +// anything. Spending it at redemption instead would let an attacker +// collect a day's worth of base-difficulty challenges up front and solve +// them at leisure, which is exactly the escalation this exists to prevent. +// The cost is that an address which requests challenges and never redeems +// them escalates itself — which only hurts someone deliberately doing that. func (m *Mint) Challenge(remoteIP string) (*token.Challenge, error) { bucket := MintBucket(remoteIP) m.mu.Lock() m.rollLocked() - n := m.mints[bucket] + k, tracked := m.nextLocked(bucket) + if tracked { + m.issued[bucket] = k + m.saveStateLocked() + } m.mu.Unlock() - return m.signer.NewChallenge(m.difficulty(n)) + return m.signer.NewChallenge(m.difficulty(k), token.Bind(bucket)) } -// difficulty is the required leading-zero-bit count for an address that -// has already minted n tokens today. -func (m *Mint) difficulty(n int) uint8 { - extra := n - m.cfg.FreeMints +// nextLocked is the 1-indexed number of the challenge this bucket would +// be issued next, reporting whether the bucket can still be tracked. When +// the map is at its bound, unseen buckets are priced at maximum: during +// an address-distributed attack that is the right answer, and honest +// newcomers recover at the daily reset. +func (m *Mint) nextLocked(bucket string) (int, bool) { + if n, seen := m.issued[bucket]; seen { + return n + 1, true + } + if len(m.issued) >= maxTrackedBuckets { + return m.cfg.FreeMints + maxEscalation/escalationBits + 1, false + } + return 1, true +} + +// difficulty is the required leading-zero-bit count for an address's k-th +// challenge of the day. The first FreeMints are at base; every one past +// that costs escalationBits more. +func (m *Mint) difficulty(k int) uint8 { + extra := k - m.cfg.FreeMints if extra < 0 { extra = 0 } @@ -122,8 +174,11 @@ func (m *Mint) Redeem(remoteIP, challenge string, nonce uint64) (*token.Token, e if err := token.Verify(challenge, c.Difficulty, nonce); err != nil { return nil, err } - - bucket := MintBucket(remoteIP) + if c.Binding != token.Bind(MintBucket(remoteIP)) { + // The challenge was issued to a different address bucket. Honouring + // it would let one cheap address farm challenges for a fleet. + return nil, fmt.Errorf("%w: challenge was issued to a different address", token.ErrMalformed) + } m.mu.Lock() m.rollLocked() @@ -134,8 +189,11 @@ func (m *Mint) Redeem(remoteIP, challenge string, nonce uint64) (*token.Token, e // of work would mint an unlimited supply. return nil, fmt.Errorf("%w: challenge already redeemed", token.ErrMalformed) } + if len(m.redeemed) >= maxOutstanding { + m.mu.Unlock() + return nil, fmt.Errorf("%w: the mint is saturated; retry shortly", token.ErrExpired) + } m.redeemed[c.ID] = m.now() - m.mints[bucket]++ m.mu.Unlock() return m.signer.NewToken(token.TierAnon) @@ -148,7 +206,8 @@ func (m *Mint) Difficulty(remoteIP string) uint8 { m.mu.Lock() defer m.mu.Unlock() m.rollLocked() - return m.difficulty(m.mints[MintBucket(remoteIP)]) + k, _ := m.nextLocked(MintBucket(remoteIP)) + return m.difficulty(k) } func (m *Mint) rollLocked() { @@ -157,7 +216,54 @@ func (m *Mint) rollLocked() { return } m.day = day - m.mints = map[string]int{} + m.issued = map[string]int{} + m.saveStateLocked() +} + +// mintState is the persisted shape of the issuance counts. +type mintState struct { + Day string `json:"day"` + Issued map[string]int `json:"issued"` +} + +// loadState restores the day's issuance counts, so a restart is not a +// difficulty amnesty. Only counts from the current UTC day are honoured. +// Errors are deliberately soft: mint state is an anti-abuse position, not +// money, and refusing to boot over a corrupt count file would be a worse +// trade than starting the day over. +func (m *Mint) loadState() { + if m.cfg.StatePath == "" { + return + } + b, err := os.ReadFile(m.cfg.StatePath) + if err != nil { + return + } + var st mintState + if json.Unmarshal(b, &st) != nil || st.Day != m.day || st.Issued == nil { + return + } + if len(st.Issued) > maxTrackedBuckets { + return + } + m.issued = st.Issued +} + +// saveStateLocked persists the issuance counts. Atomic rename, so a crash +// mid-write leaves the previous snapshot rather than a truncated one. +func (m *Mint) saveStateLocked() { + if m.cfg.StatePath == "" { + return + } + b, err := json.Marshal(mintState{Day: m.day, Issued: m.issued}) + if err != nil { + return + } + tmp := m.cfg.StatePath + ".tmp" + if os.WriteFile(tmp, b, 0o600) != nil { + return + } + os.Rename(tmp, m.cfg.StatePath) } // sweepLocked drops challenge IDs that can no longer be redeemed anyway, diff --git a/gateway/internal/mint/mint_test.go b/gateway/internal/mint/mint_test.go index 8f5a6ab..6bafefc 100644 --- a/gateway/internal/mint/mint_test.go +++ b/gateway/internal/mint/mint_test.go @@ -105,7 +105,7 @@ func TestDifficultyEscalatesPerAddress(t *testing.T) { if base != 4 { t.Fatalf("first challenge difficulty = %d, want the base 4", base) } - for i := 0; i < 3; i++ { + for i := 0; i < 2; i++ { mintOne(t, m, ip) } if got := m.Difficulty(ip); got != base { @@ -129,13 +129,71 @@ func TestDifficultyIsCapped(t *testing.T) { m := newTestMint(t, fastConfig()) const ip = "203.0.113.7" for i := 0; i < 100; i++ { - m.mints[MintBucket(ip)] = i + m.issued[MintBucket(ip)] = i if got := m.Difficulty(ip); got > 4+maxEscalation { t.Fatalf("difficulty reached %d after %d mints; a busy NAT would be locked out for good", got, i) } } } +// The escalation must price the challenge you are ASKING for, not the +// ones you have redeemed. Otherwise an attacker collects a batch of +// base-difficulty challenges first and solves them at leisure. +func TestBatchedChallengesEscalate(t *testing.T) { + m := newTestMint(t, fastConfig()) + const ip = "203.0.113.7" + + var last uint8 + for i := 0; i < 4; i++ { + c, err := m.Challenge(ip) + if err != nil { + t.Fatal(err) + } + last = c.Difficulty + } + if last <= 4 { + t.Fatalf("fourth unredeemed challenge still at difficulty %d; batching bypasses escalation", last) + } +} + +// A challenge issued to one address must not be redeemable from another, +// or one cheap address farms base-difficulty challenges for a fleet. +func TestChallengeIsBoundToItsAddress(t *testing.T) { + m := newTestMint(t, fastConfig()) + c, _ := m.Challenge("203.0.113.7") + nonce, _ := token.Solve(c.String, c.Difficulty, 1<<24) + + if _, err := m.Redeem("198.51.100.9", c.String, nonce); err == nil { + t.Fatal("a challenge issued to 203.0.113.7 was redeemed from 198.51.100.9") + } + if _, err := m.Redeem("203.0.113.7", c.String, nonce); err != nil { + t.Fatalf("the issuing address could not redeem its own challenge: %v", err) + } +} + +// A restart must not be a difficulty amnesty. +func TestIssuanceCountsSurviveRestart(t *testing.T) { + cfg := fastConfig() + cfg.StatePath = t.TempDir() + "/mint.json" + m := newTestMint(t, cfg) + const ip = "203.0.113.7" + + for i := 0; i < 6; i++ { + if _, err := m.Challenge(ip); err != nil { + t.Fatal(err) + } + } + escalated := m.Difficulty(ip) + if escalated <= 4 { + t.Fatal("difficulty did not escalate before the restart") + } + + m2 := newTestMint(t, cfg) + if got := m2.Difficulty(ip); got != escalated { + t.Errorf("after restart difficulty = %d, want %d; a deploy resets escalation", got, escalated) + } +} + func TestEscalationResetsDaily(t *testing.T) { m := newTestMint(t, fastConfig()) now := time.Date(2026, 8, 5, 23, 59, 0, 0, time.UTC) diff --git a/gateway/internal/policy/policy.go b/gateway/internal/policy/policy.go index d573ee9..67484ab 100644 --- a/gateway/internal/policy/policy.go +++ b/gateway/internal/policy/policy.go @@ -131,6 +131,9 @@ func Apply(route Route, body []byte, subject string, lim Limits) (*Decision, err if err := forbidFanOut(obj); err != nil { return nil, err } + if err := forbidServerTools(obj, route.Format); err != nil { + return nil, err + } setIdentity(obj, route.Format, subject) out, err := json.Marshal(obj) @@ -230,6 +233,29 @@ func forbidFanOut(obj map[string]any) error { return nil } +// forbidServerTools refuses tools that run on DeepSeek's side. Client +// tools ("function") only declare a schema and cost nothing extra; a +// server-side tool like web_search performs billed work that never +// appears in the usage object, which would put its cost outside every +// ceiling this gateway enforces. Only the Responses format offers them. +func forbidServerTools(obj map[string]any, f Format) error { + if f != FormatResponses { + return nil + } + tools, _ := obj["tools"].([]any) + for _, t := range tools { + tool, _ := t.(map[string]any) + kind, _ := tool["type"].(string) + if kind != "" && kind != "function" { + return &Reject{ + Message: fmt.Sprintf("the free tier does not serve server-side tools (%q)", kind), + Hint: "bring your own key for web search: https://platform.deepseek.com/api_keys", + } + } + } + return nil +} + // setIdentity stamps the subject onto the request in whichever field the // format actually reads. // diff --git a/gateway/internal/policy/policy_test.go b/gateway/internal/policy/policy_test.go index f3a8149..1e125c5 100644 --- a/gateway/internal/policy/policy_test.go +++ b/gateway/internal/policy/policy_test.go @@ -274,3 +274,26 @@ func asReject(err error, target **Reject) bool { } return ok } + +// Server-side tools perform billed work that never appears in the usage +// object — outside every ceiling this gateway enforces. Client function +// tools only declare a schema and stay allowed. +func TestServerSideToolsAreRefused(t *testing.T) { + route, _ := Lookup("POST", "/responses") + lim := Limits{MaxTokens: 100, Model: "deepseek-v4-flash"} + + _, err := Apply(route, []byte(`{"input":"hi","tools":[{"type":"web_search"}]}`), "sub", lim) + if err == nil { + t.Fatal("web_search passed policy; its per-search cost has no ceiling") + } + + if _, err := Apply(route, []byte(`{"input":"hi","tools":[{"type":"function","name":"f"}]}`), "sub", lim); err != nil { + t.Errorf("a client function tool was refused: %v", err) + } + + // The other formats have no server-side tools; their tools stay open. + chat, _ := Lookup("POST", "/chat/completions") + if _, err := Apply(chat, []byte(`{"messages":[],"tools":[{"type":"function"}]}`), "sub", lim); err != nil { + t.Errorf("chat function tools were refused: %v", err) + } +} diff --git a/gateway/internal/quota/lifetime_test.go b/gateway/internal/quota/lifetime_test.go index 574d0d4..9c8494c 100644 --- a/gateway/internal/quota/lifetime_test.go +++ b/gateway/internal/quota/lifetime_test.go @@ -171,7 +171,7 @@ func TestCreditPoolStaysExhaustedAcrossARestart(t *testing.T) { l, done := open(t, dir, poolLimits()) defer done() - err := l.Admit("a-brand-new-subject") + err := l.Admit("a-brand-new-subject", 0) if err == nil { t.Fatal("admitted a request with the credit pool already empty") } diff --git a/gateway/internal/quota/quota.go b/gateway/internal/quota/quota.go index 922f9ef..de69795 100644 --- a/gateway/internal/quota/quota.go +++ b/gateway/internal/quota/quota.go @@ -14,6 +14,8 @@ package quota import ( + "bufio" + "bytes" "encoding/json" "fmt" "os" @@ -81,6 +83,11 @@ const ( ReasonDailyBudget Reason = "daily_budget" ReasonCredits Reason = "credits_exhausted" ReasonRevoked Reason = "revoked" + // ReasonUnavailable means the ledger cannot durably record spend right + // now. Refusing is deliberate: admitting requests that cannot be + // journalled means a restart silently refunds them, which is exactly + // the fail-open hole a budget breaker must not have. + ReasonUnavailable Reason = "ledger_unavailable" ) // LimitError is a refusal. It always says when the caller may try again, @@ -99,6 +106,8 @@ func (e *LimitError) Error() string { return "the free tier has spent today's budget" case ReasonRevoked: return "this token has been revoked" + case ReasonUnavailable: + return "the free tier cannot record spend right now" default: return fmt.Sprintf("daily %s limit reached", string(e.Reason)) } @@ -133,8 +142,17 @@ type Ledger struct { // persisted alongside priorSpend so a restart can tell which journals // have already been folded in and which still have to be. through string - revoked map[string]bool - journal *os.File + // reserved is the projected worst-case cost of every admitted request + // that has not yet been charged or refunded. It counts against both + // budgets at admission, which is what makes them ceilings rather than + // horizons: without it, MAX_INFLIGHT requests could all be admitted a + // dollar before the breaker and each overshoot it. + reserved float64 + revoked map[string]bool + journal *os.File + // journalErr is the last durability failure. While set, Admit refuses: + // spend that cannot be journalled is spend a restart would refund. + journalErr error now func() time.Time } @@ -302,17 +320,28 @@ func (l *Ledger) sumJournal(day string) (float64, error) { defer f.Close() var total float64 - dec := json.NewDecoder(f) - for { + scanJournal(f, func(e entry) { total += e.USD }) + return total, nil +} + +// scanJournal feeds every parseable entry to fn, line by line. A line +// that does not parse — a crash mid-write, a disk hiccup — is skipped +// rather than treated as the end of the file: every line after it is +// still real spend, and dropping it would refund that spend on restart. +func scanJournal(f *os.File, fn func(entry)) { + sc := bufio.NewScanner(f) + sc.Buffer(make([]byte, 0, 64<<10), 1<<20) + for sc.Scan() { + line := bytes.TrimSpace(sc.Bytes()) + if len(line) == 0 { + continue + } var e entry - if err := dec.Decode(&e); err != nil { - // Same rule as replay: a truncated final line is a crash - // mid-write, and everything before it still counts. - break + if json.Unmarshal(line, &e) != nil { + continue } - total += e.USD + fn(e) } - return total, nil } func (l *Ledger) saveStateLocked(prev error) error { @@ -341,21 +370,14 @@ func (l *Ledger) replay(day string) error { } defer f.Close() - dec := json.NewDecoder(f) - for { - var e entry - if err := dec.Decode(&e); err != nil { - // A truncated final line is what a crash mid-write looks like. - // Everything before it is still good, so stop rather than fail. - break - } + scanJournal(f, func(e entry) { a := l.accountLocked(e.Subject) a.Requests++ a.InputTokens += e.InputTokens a.OutputTokens += e.OutputTokens a.SpentUSD += e.USD l.daySpend += e.USD - } + }) return nil } @@ -385,25 +407,58 @@ func (l *Ledger) rollLocked() { } if f, err := os.OpenFile(l.journalPath(day), os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o600); err == nil { l.journal = f + l.journalErr = nil } else { l.journal = nil + l.journalErr = err + } + if err := l.saveStateLocked(nil); err != nil && l.journalErr == nil { + l.journalErr = err } - l.saveStateLocked(nil) } -// Admit debits one request against a subject and reports whether it may -// proceed. +// reopenLocked retries the journal after a durability failure, so a full +// disk that has been cleared heals without a restart. +func (l *Ledger) reopenLocked() { + if l.journalErr == nil { + return + } + if l.journal != nil { + l.journal.Close() + } + f, err := os.OpenFile(l.journalPath(l.day), os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o600) + if err != nil { + l.journal = nil + return + } + // A failed write may have left a partial line behind. Terminating it + // costs one blank line, which replay skips; not terminating it would + // glue the next entry onto the partial one and lose both. + if _, err := f.Write([]byte("\n")); err != nil { + f.Close() + l.journal = nil + return + } + l.journal = f + l.journalErr = nil +} + +// Admit debits one request against a subject and reserves its worst-case +// cost, reporting whether it may proceed. // -// The request count is spent up front because it is knowable up front; -// token counts are only settled in Charge, once the model has answered. -// That leaves a bounded window where admitted-but-unbilled requests can -// overshoot a limit, which is why the server caps how many may be in -// flight at once. -func (l *Ledger) Admit(subject string) error { +// The reservation is what makes the budgets hard ceilings: a request is +// admitted only if, priced at its absolute maximum, it still fits under +// both. Charge and Refund release the reservation, so the actual (almost +// always much smaller) cost is what sticks. +func (l *Ledger) Admit(subject string, reserveUSD float64) error { l.mu.Lock() defer l.mu.Unlock() l.rollLocked() + l.reopenLocked() + if l.journalErr != nil { + return &LimitError{Reason: ReasonUnavailable} + } if l.revoked[subject] { return &LimitError{Reason: ReasonRevoked} } @@ -412,10 +467,16 @@ func (l *Ledger) Admit(subject string) error { // Service-wide limits first: when the service is out of money, the // caller's own remaining quota is irrelevant and saying "you have 28 // requests left" would be a lie. - if spent := l.priorSpend + l.daySpend; spent >= l.limits.TotalBudgetUSD { + // Refused when already spent out, and also when this request's + // worst case would break the ceiling — both conditions, because a + // pool that is exactly empty must refuse even a zero-cost admit. + projected := l.reserved + reserveUSD + if spent := l.priorSpend + l.daySpend; spent >= l.limits.TotalBudgetUSD || + spent+projected > l.limits.TotalBudgetUSD { return &LimitError{Reason: ReasonCredits} } - if l.daySpend >= l.limits.DailyBudgetUSD { + if l.daySpend >= l.limits.DailyBudgetUSD || + l.daySpend+projected > l.limits.DailyBudgetUSD { return &LimitError{Reason: ReasonDailyBudget, ResetsAt: reset} } @@ -430,28 +491,48 @@ func (l *Ledger) Admit(subject string) error { } a.Requests++ + l.reserved += reserveUSD return nil } -// Refund returns a request allowance taken by Admit, for a call that -// never reached the model. +// Refund returns everything Admit took — the request allowance and the +// reservation — for a call that never reached the model. // // Only failures the caller cannot provoke qualify — transport errors and // upstream 429/5xx. Refunding on anything the client controls, such as a // malformed body, would turn the request counter into a free retry loop. -func (l *Ledger) Refund(subject string) { +func (l *Ledger) Refund(subject string, reserveUSD float64) { l.mu.Lock() defer l.mu.Unlock() if a, ok := l.accounts[subject]; ok && a.Requests > 0 { a.Requests-- } + l.releaseLocked(reserveUSD) } -// Charge records what a completed request cost. -func (l *Ledger) Charge(subject, endpoint, model string, in, cacheHit, out int, usd float64, estimated bool) { +// Release gives back a reservation while keeping the request debit, for +// a call that reached the model but generated nothing billable — an +// upstream 4xx that was the caller's own doing. +func (l *Ledger) Release(reserveUSD float64) { + l.mu.Lock() + defer l.mu.Unlock() + l.releaseLocked(reserveUSD) +} + +func (l *Ledger) releaseLocked(reserveUSD float64) { + l.reserved -= reserveUSD + if l.reserved < 0 { + l.reserved = 0 + } +} + +// Charge settles what a completed request actually cost, releasing its +// reservation. +func (l *Ledger) Charge(subject, endpoint, model string, in, cacheHit, out int, usd float64, reserveUSD float64, estimated bool) { l.mu.Lock() defer l.mu.Unlock() l.rollLocked() + l.releaseLocked(reserveUSD) a := l.accountLocked(subject) a.InputTokens += in @@ -463,19 +544,29 @@ func (l *Ledger) Charge(subject, endpoint, model string, in, cacheHit, out int, Time: l.now().UTC(), Subject: subject, Endpoint: endpoint, Model: model, InputTokens: in, CacheHit: cacheHit, OutputTokens: out, USD: usd, Estimated: estimated, } + // The in-memory counters above are already debited, so a journal + // failure here loses no money now — but it would on restart. Recording + // the failure makes Admit refuse until the journal writes again. if l.journal == nil { + if l.journalErr == nil { + l.journalErr = fmt.Errorf("journal is not open") + } return } b, err := json.Marshal(e) if err != nil { + l.journalErr = err return } if _, err := l.journal.Write(append(b, '\n')); err != nil { + l.journalErr = err return } // Durable per debit. At this service's volume the fsync is free, and // the alternative is that a crash refunds everybody. - l.journal.Sync() + if err := l.journal.Sync(); err != nil { + l.journalErr = err + } } // Status reports one subject's standing. @@ -507,9 +598,14 @@ type Health struct { Day string `json:"day"` Subjects int `json:"subjects_today"` DaySpendUSD float64 `json:"day_spend_usd"` + ReservedUSD float64 `json:"reserved_usd"` TotalSpendUSD float64 `json:"total_spend_usd"` DailyBudgetUSD float64 `json:"daily_budget_usd"` TotalBudgetUSD float64 `json:"total_budget_usd"` + // JournalOK is false while the ledger is refusing admissions because + // it cannot write. It is the first thing to check when everything is + // suddenly 503. + JournalOK bool `json:"journal_ok"` } func (l *Ledger) Health() Health { @@ -518,8 +614,10 @@ func (l *Ledger) Health() Health { l.rollLocked() return Health{ Day: l.day, Subjects: len(l.accounts), - DaySpendUSD: l.daySpend, TotalSpendUSD: l.priorSpend + l.daySpend, + DaySpendUSD: l.daySpend, ReservedUSD: l.reserved, + TotalSpendUSD: l.priorSpend + l.daySpend, DailyBudgetUSD: l.limits.DailyBudgetUSD, TotalBudgetUSD: l.limits.TotalBudgetUSD, + JournalOK: l.journalErr == nil, } } diff --git a/gateway/internal/quota/quota_test.go b/gateway/internal/quota/quota_test.go index 75129e2..1d3e368 100644 --- a/gateway/internal/quota/quota_test.go +++ b/gateway/internal/quota/quota_test.go @@ -42,11 +42,11 @@ func TestRequestsAreCappedPerDay(t *testing.T) { defer done() for i := 0; i < 3; i++ { - if err := l.Admit("alice"); err != nil { + if err := l.Admit("alice", 0); err != nil { t.Fatalf("request %d refused: %v", i+1, err) } } - err := l.Admit("alice") + err := l.Admit("alice", 0) if err == nil { t.Fatal("a fourth request was admitted against a limit of three") } @@ -55,7 +55,7 @@ func TestRequestsAreCappedPerDay(t *testing.T) { } // One subject hitting its limit must not affect another. - if err := l.Admit("bob"); err != nil { + if err := l.Admit("bob", 0); err != nil { t.Errorf("bob was refused because alice ran out: %v", err) } } @@ -66,12 +66,12 @@ func TestTokenLimitsAreEnforcedAfterCharging(t *testing.T) { l, done := open(t, t.TempDir(), lim) defer done() - if err := l.Admit("alice"); err != nil { + if err := l.Admit("alice", 0); err != nil { t.Fatal(err) } - l.Charge("alice", "chat", "deepseek-v4-flash", 0, 0, 600, 0.0001, false) + l.Charge("alice", "chat", "deepseek-v4-flash", 0, 0, 600, 0.0001, 0, false) - err := l.Admit("alice") + err := l.Admit("alice", 0) if err == nil { t.Fatal("admitted after the output token cap was passed") } @@ -94,11 +94,11 @@ func TestDailyBudgetStopsEveryone(t *testing.T) { spenders := []string{"a", "b", "c", "d", "e"} tripped := "" for _, who := range spenders { - if err := l.Admit(who); err != nil { + if err := l.Admit(who, 0); err != nil { tripped = who break } - l.Charge(who, "chat", "deepseek-v4-flash", 100, 0, 100, 0.003, false) + l.Charge(who, "chat", "deepseek-v4-flash", 100, 0, 100, 0.003, 0, false) } if tripped == "" { t.Fatal("five subjects spent $0.015 against a $0.01 budget and none was refused") @@ -106,7 +106,7 @@ func TestDailyBudgetStopsEveryone(t *testing.T) { // And it holds for someone who has spent nothing at all: the breaker // is about the service's money, not the caller's behaviour. - err := l.Admit("someone-brand-new") + err := l.Admit("someone-brand-new", 0) if err == nil { t.Fatal("a fresh subject was admitted after the daily budget was spent") } @@ -122,10 +122,10 @@ func TestCreditsExhaustedOutranksTheDailyBudget(t *testing.T) { l, done := open(t, t.TempDir(), lim) defer done() - l.Admit("a") - l.Charge("a", "chat", "deepseek-v4-flash", 0, 0, 0, 0.10, false) + l.Admit("a", 0) + l.Charge("a", "chat", "deepseek-v4-flash", 0, 0, 0, 0.10, 0, false) - err := l.Admit("b") + err := l.Admit("b", 0) if got := reasonOf(t, err); got != ReasonCredits { t.Fatalf("reason = %q, want %q", got, ReasonCredits) } @@ -144,10 +144,10 @@ func TestCountersSurviveARestart(t *testing.T) { dir := t.TempDir() l, _ := open(t, dir, testLimits()) - l.Admit("alice") - l.Charge("alice", "chat", "deepseek-v4-flash", 400, 0, 200, 0.002, false) - l.Admit("alice") - l.Charge("alice", "chat", "deepseek-v4-flash", 100, 0, 50, 0.001, false) + l.Admit("alice", 0) + l.Charge("alice", "chat", "deepseek-v4-flash", 400, 0, 200, 0.002, 0, false) + l.Admit("alice", 0) + l.Charge("alice", "chat", "deepseek-v4-flash", 100, 0, 50, 0.001, 0, false) l.Close() again, done := open(t, dir, testLimits()) @@ -170,8 +170,8 @@ func TestCountersSurviveARestart(t *testing.T) { func TestTruncatedJournalKeepsWhatItCan(t *testing.T) { dir := t.TempDir() l, _ := open(t, dir, testLimits()) - l.Admit("alice") - l.Charge("alice", "chat", "deepseek-v4-flash", 400, 0, 200, 0.002, false) + l.Admit("alice", 0) + l.Charge("alice", "chat", "deepseek-v4-flash", 400, 0, 200, 0.002, 0, false) day := l.day l.Close() @@ -199,8 +199,8 @@ func TestLifetimeSpendSurvivesTheDayRolling(t *testing.T) { l, _ := open(t, dir, testLimits()) l.SetClock(func() time.Time { return now }) - l.Admit("alice") - l.Charge("alice", "chat", "deepseek-v4-flash", 100, 0, 100, 0.05, false) + l.Admit("alice", 0) + l.Charge("alice", "chat", "deepseek-v4-flash", 100, 0, 100, 0.05, 0, false) now = now.Add(2 * time.Minute) // past midnight UTC @@ -241,14 +241,14 @@ func TestRefundReturnsTheRequestAllowance(t *testing.T) { l, done := open(t, t.TempDir(), testLimits()) defer done() - l.Admit("alice") - l.Refund("alice") + l.Admit("alice", 0) + l.Refund("alice", 0) if st := l.Status("alice", "anon"); st.Used.Requests != 0 { t.Errorf("requests after refund = %d, want 0", st.Used.Requests) } // Refunding more than was taken must not create allowance. - l.Refund("alice") - l.Refund("alice") + l.Refund("alice", 0) + l.Refund("alice", 0) if st := l.Status("alice", "anon"); st.Used.Requests != 0 { t.Errorf("over-refunding produced %d requests", st.Used.Requests) } @@ -259,7 +259,7 @@ func TestRevocation(t *testing.T) { l, done := open(t, dir, testLimits()) defer done() - if err := l.Admit("spammer"); err != nil { + if err := l.Admit("spammer", 0); err != nil { t.Fatal(err) } if err := os.WriteFile(filepath.Join(dir, "revoked.txt"), @@ -270,10 +270,10 @@ func TestRevocation(t *testing.T) { t.Fatal(err) } - if got := reasonOf(t, l.Admit("spammer")); got != ReasonRevoked { + if got := reasonOf(t, l.Admit("spammer", 0)); got != ReasonRevoked { t.Errorf("reason = %q, want %q", got, ReasonRevoked) } - if err := l.Admit("alice"); err != nil { + if err := l.Admit("alice", 0); err != nil { t.Errorf("revoking one subject blocked another: %v", err) } } @@ -284,8 +284,8 @@ func TestRevocation(t *testing.T) { func TestJournalRecordsCountsAndNothingElse(t *testing.T) { dir := t.TempDir() l, _ := open(t, dir, testLimits()) - l.Admit("alice") - l.Charge("alice", "chat", "deepseek-v4-flash", 400, 120, 200, 0.002, false) + l.Admit("alice", 0) + l.Charge("alice", "chat", "deepseek-v4-flash", 400, 120, 200, 0.002, 0, false) day := l.day l.Close() @@ -309,8 +309,8 @@ func TestJournalRecordsCountsAndNothingElse(t *testing.T) { func TestStatusDoesNotLeakServiceFinances(t *testing.T) { l, done := open(t, t.TempDir(), testLimits()) defer done() - l.Admit("alice") - l.Charge("alice", "chat", "deepseek-v4-flash", 10, 0, 10, 0.005, false) + l.Admit("alice", 0) + l.Charge("alice", "chat", "deepseek-v4-flash", 10, 0, 10, 0.005, 0, false) st := l.Status("alice", "anon") if st.Used.SpentUSD == 0 { @@ -323,3 +323,114 @@ func TestStatusDoesNotLeakServiceFinances(t *testing.T) { t.Error("status has no reset horizon") } } + +// The reservation is what turns the budgets from horizons into ceilings: +// a request whose worst case does not fit is refused before it is +// forwarded, not billed after it overshoots. +func TestReservationIsACeiling(t *testing.T) { + lim := testLimits() // $0.01/day + lim.DailyRequests = 100 + l, done := open(t, t.TempDir(), lim) + defer done() + + if got := reasonOf(t, l.Admit("a", 0.02)); got != ReasonDailyBudget { + t.Fatalf("a $0.02 worst case fit under a $0.01 budget: %v", got) + } + if err := l.Admit("a", 0.005); err != nil { + t.Fatalf("a fitting reservation was refused: %v", err) + } + // A second request that fits the budget alone, but not alongside the + // first one's reservation, must wait. + if got := reasonOf(t, l.Admit("b", 0.006)); got != ReasonDailyBudget { + t.Fatalf("overlapping reservations overshot the budget: %v", got) + } + // Settling the first request at its (small) real cost frees the room. + l.Charge("a", "chat", "deepseek-v4-flash", 10, 0, 10, 0.001, 0.005, false) + if err := l.Admit("b", 0.006); err != nil { + t.Fatalf("room was not released at Charge: %v", err) + } +} + +func TestRefundReleasesTheReservation(t *testing.T) { + lim := testLimits() + lim.DailyRequests = 100 + l, done := open(t, t.TempDir(), lim) + defer done() + + if err := l.Admit("a", 0.009); err != nil { + t.Fatal(err) + } + l.Refund("a", 0.009) + if err := l.Admit("b", 0.009); err != nil { + t.Fatalf("a refunded reservation still held the budget: %v", err) + } + l.Release(0.009) + if err := l.Admit("c", 0.009); err != nil { + t.Fatalf("Release did not free the room: %v", err) + } +} + +// Spend that cannot be journalled is spend a restart refunds, so a +// broken journal must stop admissions rather than being shrugged off. +func TestJournalFailureFailsClosed(t *testing.T) { + l, done := open(t, t.TempDir(), testLimits()) + defer done() + + if err := l.Admit("a", 0); err != nil { + t.Fatal(err) + } + + // Break durability: close the fd out from under the ledger and make + // the file unwritable so reopening cannot silently heal. + path := l.journalPath(l.day) + l.mu.Lock() + l.journal.Close() + l.mu.Unlock() + if err := os.Chmod(path, 0o400); err != nil { + t.Fatal(err) + } + + l.Charge("a", "chat", "deepseek-v4-flash", 10, 0, 10, 0.0001, 0, false) + err := l.Admit("b", 0) + if got := reasonOf(t, err); got != ReasonUnavailable { + t.Fatalf("admissions continued with a dead journal: %v", err) + } + + // And it heals without a restart once the disk does. + if err := os.Chmod(path, 0o600); err != nil { + t.Fatal(err) + } + if err := l.Admit("b", 0); err != nil { + t.Fatalf("the ledger did not recover after the journal came back: %v", err) + } +} + +// A corrupt line in the middle of a journal must cost us that line, not +// every line after it: dropping the tail would refund real spend. +func TestReplaySkipsACorruptLine(t *testing.T) { + dir := t.TempDir() + l, done := open(t, dir, testLimits()) + l.Admit("a", 0) + l.Charge("a", "chat", "deepseek-v4-flash", 10, 0, 10, 0.002, 0, false) + day := l.day + done() + + f, err := os.OpenFile(filepath.Join(dir, "journal-"+day+".jsonl"), os.O_WRONLY|os.O_APPEND, 0o600) + if err != nil { + t.Fatal(err) + } + f.WriteString(`{"truncated`) + f.WriteString("\n") + f.Close() + + l2, done2 := open(t, dir, testLimits()) + l2.Admit("a", 0) + l2.Charge("a", "chat", "deepseek-v4-flash", 10, 0, 10, 0.003, 0, false) + done2() + + l3, done3 := open(t, dir, testLimits()) + defer done3() + if got := l3.Health().DaySpendUSD; got < 0.0049 { + t.Errorf("day spend replayed as $%.4f; the corrupt line ate the entries after it", got) + } +} diff --git a/gateway/internal/server/anon.go b/gateway/internal/server/anon.go index 7ff7cb3..aeb2b8b 100644 --- a/gateway/internal/server/anon.go +++ b/gateway/internal/server/anon.go @@ -8,9 +8,23 @@ import ( "time" "github.com/thevibeworks/deepseek-cli/gateway/internal/meter" + "github.com/thevibeworks/deepseek-cli/gateway/internal/mint" "github.com/thevibeworks/deepseek-cli/gateway/internal/quota" ) +// limitMeta throttles the read-only endpoints — info, quota, balance. +// They cost no money, but unthrottled they are still free CPU and a +// probe surface, and every other endpoint already pays a toll. +func (s *Server) limitMeta(w http.ResponseWriter, r *http.Request) bool { + ok, wait := s.limiter.Allow("meta:" + mint.RequestBucket(s.clientIP(r))) + if !ok { + retryAfter(w, wait) + writeError(w, http.StatusTooManyRequests, typeQuota, + fmt.Sprintf("slow down — retry in %s", wait.Round(time.Second))) + } + return ok +} + // challengeTTL must match the value the mint was built with; it is // echoed to the client so a solver knows how long it has. const challengeTTL = 5 * time.Minute @@ -32,7 +46,10 @@ type ChallengeResponse struct { func (s *Server) handleChallenge(w http.ResponseWriter, r *http.Request) { ip := s.clientIP(r) - if ok, wait := s.limiter.Allow("mint:" + ip); !ok { + // Throttled by the same /48 bucket that difficulty escalates on. Keyed + // on the raw address, one IPv6 /48 would be 65,536 independent + // throttles — the exact fan-out MintBucket exists to prevent. + if ok, wait := s.limiter.Allow("mint:" + mint.MintBucket(ip)); !ok { retryAfter(w, wait) writeError(w, http.StatusTooManyRequests, typeQuota, fmt.Sprintf("too many mint attempts; retry in %s", wait.Round(time.Second))) @@ -73,7 +90,7 @@ type TokenResponse struct { func (s *Server) handleToken(w http.ResponseWriter, r *http.Request) { ip := s.clientIP(r) - if ok, wait := s.limiter.Allow("mint:" + ip); !ok { + if ok, wait := s.limiter.Allow("mint:" + mint.MintBucket(ip)); !ok { retryAfter(w, wait) writeError(w, http.StatusTooManyRequests, typeQuota, fmt.Sprintf("too many mint attempts; retry in %s", wait.Round(time.Second))) @@ -105,6 +122,9 @@ func (s *Server) handleToken(w http.ResponseWriter, r *http.Request) { } func (s *Server) handleQuota(w http.ResponseWriter, r *http.Request) { + if !s.limitMeta(w, r) { + return + } t, err := s.authenticate(r) if err != nil { writeError(w, http.StatusUnauthorized, typeAuth, err.Error()) @@ -140,6 +160,9 @@ type balanceInfo struct { } func (s *Server) handleBalance(w http.ResponseWriter, r *http.Request) { + if !s.limitMeta(w, r) { + return + } t, err := s.authenticate(r) if err != nil { writeError(w, http.StatusUnauthorized, typeAuth, err.Error()) diff --git a/gateway/internal/server/proxy.go b/gateway/internal/server/proxy.go index 54b558d..2017ca6 100644 --- a/gateway/internal/server/proxy.go +++ b/gateway/internal/server/proxy.go @@ -56,6 +56,35 @@ func (s *Server) handleProxy(w http.ResponseWriter, r *http.Request) { } subject := tok.Subject.String() + // The subject's own valves, independent of the address ones: a token + // used from many addresses and many tokens behind one address are + // different attacks. + if ok, wait := s.subjLimiter.Allow(subject); !ok { + retryAfter(w, wait) + writeError(w, http.StatusTooManyRequests, typeQuota, + fmt.Sprintf("this token is sending too fast — retry in %s", wait.Round(time.Second))) + return + } + if !s.acquireSubject(subject) { + w.Header().Set("Retry-After", "2") + writeError(w, http.StatusTooManyRequests, typeQuota, + fmt.Sprintf("this token already has %d request(s) in flight — wait for one to finish", s.cfg.SubjectInflight)) + return + } + defer s.releaseSubject(subject) + + // The model list barely changes and is deliberately uncharged, so it + // is answered from a short cache when possible — otherwise the one + // free endpoint would burn in-flight slots and upstream round trips. + if route.Name == "models" { + if body, ok := s.cachedModels(); ok { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.Write(body) + return + } + } + body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, s.cfg.MaxBodyBytes)) if err != nil { writeError(w, http.StatusRequestEntityTooLarge, typeRejected, fmt.Sprintf( @@ -82,8 +111,21 @@ func (s *Server) handleProxy(w http.ResponseWriter, r *http.Request) { // `deepseek status` free and safe in a loop against the free tier, // exactly as it is against the real API. billable := route.Format != policy.FormatNone + var reserve float64 if billable { - if err := s.ledger.Admit(subject); err != nil { + if s.upstreamDry.Load() { + // DeepSeek says the account is unusable. The local ledger's + // opinion is irrelevant; honest 402 beats a confusing relay of + // upstream's insufficient-balance error. + s.writeLimit(w, "a.LimitError{Reason: quota.ReasonCredits}) + return + } + // The worst this request could cost is reserved before it is + // forwarded. This is what makes the budget a ceiling: without the + // reservation, every in-flight request is unbilled and the breaker + // only notices after the money is spent. + reserve = meter.Estimate(decision.Model, len(decision.Body), decision.MaxTokens) + if err := s.ledger.Admit(subject, reserve); err != nil { s.writeLimit(w, err) return } @@ -94,7 +136,7 @@ func (s *Server) handleProxy(w http.ResponseWriter, r *http.Request) { if err := s.acquire(r); err != nil { if billable { - s.ledger.Refund(subject) + s.ledger.Refund(subject, reserve) } w.Header().Set("Retry-After", "5") writeError(w, http.StatusServiceUnavailable, typeQuota, @@ -103,7 +145,7 @@ func (s *Server) handleProxy(w http.ResponseWriter, r *http.Request) { } defer func() { <-s.inflight }() - s.forward(w, r, route, decision, subject, billable) + s.forward(w, r, route, decision, subject, billable, reserve) } // acquire takes an in-flight slot, or gives up. @@ -123,7 +165,7 @@ func (s *Server) acquire(r *http.Request) error { } } -func (s *Server) forward(w http.ResponseWriter, r *http.Request, route policy.Route, d *policy.Decision, subject string, billable bool) { +func (s *Server) forward(w http.ResponseWriter, r *http.Request, route policy.Route, d *policy.Decision, subject string, billable bool, reserve float64) { url := strings.TrimRight(s.cfg.UpstreamBaseURL, "/") + route.Upstream var payload io.Reader @@ -133,7 +175,7 @@ func (s *Server) forward(w http.ResponseWriter, r *http.Request, route policy.Ro up, err := http.NewRequestWithContext(r.Context(), route.Method, url, payload) if err != nil { if billable { - s.ledger.Refund(subject) + s.ledger.Refund(subject, reserve) } writeError(w, http.StatusInternalServerError, typeInternal, "could not build the upstream request") return @@ -162,13 +204,23 @@ func (s *Server) forward(w http.ResponseWriter, r *http.Request, route policy.Ro resp, err := s.http.Do(up) if err != nil { + if r.Context().Err() != nil { + // The caller hung up before upstream answered. The prompt very + // likely reached DeepSeek, and whether they bill the aborted + // prefill is undocumented — so this is charged as an input-side + // estimate rather than refunded. Refunding here was a drain: an + // attacker could send-and-abort in a loop and the ledger would + // record nothing while our key paid for every prefill. + if billable { + cost := meter.Cost(d.Model, meter.Usage{InputTokens: len(d.Body) + 1}) + s.ledger.Charge(subject, route.Name, d.Model, len(d.Body)/4+1, 0, 0, cost, reserve, true) + } + return // there is nobody to tell + } // Never reached the model, so it cost nothing and the caller keeps // their request allowance. if billable { - s.ledger.Refund(subject) - } - if r.Context().Err() != nil { - return // the caller hung up; there is nobody to tell + s.ledger.Refund(subject, reserve) } writeError(w, http.StatusBadGateway, typeUpstream, "could not reach DeepSeek: "+err.Error()) return @@ -221,23 +273,27 @@ func (s *Server) forward(w http.ResponseWriter, r *http.Request, route policy.Ro if !usage.Found { if upstreamFault || resp.StatusCode >= 400 { - // No tokens were generated. Give the request allowance back - // only when the fault was not the caller's. if upstreamFault { - s.ledger.Refund(subject) + // No tokens were generated and the fault was not the + // caller's: give everything back. + s.ledger.Refund(subject, reserve) + } else { + // The caller's own 4xx keeps its request debit, but the + // money reserved for it goes back to the pool. + s.ledger.Release(reserve) } return } - // A 2xx we could not read. Charge the most it could have cost: - // unbillable must never mean free, or it becomes the way in. - cost := meter.Estimate(model, len(d.Body), d.MaxTokens) - s.ledger.Charge(subject, route.Name, model, len(d.Body)/4+1, 0, d.MaxTokens, cost, true) + // A 2xx we could not read. Charge the whole reservation — the most + // it could have cost. Unbillable must never mean free, or it + // becomes the way in. + s.ledger.Charge(subject, route.Name, model, len(d.Body)/4+1, 0, d.MaxTokens, reserve, reserve, true) return } s.ledger.Charge(subject, route.Name, model, usage.InputTokens, usage.CacheHitTokens, usage.OutputTokens, - meter.Cost(model, usage), false) + meter.Cost(model, usage), reserve, false) } // relayModels forwards the model list, minus the models this gateway @@ -281,12 +337,35 @@ func (s *Server) relayModels(w http.ResponseWriter, resp *http.Response) { out, err := json.Marshal(list) if err != nil { out = raw + } else { + s.storeModels(out) } w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) w.Write(out) } +// modelsCacheTTL is short on purpose: /models doubles as `deepseek +// status`'s reachability probe, and a long cache would keep answering +// "up" after DeepSeek went down. +const modelsCacheTTL = 30 * time.Second + +func (s *Server) cachedModels() ([]byte, bool) { + s.modelsMu.Lock() + defer s.modelsMu.Unlock() + if s.modelsBody == nil || time.Since(s.modelsAt) > modelsCacheTTL { + return nil, false + } + return s.modelsBody, true +} + +func (s *Server) storeModels(body []byte) { + s.modelsMu.Lock() + s.modelsBody = body + s.modelsAt = time.Now() + s.modelsMu.Unlock() +} + // relay copies the upstream body to the client and to the meter, // flushing as it goes. // @@ -354,6 +433,12 @@ func (s *Server) writeLimit(w http.ResponseWriter, err error) { case quota.ReasonRevoked: writeError(w, http.StatusForbidden, typeAuth, "this free-tier token has been revoked") + case quota.ReasonUnavailable: + // The ledger cannot record spend, so nothing is allowed to spend. + // A 503 is honest: this is our outage, not the caller's quota. + w.Header().Set("Retry-After", "30") + writeError(w, http.StatusServiceUnavailable, typeInternal, + "the free tier is temporarily unavailable; retry shortly") case quota.ReasonDailyBudget: retryAfter(w, lim.RetryAfter(time.Now())) writeError(w, http.StatusTooManyRequests, typeQuota, diff --git a/gateway/internal/server/ratelimit.go b/gateway/internal/server/ratelimit.go index 34a1fa1..1dde407 100644 --- a/gateway/internal/server/ratelimit.go +++ b/gateway/internal/server/ratelimit.go @@ -12,12 +12,20 @@ import ( // the cheap reflex that stops a runaway loop from reaching that code at // all. Neither one protects the budget — that is the breaker's job — so // this can stay as simple as it looks. +// maxLimiterBuckets is the hard size bound. When even a sweep cannot get +// under it — an address-distributed flood arriving faster than buckets +// idle out — new keys are refused rather than allocated. Fail closed: +// the alternative is the map growing until the container is OOM-killed, +// which turns a rate-limit evasion into a whole-service restart. +const maxLimiterBuckets = 1 << 16 + type limiter struct { - mu sync.Mutex - burst float64 - perSec float64 - buckets map[string]*bucketState - now func() time.Time + mu sync.Mutex + burst float64 + perSec float64 + buckets map[string]*bucketState + lastSweep time.Time + now func() time.Time } type bucketState struct { @@ -46,7 +54,12 @@ func (l *limiter) Allow(key string) (bool, time.Duration) { now := l.now() b, ok := l.buckets[key] if !ok { - l.sweepLocked(now) + if len(l.buckets) >= 4096 { + l.sweepLocked(now) + } + if len(l.buckets) >= maxLimiterBuckets { + return false, time.Second + } b = &bucketState{tokens: l.burst, last: now} l.buckets[key] = b } @@ -67,11 +80,13 @@ func (l *limiter) Allow(key string) (bool, time.Duration) { // sweepLocked drops buckets that have been idle long enough to have // refilled completely, since a full bucket is indistinguishable from a -// missing one. +// missing one. At most once a second: under a distinct-key flood every +// insert would otherwise pay an O(n) scan that frees nothing. func (l *limiter) sweepLocked(now time.Time) { - if len(l.buckets) < 4096 { + if now.Sub(l.lastSweep) < time.Second { return } + l.lastSweep = now full := time.Duration(l.burst / l.perSec * float64(time.Second)) for k, b := range l.buckets { if now.Sub(b.last) > full { diff --git a/gateway/internal/server/server.go b/gateway/internal/server/server.go index f1bafda..e07c99d 100644 --- a/gateway/internal/server/server.go +++ b/gateway/internal/server/server.go @@ -4,10 +4,14 @@ package server import ( + "context" "encoding/json" "fmt" + "io" "net/http" "strings" + "sync" + "sync/atomic" "time" "github.com/thevibeworks/deepseek-cli/gateway/internal/mint" @@ -34,6 +38,23 @@ type Config struct { // quota so that a runaway loop is cheap to refuse. RequestsPerMinute int + // SubjectRequestsPerMinute is the per-token burst limit. The address + // limit alone is not enough: one subject spread across addresses, or + // many subjects behind one address, are different attacks and need + // separate valves. + SubjectRequestsPerMinute int + + // SubjectInflight caps concurrent requests per token. Without it, one + // token opening MaxInflight never-reading streams parks the whole + // service behind the global cap. + SubjectInflight int + + // TokenTTL is how long a minted token stays valid. Re-enrolment is a + // second of CPU, so expiry costs honest users almost nothing — and + // stops an attacker stockpiling identities for months and spending + // them together. + TokenTTL time.Duration + // TrustProxy makes X-Forwarded-For authoritative. Set it only when // something we control terminates TLS in front of this process: // facing the internet directly, the header is attacker-supplied and @@ -59,9 +80,29 @@ type Server struct { signer *token.Signer ledger *quota.Ledger - http *http.Client - inflight chan struct{} - limiter *limiter + http *http.Client + inflight chan struct{} + limiter *limiter + subjLimiter *limiter + + subjMu sync.Mutex + subjInflight map[string]int + + // modelsCache holds the last filtered /models answer. The list changes + // on the timescale of DeepSeek launches, and without a cache the one + // deliberately-uncharged endpoint would consume an in-flight slot and + // an upstream round trip per poll. + modelsMu sync.Mutex + modelsBody []byte + modelsAt time.Time + + // upstreamDry is set when DeepSeek itself reports our account + // unusable. The local ledger only knows what this gateway spent — the + // account can empty underneath it (other spenders, a price change), + // and a gateway that keeps promising credit it does not have would + // fail every request with a confusing upstream error instead of an + // honest 402. + upstreamDry atomic.Bool origins map[string]bool started time.Time @@ -91,14 +132,46 @@ func New(cfg Config, signer *token.Signer, m *mint.Mint, ledger *quota.Ledger) * DisableCompression: false, ForceAttemptHTTP2: true, }, + // Never follow an upstream redirect. Our key rides on these + // requests, and Go strips Authorization across hosts but not + // x-api-key — a redirecting upstream would be handed the key. + CheckRedirect: func(*http.Request, []*http.Request) error { + return http.ErrUseLastResponse + }, }, - inflight: make(chan struct{}, cfg.MaxInflight), - limiter: newLimiter(cfg.RequestsPerMinute, time.Minute), - origins: origins, - started: time.Now(), + inflight: make(chan struct{}, cfg.MaxInflight), + limiter: newLimiter(cfg.RequestsPerMinute, time.Minute), + subjLimiter: newLimiter(cfg.SubjectRequestsPerMinute, time.Minute), + subjInflight: map[string]int{}, + origins: origins, + started: time.Now(), } } +// acquireSubject takes one of a token's concurrency slots, or reports +// that they are all in use. +func (s *Server) acquireSubject(subject string) bool { + s.subjMu.Lock() + defer s.subjMu.Unlock() + if s.subjInflight[subject] >= s.cfg.SubjectInflight { + return false + } + s.subjInflight[subject]++ + return true +} + +func (s *Server) releaseSubject(subject string) { + s.subjMu.Lock() + defer s.subjMu.Unlock() + if s.subjInflight[subject] <= 1 { + // Deleting at zero keeps the map's size bounded by the subjects + // actually in flight, which the global cap already bounds. + delete(s.subjInflight, subject) + return + } + s.subjInflight[subject]-- +} + // Handler builds the routing table. func (s *Server) Handler() http.Handler { mux := http.NewServeMux() @@ -221,6 +294,11 @@ func (s *Server) authenticate(r *http.Request) (*token.Token, error) { if err != nil { return nil, fmt.Errorf("token not valid: %w — run `deepseek free` to mint a new one", err) } + if s.cfg.TokenTTL > 0 && time.Since(t.Issued) > s.cfg.TokenTTL { + // Enforced here rather than in the codec: expiry is service + // policy, and the codec's job is only to say whose token it is. + return nil, fmt.Errorf("this free-tier token has expired — run `deepseek free` to mint a new one (about a second of CPU)") + } return t, nil } @@ -228,6 +306,60 @@ func (s *Server) clientIP(r *http.Request) string { return mint.ClientIP(r.RemoteAddr, r.Header.Get("X-Forwarded-For"), s.cfg.TrustProxy) } +// --- upstream balance ---------------------------------------------------- + +// StartBalanceWatch polls DeepSeek's balance endpoint so the gateway's +// idea of "we have credit" is checked against the account that actually +// pays. The poll costs nothing — /user/balance is unbilled. +func (s *Server) StartBalanceWatch(ctx context.Context, interval time.Duration) { + if interval <= 0 { + return + } + go func() { + s.checkBalance(ctx) + t := time.NewTicker(interval) + defer t.Stop() + for { + select { + case <-ctx.Done(): + return + case <-t.C: + s.checkBalance(ctx) + } + } + }() +} + +func (s *Server) checkBalance(ctx context.Context) { + ctx, cancel := context.WithTimeout(ctx, 30*time.Second) + defer cancel() + req, err := http.NewRequestWithContext(ctx, http.MethodGet, + strings.TrimRight(s.cfg.UpstreamBaseURL, "/")+"/user/balance", nil) + if err != nil { + return + } + req.Header.Set("Authorization", "Bearer "+s.cfg.UpstreamKey) + req.Header.Set("User-Agent", "dsgate") + + resp, err := s.http.Do(req) + if err != nil { + // Network trouble is not "out of money". The last known state + // stands until DeepSeek says otherwise. + return + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return + } + var b struct { + IsAvailable bool `json:"is_available"` + } + if json.NewDecoder(io.LimitReader(resp.Body, 1<<20)).Decode(&b) != nil { + return + } + s.upstreamDry.Store(!b.IsAvailable) +} + // --- simple endpoints --------------------------------------------------- func (s *Server) handleHealthz(w http.ResponseWriter, r *http.Request) { @@ -242,7 +374,10 @@ func (s *Server) handleAdminHealth(w http.ResponseWriter, r *http.Request) { writeError(w, http.StatusNotFound, typeRejected, "not found") return } - writeJSON(w, http.StatusOK, s.ledger.Health()) + writeJSON(w, http.StatusOK, struct { + quota.Health + UpstreamAvailable bool `json:"upstream_available"` + }{s.ledger.Health(), !s.upstreamDry.Load()}) } // Info is the unauthenticated description of the service, so a client can @@ -261,6 +396,9 @@ type Info struct { } func (s *Server) handleInfo(w http.ResponseWriter, r *http.Request) { + if !s.limitMeta(w, r) { + return + } h := s.ledger.Health() writeJSON(w, http.StatusOK, Info{ Service: "dsgate", @@ -275,6 +413,6 @@ func (s *Server) handleInfo(w http.ResponseWriter, r *http.Request) { "/anthropic/v1/messages", "/responses", "/models", }, Privacy: "prompts and completions are relayed to DeepSeek and are not stored or logged by this gateway; only token counts and cost are recorded", - Exhausted: h.TotalSpendUSD >= h.TotalBudgetUSD, + Exhausted: h.TotalSpendUSD >= h.TotalBudgetUSD || s.upstreamDry.Load(), }) } diff --git a/gateway/internal/server/server_test.go b/gateway/internal/server/server_test.go index 0536127..7c737d1 100644 --- a/gateway/internal/server/server_test.go +++ b/gateway/internal/server/server_test.go @@ -1,6 +1,7 @@ package server import ( + "context" "encoding/json" "fmt" "io" @@ -92,14 +93,17 @@ func newHarness(t *testing.T, up *upstream, tune func(*Config, *quota.Limits)) * TotalBudgetUSD: 10, } cfg := Config{ - UpstreamBaseURL: up.server.URL, - UpstreamKey: upstreamKey, - Model: "deepseek-v4-flash", - MaxBodyBytes: 4096, - MaxTokens: 256, - MaxInflight: 4, - RequestsPerMinute: 1000, - Origins: []string{"https://deepseek-cli.example"}, + UpstreamBaseURL: up.server.URL, + UpstreamKey: upstreamKey, + Model: "deepseek-v4-flash", + MaxBodyBytes: 4096, + MaxTokens: 256, + MaxInflight: 4, + RequestsPerMinute: 1000, + SubjectRequestsPerMinute: 1000, + SubjectInflight: 4, + TokenTTL: 7 * 24 * time.Hour, + Origins: []string{"https://deepseek-cli.example"}, } if tune != nil { tune(&cfg, &limits) @@ -759,3 +763,141 @@ func subjectOf(t *testing.T, raw string) string { } return tok.Subject.String() } + +// One token must not be able to park the whole service behind the global +// in-flight cap. +func TestSubjectConcurrencyIsCapped(t *testing.T) { + release := make(chan struct{}) + up := newUpstream(t, func(w http.ResponseWriter, r *http.Request) { + <-release + io.WriteString(w, chatReply(10, 10)) + }) + defer close(release) + h := newHarness(t, up, func(cfg *Config, lim *quota.Limits) { + cfg.SubjectInflight = 1 + lim.DailyRequests = 100 + }) + tok := h.enrol(t) + + started := make(chan *http.Response, 1) + go func() { + resp := h.do(t, "POST", "/chat/completions", tok, `{"messages":[]}`) + started <- resp + }() + + // Wait until the first request holds the subject's only slot. + deadline := time.Now().Add(4 * time.Second) + for h.upstream.count() == 0 { + if time.Now().After(deadline) { + t.Fatal("the first request never reached upstream") + } + time.Sleep(time.Millisecond) + } + + resp := h.do(t, "POST", "/chat/completions", tok, `{"messages":[]}`) + defer resp.Body.Close() + if resp.StatusCode != http.StatusTooManyRequests { + t.Fatalf("second concurrent request: HTTP %d, want 429", resp.StatusCode) + } + + release <- struct{}{} + (<-started).Body.Close() + h.settle(t) + + // With the slot free again the same token proceeds. + go func() { release <- struct{}{} }() + resp2 := h.do(t, "POST", "/chat/completions", tok, `{"messages":[]}`) + defer resp2.Body.Close() + if resp2.StatusCode != 200 { + t.Errorf("request after the slot freed: HTTP %d, want 200", resp2.StatusCode) + } +} + +// A token past its TTL is a 401 pointing at re-enrolment, not a working +// credential. Stockpiled identities must age out. +func TestExpiredTokenIsRefused(t *testing.T) { + up := newUpstream(t, func(w http.ResponseWriter, r *http.Request) { + io.WriteString(w, chatReply(10, 10)) + }) + h := newHarness(t, up, func(cfg *Config, lim *quota.Limits) { + cfg.TokenTTL = time.Millisecond + }) + tok := h.enrol(t) + time.Sleep(5 * time.Millisecond) + + resp := h.do(t, "POST", "/chat/completions", tok, `{"messages":[]}`) + defer resp.Body.Close() + if resp.StatusCode != http.StatusUnauthorized { + t.Fatalf("expired token: HTTP %d, want 401", resp.StatusCode) + } + if e := decodeError(t, resp); !strings.Contains(e.Message, "deepseek free") { + t.Errorf("the error does not say how to recover: %q", e.Message) + } + if up.count() != 0 { + t.Error("an expired token's request reached upstream") + } +} + +// The second /models inside the cache window is answered locally: the +// endpoint is deliberately uncharged, so it must not cost upstream trips. +func TestModelsIsCached(t *testing.T) { + up := newUpstream(t, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + io.WriteString(w, `{"object":"list","data":[{"id":"deepseek-v4-flash","object":"model"}]}`) + }) + h := newHarness(t, up, nil) + tok := h.enrol(t) + + for i := 0; i < 3; i++ { + resp := h.do(t, "GET", "/models", tok, "") + raw, _ := io.ReadAll(resp.Body) + resp.Body.Close() + if resp.StatusCode != 200 || !strings.Contains(string(raw), "deepseek-v4-flash") { + t.Fatalf("models call %d: HTTP %d: %s", i, resp.StatusCode, raw) + } + } + if got := up.count(); got != 1 { + t.Errorf("3 /models calls made %d upstream trips, want 1", got) + } +} + +// When DeepSeek itself reports the account unusable, the gateway must +// say 402 — not relay a confusing upstream error while its own ledger +// still claims there is credit. +func TestUpstreamDryBalanceStopsAdmissions(t *testing.T) { + up := newUpstream(t, func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/user/balance" { + w.Header().Set("Content-Type", "application/json") + io.WriteString(w, `{"is_available":false,"balance_infos":[]}`) + return + } + io.WriteString(w, chatReply(10, 10)) + }) + h := newHarness(t, up, nil) + tok := h.enrol(t) + + h.checkBalance(context.Background()) + + resp := h.do(t, "POST", "/chat/completions", tok, `{"messages":[]}`) + defer resp.Body.Close() + if resp.StatusCode != http.StatusPaymentRequired { + t.Fatalf("HTTP %d with a dry upstream account, want 402", resp.StatusCode) + } + + // And it heals when the account is topped up. + up.mu.Lock() + up.handler = func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/user/balance" { + io.WriteString(w, `{"is_available":true,"balance_infos":[]}`) + return + } + io.WriteString(w, chatReply(10, 10)) + } + up.mu.Unlock() + h.checkBalance(context.Background()) + resp2 := h.do(t, "POST", "/chat/completions", tok, `{"messages":[]}`) + defer resp2.Body.Close() + if resp2.StatusCode != 200 { + t.Errorf("HTTP %d after the account recovered, want 200", resp2.StatusCode) + } +} diff --git a/gateway/internal/token/forgery_test.go b/gateway/internal/token/forgery_test.go index d19213a..e450fa6 100644 --- a/gateway/internal/token/forgery_test.go +++ b/gateway/internal/token/forgery_test.go @@ -53,7 +53,7 @@ func TestTokenRejectsTruncatedMAC(t *testing.T) { func TestChallengeRejectsTruncatedMAC(t *testing.T) { s := testSigner(t) - c, err := s.NewChallenge(8) + c, err := s.NewChallenge(8, Bind("test")) if err != nil { t.Fatal(err) } @@ -112,7 +112,7 @@ func TestTokenRejectsOverlongMAC(t *testing.T) { // even now that both lengths are pinned. func TestDomainSeparationHolds(t *testing.T) { s := testSigner(t) - c, err := s.NewChallenge(8) + c, err := s.NewChallenge(8, Bind("test")) if err != nil { t.Fatal(err) } diff --git a/gateway/internal/token/token.go b/gateway/internal/token/token.go index d3dbddb..7c85ae5 100644 --- a/gateway/internal/token/token.go +++ b/gateway/internal/token/token.go @@ -26,8 +26,11 @@ import ( // Wire format versions. A bump invalidates every outstanding credential, // which is the intended blast radius for a format change. +// +// Challenge v2 added the 8-byte binding. Outstanding v1 challenges died +// with the bump, which cost at most one five-minute TTL of solvers. const ( - challengeVersion = 1 + challengeVersion = 2 tokenVersion = 1 ) @@ -104,8 +107,12 @@ func (s *Signer) mac(domain string, payload []byte) []byte { // --- challenges --------------------------------------------------------- -// challengePayload is version(1) | difficulty(1) | issued(4) | nonce(16). -const challengePayloadLen = 1 + 1 + 4 + 16 +// challengePayload is version(1) | difficulty(1) | issued(4) | nonce(16) | +// binding(8). +const challengePayloadLen = 1 + 1 + 4 + 16 + 8 + +// BindingLen is the size of a challenge's caller binding. +const BindingLen = 8 // challengeMACLen is how much of the HMAC a challenge carries. A challenge // is short-lived and single-use, so 128 bits is plenty and keeps the @@ -126,10 +133,23 @@ type Challenge struct { // Difficulty is the required number of leading zero bits. Difficulty uint8 Issued time.Time + // Binding ties the challenge to whoever asked for it — in practice a + // hash of the caller's address bucket. It travels inside the signed + // payload, so a challenge issued cheaply to one address cannot be + // redeemed from another after that address's difficulty has climbed. + Binding [BindingLen]byte +} + +// Bind derives a challenge binding from a caller identity string. +func Bind(identity string) [BindingLen]byte { + sum := sha256.Sum256([]byte(identity)) + var b [BindingLen]byte + copy(b[:], sum[:]) + return b } -// NewChallenge issues a puzzle at the given difficulty. -func (s *Signer) NewChallenge(difficulty uint8) (*Challenge, error) { +// NewChallenge issues a puzzle at the given difficulty, bound to a caller. +func (s *Signer) NewChallenge(difficulty uint8, binding [BindingLen]byte) (*Challenge, error) { if difficulty > 40 { // 2^40 hashes is hours of CPU. A difficulty this high is a bug in // the escalation policy, and shipping it would silently lock out @@ -139,6 +159,7 @@ func (s *Signer) NewChallenge(difficulty uint8) (*Challenge, error) { var c Challenge c.Difficulty = difficulty c.Issued = s.now().UTC() + c.Binding = binding if _, err := rand.Read(c.ID[:]); err != nil { return nil, err } @@ -147,7 +168,8 @@ func (s *Signer) NewChallenge(difficulty uint8) (*Challenge, error) { payload[0] = challengeVersion payload[1] = difficulty putUint32(payload[2:6], uint32(c.Issued.Unix())) - copy(payload[6:], c.ID[:]) + copy(payload[6:22], c.ID[:]) + copy(payload[22:], binding[:]) c.String = enc.EncodeToString(payload) + "." + enc.EncodeToString(s.mac("challenge", payload)[:challengeMACLen]) return &c, nil @@ -170,7 +192,8 @@ func (s *Signer) ParseChallenge(raw string, ttl time.Duration) (*Challenge, erro Difficulty: payload[1], Issued: time.Unix(int64(uint32From(payload[2:6])), 0).UTC(), } - copy(c.ID[:], payload[6:]) + copy(c.ID[:], payload[6:22]) + copy(c.Binding[:], payload[22:]) age := s.now().Sub(c.Issued) if age > ttl { @@ -293,10 +316,11 @@ func (s *Signer) tokenFor(sub Subject, tier Tier, issued time.Time) (*Token, err // ParseToken verifies a bearer value and returns what it vouches for. // -// There is no expiry check. A free-tier token is a name, not a lease — -// quota resets daily and is enforced by the counters, so ageing tokens -// out would only make honest users re-mint for no security gain. Ending -// a token's life is what revocation is for. +// There is no expiry check here: this codec's job is to say whose token +// it is and when it was issued. How old is too old is service policy, +// enforced by the gateway against the Issued timestamp — it changed once +// already (tokens were leases-for-life before 2026-08-06) and the wire +// format did not have to. func (s *Signer) ParseToken(raw string) (*Token, error) { if len(raw) <= len(TokenPrefix) || raw[:len(TokenPrefix)] != TokenPrefix { return nil, fmt.Errorf("%w: not a free-tier token", ErrMalformed) diff --git a/gateway/internal/token/token_test.go b/gateway/internal/token/token_test.go index dc951f4..5eabb7b 100644 --- a/gateway/internal/token/token_test.go +++ b/gateway/internal/token/token_test.go @@ -109,7 +109,7 @@ func TestTokenIsBoundToItsSecret(t *testing.T) { // token — minting an identity with no proof of work at all. func TestChallengeCannotBePresentedAsAToken(t *testing.T) { s := testSigner(t) - c, err := s.NewChallenge(8) + c, err := s.NewChallenge(8, Bind("test")) if err != nil { t.Fatal(err) } @@ -124,7 +124,7 @@ func TestChallengeCannotBePresentedAsAToken(t *testing.T) { func TestChallengeCarriesItsOwnDifficulty(t *testing.T) { s := testSigner(t) - c, err := s.NewChallenge(14) + c, err := s.NewChallenge(14, Bind("test")) if err != nil { t.Fatal(err) } @@ -144,7 +144,7 @@ func TestChallengeCarriesItsOwnDifficulty(t *testing.T) { // cannot ask for an easier puzzle than the one it was given. func TestChallengeDifficultyCannotBeEdited(t *testing.T) { s := testSigner(t) - c, _ := s.NewChallenge(20) + c, _ := s.NewChallenge(20, Bind("test")) dot := strings.LastIndex(c.String, ".") forged := flip(c.String[:dot]) + "." + c.String[dot+1:] @@ -158,7 +158,7 @@ func TestChallengeExpires(t *testing.T) { now := time.Now() s.SetClock(func() time.Time { return now }) - c, _ := s.NewChallenge(8) + c, _ := s.NewChallenge(8, Bind("test")) if _, err := s.ParseChallenge(c.String, time.Minute); err != nil { t.Fatalf("fresh challenge rejected: %v", err) } @@ -171,7 +171,7 @@ func TestChallengeExpires(t *testing.T) { func TestProofOfWork(t *testing.T) { s := testSigner(t) - c, _ := s.NewChallenge(12) + c, _ := s.NewChallenge(12, Bind("test")) nonce, err := Solve(c.String, c.Difficulty, 1<<22) if err != nil { @@ -182,7 +182,7 @@ func TestProofOfWork(t *testing.T) { } // A solution to one challenge must not satisfy another, or one proof // of work would mint an unlimited supply. - other, _ := s.NewChallenge(12) + other, _ := s.NewChallenge(12, Bind("test")) if err := Verify(other.String, other.Difficulty, nonce); err == nil { t.Error("a nonce solved for one challenge also solved another") } @@ -250,7 +250,7 @@ func TestLeadingZeroBits(t *testing.T) { func TestNewChallengeRefusesAbsurdDifficulty(t *testing.T) { s := testSigner(t) - if _, err := s.NewChallenge(64); err == nil { + if _, err := s.NewChallenge(64, Bind("test")); err == nil { t.Fatal("issued a 64-bit challenge; nobody could ever solve it") } } diff --git a/internal/cli/root.go b/internal/cli/root.go index feae7d5..c33c8d4 100644 --- a/internal/cli/root.go +++ b/internal/cli/root.go @@ -88,9 +88,36 @@ func (o *Options) resolveAuth() (key, base string, free *deepseek.FreeTier, err if !ok { return "", "", nil, err } + if time.Since(f.Enrolled) > freeRenewAfter { + if renewed := o.renewFree(f); renewed != nil { + f = renewed + } + } return f.Token, f.BaseURL, f, nil } +// freeRenewAfter is the age at which an enrolment is quietly renewed. +// The hosted gateway expires tokens after 7 days; renewing at 6 means a +// regular user never sees the expiry error at all. +const freeRenewAfter = 6 * 24 * time.Hour + +// renewFree re-runs the enrolment — about a second of CPU — and saves +// the fresh token. Failure is not an error: the old token is kept, and +// if it has actually expired the gateway's own 401 says what to do. +func (o *Options) renewFree(old *deepseek.FreeTier) *deepseek.FreeTier { + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + defer cancel() + f, err := deepseek.NewFreeGateway(old.BaseURL, 30*time.Second).Enrol(ctx, nil) + if err != nil { + return nil + } + o.verbosef("free-tier enrolment renewed (subject %s)", f.Subject) + if err := f.Save(); err != nil { + o.verbosef("could not save the renewed enrolment: %v", err) + } + return f +} + // usingFree reports whether this run would go through the free tier, // without building a client. Commands use it to phrase their output for // someone who has no key at all.