http-signature-normalization/actix/src/digest/middleware.rs

//! Types for setting up Digest middleware verification

use crate::{Canceled, DefaultSpawner, Spawn};

use super::{DigestPart, DigestVerify};
use actix_web::{
    body::MessageBody,
    dev::{Payload, Service, ServiceRequest, ServiceResponse, Transform},
    error::PayloadError,
    http::{header::HeaderValue, StatusCode},
    web, FromRequest, HttpMessage, HttpRequest, HttpResponse, ResponseError,
};
use futures_core::{future::LocalBoxFuture, Stream};
use std::{
    future::{ready, Ready},
    pin::Pin,
    task::{Context, Poll},
};
use streem::{from_fn::Yielder, IntoStreamer};
use tokio::sync::{mpsc, oneshot};
use tracing::{debug, Span};
use tracing_error::SpanTrace;

#[derive(Copy, Clone, Debug)]
/// A type implementing FromRequest that can be used in route handler to guard for verified
/// digests
///
/// This is only required when the [`VerifyDigest`] middleware is set to optional
pub struct DigestVerified;

#[derive(Clone, Debug)]
/// The VerifyDigest middleware
///
/// ```rust,ignore
/// let middleware = VerifyDigest::new(MyVerify::new())
///     .optional();
///
/// HttpServer::new(move || {
///     App::new()
///         .wrap(middleware.clone())
///         .route("/protected", web::post().to(|_: DigestVerified| "Verified Digest Header"))
///         .route("/unprotected", web::post().to(|| "No verification required"))
/// })
/// ```
pub struct VerifyDigest<T, Spawner = DefaultSpawner>(Spawner, bool, T);

#[doc(hidden)]
pub struct VerifyMiddleware<T, Spawner, S>(S, Spawner, bool, T);

#[derive(Debug, thiserror::Error)]
#[error("Error verifying digest")]
#[doc(hidden)]
pub struct VerifyError {
    context: String,
    kind: VerifyErrorKind,
}

impl VerifyError {
    fn new(span: &Span, kind: VerifyErrorKind) -> Self {
        span.in_scope(|| VerifyError {
            context: SpanTrace::capture().to_string(),
            kind,
        })
    }
}

#[derive(Debug, thiserror::Error)]
enum VerifyErrorKind {
    #[error("Missing request extension")]
    Extension,

    #[error("Digest header missing")]
    MissingDigest,

    #[error("Digest header is empty")]
    Empty,

    #[error("Failed to verify digest")]
    Verify,

    #[error("Payload dropped. If this was unexpected, it could be that the payload isn't required in the route this middleware is guarding")]
    Dropped,
}

struct RxStream<T>(mpsc::Receiver<T>);

impl<T> Stream for RxStream<T> {
    type Item = Result<T, PayloadError>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        Pin::new(&mut self.0).poll_recv(cx).map(|opt| opt.map(Ok))
    }
}

impl<T> VerifyDigest<T>
where
    T: DigestVerify + Clone,
{
    /// Produce a new VerifyDigest with a user-provided [`Digestverify`] type
    pub fn new(verify_digest: T) -> Self {
        VerifyDigest(DefaultSpawner, true, verify_digest)
    }
}

impl<T, Spawner> VerifyDigest<T, Spawner>
where
    T: DigestVerify + Clone,
{
    /// Set the spawner used for verifying bytes in the request
    ///
    /// By default this value uses `actix_web::web::block`
    pub fn spawner<NewSpawner>(self, spawner: NewSpawner) -> VerifyDigest<T, NewSpawner>
    where
        NewSpawner: Spawn,
    {
        VerifyDigest(spawner, self.1, self.2)
    }

    /// Mark verifying the Digest as optional
    ///
    /// If a digest is present in the request, it will be verified, but it is not required to be
    /// present
    pub fn optional(self) -> Self {
        VerifyDigest(self.0, false, self.2)
    }
}

struct VerifiedReceiver {
    rx: Option<oneshot::Receiver<()>>,
}

impl FromRequest for DigestVerified {
    type Error = VerifyError;
    type Future = LocalBoxFuture<'static, Result<Self, Self::Error>>;

    fn from_request(req: &HttpRequest, _: &mut Payload) -> Self::Future {
        let res = req
            .extensions_mut()
            .get_mut::<VerifiedReceiver>()
            .and_then(|r| r.rx.take())
            .ok_or_else(|| VerifyError::new(&Span::current(), VerifyErrorKind::Extension));

        if res.is_err() {
            debug!("Failed to fetch DigestVerified from request");
        }

        Box::pin(async move {
            res?.await
                .map_err(|_| VerifyError::new(&Span::current(), VerifyErrorKind::Dropped))
                .map(|()| DigestVerified)
        })
    }
}

impl<T, Spawner, S, B> Transform<S, ServiceRequest> for VerifyDigest<T, Spawner>
where
    T: DigestVerify + Clone + Send + 'static,
    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = actix_web::Error> + 'static,
    S::Error: 'static,
    B: MessageBody + 'static,
    Spawner: Spawn + Clone + 'static,
{
    type Response = ServiceResponse<B>;
    type Error = actix_web::Error;
    type Transform = VerifyMiddleware<T, Spawner, S>;
    type InitError = ();
    type Future = Ready<Result<Self::Transform, Self::InitError>>;

    fn new_transform(&self, service: S) -> Self::Future {
        ready(Ok(VerifyMiddleware(
            service,
            self.0.clone(),
            self.1,
            self.2.clone(),
        )))
    }
}

