datatrash/src/rate_limit.rs

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() {
            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(),
        )
    }

    fn response_error(&self, err: Self::KeyExtractionError) -> actix_web::Error {
        actix_web::error::ErrorUnauthorized(err.to_string())
    }
}