Skip to content

Commit 6046848

Browse files
ShubhamBhutShubhamBhut
committed
Add optional boundaries to getregistrationreceipt
1 parent 8bf3194 commit 6046848

4 files changed

Lines changed: 45 additions & 30 deletions

File tree

‎watchtower-plugin/src/convert.rs‎

Lines changed: 29 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -289,44 +289,56 @@ impl TryFrom<serde_json::Value> for GetRegistrationReceiptParams {
289289
match value {
290290
serde_json::Value::Array(a) => {
291291
let param_count = a.len();
292-
if param_count == 2{
293-
Err(GetRegistrationReceiptError::InvalidFormat(("Both ends of boundary (subscription_start and subscription_expiry) are required.").to_string()))
292+
if param_count == 2 {
293+
Err(GetRegistrationReceiptError::InvalidFormat((
294+
"Both ends of boundary (subscription_start and subscription_expiry) are required.").to_string()
295+
))
294296
} else if param_count != 1 && param_count != 3 {
295297
Err(GetRegistrationReceiptError::InvalidFormat(format!(
296298
"Unexpected request format. The request needs 1 or 3 parameter. Received: {param_count}"
297299
)))
298-
} else{
300+
} else {
299301
let tower_id = if let Some(s) = a.get(0).unwrap().as_str() {
300302
TowerId::from_str(s).map_err(|_| {
301-
GetRegistrationReceiptError::InvalidId("Invalid tower id".to_owned())
303+
GetRegistrationReceiptError::InvalidId("Invalid tower id".to_owned())
302304
})
303305
} else {
304306
Err(GetRegistrationReceiptError::InvalidId(
305-
"tower_id must be a hex encoded string".to_owned(),
307+
"tower_id must be a hex encoded string".to_owned(),
306308
))
307309
}?;
308310

309-
let (subscription_start, subscription_expiry) = if let (Some(start), Some(expire)) = (a.get(1).and_then(|v| v.as_i64()), a.get(2).and_then(|v| v.as_i64())) {
311+
let (subscription_start, subscription_expiry) = if let (Some(start), Some(expire)) = (a.get(1), a.get(2)){
312+
let start = start.as_i64().ok_or_else(|| {
313+
GetRegistrationReceiptError::InvalidFormat(
314+
"Subscription_start must be a positive integer".to_owned(),
315+
)
316+
})?;
317+
318+
let expire = expire.as_i64().ok_or_else(|| {
319+
GetRegistrationReceiptError::InvalidFormat(
320+
"Subscription_expire must be a positive integer".to_owned(),
321+
)
322+
})?;
323+
310324
if start >= 0 && expire > start {
311325
(Some(start as u32), Some(expire as u32))
312326
} else {
313327
return Err(GetRegistrationReceiptError::InvalidFormat(
314-
"Subscription_start must be a positive integer and subscription_expire must be a positive integer greater than subscription_start".to_owned(),
315-
));
328+
"subscription_start must be a positive integer and subscription_expire must be a positive integer greater than subscription_start".to_owned(),
329+
));
316330
}
317-
} else if a.get(1).is_some() || a.get(2).is_some() {
318-
return Err(GetRegistrationReceiptError::InvalidFormat(
319-
"Subscription_start and subscription_expiry must be provided together as positive integers".to_owned(),
320-
));
321331
} else {
322332
(None, None)
323333
};
324334

325-
Ok(Self {
326-
tower_id,
327-
subscription_start,
328-
subscription_expiry,
329-
})
335+
Ok(
336+
Self {
337+
tower_id,
338+
subscription_start,
339+
subscription_expiry,
340+
}
341+
)
330342
}
331343
},
332344
serde_json::Value::Object(mut m) => {

‎watchtower-plugin/src/dbm.rs‎

Lines changed: 14 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -218,7 +218,7 @@ impl DBM {
218218
user_id: UserId,
219219
subscription_start: Option<u32>,
220220
subscription_expiry: Option<u32>,
221-
) -> Vec<RegistrationReceipt> {
221+
) -> Option<Vec<RegistrationReceipt>> {
222222
let mut query = "SELECT available_slots, subscription_start, subscription_expiry, signature FROM registration_receipts WHERE tower_id = ?1".to_string();
223223

224224
let tower_id_encoded = tower_id.to_vec();
@@ -233,8 +233,8 @@ impl DBM {
233233
}
234234
let mut stmt = self.connection.prepare(&query).unwrap();
235235

236-
let receipts = stmt
237-
.query_map(params.as_slice(), |row| {
236+
Some(
237+
stmt.query_map(params.as_slice(), |row| {
238238
let slots: u32 = row.get(0)?;
239239
let start: u32 = row.get(1)?;
240240
let expiry: u32 = row.get(2)?;
@@ -245,10 +245,9 @@ impl DBM {
245245
))
246246
})
247247
.unwrap()
248-
.collect::<Result<Vec<_>, _>>()
249-
.unwrap_or_default();
250-
251-
receipts
248+
.map(|r| r.unwrap())
249+
.collect(),
250+
)
252251
}
253252

254253
/// Removes a tower record from the database.
@@ -747,7 +746,8 @@ mod tests {
747746
receipt.user_id(),
748747
subscription_start,
749748
subscription_expiry
750-
)[0],
749+
)
750+
.unwrap()[0],
751751
receipt
752752
);
753753

@@ -765,7 +765,8 @@ mod tests {
765765
latest_receipt.user_id(),
766766
subscription_start,
767767
latest_subscription_expiry
768-
),
768+
)
769+
.unwrap(),
769770
vec![receipt, latest_receipt.clone()]
770771
);
771772

@@ -774,7 +775,8 @@ mod tests {
774775
dbm.store_tower_record(tower_id, net_addr, &middle_receipt)
775776
.unwrap();
776777
assert_eq!(
777-
dbm.load_registration_receipt(tower_id, latest_receipt.user_id(), None, None)[0],
778+
dbm.load_registration_receipt(tower_id, latest_receipt.user_id(), None, None)
779+
.unwrap()[0],
778780
latest_receipt
779781
);
780782
}
@@ -799,7 +801,8 @@ mod tests {
799801
receipt.user_id(),
800802
subscription_start,
801803
subscription_expiry
802-
)[0],
804+
)
805+
.unwrap()[0],
803806
receipt
804807
);
805808

‎watchtower-plugin/src/main.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -143,7 +143,7 @@ async fn get_registration_receipt(
143143

144144
let response =
145145
state.get_registration_receipt(tower_id, subscription_start, subscription_expiry);
146-
if response.is_empty() {
146+
if response.clone().unwrap().is_empty() {
147147
if state.towers.contains_key(&tower_id) {
148148
Err(anyhow!("No registration receipt found for {tower_id}"))
149149
} else {

‎watchtower-plugin/src/wt_client.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -185,7 +185,7 @@ impl WTClient {
185185
tower_id: TowerId,
186186
subscription_start: Option<u32>,
187187
subscription_expiry: Option<u32>,
188-
) -> Vec<RegistrationReceipt> {
188+
) -> Option<Vec<RegistrationReceipt>> {
189189
self.dbm.load_registration_receipt(
190190
tower_id,
191191
self.user_id,

0 commit comments

Comments
 (0)