GlcoinCore/rust-teos
1
0
Fork 0
mirror of https://github.com/talaia-labs/rust-teos.git synced 2026-08-13 12:33:22 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge pull request #190 from mariocynicys/better-mem-cpu-usage

Browse source 
optimizing memory usage
This commit is contained in:
Sergi Delgado Segura 2023-09-14 13:11:15 -04:00 • committed by GitHub
commit 658fcca0ce
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
11 changed files with 1989 additions and 2468 deletions
Download patch file Download diff file Expand all files Collapse all files

19
teos-common/src/test_utils.rs
View file

@ -5,10 +5,9 @@ use rand::distributions::Standard;
use rand::prelude::Distribution;
use rand::Rng;
use bitcoin::consensus;
use bitcoin::hashes::Hash;
use bitcoin::secp256k1::SecretKey;
use bitcoin::Txid;
use bitcoin::{consensus, Script, Transaction, TxOut, Txid};
use crate::appointment::{Appointment, Locator};
use crate::cryptography;
@ -32,6 +31,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,
@ -42,7 +47,15 @@ pub fn generate_random_appointment(dispute_txid: Option<&Txid>) -> Appointment {
};
let tx_bytes = Vec::from_hex(TX_HEX).unwrap();
let penalty_tx = consensus::deserialize(&tx_bytes).unwrap();
let mut penalty_tx: Transaction = consensus::deserialize(&tx_bytes).unwrap();
// Append a random-sized OP_RETURN to make each transcation random in size.
penalty_tx.output.push(TxOut {
value: 0,
script_pubkey: Script::new_op_return(&cryptography::get_random_bytes(
get_random_int::<usize>() % 81,
)),
});
let mut raw_locator: [u8; 16] = cryptography::get_random_bytes(16).try_into().unwrap();
raw_locator.copy_from_slice(&dispute_txid[..16]);

21
teos/src/api/http.rs
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,

