Skip to content

Commit

Permalink
refactor: [torrust#1195] extract InMemoryKeyRepository
Browse files Browse the repository at this point in the history
  • Loading branch information
josecelano committed Jan 21, 2025
1 parent f4c7b97 commit 12a62ce
Show file tree
Hide file tree
Showing 3 changed files with 52 additions and 20 deletions.
30 changes: 30 additions & 0 deletions src/core/authentication/key/repository/in_memory.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
use crate::core::authentication::key::{Key, PeerKey};

/// In-memory implementation of the authentication key repository.
#[derive(Debug, Default)]
pub struct InMemoryKeyRepository {
/// Tracker users' keys. Only for private trackers.
keys: tokio::sync::RwLock<std::collections::HashMap<Key, PeerKey>>,
}

impl InMemoryKeyRepository {
/// It adds a new authentication key.
pub async fn insert(&self, auth_key: &PeerKey) {
self.keys.write().await.insert(auth_key.key.clone(), auth_key.clone());
}

/// It removes an authentication key.
pub async fn remove(&self, key: &Key) {
self.keys.write().await.remove(key);
}

pub async fn get(&self, key: &Key) -> Option<PeerKey> {
self.keys.read().await.get(key).cloned()
}

/// It clears all the authentication keys.
pub async fn clear(&self) {
let mut keys = self.keys.write().await;
keys.clear();
}
}
1 change: 1 addition & 0 deletions src/core/authentication/key/repository/mod.rs
Original file line number Diff line number Diff line change
@@ -1 +1,2 @@
pub mod in_memory;
pub mod persisted;
41 changes: 21 additions & 20 deletions src/core/authentication/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ use std::panic::Location;
use std::sync::Arc;
use std::time::Duration;

use key::repository::in_memory::InMemoryKeyRepository;
use key::repository::persisted::DatabaseKeyRepository;
use torrust_tracker_clock::clock::Time;
use torrust_tracker_configuration::Core;
Expand Down Expand Up @@ -36,20 +37,20 @@ pub struct Facade {
/// The tracker configuration.
config: Core,

/// Tracker users' keys. Only for private trackers.
keys: tokio::sync::RwLock<std::collections::HashMap<Key, PeerKey>>,

/// The database repository for the authentication keys.
db_key_repository: DatabaseKeyRepository,

/// In-memory implementation of the authentication key repository.
in_memory_key_repository: InMemoryKeyRepository,
}

impl Facade {
#[must_use]
pub fn new(config: &Core, database: &Arc<Box<dyn Database>>) -> Self {
Self {
config: config.clone(),
keys: tokio::sync::RwLock::new(std::collections::HashMap::new()),
db_key_repository: DatabaseKeyRepository::new(database),
in_memory_key_repository: InMemoryKeyRepository::default(),
}
}

Expand Down Expand Up @@ -82,20 +83,20 @@ impl Facade {
///
/// Will return a `key::Error` if unable to get any `auth_key`.
pub async fn verify_auth_key(&self, key: &Key) -> Result<(), Error> {
match self.keys.read().await.get(key) {
match self.in_memory_key_repository.get(key).await {
None => Err(Error::UnableToReadKey {
location: Location::caller(),
key: Box::new(key.clone()),
}),
Some(key) => match self.config.private_mode {
Some(private_mode) => {
if private_mode.check_keys_expiration {
return key::verify_key_expiration(key);
return key::verify_key_expiration(&key);
}

Ok(())
}
None => key::verify_key_expiration(key),
None => key::verify_key_expiration(&key),
},
}
}
Expand Down Expand Up @@ -203,12 +204,13 @@ impl Facade {
/// * `lifetime` - The duration in seconds for the new key. The key will be
/// no longer valid after `lifetime` seconds.
pub async fn generate_auth_key(&self, lifetime: Option<Duration>) -> Result<PeerKey, databases::error::Error> {
let auth_key = key::generate_key(lifetime);
let peer_key = key::generate_key(lifetime);

self.db_key_repository.add(&peer_key)?;

self.db_key_repository.add(&auth_key)?;
self.in_memory_key_repository.insert(&peer_key).await;

self.keys.write().await.insert(auth_key.key.clone(), auth_key.clone());
Ok(auth_key)
Ok(peer_key)
}

/// It adds a pre-generated permanent authentication key.
Expand Down Expand Up @@ -250,15 +252,16 @@ impl Facade {
key: Key,
valid_until: Option<DurationSinceUnixEpoch>,
) -> Result<PeerKey, databases::error::Error> {
let auth_key = PeerKey { key, valid_until };
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(&auth_key)?;
self.db_key_repository.add(&peer_key)?;

self.keys.write().await.insert(auth_key.key.clone(), auth_key.clone());
Ok(auth_key)
self.in_memory_key_repository.insert(&peer_key).await;

Ok(peer_key)
}

/// It removes an authentication key.
Expand All @@ -280,7 +283,7 @@ impl Facade {
///
/// # Context: Authentication
pub async fn remove_in_memory_auth_key(&self, key: &Key) {
self.keys.write().await.remove(key);
self.in_memory_key_repository.remove(key).await;
}

/// The `Tracker` stores the authentication keys in memory and in the database.
Expand All @@ -296,12 +299,10 @@ impl Facade {
pub async fn load_keys_from_database(&self) -> Result<(), databases::error::Error> {
let keys_from_database = self.db_key_repository.load_keys()?;

let mut keys = self.keys.write().await;

keys.clear();
self.in_memory_key_repository.clear().await;

for key in keys_from_database {
keys.insert(key.key.clone(), key);
self.in_memory_key_repository.insert(&key).await;
}

Ok(())
Expand Down

0 comments on commit 12a62ce

Please sign in to comment.