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 let peak = self.peak_height()?;
85 Ok(self
86 .coin_state(coin_id)?
87 .map(CoinRecord::from_coin_state)
88 .map(|record| clamp_record_heights_to_peak(record, peak)))
89 }
90
91 fn coin_records_by_puzzle_hash(
92 &self,
93 puzzle_hash: Bytes32,
94 include_spent: bool,
95 ) -> Result<Vec<CoinRecord>, Self::Error> {
96 let fetcher = self.fetcher.clone();
97 let filters = CoinStateFilters {
98 include_spent,
99 include_unspent: true,
100 include_hinted: true,
101 min_amount: 0,
102 };
103 let states = run_blocking(&self.handle, async move {
104 fetcher
105 .puzzle_states(vec![puzzle_hash], filters, false)
106 .await
107 })?
108 .map_err(ChainSourceError::from)?;
109 let peak = self.peak_height()?;
110 Ok(states
111 .into_iter()
112 .map(CoinRecord::from_coin_state)
113 .map(|record| clamp_record_heights_to_peak(record, peak))
114 .collect())
115 }
116
117 fn coin_records_by_parent(
118 &self,
119 parent_coin_id: Bytes32,
120 ) -> Result<Vec<CoinRecord>, Self::Error> {
121 let fetcher = self.fetcher.clone();
122 let states = run_blocking(&self.handle, async move {
123 fetcher.children(parent_coin_id).await
124 })?
125 .map_err(ChainSourceError::from)?;
126 let peak = self.peak_height()?;
127 Ok(states
128 .into_iter()
129 .map(CoinRecord::from_coin_state)
130 .map(|record| clamp_record_heights_to_peak(record, peak))
131 .collect())
132 }
133
134 fn coin_spend(&self, coin_id: Bytes32) -> Result<Option<CoinSpend>, Self::Error> {
135 let Some(state) = self.coin_state(coin_id)? else {
138 return Ok(None);
139 };
140 let Some(spent_height) = state.spent_height else {
141 return Ok(None);
142 };
143 let fetcher = self.fetcher.clone();
144 let (puzzle, solution) = run_blocking(&self.handle, async move {
145 fetcher.puzzle_and_solution(coin_id, spent_height).await
146 })?
147 .map_err(ChainSourceError::from)?;
148
149 verify_reveal_matches(&puzzle, state.coin.puzzle_hash)?;
152 Ok(Some(CoinSpend::new(state.coin, puzzle, solution)))
153 }
154
155 fn resolve_singleton_lineage(
156 &self,
157 _launcher_id: Bytes32,
158 ) -> Result<Option<SingletonLineage>, Self::Error> {
159 Err(ChainSourceError::Unsupported(
160 "singleton lineage resolution is not provided by the light-client source; \
161 use an aggregating chain source",
162 ))
163 }
164
165 fn peak_height(&self) -> Result<Option<u32>, Self::Error> {
166 let cache = self.cache.clone();
167 let peak = run_blocking(&self.handle, async move { cache.read().await.peak() })?;
168 Ok(peak.map(|(height, _)| height))
169 }
170
171 fn block_timestamp(&self, _height: u32) -> Result<Option<u64>, Self::Error> {
172 Err(ChainSourceError::Unsupported(
173 "block timestamps are not indexed by the light-client source",
174 ))
175 }
176}
177
178impl ChainSourceProvider for ChiaPeerProvider {
179 fn provider_info(&self) -> ProviderInfo {
180 self.info.clone()
181 }
182}
183
184fn clamp_record_heights_to_peak(mut record: CoinRecord, peak: Option<u32>) -> CoinRecord {
202 let Some(peak) = peak else { return record };
203 if let Some(confirmed) = record.confirmed_height {
204 record.confirmed_height = Some(confirmed.min(peak));
205 }
206 if let Some(spent) = record.spent_height {
207 record.spent_height = Some(spent.min(peak));
208 }
209 record
210}
211
212fn verify_reveal_matches(puzzle: &Program, expected: Bytes32) -> Result<(), ChainSourceError> {
215 let actual: Bytes32 = chia::clvm_utils::tree_hash_from_bytes(puzzle.as_ref())
216 .map_err(|e| ChainSourceError::Malformed(format!("undecodable puzzle reveal: {e}")))?
217 .into();
218 if actual != expected {
219 return Err(ChainSourceError::Malformed(
220 "puzzle reveal does not hash to the coin's puzzle hash".into(),
221 ));
222 }
223 Ok(())
224}
225
226#[cfg(test)]
227mod tests {
228 use super::*;
229 use crate::error::ChiaPeerError;
230 use async_trait::async_trait;
231 use chia_protocol::{Coin, Program};
232 use dig_chainsource_interface::{ProviderId, ProviderKind};
233 use std::borrow::Cow;
234
235 #[derive(Default, Clone)]
238 struct MockFetcher {
239 coin_states: Vec<CoinState>,
240 fail: Option<ChiaPeerError>,
241 children: Vec<CoinState>,
242 puzzle_states: Vec<CoinState>,
243 reveal: Option<(Program, Program)>,
244 }
245
246 #[async_trait]
247 impl CoinStateFetcher for MockFetcher {
248 async fn coin_states(
249 &self,
250 _coin_ids: Vec<Bytes32>,
251 _subscribe: bool,
252 ) -> Result<Vec<CoinState>, ChiaPeerError> {
253 match &self.fail {
254 Some(e) => Err(e.clone()),
255 None => Ok(self.coin_states.clone()),
256 }
257 }
258 async fn puzzle_states(
259 &self,
260 _puzzle_hashes: Vec<Bytes32>,
261 _filters: CoinStateFilters,
262 _subscribe: bool,
263 ) -> Result<Vec<CoinState>, ChiaPeerError> {
264 match &self.fail {
265 Some(e) => Err(e.clone()),
266 None => Ok(self.puzzle_states.clone()),
267 }
268 }
269 async fn children(&self, _coin_id: Bytes32) -> Result<Vec<CoinState>, ChiaPeerError> {
270 match &self.fail {
271 Some(e) => Err(e.clone()),
272 None => Ok(self.children.clone()),
273 }
274 }
275 async fn puzzle_and_solution(
276 &self,
277 _coin_id: Bytes32,
278 _height: u32,
279 ) -> Result<(Program, Program), ChiaPeerError> {
280 if let Some(e) = &self.fail {
281 return Err(e.clone());
282 }
283 match &self.reveal {
284 Some(reveal) => Ok(reveal.clone()),
285 None => Err(ChiaPeerError::Rejected("no reveal".into())),
287 }
288 }
289 }
290
291 fn reveal_and_matching_puzzle_hash() -> (Program, Bytes32) {
294 let puzzle = Program::from(vec![1u8]);
295 let ph: Bytes32 = chia::clvm_utils::tree_hash_from_bytes(puzzle.as_ref())
296 .unwrap()
297 .into();
298 (puzzle, ph)
299 }
300
301 fn info() -> ProviderInfo {
302 ProviderInfo {
303 id: ProviderId(Cow::Borrowed("chia-peer-test")),
304 kind: ProviderKind::Custom,
305 priority: 20,
306 trustless: false,
307 }
308 }
309
310 fn provider_with(fetcher: MockFetcher) -> (tokio::runtime::Runtime, ChiaPeerProvider) {
311 provider_with_peak(fetcher, None)
312 }
313
314 fn provider_with_peak(
317 fetcher: MockFetcher,
318 peak: Option<u32>,
319 ) -> (tokio::runtime::Runtime, ChiaPeerProvider) {
320 let rt = tokio::runtime::Builder::new_multi_thread()
321 .worker_threads(1)
322 .enable_all()
323 .build()
324 .expect("multi-thread runtime");
325 let mut cache = CoinStateCache::new();
326 if let Some(height) = peak {
327 cache.set_peak(height, Bytes32::new([0xAB; 32]));
328 }
329 let provider = ChiaPeerProvider::new(
330 Arc::new(fetcher),
331 Arc::new(RwLock::new(cache)),
332 rt.handle().clone(),
333 info(),
334 );
335 (rt, provider)
336 }
337
338 fn call<T: Send>(f: impl FnOnce() -> T + Send) -> T {
340 std::thread::scope(|s| s.spawn(f).join().expect("thread panicked"))
341 }
342
343 fn coin(seed: u8) -> Coin {
344 Coin::new(Bytes32::new([seed; 32]), Bytes32::new([seed ^ 1; 32]), 1)
345 }
346
347 #[test]
350 fn coin_record_returns_some_for_a_known_coin() {
351 let c = coin(7);
352 let id = c.coin_id();
353 let fetcher = MockFetcher {
354 coin_states: vec![CoinState {
355 coin: c,
356 created_height: Some(100),
357 spent_height: None,
358 }],
359 ..Default::default()
360 };
361 let (_rt, provider) = provider_with(fetcher);
362 let record = call(move || provider.coin_record(id)).expect("read ok");
363 assert!(record.is_some());
364 assert_eq!(record.unwrap().confirmed_height, Some(100));
365 }
366
367 #[test]
372 fn live_fetched_coin_above_peak_reports_clamped_confirmed_height() {
373 let c = coin(11);
374 let id = c.coin_id();
375 let fetcher = MockFetcher {
376 coin_states: vec![CoinState {
377 coin: c,
378 created_height: Some(1_000_001), spent_height: None,
380 }],
381 ..Default::default()
382 };
383 let (_rt, provider) = provider_with_peak(fetcher, Some(1_000_000));
384 let record = call(move || provider.coin_record(id))
385 .expect("read ok")
386 .expect("coin present, never omitted");
387 assert_eq!(
388 record.confirmed_height,
389 Some(1_000_000),
390 "an above-peak live coin must clamp to the peak (0 confirmations), never overstate"
391 );
392 }
393
394 #[test]
397 fn live_fetched_coin_at_or_below_peak_is_unaffected() {
398 let c = coin(12);
399 let id = c.coin_id();
400 let fetcher = MockFetcher {
401 coin_states: vec![CoinState {
402 coin: c,
403 created_height: Some(900_000),
404 spent_height: None,
405 }],
406 ..Default::default()
407 };
408 let (_rt, provider) = provider_with_peak(fetcher, Some(1_000_000));
409 let record = call(move || provider.coin_record(id))
410 .expect("read ok")
411 .expect("coin present");
412 assert_eq!(record.confirmed_height, Some(900_000));
413 }
414
415 #[test]
417 fn discovery_reads_clamp_above_peak_confirmed_height() {
418 let fetcher = MockFetcher {
419 puzzle_states: vec![CoinState {
420 coin: coin(13),
421 created_height: Some(2_000_000),
422 spent_height: Some(2_000_000),
423 }],
424 children: vec![CoinState {
425 coin: coin(14),
426 created_height: Some(2_000_000),
427 spent_height: Some(2_000_000),
428 }],
429 ..Default::default()
430 };
431 let (_rt, provider) = provider_with_peak(fetcher, Some(1_000_000));
432 let ph = Bytes32::new([8; 32]);
433 let parent = Bytes32::new([9; 32]);
434 let p = provider.clone();
435 let by_ph = call(move || p.coin_records_by_puzzle_hash(ph, true)).unwrap();
436 assert_eq!(by_ph[0].confirmed_height, Some(1_000_000));
437 assert_eq!(by_ph[0].spent_height, Some(1_000_000));
438 let by_parent = call(move || provider.coin_records_by_parent(parent)).unwrap();
439 assert_eq!(by_parent[0].confirmed_height, Some(1_000_000));
440 assert_eq!(by_parent[0].spent_height, Some(1_000_000));
441 }
442
443 #[test]
449 fn live_fetched_coin_spent_above_peak_reports_clamped_spent_height() {
450 let c = coin(15);
451 let id = c.coin_id();
452 let fetcher = MockFetcher {
453 coin_states: vec![CoinState {
454 coin: c,
455 created_height: Some(999_999),
456 spent_height: Some(1_000_001), }],
458 ..Default::default()
459 };
460 let (_rt, provider) = provider_with_peak(fetcher, Some(1_000_000));
461 let record = call(move || provider.coin_record(id))
462 .expect("read ok")
463 .expect("coin present, never omitted");
464 assert_eq!(
465 record.spent_height,
466 Some(1_000_000),
467 "an above-peak spent coin must clamp spent_height to the peak (0 spend-depth), never overstate"
468 );
469 assert!(
470 record.is_spent(),
471 "the coin IS spent — clamping the reported height must never drop the spent flag"
472 );
473 }
474
475 #[test]
478 fn live_fetched_coin_spent_at_or_below_peak_is_unaffected() {
479 let c = coin(16);
480 let id = c.coin_id();
481 let fetcher = MockFetcher {
482 coin_states: vec![CoinState {
483 coin: c,
484 created_height: Some(800_000),
485 spent_height: Some(900_000),
486 }],
487 ..Default::default()
488 };
489 let (_rt, provider) = provider_with_peak(fetcher, Some(1_000_000));
490 let record = call(move || provider.coin_record(id))
491 .expect("read ok")
492 .expect("coin present");
493 assert_eq!(record.spent_height, Some(900_000));
494 }
495
496 #[test]
499 fn coin_spend_of_coin_spent_above_peak_still_identified_as_spent() {
500 let (puzzle, ph) = reveal_and_matching_puzzle_hash();
501 let c = Coin::new(Bytes32::new([17; 32]), ph, 1);
502 let id = c.coin_id();
503 let fetcher = MockFetcher {
504 coin_states: vec![CoinState {
505 coin: c,
506 created_height: Some(999_999),
507 spent_height: Some(1_000_001), }],
509 reveal: Some((puzzle, Program::from(vec![2u8]))),
510 ..Default::default()
511 };
512 let (_rt, provider) = provider_with_peak(fetcher, Some(1_000_000));
513 let spend = call(move || provider.coin_spend(id))
514 .unwrap()
515 .expect("a coin spent above the lagged peak is still spent");
516 assert_eq!(spend.coin, c);
517 }
518
519 #[test]
520 fn coin_record_returns_none_for_provable_absence() {
521 let (_rt, provider) = provider_with(MockFetcher::default());
522 let id = coin(9).coin_id();
523 let record = call(move || provider.coin_record(id)).expect("read ok");
524 assert_eq!(record, None);
525 }
526
527 #[test]
528 fn transport_failure_is_err_never_false_absence() {
529 let fetcher = MockFetcher {
530 fail: Some(ChiaPeerError::Transport("socket reset".into())),
531 ..Default::default()
532 };
533 let (_rt, provider) = provider_with(fetcher);
534 let id = coin(3).coin_id();
535 let result = call(move || provider.coin_record(id));
536 assert!(
537 matches!(result, Err(ChainSourceError::Transport(_))),
538 "a transport failure MUST be Err, never Ok(None): {result:?}"
539 );
540 }
541
542 #[test]
543 fn coin_spend_of_unspent_coin_is_none() {
544 let c = coin(4);
545 let id = c.coin_id();
546 let fetcher = MockFetcher {
547 coin_states: vec![CoinState {
548 coin: c,
549 created_height: Some(10),
550 spent_height: None,
551 }],
552 ..Default::default()
553 };
554 let (_rt, provider) = provider_with(fetcher);
555 assert_eq!(call(move || provider.coin_spend(id)).unwrap(), None);
556 }
557
558 #[test]
559 fn coin_spend_of_spent_coin_assembles_from_real_coin() {
560 let (puzzle, ph) = reveal_and_matching_puzzle_hash();
561 let c = Coin::new(Bytes32::new([5; 32]), ph, 1);
562 let id = c.coin_id();
563 let fetcher = MockFetcher {
564 coin_states: vec![CoinState {
565 coin: c,
566 created_height: Some(10),
567 spent_height: Some(20),
568 }],
569 reveal: Some((puzzle, Program::from(vec![2u8]))),
570 ..Default::default()
571 };
572 let (_rt, provider) = provider_with(fetcher);
573 let spend = call(move || provider.coin_spend(id))
574 .unwrap()
575 .expect("spend");
576 assert_eq!(spend.coin, c);
577 }
578
579 #[test]
582 fn coin_spend_of_spent_coin_with_rejected_reveal_is_err_never_none() {
583 let c = coin(6);
584 let id = c.coin_id();
585 let fetcher = MockFetcher {
586 coin_states: vec![CoinState {
587 coin: c,
588 created_height: Some(10),
589 spent_height: Some(20),
590 }],
591 reveal: None, ..Default::default()
593 };
594 let (_rt, provider) = provider_with(fetcher);
595 let result = call(move || provider.coin_spend(id));
596 assert!(
597 matches!(result, Err(ChainSourceError::Transport(_))),
598 "a rejected reveal for a spent coin must be Err, never Ok(None): {result:?}"
599 );
600 }
601
602 #[test]
605 fn coin_spend_rejects_a_reveal_that_does_not_hash_to_the_coin() {
606 let c = coin(8); let id = c.coin_id();
608 let fetcher = MockFetcher {
609 coin_states: vec![CoinState {
610 coin: c,
611 created_height: Some(10),
612 spent_height: Some(20),
613 }],
614 reveal: Some((Program::from(vec![1u8]), Program::from(vec![2u8]))),
615 ..Default::default()
616 };
617 let (_rt, provider) = provider_with(fetcher);
618 let result = call(move || provider.coin_spend(id));
619 assert!(
620 matches!(result, Err(ChainSourceError::Malformed(_))),
621 "a mismatched reveal must be Malformed: {result:?}"
622 );
623 }
624
625 #[test]
626 fn records_by_puzzle_hash_and_parent_map_states() {
627 let fetcher = MockFetcher {
628 puzzle_states: vec![CoinState {
629 coin: coin(1),
630 created_height: Some(1),
631 spent_height: None,
632 }],
633 children: vec![CoinState {
634 coin: coin(2),
635 created_height: Some(2),
636 spent_height: None,
637 }],
638 ..Default::default()
639 };
640 let (_rt, provider) = provider_with(fetcher);
641 let ph = Bytes32::new([8; 32]);
642 let parent = Bytes32::new([9; 32]);
643 let p = provider.clone();
644 assert_eq!(
645 call(move || p.coin_records_by_puzzle_hash(ph, true))
646 .unwrap()
647 .len(),
648 1
649 );
650 assert_eq!(
651 call(move || provider.coin_records_by_parent(parent))
652 .unwrap()
653 .len(),
654 1
655 );
656 }
657
658 #[test]
659 fn lineage_and_timestamp_are_unsupported_not_false_absence() {
660 let (_rt, provider) = provider_with(MockFetcher::default());
661 let p = provider.clone();
662 assert!(matches!(
663 call(move || p.resolve_singleton_lineage(Bytes32::new([1; 32]))),
664 Err(ChainSourceError::Unsupported(_))
665 ));
666 assert!(matches!(
667 call(move || provider.block_timestamp(1)),
668 Err(ChainSourceError::Unsupported(_))
669 ));
670 }
671
672 #[test]
673 fn provider_info_is_reported() {
674 let (_rt, provider) = provider_with(MockFetcher::default());
675 assert_eq!(provider.provider_info().priority, 20);
676 assert_eq!(provider.peak_height().unwrap(), None);
677 }
678}