handler.rs

//! This module implements the `KeysHandler` service
//!
//! It's responsible for managing authentication keys for the `BitTorrent` tracker.
//!
//! The service handles both persistent and in-memory storage of peer keys, and
//! supports adding new keys (either pre-generated or randomly created),
//! removing keys, and loading keys from the database into memory. Keys can be
//! either permanent or expire after a configurable duration per key.
use std::sync::Arc;
use std::time::Duration;

use torrust_tracker_clock::clock::Time;
use torrust_tracker_located_error::Located;
use torrust_tracker_primitives::DurationSinceUnixEpoch;

use super::key::repository::in_memory::InMemoryKeyRepository;
use super::key::repository::persisted::DatabaseKeyRepository;
use super::{key, CurrentClock, Key, PeerKey};
use crate::databases;
use crate::error::PeerKeyError;

/// Contains the information needed to add a new tracker key.
///
/// A new key can either be a pre-generated key provided by the user or can be
/// randomly generated by the application. Additionally, the key may be set to
/// expire after a certain number of seconds, or be permanent (if no expiration
/// is specified).
#[derive(Debug)]
pub struct AddKeyRequest {
    /// The pre-generated key as a string. If `None` the service will generate a
    ///  random key.
    pub opt_key: Option<String>,

    /// The duration (in seconds) for which the key is valid. Use `None` for
    /// permanent keys.
    pub opt_seconds_valid: Option<u64>,
}

/// The `KeysHandler` service manages the creation, addition, removal, and loading
///  of authentication keys for the tracker.
///
/// It uses both a persistent (database) repository and an in-memory repository
/// to manage keys.
pub struct KeysHandler {
    /// The database repository for storing authentication keys persistently.
    db_key_repository: Arc<DatabaseKeyRepository>,

    /// The in-memory repository for caching authentication keys.
    in_memory_key_repository: Arc<InMemoryKeyRepository>,
}

impl KeysHandler {
    /// Creates a new instance of the `KeysHandler` service.
    ///
    /// # Parameters
    ///
    /// - `db_key_repository`: A shared reference to the database key repository.
    /// - `in_memory_key_repository`: A shared reference to the in-memory key
    ///   repository.
    #[must_use]
    pub fn new(db_key_repository: &Arc<DatabaseKeyRepository>, in_memory_key_repository: &Arc<InMemoryKeyRepository>) -> Self {
        Self {
            db_key_repository: db_key_repository.clone(),
            in_memory_key_repository: in_memory_key_repository.clone(),
        }
    }

    /// Adds a new peer key to the tracker.
    ///
    /// The key may be pre-generated or generated on-the-fly.
    ///
    /// Depending on whether an expiration duration is specified, the key will
    /// be either expiring or permanent.
    ///
    /// # Parameters
    ///
    /// - `add_key_req`: The request containing options for key creation.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    ///
    /// - The provided key duration exceeds the maximum allowed value.
    /// - The provided pre-generated key is invalid.
    /// - There is an error persisting the key in the database.
    pub async fn add_peer_key(&self, add_key_req: AddKeyRequest) -> Result<PeerKey, PeerKeyError> {
        if let Some(pre_existing_key) = add_key_req.opt_key {
            // Pre-generated key

            if let Some(seconds_valid) = add_key_req.opt_seconds_valid {
                // Expiring key

                let Some(valid_until) = CurrentClock::now_add(&Duration::from_secs(seconds_valid)) else {
                    return Err(PeerKeyError::DurationOverflow { seconds_valid });
                };

                let key = pre_existing_key.parse::<Key>();

                match key {
                    Ok(key) => match self.add_expiring_peer_key(key, Some(valid_until)).await {
                        Ok(auth_key) => Ok(auth_key),
                        Err(err) => Err(PeerKeyError::DatabaseError {
                            source: Located(err).into(),
                        }),
                    },
                    Err(err) => Err(PeerKeyError::InvalidKey {
                        key: pre_existing_key,
                        source: Located(err).into(),
                    }),
                }
            } else {
                // Permanent key

                let key = pre_existing_key.parse::<Key>();

                match key {
                    Ok(key) => match self.add_permanent_peer_key(key).await {
                        Ok(auth_key) => Ok(auth_key),
                        Err(err) => Err(PeerKeyError::DatabaseError {
                            source: Located(err).into(),
                        }),
                    },
                    Err(err) => Err(PeerKeyError::InvalidKey {
                        key: pre_existing_key,
                        source: Located(err).into(),
                    }),
                }
            }
        } else {
            // New randomly generate key

            if let Some(seconds_valid) = add_key_req.opt_seconds_valid {
                // Expiring key

                match self
                    .generate_expiring_peer_key(Some(Duration::from_secs(seconds_valid)))
                    .await
                {
                    Ok(auth_key) => Ok(auth_key),
                    Err(err) => Err(PeerKeyError::DatabaseError {
                        source: Located(err).into(),
                    }),
                }
            } else {
                // Permanent key

                match self.generate_permanent_peer_key().await {
                    Ok(auth_key) => Ok(auth_key),
                    Err(err) => Err(PeerKeyError::DatabaseError {
                        source: Located(err).into(),
                    }),
                }
            }
        }
    }

