diff --git a/.gitignore b/.gitignore index 0fcb505..f3eba0a 100644 --- a/.gitignore +++ b/.gitignore @@ -5,4 +5,5 @@ runtime/*.jsonl runtime/*.log runtime/attack_events.jsonl +configs/runtime/ diff --git a/src/adaptive.rs b/src/adaptive.rs index 9129b06..04026b1 100644 --- a/src/adaptive.rs +++ b/src/adaptive.rs @@ -228,7 +228,9 @@ impl AdaptiveDetector { poisoned.into_inner() } }; - if !windows.entries.contains_key(signature) { + let window = if let Some(window) = windows.entries.get_mut(signature) { + window + } else { if !ensure_signature_window_capacity( &mut windows, self.max_signature_windows_per_shard, @@ -240,18 +242,19 @@ impl AdaptiveDetector { emit: false, }; } - windows.order.push_back(signature.to_string()); - } - let window = windows - .entries - .entry(signature.to_string()) - .or_insert_with(|| SignatureWindow { - bucket: adaptive_bucket(self.threshold_per_second, now), - observed_count: 0, - last_event: None, - last_strong_event: None, - last_seen: now, - }); + let signature_owned = signature.to_string(); + windows.order.push_back(signature_owned.clone()); + windows + .entries + .entry(signature_owned) + .or_insert_with(|| SignatureWindow { + bucket: adaptive_bucket(self.threshold_per_second, now), + observed_count: 0, + last_event: None, + last_strong_event: None, + last_seen: now, + }) + }; window.last_seen = now; window.observed_count = window.observed_count.saturating_add(1); let count = window.observed_count; @@ -294,7 +297,9 @@ impl AdaptiveDetector { poisoned.into_inner() } }; - if !windows.entries.contains_key(path_shape) { + let window = if let Some(window) = windows.entries.get_mut(path_shape) { + window + } else { if !ensure_shape_window_capacity( &mut windows, self.max_path_shape_windows_per_shard, @@ -302,20 +307,21 @@ impl AdaptiveDetector { ) { return None; } - windows.order.push_back(path_shape.to_string()); - } - let window = windows - .entries - .entry(path_shape.to_string()) - .or_insert_with(|| ShapeWindow { - bucket: adaptive_bucket(self.threshold_per_second, now), - strong_bucket: adaptive_bucket(self.threshold_per_second, now), - observed_count: 0, - strong_count: 0, - sample_window_start: now, - emitted_samples: 0, - last_seen: now, - }); + let path_shape_owned = path_shape.to_string(); + windows.order.push_back(path_shape_owned.clone()); + windows + .entries + .entry(path_shape_owned) + .or_insert_with(|| ShapeWindow { + bucket: adaptive_bucket(self.threshold_per_second, now), + strong_bucket: adaptive_bucket(self.threshold_per_second, now), + observed_count: 0, + strong_count: 0, + sample_window_start: now, + emitted_samples: 0, + last_seen: now, + }) + }; window.last_seen = now; window.observed_count = window.observed_count.saturating_add(1); if now.saturating_duration_since(window.sample_window_start) >= Duration::from_secs(1) { @@ -1266,6 +1272,69 @@ mod tests { } } + fn signature_observed_count(detector: &AdaptiveDetector, signature: &str) -> u64 { + let shard_idx = shard_for(signature); + let windows = match detector.windows[shard_idx].lock() { + Ok(windows) => windows, + Err(poisoned) => poisoned.into_inner(), + }; + windows + .entries + .get(signature) + .map(|window| window.observed_count) + .unwrap_or(0) + } + + #[tokio::test] + async fn body_too_large_hot_path_observes_before_strong_evidence() { + let engine = FilterEngine::new( + vec![], + PathBuf::from("/tmp/altura-prot-nonexistent-body-too-large-filters.json"), + Duration::from_secs(30), + ) + .await; + let logger = Arc::new( + EventLogger::new("/tmp/altura-prot-test-body-too-large-events.jsonl").unwrap(), + ); + let detector = AdaptiveDetector::new( + test_detector_config(60, 8_192, 8_192), + Arc::clone(&engine), + logger, + ); + let headers = HeaderMap::new(); + let signature = crate::filter::request_signature("POST", "/upload", None, &headers); + let ctx = RequestContext { + client_ip: "127.0.0.1".parse().unwrap(), + method: "POST", + path: "/upload", + query: None, + headers: &headers, + signature: signature.clone(), + }; + let path_shape = request_path_shape("/upload"); + + detector.observe_with_path_shape(&ctx, "body_too_large", &path_shape); + assert_eq!(signature_observed_count(&detector, &signature), 1); + + let engine = FilterEngine::new( + vec![], + PathBuf::from("/tmp/altura-prot-nonexistent-body-too-large-hot-path-filters.json"), + Duration::from_secs(30), + ) + .await; + let logger = Arc::new( + EventLogger::new("/tmp/altura-prot-test-body-too-large-hot-path-events.jsonl").unwrap(), + ); + let detector = AdaptiveDetector::new( + test_detector_config(60, 8_192, 8_192), + Arc::clone(&engine), + logger, + ); + detector.observe_with_path_shape(&ctx, "observed", &path_shape); + detector.observe_with_path_shape(&ctx, "body_too_large", &path_shape); + assert_eq!(signature_observed_count(&detector, &signature), 2); + } + fn same_signature_shard_keys(prefix: &str, count: usize) -> Vec { let mut selected = Vec::new(); let mut target_shard = None; diff --git a/src/filter.rs b/src/filter.rs index 01abd2c..86bf416 100644 --- a/src/filter.rs +++ b/src/filter.rs @@ -405,8 +405,9 @@ impl FilterEvalScratch { .as_str() } - fn matches_path_shape(&mut self, path: &str, expected: &str) -> bool { - if self.path_shape(path) == expected { + fn matches_path_shape(&mut self, path: &str, expected: &str, hint: Option<&str>) -> bool { + let shape = hint.unwrap_or_else(|| self.path_shape(path)); + if shape == expected { return true; } let parent = self @@ -613,6 +614,14 @@ impl FilterEngine { } pub fn evaluate(&self, ctx: &RequestContext<'_>) -> Option { + self.evaluate_with_path_shape(ctx, None) + } + + pub fn evaluate_with_path_shape( + &self, + ctx: &RequestContext<'_>, + path_shape_hint: Option<&str>, + ) -> Option { let unix_ms = unix_time_ms(); let rules = self.rules_snapshot("evaluate"); let mut scratch = FilterEvalScratch::default(); @@ -624,7 +633,7 @@ impl FilterEngine { if rule.adaptive && runtime.active_until_ms.load(Ordering::Relaxed) <= unix_ms { continue; } - if rule_matches(runtime, ctx, &mut scratch) { + if rule_matches(runtime, ctx, &mut scratch, path_shape_hint) { return Some(FilterDecision { rule_id: rule.id.clone(), status: rule.action.status, @@ -722,11 +731,9 @@ pub fn request_signature( query: Option<&str>, headers: &HeaderMap, ) -> String { - let basis = signature_basis(method, path, query, headers); - hex_prefix( - blake3::hash(basis.as_bytes()).as_bytes(), - REQUEST_SIGNATURE_HASH_BYTES, - ) + let mut hasher = blake3::Hasher::new(); + append_signature_basis(&mut hasher, method, path, query, headers); + hex_prefix(hasher.finalize().as_bytes(), REQUEST_SIGNATURE_HASH_BYTES) } pub fn legacy_request_signature( @@ -735,7 +742,8 @@ pub fn legacy_request_signature( query: Option<&str>, headers: &HeaderMap, ) -> String { - let basis = signature_basis(method, path, query, headers); + let mut basis = String::with_capacity(128); + append_signature_basis_string(&mut basis, method, path, query, headers); format!("{:016x}", fnv1a64(basis.as_bytes())) } @@ -745,26 +753,72 @@ pub fn signature_basis( query: Option<&str>, headers: &HeaderMap, ) -> String { - let ua = combined_header_value(headers, &http::header::USER_AGENT) - .map(|value| user_agent_family(value.as_ref())) - .unwrap_or_else(|| "empty".to_string()); - let accept = combined_header_value(headers, &http::header::ACCEPT) - .map(|value| header_class(value.as_ref())) - .unwrap_or_else(|| "empty".to_string()); - format!( - "{}|{}|{}|{}|{}", - method.to_ascii_uppercase(), - normalize_path(path), - query_shape(query.unwrap_or("")), - ua, - accept - ) + let mut basis = String::with_capacity(128); + append_signature_basis_string(&mut basis, method, path, query, headers); + basis +} + +fn append_signature_basis_string( + out: &mut String, + method: &str, + path: &str, + query: Option<&str>, + headers: &HeaderMap, +) { + for byte in method.bytes() { + out.push(byte.to_ascii_uppercase() as char); + } + out.push('|'); + out.push_str(&normalize_path(path)); + out.push('|'); + out.push_str(&query_shape(query.unwrap_or(""))); + out.push('|'); + match combined_header_value(headers, &http::header::USER_AGENT) { + Some(value) => out.push_str(&user_agent_family(value.as_ref())), + None => out.push_str("empty"), + } + out.push('|'); + match combined_header_value(headers, &http::header::ACCEPT) { + Some(value) => out.push_str(&header_class(value.as_ref())), + None => out.push_str("empty"), + } +} + +fn append_signature_basis( + hasher: &mut blake3::Hasher, + method: &str, + path: &str, + query: Option<&str>, + headers: &HeaderMap, +) { + for byte in method.bytes() { + hasher.update(&[byte.to_ascii_uppercase()]); + } + hasher.update(b"|"); + let normalized_path = normalize_path(path); + hasher.update(normalized_path.as_bytes()); + hasher.update(b"|"); + let query = query_shape(query.unwrap_or("")); + hasher.update(query.as_bytes()); + hasher.update(b"|"); + if let Some(value) = combined_header_value(headers, &http::header::USER_AGENT) { + hasher.update(user_agent_family(value.as_ref()).as_bytes()); + } else { + hasher.update(b"empty"); + } + hasher.update(b"|"); + if let Some(value) = combined_header_value(headers, &http::header::ACCEPT) { + hasher.update(header_class(value.as_ref()).as_bytes()); + } else { + hasher.update(b"empty"); + } } fn rule_matches( runtime: &RuntimeRule, ctx: &RequestContext<'_>, scratch: &mut FilterEvalScratch, + path_shape_hint: Option<&str>, ) -> bool { let rule = &runtime.rule; let compiled = &runtime.condition; @@ -801,7 +855,7 @@ fn rule_matches( } } if let Some(path_shape) = &rule.condition.path_shape { - if !scratch.matches_path_shape(ctx.path, path_shape) { + if !scratch.matches_path_shape(ctx.path, path_shape, path_shape_hint) { return false; } } @@ -2218,6 +2272,13 @@ mod tests { }; assert!(engine.evaluate(&matching).is_some()); assert!(engine.evaluate(&benign).is_none()); + let hinted_shape = request_path_shape(matching.path); + assert_eq!( + engine + .evaluate_with_path_shape(&matching, Some(&hinted_shape)) + .map(|decision| decision.rule_id), + engine.evaluate(&matching).map(|decision| decision.rule_id) + ); } #[tokio::test] diff --git a/src/http_proxy.rs b/src/http_proxy.rs index ce84b66..137b29f 100644 --- a/src/http_proxy.rs +++ b/src/http_proxy.rs @@ -162,22 +162,23 @@ impl UpstreamFailureCircuitBreaker { Ok(shard) => shard, Err(poisoned) => poisoned.into_inner(), }; - if !shard.entries.contains_key(path_shape) { + let entry = if let Some(entry) = shard.entries.get_mut(path_shape) { + entry + } else { if !ensure_upstream_circuit_shard_capacity(&mut shard, self.shard_capacity, now_ms) { return; } - shard.order.push_back(path_shape.to_string()); - shard.entries.insert( - path_shape.to_string(), - UpstreamFailureCircuitEntry::default(), - ); - } - if let Some(entry) = shard.entries.get_mut(path_shape) { - entry.consecutive_failures = entry.consecutive_failures.saturating_add(1); - if entry.consecutive_failures >= self.threshold { - entry.open_until_ms = now_ms.saturating_add(self.open_ms); - entry.consecutive_failures = 0; - } + let path_shape_owned = path_shape.to_string(); + shard.order.push_back(path_shape_owned.clone()); + shard + .entries + .entry(path_shape_owned) + .or_insert_with(UpstreamFailureCircuitEntry::default) + }; + entry.consecutive_failures = entry.consecutive_failures.saturating_add(1); + if entry.consecutive_failures >= self.threshold { + entry.open_until_ms = now_ms.saturating_add(self.open_ms); + entry.consecutive_failures = 0; } } @@ -1095,7 +1096,6 @@ async fn handle_http( return Ok(early_rejection_response(400, format!("{reason}\n"), None)); } }; - let method = req.method().as_str().to_string(); if !method_allowed(req.method(), &state.cfg) { Stats::inc(&state.stats.http_method_rejected); return Ok(method_not_allowed_response(&state.cfg)); @@ -1137,24 +1137,28 @@ async fn handle_http( Stats::inc(&state.stats.http_uri_rejected); return Ok(request_target_rejected_response(format!("{reason}\n"))); } - let path = req.uri().path().to_string(); - let query = req.uri().query().map(ToString::to_string); - let signature = request_signature(&method, &path, query.as_deref(), req.headers()); - let path_shape = request_path_shape(&path); - let short_token_parent = short_token_parent_shape(&path); - let ctx = RequestContext { - client_ip, - method: &method, - path: &path, - query: query.as_deref(), - headers: req.headers(), - signature, - }; + let method = req.method().as_str(); + let path = req.uri().path(); + let query = req.uri().query(); - if admin_endpoint(&state.cfg.admin_path_prefix, &method, &path).is_some() { + if admin_endpoint(&state.cfg.admin_path_prefix, method, path).is_some() { + let method = method.to_string(); + let path = path.to_string(); + let query = query.map(ToString::to_string); + let signature = request_signature(&method, &path, query.as_deref(), req.headers()); + let path_shape = request_path_shape(&path); + let short_token_parent = short_token_parent_shape(&path); + let ctx = RequestContext { + client_ip, + method: &method, + path: &path, + query: query.as_deref(), + headers: req.headers(), + signature, + }; if let Some(response) = - maybe_rate_limit_response(&state.limiter, &state.stats, client_ip, None, &ctx) + maybe_rate_limit_response(&state.limiter, &state.stats, client_ip, None, &ctx, None) { return Ok(response); } @@ -1165,12 +1169,17 @@ async fn handle_http( peer_addr.ip(), None, &ctx, + None, ) { return Ok(response); } - if let Some(response) = - maybe_signature_rate_limit_response(&state.signature_limiter, &state.stats, None, &ctx) - { + if let Some(response) = maybe_signature_rate_limit_response( + &state.signature_limiter, + &state.stats, + None, + &ctx, + None, + ) { return Ok(response); } if let Some(response) = maybe_path_shape_rate_limit_response( @@ -1196,21 +1205,61 @@ async fn handle_http( } } - state.detector.observe(&ctx, "observed"); - if let Some(length) = content_length(req.headers()) { if state.cfg.max_body_bytes > 0 && length > state.cfg.max_body_bytes { Stats::inc(&state.stats.http_body_rejected); - state.detector.observe(&ctx, "body_too_large"); + let path_owned = path.to_string(); + let path_shape = request_path_shape(&path_owned); + let method_owned = method.to_string(); + let query_owned = query.map(ToString::to_string); + let ctx = RequestContext { + client_ip, + method: &method_owned, + path: &path_owned, + query: query_owned.as_deref(), + headers: req.headers(), + signature: request_signature( + &method_owned, + &path_owned, + query_owned.as_deref(), + req.headers(), + ), + }; + state + .detector + .observe_with_path_shape(&ctx, "observed", &path_shape); + state + .detector + .observe_with_path_shape(&ctx, "body_too_large", &path_shape); return Ok(content_too_large_response("request body too large\n")); } } + let method = method.to_string(); + let path = path.to_string(); + let query = query.map(ToString::to_string); + let signature = request_signature(&method, &path, query.as_deref(), req.headers()); + let path_shape = request_path_shape(&path); + let short_token_parent = short_token_parent_shape(&path); + let ctx = RequestContext { + client_ip, + method: &method, + path: &path, + query: query.as_deref(), + headers: req.headers(), + signature, + }; + + state + .detector + .observe_with_path_shape(&ctx, "observed", &path_shape); + if let Some(response) = maybe_signature_rate_limit_response( &state.signature_limiter, &state.stats, Some(&state.detector), &ctx, + Some(&path_shape), ) { return Ok(response); } @@ -1240,6 +1289,7 @@ async fn handle_http( client_ip, Some(&state.detector), &ctx, + Some(&path_shape), ) { return Ok(response); } @@ -1250,13 +1300,19 @@ async fn handle_http( peer_addr.ip(), Some(&state.detector), &ctx, + Some(&path_shape), ) { return Ok(response); } - if let Some(decision) = state.engine.evaluate(&ctx) { + if let Some(decision) = state + .engine + .evaluate_with_path_shape(&ctx, Some(&path_shape)) + { Stats::inc(&state.stats.http_blocked); - state.detector.observe(&ctx, "filter_block"); + state + .detector + .observe_with_path_shape(&ctx, "filter_block", &path_shape); return Ok(early_rejection_response( decision.status, decision.body, @@ -1958,6 +2014,7 @@ fn maybe_rate_limit_response( client_ip: IpAddr, detector: Option<&AdaptiveDetector>, ctx: &RequestContext<'_>, + path_shape: Option<&str>, ) -> Option> { let limit = limiter.check(client_ip); if limit.allowed { @@ -1965,14 +2022,16 @@ fn maybe_rate_limit_response( } Stats::inc(&stats.http_rate_limited); if let Some(detector) = detector { - detector.observe( - ctx, - match limit.reason { - Some(LimitReason::GlobalRate) => "global_rate_limited", - Some(LimitReason::PerIpRate) => "per_ip_rate_limited", - _ => "rate_limited", - }, - ); + let reason = match limit.reason { + Some(LimitReason::GlobalRate) => "global_rate_limited", + Some(LimitReason::PerIpRate) => "per_ip_rate_limited", + _ => "rate_limited", + }; + if let Some(path_shape) = path_shape { + detector.observe_with_path_shape(ctx, reason, path_shape); + } else { + detector.observe(ctx, reason); + } } Some(rate_limit_response("rate limited\n")) } @@ -1982,6 +2041,7 @@ fn maybe_signature_rate_limit_response( stats: &Stats, detector: Option<&AdaptiveDetector>, ctx: &RequestContext<'_>, + path_shape: Option<&str>, ) -> Option> { if limiter.check(&ctx.signature) { return None; @@ -1989,7 +2049,11 @@ fn maybe_signature_rate_limit_response( Stats::inc(&stats.http_rate_limited); Stats::inc(&stats.http_signature_rate_limited); if let Some(detector) = detector { - detector.observe(ctx, "signature_rate_limited"); + if let Some(path_shape) = path_shape { + detector.observe_with_path_shape(ctx, "signature_rate_limited", path_shape); + } else { + detector.observe(ctx, "signature_rate_limited"); + } } Some(rate_limit_response("signature rate limited\n")) } @@ -2038,6 +2102,7 @@ fn maybe_trusted_proxy_rate_limit_response( peer_ip: IpAddr, detector: Option<&AdaptiveDetector>, ctx: &RequestContext<'_>, + path_shape: Option<&str>, ) -> Option> { if client_ip == peer_ip { return None; @@ -2049,7 +2114,11 @@ fn maybe_trusted_proxy_rate_limit_response( Stats::inc(&stats.http_rate_limited); Stats::inc(&stats.http_trusted_proxy_rate_limited); if let Some(detector) = detector { - detector.observe(ctx, "trusted_proxy_rate_limited"); + if let Some(path_shape) = path_shape { + detector.observe_with_path_shape(ctx, "trusted_proxy_rate_limited", path_shape); + } else { + detector.observe(ctx, "trusted_proxy_rate_limited"); + } } Some(rate_limit_response("trusted proxy rate limited\n")) } @@ -5470,8 +5539,8 @@ mod tests { signature, }; - assert!(maybe_rate_limit_response(&limiter, &stats, client_ip, None, &ctx).is_none()); - let response = maybe_rate_limit_response(&limiter, &stats, client_ip, None, &ctx) + assert!(maybe_rate_limit_response(&limiter, &stats, client_ip, None, &ctx, None).is_none()); + let response = maybe_rate_limit_response(&limiter, &stats, client_ip, None, &ctx, None) .expect("second request should be rate limited"); assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); assert_retry_after_header(&response); @@ -5529,14 +5598,16 @@ mod tests { signature: "cold-signature".to_string(), }; - assert!(maybe_signature_rate_limit_response(&limiter, &stats, None, &ctx).is_none()); - let response = maybe_signature_rate_limit_response(&limiter, &stats, None, &ctx) + assert!(maybe_signature_rate_limit_response(&limiter, &stats, None, &ctx, None).is_none()); + let response = maybe_signature_rate_limit_response(&limiter, &stats, None, &ctx, None) .expect("second hot signature request should be limited"); assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); assert_retry_after_header(&response); assert_generated_no_store_header(&response); assert_connection_close_header(&response); - assert!(maybe_signature_rate_limit_response(&limiter, &stats, None, &other_ctx).is_none()); + assert!( + maybe_signature_rate_limit_response(&limiter, &stats, None, &other_ctx, None).is_none() + ); assert_eq!( stats .http_signature_rate_limited @@ -5714,11 +5785,11 @@ mod tests { }; assert!(maybe_trusted_proxy_rate_limit_response( - &limiter, &stats, client_ip, peer_ip, None, &ctx + &limiter, &stats, client_ip, peer_ip, None, &ctx, None ) .is_none()); let response = maybe_trusted_proxy_rate_limit_response( - &limiter, &stats, client_ip, peer_ip, None, &ctx, + &limiter, &stats, client_ip, peer_ip, None, &ctx, None, ) .expect("second trusted-proxy request should be rate limited"); assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); @@ -5732,7 +5803,7 @@ mod tests { 1 ); assert!(maybe_trusted_proxy_rate_limit_response( - &limiter, &stats, peer_ip, peer_ip, None, &ctx + &limiter, &stats, peer_ip, peer_ip, None, &ctx, None ) .is_none()); } diff --git a/src/limiter.rs b/src/limiter.rs index b74408a..da80831 100644 --- a/src/limiter.rs +++ b/src/limiter.rs @@ -165,21 +165,19 @@ impl KeyRateLimiter { let now = Instant::now(); let shard_idx = shard_for_key(key); let mut shard = lock_or_recover(&self.shards[shard_idx], "key rate limiter"); - if !shard.entries.contains_key(key) - && !ensure_key_shard_capacity(&mut shard, self.max_tracked_keys, now) - { - return false; + if let Some(entry) = shard.entries.get_mut(key) { + entry.last_seen = now; + return entry.bucket.allow(now, 1.0); } - if !shard.entries.contains_key(key) { - shard.order.push_back(key.to_string()); + if !ensure_key_shard_capacity(&mut shard, self.max_tracked_keys, now) { + return false; } - let entry = shard - .entries - .entry(key.to_string()) - .or_insert_with(|| KeyBucket { - bucket: TokenBucket::new(self.rps, self.burst, now), - last_seen: now, - }); + let key_owned = key.to_string(); + shard.order.push_back(key_owned.clone()); + let entry = shard.entries.entry(key_owned).or_insert_with(|| KeyBucket { + bucket: TokenBucket::new(self.rps, self.burst, now), + last_seen: now, + }); entry.last_seen = now; entry.bucket.allow(now, 1.0) } @@ -263,26 +261,30 @@ impl ShortTokenSiblingRateLimiter { let now = Instant::now(); let shard_idx = shard_for_key(parent_shape); let mut shard = lock_or_recover(&self.shards[shard_idx], "short-token sibling limiter"); - if !shard.entries.contains_key(parent_shape) - && !ensure_short_token_parent_shard_capacity(&mut shard, self.max_tracked_parents, now) - { - return false; - } - if !shard.entries.contains_key(parent_shape) { - shard.order.push_back(parent_shape.to_string()); - } - let entry = shard - .entries - .entry(parent_shape.to_string()) - .or_insert_with(|| ShortTokenSiblingEntry { - bucket: TokenBucket::new(self.rps, self.burst, now), - distinct_tokens: HashSet::new(), - token_order: VecDeque::new(), - last_seen: now, - }); + let entry = if let Some(entry) = shard.entries.get_mut(parent_shape) { + entry + } else { + if !ensure_short_token_parent_shard_capacity(&mut shard, self.max_tracked_parents, now) + { + return false; + } + let parent_owned = parent_shape.to_string(); + shard.order.push_back(parent_owned.clone()); + shard + .entries + .entry(parent_owned) + .or_insert_with(|| ShortTokenSiblingEntry { + bucket: TokenBucket::new(self.rps, self.burst, now), + distinct_tokens: HashSet::new(), + token_order: VecDeque::new(), + last_seen: now, + }) + }; entry.last_seen = now; - if entry.distinct_tokens.insert(token.to_string()) { - entry.token_order.push_back(token.to_string()); + if !entry.distinct_tokens.contains(token) { + let token_owned = token.to_string(); + entry.distinct_tokens.insert(token_owned.clone()); + entry.token_order.push_back(token_owned); while entry.distinct_tokens.len() > SHORT_TOKEN_SIBLING_DISTINCT_CAP { let Some(oldest) = entry.token_order.pop_front() else { break;