| 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 | #[derive(Clone, Copy)] |
| 25 | pub struct Policy { signed_in: bool, limited: bool } |
| 26 | |
| 27 | impl Policy { |
| 28 | fn name(self) -> &'static str { if self.signed_in { "token" } else { "anonymous" } } |
| 29 | fn header(self) -> String { |
| 30 | format!("\"{}\";q={};w={PERIOD_SECONDS}", self.name(), if self.signed_in { 1000 } else { 60 }) |
| 31 | } |
| 32 | fn remaining(self) -> Option<String> { |
| 33 | self.limited.then(|| format!("\"{}\";r=0;t={PERIOD_SECONDS}", self.name())) |
| 34 | } |
| 35 | pub fn apply(self, response: &mut Response) -> Result<()> { |
| 36 | response.headers_mut().set("ratelimit-policy", &self.header())?; |
| 37 | // The binding exposes a verdict, not remaining quota or a reset clock. |
| 38 | // Only advertise remaining/reset information when actually limited. |
| 39 | if let Some(remaining) = self.remaining() { |
| 40 | response.headers_mut().set("ratelimit", &remaining)?; |
| 41 | } |
| 42 | Ok(()) |
| 43 | } |
| 44 | } |
| 45 | |
| 46 | fn policy(verdict: limits::Verdict, signed_in: bool) -> Option<Policy> { |
| 47 | match verdict { |
| 48 | limits::Verdict::Unavailable => None, |
| 49 | _ => Some(Policy { signed_in, limited: verdict.limited() }), |
| 50 | } |
| 51 | } |
| 52 | |
| 53 | /// Paths a sandbox or runner reports to with its own token in the body: |
| 54 | /// many sandboxes share an address, and none of them is a person. |
| 55 | const REPORTS: &[&str] = &[ |
| 56 | "/mergechecks/", |
| 57 | "/backups/", |
| 58 | "/queue/", |
| 59 | "/actions/jobs/", |
| 60 | "/agent-runs/", |
| 61 | "/checks/", |
| 62 | "/runs/", |
| 63 | "/plans/", |
| 64 | "/reviews/", |
| 65 | "/runners/", |
| 66 | ]; |
| 67 | |
| 68 | /// Whether a request counts against a limit at all. |
| 69 | pub fn counts(method: &str, path: &str, has_token: bool) -> bool { |
| 70 | has_token || method != "POST" || !REPORTS.iter().any(|prefix| path.starts_with(prefix)) |
| 71 | } |
| 72 | |
| 73 | /// The binding a request counts against and its key there. |
| 74 | pub fn key(token: Option<&str>, address: Option<&str>, on_mcp: bool) -> (&'static str, String) { |
| 75 | let surface = if on_mcp { "mcp" } else { "rest" }; |
| 76 | match token.filter(|token| !token.is_empty()) { |
| 77 | Some(token) => (TOKEN, format!("{surface}:tok:{}", token_hash(token))), |
| 78 | None => (ANONYMOUS, format!("{surface}:{}", limits::address_key(address))), |
| 79 | } |
| 80 | } |
| 81 | |
| 82 | /// The first 16 hex digits of the token's SHA-256. |
| 83 | fn token_hash(token: &str) -> String { |
| 84 | Sha256::digest(token.as_bytes())[..8].iter().map(|byte| format!("{byte:02x}")).collect() |
| 85 | } |
| 86 | |
| 87 | /// The access token of an `Authorization` header, if it has one: `Bearer |
| 88 | /// <token>`, or `token <token>`, the form GitHub also takes and Octokit |
| 89 | /// (`@actions/github`, so `actions/checkout` and `actions/github-script`) |
| 90 | /// sends. |
| 91 | pub fn bearer(header: &str) -> Option<&str> { |
| 92 | match header.split_once(' ') { |
| 93 | Some((scheme, token)) if scheme.eq_ignore_ascii_case("bearer") || scheme.eq_ignore_ascii_case("token") => { |
| 94 | Some(token.trim()).filter(|t| !t.is_empty()) |
| 95 | } |
| 96 | _ => None, |
| 97 | } |
| 98 | } |
| 99 | |
| 100 | /// The 429 every limit answers, in the shape every error takes. |
| 101 | pub fn too_many(signed_in: bool) -> Result<Response> { |
| 102 | let message = if signed_in { |
| 103 | "Too many requests with this token. Wait a minute and try again: https://docs.g1t.sh/reference/rate-limits/" |
| 104 | } else { |
| 105 | "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/" |
| 106 | }; |
| 107 | let mut response = Response::from_json(&json!({ "error": { "code": "rate_limited", "message": message } }))?.with_status(429); |
| 108 | response.headers_mut().set("retry-after", &PERIOD_SECONDS.to_string())?; |
| 109 | Ok(response) |
| 110 | } |
| 111 | |
| 112 | /// The 429 for a request past its limit, or `None` to go on. |
| 113 | pub async fn limited(request: &Request, env: &Env, method: &str, path: &str, on_mcp: bool, advertised: &mut Option<Policy>) -> Result<Option<Response>> { |
| 114 | let header = request.headers().get("authorization")?.unwrap_or_default(); |
| 115 | let token = bearer(&header); |
| 116 | if !counts(method, path, token.is_some()) { |
| 117 | return Ok(None); |
| 118 | } |
| 119 | let address = request.headers().get("cf-connecting-ip")?; |
| 120 | let (binding, key) = key(token, address.as_deref(), on_mcp); |
| 121 | let verdict = limits::check(env, binding, key).await; |
| 122 | *advertised = policy(verdict, token.is_some()); |
| 123 | if verdict.limited() { |
| 124 | return too_many(token.is_some()).map(Some); |
| 125 | } |
| 126 | Ok(None) |
| 127 | } |
| 128 | |
| 129 | /// A request whose token was wrong also counts against its address. |
| 130 | pub async fn wrong_token(request: &Request, env: &Env, on_mcp: bool, advertised: &mut Option<Policy>) -> Result<Option<Response>> { |
| 131 | let address = request.headers().get("cf-connecting-ip")?; |
| 132 | let (binding, key) = key(None, address.as_deref(), on_mcp); |
| 133 | let verdict = limits::check(env, binding, key).await; |
| 134 | *advertised = policy(verdict, false); |
| 135 | if verdict.limited() { |
| 136 | return too_many(false).map(Some); |
| 137 | } |
| 138 | Ok(None) |
| 139 | } |
| 140 | |
| 141 | #[cfg(test)] |
| 142 | mod tests { |
| 143 | use super::*; |
| 144 | |
| 145 | #[test] |
| 146 | fn advertised_limits_match_bindings_without_inventing_remaining_counts() { |
| 147 | let config: String = include_str!("../wrangler.jsonc").chars().filter(|c| !c.is_whitespace()).collect(); |
| 148 | for (signed_in, binding, quota) in [(false, ANONYMOUS, 60), (true, TOKEN, 1000)] { |
| 149 | let configured = config.split(&format!("\"name\":\"{binding}\"")).nth(1).unwrap().split('}').next().unwrap(); |
| 150 | assert!(configured.contains(&format!("\"limit\":{quota},\"period\":{PERIOD_SECONDS}"))); |
| 151 | let allowed = policy(limits::Verdict::Allowed, signed_in).unwrap(); |
| 152 | assert!(allowed.header().ends_with(&format!(";q={quota};w=60"))); |
| 153 | assert_eq!(allowed.remaining(), None); |
| 154 | assert!(policy(limits::Verdict::Limited, signed_in).unwrap().remaining().unwrap().ends_with(";r=0;t=60")); |
| 155 | } |
| 156 | assert!(policy(limits::Verdict::Unavailable, false).is_none()); |
| 157 | } |
| 158 | |
| 159 | #[test] |
| 160 | fn tokens_are_counted_by_their_hash_and_others_by_address() { |
| 161 | let (binding, tok) = key(Some("g1t_secret"), Some("203.0.113.9"), false); |
| 162 | assert_eq!(binding, TOKEN); |
| 163 | assert!(tok.starts_with("rest:tok:") && tok.len() == "rest:tok:".len() + 16); |
| 164 | assert!(!tok.contains("g1t_secret"), "the token never reaches the limiter"); |
| 165 | assert_eq!(key(Some("g1t_secret"), None, false).1, tok, "the same token, the same key"); |
| 166 | assert_ne!(key(Some("g1t_other"), None, false).1, tok); |
| 167 | assert_eq!(key(None, Some("203.0.113.9"), false), (ANONYMOUS, "rest:ip:203.0.113.9".to_owned())); |
| 168 | assert_eq!(key(Some(""), None, false), (ANONYMOUS, "rest:ip:unknown".to_owned())); |
| 169 | } |
| 170 | |
| 171 | #[test] |
| 172 | fn rest_and_mcp_count_apart() { |
| 173 | assert_eq!(key(None, Some("203.0.113.9"), true).1, "mcp:ip:203.0.113.9"); |
| 174 | assert!(key(Some("g1t_secret"), None, true).1.starts_with("mcp:tok:")); |
| 175 | } |
| 176 | |
| 177 | #[test] |
| 178 | fn sandbox_reports_are_not_counted_but_everything_else_is() { |
| 179 | assert!(!counts("POST", "/checks/run_1", false)); |
| 180 | assert!(!counts("POST", "/agent-runs/run_1/report", false)); |
| 181 | assert!(!counts("POST", "/runs/run_1/usage", false)); |
| 182 | assert!(counts("GET", "/repos/acme/rocket", false)); |
| 183 | assert!(counts("POST", "/device/code", false)); |
| 184 | assert!(counts("POST", "/checks/run_1", true), "with a bearer token it counts as that token"); |
| 185 | } |
| 186 | |
| 187 | #[test] |
| 188 | fn bearer_and_token_schemes_are_read() { |
| 189 | assert_eq!(bearer("Bearer g1t_abc "), Some("g1t_abc")); |
| 190 | assert_eq!(bearer("bearer g1t_abc"), Some("g1t_abc")); |
| 191 | // Octokit's form, as `actions/checkout` sends it for a default branch. |
| 192 | assert_eq!(bearer("token g1t_abc"), Some("g1t_abc")); |
| 193 | assert_eq!(bearer("Token g1t_abc"), Some("g1t_abc")); |
| 194 | assert_eq!(bearer("token "), None); |
| 195 | assert_eq!(bearer("Basic dXNlcjpwYXNz"), None); |
| 196 | assert_eq!(bearer("Bearer "), None); |
| 197 | assert_eq!(bearer(""), None); |
| 198 | } |
| 199 | } |