1use std::sync::Arc;
22
23use chia_protocol::{Bytes32, CoinSpend, CoinState, CoinStateFilters, Program};
24use dig_chainsource_interface::{
25 ChainSource, ChainSourceError, ChainSourceProvider, CoinRecord, ProviderInfo, SingletonLineage,
26};
27use tokio::runtime::Handle;
28use tokio::sync::RwLock;
29
30use crate::bridge::run_blocking;
31use crate::cache::CoinStateCache;
32use crate::fetcher::CoinStateFetcher;
33
34#[derive(Clone)]
38pub struct ChiaPeerProvider {
39 fetcher: Arc<dyn CoinStateFetcher>,
40 cache: Arc<RwLock<CoinStateCache>>,
41 handle: Handle,
42 info: ProviderInfo,
43}
44
45impl ChiaPeerProvider {
46 pub fn new(
50 fetcher: Arc<dyn CoinStateFetcher>,
51 cache: Arc<RwLock<CoinStateCache>>,
52 handle: Handle,
53 info: ProviderInfo,
54 ) -> Self {
55 Self {
56 fetcher,
57 cache,
58 handle,
59 info,
60 }
61 }
62
63 fn coin_state(&self, coin_id: Bytes32) -> Result<Option<CoinState>, ChainSourceError> {
65 let fetcher = self.fetcher.clone();
66 let cache = self.cache.clone();
67 run_blocking(&self.handle, async move {
68 if let Some(cached) = cache.read().await.get(coin_id) {
69 return Ok(Some(cached));
70 }
71 let states = fetcher.coin_states(vec![coin_id], false).await?;
72 Ok::<_, crate::error::ChiaPeerError>(
73 states.into_iter().find(|s| s.coin.coin_id() == coin_id),
74 )
75 })?
76 .map_err(ChainSourceError::from)
77 }
78}
79
80impl ChainSource for ChiaPeerProvider {
81 type Error = ChainSourceError;
82
83 fn coin_record(&self, coin_id: Bytes32) -> Result<Option<CoinRecord>, Self::Error> {
84 Ok(self.coin_state(coin_id)?.map(CoinRecord::from_coin_state))
85 }
86
87 fn coin_records_by_puzzle_hash(
88 &self,
89 puzzle_hash: Bytes32,
90 include_spent: bool,
91 ) -> Result<Vec<CoinRecord>, Self::Error> {
92 let fetcher = self.fetcher.clone();
93 let filters = CoinStateFilters {
94 include_spent,
95 include_unspent: true,
96 include_hinted: true,
97 min_amount: 0,
98 };
99 let states = run_blocking(&self.handle, async move {
100 fetcher
101 .puzzle_states(vec![puzzle_hash], filters, false)
102 .await
103 })?
104 .map_err(ChainSourceError::from)?;
105 Ok(states
106 .into_iter()
107 .map(CoinRecord::from_coin_state)
108 .collect())
109 }
110
111 fn coin_records_by_parent(
112 &self,
113 parent_coin_id: Bytes32,
114 ) -> Result<Vec<CoinRecord>, Self::Error> {
115 let fetcher = self.fetcher.clone();
116 let states = run_blocking(&self.handle, async move {
117 fetcher.children(parent_coin_id).await
118 })?
119 .map_err(ChainSourceError::from)?;
120 Ok(states
121 .into_iter()
122 .map(CoinRecord::from_coin_state)
123 .collect())
124 }
125
126 fn coin_spend(&self, coin_id: Bytes32) -> Result<Option<CoinSpend>, Self::Error> {
127 let Some(state) = self.coin_state(coin_id)? else {
130 return Ok(None);
131 };
132 let Some(spent_height) = state.spent_height else {
133 return Ok(None);
134 };
135 let fetcher = self.fetcher.clone();
136 let (puzzle, solution) = run_blocking(&self.handle, async move {
137 fetcher.puzzle_and_solution(coin_id, spent_height).await
138 })?
139 .map_err(ChainSourceError::from)?;
140
141 verify_reveal_matches(&puzzle, state.coin.puzzle_hash)?;
144 Ok(Some(CoinSpend::new(state.coin, puzzle, solution)))
145 }
146
147 fn resolve_singleton_lineage(
148 &self,
149 _launcher_id: Bytes32,
150 ) -> Result<Option<SingletonLineage>, Self::Error> {
151 Err(ChainSourceError::Unsupported(
152 "singleton lineage resolution is not provided by the light-client source; \
153 use an aggregating chain source",
154 ))
155 }
156
157 fn peak_height(&self) -> Result<Option<u32>, Self::Error> {
158 let cache = self.cache.clone();
159 let peak = run_blocking(&self.handle, async move { cache.read().await.peak() })?;
160 Ok(peak.map(|(height, _)| height))
161 }
162
163 fn block_timestamp(&self, _height: u32) -> Result<Option<u64>, Self::Error> {
164 Err(ChainSourceError::Unsupported(
165 "block timestamps are not indexed by the light-client source",
166 ))
167 }
168}
169
170impl ChainSourceProvider for ChiaPeerProvider {
171 fn provider_info(&self) -> ProviderInfo {
172 self.info.clone()
173 }
174}
175
176fn verify_reveal_matches(puzzle: &Program, expected: Bytes32) -> Result<(), ChainSourceError> {
179 let actual: Bytes32 = chia::clvm_utils::tree_hash_from_bytes(puzzle.as_ref())
180 .map_err(|e| ChainSourceError::Malformed(format!("undecodable puzzle reveal: {e}")))?
181 .into();
182 if actual != expected {
183 return Err(ChainSourceError::Malformed(
184 "puzzle reveal does not hash to the coin's puzzle hash".into(),
185 ));
186 }
187 Ok(())
188}
189
190#[cfg(test)]
191mod tests {
192 use super::*;
193 use crate::error::ChiaPeerError;
194 use async_trait::async_trait;
195 use chia_protocol::{Coin, Program};
196 use dig_chainsource_interface::{ProviderId, ProviderKind};
197 use std::borrow::Cow;
198
199 #[derive(Default, Clone)]
202 struct MockFetcher {
203 coin_states: Vec<CoinState>,
204 fail: Option<ChiaPeerError>,
205 children: Vec<CoinState>,
206 puzzle_states: Vec<CoinState>,
207 reveal: Option<(Program, Program)>,
208 }
209
210 #[async_trait]
211 impl CoinStateFetcher for MockFetcher {
212 async fn coin_states(
213 &self,
214 _coin_ids: Vec<Bytes32>,
215 _subscribe: bool,
216 ) -> Result<Vec<CoinState>, ChiaPeerError> {
217 match &self.fail {
218 Some(e) => Err(e.clone()),
219 None => Ok(self.coin_states.clone()),
220 }
221 }
222 async fn puzzle_states(
223 &self,
224 _puzzle_hashes: Vec<Bytes32>,
225 _filters: CoinStateFilters,
226 _subscribe: bool,
227 ) -> Result<Vec<CoinState>, ChiaPeerError> {
228 match &self.fail {
229 Some(e) => Err(e.clone()),
230 None => Ok(self.puzzle_states.clone()),
231 }
232 }
233 async fn children(&self, _coin_id: Bytes32) -> Result<Vec<CoinState>, ChiaPeerError> {
234 match &self.fail {
235 Some(e) => Err(e.clone()),
236 None => Ok(self.children.clone()),
237 }
238 }
239 async fn puzzle_and_solution(
240 &self,
241 _coin_id: Bytes32,
242 _height: u32,
243 ) -> Result<(Program, Program), ChiaPeerError> {
244 if let Some(e) = &self.fail {
245 return Err(e.clone());
246 }
247 match &self.reveal {
248 Some(reveal) => Ok(reveal.clone()),
249 None => Err(ChiaPeerError::Rejected("no reveal".into())),
251 }
252 }
253 }
254
255 fn reveal_and_matching_puzzle_hash() -> (Program, Bytes32) {
258 let puzzle = Program::from(vec![1u8]);
259 let ph: Bytes32 = chia::clvm_utils::tree_hash_from_bytes(puzzle.as_ref())
260 .unwrap()
261 .into();
262 (puzzle, ph)
263 }
264
265 fn info() -> ProviderInfo {
266 ProviderInfo {
267 id: ProviderId(Cow::Borrowed("chia-peer-test")),
268 kind: ProviderKind::Custom,
269 priority: 20,
270 trustless: false,
271 }
272 }
273
274 fn provider_with(fetcher: MockFetcher) -> (tokio::runtime::Runtime, ChiaPeerProvider) {
275 let rt = tokio::runtime::Builder::new_multi_thread()
276 .worker_threads(1)
277 .enable_all()
278 .build()
279 .expect("multi-thread runtime");
280 let provider = ChiaPeerProvider::new(
281 Arc::new(fetcher),
282 Arc::new(RwLock::new(CoinStateCache::new())),
283 rt.handle().clone(),
284 info(),
285 );
286 (rt, provider)
287 }
288
289 fn call<T: Send>(f: impl FnOnce() -> T + Send) -> T {
291 std::thread::scope(|s| s.spawn(f).join().expect("thread panicked"))
292 }
293
294 fn coin(seed: u8) -> Coin {
295 Coin::new(Bytes32::new([seed; 32]), Bytes32::new([seed ^ 1; 32]), 1)
296 }
297
298 #[test]
301 fn coin_record_returns_some_for_a_known_coin() {
302 let c = coin(7);
303 let id = c.coin_id();
304 let fetcher = MockFetcher {
305 coin_states: vec![CoinState {
306 coin: c,
307 created_height: Some(100),
308 spent_height: None,
309 }],
310 ..Default::default()
311 };
312 let (_rt, provider) = provider_with(fetcher);
313 let record = call(move || provider.coin_record(id)).expect("read ok");
314 assert!(record.is_some());
315 assert_eq!(record.unwrap().confirmed_height, Some(100));
316 }
317
318 #[test]
319 fn coin_record_returns_none_for_provable_absence() {
320 let (_rt, provider) = provider_with(MockFetcher::default());
321 let id = coin(9).coin_id();
322 let record = call(move || provider.coin_record(id)).expect("read ok");
323 assert_eq!(record, None);
324 }
325
326 #[test]
327 fn transport_failure_is_err_never_false_absence() {
328 let fetcher = MockFetcher {
329 fail: Some(ChiaPeerError::Transport("socket reset".into())),
330 ..Default::default()
331 };
332 let (_rt, provider) = provider_with(fetcher);
333 let id = coin(3).coin_id();
334 let result = call(move || provider.coin_record(id));
335 assert!(
336 matches!(result, Err(ChainSourceError::Transport(_))),
337 "a transport failure MUST be Err, never Ok(None): {result:?}"
338 );
339 }
340
341 #[test]
342 fn coin_spend_of_unspent_coin_is_none() {
343 let c = coin(4);
344 let id = c.coin_id();
345 let fetcher = MockFetcher {
346 coin_states: vec![CoinState {
347 coin: c,
348 created_height: Some(10),
349 spent_height: None,
350 }],
351 ..Default::default()
352 };
353 let (_rt, provider) = provider_with(fetcher);
354 assert_eq!(call(move || provider.coin_spend(id)).unwrap(), None);
355 }
356
357 #[test]
358 fn coin_spend_of_spent_coin_assembles_from_real_coin() {
359 let (puzzle, ph) = reveal_and_matching_puzzle_hash();
360 let c = Coin::new(Bytes32::new([5; 32]), ph, 1);
361 let id = c.coin_id();
362 let fetcher = MockFetcher {
363 coin_states: vec![CoinState {
364 coin: c,
365 created_height: Some(10),
366 spent_height: Some(20),
367 }],
368 reveal: Some((puzzle, Program::from(vec![2u8]))),
369 ..Default::default()
370 };
371 let (_rt, provider) = provider_with(fetcher);
372 let spend = call(move || provider.coin_spend(id))
373 .unwrap()
374 .expect("spend");
375 assert_eq!(spend.coin, c);
376 }
377
378 #[test]
381 fn coin_spend_of_spent_coin_with_rejected_reveal_is_err_never_none() {
382 let c = coin(6);
383 let id = c.coin_id();
384 let fetcher = MockFetcher {
385 coin_states: vec![CoinState {
386 coin: c,
387 created_height: Some(10),
388 spent_height: Some(20),
389 }],
390 reveal: None, ..Default::default()
392 };
393 let (_rt, provider) = provider_with(fetcher);
394 let result = call(move || provider.coin_spend(id));
395 assert!(
396 matches!(result, Err(ChainSourceError::Transport(_))),
397 "a rejected reveal for a spent coin must be Err, never Ok(None): {result:?}"
398 );
399 }
400
401 #[test]
404 fn coin_spend_rejects_a_reveal_that_does_not_hash_to_the_coin() {
405 let c = coin(8); let id = c.coin_id();
407 let fetcher = MockFetcher {
408 coin_states: vec![CoinState {
409 coin: c,
410 created_height: Some(10),
411 spent_height: Some(20),
412 }],
413 reveal: Some((Program::from(vec![1u8]), Program::from(vec![2u8]))),
414 ..Default::default()
415 };
416 let (_rt, provider) = provider_with(fetcher);
417 let result = call(move || provider.coin_spend(id));
418 assert!(
419 matches!(result, Err(ChainSourceError::Malformed(_))),
420 "a mismatched reveal must be Malformed: {result:?}"
421 );
422 }
423
424 #[test]
425 fn records_by_puzzle_hash_and_parent_map_states() {
426 let fetcher = MockFetcher {
427 puzzle_states: vec![CoinState {
428 coin: coin(1),
429 created_height: Some(1),
430 spent_height: None,
431 }],
432 children: vec![CoinState {
433 coin: coin(2),
434 created_height: Some(2),
435 spent_height: None,
436 }],
437 ..Default::default()
438 };
439 let (_rt, provider) = provider_with(fetcher);
440 let ph = Bytes32::new([8; 32]);
441 let parent = Bytes32::new([9; 32]);
442 let p = provider.clone();
443 assert_eq!(
444 call(move || p.coin_records_by_puzzle_hash(ph, true))
445 .unwrap()
446 .len(),
447 1
448 );
449 assert_eq!(
450 call(move || provider.coin_records_by_parent(parent))
451 .unwrap()
452 .len(),
453 1
454 );
455 }
456
457 #[test]
458 fn lineage_and_timestamp_are_unsupported_not_false_absence() {
459 let (_rt, provider) = provider_with(MockFetcher::default());
460 let p = provider.clone();
461 assert!(matches!(
462 call(move || p.resolve_singleton_lineage(Bytes32::new([1; 32]))),
463 Err(ChainSourceError::Unsupported(_))
464 ));
465 assert!(matches!(
466 call(move || provider.block_timestamp(1)),
467 Err(ChainSourceError::Unsupported(_))
468 ));
469 }
470
471 #[test]
472 fn provider_info_is_reported() {
473 let (_rt, provider) = provider_with(MockFetcher::default());
474 assert_eq!(provider.provider_info().priority, 20);
475 assert_eq!(provider.peak_height().unwrap(), None);
476 }
477}