impl<T, Spawner, S, B> Service<ServiceRequest> for VerifyMiddleware<T, Spawner, S>
where
    T: DigestVerify + Clone + Send + 'static,
    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = actix_web::Error> + 'static,
    S::Error: 'static,
    B: MessageBody + 'static,
    Spawner: Spawn + Clone + 'static,
{
    type Response = ServiceResponse<B>;
    type Error = actix_web::Error;
    type Future = LocalBoxFuture<'static, Result<Self::Response, Self::Error>>;

    fn poll_ready(&self, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
        self.0.poll_ready(cx)
    }

    fn call(&self, mut req: ServiceRequest) -> Self::Future {
        let span = tracing::info_span!(
            "Verify digest",
            digest.required = tracing::field::display(&self.2),
        );

        if let Some(digest) = req.headers().get("Digest") {
            let vec = match parse_digest(digest) {
                Some(vec) => vec,
                None => {
                    return Box::pin(ready(Err(
                        VerifyError::new(&span, VerifyErrorKind::Empty).into()
                    )));
                }
            };

            let spawner = self.1.clone();
            let digest = self.3.clone();
            let (verify_tx, verify_rx) = oneshot::channel();

            let payload = req.take_payload();
            let payload: Pin<Box<dyn Stream<Item = Result<web::Bytes, PayloadError>> + 'static>> =
                Box::pin(streem::try_from_fn(|yielder| async move {
                    verify_payload(yielder, spawner, vec, digest, payload, verify_tx).await
                }));
            req.set_payload(payload.into());

            req.extensions_mut().insert(VerifiedReceiver {
                rx: Some(verify_rx),
            });

            Box::pin(self.0.call(req))
        } else if self.2 {
            Box::pin(ready(Err(VerifyError::new(
                &span,
                VerifyErrorKind::MissingDigest,
            )
            .into())))
        } else {
            Box::pin(self.0.call(req))
        }
    }
}

fn canceled_error(error: Canceled) -> PayloadError {
    PayloadError::Io(std::io::Error::new(std::io::ErrorKind::Other, error))
}

fn verified_error(error: VerifyError) -> PayloadError {
    PayloadError::Io(std::io::Error::new(std::io::ErrorKind::Other, error))
}

async fn verify_payload<T, Spawner>(
    yielder: Yielder<Result<web::Bytes, PayloadError>>,
    spawner: Spawner,
    vec: Vec<DigestPart>,
    mut verify_digest: T,
    payload: Payload,
    verify_tx: oneshot::Sender<()>,
) -> Result<(), PayloadError>
where
    T: DigestVerify + Clone + Send + 'static,
    Spawner: Spawn,
{
    let mut payload = payload.into_streamer();

    let mut error = None;

    while let Some(bytes) = payload.try_next().await? {
        if error.is_none() {
            let bytes2 = bytes.clone();
            let mut verify_digest2 = verify_digest.clone();

            let task = spawner.spawn_blocking(move || {
                verify_digest2.update(bytes2.as_ref());
                Ok(verify_digest2) as Result<T, VerifyError>
            });

            yielder.yield_ok(bytes).await;

            match task.await {
                Ok(Ok(digest)) => {
                    verify_digest = digest;
                }
                Ok(Err(e)) => {
                    error = Some(verified_error(e));
                }
                Err(e) => {
                    error = Some(canceled_error(e));
                }
            }
        } else {
            yielder.yield_ok(bytes).await;
        }
    }

    if let Some(error) = error {
        return Err(error);
    }

    let verified = spawner
        .spawn_blocking(move || Ok(verify_digest.verify(&vec)) as Result<_, VerifyError>)
        .await
        .map_err(canceled_error)?
        .map_err(verified_error)?;

    if verified {
        if verify_tx.send(()).is_err() {
            debug!("handler dropped");
        }

        Ok(())
    } else {
        Err(verified_error(VerifyError::new(
            &Span::current(),
            VerifyErrorKind::Verify,
        )))
    }
}

fn parse_digest(h: &HeaderValue) -> Option<Vec<DigestPart>> {
    let h = h.to_str().ok()?.split(';').next()?;
    let v: Vec<_> = h
        .split(',')
        .filter_map(|p| {
            let mut iter = p.splitn(2, '=');
            iter.next()
                .and_then(|alg| iter.next().map(|value| (alg, value)))
        })
        .map(|(alg, value)| DigestPart {
            algorithm: alg.to_owned(),
            digest: value.to_owned(),
        })
        .collect();

    if v.is_empty() {
        None
    } else {
        Some(v)
    }
}

impl ResponseError for VerifyError {
    fn status_code(&self) -> StatusCode {
        StatusCode::BAD_REQUEST
    }

    fn error_response(&self) -> HttpResponse {
        HttpResponse::new(self.status_code())
    }
}