Skip to content
Sign up

flagon-io/g1t

public
0
143 linesCodeBlameRaw
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
16use g1t_kit::limits::{self, PERIOD_SECONDS};
17use serde_json::json;
18use sha2::{Digest, Sha256};
19use worker::{Env, Request, Response, Result};
20
21pub const ANONYMOUS: &str = "API_ANONYMOUS_LIMIT";
22pub 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.
26const 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.
40pub 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.
45pub 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.
54fn 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.
59pub 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.
67pub 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.
79pub 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.
94pub 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)]
104mod 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}