order.rs

use std::borrow::Cow;
use std::ops::{ControlFlow, Deref};
use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime};
use std::{fmt, slice};

use base64::prelude::{BASE64_URL_SAFE_NO_PAD, Engine};
#[cfg(feature = "rcgen")]
use rcgen::{CertificateParams, DistinguishedName, KeyPair};
use serde::Serialize;
use tokio::time::sleep;

use crate::account::AccountInner;
use crate::types::{
    Authorization, AuthorizationState, AuthorizationStatus, AuthorizedIdentifier, Challenge,
    ChallengeType, DeviceAttestation, Empty, FinalizeRequest, OrderState, OrderStatus, Problem,
};
use crate::{ChallengeStatus, Error, Key, crypto, nonce_from_response, retry_after};

/// An ACME order as described in RFC 8555 (section 7.1.3)
///
/// An order is created from an [`Account`][crate::Account] by calling
/// [`Account::new_order()`][crate::Account::new_order()]. The `Order` type represents the stable
/// identity of an order, while the [`Order::state()`] method gives you access to the current
/// state of the order according to the server.
///
/// <https://datatracker.ietf.org/doc/html/rfc8555#section-7.1.3>
pub struct Order {
    pub(crate) account: Arc<AccountInner>,
    pub(crate) nonce: Option<String>,
    pub(crate) retry_after: Option<SystemTime>,
    pub(crate) url: String,
    pub(crate) state: OrderState,
}

