actix_ipquery/
lib.rs

1/// # IPQuery Actix Web Middleware
2///
3/// This middleware allows you to query IP information using the `ipapi` crate and store the results
4/// using a custom store that implements the `IPQueryStore` trait. It supports querying the IP address
5/// from either the `X-Forwarded-For` header or the peer address of the request.
6///
7/// ## Features
8/// - Query IP information using a specified endpoint.
9/// - Store IP information using a custom store.
10/// - Option to use the `X-Forwarded-For` header for IP address extraction.
11///
12/// ## Usage Example
13/// ```rust
14/// use actix_ipquery::{IPInfo, IPQuery, IPQueryStore};
15/// use actix_web::{App, HttpServer};
16///
17/// #[actix_web::main]
18/// async fn main() {
19///     HttpServer::new(|| App::new().wrap(IPQuery::new(Store).finish()))
20///         .bind("127.0.0.1:8080")
21///         .unwrap()
22///         .run()
23///         .await
24///         .unwrap()
25/// }
26///
27/// #[derive(Clone)]
28/// struct Store;
29/// impl IPQueryStore for Store {
30///     fn store(
31///         &self,
32///         ip_info: IPInfo,
33///     ) -> std::pin::Pin<
34///         Box<dyn std::future::Future<Output = Result<(), std::io::Error>> + Send + 'static>,
35///     > {
36///         println!("{:?}", ip_info);
37///         Box::pin(async { Ok(()) })
38///     }
39/// }
40/// ```
41use ipapi::query_ip;
42pub use ipapi::IPInfo;
43use std::future::{ready, Ready};
44
45use actix_web::{
46    dev::{forward_ready, Service, ServiceRequest, ServiceResponse, Transform},
47    Error,
48};
49use futures_util::future::LocalBoxFuture;
50
51/// The IPQuery struct that implements actix-web's middleware.
52#[derive(Clone)]
53pub struct IPQuery<T: IPQueryStore> {
54    store: T,
55    forwarded_for: bool,
56}
57impl<T: IPQueryStore + 'static> IPQuery<T> {
58    /// Create a new IPQuery middleware
59    pub fn new(store: T) -> IPQuery<T> {
60        IPQuery {
61            store,
62            forwarded_for: false,
63        }
64    }
65
66    /// Use the `X-Forwarded-For` header for IP address extraction
67    pub fn forwarded_for(&mut self, y: bool) -> &mut Self {
68        self.forwarded_for = y;
69        self
70    }
71    /// Finish the configuration and return the middleware
72    pub fn finish(&self) -> IPQuery<T> {
73        self.clone()
74    }
75    async fn query_ip(&self, ip: &str) -> Result<IPInfo, Box<dyn std::error::Error + Send + Sync>> {
76        Ok(query_ip(ip).await?)
77    }
78}
79impl<S, B, T> Transform<S, ServiceRequest> for IPQuery<T>
80where
81    T: IPQueryStore + 'static,
82    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error>,
83    S::Future: 'static,
84    B: 'static,
85    T: IPQueryStore + Clone,
86{
87    type Response = ServiceResponse<B>;
88    type Error = Error;
89    type InitError = ();
90    type Transform = IPQueryMiddleware<S, T>;
91    type Future = Ready<Result<Self::Transform, Self::InitError>>;
92
93    fn new_transform(&self, service: S) -> Self::Future {
94        ready(Ok(IPQueryMiddleware {
95            service,
96            ip_query: std::sync::Arc::new(self.clone()),
97        }))
98    }
99}
100
101pub struct IPQueryMiddleware<S, T>
102where
103    T: IPQueryStore,
104{
105    service: S,
106    ip_query: std::sync::Arc<IPQuery<T>>,
107}
108
109impl<S, B, T> Service<ServiceRequest> for IPQueryMiddleware<S, T>
110where
111    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error>,
112    S::Future: 'static,
113    B: 'static,
114    T: IPQueryStore + Clone + 'static,
115{
116    type Response = ServiceResponse<B>;
117    type Error = Error;
118    type Future = LocalBoxFuture<'static, Result<Self::Response, Self::Error>>;
119    forward_ready!(service);
120
121    fn call(&self, req: ServiceRequest) -> Self::Future {
122        let ip = if self.ip_query.forwarded_for {
123            req.connection_info()
124                .realip_remote_addr()
125                .unwrap()
126                .to_string()
127        } else {
128            match req.peer_addr() {
129                Some(addr) => addr.ip().to_string(),
130                None => {
131                    return Box::pin(async {
132                        Err(Error::from(actix_web::error::ErrorInternalServerError(
133                            "No peer address",
134                        )))
135                    })
136                }
137            }
138        };
139
140        let fut = self.service.call(req);
141        let ip_query_clone = self.ip_query.clone();
142        Box::pin(async move {
143            let res = fut.await?;
144            let ip_info = match ip_query_clone.query_ip(&ip).await {
145                Ok(info) => info,
146                Err(e) => {
147                    return Err(Error::from(actix_web::error::ErrorInternalServerError(
148                        e.to_string(),
149                    )))
150                }
151            };
152            ip_query_clone.store.store(ip_info).await?;
153            Ok(res)
154        })
155    }
156}
157
158/// Define the IPQueryStore trait
159pub trait IPQueryStore: Send + Sync + Clone {
160    fn store(
161        &self,
162        ip_info: IPInfo,
163    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<(), std::io::Error>> + Send>>;
164}
165
166#[cfg(test)]
167mod tests {
168    #[tokio::test]
169    async fn my_ip() {
170        let ip = ipapi::query_own_ip().await.unwrap();
171        println!("{:?}", ip);
172    }
173}
actix_ipquery/lib.rs

actix_ipquery/
lib.rs
