krill 0.16.0

use std::{
    marker::PhantomData,
    sync::{Arc, RwLock, RwLockReadGuard},
    time::{Duration, Instant},
};
use log::info;

use super::error::SignerError;

#[derive(Debug)]
pub enum ProbeError<E> {
    WrongState,
    AwaitingNextProbe,
    CompletedUnusable,
    CallbackFailed(E),
}

/// Probe status based access to the PKCS#11 server.
///
/// To avoid blocking Krill startup due to HSM connection timeout or failure
/// we start in a `AwaitingNextProbe` status which signifies that we haven't
/// yet verified that we can connect to the HSM or that it supports the
/// capabilities that we require.
///
/// At some point later once an initial connection has been established the
/// PKCS#11 signer changes status to either `Usable` or `Unusable` based on
/// what was discovered about the PKCS#11 server.
#[derive(Debug)]
pub struct StatefulProbe<C, E, S> {
    name: String,

    status: RwLock<ProbeStatus<C, E, S>>,

    probe_interval: Duration,
}

pub enum ProbeStatus<C, E, S> {
    /// We haven't yet been able to connect to the HSM. If there was already
    /// a failed attempt to connect the timestamp of the attempt is
    /// remembered so that we can choose to space out connection attempts
    /// rather than attempt to connect every time Krill tries to use the
    /// signer.
    Probing {
        config: Arc<C>,
        last_probe_time: Option<Instant>,
        phantom: PhantomData<E>,
    },

    /// The HSM was successfully probed but found to be lacking required
    /// capabilities and is thus unusable by Krill.
    Unusable,

    /// The HSM was successfully probed and confirmed to have the required
    /// capabilities.
    ///
    /// Note that this does not mean that the HSM is currently contactable,
    /// only that we were able to contact it at least once since Krill
    /// was started. If the domain name/IP address used to connect to Krill
    /// now point to a different HSM instance the previously determined
    /// conclusion that the HSM is usable may no longer be valid.
    ///
    /// In this status we keep state concerning our relationship with the
    /// HSM.
    Usable(S),
}

impl<C, E, S> std::fmt::Debug for ProbeStatus<C, E, S> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Probing { .. } => write!(f, "Probing"),
            Self::Unusable => write!(f, "Unusable"),
            Self::Usable(_) => write!(f, "Usable"),
        }
    }
}

impl<C, E, S> ProbeStatus<C, E, S> {
    /// Marks now as the last probe attempt timestamp.
    ///
    /// Calling this function while not in the Probing state will result in a
    /// panic.
    pub fn mark(&mut self) -> Result<(), ProbeError<E>> {
        match self {
            #[rustfmt::skip]
            ProbeStatus::Probing { last_probe_time, .. } => {
                last_probe_time.replace(Instant::now());
                Ok(())
            }
            _ => Err(ProbeError::WrongState),
        }
    }

    pub fn config(&self) -> Result<Arc<C>, ProbeError<E>> {
        match self {
            ProbeStatus::Probing { config, .. } => Ok(config.clone()),
            _ => Err(ProbeError::WrongState),
        }
    }

    #[cfg(test)]
    pub fn last_probe_time(&self) -> Result<Option<Instant>, ProbeError<E>> {
        match self {
            ProbeStatus::Probing {
                last_probe_time, ..
            } => Ok(*last_probe_time),
            _ => Err(ProbeError::WrongState),
        }
    }

    /// Helper function to retrieve the state associated with status Usable.
    /// Only callable when in status `Usable`. Calling this function while
    /// in another state will result in a panic.
    pub fn state(&self) -> Result<&S, ProbeError<E>> {
        match self {
            ProbeStatus::Usable(state) => Ok(state),
            _ => Err(ProbeError::WrongState),
        }
    }
}

impl<C, E, S> StatefulProbe<C, E, S> {
    /// Create a new connector to a server that hasn't been probed yet.
    pub fn new(
        name: String,
        config: Arc<C>,
        probe_interval: Duration,
    ) -> Self {
        let status = RwLock::new(ProbeStatus::Probing {
            config,
            last_probe_time: None,
            phantom: PhantomData,
        });
        StatefulProbe {
            name,
            status,
            probe_interval,
        }
    }

    #[cfg(test)]
    pub fn last_probe_time(&self) -> Result<Option<Instant>, ProbeError<E>> {
        self.status.read().unwrap().last_probe_time()
    }

