use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use bytes::Bytes;
use http::{Request, Response, StatusCode};
use tower::Service;
use crate::http::{Body, BoxError, HttpService, full_body};
#[derive(Clone)]
pub struct SetResponse {
status: StatusCode,
body: Bytes,
}
impl SetResponse {
pub fn ok(body: impl Into<Bytes>) -> Self {
Self {
status: StatusCode::OK,
body: body.into(),
}
}
pub fn new(status: StatusCode, body: impl Into<Bytes>) -> Self {
Self {
status,
body: body.into(),
}
}
}
impl tower::Layer<HttpService> for SetResponse {
type Service = SetResponseService;
fn layer(&self, _inner: HttpService) -> Self::Service {
SetResponseService {
status: self.status,
body: self.body.clone(),
}
}
}
pub struct SetResponseService {
status: StatusCode,
body: Bytes,
}
impl Service<Request<Body>> for SetResponseService {
type Response = Response<Body>;
type Error = BoxError;
type Future = Pin<Box<dyn Future<Output = Result<Response<Body>, BoxError>> + Send>>;
fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn call(&mut self, _req: Request<Body>) -> Self::Future {
let status = self.status;
let body = self.body.clone();
Box::pin(async move {
Ok(Response::builder()
.status(status)
.body(full_body(body))
.unwrap())
})
}
}