libp2p_kad/kbucket/
key.rs1use crate::record;
22use libp2p_core::multihash::Multihash;
23use libp2p_identity::PeerId;
24use sha2::digest::generic_array::{typenum::U32, GenericArray};
25use sha2::{Digest, Sha256};
26use std::borrow::Borrow;
27use std::hash::{Hash, Hasher};
28use uint::*;
29
30construct_uint! {
31 pub(super) struct U256(4);
33}
34
35#[derive(Clone, Copy, Debug)]
43pub struct Key<T> {
44 preimage: T,
45 bytes: KeyBytes,
46}
47
48impl<T> Key<T> {
49 pub fn new(preimage: T) -> Key<T>
55 where
56 T: Borrow<[u8]>,
57 {
58 let bytes = KeyBytes::new(preimage.borrow());
59 Key { preimage, bytes }
60 }
61
62 pub fn preimage(&self) -> &T {
64 &self.preimage
65 }
66
67 pub fn into_preimage(self) -> T {
69 self.preimage
70 }
71
72 pub fn distance<U>(&self, other: &U) -> Distance
74 where
75 U: AsRef<KeyBytes>,
76 {
77 self.bytes.distance(other)
78 }
79
80 pub fn hashed_bytes(&self) -> &[u8] {
82 &self.bytes.0
83 }
84
85 pub fn for_distance(&self, d: Distance) -> KeyBytes {
91 self.bytes.for_distance(d)
92 }
93}
94
95impl<T> From<Key<T>> for KeyBytes {
96 fn from(key: Key<T>) -> KeyBytes {
97 key.bytes
98 }
99}
100
101impl<const S: usize> From<Multihash<S>> for Key<Multihash<S>> {
102 fn from(m: Multihash<S>) -> Self {
103 let bytes = KeyBytes(Sha256::digest(m.to_bytes()));
104 Key { preimage: m, bytes }
105 }
106}
107
108impl From<PeerId> for Key<PeerId> {
109 fn from(p: PeerId) -> Self {
110 let bytes = KeyBytes(Sha256::digest(p.to_bytes()));
111 Key { preimage: p, bytes }
112 }
113}
114
115impl From<Vec<u8>> for Key<Vec<u8>> {
116 fn from(b: Vec<u8>) -> Self {
117 Key::new(b)
118 }
119}
120
121impl From<record::Key> for Key<record::Key> {
122 fn from(k: record::Key) -> Self {
123 Key::new(k)
124 }
125}
126
127impl<T> AsRef<KeyBytes> for Key<T> {
128 fn as_ref(&self) -> &KeyBytes {
129 &self.bytes
130 }
131}
132
133impl<T, U> PartialEq<Key<U>> for Key<T> {
134 fn eq(&self, other: &Key<U>) -> bool {
135 self.bytes == other.bytes
136 }
137}
138
139impl<T> Eq for Key<T> {}
140
141impl<T> Hash for Key<T> {
142 fn hash<H: Hasher>(&self, state: &mut H) {
143 self.bytes.0.hash(state);
144 }
145}
146
147#[derive(PartialEq, Eq, Clone, Copy, Debug)]
149pub struct KeyBytes(GenericArray<u8, U32>);
150
151impl KeyBytes {
152 pub fn new<T>(value: T) -> Self
155 where
156 T: Borrow<[u8]>,
157 {
158 KeyBytes(Sha256::digest(value.borrow()))
159 }
160
161 pub fn distance<U>(&self, other: &U) -> Distance
163 where
164 U: AsRef<KeyBytes>,
165 {
166 let a = U256::from(self.0.as_slice());
167 let b = U256::from(other.as_ref().0.as_slice());
168 Distance(a ^ b)
169 }
170
171 pub fn for_distance(&self, d: Distance) -> KeyBytes {
177 let key_int = U256::from(self.0.as_slice()) ^ d.0;
178 KeyBytes(GenericArray::from(<[u8; 32]>::from(key_int)))
179 }
180}
181
182impl AsRef<KeyBytes> for KeyBytes {
183 fn as_ref(&self) -> &KeyBytes {
184 self
185 }
186}
187
188#[derive(Copy, Clone, PartialEq, Eq, Default, PartialOrd, Ord, Debug)]
190pub struct Distance(pub(super) U256);
191
192impl Distance {
193 pub fn ilog2(&self) -> Option<u32> {
197 (256 - self.0.leading_zeros()).checked_sub(1)
198 }
199}
200
201#[cfg(test)]
202mod tests {
203 use super::*;
204 use crate::SHA_256_MH;
205 use quickcheck::*;
206
207 impl Arbitrary for Key<PeerId> {
208 fn arbitrary(_: &mut Gen) -> Key<PeerId> {
209 Key::from(PeerId::random())
210 }
211 }
212
213 impl Arbitrary for Key<Multihash<64>> {
214 fn arbitrary(g: &mut Gen) -> Key<Multihash<64>> {
215 let hash: [u8; 32] = core::array::from_fn(|_| u8::arbitrary(g));
216 Key::from(Multihash::wrap(SHA_256_MH, &hash).unwrap())
217 }
218 }
219
220 #[test]
221 fn identity() {
222 fn prop(a: Key<PeerId>) -> bool {
223 a.distance(&a) == Distance::default()
224 }
225 quickcheck(prop as fn(_) -> _)
226 }
227
228 #[test]
229 fn symmetry() {
230 fn prop(a: Key<PeerId>, b: Key<PeerId>) -> bool {
231 a.distance(&b) == b.distance(&a)
232 }
233 quickcheck(prop as fn(_, _) -> _)
234 }
235
236 #[test]
237 fn triangle_inequality() {
238 fn prop(a: Key<PeerId>, b: Key<PeerId>, c: Key<PeerId>) -> TestResult {
239 let ab = a.distance(&b);
240 let bc = b.distance(&c);
241 let (ab_plus_bc, overflow) = ab.0.overflowing_add(bc.0);
242 if overflow {
243 TestResult::discard()
244 } else {
245 TestResult::from_bool(a.distance(&c) <= Distance(ab_plus_bc))
246 }
247 }
248 quickcheck(prop as fn(_, _, _) -> _)
249 }
250
251 #[test]
252 fn unidirectionality() {
253 fn prop(a: Key<PeerId>, b: Key<PeerId>) -> bool {
254 let d = a.distance(&b);
255 (0..100).all(|_| {
256 let c = Key::from(PeerId::random());
257 a.distance(&c) != d || b == c
258 })
259 }
260 quickcheck(prop as fn(_, _) -> _)
261 }
262}