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

featherbit/plugins/native/
ua_restriction.rs
