use crate::service::{Layer, Service};
use http::header::HeaderName;
use http::{HeaderValue, Request, Response, Version};
use witchcraft_server_config::install::InstallConfig;
#[allow(clippy::declare_interior_mutable_const)]
const KEEP_ALIVE: HeaderName = HeaderName::from_static("keep-alive");
pub struct KeepAliveHeaderLayer {
value: Option<HeaderValue>,
}
impl KeepAliveHeaderLayer {
pub fn new(config: &InstallConfig) -> Self {
KeepAliveHeaderLayer {
value: config
.server()
.idle_connection_timeout()
.map(|d| HeaderValue::try_from(format!("timeout={}", d.as_secs())).unwrap()),
}
}
}
impl<S> Layer<S> for KeepAliveHeaderLayer {
type Service = KeepAliveHeaderService<S>;
fn layer(self, inner: S) -> Self::Service {
KeepAliveHeaderService {
inner,
value: self.value,
}
}
}
pub struct KeepAliveHeaderService<S> {
inner: S,
value: Option<HeaderValue>,
}
impl<S, B1, B2> Service<Request<B1>> for KeepAliveHeaderService<S>
where
S: Service<Request<B1>, Response = Response<B2>> + Sync,
B1: Send,
{
type Response = S::Response;
async fn call(&self, req: Request<B1>) -> Self::Response {
let value = match req.version() {
Version::HTTP_09 | Version::HTTP_10 | Version::HTTP_11 => self.value.clone(),
_ => None,
};
let mut response = self.inner.call(req).await;
if let Some(value) = value {
response.headers_mut().insert(KEEP_ALIVE, value);
}
response
}
}