reth_dns_discovery/
sync.rs

1use crate::tree::{LinkEntry, TreeRootEntry};
2use enr::EnrKeyUnambiguous;
3use linked_hash_set::LinkedHashSet;
4use secp256k1::SecretKey;
5use std::time::{Duration, Instant};
6
7/// A sync-able tree
8pub(crate) struct SyncTree<K: EnrKeyUnambiguous = SecretKey> {
9    /// Root of the tree
10    root: TreeRootEntry,
11    /// Link to this tree
12    link: LinkEntry<K>,
13    /// Timestamp when the root was updated
14    root_updated: Instant,
15    /// The state of the tree sync progress.
16    sync_state: SyncState,
17    /// Unresolved links of the tree
18    unresolved_links: LinkedHashSet<String>,
19    /// Unresolved nodes of the tree
20    unresolved_nodes: LinkedHashSet<String>,
21}
22
23// === impl SyncTree ===
24
25impl<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    /// Advances the state of the tree by returning actions to perform
62    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    /// Updates the root and returns what changed
94    pub(crate) fn update_root(&mut self, root: TreeRootEntry) {
95        let enr_unchanged = root.enr_root == self.root.enr_root;
96        let link_unchanged = root.link_root == self.root.link_root;
97
98        self.root = root;
99        self.root_updated = Instant::now();
100
101        let state = match (enr_unchanged, link_unchanged) {
102            // both unchanged — no resync needed
103            (true, true) => return,
104            // only ENR changed
105            (false, true) => {
106                self.unresolved_nodes.clear();
107                SyncState::Enr
108            }
109            // only LINK changed
110            (true, false) => {
111                self.unresolved_links.clear();
112                SyncState::Link
113            }
114            // both changed
115            (false, false) => {
116                self.unresolved_nodes.clear();
117                self.unresolved_links.clear();
118                SyncState::Pending
119            }
120        };
121        self.sync_state = state;
122    }
123}
124
125/// The action to perform by the service
126#[derive(Debug)]
127pub(crate) enum SyncAction {
128    UpdateRoot,
129    Enr(String),
130    Link(String),
131}
132
133/// How the [`SyncTree::update_root`] changed the root
134enum SyncState {
135    RootUpdate,
136    Pending,
137    Enr,
138    Link,
139    Active,
140}
141
142/// What kind of hash to resolve
143pub(crate) enum ResolveKind {
144    Enr,
145    Link,
146}
147
148// === impl ResolveKind ===
149
150impl ResolveKind {
151    pub(crate) const fn is_link(&self) -> bool {
152        matches!(self, Self::Link)
153    }
154}
155
156#[cfg(test)]
157mod tests {
158    use super::*;
159    use enr::EnrKey;
160    use secp256k1::rand::thread_rng;
161
162    fn base_root() -> TreeRootEntry {
163        // taken from existing tests to ensure valid formatting
164        let s = "enrtree-root:v1 e=QFT4PBCRX4XQCV3VUYJ6BTCEPU l=JGUFMSAGI7KZYB3P7IZW4S5Y3A seq=3 sig=3FmXuVwpa8Y7OstZTx9PIb1mt8FrW7VpDOFv4AaGCsZ2EIHmhraWhe4NxYhQDlw5MjeFXYMbJjsPeKlHzmJREQE";
165        s.parse::<TreeRootEntry>().unwrap()
166    }
167
168    fn make_tree() -> SyncTree {
169        let secret_key = SecretKey::new(&mut thread_rng());
170        let link =
171            LinkEntry { domain: "nodes.example.org".to_string(), pubkey: secret_key.public() };
172        SyncTree::new(base_root(), link)
173    }
174
175    fn advance_to_active(tree: &mut SyncTree) {
176        // Move Pending -> (emit Link) -> Enr, then Enr -> (emit Enr) -> Active
177        let now = Instant::now();
178        let timeout = Duration::from_secs(60 * 60 * 24);
179        let _ = tree.poll(now, timeout);
180        let _ = tree.poll(now, timeout);
181    }
182
183    #[test]
184    fn update_root_unchanged_no_action_from_active() {
185        let mut tree = make_tree();
186        let now = Instant::now();
187        let timeout = Duration::from_secs(60 * 60 * 24);
188        advance_to_active(&mut tree);
189
190        // same root -> no resync
191        let same = base_root();
192        tree.update_root(same);
193        assert!(tree.poll(now, timeout).is_none());
194    }
195
196    #[test]
197    fn update_root_only_enr_changed_triggers_enr() {
198        let mut tree = make_tree();
199        advance_to_active(&mut tree);
200        let mut new_root = base_root();
201        new_root.enr_root = "NEW_ENR_ROOT".to_string();
202        let now = Instant::now();
203        let timeout = Duration::from_secs(60 * 60 * 24);
204
205        tree.update_root(new_root.clone());
206        match tree.poll(now, timeout) {
207            Some(SyncAction::Enr(hash)) => assert_eq!(hash, new_root.enr_root),
208            other => panic!("expected Enr action, got {:?}", other),
209        }
210    }
211
212    #[test]
213    fn update_root_only_link_changed_triggers_link() {
214        let mut tree = make_tree();
215        advance_to_active(&mut tree);
216        let mut new_root = base_root();
217        new_root.link_root = "NEW_LINK_ROOT".to_string();
218        let now = Instant::now();
219        let timeout = Duration::from_secs(60 * 60 * 24);
220
221        tree.update_root(new_root.clone());
222        match tree.poll(now, timeout) {
223            Some(SyncAction::Link(hash)) => assert_eq!(hash, new_root.link_root),
224            other => panic!("expected Link action, got {:?}", other),
225        }
226    }
227
228    #[test]
229    fn update_root_both_changed_triggers_link_then_enr() {
230        let mut tree = make_tree();
231        advance_to_active(&mut tree);
232        let mut new_root = base_root();
233        new_root.enr_root = "NEW_ENR_ROOT".to_string();
234        new_root.link_root = "NEW_LINK_ROOT".to_string();
235        let now = Instant::now();
236        let timeout = Duration::from_secs(60 * 60 * 24);
237
238        tree.update_root(new_root.clone());
239        match tree.poll(now, timeout) {
240            Some(SyncAction::Link(hash)) => assert_eq!(hash, new_root.link_root),
241            other => panic!("expected first Link action, got {:?}", other),
242        }
243        match tree.poll(now, timeout) {
244            Some(SyncAction::Enr(hash)) => assert_eq!(hash, new_root.enr_root),
245            other => panic!("expected second Enr action, got {:?}", other),
246        }
247    }
248}
reth_dns_discovery/sync.rs

reth_dns_discovery/
sync.rs
