1use std::collections::HashSet;
7
8#[derive(Debug, Clone, Default)]
10pub struct ProviderCapabilities {
11 pub archive: bool,
13 pub trace: bool,
15 pub max_block_range: u64,
17 pub max_batch_size: usize,
19 pub supported_methods: HashSet<String>,
21}
22
23impl ProviderCapabilities {
24 pub fn full_node() -> Self {
26 Self {
27 archive: false,
28 trace: false,
29 max_block_range: 10_000,
30 max_batch_size: 100,
31 supported_methods: HashSet::new(),
32 }
33 }
34
35 pub fn archive_node() -> Self {
37 Self {
38 archive: true,
39 trace: true,
40 max_block_range: 0,
41 max_batch_size: 100,
42 supported_methods: HashSet::new(),
43 }
44 }
45
46 pub fn can_handle(&self, requirement: &RequestRequirement) -> bool {
48 if requirement.needs_archive && !self.archive {
49 return false;
50 }
51 if requirement.needs_trace && !self.trace {
52 return false;
53 }
54 if !self.supported_methods.is_empty() {
55 if let Some(ref method) = requirement.method {
56 if !self.supported_methods.contains(method.as_str()) {
57 return false;
58 }
59 }
60 }
61 true
62 }
63}
64
65#[derive(Debug, Clone, Default)]
67pub struct RequestRequirement {
68 pub needs_archive: bool,
70 pub needs_trace: bool,
72 pub method: Option<String>,
74}
75
76pub fn analyze_request(
81 method: &str,
82 block_param: Option<&str>,
83 current_block: u64,
84) -> RequestRequirement {
85 let mut req = RequestRequirement {
86 method: Some(method.to_string()),
87 ..Default::default()
88 };
89
90 if method.starts_with("debug_") || method.starts_with("trace_") {
92 req.needs_trace = true;
93 req.needs_archive = true; return req;
95 }
96
97 if let Some(block) = block_param {
99 if is_historical_block(block, current_block) {
100 req.needs_archive = true;
101 }
102 }
103
104 req
105}
106
107fn is_historical_block(block: &str, current_block: u64) -> bool {
112 match block {
113 "latest" | "pending" | "safe" | "finalized" => false,
114 "earliest" => true,
115 hex if hex.starts_with("0x") => {
116 if let Ok(num) = u64::from_str_radix(hex.trim_start_matches("0x"), 16) {
117 current_block.saturating_sub(num) > 256
119 } else {
120 false
121 }
122 }
123 _ => false,
124 }
125}
126
127pub fn select_capable_provider(
132 capabilities: &[ProviderCapabilities],
133 allowed: &[bool],
134 requirement: &RequestRequirement,
135) -> Option<usize> {
136 for (idx, (cap, &ok)) in capabilities.iter().zip(allowed.iter()).enumerate() {
138 if ok && cap.can_handle(requirement) {
139 return Some(idx);
140 }
141 }
142 None
143}
144
145#[cfg(test)]
146mod tests {
147 use super::*;
148
149 #[test]
150 fn full_node_handles_recent() {
151 let cap = ProviderCapabilities::full_node();
152 let req = RequestRequirement::default();
153 assert!(cap.can_handle(&req));
154 }
155
156 #[test]
157 fn full_node_rejects_archive() {
158 let cap = ProviderCapabilities::full_node();
159 let req = RequestRequirement {
160 needs_archive: true,
161 ..Default::default()
162 };
163 assert!(!cap.can_handle(&req));
164 }
165
166 #[test]
167 fn archive_node_handles_everything() {
168 let cap = ProviderCapabilities::archive_node();
169 assert!(cap.can_handle(&RequestRequirement {
170 needs_archive: true,
171 ..Default::default()
172 }));
173 assert!(cap.can_handle(&RequestRequirement {
174 needs_trace: true,
175 ..Default::default()
176 }));
177 assert!(cap.can_handle(&RequestRequirement::default()));
178 }
179
180 #[test]
181 fn analyze_trace_method() {
182 let req = analyze_request("debug_traceTransaction", None, 1000);
183 assert!(req.needs_trace);
184 assert!(req.needs_archive);
185 }
186
187 #[test]
188 fn analyze_historical_block() {
189 let req = analyze_request("eth_getBalance", Some("0x1"), 1_000_000);
190 assert!(req.needs_archive);
191 }
192
193 #[test]
194 fn analyze_recent_block() {
195 let req = analyze_request("eth_getBalance", Some("latest"), 1_000_000);
196 assert!(!req.needs_archive);
197 }
198
199 #[test]
200 fn analyze_earliest() {
201 let req = analyze_request("eth_getBalance", Some("earliest"), 1_000_000);
202 assert!(req.needs_archive);
203 }
204
205 #[test]
206 fn analyze_close_to_head() {
207 let current = 1_000_000u64;
208 let req = analyze_request(
209 "eth_getBalance",
210 Some(&format!("0x{:x}", current - 10)),
211 current,
212 );
213 assert!(!req.needs_archive); }
215
216 #[test]
217 fn select_capable() {
218 let caps = vec![
219 ProviderCapabilities::full_node(),
220 ProviderCapabilities::archive_node(),
221 ];
222 let allowed = [true, true];
223
224 let req = RequestRequirement {
226 needs_archive: true,
227 ..Default::default()
228 };
229 assert_eq!(select_capable_provider(&caps, &allowed, &req), Some(1));
230
231 let req = RequestRequirement::default();
233 assert_eq!(select_capable_provider(&caps, &allowed, &req), Some(0));
234 }
235
236 #[test]
237 fn select_when_no_capable() {
238 let caps = vec![ProviderCapabilities::full_node()];
239 let allowed = [true];
240
241 let req = RequestRequirement {
242 needs_archive: true,
243 ..Default::default()
244 };
245 assert_eq!(select_capable_provider(&caps, &allowed, &req), None);
246 }
247}