use std::{
collections::HashMap,
net::{IpAddr, Ipv4Addr},
str::FromStr,
};
use itertools::Itertools;
use tardis::{
basic::{dto::TardisContext, result::TardisResult},
web::poem::{http::header::FORWARDED, Request},
};
pub const REMOTE_ADDR: &str = "remote-addr";
pub async fn add_ip(ip: Option<String>, ctx: &TardisContext) -> TardisResult<()> {
if let Some(ip) = ip {
ctx.add_ext(REMOTE_ADDR, &ip).await?;
}
Ok(())
}
pub async fn try_set_real_ip_from_req_to_ctx(request: &Request, ctx: &TardisContext) -> TardisResult<()> {
ctx.add_ext(REMOTE_ADDR, &try_get_real_ip_from_req(request).await?.unwrap_or_default()).await?;
Ok(())
}
pub fn parse_forwarded_ip(forwarded_value: &str) -> Option<IpAddr> {
forwarded_value.strip_prefix("Forwarded: ").and_then(|forwarded_value| {
forwarded_value
.split(';')
.find(|part| part.trim().starts_with("for="))
.and_then(|part| part.trim()[4..].split(',').next().and_then(|ip_str| IpAddr::from_str(ip_str).ok()))
})
}
pub fn mapped_ipv6_to_ipv4(ip_addr: IpAddr) -> IpAddr {
match ip_addr {
IpAddr::V6(ip_addr) => {
let segments = ip_addr.segments();
if segments[0] == 0 && segments[1] == 0 && segments[2] == 0 && segments[3] == 0 && segments[4] == 0 && segments[5] == 0xffff {
IpAddr::V4(Ipv4Addr::new(ip_addr.octets()[12], ip_addr.octets()[13], ip_addr.octets()[14], ip_addr.octets()[15]))
} else {
IpAddr::V6(ip_addr)
}
}
IpAddr::V4(ip_addr) => IpAddr::V4(ip_addr),
}
}
pub async fn try_get_real_ip_from_req(request: &Request) -> TardisResult<Option<String>> {
if let Some(forwarded_header) = request.headers().get(FORWARDED) {
if let Ok(forwarded_value) = forwarded_header.to_str() {
if let Some(ip) = parse_forwarded_ip(forwarded_value.trim()) {
return Ok(Some(mapped_ipv6_to_ipv4(ip).to_string()));
}
}
}
if let Some(xff_header) = request.headers().get("X-Forwarded-For") {
if let Ok(xff_value) = xff_header.to_str() {
if let Some(ip) = xff_value.split(',').next().and_then(|s| IpAddr::from_str(s.trim()).ok()) {
return Ok(Some(mapped_ipv6_to_ipv4(ip).to_string()));
}
}
}
if let Some(xrp_header) = request.headers().get("X-Real-IP") {
if let Ok(xrp_value) = xrp_header.to_str() {
if let Some(ip) = xrp_value.split(',').next().and_then(|s| IpAddr::from_str(s.trim()).ok()) {
return Ok(Some(mapped_ipv6_to_ipv4(ip).to_string()));
}
}
}
Ok(request.remote_addr().as_socket_addr().map(|addr| mapped_ipv6_to_ipv4(addr.ip()).to_string()))
}
pub async fn get_real_ip_from_ctx(ctx: &TardisContext) -> TardisResult<Option<String>> {
ctx.get_ext(REMOTE_ADDR).await
}
pub fn sort_query(query: &str) -> String {
if query.is_empty() {
return "".to_string();
}
query.split('&').sorted_by(|a, b| Ord::cmp(&a.to_lowercase(), &b.to_lowercase())).join("&")
}
pub fn sort_hashmap_query(query: HashMap<String, String>) -> String {
if query.is_empty() {
return "".to_string();
}
query.iter().map(|a| format!("{}={}", a.0, a.1)).sorted_by(|a, b| Ord::cmp(&a.to_lowercase(), &b.to_lowercase())).join("&")
}