datatrash/src/rate_limit.rs

46 lines
1.5 KiB
Rust
Raw Normal View History

2022-08-21 16:44:12 +00:00
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;
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() {
2022-08-21 16:44:12 +00:00
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();
(
2022-10-07 13:52:12 +00:00
format!("too many requests, retry in {wait_time}s"),
2022-08-21 16:44:12 +00:00
ContentType::plaintext(),
)
}
2022-09-30 12:38:33 +00:00
fn response_error(&self, err: Self::KeyExtractionError) -> actix_web::Error {
actix_web::error::ErrorUnauthorized(err.to_string())
2022-08-21 16:44:12 +00:00
}
}