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