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

//! Types for setting up Digest middleware verification

use super::{DigestPart, DigestVerify};
use actix_web::{
    dev::{MessageBody, Payload, Service, ServiceRequest, ServiceResponse, Transform},
    error::PayloadError,
    http::{header::HeaderValue, StatusCode},
    web, FromRequest, HttpMessage, HttpRequest, HttpResponse, ResponseError,
};
use futures_util::{
    future::LocalBoxFuture,
    stream::{Stream, StreamExt},
};
use std::{
    future::{ready, Ready},
    pin::Pin,
    rc::Rc,
    task::{Context, Poll},
};
use tracing::{debug, Instrument, 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>(bool, T);

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

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

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

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(true, verify_digest)
    }

    /// 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(false, self.1)
    }
}

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

    fn from_request(req: &HttpRequest, _: &mut Payload) -> Self::Future {
        let res = req
            .extensions()
            .get::<Self>()
            .copied()
            .ok_or_else(|| VerifyError::new(&Span::current(), VerifyErrorKind::Extension));

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

        ready(res)
    }
}

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

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

impl<T, S, B> Service<ServiceRequest> for VerifyMiddleware<T, S>
where
    T: DigestVerify + Clone + Send + 'static,
    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = actix_web::Error> + 'static,
    S::Error: 'static,
    B: MessageBody + '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.1),
        );

        let verifier = self.2.clone();
        let svc = Rc::clone(&self.0);
        let required = self.1;

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

                let payload = req.take_payload();

                let payload = verify_payload(vec, verifier, payload)
                    .instrument(span)
                    .await?;

                req.set_payload(payload);
                req.extensions_mut().insert(DigestVerified);

                svc.call(req).await
            } else if required {
                Err(VerifyError::new(&span, VerifyErrorKind::MissingDigest).into())
            } else {
                svc.call(req).await
            }
        })
    }
}

#[tracing::instrument(name = "Verify Payload", skip(verify_digest, payload))]
async fn verify_payload<T>(
    vec: Vec<DigestPart>,
    mut verify_digest: T,
    mut payload: Payload,
) -> Result<Payload, actix_web::Error>
where
    T: DigestVerify + Clone + Send + 'static,
{
    let mut payload_vec = vec![];
    while let Some(res) = payload.next().await {
        let bytes = res?;
        let bytes2 = bytes.clone();
        verify_digest = web::block(move || {
            verify_digest.update(bytes2.as_ref());
            Ok(verify_digest) as Result<T, VerifyError>
        })
        .await??;

        payload_vec.push(Ok(bytes));
    }

    let verified =
        web::block(move || Ok(verify_digest.verify(&vec)) as Result<_, VerifyError>).await??;

    if verified {
        let stream: Pin<Box<dyn Stream<Item = Result<web::Bytes, PayloadError>>>> =
            Box::pin(futures_util::stream::iter(payload_vec));

        Ok(Payload::Stream(stream))
    } else {
        Err(VerifyError::new(&Span::current(), VerifyErrorKind::Verify).into())
    }
}

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())
    }
}