hive_router/
shared_state.rs

1use graphql_tools::validation::validate::ValidationPlan;
2use hive_console_sdk::agent::usage_agent::{AgentError, UsageAgent};
3use hive_router_config::HiveRouterConfig;
4use hive_router_internal::expressions::values::boolean::BooleanOrProgram;
5use hive_router_internal::expressions::ExpressionCompileError;
6use hive_router_internal::telemetry::TelemetryContext;
7use hive_router_plan_executor::headers::{
8    compile::compile_headers_plan, errors::HeaderRuleCompileError, plan::HeaderRulesPlan,
9};
10use moka::future::Cache;
11use moka::Expiry;
12use std::sync::Arc;
13use std::time::{Duration, SystemTime, UNIX_EPOCH};
14
15use crate::jwt::context::JwtTokenPayload;
16use crate::jwt::JwtAuthRuntime;
17use crate::pipeline::cors::{CORSConfigError, Cors};
18use crate::pipeline::introspection_policy::compile_introspection_policy;
19use crate::pipeline::parser::ParseCacheEntry;
20use crate::pipeline::progressive_override::{OverrideLabelsCompileError, OverrideLabelsEvaluator};
21
22pub type JwtClaimsCache = Cache<String, Arc<JwtTokenPayload>>;
23
24/// Default TTL for JWT claims cache entries (5 seconds)
25const DEFAULT_JWT_CACHE_TTL_SECS: u64 = 5;
26
27struct JwtClaimsExpiry;
28
29impl Expiry<String, Arc<JwtTokenPayload>> for JwtClaimsExpiry {
30    fn expire_after_create(
31        &self,
32        _key: &String,
33        value: &Arc<JwtTokenPayload>,
34        _created_at: std::time::Instant,
35    ) -> Option<Duration> {
36        const DEFAULT_TTL: Duration = Duration::from_secs(DEFAULT_JWT_CACHE_TTL_SECS);
37
38        // if token has no exp claim, use default TTL (avoids syscall)
39        let exp = match value.claims.exp {
40            Some(e) => e,
41            None => return Some(DEFAULT_TTL),
42        };
43
44        let now = match SystemTime::now().duration_since(UNIX_EPOCH) {
45            Ok(duration) => duration.as_secs(),
46            Err(_) => return Some(DEFAULT_TTL), // Clock error: fall back to default
47        };
48
49        // If token is already expired, return zero TTL to remove it immediately
50        if exp <= now {
51            return Some(Duration::ZERO);
52        }
53
54        // Calculate time until token expiration
55        let time_until_exp = Duration::from_secs(exp - now);
56
57        // Return the minimum of default TTL and time until expiration.
58        // Short-lived tokens (exp < 5s) are evicted when they expire
59        // Long-lived tokens still respect the 5s cache limit.
60        Some(DEFAULT_TTL.min(time_until_exp))
61    }
62}
63pub struct RouterSharedState {
64    pub validation_plan: ValidationPlan,
65    pub parse_cache: Cache<u64, ParseCacheEntry>,
66    pub router_config: Arc<HiveRouterConfig>,
67    pub headers_plan: HeaderRulesPlan,
68    pub override_labels_evaluator: OverrideLabelsEvaluator,
69    pub cors_runtime: Option<Cors>,
70    /// Cache for validated JWT claims to avoid re-parsing on every request.
71    /// The cache key is the raw JWT token string.
72    /// Stores the parsed claims payload for 5s,
73    /// but no longer than `exp` date.
74    pub jwt_claims_cache: JwtClaimsCache,
75    pub jwt_auth_runtime: Option<JwtAuthRuntime>,
76    pub hive_usage_agent: Option<UsageAgent>,
77    pub introspection_policy: BooleanOrProgram,
78    pub telemetry_context: Arc<TelemetryContext>,
79}
80
81impl RouterSharedState {
82    pub fn new(
83        router_config: Arc<HiveRouterConfig>,
84        jwt_auth_runtime: Option<JwtAuthRuntime>,
85        hive_usage_agent: Option<UsageAgent>,
86        validation_plan: ValidationPlan,
87        telemetry_context: Arc<TelemetryContext>,
88    ) -> Result<Self, SharedStateError> {
89        Ok(Self {
90            validation_plan,
91            headers_plan: compile_headers_plan(&router_config.headers).map_err(Box::new)?,
92            parse_cache: moka::future::Cache::new(1000),
93            cors_runtime: Cors::from_config(&router_config.cors).map_err(Box::new)?,
94            jwt_claims_cache: Cache::builder()
95                // High capacity due to potentially high token diversity.
96                // Capping prevents unbounded memory usage.
97                .max_capacity(10_000)
98                .expire_after(JwtClaimsExpiry)
99                .build(),
100            router_config: router_config.clone(),
101            override_labels_evaluator: OverrideLabelsEvaluator::from_config(
102                &router_config.override_labels,
103            )
104            .map_err(Box::new)?,
105            jwt_auth_runtime,
106            hive_usage_agent,
107            introspection_policy: compile_introspection_policy(&router_config.introspection)
108                .map_err(Box::new)?,
109            telemetry_context,
110        })
111    }
112}
113
114#[derive(thiserror::Error, Debug)]
115pub enum SharedStateError {
116    #[error("invalid headers config: {0}")]
117    HeaderRuleCompile(#[from] Box<HeaderRuleCompileError>),
118    #[error("invalid regex in CORS config: {0}")]
119    CORSConfig(#[from] Box<CORSConfigError>),
120    #[error("invalid override labels config: {0}")]
121    OverrideLabelsCompile(#[from] Box<OverrideLabelsCompileError>),
122    #[error("error creating hive usage agent: {0}")]
123    UsageAgent(#[from] Box<AgentError>),
124    #[error("invalid introspection config: {0}")]
125    IntrospectionPolicyCompile(#[from] Box<ExpressionCompileError>),
126}
hive_router/shared_state.rs

hive_router/
shared_state.rs
