52 lines
1.6 KiB
Rust
52 lines
1.6 KiB
Rust
use actix_governor::KeyExtractor;
|
|
use actix_governor::PeerIpKeyExtractor;
|
|
use actix_web::{dev::ServiceRequest, http::header::ContentType};
|
|
use governor::clock::{Clock, DefaultClock, QuantaInstant};
|
|
use governor::NotUntil;
|
|
use std::net::IpAddr;
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub struct ForwardedPeerIpKeyExtractor {
|
|
pub proxied: bool,
|
|
}
|
|
|
|
impl KeyExtractor for ForwardedPeerIpKeyExtractor {
|
|
type Key = IpAddr;
|
|
type KeyExtractionError = &'static str;
|
|
|
|
#[cfg(feature = "log")]
|
|
fn name(&self) -> &'static str {
|
|
"Forwarded peer IP"
|
|
}
|
|
|
|
fn extract(&self, req: &ServiceRequest) -> Result<Self::Key, Self::KeyExtractionError> {
|
|
let forwarded_for = req.headers().get("x-forwarded-for");
|
|
if !self.proxied && forwarded_for.is_some() {
|
|
let forwarded_for = forwarded_for
|
|
.unwrap()
|
|
.to_str()
|
|
.map_err(|_| "x-forwarded-for contains invalid header value")?;
|
|
forwarded_for
|
|
.parse::<IpAddr>()
|
|
.map_err(|_| "x-forwarded-for contains invalid ip adress")
|
|
} else {
|
|
PeerIpKeyExtractor.extract(req)
|
|
}
|
|
}
|
|
|
|
fn response_error_content(&self, negative: &NotUntil<QuantaInstant>) -> (String, ContentType) {
|
|
let wait_time = negative
|
|
.wait_time_from(DefaultClock::default().now())
|
|
.as_secs();
|
|
(
|
|
format!("too many requests, retry in {}s", wait_time),
|
|
ContentType::plaintext(),
|
|
)
|
|
}
|
|
|
|
#[cfg(feature = "log")]
|
|
fn key_name(&self, key: &Self::Key) -> Option<String> {
|
|
Some(key.to_string())
|
|
}
|
|
}
|