use std::collections::HashMap;
use std::pin::Pin;
use std::sync::Arc;
use futures_util::future::{self, FutureExt};
use http::{Response, StatusCode};
use log::trace;
use crate::handler::HandlerFuture;
use crate::state::{request_id, State};
use crate::helpers::http::Body;
use crate::router::response::extender::ResponseExtender;
#[derive(Clone)]
pub struct ResponseFinalizer {
data: Arc<HashMap<StatusCode, Box<dyn ResponseExtender<Body> + Send + Sync>>>,
}
pub struct ResponseFinalizerBuilder {
data: HashMap<StatusCode, Box<dyn ResponseExtender<Body> + Send + Sync>>,
}
impl ResponseFinalizerBuilder {
pub(in crate::router) fn new() -> Self {
let handlers = HashMap::new();
ResponseFinalizerBuilder { data: handlers }
}
pub fn add(
&mut self,
status_code: StatusCode,
extender: Box<dyn ResponseExtender<Body> + Send + Sync>,
) {
trace!(" adding response extender for {}", status_code);
self.data.insert(status_code, extender);
}
pub fn finalize(self) -> ResponseFinalizer {
ResponseFinalizer {
data: Arc::new(self.data),
}
}
}
impl ResponseFinalizer {
pub fn finalize(&self, mut state: State, mut res: Response<Body>) -> Pin<Box<HandlerFuture>> {
match self.data.get(&res.status()) {
Some(extender) => {
trace!(
"[{}] invoking {} response extender",
request_id(&state),
res.status()
);
extender.extend(&mut state, &mut res);
}
None => {
trace!(
"[{}] no response extender for {}",
request_id(&state),
res.status()
);
}
}
future::ok((state, res)).boxed()
}
}