    /// Generates a new permanent authentication key.
    ///
    /// Permanent keys do not expire.
    ///
    /// # Errors
    ///
    /// Returns a `databases::error::Error` if the key cannot be persisted in
    /// the database.
    pub(crate) async fn generate_permanent_peer_key(&self) -> Result<PeerKey, databases::error::Error> {
        self.generate_expiring_peer_key(None).await
    }

    /// Generates a new authentication key with an optional expiration lifetime.
    ///
    /// If a `lifetime` is provided, the generated key will expire after that
    /// duration. The new key is stored both in the database and in memory.
    ///
    /// # Parameters
    ///
    /// - `lifetime`: An optional duration specifying how long the key is valid.
    ///
    /// # Errors
    ///
    /// Returns a `databases::error::Error` if there is an issue adding the key
    /// to the database.
    pub async fn generate_expiring_peer_key(&self, lifetime: Option<Duration>) -> Result<PeerKey, databases::error::Error> {
        let peer_key = key::generate_key(lifetime);

        self.db_key_repository.add(&peer_key)?;

        self.in_memory_key_repository.insert(&peer_key).await;

        Ok(peer_key)
    }

    /// Adds a pre-generated permanent authentication key.
    ///
    /// Internally, this calls `add_expiring_peer_key` with no expiration.
    ///
    /// # Parameters
    ///
    /// - `key`: The pre-generated key.
    ///
    /// # Errors
    ///
    /// Returns a `databases::error::Error` if there is an issue persisting the
    /// key.
    pub(crate) async fn add_permanent_peer_key(&self, key: Key) -> Result<PeerKey, databases::error::Error> {
        self.add_expiring_peer_key(key, None).await
    }

    /// Adds a pre-generated authentication key with an optional expiration.
    ///
    /// The key is stored in both the database and the in-memory repository.
    ///
    /// # Parameters
    ///
    /// - `key`: The pre-generated key.
    /// - `valid_until`: An optional timestamp (as a duration since the Unix
    ///   epoch) after which the key expires.
    ///
    /// # Errors
    ///
    /// Returns a `databases::error::Error` if there is an issue adding the key
    /// to the database.
    pub(crate) async fn add_expiring_peer_key(
        &self,
        key: Key,
        valid_until: Option<DurationSinceUnixEpoch>,
    ) -> Result<PeerKey, databases::error::Error> {
        let peer_key = PeerKey { key, valid_until };

        // code-review: should we return a friendly error instead of the DB
        // constrain error when the key already exist? For now, it's returning
        // the specif error for each DB driver when a UNIQUE constrain fails.
        self.db_key_repository.add(&peer_key)?;

        self.in_memory_key_repository.insert(&peer_key).await;

        Ok(peer_key)
    }

