featherbit/plugins/native/
ua_restriction.rs

1//! User-Agent allow/deny list plugin (`ua-restriction`).
2//!
3//! Port of APISIX's `ua-restriction`: matches the request's `User-Agent`
4//! header against a list of regexes — either an allowlist (only matching
5//! agents pass) or a denylist (matching agents are rejected). Rejections are
6//! routed through the node's `denied` port.
7
8use async_trait::async_trait;
9use bytes::Bytes;
10use regex::Regex;
11use std::collections::HashMap;
12
13use crate::context::Context;
14use crate::plugins::{Plugin, PluginOutput, PluginResult};
15use crate::vars::template::Template;
16
17/// Restricts access based on the `User-Agent` request header.
18///
19/// Exactly one of `allowlist` / `denylist` is configured (APISIX `oneOf`
20/// parity). Each User-Agent value is trimmed and tested against the regexes;
21/// in allowlist mode a request passes when any value matches any rule, in
22/// denylist mode a request is rejected when any value matches any rule.
23/// A missing User-Agent is rejected unless `bypass_missing` is set.
24pub struct UaRestrictionPlugin {
25    /// Compiled allowlist regexes; non-empty means allowlist mode.
26    allowlist: Vec<Regex>,
27    /// Compiled denylist regexes; non-empty means denylist mode.
28    denylist: Vec<Regex>,
29    /// Pass requests that carry no User-Agent header (default `false`).
30    bypass_missing: bool,
31    /// HTTP status for rejections (default 403).
32    rejected_code: u16,
33    /// Body message for rejections. Supports `{{namespace.path}}` references
34    /// (no legacy `$var` interpolation — this field never supported it, so
35    /// this sweep must not start).
36    rejected_msg: Template,
37}
38
39/// Parses a config key as an array of non-empty regex strings, compiling
40/// each at config load. An absent key yields an empty list.
41fn parse_regex_list(
42    config: &HashMap<String, serde_json::Value>,
43    key: &str,
44) -> Result<Vec<Regex>, String> {
45    let Some(raw) = config.get(key) else {
46        return Ok(Vec::new());
47    };
48    let arr = raw
49        .as_array()
50        .ok_or_else(|| format!("{} must be an array of regex strings", key))?;
51    arr.iter()
52        .map(|item| {
53            let s = item
54                .as_str()
55                .ok_or_else(|| format!("{} entries must be strings", key))?;
56            if s.is_empty() {
57                return Err(format!("{} entries must be non-empty", key));
58            }
59            Regex::new(s).map_err(|e| format!("invalid regex '{}' in {}: {}", s, key, e))
60        })
61        .collect()
62}
63
64impl UaRestrictionPlugin {
65    /// Builds the plugin from node config.
66    ///
67    /// Accepted keys:
68    /// - `allowlist` (array of regex strings): only matching User-Agents pass.
69    /// - `denylist` (array of regex strings): matching User-Agents are rejected.
70    ///   Exactly **one** of `allowlist` / `denylist` must be non-empty
71    ///   (APISIX rejects both-set and neither-set configs); regexes are
72    ///   compiled here, so an invalid pattern is a config error.
73    /// - `bypass_missing` (bool, default `false`): pass requests without a
74    ///   User-Agent header instead of rejecting them.
75    /// - `rejected_code` (integer 200–599, default `403`): rejection status.
76    /// - `rejected_msg` (string, default `"Not allowed"`): rejection message,
77    ///   returned as `{"message": ...}`. Supports `{{namespace.path}}`
78    ///   references.
79    ///
80    /// ```yaml
81    /// type: ua-restriction
82    /// config:
83    ///   denylist: ["curl/.*", "(?i)spider"]
84    ///   bypass_missing: false
85    ///   rejected_msg: Not allowed
86    /// ```
87    pub fn from_config(config: &HashMap<String, serde_json::Value>) -> Result<Self, String> {
88        let allowlist = parse_regex_list(config, "allowlist")?;
89        let denylist = parse_regex_list(config, "denylist")?;
90
91        if !allowlist.is_empty() && !denylist.is_empty() {
92            return Err("ua-restriction: allowlist and denylist cannot both be set".to_string());
93        }
94        if allowlist.is_empty() && denylist.is_empty() {
95            return Err(
96                "ua-restriction: exactly one of allowlist or denylist must be non-empty"
97                    .to_string(),
98            );
99        }
100
101        let rejected_code = match config.get("rejected_code") {
102            None => 403,
103            Some(v) => {
104                let code = v
105                    .as_u64()
106                    .ok_or_else(|| "rejected_code must be an integer".to_string())?;
107                if !(200..=599).contains(&code) {
108                    return Err("rejected_code must be between 200 and 599".to_string());
109                }
110                code as u16
111            }
112        };
113
114        let rejected_msg = config
115            .get("rejected_msg")
116            .and_then(|v| v.as_str())
117            .unwrap_or("Not allowed");
118        // Discard warnings here — the compile-time walk (a later task)
119        // reports well-formed-but-unknown references; execution must not.
120        let rejected_msg = Template::parse(rejected_msg).0;
121
122        Ok(Self {
123            allowlist,
124            denylist,
125            bypass_missing: config
126                .get("bypass_missing")
127                .and_then(|v| v.as_bool())
128                .unwrap_or(false),
129            rejected_code,
130            rejected_msg,
131        })
132    }
133
134    /// Builds the 403-style rejection: JSON body on the response, routed
135    /// through the `denied` port.
136    fn reject(&self, mut ctx: Context) -> PluginResult {
137        let message = self.rejected_msg.render(&ctx).into_owned();
138        ctx.response.status_code = self.rejected_code;
139        ctx.response.body = Bytes::from(serde_json::json!({ "message": message }).to_string());
140        ctx.response.headers.insert(
141            "content-type".to_string(),
142            vec!["application/json".to_string()],
143        );
144        Ok(PluginOutput::on_port(ctx, "denied"))
145    }
146}
147
148#[async_trait]
149impl Plugin for UaRestrictionPlugin {
150    fn plugin_type(&self) -> &str {
151        "ua-restriction"
152    }
153
154    async fn execute(&self, ctx: Context) -> PluginResult {
155        let user_agents: Vec<&str> = ctx
156            .request
157            .headers
158            .get("user-agent")
159            .map(|values| values.iter().map(|v| v.trim()).collect())
160            .unwrap_or_default();
161
162        if user_agents.is_empty() {
163            if self.bypass_missing {
164                return Ok(PluginOutput::success(ctx));
165            }
166            return self.reject(ctx);
167        }
168
169        let passed = if !self.allowlist.is_empty() {
170            // Allowlist mode: any UA value matching any rule passes.
171            user_agents
172                .iter()
173                .any(|ua| self.allowlist.iter().any(|re| re.is_match(ua)))
174        } else {
175            // Denylist mode: any UA value matching any rule rejects.
176            !user_agents
177                .iter()
178                .any(|ua| self.denylist.iter().any(|re| re.is_match(ua)))
179        };
180
181        if !passed {
182            return self.reject(ctx);
183        }
184
185        Ok(PluginOutput::success(ctx))
186    }
187}
188
189#[cfg(test)]
190mod tests {
191    use super::*;
192    use crate::context::{GatewayRequest, GatewayResponse, Protocol};
193
194    fn test_context(user_agent: Option<&str>) -> Context {
195        let mut headers = HashMap::new();
196        if let Some(ua) = user_agent {
197            headers.insert("user-agent".to_string(), vec![ua.to_string()]);
198        }
199        Context {
200            request: GatewayRequest {
201                method: "GET".to_string(),
202                path: "/test".to_string(),
203                host: "localhost".to_string(),
204                scheme: "http".to_string(),
205                headers,
206                query_params: HashMap::new(),
207                body: Bytes::new(),
208                remote_addr: "127.0.0.1:12345".to_string(),
209                protocol: Protocol::Http1,
210            },
211            response: GatewayResponse {
212                status_code: 0,
213                headers: HashMap::new(),
214                body: Bytes::new(),
215                stream: None,
216            },
217            message: HashMap::new(),
218            errors: Vec::new(),
219        }
220    }
221
222    fn config(json: serde_json::Value) -> HashMap<String, serde_json::Value> {
223        serde_json::from_value(json).unwrap()
224    }
225
226    #[test]
227    fn test_config_requires_exactly_one_list() {
228        // neither list
229        assert!(UaRestrictionPlugin::from_config(&config(serde_json::json!({}))).is_err());
230        // both lists
231        assert!(UaRestrictionPlugin::from_config(&config(serde_json::json!({
232            "allowlist": ["a"], "denylist": ["b"]
233        })))
234        .is_err());
235        // one list is fine
236        assert!(UaRestrictionPlugin::from_config(&config(serde_json::json!({
237            "denylist": ["curl"]
238        })))
239        .is_ok());
240    }
241
242    #[test]
243    fn test_config_rejects_invalid_regex_and_bad_shapes() {
244        assert!(UaRestrictionPlugin::from_config(&config(serde_json::json!({
245            "denylist": ["("]
246        })))
247        .is_err());
248        assert!(UaRestrictionPlugin::from_config(&config(serde_json::json!({
249            "denylist": "curl"
250        })))
251        .is_err());
252        assert!(UaRestrictionPlugin::from_config(&config(serde_json::json!({
253            "denylist": [""]
254        })))
255        .is_err());
256        assert!(UaRestrictionPlugin::from_config(&config(serde_json::json!({
257            "denylist": ["curl"], "rejected_code": 100
258        })))
259        .is_err());
260    }
261
262    #[tokio::test]
263    async fn test_denylist_blocks_matching_ua() {
264        let plugin = UaRestrictionPlugin::from_config(&config(serde_json::json!({
265            "denylist": ["curl/.*", "(?i)spider"]
266        })))
267        .unwrap();
268
269        let out = plugin
270            .execute(test_context(Some("curl/8.1.2")))
271            .await
272            .unwrap();
273        assert_eq!(out.port, Some("denied"));
274        assert_eq!(out.context.response.status_code, 403);
275
276        // non-matching UA passes
277        assert!(plugin
278            .execute(test_context(Some("Mozilla/5.0")))
279            .await
280            .unwrap()
281            .port
282            .is_none());
283    }
284
285    #[tokio::test]
286    async fn test_allowlist_only_matching_ua_passes() {
287        let plugin = UaRestrictionPlugin::from_config(&config(serde_json::json!({
288            "allowlist": ["Mozilla.*"]
289        })))
290        .unwrap();
291
292        assert!(plugin
293            .execute(test_context(Some("Mozilla/5.0")))
294            .await
295            .unwrap()
296            .port
297            .is_none());
298        // trimmed before matching (APISIX str_strip parity)
299        assert!(plugin
300            .execute(test_context(Some("  Mozilla/5.0  ")))
301            .await
302            .unwrap()
303            .port
304            .is_none());
305        assert_eq!(
306            plugin
307                .execute(test_context(Some("curl/8.1.2")))
308                .await
309                .unwrap()
310                .port,
311            Some("denied")
312        );
313    }
314
315    #[tokio::test]
316    async fn test_missing_ua_bypass() {
317        let deny = UaRestrictionPlugin::from_config(&config(serde_json::json!({
318            "denylist": ["curl"]
319        })))
320        .unwrap();
321        // default: missing UA is rejected
322        assert_eq!(
323            deny.execute(test_context(None)).await.unwrap().port,
324            Some("denied")
325        );
326
327        let bypass = UaRestrictionPlugin::from_config(&config(serde_json::json!({
328            "denylist": ["curl"], "bypass_missing": true
329        })))
330        .unwrap();
331        assert!(bypass
332            .execute(test_context(None))
333            .await
334            .unwrap()
335            .port
336            .is_none());
337    }
338
339    #[tokio::test]
340    async fn test_custom_code_and_message() {
341        let plugin = UaRestrictionPlugin::from_config(&config(serde_json::json!({
342            "denylist": ["curl"], "rejected_code": 405, "rejected_msg": "go away"
343        })))
344        .unwrap();
345        let out = plugin
346            .execute(test_context(Some("curl/8.1.2")))
347            .await
348            .unwrap();
349        assert_eq!(out.port, Some("denied"));
350        assert_eq!(out.context.response.status_code, 405);
351        let body: serde_json::Value = serde_json::from_slice(&out.context.response.body).unwrap();
352        assert_eq!(body["message"], "go away");
353    }
354
355    #[tokio::test]
356    async fn test_rejected_msg_renders_template() {
357        let plugin = UaRestrictionPlugin::from_config(&config(serde_json::json!({
358            "denylist": ["curl"], "rejected_msg": "blocked {{request.method}}"
359        })))
360        .unwrap();
361        let out = plugin
362            .execute(test_context(Some("curl/8.1.2")))
363            .await
364            .unwrap();
365        assert_eq!(out.port, Some("denied"));
366        let body: serde_json::Value = serde_json::from_slice(&out.context.response.body).unwrap();
367        assert_eq!(body["message"], "blocked GET");
368    }
369}
featherbit/plugins/native/ua_restriction.rs

featherbit/plugins/native/
ua_restriction.rs
