Getting rid of in-memory data to avoid unbounded memory growth

Regrading the `Watcher`, fields (appointments, locator_uuid_map) has
been replaced by DB calls when needed.

For `Responder`, the field `trackers` has been replaced by DB calls when
needed, and `tx_tracker_map` wasn't actually needed for the tower to
operate, so was just dropped.

For `GateKeeper`, `registered_users::appointments` which used to hold
the uuids of every appointment the user submitted was removed so that
`registered_users` only holds meta information about users.

Also now the gatekeeper is the entity responsible for deleting appointments from the database. Instead of the watcher/responder asking the gatekeeper for the users to update and carry out the deletion and update itself, now the watcher/responder will hand the gatekeeper the uuids to delete and the gatekeeper will figure out which users it needs to update (refund the freed slots to).

Also now, like in `Watcher::store_triggered_appointment`, if the appointment is invalid or was rejected by the network in block connections, the freed slots will not be refunded to the user.

Also the block connection order starts with the gatekeeper first, this
allows the gatekeeper to delete the outdated users so that the watcher
and the responder doesn't take them into account.
This commit is contained in:
Omer Yacine 2023-03-07 18:50:54 +02:00
parent fad3ad1c08
commit cab6151ccc
No known key found for this signature in database
GPG key ID: C3BED6698142B393
11 changed files with 1947 additions and 2510 deletions

View file

