geo302 0.1.2

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

use crate::config::parse_config;
use crate::geo::{Continent, Geo};
use crate::healthcheck::check_health;
use crate::mirror::{ContinentMap, Mirror, MirrorVec};
use crate::rejects::{handle_rejection, BrokenPath, MirrorsUnavailable};
use std::net::{IpAddr, SocketAddr};
use std::str::FromStr;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::Duration;
use thiserror::Error;
use warp::http::header::{
    HeaderMap, HeaderName, HeaderValue, InvalidHeaderName, InvalidHeaderValue,
};
use warp::http::Uri;
use warp::path::FullPath;
use warp::reply::Reply;
use warp::Filter;

mod config;
mod filters;
mod geo;
mod healthcheck;
mod mirror;
mod rejects;

#[derive(Error, Debug)]
pub enum HeaderError {
    #[error("{0}")]
    InvalidHeaderName(InvalidHeaderName),
    #[error("{0}")]
    InvalidHeaderValue(InvalidHeaderValue),
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let config_path = std::env::args()
        .nth(1)
        .unwrap_or_else(|| "geo302.toml".to_owned());

    let config = parse_config(config_path)?;
    let host: SocketAddr = config.host.parse()?;
    let ip_header_recursive = config.ip_headers_recursive;
    let continent_map = ContinentMap::from_config(&config)?;
    let geo = Arc::new(Geo::from_config(&config)?);
    let check_interval = Duration::new(config.healthckeck_interval.get().into(), 0);
    let ip_header_names = config.ip_headers;
    let response_headers = config
        .response_headers
        .iter()
        .map(|(name, value)| {
            let name = match HeaderName::from_str(name) {
                Ok(name) => name,
                Err(err) => return Err(HeaderError::InvalidHeaderName(err)),
            };
            let value = match HeaderValue::from_str(value) {
                Ok(value) => value,
                Err(err) => return Err(HeaderError::InvalidHeaderValue(err)),
            };
            Ok((name, value))
        })
        .collect::<Result<HeaderMap, _>>()?;

    simple_logger::init_with_level(config.log_level)?;

    check_health(continent_map.all_mirrors(), check_interval);

    let logs = warp::log::custom(|info| {
        log::info!(
            "{} {} {} {}",
            info.remote_addr()
                .map_or_else(|| "_".into(), |addr| format!("{}", addr.ip())),
            info.method(),
            info.path(),
            info.status(),
        )
    });

    let routes = warp::get()
        .and(filters::client_ip(ip_header_names, ip_header_recursive))
        .map(
            move |ip: Option<IpAddr>| match ip.and_then(|ip| geo.try_lookup_continent(ip).ok()) {
                Some(continent) => continent_map.get(continent),
                None => continent_map.get_default(),
            },
        )
        .and(warp::path::full())
        .and(
            warp::query::raw()
                .map(Option::Some)
                .or_else(|_err| async { Ok::<_, warp::Rejection>((None,)) }),
        )
        .and_then(
            |mirrors: MirrorVec, path: FullPath, query: Option<String>| async move {
                let mirror = {
                    let mut it_mirrors = mirrors.iter();
                    loop {
                        match it_mirrors.next() {
                            Some(mirror) => {
                                if mirror.available.load(Ordering::Acquire) {
                                    break mirror;
                                }
                            }
                            None => return Err(warp::reject::custom(MirrorsUnavailable)),
                        }
                    }
                };
                let url = {
                    let mut url = mirror
                        .upstream
                        .join(path.as_str().trim_start_matches('/'))
                        .map_err(|_| warp::reject::custom(BrokenPath))?;
                    url.set_query(query.as_deref());
                    url
                };

                Ok(warp::redirect::found(url.as_str().parse::<Uri>().unwrap()).into_response())
            },
        )
        .with(warp::filters::reply::headers(response_headers))
        .recover(handle_rejection)
        .with(logs);

    warp::serve(routes).run(host).await;
    Ok(())
}