cratestack_axum/ratelimit/
subscription.rs

1//! Rate-limiting WebSocket / WebTransport subscription opens
2//! (`docs/design/subscriptions-websocket.md` §6.7). The bindings call the
3//! core trait [`SubscriptionRateLimit`]; this implements it over
4//! [`OpExecutor::admit_rate_limit`] with the HTTP layer's own pieces, so the
5//! store, token bucket, [`StoreErrorPolicy`] and lookup timeout are shared.
6
7use std::fmt;
8use std::net::IpAddr;
9use std::sync::Arc;
10use std::time::Duration;
11
12use async_trait::async_trait;
13use cratestack_core::subscription_wire::{
14    SubscriptionCharge, SubscriptionRateLimit, SubscriptionRefusal,
15};
16use cratestack_core::{BoundedOutcome, CratestackError, RateLimitDecision};
17use cratestack_exec::{OpAdmission, OpExecutor, OpInput, RateLimitAdmission, RateLimitBucket};
18
19use super::config::RateLimitConfig;
20use super::key_fn::sha256_hex;
21use super::layer::RateLimitLayer;
22use super::policy::{DEFAULT_STORE_TIMEOUT, StoreErrorPolicy, StoreErrorWarnings};
23use super::scope::bucket_address;
24use super::store::RateLimitStore;
25use super::store_error::{StoreFailure, classify_store_failure};
26
27type OpResolver = Arc<dyn Fn(&str) -> OpAdmission + Send + Sync>;
28
29/// Charges one token per subscription open. Build it with
30/// [`RateLimitLayer::for_subscriptions`] (sharing that layer's store and
31/// settings) or [`SubscriptionRateLimiter::new`], and hand it to
32/// `WebSocketConfig::rate_limit` / `WebTransportConfig::rate_limit`.
33///
34/// The bucket key follows the HTTP layer's precedence on an already
35/// verified session: `princ:<sha256 of the principal key>`, else
36/// `ip:<peer>` (IPv6 aggregated to its /64), else the open is refused
37/// `failed_precondition` (HTTP's 412 rule, cratestack#416). No `auth:` key
38/// is ever minted: identity is verified before admission.
39#[derive(Clone)]
40pub struct SubscriptionRateLimiter {
41    executor: OpExecutor,
42    op_resolver: OpResolver,
43    store_error_policy: StoreErrorPolicy,
44    store_timeout: Duration,
45    warnings: Arc<StoreErrorWarnings>,
46}
47
48impl SubscriptionRateLimiter {
49    pub fn new(store: Arc<dyn RateLimitStore>, config: RateLimitConfig) -> Self {
50        Self {
51            executor: OpExecutor::new(None, Duration::ZERO).with_rate_limit(store, config),
52            // Every open charged until a resolver says otherwise: fail closed.
53            op_resolver: Arc::new(|_| OpAdmission::unresolved()),
54            store_error_policy: StoreErrorPolicy::default(),
55            store_timeout: DEFAULT_STORE_TIMEOUT,
56            warnings: Arc::new(StoreErrorWarnings::default()),
57        }
58    }
59
60    pub fn with_store_error_policy(mut self, policy: StoreErrorPolicy) -> Self {
61        self.store_error_policy = policy;
62        self
63    }
64
65    pub fn with_store_timeout(mut self, timeout: Duration) -> Self {
66        self.store_timeout = timeout;
67        self
68    }
69
70    /// Which ops are charged, by op id. The default charges all of them,
71    /// because subscription ops cannot be `@no_rate_limit`; this hook exists
72    /// for parity with [`RateLimitLayer::with_op_resolver`].
73    pub fn with_op_resolver(
74        mut self,
75        resolver: impl Fn(&str) -> OpAdmission + Send + Sync + 'static,
76    ) -> Self {
77        self.op_resolver = Arc::new(resolver);
78        self
79    }
80}
81
82impl RateLimitLayer {
83    /// A limiter for WebSocket / WebTransport subscription opens with this
84    /// layer's store, token bucket, [`StoreErrorPolicy`], lookup timeout and
85    /// warning throttle. A principal's bucket is the one its HTTP requests
86    /// use when the application's `VerifiedPrincipal` string equals the
87    /// binding's principal key (not so for COSE's `cose:<thumbprint>`).
88    pub fn for_subscriptions(&self) -> SubscriptionRateLimiter {
89        SubscriptionRateLimiter {
90            executor: OpExecutor::new(None, Duration::ZERO)
91                .with_rate_limit(self.store.clone(), self.config),
92            op_resolver: Arc::new(|_| OpAdmission::unresolved()),
93            store_error_policy: self.store_error_policy,
94            store_timeout: self.store_timeout,
95            warnings: self.warnings.clone(),
96        }
97    }
98}
99
100impl fmt::Debug for SubscriptionRateLimiter {
101    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
102        f.debug_struct("SubscriptionRateLimiter")
103            .field("store_timeout", &self.store_timeout)
104            .field("store_error_policy", &self.store_error_policy)
105            .finish_non_exhaustive()
106    }
107}
108
109#[async_trait]
110impl SubscriptionRateLimit for SubscriptionRateLimiter {
111    async fn admit(&self, charge: SubscriptionCharge<'_>) -> Result<(), SubscriptionRefusal> {
112        let op = (self.op_resolver)(charge.op_id);
113        if !self.executor.rate_limit_applies(&op) {
114            return Ok(());
115        }
116        let key =
117            bucket_key(charge.principal, charge.remote).map_err(SubscriptionRefusal::Error)?;
118        let input = OpInput::for_rate_limit(op, RateLimitBucket::new(&key, None));
119        let admitted =
120            tokio::time::timeout(self.store_timeout, self.executor.admit_rate_limit(&input))
121                .await
122                .unwrap_or_else(|_| Err(super::policy::store_timeout_error()));
123        match admitted {
124            Ok(RateLimitAdmission::Bypass) => Ok(()),
125            Ok(RateLimitAdmission::Consumed(BoundedOutcome { decision, .. })) => match decision {
126                RateLimitDecision::Allowed { .. } => Ok(()),
127                RateLimitDecision::Throttled { retry_after_secs } => {
128                    Err(SubscriptionRefusal::Throttled { retry_after_secs })
129                }
130            },
131            // `#[non_exhaustive]`: an outcome this build does not know must
132            // refuse, never open the subscription.
133            Ok(_) => Err(SubscriptionRefusal::Error(CratestackError::Internal(
134                "rate limit: unhandled admission outcome; refusing rather than opening the \
135                 subscription"
136                    .to_owned(),
137            ))),
138            Err(error) => {
139                match classify_store_failure(error, self.store_error_policy, &self.warnings) {
140                    StoreFailure::Serve => Ok(()),
141                    StoreFailure::Refuse(error) => Err(SubscriptionRefusal::Error(error)),
142                }
143            }
144        }
145    }
146}
147
148/// HTTP's precedence (`key_fn.rs`) on a verified session: principal, then
149/// peer, else 412.
150fn bucket_key(principal: Option<&str>, remote: Option<IpAddr>) -> Result<String, CratestackError> {
151    if let Some(principal) = principal.filter(|p| !p.is_empty()) {
152        return Ok(format!("princ:{}", sha256_hex(principal.as_bytes())));
153    }
154    if let Some(ip) = remote {
155        return Ok(format!("ip:{}", bucket_address(ip)));
156    }
157    Err(CratestackError::PreconditionFailed(
158        "rate limit: an anonymous subscription session has no peer address to key its bucket \
159         on; serve the router through into_make_service_with_connect_info::<SocketAddr>()"
160            .to_owned(),
161    ))
162}
cratestack_axum/ratelimit/subscription.rs

cratestack_axum/ratelimit/
subscription.rs
