Skip to content

Commit efc5fd4

Browse files
benthecarmanclaude
andcommitted
Add PaginatedKVStore support to VssStore
Implement PaginatedKVStore and PaginatedKVStoreSync traits for VssStore, enabling paginated key listing via cursor-based pagination using PageToken. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 3aef2b3 commit efc5fd4

File tree

1 file changed

+180
-1
lines changed

1 file changed

+180
-1
lines changed

src/io/vss_store.rs

Lines changed: 180 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,9 @@ use bitcoin::Network;
2424
use lightning::impl_writeable_tlv_based_enum;
2525
use lightning::io::{self, Error, ErrorKind};
2626
use lightning::sign::{EntropySource as LdkEntropySource, RandomBytes};
27-
use lightning::util::persist::{KVStore, KVStoreSync};
27+
use lightning::util::persist::{
28+
KVStore, KVStoreSync, PageToken, PaginatedKVStore, PaginatedKVStoreSync, PaginatedListResponse,
29+
};
2830
use lightning::util::ser::{Readable, Writeable};
2931
use prost::Message;
3032
use vss_client::client::VssClient;
@@ -377,6 +379,52 @@ impl KVStore for VssStore {
377379
}
378380
}
379381

382+
impl PaginatedKVStoreSync for VssStore {
383+
fn list_paginated(
384+
&self, primary_namespace: &str, secondary_namespace: &str, page_token: Option<PageToken>,
385+
) -> io::Result<PaginatedListResponse> {
386+
let internal_runtime = self.internal_runtime.as_ref().ok_or_else(|| {
387+
debug_assert!(false, "Failed to access internal runtime");
388+
let msg = format!("Failed to access internal runtime");
389+
Error::new(ErrorKind::Other, msg)
390+
})?;
391+
let primary_namespace = primary_namespace.to_string();
392+
let secondary_namespace = secondary_namespace.to_string();
393+
let inner = Arc::clone(&self.inner);
394+
let fut = async move {
395+
inner
396+
.list_paginated_internal(
397+
&inner.blocking_client,
398+
primary_namespace,
399+
secondary_namespace,
400+
page_token,
401+
)
402+
.await
403+
};
404+
tokio::task::block_in_place(move || internal_runtime.block_on(fut))
405+
}
406+
}
407+
408+
impl PaginatedKVStore for VssStore {
409+
fn list_paginated(
410+
&self, primary_namespace: &str, secondary_namespace: &str, page_token: Option<PageToken>,
411+
) -> impl Future<Output = Result<PaginatedListResponse, io::Error>> + 'static + Send {
412+
let primary_namespace = primary_namespace.to_string();
413+
let secondary_namespace = secondary_namespace.to_string();
414+
let inner = Arc::clone(&self.inner);
415+
async move {
416+
inner
417+
.list_paginated_internal(
418+
&inner.async_client,
419+
primary_namespace,
420+
secondary_namespace,
421+
page_token,
422+
)
423+
.await
424+
}
425+
}
426+
}
427+
380428
impl Drop for VssStore {
381429
fn drop(&mut self) {
382430
let internal_runtime = self.internal_runtime.take();
@@ -638,6 +686,49 @@ impl VssStoreInner {
638686
Ok(keys)
639687
}
640688

689+
async fn list_paginated_internal(
690+
&self, client: &VssClient<CustomRetryPolicy>, primary_namespace: String,
691+
secondary_namespace: String, page_token: Option<PageToken>,
692+
) -> io::Result<PaginatedListResponse> {
693+
check_namespace_key_validity(
694+
&primary_namespace,
695+
&secondary_namespace,
696+
None,
697+
"list_paginated",
698+
)?;
699+
700+
const PAGE_SIZE: i32 = 50;
701+
702+
let key_prefix = self.build_obfuscated_prefix(&primary_namespace, &secondary_namespace);
703+
let vss_page_token = page_token.map(|t| t.as_str().to_string());
704+
705+
let request = ListKeyVersionsRequest {
706+
store_id: self.store_id.clone(),
707+
key_prefix: Some(key_prefix),
708+
page_token: vss_page_token,
709+
page_size: Some(PAGE_SIZE),
710+
};
711+
712+
let response = client.list_key_versions(&request).await.map_err(|e| {
713+
let msg = format!(
714+
"Failed to list keys in {}/{}: {}",
715+
primary_namespace, secondary_namespace, e
716+
);
717+
Error::new(ErrorKind::Other, msg)
718+
})?;
719+
720+
let mut keys = Vec::with_capacity(response.key_versions.len());
721+
for kv in response.key_versions {
722+
keys.push(self.extract_key(&kv.key)?);
723+
}
724+
725+
// VSS uses empty string to signal the last page
726+
let next_page_token =
727+
response.next_page_token.filter(|t| !t.is_empty()).map(PageToken::new);
728+
729+
Ok(PaginatedListResponse { keys, next_page_token })
730+
}
731+
641732
async fn execute_locked_write<
642733
F: Future<Output = Result<(), lightning::io::Error>>,
643734
FN: FnOnce() -> F,
@@ -1051,4 +1142,92 @@ mod tests {
10511142
do_read_write_remove_list_persist(&vss_store);
10521143
drop(vss_store)
10531144
}
1145+
1146+
fn build_vss_store() -> VssStore {
1147+
let vss_base_url = std::env::var("TEST_VSS_BASE_URL").unwrap();
1148+
let mut rng = rng();
1149+
let rand_store_id: String = (0..7).map(|_| rng.sample(Alphanumeric) as char).collect();
1150+
let mut node_seed = [0u8; 64];
1151+
rng.fill_bytes(&mut node_seed);
1152+
let entropy = NodeEntropy::from_seed_bytes(node_seed);
1153+
VssStoreBuilder::new(entropy, vss_base_url, rand_store_id, Network::Testnet)
1154+
.build_with_sigs_auth(HashMap::new())
1155+
.unwrap()
1156+
}
1157+
1158+
#[test]
1159+
fn vss_paginated_listing() {
1160+
let store = build_vss_store();
1161+
let ns = "test_paginated";
1162+
let sub = "listing";
1163+
let num_entries = 5;
1164+
1165+
for i in 0..num_entries {
1166+
let key = format!("key_{:04}", i);
1167+
let data = vec![i as u8; 32];
1168+
KVStoreSync::write(&store, ns, sub, &key, data).unwrap();
1169+
}
1170+
1171+
let mut all_keys = Vec::new();
1172+
let mut page_token = None;
1173+
1174+
loop {
1175+
let response =
1176+
PaginatedKVStoreSync::list_paginated(&store, ns, sub, page_token).unwrap();
1177+
all_keys.extend(response.keys);
1178+
match response.next_page_token {
1179+
Some(token) => page_token = Some(token),
1180+
None => break,
1181+
}
1182+
}
1183+
1184+
assert_eq!(all_keys.len(), num_entries);
1185+
1186+
// Verify no duplicates
1187+
let mut unique = all_keys.clone();
1188+
unique.sort();
1189+
unique.dedup();
1190+
assert_eq!(unique.len(), num_entries);
1191+
}
1192+
1193+
#[test]
1194+
fn vss_paginated_empty_namespace() {
1195+
let store = build_vss_store();
1196+
let response =
1197+
PaginatedKVStoreSync::list_paginated(&store, "nonexistent", "ns", None).unwrap();
1198+
assert!(response.keys.is_empty());
1199+
assert!(response.next_page_token.is_none());
1200+
}
1201+
1202+
#[test]
1203+
fn vss_paginated_removal() {
1204+
let store = build_vss_store();
1205+
let ns = "test_paginated";
1206+
let sub = "removal";
1207+
1208+
KVStoreSync::write(&store, ns, sub, "a", vec![1u8; 8]).unwrap();
1209+
KVStoreSync::write(&store, ns, sub, "b", vec![2u8; 8]).unwrap();
1210+
KVStoreSync::write(&store, ns, sub, "c", vec![3u8; 8]).unwrap();
1211+
1212+
KVStoreSync::remove(&store, ns, sub, "b", false).unwrap();
1213+
1214+
let response = PaginatedKVStoreSync::list_paginated(&store, ns, sub, None).unwrap();
1215+
assert_eq!(response.keys.len(), 2);
1216+
assert!(!response.keys.contains(&"b".to_string()));
1217+
}
1218+
1219+
#[test]
1220+
fn vss_paginated_namespace_isolation() {
1221+
let store = build_vss_store();
1222+
1223+
KVStoreSync::write(&store, "ns_a", "sub", "key_1", vec![1u8; 8]).unwrap();
1224+
KVStoreSync::write(&store, "ns_a", "sub", "key_2", vec![2u8; 8]).unwrap();
1225+
KVStoreSync::write(&store, "ns_b", "sub", "key_3", vec![3u8; 8]).unwrap();
1226+
1227+
let response = PaginatedKVStoreSync::list_paginated(&store, "ns_a", "sub", None).unwrap();
1228+
assert_eq!(response.keys.len(), 2);
1229+
1230+
let response = PaginatedKVStoreSync::list_paginated(&store, "ns_b", "sub", None).unwrap();
1231+
assert_eq!(response.keys.len(), 1);
1232+
}
10541233
}

0 commit comments

Comments
 (0)