rama_proxy/
username.rs

1use super::ProxyFilter;
2use rama_core::{
3    context::Extensions,
4    error::{OpaqueError, error},
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            ("john", String::from("john"), None),
271            (
272                "john-datacenter",
273                String::from("john"),
274                Some(ProxyFilter {
275                    datacenter: Some(true),
276                    ..Default::default()
277                }),
278            ),
279            (
280                "john-!datacenter",
281                String::from("john"),
282                Some(ProxyFilter {
283                    datacenter: Some(false),
284                    ..Default::default()
285                }),
286            ),
287            (
288                "john-country-us-datacenter",
289                String::from("john"),
290                Some(ProxyFilter {
291                    country: Some(vec!["us".into()]),
292                    datacenter: Some(true),
293                    ..Default::default()
294                }),
295            ),
296            (
297                "john-city-tokyo-residential",
298                String::from("john"),
299                Some(ProxyFilter {
300                    city: Some(vec!["tokyo".into()]),
301                    residential: Some(true),
302                    ..Default::default()
303                }),
304            ),
305            (
306                "john-country-us-datacenter-pool-1",
307                String::from("john"),
308                Some(ProxyFilter {
309                    pool_id: Some(vec![StringFilter::from("1")]),
310                    country: Some(vec![StringFilter::from("us")]),
311                    datacenter: Some(true),
312                    ..Default::default()
313                }),
314            ),
315            (
316                "john-country-us-datacenter-pool-1-residential",
317                String::from("john"),
318                Some(ProxyFilter {
319                    pool_id: Some(vec![StringFilter::from("1")]),
320                    country: Some(vec![StringFilter::from("us")]),
321                    datacenter: Some(true),
322                    residential: Some(true),
323                    ..Default::default()
324                }),
325            ),
326            (
327                "john-country-us-datacenter-pool-1-residential-mobile",
328                String::from("john"),
329                Some(ProxyFilter {
330                    pool_id: Some(vec![StringFilter::from("1")]),
331                    country: Some(vec![StringFilter::from("us")]),
332                    datacenter: Some(true),
333                    residential: Some(true),
334                    mobile: Some(true),
335                    ..Default::default()
336                }),
337            ),
338            (
339                "john-country-us-datacenter-pool-1-residential-!mobile",
340                String::from("john"),
341                Some(ProxyFilter {
342                    pool_id: Some(vec![StringFilter::from("1")]),
343                    country: Some(vec![StringFilter::from("us")]),
344                    datacenter: Some(true),
345                    residential: Some(true),
346                    mobile: Some(false),
347                    ..Default::default()
348                }),
349            ),
350            (
351                "john-country-us-city-california-datacenter-pool-1-!residential-mobile",
352                String::from("john"),
353                Some(ProxyFilter {
354                    pool_id: Some(vec![StringFilter::from("1")]),
355                    country: Some(vec![StringFilter::from("us")]),
356                    city: Some(vec![StringFilter::from("california")]),
357                    datacenter: Some(true),
358                    residential: Some(false),
359                    mobile: Some(true),
360                    ..Default::default()
361                }),
362            ),
363            (
364                "john-country-us-datacenter-pool-1-residential-mobile-id-1",
365                String::from("john"),
366                Some(ProxyFilter {
367                    id: Some(NonEmptyString::from_static("1")),
368                    pool_id: Some(vec![StringFilter::from("1")]),
369                    country: Some(vec![StringFilter::from("us")]),
370                    datacenter: Some(true),
371                    residential: Some(true),
372                    mobile: Some(true),
373                    ..Default::default()
374                }),
375            ),
376            (
377                "john-country-us-datacenter-pool-1-residential-mobile-carrier-bar-id-1",
378                String::from("john"),
379                Some(ProxyFilter {
380                    id: Some(NonEmptyString::from_static("1")),
381                    pool_id: Some(vec![StringFilter::from("1")]),
382                    country: Some(vec![StringFilter::from("us")]),
383                    datacenter: Some(true),
384                    residential: Some(true),
385                    mobile: Some(true),
386                    carrier: Some(vec![StringFilter::from("bar")]),
387                    ..Default::default()
388                }),
389            ),
390            (
391                "john-country-us-datacenter-pool-1-residential-mobile-id-1-country-uk",
392                String::from("john"),
393                Some(ProxyFilter {
394                    id: Some(NonEmptyString::from_static("1")),
395                    pool_id: Some(vec![StringFilter::from("1")]),
396                    country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
397                    datacenter: Some(true),
398                    residential: Some(true),
399                    mobile: Some(true),
400                    ..Default::default()
401                }),
402            ),
403            (
404                "john-country-us-!datacenter-pool-1-residential-mobile-id-1-country-uk",
405                String::from("john"),
406                Some(ProxyFilter {
407                    id: Some(NonEmptyString::from_static("1")),
408                    pool_id: Some(vec![StringFilter::from("1")]),
409                    country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
410                    datacenter: Some(false),
411                    residential: Some(true),
412                    mobile: Some(true),
413                    ..Default::default()
414                }),
415            ),
416            (
417                "john-country-us-datacenter-pool-1-residential-mobile-id-1-country-uk-pool-2",
418                String::from("john"),
419                Some(ProxyFilter {
420                    id: Some(NonEmptyString::from_static("1")),
421                    pool_id: Some(vec![StringFilter::from("1"), StringFilter::from("2")]),
422                    country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
423                    datacenter: Some(true),
424                    residential: Some(true),
425                    mobile: Some(true),
426                    ..Default::default()
427                }),
428            ),
429            (
430                "john-country-us-datacenter-pool-1-!residential-mobile-id-1-country-uk-pool-2",
431                String::from("john"),
432                Some(ProxyFilter {
433                    id: Some(NonEmptyString::from_static("1")),
434                    pool_id: Some(vec![StringFilter::from("1"), StringFilter::from("2")]),
435                    country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
436                    datacenter: Some(true),
437                    residential: Some(false),
438                    mobile: Some(true),
439                    ..Default::default()
440                }),
441            ),
442            (
443                "john-country-us-datacenter-pool-1-residential-mobile-id-1-country-uk-pool-2-datacenter",
444                String::from("john"),
445                Some(ProxyFilter {
446                    id: Some(NonEmptyString::from_static("1")),
447                    pool_id: Some(vec![StringFilter::from("1"), StringFilter::from("2")]),
448                    country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
449                    datacenter: Some(true),
450                    residential: Some(true),
451                    mobile: Some(true),
452                    ..Default::default()
453                }),
454            ),
455            (
456                "john-country-us-datacenter-pool-1-residential-mobile-id-1-country-uk-pool-2-datacenter-residential",
457                String::from("john"),
458                Some(ProxyFilter {
459                    id: Some(NonEmptyString::from_static("1")),
460                    pool_id: Some(vec![StringFilter::from("1"), StringFilter::from("2")]),
461                    country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
462                    datacenter: Some(true),
463                    residential: Some(true),
464                    mobile: Some(true),
465                    ..Default::default()
466                }),
467            ),
468            (
469                "john-country-us-datacenter-pool-1-residential-mobile-id-1-country-uk-pool-2-datacenter-residential-mobile",
470                String::from("john"),
471                Some(ProxyFilter {
472                    id: Some(NonEmptyString::from_static("1")),
473                    pool_id: Some(vec![StringFilter::from("1"), StringFilter::from("2")]),
474                    country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
475                    datacenter: Some(true),
476                    residential: Some(true),
477                    mobile: Some(true),
478                    ..Default::default()
479                }),
480            ),
481            (
482                "john-continent-americas-country-us-state-NY-city-ny-asn-7018",
483                String::from("john"),
484                Some(ProxyFilter {
485                    continent: Some(vec![StringFilter::from("americas")]),
486                    country: Some(vec![StringFilter::from("us")]),
487                    state: Some(vec![StringFilter::from("ny")]),
488                    city: Some(vec![StringFilter::from("ny")]),
489                    asn: Some(vec![Asn::from_static(7018)]),
490                    ..Default::default()
491                }),
492            ),
493            (
494                "john-continent-europe-continent-asia",
495                String::from("john"),
496                Some(ProxyFilter {
497                    continent: Some(vec![
498                        StringFilter::from("europe"),
499                        StringFilter::from("asia"),
500                    ]),
501                    ..Default::default()
502                }),
503            ),
504            (
505                "john-country-us-datacenter-pool-1-residential-mobile-id-1-country-uk-pool-2-!datacenter-!residential-!mobile",
506                String::from("john"),
507                Some(ProxyFilter {
508                    id: Some(NonEmptyString::from_static("1")),
509                    pool_id: Some(vec![StringFilter::from("1"), StringFilter::from("2")]),
510                    country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
511                    datacenter: Some(false),
512                    residential: Some(false),
513                    mobile: Some(false),
514                    ..Default::default()
515                }),
516            ),
517        ];
518
519        for (username, expected_username, expected_filter) in test_cases.into_iter() {
520            let mut ext = Extensions::default();
521
522            let parser = ProxyFilterUsernameParser::default();
523
524            let username = parse_username(&mut ext, parser, username).unwrap();
525            let filter = ext.get::<ProxyFilter>().cloned();
526            assert_eq!(
527                username, expected_username,
528                "username = '{}' ; expected_username = '{}'",
529                username, expected_username
530            );
531            assert_eq!(
532                filter, expected_filter,
533                "username = '{}' ; expected_username = '{}'",
534                username, expected_username
535            );
536        }
537    }
538
539    #[test]
540    fn test_username_config_error() {
541        for username in [
542            "john-country-us-datacenter-",
543            "",
544            "-",
545            "john-country-us-datacenter-pool",
546            "john-foo",
547            "john-foo-country",
548            "john-country",
549            "john-id-", // empty id is invalid
550        ] {
551            let mut ext = Extensions::default();
552
553            let parser = ProxyFilterUsernameParser::default();
554
555            assert!(
556                parse_username(&mut ext, parser, username).is_err(),
557                "username = {}",
558                username
559            );
560        }
561    }
562
563    #[test]
564    fn test_username_negation_key_failures() {
565        for username in [
566            "john-!id-a",
567            "john-!pool-b",
568            "john-!country-us",
569            "john-!city-ny",
570            "john-!carrier-c",
571        ] {
572            let mut ext = Extensions::default();
573
574            let parser = ProxyFilterUsernameParser::default();
575
576            assert!(
577                parse_username(&mut ext, parser, username).is_err(),
578                "username = {}",
579                username
580            );
581        }
582    }
583
584    #[test]
585    fn test_username_compose_parser_proxy_filter() {
586        let test_cases = [
587            ProxyFilter::default(),
588            ProxyFilter {
589                id: Some(NonEmptyString::from_static("p42")),
590                ..Default::default()
591            },
592            ProxyFilter {
593                id: Some(NonEmptyString::from_static("1")),
594                pool_id: Some(vec![StringFilter::from("1")]),
595                country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
596                datacenter: Some(false),
597                residential: Some(true),
598                mobile: Some(true),
599                ..Default::default()
600            },
601            ProxyFilter {
602                id: Some(NonEmptyString::from_static("1")),
603                pool_id: Some(vec![StringFilter::from("1"), StringFilter::from("2")]),
604                country: Some(vec![StringFilter::from("us"), StringFilter::from("uk")]),
605                datacenter: Some(false),
606                residential: Some(false),
607                mobile: Some(false),
608                ..Default::default()
609            },
610            ProxyFilter {
611                id: Some(NonEmptyString::from_static("a")),
612                pool_id: Some(vec![StringFilter::from("1"), StringFilter::from("2")]),
613                continent: Some(vec![StringFilter::from("na"), StringFilter::from("eu")]),
614                country: Some(vec![StringFilter::from("us"), StringFilter::from("be")]),
615                state: Some(vec![
616                    StringFilter::from("ca"),
617                    StringFilter::from("ny"),
618                    StringFilter::from("ovl"),
619                ]),
620                city: Some(vec![
621                    StringFilter::from("berkeley"),
622                    StringFilter::from("bruxelles"),
623                    StringFilter::from("gent"),
624                ]),
625                datacenter: Some(false),
626                residential: Some(true),
627                mobile: Some(true),
628                carrier: Some(vec![
629                    StringFilter::from("at&t"),
630                    StringFilter::from("orange"),
631                ]),
632                asn: Some(vec![Asn::from_static(7018), Asn::from_static(1)]),
633            },
634        ];
635
636        for test_case in test_cases {
637            let fmt_username = compose_username("john".to_owned(), &test_case).unwrap();
638            let mut ext = Extensions::new();
639            let username = parse_username(
640                &mut ext,
641                ProxyFilterUsernameParser::default(),
642                &fmt_username,
643            )
644            .unwrap_or_else(|_| panic!("to be ok: {fmt_username}"));
645            assert_eq!("john", username);
646            if test_case == Default::default() {
647                assert!(!ext.contains::<ProxyFilter>());
648            } else {
649                let result = ext.get::<ProxyFilter>().unwrap();
650                assert_eq!(test_case, *result);
651            }
652        }
653    }
654}
rama_proxy/username.rs

rama_proxy/
username.rs
