geo302 0.1.3

A simple geoIP-based redirect proxy
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

use hyper::HeaderMap;
use std::net::IpAddr;

pub fn client_ip(headers: &HeaderMap, header_names: &[String], recursive: bool) -> Option<IpAddr> {
    header_names
        .iter()
        .filter_map(|name| {
            let values = headers.get_all(name);
            let mut it_values = values.iter();
            if recursive {
                it_values.next()
            } else {
                it_values.next_back()
            }
        })
        .next()
        .and_then(|value| {
            let value = value.to_str().ok()?;
            let mut split = value.split(',');
            if recursive {
                split.next()
            } else {
                split.next_back()
            }
        })
        .and_then(|s| s.trim().parse::<IpAddr>().ok())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::net::{IpAddr, Ipv4Addr};

    #[test]
    fn client_ip_filter_no_addr() {
        for is_recursive in [false, true] {
            let actual = client_ip(&HeaderMap::new(), &[], is_recursive);
            assert_eq!(actual, None);
        }
    }

    #[test]
    fn client_ip_filter_single_header() {
        let ip_expected = IpAddr::V4(Ipv4Addr::new(128, 174, 199, 60));
        let header_name = "X-FORWARDED-FOR";
        let request_headers = {
            let mut headers = HeaderMap::new();
            headers.insert(header_name, ip_expected.to_string().parse().unwrap());
            headers
        };

        for is_recursive in [false, true] {
            let actual = client_ip(&request_headers, &[header_name.to_string()], is_recursive);
            assert_eq!(actual, Some(ip_expected));
        }
    }

    #[test]
    fn client_ip_filter_one_header_two_values() {
        let ip_client = IpAddr::V4(Ipv4Addr::new(128, 174, 199, 60));
        let ip_proxy = IpAddr::V4(Ipv4Addr::new(80, 94, 184, 70));
        let ip_header_name = "X-FORWARDED-FOR";
        let request_header_value = format!("{}, {}", ip_client, ip_proxy);
        let request_headers = {
            let mut headers = HeaderMap::new();
            headers.insert(ip_header_name, request_header_value.parse().unwrap());
            headers
        };

        for (ip_expected, is_recursive) in [(ip_proxy, false), (ip_client, true)] {
            let actual = client_ip(
                &request_headers,
                &[ip_header_name.to_string()],
                is_recursive,
            );
            assert_eq!(actual, Some(ip_expected));
        }
    }

    #[tokio::test]
    async fn client_ip_filter_one_header_three_values() {
        let ip_client = IpAddr::V4(Ipv4Addr::new(128, 174, 199, 60));
        let ip_proxy1 = IpAddr::V4(Ipv4Addr::new(80, 94, 184, 70));
        let ip_proxy2 = IpAddr::V4(Ipv4Addr::new(52, 0, 14, 116));
        let ip_header_name = "X-FORWARDED-FOR";
        let request_header_value = format!("{}, {}, {}", ip_client, ip_proxy1, ip_proxy2);
        let request_headers = {
            let mut headers = HeaderMap::new();
            headers.insert(ip_header_name, request_header_value.parse().unwrap());
            headers
        };

        for (ip_expected, is_recursive) in [(ip_proxy2, false), (ip_client, true)] {
            let actual = client_ip(
                &request_headers,
                &[ip_header_name.to_string()],
                is_recursive,
            );
            assert_eq!(actual, Some(ip_expected));
        }
    }

    #[test]
    fn client_ip_filter_two_headers() {
        let ip_client = IpAddr::V4(Ipv4Addr::new(128, 174, 199, 60));
        let ip_nginx_proxy = IpAddr::V4(Ipv4Addr::new(80, 94, 184, 70));
        let ip_other_proxy = IpAddr::V4(Ipv4Addr::new(80, 94, 184, 70));
        let request_headers = {
            let mut headers = HeaderMap::new();
            headers.insert(
                "X-REAL-IP",
                format!("{}, {}", ip_client, ip_nginx_proxy)
                    .parse()
                    .unwrap(),
            );
            headers.insert(
                "X-FORWARDED-FOR",
                format!("{}, {}", ip_nginx_proxy, ip_other_proxy)
                    .parse()
                    .unwrap(),
            );
            headers
        };
        let header_names: Vec<_> = request_headers
            .keys()
            .map(|name| name.as_str().to_string())
            .collect();

        for (ip_expected, is_recursive) in [(ip_nginx_proxy, false), (ip_client, true)] {
            let actual = client_ip(&request_headers, &header_names, is_recursive);
            assert_eq!(actual, Some(ip_expected));
        }
    }
}