@ -32,6 +32,12 @@ pub fn get_random_user_id() -> UserId {
UserId(pk)
}
pub fn get_random_locator() -> Locator {
let mut rng = rand::thread_rng();
Locator::from_slice(&rng.gen::<[u8; 16]>()).unwrap()
}
pub fn generate_random_appointment(dispute_txid: Option<&Txid>) -> Appointment {
let dispute_txid = match dispute_txid {
Some(l) => *l,

View file

@ -660,8 +660,11 @@ mod tests_methods {
};
use super::*;
use crate::extended_appointment::UUID;
use crate::test_utils::{generate_dummy_appointment, ApiConfig, DURATION, SLOTS};
use crate::responder::{ConfirmationStatus, TransactionTracker};
use crate::test_utils::{
generate_dummy_appointment, get_random_tx, ApiConfig, DURATION, SLOTS,
};
use crate::watcher::Breach;
use teos_common::test_utils::get_random_user_id;
use teos_common::{cryptography, UserId};
@ -830,14 +833,20 @@ mod tests_methods {
.await
.unwrap();
// Add the appointment to the Responder so it counts as triggered
let appointment = generate_dummy_appointment(None).inner;
let signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();
// Add the appointment to the Responder as a tracker so it counts as triggered
let dispute_tx = get_random_tx();
let tracker = TransactionTracker::new(
Breach::new(dispute_tx.clone(), get_random_tx()),
UserId(user_pk),
ConfirmationStatus::ConfirmedIn(100),
);
internal_api
.get_watcher()
.add_random_tracker_to_responder(UUID::new(appointment.locator, UserId(user_pk)));
.add_dummy_tracker_to_responder(&tracker);
// Try to add it via the http API
let appointment = generate_dummy_appointment(Some(&dispute_tx.txid())).inner;
let signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();
assert_eq!(
check_api_error(
Endpoint::AddAppointment,

View file

@ -2,6 +2,7 @@ use std::sync::{Arc, Condvar, Mutex};
use tonic::{Code, Request, Response, Status};
use triggered::Trigger;
use crate::extended_appointment::UUID;
use crate::protos as msgs;
use crate::protos::private_tower_services_server::PrivateTowerServices;
use crate::protos::public_tower_services_server::PublicTowerServices;
@ -386,10 +387,14 @@ impl PrivateTowerServices for Arc<InternalAPI> {
})?;
match self.watcher.get_user_info(user_id) {
Some(info) => Ok(Response::new(msgs::GetUserResponse {
Some((info, locators)) => Ok(Response::new(msgs::GetUserResponse {
available_slots: info.available_slots,
subscription_expiry: info.subscription_expiry,
appointments: info.appointments.keys().map(|uuid| uuid.to_vec()).collect(),
// TODO: Should make it return locators and make `get_appointments` queryable using the (user_id, locator) pair for consistency.
appointments: locators
.into_iter()
.map(|locator| UUID::new(locator, user_id).to_vec())
.collect(),
})),
None => Err(Status::new(Code::NotFound, "User not found")),
}
@ -429,11 +434,10 @@ mod tests_private_api {
use bitcoin::hashes::Hash;
use bitcoin::Txid;
use crate::extended_appointment::UUID;
use crate::responder::{ConfirmationStatus, TransactionTracker};
use crate::test_utils::{
create_api, generate_dummy_appointment, generate_uuid, get_random_tx, DURATION, SLOTS,
START_HEIGHT,
create_api, generate_dummy_appointment, generate_dummy_appointment_with_user,
get_random_tx, DURATION, SLOTS, START_HEIGHT,
};
use crate::watcher::Breach;
@ -486,9 +490,7 @@ mod tests_private_api {
let (internal_api, _s) = create_api().await;
// Add data to the Responser so we can retrieve it later on
internal_api
.watcher
.add_random_tracker_to_responder(generate_uuid());
internal_api.watcher.add_random_tracker_to_responder();
let response = internal_api
.get_all_appointments(Request::new(()))
@ -588,7 +590,7 @@ mod tests_private_api {
);
internal_api
.watcher
.add_dummy_tracker_to_responder(generate_uuid(), &tracker);
.add_dummy_tracker_to_responder(&tracker);
}
let locator = Locator::new(dispute_tx.txid());
@ -655,9 +657,7 @@ mod tests_private_api {
// And the Responder
for _ in 0..3 {
internal_api
.watcher
.add_random_tracker_to_responder(generate_uuid());
internal_api.watcher.add_random_tracker_to_responder();
}
let response = internal_api
@ -730,12 +730,11 @@ mod tests_private_api {
assert!(response.appointments.is_empty());
// Add an appointment and check back
let appointment = generate_dummy_appointment(None).inner;
let uuid = UUID::new(appointment.locator, user_id);
let user_signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();
let (uuid, appointment) = generate_dummy_appointment_with_user(user_id, None);
let user_signature = cryptography::sign(&appointment.inner.to_vec(), &user_sk).unwrap();
internal_api
.watcher
.add_appointment(appointment.clone(), user_signature)
.add_appointment(appointment.inner, user_signature)
.unwrap();
let response = internal_api
@ -786,10 +785,12 @@ mod tests_private_api {
mod tests_public_api {
use super::*;
use crate::extended_appointment::UUID;
use crate::responder::{ConfirmationStatus, TransactionTracker};
use crate::test_utils::{
create_api, create_api_with_config, generate_dummy_appointment, ApiConfig, DURATION, SLOTS,
create_api, create_api_with_config, generate_dummy_appointment, get_random_tx, ApiConfig,
DURATION, SLOTS,
};
use crate::watcher::Breach;
use teos_common::cryptography::{self, get_random_keypair};
#[tokio::test]
@ -900,12 +901,12 @@ mod tests_public_api {
internal_api.watcher.register(UserId(user_pk)).unwrap();
let appointment = generate_dummy_appointment(None).inner;
let user_signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();
let signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();
let response = internal_api
.add_appointment(Request::new(common_msgs::AddAppointmentRequest {
appointment: Some(appointment.clone().into()),
signature: user_signature.clone(),
signature,
}))
.await
.unwrap()
@ -925,12 +926,12 @@ mod tests_public_api {
let (user_sk, _) = get_random_keypair();
let appointment = generate_dummy_appointment(None).inner;
let user_signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();
let signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();
match internal_api
.add_appointment(Request::new(common_msgs::AddAppointmentRequest {
appointment: Some(appointment.clone().into()),
signature: user_signature.clone(),
signature,
}))
.await
{
@ -954,12 +955,12 @@ mod tests_public_api {
internal_api.watcher.register(UserId(user_pk)).unwrap();
let appointment = generate_dummy_appointment(None).inner;
let user_signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();
let signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();
match internal_api
.add_appointment(Request::new(common_msgs::AddAppointmentRequest {
appointment: Some(appointment.clone().into()),
signature: user_signature.clone(),
signature,
}))
.await
{
@ -983,12 +984,12 @@ mod tests_public_api {
internal_api.watcher.register(UserId(user_pk)).unwrap();
let appointment = generate_dummy_appointment(None).inner;
let user_signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();
let signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();
match internal_api
.add_appointment(Request::new(common_msgs::AddAppointmentRequest {
appointment: Some(appointment.clone().into()),
signature: user_signature.clone(),
signature,
}))
.await
{
@ -1008,16 +1009,24 @@ mod tests_public_api {
let user_id = UserId(user_pk);
internal_api.watcher.register(user_id).unwrap();
let appointment = generate_dummy_appointment(None).inner;
let user_signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();
// Add a tracker to the responder to simulate it being triggered.
let dispute_tx = get_random_tx();
let tracker = TransactionTracker::new(
Breach::new(dispute_tx.clone(), get_random_tx()),
user_id,
ConfirmationStatus::ConfirmedIn(100),
);
internal_api
.watcher
.add_random_tracker_to_responder(UUID::new(appointment.locator, user_id));
.get_watcher()
.add_dummy_tracker_to_responder(&tracker);
// Try to add it again using the API.
let appointment = generate_dummy_appointment(Some(&dispute_tx.txid())).inner;
let signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();
match internal_api
.add_appointment(Request::new(common_msgs::AddAppointmentRequest {
appointment: Some(appointment.clone().into()),
signature: user_signature.clone(),
appointment: Some(appointment.into()),
signature,
}))
.await
{
@ -1038,12 +1047,12 @@ mod tests_public_api {
let (user_sk, _) = get_random_keypair();
let appointment = generate_dummy_appointment(None).inner;
let user_signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();
let signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();
match internal_api
.add_appointment(Request::new(common_msgs::AddAppointmentRequest {
appointment: Some(appointment.clone().into()),
signature: user_signature.clone(),
signature,
}))
.await
{

File diff suppressed because it is too large Load diff

View file

@ -46,8 +46,6 @@ impl std::fmt::Display for UUID {
/// An extended version of the appointment hold by the tower.
///
/// The [Appointment] is extended in terms of data, that is, it provides further information only relevant to the tower.
/// Notice [ExtendedAppointment]s are not kept in memory but persisted on disk. The [Watcher](crate::watcher::Watcher)
/// keeps [AppointmentSummary] instead.
#[derive(Debug, Eq, PartialEq, Clone)]
pub(crate) struct ExtendedAppointment {
/// The underlying appointment extended by [ExtendedAppointment].
@ -60,24 +58,6 @@ pub(crate) struct ExtendedAppointment {
pub start_block: u32,
}
/// A summary of an appointment.
///
/// Contains the minimal amount of data the [Watcher](crate::watcher::Watcher) needs to keep in memory in order to
/// watch for breaches.
#[derive(Debug, Eq, PartialEq, Clone)]
pub(crate) struct AppointmentSummary {
/// The [Appointment] locator.
pub locator: Locator,
/// The user this [Appointment] belongs to.
pub user_id: UserId,
}
impl AppointmentSummary {
pub fn new(locator: Locator, user_id: UserId) -> Self {
Self { locator, user_id }
}
}
impl ExtendedAppointment {
/// Create a new [ExtendedAppointment].
pub fn new(
@ -109,12 +89,8 @@ impl ExtendedAppointment {
self.inner.to_self_delay
}
/// Computes the summary of the [ExtendedAppointment].
pub fn get_summary(&self) -> AppointmentSummary {
AppointmentSummary {
locator: self.locator(),
user_id: self.user_id,
}
pub fn uuid(&self) -> UUID {
UUID::new(self.inner.locator, self.user_id)
}
}
@ -122,22 +98,14 @@ impl ExtendedAppointment {
mod tests {
use super::*;
use teos_common::appointment::Appointment;
use teos_common::cryptography::get_random_bytes;
use teos_common::test_utils::get_random_user_id;
use crate::test_utils::generate_uuid;
#[test]
fn test_get_summary() {
let locator = Locator::from_slice(&get_random_bytes(16)).unwrap();
let user_id = get_random_user_id();
let signature = String::new();
let a = Appointment::new(locator, get_random_bytes(32), 42);
let e = ExtendedAppointment::new(a, user_id, signature, 21);
let s = e.get_summary();
assert_eq!(e.locator(), s.locator);
assert_eq!(e.user_id, s.user_id);
fn test_uuid_ser_deser() {
let original_uuid = generate_uuid();
assert_eq!(
UUID::from_slice(&original_uuid.to_vec()).unwrap(),
original_uuid
);
}
}

View file

@ -1,13 +1,11 @@
//! Logic related to the Gatekeeper, the component in charge of managing access to the tower resources.
use std::collections::{HashMap, HashSet};
use std::iter::FromIterator;
use lightning::chain;
use std::collections::HashMap;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::{Arc, Mutex};
use lightning::chain;
use teos_common::appointment::compute_appointment_slots;
use teos_common::appointment::{compute_appointment_slots, Locator};
use teos_common::constants::ENCRYPTED_BLOB_MAX_SIZE;
use teos_common::cryptography;
use teos_common::receipts::RegistrationReceipt;
@ -17,7 +15,7 @@ use crate::dbm::DBM;
use crate::extended_appointment::{ExtendedAppointment, UUID};
/// Data regarding a user subscription with the tower.
#[derive(Debug, Clone, PartialEq, Eq)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct UserInfo {
/// Number of appointment slots available for a given user.
pub(crate) available_slots: u32,
@ -25,8 +23,6 @@ pub(crate) struct UserInfo {
pub(crate) subscription_start: u32,
/// Block height where the user subscription expires.
pub(crate) subscription_expiry: u32,
/// Map of appointment ids and the how many slots they take from the subscription.
pub(crate) appointments: HashMap<UUID, u32>,
}
impl UserInfo {
@ -36,22 +32,6 @@ impl UserInfo {
available_slots,
subscription_start,
subscription_expiry,
appointments: HashMap::new(),
}
}
/// Creates a new [UserInfo] instance with some associated appointments.
pub fn with_appointments(
available_slots: u32,
subscription_start: u32,
subscription_expiry: u32,
appointments: HashMap<UUID, u32>,
) -> Self {
UserInfo {
available_slots,
subscription_start,
subscription_expiry,
appointments,
}
}
}
@ -134,8 +114,9 @@ impl Gatekeeper {
}
/// Gets the data held by the tower about a given user.
pub(crate) fn get_user_info(&self, user_id: UserId) -> Option<UserInfo> {
self.registered_users.lock().unwrap().get(&user_id).cloned()
pub(crate) fn get_user_info(&self, user_id: UserId) -> Option<(UserInfo, Vec<Locator>)> {
let info = self.registered_users.lock().unwrap().get(&user_id).cloned();
info.map(|info| (info, self.dbm.lock().unwrap().load_user_locators(user_id)))
}
/// Authenticates a user.
@ -219,7 +200,13 @@ impl Gatekeeper {
// For updates, the difference between the existing appointment size and the update is computed.
let mut registered_users = self.registered_users.lock().unwrap();
let user_info = registered_users.get_mut(&user_id).unwrap();
let used_slots = user_info.appointments.get(&uuid).map_or(0, |x| *x);
let used_blob_size = self
.dbm
.lock()
.unwrap()
.get_appointment_length(uuid)
.unwrap_or(0);
let used_slots = compute_appointment_slots(used_blob_size, ENCRYPTED_BLOB_MAX_SIZE);
let required_slots =
compute_appointment_slots(appointment.encrypted_blob().len(), ENCRYPTED_BLOB_MAX_SIZE);
@ -228,7 +215,6 @@ impl Gatekeeper {
if diff <= user_info.available_slots as i64 {
// Filling / freeing slots depending on whether this is an update or not, and if it is bigger or smaller
// than the old appointment
user_info.appointments.insert(uuid, required_slots);
user_info.available_slots = (user_info.available_slots as i64 - diff) as u32;
self.dbm.lock().unwrap().update_user(user_id, user_info);
@ -258,56 +244,50 @@ impl Gatekeeper {
/// Gets a map of outdated users. Outdated users are those whose subscription has expired and the renewal grace period
/// has already passed ([expiry_delta](Self::expiry_delta)).
pub(crate) fn get_outdated_users(&self, block_height: u32) -> HashMap<UserId, HashSet<UUID>> {
let registered_users = self.registered_users.lock().unwrap().clone();
registered_users
.into_iter()
.filter(|(_, info)| block_height == info.subscription_expiry + self.expiry_delta)
.map(|(id, info)| (id, info.appointments.keys().cloned().collect()))
pub(crate) fn get_outdated_users(&self, block_height: u32) -> Vec<UserId> {
self.registered_users
.lock()
.unwrap()
.iter()
// NOTE: Ideally there won't be a user with `block_height > subscription_expiry + expiry_delta`, but
// this might happen if we skip a couple of block connections due to a force update.
.filter(|(_, info)| block_height >= info.subscription_expiry + self.expiry_delta)
.map(|(user_id, _)| *user_id)
.collect()
}
/// Gets a set of outdated user ids.
pub(crate) fn get_outdated_user_ids(&self, block_height: u32) -> HashSet<UserId> {
self.get_outdated_users(block_height)
.keys()
.cloned()
.collect()
}
/// Get a map of outdated appointments (from any user).
pub(crate) fn get_outdated_appointments(&self, block_height: u32) -> HashSet<UUID> {
HashSet::from_iter(
self.get_outdated_users(block_height)
.into_values()
.flatten(),
)
}
/// Deletes a collection of appointments from the users' subscriptions (from memory only)
/// and updates the available_slots count for the given user.
/// Deletes these appointments from the database and updates the user's information.
///
/// Notice appointments are only de-linked from users, but not actually removed. This is because the [Gatekeeper]
/// does not actually hold any [ExtendedAppointment](crate::extended_appointment::ExtendedAppointment) data,
/// just references to them.
pub(crate) fn delete_appointments_from_memory(
&self,
appointments: &HashMap<UUID, UserId>,
) -> HashMap<UserId, UserInfo> {
let mut updated_users = HashMap::new();
let mut registered_users = self.registered_users.lock().unwrap();
/// If `refund` is set, the appointments owners will get their slots refunded back.
///
/// DISCUSS: When `refund` is `false` we don't give back the slots to the user for the deleted appointments.
/// This is to discourage misbehavior (sending bad appointments, either non-decryptable or rejected by the network).
pub(crate) fn delete_appointments(&self, appointments: Vec<UUID>, refund: bool) {
let mut dbm = self.dbm.lock().unwrap();
for (uuid, user_id) in appointments {
// Remove the appointment from the appointment list and update the available slots
if let Some(user_info) = registered_users.get_mut(user_id) {
if let Some(x) = user_info.appointments.remove(uuid) {
user_info.available_slots += x;
}
updated_users.insert(*user_id, user_info.clone());
};
let updated_users = if refund {
let mut updated_users = HashMap::new();
let mut registered_users = self.registered_users.lock().unwrap();
// Give back the consumed slots to each user.
for uuid in appointments.iter() {
let (user_id, blob_size) = dbm.get_appointment_user_and_length(*uuid).unwrap();
registered_users.get_mut(&user_id).unwrap().available_slots +=
compute_appointment_slots(blob_size, ENCRYPTED_BLOB_MAX_SIZE);
updated_users.insert(user_id, registered_users[&user_id]);
}
updated_users
} else {
// No updated users.
HashMap::new()
};
// An optimization for the case when only one appointment is being deleted without refunding.
// This avoids creating a DB transaction for a single query.
if appointments.len() == 1 && updated_users.is_empty() {
dbm.remove_appointment(appointments[0])
} else {
dbm.batch_remove_appointments(&appointments, &updated_users);
}
updated_users
}
}
@ -324,12 +304,17 @@ impl chain::Listen for Gatekeeper {
log::info!("New block received: {}", header.block_hash());
// Expired user deletion is delayed. Users are deleted when their subscription is outdated, not expired.
let outdated_users = self.get_outdated_user_ids(height);
let outdated_users = self.get_outdated_users(height);
if !outdated_users.is_empty() {
self.registered_users
.lock()
.unwrap()
.retain(|id, _| !outdated_users.contains(id));
// Remove the outdated users from memory first.
{
let mut registered_users = self.registered_users.lock().unwrap();
// Removing each outdated user in a loop is more efficient than retaining non-outdated users
// because retaining would loop over all the available users which is always more than the outdated ones.
for outdated_user in outdated_users.iter() {
registered_users.remove(outdated_user);
}
}
self.dbm.lock().unwrap().batch_remove_users(&outdated_users);
}
@ -351,13 +336,13 @@ impl chain::Listen for Gatekeeper {
mod tests {
use super::*;
use crate::test_utils::{
generate_dummy_appointment, generate_dummy_appointment_with_user, generate_uuid, Blockchain,
};
use crate::test_utils::{generate_dummy_appointment_with_user, get_random_tracker, Blockchain};
use lightning::chain::Listen;
use teos_common::cryptography::{get_random_bytes, get_random_keypair};
use teos_common::test_utils::get_random_user_id;
use crate::responder::ConfirmationStatus;
const SLOTS: u32 = 21;
const DURATION: u32 = 500;
const EXPIRY_DELTA: u32 = 42;
@ -380,21 +365,11 @@ mod tests {
&self.registered_users
}
pub(crate) fn add_outdated_user(
&self,
user_id: UserId,
outdates_at: u32,
appointments: Option<Vec<UUID>>,
) {
pub(crate) fn add_outdated_user(&self, user_id: UserId, outdates_at: u32) {
self.add_update_user(user_id).unwrap();
let mut registered_users = self.registered_users.lock().unwrap();
let user = registered_users.get_mut(&user_id).unwrap();
user.subscription_expiry = outdates_at - self.expiry_delta;
if let Some(uuids) = appointments {
for uuid in uuids.iter() {
user.appointments.insert(*uuid, 1);
}
}
}
}
@ -567,25 +542,32 @@ mod tests {
let available_slots = gatekeeper
.add_update_appointment(user_id, uuid, &appointment)
.unwrap();
// Simulate the watcher adding the appointment in the database.
gatekeeper
.dbm
.lock()
.unwrap()
.store_appointment(uuid, &appointment)
.unwrap();
assert!(gatekeeper.registered_users.lock().unwrap()[&user_id]
.appointments
.contains_key(&uuid));
let (_, user_locators) = gatekeeper.get_user_info(user_id).unwrap();
assert!(user_locators.contains(&appointment.locator()));
assert_eq!(slots_before, available_slots + 1);
// Slots should have been updated in the database too. Notice the appointment won't be there yet
// given the Watcher is responsible for adding it, and it will do so after calling this method
// Slots should have been updated in the database too.
let mut loaded_user = gatekeeper.dbm.lock().unwrap().load_user(user_id).unwrap();
assert_eq!(loaded_user.available_slots, available_slots);
// Adding the exact same appointment should leave the slots count unchanged
// Adding the exact same appointment should leave the slots count unchanged.
// We don't really need to update the appointment in the DB since it's the very same appointment.
let mut updated_slot_count = gatekeeper
.add_update_appointment(user_id, uuid, &appointment)
.unwrap();
assert!(gatekeeper.registered_users.lock().unwrap()[&user_id]
.appointments
.contains_key(&uuid));
let (_, user_locators) = gatekeeper.get_user_info(user_id).unwrap();
assert!(user_locators.contains(&appointment.locator()));
assert_eq!(updated_slot_count, available_slots);
loaded_user = gatekeeper.dbm.lock().unwrap().load_user(user_id).unwrap();
assert_eq!(loaded_user.available_slots, updated_slot_count);
@ -595,10 +577,18 @@ mod tests {
updated_slot_count = gatekeeper
.add_update_appointment(user_id, uuid, &bigger_appointment)
.unwrap();
assert!(gatekeeper.registered_users.lock().unwrap()[&user_id]
.appointments
.contains_key(&uuid));
// Simulate the watcher updating the appointment in the database.
gatekeeper
.dbm
.lock()
.unwrap()
.update_appointment(uuid, &bigger_appointment)
.unwrap();
let (_, user_locators) = gatekeeper.get_user_info(user_id).unwrap();
assert!(user_locators.contains(&appointment.locator()));
assert_eq!(updated_slot_count, available_slots - 1);
loaded_user = gatekeeper.dbm.lock().unwrap().load_user(user_id).unwrap();
assert_eq!(loaded_user.available_slots, updated_slot_count);
@ -606,26 +596,43 @@ mod tests {
updated_slot_count = gatekeeper
.add_update_appointment(user_id, uuid, &appointment)
.unwrap();
assert!(gatekeeper.registered_users.lock().unwrap()[&user_id]
.appointments
.contains_key(&uuid));
// Simulate the watcher updating the appointment in the database.
gatekeeper
.dbm
.lock()
.unwrap()
.update_appointment(uuid, &appointment)
.unwrap();
let (_, user_locators) = gatekeeper.get_user_info(user_id).unwrap();
assert!(user_locators.contains(&appointment.locator()));
assert_eq!(updated_slot_count, available_slots);
loaded_user = gatekeeper.dbm.lock().unwrap().load_user(user_id).unwrap();
assert_eq!(loaded_user.available_slots, updated_slot_count);
// Adding an appointment with a different uuid should not count as an update
let new_uuid = generate_uuid();
let (uuid, appointment) = generate_dummy_appointment_with_user(user_id, None);
updated_slot_count = gatekeeper
.add_update_appointment(user_id, new_uuid, &appointment)
.add_update_appointment(user_id, uuid, &appointment)
.unwrap();
assert!(gatekeeper.registered_users.lock().unwrap()[&user_id]
.appointments
.contains_key(&new_uuid));
// Simulate the watcher adding the appointment in the database.
gatekeeper
.dbm
.lock()
.unwrap()
.store_appointment(uuid, &appointment)
.unwrap();
let (_, user_locators) = gatekeeper.get_user_info(user_id).unwrap();
assert!(user_locators.contains(&appointment.locator()));
assert_eq!(updated_slot_count, available_slots - 1);
loaded_user = gatekeeper.dbm.lock().unwrap().load_user(user_id).unwrap();
assert_eq!(loaded_user.available_slots, updated_slot_count);
// Finally, trying to add an appointment when the user has no enough slots should fail
let (uuid, appointment) = generate_dummy_appointment_with_user(user_id, None);
gatekeeper
.registered_users
.lock()
@ -634,9 +641,10 @@ mod tests {
.unwrap()
.available_slots = 0;
assert!(matches!(
gatekeeper.add_update_appointment(user_id, generate_uuid(), &appointment),
gatekeeper.add_update_appointment(user_id, uuid, &appointment),
Err(NotEnoughSlots)
));
// The entry in the database should remain unchanged in this case
loaded_user = gatekeeper.dbm.lock().unwrap().load_user(user_id).unwrap();
assert_eq!(loaded_user.available_slots, updated_slot_count);
@ -682,144 +690,175 @@ mod tests {
// Initially, there are not outdated users, so querying any block height should return an empty map
for i in 0..start_height {
assert_eq!(gatekeeper.get_outdated_users(i).len(), 0);
assert_eq!(gatekeeper.get_outdated_users(i), vec![]);
}
// Adding a user whose subscription is outdated should return an entry
let user_id = get_random_user_id();
gatekeeper.add_update_user(user_id).unwrap();
// Add also an appointment so we can check the returned data
let appointment = generate_dummy_appointment(None);
let uuid = generate_uuid();
gatekeeper
.add_update_appointment(user_id, uuid, &appointment)
.unwrap();
// Check that data is not yet outdated
assert_eq!(gatekeeper.get_outdated_users(start_height).len(), 0);
assert_eq!(gatekeeper.get_outdated_users(start_height), vec![]);
// Add an outdated user and check again
gatekeeper.add_outdated_user(user_id, start_height, None);
let outdated_users = gatekeeper.get_outdated_users(start_height);
assert_eq!(outdated_users.len(), 1);
assert_eq!(outdated_users[&user_id], HashSet::from_iter([uuid]));
gatekeeper.add_outdated_user(user_id, start_height);
assert_eq!(gatekeeper.get_outdated_users(start_height), vec![user_id]);
}
#[test]
fn test_get_outdated_appointments() {
let start_height = START_HEIGHT as u32 + EXPIRY_DELTA;
let gatekeeper = init_gatekeeper(&Blockchain::default().with_height(start_height as usize));
// get_outdated_appointments returns a list of appointments that were outdated at a given block height, indistinguishably of their user.
// If there are no outdated users, there cannot be outdated appointments
for i in 0..start_height {
assert_eq!(gatekeeper.get_outdated_appointments(i).len(), 0);
}
// Adding data about different users and appointments should return a flattened list of appointments
let user1_id = get_random_user_id();
let user2_id = get_random_user_id();
let uuid1 = generate_uuid();
let uuid2 = generate_uuid();
// Manually set the user expiry for the test
for (user_id, uuid) in [(user1_id, uuid1), (user2_id, uuid2)] {
gatekeeper.add_outdated_user(user_id, start_height, Some(Vec::from_iter([uuid])));
}
let outdated_appointments = gatekeeper.get_outdated_appointments(start_height);
assert_eq!(outdated_appointments.len(), 2);
assert!(outdated_appointments.contains(&uuid1));
assert!(outdated_appointments.contains(&uuid2));
}
#[test]
fn test_delete_appointments_from_memory() {
fn test_delete_appointments_without_refund() {
let gatekeeper = init_gatekeeper(&Blockchain::default().with_height(START_HEIGHT));
let n_users = 100;
let n_apps = 10;
let mut uuids_to_delete = Vec::new();
let mut rest = Vec::new();
let mut trackers = Vec::new();
let mut users_info = HashMap::new();
// delete_appointments will remove a list of appointments from the Gatekeeper (as long as they exist)
let mut all_appointments = HashMap::new();
let mut to_be_deleted = HashMap::new();
let mut rest = HashMap::new();
for i in 1..11 {
for _ in 0..n_users {
let user_id = get_random_user_id();
let uuid = generate_uuid();
all_appointments.insert(uuid, user_id);
if i % 2 == 0 {
to_be_deleted.insert(uuid, user_id);
} else {
rest.insert(uuid, user_id);
}
}
// Calling the method with unknown data should work but do nothing
assert!(gatekeeper.registered_users.lock().unwrap().is_empty());
assert!(gatekeeper
.delete_appointments_from_memory(&all_appointments)
.is_empty());
// If there's matching data in the gatekeeper it should be deleted
for (uuid, user_id) in to_be_deleted.iter() {
gatekeeper.add_update_user(*user_id).unwrap();
gatekeeper
.add_update_appointment(*user_id, *uuid, &generate_dummy_appointment(None))
.unwrap();
}
// Check before deleting
assert_eq!(gatekeeper.registered_users.lock().unwrap().len(), 5);
for (uuid, user_id) in to_be_deleted.iter() {
assert!(gatekeeper.registered_users.lock().unwrap()[user_id]
.appointments
.contains_key(uuid));
// The slot count should be decreased now too (both in memory and in the database)
assert_ne!(
gatekeeper.registered_users.lock().unwrap()[user_id].available_slots,
gatekeeper.subscription_slots
);
assert_ne!(
gatekeeper.add_update_user(user_id).unwrap();
for i in 0..n_apps {
let (uuid, appointment) = generate_dummy_appointment_with_user(user_id, None);
gatekeeper
.add_update_appointment(user_id, uuid, &appointment)
.unwrap();
// Add the appointment to the database. This is normally done by the Watcher.
gatekeeper
.dbm
.lock()
.unwrap()
.load_user(*user_id)
.store_appointment(uuid, &appointment)
.unwrap();
if i % 2 == 0 {
uuids_to_delete.push(uuid);
} else {
rest.push(uuid);
}
// Also trigger some of these appointments as trackers.
if i % 5 == 0 {
gatekeeper
.dbm
.lock()
.unwrap()
.store_tracker(
uuid,
&get_random_tracker(user_id, ConfirmationStatus::ConfirmedIn(42)),
)
.unwrap();
trackers.push(uuid);
}
}
users_info.insert(user_id, gatekeeper.get_user_info(user_id).unwrap().0);
}
// Delete these appointments without refunding their owners.
gatekeeper.delete_appointments(uuids_to_delete.clone(), false);
for uuid in uuids_to_delete.clone() {
assert!(!gatekeeper.dbm.lock().unwrap().appointment_exists(uuid));
}
for uuid in rest {
assert!(gatekeeper.dbm.lock().unwrap().appointment_exists(uuid));
}
for uuid in trackers {
if uuids_to_delete.contains(&uuid) {
// The tracker should be deleted as well.
assert!(!gatekeeper.dbm.lock().unwrap().tracker_exists(uuid));
} else {
assert!(gatekeeper.dbm.lock().unwrap().tracker_exists(uuid));
}
}
for (user_id, user_info_before_deletion) in users_info {
// Since `refund` was false, the users' slots should not have changed after deleting appointments.
let (user_info_after_deletion, _) = gatekeeper.get_user_info(user_id).unwrap();
assert_eq!(user_info_after_deletion, user_info_before_deletion);
}
}
#[test]
fn test_delete_appointments_with_refund() {
let gatekeeper = init_gatekeeper(&Blockchain::default().with_height(START_HEIGHT));
let n_users = 100;
let n_apps = 10;
let mut uuids_to_delete = Vec::new();
let mut rest = Vec::new();
let mut trackers = Vec::new();
let mut users_remaining_slots = HashMap::new();
for _ in 0..n_users {
let user_id = get_random_user_id();
gatekeeper.add_update_user(user_id).unwrap();
let mut user_remaining_slots =
gatekeeper.get_user_info(user_id).unwrap().0.available_slots;
for i in 0..n_apps {
let (uuid, appointment) = generate_dummy_appointment_with_user(user_id, None);
gatekeeper
.add_update_appointment(user_id, uuid, &appointment)
.unwrap();
// Add the appointment to the database. This is normally done by the Watcher.
gatekeeper
.dbm
.lock()
.unwrap()
.available_slots,
gatekeeper.subscription_slots
);
}
for (_, user_id) in rest.iter() {
assert!(!gatekeeper
.registered_users
.lock()
.unwrap()
.contains_key(user_id));
.store_appointment(uuid, &appointment)
.unwrap();
if i % 2 == 0 {
// We don't reduce the remaining slots for the appointments which are
// going to delete since we will refund their owners.
uuids_to_delete.push(uuid);
} else {
rest.push(uuid);
user_remaining_slots -= compute_appointment_slots(
appointment.encrypted_blob().len(),
ENCRYPTED_BLOB_MAX_SIZE,
);
}
// Also trigger some of these appointments as trackers.
if i % 5 == 0 {
gatekeeper
.dbm
.lock()
.unwrap()
.store_tracker(
uuid,
&get_random_tracker(user_id, ConfirmationStatus::ConfirmedIn(42)),
)
.unwrap();
trackers.push(uuid);
}
}
users_remaining_slots.insert(user_id, user_remaining_slots);
}
// And after
gatekeeper.delete_appointments_from_memory(&all_appointments);
for (uuid, user_id) in to_be_deleted.iter() {
assert!(!gatekeeper.registered_users.lock().unwrap()[user_id]
.appointments
.contains_key(uuid));
// Delete these appointments and refund their owners their slots back.
gatekeeper.delete_appointments(uuids_to_delete.clone(), true);
// The slot count is back to default
for uuid in uuids_to_delete.clone() {
assert!(!gatekeeper.dbm.lock().unwrap().appointment_exists(uuid));
}
for uuid in rest {
assert!(gatekeeper.dbm.lock().unwrap().appointment_exists(uuid));
}
for uuid in trackers {
if uuids_to_delete.contains(&uuid) {
// The tracker should be deleted as well.
assert!(!gatekeeper.dbm.lock().unwrap().tracker_exists(uuid));
} else {
assert!(gatekeeper.dbm.lock().unwrap().tracker_exists(uuid));
}
}
for (user_id, correct_remaining_slots) in users_remaining_slots {
let remaining_slots_from_db =
gatekeeper.get_user_info(user_id).unwrap().0.available_slots;
assert_eq!(remaining_slots_from_db, correct_remaining_slots);
assert_eq!(
gatekeeper.registered_users.lock().unwrap()[user_id].available_slots,
gatekeeper.subscription_slots
gatekeeper.registered_users.lock().unwrap()[&user_id].available_slots,
correct_remaining_slots
);
}
for (_, user_id) in rest.iter() {
assert!(!gatekeeper
.registered_users
.lock()
.unwrap()
.contains_key(user_id));
}
}
#[test]
@ -835,7 +874,7 @@ mod tests {
let user3_id = get_random_user_id();
for user_id in &[user1_id, user2_id, user3_id] {
gatekeeper.add_outdated_user(*user_id, chain.tip().height + 1, None)
gatekeeper.add_outdated_user(*user_id, chain.tip().height + 1)
}
// Connect a new block. Outdated users are deleted

View file

@ -48,6 +48,7 @@ where
{
let mut last_n_blocks = Vec::with_capacity(n);
for _ in 0..n {
log::debug!("Fetching block #{}", last_known_block.height);
let block = poller.fetch_block(&last_known_block).await?;
last_known_block = poller.look_up_previous_header(&last_known_block).await?;
last_n_blocks.push(block);
@ -309,8 +310,8 @@ async fn main() {
let shutdown_signal_tor = shutdown_signal_rpc_api.clone();
// The ordering here actually matters. Listeners are called by order, and we want the gatekeeper to be called
// last, so both the Watcher and the Responder can query the necessary data from it during data deletion.
let listener = &(watcher.clone(), &(responder, gatekeeper));
// first so it updates the users' states and both the Watcher and the Responder operate only on registered users.
let listener = &(gatekeeper, &(watcher.clone(), responder));
let cache = &mut UnboundedCache::new();
let spv_client = SpvClient::new(tip, poller, cache, listener);
let mut chain_monitor = ChainMonitor::new(

File diff suppressed because it is too large Load diff

View file

@ -323,7 +323,7 @@ pub(crate) fn generate_dummy_appointment_with_user(
let mut app = generate_dummy_appointment(dispute_txid);
app.user_id = user_id;
(UUID::new(app.locator(), user_id), app)
(app.uuid(), app)
}
pub(crate) fn get_random_breach() -> Breach {
@ -341,17 +341,15 @@ pub(crate) fn get_random_tracker(
TransactionTracker::new(breach, user_id, status)
}
pub(crate) fn store_appointment_and_fks_to_db(
dbm: &DBM,
uuid: UUID,
appointment: &ExtendedAppointment,
) {
pub(crate) fn store_appointment_and_its_user(dbm: &DBM, appointment: &ExtendedAppointment) {
dbm.store_user(
appointment.user_id,
&UserInfo::new(AVAILABLE_SLOTS, SUBSCRIPTION_START, SUBSCRIPTION_EXPIRY),
)
.unwrap();
dbm.store_appointment(uuid, appointment).unwrap();
// It's ok if the user is already stored.
.ok();
dbm.store_appointment(appointment.uuid(), appointment)
.unwrap();
}
pub(crate) async fn get_last_n_blocks(chain: &mut Blockchain, n: usize) -> Vec<ValidatedBlock> {

View file

@ -9,7 +9,7 @@ use lightning_block_sync::poll::ValidatedBlock;
use teos_common::appointment::Locator;
/// A trait implemented by types that can be used as key in a [TxIndex].
pub trait Key: Hash {
pub trait Key: Hash + Eq {
fn from_txid(txid: Txid) -> Self;
}
@ -79,8 +79,8 @@ impl Value for Transaction {
/// Data structure used to index locators computed from parsed blocks.
///
/// Holds up to `size` blocks with their corresponding computed [Locator]s.
#[derive(Debug)]
pub struct TxIndex<K, V> {
#[derive(Debug, PartialEq, Eq)]
pub struct TxIndex<K: Key, V: Value> {
/// A [K]:[V] map.
index: HashMap<K, V>,
/// Vector of block hashes covered by the index.
@ -95,7 +95,7 @@ pub struct TxIndex<K, V> {
impl<K, V> TxIndex<K, V>
where
K: Key + std::cmp::Eq + Copy,
K: Key + Copy,
V: Value + Clone,
Self: Sized,
{
@ -143,11 +143,6 @@ where
self.index.get(k)
}
/// Checks whether the index contains a certain key.
pub fn contains_key(&self, k: &K) -> bool {
self.index.contains_key(k)
}
/// Checks if the index if full.
pub fn is_full(&self) -> bool {
self.blocks.len() > self.size
@ -175,7 +170,7 @@ where
if self.is_full() {
// Avoid logging during bootstrap
log::info!("New block added to index: {}", block_header.block_hash());
log::debug!("New block added to index: {}", block_header.block_hash());
self.tip += 1;
self.remove_oldest_block();
}
@ -204,11 +199,11 @@ where
let ks = self.tx_in_block.remove(&h).unwrap();
self.index.retain(|k, _| !ks.contains(k));
log::info!("Oldest block removed from index: {h}");
log::debug!("Oldest block removed from index: {h}");
}
}
impl<K: std::fmt::Debug, V: std::fmt::Debug> fmt::Display for TxIndex<K, V> {
impl<K: std::fmt::Debug + Key, V: std::fmt::Debug + Value> fmt::Display for TxIndex<K, V> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(
f,
@ -240,6 +235,10 @@ mod tests {
pub fn blocks(&self) -> &VecDeque<BlockHash> {
&self.blocks
}
pub fn contains_key(&self, k: &K) -> bool {
self.index.contains_key(k)
}
}
#[tokio::test]
@ -304,7 +303,7 @@ mod tests {
);
let fake_hash = BlockHash::default();
assert!(matches!(cache.get_height(&fake_hash), None));
assert!(cache.get_height(&fake_hash).is_none());
}
#[tokio::test]

File diff suppressed because it is too large Load diff