    /// Removes an authentication key.
    ///
    /// The key is removed from both the database and the in-memory repository.
    ///
    /// # Parameters
    ///
    /// - `key`: A reference to the key to be removed.
    ///
    /// # Errors
    ///
    /// Returns a `databases::error::Error` if the key cannot be removed from
    /// the database.
    pub async fn remove_peer_key(&self, key: &Key) -> Result<(), databases::error::Error> {
        self.db_key_repository.remove(key)?;

        self.remove_in_memory_auth_key(key).await;

        Ok(())
    }

    /// Removes an authentication key from the in-memory repository.
    ///
    /// This function does not interact with the database.
    ///
    /// # Parameters
    ///
    /// - `key`: A reference to the key to be removed.
    pub(crate) async fn remove_in_memory_auth_key(&self, key: &Key) {
        self.in_memory_key_repository.remove(key).await;
    }

    /// Loads all authentication keys from the database into the in-memory
    /// repository.
    ///
    /// This is useful during tracker startup to ensure that all persisted keys
    /// are available in memory.
    ///
    /// # Errors
    ///
    /// Returns a `databases::error::Error` if there is an issue loading the keys from the database.
    pub async fn load_peer_keys_from_database(&self) -> Result<(), databases::error::Error> {
        let keys_from_database = self.db_key_repository.load_keys()?;

        self.in_memory_key_repository.reset_with(keys_from_database).await;

        Ok(())
    }
}

#[cfg(test)]
mod tests {

    mod the_keys_handler_when_the_tracker_is_configured_as_private {

        use std::sync::Arc;

        use torrust_tracker_configuration::Configuration;
        use torrust_tracker_test_helpers::configuration;

        use crate::authentication::handler::KeysHandler;
        use crate::authentication::key::repository::in_memory::InMemoryKeyRepository;
        use crate::authentication::key::repository::persisted::DatabaseKeyRepository;
        use crate::databases::setup::initialize_database;
        use crate::databases::Database;

        fn instantiate_keys_handler() -> KeysHandler {
            let config = configuration::ephemeral_private();

            instantiate_keys_handler_with_configuration(&config)
        }

        fn instantiate_keys_handler_with_database(database: &Arc<Box<dyn Database>>) -> KeysHandler {
            let db_key_repository = Arc::new(DatabaseKeyRepository::new(database));
            let in_memory_key_repository = Arc::new(InMemoryKeyRepository::default());

            KeysHandler::new(&db_key_repository, &in_memory_key_repository)
        }

        fn instantiate_keys_handler_with_configuration(config: &Configuration) -> KeysHandler {
            // todo: pass only Core configuration

            let database = initialize_database(&config.core);
            let db_key_repository = Arc::new(DatabaseKeyRepository::new(&database));
            let in_memory_key_repository = Arc::new(InMemoryKeyRepository::default());

            KeysHandler::new(&db_key_repository, &in_memory_key_repository)
        }

        mod handling_expiring_peer_keys {

            use std::time::Duration;

            use torrust_tracker_clock::clock::Time;

            use crate::authentication::handler::tests::the_keys_handler_when_the_tracker_is_configured_as_private::instantiate_keys_handler;
            use crate::CurrentClock;

            #[tokio::test]
            async fn it_should_generate_the_key() {
                let keys_handler = instantiate_keys_handler();

                let peer_key = keys_handler
                    .generate_expiring_peer_key(Some(Duration::from_secs(100)))
                    .await
                    .unwrap();

                assert_eq!(
                    peer_key.valid_until,
                    Some(CurrentClock::now_add(&Duration::from_secs(100)).unwrap())
                );
            }

            mod randomly_generated {
                use std::panic::Location;
                use std::sync::Arc;
                use std::time::Duration;

