rama_core/layer/limit/
layer.rs

1use std::fmt;
2
3use super::{Limit, into_response::ErrorIntoResponseFn, policy::UnlimitedPolicy};
4use crate::Layer;
5
6/// Limit requests based on a [`Policy`].
7///
8/// [`Policy`]: crate::layer::limit::Policy
9pub struct LimitLayer<P, F = ()> {
10    policy: P,
11    error_into_response: F,
12}
13
14impl<P> LimitLayer<P> {
15    /// Creates a new [`LimitLayer`] from a [`crate::layer::limit::Policy`].
16    pub const fn new(policy: P) -> Self {
17        LimitLayer {
18            policy,
19            error_into_response: (),
20        }
21    }
22
23    /// Attach a function to this [`LimitLayer`] to allow you to turn the Policy error
24    /// into a Result fully compatible with the inner `Service` Result.
25    pub fn with_error_into_response_fn<F>(self, f: F) -> LimitLayer<P, ErrorIntoResponseFn<F>> {
26        LimitLayer {
27            policy: self.policy,
28            error_into_response: ErrorIntoResponseFn(f),
29        }
30    }
31}
32
33impl LimitLayer<UnlimitedPolicy> {
34    /// Creates a new [`LimitLayer`] with an unlimited policy.
35    ///
36    /// Meaning that all requests are allowed to proceed.
37    pub fn unlimited() -> Self {
38        Self::new(UnlimitedPolicy::default())
39    }
40}
41
42impl<P: fmt::Debug, F: fmt::Debug> std::fmt::Debug for LimitLayer<P, F> {
43    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
44        f.debug_struct("LimitLayer")
45            .field("policy", &self.policy)
46            .field("self.error_into_response", &self.error_into_response)
47            .finish()
48    }
49}
50
51impl<P, F> Clone for LimitLayer<P, F>
52where
53    P: Clone,
54    F: Clone,
55{
56    fn clone(&self) -> Self {
57        Self {
58            policy: self.policy.clone(),
59            error_into_response: self.error_into_response.clone(),
60        }
61    }
62}
63
64impl<T, P, F> Layer<T> for LimitLayer<P, F>
65where
66    P: Clone,
67    F: Clone,
68{
69    type Service = Limit<T, P, F>;
70
71    fn layer(&self, service: T) -> Self::Service {
72        Limit {
73            inner: service,
74            policy: self.policy.clone(),
75            error_into_response: self.error_into_response.clone(),
76        }
77    }
78
79    fn into_layer(self, service: T) -> Self::Service {
80        Limit {
81            inner: service,
82            policy: self.policy,
83            error_into_response: self.error_into_response,
84        }
85    }
86}