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