                use mockall::predicate::function;
                use torrust_tracker_clock::clock::stopped::Stopped;
                use torrust_tracker_clock::clock::{self, Time};

                use crate::authentication::handler::tests::the_keys_handler_when_the_tracker_is_configured_as_private::{
                    instantiate_keys_handler, instantiate_keys_handler_with_database,
                };
                use crate::authentication::handler::AddKeyRequest;
                use crate::authentication::PeerKey;
                use crate::databases::driver::Driver;
                use crate::databases::{self, Database, MockDatabase};
                use crate::error::PeerKeyError;
                use crate::CurrentClock;

                #[tokio::test]
                async fn it_should_add_a_randomly_generated_key() {
                    let keys_handler = instantiate_keys_handler();

                    let peer_key = keys_handler
                        .add_peer_key(AddKeyRequest {
                            opt_key: None,
                            opt_seconds_valid: Some(100),
                        })
                        .await
                        .unwrap();

                    assert_eq!(
                        peer_key.valid_until,
                        Some(CurrentClock::now_add(&Duration::from_secs(100)).unwrap())
                    );
                }

                #[tokio::test]
                async fn it_should_fail_adding_a_randomly_generated_key_when_there_is_a_database_error() {
                    clock::Stopped::local_set(&Duration::from_secs(0));

                    // The key should be valid the next 60 seconds.
                    let expected_valid_until = clock::Stopped::now_add(&Duration::from_secs(60)).unwrap();

                    let mut database_mock = MockDatabase::default();
                    database_mock
                        .expect_add_key_to_keys()
                        .with(function(move |peer_key: &PeerKey| {
                            peer_key.valid_until == Some(expected_valid_until)
                        }))
                        .times(1)
                        .returning(|_peer_key| {
                            Err(databases::error::Error::InsertFailed {
                                location: Location::caller(),
                                driver: Driver::Sqlite3,
                            })
                        });
                    let database_mock: Arc<Box<dyn Database>> = Arc::new(Box::new(database_mock));

                    let keys_handler = instantiate_keys_handler_with_database(&database_mock);

                    let result = keys_handler
                        .add_peer_key(AddKeyRequest {
                            opt_key: None,
                            opt_seconds_valid: Some(60), // The key is valid for 60 seconds.
                        })
                        .await;

                    assert!(matches!(result.unwrap_err(), PeerKeyError::DatabaseError { .. }));
                }
            }

            mod pre_generated {
                use std::panic::Location;
                use std::sync::Arc;
                use std::time::Duration;

                use mockall::predicate;
                use torrust_tracker_clock::clock::stopped::Stopped;
                use torrust_tracker_clock::clock::{self, Time};

                use crate::authentication::handler::tests::the_keys_handler_when_the_tracker_is_configured_as_private::{
                    instantiate_keys_handler, instantiate_keys_handler_with_database,
                };
                use crate::authentication::handler::AddKeyRequest;
                use crate::authentication::{Key, PeerKey};
                use crate::databases::driver::Driver;
                use crate::databases::{self, Database, MockDatabase};
                use crate::error::PeerKeyError;
                use crate::CurrentClock;

                #[tokio::test]
                async fn it_should_add_a_pre_generated_key() {
                    let keys_handler = instantiate_keys_handler();

                    let peer_key = keys_handler
                        .add_peer_key(AddKeyRequest {
                            opt_key: Some(Key::new("YZSl4lMZupRuOpSRC3krIKR5BPB14nrJ").unwrap().to_string()),
                            opt_seconds_valid: Some(100),
                        })
                        .await
                        .unwrap();

                    assert_eq!(
                        peer_key,
                        PeerKey {
                            key: Key::new("YZSl4lMZupRuOpSRC3krIKR5BPB14nrJ").unwrap(),
                            valid_until: Some(CurrentClock::now_add(&Duration::from_secs(100)).unwrap()),
                        }
                    );
                }