impl Order {
    /// Retrieve the authorizations for this order
    ///
    /// An order contains one authorization to complete per identifier in the order.
    /// After creating an order, you'll need to retrieve the authorizations so that
    /// you can set up a challenge response for each authorization.
    ///
    /// This method will retrieve the authorizations attached to this order if they have not
    /// been retrieved yet. If you have already consumed the stream generated by this method
    /// before, processing it again will not involve any network activity.
    pub fn authorizations(&mut self) -> Authorizations<'_> {
        Authorizations {
            inner: AuthStream {
                iter: self.state.authorizations.iter_mut(),
                nonce: &mut self.nonce,
                account: &self.account,
            },
        }
    }

    /// Generate a Certificate Signing Request for the order's identifiers and request finalization
    ///
    /// Uses the rcgen crate to generate a Certificate Signing Request (CSR) converting the order's
    /// identifiers to Subject Alternative Names, then calls [`Order::finalize_csr()`] with it.
    /// Returns the generated private key, serialized as PEM.
    ///
    /// After this succeeds, call [`Order::certificate()`] to retrieve the certificate chain once
    /// the order is in the appropriate state.
    #[cfg(feature = "rcgen")]
    pub async fn finalize(&mut self) -> Result<String, Error> {
        let mut names = Vec::with_capacity(self.state.authorizations.len());
        let mut identifiers = self.identifiers();
        while let Some(result) = identifiers.next().await {
            names.push(result?.to_string());
        }

        let mut params = CertificateParams::new(names).map_err(Error::from_rcgen)?;
        params.distinguished_name = DistinguishedName::new();
        let private_key = KeyPair::generate().map_err(Error::from_rcgen)?;
        let csr = params
            .serialize_request(&private_key)
            .map_err(Error::from_rcgen)?;

        self.finalize_csr(csr.der()).await?;
        Ok(private_key.serialize_pem())
    }

    /// Request a certificate from the given Certificate Signing Request (CSR)
    ///
    /// `csr_der` contains the CSR representation serialized in DER encoding. If you don't need
    /// custom certificate parameters, [`Order::finalize()`] can generate the CSR for you.
    ///
    /// After this succeeds, call [`Order::certificate()`] to retrieve the certificate chain once
    /// the order is in the appropriate state.
    pub async fn finalize_csr(&mut self, csr_der: &[u8]) -> Result<(), Error> {
        let rsp = self
            .account
            .post(
                Some(&FinalizeRequest::new(csr_der)),
                self.nonce.take(),
                &self.state.finalize,
            )
            .await?;

        self.nonce = nonce_from_response(&rsp);
        self.state = Problem::check::<OrderState>(rsp).await?;
        Ok(())
    }

    /// Get the certificate for this order
    ///
    /// If the cached order state is in `ready` or `processing` state, this will poll the server
    /// for the latest state. If the order is still in `processing` state after that, this will
    /// return `Ok(None)`. If the order is in `valid` state, this will attempt to retrieve
    /// the certificate from the server and return it as a `String`. If the order contains
    /// an error or ends up in any state other than `valid` or `processing`, return an error.
    pub async fn certificate(&mut self) -> Result<Option<String>, Error> {
        if matches!(self.state.status, OrderStatus::Processing) {
            let rsp = self
                .account
                .post(None::<&Empty>, self.nonce.take(), &self.url)
                .await?;
            self.nonce = nonce_from_response(&rsp);
            self.state = Problem::check::<OrderState>(rsp).await?;
        }

        if let Some(error) = &self.state.error {
            return Err(Error::Api(error.clone()));
        } else if self.state.status == OrderStatus::Processing {
            return Ok(None);
        } else if self.state.status != OrderStatus::Valid {
            return Err(Error::Str("invalid order state"));
        }

        let Some(cert_url) = &self.state.certificate else {
            return Err(Error::Str("no certificate URL found"));
        };

        let rsp = self
            .account
            .post(None::<&Empty>, self.nonce.take(), cert_url)
            .await?;

        self.nonce = nonce_from_response(&rsp);
        let body = Problem::from_response(rsp).await?;
        Ok(Some(
            String::from_utf8(body.to_vec())
                .map_err(|_| "unable to decode certificate as UTF-8")?,
        ))
    }

    /// Retrieve the identifiers for this order
    ///
    /// This method will retrieve the identifiers attached to the authorizations for this order
    /// if they have not been retrieved yet. If you have already consumed the stream generated
    /// by [`Order::authorizations()`], this will not involve any network activity.
    pub fn identifiers(&mut self) -> Identifiers<'_> {
        Identifiers {
            inner: AuthStream {
                iter: self.state.authorizations.iter_mut(),
                nonce: &mut self.nonce,
                account: &self.account,
            },
        }
    }

    /// Poll the order with the given [`RetryPolicy`]
    ///
    /// Yields the [`OrderStatus`] immediately if `Ready` or `Invalid`, or yields an
    /// [`Error::Timeout`] if the [`RetryPolicy::timeout`] has been reached.
    pub async fn poll_ready(&mut self, retries: &RetryPolicy) -> Result<OrderStatus, Error> {
        let mut retrying = retries.state();
        self.retry_after = None;
        loop {
            if let ControlFlow::Break(err) = retrying.wait(self.retry_after.take()).await {
                return Err(err);
            }

            let state = self.refresh().await?;
            if let Some(error) = &state.error {
                return Err(Error::Api(error.clone()));
            } else if let OrderStatus::Ready | OrderStatus::Invalid = state.status {
                return Ok(state.status);
            }
        }
    }

    /// Poll the certificate with the given [`RetryPolicy`]
    ///
    /// Yields the PEM encoded certificate chain for this order if the order state becomes
    /// `Valid`. The function keeps polling as long as the order state is `Processing`.
    /// An error is returned immediately: if the order state is `Invalid`, if polling runs
    /// into a timeout, or if the ACME CA suggest to retry at a later time.
    pub async fn poll_certificate(&mut self, retries: &RetryPolicy) -> Result<String, Error> {
        let mut retrying = retries.state();
        self.retry_after = None;
        loop {
            if let ControlFlow::Break(err) = retrying.wait(self.retry_after.take()).await {
                return Err(err);
            }

            let state = self.refresh().await?;
            if let Some(error) = &state.error {
                return Err(Error::Api(error.clone()));
            } else if let OrderStatus::Valid | OrderStatus::Invalid = state.status {
                return self
                    .certificate()
                    .await?
                    .ok_or(Error::Str("no certificates received from ACME CA"));
            }
        }
    }

    /// Refresh the current state of the order
    pub async fn refresh(&mut self) -> Result<&OrderState, Error> {
        let rsp = self
            .account
            .post(None::<&Empty>, self.nonce.take(), &self.url)
            .await?;

        self.nonce = nonce_from_response(&rsp);
        self.retry_after = retry_after(&rsp);
        self.state = Problem::check::<OrderState>(rsp).await?;
        Ok(&self.state)
    }

    /// Extract the URL and last known state from the `Order`
    pub fn into_parts(self) -> (String, OrderState) {
        (self.url, self.state)
    }

    /// Get the last known state of the order
    ///
    /// Call `refresh()` to get the latest state from the server.
    pub fn state(&mut self) -> &OrderState {
        &self.state
    }

    /// Get the URL of the order
    pub fn url(&self) -> &str {
        &self.url
    }
}

/// An stream-like interface that yields an [`Order`]'s authorizations
///
/// Call [`next()`] to get the next authorization in the order. If the order state
/// does not yet contain the state of the authorization, it will be fetched from the server.
///
/// [`next()`]: Authorizations::next()
pub struct Authorizations<'a> {
    inner: AuthStream<'a>,
}

