axum_analytics/
analytics.rs

1use axum::{body::Body, extract::ConnectInfo, http::Request, response::Response};
2use chrono::Utc;
3use std::future::Future;
4use http::{
5    header::{HeaderValue, HOST, USER_AGENT},
6    Extensions, HeaderMap,
7};
8use reqwest::Client;
9use serde::Serialize;
10use std::{
11    net::{IpAddr, SocketAddr},
12    pin::Pin,
13    sync::Arc,
14    task::{Context, Poll},
15    time::{Duration, Instant},
16};
17use tokio::sync::RwLock;
18use tokio::time::interval;
19use tower::{Layer, Service};
20
21#[derive(Debug, Clone, Serialize)]
22struct RequestData {
23    hostname: String,
24    ip_address: Option<String>,
25    path: String,
26    user_agent: String,
27    method: String,
28    response_time: u32,
29    status: u16,
30    user_id: Option<String>,
31    created_at: String,
32}
33
34type StringMapper = dyn for<'a> Fn(&Request<Body>) -> String + Send + Sync;
35
36#[derive(Clone)]
37struct Config {
38    privacy_level: i32,
39    server_url: String,
40    get_hostname: Arc<StringMapper>,
41    get_ip_address: Arc<StringMapper>,
42    get_path: Arc<StringMapper>,
43    get_user_agent: Arc<StringMapper>,
44    get_user_id: Arc<StringMapper>,
45}
46
47impl Default for Config {
48    fn default() -> Self {
49        Self {
50            privacy_level: 0,
51            server_url: "https://www.apianalytics-server.com/".to_string(),
52            get_hostname: Arc::new(get_hostname),
53            get_ip_address: Arc::new(get_ip_address),
54            get_path: Arc::new(get_path),
55            get_user_agent: Arc::new(get_user_agent),
56            get_user_id: Arc::new(get_user_id),
57        }
58    }
59}
60
61pub trait HeaderValueExt {
62    fn to_string(&self) -> String;
63}
64
65impl HeaderValueExt for HeaderValue {
66    fn to_string(&self) -> String {
67        self.to_str().unwrap_or_default().to_string()
68    }
69}
70
71fn get_hostname(req: &Request<Body>) -> String {
72    req.headers()
73        .get(HOST)
74        .map(|x| x.to_string())
75        .unwrap_or_default()
76}
77
78fn get_ip_address(req: &Request<Body>) -> String {
79    let headers = req.headers();
80    let extensions = req.extensions();
81    if let Some(val) = ip_from_cf_connecting_ip(headers) {
82        return val.to_string();
83    }
84    if let Some(val) = ip_from_x_forwarded_for(headers) {
85        return val.to_string();
86    }
87    if let Some(val) = ip_from_x_real_ip(headers) {
88        return val.to_string();
89    }
90    if let Some(val) = ip_from_connect_info(extensions) {
91        return val.to_string();
92    }
93    String::new()
94}
95
96fn ip_from_cf_connecting_ip(headers: &HeaderMap) -> Option<IpAddr> {
97    headers
98        .get("cf-connecting-ip")
99        .and_then(|hv| hv.to_str().ok())
100        .and_then(|s| s.trim().parse::<IpAddr>().ok())
101}
102
103fn ip_from_x_forwarded_for(headers: &HeaderMap) -> Option<IpAddr> {
104    headers
105        .get("x-forwarded-for")
106        .and_then(|hv| hv.to_str().ok())
107        .and_then(|s| s.split(',').find_map(|s| s.trim().parse::<IpAddr>().ok()))
108}
109
110fn ip_from_x_real_ip(headers: &HeaderMap) -> Option<IpAddr> {
111    headers
112        .get("x-real-ip")
113        .and_then(|hv| hv.to_str().ok())
114        .and_then(|s| s.parse::<IpAddr>().ok())
115}
116
117fn ip_from_connect_info(extensions: &Extensions) -> Option<IpAddr> {
118    extensions
119        .get::<ConnectInfo<SocketAddr>>()
120        .map(|ConnectInfo(addr)| addr.ip())
121}
122
123fn get_path(req: &Request<Body>) -> String {
124    req.uri().path().to_owned()
125}
126
127fn get_user_agent(req: &Request<Body>) -> String {
128    req.headers()
129        .get(USER_AGENT)
130        .map(|x| x.to_string())
131        .unwrap_or_default()
132}
133
134fn get_user_id(_req: &Request<Body>) -> String {
135    String::new()
136}
137
138#[derive(Clone)]
139pub struct Analytics {
140    config: Config,
141    requests: Arc<RwLock<Vec<RequestData>>>,
142}
143
144impl Analytics {
145    pub fn new(api_key: String) -> Self {
146        let requests: Arc<RwLock<Vec<RequestData>>> = Arc::new(RwLock::new(vec![]));
147        let requests_clone = Arc::clone(&requests);
148        let config = Config::default();
149        let privacy_level = config.privacy_level;
150        let server_url = config.server_url.clone();
151
152        tokio::spawn(async move {
153            let client = Client::builder()
154                .timeout(Duration::from_secs(10))
155                .build()
156                .unwrap_or_else(|_| Client::new());
157
158            let mut ticker = interval(Duration::from_secs(60));
159            ticker.tick().await; // skip the immediate first tick
160            loop {
161                ticker.tick().await;
162                let batch = {
163                    let mut reqs = requests_clone.write().await;
164                    std::mem::take(&mut *reqs)
165                };
166                if !batch.is_empty() {
167                    let payload = Payload::new(api_key.clone(), batch, privacy_level);
168                    post_requests(&client, payload, &server_url).await;
169                }
170            }
171        });
172
173        Self { config, requests }
174    }
175
176    pub fn with_privacy_level(mut self, privacy_level: i32) -> Self {
177        self.config.privacy_level = privacy_level;
178        self
179    }
180
181    pub fn with_server_url(mut self, server_url: String) -> Self {
182        self.config.server_url = if server_url.ends_with('/') {
183            server_url
184        } else {
185            server_url + "/"
186        };
187        self
188    }
189
190    pub fn with_hostname_mapper<F>(mut self, mapper: F) -> Self
191    where
192        F: Fn(&Request<Body>) -> String + Send + Sync + 'static,
193    {
194        self.config.get_hostname = Arc::new(mapper);
195        self
196    }
197
198    pub fn with_ip_address_mapper<F>(mut self, mapper: F) -> Self
199    where
200        F: Fn(&Request<Body>) -> String + Send + Sync + 'static,
201    {
202        self.config.get_ip_address = Arc::new(mapper);
203        self
204    }
205
206    pub fn with_path_mapper<F>(mut self, mapper: F) -> Self
207    where
208        F: Fn(&Request<Body>) -> String + Send + Sync + 'static,
209    {
210        self.config.get_path = Arc::new(mapper);
211        self
212    }
213
214    pub fn with_user_agent_mapper<F>(mut self, mapper: F) -> Self
215    where
216        F: Fn(&Request<Body>) -> String + Send + Sync + 'static,
217    {
218        self.config.get_user_agent = Arc::new(mapper);
219        self
220    }
221
222    pub fn with_user_id_mapper<F>(mut self, mapper: F) -> Self
223    where
224        F: Fn(&Request<Body>) -> String + Send + Sync + 'static,
225    {
226        self.config.get_user_id = Arc::new(mapper);
227        self
228    }
229}
230
231impl<S> Layer<S> for Analytics {
232    type Service = AnalyticsMiddleware<S>;
233
234    fn layer(&self, inner: S) -> Self::Service {
235        AnalyticsMiddleware {
236            config: Arc::new(self.config.clone()),
237            requests: Arc::clone(&self.requests),
238            inner,
239        }
240    }
241}
242
243#[derive(Clone)]
244pub struct AnalyticsMiddleware<S> {
245    config: Arc<Config>,
246    requests: Arc<RwLock<Vec<RequestData>>>,
247    inner: S,
248}
249
250#[derive(Debug, Clone, Serialize)]
251struct Payload {
252    api_key: String,
253    requests: Vec<RequestData>,
254    framework: String,
255    privacy_level: i32,
256}
257
258impl Payload {
259    pub fn new(api_key: String, requests: Vec<RequestData>, privacy_level: i32) -> Self {
260        Self {
261            api_key,
262            requests,
263            framework: "Axum".to_string(),
264            privacy_level,
265        }
266    }
267}
268
269async fn post_requests(client: &Client, data: Payload, server_url: &str) {
270    let url = format!("{}api/log-request", server_url);
271    match client.post(url).json(&data).send().await {
272        Ok(resp) if resp.status().is_success() => {}
273        Ok(resp) => {
274            eprintln!(
275                "Failed to send analytics data. Server responded with status: {}",
276                resp.status()
277            );
278        }
279        Err(err) => {
280            eprintln!("Error sending analytics data: {}", err);
281        }
282    }
283}
284
285impl<S> Service<Request<Body>> for AnalyticsMiddleware<S>
286where
287    S: Service<Request<Body>, Response = Response> + Clone + Send + Sync + 'static,
288    S::Future: Send + 'static,
289{
290    type Response = S::Response;
291    type Error = S::Error;
292    type Future =
293        Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'static>>;
294
295    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
296        self.inner.poll_ready(cx)
297    }
298
299    fn call(&mut self, req: Request<Body>) -> Self::Future {
300        let now = Instant::now();
301
302        let hostname = (self.config.get_hostname)(&req);
303        let ip_address = if self.config.privacy_level >= 2 {
304            None
305        } else {
306            let val = (self.config.get_ip_address)(&req);
307            if val.is_empty() { None } else { Some(val) }
308        };
309        let path = (self.config.get_path)(&req);
310        let method = req.method().to_string();
311        let user_agent = (self.config.get_user_agent)(&req);
312        let user_id = {
313            let val = (self.config.get_user_id)(&req);
314            if val.is_empty() { None } else { Some(val) }
315        };
316        let requests = Arc::clone(&self.requests);
317
318        let future = self.inner.call(req);
319
320        Box::pin(async move {
321            let res: Response = future.await?;
322            let response_time = now.elapsed().as_millis().min(u32::MAX as u128) as u32;
323
324            let request_data = RequestData {
325                hostname,
326                ip_address,
327                path,
328                user_agent,
329                method,
330                response_time,
331                status: res.status().as_u16(),
332                user_id,
333                created_at: Utc::now().to_rfc3339(),
334            };
335
336            requests.write().await.push(request_data);
337
338            Ok(res)
339        })
340    }
341}
342
343#[cfg(test)]
344mod tests {
345    use super::*;
346    use std::str::FromStr;
347
348    fn make_headers(pairs: &[(&str, &str)]) -> HeaderMap {
349        let mut map = HeaderMap::new();
350        for (name, value) in pairs {
351            map.insert(
352                http::header::HeaderName::from_str(name).unwrap(),
353                http::header::HeaderValue::from_str(value).unwrap(),
354            );
355        }
356        map
357    }
358
359    #[test]
360    fn test_cf_connecting_ip() {
361        let headers = make_headers(&[("cf-connecting-ip", "203.0.113.1")]);
362        assert_eq!(ip_from_cf_connecting_ip(&headers), Some("203.0.113.1".parse().unwrap()));
363    }
364
365    #[test]
366    fn test_x_forwarded_for_single() {
367        let headers = make_headers(&[("x-forwarded-for", "203.0.113.1")]);
368        assert_eq!(ip_from_x_forwarded_for(&headers), Some("203.0.113.1".parse().unwrap()));
369    }
370
371    #[test]
372    fn test_x_forwarded_for_returns_first_ip() {
373        // First IP is the client; last is the closest proxy — must NOT use .rev()
374        let headers = make_headers(&[("x-forwarded-for", "203.0.113.1, 10.0.0.1, 192.168.1.1")]);
375        assert_eq!(ip_from_x_forwarded_for(&headers), Some("203.0.113.1".parse().unwrap()));
376    }
377
378    #[test]
379    fn test_x_real_ip() {
380        let headers = make_headers(&[("x-real-ip", "203.0.113.1")]);
381        assert_eq!(ip_from_x_real_ip(&headers), Some("203.0.113.1".parse().unwrap()));
382    }
383
384    #[test]
385    fn test_invalid_ip_returns_none() {
386        let headers = make_headers(&[("x-real-ip", "not-an-ip")]);
387        assert_eq!(ip_from_x_real_ip(&headers), None);
388    }
389
390    #[test]
391    fn test_cf_takes_priority_over_forwarded() {
392        let headers = make_headers(&[
393            ("cf-connecting-ip", "1.1.1.1"),
394            ("x-forwarded-for", "2.2.2.2"),
395        ]);
396        assert_eq!(ip_from_cf_connecting_ip(&headers), Some("1.1.1.1".parse().unwrap()));
397        assert_eq!(ip_from_x_forwarded_for(&headers), Some("2.2.2.2".parse().unwrap()));
398    }
399
400    #[test]
401    fn test_ipv6_address() {
402        let headers = make_headers(&[("x-real-ip", "2001:db8::1")]);
403        assert_eq!(ip_from_x_real_ip(&headers), Some("2001:db8::1".parse().unwrap()));
404    }
405
406    #[test]
407    fn test_missing_header_returns_none() {
408        let headers = HeaderMap::new();
409        assert_eq!(ip_from_cf_connecting_ip(&headers), None);
410        assert_eq!(ip_from_x_forwarded_for(&headers), None);
411        assert_eq!(ip_from_x_real_ip(&headers), None);
412    }
413}
axum_analytics/analytics.rs

axum_analytics/
analytics.rs