77
teos/src/api/internal.rs
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
{

717
teos/src/dbm.rs
View file

@ -1,7 +1,7 @@
//! Logic related to the tower database manager (DBM), component in charge of persisting data on disk.
//!
use std::collections::{HashMap, HashSet};
use std::collections::HashMap;
use std::iter::FromIterator;
use std::path::PathBuf;
use std::str::FromStr;
@ -14,16 +14,15 @@ use bitcoin::hashes::Hash;
use bitcoin::secp256k1::SecretKey;
use bitcoin::BlockHash;
use teos_common::appointment::{compute_appointment_slots, Appointment, Locator};
use teos_common::constants::ENCRYPTED_BLOB_MAX_SIZE;
use teos_common::appointment::{Appointment, Locator};
use teos_common::dbm::{DatabaseConnection, DatabaseManager, Error};
use teos_common::UserId;
use crate::extended_appointment::{ExtendedAppointment, UUID};
use crate::gatekeeper::UserInfo;
use crate::responder::{ConfirmationStatus, TransactionTracker};
use crate::responder::{ConfirmationStatus, PenaltySummary, TransactionTracker};
const TABLES: [&str; 5] = [
const TABLES: [&str; 6] = [
"CREATE TABLE IF NOT EXISTS users (
user_id INT PRIMARY KEY,
available_slots INT NOT NULL,
@ -59,6 +58,9 @@ const TABLES: [&str; 5] = [
"CREATE TABLE IF NOT EXISTS keys (
id INTEGER PRIMARY KEY AUTOINCREMENT,
key INT NOT NULL
)",
"CREATE INDEX IF NOT EXISTS locators_index ON appointments (
locator
)",
];
@ -139,27 +141,21 @@ impl DBM {
}
}
/// Loads the associated appointments ([Appointment]) of a given user ([UserInfo]).
pub(crate) fn load_user_appointments(&self, user_id: UserId) -> HashMap<UUID, u32> {
/// Loads the associated locators ([Locator]) of a given user ([UserId]).
pub(crate) fn load_user_locators(&self, user_id: UserId) -> Vec<Locator> {
let mut stmt = self
.connection
.prepare("SELECT UUID, encrypted_blob FROM appointments WHERE user_id=(?)")
.prepare("SELECT locator FROM appointments WHERE user_id=(?)")
.unwrap();
let mut rows = stmt.query([user_id.to_vec()]).unwrap();
let mut appointments = HashMap::new();
while let Ok(Some(inner_row)) = rows.next() {
let raw_uuid: Vec<u8> = inner_row.get(0).unwrap();
let uuid = UUID::from_slice(&raw_uuid[0..20]).unwrap();
let e_blob: Vec<u8> = inner_row.get(1).unwrap();
appointments.insert(
uuid,
compute_appointment_slots(e_blob.len(), ENCRYPTED_BLOB_MAX_SIZE),
);
}
appointments
stmt.query_map([user_id.to_vec()], |row| {
let raw_locator: Vec<u8> = row.get(0).unwrap();
let locator = Locator::from_slice(&raw_locator).unwrap();
Ok(locator)
})
.unwrap()
.map(|res| res.unwrap())
.collect()
}
/// Loads all users from the database.
@ -178,22 +174,14 @@ impl DBM {
let start = row.get(2).unwrap();
let expiry = row.get(3).unwrap();
users.insert(
user_id,
UserInfo::with_appointments(
slots,
start,
expiry,
self.load_user_appointments(user_id),
),
);
users.insert(user_id, UserInfo::new(slots, start, expiry));
}
users
}
/// Removes some users from the database in batch.
pub(crate) fn batch_remove_users(&mut self, users: &HashSet<UserId>) -> usize {
pub(crate) fn batch_remove_users(&mut self, users: &Vec<UserId>) -> usize {
let limit = self.connection.limit(Limit::SQLITE_LIMIT_VARIABLE_NUMBER) as usize;
let tx = self.connection.transaction().unwrap();
let iter = users
@ -219,6 +207,24 @@ impl DBM {
(users.len() as f64 / limit as f64).ceil() as usize
}
/// Get the number of stored appointments.
pub(crate) fn get_appointments_count(&self) -> usize {
let mut stmt = self
.connection
.prepare("SELECT COUNT(*) FROM appointments as a LEFT JOIN trackers as t ON a.UUID=t.UUID WHERE t.UUID IS NULL")
.unwrap();
stmt.query_row([], |row| row.get(0)).unwrap()
}
/// Get the number of stored trackers.
pub(crate) fn get_trackers_count(&self) -> usize {
let mut stmt = self
.connection
.prepare("SELECT COUNT(*) FROM trackers")
.unwrap();
stmt.query_row([], |row| row.get(0)).unwrap()
}
/// Stores an [Appointment] into the database.
pub(crate) fn store_appointment(
&self,
@ -250,7 +256,11 @@ impl DBM {
}
/// Updates an existing [Appointment] in the database.
pub(crate) fn update_appointment(&self, uuid: UUID, appointment: &ExtendedAppointment) {
pub(crate) fn update_appointment(
&self,
uuid: UUID,
appointment: &ExtendedAppointment,
) -> Result<(), Error> {
// DISCUSS: Check what fields we'd like to make updatable. e_blob and signature are the obvious, to_self_delay and start_block may not be necessary (or even risky)
let query =
"UPDATE appointments SET encrypted_blob=(?1), to_self_delay=(?2), user_signature=(?3), start_block=(?4) WHERE UUID=(?5)";
@ -266,9 +276,11 @@ impl DBM {
) {
Ok(_) => {
log::debug!("Appointment successfully updated: {uuid}");
Ok(())
}
Err(_) => {
log::error!("Appointment not found, data cannot be updated: {uuid}");
Err(e) => {
log::error!("Appointment not found, data cannot be updated: {uuid}. Error: {e:?}");
Err(e)
}
}
}
@ -305,6 +317,15 @@ impl DBM {
.ok()
}
/// Check if an appointment with `uuid` exists.
pub(crate) fn appointment_exists(&self, uuid: UUID) -> bool {
self.connection
.prepare("SELECT UUID FROM appointments WHERE UUID=(?)")
.unwrap()
.exists([uuid.to_vec()])
.unwrap()
}
/// Loads appointments from the database. If a locator is given, this method loads only the appointments
/// matching this locator. If no locator is given, all the appointments in the database would be returned.
pub(crate) fn load_appointments(
@ -352,6 +373,32 @@ impl DBM {
appointments
}
/// Gets the length of an appointment (the length of `appointment.encrypted_blob`).
pub(crate) fn get_appointment_length(&self, uuid: UUID) -> Option<usize> {
let mut stmt = self
.connection
.prepare("SELECT length(encrypted_blob) FROM appointments WHERE UUID=(?)")
.unwrap();
stmt.query_row([uuid.to_vec()], |row| row.get(0)).ok()
}
/// Gets the [`UserId`] of the owner of the appointment along with the appointment
/// length (same as [DBM::get_appointment_length]) for `uuid`.
pub(crate) fn get_appointment_user_and_length(&self, uuid: UUID) -> Option<(UserId, usize)> {
let mut stmt = self
.connection
.prepare("SELECT user_id, length(encrypted_blob) FROM appointments WHERE UUID=(?)")
.unwrap();
stmt.query_row([uuid.to_vec()], |row| {
let raw_userid: Vec<u8> = row.get(0).unwrap();
let length = row.get(1).unwrap();
Ok((UserId::from_slice(&raw_userid).unwrap(), length))
})
.ok()
}
/// Removes an [Appointment] from the database.
pub(crate) fn remove_appointment(&self, uuid: UUID) {
let query = "DELETE FROM appointments WHERE UUID=(?)";
@ -365,11 +412,12 @@ impl DBM {
}
}
/// Removes some appointments from the database in batch and updates the associated users giving back
/// the freed appointment slots
/// Removes some appointments from the database in batch and updates the associated users
/// (giving back freed appointment slots) in one transaction so that the deletion and the
/// update is atomic.
pub(crate) fn batch_remove_appointments(
&mut self,
appointments: &HashSet<UUID>,
appointments: &Vec<UUID>,
updated_users: &HashMap<UserId, UserInfo>,
) -> usize {
let limit = self.connection.limit(Limit::SQLITE_LIMIT_VARIABLE_NUMBER) as usize;
@ -405,18 +453,49 @@ impl DBM {
(appointments.len() as f64 / limit as f64).ceil() as usize
}
/// Loads the locator associated to a given UUID
pub(crate) fn load_locator(&self, uuid: UUID) -> Option<Locator> {
/// Loads the [`UUID`]s of appointments triggered by `locator`.
pub(crate) fn load_uuids(&self, locator: Locator) -> Vec<UUID> {
let mut stmt = self
.connection
.prepare("SELECT locator FROM appointments WHERE UUID=(?)")
.prepare("SELECT UUID from appointments WHERE locator=(?)")
.unwrap();
stmt.query_row([uuid.to_vec()], |row| {
let raw_locator: Vec<u8> = row.get(0).unwrap();
Ok(Locator::from_slice(&raw_locator).unwrap())
stmt.query_map([locator.to_vec()], |row| {
let raw_uuid: Vec<u8> = row.get(0).unwrap();
let uuid = UUID::from_slice(&raw_uuid).unwrap();
Ok(uuid)
})
.ok()
.unwrap()
.map(|uuid_res| uuid_res.unwrap())
.collect()
}
/// Filters the given set of [`Locator`]s by including only the ones which trigger any of our stored appointments.
pub(crate) fn batch_check_locators_exist(&self, locators: Vec<&Locator>) -> Vec<Locator> {
let mut registered_locators = Vec::new();
let locators: Vec<Vec<u8>> = locators.iter().map(|l| l.to_vec()).collect();
let limit = self.connection.limit(Limit::SQLITE_LIMIT_VARIABLE_NUMBER) as usize;
for chunk in locators.chunks(limit) {
let query = "SELECT locator FROM appointments WHERE locator IN ".to_owned();
let placeholders = format!("(?{})", (", ?").repeat(chunk.len() - 1));
let mut stmt = self
.connection
.prepare(&format!("{query}{placeholders}"))
.unwrap();
let known_locators = stmt
.query_map(params_from_iter(chunk), |row| {
let raw_locator: Vec<u8> = row.get(0).unwrap();
let locator = Locator::from_slice(&raw_locator).unwrap();
Ok(locator)
})
.unwrap()
.map(|locator_res| locator_res.unwrap());
registered_locators.extend(known_locators);
}
registered_locators
}
/// Stores a [TransactionTracker] into the database.
@ -450,6 +529,29 @@ impl DBM {
}
}
/// Updates the tracker status in the database.
///
/// The only updatable fields are `height` and `confirmed`.
pub(crate) fn update_tracker_status(
&self,
uuid: UUID,
status: &ConfirmationStatus,
) -> Result<(), Error> {
let (height, confirmed) = status.to_db_data().ok_or(Error::MissingField)?;
let query = "UPDATE trackers SET height=(?1), confirmed=(?2) WHERE UUID=(?3)";
match self.update_data(query, params![height, confirmed, uuid.to_vec(),]) {
Ok(x) => {
log::debug!("Tracker successfully updated: {uuid}");
Ok(x)
}
Err(e) => {
log::error!("Couldn't update tracker: {uuid}. Error: {e:?}");
Err(e)
}
}
}
/// Loads a [TransactionTracker] from the database.
pub(crate) fn load_tracker(&self, uuid: UUID) -> Option<TransactionTracker> {
let key = uuid.to_vec();
@ -481,6 +583,15 @@ impl DBM {
.ok()
}
/// Check if a tracker with `uuid` exists.
pub(crate) fn tracker_exists(&self, uuid: UUID) -> bool {
self.connection
.prepare("SELECT UUID FROM trackers WHERE UUID=(?)")
.unwrap()
.exists([uuid.to_vec()])
.unwrap()
}
/// Loads trackers from the database. If a locator is given, this method loads only the trackers
/// matching this locator. If no locator is given, all the trackers in the database would be returned.
pub(crate) fn load_trackers(
@ -530,6 +641,66 @@ impl DBM {
trackers
}
/// Loads trackers with the given confirmation status.
///
/// Note that for [`ConfirmationStatus::InMempoolSince(height)`] variant, this pulls trackers
/// with `h <= height` and not just `h = height`.
pub(crate) fn load_trackers_with_confirmation_status(
&self,
status: ConfirmationStatus,
) -> Result<Vec<UUID>, Error> {
let (height, confirmed) = status.to_db_data().ok_or(Error::MissingField)?;
let sql = format!(
"SELECT UUID FROM trackers WHERE confirmed=(?1) AND height{}(?2)",
if confirmed { "=" } else { "<=" }
);
let mut stmt = self.connection.prepare(&sql).unwrap();
Ok(stmt
.query_map(params![confirmed, height], |row| {
let raw_uuid: Vec<u8> = row.get(0).unwrap();
let uuid = UUID::from_slice(&raw_uuid).unwrap();
Ok(uuid)
})
.unwrap()
.map(|uuid_res| uuid_res.unwrap())
.collect())
}
/// Loads the transaction IDs of all the penalties and their status from the database.
pub(crate) fn load_penalties_summaries(&self) -> HashMap<UUID, PenaltySummary> {
let mut summaries = HashMap::new();
let mut stmt = self
.connection
.prepare(
"SELECT t.UUID, t.penalty_tx, t.height, t.confirmed
FROM trackers as t INNER JOIN appointments as a ON t.UUID=a.UUID",
)
.unwrap();
let mut rows = stmt.query([]).unwrap();
while let Ok(Some(row)) = rows.next() {
let raw_uuid: Vec<u8> = row.get(0).unwrap();
let raw_penalty_tx: Vec<u8> = row.get(1).unwrap();
let height: u32 = row.get(2).unwrap();
let confirmed: bool = row.get(3).unwrap();
// DISCUSS: Should we store the txids to avoid pulling raw txs and deserializing then hashing them.
let penalty_txid = consensus::deserialize::<bitcoin::Transaction>(&raw_penalty_tx)
.unwrap()
.txid();
summaries.insert(
UUID::from_slice(&raw_uuid).unwrap(),
PenaltySummary::new(
penalty_txid,
ConfirmationStatus::from_db_data(height, confirmed),
),
);
}
summaries
}
/// Stores the last known block into the database.
pub(crate) fn store_last_known_block(&self, block_hash: &BlockHash) -> Result<(), Error> {
let query = "INSERT OR REPLACE INTO last_known_block (id, block_hash) VALUES (0, ?)";
@ -581,11 +752,13 @@ impl DBM {
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashSet;
use std::iter::FromIterator;
use teos_common::cryptography::{get_random_bytes, get_random_keypair};
use teos_common::test_utils::get_random_user_id;
use teos_common::test_utils::{get_random_locator, get_random_user_id};
use crate::rpc_errors;
use crate::test_utils::{
generate_dummy_appointment, generate_dummy_appointment_with_user, generate_uuid,
get_random_tracker, get_random_tx, AVAILABLE_SLOTS, SUBSCRIPTION_EXPIRY,
@ -607,20 +780,15 @@ mod tests {
let mut stmt = self
.connection
.prepare(
"SELECT user_id, available_slots, subscription_start, subscription_expiry
"SELECT available_slots, subscription_start, subscription_expiry
FROM users WHERE user_id=(?)",
)
.unwrap();
stmt.query_row([&key], |row| {
let slots = row.get(1).unwrap();
let start = row.get(2).unwrap();
let expiry = row.get(3).unwrap();
Ok(UserInfo::with_appointments(
slots,
start,
expiry,
self.load_user_appointments(user_id),
))
let slots = row.get(0).unwrap();
let start = row.get(1).unwrap();
let expiry = row.get(2).unwrap();
Ok(UserInfo::new(slots, start, expiry))
})
.ok()
}
@ -651,27 +819,6 @@ mod tests {
));
}
#[test]
fn test_store_load_user_with_appointments() {
let dbm = DBM::in_memory().unwrap();
let user_id = get_random_user_id();
let mut user = UserInfo::new(AVAILABLE_SLOTS, SUBSCRIPTION_START, SUBSCRIPTION_EXPIRY);
dbm.store_user(user_id, &user).unwrap();
// Add some appointments to the user
for _ in 0..10 {
let (uuid, appointment) = generate_dummy_appointment_with_user(user_id, None);
dbm.store_appointment(uuid, &appointment).unwrap();
user.appointments.insert(uuid, 1);
}
// Check both loading the whole user info or only the associated appointments
assert_eq!(dbm.load_user(user_id).unwrap(), user);
assert_eq!(dbm.load_user_appointments(user_id), user.appointments);
}
#[test]
fn test_load_nonexistent_user() {
let dbm = DBM::in_memory().unwrap();
@ -695,6 +842,30 @@ mod tests {
assert_eq!(dbm.load_user(user_id).unwrap(), user);
}
#[test]
fn test_load_user_locators() {
let dbm = DBM::in_memory().unwrap();
let user_id = get_random_user_id();
let user = UserInfo::new(AVAILABLE_SLOTS, SUBSCRIPTION_START, SUBSCRIPTION_EXPIRY);
dbm.store_user(user_id, &user).unwrap();
let mut locators = HashSet::new();
// Add some appointments to the user
for _ in 0..10 {
let (uuid, appointment) = generate_dummy_appointment_with_user(user_id, None);
dbm.store_appointment(uuid, &appointment).unwrap();
locators.insert(appointment.locator());
}
assert_eq!(dbm.load_user(user_id).unwrap(), user);
assert_eq!(
HashSet::from_iter(dbm.load_user_locators(user_id)),
locators
);
}
#[test]
fn test_load_all_users() {
let dbm = DBM::in_memory().unwrap();
@ -707,19 +878,8 @@ mod tests {
SUBSCRIPTION_START + i,
SUBSCRIPTION_EXPIRY + i,
);
users.insert(user_id, user.clone());
users.insert(user_id, user);
dbm.store_user(user_id, &user).unwrap();
// Add appointments to some of the users
if i % 2 == 0 {
let (uuid, appointment) = generate_dummy_appointment_with_user(user_id, None);
dbm.store_appointment(uuid, &appointment).unwrap();
users
.get_mut(&user_id)
.unwrap()
.appointments
.insert(uuid, 1);
}
}
assert_eq!(dbm.load_all_users(), users);
@ -735,7 +895,7 @@ mod tests {
dbm.connection
.set_limit(Limit::SQLITE_LIMIT_VARIABLE_NUMBER, limit);
let mut to_be_deleted = HashSet::new();
let mut to_be_deleted = Vec::new();
let mut rest = HashSet::new();
for i in 1..100 {
let user_id = get_random_user_id();
@ -743,7 +903,7 @@ mod tests {
dbm.store_user(user_id, &user).unwrap();
if i % 2 == 0 {
to_be_deleted.insert(user_id);
to_be_deleted.push(user_id);
} else {
rest.insert(user_id);
}
@ -775,7 +935,7 @@ mod tests {
Ok { .. }
));
dbm.batch_remove_users(&HashSet::from_iter(vec![appointment.user_id]));
dbm.batch_remove_users(&vec![appointment.user_id]);
assert!(dbm.load_user(appointment.user_id).is_none());
assert!(dbm.load_appointment(uuid).is_none());
@ -787,7 +947,7 @@ mod tests {
));
assert!(matches!(dbm.store_tracker(uuid, &tracker), Ok { .. }));
dbm.batch_remove_users(&HashSet::from_iter(vec![appointment.user_id]));
dbm.batch_remove_users(&vec![appointment.user_id]);
assert!(dbm.load_user(appointment.user_id).is_none());
assert!(dbm.load_appointment(uuid).is_none());
assert!(dbm.load_tracker(uuid).is_none());
@ -802,6 +962,37 @@ mod tests {
dbm.batch_remove_users(&users);
}
#[test]
fn test_get_appointments_trackers_count() {
let dbm = DBM::in_memory().unwrap();
let n_users = 100;
let n_app_per_user = 4;
let n_trk_per_user = 6;
for _ in 0..n_users {
let user_id = get_random_user_id();
let user = UserInfo::new(AVAILABLE_SLOTS, SUBSCRIPTION_START, SUBSCRIPTION_EXPIRY);
dbm.store_user(user_id, &user).unwrap();
// These are un-triggered appointments.
for _ in 0..n_app_per_user {
let (uuid, appointment) = generate_dummy_appointment_with_user(user_id, None);
dbm.store_appointment(uuid, &appointment).unwrap();
}
// And these are triggered ones (trackers).
for _ in 0..n_trk_per_user {
let (uuid, appointment) = generate_dummy_appointment_with_user(user_id, None);
dbm.store_appointment(uuid, &appointment).unwrap();
let tracker = get_random_tracker(user_id, ConfirmationStatus::ConfirmedIn(42));
dbm.store_tracker(uuid, &tracker).unwrap();
}
}
assert_eq!(dbm.get_appointments_count(), n_users * n_app_per_user);
assert_eq!(dbm.get_trackers_count(), n_users * n_trk_per_user);
}
#[test]
fn test_store_load_appointment() {
let dbm = DBM::in_memory().unwrap();
@ -848,6 +1039,22 @@ mod tests {
assert!(dbm.load_appointment(uuid).is_none());
}
#[test]
fn test_appointment_exists() {
let dbm = DBM::in_memory().unwrap();
let user_id = get_random_user_id();
let user = UserInfo::new(AVAILABLE_SLOTS, SUBSCRIPTION_START, SUBSCRIPTION_EXPIRY);
let (uuid, appointment) = generate_dummy_appointment_with_user(user_id, None);
assert!(!dbm.appointment_exists(uuid));
dbm.store_user(user_id, &user).unwrap();
dbm.store_appointment(uuid, &appointment).unwrap();
assert!(dbm.appointment_exists(uuid));
}
#[test]
fn test_update_appointment() {
let dbm = DBM::in_memory().unwrap();
@ -871,7 +1078,8 @@ mod tests {
another_modified_appointment.user_id = get_random_user_id();
// Check how only the modifiable fields have been updated
dbm.update_appointment(uuid, &another_modified_appointment);
dbm.update_appointment(uuid, &another_modified_appointment)
.unwrap();
assert_eq!(dbm.load_appointment(uuid).unwrap(), modified_appointment);
assert_ne!(
dbm.load_appointment(uuid).unwrap(),
@ -970,6 +1178,44 @@ mod tests {
assert_eq!(dbm.load_appointments(Some(locator)), appointments);
}
#[test]
fn test_get_appointment_length() {
let dbm = DBM::in_memory().unwrap();
let user_id = get_random_user_id();
let user = UserInfo::new(AVAILABLE_SLOTS, SUBSCRIPTION_START, SUBSCRIPTION_EXPIRY);
let (uuid, appointment) = generate_dummy_appointment_with_user(user_id, None);
dbm.store_user(user_id, &user).unwrap();
dbm.store_appointment(uuid, &appointment).unwrap();
assert_eq!(
dbm.get_appointment_length(uuid).unwrap(),
appointment.inner.encrypted_blob.len()
);
assert!(dbm.get_appointment_length(generate_uuid()).is_none());
}
#[test]
fn test_get_appointment_user_and_length() {
let dbm = DBM::in_memory().unwrap();
let user_id = get_random_user_id();
let user = UserInfo::new(AVAILABLE_SLOTS, SUBSCRIPTION_START, SUBSCRIPTION_EXPIRY);
let (uuid, appointment) = generate_dummy_appointment_with_user(user_id, None);
dbm.store_user(user_id, &user).unwrap();
dbm.store_appointment(uuid, &appointment).unwrap();
assert_eq!(
dbm.get_appointment_user_and_length(uuid).unwrap(),
(user_id, appointment.encrypted_blob().len())
);
assert!(dbm
.get_appointment_user_and_length(generate_uuid())
.is_none());
}
#[test]
fn test_batch_remove_appointments() {
let mut dbm = DBM::in_memory().unwrap();
@ -990,13 +1236,13 @@ mod tests {
let mut rest = HashSet::new();
for i in 1..6 {
let mut to_be_deleted = HashSet::new();
let mut to_be_deleted = Vec::new();
for j in 0..limit * 2 * i {
let (uuid, appointment) = generate_dummy_appointment_with_user(user_id, None);
dbm.store_appointment(uuid, &appointment).unwrap();
if j % 2 == 0 {
to_be_deleted.insert(uuid);
to_be_deleted.push(uuid);
} else {
rest.insert(uuid);
}
@ -1005,7 +1251,7 @@ mod tests {
// When the appointment are deleted, the user will get back slots based on the deleted data.
// Here we can just make a number up to make sure it matches.
user.available_slots = i as u32;
let updated_users = HashMap::from_iter([(user_id, user.clone())]);
let updated_users = HashMap::from_iter([(user_id, user)]);
// Check that the db transaction had i queries on it
assert_eq!(
@ -1041,8 +1287,8 @@ mod tests {
));
dbm.batch_remove_appointments(
&HashSet::from_iter(vec![uuid]),
&HashMap::from_iter([(appointment.user_id, info.clone())]),
&vec![uuid],
&HashMap::from_iter([(appointment.user_id, info)]),
);
assert!(dbm.load_appointment(uuid).is_none());
@ -1054,7 +1300,7 @@ mod tests {
assert!(matches!(dbm.store_tracker(uuid, &tracker), Ok { .. }));
dbm.batch_remove_appointments(
&HashSet::from_iter(vec![uuid]),
&vec![uuid],
&HashMap::from_iter([(appointment.user_id, info)]),
);
assert!(dbm.load_appointment(uuid).is_none());
@ -1069,32 +1315,83 @@ mod tests {
// Test it does not fail even if the user does not exist (it will log though)
dbm.batch_remove_appointments(&appointments, &HashMap::new());
}
#[test]
fn test_load_locator() {
fn test_load_uuids() {
let dbm = DBM::in_memory().unwrap();
// In order to add an appointment we need the associated user to be present
let user_id = get_random_user_id();
let user = UserInfo::new(AVAILABLE_SLOTS, SUBSCRIPTION_START, SUBSCRIPTION_EXPIRY);
dbm.store_user(user_id, &user).unwrap();
let dispute_tx = get_random_tx();
let dispute_txid = dispute_tx.txid();
let mut uuids = HashSet::new();
let (uuid, appointment) = generate_dummy_appointment_with_user(user_id, None);
// Add ten appointments triggered by the same locator.
for _ in 0..10 {
let user_id = get_random_user_id();
dbm.store_user(user_id, &user).unwrap();
assert!(matches!(
dbm.store_appointment(uuid, &appointment),
Ok { .. }
));
let (uuid, appointment) =
generate_dummy_appointment_with_user(user_id, Some(&dispute_txid));
dbm.store_appointment(uuid, &appointment).unwrap();
// We should be able to load the locator now the appointment exists
assert_eq!(dbm.load_locator(uuid).unwrap(), appointment.locator());
uuids.insert(uuid);
}
// Add ten more appointments triggered by different locators.
for _ in 0..10 {
let user_id = get_random_user_id();
dbm.store_user(user_id, &user).unwrap();
let dispute_txid = get_random_tx().txid();
let (uuid, appointment) =
generate_dummy_appointment_with_user(user_id, Some(&dispute_txid));
dbm.store_appointment(uuid, &appointment).unwrap();
}
assert_eq!(
HashSet::from_iter(dbm.load_uuids(Locator::new(dispute_txid))),
uuids
);
}
#[test]
fn test_load_nonexistent_locator() {
fn test_batch_check_locators_exist() {
let dbm = DBM::in_memory().unwrap();
// Generate `n_app` appointments which we will store in the DB.
let n_app = 100;
let appointments: Vec<_> = (0..n_app)
.map(|_| generate_dummy_appointment(None))
.collect();
let (uuid, _) = generate_dummy_appointment_with_user(get_random_user_id(), None);
assert!(dbm.load_locator(uuid).is_none());
// Register all the users beforehand.
for user_id in appointments.iter().map(|a| a.user_id) {
let user = UserInfo::new(AVAILABLE_SLOTS, SUBSCRIPTION_START, SUBSCRIPTION_EXPIRY);
dbm.store_user(user_id, &user).unwrap();
}
// Store all the `n_app` appointments.
for appointment in appointments.iter() {
dbm.store_appointment(appointment.uuid(), appointment)
.unwrap();
}
// Select `n_app / 5` locators as if they appeared in a new block.
let known_locators: HashSet<_> = appointments
.iter()
.take(n_app / 5)
.map(|a| a.locator())
.collect();
// And extra `n_app / 5` unknown locators.
let unknown_locators: HashSet<_> = (0..n_app / 5).map(|_| get_random_locator()).collect();
let all_locators = known_locators
.iter()
.chain(unknown_locators.iter())
.collect();
assert_eq!(
HashSet::from_iter(dbm.batch_check_locators_exist(all_locators)),
known_locators
);
}
#[test]
@ -1154,6 +1451,38 @@ mod tests {
));
}
#[test]
fn test_update_tracker_status() {
let dbm = DBM::in_memory().unwrap();
let user_id = get_random_user_id();
let user = UserInfo::new(AVAILABLE_SLOTS, SUBSCRIPTION_START, SUBSCRIPTION_EXPIRY);
dbm.store_user(user_id, &user).unwrap();
let (uuid, appointment) = generate_dummy_appointment_with_user(user_id, None);
dbm.store_appointment(uuid, &appointment).unwrap();
let tracker = get_random_tracker(user_id, ConfirmationStatus::InMempoolSince(42));
dbm.store_tracker(uuid, &tracker).unwrap();
// Update the status and check if it's actually updated.
dbm.update_tracker_status(uuid, &ConfirmationStatus::ConfirmedIn(100))
.unwrap();
assert_eq!(
dbm.load_tracker(uuid).unwrap().status,
ConfirmationStatus::ConfirmedIn(100)
);
// Rejected status doesn't have a persistent DB representation.
assert!(matches!(
dbm.update_tracker_status(
uuid,
&ConfirmationStatus::Rejected(rpc_errors::RPC_VERIFY_REJECTED)
),
Err(Error::MissingField)
));
}
#[test]
fn test_load_nonexistent_tracker() {
let dbm = DBM::in_memory().unwrap();
@ -1224,6 +1553,166 @@ mod tests {
assert_eq!(dbm.load_trackers(Some(locator)), trackers);
}
#[test]
fn test_load_trackers_with_confirmation_status_in_mempool() {
let dbm = DBM::in_memory().unwrap();
let n_trackers = 100;
let mut tracker_statuses = HashMap::new();
// Store a bunch of trackers.
for i in 0..n_trackers {
let user_id = get_random_user_id();
let user = UserInfo::new(
AVAILABLE_SLOTS + i,
SUBSCRIPTION_START + i,
SUBSCRIPTION_EXPIRY + i,
);
dbm.store_user(user_id, &user).unwrap();
let (uuid, appointment) = generate_dummy_appointment_with_user(user_id, None);
dbm.store_appointment(uuid, &appointment).unwrap();
// Some trackers confirmed and some aren't.
let status = if i % 2 == 0 {
ConfirmationStatus::InMempoolSince(i)
} else {
ConfirmationStatus::ConfirmedIn(i)
};
let tracker = get_random_tracker(user_id, status);
dbm.store_tracker(uuid, &tracker).unwrap();
tracker_statuses.insert(uuid, status);
}
for i in 0..n_trackers + 10 {
let in_mempool_since_i: HashSet<UUID> = tracker_statuses
.iter()
.filter_map(|(&uuid, &status)| {
if let ConfirmationStatus::InMempoolSince(x) = status {
// If a tracker was in mempool since x, then it's also in mempool since x + 1, x + 2, etc...
return (x <= i).then_some(uuid);
}
None
})
.collect();
assert_eq!(
HashSet::from_iter(
dbm.load_trackers_with_confirmation_status(ConfirmationStatus::InMempoolSince(
i
))
.unwrap()
),
in_mempool_since_i,
);
}
}
#[test]
fn test_load_trackers_with_confirmation_status_confirmed() {
let dbm = DBM::in_memory().unwrap();
let n_blocks = 100;
let n_trackers = 30;
let mut tracker_statuses = HashMap::new();
// Loop over a bunch of blocks.
for i in 0..n_blocks {
// Store a bunch of trackers in each block.
for j in 0..n_trackers {
let user_id = get_random_user_id();
let user = UserInfo::new(
AVAILABLE_SLOTS + i,
SUBSCRIPTION_START + i,
SUBSCRIPTION_EXPIRY + i,
);
dbm.store_user(user_id, &user).unwrap();
let (uuid, appointment) = generate_dummy_appointment_with_user(user_id, None);
dbm.store_appointment(uuid, &appointment).unwrap();
// Some trackers confirmed and some aren't.
let status = if j % 2 == 0 {
ConfirmationStatus::InMempoolSince(i)
} else {
ConfirmationStatus::ConfirmedIn(i)
};
let tracker = get_random_tracker(user_id, status);
dbm.store_tracker(uuid, &tracker).unwrap();
tracker_statuses.insert(uuid, status);
}
}
for i in 0..n_blocks + 10 {
let confirmed_in_i: HashSet<UUID> = tracker_statuses
.iter()
.filter_map(|(&uuid, &status)| {
if let ConfirmationStatus::ConfirmedIn(x) = status {
return (x == i).then_some(uuid);
}
None
})
.collect();
assert_eq!(
HashSet::from_iter(
dbm.load_trackers_with_confirmation_status(ConfirmationStatus::ConfirmedIn(i))
.unwrap()
),
confirmed_in_i,
);
}
}
#[test]
fn test_load_trackers_with_confirmation_status_bad_status() {
let dbm = DBM::in_memory().unwrap();
assert!(matches!(
dbm.load_trackers_with_confirmation_status(ConfirmationStatus::Rejected(
rpc_errors::RPC_VERIFY_REJECTED
)),
Err(Error::MissingField)
));
assert!(matches!(
dbm.load_trackers_with_confirmation_status(ConfirmationStatus::IrrevocablyResolved),
Err(Error::MissingField)
));
}
#[test]
fn test_load_penalties_summaries() {
let dbm = DBM::in_memory().unwrap();
let n_trackers = 100;
let mut penalties_summaries = HashMap::new();
for i in 0..n_trackers {
let user_id = get_random_user_id();
let user = UserInfo::new(
AVAILABLE_SLOTS + i,
SUBSCRIPTION_START + i,
SUBSCRIPTION_EXPIRY + i,
);
dbm.store_user(user_id, &user).unwrap();
let (uuid, appointment) = generate_dummy_appointment_with_user(user_id, None);
dbm.store_appointment(uuid, &appointment).unwrap();
let status = if i % 2 == 0 {
ConfirmationStatus::InMempoolSince(i)
} else {
ConfirmationStatus::ConfirmedIn(i)
};
let tracker = get_random_tracker(user_id, status);
dbm.store_tracker(uuid, &tracker).unwrap();
penalties_summaries
.insert(uuid, PenaltySummary::new(tracker.penalty_tx.txid(), status));
}
assert_eq!(dbm.load_penalties_summaries(), penalties_summaries);
}
#[test]
fn test_store_load_last_known_block() {
let dbm = DBM::in_memory().unwrap();

44
teos/src/extended_appointment.rs
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,18 +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 ExtendedAppointment {
/// Create a new [ExtendedAppointment].
pub fn new(
@ -103,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)
}
}
@ -116,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
);
}
}

489
teos/src/gatekeeper.rs
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

61
teos/src/main.rs
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);
@ -257,16 +258,6 @@ async fn main() {
any => any,
};
let mut poller = ChainPoller::new(&mut derefed, Network::from_str(btc_network).unwrap());
let last_n_blocks = get_last_n_blocks(&mut poller, tip, IRREVOCABLY_RESOLVED as usize)
.await.unwrap_or_else(|e| {
// I'm pretty sure this can only happen if we are pulling blocks from the target to the prune height, and by the time we get to
// the end at least one has been pruned.
log::error!("Couldn't load the latest {IRREVOCABLY_RESOLVED} blocks. Please try again (Error: {})", e.into_inner());
std::process::exit(1);
}
);
// Build components
let gatekeeper = Arc::new(Gatekeeper::new(
tip.height,
@ -276,23 +267,35 @@ async fn main() {
dbm.clone(),
));
let carrier = Carrier::new(rpc, bitcoind_reachable.clone(), tip.height);
let responder = Arc::new(Responder::new(
&last_n_blocks,
tip.height,
carrier,
gatekeeper.clone(),
dbm.clone(),
));
let watcher = Arc::new(Watcher::new(
gatekeeper.clone(),
responder.clone(),
&last_n_blocks[0..6],
tip.height,
tower_sk,
TowerId(tower_pk),
dbm.clone(),
));
let mut poller = ChainPoller::new(&mut derefed, Network::from_str(btc_network).unwrap());
let (responder, watcher) = {
let last_n_blocks = get_last_n_blocks(&mut poller, tip, IRREVOCABLY_RESOLVED as usize)
.await.unwrap_or_else(|e| {
// I'm pretty sure this can only happen if we are pulling blocks from the target to the prune height, and by the time we get to
// the end at least one has been pruned.
log::error!("Couldn't load the latest {IRREVOCABLY_RESOLVED} blocks. Please try again (Error: {})", e.into_inner());
std::process::exit(1);
}
);
let responder = Arc::new(Responder::new(
&last_n_blocks,
tip.height,
Carrier::new(rpc, bitcoind_reachable.clone(), tip.height),
gatekeeper.clone(),
dbm.clone(),
));
let watcher = Arc::new(Watcher::new(
gatekeeper.clone(),
responder.clone(),
&last_n_blocks[0..6],
tip.height,
tower_sk,
TowerId(tower_pk),
dbm.clone(),
));
(responder, watcher)
};
if watcher.is_fresh() & responder.is_fresh() & gatekeeper.is_fresh() {
log::info!("Fresh bootstrap");
@ -307,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(

1822
teos/src/responder.rs
View file

File diff suppressed because it is too large Load diff

14
teos/src/test_utils.rs
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> {

25
teos/src/tx_index.rs
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]

1168
teos/src/watcher.rs
View file

File diff suppressed because it is too large Load diff