    /// Get a read lock on the Usable server status, if the server is usable.
    ///
    /// Returns `Ok` with the status read lock if the server is usable,
    /// otherwise returns an `Err` because the server is unusable or we
    /// haven't yet been able to establish if it is usable or not.
    ///
    /// Will try probing again if we didn't already manage to connect to the
    /// server and the delay period between probes has elapsed.
    pub fn status<F>(
        &self,
        probe: F,
    ) -> Result<RwLockReadGuard<'_, ProbeStatus<C, E, S>>, ProbeError<E>>
    where
        F: Fn(String, &ProbeStatus<C, E, S>) -> Result<S, ProbeError<E>>,
    {
        fn is_time_to_check(
            time_between_probes: Duration,
            last_probe_time: Option<Instant>,
        ) -> bool {
            match last_probe_time {
                None => true,
                Some(instant) => {
                    Instant::now().saturating_duration_since(instant)
                        > time_between_probes
                }
            }
        }

        #[allow(clippy::type_complexity)]
        fn get_if_usable<C, E, S>(
            name: String,
            status: RwLockReadGuard<ProbeStatus<C, E, S>>,
            retry_interval: Duration,
        ) -> Option<
            Result<RwLockReadGuard<ProbeStatus<C, E, S>>, ProbeError<E>>,
        > {
            // Check the status through the unlocked read lock
            match &*status {
                ProbeStatus::Usable(_) => {
                    // The server has been confirmed as usable, return the
                    // read-lock granting access to the current
                    // status and via it the current state of our relationship
                    // with the server.
                    Some(Ok(status))
                }

                ProbeStatus::Unusable => {
                    // The server has been confirmed as unusable, fail.
                    Some(Err(ProbeError::CompletedUnusable))
                }

                ProbeStatus::Probing {
                    last_probe_time, ..
                } => {
                    // We haven't yet established whether the  server is
                    // usable or not. If we haven't yet checked or we
                    // haven't tried checking again for a while, then try
                    // contacting it again. If we can't establish
                    // whether or not the server is usable, return an error.
                    if !is_time_to_check(retry_interval, *last_probe_time) {
                        if let Some(instant) = last_probe_time {
                            let until = retry_interval
                                .checked_sub(instant.elapsed())
                                .unwrap(); // TODO
                            info!(
                                "Signer '{}' is currently unavailable. {:?}s until next check.",
                                name,
                                until.as_secs()
                            );
                        } else {
                            // This should be unreachable
                            info!(
                                "Signer '{}' is currently unavailable. Checking every {:?}s",
                                name,
                                retry_interval.as_secs()
                            );
                        }
                        Some(Err(ProbeError::AwaitingNextProbe))
                    } else {
                        None
                    }
                }
            }
        }

        /// Verify if the configured server is contactable and supports the
        /// required capabilities.
        fn send_probe<C, E, F, S>(
            probe: &StatefulProbe<C, E, S>,
            probe_cb: F,
        ) -> Result<(), ProbeError<E>>
        where
            F: Fn(String, &ProbeStatus<C, E, S>) -> Result<S, ProbeError<E>>,
        {
            // Hold a write lock for the duration of our attempt to verify the
            // server so that no other attempt occurs at the same
            // time. Bail out if another thread is performing a probe and has
            // the lock. This is the same result
            // as when attempting to use the server between probe retries.
            let mut status = probe
                .status
                .try_write()
                .map_err(|_| ProbeError::AwaitingNextProbe)?;

            // Update the timestamp of our last attempt to contact the server.
            // This is used above to know when we have waited long
            // enough before attempting to contact the server again. This also
            // guards against attempts to probe when probing has
            // already finished as mark() will fail in that case.
            status.mark()?;

            match (probe_cb)(probe.name.clone(), &*status) {
                Ok(usable_state) => {
                    *status = ProbeStatus::Usable(usable_state);
                    Ok(())
                }
                Err(err) => {
                    if matches!(err, ProbeError::CompletedUnusable) {
                        *status = ProbeStatus::Unusable;
                    }
                    Err(err)
                }
            }
        }

        // Return the current status or attempt to set it by probing the
        // server
        let status = self.status.read().unwrap();
        get_if_usable(self.name.clone(), status, self.probe_interval)
            .unwrap_or_else(|| {
                send_probe(self, probe)
                    .map(|_| self.status.read().unwrap())
                    .map_err(|err| match err {
                        ProbeError::CompletedUnusable => err,
                        _ => ProbeError::AwaitingNextProbe,
                    })
            })
    }
}

impl From<ProbeError<SignerError>> for SignerError {
    fn from(err: ProbeError<SignerError>) -> Self {
        match err {
            ProbeError::WrongState => SignerError::Other(
                "Internal error: probe is not in the expected state"
                    .to_string(),
            ),
            ProbeError::AwaitingNextProbe => {
                SignerError::TemporarilyUnavailable
            }
            ProbeError::CompletedUnusable => SignerError::PermanentlyUnusable,
            ProbeError::CallbackFailed(err) => err,
        }
    }
}

#[cfg(test)]
pub mod tests {
    use std::time::Duration;

    use super::*;

    #[derive(Debug, Default)]
    struct Config {
        _hostname: String,

        _port: u64,
    }

    #[derive(Copy, Clone, Debug, Default)]
    struct State {
        some_state: u8,
    }

    impl State {
        fn some_func(&self) -> u8 {
            self.some_state
        }
    }

