//! Types for verifying requests with Actix Web

use crate::{Config, PrepareVerifyError, SignatureVerify};
use actix_web::{
    body::MessageBody,
    dev::{Payload, Service, ServiceRequest, ServiceResponse, Transform},
    http::StatusCode,
    Error, FromRequest, HttpMessage, HttpRequest, HttpResponse, ResponseError,
};
use futures_util::future::LocalBoxFuture;
use std::{
    collections::HashSet,
    future::{ready, Ready},
    task::{Context, Poll},
};
use tracing::{debug, Span};
use tracing_error::SpanTrace;
use tracing_futures::Instrument;

#[derive(Clone, Debug)]
/// A marker type that can be used to guard routes when the signature middleware is set to
/// 'optional'
pub struct SignatureVerified(String);

impl SignatureVerified {
    /// Return the Key ID used to verify the request
    ///
    /// It might be important for an application to verify that the payload being processed indeed
    /// belongs to the owner of the key used to sign the request.
    pub fn key_id(&self) -> &str {
        &self.0
    }
}

#[derive(Clone, Debug)]
/// The Verify signature middleware
///
/// ```rust,ignore
/// let middleware = VerifySignature::new(MyVerifier::new(), Config::default())
///     .authorization()
///     .optional();
///
/// HttpServer::new(move || {
///     App::new()
///         .wrap(middleware.clone())
///         .route("/protected", web::post().to(|_: SignatureVerified| "Verified Authorization Header"))
///         .route("/unprotected", web::post().to(|| "No verification required"))
/// })
/// ```
pub struct VerifySignature<T>(T, Config, HeaderKind, bool);

#[derive(Clone, Debug)]
#[doc(hidden)]
pub struct VerifyMiddleware<T, S>(S, Config, HeaderKind, bool, T);

#[derive(Copy, Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
enum HeaderKind {
    Authorization,
    Signature,
}

impl std::fmt::Display for HeaderKind {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Authorization => {
                write!(f, "Authorization")
            }
            Self::Signature => {
                write!(f, "Signature")
            }
        }
    }
}

#[derive(Clone, Debug, thiserror::Error)]
#[error("Failed to verify http signature, {kind}")]
#[doc(hidden)]
pub struct VerifyError {
    context: SpanTrace,
    kind: VerifyErrorKind,
}

#[derive(Clone, Debug, thiserror::Error)]
enum VerifyErrorKind {
    #[error("Signature or Authorization header is missing")]
    MissingSignature,

    #[error("{0}")]
    ExpiredSignature(String),

