rocket_analytics/
analytics.rs

1use chrono::Utc;
2use lazy_static::lazy_static;
3use reqwest::blocking::Client;
4use rocket::fairing::{Fairing, Info, Kind};
5use rocket::http::Status;
6use rocket::request::{FromRequest, Outcome};
7use rocket::{Data, Request, Response};
8use serde::Serialize;
9use std::sync::Mutex;
10use std::thread::spawn;
11use std::time::Instant;
12
13#[derive(Debug, Clone, Serialize)]
14struct RequestData {
15    hostname: String,
16    ip_address: String,
17    path: String,
18    user_agent: String,
19    method: String,
20    response_time: u32,
21    status: u16,
22    user_id: String,
23    created_at: String,
24}
25
26impl RequestData {
27    pub fn new(
28        hostname: String,
29        ip_address: String,
30        path: String,
31        user_agent: String,
32        method: String,
33        response_time: u32,
34        status: u16,
35        user_id: String,
36        created_at: String,
37    ) -> Self {
38        Self {
39            hostname,
40            ip_address,
41            path,
42            user_agent,
43            method,
44            response_time,
45            status,
46            user_id,
47            created_at,
48        }
49    }
50}
51
52type StringMapper = dyn for<'a> Fn(&Request<'a>) -> String + Send + Sync;
53
54struct Config {
55    privacy_level: i32,
56    server_url: String,
57    get_hostname: Box<StringMapper>,
58    get_ip_address: Box<StringMapper>,
59    get_path: Box<StringMapper>,
60    get_user_agent: Box<StringMapper>,
61    get_user_id: Box<StringMapper>,
62}
63
64impl Default for Config {
65    fn default() -> Self {
66        Self {
67            privacy_level: 0,
68            server_url: String::from("https://www.apianalytics-server.com/"),
69            get_hostname: Box::new(get_hostname),
70            get_ip_address: Box::new(get_ip_address),
71            get_path: Box::new(get_path),
72            get_user_agent: Box::new(get_user_agent),
73            get_user_id: Box::new(get_user_id),
74        }
75    }
76}
77
78fn get_hostname(req: &Request) -> String {
79    req.host().unwrap().to_string()
80}
81
82fn get_ip_address(req: &Request) -> String {
83    req.client_ip().unwrap().to_string()
84}
85
86fn get_path(req: &Request) -> String {
87    req.uri().path().to_string()
88}
89
90fn get_user_agent(req: &Request) -> String {
91    req.headers()
92        .get_one("User-Agent")
93        .unwrap_or_default()
94        .to_owned()
95}
96
97fn get_user_id(_req: &Request) -> String {
98    String::new()
99}
100
101#[derive(Default)]
102pub struct Analytics {
103    api_key: String,
104    config: Config,
105}
106
107impl Analytics {
108    pub fn new(api_key: String) -> Self {
109        Self {
110            api_key,
111            config: Config::default(),
112        }
113    }
114
115    pub fn with_privacy_level(mut self, privacy_level: i32) -> Self {
116        self.config.privacy_level = privacy_level;
117        self
118    }
119
120    pub fn with_server_url(mut self, server_url: String) -> Self {
121        if server_url.ends_with("/") {
122            self.config.server_url = server_url;
123        } else {
124            self.config.server_url = server_url + "/";
125        }
126        self
127    }
128
129    pub fn with_hostname_mapper<F>(mut self, mapper: F) -> Self
130    where
131        F: for<'a> Fn(&Request<'a>) -> String + Send + Sync + 'static,
132    {
133        self.config.get_hostname = Box::new(mapper);
134        self
135    }
136
137    pub fn with_ip_address_mapper<F>(mut self, mapper: F) -> Self
138    where
139        F: for<'a> Fn(&Request<'a>) -> String + Send + Sync + 'static,
140    {
141        self.config.get_ip_address = Box::new(mapper);
142        self
143    }
144
145    pub fn with_path_mapper<F>(mut self, mapper: F) -> Self
146    where
147        F: for<'a> Fn(&Request<'a>) -> String + Send + Sync + 'static,
148    {
149        self.config.get_path = Box::new(mapper);
150        self
151    }
152
153    pub fn with_user_agent_mapper<F>(mut self, mapper: F) -> Self
154    where
155        F: for<'a> Fn(&Request<'a>) -> String + Send + Sync + 'static,
156    {
157        self.config.get_user_agent = Box::new(mapper);
158        self
159    }
160}
161
162#[derive(Clone)]
163pub struct Start<T = Instant>(T);
164
165lazy_static! {
166    static ref REQUESTS: Mutex<Vec<RequestData>> = Mutex::new(vec![]);
167    static ref LAST_POSTED: Mutex<Instant> = Mutex::new(Instant::now());
168}
169
170#[derive(Debug, Clone, Serialize)]
171struct Payload {
172    api_key: String,
173    requests: Vec<RequestData>,
174    framework: String,
175    privacy_level: i32,
176}
177
178impl Payload {
179    pub fn new(api_key: String, requests: Vec<RequestData>, privacy_level: i32) -> Self {
180        Self {
181            api_key,
182            requests,
183            framework: String::from("Rocket"),
184            privacy_level,
185        }
186    }
187}
188
189fn post_requests(data: Payload, server_url: String) {
190    let _ = Client::new()
191        .post(server_url + "api/log-request")
192        .json(&data)
193        .send();
194}
195
196fn log_request(api_key: &str, request_data: RequestData, config: &Config) {
197    REQUESTS.lock().unwrap().push(request_data);
198    if LAST_POSTED.lock().unwrap().elapsed().as_secs_f64() > 60.0 {
199        let payload = Payload::new(
200            api_key.to_string(),
201            REQUESTS.lock().unwrap().to_vec(),
202            config.privacy_level,
203        );
204        let server_url = config.server_url.to_owned();
205        REQUESTS.lock().unwrap().clear();
206        spawn(|| post_requests(payload, server_url));
207        *LAST_POSTED.lock().unwrap() = Instant::now();
208    }
209}
210
211#[rocket::async_trait]
212impl Fairing for Analytics {
213    fn info(&self) -> Info {
214        Info {
215            name: "API Analytics",
216            kind: Kind::Request | Kind::Response,
217        }
218    }
219
220    async fn on_request(&self, req: &mut Request<'_>, _data: &mut Data<'_>) {
221        req.local_cache(|| Start::<Option<Instant>>(Some(Instant::now())));
222    }
223
224    async fn on_response<'r>(&self, req: &'r Request<'_>, res: &mut Response<'r>) {
225        let start = &req.local_cache(|| Start::<Option<Instant>>(None)).0;
226        let hostname = (self.config.get_hostname)(req);
227        let ip_address = (self.config.get_ip_address)(req);
228        let method = req.method().to_string();
229        let user_agent = (self.config.get_user_agent)(req);
230        let path = (self.config.get_path)(req);
231        let user_id = (self.config.get_user_id)(req);
232
233        let request_data = RequestData::new(
234            hostname,
235            ip_address,
236            path,
237            user_agent,
238            method,
239            start.unwrap().elapsed().as_millis() as u32,
240            res.status().code,
241            user_id,
242            Utc::now().to_rfc3339(),
243        );
244
245        log_request(&self.api_key, request_data, &self.config);
246    }
247}
248
249// Allows a route to access the start time
250#[rocket::async_trait]
251impl<'r> FromRequest<'r> for Start {
252    type Error = ();
253
254    async fn from_request(request: &'r Request<'_>) -> Outcome<Self, ()> {
255        match &*request.local_cache(|| Start::<Option<Instant>>(None)) {
256            Start(Some(start)) => Outcome::Success(Start(start.to_owned())),
257            Start(None) => Outcome::Error((Status::InternalServerError, ())),
258        }
259    }
260}