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.