    #[error("Signature field could not be parsed")]
    ParseField(&'static str),

    #[error("Signature is not a valid string")]
    ParseSignature,

    #[error("Signature is invalid")]
    Validate,

    #[error("Request extension not present")]
    Extension,

    #[error("Required headers are missing")]
    MissingHeader(HashSet<String>),
}

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

impl<T> VerifySignature<T>
where
    T: SignatureVerify,
{
    /// Create a new middleware for verifying HTTP Signatures. A type implementing
    /// [`SignatureVerify`] is required, as well as a Config
    ///
    /// By default, this middleware expects to verify Signature headers, and requires the presence
    /// of the header
    pub fn new(verify_signature: T, config: Config) -> Self {
        VerifySignature(verify_signature, config, HeaderKind::Signature, false)
    }

    /// Verify Authorization headers instead of Signature headers
    pub fn authorization(self) -> Self {
        VerifySignature(self.0, self.1, HeaderKind::Authorization, self.3)
    }

    /// Mark the presence of a Signature or Authorization header as optional
    ///
    /// If a header is present, it will be verified, but if there is not one present, the request
    /// is passed through. This can be used to set a global middleware, and then guard each route
    /// handler with the [`SignatureVerified`] type.
    pub fn optional(self) -> Self {
        VerifySignature(self.0, self.1, self.2, true)
    }
}

impl<T, S, B> VerifyMiddleware<T, S>
where
    T: SignatureVerify + Clone + 'static,
    T::Future: 'static,
    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error> + 'static,
    B: MessageBody + 'static,
{
    fn handle(
        &self,
        span: Span,
        req: ServiceRequest,
    ) -> LocalBoxFuture<'static, Result<ServiceResponse<B>, Error>> {
        let res = self.1.begin_verify(
            req.method(),
            req.uri().path_and_query(),
            req.headers().clone(),
        );

        let unverified = match res {
            Ok(unverified) => unverified,
            Err(PrepareVerifyError::Expired(reason)) => {
                return Box::pin(ready(Err(VerifyError::new(
                    &span,
                    VerifyErrorKind::ExpiredSignature(reason),
                )
                .into())));
            }
            Err(PrepareVerifyError::Missing) => {
                return Box::pin(ready(Err(VerifyError::new(
                    &span,
                    VerifyErrorKind::MissingSignature,
                )
                .into())));
            }
            Err(PrepareVerifyError::ParseField(field)) => {
                return Box::pin(ready(Err(VerifyError::new(
                    &span,
                    VerifyErrorKind::ParseField(field),
                )
                .into())));
            }
            Err(PrepareVerifyError::Header(_)) => {
                return Box::pin(ready(Err(VerifyError::new(
                    &span,
                    VerifyErrorKind::ParseSignature,
                )
                .into())));
            }
            Err(PrepareVerifyError::Required(mut req)) => {
                return Box::pin(ready(Err(VerifyError::new(
                    &span,
                    VerifyErrorKind::MissingHeader(req.take_headers()),
                )
                .into())));
            }
        };

        let algorithm = unverified.algorithm().cloned();
        let key_id = unverified.key_id().to_owned();

        let f1 = unverified.verify(|signature, signing_string| {
            let fut = span.in_scope(|| {
                self.4.clone().signature_verify(
                    algorithm,
                    key_id.clone(),
                    signature.to_string(),
                    signing_string.to_string(),
                )
            });

            fut.instrument(span.clone())
        });

        req.extensions_mut().insert(SignatureVerified(key_id));

        let f2 = self.0.call(req);

        Box::pin(async move {
            let span = span;

            if f1.await? {
                f2.await
            } else {
                Err(VerifyError::new(&span, VerifyErrorKind::Validate).into())
            }
        })
    }
}

impl HeaderKind {
    pub fn is_authorization(self) -> bool {
        HeaderKind::Authorization == self
    }

    pub fn is_signature(self) -> bool {
        HeaderKind::Signature == self
    }
}

impl FromRequest for SignatureVerified {
    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>()
            .cloned()
            .ok_or_else(|| VerifyError::new(&Span::current(), VerifyErrorKind::Extension));

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

        ready(res)
    }
}

impl<T, S, B> Transform<S, ServiceRequest> for VerifySignature<T>
where
    T: SignatureVerify + Clone + '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(
            service,
            self.1.clone(),
            self.2,
            self.3,
            self.0.clone(),
        )))
    }
}

impl<T, S, B> Service<ServiceRequest> for VerifyMiddleware<T, S>
where
    T: SignatureVerify + Clone + '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, req: ServiceRequest) -> Self::Future {
        let span = tracing::info_span!(
            "Signature Verification",
            signature.kind = tracing::field::Empty,
            signature.expected_kind = tracing::field::display(&self.2),
            signature.optional = tracing::field::display(&self.3),
        );
        let authorization = req.headers().get("Authorization").is_some();
        let signature = req.headers().get("Signature").is_some();

        if authorization {
            span.record("signature.kind", &tracing::field::display("Authorization"));

            if self.2.is_authorization() {
                return self.handle(span, req);
            }
        } else if signature {
            span.record("signature.kind", &tracing::field::display("Signature"));

            if self.2.is_signature() {
                return self.handle(span, req);
            }
        } else {
            span.record("signature.kind", &tracing::field::display("None"));

            if self.3 {
                return Box::pin(self.0.call(req));
            }
        }

        Box::pin(ready(Err(VerifyError::new(
            &span,
            VerifyErrorKind::MissingSignature,
        )
        .into())))
    }
}

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

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