From efc5fd4f5d0dc55ace2b56bdeb346f21efeccd77 Mon Sep 17 00:00:00 2001 From: benthecarman Date: Wed, 1 Apr 2026 15:25:29 -0500 Subject: [PATCH] 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) --- src/io/vss_store.rs | 181 +++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 180 insertions(+), 1 deletion(-) diff --git a/src/io/vss_store.rs b/src/io/vss_store.rs index 2f7a689b2..da41fe767 100644 --- a/src/io/vss_store.rs +++ b/src/io/vss_store.rs @@ -24,7 +24,9 @@ use bitcoin::Network; use lightning::impl_writeable_tlv_based_enum; use lightning::io::{self, Error, ErrorKind}; use lightning::sign::{EntropySource as LdkEntropySource, RandomBytes}; -use lightning::util::persist::{KVStore, KVStoreSync}; +use lightning::util::persist::{ + KVStore, KVStoreSync, PageToken, PaginatedKVStore, PaginatedKVStoreSync, PaginatedListResponse, +}; use lightning::util::ser::{Readable, Writeable}; use prost::Message; use vss_client::client::VssClient; @@ -377,6 +379,52 @@ impl KVStore for VssStore { } } +impl PaginatedKVStoreSync for VssStore { + fn list_paginated( + &self, primary_namespace: &str, secondary_namespace: &str, page_token: Option, + ) -> io::Result { + let internal_runtime = self.internal_runtime.as_ref().ok_or_else(|| { + debug_assert!(false, "Failed to access internal runtime"); + let msg = format!("Failed to access internal runtime"); + Error::new(ErrorKind::Other, msg) + })?; + let primary_namespace = primary_namespace.to_string(); + let secondary_namespace = secondary_namespace.to_string(); + let inner = Arc::clone(&self.inner); + let fut = async move { + inner + .list_paginated_internal( + &inner.blocking_client, + primary_namespace, + secondary_namespace, + page_token, + ) + .await + }; + tokio::task::block_in_place(move || internal_runtime.block_on(fut)) + } +} + +impl PaginatedKVStore for VssStore { + fn list_paginated( + &self, primary_namespace: &str, secondary_namespace: &str, page_token: Option, + ) -> impl Future> + 'static + Send { + let primary_namespace = primary_namespace.to_string(); + let secondary_namespace = secondary_namespace.to_string(); + let inner = Arc::clone(&self.inner); + async move { + inner + .list_paginated_internal( + &inner.async_client, + primary_namespace, + secondary_namespace, + page_token, + ) + .await + } + } +} + impl Drop for VssStore { fn drop(&mut self) { let internal_runtime = self.internal_runtime.take(); @@ -638,6 +686,49 @@ impl VssStoreInner { Ok(keys) } + async fn list_paginated_internal( + &self, client: &VssClient, primary_namespace: String, + secondary_namespace: String, page_token: Option, + ) -> io::Result { + check_namespace_key_validity( + &primary_namespace, + &secondary_namespace, + None, + "list_paginated", + )?; + + const PAGE_SIZE: i32 = 50; + + let key_prefix = self.build_obfuscated_prefix(&primary_namespace, &secondary_namespace); + let vss_page_token = page_token.map(|t| t.as_str().to_string()); + + let request = ListKeyVersionsRequest { + store_id: self.store_id.clone(), + key_prefix: Some(key_prefix), + page_token: vss_page_token, + page_size: Some(PAGE_SIZE), + }; + + let response = client.list_key_versions(&request).await.map_err(|e| { + let msg = format!( + "Failed to list keys in {}/{}: {}", + primary_namespace, secondary_namespace, e + ); + Error::new(ErrorKind::Other, msg) + })?; + + let mut keys = Vec::with_capacity(response.key_versions.len()); + for kv in response.key_versions { + keys.push(self.extract_key(&kv.key)?); + } + + // VSS uses empty string to signal the last page + let next_page_token = + response.next_page_token.filter(|t| !t.is_empty()).map(PageToken::new); + + Ok(PaginatedListResponse { keys, next_page_token }) + } + async fn execute_locked_write< F: Future>, FN: FnOnce() -> F, @@ -1051,4 +1142,92 @@ mod tests { do_read_write_remove_list_persist(&vss_store); drop(vss_store) } + + fn build_vss_store() -> VssStore { + let vss_base_url = std::env::var("TEST_VSS_BASE_URL").unwrap(); + let mut rng = rng(); + let rand_store_id: String = (0..7).map(|_| rng.sample(Alphanumeric) as char).collect(); + let mut node_seed = [0u8; 64]; + rng.fill_bytes(&mut node_seed); + let entropy = NodeEntropy::from_seed_bytes(node_seed); + VssStoreBuilder::new(entropy, vss_base_url, rand_store_id, Network::Testnet) + .build_with_sigs_auth(HashMap::new()) + .unwrap() + } + + #[test] + fn vss_paginated_listing() { + let store = build_vss_store(); + let ns = "test_paginated"; + let sub = "listing"; + let num_entries = 5; + + for i in 0..num_entries { + let key = format!("key_{:04}", i); + let data = vec![i as u8; 32]; + KVStoreSync::write(&store, ns, sub, &key, data).unwrap(); + } + + let mut all_keys = Vec::new(); + let mut page_token = None; + + loop { + let response = + PaginatedKVStoreSync::list_paginated(&store, ns, sub, page_token).unwrap(); + all_keys.extend(response.keys); + match response.next_page_token { + Some(token) => page_token = Some(token), + None => break, + } + } + + assert_eq!(all_keys.len(), num_entries); + + // Verify no duplicates + let mut unique = all_keys.clone(); + unique.sort(); + unique.dedup(); + assert_eq!(unique.len(), num_entries); + } + + #[test] + fn vss_paginated_empty_namespace() { + let store = build_vss_store(); + let response = + PaginatedKVStoreSync::list_paginated(&store, "nonexistent", "ns", None).unwrap(); + assert!(response.keys.is_empty()); + assert!(response.next_page_token.is_none()); + } + + #[test] + fn vss_paginated_removal() { + let store = build_vss_store(); + let ns = "test_paginated"; + let sub = "removal"; + + KVStoreSync::write(&store, ns, sub, "a", vec![1u8; 8]).unwrap(); + KVStoreSync::write(&store, ns, sub, "b", vec![2u8; 8]).unwrap(); + KVStoreSync::write(&store, ns, sub, "c", vec![3u8; 8]).unwrap(); + + KVStoreSync::remove(&store, ns, sub, "b", false).unwrap(); + + let response = PaginatedKVStoreSync::list_paginated(&store, ns, sub, None).unwrap(); + assert_eq!(response.keys.len(), 2); + assert!(!response.keys.contains(&"b".to_string())); + } + + #[test] + fn vss_paginated_namespace_isolation() { + let store = build_vss_store(); + + KVStoreSync::write(&store, "ns_a", "sub", "key_1", vec![1u8; 8]).unwrap(); + KVStoreSync::write(&store, "ns_a", "sub", "key_2", vec![2u8; 8]).unwrap(); + KVStoreSync::write(&store, "ns_b", "sub", "key_3", vec![3u8; 8]).unwrap(); + + let response = PaginatedKVStoreSync::list_paginated(&store, "ns_a", "sub", None).unwrap(); + assert_eq!(response.keys.len(), 2); + + let response = PaginatedKVStoreSync::list_paginated(&store, "ns_b", "sub", None).unwrap(); + assert_eq!(response.keys.len(), 1); + } }