-
Notifications
You must be signed in to change notification settings - Fork 65
/
query_info.rs
134 lines (113 loc) · 4.25 KB
/
query_info.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
use crate::{kbucket::Key, rpc::RequestBody, Enr};
use enr::NodeId;
use sha2::digest::generic_array::GenericArray;
use smallvec::SmallVec;
use tokio::sync::oneshot;
/// Information about a query.
#[derive(Debug)]
pub struct QueryInfo {
/// What we are querying and why.
pub query_type: QueryType,
/// Temporary ENRs used when trying to reach nodes.
pub untrusted_enrs: SmallVec<[Enr; 16]>,
/// A callback channel for the service that requested the query.
pub callback: oneshot::Sender<Vec<Enr>>,
/// The number of distances we request for each peer.
pub distances_to_request: usize,
}
/// Additional information about the query.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum QueryType {
/// The user requested a `FIND_NODE` query to be performed. It should be reported when finished.
FindNode(NodeId),
}
impl QueryInfo {
/// Builds an RPC Request, given the QueryInfo
pub(crate) fn rpc_request(&self, peer: NodeId) -> Result<RequestBody, &'static str> {
let request = match self.query_type {
QueryType::FindNode(node_id) => {
let distances = findnode_log2distance(node_id, peer, self.distances_to_request)
.ok_or("Requested a node find itself")?;
RequestBody::FindNode { distances }
}
};
Ok(request)
}
}
impl crate::query_pool::TargetKey<NodeId> for QueryInfo {
fn key(&self) -> Key<NodeId> {
match self.query_type {
QueryType::FindNode(ref node_id) => {
Key::new_raw(*node_id, *GenericArray::from_slice(&node_id.raw()))
}
}
}
}
/// Calculates the log2 distance for a destination peer given a target and the size (number of
/// distances to request).
///
/// As the iteration increases, FINDNODE requests adjacent distances from the exact peer distance.
///
/// As an example, if the target has a distance of 12 from the remote peer, the sequence of distances that are sent for increasing iterations would be [12, 13, 11, 14, 10, .. ].
fn findnode_log2distance(target: NodeId, peer: NodeId, size: usize) -> Option<Vec<u64>> {
if size > 127 {
// invoke and endless loop - coding error
panic!("Iterations cannot be greater than 127");
}
let dst_key: Key<NodeId> = peer.into();
let distance = dst_key.log2_distance(&target.into())?;
let mut result_list = vec![distance];
let mut difference = 1;
while result_list.len() < size {
if distance + difference <= 256 {
result_list.push(distance + difference);
}
if result_list.len() < size {
if let Some(d) = distance.checked_sub(difference) {
result_list.push(d);
}
}
difference += 1;
}
Some(result_list[..size].to_vec())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_log2distance() {
let target = NodeId::new(&[0u8; 32]);
let mut destination = [0u8; 32];
destination[10] = 1; // gives a log2 distance of 169
let destination = NodeId::new(&destination);
let expected_distances = vec![169, 170, 168, 171, 167, 172, 166, 173, 165];
assert_eq!(
findnode_log2distance(target, destination, expected_distances.len()).unwrap(),
expected_distances
);
}
#[test]
fn test_log2distance_lower() {
let target = NodeId::new(&[0u8; 32]);
let mut destination = [0u8; 32];
destination[31] = 8; // gives a log2 distance of 5
let destination = NodeId::new(&destination);
let expected_distances = vec![4, 5, 3, 6, 2, 7, 1, 8, 0, 9, 10];
assert_eq!(
findnode_log2distance(target, destination, expected_distances.len()).unwrap(),
expected_distances
);
}
#[test]
fn test_log2distance_upper() {
let target = NodeId::new(&[0u8; 32]);
let mut destination = [0u8; 32];
destination[0] = 8; // gives a log2 distance of 252
let destination = NodeId::new(&destination);
let expected_distances = vec![252, 253, 251, 254, 250, 255, 249, 256, 248, 247, 246];
assert_eq!(
findnode_log2distance(target, destination, expected_distances.len()).unwrap(),
expected_distances
);
}
}