This repository has no description
1use std::convert::Infallible;
2use std::net::{IpAddr, SocketAddr};
3
4use axum::extract::{ConnectInfo, FromRequestParts};
5use http::request::Parts;
6
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8pub struct SocketPeer(Option<IpAddr>);
9
10impl SocketPeer {
11 pub fn ip(self) -> Option<IpAddr> {
12 self.0
13 }
14}
15
16impl<S: Send + Sync> FromRequestParts<S> for SocketPeer {
17 type Rejection = Infallible;
18
19 async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
20 Ok(Self(
21 parts
22 .extensions
23 .get::<ConnectInfo<SocketAddr>>()
24 .map(|connect| connect.0.ip()),
25 ))
26 }
27}
28
29#[cfg(test)]
30mod tests {
31 use super::*;
32 use axum::body::Body;
33 use std::net::Ipv4Addr;
34
35 #[tokio::test]
36 async fn the_extractor_reads_connect_info_and_tolerates_its_absence() {
37 let with = {
38 let mut request = http::Request::builder().body(Body::empty()).unwrap();
39 request
40 .extensions_mut()
41 .insert(ConnectInfo(SocketAddr::from(([203, 0, 113, 7], 443))));
42 let (mut parts, _) = request.into_parts();
43 SocketPeer::from_request_parts(&mut parts, &())
44 .await
45 .unwrap()
46 };
47 assert_eq!(with.ip(), Some(IpAddr::V4(Ipv4Addr::new(203, 0, 113, 7))));
48
49 let without = {
50 let request = http::Request::builder().body(Body::empty()).unwrap();
51 let (mut parts, _) = request.into_parts();
52 SocketPeer::from_request_parts(&mut parts, &())
53 .await
54 .unwrap()
55 };
56 assert_eq!(without.ip(), None);
57 }
58}