use http::{header, HeaderValue, Response, StatusCode, Uri};
use std::{
convert::{Infallible, TryFrom},
fmt,
future::Future,
marker::PhantomData,
pin::Pin,
task::{Context, Poll},
};
use tower_service::Service;
pub struct Redirect<ResBody> {
status_code: StatusCode,
location: HeaderValue,
_marker: PhantomData<fn() -> ResBody>,
}
impl<ResBody> Redirect<ResBody> {
pub fn temporary(uri: Uri) -> Self {
Self::with_status_code(StatusCode::TEMPORARY_REDIRECT, uri)
}
pub fn permanent(uri: Uri) -> Self {
Self::with_status_code(StatusCode::PERMANENT_REDIRECT, uri)
}
pub fn with_status_code(status_code: StatusCode, uri: Uri) -> Self {
assert!(
status_code.is_redirection(),
"not a redirection status code"
);
Self {
status_code,
location: HeaderValue::try_from(uri.to_string())
.expect("URI isn't a valid header value"),
_marker: PhantomData,
}
}
}
impl<R, ResBody> Service<R> for Redirect<ResBody>
where
ResBody: Default,
{
type Response = Response<ResBody>;
type Error = Infallible;
type Future = ResponseFuture<ResBody>;
#[inline]
fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn call(&mut self, _req: R) -> Self::Future {
ResponseFuture {
status_code: self.status_code,
location: Some(self.location.clone()),
_marker: PhantomData,
}
}
}
impl<ResBody> fmt::Debug for Redirect<ResBody> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Redirect")
.field("status_code", &self.status_code)
.field("location", &self.location)
.finish()
}
}
impl<ResBody> Clone for Redirect<ResBody> {
fn clone(&self) -> Self {
Self {
status_code: self.status_code,
location: self.location.clone(),
_marker: PhantomData,
}
}
}
#[derive(Debug)]
pub struct ResponseFuture<ResBody> {
location: Option<HeaderValue>,
status_code: StatusCode,
_marker: PhantomData<fn() -> ResBody>,
}
impl<ResBody> Future for ResponseFuture<ResBody>
where
ResBody: Default,
{
type Output = Result<Response<ResBody>, Infallible>;
fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
let mut res = Response::default();
*res.status_mut() = self.status_code;
res.headers_mut()
.insert(header::LOCATION, self.location.take().unwrap());
Poll::Ready(Ok(res))
}
}