hotshot_libp2p_networking/network/behaviours/dht/store/
validated.rs

1//! This file contains the `ValidatedStore` struct, which is a wrapper around a `RecordStore` that
2//! validates records before storing them.
3//!
4//! The `ValidatedStore` struct is used to ensure that only valid records are stored in the DHT.
5
6use std::marker::PhantomData;
7
8use delegate::delegate;
9use hotshot_types::traits::signature_key::SignatureKey;
10use libp2p::kad::store::{Error, RecordStore, Result};
11use tracing::warn;
12
13use crate::network::behaviours::dht::record::{RecordKey, RecordValue};
14
15/// A `RecordStore` wrapper that validates records before storing them.
16pub struct ValidatedStore<R: RecordStore, K: SignatureKey> {
17    /// The underlying store
18    store: R,
19
20    /// Phantom type for the key
21    phantom: std::marker::PhantomData<K>,
22}
23
24impl<R: RecordStore, K: SignatureKey> ValidatedStore<R, K> {
25    /// Create a new `ValidatedStore` with the given underlying store
26    pub fn new(store: R) -> Self {
27        ValidatedStore {
28            store,
29            phantom: PhantomData,
30        }
31    }
32}
33
34/// Implement the `RecordStore` trait for `ValidatedStore`
35impl<R: RecordStore, K: SignatureKey> RecordStore for ValidatedStore<R, K>
36where
37    K: 'static,
38{
39    type ProvidedIter<'a>
40        = R::ProvidedIter<'a>
41    where
42        R: 'a,
43        K: 'a;
44    type RecordsIter<'a>
45        = R::RecordsIter<'a>
46    where
47        R: 'a,
48        K: 'a;
49
50    // Delegate all `RecordStore` methods except `put` to the inner store
51    delegate! {
52        to self.store{
53            fn add_provider(&mut self, record: libp2p::kad::ProviderRecord) -> libp2p::kad::store::Result<()>;
54            fn get(&self, k: &libp2p::kad::RecordKey) -> Option<std::borrow::Cow<'_, libp2p::kad::Record>>;
55            fn provided(&self) -> Self::ProvidedIter<'_>;
56            fn providers(&self, key: &libp2p::kad::RecordKey) -> Vec<libp2p::kad::ProviderRecord>;
57            fn records(&self) -> Self::RecordsIter<'_>;
58            fn remove(&mut self, k: &libp2p::kad::RecordKey);
59            fn remove_provider(&mut self, k: &libp2p::kad::RecordKey, p: &libp2p::PeerId);
60        }
61    }
62
63    /// Overwrite the `put` method to validate the record before storing it
64    fn put(&mut self, record: libp2p::kad::Record) -> Result<()> {
65        // Convert the record to the correct type
66        if let Ok(record_value) = RecordValue::<K>::try_from(record.clone()) {
67            // Convert the key to the correct type
68            let Ok(record_key) = RecordKey::try_from_bytes(&record.key.to_vec()) else {
69                warn!("Failed to convert record key");
70                return Err(Error::MaxRecords);
71            };
72
73            // If the record is signed by the correct key,
74            if record_value.validate(&record_key) {
75                // Store the record
76                if let Err(err) = self.store.put(record.clone()) {
77                    warn!("Failed to store record: {err:?}");
78                    return Err(Error::MaxRecords);
79                }
80            } else {
81                warn!("Failed to validate record");
82                return Err(Error::MaxRecords);
83            }
84        }
85
86        Ok(())
87    }
88}
89
90#[cfg(test)]
91mod test {
92    use hotshot_types::signature_key::BLSPubKey;
93    use libp2p::{
94        kad::{store::MemoryStore, Record},
95        PeerId,
96    };
97
98    use super::*;
99    use crate::network::behaviours::dht::record::Namespace;
100
101    /// Test that a valid record is stored
102    #[test]
103    fn test_valid_stored() {
104        // Generate a staking keypair
105        let (public_key, private_key) = BLSPubKey::generated_from_seed_indexed([1; 32], 1337);
106
107        // Create a value. The key is the public key
108        let value = vec![5, 6, 7, 8];
109
110        // Create a record key (as we need to sign both the key and the value)
111        let record_key = RecordKey::new(Namespace::Lookup, public_key.to_bytes());
112
113        // Sign the record and value with the private key
114        let record_value: RecordValue<BLSPubKey> =
115            RecordValue::new_signed(&record_key, value.clone(), &private_key).unwrap();
116
117        // Initialize the store
118        let mut store: ValidatedStore<MemoryStore, BLSPubKey> =
119            ValidatedStore::new(MemoryStore::new(PeerId::random()));
120
121        // Serialize the record value
122        let record_value_bytes =
123            bincode::serialize(&record_value).expect("Failed to serialize record value");
124
125        // Create and store the record
126        let record = Record::new(record_key.to_bytes(), record_value_bytes);
127        store.put(record).expect("Failed to store record");
128
129        // Check that the record is stored
130        let libp2p_record_key = libp2p::kad::RecordKey::new(&record_key.to_bytes());
131        let stored_record = store.get(&libp2p_record_key).expect("Failed to get record");
132        let stored_record_value: RecordValue<BLSPubKey> =
133            bincode::deserialize(&stored_record.value).expect("Failed to deserialize record value");
134
135        // Make sure the stored record is the same as the original record
136        assert_eq!(
137            record_value, stored_record_value,
138            "Stored record is not the same as original"
139        );
140    }
141
142    /// Test that an invalid record is not stored
143    #[test]
144    fn test_invalid_not_stored() {
145        // Generate a staking keypair
146        let (public_key, _) = BLSPubKey::generated_from_seed_indexed([1; 32], 1337);
147
148        // Create a record key (as we need to sign both the key and the value)
149        let record_key = RecordKey::new(Namespace::Lookup, public_key.to_bytes());
150
151        // Create a new (unsigned, invalid) record value
152        let record_value: RecordValue<BLSPubKey> = RecordValue::new(vec![2, 3]);
153
154        // Initialize the store
155        let mut store: ValidatedStore<MemoryStore, BLSPubKey> =
156            ValidatedStore::new(MemoryStore::new(PeerId::random()));
157
158        // Serialize the record value
159        let record_value_bytes =
160            bincode::serialize(&record_value).expect("Failed to serialize record value");
161
162        // Make sure we are unable to store the record
163        let record = Record::new(record_key.to_bytes(), record_value_bytes);
164        assert!(store.put(record).is_err(), "Should not have stored record");
165
166        // Check that the record is not stored
167        let libp2p_record_key = libp2p::kad::RecordKey::new(&record_key.to_bytes());
168        assert!(
169            store.get(&libp2p_record_key).is_none(),
170            "Should not have stored record"
171        );
172    }
173}
hotshot_libp2p_networking/network/behaviours/dht/store/validated.rs

hotshot_libp2p_networking/network/behaviours/dht/store/
validated.rs