    #[derive(Debug)]
    enum SomeError {
        SomeErrorCode,
    }

    fn probe_func(
        _name: String,
        _status: &ProbeStatus<Config, SomeError, State>,
    ) -> Result<State, ProbeError<SomeError>> {
        Err(ProbeError::CompletedUnusable)
    }

    #[test]
    fn probe_should_be_permanently_unavailable_with_closure() {
        let config = Arc::new(Config::default());
        let conn = StatefulProbe::<_, SomeError, State>::new(
            "dummy".to_string(),
            config,
            Duration::from_secs(0),
        );
        let res = conn.status(|_, _| Err(ProbeError::CompletedUnusable));
        match res {
            Err(ProbeError::CompletedUnusable) => {}
            other => panic!(
                "Expected Err(ProbeError::PermanentlyUnusable) but got {other:?}"
            ),
        }
    }

    #[test]
    fn probe_should_be_permanently_unavailable_with_fn() {
        let config = Arc::new(Config::default());
        let conn = StatefulProbe::<_, SomeError, State>::new(
            "dummy".to_string(),
            config,
            Duration::from_secs(0),
        );
        let res = conn.status(probe_func);
        match res {
            Err(ProbeError::CompletedUnusable) => {}
            other => panic!(
                "Expected Err(ProbeError::PermanentlyUnusable) but got {other:?}"
            ),
        }
    }

    #[test]
    fn probe_should_be_permanently_unavailable() {
        let config = Arc::new(Config::default());
        let conn = StatefulProbe::<_, SomeError, State>::new(
            "dummy".to_string(),
            config,
            Duration::from_secs(0),
        );
        let res = conn.status(|_, _| Err(ProbeError::CompletedUnusable));
        match res {
            Err(ProbeError::CompletedUnusable) => {}
            other => panic!(
                "Expected Err(ProbeError::PermanentlyUnusable) but got {other:?}"
            ),
        }
    }

    #[test]
    fn probe_should_be_temporarily_unavailable() {
        let config = Arc::new(Config::default());
        let conn = StatefulProbe::<_, SomeError, State>::new(
            "dummy".to_string(),
            config,
            Duration::from_secs(0),
        );
        let res = conn.status(|_, _| Err(ProbeError::AwaitingNextProbe));
        match res {
            Err(ProbeError::AwaitingNextProbe) => {}
            other => panic!(
                "Expected Err(ProbeError::AwaitingNextProbe) but got {other:?}"
            ),
        }
    }

    #[test]
    fn probe_should_be_temporarily_unavailable_on_custom_error() {
        let config = Arc::new(Config::default());
        let conn = StatefulProbe::<_, SomeError, State>::new(
            "dummy".to_string(),
            config,
            Duration::from_secs(0),
        );
        let res = conn.status(|_, _| {
            Err(ProbeError::CallbackFailed(SomeError::SomeErrorCode))
        });
        match res {
            Err(ProbeError::AwaitingNextProbe) => {}
            other => panic!(
                "Expected Err(ProbeError::AwaitingNextProbe) but got {other:?}"
            ),
        }
    }

    #[test]
    fn last_probe_time_should_advance() -> Result<(), ProbeError<SomeError>> {
        let config = Arc::new(Config::default());

        // Probing is only done when .get() is called
        let conn = StatefulProbe::<_, SomeError, State>::new(
            "dummy".to_string(),
            config,
            Duration::from_millis(200),
        );
        assert_eq!(None, conn.last_probe_time()?);

        // The first call to .get() should trigger a probe
        let _probe1 = conn.status(|_, _| Err(ProbeError::AwaitingNextProbe));
        let t1 = conn.last_probe_time()?;
        assert!(t1.is_some());

        // A call to .get() before the next probe interval should NOT result
        // in an updated last probe time
        std::thread::sleep(Duration::from_millis(10));
        let _probe2 = conn.status(|_, _| Err(ProbeError::AwaitingNextProbe));
        let t2 = conn.last_probe_time()?;
        assert!(t2 == t1);

        // A call to .get() after the next probe interval SHOULD result in an
        // updated last probe time
        std::thread::sleep(Duration::from_millis(191)); // plus 10 earlier is >200
        let _probe3 = conn.status(|_, _| Err(ProbeError::AwaitingNextProbe));
        let t3 = conn.last_probe_time()?;
        assert!(t3 > t1);

        Ok(())
    }

    #[test]
    fn probe_should_change_state_when_usable(
    ) -> Result<(), ProbeError<SomeError>> {
        let config = Arc::new(Config::default());
        let new_state = State { some_state: 1 };

        // Probing only happens when .get() is called
        let conn = StatefulProbe::<_, SomeError, State>::new(
            "dummy".to_string(),
            config,
            Duration::from_millis(0),
        );
        let new_status = conn.status(|_, _| Ok(new_state))?;
        assert_eq!(1, new_status.state()?.some_state);
        assert_eq!(1, new_status.state()?.some_func());

        Ok(())
    }
}