Skip to main content

rama_proxy/
username.rs

1use super::ProxyFilter;
2use rama_core::error::BoxErrorExt as _;
3use rama_core::{
4    error::{BoxError, ErrorExt},
5    extensions::Extensions,
6    telemetry::tracing,
7    username::{UsernameLabelParser, UsernameLabelState, UsernameLabelWriter},
8};
9use rama_utils::macros::match_ignore_ascii_case_str;
10
11#[derive(Debug, Clone, Default)]
12#[non_exhaustive]
13/// A parser which parses [`ProxyFilter`]s from username labels
14/// and adds it to the input [`Extensions`].
15///
16/// [`Extensions`]: rama_core::extensions::Extensions
17pub struct ProxyFilterUsernameParser {
18    key: Option<ProxyFilterKey>,
19    proxy_filter: ProxyFilter,
20}
21
22#[derive(Debug, Clone)]
23enum ProxyFilterKey {
24    Id,
25    Pool,
26    Continent,
27    Country,
28    State,
29    City,
30    Carrier,
31    Asn,
32}
33
34impl ProxyFilterUsernameParser {
35    /// Create a new [`ProxyFilterUsernameParser`].
36    #[must_use]
37    pub fn new() -> Self {
38        Self::default()
39    }
40}
41
42impl UsernameLabelParser for ProxyFilterUsernameParser {
43    type Error = BoxError;
44
45    fn parse_label(&mut self, label: &str) -> UsernameLabelState {
46        if let Some(key) = self.key.take() {
47            match key {
48                ProxyFilterKey::Id => {
49                    self.proxy_filter.id = Some(match label.try_into() {
50                        Ok(id) => id,
51                        Err(err) => {
52                            tracing::trace!(
53                                "abort username label parsing: invalid parse label: {err:?}"
54                            );
55                            return UsernameLabelState::Abort;
56                        }
57                    })
58                }
59                ProxyFilterKey::Pool => {
60                    self.proxy_filter.pool_id = match self.proxy_filter.pool_id.take() {
61                        Some(mut pool_ids) => {
62                            pool_ids.push(label.into());
63                            Some(pool_ids)
64                        }
65                        None => Some(vec![label.into()]),
66                    }
67                }
68                ProxyFilterKey::Continent => {
69                    self.proxy_filter.continent = match self.proxy_filter.continent.take() {
70                        Some(mut continents) => {
71                            continents.push(label.into());
72                            Some(continents)
73                        }
74                        None => Some(vec![label.into()]),
75                    }
76                }
77                ProxyFilterKey::Country => {
78                    self.proxy_filter.country = match self.proxy_filter.country.take() {
79                        Some(mut countries) => {
80                            countries.push(label.into());
81                            Some(countries)
82                        }
83                        None => Some(vec![label.into()]),
84                    }
85                }
86                ProxyFilterKey::State => {
87                    self.proxy_filter.state = match self.proxy_filter.state.take() {
88                        Some(mut states) => {
89                            states.push(label.into());
90                            Some(states)
91                        }
92                        None => Some(vec![label.into()]),
93                    }
94                }
95                ProxyFilterKey::City => {
96                    self.proxy_filter.city = match self.proxy_filter.city.take() {
97                        Some(mut cities) => {
98                            cities.push(label.into());
99                            Some(cities)
100                        }
101                        None => Some(vec![label.into()]),
102                    }
103                }
104                ProxyFilterKey::Carrier => {
105                    self.proxy_filter.carrier = match self.proxy_filter.carrier.take() {
106                        Some(mut carriers) => {
107                            carriers.push(label.into());
108                            Some(carriers)
109                        }
110                        None => Some(vec![label.into()]),
111                    }
112                }
113                ProxyFilterKey::Asn => {
114                    let asn = match label.try_into() {
115                        Ok(asn) => asn,
116                        Err(err) => {
117                            tracing::trace!(
118                                "failed to parse asn username label; abort username parsing: {err:?}"
119                            );
120                            return UsernameLabelState::Abort;
121                        }
122                    };
123                    self.proxy_filter.asn = match self.proxy_filter.asn.take() {
124                        Some(mut asns) => {
125                            asns.push(asn);
126                            Some(asns)
127                        }
128                        None => Some(vec![asn]),
129                    }
130                }
131            }
132        } else {
133            // allow bool-keys to be negated
134            let (key, bval) = if let Some(key) = label.strip_prefix('!') {
135                (key, false)
136            } else {
137                (label, true)
138            };
139
140            match_ignore_ascii_case_str! {
141                match(key) {
142                    "datacenter" => self.proxy_filter.datacenter = Some(bval),
143                    "residential" => self.proxy_filter.residential = Some(bval),
144                    "mobile" => self.proxy_filter.mobile = Some(bval),
145                    "id" => self.key = Some(ProxyFilterKey::Id),
146                    "pool" => self.key = Some(ProxyFilterKey::Pool),
147                    "continent" => self.key = Some(ProxyFilterKey::Continent),
148                    "country" => self.key = Some(ProxyFilterKey::Country),
149                    "state" => self.key = Some(ProxyFilterKey::State),
150                    "city" => self.key = Some(ProxyFilterKey::City),
151                    "carrier" => self.key = Some(ProxyFilterKey::Carrier),
152                    "asn" => self.key = Some(ProxyFilterKey::Asn),
153                    _ => return UsernameLabelState::Ignored,
154                }
155            }
156
157            if !bval && self.key.take().is_some() {
158                // negation only possible for standalone labels
159                return UsernameLabelState::Ignored;
160            }
161        }
162
163        UsernameLabelState::Used
164    }
165
166    fn build(self, ext: &Extensions) -> Result<(), Self::Error> {
167        if let Some(key) = self.key {
168            return Err(
169                BoxError::from_static_str("unused proxy filter username key")
170                    .context_debug_field("key", key),
171            );
172        }
173        if self.proxy_filter != ProxyFilter::default() {
174            ext.insert(self.proxy_filter);
175        }
176        Ok(())
177    }
178}
179
180impl<const SEPARATOR: char> UsernameLabelWriter<SEPARATOR> for ProxyFilter {
181    fn write_labels(
182        &self,
183        composer: &mut rama_core::username::Composer<SEPARATOR>,
184    ) -> Result<(), rama_core::username::ComposeError> {
185        if let Some(id) = &self.id {
186            composer.write_label("id")?;
187            composer.write_label(id)?;
188        }
189
190        if let Some(pool_id_vec) = &self.pool_id {
191            for pool_id in pool_id_vec {
192                composer.write_label("pool")?;
193                composer.write_label(pool_id.as_ref())?;
194            }
195        }
196
197        if let Some(continent_vec) = &self.continent {
198            for continent in continent_vec {
199                composer.write_label("continent")?;
200                composer.write_label(continent.as_ref())?;
201            }
202        }
203
204        if let Some(country_vec) = &self.country {
205            for country in country_vec {
206                composer.write_label("country")?;
207                composer.write_label(country.as_ref())?;
208            }
209        }
210
211        if let Some(state_vec) = &self.state {
212            for state in state_vec {
213                composer.write_label("state")?;
214                composer.write_label(state.as_ref())?;
215            }
216        }
217
218        if let Some(city_vec) = &self.city {
219            for city in city_vec {
220                composer.write_label("city")?;
221                composer.write_label(city.as_ref())?;
222            }
223        }
224
225        if let Some(datacenter) = &self.datacenter {
226            if *datacenter {
227                composer.write_label("datacenter")?;
228            } else {
229                composer.write_label("!datacenter")?;
230            }
231        }
232
233        if let Some(residential) = &self.residential {
234            if *residential {
235                composer.write_label("residential")?;
236            } else {
237                composer.write_label("!residential")?;
238            }
239        }
240
241        if let Some(mobile) = &self.mobile {
242            if *mobile {
243                composer.write_label("mobile")?;
244            } else {
245                composer.write_label("!mobile")?;
246            }
247        }
248
249        if let Some(carrier_vec) = &self.carrier {
250            for carrier in carrier_vec {
251                composer.write_label("carrier")?;
252                composer.write_label(carrier.as_ref())?;
253            }
254        }
255
256        if let Some(asn_vec) = &self.asn {
257            for asn in asn_vec {
258                composer.write_label("asn")?;
259                composer.write_label(asn.as_u32().to_string())?;
260            }
261        }
262
263        Ok(())
264    }
265}
266
267#[cfg(test)]
268mod tests {
269    use super::*;
270    use crate::StringFilter;
271    use rama_core::username::{compose_username, parse_username};
272    use rama_net::asn::Asn;
273    use rama_utils::str::non_empty_str;
274
275    #[test]
276    fn test_username_config() {
277        let test_cases = [
278            ("john", String::from("john"), None),
279            (
280                "john-datacenter",
281                String::from("john"),
282                Some(ProxyFilter {
283                    datacenter: Some(true),
284                    ..Default::default()
285                }),
286            ),
287            (
288                "john-!datacenter",
289                String::from("john"),
290                Some(ProxyFilter {
291                    datacenter: Some(false),
292                    ..Default::default()
293                }),
294            ),
295            (
296                "john-country-us-datacenter",
297                String::from("john"),
298                Some(ProxyFilter {
299                    country: Some(vec!["us".into()]),
300                    datacenter: Some(true),
301                    ..Default::default()
302                }),
303            ),
304            (
305                "john-city-tokyo-residential",
306                String::from("john"),
307                Some(ProxyFilter {
308                    city: Some(vec!["tokyo".into()]),
309                    residential: Some(true),
310                    ..Default::default()
311                }),
312            ),
313            (
314                "john-country-us-datacenter-pool-1",
315                String::from("john"),
316                Some(ProxyFilter {
317                    pool_id: Some(vec![StringFilter::from("1")]),
318                    country: Some(vec![StringFilter::from("us")]),
319                    datacenter: Some(true),
320                    ..Default::default()
321                }),
322            ),
323            (
324                "john-country-us-datacenter-pool-1-residential",
325                String::from("john"),
326                Some(ProxyFilter {
327                    pool_id: Some(vec![StringFilter::from("1")]),
328                    country: Some(vec![StringFilter::from("us")]),
329                    datacenter: Some(true),
330                    residential: Some(true),
331                    ..Default::default()
332                }),
333            ),
334            (
335                "john-country-us-datacenter-pool-1-residential-mobile",
336                String::from("john"),
337                Some(ProxyFilter {
338                    pool_id: Some(vec![StringFilter::from("1")]),
339                    country: Some(vec![StringFilter::from("us")]),
340                    datacenter: Some(true),
341                    residential: Some(true),
342                    mobile: Some(true),
343                    ..Default::default()
344                }),
345            ),
346            (
347                "john-country-us-datacenter-pool-1-residential-!mobile",
348                String::from("john"),
349                Some(ProxyFilter {
350                    pool_id: Some(vec![StringFilter::from("1")]),
351                    country: Some(vec![StringFilter::from("us")]),
352                    datacenter: Some(true),
353                    residential: Some(true),
354                    mobile: Some(false),
355                    ..Default::default()
356                }),
357            ),
358            (
359                "john-country-us-city-california-datacenter-pool-1-!residential-mobile",
360                String::from("john"),
361                Some(ProxyFilter {
362                    pool_id: Some(vec![StringFilter::from("1")]),
363                    country: Some(vec![StringFilter::from("us")]),
364                    city: Some(vec![StringFilter::from("california")]),
365                    datacenter: Some(true),
366                    residential: Some(false),
367                    mobile: Some(true),
368                    ..Default::default()
369                }),
370            ),
371            (
372                "john-country-us-datacenter-pool-1-residential-mobile-id-1",
373                String::from("john"),
374                Some(ProxyFilter {
375                    id: Some(non_empty_str!("1")),
376                    pool_id: Some(vec![StringFilter::from("1")]),
377                    country: Some(vec![StringFilter::from("us")]),
378                    datacenter: Some(true),
379                    residential: Some(true),
380                    mobile: Some(true),
381                    ..Default::default()
382                }),
383            ),
384            (
385                "john-country-us-datacenter-pool-1-residential-mobile-carrier-bar-id-1",
386                String::from("john"),
387                Some(ProxyFilter {
388                    id: Some(non_empty_str!("1")),
389                    pool_id: Some(vec![StringFilter::from("1")]),
390                    country: Some(vec![StringFilter::from("us")]),
391                    datacenter: Some(true),
392                    residential: Some(true),
393                    mobile: Some(true),
394                    carrier: Some(vec![StringFilter::from("bar")]),
395                    ..Default::default()
396                }),
397            ),
398            (
399                "john-country-us-datacenter-pool-1-residential-mobile-id-1-country-uk",
400                String::from("john"),
401                Some(ProxyFilter {
402                    id: Some(non_empty_str!("1")),
403                    pool_id: Some(vec![StringFilter::from("1")]),
404                    country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
405                    datacenter: Some(true),
406                    residential: Some(true),
407                    mobile: Some(true),
408                    ..Default::default()
409                }),
410            ),
411            (
412                "john-country-us-!datacenter-pool-1-residential-mobile-id-1-country-uk",
413                String::from("john"),
414                Some(ProxyFilter {
415                    id: Some(non_empty_str!("1")),
416                    pool_id: Some(vec![StringFilter::from("1")]),
417                    country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
418                    datacenter: Some(false),
419                    residential: Some(true),
420                    mobile: Some(true),
421                    ..Default::default()
422                }),
423            ),
424            (
425                "john-country-us-datacenter-pool-1-residential-mobile-id-1-country-uk-pool-2",
426                String::from("john"),
427                Some(ProxyFilter {
428                    id: Some(non_empty_str!("1")),
429                    pool_id: Some(vec![StringFilter::from("1"), StringFilter::from("2")]),
430                    country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
431                    datacenter: Some(true),
432                    residential: Some(true),
433                    mobile: Some(true),
434                    ..Default::default()
435                }),
436            ),
437            (
438                "john-country-us-datacenter-pool-1-!residential-mobile-id-1-country-uk-pool-2",
439                String::from("john"),
440                Some(ProxyFilter {
441                    id: Some(non_empty_str!("1")),
442                    pool_id: Some(vec![StringFilter::from("1"), StringFilter::from("2")]),
443                    country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
444                    datacenter: Some(true),
445                    residential: Some(false),
446                    mobile: Some(true),
447                    ..Default::default()
448                }),
449            ),
450            (
451                "john-country-us-datacenter-pool-1-residential-mobile-id-1-country-uk-pool-2-datacenter",
452                String::from("john"),
453                Some(ProxyFilter {
454                    id: Some(non_empty_str!("1")),
455                    pool_id: Some(vec![StringFilter::from("1"), StringFilter::from("2")]),
456                    country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
457                    datacenter: Some(true),
458                    residential: Some(true),
459                    mobile: Some(true),
460                    ..Default::default()
461                }),
462            ),
463            (
464                "john-country-us-datacenter-pool-1-residential-mobile-id-1-country-uk-pool-2-datacenter-residential",
465                String::from("john"),
466                Some(ProxyFilter {
467                    id: Some(non_empty_str!("1")),
468                    pool_id: Some(vec![StringFilter::from("1"), StringFilter::from("2")]),
469                    country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
470                    datacenter: Some(true),
471                    residential: Some(true),
472                    mobile: Some(true),
473                    ..Default::default()
474                }),
475            ),
476            (
477                "john-country-us-datacenter-pool-1-residential-mobile-id-1-country-uk-pool-2-datacenter-residential-mobile",
478                String::from("john"),
479                Some(ProxyFilter {
480                    id: Some(non_empty_str!("1")),
481                    pool_id: Some(vec![StringFilter::from("1"), StringFilter::from("2")]),
482                    country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
483                    datacenter: Some(true),
484                    residential: Some(true),
485                    mobile: Some(true),
486                    ..Default::default()
487                }),
488            ),
489            (
490                "john-continent-americas-country-us-state-NY-city-ny-asn-7018",
491                String::from("john"),
492                Some(ProxyFilter {
493                    continent: Some(vec![StringFilter::from("americas")]),
494                    country: Some(vec![StringFilter::from("us")]),
495                    state: Some(vec![StringFilter::from("ny")]),
496                    city: Some(vec![StringFilter::from("ny")]),
497                    asn: Some(vec![Asn::from_static(7018)]),
498                    ..Default::default()
499                }),
500            ),
501            (
502                "john-continent-europe-continent-asia",
503                String::from("john"),
504                Some(ProxyFilter {
505                    continent: Some(vec![
506                        StringFilter::from("europe"),
507                        StringFilter::from("asia"),
508                    ]),
509                    ..Default::default()
510                }),
511            ),
512            (
513                "john-country-us-datacenter-pool-1-residential-mobile-id-1-country-uk-pool-2-!datacenter-!residential-!mobile",
514                String::from("john"),
515                Some(ProxyFilter {
516                    id: Some(non_empty_str!("1")),
517                    pool_id: Some(vec![StringFilter::from("1"), StringFilter::from("2")]),
518                    country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
519                    datacenter: Some(false),
520                    residential: Some(false),
521                    mobile: Some(false),
522                    ..Default::default()
523                }),
524            ),
525        ];
526
527        for (username, expected_username, expected_filter) in test_cases.into_iter() {
528            let ext = Extensions::default();
529
530            let parser = ProxyFilterUsernameParser::default();
531
532            let username = parse_username(&ext, parser, username).unwrap();
533            let filter = ext.get_ref::<ProxyFilter>().cloned();
534            assert_eq!(
535                username, expected_username,
536                "username = '{username}' ; expected_username = '{expected_username}'",
537            );
538            assert_eq!(
539                filter, expected_filter,
540                "username = '{username}' ; expected_username = '{expected_username}'",
541            );
542        }
543    }
544
545    #[test]
546    fn test_username_config_error() {
547        for username in [
548            "john-country-us-datacenter-",
549            "",
550            "-",
551            "john-country-us-datacenter-pool",
552            "john-foo",
553            "john-foo-country",
554            "john-country",
555            "john-id-", // empty id is invalid
556        ] {
557            let ext = Extensions::default();
558
559            let parser = ProxyFilterUsernameParser::default();
560
561            assert!(
562                parse_username(&ext, parser, username).is_err(),
563                "username = {username}",
564            );
565        }
566    }
567
568    #[test]
569    fn test_username_negation_key_failures() {
570        for username in [
571            "john-!id-a",
572            "john-!pool-b",
573            "john-!country-us",
574            "john-!city-ny",
575            "john-!carrier-c",
576        ] {
577            let ext = Extensions::default();
578
579            let parser = ProxyFilterUsernameParser::default();
580
581            assert!(
582                parse_username(&ext, parser, username).is_err(),
583                "username = {username}",
584            );
585        }
586    }
587
588    #[test]
589    fn test_username_compose_parser_proxy_filter() {
590        let test_cases = [
591            ProxyFilter::default(),
592            ProxyFilter {
593                id: Some(non_empty_str!("p42")),
594                ..Default::default()
595            },
596            ProxyFilter {
597                id: Some(non_empty_str!("1")),
598                pool_id: Some(vec![StringFilter::from("1")]),
599                country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
600                datacenter: Some(false),
601                residential: Some(true),
602                mobile: Some(true),
603                ..Default::default()
604            },
605            ProxyFilter {
606                id: Some(non_empty_str!("1")),
607                pool_id: Some(vec![StringFilter::from("1"), StringFilter::from("2")]),
608                country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
609                datacenter: Some(false),
610                residential: Some(false),
611                mobile: Some(false),
612                ..Default::default()
613            },
614            ProxyFilter {
615                id: Some(non_empty_str!("a")),
616                pool_id: Some(vec![StringFilter::from("1"), StringFilter::from("2")]),
617                continent: Some(vec![StringFilter::from("na"), StringFilter::from("eu")]),
618                country: Some(vec![StringFilter::from("us"), StringFilter::from("be")]),
619                state: Some(vec![
620                    StringFilter::from("ca"),
621                    StringFilter::from("ny"),
622                    StringFilter::from("ovl"),
623                ]),
624                city: Some(vec![
625                    StringFilter::from("berkeley"),
626                    StringFilter::from("bruxelles"),
627                    StringFilter::from("gent"),
628                ]),
629                datacenter: Some(false),
630                residential: Some(true),
631                mobile: Some(true),
632                carrier: Some(vec![
633                    StringFilter::from("at&t"),
634                    StringFilter::from("orange"),
635                ]),
636                asn: Some(vec![Asn::from_static(7018), Asn::from_static(1)]),
637            },
638        ];
639
640        for test_case in test_cases {
641            let fmt_username = compose_username("john".to_owned(), &test_case).unwrap();
642            let ext = Extensions::new();
643            let username =
644                parse_username(&ext, ProxyFilterUsernameParser::default(), &fmt_username)
645                    .unwrap_or_else(|_| panic!("to be ok: {fmt_username}"));
646            assert_eq!("john", username);
647            if test_case == Default::default() {
648                assert!(!ext.contains::<ProxyFilter>());
649            } else {
650                let result = ext.get_ref::<ProxyFilter>().unwrap();
651                assert_eq!(test_case, *result);
652            }
653        }
654    }
655}