impl Authorizations<'_> {
    /// Yield the next [`AuthorizationHandle`], fetching its state if we don't have it yet
    pub async fn next(&mut self) -> Option<Result<AuthorizationHandle<'_>, Error>> {
        let (url, state) = match self.inner.next().await? {
            Ok((url, state)) => (url, state),
            Err(err) => return Some(Err(err)),
        };

        Some(Ok(AuthorizationHandle {
            state,
            url,
            nonce: self.inner.nonce,
            account: self.inner.account,
        }))
    }
}

/// An stream-like interface that yields an [`Order`]'s identifiers
///
/// Call [`next()`] to get the next authorization in the order. If the order state
/// does not yet contain the state of the authorization, it will be fetched from the server.
///
/// [`next()`]: Identifiers::next()
pub struct Identifiers<'a> {
    inner: AuthStream<'a>,
}

impl<'a> Identifiers<'a> {
    /// Yield the next [`Identifier`][crate::Identifier], fetching the authorization's state if
    /// we don't have it yet
    pub async fn next(&mut self) -> Option<Result<AuthorizedIdentifier<'a>, Error>> {
        Some(match self.inner.next().await? {
            Ok((_, state)) => Ok(state.identifier()),
            Err(err) => Err(err),
        })
    }
}

struct AuthStream<'a> {
    iter: slice::IterMut<'a, Authorization>,
    nonce: &'a mut Option<String>,
    account: &'a AccountInner,
}

impl<'a> AuthStream<'a> {
    async fn next(&mut self) -> Option<Result<(&'a str, &'a mut AuthorizationState), Error>> {
        let authz = self.iter.next()?;
        if authz.state.is_none() {
            match self.account.get(self.nonce, &authz.url).await {
                Ok(state) => authz.state = Some(state),
                Err(e) => return Some(Err(e)),
            }
        }

        // The `unwrap()` here is safe: the code above will either set it to `Some` or yield
        // an error to the caller if it was `None` upon entering this method. I attempted to
        // use `Option::insert()` which did not pass the borrow checker for reasons that I
        // think have to do with the let scope extension that got fixed for 2024 edition.
        // For now, our MSRV does not allow the use of the new edition.
        let state = authz.state.as_mut().unwrap();
        Some(Ok((&authz.url, state)))
    }
}

/// An ACME authorization as described in RFC 8555 (section 7.1.4)
///
/// Authorizations are retrieved from an associated [`Order`] by calling
/// [`Order::authorizations()`]. This type dereferences to the underlying
/// [`AuthorizationState`] for easy access to the authorization's state.
///
/// For each authorization, you'll need to:
///
/// * Select which [`ChallengeType`] you want to complete
/// * Call [`AuthorizationHandle::challenge()`] to get a [`ChallengeHandle`]
/// * Use the `ChallengeHandle` to complete the authorization's challenge
///
/// <https://datatracker.ietf.org/doc/html/rfc8555#section-7.1.3>
pub struct AuthorizationHandle<'a> {
    state: &'a mut AuthorizationState,
    url: &'a str,
    nonce: &'a mut Option<String>,
    account: &'a AccountInner,
}

impl<'a> AuthorizationHandle<'a> {
    /// Refresh the current state of the authorization
    pub async fn refresh(&mut self) -> Result<&AuthorizationState, Error> {
        let rsp = self
            .account
            .post(None::<&Empty>, self.nonce.take(), self.url)
            .await?;

        *self.nonce = nonce_from_response(&rsp);
        *self.state = Problem::check::<AuthorizationState>(rsp).await?;
        Ok(self.state)
    }

    /// Deactivate a pending or valid authorization
    ///
    /// Returns the updated [`AuthorizationState`] if the deactivation was successful.
    /// If the authorization was not pending or valid, an error is returned.
    ///
    /// Once deactivated the authorization and associated challenges can not be updated
    /// further.
    ///
    /// This is useful when you want to cancel a pending authorization attempt you wish
    /// to abandon, or if you wish to revoke valid authorization for an identifier to
    /// force future uses of the identifier by the same ACME account to require
    /// re-verification with fresh authorizations/challenges.
    pub async fn deactivate(&mut self) -> Result<&AuthorizationState, Error> {
        if !matches!(
            self.state.status,
            AuthorizationStatus::Pending | AuthorizationStatus::Valid
        ) {
            return Err(Error::Other("authorization not pending or valid".into()));
        }

        #[derive(Serialize)]
        struct DeactivateRequest {
            status: AuthorizationStatus,
        }

        let rsp = self
            .account
            .post(
                Some(&DeactivateRequest {
                    status: AuthorizationStatus::Deactivated,
                }),
                self.nonce.take(),
                self.url,
            )
            .await?;

        *self.nonce = nonce_from_response(&rsp);
        *self.state = Problem::check::<AuthorizationState>(rsp).await?;
        match self.state.status {
            AuthorizationStatus::Deactivated => Ok(self.state),
            _ => Err(Error::Other(
                "authorization was not deactivated by ACME server".into(),
            )),
        }
    }

