diff --git a/Cargo.lock b/Cargo.lock index 13eb343..6a0c79c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -944,7 +944,7 @@ dependencies = [ [[package]] name = "ghpool" -version = "0.3.0" +version = "0.3.3" dependencies = [ "aws-config", "aws-sdk-secretsmanager", diff --git a/config.example.toml b/config.example.toml index b3ae500..e16d0cb 100644 --- a/config.example.toml +++ b/config.example.toml @@ -28,3 +28,8 @@ run_list_ttl_secs = 15 commit_list_ttl_secs = 120 repo_view_ttl_secs = 300 default_ttl_secs = 60 +# Raw (non-JSON) response cache — used for `gh pr diff`, patches, etc. +# Scoped per (path, query, Accept header, identity) to avoid cross-identity +# leakage and to correctly distinguish paginated/variant requests. +raw_ttl_secs = 30 +raw_max_bytes = 268435456 # 256 MiB — total byte budget, enforced via weigher diff --git a/src/cache.rs b/src/cache.rs index 13be6bb..ad880ae 100644 --- a/src/cache.rs +++ b/src/cache.rs @@ -20,6 +20,7 @@ pub enum RouteKind { pub struct Cache { store: MokaCache, + raw_store: MokaCache, ttls: CacheTtls, hits: AtomicU64, misses: AtomicU64, @@ -40,8 +41,25 @@ impl Cache { .max_capacity(config.max_entries) .time_to_live(Duration::from_secs(config.default_ttl_secs)) .build(); + // Raw (diff/patch) cache: bounded primarily by total bytes via a + // weigher, since diff bodies can be multi-MB (unlike JSON metadata + // responses). moka's max_capacity here tracks weighted bytes, not + // entry count. + let raw_store = MokaCache::builder() + .max_capacity(config.raw_max_bytes) + .weigher(move |key: &String, value: &String| -> u32 { + (key.len() + value.len()).min(u32::MAX as usize) as u32 + }) + .time_to_live(Duration::from_secs(config.raw_ttl_secs)) + .build(); + tracing::debug!( + raw_max_bytes = config.raw_max_bytes, + raw_ttl_secs = config.raw_ttl_secs, + "raw cache configured" + ); Self { store, + raw_store, ttls: CacheTtls { pr_view: Duration::from_secs(config.pr_view_ttl_secs), issue_list: Duration::from_secs(config.issue_list_ttl_secs), @@ -88,7 +106,36 @@ impl Cache { CacheStats { hits: self.hits.load(Ordering::Relaxed), misses: self.misses.load(Ordering::Relaxed), - entries: self.store.entry_count(), + entries: self.store.entry_count() + self.raw_store.entry_count(), + } + } + + /// Fetch-or-compute with stampede protection: concurrent callers with the + /// same key share a single in-flight future via moka's `entry` API, + /// so N simultaneous cache misses only trigger one upstream request. + /// Uses `is_fresh()` to correctly distinguish hits from misses. + pub async fn get_or_insert_raw( + &self, + key: &str, + init: F, + ) -> Result + where + F: std::future::Future>, + E: Clone + std::fmt::Debug + Send + Sync + 'static, + { + match self.raw_store.entry_by_ref(key).or_try_insert_with(init).await { + Ok(entry) => { + if entry.is_fresh() { + self.misses.fetch_add(1, Ordering::Relaxed); + } else { + self.hits.fetch_add(1, Ordering::Relaxed); + } + Ok(entry.into_value()) + } + Err(e) => { + self.misses.fetch_add(1, Ordering::Relaxed); + Err((*e).clone()) + } } } } @@ -107,6 +154,30 @@ pub fn build_key(path: &str, query: &HashMap) -> String { if qs.is_empty() { path.to_string() } else { format!("{}?{}", path, qs) } } +/// Build a cache key for raw (non-JSON) responses. +/// +/// Includes: +/// - `path` + sorted `query` params (consistent with `build_key`, fixes a gap +/// where paginated raw requests would previously collide on the same key) +/// - normalized `accept` header (lowercased, so `Application/Vnd.Github.V3.Diff` +/// and `application/vnd.github.v3.diff` share one cache entry) +/// - `identity_scope`: a caller-supplied token identifying which PAT/identity's +/// access scope produced the response. This prevents cross-identity leakage +/// when the pool holds PATs with different repo access — a response fetched +/// with a broadly-scoped PAT must never be served to a caller whose own PAT +/// would have been denied access. +pub fn build_raw_key(path: &str, query: &HashMap, accept: &str, identity_scope: &str) -> String { + let mut parts: Vec<(&str, &str)> = query.iter().map(|(k, v)| (k.as_str(), v.as_str())).collect(); + parts.sort(); + let qs: String = parts.iter().map(|(k, v)| format!("{}={}", k, v)).collect::>().join("&"); + let normalized_accept = accept.trim().to_ascii_lowercase(); + if qs.is_empty() { + format!("raw:{}:{}:{}", path, normalized_accept, identity_scope) + } else { + format!("raw:{}?{}:{}:{}", path, qs, normalized_accept, identity_scope) + } +} + pub fn build_graphql_key(body: &[u8]) -> String { use std::hash::{Hash, Hasher}; let mut hasher = std::collections::hash_map::DefaultHasher::new(); @@ -129,3 +200,126 @@ pub fn classify_route(path: &str) -> RouteKind { _ => RouteKind::Other, } } + +#[cfg(test)] +mod tests { + use super::*; + + fn empty_query() -> HashMap { + HashMap::new() + } + + #[test] + fn test_build_raw_key_no_query() { + let key = build_raw_key("/repos/o/r/pulls/1", &empty_query(), "application/vnd.github.v3.diff", "id1"); + assert_eq!(key, "raw:/repos/o/r/pulls/1:application/vnd.github.v3.diff:id1"); + } + + #[test] + fn test_build_raw_key_normalizes_accept_case() { + let a = build_raw_key("/repos/o/r/pulls/1", &empty_query(), "Application/Vnd.Github.V3.Diff", "id1"); + let b = build_raw_key("/repos/o/r/pulls/1", &empty_query(), "application/vnd.github.v3.diff", "id1"); + assert_eq!(a, b); + } + + #[test] + fn test_build_raw_key_includes_query_params() { + let mut q = HashMap::new(); + q.insert("page".to_string(), "2".to_string()); + let key = build_raw_key("/repos/o/r/pulls/1", &q, "application/vnd.github.v3.diff", "id1"); + assert_eq!(key, "raw:/repos/o/r/pulls/1?page=2:application/vnd.github.v3.diff:id1"); + } + + #[test] + fn test_build_raw_key_differs_by_query() { + let mut q1 = HashMap::new(); + q1.insert("page".to_string(), "1".to_string()); + let mut q2 = HashMap::new(); + q2.insert("page".to_string(), "2".to_string()); + let k1 = build_raw_key("/repos/o/r/pulls/1", &q1, "application/vnd.github.v3.diff", "id1"); + let k2 = build_raw_key("/repos/o/r/pulls/1", &q2, "application/vnd.github.v3.diff", "id1"); + assert_ne!(k1, k2); + } + + #[test] + fn test_build_raw_key_differs_by_identity_scope() { + let k1 = build_raw_key("/repos/o/r/pulls/1", &empty_query(), "application/vnd.github.v3.diff", "id1"); + let k2 = build_raw_key("/repos/o/r/pulls/1", &empty_query(), "application/vnd.github.v3.diff", "id2"); + assert_ne!(k1, k2, "responses from different identity scopes must not share a cache entry"); + } + + #[test] + fn test_build_raw_key_differs_by_accept() { + let k1 = build_raw_key("/repos/o/r/pulls/1", &empty_query(), "application/vnd.github.v3.diff", "id1"); + let k2 = build_raw_key("/repos/o/r/pulls/1", &empty_query(), "application/json", "id1"); + assert_ne!(k1, k2); + } + + #[tokio::test] + async fn test_cache_raw_get_insert_roundtrip() { + let cfg = CacheConfig::default(); + let cache = Cache::new(&cfg); + let key = build_raw_key("/repos/o/r/pulls/1", &empty_query(), "application/vnd.github.v3.diff", "id1"); + let call_count = std::sync::Arc::new(AtomicU64::new(0)); + let cc1 = call_count.clone(); + let v1 = cache + .get_or_insert_raw(&key, async move { + cc1.fetch_add(1, Ordering::SeqCst); + Ok::("diff content".to_string()) + }) + .await + .unwrap(); + assert_eq!(v1, "diff content"); + // First call should be a miss (fresh insert). + assert_eq!(cache.stats().misses, 1); + assert_eq!(cache.stats().hits, 0); + // Second call should hit cache, not invoke init again. + let cc2 = call_count.clone(); + let v2 = cache + .get_or_insert_raw(&key, async move { + cc2.fetch_add(1, Ordering::SeqCst); + Ok::("should not be used".to_string()) + }) + .await + .unwrap(); + assert_eq!(v2, "diff content"); + assert_eq!(call_count.load(Ordering::SeqCst), 1); + assert_eq!(cache.stats().hits, 1); + assert_eq!(cache.stats().misses, 1); + } + + #[tokio::test] + async fn test_cache_raw_stampede_dedup() { + // Concurrent callers requesting the same uncached key should only + // trigger the init closure once (moka's entry API coalesces). + let cfg = CacheConfig::default(); + let cache = std::sync::Arc::new(Cache::new(&cfg)); + let call_count = std::sync::Arc::new(AtomicU64::new(0)); + let key = build_raw_key("/repos/o/r/pulls/1", &empty_query(), "application/vnd.github.v3.diff", "id1"); + + let mut handles = Vec::new(); + for _ in 0..10 { + let cache = cache.clone(); + let call_count = call_count.clone(); + let key = key.clone(); + handles.push(tokio::spawn(async move { + cache + .get_or_insert_raw(&key, async move { + call_count.fetch_add(1, Ordering::SeqCst); + tokio::time::sleep(Duration::from_millis(20)).await; + Ok::("diff".to_string()) + }) + .await + })); + } + for h in handles { + h.await.unwrap().unwrap(); + } + // With entry API, all concurrent calls coalesce: only 1 init runs. + // The one that evaluated is_fresh()=true counts as a miss. + // The rest see is_fresh()=false and count as hits. + let stats = cache.stats(); + assert_eq!(stats.misses, 1, "only one call should be a miss (the one that ran init)"); + assert_eq!(stats.hits, 9, "the rest should be hits (served from cache after init completed)"); + } +} diff --git a/src/config.rs b/src/config.rs index 392d3d1..56bd965 100644 --- a/src/config.rs +++ b/src/config.rs @@ -31,6 +31,12 @@ pub struct CacheConfig { pub repo_view_ttl_secs: u64, #[serde(default = "default_ttl")] pub default_ttl_secs: u64, + /// TTL for raw (non-JSON, e.g. diff/patch) responses. + #[serde(default = "default_raw_ttl")] + pub raw_ttl_secs: u64, + /// Max total bytes held by the raw response cache (weigher-enforced). + #[serde(default = "default_raw_max_bytes")] + pub raw_max_bytes: u64, } impl Default for CacheConfig { @@ -43,6 +49,8 @@ impl Default for CacheConfig { commit_list_ttl_secs: default_commit_ttl(), repo_view_ttl_secs: default_repo_ttl(), default_ttl_secs: default_ttl(), + raw_ttl_secs: default_raw_ttl(), + raw_max_bytes: default_raw_max_bytes(), } } } @@ -51,6 +59,8 @@ fn default_port() -> u16 { 8080 } fn default_max_entries() -> u64 { 10000 } fn default_pr_ttl() -> u64 { 30 } fn default_run_ttl() -> u64 { 15 } +fn default_raw_ttl() -> u64 { 30 } +fn default_raw_max_bytes() -> u64 { 256 * 1024 * 1024 } // 256 MiB fn default_commit_ttl() -> u64 { 120 } fn default_repo_ttl() -> u64 { 300 } fn default_ttl() -> u64 { 60 } diff --git a/src/main.rs b/src/main.rs index b7c61f0..cc59502 100644 --- a/src/main.rs +++ b/src/main.rs @@ -105,7 +105,7 @@ async fn proxy( // Forward request let mut req = state.http.get(&url) .header("Authorization", format!("Bearer {}", identity.token)) - .header("User-Agent", "ghpool/0.1.0") + .header("User-Agent", concat!("ghpool/", env!("CARGO_PKG_VERSION"))) .header("Accept", "application/vnd.github+json"); if let Some(version) = headers.get("x-github-api-version") { @@ -160,7 +160,18 @@ async fn proxy_raw( return Err(StatusCode::FORBIDDEN); } + let accept = headers.get("accept") + .and_then(|v| v.to_str().ok()) + .unwrap_or("application/vnd.github.v3.diff") + .to_string(); + + // Select the identity BEFORE checking the cache. The cache key is scoped + // to this identity so a response fetched under one PAT's access scope is + // never served to a caller that would resolve to a different identity + // (prevents cross-identity leakage when the pool holds PATs with + // different repo access). let identity = state.pool.select().map_err(|_| StatusCode::SERVICE_UNAVAILABLE)?; + let cache_key = cache::build_raw_key(&api_path, &query, &accept, &identity.id); let mut url = format!("https://api.github.com{}", api_path); if !query.is_empty() { @@ -168,43 +179,59 @@ async fn proxy_raw( url = format!("{}?{}", url, qs.join("&")); } - let accept = headers.get("accept") - .and_then(|v| v.to_str().ok()) - .unwrap_or("application/vnd.github.v3.diff"); + let state_for_fetch = state.clone(); + let token = identity.token.clone(); + let identity_id = identity.id.clone(); + let api_path_for_log = api_path.clone(); + + let result = state.cache.get_or_insert_raw(&cache_key, async move { + let resp = state_for_fetch.http.get(&url) + .header("Authorization", format!("Bearer {}", token)) + .header("User-Agent", concat!("ghpool/", env!("CARGO_PKG_VERSION"))) + .header("Accept", &accept) + .send() + .await + .map_err(|e| format!("github request failed: {e}"))?; - let resp = state.http.get(&url) - .header("Authorization", format!("Bearer {}", identity.token)) - .header("User-Agent", "ghpool/0.1.0") - .header("Accept", accept) - .send() - .await - .map_err(|e| { - tracing::error!("github request failed: {}", e); - StatusCode::BAD_GATEWAY - })?; + let rate_remaining = resp.headers() + .get("x-ratelimit-remaining") + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.parse::().ok()); + let rate_reset = resp.headers() + .get("x-ratelimit-reset") + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.parse::().ok()); + state_for_fetch.pool.update_rate(&identity_id, rate_remaining, rate_reset); - let rate_remaining = resp.headers() - .get("x-ratelimit-remaining") - .and_then(|v| v.to_str().ok()) - .and_then(|v| v.parse::().ok()); - let rate_reset = resp.headers() - .get("x-ratelimit-reset") - .and_then(|v| v.to_str().ok()) - .and_then(|v| v.parse::().ok()); - state.pool.update_rate(&identity.id, rate_remaining, rate_reset); + let status = resp.status(); + let body = resp.text().await.map_err(|_| "failed to read response body".to_string())?; - let status = resp.status(); - let body = resp.text().await.map_err(|_| StatusCode::BAD_GATEWAY)?; + if !status.is_success() { + tracing::warn!("github returned {}: {}", status, api_path_for_log); + // Encode the status so the caller can map it back to an HTTP error. + return Err(format!("upstream_status:{}", status.as_u16())); + } - if !status.is_success() { - tracing::warn!("github returned {}: {}", status, api_path); - return Err(StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::BAD_GATEWAY)); - } + Ok(body) + }).await; - tracing::info!("200 OK {} [raw, via {}]", api_path, identity.id); - let mut resp_headers = HeaderMap::new(); - resp_headers.insert("content-type", "text/plain".parse().unwrap()); - Ok((StatusCode::OK, resp_headers, body)) + match result { + Ok(body) => { + tracing::info!("200 OK {} [raw, via {}]", api_path, identity.id); + let mut resp_headers = HeaderMap::new(); + resp_headers.insert("content-type", "text/plain".parse().unwrap()); + Ok((StatusCode::OK, resp_headers, body)) + } + Err(e) => { + if let Some(code_str) = e.strip_prefix("upstream_status:") { + if let Ok(code) = code_str.parse::() { + return Err(StatusCode::from_u16(code).unwrap_or(StatusCode::BAD_GATEWAY)); + } + } + tracing::error!("raw proxy failed: {}", e); + Err(StatusCode::BAD_GATEWAY) + } + } } fn is_allowed_path(path: &str, allowed_owners: &[String]) -> bool { @@ -254,7 +281,7 @@ async fn graphql_proxy( let resp = state.http.post("https://api.github.com/graphql") .header("Authorization", &auth_header) - .header("User-Agent", "ghpool/0.1.0") + .header("User-Agent", concat!("ghpool/", env!("CARGO_PKG_VERSION"))) .header("Content-Type", "application/json") .body(body.to_vec()) .send() @@ -302,7 +329,7 @@ async fn resolve_token_user(state: &AppState, auth_header: &str) -> String { } let user = match state.http.get("https://api.github.com/user") .header("Authorization", auth_header) - .header("User-Agent", "ghpool/0.1.0") + .header("User-Agent", concat!("ghpool/", env!("CARGO_PKG_VERSION"))) .send() .await {