summaryrefslogtreecommitdiff
path: root/exes/rest/src/proxy/mod.rs
blob: 65d77aaec907d5a51fe6f40577f2592378008ad4 (plain)
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
135
136
137
138
use crate::{config::Config, ratelimit::Ratelimiter};
use hyper::{
    client::HttpConnector, header::HeaderValue, http::uri::Parts, service::Service, Body, Client,
    Request, Response, Uri,
};
use hyper_tls::HttpsConnector;
use shared::{
    log::debug,
    prometheus::{labels, opts, register_counter, register_histogram_vec, Counter, HistogramVec},
};
use std::{future::Future, pin::Pin, sync::Arc, task::Poll};
use tokio::sync::Mutex;

lazy_static::lazy_static! {
    static ref HTTP_COUNTER: Counter = register_counter!(opts!(
        "nova_rest_http_requests_total",
        "Number of HTTP requests made.",
        labels! {"handler" => "all",}
    ))
    .unwrap();

    static ref HTTP_REQ_HISTOGRAM: HistogramVec = register_histogram_vec!(
        "nova_rest_http_request_duration_seconds",
        "The HTTP request latencies in seconds.",
        &["handler"]
    )
    .unwrap();

    static ref HTTP_COUNTER_STATUS: Counter = register_counter!(opts!(
        "nova_rest_http_requests_status",
        "Number of HTTP requests made by status",
        labels! {"" => ""}
    ))
    .unwrap();
}

#[derive(Clone)]
pub struct ServiceProxy {
    client: Client<HttpsConnector<HttpConnector>>,
    ratelimiter: Arc<Ratelimiter>,
    config: Arc<Config>,
    fail: Arc<Mutex<i32>>,
}

impl Service<Request<Body>> for ServiceProxy {
    type Response = Response<Body>;
    type Error = hyper::Error;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

    fn poll_ready(
        &mut self,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Result<(), Self::Error>> {
        match self.client.poll_ready(cx) {
            Poll::Ready(Ok(())) => Poll::Ready(Ok(())),
            Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
            Poll::Pending => Poll::Pending,
        }
    }

    fn call(&mut self, mut req: Request<hyper::Body>) -> Self::Future {
        HTTP_COUNTER.inc();

        let timer = HTTP_REQ_HISTOGRAM.with_label_values(&["all"]).start_timer();
        let host = "discord.com";
        let mut new_parts = Parts::default();

        let path = req.uri().path().to_string();

        new_parts.scheme = Some("https".parse().unwrap());
        new_parts.authority = Some(host.parse().unwrap());
        new_parts.path_and_query = Some(path.parse().unwrap());

        *req.uri_mut() = Uri::from_parts(new_parts).unwrap();

        let headers = req.headers_mut();
        headers.remove("user-agent");
        headers.insert("Host", HeaderValue::from_str("discord.com").unwrap());
        headers.insert(
            "Authorization",
            HeaderValue::from_str(&format!("Bot {}", self.config.discord.token)).unwrap(),
        );

        println!("{:?}", headers);

        let client = self.client.clone();
        let ratelimiter = self.ratelimiter.clone();
        let fail = self.fail.clone();

        return Box::pin(async move {
            let resp = match ratelimiter.before_request(&req).await {
                Ok(allowed) => match allowed {
                    crate::ratelimit::RatelimiterResponse::Ratelimited => {
                        debug!("ratelimited");
                        Ok(Response::builder().body("ratelimited".into()).unwrap())
                    }
                    _ => {
                        debug!("forwarding request");
                        match client.request(req).await {
                            Ok(mut response) => {
                                ratelimiter.after_request(&path, &response).await;
                                if response.status() != 200 {
                                    *fail.lock().await += 1
                                }
                                response.headers_mut().insert(
                                    "x-fails",
                                    HeaderValue::from_str(&format!("{}", fail.lock().await))
                                        .unwrap(),
                                );
                                Ok(response)
                            }
                            Err(e) => Err(e),
                        }
                    }
                },
                Err(e) => Ok(Response::builder()
                    .body(format!("server error: {}", e).into())
                    .unwrap()),
            };
            timer.observe_duration();
            resp
        });
    }
}

impl ServiceProxy {
    pub fn new(config: Arc<Config>, ratelimiter: Arc<Ratelimiter>) -> Self {
        let https = HttpsConnector::new();
        let client = Client::builder().build::<_, hyper::Body>(https);
        let fail = Arc::new(Mutex::new(0));
        ServiceProxy {
            client,
            config,
            ratelimiter,
            fail,
        }
    }
}