rama_proxy/
username.rs

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

rama_proxy/
username.rs
