rama_proxy/
username.rs

1use super::ProxyFilter;
2use rama_core::{
3    context::Extensions,
4    error::{OpaqueError, error},
5    telemetry::tracing,
6    username::{UsernameLabelParser, UsernameLabelState, UsernameLabelWriter},
7};
8use rama_utils::macros::match_ignore_ascii_case_str;
9
10#[derive(Debug, Clone, Default)]
11#[non_exhaustive]
12/// A parser which parses [`ProxyFilter`]s from username labels
13/// and adds it to the [`Context`]'s [`Extensions`].
14///
15/// [`Context`]: rama_core::Context
16/// [`Extensions`]: rama_core::context::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 = OpaqueError;
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: &mut Extensions) -> Result<(), Self::Error> {
167        if let Some(key) = self.key {
168            return Err(error!("unused proxy filter username key: {:?}", key));
169        }
170        if self.proxy_filter != ProxyFilter::default() {
171            ext.insert(self.proxy_filter);
172        }
173        Ok(())
174    }
175}
176
177impl<const SEPARATOR: char> UsernameLabelWriter<SEPARATOR> for ProxyFilter {
178    fn write_labels(
179        &self,
180        composer: &mut rama_core::username::Composer<SEPARATOR>,
181    ) -> Result<(), rama_core::username::ComposeError> {
182        if let Some(id) = &self.id {
183            composer.write_label("id")?;
184            composer.write_label(id.as_str())?;
185        }
186
187        if let Some(pool_id_vec) = &self.pool_id {
188            for pool_id in pool_id_vec {
189                composer.write_label("pool")?;
190                composer.write_label(pool_id.as_ref())?;
191            }
192        }
193
194        if let Some(continent_vec) = &self.continent {
195            for continent in continent_vec {
196                composer.write_label("continent")?;
197                composer.write_label(continent.as_ref())?;
198            }
199        }
200
201        if let Some(country_vec) = &self.country {
202            for country in country_vec {
203                composer.write_label("country")?;
204                composer.write_label(country.as_ref())?;
205            }
206        }
207
208        if let Some(state_vec) = &self.state {
209            for state in state_vec {
210                composer.write_label("state")?;
211                composer.write_label(state.as_ref())?;
212            }
213        }
214
215        if let Some(city_vec) = &self.city {
216            for city in city_vec {
217                composer.write_label("city")?;
218                composer.write_label(city.as_ref())?;
219            }
220        }
221
222        if let Some(datacenter) = &self.datacenter {
223            if *datacenter {
224                composer.write_label("datacenter")?;
225            } else {
226                composer.write_label("!datacenter")?;
227            }
228        }
229
230        if let Some(residential) = &self.residential {
231            if *residential {
232                composer.write_label("residential")?;
233            } else {
234                composer.write_label("!residential")?;
235            }
236        }
237
238        if let Some(mobile) = &self.mobile {
239            if *mobile {
240                composer.write_label("mobile")?;
241            } else {
242                composer.write_label("!mobile")?;
243            }
244        }
245
246        if let Some(carrier_vec) = &self.carrier {
247            for carrier in carrier_vec {
248                composer.write_label("carrier")?;
249                composer.write_label(carrier.as_ref())?;
250            }
251        }
252
253        if let Some(asn_vec) = &self.asn {
254            for asn in asn_vec {
255                composer.write_label("asn")?;
256                composer.write_label(asn.as_u32().to_string())?;
257            }
258        }
259
260        Ok(())
261    }
262}
263
264#[cfg(test)]
265mod tests {
266    use super::*;
267    use crate::StringFilter;
268    use rama_core::username::{compose_username, parse_username};
269    use rama_net::asn::Asn;
270    use rama_utils::str::NonEmptyString;
271
272    #[test]
273    fn test_username_config() {
274        let test_cases = [
275            ("john", String::from("john"), None),
276            (
277                "john-datacenter",
278                String::from("john"),
279                Some(ProxyFilter {
280                    datacenter: Some(true),
281                    ..Default::default()
282                }),
283            ),
284            (
285                "john-!datacenter",
286                String::from("john"),
287                Some(ProxyFilter {
288                    datacenter: Some(false),
289                    ..Default::default()
290                }),
291            ),
292            (
293                "john-country-us-datacenter",
294                String::from("john"),
295                Some(ProxyFilter {
296                    country: Some(vec!["us".into()]),
297                    datacenter: Some(true),
298                    ..Default::default()
299                }),
300            ),
301            (
302                "john-city-tokyo-residential",
303                String::from("john"),
304                Some(ProxyFilter {
305                    city: Some(vec!["tokyo".into()]),
306                    residential: Some(true),
307                    ..Default::default()
308                }),
309            ),
310            (
311                "john-country-us-datacenter-pool-1",
312                String::from("john"),
313                Some(ProxyFilter {
314                    pool_id: Some(vec![StringFilter::from("1")]),
315                    country: Some(vec![StringFilter::from("us")]),
316                    datacenter: Some(true),
317                    ..Default::default()
318                }),
319            ),
320            (
321                "john-country-us-datacenter-pool-1-residential",
322                String::from("john"),
323                Some(ProxyFilter {
324                    pool_id: Some(vec![StringFilter::from("1")]),
325                    country: Some(vec![StringFilter::from("us")]),
326                    datacenter: Some(true),
327                    residential: Some(true),
328                    ..Default::default()
329                }),
330            ),
331            (
332                "john-country-us-datacenter-pool-1-residential-mobile",
333                String::from("john"),
334                Some(ProxyFilter {
335                    pool_id: Some(vec![StringFilter::from("1")]),
336                    country: Some(vec![StringFilter::from("us")]),
337                    datacenter: Some(true),
338                    residential: Some(true),
339                    mobile: Some(true),
340                    ..Default::default()
341                }),
342            ),
343            (
344                "john-country-us-datacenter-pool-1-residential-!mobile",
345                String::from("john"),
346                Some(ProxyFilter {
347                    pool_id: Some(vec![StringFilter::from("1")]),
348                    country: Some(vec![StringFilter::from("us")]),
349                    datacenter: Some(true),
350                    residential: Some(true),
351                    mobile: Some(false),
352                    ..Default::default()
353                }),
354            ),
355            (
356                "john-country-us-city-california-datacenter-pool-1-!residential-mobile",
357                String::from("john"),
358                Some(ProxyFilter {
359                    pool_id: Some(vec![StringFilter::from("1")]),
360                    country: Some(vec![StringFilter::from("us")]),
361                    city: Some(vec![StringFilter::from("california")]),
362                    datacenter: Some(true),
363                    residential: Some(false),
364                    mobile: Some(true),
365                    ..Default::default()
366                }),
367            ),
368            (
369                "john-country-us-datacenter-pool-1-residential-mobile-id-1",
370                String::from("john"),
371                Some(ProxyFilter {
372                    id: Some(NonEmptyString::from_static("1")),
373                    pool_id: Some(vec![StringFilter::from("1")]),
374                    country: Some(vec![StringFilter::from("us")]),
375                    datacenter: Some(true),
376                    residential: Some(true),
377                    mobile: Some(true),
378                    ..Default::default()
379                }),
380            ),
381            (
382                "john-country-us-datacenter-pool-1-residential-mobile-carrier-bar-id-1",
383                String::from("john"),
384                Some(ProxyFilter {
385                    id: Some(NonEmptyString::from_static("1")),
386                    pool_id: Some(vec![StringFilter::from("1")]),
387                    country: Some(vec![StringFilter::from("us")]),
388                    datacenter: Some(true),
389                    residential: Some(true),
390                    mobile: Some(true),
391                    carrier: Some(vec![StringFilter::from("bar")]),
392                    ..Default::default()
393                }),
394            ),
395            (
396                "john-country-us-datacenter-pool-1-residential-mobile-id-1-country-uk",
397                String::from("john"),
398                Some(ProxyFilter {
399                    id: Some(NonEmptyString::from_static("1")),
400                    pool_id: Some(vec![StringFilter::from("1")]),
401                    country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
402                    datacenter: Some(true),
403                    residential: Some(true),
404                    mobile: Some(true),
405                    ..Default::default()
406                }),
407            ),
408            (
409                "john-country-us-!datacenter-pool-1-residential-mobile-id-1-country-uk",
410                String::from("john"),
411                Some(ProxyFilter {
412                    id: Some(NonEmptyString::from_static("1")),
413                    pool_id: Some(vec![StringFilter::from("1")]),
414                    country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
415                    datacenter: Some(false),
416                    residential: Some(true),
417                    mobile: Some(true),
418                    ..Default::default()
419                }),
420            ),
421            (
422                "john-country-us-datacenter-pool-1-residential-mobile-id-1-country-uk-pool-2",
423                String::from("john"),
424                Some(ProxyFilter {
425                    id: Some(NonEmptyString::from_static("1")),
426                    pool_id: Some(vec![StringFilter::from("1"), StringFilter::from("2")]),
427                    country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
428                    datacenter: Some(true),
429                    residential: Some(true),
430                    mobile: Some(true),
431                    ..Default::default()
432                }),
433            ),
434            (
435                "john-country-us-datacenter-pool-1-!residential-mobile-id-1-country-uk-pool-2",
436                String::from("john"),
437                Some(ProxyFilter {
438                    id: Some(NonEmptyString::from_static("1")),
439                    pool_id: Some(vec![StringFilter::from("1"), StringFilter::from("2")]),
440                    country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
441                    datacenter: Some(true),
442                    residential: Some(false),
443                    mobile: Some(true),
444                    ..Default::default()
445                }),
446            ),
447            (
448                "john-country-us-datacenter-pool-1-residential-mobile-id-1-country-uk-pool-2-datacenter",
449                String::from("john"),
450                Some(ProxyFilter {
451                    id: Some(NonEmptyString::from_static("1")),
452                    pool_id: Some(vec![StringFilter::from("1"), StringFilter::from("2")]),
453                    country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
454                    datacenter: Some(true),
455                    residential: Some(true),
456                    mobile: Some(true),
457                    ..Default::default()
458                }),
459            ),
460            (
461                "john-country-us-datacenter-pool-1-residential-mobile-id-1-country-uk-pool-2-datacenter-residential",
462                String::from("john"),
463                Some(ProxyFilter {
464                    id: Some(NonEmptyString::from_static("1")),
465                    pool_id: Some(vec![StringFilter::from("1"), StringFilter::from("2")]),
466                    country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
467                    datacenter: Some(true),
468                    residential: Some(true),
469                    mobile: Some(true),
470                    ..Default::default()
471                }),
472            ),
473            (
474                "john-country-us-datacenter-pool-1-residential-mobile-id-1-country-uk-pool-2-datacenter-residential-mobile",
475                String::from("john"),
476                Some(ProxyFilter {
477                    id: Some(NonEmptyString::from_static("1")),
478                    pool_id: Some(vec![StringFilter::from("1"), StringFilter::from("2")]),
479                    country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
480                    datacenter: Some(true),
481                    residential: Some(true),
482                    mobile: Some(true),
483                    ..Default::default()
484                }),
485            ),
486            (
487                "john-continent-americas-country-us-state-NY-city-ny-asn-7018",
488                String::from("john"),
489                Some(ProxyFilter {
490                    continent: Some(vec![StringFilter::from("americas")]),
491                    country: Some(vec![StringFilter::from("us")]),
492                    state: Some(vec![StringFilter::from("ny")]),
493                    city: Some(vec![StringFilter::from("ny")]),
494                    asn: Some(vec![Asn::from_static(7018)]),
495                    ..Default::default()
496                }),
497            ),
498            (
499                "john-continent-europe-continent-asia",
500                String::from("john"),
501                Some(ProxyFilter {
502                    continent: Some(vec![
503                        StringFilter::from("europe"),
504                        StringFilter::from("asia"),
505                    ]),
506                    ..Default::default()
507                }),
508            ),
509            (
510                "john-country-us-datacenter-pool-1-residential-mobile-id-1-country-uk-pool-2-!datacenter-!residential-!mobile",
511                String::from("john"),
512                Some(ProxyFilter {
513                    id: Some(NonEmptyString::from_static("1")),
514                    pool_id: Some(vec![StringFilter::from("1"), StringFilter::from("2")]),
515                    country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
516                    datacenter: Some(false),
517                    residential: Some(false),
518                    mobile: Some(false),
519                    ..Default::default()
520                }),
521            ),
522        ];
523
524        for (username, expected_username, expected_filter) in test_cases.into_iter() {
525            let mut ext = Extensions::default();
526
527            let parser = ProxyFilterUsernameParser::default();
528
529            let username = parse_username(&mut ext, parser, username).unwrap();
530            let filter = ext.get::<ProxyFilter>().cloned();
531            assert_eq!(
532                username, expected_username,
533                "username = '{username}' ; expected_username = '{expected_username}'",
534            );
535            assert_eq!(
536                filter, expected_filter,
537                "username = '{username}' ; expected_username = '{expected_username}'",
538            );
539        }
540    }
541
542    #[test]
543    fn test_username_config_error() {
544        for username in [
545            "john-country-us-datacenter-",
546            "",
547            "-",
548            "john-country-us-datacenter-pool",
549            "john-foo",
550            "john-foo-country",
551            "john-country",
552            "john-id-", // empty id is invalid
553        ] {
554            let mut ext = Extensions::default();
555
556            let parser = ProxyFilterUsernameParser::default();
557
558            assert!(
559                parse_username(&mut ext, parser, username).is_err(),
560                "username = {username}",
561            );
562        }
563    }
564
565    #[test]
566    fn test_username_negation_key_failures() {
567        for username in [
568            "john-!id-a",
569            "john-!pool-b",
570            "john-!country-us",
571            "john-!city-ny",
572            "john-!carrier-c",
573        ] {
574            let mut ext = Extensions::default();
575
576            let parser = ProxyFilterUsernameParser::default();
577
578            assert!(
579                parse_username(&mut ext, parser, username).is_err(),
580                "username = {username}",
581            );
582        }
583    }
584
585    #[test]
586    fn test_username_compose_parser_proxy_filter() {
587        let test_cases = [
588            ProxyFilter::default(),
589            ProxyFilter {
590                id: Some(NonEmptyString::from_static("p42")),
591                ..Default::default()
592            },
593            ProxyFilter {
594                id: Some(NonEmptyString::from_static("1")),
595                pool_id: Some(vec![StringFilter::from("1")]),
596                country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
597                datacenter: Some(false),
598                residential: Some(true),
599                mobile: Some(true),
600                ..Default::default()
601            },
602            ProxyFilter {
603                id: Some(NonEmptyString::from_static("1")),
604                pool_id: Some(vec![StringFilter::from("1"), StringFilter::from("2")]),
605                country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
606                datacenter: Some(false),
607                residential: Some(false),
608                mobile: Some(false),
609                ..Default::default()
610            },
611            ProxyFilter {
612                id: Some(NonEmptyString::from_static("a")),
613                pool_id: Some(vec![StringFilter::from("1"), StringFilter::from("2")]),
614                continent: Some(vec![StringFilter::from("na"), StringFilter::from("eu")]),
615                country: Some(vec![StringFilter::from("us"), StringFilter::from("be")]),
616                state: Some(vec![
617                    StringFilter::from("ca"),
618                    StringFilter::from("ny"),
619                    StringFilter::from("ovl"),
620                ]),
621                city: Some(vec![
622                    StringFilter::from("berkeley"),
623                    StringFilter::from("bruxelles"),
624                    StringFilter::from("gent"),
625                ]),
626                datacenter: Some(false),
627                residential: Some(true),
628                mobile: Some(true),
629                carrier: Some(vec![
630                    StringFilter::from("at&t"),
631                    StringFilter::from("orange"),
632                ]),
633                asn: Some(vec![Asn::from_static(7018), Asn::from_static(1)]),
634            },
635        ];
636
637        for test_case in test_cases {
638            let fmt_username = compose_username("john".to_owned(), &test_case).unwrap();
639            let mut ext = Extensions::new();
640            let username = parse_username(
641                &mut ext,
642                ProxyFilterUsernameParser::default(),
643                &fmt_username,
644            )
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::<ProxyFilter>().unwrap();
651                assert_eq!(test_case, *result);
652            }
653        }
654    }
655}