Test panic when return error

[DimitriTimoz]

hello,
When I try to test my spam system and when my middleware return HttpResponse::TooManyRequests() the test panic with this error:

thread 'middlewares::firewall::tests::test_spam_protection' panicked at 'called `Result::unwrap()` on an `Err` value: ""', /Users/myname/.cargo/registry/src/github.com-1ecc6299db9ec823/actix-web-3.3.2/src/test.rs:144:25
stack backtrace:

    #[actix_rt::test]
    async fn test_spam_protection() {

        let user_id = DATABASE.create_user("test_spam_user_blocked", None).unwrap();
        let api_key = DATABASE.create_api_key(user_id).unwrap();
        let mut app = init_service(
            App::new()
                .service(get_schedule)
                .wrap(super::Auth::enabled())
                .data(EventDb::new_passive()),
        )
        .await;

        for i in 1..100 {
            let code = gen_code(api_key, i);
            let api_key_header_value = format!("{}-{}-{}", api_key, i, code);
            let req = TestRequest::get()
                .uri("/schedule/STPI1/E-1/ALL/0-99999999999")
                .header("Api-Key", api_key_header_value)
                .to_request();
            let resp = call_service(&mut app, req).await;
            assert_eq!(resp.status(), StatusCode::OK);
        }
    }
}

How can I fix this problem ?

[robjtede]

Would need to see code in the middleware to help properly. Have you stepped through the backtrace lines and searched for unwraps and expects in your code?

[DimitriTimoz]

yes sorry,

use crate::prelude::*;
use std::{
    future::{ready, Future, Ready}, pin::Pin, task::{Context, Poll}
};

use actix_web::{
    dev::{Service, ServiceRequest, ServiceResponse, Transform}, http::HeaderMap, Error, HttpResponse
};

#[derive(Debug, Clone)]
pub struct Firewall;

#[doc(hidden)]
pub struct FirewallService<S> {
    service: S,
    enabled: bool,
}

fn check_header_api_key(headers: &HeaderMap) -> Result<Option<u64>, Error> {
    let key = match headers.get("Api-Key") {
        Some(V) => V.to_str(),
        None => return Ok(None),
    };
    let key = key.unwrap();
    if key.contains('-') {
        let args: Vec<_> = key.split('-').collect();
        if args.len() == 3 {
            let api_key = args[0].parse::<u64>();
            let counter = args[1].parse::<u64>();
            let code = args[2].parse::<u64>();
            if api_key.is_err() || counter.is_err() || code.is_err() {
                Err(HttpResponse::BadRequest().body("The api key format is wrong").into())
            } else if let Ok(v) = DATABASE.check_api_key(api_key.unwrap(), counter.unwrap(), code.unwrap()) {
                Ok(Some(v))
            } else {
                Ok(None)
            }
        } else {
            Err(HttpResponse::BadRequest().body("The api key format is wrong").into())
        }
    } else {
        Err(HttpResponse::BadRequest().body("The api key format is wrong").into())
    }
}

impl<S, B> Service for FirewallService<S>
where
    S: Service<Request = ServiceRequest, Response = ServiceResponse<B>, Error = actix_web::Error>,
    S::Future: 'static,
    B: 'static,
{
    type Error = Error;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Error>>>>;
    type Request = ServiceRequest;
    type Response = ServiceResponse<B>;

    fn poll_ready(&mut self, ctx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { self.service.poll_ready(ctx) }

    fn call(&mut self, req: Self::Request) -> Self::Future {
        println!("request is passing through the filter middleware");

        // get mut HttpRequest from ServiceRequest
        let (request, pl) = req.into_parts();
        let headers = request.headers();

        match check_header_api_key(headers) {
            Ok(v) => {
                request.extensions_mut().insert(v);
                if let Some(uid) = v {
                    if !anti_span::is_authorised(uid) {
                        error!("User: {} is trying to access to the api but is rate limited.", uid);
                        println!("too many");
                        return Box::pin(ready(Err(HttpResponse::TooManyRequests()
                            .body("Too many requests, you are temporarily rate limited.")
                            .into())));
                    }
                }

                match ServiceRequest::from_parts(request, pl) {
                    Ok(req) => Box::pin(self.service.call(req)),
                    Err(_) => Box::pin(ready(Err(Error::from(())))),
                }
            },
            Err(e) => Box::pin(ready(Err(e))),
        }
    }
}

#[derive(Clone, Debug)]
pub struct Auth {
    enabled: bool,
}

impl Auth {
    pub fn enabled() -> Self { Self { enabled: true } }

    pub fn disabled() -> Self { Self { enabled: false } }
}

impl<S, B> Transform<S> for Auth
where
    S: Service<Request = ServiceRequest, Response = ServiceResponse<B>, Error = actix_web::Error>,
    S::Future: 'static,
    B: 'static,
{
    type Error = Error;
    type Future = Ready<Result<Self::Transform, Self::InitError>>;
    type InitError = ();
    type Request = ServiceRequest;
    type Response = ServiceResponse<B>;
    type Transform = FirewallService<S>;

    fn new_transform(&self, service: S) -> Self::Future {
        ready(Ok(FirewallService {
            service,
            enabled: self.enabled,
        }))
    }
}

mod anti_span {
    const THRESHOLD: u64 = 600;
    const HIT: u64 = 50;

    use crate::database::FirewallEntry;

    use super::*;

    #[inline]
    fn calcul_heat(time_difference: u32, heat: u64) -> u64 {
        let heat = heat + HIT;
        heat * (((heat - 1) / heat) as u64).pow(time_difference / 10)
    }
    pub(crate) fn is_authorised(uid: u64) -> bool {
        let timestamp = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_secs();
        if let Ok(v) = DATABASE.get_firewall_entry(uid) {
            let mut time_difference = timestamp - v.last_request_time;
            if time_difference > 4294967295 {
                time_difference = 4294967295;
            }
            if v.heat > THRESHOLD {
                let new_entry = FirewallEntry {
                    heat: calcul_heat(time_difference as u32, v.heat),
                    last_request_time: timestamp,
                    recidivism_counter: v.recidivism_counter,
                };
                if let Err(e) = DATABASE.set_firewall_entry(uid, new_entry) {
                    error!(
                        "Impossible to save the new FirewallEntry for the user id: {}. Message: {}",
                        uid, e
                    )
                }
                return false;
            }
            let mut heat = calcul_heat(time_difference as u32, v.heat);
            let mut recidivism_counter = v.recidivism_counter;

            if heat > THRESHOLD {
                recidivism_counter += 1;
                heat = (THRESHOLD * 2) * (recidivism_counter as u64).pow(2);
            }
            let new_entry = FirewallEntry {
                heat,
                last_request_time: timestamp,
                recidivism_counter,
            };
            if let Err(e) = DATABASE.set_firewall_entry(uid, new_entry) {
                error!(
                    "Impossible to save the new FirewallEntry for the user id: {}. Message: {}",
                    uid, e
                )
            }
            true
        } else {
            if let Err(e) = DATABASE.create_firewall_entry(uid) {
                error!(
                    "Impossible to add the new FirewallEntry for the user id: {}. Message: {}",
                    uid, e
                )
            } else {
                let new_entry = FirewallEntry {
                    heat: HIT,
                    last_request_time: timestamp,
                    recidivism_counter: 0,
                };
                if let Err(e) = DATABASE.set_firewall_entry(uid, new_entry) {
                    error!(
                        "Impossible to save the new FirewallEntry for the user id: {}. Message: {}",
                        uid, e
                    )
                }
            }

            true
        }
    }
}```