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

Replaces Result<T, E> for Option<T> in DBM load methods

Browse source 
The `DBM` methods related to loading data return `Result<T, E>` where `E` is always `dbm::Error::NotFound`.

It may makes more sense make them return `Option<T>`.
This commit is contained in:
Sergi Delgado Segura 2023-01-25 12:18:06 -05:00
parent 43f9971315
commit f628b358db
No known key found for this signature in database
GPG key ID: 35DDB7126CCB7618
9 changed files with 170 additions and 274 deletions
Download patch file Download diff file Expand all files Collapse all files

5
teos/src/chain_monitor.rs
View file

@ -264,10 +264,7 @@ mod tests {
// If a new (worse, just one) block gets mined, nothing gets connected nor disconnected
cm.poll_best_tip().await;
assert_eq!(cm.last_known_block_header, best_tip);
assert!(matches!(
cm.dbm.lock().unwrap().load_last_known_block(),
Err { .. }
));
assert!(cm.dbm.lock().unwrap().load_last_known_block().is_none());
assert!(listener.connected_blocks.borrow().is_empty());
assert!(listener.disconnected_blocks.borrow().is_empty());
}

85
teos/src/dbm.rs
View file

@ -274,7 +274,7 @@ impl DBM {
}
/// Loads an [Appointment] from the database.
pub(crate) fn load_appointment(&self, uuid: UUID) -> Result<ExtendedAppointment, Error> {
pub(crate) fn load_appointment(&self, uuid: UUID) -> Option<ExtendedAppointment> {
let key = uuid.to_vec();
let mut stmt = self
.connection
@ -302,7 +302,7 @@ impl DBM {
start_block,
))
})
.map_err(|_| Error::NotFound)
.ok()
}
/// Loads appointments from the database. If a locator is given, this method loads only the appointments
@ -406,7 +406,7 @@ impl DBM {
}
/// Loads the locator associated to a given UUID
pub(crate) fn load_locator(&self, uuid: UUID) -> Result<Locator, Error> {
pub(crate) fn load_locator(&self, uuid: UUID) -> Option<Locator> {
let mut stmt = self
.connection
.prepare("SELECT locator FROM appointments WHERE UUID=(?)")
@ -416,7 +416,7 @@ impl DBM {
let raw_locator: Vec<u8> = row.get(0).unwrap();
Ok(Locator::from_slice(&raw_locator).unwrap())
})
.map_err(|_| Error::NotFound)
.ok()
}
/// Stores a [TransactionTracker] into the database.
@ -451,7 +451,7 @@ impl DBM {
}
/// Loads a [TransactionTracker] from the database.
pub(crate) fn load_tracker(&self, uuid: UUID) -> Result<TransactionTracker, Error> {
pub(crate) fn load_tracker(&self, uuid: UUID) -> Option<TransactionTracker> {
let key = uuid.to_vec();
let mut stmt = self
.connection.prepare(
@ -478,7 +478,7 @@ impl DBM {
user_id,
})
})
.map_err(|_| Error::NotFound)
.ok()
}
/// Loads trackers from the database. If a locator is given, this method loads only the trackers
@ -537,7 +537,7 @@ impl DBM {
}
/// Loads the last known block from the database.
pub fn load_last_known_block(&self) -> Result<BlockHash, Error> {
pub fn load_last_known_block(&self) -> Option<BlockHash> {
let mut stmt = self
.connection
.prepare("SELECT block_hash FROM last_known_block WHERE id=0")
@ -547,7 +547,7 @@ impl DBM {
let raw_hash: Vec<u8> = row.get(0).unwrap();
Ok(BlockHash::from_slice(&raw_hash).unwrap())
})
.map_err(|_| Error::NotFound)
.ok()
}
/// Stores the tower secret key into the database.
@ -562,7 +562,7 @@ impl DBM {
///
/// Loads the key with higher id from the database. Old keys are not overwritten just in case a recovery is needed,
/// but they are not accessible from the API either.
pub fn load_tower_key(&self) -> Result<SecretKey, Error> {
pub fn load_tower_key(&self) -> Option<SecretKey> {
let mut stmt = self
.connection
.prepare(
@ -574,7 +574,7 @@ impl DBM {
let sk: String = row.get(0).unwrap();
Ok(SecretKey::from_str(&sk).unwrap())
})
.map_err(|_| Error::NotFound)
.ok()
}
}
@ -602,7 +602,7 @@ mod tests {
Ok(dbm)
}
pub(crate) fn load_user(&self, user_id: UserId) -> Result<UserInfo, Error> {
pub(crate) fn load_user(&self, user_id: UserId) -> Option<UserInfo> {
let key = user_id.to_vec();
let mut stmt = self
.connection
@ -611,21 +611,18 @@ mod tests {
FROM users WHERE user_id=(?)",
)
.unwrap();
let user = 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),
))
})
.map_err(|_| Error::NotFound)?;
Ok(user)
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),
))
})
.ok()
}
}
@ -680,7 +677,7 @@ mod tests {
let dbm = DBM::in_memory().unwrap();
let user_id = get_random_user_id();
assert!(matches!(dbm.load_user(user_id), Err(Error::NotFound)));
assert!(dbm.load_user(user_id).is_none());
}
#[test]
@ -779,11 +776,8 @@ mod tests {
));
dbm.batch_remove_users(&HashSet::from_iter(vec![appointment.user_id]));
assert!(matches!(
dbm.load_user(appointment.user_id),
Err(Error::NotFound)
));
assert!(matches!(dbm.load_appointment(uuid), Err(Error::NotFound)));
assert!(dbm.load_user(appointment.user_id).is_none());
assert!(dbm.load_appointment(uuid).is_none());
// Appointment + Tracker
dbm.store_user(appointment.user_id, &info).unwrap();
@ -794,12 +788,9 @@ mod tests {
assert!(matches!(dbm.store_tracker(uuid, &tracker), Ok { .. }));
dbm.batch_remove_users(&HashSet::from_iter(vec![appointment.user_id]));
assert!(matches!(
dbm.load_user(appointment.user_id),
Err(Error::NotFound)
));
assert!(matches!(dbm.load_appointment(uuid), Err(Error::NotFound)));
assert!(matches!(dbm.load_tracker(uuid), Err(Error::NotFound)));
assert!(dbm.load_user(appointment.user_id).is_none());
assert!(dbm.load_appointment(uuid).is_none());
assert!(dbm.load_tracker(uuid).is_none());
}
#[test]
@ -846,7 +837,7 @@ mod tests {
dbm.store_appointment(uuid, &appointment),
Err(Error::MissingForeignKey)
));
assert!(matches!(dbm.load_tracker(uuid), Err(Error::NotFound)));
assert!((dbm.load_tracker(uuid).is_none()));
}
#[test]
@ -854,7 +845,7 @@ mod tests {
let dbm = DBM::in_memory().unwrap();
let uuid = generate_uuid();
assert!(matches!(dbm.load_appointment(uuid), Err(Error::NotFound)));
assert!(dbm.load_appointment(uuid).is_none());
}
#[test]
@ -1053,7 +1044,7 @@ mod tests {
&HashSet::from_iter(vec![uuid]),
&HashMap::from_iter([(appointment.user_id, info.clone())]),
);
assert!(matches!(dbm.load_appointment(uuid), Err(Error::NotFound)));
assert!(dbm.load_appointment(uuid).is_none());
// Appointment + Tracker
assert!(matches!(
@ -1066,8 +1057,8 @@ mod tests {
&HashSet::from_iter(vec![uuid]),
&HashMap::from_iter([(appointment.user_id, info)]),
);
assert!(matches!(dbm.load_appointment(uuid), Err(Error::NotFound)));
assert!(matches!(dbm.load_tracker(uuid), Err(Error::NotFound)));
assert!(dbm.load_appointment(uuid).is_none());
assert!(dbm.load_tracker(uuid).is_none());
}
#[test]
@ -1103,7 +1094,7 @@ mod tests {
let dbm = DBM::in_memory().unwrap();
let (uuid, _) = generate_dummy_appointment_with_user(get_random_user_id(), None);
assert!(matches!(dbm.load_locator(uuid), Err(Error::NotFound)));
assert!(dbm.load_locator(uuid).is_none());
}
#[test]
@ -1168,7 +1159,7 @@ mod tests {
let dbm = DBM::in_memory().unwrap();
let uuid = generate_uuid();
assert!(matches!(dbm.load_tracker(uuid), Err(Error::NotFound)));
assert!(dbm.load_tracker(uuid).is_none());
}
#[test]
@ -1251,14 +1242,14 @@ mod tests {
fn test_store_load_nonexistent_last_known_block() {
let dbm = DBM::in_memory().unwrap();
assert!(matches!(dbm.load_last_known_block(), Err(Error::NotFound)));
assert!(dbm.load_last_known_block().is_none());
}
#[test]
fn test_store_load_tower_key() {
let dbm = DBM::in_memory().unwrap();
assert!(matches!(dbm.load_tower_key(), Err(Error::NotFound)));
assert!(dbm.load_tower_key().is_none());
for _ in 0..7 {
let sk = get_random_keypair().0;
dbm.store_tower_key(&sk).unwrap();

6
teos/src/gatekeeper.rs
View file

@ -356,7 +356,6 @@ mod tests {
};
use lightning::chain::Listen;
use teos_common::cryptography::{get_random_bytes, get_random_keypair};
use teos_common::dbm::Error as DBError;
use teos_common::test_utils::get_random_user_id;
const SLOTS: u32 = 21;
@ -849,10 +848,7 @@ mod tests {
.lock()
.unwrap()
.contains_key(user_id));
assert!(matches!(
gatekeeper.dbm.lock().unwrap().load_user(*user_id),
Err(DBError::NotFound)
));
assert!(gatekeeper.dbm.lock().unwrap().load_user(*user_id).is_none());
}
// Check that the last_known_block_header has been properly updated

13
teos/src/main.rs
View file

@ -126,14 +126,11 @@ async fn main() {
if conf.overwrite_key {
log::info!("Overwriting tower keys");
create_new_tower_keypair(&locked_db)
} else if let Some(sk) = locked_db.load_tower_key() {
(sk, PublicKey::from_secret_key(&Secp256k1::new(), &sk))
} else {
match locked_db.load_tower_key() {
Ok(sk) => (sk, PublicKey::from_secret_key(&Secp256k1::new(), &sk)),
Err(_) => {
log::info!("Tower keys not found. Creating a fresh set");
create_new_tower_keypair(&locked_db)
}
}
log::info!("Tower keys not found. Creating a fresh set");
create_new_tower_keypair(&locked_db)
}
};
log::info!("tower_id: {tower_pk}");
@ -179,7 +176,7 @@ async fn main() {
let mut derefed = bitcoin_cli.deref();
// Load last known block from DB if found. Poll it from Bitcoind otherwise.
let last_known_block = dbm.lock().unwrap().load_last_known_block();
let tip = if let Ok(block_hash) = last_known_block {
let tip = if let Some(block_hash) = last_known_block {
derefed
.get_header(&block_hash, None)
.await

18
teos/src/responder.rs
View file

@ -297,7 +297,7 @@ impl Responder {
/// The [TransactionTracker] is queried to the [DBM].
pub(crate) fn get_tracker(&self, uuid: UUID) -> Option<TransactionTracker> {
if self.trackers.lock().unwrap().contains_key(&uuid) {
self.dbm.lock().unwrap().load_tracker(uuid).ok()
self.dbm.lock().unwrap().load_tracker(uuid)
} else {
None
}
@ -623,7 +623,6 @@ mod tests {
};
use teos_common::constants::IRREVOCABLY_RESOLVED;
use teos_common::dbm::Error as DBError;
use teos_common::test_utils::get_random_user_id;
impl PartialEq for Responder {
@ -1581,10 +1580,7 @@ mod tests {
assert!(!responder.tx_tracker_map.lock().unwrap().contains_key(&txid));
// But it can be found in the database
assert!(matches!(
responder.dbm.lock().unwrap().load_tracker(uuid),
Ok(TransactionTracker { .. })
));
assert!(responder.dbm.lock().unwrap().load_tracker(uuid).is_some());
}
}
@ -1669,10 +1665,7 @@ mod tests {
for uuid in all_trackers {
if target_trackers.contains(&uuid) {
assert!(!responder.trackers.lock().unwrap().contains_key(&uuid));
assert!(matches!(
responder.dbm.lock().unwrap().load_tracker(uuid),
Err(DBError::NotFound)
));
assert!(responder.dbm.lock().unwrap().load_tracker(uuid).is_none());
let penalty_txid = &uuid_txid_map[&uuid];
// If the penalty had more than one associated uuid, only one has been deleted
// (because that's how the test has been designed)
@ -1702,10 +1695,7 @@ mod tests {
.lock()
.unwrap()
.contains_key(&uuid_txid_map[&uuid]));
assert!(matches!(
responder.dbm.lock().unwrap().load_tracker(uuid),
Ok(TransactionTracker { .. })
));
assert!(responder.dbm.lock().unwrap().load_tracker(uuid).is_some());
}
}

135
teos/src/watcher.rs
View file

@ -636,11 +636,10 @@ impl Watcher {
Some(a) => locators.push(a.locator),
None => {
if self.responder.has_tracker(*uuid) {
match dbm.load_locator(*uuid) {
Ok(locator) => locators.push(locator),
Err(_) => log::error!(
"Tracker found in Responder but not in DB (uuid = {uuid})"
),
if let Some(locator) = dbm.load_locator(*uuid) {
locators.push(locator)
} else {
log::error!("Tracker found in Responder but not in DB (uuid = {uuid})")
}
} else {
log::error!("Appointment found in the Gatekeeper but not in the Watcher nor the Responder (uuid = {uuid})")
@ -770,7 +769,6 @@ mod tests {
SUBSCRIPTION_EXPIRY, SUBSCRIPTION_START,
};
use teos_common::cryptography::{get_random_bytes, get_random_keypair};
use teos_common::dbm::Error as DBError;
use bitcoin::hash_types::Txid;
use bitcoin::hashes::Hash;
@ -961,10 +959,12 @@ mod tests {
// Check data was added to the database
for uuid in watcher.appointments.lock().unwrap().keys() {
assert!(matches!(
watcher.dbm.lock().unwrap().load_appointment(*uuid),
Ok(ExtendedAppointment { .. })
));
assert!(watcher
.dbm
.lock()
.unwrap()
.load_appointment(*uuid)
.is_some());
}
// If an appointment is already in the Responder, it should bounce
@ -1010,14 +1010,8 @@ mod tests {
assert!(watcher.responder.has_tracker(uuid));
// Check data was added to the database
assert!(matches!(
watcher.dbm.lock().unwrap().load_appointment(uuid),
Ok(ExtendedAppointment { .. })
));
assert!(matches!(
watcher.dbm.lock().unwrap().load_tracker(uuid),
Ok(TransactionTracker { .. })
));
assert!(watcher.dbm.lock().unwrap().load_appointment(uuid).is_some());
assert!(watcher.dbm.lock().unwrap().load_tracker(uuid).is_some());
// If an appointment is rejected by the Responder, it is considered misbehavior and the slot count is kept
// Wrong penalty
@ -1034,14 +1028,8 @@ mod tests {
assert_eq!(watcher.appointments.lock().unwrap().len(), 3);
// Data should not be in the database
assert!(matches!(
watcher.dbm.lock().unwrap().load_appointment(uuid),
Err(DBError::NotFound)
));
assert!(matches!(
watcher.dbm.lock().unwrap().load_tracker(uuid),
Err(DBError::NotFound)
));
assert!(watcher.dbm.lock().unwrap().load_appointment(uuid).is_none());
assert!(watcher.dbm.lock().unwrap().load_tracker(uuid).is_none());
// Transaction rejected
// Update the Responder with a new Carrier
@ -1062,10 +1050,7 @@ mod tests {
assert_eq!(watcher.appointments.lock().unwrap().len(), 3);
// Data should not be in the database
assert!(matches!(
watcher.dbm.lock().unwrap().load_appointment(uuid),
Err(DBError::NotFound)
));
assert!(watcher.dbm.lock().unwrap().load_appointment(uuid).is_none());
// FAIL cases (non-registered, subscription expired and not enough slots)
@ -1078,10 +1063,7 @@ mod tests {
Err(AddAppointmentFailure::AuthenticationFailure)
));
// Data should not be in the database
assert!(matches!(
watcher.dbm.lock().unwrap().load_appointment(uuid),
Err(DBError::NotFound)
));
assert!(watcher.dbm.lock().unwrap().load_appointment(uuid).is_none());
// If the user has no enough slots, the appointment is rejected. We do not test all possible cases since updates are
// already tested int he Gatekeeper. Testing that it is rejected if the condition is met should suffice.
@ -1103,10 +1085,7 @@ mod tests {
Err(AddAppointmentFailure::NotEnoughSlots)
));
// Data should not be in the database
assert!(matches!(
watcher.dbm.lock().unwrap().load_appointment(uuid),
Err(DBError::NotFound)
));
assert!(watcher.dbm.lock().unwrap().load_appointment(uuid).is_none());
// If the user subscription has expired, the appointment should be rejected.
watcher
@ -1123,10 +1102,7 @@ mod tests {
Err(AddAppointmentFailure::SubscriptionExpired { .. })
));
// Data should not be in the database
assert!(matches!(
watcher.dbm.lock().unwrap().load_appointment(uuid),
Err(DBError::NotFound)
));
assert!(watcher.dbm.lock().unwrap().load_appointment(uuid).is_none());
}
#[tokio::test]
@ -1211,10 +1187,7 @@ mod tests {
);
// In this case the appointment is kept in the Responder and, therefore, in the database
assert!(watcher.responder.has_tracker(uuid));
assert!(matches!(
watcher.dbm.lock().unwrap().load_appointment(uuid),
Ok(ExtendedAppointment { .. })
));
assert!(watcher.dbm.lock().unwrap().load_appointment(uuid).is_some());
// A properly formatted but invalid transaction should be rejected by the Responder
// Update the Responder with a new Carrier that will reject the transaction
@ -1232,10 +1205,7 @@ mod tests {
);
// In this case the appointment is not kept in the Responder nor in the database
assert!(!watcher.responder.has_tracker(uuid));
assert!(matches!(
watcher.dbm.lock().unwrap().load_appointment(uuid),
Err { .. }
));
assert!(watcher.dbm.lock().unwrap().load_appointment(uuid).is_none());
// Invalid triggered appointments should not be passed to the Responder
// Use a dispute_tx that does not match the appointment to replicate a decryption error
@ -1247,10 +1217,7 @@ mod tests {
);
// The appointment is not kept anywhere
assert!(!watcher.responder.has_tracker(uuid));
assert!(matches!(
watcher.dbm.lock().unwrap().load_appointment(uuid),
Err { .. }
));
assert!(watcher.dbm.lock().unwrap().load_appointment(uuid).is_none());
}
#[tokio::test]
@ -1502,10 +1469,7 @@ mod tests {
.contains_key(&locator));
// But it can be found in the database
assert!(matches!(
watcher.dbm.lock().unwrap().load_appointment(uuid),
Ok(ExtendedAppointment { .. })
));
assert!(watcher.dbm.lock().unwrap().load_appointment(uuid).is_some());
}
}
@ -1586,10 +1550,7 @@ mod tests {
for uuid in all_appointments {
if target_appointments.contains(&uuid) {
assert!(!watcher.appointments.lock().unwrap().contains_key(&uuid));
assert!(matches!(
watcher.dbm.lock().unwrap().load_appointment(uuid),
Err(DBError::NotFound)
));
assert!(watcher.dbm.lock().unwrap().load_appointment(uuid).is_none());
let locator = &uuid_locator_map[&uuid];
// If the penalty had more than one associated uuid, only one has been deleted
@ -1620,10 +1581,7 @@ mod tests {
.lock()
.unwrap()
.contains_key(&uuid_locator_map[&uuid]));
assert!(matches!(
watcher.dbm.lock().unwrap().load_appointment(uuid),
Ok(ExtendedAppointment { .. })
));
assert!(watcher.dbm.lock().unwrap().load_appointment(uuid).is_some());
}
}
@ -1728,10 +1686,12 @@ mod tests {
.appointments
.contains_key(&uuid1)
);
assert!(matches!(
watcher.dbm.lock().unwrap().load_appointment(uuid1),
Ok(ExtendedAppointment { .. })
));
assert!(watcher
.dbm
.lock()
.unwrap()
.load_appointment(uuid1)
.is_some());
assert!(watcher.appointments.lock().unwrap().contains_key(&uuid2));
assert!(watcher.locator_uuid_map.lock().unwrap()[&appointment.locator()].contains(&uuid2));
@ -1740,10 +1700,12 @@ mod tests {
.appointments
.contains_key(&uuid2)
);
assert!(matches!(
watcher.dbm.lock().unwrap().load_appointment(uuid2),
Ok(ExtendedAppointment { .. })
));
assert!(watcher
.dbm
.lock()
.unwrap()
.load_appointment(uuid2)
.is_some());
// Check triggers. Add a new appointment and trigger it with valid data.
let dispute_tx = get_random_tx();
@ -1774,14 +1736,8 @@ mod tests {
);
// Data should have been kept in the database
assert!(matches!(
watcher.dbm.lock().unwrap().load_appointment(uuid),
Ok(ExtendedAppointment { .. })
));
assert!(matches!(
watcher.dbm.lock().unwrap().load_tracker(uuid),
Ok(TransactionTracker { .. })
));
assert!(watcher.dbm.lock().unwrap().load_appointment(uuid).is_some());
assert!(watcher.dbm.lock().unwrap().load_tracker(uuid).is_some());
// Check triggering with a valid formatted transaction but that is rejected by the Responder.
let dispute_tx = get_random_tx();
@ -1816,14 +1772,8 @@ mod tests {
.contains_key(&uuid)
);
// Data should also have been deleted from the database
assert!(matches!(
watcher.dbm.lock().unwrap().load_appointment(uuid),
Err(DBError::NotFound)
));
assert!(matches!(
watcher.dbm.lock().unwrap().load_tracker(uuid),
Err(DBError::NotFound)
));
assert!(watcher.dbm.lock().unwrap().load_appointment(uuid).is_none());
assert!(watcher.dbm.lock().unwrap().load_tracker(uuid).is_none());
// Checks invalid triggers. Add a new appointment and trigger it with invalid data.
let dispute_tx = get_random_tx();
@ -1855,10 +1805,7 @@ mod tests {
.appointments
.contains_key(&uuid)
);
assert!(matches!(
watcher.dbm.lock().unwrap().load_appointment(uuid),
Err(DBError::NotFound)
));
assert!(watcher.dbm.lock().unwrap().load_appointment(uuid).is_none());
}
#[tokio::test]

77
watchtower-plugin/src/dbm.rs
View file

@ -125,7 +125,7 @@ impl DBM {
///
/// Loads the key with higher id from the database. Old keys are not overwritten just in case a recovery is needed,
/// but they are not accessible from the API either.
pub fn load_client_key(&self) -> Result<SecretKey, Error> {
pub fn load_client_key(&self) -> Option<SecretKey> {
let mut stmt = self
.connection
.prepare(
@ -137,7 +137,7 @@ impl DBM {
let sk: String = row.get(0).unwrap();
Ok(SecretKey::from_str(&sk).unwrap())
})
.map_err(|_| Error::NotFound)
.ok()
}
/// Stores a tower record into the database alongside the corresponding registration receipt.
@ -171,7 +171,7 @@ impl DBM {
/// Tower records are composed from the tower information and the appointment data. The latter is split in:
/// accepted appointments (represented by appointment receipts), pending appointments and invalid appointments.
/// In the case that the tower has misbehaved, then a misbehaving proof is also attached to the record.
pub fn load_tower_record(&self, tower_id: TowerId) -> Result<TowerInfo, Error> {
pub fn load_tower_record(&self, tower_id: TowerId) -> Option<TowerInfo> {
let mut stmt = self
.connection
.prepare("SELECT t.net_addr, t.available_slots, r.subscription_start, r.subscription_expiry
@ -197,16 +197,16 @@ impl DBM {
self.load_appointments(tower_id, AppointmentStatus::Invalid),
))
})
.map_err(|_| Error::NotFound)?;
.ok()?;
if let Ok(proof) = self.load_misbehaving_proof(tower_id) {
if let Some(proof) = self.load_misbehaving_proof(tower_id) {
tower.status = TowerStatus::Misbehaving;
tower.set_misbehaving_proof(proof);
} else if !tower.pending_appointments.is_empty() {
tower.status = TowerStatus::TemporaryUnreachable;
}
Ok(tower)
Some(tower)
}
/// Loads the latest registration receipt for a given tower.
@ -216,7 +216,7 @@ impl DBM {
&self,
tower_id: TowerId,
user_id: UserId,
) -> Result<RegistrationReceipt, Error> {
) -> Option<RegistrationReceipt> {
let mut stmt = self
.connection
.prepare(
@ -228,20 +228,17 @@ impl DBM {
)
.unwrap();
let receipt = stmt
.query_row([tower_id.to_vec()], |row| {
let slots: u32 = row.get(0).unwrap();
let start: u32 = row.get(1).unwrap();
let expiry: u32 = row.get(2).unwrap();
let signature: String = row.get(3).unwrap();
stmt.query_row([tower_id.to_vec()], |row| {
let slots: u32 = row.get(0).unwrap();
let start: u32 = row.get(1).unwrap();
let expiry: u32 = row.get(2).unwrap();
let signature: String = row.get(3).unwrap();
Ok(RegistrationReceipt::with_signature(
user_id, slots, start, expiry, signature,
))
})
.map_err(|_| Error::NotFound)?;
Ok(receipt)
Ok(RegistrationReceipt::with_signature(
user_id, slots, start, expiry, signature,
))
})
.ok()
}
/// Removes a tower record from the database.
@ -333,7 +330,7 @@ impl DBM {
&self,
tower_id: TowerId,
locator: Locator,
) -> Result<AppointmentReceipt, Error> {
) -> Option<AppointmentReceipt> {
let mut stmt = self
.connection
.prepare("SELECT start_block, user_signature, tower_signature FROM appointment_receipts WHERE tower_id = ?1 and locator = ?2")
@ -350,7 +347,7 @@ impl DBM {
tower_sig,
))
})
.map_err(|_| Error::NotFound)
.ok()
}
/// Loads the appointment receipts associated to a given tower.
@ -406,7 +403,7 @@ impl DBM {
}
/// Loads an appointment from the database.
pub fn load_appointment(&self, locator: Locator) -> Result<Appointment, Error> {
pub fn load_appointment(&self, locator: Locator) -> Option<Appointment> {
let mut stmt = self
.connection
.prepare("SELECT encrypted_blob, to_self_delay FROM appointments WHERE locator = ?")
@ -418,7 +415,7 @@ impl DBM {
Ok(Appointment::new(locator, encrypted_blob, to_self_delay))
})
.map_err(|_| Error::NotFound)
.ok()
}
/// Stores an appointment into the database.
@ -598,7 +595,7 @@ impl DBM {
}
/// Loads the misbehaving proof for a given tower from the database (if found).
fn load_misbehaving_proof(&self, tower_id: TowerId) -> Result<MisbehaviorProof, Error> {
fn load_misbehaving_proof(&self, tower_id: TowerId) -> Option<MisbehaviorProof> {
let mut misbehaving_stmt = self
.connection
.prepare("SELECT locator, recovered_id FROM misbehaving_proofs WHERE tower_id = ?")
@ -633,7 +630,7 @@ impl DBM {
.unwrap();
MisbehaviorProof::new(locator, receipt, recovered_id)
})
.map_err(|_| Error::NotFound)
.ok()
}
/// Checks whether a misbehaving proof exists for a given tower.
@ -792,10 +789,7 @@ mod tests {
// If the tower does not exists, `load_tower` will fail.
let tower_id = get_random_user_id();
assert!(matches!(
dbm.load_tower_record(tower_id),
Err(Error::NotFound)
));
assert!(dbm.load_tower_record(tower_id).is_none());
}
#[test]
@ -917,10 +911,9 @@ mod tests {
// If there is no appointment receipt for the given (locator, tower_id) pair, Error::NotFound is returned
// Try first with both being unknown
assert!(matches!(
dbm.load_appointment_receipt(tower_id, appointment.locator),
Err(Error::NotFound)
));
assert!(dbm
.load_appointment_receipt(tower_id, appointment.locator)
.is_none());
// Add the tower but not the appointment and try again
let net_addr = "talaia.watch";
@ -928,10 +921,9 @@ mod tests {
dbm.store_tower_record(tower_id, net_addr, &receipt)
.unwrap();
assert!(matches!(
dbm.load_appointment_receipt(tower_id, appointment.locator),
Err(Error::NotFound)
));
assert!(dbm
.load_appointment_receipt(tower_id, appointment.locator)
.is_none());
// Add both
let tower_summary = TowerSummary::new(
@ -1045,7 +1037,7 @@ mod tests {
let locator = generate_random_appointment(None).locator;
let loaded_appointment = dbm.load_appointment(locator);
assert!(matches!(loaded_appointment, Err(Error::NotFound)));
assert!(loaded_appointment.is_none());
}
#[test]
@ -1284,10 +1276,7 @@ mod tests {
#[test]
fn test_store_load_non_existing_misbehaving_proof() {
let dbm = DBM::in_memory().unwrap();
assert!(matches!(
dbm.load_misbehaving_proof(get_random_user_id()),
Err(Error::NotFound)
));
assert!(dbm.load_misbehaving_proof(get_random_user_id()).is_none());
}
#[test]
@ -1340,7 +1329,7 @@ mod tests {
fn test_store_load_client_key() {
let dbm = DBM::in_memory().unwrap();
assert!(matches!(dbm.load_client_key(), Err(Error::NotFound)));
assert!(dbm.load_client_key().is_none());
for _ in 0..7 {
let sk = get_random_keypair().0;
dbm.store_client_key(&sk).unwrap();

63
watchtower-plugin/src/main.rs
View file

@ -136,11 +136,13 @@ async fn get_registration_receipt(
let tower_id = TowerId::try_from(v).map_err(|x| anyhow!(x))?;
let state = plugin.state().lock().unwrap();
let response = state.get_registration_receipt(tower_id).map_err(|_| {
anyhow!("Cannot find {tower_id} within the known towers. Have you registered?")
})?;
Ok(json!(response))
if let Some(response) = state.get_registration_receipt(tower_id) {
Ok(json!(response))
} else {
Err(anyhow!(
"Cannot find {tower_id} within the known towers. Have you registered?"
))
}
}
/// Gets the subscription information directly form the tower.
@ -244,24 +246,20 @@ async fn get_appointment_receipt(
let params = GetAppointmentParams::try_from(v).map_err(|x| anyhow!(x))?;
let state = plugin.state().lock().unwrap();
let response = state
.get_appointment_receipt(params.tower_id, params.locator)
.map_err(|_| {
if state.towers.contains_key(&params.tower_id) {
anyhow!(
"Cannot find {} within {}. Did you send that appointment?",
params.locator,
params.tower_id
)
} else {
anyhow!(
"Cannot find {} within the known towers. Have you registered?",
params.tower_id
)
}
})?;
Ok(json!(response))
if let Some(r) = state.get_appointment_receipt(params.tower_id, params.locator) {
Ok(json!(r))
} else if state.towers.contains_key(&params.tower_id) {
Err(anyhow!(
"Cannot find {} within {}. Did you send that appointment?",
params.locator,
params.tower_id
))
} else {
Err(anyhow!(
"Cannot find {} within the known towers. Have you registered?",
params.tower_id
))
}
}
/// Lists all the registered towers.
@ -283,15 +281,18 @@ async fn get_tower_info(
) -> Result<serde_json::Value, Error> {
let state = plugin.state().lock().unwrap();
let tower_id = TowerId::try_from(v).map_err(|e| anyhow!(e))?;
let tower_info = state.load_tower_info(tower_id).map_err(|_| {
anyhow!("Cannot find {tower_id} within the known towers. Have you registered?")
})?;
// Notice we need to check the status in memory since we cannot distinguish between unreachable and temporary unreachable
// by just checking the data in the database.
Ok(json!(
tower_info.with_status(state.get_tower_status(&tower_id).unwrap())
))
if let Some(tower_info) = state.load_tower_info(tower_id) {
// Notice we need to check the status in memory since we cannot distinguish between unreachable and temporary unreachable
// by just checking the data in the database.
Ok(json!(
tower_info.with_status(state.get_tower_status(&tower_id).unwrap())
))
} else {
Err(anyhow!(
"Cannot find {tower_id} within the known towers. Have you registered?",
))
}
}
/// Triggers a manual retry of a tower, tries to send all pending appointments to it.

42
watchtower-plugin/src/wt_client.rs
View file

@ -95,17 +95,17 @@ impl WTClient {
});
let dbm = DBM::new(&data_dir.join("watchtowers_db.sql3")).unwrap();
let (user_sk, user_id) = match dbm.load_client_key() {
Ok(sk) => (
let (user_sk, user_id) = if let Some(sk) = dbm.load_client_key() {
(
sk,
UserId(PublicKey::from_secret_key(&Secp256k1::new(), &sk)),
),
Err(_) => {
log::info!("Watchtower client keys not found. Creating a fresh set");
let (sk, pk) = cryptography::get_random_keypair();
dbm.store_client_key(&sk).unwrap();
(sk, UserId(pk))
}
)
} else {
log::info!("Watchtower client keys not found. Creating a fresh set");
let (sk, pk) = cryptography::get_random_keypair();
dbm.store_client_key(&sk).unwrap();
(sk, UserId(pk))
};
let towers = dbm.load_towers();
@ -180,15 +180,12 @@ impl WTClient {
}
/// Gets the latest registration receipt of a given tower.
pub fn get_registration_receipt(
&self,
tower_id: TowerId,
) -> Result<RegistrationReceipt, DBError> {
pub fn get_registration_receipt(&self, tower_id: TowerId) -> Option<RegistrationReceipt> {
self.dbm.load_registration_receipt(tower_id, self.user_id)
}
/// Loads a tower record from the database.
pub fn load_tower_info(&self, tower_id: TowerId) -> Result<TowerInfo, DBError> {
pub fn load_tower_info(&self, tower_id: TowerId) -> Option<TowerInfo> {
self.dbm.load_tower_record(tower_id)
}
@ -240,7 +237,7 @@ impl WTClient {
&self,
tower_id: TowerId,
locator: Locator,
) -> Result<AppointmentReceipt, DBError> {
) -> Option<AppointmentReceipt> {
self.dbm.load_appointment_receipt(tower_id, locator)
}
@ -810,10 +807,7 @@ mod tests {
// Remove the tower and check it is not there anymore
wt_client.remove_tower(tower_id).unwrap();
assert!(matches!(
wt_client.load_tower_info(tower_id),
Err(DBError::NotFound)
));
assert!(wt_client.load_tower_info(tower_id).is_none());
assert!(!wt_client.towers.contains_key(&tower_id));
// Try again but this time with an associated appointment to check that it also gets removed
@ -836,10 +830,7 @@ mod tests {
// Remove and check both the tower and the appointment
wt_client.remove_tower(tower_id).unwrap();
assert!(matches!(
wt_client.load_tower_info(tower_id),
Err(DBError::NotFound)
));
assert!(wt_client.load_tower_info(tower_id).is_none());
assert!(!wt_client.towers.contains_key(&tower_id));
assert!(!wt_client.dbm.appointment_receipt_exists(locator, tower_id));
}
@ -890,10 +881,7 @@ mod tests {
// Remove tower1 and check that the appointment receipt can still be found for tower2
wt_client.remove_tower(tower1_id).unwrap();
assert!(matches!(
wt_client.load_tower_info(tower1_id),
Err(DBError::NotFound)
));
assert!(wt_client.load_tower_info(tower1_id).is_none());
assert!(!wt_client.dbm.appointment_receipt_exists(locator, tower1_id));
assert!(wt_client.dbm.appointment_receipt_exists(locator, tower2_id));