apisix_admin_client/models/plugins/
limit_count.rs

1use serde::{Deserialize, Serialize};
2use strum_macros::{Display, EnumString};
3use validator::Validate;
4use crate::error::ApisixClientError;
5use crate::{Result};
6use crate::models::Plugin;
7
8/// Builder to create a LimitCount
9#[serde_with::skip_serializing_none]
10#[derive(Validate, Default, Debug, Clone, PartialEq, Serialize, Deserialize)]
11pub struct LimitCountBuilder {
12    #[validate(range(min = 1))]
13    pub count: Option<i64>,
14    #[validate(range(min = 1))]
15    pub time_window: Option<i64>,
16    pub key_type: Option<LimitCountKeyType>,
17    pub key: Option<String>,
18    #[validate(range(min = 200, max = 599))]
19    pub rejected_code: Option<i64>,
20    #[validate(length(min = 1))]
21    pub rejected_msg: Option<String>,
22    pub policy: Option<LimitCountPolicy>,
23    pub allow_degradation: Option<bool>,
24    pub show_limit_quota_headers: Option<bool>,
25    #[validate(length(min = 1))]
26    pub group: Option<String>,
27    pub redis_host: Option<String>,
28    pub redis_port: Option<i64>,
29    pub redis_username: Option<String>,
30    pub redis_password: Option<String>,
31    pub redis_ssl: Option<bool>,
32    pub redis_ssl_verify: Option<bool>,
33    pub redis_database: Option<i64>,
34    #[validate(range(min = 1))]
35    pub redis_timeout: Option<i64>,
36    pub redis_cluster_nodes: Option<Vec<String>>,
37    pub redis_cluster_name: Option<String>,
38    pub redis_cluster_ssl: Option<bool>,
39    pub redis_cluster_ssl_verify: Option<bool>,
40}
41
42impl LimitCountBuilder {
43    pub fn new() -> Self {
44        LimitCount::default().into()
45    }
46
47    /// Maximum number of requests to allow.
48    pub fn with_count(mut self, count: i64) -> Self {
49        self.count = Some(count);
50        self
51    }
52
53    pub fn with_time_window(mut self, time_window: i64) -> Self {
54        self.time_window = Some(time_window);
55        self
56    }
57
58    pub fn with_key_type(mut self, key_type: LimitCountKeyType) -> Self {
59        self.key_type = Some(key_type);
60        self
61    }
62
63    pub fn with_key(mut self, key: impl Into<String>) -> Self {
64        self.key = Some(key.into());
65        self
66    }
67
68    pub fn with_rejected_code(mut self, rejected_code: i64) -> Self {
69        self.rejected_code = Some(rejected_code);
70        self
71    }
72
73    pub fn with_rejected_msg(mut self, rejected_msg: impl Into<String>) -> Self {
74        self.rejected_msg = Some(rejected_msg.into());
75        self
76    }
77
78    pub fn with_policy(mut self, policy: LimitCountPolicy) -> Self {
79        self.policy = Some(policy);
80        self
81    }
82
83    pub fn with_allow_degradation(mut self, allow_degradation: bool) -> Self {
84        self.allow_degradation = Some(allow_degradation);
85        self
86    }
87
88    pub fn with_show_limit_quota_headers(mut self, show_limit_quota_headers: bool) -> Self {
89        self.show_limit_quota_headers = Some(show_limit_quota_headers);
90        self
91    }
92
93    pub fn with_group(mut self, group: impl Into<String>) -> Self {
94        self.group = Some(group.into());
95        self
96    }
97
98    pub fn with_redis_host(mut self, redis_host: impl Into<String>) -> Self {
99        self.redis_host = Some(redis_host.into());
100        self
101    }
102
103    pub fn with_redis_port(mut self, redis_port: i64) -> Self {
104        self.redis_port = Some(redis_port);
105        self
106    }
107
108    pub fn with_redis_username(mut self, redis_username: impl Into<String>) -> Self {
109        self.redis_username = Some(redis_username.into());
110        self
111    }
112
113    pub fn with_redis_password(mut self, redis_password: impl Into<String>) -> Self {
114        self.redis_password = Some(redis_password.into());
115        self
116    }
117
118    pub fn with_redis_ssl(mut self, redis_ssl: bool) -> Self {
119        self.redis_ssl = Some(redis_ssl);
120        self
121    }
122
123    pub fn with_redis_ssl_verify(mut self, redis_ssl_verify: bool) -> Self {
124        self.redis_ssl_verify = Some(redis_ssl_verify);
125        self
126    }
127
128    pub fn with_redis_database(mut self, redis_database: i64) -> Self {
129        self.redis_database = Some(redis_database);
130        self
131    }
132
133    pub fn with_redis_timeout(mut self, redis_timeout: i64) -> Self {
134        self.redis_timeout = Some(redis_timeout);
135        self
136    }
137
138    pub fn with_redis_cluster_nodes(mut self, redis_cluster_nodes: Vec<String>) -> Self {
139        self.redis_cluster_nodes = Some(redis_cluster_nodes);
140        self
141    }
142
143    pub fn with_redis_cluster_name(mut self, redis_cluster_name: impl Into<String>) -> Self {
144        self.redis_cluster_name = Some(redis_cluster_name.into());
145        self
146    }
147
148    pub fn with_redis_cluster_ssl(mut self, redis_cluster_ssl: bool) -> Self {
149        self.redis_cluster_ssl = Some(redis_cluster_ssl);
150        self
151    }
152
153    pub fn with_redis_cluster_ssl_verify(mut self, redis_cluster_ssl_verify: bool) -> Self {
154        self.redis_cluster_ssl_verify = Some(redis_cluster_ssl_verify);
155        self
156    }
157
158    pub fn build(self) -> Result<LimitCount> {
159        self.validate().map_err(|v| ApisixClientError::PluginConfigException(v.to_string()))?;
160        Ok(LimitCount {
161            count: self.count,
162            time_window: self.time_window,
163            key_type: self.key_type,
164            key: self.key,
165            rejected_code: self.rejected_code,
166            rejected_msg: self.rejected_msg,
167            policy: self.policy,
168            allow_degradation: self.allow_degradation,
169            show_limit_quota_headers: self.show_limit_quota_headers,
170            group: self.group,
171            redis_host: self.redis_host,
172            redis_port: self.redis_port,
173            redis_username: self.redis_username,
174            redis_password: self.redis_password,
175            redis_ssl: self.redis_ssl,
176            redis_ssl_verify: self.redis_ssl_verify,
177            redis_database: self.redis_database,
178            redis_timeout: self.redis_timeout,
179            redis_cluster_nodes: self.redis_cluster_nodes,
180            redis_cluster_name: self.redis_cluster_name,
181            redis_cluster_ssl: self.redis_cluster_ssl,
182            redis_cluster_ssl_verify: self.redis_cluster_ssl_verify,
183        })
184    }
185}
186
187/// The limit-count Plugin limits the number of requests to your service by a given count per time.
188/// The plugin is using Fixed Window algorithm.
189/// [Documentation](https://apisix.apache.org/docs/apisix/plugins/limit-count/)
190#[serde_with::skip_serializing_none]
191#[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)]
192pub struct LimitCount {
193    pub count: Option<i64>,
194    pub time_window: Option<i64>,
195    pub key_type: Option<LimitCountKeyType>,
196    pub key: Option<String>,
197    pub rejected_code: Option<i64>,
198    pub rejected_msg: Option<String>,
199    pub policy: Option<LimitCountPolicy>,
200    pub allow_degradation: Option<bool>,
201    pub show_limit_quota_headers: Option<bool>,
202    pub group: Option<String>,
203    pub redis_host: Option<String>,
204    pub redis_port: Option<i64>,
205    pub redis_username: Option<String>,
206    pub redis_password: Option<String>,
207    pub redis_ssl: Option<bool>,
208    pub redis_ssl_verify: Option<bool>,
209    pub redis_database: Option<i64>,
210    pub redis_timeout: Option<i64>,
211    pub redis_cluster_nodes: Option<Vec<String>>,
212    pub redis_cluster_name: Option<String>,
213    pub redis_cluster_ssl: Option<bool>,
214    pub redis_cluster_ssl_verify: Option<bool>,
215}
216
217impl From<LimitCount> for LimitCountBuilder {
218    fn from(item: LimitCount) -> Self {
219        LimitCountBuilder {
220            count: item.count,
221            time_window: item.time_window,
222            key_type: item.key_type,
223            key: item.key,
224            rejected_code: item.rejected_code,
225            rejected_msg: item.rejected_msg,
226            policy: item.policy,
227            allow_degradation: item.allow_degradation,
228            show_limit_quota_headers: item.show_limit_quota_headers,
229            group: item.group,
230            redis_host: item.redis_host,
231            redis_port: item.redis_port,
232            redis_username: item.redis_username,
233            redis_password: item.redis_password,
234            redis_ssl: item.redis_ssl,
235            redis_ssl_verify: item.redis_ssl_verify,
236            redis_database: item.redis_database,
237            redis_timeout: item.redis_timeout,
238            redis_cluster_nodes: item.redis_cluster_nodes,
239            redis_cluster_name: item.redis_cluster_name,
240            redis_cluster_ssl: item.redis_cluster_ssl,
241            redis_cluster_ssl_verify: item.redis_cluster_ssl_verify,
242        }
243    }
244}
245
246impl Plugin for LimitCount {}
247
248/// Type of user specified key to use
249#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize, Display, EnumString)]
250#[allow(non_camel_case_types)]
251#[strum(ascii_case_insensitive)]
252#[non_exhaustive]
253pub enum LimitCountKeyType {
254    var,
255    var_combination,
256    constant,
257}
258
259/// Rate-limiting policies to use for retrieving and increment the limit count
260#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize, Display)]
261#[allow(non_camel_case_types)]
262#[non_exhaustive]
263pub enum LimitCountPolicy {
264    local,
265    redis,
266    #[strum(serialize = "redis-cluster")]
267    redis_cluster,
268}
269
270// region: tests
271#[cfg(test)]
272mod tests {
273    use serde_json::{to_string, to_string_pretty};
274    use super::*;
275    use tracing::{error, info};
276    use tracing_test::traced_test;
277    use crate::models::admin_upstream_requests::UpstreamType;
278    use crate::models::common::TypedItem;
279
280    #[traced_test]
281    #[tokio::test]
282    async fn test_parse_limit_count_empty_response() {
283        let nodes = r#"{}"#;
284        let nodes: LimitCount = serde_json::from_str(nodes).unwrap();
285        assert_eq!(nodes.count, None);
286        assert_eq!(nodes.time_window, None);
287        assert_eq!(nodes.key_type, None);
288        assert_eq!(nodes.key, None);
289        assert_eq!(nodes.rejected_code, None);
290    }
291
292    #[traced_test]
293    #[tokio::test]
294    async fn test_parse_limit_count_response() {
295        let nodes = r#"
296        {
297            "count": 2,
298            "time_window": 60,
299            "rejected_code": 503,
300            "key_type": "var",
301            "key": "remote_addr",
302            "policy": "redis_cluster"
303        }"#;
304        let nodes: LimitCount = serde_json::from_str(nodes).unwrap();
305        assert_eq!(nodes.count.unwrap(), 2);
306        assert_eq!(nodes.time_window.unwrap(), 60);
307        assert_eq!(nodes.rejected_code.unwrap(), 503);
308        assert_eq!(nodes.key_type.unwrap(), LimitCountKeyType::var);
309        assert_eq!(nodes.key.unwrap(), "remote_addr");
310        assert_eq!(nodes.policy.unwrap(), LimitCountPolicy::redis_cluster);
311    }
312
313    #[traced_test]
314    #[tokio::test]
315    async fn test_validate() {
316        let mut test = LimitCountBuilder::new()
317            .with_count(0) // count min > 0
318            .build();
319        assert!(test.is_err());
320        test = LimitCountBuilder::new()
321            .with_rejected_code(199) // rejected_code min > 200
322            .build();
323        assert!(test.is_err());
324        test = LimitCountBuilder::new()
325            .with_rejected_code(600) // rejected_code min > 200
326            .build();
327        assert!(test.is_err());
328        test = LimitCountBuilder::new()
329            .with_rejected_code(200) // rejected_code min > 200
330            .build();
331        assert!(test.is_ok());
332    }
333
334}
335// endregion: tests