                #[tokio::test]
                async fn it_should_fail_adding_a_pre_generated_key_when_the_key_duration_exceeds_the_maximum_duration() {
                    let keys_handler = instantiate_keys_handler();

                    let result = keys_handler
                        .add_peer_key(AddKeyRequest {
                            opt_key: Some(Key::new("YZSl4lMZupRuOpSRC3krIKR5BPB14nrJ").unwrap().to_string()),
                            opt_seconds_valid: Some(u64::MAX),
                        })
                        .await;

                    assert!(matches!(result.unwrap_err(), PeerKeyError::DurationOverflow { .. }));
                }

                #[tokio::test]
                async fn it_should_fail_adding_a_pre_generated_key_when_the_key_is_invalid() {
                    let keys_handler = instantiate_keys_handler();

                    let result = keys_handler
                        .add_peer_key(AddKeyRequest {
                            opt_key: Some("INVALID KEY".to_string()),
                            opt_seconds_valid: Some(100),
                        })
                        .await;

                    assert!(matches!(result.unwrap_err(), PeerKeyError::InvalidKey { .. }));
                }

                #[tokio::test]
                async fn it_should_fail_adding_a_pre_generated_key_when_there_is_a_database_error() {
                    clock::Stopped::local_set(&Duration::from_secs(0));

                    // The key should be valid the next 60 seconds.
                    let expected_valid_until = clock::Stopped::now_add(&Duration::from_secs(60)).unwrap();
                    let expected_peer_key = PeerKey {
                        key: Key::new("YZSl4lMZupRuOpSRC3krIKR5BPB14nrJ").unwrap(),
                        valid_until: Some(expected_valid_until),
                    };

                    let mut database_mock = MockDatabase::default();
                    database_mock
                        .expect_add_key_to_keys()
                        .with(predicate::eq(expected_peer_key))
                        .times(1)
                        .returning(|_peer_key| {
                            Err(databases::error::Error::InsertFailed {
                                location: Location::caller(),
                                driver: Driver::Sqlite3,
                            })
                        });
                    let database_mock: Arc<Box<dyn Database>> = Arc::new(Box::new(database_mock));

                    let keys_handler = instantiate_keys_handler_with_database(&database_mock);

                    let result = keys_handler
                        .add_peer_key(AddKeyRequest {
                            opt_key: Some(Key::new("YZSl4lMZupRuOpSRC3krIKR5BPB14nrJ").unwrap().to_string()),
                            opt_seconds_valid: Some(60), // The key is valid for 60 seconds.
                        })
                        .await;

                    assert!(matches!(result.unwrap_err(), PeerKeyError::DatabaseError { .. }));
                }
            }
        }

        mod handling_permanent_peer_keys {

            mod randomly_generated {

                use std::panic::Location;
                use std::sync::Arc;

                use mockall::predicate::function;

                use crate::authentication::handler::tests::the_keys_handler_when_the_tracker_is_configured_as_private::{
                    instantiate_keys_handler, instantiate_keys_handler_with_database,
                };
                use crate::authentication::handler::AddKeyRequest;
                use crate::authentication::PeerKey;
                use crate::databases::driver::Driver;
                use crate::databases::{self, Database, MockDatabase};
                use crate::error::PeerKeyError;

                #[tokio::test]
                async fn it_should_generate_the_key() {
                    let keys_handler = instantiate_keys_handler();

                    let peer_key = keys_handler.generate_permanent_peer_key().await.unwrap();

                    assert_eq!(peer_key.valid_until, None);
                }

                #[tokio::test]
                async fn it_should_add_a_randomly_generated_key() {
                    let keys_handler = instantiate_keys_handler();

                    let peer_key = keys_handler
                        .add_peer_key(AddKeyRequest {
                            opt_key: None,
                            opt_seconds_valid: None,
                        })
                        .await
                        .unwrap();

                    assert_eq!(peer_key.valid_until, None);
                }

