Pick any line to see why it is the way it is: the commit, the pull request and issue it came from, and what the agent was thinking.
| Merge branch 'worktree-agent-a8752162fea25f63f' into spend-guardrails | 1 | //! How often a client may call the API and the MCP server. |
| 2 | //! | |
| 3 | //! A request with a token counts against API_TOKEN_LIMIT under a hash of | |
| 4 | //! the token (never the token itself); one without counts against | |
| 5 | //! API_ANONYMOUS_LIMIT under the client's address (`CF-Connecting-IP`), as | |
| 6 | //! does one whose token turns out wrong, so guessing is limited by address. | |
| 7 | //! REST and MCP count apart. The limits are in `RATE_LIMITS` | |
| 8 | //! (packages/contracts/src/rate-limits.ts) and the docs' rate limits page. | |
| 9 | //! | |
| 10 | //! Not counted: what sandboxes, runners and outside systems send with | |
| 11 | //! credentials of their own (Stripe, connections' hooks, job tokens, | |
| 12 | //! report routes), which are answered before this or listed in | |
| 13 | //! [`counts`]. Without the bindings (self-hosted) nothing is limited, and | |
| 14 | //! a binding that fails lets the request through. | |
| 15 | ||
| 16 | use g1t_kit::limits::{self, PERIOD_SECONDS}; | |
| 17 | use serde_json::json; | |
| 18 | use sha2::{Digest, Sha256}; | |
| 19 | use worker::{Env, Request, Response, Result}; | |
| 20 | ||
| 21 | pub const ANONYMOUS: &str = "API_ANONYMOUS_LIMIT"; | |
| 22 | pub const TOKEN: &str = "API_TOKEN_LIMIT"; | |
| 23 | ||
| 24 | /// Paths a sandbox or runner reports to with its own token in the body: | |
| 25 | /// many sandboxes share an address, and none of them is a person. | |
| 26 | const REPORTS: &[&str] = &[ | |
| 27 | "/mergechecks/", | |
| 28 | "/backups/", | |
| 29 | "/queue/", | |
| 30 | "/actions/jobs/", | |
| 31 | "/agent-runs/", | |
| 32 | "/checks/", | |
| 33 | "/runs/", | |
| 34 | "/plans/", | |
| 35 | "/reviews/", | |
| 36 | "/runners/", | |
| 37 | ]; | |
| 38 | ||
| 39 | /// Whether a request counts against a limit at all. | |
| 40 | pub fn counts(method: &str, path: &str, has_token: bool) -> bool { | |
| 41 | has_token || method != "POST" || !REPORTS.iter().any(|prefix| path.starts_with(prefix)) | |
| 42 | } | |
| 43 | ||
| 44 | /// The binding a request counts against and its key there. | |
| 45 | pub fn key(token: Option<&str>, address: Option<&str>, on_mcp: bool) -> (&'static str, String) { | |
| 46 | let surface = if on_mcp { "mcp" } else { "rest" }; | |
| 47 | match token.filter(|token| !token.is_empty()) { | |
| 48 | Some(token) => (TOKEN, format!("{surface}:tok:{}", token_hash(token))), | |
| 49 | None => (ANONYMOUS, format!("{surface}:{}", limits::address_key(address))), | |
| 50 | } | |
| 51 | } | |
| 52 | ||
| 53 | /// The first 16 hex digits of the token's SHA-256. | |
| 54 | fn token_hash(token: &str) -> String { | |
| 55 | Sha256::digest(token.as_bytes())[..8].iter().map(|byte| format!("{byte:02x}")).collect() | |
| 56 | } | |
| 57 | ||
| 58 | /// The bearer token of an `Authorization` header, if it has one. | |
| 59 | pub fn bearer(header: &str) -> Option<&str> { | |
| 60 | match header.split_once(' ') { | |
| 61 | Some((scheme, token)) if scheme.eq_ignore_ascii_case("bearer") => Some(token.trim()).filter(|t| !t.is_empty()), | |
| 62 | _ => None, | |
| 63 | } | |
| 64 | } | |
| 65 | ||
| 66 | /// The 429 every limit answers, in the shape every error takes. | |
| 67 | pub fn too_many(signed_in: bool) -> Result<Response> { | |
| 68 | let message = if signed_in { | |
| 69 | "Too many requests with this token. Wait a minute and try again: https://docs.g1t.sh/reference/rate-limits/" | |
| 70 | } else { | |
| 71 | "Too many requests from this address. Wait a minute, or use an access token for a higher limit: https://docs.g1t.sh/reference/rate-limits/" | |
| 72 | }; | |
| 73 | let mut response = Response::from_json(&json!({ "error": { "code": "rate_limited", "message": message } }))?.with_status(429); | |
| 74 | response.headers_mut().set("retry-after", &PERIOD_SECONDS.to_string())?; | |
| 75 | Ok(response) | |
| 76 | } | |
| 77 | ||
| 78 | /// The 429 for a request past its limit, or `None` to go on. | |
| 79 | pub async fn limited(request: &Request, env: &Env, method: &str, path: &str, on_mcp: bool) -> Result<Option<Response>> { | |
| 80 | let header = request.headers().get("authorization")?.unwrap_or_default(); | |
| 81 | let token = bearer(&header); | |
| 82 | if !counts(method, path, token.is_some()) { | |
| 83 | return Ok(None); | |
| 84 | } | |
| 85 | let address = request.headers().get("cf-connecting-ip")?; | |
| 86 | let (binding, key) = key(token, address.as_deref(), on_mcp); | |
| 87 | if limits::check(env, binding, key).await.limited() { | |
| 88 | return too_many(token.is_some()).map(Some); | |
| 89 | } | |
| 90 | Ok(None) | |
| 91 | } | |
| 92 | ||
| 93 | /// A request whose token was wrong also counts against its address. | |
| 94 | pub async fn wrong_token(request: &Request, env: &Env, on_mcp: bool) -> Result<Option<Response>> { | |
| 95 | let address = request.headers().get("cf-connecting-ip")?; | |
| 96 | let (binding, key) = key(None, address.as_deref(), on_mcp); | |
| 97 | if limits::check(env, binding, key).await.limited() { | |
| 98 | return too_many(false).map(Some); | |
| 99 | } | |
| 100 | Ok(None) | |
| 101 | } | |
| 102 | ||
| 103 | #[cfg(test)] | |
| 104 | mod tests { | |
| 105 | use super::*; | |
| 106 | ||
| 107 | #[test] | |
| 108 | fn tokens_are_counted_by_their_hash_and_others_by_address() { | |
| 109 | let (binding, tok) = key(Some("g1t_secret"), Some("203.0.113.9"), false); | |
| 110 | assert_eq!(binding, TOKEN); | |
| 111 | assert!(tok.starts_with("rest:tok:") && tok.len() == "rest:tok:".len() + 16); | |
| 112 | assert!(!tok.contains("g1t_secret"), "the token never reaches the limiter"); | |
| 113 | assert_eq!(key(Some("g1t_secret"), None, false).1, tok, "the same token, the same key"); | |
| 114 | assert_ne!(key(Some("g1t_other"), None, false).1, tok); | |
| 115 | assert_eq!(key(None, Some("203.0.113.9"), false), (ANONYMOUS, "rest:ip:203.0.113.9".to_owned())); | |
| 116 | assert_eq!(key(Some(""), None, false), (ANONYMOUS, "rest:ip:unknown".to_owned())); | |
| 117 | } | |
| 118 | ||
| 119 | #[test] | |
| 120 | fn rest_and_mcp_count_apart() { | |
| 121 | assert_eq!(key(None, Some("203.0.113.9"), true).1, "mcp:ip:203.0.113.9"); | |
| 122 | assert!(key(Some("g1t_secret"), None, true).1.starts_with("mcp:tok:")); | |
| 123 | } | |
| 124 | ||
| 125 | #[test] | |
| 126 | fn sandbox_reports_are_not_counted_but_everything_else_is() { | |
| 127 | assert!(!counts("POST", "/checks/run_1", false)); | |
| 128 | assert!(!counts("POST", "/agent-runs/run_1/report", false)); | |
| 129 | assert!(!counts("POST", "/runs/run_1/usage", false)); | |
| 130 | assert!(counts("GET", "/repos/acme/rocket", false)); | |
| 131 | assert!(counts("POST", "/device/code", false)); | |
| 132 | assert!(counts("POST", "/checks/run_1", true), "with a bearer token it counts as that token"); | |
| 133 | } | |
| 134 | ||
| 135 | #[test] | |
| 136 | fn only_bearer_tokens_are_read() { | |
| 137 | assert_eq!(bearer("Bearer g1t_abc "), Some("g1t_abc")); | |
| 138 | assert_eq!(bearer("bearer g1t_abc"), Some("g1t_abc")); | |
| 139 | assert_eq!(bearer("Basic dXNlcjpwYXNz"), None); | |
| 140 | assert_eq!(bearer("Bearer "), None); | |
| 141 | assert_eq!(bearer(""), None); | |
| 142 | } | |
| 143 | } |
This file's history is long; its oldest lines are credited to the oldest commit read.