Skip to main content

rama_proxy/proxydb/
csv.rs

1use super::{Proxy, StringFilter};
2use base64::Engine;
3use base64::engine::general_purpose::STANDARD as ENGINE;
4use rama_net::{
5    address::ProxyAddress,
6    asn::{Asn, InvalidAsn},
7    user::ProxyCredential,
8};
9use std::path::Path;
10use tokio::{
11    fs::File,
12    io::{AsyncBufReadExt, BufReader, Lines},
13};
14
15#[derive(Debug)]
16/// A CSV Reader that can be used to create a [`Proxy`] database from a CSV file or raw data.
17pub struct ProxyCsvRowReader {
18    data: ProxyCsvRowReaderData,
19}
20
21impl ProxyCsvRowReader {
22    /// Create a new [`ProxyCsvRowReader`] from the given CSV file.
23    pub async fn open(path: impl AsRef<Path>) -> Result<Self, ProxyCsvRowReaderError> {
24        let file = rama_utils::fs::safe_open(path).await?;
25        let reader = BufReader::new(file);
26        let lines = reader.lines();
27        Ok(Self {
28            data: ProxyCsvRowReaderData::File(lines),
29        })
30    }
31
32    /// Create a new [`ProxyCsvRowReader`] from the given CSV data.
33    pub fn raw(data: impl AsRef<str>) -> Self {
34        let lines: Vec<_> = data.as_ref().lines().rev().map(str::to_owned).collect();
35        Self {
36            data: ProxyCsvRowReaderData::Raw(lines),
37        }
38    }
39
40    /// Read the next row from the CSV file.
41    pub async fn next(&mut self) -> Result<Option<Proxy>, ProxyCsvRowReaderError> {
42        match &mut self.data {
43            ProxyCsvRowReaderData::File(lines) => {
44                let line = lines.next_line().await?;
45                match line {
46                    Some(line) => Ok(Some(match parse_csv_row(&line) {
47                        Some(proxy) => proxy,
48                        None => {
49                            return Err(ProxyCsvRowReaderError {
50                                kind: ProxyCsvRowReaderErrorKind::InvalidRow(line),
51                            });
52                        }
53                    })),
54                    None => Ok(None),
55                }
56            }
57            ProxyCsvRowReaderData::Raw(lines) => match lines.pop() {
58                Some(line) => Ok(Some(match parse_csv_row(&line) {
59                    Some(proxy) => proxy,
60                    None => {
61                        return Err(ProxyCsvRowReaderError {
62                            kind: ProxyCsvRowReaderErrorKind::InvalidRow(line),
63                        });
64                    }
65                })),
66                None => Ok(None),
67            },
68        }
69    }
70}
71
72fn strip_csv_quotes(p: &str) -> &str {
73    p.strip_prefix('"')
74        .and_then(|p| p.strip_suffix('"'))
75        .unwrap_or(p)
76}
77
78pub(crate) fn parse_csv_row(row: &str) -> Option<Proxy> {
79    let mut iter = row.split(',').map(strip_csv_quotes);
80
81    let id = iter.next().and_then(|s| s.try_into().ok())?;
82
83    let tcp = iter.next().and_then(parse_csv_bool)?;
84    let udp = iter.next().and_then(parse_csv_bool)?;
85    let http = iter.next().and_then(parse_csv_bool)?;
86    let https = iter.next().and_then(parse_csv_bool)?;
87    let socks5 = iter.next().and_then(parse_csv_bool)?;
88    let socks5h = iter.next().and_then(parse_csv_bool)?;
89    let datacenter = iter.next().and_then(parse_csv_bool)?;
90    let residential = iter.next().and_then(parse_csv_bool)?;
91    let mobile = iter.next().and_then(parse_csv_bool)?;
92    let mut address = iter.next().and_then(|s| {
93        if s.is_empty() {
94            None
95        } else {
96            ProxyAddress::try_from(s).ok()
97        }
98    })?;
99    let pool_id = parse_csv_opt_string_filter(iter.next()?);
100    let continent = parse_csv_opt_string_filter(iter.next()?);
101    let country = parse_csv_opt_string_filter(iter.next()?);
102    let state = parse_csv_opt_string_filter(iter.next()?);
103    let city = parse_csv_opt_string_filter(iter.next()?);
104    let carrier = parse_csv_opt_string_filter(iter.next()?);
105    let asn = parse_csv_opt_asn(iter.next()?).ok()?;
106
107    // support header format or cleartext format
108    if let Some(value) = iter.next()
109        && !value.is_empty()
110    {
111        address.credential = Some(match value.split_once(' ') {
112            Some((t, v)) => {
113                if t.eq_ignore_ascii_case("basic") {
114                    let bytes = ENGINE.decode(v).ok()?;
115                    let decoded = String::from_utf8(bytes).ok()?;
116                    ProxyCredential::Basic(decoded.parse().ok()?)
117                } else if t.eq_ignore_ascii_case("bearer") {
118                    ProxyCredential::Bearer(v.parse().ok()?)
119                } else {
120                    ProxyCredential::Basic(value.parse().ok()?)
121                }
122            }
123            None => ProxyCredential::Basic(value.parse().ok()?),
124        });
125    }
126
127    // Ensure there are no more values in the row
128    if iter.next().is_some() {
129        return None;
130    }
131
132    Some(Proxy {
133        id,
134        address,
135        tcp,
136        udp,
137        http,
138        https,
139        socks5,
140        socks5h,
141        datacenter,
142        residential,
143        mobile,
144        pool_id,
145        continent,
146        country,
147        state,
148        city,
149        carrier,
150        asn,
151    })
152}
153
154fn parse_csv_bool(value: &str) -> Option<bool> {
155    rama_utils::macros::match_ignore_ascii_case_str! {
156        match(value) {
157            "true" | "1" => Some(true),
158            "" | "false" | "0" | "null" | "nil" => Some(false),
159            _ => None,
160        }
161    }
162}
163
164fn parse_csv_opt_string_filter(value: &str) -> Option<StringFilter> {
165    if value.is_empty() {
166        None
167    } else {
168        Some(StringFilter::from(value))
169    }
170}
171
172fn parse_csv_opt_asn(value: &str) -> Result<Option<Asn>, InvalidAsn> {
173    if value.is_empty() {
174        Ok(None)
175    } else {
176        Asn::try_from(value).map(Some)
177    }
178}
179
180#[derive(Debug)]
181enum ProxyCsvRowReaderData {
182    File(Lines<BufReader<File>>),
183    Raw(Vec<String>),
184}
185
186#[derive(Debug)]
187/// An error that can occur when reading a Proxy CSV row.
188pub struct ProxyCsvRowReaderError {
189    kind: ProxyCsvRowReaderErrorKind,
190}
191
192#[derive(Debug)]
193/// The kind of error that can occur when reading a Proxy CSV row.
194pub enum ProxyCsvRowReaderErrorKind {
195    /// An I/O error occurred while reading the CSV row.
196    IoError(std::io::Error),
197    /// The CSV row is invalid, and could not be parsed.
198    InvalidRow(String),
199}
200
201impl std::fmt::Display for ProxyCsvRowReaderError {
202    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
203        match &self.kind {
204            ProxyCsvRowReaderErrorKind::IoError(err) => write!(f, "I/O error: {err}"),
205            ProxyCsvRowReaderErrorKind::InvalidRow(row) => write!(f, "Invalid row: {row}"),
206        }
207    }
208}
209
210impl std::error::Error for ProxyCsvRowReaderError {
211    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
212        match &self.kind {
213            ProxyCsvRowReaderErrorKind::IoError(err) => Some(err),
214            ProxyCsvRowReaderErrorKind::InvalidRow(_) => None,
215        }
216    }
217}
218
219impl From<std::io::Error> for ProxyCsvRowReaderError {
220    fn from(err: std::io::Error) -> Self {
221        Self {
222            kind: ProxyCsvRowReaderErrorKind::IoError(err),
223        }
224    }
225}
226
227#[cfg(test)]
228mod tests {
229    use super::*;
230    use crate::{ProxyFilter, proxydb::ProxyContext};
231    use rama_net::transport::TransportProtocol;
232    use rama_utils::str::non_empty_str;
233    use std::str::FromStr;
234
235    #[test]
236    fn test_parse_csv_bool() {
237        for (input, output) in &[
238            ("1", Some(true)),
239            ("true", Some(true)),
240            ("True", Some(true)),
241            ("TRUE", Some(true)),
242            ("0", Some(false)),
243            ("false", Some(false)),
244            ("False", Some(false)),
245            ("FALSE", Some(false)),
246            ("null", Some(false)),
247            ("nil", Some(false)),
248            ("NULL", Some(false)),
249            ("NIL", Some(false)),
250            ("", Some(false)),
251            ("invalid", None),
252        ] {
253            assert_eq!(parse_csv_bool(input), *output);
254        }
255    }
256
257    #[test]
258    fn test_parse_csv_opt_string_filter() {
259        for (input, output) in [
260            ("", None),
261            ("value", Some("value")),
262            ("*", Some("*")),
263            ("Foo", Some("foo")),
264            ("  ok ", Some("ok")),
265            (" NO  ", Some("no")),
266        ] {
267            assert_eq!(
268                parse_csv_opt_string_filter(input)
269                    .as_ref()
270                    .map(|f| f.as_ref()),
271                output,
272            );
273        }
274    }
275
276    #[test]
277    fn test_parse_csv_opt_string_filter_is_any() {
278        let filter = parse_csv_opt_string_filter("*").unwrap();
279        assert!(venndb::Any::is_any(&filter));
280    }
281
282    #[test]
283    fn test_parse_csv_row_happy_path() {
284        for (input, output) in [
285            // most minimal
286            (
287                "id,,,,,,,,,,authority:80,,,,,,,,",
288                Proxy {
289                    id: non_empty_str!("id"),
290                    address: ProxyAddress::from_str("authority:80").unwrap(),
291                    tcp: false,
292                    udp: false,
293                    http: false,
294                    https: false,
295                    socks5: false,
296                    socks5h: false,
297                    datacenter: false,
298                    residential: false,
299                    mobile: false,
300                    pool_id: None,
301                    continent: None,
302                    country: None,
303                    state: None,
304                    city: None,
305                    carrier: None,
306                    asn: None,
307                },
308            ),
309            // more happy row tests
310            (
311                "id,true,false,true,,false,,true,false,true,authority:80,pool_id,,country,,city,carrier,,Basic dXNlcm5hbWU6cGFzc3dvcmQ=",
312                Proxy {
313                    id: non_empty_str!("id"),
314                    address: ProxyAddress::from_str("username:password@authority:80").unwrap(),
315                    tcp: true,
316                    udp: false,
317                    http: true,
318                    https: false,
319                    socks5: false,
320                    socks5h: false,
321                    datacenter: true,
322                    residential: false,
323                    mobile: true,
324                    pool_id: Some("pool_id".into()),
325                    continent: None,
326                    country: Some("country".into()),
327                    state: None,
328                    city: Some("city".into()),
329                    carrier: Some("carrier".into()),
330                    asn: None,
331                },
332            ),
333            (
334                "123,1,0,False,,True,,null,false,true,host:1234,,americas,*,*,*,carrier,13335,",
335                Proxy {
336                    id: non_empty_str!("123"),
337                    address: ProxyAddress::from_str("host:1234").unwrap(),
338                    tcp: true,
339                    udp: false,
340                    http: false,
341                    https: false,
342                    socks5: true,
343                    socks5h: false,
344                    datacenter: false,
345                    residential: false,
346                    mobile: true,
347                    pool_id: None,
348                    continent: Some("americas".into()),
349                    country: Some("*".into()),
350                    state: Some("*".into()),
351                    city: Some("*".into()),
352                    carrier: Some("carrier".into()),
353                    asn: Some(Asn::from_static(13335)),
354                },
355            ),
356            (
357                "123,1,0,False,,True,,null,false,true,host:1234,,europe,*,,*,carrier,0",
358                Proxy {
359                    id: non_empty_str!("123"),
360                    address: ProxyAddress::from_str("host:1234").unwrap(),
361                    tcp: true,
362                    udp: false,
363                    http: false,
364                    https: false,
365                    socks5: true,
366                    socks5h: false,
367                    datacenter: false,
368                    residential: false,
369                    mobile: true,
370                    pool_id: None,
371                    continent: Some("europe".into()),
372                    country: Some("*".into()),
373                    state: None,
374                    city: Some("*".into()),
375                    carrier: Some("carrier".into()),
376                    asn: Some(Asn::unspecified()),
377                },
378            ),
379            (
380                "foo,1,0,1,,0,,1,0,0,http://bar,baz,,US,,,,",
381                Proxy {
382                    id: non_empty_str!("foo"),
383                    address: ProxyAddress::from_str("http://bar").unwrap(),
384                    tcp: true,
385                    udp: false,
386                    http: true,
387                    https: false,
388                    socks5: false,
389                    socks5h: false,
390                    datacenter: true,
391                    residential: false,
392                    mobile: false,
393                    pool_id: Some("baz".into()),
394                    continent: None,
395                    country: Some("us".into()),
396                    state: None,
397                    city: None,
398                    carrier: None,
399                    asn: None,
400                },
401            ),
402        ] {
403            let proxy = parse_csv_row(input).unwrap();
404            assert_eq!(proxy.id, output.id);
405            assert_eq!(proxy.address, output.address);
406            assert_eq!(proxy.tcp, output.tcp);
407            assert_eq!(proxy.udp, output.udp);
408            assert_eq!(proxy.http, output.http);
409            assert_eq!(proxy.socks5, output.socks5);
410            assert_eq!(proxy.datacenter, output.datacenter);
411            assert_eq!(proxy.residential, output.residential);
412            assert_eq!(proxy.mobile, output.mobile);
413            assert_eq!(proxy.pool_id, output.pool_id);
414            assert_eq!(proxy.continent, output.continent);
415            assert_eq!(proxy.country, output.country);
416            assert_eq!(proxy.state, output.state);
417            assert_eq!(proxy.city, output.city);
418            assert_eq!(proxy.carrier, output.carrier);
419            assert_eq!(proxy.asn, output.asn);
420        }
421    }
422
423    #[test]
424    fn test_parse_csv_row_mistakes() {
425        for input in [
426            // garbage rows
427            "",
428            ",",
429            ",,,,,,",
430            ",,,,,,,,,,,,,,,,,,,,",
431            ",,,,,,,,,,,,,,,,,,,,,,",
432            ",,,,,,,,,,,,,,,,,,,,,,,",
433            // too many rows
434            "id,true,false,true,false,true,false,true,authority:80,pool_id,continent,country,state,city,carrier,15169,Basic dXNlcm5hbWU6cGFzc3dvcmQ=,",
435            // missing authority
436            "id,,,,,,,,,,,,,,,,",
437            // missing proxy id
438            ",,,,,,,,authority:80,,,,,,,,",
439            // invalid bool values
440            "id,foo,,,,,,,,,authority:80,,,,,,,,",
441            "id,,foo,,,,,,,,authority:80,,,,,,,,",
442            "id,,,foo,,,,,,,authority:80,,,,,,,,",
443            "id,,,,,foo,,,,,authority:80,,,,,,,,",
444            "id,,,,,,foo,,,,authority:80,,,,,,,,",
445            "id,,,,,,,,foo,,authority:80,,,,,,,,",
446            "id,,,,,,,foo,authority:80,,,,,,,,",
447            // invalid credentials
448            "id,,,,,,,,authority:80,,,,,:foo",
449        ] {
450            assert!(parse_csv_row(input).is_none(), "input: {input}");
451        }
452    }
453
454    #[tokio::test]
455    async fn test_proxy_csv_row_reader_happy_one_row() {
456        let mut reader = ProxyCsvRowReader::raw(
457            "id,true,false,true,,false,,true,false,true,authority:80,pool_id,continent,country,state,city,carrier,13335,Basic dXNlcm5hbWU6cGFzc3dvcmQ=",
458        );
459        let proxy = reader.next().await.unwrap().unwrap();
460
461        assert_eq!(proxy.id, "id");
462        assert_eq!(
463            proxy.address,
464            ProxyAddress::from_str("username:password@authority:80").unwrap()
465        );
466        assert!(proxy.tcp);
467        assert!(!proxy.udp);
468        assert!(proxy.http);
469        assert!(!proxy.socks5);
470        assert!(proxy.datacenter);
471        assert!(!proxy.residential);
472        assert!(proxy.mobile);
473        assert_eq!(proxy.pool_id, Some("pool_id".into()));
474        assert_eq!(proxy.continent, Some("continent".into()));
475        assert_eq!(proxy.country, Some("country".into()));
476        assert_eq!(proxy.state, Some("state".into()));
477        assert_eq!(proxy.city, Some("city".into()));
478        assert_eq!(proxy.carrier, Some("carrier".into()));
479        assert_eq!(proxy.asn, Some(Asn::from_static(13335)));
480
481        // no more rows to read
482        assert!(reader.next().await.unwrap().is_none());
483    }
484
485    #[tokio::test]
486    async fn test_proxy_csv_row_reader_happy_multi_row() {
487        let mut reader = ProxyCsvRowReader::raw(
488            "id,true,false,false,true,true,false,true,false,true,authority:80,pool_id,continent,country,state,city,carrier,42,Basic dXNlcm5hbWU6cGFzc3dvcmQ=\nid2,1,0,0,0,0,0,1,0,0,authority2:80,pool_id2,continent2,country2,state2,city2,carrier2,1",
489        );
490
491        let proxy = reader.next().await.unwrap().unwrap();
492        assert_eq!(proxy.id, "id");
493        assert_eq!(
494            proxy.address,
495            ProxyAddress::from_str("username:password@authority:80").unwrap()
496        );
497        assert!(proxy.tcp);
498        assert!(!proxy.udp);
499        assert!(!proxy.http);
500        assert!(proxy.https);
501        assert!(proxy.socks5);
502        assert!(!proxy.socks5h);
503        assert!(proxy.datacenter);
504        assert!(!proxy.residential);
505        assert!(proxy.mobile);
506        assert_eq!(proxy.pool_id, Some("pool_id".into()));
507        assert_eq!(proxy.continent, Some("continent".into()));
508        assert_eq!(proxy.country, Some("country".into()));
509        assert_eq!(proxy.state, Some("state".into()));
510        assert_eq!(proxy.city, Some("city".into()));
511        assert_eq!(proxy.carrier, Some("carrier".into()));
512        assert_eq!(proxy.asn, Some(Asn::from_static(42)));
513
514        let proxy = reader.next().await.unwrap().unwrap();
515
516        assert_eq!(proxy.id, "id2");
517        assert_eq!(
518            proxy.address,
519            ProxyAddress::from_str("authority2:80").unwrap()
520        );
521        assert!(proxy.tcp);
522        assert!(!proxy.udp);
523        assert!(!proxy.http);
524        assert!(!proxy.https);
525        assert!(!proxy.socks5);
526        assert!(!proxy.socks5h);
527        assert!(proxy.datacenter);
528        assert!(!proxy.residential);
529        assert!(!proxy.mobile);
530        assert_eq!(proxy.pool_id, Some("pool_id2".into()));
531        assert_eq!(proxy.continent, Some("continent2".into()));
532        assert_eq!(proxy.country, Some("country2".into()));
533        assert_eq!(proxy.city, Some("city2".into()));
534        assert_eq!(proxy.state, Some("state2".into()));
535        assert_eq!(proxy.carrier, Some("carrier2".into()));
536        assert_eq!(proxy.asn, Some(Asn::from_static(1)));
537
538        // no more rows to read
539        assert!(reader.next().await.unwrap().is_none());
540    }
541
542    #[tokio::test]
543    async fn test_proxy_csv_row_reader_failure_empty_data() {
544        let mut reader = ProxyCsvRowReader::raw("");
545        assert!(reader.next().await.unwrap().is_none());
546    }
547
548    #[tokio::test]
549    async fn test_proxy_csv_row_reader_failure_invalid_row() {
550        let mut reader = ProxyCsvRowReader::raw(",,,,,,,,,,,");
551        reader.next().await.unwrap_err();
552    }
553
554    #[test]
555    fn test_proxy_is_match_happy_path_proxy_with_any_filter_string_cases() {
556        let proxy = parse_csv_row("id,1,,1,,,,,,,authority:80,*,*,*,*,*,*,0").unwrap();
557        let ctx = ProxyContext {
558            protocol: TransportProtocol::Tcp,
559        };
560
561        for filter in [
562            ProxyFilter::default(),
563            ProxyFilter {
564                pool_id: Some(vec![StringFilter::new("pool_a")]),
565                country: Some(vec![StringFilter::new("country_a")]),
566                city: Some(vec![StringFilter::new("city_a")]),
567                carrier: Some(vec![StringFilter::new("carrier_a")]),
568                ..Default::default()
569            },
570            ProxyFilter {
571                pool_id: Some(vec![StringFilter::new("pool_a")]),
572                ..Default::default()
573            },
574            ProxyFilter {
575                continent: Some(vec![StringFilter::new("continent_a")]),
576                ..Default::default()
577            },
578            ProxyFilter {
579                country: Some(vec![StringFilter::new("country_a")]),
580                ..Default::default()
581            },
582            ProxyFilter {
583                state: Some(vec![StringFilter::new("state_a")]),
584                ..Default::default()
585            },
586            ProxyFilter {
587                city: Some(vec![StringFilter::new("city_a")]),
588                carrier: Some(vec![StringFilter::new("carrier_a")]),
589                ..Default::default()
590            },
591            ProxyFilter {
592                carrier: Some(vec![StringFilter::new("carrier_a")]),
593                ..Default::default()
594            },
595        ] {
596            assert!(proxy.is_match(&ctx, &filter), "filter: {filter:?}");
597        }
598    }
599
600    #[test]
601    fn test_proxy_is_match_happy_path_proxy_with_any_filters_cases() {
602        let proxy = parse_csv_row(
603            "id,1,,1,,,,,,,authority:80,pool,continent,country,state,city,carrier,42",
604        )
605        .unwrap();
606        let ctx = ProxyContext {
607            protocol: TransportProtocol::Tcp,
608        };
609
610        for filter in [
611            ProxyFilter::default(),
612            ProxyFilter {
613                pool_id: Some(vec![StringFilter::new("*")]),
614                ..Default::default()
615            },
616            ProxyFilter {
617                continent: Some(vec![StringFilter::new("*")]),
618                ..Default::default()
619            },
620            ProxyFilter {
621                country: Some(vec![StringFilter::new("*")]),
622                ..Default::default()
623            },
624            ProxyFilter {
625                state: Some(vec![StringFilter::new("*")]),
626                ..Default::default()
627            },
628            ProxyFilter {
629                city: Some(vec![StringFilter::new("*")]),
630                ..Default::default()
631            },
632            ProxyFilter {
633                carrier: Some(vec![StringFilter::new("*")]),
634                ..Default::default()
635            },
636            ProxyFilter {
637                pool_id: Some(vec![StringFilter::new("pool")]),
638                continent: Some(vec![StringFilter::new("continent")]),
639                country: Some(vec![StringFilter::new("country")]),
640                state: Some(vec![StringFilter::new("state")]),
641                city: Some(vec![StringFilter::new("city")]),
642                carrier: Some(vec![StringFilter::new("carrier")]),
643                asn: Some(vec![Asn::from_static(42)]),
644                ..Default::default()
645            },
646            ProxyFilter {
647                pool_id: Some(vec![StringFilter::new("*")]),
648                country: Some(vec![StringFilter::new("country")]),
649                city: Some(vec![StringFilter::new("city")]),
650                carrier: Some(vec![StringFilter::new("carrier")]),
651                ..Default::default()
652            },
653            ProxyFilter {
654                pool_id: Some(vec![StringFilter::new("pool")]),
655                country: Some(vec![StringFilter::new("*")]),
656                city: Some(vec![StringFilter::new("city")]),
657                carrier: Some(vec![StringFilter::new("carrier")]),
658                ..Default::default()
659            },
660            ProxyFilter {
661                pool_id: Some(vec![StringFilter::new("pool")]),
662                country: Some(vec![StringFilter::new("country")]),
663                city: Some(vec![StringFilter::new("*")]),
664                carrier: Some(vec![StringFilter::new("carrier")]),
665                ..Default::default()
666            },
667            ProxyFilter {
668                pool_id: Some(vec![StringFilter::new("pool")]),
669                country: Some(vec![StringFilter::new("country")]),
670                city: Some(vec![StringFilter::new("city")]),
671                carrier: Some(vec![StringFilter::new("*")]),
672                ..Default::default()
673            },
674            ProxyFilter {
675                continent: Some(vec![StringFilter::new("*")]),
676                ..Default::default()
677            },
678        ] {
679            assert!(proxy.is_match(&ctx, &filter), "filter: {filter:?}");
680        }
681    }
682}