iron_control_api/routes/budget/
state.rs

1//! Budget protocol shared state
2//!
3//! Provides `BudgetState` which holds all managers and dependencies needed
4//! for budget protocol endpoints.
5
6use crate::{ ic_token::IcTokenManager, ip_token::IpTokenCrypto, jwt_auth::JwtSecret, routes::auth::AuthState };
7use axum::extract::FromRef;
8use iron_secrets::crypto::CryptoService;
9use iron_token_manager::
10{
11  agent_budget::AgentBudgetManager,
12  lease_manager::LeaseManager,
13  provider_key_storage::ProviderKeyStorage,
14};
15use sqlx::SqlitePool;
16use std::sync::Arc;
17
18/// Budget protocol shared state
19#[ derive( Clone ) ]
20pub struct BudgetState
21{
22  pub ic_token_manager: Arc< IcTokenManager >,
23  pub ip_token_crypto: Arc< IpTokenCrypto >,
24  pub lease_manager: Arc< LeaseManager >,
25  pub agent_budget_manager: Arc< AgentBudgetManager >,
26  pub provider_key_storage: Arc< ProviderKeyStorage >,
27  pub provider_key_crypto: Arc< CryptoService >,
28  pub db_pool: SqlitePool,
29  pub jwt_secret: Arc< JwtSecret >,
30  /// Crypto service for decrypting provider keys (Feature 014)
31  pub crypto_service: Option< Arc< CryptoService > >,
32}
33
34/// Enable AuthState extraction from BudgetState
35impl FromRef< BudgetState > for AuthState
36{
37  fn from_ref( state: &BudgetState ) -> Self
38  {
39    AuthState
40    {
41      jwt_secret: state.jwt_secret.clone(),
42      db_pool: state.db_pool.clone(),
43      rate_limiter: crate::rate_limiter::LoginRateLimiter::new(),
44      // Budget endpoints use IC Token auth, not login rate limiting
45      rate_limiting_enabled: false,
46    }
47  }
48}
49
50impl BudgetState
51{
52  /// Create new budget state
53  ///
54  /// # Arguments
55  ///
56  /// * `ic_token_secret` - Secret key for IC Token JWT signing
57  /// * `ip_token_key` - 32-byte encryption key for IP Token AES-256-GCM
58  /// * `provider_key_master` - 32-byte master key for provider key encryption/decryption
59  /// * `jwt_secret` - Secret key for JWT access token signing/verification
60  /// * `database_url` - Database connection string
61  /// * `crypto_service` - Optional crypto service for provider key decryption
62  ///
63  /// # Errors
64  ///
65  /// Returns error if database connection or crypto initialization fails
66  pub async fn new(
67    ic_token_secret: String,
68    ip_token_key: &[ u8 ],
69    provider_key_master: &[ u8 ],
70    jwt_secret: Arc< JwtSecret >,
71    database_url: &str,
72    crypto_service: Option< Arc< CryptoService > >,
73  ) -> Result< Self, Box< dyn std::error::Error > >
74  {
75    let db_pool = SqlitePool::connect( database_url ).await?;
76    let ic_token_manager = Arc::new( IcTokenManager::new( ic_token_secret ) );
77    let ip_token_crypto = Arc::new( IpTokenCrypto::new( ip_token_key )? );
78    let provider_key_crypto = Arc::new( CryptoService::new( provider_key_master )? );
79    let lease_manager = Arc::new( LeaseManager::from_pool( db_pool.clone() ) );
80    let agent_budget_manager = Arc::new( AgentBudgetManager::from_pool( db_pool.clone() ) );
81    let provider_key_storage = Arc::new( ProviderKeyStorage::new( db_pool.clone() ) );
82
83    Ok( Self
84    {
85      ic_token_manager,
86      ip_token_crypto,
87      lease_manager,
88      agent_budget_manager,
89      provider_key_storage,
90      provider_key_crypto,
91      db_pool,
92      jwt_secret,
93      crypto_service,
94    } )
95  }
96}