1use crate::tree::{LinkEntry, TreeRootEntry};
2use enr::EnrKeyUnambiguous;
3use linked_hash_set::LinkedHashSet;
4use secp256k1::SecretKey;
5use std::time::{Duration, Instant};
6
7pub(crate) struct SyncTree<K: EnrKeyUnambiguous = SecretKey> {
9 root: TreeRootEntry,
11 link: LinkEntry<K>,
13 root_updated: Instant,
15 sync_state: SyncState,
17 unresolved_links: LinkedHashSet<String>,
19 unresolved_nodes: LinkedHashSet<String>,
21}
22
23impl<K: EnrKeyUnambiguous> SyncTree<K> {
26 pub(crate) fn new(root: TreeRootEntry, link: LinkEntry<K>) -> Self {
27 Self {
28 root,
29 link,
30 root_updated: Instant::now(),
31 sync_state: SyncState::Pending,
32 unresolved_links: Default::default(),
33 unresolved_nodes: Default::default(),
34 }
35 }
36
37 #[cfg(test)]
38 pub(crate) const fn root(&self) -> &TreeRootEntry {
39 &self.root
40 }
41
42 pub(crate) const fn link(&self) -> &LinkEntry<K> {
43 &self.link
44 }
45
46 pub(crate) fn extend_children(
47 &mut self,
48 kind: ResolveKind,
49 children: impl IntoIterator<Item = String>,
50 ) {
51 match kind {
52 ResolveKind::Enr => {
53 self.unresolved_nodes.extend(children);
54 }
55 ResolveKind::Link => {
56 self.unresolved_links.extend(children);
57 }
58 }
59 }
60
61 pub(crate) fn poll(&mut self, now: Instant, update_timeout: Duration) -> Option<SyncAction> {
63 match self.sync_state {
64 SyncState::Pending => {
65 self.sync_state = SyncState::Enr;
66 return Some(SyncAction::Link(self.root.link_root.clone()))
67 }
68 SyncState::Enr => {
69 self.sync_state = SyncState::Active;
70 return Some(SyncAction::Enr(self.root.enr_root.clone()))
71 }
72 SyncState::Link => {
73 self.sync_state = SyncState::Active;
74 return Some(SyncAction::Link(self.root.link_root.clone()))
75 }
76 SyncState::Active => {
77 if now > self.root_updated + update_timeout {
78 self.sync_state = SyncState::RootUpdate;
79 return Some(SyncAction::UpdateRoot)
80 }
81 }
82 SyncState::RootUpdate => return None,
83 }
84
85 if let Some(link) = self.unresolved_links.pop_front() {
86 return Some(SyncAction::Link(link))
87 }
88
89 let enr = self.unresolved_nodes.pop_front()?;
90 Some(SyncAction::Enr(enr))
91 }
92
93 pub(crate) fn root_update_failed(&mut self) {
99 if matches!(self.sync_state, SyncState::RootUpdate) {
104 self.root_updated = Instant::now();
105 self.sync_state = SyncState::Active;
106 }
107 }
108
109 pub(crate) fn update_root(&mut self, root: TreeRootEntry) {
111 let enr_unchanged = root.enr_root == self.root.enr_root;
112 let link_unchanged = root.link_root == self.root.link_root;
113
114 self.root = root;
115 self.root_updated = Instant::now();
116
117 let state = match (enr_unchanged, link_unchanged) {
118 (true, true) => {
122 if !matches!(self.sync_state, SyncState::RootUpdate) {
123 return
124 }
125 SyncState::Active
126 }
127 (false, true) => {
129 self.unresolved_nodes.clear();
130 SyncState::Enr
131 }
132 (true, false) => {
134 self.unresolved_links.clear();
135 SyncState::Link
136 }
137 (false, false) => {
139 self.unresolved_nodes.clear();
140 self.unresolved_links.clear();
141 SyncState::Pending
142 }
143 };
144 self.sync_state = state;
145 }
146}
147
148#[derive(Debug)]
150pub(crate) enum SyncAction {
151 UpdateRoot,
152 Enr(String),
153 Link(String),
154}
155
156enum SyncState {
158 RootUpdate,
159 Pending,
160 Enr,
161 Link,
162 Active,
163}
164
165pub(crate) enum ResolveKind {
167 Enr,
168 Link,
169}
170
171impl ResolveKind {
174 pub(crate) const fn is_link(&self) -> bool {
175 matches!(self, Self::Link)
176 }
177}
178
179#[cfg(test)]
180mod tests {
181 use super::*;
182 use enr::EnrKey;
183 use secp256k1::rand::thread_rng;
184
185 fn base_root() -> TreeRootEntry {
186 let s = "enrtree-root:v1 e=QFT4PBCRX4XQCV3VUYJ6BTCEPU l=JGUFMSAGI7KZYB3P7IZW4S5Y3A seq=3 sig=3FmXuVwpa8Y7OstZTx9PIb1mt8FrW7VpDOFv4AaGCsZ2EIHmhraWhe4NxYhQDlw5MjeFXYMbJjsPeKlHzmJREQE";
188 s.parse::<TreeRootEntry>().unwrap()
189 }
190
191 fn make_tree() -> SyncTree {
192 let secret_key = SecretKey::new(&mut thread_rng());
193 let link =
194 LinkEntry { domain: "nodes.example.org".to_string(), pubkey: secret_key.public() };
195 SyncTree::new(base_root(), link)
196 }
197
198 fn advance_to_active(tree: &mut SyncTree) {
199 let now = Instant::now();
201 let timeout = Duration::from_secs(60 * 60 * 24);
202 let _ = tree.poll(now, timeout);
203 let _ = tree.poll(now, timeout);
204 }
205
206 #[test]
207 fn update_root_unchanged_no_action_from_active() {
208 let mut tree = make_tree();
209 let now = Instant::now();
210 let timeout = Duration::from_secs(60 * 60 * 24);
211 advance_to_active(&mut tree);
212
213 let same = base_root();
215 tree.update_root(same);
216 assert!(tree.poll(now, timeout).is_none());
217 }
218
219 #[test]
220 fn unsolicited_root_outcome_does_not_cancel_a_walk() {
221 let timeout = Duration::from_secs(60);
222
223 for unsolicited in [true, false] {
227 let mut tree = make_tree();
228 assert!(matches!(tree.poll(Instant::now(), timeout), Some(SyncAction::Link(_))));
230
231 if unsolicited {
232 tree.update_root(base_root());
233 } else {
234 tree.root_update_failed();
235 }
236
237 assert!(
238 matches!(tree.poll(Instant::now(), timeout), Some(SyncAction::Enr(_))),
239 "walk was cancelled by an unsolicited root outcome"
240 );
241 }
242 }
243
244 #[test]
245 fn update_root_unchanged_still_rechecks_later() {
246 let mut tree = make_tree();
247 let timeout = Duration::from_secs(60);
248 advance_to_active(&mut tree);
249
250 let due = Instant::now() + timeout * 2;
252 assert!(matches!(tree.poll(due, timeout), Some(SyncAction::UpdateRoot)));
253
254 tree.update_root(base_root());
256 assert!(tree.poll(Instant::now(), timeout).is_none());
257
258 let later = Instant::now() + timeout * 4;
260 assert!(
261 matches!(tree.poll(later, timeout), Some(SyncAction::UpdateRoot)),
262 "tree stopped rechecking its root after an unchanged update"
263 );
264 }
265
266 #[test]
267 fn failed_root_update_still_rechecks_later() {
268 let mut tree = make_tree();
269 let timeout = Duration::from_secs(60);
270 advance_to_active(&mut tree);
271
272 let due = Instant::now() + timeout * 2;
273 assert!(matches!(tree.poll(due, timeout), Some(SyncAction::UpdateRoot)));
274
275 tree.root_update_failed();
277 assert!(tree.poll(Instant::now(), timeout).is_none());
278
279 let later = Instant::now() + timeout * 4;
280 assert!(
281 matches!(tree.poll(later, timeout), Some(SyncAction::UpdateRoot)),
282 "tree stopped rechecking its root after a failed lookup"
283 );
284 }
285
286 #[test]
287 fn update_root_only_enr_changed_triggers_enr() {
288 let mut tree = make_tree();
289 advance_to_active(&mut tree);
290 let mut new_root = base_root();
291 new_root.enr_root = "NEW_ENR_ROOT".to_string();
292 let now = Instant::now();
293 let timeout = Duration::from_secs(60 * 60 * 24);
294
295 tree.update_root(new_root.clone());
296 match tree.poll(now, timeout) {
297 Some(SyncAction::Enr(hash)) => assert_eq!(hash, new_root.enr_root),
298 other => panic!("expected Enr action, got {:?}", other),
299 }
300 }
301
302 #[test]
303 fn update_root_only_link_changed_triggers_link() {
304 let mut tree = make_tree();
305 advance_to_active(&mut tree);
306 let mut new_root = base_root();
307 new_root.link_root = "NEW_LINK_ROOT".to_string();
308 let now = Instant::now();
309 let timeout = Duration::from_secs(60 * 60 * 24);
310
311 tree.update_root(new_root.clone());
312 match tree.poll(now, timeout) {
313 Some(SyncAction::Link(hash)) => assert_eq!(hash, new_root.link_root),
314 other => panic!("expected Link action, got {:?}", other),
315 }
316 }
317
318 #[test]
319 fn update_root_both_changed_triggers_link_then_enr() {
320 let mut tree = make_tree();
321 advance_to_active(&mut tree);
322 let mut new_root = base_root();
323 new_root.enr_root = "NEW_ENR_ROOT".to_string();
324 new_root.link_root = "NEW_LINK_ROOT".to_string();
325 let now = Instant::now();
326 let timeout = Duration::from_secs(60 * 60 * 24);
327
328 tree.update_root(new_root.clone());
329 match tree.poll(now, timeout) {
330 Some(SyncAction::Link(hash)) => assert_eq!(hash, new_root.link_root),
331 other => panic!("expected first Link action, got {:?}", other),
332 }
333 match tree.poll(now, timeout) {
334 Some(SyncAction::Enr(hash)) => assert_eq!(hash, new_root.enr_root),
335 other => panic!("expected second Enr action, got {:?}", other),
336 }
337 }
338}