#![cfg(all(feature = "lineage-walk", feature = "testing"))]
use std::cell::{Cell, RefCell};
use std::collections::HashMap;
use std::time::{Duration, Instant};
use anyhow::Result;
use chia_protocol::{Bytes32, Coin, CoinSpend, Program};
use chia_puzzle_types::singleton::{SingletonArgs, SingletonSolution};
use chia_puzzle_types::{EveProof, Proof};
use chia_sdk_driver::SpendContext;
use clvm_utils::{tree_hash, tree_hash_from_bytes};
use clvmr::serde::{node_from_bytes, node_from_bytes_backrefs, node_to_bytes_backrefs};
use clvmr::{Allocator, NodePtr};
use dig_chainsource_interface::{
walk_singleton_lineage_within, ChainSource, ChainSourceError, CoinRecord, LineageWalkError,
SingletonLineage, WalkBounds, MAX_REVEAL_EXPANDED_BYTES,
};
const CONS: i64 = 4;
const QUOTE: i64 = 1;
const CREATE_COIN: i64 = 51;
const FIRST_SOLUTION_ARG: i64 = 2;
const BACKREF_TAG: u8 = 0xFE;
fn compress(ctx: &mut SpendContext, program: &Program) -> Result<Program> {
let node = node_from_bytes(ctx, program.as_ref())?;
Ok(Program::from(node_to_bytes_backrefs(ctx, node)?))
}
fn endless_inner_puzzle(ctx: &mut SpendContext) -> Result<NodePtr> {
let quoted_amount = ctx.alloc(&(QUOTE, QUOTE))?;
let quoted_opcode = ctx.alloc(&(QUOTE, CREATE_COIN))?;
let cons = ctx.alloc(&CONS)?;
let target = ctx.alloc(&FIRST_SOLUTION_ARG)?;
let amount_tail = ctx.alloc(&vec![cons, quoted_amount, NodePtr::NIL])?;
let arguments = ctx.alloc(&vec![cons, target, amount_tail])?;
let condition = ctx.alloc(&vec![cons, quoted_opcode, arguments])?;
Ok(ctx.alloc(&vec![cons, condition, NodePtr::NIL])?)
}
struct HostileSource {
launcher: Coin,
outer_puzzle_hash: Bytes32,
launcher_reveal: Program,
launcher_solution: Program,
singleton_reveal: Program,
singleton_solution: Program,
known: RefCell<HashMap<Bytes32, Coin>>,
frontier: RefCell<Coin>,
reads: Cell<usize>,
}
impl HostileSource {
fn new(ctx: &mut SpendContext) -> Result<Self> {
let launcher = Coin::new(
Bytes32::new([0xA1; 32]),
Bytes32::new(chia_puzzles::SINGLETON_LAUNCHER_HASH),
1,
);
let launcher_id = launcher.coin_id();
let inner = endless_inner_puzzle(ctx)?;
let inner_puzzle_hash = Bytes32::from(tree_hash(ctx, inner));
let singleton = ctx.curry(SingletonArgs::new(launcher_id, inner))?;
let outer_puzzle_hash = Bytes32::from(tree_hash(ctx, singleton));
let inner_solution = ctx.alloc(&vec![inner_puzzle_hash])?;
let singleton_solution = ctx.alloc(&SingletonSolution {
lineage_proof: Proof::Eve(EveProof {
parent_parent_coin_info: launcher.parent_coin_info,
parent_amount: 1,
}),
amount: 1,
inner_solution,
})?;
let launcher_reveal =
ctx.alloc(&Program::from(chia_puzzles::SINGLETON_LAUNCHER.to_vec()))?;
let launcher_solution =
ctx.alloc(&(outer_puzzle_hash, (1, (Vec::<Bytes32>::new(), ()))))?;
let eve = Coin::new(launcher_id, outer_puzzle_hash, 1);
Ok(Self {
launcher,
outer_puzzle_hash,
launcher_reveal: ctx.serialize(&launcher_reveal)?,
launcher_solution: ctx.serialize(&launcher_solution)?,
singleton_reveal: ctx.serialize(&singleton)?,
singleton_solution: ctx.serialize(&singleton_solution)?,
known: RefCell::new(HashMap::from([(eve.coin_id(), eve)])),
frontier: RefCell::new(eve),
reads: Cell::new(0),
})
}
fn with_backref_compressed_reveals(mut self, ctx: &mut SpendContext) -> Result<Self> {
self.singleton_reveal = compress(ctx, &self.singleton_reveal)?;
assert!(
self.singleton_reveal.as_ref().contains(&BACKREF_TAG),
"the control is only load-bearing if the reveal really carries a back-reference"
);
self.launcher_reveal = compress(ctx, &self.launcher_reveal)?;
Ok(self)
}
fn launcher_id(&self) -> Bytes32 {
self.launcher.coin_id()
}
fn coin(&self, coin_id: Bytes32) -> Option<Coin> {
for _ in 0..2 {
if let Some(coin) = self.known.borrow().get(&coin_id) {
return Some(*coin);
}
let mut frontier = self.frontier.borrow_mut();
let next = Coin::new(frontier.coin_id(), self.outer_puzzle_hash, 1);
*frontier = next;
self.known.borrow_mut().insert(next.coin_id(), next);
}
None
}
fn record(&self, coin: Coin) -> CoinRecord {
CoinRecord {
coin,
confirmed_height: Some(1),
spent_height: Some(2),
timestamp: None,
coinbase: false,
}
}
}
impl ChainSource for HostileSource {
type Error = ChainSourceError;
fn coin_record(&self, coin_id: Bytes32) -> Result<Option<CoinRecord>, Self::Error> {
self.reads.set(self.reads.get() + 1);
if coin_id == self.launcher_id() {
return Ok(Some(self.record(self.launcher)));
}
Ok(self.coin(coin_id).map(|coin| self.record(coin)))
}
fn coin_records_by_puzzle_hash(
&self,
_puzzle_hash: Bytes32,
_include_spent: bool,
) -> Result<Vec<CoinRecord>, Self::Error> {
Ok(Vec::new())
}
fn coin_records_by_parent(
&self,
_parent_coin_id: Bytes32,
) -> Result<Vec<CoinRecord>, Self::Error> {
Ok(Vec::new())
}
fn coin_spend(&self, coin_id: Bytes32) -> Result<Option<CoinSpend>, Self::Error> {
self.reads.set(self.reads.get() + 1);
if coin_id == self.launcher_id() {
return Ok(Some(CoinSpend::new(
self.launcher,
self.launcher_reveal.clone(),
self.launcher_solution.clone(),
)));
}
Ok(self.coin(coin_id).map(|coin| {
CoinSpend::new(
coin,
self.singleton_reveal.clone(),
self.singleton_solution.clone(),
)
}))
}
fn resolve_singleton_lineage(
&self,
launcher_id: Bytes32,
) -> Result<Option<SingletonLineage>, Self::Error> {
dig_chainsource_interface::resolve_singleton_lineage_via_walk(self, launcher_id)
}
fn peak_height(&self) -> Result<Option<u32>, Self::Error> {
Ok(None)
}
fn block_timestamp(&self, _height: u32) -> Result<Option<u64>, Self::Error> {
Ok(None)
}
}
#[test]
fn the_hostile_chain_is_genuinely_endless_and_genuinely_well_formed() -> Result<()> {
let ctx = &mut SpendContext::new();
let source = HostileSource::new(ctx)?;
let error = walk_singleton_lineage_within(&source, source.launcher_id(), WalkBounds::hops(64))
.expect_err("an endless chain never reaches a tip");
assert_eq!(error, LineageWalkError::TooDeep { limit: 64 });
Ok(())
}
#[test]
fn an_endless_chain_is_refused_on_the_wall_clock_budget() -> Result<()> {
let ctx = &mut SpendContext::new();
let source = HostileSource::new(ctx)?;
let budget = Duration::from_millis(50);
let started = Instant::now();
let error = walk_singleton_lineage_within(
&source,
source.launcher_id(),
WalkBounds::default().within(budget),
)
.expect_err("an endless chain must not resolve");
let elapsed = started.elapsed();
assert_eq!(error, LineageWalkError::DeadlineExceeded { budget });
assert!(
elapsed < Duration::from_secs(10),
"the budget must actually stop the walk; it ran for {elapsed:?}"
);
assert_eq!(
ChainSourceError::from(error),
ChainSourceError::Timeout,
"running out of time is a timeout, distinct from malformed chain data and from too-deep"
);
Ok(())
}
#[test]
fn a_short_honest_walk_finishes_well_inside_its_budget() -> Result<()> {
let ctx = &mut SpendContext::new();
let source = HostileSource::new(ctx)?;
let error = walk_singleton_lineage_within(
&source,
source.launcher_id(),
WalkBounds::hops(4).within(Duration::from_secs(30)),
)
.expect_err("four hops of an endless chain still exhaust the cap");
assert_eq!(
error,
LineageWalkError::TooDeep { limit: 4 },
"a generous budget must leave the hop cap in charge"
);
Ok(())
}
#[test]
#[ignore = "a multi-second DoS measurement, not a correctness gate"]
fn measure_the_cost_of_the_default_hop_bound() -> Result<()> {
let ctx = &mut SpendContext::new();
let source = HostileSource::new(ctx)?;
let started = Instant::now();
let outcome = walk_singleton_lineage_within(
&source,
source.launcher_id(),
WalkBounds::default().within(Duration::from_secs(3600)),
);
println!(
"hops={} elapsed={:.2}s reads={} peak_working_set={:.1} MB outcome={:?}",
dig_chainsource_interface::MAX_LINEAGE_DEPTH,
started.elapsed().as_secs_f64(),
source.reads.get(),
peak_working_set_bytes() as f64 / (1024.0 * 1024.0),
outcome
);
Ok(())
}
fn backref_bomb(depth: u32) -> Program {
let mut allocator = Allocator::new();
let mut node = allocator
.new_atom(&[1])
.expect("a one-byte atom always fits");
for _ in 0..depth {
node = allocator
.new_pair(node, node)
.expect("a self-cons adds one pair");
}
let bytes = node_to_bytes_backrefs(&allocator, node).expect("the DAG serializes");
Program::from(bytes)
}
struct BombSource {
launcher: Coin,
bomb: Program,
}
impl BombSource {
fn at_depth(depth: u32) -> Self {
Self {
launcher: Coin::new(
Bytes32::new([0xB0; 32]),
Bytes32::new(chia_puzzles::SINGLETON_LAUNCHER_HASH),
1,
),
bomb: backref_bomb(depth),
}
}
}
impl ChainSource for BombSource {
type Error = ChainSourceError;
fn coin_record(&self, coin_id: Bytes32) -> Result<Option<CoinRecord>, Self::Error> {
if coin_id != self.launcher.coin_id() {
return Ok(None);
}
Ok(Some(CoinRecord {
coin: self.launcher,
confirmed_height: Some(1),
spent_height: Some(2),
timestamp: None,
coinbase: false,
}))
}
fn coin_records_by_puzzle_hash(
&self,
_puzzle_hash: Bytes32,
_include_spent: bool,
) -> Result<Vec<CoinRecord>, Self::Error> {
Ok(Vec::new())
}
fn coin_records_by_parent(
&self,
_parent_coin_id: Bytes32,
) -> Result<Vec<CoinRecord>, Self::Error> {
Ok(Vec::new())
}
fn coin_spend(&self, coin_id: Bytes32) -> Result<Option<CoinSpend>, Self::Error> {
if coin_id != self.launcher.coin_id() {
return Ok(None);
}
Ok(Some(CoinSpend::new(
self.launcher,
self.bomb.clone(),
Program::from(vec![0x80]),
)))
}
fn resolve_singleton_lineage(
&self,
launcher_id: Bytes32,
) -> Result<Option<SingletonLineage>, Self::Error> {
dig_chainsource_interface::resolve_singleton_lineage_via_walk(self, launcher_id)
}
fn peak_height(&self) -> Result<Option<u32>, Self::Error> {
Ok(None)
}
fn block_timestamp(&self, _height: u32) -> Result<Option<u64>, Self::Error> {
Ok(None)
}
}
#[test]
fn a_backref_decompression_bomb_in_a_puzzle_reveal_is_refused_in_bounded_time() -> Result<()> {
let source = BombSource::at_depth(30);
let started = Instant::now();
let error = walk_singleton_lineage_within(
&source,
source.launcher.coin_id(),
WalkBounds::hops(4).within(Duration::from_secs(600)),
)
.expect_err("a reveal that expands beyond the bound is refused");
let elapsed = started.elapsed();
assert_eq!(
error,
LineageWalkError::RevealTooLarge {
coin_id: source.launcher.coin_id(),
limit: MAX_REVEAL_EXPANDED_BYTES,
},
"the bomb is refused for its EXPANSION, not as malformed chain data: {error:?}"
);
assert!(
elapsed < Duration::from_secs(5),
"a {}-byte reveal must be bounded before it is hashed; it took {elapsed:?}",
source.bomb.as_ref().len()
);
Ok(())
}
struct EveInnerBombSource {
launcher: Coin,
launcher_solution: Program,
eve: Coin,
eve_reveal: Program,
records: HashMap<Bytes32, CoinRecord>,
}
impl EveInnerBombSource {
fn at_depth(depth: u32) -> Result<Self> {
let ctx = &mut SpendContext::new();
let launcher = Coin::new(
Bytes32::new([0xC0; 32]),
Bytes32::new(chia_puzzles::SINGLETON_LAUNCHER_HASH),
1,
);
let launcher_id = launcher.coin_id();
let bomb = backref_bomb(depth);
let inner_hash = tree_hash_from_bytes(bomb.as_ref()).expect("the bomb decodes");
let inner = node_from_bytes_backrefs(ctx, bomb.as_ref())?;
let outer = ctx.curry(SingletonArgs::new(launcher_id, inner))?;
let eve_reveal = Program::from(node_to_bytes_backrefs(ctx, outer)?);
let eve_puzzle_hash =
Bytes32::from(SingletonArgs::curry_tree_hash(launcher_id, inner_hash));
let eve = Coin::new(launcher_id, eve_puzzle_hash, 1);
let solution = ctx.alloc(&(eve_puzzle_hash, (1u64, (Vec::<Bytes32>::new(), ()))))?;
let spent_at = |coin: Coin, height: u32| CoinRecord {
coin,
confirmed_height: Some(1),
spent_height: Some(height),
timestamp: None,
coinbase: false,
};
Ok(Self {
launcher,
launcher_solution: ctx.serialize(&solution)?,
eve,
eve_reveal,
records: HashMap::from([
(launcher_id, spent_at(launcher, 2)),
(eve.coin_id(), spent_at(eve, 3)),
]),
})
}
}
impl ChainSource for EveInnerBombSource {
type Error = ChainSourceError;
fn coin_record(&self, coin_id: Bytes32) -> Result<Option<CoinRecord>, Self::Error> {
Ok(self.records.get(&coin_id).cloned())
}
fn coin_records_by_puzzle_hash(
&self,
_puzzle_hash: Bytes32,
_include_spent: bool,
) -> Result<Vec<CoinRecord>, Self::Error> {
Ok(Vec::new())
}
fn coin_records_by_parent(
&self,
_parent_coin_id: Bytes32,
) -> Result<Vec<CoinRecord>, Self::Error> {
Ok(Vec::new())
}
fn coin_spend(&self, coin_id: Bytes32) -> Result<Option<CoinSpend>, Self::Error> {
if coin_id == self.launcher.coin_id() {
return Ok(Some(CoinSpend::new(
self.launcher,
Program::from(chia_puzzles::SINGLETON_LAUNCHER.to_vec()),
self.launcher_solution.clone(),
)));
}
if coin_id == self.eve.coin_id() {
return Ok(Some(CoinSpend::new(
self.eve,
self.eve_reveal.clone(),
Program::from(vec![0x80]),
)));
}
Ok(None)
}
fn resolve_singleton_lineage(
&self,
launcher_id: Bytes32,
) -> Result<Option<SingletonLineage>, Self::Error> {
dig_chainsource_interface::resolve_singleton_lineage_via_walk(self, launcher_id)
}
fn peak_height(&self) -> Result<Option<u32>, Self::Error> {
Ok(None)
}
fn block_timestamp(&self, _height: u32) -> Result<Option<u64>, Self::Error> {
Ok(None)
}
}
#[test]
fn a_bomb_curried_inside_a_genuine_singleton_reveal_is_refused_before_it_is_parsed() -> Result<()> {
let source = EveInnerBombSource::at_depth(30)?;
let started = Instant::now();
let error = walk_singleton_lineage_within(
&source,
source.launcher.coin_id(),
WalkBounds::hops(4).within(Duration::from_secs(600)),
)
.expect_err("a reveal that expands beyond the bound is refused");
let elapsed = started.elapsed();
assert_eq!(
error,
LineageWalkError::RevealTooLarge {
coin_id: source.eve.coin_id(),
limit: MAX_REVEAL_EXPANDED_BYTES,
},
"the inner bomb is refused for its EXPANSION, at the eve hop: {error:?}"
);
assert!(
elapsed < Duration::from_secs(5),
"a {}-byte reveal must be bounded before it is parsed; it took {elapsed:?}",
source.eve_reveal.as_ref().len()
);
Ok(())
}
#[test]
fn a_genuinely_backref_compressed_reveal_is_still_accepted() -> Result<()> {
let ctx = &mut SpendContext::new();
let source = HostileSource::new(ctx)?.with_backref_compressed_reveals(ctx)?;
let error = walk_singleton_lineage_within(
&source,
source.launcher_id(),
WalkBounds::hops(4).within(Duration::from_secs(60)),
)
.expect_err("four hops of an endless chain still exhaust the cap");
assert_eq!(
error,
LineageWalkError::TooDeep { limit: 4 },
"a back-reference-compressed reveal must be accepted, not reported as malformed"
);
Ok(())
}
#[cfg(windows)]
fn peak_working_set_bytes() -> u64 {
#[repr(C)]
#[derive(Default)]
struct ProcessMemoryCounters {
cb: u32,
page_fault_count: u32,
peak_working_set_size: usize,
working_set_size: usize,
quota_peak_paged_pool_usage: usize,
quota_paged_pool_usage: usize,
quota_peak_non_paged_pool_usage: usize,
quota_non_paged_pool_usage: usize,
pagefile_usage: usize,
peak_pagefile_usage: usize,
}
extern "system" {
fn GetCurrentProcess() -> isize;
fn K32GetProcessMemoryInfo(
process: isize,
counters: *mut ProcessMemoryCounters,
size: u32,
) -> i32;
}
let mut counters = ProcessMemoryCounters {
cb: std::mem::size_of::<ProcessMemoryCounters>() as u32,
..Default::default()
};
unsafe {
K32GetProcessMemoryInfo(GetCurrentProcess(), &mut counters, counters.cb);
}
counters.peak_working_set_size as u64
}
#[cfg(not(windows))]
fn peak_working_set_bytes() -> u64 {
0
}