cht proof checker, use it in on_demand

ethcore/light/src/cht.rs

@@ -23,7 +23,7 @@
 use ethcore::ids::BlockId;
 use util::{Bytes, H256, U256, HashDB, MemoryDB};
 use util::trie::{self, TrieMut, TrieDBMut, Trie, TrieDB, Recorder};
-use rlp::{Stream, RlpStream};
+use rlp::{Stream, RlpStream, UntrustedRlp, View};
 
 // encode a key.
 macro_rules! key {
@@ -136,6 +136,30 @@ pub fn compute_root<I>(cht_num: u64, iterable: I) -> Option<H256>
 	}
 }
 
+/// Check a proof for a CHT.
+/// Given a set of a trie nodes, a number to query, and a trie root,
+/// verify the given trie branch and extract the canonical hash and total difficulty.
+// TODO: better support for partially-checked queries.
+pub fn check_proof(proof: &[Bytes], num: u64, root: H256) -> Option<(H256, U256)> {
+	let mut db = MemoryDB::new();
+
+	for node in proof { db.insert(&node[..]); }
+	let res = match TrieDB::new(&db, &root) {
+		Err(_) => return None,
+		Ok(trie) => trie.get_with(&key!(num), |val: &[u8]| {
+			let rlp = UntrustedRlp::new(val);
+			rlp.val_at::<H256>(0)
+				.and_then(|h| rlp.val_at::<U256>(1).map(|td| (h, td)))
+				.ok()
+		})
+	};
+
+	match res {
+		Ok(Some(Some((hash, td)))) => Some((hash, td)),
+		_ => None,
+	}
+}
+
 /// Convert a block number to a CHT number.
 /// Returns `None` for `block_num` == 0, `Some` otherwise.
 pub fn block_to_cht_number(block_num: u64) -> Option<u64> {

ethcore/light/src/on_demand/mod.rs

@@ -29,7 +29,7 @@ use futures::sync::oneshot;
 use network::PeerId;
 
 use net::{Handler, Status, Capabilities, Announcement, EventContext, BasicContext, ReqId};
-use util::{Bytes, RwLock};
+use util::{Bytes, RwLock, U256};
 use types::les_request::{self as les_request, Request as LesRequest};
 
 pub mod request;
@@ -79,7 +79,7 @@ struct Peer {
 
 // Attempted request info and sender to put received value.
 enum Pending {
-	HeaderByNumber(request::HeaderByNumber, Sender<encoded::Header>), // num + CHT root
+	HeaderByNumber(request::HeaderByNumber, Sender<(encoded::Header, U256)>), // num + CHT root
 	HeaderByHash(request::HeaderByHash, Sender<encoded::Header>),
 	Block(request::Body, Sender<encoded::Block>),
 	BlockReceipts(request::BlockReceipts, Sender<Vec<Receipt>>),
@@ -105,14 +105,15 @@ impl Default for OnDemand {
 
 impl OnDemand {
 	/// Request a header by block number and CHT root hash.
-	pub fn header_by_number(&self, ctx: &BasicContext, req: request::HeaderByNumber) -> Response<encoded::Header> {
+	/// Returns the header and the total difficulty.
+	pub fn header_by_number(&self, ctx: &BasicContext, req: request::HeaderByNumber) -> Response<(encoded::Header, U256)> {
 		let (sender, receiver) = oneshot::channel();
 		self.dispatch_header_by_number(ctx, req, sender);
 		Response(receiver)
 	}
 
 	// dispatch the request, completing the request if no peers available.
-	fn dispatch_header_by_number(&self, ctx: &BasicContext, req: request::HeaderByNumber, sender: Sender<encoded::Header>) {
+	fn dispatch_header_by_number(&self, ctx: &BasicContext, req: request::HeaderByNumber, sender: Sender<(encoded::Header, U256)>) {
 		let num = req.num;
 		let cht_num = match ::cht::block_to_cht_number(req.num) {
 			Some(cht_num) => cht_num,

ethcore/light/src/on_demand/request.rs

@@ -21,7 +21,7 @@ use ethcore::encoded;
 use ethcore::receipt::Receipt;
 
 use rlp::{RlpStream, Stream, UntrustedRlp, View};
-use util::{Address, Bytes, HashDB, H256};
+use util::{Address, Bytes, HashDB, H256, U256};
 use util::memorydb::MemoryDB;
 use util::sha3::Hashable;
 use util::trie::{Trie, TrieDB, TrieError};
@@ -66,24 +66,16 @@ pub struct HeaderByNumber {
 
 impl HeaderByNumber {
 	/// Check a response with a header and cht proof.
-	pub fn check_response(&self, header: &[u8], proof: &[Bytes]) -> Result<encoded::Header, Error> {
-		use util::trie::{Trie, TrieDB};
-
-		// check the proof
-		let mut db = MemoryDB::new();
-
-		for node in proof { db.insert(&node[..]); }
-		let key = ::rlp::encode(&self.num);
-
-		let expected_hash: H256 = match TrieDB::new(&db, &self.cht_root).and_then(|t| t.get(&*key))? {
-			Some(val) => ::rlp::decode(&val),
-			None => return Err(Error::BadProof)
+	pub fn check_response(&self, header: &[u8], proof: &[Bytes]) -> Result<(encoded::Header, U256), Error> {
+		let (expected_hash, td) = match ::cht::check_proof(proof, self.num, self.cht_root) {
+			Some((expected_hash, td)) => (expected_hash, td),
+			None => return Err(Error::BadProof),
 		};
 
 		// and compare the hash to the found header.
 		let found_hash = header.sha3();
 		match expected_hash == found_hash {
-			true => Ok(encoded::Header::new(header.to_vec())),
+			true => Ok((encoded::Header::new(header.to_vec()), td)),
 			false => Err(Error::WrongHash(expected_hash, found_hash)),
 		}
 	}
@@ -191,51 +183,44 @@ impl Account {
 mod tests {
 	use super::*;
 	use util::{MemoryDB, Address, H256, FixedHash};
-	use util::trie::{Trie, TrieMut, TrieDB, SecTrieDB, TrieDBMut, SecTrieDBMut};
+	use util::trie::{Trie, TrieMut, SecTrieDB, SecTrieDBMut};
 	use util::trie::recorder::Recorder;
 
+	use ethcore::client::{BlockChainClient, TestBlockChainClient, EachBlockWith};
 	use ethcore::header::Header;
 	use ethcore::encoded;
 	use ethcore::receipt::Receipt;
 
 	#[test]
 	fn check_header_by_number() {
-		let mut root = H256::default();
-		let mut db = MemoryDB::new();
-		let mut header = Header::new();
-		header.set_number(10_000);
-		header.set_extra_data(b"test_header".to_vec());
+		use ::cht;
 
-		{
-			let mut trie = TrieDBMut::new(&mut db, &mut root);
-			for i in (0..2048u64).map(|x| x + 8192) {
-				let hash = if i == 10_000 {
-					header.hash()
-				} else {
-					H256::random()
-				};
-				trie.insert(&*::rlp::encode(&i), &*::rlp::encode(&hash)).unwrap();
-			}
-		}
+		let test_client = TestBlockChainClient::new();
+		test_client.add_blocks(10500, EachBlockWith::Nothing);
 
-		let proof = {
-			let trie = TrieDB::new(&db, &root).unwrap();
-			let key = ::rlp::encode(&10_000u64);
-			let mut recorder = Recorder::new();
+		let cht = {
+			let fetcher = |id| {
+				let hdr = test_client.block_header(id).unwrap();
+				let td = test_client.block_total_difficulty(id).unwrap();
+				Some(cht::BlockInfo {
+					hash: hdr.hash(),
+					parent_hash: hdr.parent_hash(),
+					total_difficulty: td,
+				})
+			};
 
-			trie.get_with(&*key, &mut recorder).unwrap().unwrap();
-
-			recorder.drain().into_iter().map(|r| r.data).collect::<Vec<_>>()
+			cht::build(cht::block_to_cht_number(10_000).unwrap(), fetcher).unwrap()
 		};
 
+		let proof = cht.prove(10_000, 0).unwrap().unwrap();
 		let req = HeaderByNumber {
 			num: 10_000,
-			cht_root: root,
+			cht_root: cht.root(),
 		};
 
-		let raw_header = ::rlp::encode(&header);
+		let raw_header = test_client.block_header(::ethcore::ids::BlockId::Number(10_000)).unwrap();
 
-		assert!(req.check_response(&*raw_header, &proof[..]).is_ok());
+		assert!(req.check_response(&raw_header.into_inner(), &proof[..]).is_ok());
 	}
 
 	#[test]