From 30d1b7f77c6863177efa5898a76ed7350a02d485 Mon Sep 17 00:00:00 2001
From: Neil Booth <kyuupichan@gmail.com>
Date: Mon, 17 Oct 2016 23:58:06 +0900
Subject: [PATCH] Add proper block chaining check

---
 lib/coins.py         | 13 ++++++++++---
 server/controller.py |  2 +-
 server/db.py         | 24 +++++++++++++++++-------
 3 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/lib/coins.py b/lib/coins.py
index b04bec2..30a0b12 100644
--- a/lib/coins.py
+++ b/lib/coins.py
@@ -6,7 +6,7 @@ from decimal import Decimal
 import inspect
 import sys
 
-from lib.hash import Base58, hash160
+from lib.hash import Base58, hash160, double_sha256
 from lib.script import ScriptPubKey
 from lib.tx import Deserializer
 
@@ -135,10 +135,17 @@ class Coin(object):
             payload.append(0x01)
         return Base58.encode_check(payload)
 
+    @classmethod
+    def header_hashes(cls, header):
+        '''Given a header return the previous block hash and the current block
+        hash.'''
+        return header[4:36], double_sha256(header)
+
     @classmethod
     def read_block(cls, block):
-        d = Deserializer(block[cls.HEADER_LEN:])
-        return d.read_block()
+        '''Read a block and return (header, tx_hashes, txs)'''
+        header, rest = block[:cls.HEADER_LEN], block[cls.HEADER_LEN:]
+        return (header, ) + Deserializer(rest).read_block()
 
     @classmethod
     def decimal_value(cls, value):
diff --git a/server/controller.py b/server/controller.py
index 9925335..ca013da 100644
--- a/server/controller.py
+++ b/server/controller.py
@@ -260,7 +260,7 @@ class BlockCache(LoggedClass):
         data = json.dumps(payload)
         while True:
             try:
-                async with aiohttp.post(self.daemon_url, data = data) as resp:
+                async with aiohttp.post(self.daemon_url, data=data) as resp:
                     result = await resp.json()
             except asyncio.CancelledError:
                 raise
diff --git a/server/db.py b/server/db.py
index 26f044a..4f6b6cf 100644
--- a/server/db.py
+++ b/server/db.py
@@ -310,17 +310,17 @@ class FSCache(LoggedClass):
             raise
 
     def process_block(self, block):
-        '''Process a new block.'''
+        '''Process a new block and return (header, tx_hashes, txs)'''
         assert len(self.tx_counts) == self.height + 1 + len(self.headers)
 
-        tx_hashes, txs = self.coin.read_block(block)
+        triple = header, tx_hashes, txs = self.coin.read_block(block)
 
         # Cache the new header, tx hashes and cumulative tx count
-        self.headers.append(block[:self.coin.HEADER_LEN])
+        self.headers.append(header)
         self.tx_hashes.append(tx_hashes)
         self.tx_counts.append(self.tx_count + len(txs))
 
-        return tx_hashes, txs
+        return triple
 
     def flush(self, new_height, new_tx_count):
         '''Flush the things stored on the filesystem.'''
@@ -431,6 +431,9 @@ class DB(LoggedClass):
     class Error(Exception):
         pass
 
+    class ChainError(Exception):
+        pass
+
     def __init__(self, env):
         super().__init__()
 
@@ -447,7 +450,7 @@ class DB(LoggedClass):
         self.flush_count = 0
         self.utxo_flush_count = 0
         self.wall_time = 0
-        self.tip = self.coin.GENESIS_HASH
+        self.tip = b'\0' * 32
 
         # Open DB and metadata files.  Record some of its state.
         self.db = self.open_db(self.coin)
@@ -654,8 +657,15 @@ class DB(LoggedClass):
     def process_block(self, block, daemon_height):
         # We must update the fs_cache before calling process_tx() as
         # it uses the fs_cache for tx hash lookup
-        tx_hashes, txs = self.fs_cache.process_block(block)
-
+        header, tx_hashes, txs = self.fs_cache.process_block(block)
+        prev_hash, header_hash = self.coin.header_hashes(header)
+        if prev_hash != self.tip:
+            raise self.ChainError('trying to build header with prev_hash {} '
+                                  'on top of tip with hash {}'
+                                  .format(hash_to_str(prev_hash),
+                                          hash_to_str(self.tip)))
+
+        self.tip = header_hash
         self.height += 1
         for tx_hash, tx in zip(tx_hashes, txs):
             self.process_tx(tx_hash, tx)