    /// Get a [`ChallengeHandle`] for the given `type`
    ///
    /// Yields an object to interact with the challenge for the given type, if available.
    pub fn challenge(&'a mut self, r#type: ChallengeType) -> Option<ChallengeHandle<'a>> {
        let challenge = self
            .state
            .challenges
            .iter()
            .find(|c| c.state.r#type() == r#type)?;
        Some(ChallengeHandle {
            identifier: self.state.identifier(),
            challenge,
            nonce: self.nonce,
            account: self.account,
        })
    }

    /// Get the URL of the authorization
    pub fn url(&self) -> &str {
        self.url
    }
}

impl Deref for AuthorizationHandle<'_> {
    type Target = AuthorizationState;

    fn deref(&self) -> &Self::Target {
        self.state
    }
}

/// Wrapper type for interacting with a [`Challenge`]'s state
///
/// For each challenge, you'll need to:
///
/// * Obtain the [`ChallengeHandle::key_authorization()`] for the challenge response
/// * Set up the challenge response in your infrastructure (details vary by challenge type)
/// * Call [`ChallengeHandle::set_ready()`] for that challenge after setup is complete
///
/// After the challenges have been set to ready, call [`Order::poll_ready()`] to wait until the
/// order is ready to be finalized (or to learn if it becomes invalid). Once it is ready, call
/// [`Order::finalize()`] to get the certificate.
///
/// Dereferences to the underlying [`Challenge`] for easy access to the challenge's state.
pub struct ChallengeHandle<'a> {
    identifier: AuthorizedIdentifier<'a>,
    challenge: &'a Challenge,
    nonce: &'a mut Option<String>,
    account: &'a AccountInner,
}

impl ChallengeHandle<'_> {
    /// Notify the server that the given challenge is ready to be completed
    pub async fn set_ready(&mut self) -> Result<(), Error> {
        let rsp = self
            .account
            .post(Some(&Empty {}), self.nonce.take(), &self.challenge.url)
            .await?;

        *self.nonce = nonce_from_response(&rsp);
        let response = Problem::check::<Challenge>(rsp).await?;
        match response.error {
            Some(details) => Err(Error::Api(details)),
            None => Ok(()),
        }
    }

    /// Notify the server that the challenge is ready by sending a device attestation
    ///
    /// This function is for the ACME challenge device-attest-01. It should not be used
    /// with other challenge types.
    /// See <https://datatracker.ietf.org/doc/draft-acme-device-attest/> for details.
    ///
    /// `payload` is the device attestation object as defined in link. Provide the attestation
    /// object as a raw blob. Base64 encoding of the attestation object `payload.att_obj`
    /// is done by this function.
    ///
    /// The function yields the challenge status from the ACME server that validated the
    /// attestation challenge.
    ///
    /// Note: Device attestation support is experimental.
    pub async fn send_device_attestation(
        &mut self,
        payload: &DeviceAttestation<'_>,
    ) -> Result<ChallengeStatus, Error> {
        if self.challenge.state.r#type() != ChallengeType::DeviceAttest01 {
            return Err(Error::Str("challenge type should be device-attest-01"));
        }

        #[derive(Serialize)]
        #[serde(rename_all = "camelCase")]
        struct DeviceAttestationBase64<'a> {
            att_obj: Cow<'a, str>,
        }

        let payload = DeviceAttestationBase64 {
            att_obj: Cow::Owned(BASE64_URL_SAFE_NO_PAD.encode(&payload.att_obj)),
        };

        let rsp = self
            .account
            .post(Some(&payload), self.nonce.take(), &self.challenge.url)
            .await?;

        *self.nonce = nonce_from_response(&rsp);
        let response = Problem::check::<Challenge>(rsp).await?;
        match response.error {
            Some(details) => Err(Error::Api(details)),
            None => Ok(response.status),
        }
    }

    /// Create a [`KeyAuthorization`] for this challenge
    ///
    /// Combines a challenge's token with the thumbprint of the account's public key to compute
    /// the challenge's `KeyAuthorization`. The `KeyAuthorization` must be used to provision the
    /// expected challenge response based on the challenge type in use.
    pub fn key_authorization(&self) -> Option<KeyAuthorization> {
        self.challenge
            .token()
            .map(|token| KeyAuthorization::new(token, &self.account.key))
    }

    /// The identifier for this challenge's authorization
    pub fn identifier(&self) -> &AuthorizedIdentifier<'_> {
        &self.identifier
    }
}

