systemprompt_api/services/middleware/
client_addr.rs1use std::convert::Infallible;
29use std::future::{Future, ready};
30use std::net::{IpAddr, SocketAddr};
31
32use axum::extract::{ConnectInfo, FromRequestParts};
33use axum::http::HeaderMap;
34use axum::http::request::Parts;
35use ipnet::IpNet;
36
37fn is_trusted(addr: IpAddr, trusted: &[IpNet]) -> bool {
38 trusted.iter().any(|net| net.contains(&addr))
39}
40
41#[must_use]
42pub fn resolve_client_ip(
43 headers: &HeaderMap,
44 connect_info: Option<&ConnectInfo<SocketAddr>>,
45 trusted: &[IpNet],
46) -> Option<IpAddr> {
47 let peer_ip = connect_info.map(|c| c.0.ip())?;
48
49 if !is_trusted(peer_ip, trusted) {
50 return Some(peer_ip);
51 }
52
53 if let Some(xff) = headers.get("x-forwarded-for").and_then(|v| v.to_str().ok()) {
54 let hops: Vec<&str> = xff
55 .split(',')
56 .map(str::trim)
57 .filter(|s| !s.is_empty())
58 .collect();
59 for hop in hops.iter().rev() {
60 if let Ok(addr) = hop.parse::<IpAddr>()
61 && !is_trusted(addr, trusted)
62 {
63 return Some(addr);
64 }
65 }
66 }
67
68 for header in ["x-real-ip", "cf-connecting-ip"] {
69 if let Some(raw) = headers.get(header).and_then(|v| v.to_str().ok())
70 && let Ok(addr) = raw.trim().parse::<IpAddr>()
71 && !is_trusted(addr, trusted)
72 {
73 return Some(addr);
74 }
75 }
76
77 Some(peer_ip)
78}
79
80#[must_use]
81pub fn resolve_client_ip_from_config(
82 headers: &HeaderMap,
83 connect_info: Option<&ConnectInfo<SocketAddr>>,
84) -> Option<IpAddr> {
85 let trusted = systemprompt_models::Config::get()
86 .map(|c| c.trusted_proxies.clone())
87 .unwrap_or_default();
88 resolve_client_ip(headers, connect_info, &trusted)
89}
90
91#[must_use]
92pub fn client_ip_from_request(request: &axum::extract::Request) -> Option<IpAddr> {
93 resolve_client_ip_from_config(
94 request.headers(),
95 request.extensions().get::<ConnectInfo<SocketAddr>>(),
96 )
97}
98
99#[derive(Debug, Clone, Copy)]
100pub struct ClientIp(pub Option<IpAddr>);
101
102impl<S: Sync> FromRequestParts<S> for ClientIp {
103 type Rejection = Infallible;
104
105 fn from_request_parts(
106 parts: &mut Parts,
107 _state: &S,
108 ) -> impl Future<Output = Result<Self, Infallible>> + Send {
109 let resolved = resolve_client_ip_from_config(
110 &parts.headers,
111 parts.extensions.get::<ConnectInfo<SocketAddr>>(),
112 );
113 ready(Ok(Self(resolved)))
114 }
115}