apollo-errors 0.7.0

//! Error service middleware

use super::{NegotiationConfig, body::ErrorBody, future::ErrorServiceFuture};
use bytes::Buf;
use http::{Request, Response, header};
use http_body::Body;
use std::{
    marker::PhantomData,
    task::{Context, Poll},
};
use tower::{BoxError, Service};

/// Service that wraps another service and renders errors as HTTP responses
#[derive(Debug)]
pub struct ErrorService<S> {
    inner: S,
    config: NegotiationConfig,
    poll_ready_error: Option<BoxError>,
}

impl<S: Clone> Clone for ErrorService<S> {
    fn clone(&self) -> Self {
        Self {
            inner: self.inner.clone(),
            config: self.config.clone(),
            // Don't clone the error - it will be consumed on first call anyway
            poll_ready_error: None,
        }
    }
}

impl<S> ErrorService<S> {
    /// Create a new error service
    pub(super) fn new(inner: S, config: NegotiationConfig) -> Self {
        Self {
            inner,
            config,
            poll_ready_error: None,
        }
    }
}

impl<S, ReqBody, ResBody> Service<Request<ReqBody>> for ErrorService<S>
where
    S: Service<Request<ReqBody>, Response = Response<ResBody>>,
    S::Error: Into<BoxError>,
    ResBody: Body,
    ResBody::Data: Buf,
    ResBody::Error: Into<BoxError>,
{
    type Response = Response<ErrorBody<ResBody>>;
    type Error = std::convert::Infallible;
    type Future = ErrorServiceFuture<S::Future, ResBody>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        match self.inner.poll_ready(cx) {
            Poll::Ready(Ok(())) => Poll::Ready(Ok(())),
            Poll::Ready(Err(e)) => {
                // Store the error to return it in the future
                self.poll_ready_error = Some(e.into());
                Poll::Ready(Ok(()))
            }
            Poll::Pending => Poll::Pending,
        }
    }

    fn call(&mut self, req: Request<ReqBody>) -> Self::Future {
        // Extract Accept header for content negotiation
        let accept = req
            .headers()
            .get(header::ACCEPT)
            .and_then(|v| v.to_str().ok())
            .map(String::from);

        // Use mem::take to move the error out if present
        let poll_ready_error = std::mem::take(&mut self.poll_ready_error);

        ErrorServiceFuture {
            inner: self.inner.call(req),
            config: self.config.clone(),
            accept,
            poll_ready_error,
            _marker: PhantomData,
        }
    }
}