impl Deref for ChallengeHandle<'_> {
    type Target = Challenge;

    fn deref(&self) -> &Self::Target {
        self.challenge
    }
}

/// The response value to use for challenge responses
///
/// Refer to the methods below to see which encoding to use for your challenge type.
///
/// <https://datatracker.ietf.org/doc/html/rfc8555#section-8.1>
pub struct KeyAuthorization(String);

impl KeyAuthorization {
    fn new(token: &str, key: &Key) -> Self {
        Self(format!("{token}.{}", &key.thumb))
    }

    /// Get the key authorization value
    ///
    /// This can be used for HTTP-01 challenge responses.
    pub fn as_str(&self) -> &str {
        &self.0
    }

    /// Get the SHA-256 digest of the key authorization
    ///
    /// This can be used for TLS-ALPN-01 challenge responses.
    ///
    /// <https://datatracker.ietf.org/doc/html/rfc8737#section-3>
    pub fn digest(&self) -> impl AsRef<[u8]> {
        crypto::digest(&crypto::SHA256, self.0.as_bytes())
    }

    /// Get the base64-encoded SHA256 digest of the key authorization
    ///
    /// This can be used for DNS-01 challenge responses.
    pub fn dns_value(&self) -> String {
        BASE64_URL_SAFE_NO_PAD.encode(self.digest())
    }
}

impl fmt::Debug for KeyAuthorization {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_tuple("KeyAuthorization").finish()
    }
}

/// A policy for retrying API requests
///
/// Refresh the order state repeatedly, waiting `delay` before the first attempt and increasing
/// the delay by a factor of `backoff` after each attempt, until the `timeout` is reached.
#[derive(Debug, Clone, Copy)]
pub struct RetryPolicy {
    delay: Duration,
    backoff: f32,
    timeout: Duration,
}

impl RetryPolicy {
    /// A constructor for the default `RetryPolicy`
    ///
    /// Will retry for 5s with an initial delay of 250ms and a backoff factor of 2.0.
    pub const fn new() -> Self {
        Self {
            delay: Duration::from_millis(250),
            backoff: 2.0,
            timeout: Duration::from_secs(30),
        }
    }

    /// Set the initial delay
    ///
    /// This is the delay before the first retry attempt. The delay will be multiplied by the
    /// backoff factor after each attempt.
    pub const fn initial_delay(mut self, delay: Duration) -> Self {
        self.delay = delay;
        self
    }

    /// Set the backoff factor
    ///
    /// The delay will be multiplied by this factor after each retry attempt.
    pub const fn backoff(mut self, backoff: f32) -> Self {
        self.backoff = backoff;
        self
    }

    /// Set the timeout for retries
    ///
    /// After this duration has passed, no more retries will be attempted.
    pub const fn timeout(mut self, timeout: Duration) -> Self {
        self.timeout = timeout;
        self
    }

    fn state(&self) -> RetryState {
        RetryState {
            delay: self.delay,
            backoff: self.backoff,
            deadline: Instant::now() + self.timeout,
        }
    }
}

impl Default for RetryPolicy {
    fn default() -> Self {
        Self::new()
    }
}

struct RetryState {
    delay: Duration,
    backoff: f32,
    deadline: Instant,
}

impl RetryState {
    async fn wait(&mut self, after: Option<SystemTime>) -> ControlFlow<Error, ()> {
        if let Some(after) = after {
            let now = SystemTime::now();
            if let Ok(delay) = after.duration_since(now) {
                let next = Instant::now() + delay;
                if next > self.deadline {
                    return ControlFlow::Break(Error::Timeout(Some(next)));
                } else {
                    sleep(delay).await;
                    return ControlFlow::Continue(());
                }
            }
        }

        sleep(self.delay).await;
        self.delay = self.delay.mul_f32(self.backoff);
        match Instant::now() + self.delay > self.deadline {
            true => ControlFlow::Break(Error::Timeout(None)),
            false => ControlFlow::Continue(()),
        }
    }
}