Skip to main content

rama_proxy/proxydb/
mod.rs

1use rama_core::error::BoxErrorExt as _;
2use rama_core::error::{BoxError, ErrorContext};
3use rama_core::extensions::Extension;
4use rama_net::asn::Asn;
5use rama_utils::collections::NonEmptyVec;
6use rama_utils::str::NonEmptyStr;
7use serde::{Deserialize, Serialize};
8use std::fmt;
9use std::num::NonZeroUsize;
10
11#[cfg(feature = "live-update")]
12mod update;
13#[cfg(feature = "live-update")]
14#[cfg_attr(docsrs, doc(cfg(feature = "live-update")))]
15#[doc(inline)]
16pub use update::{LiveUpdateProxyDB, LiveUpdateProxyDBSetter, proxy_db_updater};
17
18mod context;
19pub use context::ProxyContext;
20
21mod internal;
22#[doc(inline)]
23pub use internal::Proxy;
24
25#[cfg(feature = "csv")]
26mod csv;
27
28#[cfg(feature = "csv")]
29#[cfg_attr(docsrs, doc(cfg(feature = "csv")))]
30#[doc(inline)]
31pub use csv::{ProxyCsvRowReader, ProxyCsvRowReaderError, ProxyCsvRowReaderErrorKind};
32
33pub(super) mod layer;
34
35mod str;
36#[doc(inline)]
37pub use str::StringFilter;
38
39#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Extension)]
40#[extension(tags(proxy))]
41/// ID of a proxy selected for a connection attempt.
42///
43/// In plural routing mode this is attached to the corresponding route and
44/// installed on that route's isolated input by `ProxyRoutesConnector`.
45pub struct ProxyID(NonEmptyStr);
46
47impl ProxyID {
48    /// View  this [`ProxyID`] as a `str`.
49    #[must_use]
50    pub fn as_str(&self) -> &str {
51        self.0.as_ref()
52    }
53}
54
55impl AsRef<str> for ProxyID {
56    fn as_ref(&self) -> &str {
57        self.0.as_ref()
58    }
59}
60
61impl fmt::Display for ProxyID {
62    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
63        self.0.fmt(f)
64    }
65}
66
67impl From<NonEmptyStr> for ProxyID {
68    fn from(value: NonEmptyStr) -> Self {
69        Self(value)
70    }
71}
72
73#[derive(Debug, Default, Clone, Serialize, Deserialize, PartialEq, Eq, Hash, Extension)]
74#[extension(tags(proxy))]
75/// Filter to select a specific kind of proxy.
76///
77/// If the `id` is specified the other fields are used
78/// as a validator to see if the only possible matching proxy
79/// matches these fields.
80///
81/// If the `id` is not specified, the other fields select matching proxy
82/// candidates. The database defines their preference order; [`MemoryProxyDB`]
83/// randomizes its complete matching set.
84///
85/// Filters can be combined to make combinations with special meaning.
86/// E.g. `datacenter:true, residential:true` is essentially an ISP proxy.
87///
88/// ## Usage
89///
90/// - Use `HeaderConfigLayer` (`rama-http`) to have this proxy filter be given by the http request headers,
91///   which will add the extracted and parsed [`ProxyFilter`] to the input [`Extensions`].
92/// - Or extract yourself from the username/token validated in the `ProxyAuthLayer` (`rama-http`)
93///   to add it manually to the input [`Extensions`].
94///
95/// [`Extensions`]: rama_core::extensions::Extensions
96pub struct ProxyFilter {
97    /// The ID of the proxy to select.
98    pub id: Option<NonEmptyStr>,
99
100    /// The ID of the pool from which to select the proxy.
101    #[serde(alias = "pool")]
102    pub pool_id: Option<Vec<StringFilter>>,
103
104    /// The continent of the proxy.
105    pub continent: Option<Vec<StringFilter>>,
106
107    /// The country of the proxy.
108    pub country: Option<Vec<StringFilter>>,
109
110    /// The state of the proxy.
111    pub state: Option<Vec<StringFilter>>,
112
113    /// The city of the proxy.
114    pub city: Option<Vec<StringFilter>>,
115
116    /// Set explicitly to `true` to select a datacenter proxy.
117    pub datacenter: Option<bool>,
118
119    /// Set explicitly to `true` to select a residential proxy.
120    pub residential: Option<bool>,
121
122    /// Set explicitly to `true` to select a mobile proxy.
123    pub mobile: Option<bool>,
124
125    /// The mobile carrier desired.
126    pub carrier: Option<Vec<StringFilter>>,
127
128    ///  Autonomous System Number (ASN).
129    pub asn: Option<Vec<Asn>>,
130}
131
132/// The trait to implement to provide a proxy database to other facilities,
133/// such as connection pools, to provide proxy candidates based on the given
134/// [`ProxyContext`] and [`ProxyFilter`].
135pub trait ProxyDB: Send + Sync + 'static {
136    /// The error type that can be returned by the proxy database
137    ///
138    /// Examples are generic I/O issues or
139    /// even more common if no proxy match could be found.
140    type Error: Send + 'static;
141
142    /// Return matching [`Proxy`] values in database-defined preference order,
143    /// after applying an additional predicate and optional result limit.
144    fn get_proxies_if(
145        &self,
146        ctx: ProxyContext,
147        filter: ProxyFilter,
148        predicate: impl ProxyQueryPredicate,
149        limit: Option<NonZeroUsize>,
150    ) -> impl Future<Output = Result<NonEmptyVec<Proxy>, Self::Error>> + Send + '_;
151
152    /// Return matching [`Proxy`] values in database-defined preference order,
153    /// up to the optional result limit.
154    fn get_proxies(
155        &self,
156        ctx: ProxyContext,
157        filter: ProxyFilter,
158        limit: Option<NonZeroUsize>,
159    ) -> impl Future<Output = Result<NonEmptyVec<Proxy>, Self::Error>> + Send + '_ {
160        self.get_proxies_if(ctx, filter, true, limit)
161    }
162
163    /// Return one matching proxy after applying an additional predicate.
164    ///
165    /// Implementations may override this to provide a more efficient random
166    /// selection path. By default the first plural result is returned.
167    fn get_proxy_if(
168        &self,
169        ctx: ProxyContext,
170        filter: ProxyFilter,
171        predicate: impl ProxyQueryPredicate,
172    ) -> impl Future<Output = Result<Proxy, Self::Error>> + Send + '_ {
173        async move {
174            let proxies = self
175                .get_proxies_if(ctx, filter, predicate, Some(NonZeroUsize::MIN))
176                .await?;
177            Ok(proxies.head)
178        }
179    }
180
181    /// Get a [`Proxy`] based on the given [`ProxyContext`] and [`ProxyFilter`],
182    /// or return an error in case no [`Proxy`] could be returned.
183    fn get_proxy(
184        &self,
185        ctx: ProxyContext,
186        filter: ProxyFilter,
187    ) -> impl Future<Output = Result<Proxy, Self::Error>> + Send + '_ {
188        self.get_proxy_if(ctx, filter, true)
189    }
190}
191
192impl ProxyDB for () {
193    type Error = BoxError;
194
195    #[inline]
196    async fn get_proxies_if(
197        &self,
198        _ctx: ProxyContext,
199        _filter: ProxyFilter,
200        _predicate: impl ProxyQueryPredicate,
201        _limit: Option<NonZeroUsize>,
202    ) -> Result<NonEmptyVec<Proxy>, Self::Error> {
203        Err(BoxError::from_static_str(
204            "()::get_proxies_if: no ProxyDB defined",
205        ))
206    }
207
208    #[inline]
209    async fn get_proxy_if(
210        &self,
211        _ctx: ProxyContext,
212        _filter: ProxyFilter,
213        _predicate: impl ProxyQueryPredicate,
214    ) -> Result<Proxy, Self::Error> {
215        Err(BoxError::from_static_str(
216            "()::get_proxy_if: no ProxyDB defined",
217        ))
218    }
219
220    #[inline]
221    async fn get_proxy(
222        &self,
223        _ctx: ProxyContext,
224        _filter: ProxyFilter,
225    ) -> Result<Proxy, Self::Error> {
226        Err(BoxError::from_static_str(
227            "()::get_proxy: no ProxyDB defined",
228        ))
229    }
230}
231
232impl<T> ProxyDB for Option<T>
233where
234    T: ProxyDB<Error: Into<BoxError>>,
235{
236    type Error = BoxError;
237
238    #[inline]
239    async fn get_proxies_if(
240        &self,
241        ctx: ProxyContext,
242        filter: ProxyFilter,
243        predicate: impl ProxyQueryPredicate,
244        limit: Option<NonZeroUsize>,
245    ) -> Result<NonEmptyVec<Proxy>, Self::Error> {
246        match self {
247            Some(db) => db
248                .get_proxies_if(ctx, filter, predicate, limit)
249                .await
250                .context("Some::get_proxies_if"),
251            None => Err(BoxError::from_static_str(
252                "None::get_proxies_if: no ProxyDB defined",
253            )),
254        }
255    }
256
257    #[inline]
258    async fn get_proxies(
259        &self,
260        ctx: ProxyContext,
261        filter: ProxyFilter,
262        limit: Option<NonZeroUsize>,
263    ) -> Result<NonEmptyVec<Proxy>, Self::Error> {
264        match self {
265            Some(db) => db
266                .get_proxies(ctx, filter, limit)
267                .await
268                .context("Some::get_proxies"),
269            None => Err(BoxError::from_static_str(
270                "None::get_proxies: no ProxyDB defined",
271            )),
272        }
273    }
274
275    #[inline]
276    async fn get_proxy_if(
277        &self,
278        ctx: ProxyContext,
279        filter: ProxyFilter,
280        predicate: impl ProxyQueryPredicate,
281    ) -> Result<Proxy, Self::Error> {
282        match self {
283            Some(db) => db
284                .get_proxy_if(ctx, filter, predicate)
285                .await
286                .context("Some::get_proxy_if"),
287            None => Err(BoxError::from_static_str(
288                "None::get_proxy_if: no ProxyDB defined",
289            )),
290        }
291    }
292
293    #[inline]
294    async fn get_proxy(
295        &self,
296        ctx: ProxyContext,
297        filter: ProxyFilter,
298    ) -> Result<Proxy, Self::Error> {
299        match self {
300            Some(db) => db.get_proxy(ctx, filter).await.context("Some::get_proxy"),
301            None => Err(BoxError::from_static_str(
302                "None::get_proxy: no ProxyDB defined",
303            )),
304        }
305    }
306}
307
308impl<T> ProxyDB for std::sync::Arc<T>
309where
310    T: ProxyDB,
311{
312    type Error = T::Error;
313
314    #[inline]
315    fn get_proxies_if(
316        &self,
317        ctx: ProxyContext,
318        filter: ProxyFilter,
319        predicate: impl ProxyQueryPredicate,
320        limit: Option<NonZeroUsize>,
321    ) -> impl Future<Output = Result<NonEmptyVec<Proxy>, Self::Error>> + Send + '_ {
322        (**self).get_proxies_if(ctx, filter, predicate, limit)
323    }
324
325    #[inline]
326    fn get_proxies(
327        &self,
328        ctx: ProxyContext,
329        filter: ProxyFilter,
330        limit: Option<NonZeroUsize>,
331    ) -> impl Future<Output = Result<NonEmptyVec<Proxy>, Self::Error>> + Send + '_ {
332        (**self).get_proxies(ctx, filter, limit)
333    }
334
335    #[inline]
336    fn get_proxy_if(
337        &self,
338        ctx: ProxyContext,
339        filter: ProxyFilter,
340        predicate: impl ProxyQueryPredicate,
341    ) -> impl Future<Output = Result<Proxy, Self::Error>> + Send + '_ {
342        (**self).get_proxy_if(ctx, filter, predicate)
343    }
344
345    #[inline]
346    fn get_proxy(
347        &self,
348        ctx: ProxyContext,
349        filter: ProxyFilter,
350    ) -> impl Future<Output = Result<Proxy, Self::Error>> + Send + '_ {
351        (**self).get_proxy(ctx, filter)
352    }
353}
354
355macro_rules! impl_proxydb_either {
356    ($id:ident, $($param:ident),+ $(,)?) => {
357        impl<$($param),+> ProxyDB for rama_core::combinators::$id<$($param),+>
358        where
359            $(
360                $param: ProxyDB<Error: Into<BoxError>>,
361            )+
362    {
363        type Error = BoxError;
364
365        #[inline]
366        async fn get_proxies_if(
367            &self,
368            ctx: ProxyContext,
369            filter: ProxyFilter,
370            predicate: impl ProxyQueryPredicate,
371            limit: Option<NonZeroUsize>,
372        ) -> Result<NonEmptyVec<Proxy>, Self::Error> {
373            match self {
374                $(
375                    rama_core::combinators::$id::$param(s) => s.get_proxies_if(ctx, filter, predicate, limit).await.into_box_error(),
376                )+
377            }
378        }
379
380        #[inline]
381        async fn get_proxies(
382            &self,
383            ctx: ProxyContext,
384            filter: ProxyFilter,
385            limit: Option<NonZeroUsize>,
386        ) -> Result<NonEmptyVec<Proxy>, Self::Error> {
387            match self {
388                $(
389                    rama_core::combinators::$id::$param(s) => s.get_proxies(ctx, filter, limit).await.into_box_error(),
390                )+
391            }
392        }
393
394        #[inline]
395        async fn get_proxy_if(
396            &self,
397            ctx: ProxyContext,
398            filter: ProxyFilter,
399            predicate: impl ProxyQueryPredicate,
400        ) -> Result<Proxy, Self::Error> {
401            match self {
402                $(
403                    rama_core::combinators::$id::$param(s) => s.get_proxy_if(ctx, filter, predicate).await.into_box_error(),
404                )+
405            }
406        }
407
408        #[inline]
409        async fn get_proxy(
410            &self,
411            ctx: ProxyContext,
412            filter: ProxyFilter,
413        ) -> Result<Proxy, Self::Error> {
414            match self {
415                $(
416                    rama_core::combinators::$id::$param(s) => s.get_proxy(ctx, filter).await.into_box_error(),
417                )+
418            }
419        }
420        }
421    };
422}
423
424rama_core::combinators::impl_either!(impl_proxydb_either);
425
426/// Trait that is used by the [`ProxyDB`] for providing an optional
427/// filter predicate to rule out returned results.
428pub trait ProxyQueryPredicate: Clone + Send + Sync + 'static {
429    /// Execute the predicate.
430    fn execute(&self, proxy: &Proxy) -> bool;
431}
432
433impl ProxyQueryPredicate for bool {
434    fn execute(&self, _proxy: &Proxy) -> bool {
435        *self
436    }
437}
438
439impl<F> ProxyQueryPredicate for F
440where
441    F: Fn(&Proxy) -> bool + Clone + Send + Sync + 'static,
442{
443    fn execute(&self, proxy: &Proxy) -> bool {
444        (self)(proxy)
445    }
446}
447
448impl ProxyDB for Proxy {
449    type Error = BoxError;
450
451    async fn get_proxies_if(
452        &self,
453        ctx: ProxyContext,
454        filter: ProxyFilter,
455        predicate: impl ProxyQueryPredicate,
456        _limit: Option<NonZeroUsize>,
457    ) -> Result<NonEmptyVec<Self>, Self::Error> {
458        (self.is_match(&ctx, &filter) && predicate.execute(self))
459            .then(|| NonEmptyVec::new(self.clone()))
460            .context("hardcoded proxy no match")
461    }
462
463    async fn get_proxy_if(
464        &self,
465        ctx: ProxyContext,
466        filter: ProxyFilter,
467        predicate: impl ProxyQueryPredicate,
468    ) -> Result<Self, Self::Error> {
469        (self.is_match(&ctx, &filter) && predicate.execute(self))
470            .then(|| self.clone())
471            .context("hardcoded proxy no match")
472    }
473}
474
475#[cfg(test)]
476mod trait_tests {
477    use super::*;
478    use rama_core::combinators::Either;
479    use rama_net::{address::ProxyAddress, transport::TransportProtocol};
480    use rama_utils::str::non_empty_str;
481    use std::{str::FromStr, sync::Arc};
482
483    fn proxy() -> Proxy {
484        Proxy {
485            id: non_empty_str!("proxy"),
486            address: ProxyAddress::from_str("proxy.example:8080").unwrap(),
487            tcp: true,
488            udp: false,
489            http: true,
490            https: false,
491            socks5: false,
492            socks5h: false,
493            datacenter: true,
494            residential: false,
495            mobile: false,
496            pool_id: None,
497            continent: None,
498            country: None,
499            state: None,
500            city: None,
501            carrier: None,
502            asn: None,
503        }
504    }
505
506    fn context() -> ProxyContext {
507        ProxyContext {
508            protocol: TransportProtocol::Tcp,
509        }
510    }
511
512    #[tokio::test]
513    async fn wrappers_forward_plural_queries() {
514        let optional = Some(Arc::new(proxy()));
515        let proxies = optional
516            .get_proxies(context(), ProxyFilter::default(), None)
517            .await
518            .unwrap();
519        assert_eq!(proxies.len(), 1);
520        assert_eq!(proxies.head.id, "proxy");
521
522        let either: Either<Proxy, Proxy> = Either::B(proxy());
523        let proxies = either
524            .get_proxies_if(context(), ProxyFilter::default(), true, None)
525            .await
526            .unwrap();
527        assert_eq!(proxies.len(), 1);
528        assert_eq!(proxies.head.id, "proxy");
529    }
530
531    #[tokio::test]
532    async fn hardcoded_proxy_applies_plural_predicate() {
533        let error = proxy()
534            .get_proxies_if(context(), ProxyFilter::default(), false, None)
535            .await
536            .unwrap_err();
537
538        assert!(error.to_string().contains("hardcoded proxy no match"));
539    }
540}
541
542#[cfg(feature = "memory-db")]
543mod memdb {
544    use super::*;
545    use crate::proxydb::internal::ProxyDBErrorKind;
546    use rama_net::transport::TransportProtocol;
547    use rand::seq::{IteratorRandom as _, SliceRandom as _};
548
549    /// A fast in-memory ProxyDatabase that is the default choice for Rama.
550    ///
551    /// Plural queries return a uniformly shuffled sample up to the requested
552    /// limit, or every match when unbounded. Singular queries retain the
553    /// allocation-efficient random-selection path.
554    #[derive(Debug)]
555    pub struct MemoryProxyDB {
556        data: internal::ProxyDB,
557    }
558
559    impl MemoryProxyDB {
560        /// Create a new in-memory proxy database with the given proxies.
561        pub fn try_from_rows(proxies: Vec<Proxy>) -> Result<Self, MemoryProxyDBInsertError> {
562            Ok(Self {
563                data: internal::ProxyDB::from_rows(proxies).map_err(|err| match err.kind() {
564                    ProxyDBErrorKind::DuplicateKey => {
565                        MemoryProxyDBInsertError::duplicate_key(err.into_input())
566                    }
567                    ProxyDBErrorKind::InvalidRow => {
568                        MemoryProxyDBInsertError::invalid_proxy(err.into_input())
569                    }
570                })?,
571            })
572        }
573
574        /// Create a new in-memory proxy database with the given proxies from an iterator.
575        pub fn try_from_iter<I>(proxies: I) -> Result<Self, MemoryProxyDBInsertError>
576        where
577            I: IntoIterator<Item = Proxy>,
578        {
579            Ok(Self {
580                data: internal::ProxyDB::from_iter(proxies).map_err(|err| match err.kind() {
581                    ProxyDBErrorKind::DuplicateKey => {
582                        MemoryProxyDBInsertError::duplicate_key(err.into_input())
583                    }
584                    ProxyDBErrorKind::InvalidRow => {
585                        MemoryProxyDBInsertError::invalid_proxy(err.into_input())
586                    }
587                })?,
588            })
589        }
590
591        /// Return the number of proxies in the database.
592        #[must_use]
593        pub fn len(&self) -> usize {
594            self.data.len()
595        }
596
597        /// Rerturns if the database is empty.
598        #[must_use]
599        pub fn is_empty(&self) -> bool {
600            self.data.is_empty()
601        }
602
603        #[expect(clippy::needless_pass_by_value)]
604        fn query_from_filter(
605            &self,
606            ctx: ProxyContext,
607            filter: ProxyFilter,
608        ) -> internal::ProxyDBQuery<'_> {
609            let mut query = self.data.query();
610
611            for pool_id in filter.pool_id.into_iter().flatten() {
612                query.pool_id(pool_id);
613            }
614            for continent in filter.continent.into_iter().flatten() {
615                query.continent(continent);
616            }
617            for country in filter.country.into_iter().flatten() {
618                query.country(country);
619            }
620            for state in filter.state.into_iter().flatten() {
621                query.state(state);
622            }
623            for city in filter.city.into_iter().flatten() {
624                query.city(city);
625            }
626            for carrier in filter.carrier.into_iter().flatten() {
627                query.carrier(carrier);
628            }
629            for asn in filter.asn.into_iter().flatten() {
630                query.asn(asn);
631            }
632
633            if let Some(value) = filter.datacenter {
634                query.datacenter(value);
635            }
636            if let Some(value) = filter.residential {
637                query.residential(value);
638            }
639            if let Some(value) = filter.mobile {
640                query.mobile(value);
641            }
642
643            match ctx.protocol {
644                TransportProtocol::Tcp => {
645                    query.tcp(true);
646                }
647                TransportProtocol::Udp => {
648                    query.udp(true).socks5(true);
649                }
650            }
651
652            query
653        }
654    }
655
656    // TODO: custom query filters using ProxyQueryPredicate
657    // might be a lot faster for cases where we want to filter a big batch of proxies,
658    // in which case a bitmap could be supported by a future VennDB version...
659    //
660    // Would just need to figure out how to allow this to happen.
661
662    impl ProxyDB for MemoryProxyDB {
663        type Error = MemoryProxyDBQueryError;
664
665        async fn get_proxies_if(
666            &self,
667            ctx: ProxyContext,
668            filter: ProxyFilter,
669            predicate: impl ProxyQueryPredicate,
670            limit: Option<NonZeroUsize>,
671        ) -> Result<NonEmptyVec<Proxy>, Self::Error> {
672            if let Some(id) = &filter.id {
673                match self.data.get_by_id(id) {
674                    None => Err(MemoryProxyDBQueryError::not_found()),
675                    Some(proxy) => {
676                        if proxy.is_match(&ctx, &filter) && predicate.execute(proxy) {
677                            Ok(NonEmptyVec::new(proxy.clone()))
678                        } else {
679                            Err(MemoryProxyDBQueryError::mismatch())
680                        }
681                    }
682                }
683            } else {
684                let query = self.query_from_filter(ctx, filter);
685                let result = query
686                    .execute()
687                    .and_then(|result| result.filter(|proxy| predicate.execute(proxy)))
688                    .ok_or_else(MemoryProxyDBQueryError::not_found)?;
689                let mut rng = rand::rng();
690                let mut proxies = match limit {
691                    Some(limit) => result
692                        .iter()
693                        .sample(&mut rng, limit.get())
694                        .into_iter()
695                        .cloned()
696                        .collect(),
697                    None => result.iter().cloned().collect::<Vec<_>>(),
698                };
699                if limit.is_none() {
700                    proxies.shuffle(&mut rng);
701                }
702                NonEmptyVec::collect(proxies).ok_or_else(MemoryProxyDBQueryError::not_found)
703            }
704        }
705
706        async fn get_proxy_if(
707            &self,
708            ctx: ProxyContext,
709            filter: ProxyFilter,
710            predicate: impl ProxyQueryPredicate,
711        ) -> Result<Proxy, Self::Error> {
712            if let Some(id) = &filter.id {
713                match self.data.get_by_id(id) {
714                    None => Err(MemoryProxyDBQueryError::not_found()),
715                    Some(proxy) => {
716                        if proxy.is_match(&ctx, &filter) && predicate.execute(proxy) {
717                            Ok(proxy.clone())
718                        } else {
719                            Err(MemoryProxyDBQueryError::mismatch())
720                        }
721                    }
722                }
723            } else {
724                let query = self.query_from_filter(ctx, filter);
725                match query
726                    .execute()
727                    .and_then(|result| result.filter(|proxy| predicate.execute(proxy)))
728                    .map(|result| result.any())
729                {
730                    None => Err(MemoryProxyDBQueryError::not_found()),
731                    Some(proxy) => Ok(proxy.clone()),
732                }
733            }
734        }
735    }
736
737    /// The error type that can be returned by [`MemoryProxyDB`] when some of the proxies
738    /// could not be inserted due to a proxy that had a duplicate key or was invalid for some other reason.
739    #[derive(Debug)]
740    pub struct MemoryProxyDBInsertError {
741        kind: MemoryProxyDBInsertErrorKind,
742        proxies: Vec<Proxy>,
743    }
744
745    impl std::fmt::Display for MemoryProxyDBInsertError {
746        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
747            match self.kind {
748                MemoryProxyDBInsertErrorKind::DuplicateKey => write!(
749                    f,
750                    "A proxy with the same key already exists in the database"
751                ),
752                MemoryProxyDBInsertErrorKind::InvalidProxy => {
753                    write!(f, "A proxy in the list is invalid for some reason")
754                }
755            }
756        }
757    }
758
759    impl std::error::Error for MemoryProxyDBInsertError {}
760
761    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
762    /// The kind of error that [`MemoryProxyDBInsertError`] represents.
763    pub enum MemoryProxyDBInsertErrorKind {
764        /// Duplicate key found in the proxies.
765        DuplicateKey,
766        /// Invalid proxy found in the proxies.
767        ///
768        /// This could be due to a proxy that is not valid for some reason.
769        /// E.g. a proxy that neither supports http or socks5.
770        InvalidProxy,
771    }
772
773    impl MemoryProxyDBInsertError {
774        fn duplicate_key(proxies: Vec<Proxy>) -> Self {
775            Self {
776                kind: MemoryProxyDBInsertErrorKind::DuplicateKey,
777                proxies,
778            }
779        }
780
781        fn invalid_proxy(proxies: Vec<Proxy>) -> Self {
782            Self {
783                kind: MemoryProxyDBInsertErrorKind::InvalidProxy,
784                proxies,
785            }
786        }
787
788        /// Returns the kind of error that [`MemoryProxyDBInsertError`] represents.
789        #[must_use]
790        pub fn kind(&self) -> MemoryProxyDBInsertErrorKind {
791            self.kind
792        }
793
794        /// Returns the proxies that were not inserted.
795        #[must_use]
796        pub fn proxies(&self) -> &[Proxy] {
797            &self.proxies
798        }
799
800        /// Consumes the error and returns the proxies that were not inserted.
801        #[must_use]
802        pub fn into_proxies(self) -> Vec<Proxy> {
803            self.proxies
804        }
805    }
806
807    /// The error type that can be returned by [`MemoryProxyDB`] when no proxy could be returned.
808    #[derive(Debug)]
809    pub struct MemoryProxyDBQueryError {
810        kind: MemoryProxyDBQueryErrorKind,
811    }
812
813    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
814    /// The kind of error that [`MemoryProxyDBQueryError`] represents.
815    pub enum MemoryProxyDBQueryErrorKind {
816        /// No proxy match could be found.
817        NotFound,
818        /// A proxy looked up by key had a config that did not match the given filters/requirements.
819        Mismatch,
820    }
821
822    impl std::fmt::Display for MemoryProxyDBQueryError {
823        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
824            match self.kind {
825                MemoryProxyDBQueryErrorKind::NotFound => write!(f, "No proxy match could be found"),
826                MemoryProxyDBQueryErrorKind::Mismatch => write!(
827                    f,
828                    "Proxy config did not match the given filters/requirements"
829                ),
830            }
831        }
832    }
833
834    impl std::error::Error for MemoryProxyDBQueryError {}
835
836    impl MemoryProxyDBQueryError {
837        /// Create a new error that indicates no proxy match could be found.
838        #[must_use]
839        pub fn not_found() -> Self {
840            Self {
841                kind: MemoryProxyDBQueryErrorKind::NotFound,
842            }
843        }
844
845        /// Create a new error that indicates a proxy looked up by key had a config that did not match the given filters/requirements.
846        #[must_use]
847        pub fn mismatch() -> Self {
848            Self {
849                kind: MemoryProxyDBQueryErrorKind::Mismatch,
850            }
851        }
852
853        /// Returns the kind of error that [`MemoryProxyDBQueryError`] represents.
854        #[must_use]
855        pub fn kind(&self) -> MemoryProxyDBQueryErrorKind {
856            self.kind
857        }
858    }
859
860    #[cfg(test)]
861    mod tests {
862        use super::*;
863        use itertools::Itertools;
864        use rama_net::address::ProxyAddress;
865        use rama_utils::str::non_empty_str;
866        use std::str::FromStr;
867
868        const RAW_CSV_DATA: &str = include_str!("./test_proxydb_rows.csv");
869
870        async fn memproxydb() -> MemoryProxyDB {
871            let mut reader = ProxyCsvRowReader::raw(RAW_CSV_DATA);
872            let mut rows = Vec::new();
873            while let Some(proxy) = reader.next().await.unwrap() {
874                rows.push(proxy);
875            }
876            MemoryProxyDB::try_from_rows(rows).unwrap()
877        }
878
879        #[tokio::test]
880        async fn test_load_memproxydb_from_rows() {
881            let db = memproxydb().await;
882            assert_eq!(db.len(), 64);
883        }
884
885        fn h2_proxy_context() -> ProxyContext {
886            ProxyContext {
887                protocol: TransportProtocol::Tcp,
888            }
889        }
890
891        #[tokio::test]
892        async fn test_memproxydb_get_proxy_by_id_found() {
893            let db = memproxydb().await;
894            let ctx = h2_proxy_context();
895            let filter = ProxyFilter {
896                id: Some(non_empty_str!("3031533634")),
897                ..Default::default()
898            };
899            let proxies = db
900                .get_proxies(ctx.clone(), filter.clone(), None)
901                .await
902                .unwrap();
903            assert_eq!(proxies.len(), 1);
904            assert_eq!(proxies.head.id, "3031533634");
905
906            let proxy = db.get_proxy(ctx, filter).await.unwrap();
907            assert_eq!(proxy.id, "3031533634");
908        }
909
910        #[tokio::test]
911        async fn test_memproxydb_get_proxy_by_id_found_correct_filters() {
912            let db = memproxydb().await;
913            let ctx = h2_proxy_context();
914            let filter = ProxyFilter {
915                id: Some(non_empty_str!("3031533634")),
916                pool_id: Some(vec![StringFilter::new("poolF")]),
917                country: Some(vec![StringFilter::new("JP")]),
918                city: Some(vec![StringFilter::new("Yokohama")]),
919                datacenter: Some(true),
920                residential: Some(false),
921                mobile: Some(true),
922                carrier: Some(vec![StringFilter::new("Verizon")]),
923                ..Default::default()
924            };
925            let proxy = db.get_proxy(ctx, filter).await.unwrap();
926            assert_eq!(proxy.id, "3031533634");
927        }
928
929        #[tokio::test]
930        async fn test_memproxydb_get_proxy_by_id_not_found() {
931            let db = memproxydb().await;
932            let ctx = h2_proxy_context();
933            let filter = ProxyFilter {
934                id: Some(non_empty_str!("notfound")),
935                ..Default::default()
936            };
937            let err = db
938                .get_proxies(ctx.clone(), filter.clone(), None)
939                .await
940                .unwrap_err();
941            assert_eq!(err.kind(), MemoryProxyDBQueryErrorKind::NotFound);
942
943            let err = db.get_proxy(ctx, filter).await.unwrap_err();
944            assert_eq!(err.kind(), MemoryProxyDBQueryErrorKind::NotFound);
945        }
946
947        #[tokio::test]
948        async fn test_memproxydb_get_proxy_by_id_mismatch_filter() {
949            let db = memproxydb().await;
950            let ctx = h2_proxy_context();
951            let filters = [
952                ProxyFilter {
953                    id: Some(non_empty_str!("3031533634")),
954                    pool_id: Some(vec![StringFilter::new("poolB")]),
955                    ..Default::default()
956                },
957                ProxyFilter {
958                    id: Some(non_empty_str!("3031533634")),
959                    country: Some(vec![StringFilter::new("US")]),
960                    ..Default::default()
961                },
962                ProxyFilter {
963                    id: Some(non_empty_str!("3031533634")),
964                    city: Some(vec![StringFilter::new("New York")]),
965                    ..Default::default()
966                },
967                ProxyFilter {
968                    id: Some(non_empty_str!("3031533634")),
969                    continent: Some(vec![StringFilter::new("americas")]),
970                    ..Default::default()
971                },
972                ProxyFilter {
973                    id: Some(non_empty_str!("3732488183")),
974                    state: Some(vec![StringFilter::new("Texas")]),
975                    ..Default::default()
976                },
977                ProxyFilter {
978                    id: Some(non_empty_str!("3031533634")),
979                    datacenter: Some(false),
980                    ..Default::default()
981                },
982                ProxyFilter {
983                    id: Some(non_empty_str!("3031533634")),
984                    residential: Some(true),
985                    ..Default::default()
986                },
987                ProxyFilter {
988                    id: Some(non_empty_str!("3031533634")),
989                    mobile: Some(false),
990                    ..Default::default()
991                },
992                ProxyFilter {
993                    id: Some(non_empty_str!("3031533634")),
994                    carrier: Some(vec![StringFilter::new("AT&T")]),
995                    ..Default::default()
996                },
997                ProxyFilter {
998                    id: Some(non_empty_str!("292096733")),
999                    asn: Some(vec![Asn::from_static(1)]),
1000                    ..Default::default()
1001                },
1002            ];
1003            for filter in filters.iter() {
1004                let err = db
1005                    .get_proxies(ctx.clone(), filter.clone(), None)
1006                    .await
1007                    .unwrap_err();
1008                assert_eq!(err.kind(), MemoryProxyDBQueryErrorKind::Mismatch);
1009
1010                let err = db.get_proxy(ctx.clone(), filter.clone()).await.unwrap_err();
1011                assert_eq!(err.kind(), MemoryProxyDBQueryErrorKind::Mismatch);
1012            }
1013        }
1014
1015        fn h3_proxy_context() -> ProxyContext {
1016            ProxyContext {
1017                protocol: TransportProtocol::Udp,
1018            }
1019        }
1020
1021        #[tokio::test]
1022        async fn test_memproxydb_get_proxy_by_id_mismatch_req_context() {
1023            let db = memproxydb().await;
1024            let ctx = h3_proxy_context();
1025            let filter = ProxyFilter {
1026                id: Some(non_empty_str!("3031533634")),
1027                ..Default::default()
1028            };
1029            // this proxy does not support socks5 UDP, which is what we need
1030            let err = db
1031                .get_proxies(ctx.clone(), filter.clone(), None)
1032                .await
1033                .unwrap_err();
1034            assert_eq!(err.kind(), MemoryProxyDBQueryErrorKind::Mismatch);
1035
1036            let err = db.get_proxy(ctx, filter).await.unwrap_err();
1037            assert_eq!(err.kind(), MemoryProxyDBQueryErrorKind::Mismatch);
1038        }
1039
1040        #[tokio::test]
1041        async fn test_memorydb_get_h3_capable_proxies() {
1042            let db = memproxydb().await;
1043            let ctx = h3_proxy_context();
1044            let filter = ProxyFilter::default();
1045            let mut found_ids = Vec::new();
1046            for _ in 0..5000 {
1047                let proxy = db.get_proxy(ctx.clone(), filter.clone()).await.unwrap();
1048                if found_ids.contains(&proxy.id) {
1049                    continue;
1050                }
1051                assert!(proxy.udp);
1052                assert!(proxy.socks5);
1053                found_ids.push(proxy.id);
1054            }
1055            assert_eq!(found_ids.len(), 40);
1056            assert_eq!(
1057                found_ids.iter().sorted().join(","),
1058                r##"1125300915,1259341971,1316455915,153202126,1571861931,1684342915,1742367441,1844412609,1916851007,20647117,2107229589,2261612122,2497865606,2521901221,2560727338,2593294918,2596743625,2745456299,2880295577,2909724448,2950022859,2951529660,3187902553,3269411602,3269465574,3269921904,3481200027,3498810974,362091157,3679054656,3732488183,3836943127,39048766,3951672504,3976711563,4187178960,56402588,724884866,738626121,906390012"##
1059            );
1060        }
1061
1062        #[tokio::test]
1063        async fn test_memorydb_get_h2_capable_proxies() {
1064            let db = memproxydb().await;
1065            let ctx = h2_proxy_context();
1066            let filter = ProxyFilter::default();
1067            let mut found_ids = Vec::new();
1068            for _ in 0..5000 {
1069                let proxy = db.get_proxy(ctx.clone(), filter.clone()).await.unwrap();
1070                if found_ids.contains(&proxy.id) {
1071                    continue;
1072                }
1073                assert!(proxy.tcp);
1074                found_ids.push(proxy.id);
1075            }
1076            assert_eq!(found_ids.len(), 50);
1077            assert_eq!(
1078                found_ids.iter().sorted().join(","),
1079                r#"1125300915,1259341971,1264821985,129108927,1316455915,1425588737,1571861931,1810781137,1836040682,1844412609,1885107293,2021561518,2079461709,2107229589,2141152822,2438596154,2497865606,2521901221,2551759475,2560727338,2593294918,2798907087,2854473221,2880295577,2909724448,2912880381,292096733,2951529660,3031533634,3187902553,3269411602,3269465574,339020035,3481200027,3498810974,3503691556,362091157,3679054656,371209663,3861736957,39048766,3976711563,4062553709,49590203,56402588,724884866,738626121,767809962,846528631,906390012"#,
1080            );
1081        }
1082
1083        #[tokio::test]
1084        async fn plural_query_returns_every_matching_proxy_once() {
1085            let db = memproxydb().await;
1086
1087            let proxies = db
1088                .get_proxies(h2_proxy_context(), ProxyFilter::default(), None)
1089                .await
1090                .unwrap();
1091            let ids = proxies
1092                .iter()
1093                .map(|proxy| proxy.id.as_ref())
1094                .sorted()
1095                .collect::<Vec<_>>();
1096
1097            assert_eq!(ids.len(), 50);
1098            assert_eq!(ids.iter().unique().count(), ids.len());
1099            assert_eq!(
1100                ids.join(","),
1101                r#"1125300915,1259341971,1264821985,129108927,1316455915,1425588737,1571861931,1810781137,1836040682,1844412609,1885107293,2021561518,2079461709,2107229589,2141152822,2438596154,2497865606,2521901221,2551759475,2560727338,2593294918,2798907087,2854473221,2880295577,2909724448,2912880381,292096733,2951529660,3031533634,3187902553,3269411602,3269465574,339020035,3481200027,3498810974,3503691556,362091157,3679054656,371209663,3861736957,39048766,3976711563,4062553709,49590203,56402588,724884866,738626121,767809962,846528631,906390012"#,
1102            );
1103        }
1104
1105        #[tokio::test]
1106        async fn plural_query_honors_result_limit() {
1107            let db = memproxydb().await;
1108
1109            let proxies = db
1110                .get_proxies(
1111                    h2_proxy_context(),
1112                    ProxyFilter::default(),
1113                    NonZeroUsize::new(5),
1114                )
1115                .await
1116                .unwrap();
1117
1118            assert_eq!(proxies.len(), 5);
1119            assert_eq!(proxies.iter().map(|proxy| &proxy.id).unique().count(), 5);
1120            assert!(proxies.iter().all(|proxy| proxy.tcp));
1121        }
1122
1123        #[tokio::test]
1124        async fn plural_query_applies_predicate_to_every_candidate() {
1125            let db = memproxydb().await;
1126
1127            let proxies = db
1128                .get_proxies_if(
1129                    h2_proxy_context(),
1130                    ProxyFilter::default(),
1131                    |proxy: &Proxy| proxy.mobile,
1132                    None,
1133                )
1134                .await
1135                .unwrap();
1136
1137            assert!(proxies.iter().all(|proxy| proxy.mobile));
1138            assert!(proxies.len() < 50);
1139        }
1140
1141        #[tokio::test]
1142        async fn test_memorydb_get_any_country_proxies() {
1143            let db = memproxydb().await;
1144            let ctx = h2_proxy_context();
1145            let filter = ProxyFilter {
1146                // there are no explicit BE proxies,
1147                // so these will only match the proxies that have a wildcard country
1148                country: Some(vec!["BE".into()]),
1149                ..Default::default()
1150            };
1151            let mut found_ids = Vec::new();
1152            for _ in 0..5000 {
1153                let proxy = db.get_proxy(ctx.clone(), filter.clone()).await.unwrap();
1154                if found_ids.contains(&proxy.id) {
1155                    continue;
1156                }
1157                found_ids.push(proxy.id);
1158            }
1159            assert_eq!(found_ids.len(), 5);
1160            assert_eq!(
1161                found_ids.iter().sorted().join(","),
1162                r#"2141152822,2593294918,2912880381,371209663,767809962"#,
1163            );
1164        }
1165
1166        #[tokio::test]
1167        async fn test_memorydb_get_illinois_proxies() {
1168            let db = memproxydb().await;
1169            let ctx = h2_proxy_context();
1170            let filter = ProxyFilter {
1171                // this will also work for proxies that have 'any' state
1172                state: Some(vec!["illinois".into()]),
1173                ..Default::default()
1174            };
1175            let mut found_ids = Vec::new();
1176            for _ in 0..5000 {
1177                let proxy = db.get_proxy(ctx.clone(), filter.clone()).await.unwrap();
1178                if found_ids.contains(&proxy.id) {
1179                    continue;
1180                }
1181                found_ids.push(proxy.id);
1182            }
1183            assert_eq!(found_ids.len(), 9);
1184            assert_eq!(
1185                found_ids.iter().sorted().join(","),
1186                r#"2141152822,2521901221,2560727338,2593294918,2912880381,292096733,371209663,39048766,767809962"#,
1187            );
1188        }
1189
1190        #[tokio::test]
1191        async fn test_memorydb_get_asn_proxies() {
1192            let db = memproxydb().await;
1193            let ctx = h2_proxy_context();
1194            let filter = ProxyFilter {
1195                // this will also work for proxies that have 'any' ASN
1196                asn: Some(vec![Asn::from_static(42)]),
1197                ..Default::default()
1198            };
1199            let mut found_ids = Vec::new();
1200            for _ in 0..5000 {
1201                let proxy = db.get_proxy(ctx.clone(), filter.clone()).await.unwrap();
1202                if found_ids.contains(&proxy.id) {
1203                    continue;
1204                }
1205                found_ids.push(proxy.id);
1206            }
1207            assert_eq!(found_ids.len(), 4);
1208            assert_eq!(
1209                found_ids.iter().sorted().join(","),
1210                r#"2141152822,2912880381,292096733,3481200027"#,
1211            );
1212        }
1213
1214        #[tokio::test]
1215        async fn test_memorydb_get_h3_capable_mobile_residential_be_asterix_proxies() {
1216            let db = memproxydb().await;
1217            let ctx = h3_proxy_context();
1218            let filter = ProxyFilter {
1219                country: Some(vec!["BE".into()]),
1220                mobile: Some(true),
1221                residential: Some(true),
1222                ..Default::default()
1223            };
1224            for _ in 0..50 {
1225                let proxy = db.get_proxy(ctx.clone(), filter.clone()).await.unwrap();
1226                assert_eq!(proxy.id, "2593294918");
1227            }
1228        }
1229
1230        #[tokio::test]
1231        async fn test_memorydb_get_blocked_proxies() {
1232            let db = memproxydb().await;
1233            let ctx = h2_proxy_context();
1234            let filter = ProxyFilter::default();
1235
1236            let mut blocked_proxies = vec![
1237                "1125300915",
1238                "1259341971",
1239                "1264821985",
1240                "129108927",
1241                "1316455915",
1242                "1425588737",
1243                "1571861931",
1244                "1810781137",
1245                "1836040682",
1246                "1844412609",
1247                "1885107293",
1248                "2021561518",
1249                "2079461709",
1250                "2107229589",
1251                "2141152822",
1252                "2438596154",
1253                "2497865606",
1254                "2521901221",
1255                "2551759475",
1256                "2560727338",
1257                "2593294918",
1258                "2798907087",
1259                "2854473221",
1260                "2880295577",
1261                "2909724448",
1262                "2912880381",
1263                "292096733",
1264                "2951529660",
1265                "3031533634",
1266                "3187902553",
1267                "3269411602",
1268                "3269465574",
1269                "339020035",
1270                "3481200027",
1271                "3498810974",
1272                "3503691556",
1273                "362091157",
1274                "3679054656",
1275                "371209663",
1276                "3861736957",
1277                "39048766",
1278                "3976711563",
1279                "4062553709",
1280                "49590203",
1281                "56402588",
1282                "724884866",
1283                "738626121",
1284                "767809962",
1285                "846528631",
1286                "906390012",
1287            ];
1288
1289            {
1290                let blocked_proxies = blocked_proxies.clone();
1291
1292                assert_eq!(
1293                    MemoryProxyDBQueryErrorKind::NotFound,
1294                    db.get_proxy_if(ctx.clone(), filter.clone(), move |proxy: &Proxy| {
1295                        !blocked_proxies.contains(&proxy.id.as_ref())
1296                    })
1297                    .await
1298                    .unwrap_err()
1299                    .kind()
1300                );
1301            }
1302
1303            let last_proxy_id = blocked_proxies.pop().unwrap();
1304
1305            let proxy = db
1306                .get_proxy_if(ctx, filter.clone(), move |proxy: &Proxy| {
1307                    !blocked_proxies.contains(&proxy.id.as_ref())
1308                })
1309                .await
1310                .unwrap();
1311            assert_eq!(proxy.id, last_proxy_id);
1312        }
1313
1314        #[tokio::test]
1315        async fn test_db_proxy_filter_any_use_filter_property() {
1316            let db = MemoryProxyDB::try_from_iter([Proxy {
1317                id: non_empty_str!("1"),
1318                address: ProxyAddress::from_str("example.com:80").unwrap(),
1319                tcp: true,
1320                udp: true,
1321                http: true,
1322                https: true,
1323                socks5: true,
1324                socks5h: true,
1325                datacenter: true,
1326                residential: true,
1327                mobile: true,
1328                pool_id: Some("*".into()),
1329                continent: Some("*".into()),
1330                country: Some("*".into()),
1331                state: Some("*".into()),
1332                city: Some("*".into()),
1333                carrier: Some("*".into()),
1334                asn: Some(Asn::unspecified()),
1335            }])
1336            .unwrap();
1337
1338            let ctx = h2_proxy_context();
1339
1340            for filter in [
1341                ProxyFilter {
1342                    id: Some(non_empty_str!("1")),
1343                    ..Default::default()
1344                },
1345                ProxyFilter {
1346                    pool_id: Some(vec![StringFilter::new("*")]),
1347                    ..Default::default()
1348                },
1349                ProxyFilter {
1350                    pool_id: Some(vec![StringFilter::new("hq")]),
1351                    ..Default::default()
1352                },
1353                ProxyFilter {
1354                    country: Some(vec![StringFilter::new("*")]),
1355                    ..Default::default()
1356                },
1357                ProxyFilter {
1358                    country: Some(vec![StringFilter::new("US")]),
1359                    ..Default::default()
1360                },
1361                ProxyFilter {
1362                    city: Some(vec![StringFilter::new("*")]),
1363                    ..Default::default()
1364                },
1365                ProxyFilter {
1366                    city: Some(vec![StringFilter::new("NY")]),
1367                    ..Default::default()
1368                },
1369                ProxyFilter {
1370                    carrier: Some(vec![StringFilter::new("*")]),
1371                    ..Default::default()
1372                },
1373                ProxyFilter {
1374                    carrier: Some(vec![StringFilter::new("Telenet")]),
1375                    ..Default::default()
1376                },
1377                ProxyFilter {
1378                    pool_id: Some(vec![StringFilter::new("hq")]),
1379                    country: Some(vec![StringFilter::new("US")]),
1380                    city: Some(vec![StringFilter::new("NY")]),
1381                    carrier: Some(vec![StringFilter::new("AT&T")]),
1382                    ..Default::default()
1383                },
1384            ] {
1385                let proxy = db.get_proxy(ctx.clone(), filter.clone()).await.unwrap();
1386                assert!(filter.id.map(|id| proxy.id == id).unwrap_or(true));
1387                assert!(
1388                    filter
1389                        .pool_id
1390                        .map(|pool_id| pool_id.contains(proxy.pool_id.as_ref().unwrap()))
1391                        .unwrap_or(true)
1392                );
1393                assert!(
1394                    filter
1395                        .country
1396                        .map(|country| country.contains(proxy.country.as_ref().unwrap()))
1397                        .unwrap_or(true)
1398                );
1399                assert!(
1400                    filter
1401                        .city
1402                        .map(|city| city.contains(proxy.city.as_ref().unwrap()))
1403                        .unwrap_or(true)
1404                );
1405                assert!(
1406                    filter
1407                        .carrier
1408                        .map(|carrier| carrier.contains(proxy.carrier.as_ref().unwrap()))
1409                        .unwrap_or(true)
1410                );
1411            }
1412        }
1413
1414        #[tokio::test]
1415        async fn test_db_proxy_filter_any_only_matches_any_value() {
1416            let db = MemoryProxyDB::try_from_iter([Proxy {
1417                id: non_empty_str!("1"),
1418                address: ProxyAddress::from_str("example.com:80").unwrap(),
1419                tcp: true,
1420                udp: true,
1421                http: true,
1422                https: true,
1423                socks5: true,
1424                socks5h: true,
1425                datacenter: true,
1426                residential: true,
1427                mobile: true,
1428                pool_id: Some("hq".into()),
1429                continent: Some("americas".into()),
1430                country: Some("US".into()),
1431                state: Some("NY".into()),
1432                city: Some("NY".into()),
1433                carrier: Some("AT&T".into()),
1434                asn: Some(Asn::from_static(7018)),
1435            }])
1436            .unwrap();
1437
1438            let ctx = h2_proxy_context();
1439
1440            for filter in [
1441                ProxyFilter {
1442                    pool_id: Some(vec![StringFilter::new("*")]),
1443                    ..Default::default()
1444                },
1445                ProxyFilter {
1446                    continent: Some(vec![StringFilter::new("*")]),
1447                    ..Default::default()
1448                },
1449                ProxyFilter {
1450                    country: Some(vec![StringFilter::new("*")]),
1451                    ..Default::default()
1452                },
1453                ProxyFilter {
1454                    state: Some(vec![StringFilter::new("*")]),
1455                    ..Default::default()
1456                },
1457                ProxyFilter {
1458                    city: Some(vec![StringFilter::new("*")]),
1459                    ..Default::default()
1460                },
1461                ProxyFilter {
1462                    carrier: Some(vec![StringFilter::new("*")]),
1463                    ..Default::default()
1464                },
1465                ProxyFilter {
1466                    asn: Some(vec![Asn::unspecified()]),
1467                    ..Default::default()
1468                },
1469                ProxyFilter {
1470                    pool_id: Some(vec![StringFilter::new("*")]),
1471                    continent: Some(vec![StringFilter::new("*")]),
1472                    country: Some(vec![StringFilter::new("*")]),
1473                    state: Some(vec![StringFilter::new("*")]),
1474                    city: Some(vec![StringFilter::new("*")]),
1475                    carrier: Some(vec![StringFilter::new("*")]),
1476                    asn: Some(vec![Asn::unspecified()]),
1477                    ..Default::default()
1478                },
1479            ] {
1480                let err = match db.get_proxy(ctx.clone(), filter.clone()).await {
1481                    Ok(proxy) => {
1482                        panic!("expected error for filter {filter:?}, not found proxy: {proxy:?}");
1483                    }
1484                    Err(err) => err,
1485                };
1486                assert_eq!(
1487                    MemoryProxyDBQueryErrorKind::NotFound,
1488                    err.kind(),
1489                    "filter: {filter:?}",
1490                );
1491            }
1492        }
1493
1494        #[tokio::test]
1495        async fn test_search_proxy_for_any_of_given_pools() {
1496            let db = MemoryProxyDB::try_from_iter([
1497                Proxy {
1498                    id: non_empty_str!("1"),
1499                    address: ProxyAddress::from_str("example.com:80").unwrap(),
1500                    tcp: true,
1501                    udp: true,
1502                    http: true,
1503                    https: true,
1504                    socks5: true,
1505                    socks5h: true,
1506                    datacenter: true,
1507                    residential: true,
1508                    mobile: true,
1509                    pool_id: Some("a".into()),
1510                    continent: Some("americas".into()),
1511                    country: Some("US".into()),
1512                    state: Some("NY".into()),
1513                    city: Some("NY".into()),
1514                    carrier: Some("AT&T".into()),
1515                    asn: Some(Asn::from_static(7018)),
1516                },
1517                Proxy {
1518                    id: non_empty_str!("2"),
1519                    address: ProxyAddress::from_str("example.com:80").unwrap(),
1520                    tcp: true,
1521                    udp: true,
1522                    http: true,
1523                    https: true,
1524                    socks5: true,
1525                    socks5h: true,
1526                    datacenter: true,
1527                    residential: true,
1528                    mobile: true,
1529                    pool_id: Some("b".into()),
1530                    continent: Some("americas".into()),
1531                    country: Some("US".into()),
1532                    state: Some("NY".into()),
1533                    city: Some("NY".into()),
1534                    carrier: Some("AT&T".into()),
1535                    asn: Some(Asn::from_static(7018)),
1536                },
1537                Proxy {
1538                    id: non_empty_str!("3"),
1539                    address: ProxyAddress::from_str("example.com:80").unwrap(),
1540                    tcp: true,
1541                    udp: true,
1542                    http: true,
1543                    https: true,
1544                    socks5: true,
1545                    socks5h: true,
1546                    datacenter: true,
1547                    residential: true,
1548                    mobile: true,
1549                    pool_id: Some("b".into()),
1550                    continent: Some("americas".into()),
1551                    country: Some("US".into()),
1552                    state: Some("NY".into()),
1553                    city: Some("NY".into()),
1554                    carrier: Some("AT&T".into()),
1555                    asn: Some(Asn::from_static(7018)),
1556                },
1557                Proxy {
1558                    id: non_empty_str!("4"),
1559                    address: ProxyAddress::from_str("example.com:80").unwrap(),
1560                    tcp: true,
1561                    udp: true,
1562                    http: true,
1563                    https: true,
1564                    socks5: true,
1565                    socks5h: true,
1566                    datacenter: true,
1567                    residential: true,
1568                    mobile: true,
1569                    pool_id: Some("c".into()),
1570                    continent: Some("americas".into()),
1571                    country: Some("US".into()),
1572                    state: Some("NY".into()),
1573                    city: Some("NY".into()),
1574                    carrier: Some("AT&T".into()),
1575                    asn: Some(Asn::from_static(7018)),
1576                },
1577            ])
1578            .unwrap();
1579
1580            let ctx = h2_proxy_context();
1581
1582            let filter = ProxyFilter {
1583                pool_id: Some(vec![StringFilter::new("a"), StringFilter::new("c")]),
1584                ..Default::default()
1585            };
1586
1587            let mut seen_1 = false;
1588            let mut seen_4 = false;
1589            for _ in 0..100 {
1590                let proxy = db.get_proxy(ctx.clone(), filter.clone()).await.unwrap();
1591                match proxy.id.as_ref() {
1592                    "1" => seen_1 = true,
1593                    "4" => seen_4 = true,
1594                    _ => panic!("unexpected pool id"),
1595                }
1596            }
1597            assert!(seen_1);
1598            assert!(seen_4);
1599        }
1600
1601        #[tokio::test]
1602        async fn test_deserialize_url_proxy_filter() {
1603            for (input, expected_output) in [
1604                (
1605                    "id=1",
1606                    ProxyFilter {
1607                        id: Some(non_empty_str!("1")),
1608                        ..Default::default()
1609                    },
1610                ),
1611                (
1612                    "pool=hq&country=us",
1613                    ProxyFilter {
1614                        pool_id: Some(vec![StringFilter::new("hq")]),
1615                        country: Some(vec![StringFilter::new("us")]),
1616                        ..Default::default()
1617                    },
1618                ),
1619                (
1620                    "pool=hq&country=us&country=be",
1621                    ProxyFilter {
1622                        pool_id: Some(vec![StringFilter::new("hq")]),
1623                        country: Some(vec![StringFilter::new("us"), StringFilter::new("be")]),
1624                        ..Default::default()
1625                    },
1626                ),
1627                (
1628                    "pool=a&country=uk&pool=b",
1629                    ProxyFilter {
1630                        pool_id: Some(vec![StringFilter::new("a"), StringFilter::new("b")]),
1631                        country: Some(vec![StringFilter::new("uk")]),
1632                        ..Default::default()
1633                    },
1634                ),
1635                (
1636                    "continent=europe&continent=asia",
1637                    ProxyFilter {
1638                        continent: Some(vec![
1639                            StringFilter::new("europe"),
1640                            StringFilter::new("asia"),
1641                        ]),
1642                        ..Default::default()
1643                    },
1644                ),
1645                (
1646                    "continent=americas&country=us&state=NY&city=buffalo&carrier=AT%26T&asn=7018",
1647                    ProxyFilter {
1648                        continent: Some(vec![StringFilter::new("americas")]),
1649                        country: Some(vec![StringFilter::new("us")]),
1650                        state: Some(vec![StringFilter::new("ny")]),
1651                        city: Some(vec![StringFilter::new("buffalo")]),
1652                        carrier: Some(vec![StringFilter::new("at&t")]),
1653                        asn: Some(vec![Asn::from_static(7018)]),
1654                        ..Default::default()
1655                    },
1656                ),
1657                (
1658                    "asn=1&asn=2",
1659                    ProxyFilter {
1660                        asn: Some(vec![Asn::from_static(1), Asn::from_static(2)]),
1661                        ..Default::default()
1662                    },
1663                ),
1664            ] {
1665                let filter: ProxyFilter = serde_html_form::from_str(input).unwrap();
1666                assert_eq!(filter, expected_output);
1667            }
1668        }
1669    }
1670}
1671
1672#[cfg(feature = "memory-db")]
1673#[cfg_attr(docsrs, doc(cfg(feature = "memory-db")))]
1674pub use memdb::{
1675    MemoryProxyDB, MemoryProxyDBInsertError, MemoryProxyDBInsertErrorKind, MemoryProxyDBQueryError,
1676    MemoryProxyDBQueryErrorKind,
1677};