datatrash/src/rate_limit.rs

use actix_governor::KeyExtractor;
use actix_governor::PeerIpKeyExtractor;
use actix_governor::SimpleKeyExtractionError;
use actix_web::HttpResponse;
use actix_web::HttpResponseBuilder;
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 = SimpleKeyExtractionError<&'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() {
            read_forwareded_for(forwarded_for).map_err(SimpleKeyExtractionError::new)
        } else {
            PeerIpKeyExtractor.extract(req)
        }
    }

    fn exceed_rate_limit_response(
        &self,
        negative: &NotUntil<QuantaInstant>,
        mut response: HttpResponseBuilder,
    ) -> HttpResponse {
        let wait_time = negative
            .wait_time_from(DefaultClock::default().now())
            .as_secs();
        response
            .content_type(ContentType::plaintext())
            .body(format!("too many requests, retry in {wait_time}s"))
    }
}

fn read_forwareded_for(
    forwarded_for: Option<&actix_web::http::header::HeaderValue>,
) -> Result<IpAddr, &'static str> {
    forwarded_for
        .unwrap()
        .to_str()
        .map_err(|_| "x-forwarded-for contains invalid header value")?
        .parse::<IpAddr>()
        .map_err(|_| "x-forwarded-for contains invalid ip adress")
}