                #[tokio::test]
                async fn it_should_fail_adding_a_randomly_generated_key_when_there_is_a_database_error() {
                    let mut database_mock = MockDatabase::default();
                    database_mock
                        .expect_add_key_to_keys()
                        .with(function(move |peer_key: &PeerKey| peer_key.valid_until.is_none()))
                        .times(1)
                        .returning(|_peer_key| {
                            Err(databases::error::Error::InsertFailed {
                                location: Location::caller(),
                                driver: Driver::Sqlite3,
                            })
                        });
                    let database_mock: Arc<Box<dyn Database>> = Arc::new(Box::new(database_mock));

                    let keys_handler = instantiate_keys_handler_with_database(&database_mock);

                    let result = keys_handler
                        .add_peer_key(AddKeyRequest {
                            opt_key: None,
                            opt_seconds_valid: None,
                        })
                        .await;

                    assert!(matches!(result.unwrap_err(), PeerKeyError::DatabaseError { .. }));
                }
            }

            mod pre_generated_keys {

                use std::panic::Location;
                use std::sync::Arc;

                use mockall::predicate;

                use crate::authentication::handler::tests::the_keys_handler_when_the_tracker_is_configured_as_private::{
                    instantiate_keys_handler, instantiate_keys_handler_with_database,
                };
                use crate::authentication::handler::AddKeyRequest;
                use crate::authentication::{Key, PeerKey};
                use crate::databases::driver::Driver;
                use crate::databases::{self, Database, MockDatabase};
                use crate::error::PeerKeyError;

                #[tokio::test]
                async fn it_should_add_a_pre_generated_key() {
                    let keys_handler = instantiate_keys_handler();

                    let peer_key = keys_handler
                        .add_peer_key(AddKeyRequest {
                            opt_key: Some(Key::new("YZSl4lMZupRuOpSRC3krIKR5BPB14nrJ").unwrap().to_string()),
                            opt_seconds_valid: None,
                        })
                        .await
                        .unwrap();

                    assert_eq!(
                        peer_key,
                        PeerKey {
                            key: Key::new("YZSl4lMZupRuOpSRC3krIKR5BPB14nrJ").unwrap(),
                            valid_until: None,
                        }
                    );
                }

                #[tokio::test]
                async fn it_should_fail_adding_a_pre_generated_key_when_the_key_is_invalid() {
                    let keys_handler = instantiate_keys_handler();

                    let result = keys_handler
                        .add_peer_key(AddKeyRequest {
                            opt_key: Some("INVALID KEY".to_string()),
                            opt_seconds_valid: None,
                        })
                        .await;

                    assert!(matches!(result.unwrap_err(), PeerKeyError::InvalidKey { .. }));
                }

                #[tokio::test]
                async fn it_should_fail_adding_a_pre_generated_key_when_there_is_a_database_error() {
                    let expected_peer_key = PeerKey {
                        key: Key::new("YZSl4lMZupRuOpSRC3krIKR5BPB14nrJ").unwrap(),
                        valid_until: None,
                    };

                    let mut database_mock = MockDatabase::default();
                    database_mock
                        .expect_add_key_to_keys()
                        .with(predicate::eq(expected_peer_key))
                        .times(1)
                        .returning(|_peer_key| {
                            Err(databases::error::Error::InsertFailed {
                                location: Location::caller(),
                                driver: Driver::Sqlite3,
                            })
                        });
                    let database_mock: Arc<Box<dyn Database>> = Arc::new(Box::new(database_mock));

                    let keys_handler = instantiate_keys_handler_with_database(&database_mock);

                    let result = keys_handler
                        .add_peer_key(AddKeyRequest {
                            opt_key: Some(Key::new("YZSl4lMZupRuOpSRC3krIKR5BPB14nrJ").unwrap().to_string()),
                            opt_seconds_valid: None,
                        })
                        .await;

                    assert!(matches!(result.unwrap_err(), PeerKeyError::DatabaseError { .. }));
                }
            }
        }
    }
}