Add segwit block (de)serialization. (#9)

Add witness merkle tree;
Add parameters to stream_(de)serialize and (de)serialize methods;
Check witness merkle tree.
CheckBlock:
 - check MAX_BLOCK_WEIGHT;
 - check witness witnessScript;
 - check witness commitment.

Author: the9ull

bitcoin/core/__init__.py

@@ -16,14 +16,21 @@
 import sys
 import time
 
-from .script import CScript, CScriptWitness
+from .script import CScript, CScriptWitness, CScriptOp, OP_RETURN
 
 from .serialize import *
 
+if sys.version > '3':
+    _bytes = bytes
+else:
+    _bytes = lambda x: bytes(bytearray(x))
+
 # Core definitions
 COIN = 100000000
 MAX_BLOCK_SIZE = 1000000
+MAX_BLOCK_WEIGHT = 4000000
 MAX_BLOCK_SIGOPS = MAX_BLOCK_SIZE/50
+WITNESS_COINBASE_SCRIPTPUBKEY_MAGIC = _bytes([OP_RETURN, 0x24, 0xaa, 0x21, 0xa9, 0xed])
 
 def MoneyRange(nValue, params=None):
     global coreparams
@@ -426,9 +433,9 @@ def stream_deserialize(cls, f):
             return cls(vin, vout, nLockTime, nVersion)
 
 
-    def stream_serialize(self, f):
+    def stream_serialize(self, f, include_witness=True):
         f.write(struct.pack(b"<i", self.nVersion))
-        if not self.wit.is_null():
+        if include_witness and not self.wit.is_null():
             assert(len(self.wit.vtxinwit) <= len(self.vin))
             f.write(b'\x00') # Marker
             f.write(b'\x01') # Flag
@@ -443,6 +450,10 @@ def stream_serialize(self, f):
     def is_coinbase(self):
         return len(self.vin) == 1 and self.vin[0].prevout.is_null()
 
+    def has_witness(self):
+        """True if witness"""
+        return not self.wit.is_null()
+
     def __repr__(self):
         return "CTransaction(%r, %r, %i, %i, %r)" % (self.vin, self.vout,
                 self.nLockTime, self.nVersion, self.wit)
@@ -561,9 +572,12 @@ def __repr__(self):
                 (self.__class__.__name__, self.nVersion, b2lx(self.hashPrevBlock), b2lx(self.hashMerkleRoot),
                  self.nTime, self.nBits, self.nNonce)
 
+class NoWitnessData(Exception):
+    """The block does not have witness data"""
+
 class CBlock(CBlockHeader):
     """A block including all transactions in it"""
-    __slots__ = ['vtx', 'vMerkleTree']
+    __slots__ = ['vtx', 'vMerkleTree', 'vWitnessMerkleTree']
 
     @staticmethod
     def build_merkle_tree_from_txids(txids):
@@ -612,12 +626,60 @@ def calc_merkle_root(self):
             raise ValueError('Block contains no transactions')
         return self.build_merkle_tree_from_txs(self.vtx)[-1]
 
+    @staticmethod
+    def build_witness_merkle_tree_from_txs(txs):
+        """Calculate the witness merkle tree from transactions"""
+        has_witness = False
+        hashes = []
+        for tx in txs:
+            hashes.append(tx.GetHash())
+            has_witness |= tx.has_witness()
+        if not has_witness:
+            raise NoWitnessData
+        hashes[0] = b'\x00' * 32
+        return CBlock.build_merkle_tree_from_txids(hashes)
+
+    def calc_witness_merkle_root(self):
+        """Calculate the witness merkle root
+
+        The calculated merkle root is not cached; every invocation
+        re-calculates it from scratch.
+        """
+        if not len(self.vtx):
+            raise ValueError('Block contains no transactions')
+        return self.build_witness_merkle_tree_from_txs(self.vtx)[-1]
+
+    def get_witness_commitment_index(self):
+        """Return None or an index"""
+        if not len(self.vtx):
+            raise ValueError('Block contains no transactions')
+        commit_pos = None
+        for index, out in enumerate(self.vtx[0].vout):
+            script = out.scriptPubKey
+            if len(script) >= 38 and script[:6] == WITNESS_COINBASE_SCRIPTPUBKEY_MAGIC:
+                commit_pos = index
+        if commit_pos is None:
+            raise ValueError('The witness commitment is missed')
+        return commit_pos
+
     def __init__(self, nVersion=2, hashPrevBlock=b'\x00'*32, hashMerkleRoot=b'\x00'*32, nTime=0, nBits=0, nNonce=0, vtx=()):
         """Create a new block"""
+        if vtx:
+            vMerkleTree = tuple(CBlock.build_merkle_tree_from_txs(vtx))
+            if hashMerkleRoot == b'\x00'*32:
+                hashMerkleRoot = vMerkleTree[-1]
+            elif hashMerkleRoot != vMerkleTree[-1]:
+                raise CheckBlockError("CBlock : hashMerkleRoot is not compatible with vtx")
+        else:
+            vMerkleTree = ()
         super(CBlock, self).__init__(nVersion, hashPrevBlock, hashMerkleRoot, nTime, nBits, nNonce)
 
-        vMerkleTree = tuple(CBlock.build_merkle_tree_from_txs(vtx))
         object.__setattr__(self, 'vMerkleTree', vMerkleTree)
+        try:
+            vWitnessMerkleTree = tuple(CBlock.build_witness_merkle_tree_from_txs(vtx))
+        except NoWitnessData:
+            vWitnessMerkleTree = ()
+        object.__setattr__(self, 'vWitnessMerkleTree', vWitnessMerkleTree)
         object.__setattr__(self, 'vtx', tuple(CTransaction.from_tx(tx) for tx in vtx))
 
     @classmethod
@@ -627,13 +689,18 @@ def stream_deserialize(cls, f):
         vtx = VectorSerializer.stream_deserialize(CTransaction, f)
         vMerkleTree = tuple(CBlock.build_merkle_tree_from_txs(vtx))
         object.__setattr__(self, 'vMerkleTree', vMerkleTree)
+        try:
+            vWitnessMerkleTree = tuple(CBlock.build_witness_merkle_tree_from_txs(vtx))
+        except NoWitnessData:
+            vWitnessMerkleTree = ()
+        object.__setattr__(self, 'vWitnessMerkleTree', vWitnessMerkleTree)
         object.__setattr__(self, 'vtx', tuple(vtx))
 
         return self
 
-    def stream_serialize(self, f):
+    def stream_serialize(self, f, include_witness=True):
         super(CBlock, self).stream_serialize(f)
-        VectorSerializer.stream_serialize(CTransaction, self.vtx, f)
+        VectorSerializer.stream_serialize(CTransaction, self.vtx, f, dict(include_witness=include_witness))
 
     def get_header(self):
         """Return the block header
@@ -660,6 +727,10 @@ def GetHash(self):
             object.__setattr__(self, '_cached_GetHash', _cached_GetHash)
             return _cached_GetHash
 
+    def GetWeight(self):
+        """Return the block weight: (stripped_size * 3) + total_size"""
+        return len(self.serialize(dict(include_witness=False))) * 3 + len(self.serialize())
+
 class CoreChainParams(object):
     """Define consensus-critical parameters of a given instance of the Bitcoin system"""
     MAX_MONEY = None
@@ -721,7 +792,8 @@ def CheckTransaction(tx):
         raise CheckTransactionError("CheckTransaction() : vout empty")
 
     # Size limits
-    if len(tx.serialize()) > MAX_BLOCK_SIZE:
+    base_tx = CTransaction(tx.vin, tx.vout, tx.nLockTime, tx.nVersion)
+    if len(base_tx.serialize()) > MAX_BLOCK_SIZE:
         raise CheckTransactionError("CheckTransaction() : size limits failed")
 
     # Check for negative or overflow output values
@@ -820,7 +892,8 @@ def CheckBlock(block, fCheckPoW = True, fCheckMerkleRoot = True, cur_time=None):
 
     fCheckPoW        - Check proof-of-work.
 
-    fCheckMerkleRoot - Check merkle root matches transactions.
+    fCheckMerkleRoot - Check merkle root and witness merkle root matches transactions.
+                     - Check witness commitment in coinbase
 
     cur_time         - Current time. Defaults to time.time()
     """
@@ -831,9 +904,12 @@ def CheckBlock(block, fCheckPoW = True, fCheckMerkleRoot = True, cur_time=None):
     # Size limits
     if not block.vtx:
         raise CheckBlockError("CheckBlock() : vtx empty")
-    if len(block.serialize()) > MAX_BLOCK_SIZE:
+    if len(block.serialize(dict(include_witness=False))) > MAX_BLOCK_SIZE:
         raise CheckBlockError("CheckBlock() : block larger than MAX_BLOCK_SIZE")
 
+    if block.GetWeight() > MAX_BLOCK_WEIGHT:
+        raise CheckBlockError("CheckBlock() : block larger than MAX_BLOCK_WEIGHT")
+
     # First transaction must be coinbase
     if not block.vtx[0].is_coinbase():
         raise CheckBlockError("CheckBlock() : first tx is not coinbase")
@@ -861,8 +937,29 @@ def CheckBlock(block, fCheckPoW = True, fCheckMerkleRoot = True, cur_time=None):
             raise CheckBlockError("CheckBlock() : out-of-bounds SigOpCount")
 
     # Check merkle root
-    if fCheckMerkleRoot and block.hashMerkleRoot != block.calc_merkle_root():
-        raise CheckBlockError("CheckBlock() : hashMerkleRoot mismatch")
+    if fCheckMerkleRoot:
+        if block.hashMerkleRoot != block.calc_merkle_root():
+            raise CheckBlockError("CheckBlock() : hashMerkleRoot mismatch")
+        if len(block.vWitnessMerkleTree):
+            # At least 1 tx has witness: check witness merkle tree
+            root = block.vWitnessMerkleTree[-1]
+            # vtx[0]: coinbase
+            # vtxinwit[0]: first input
+            nonce_script = block.vtx[0].wit.vtxinwit[0].scriptWitness
+            nonce = nonce_script.stack[0]
+            if len(nonce_script.stack) != 1 or len(nonce) != 32:
+                raise CheckBlockError("CheckBlock() : invalid coinbase witnessScript")
+            try:
+                index = block.get_witness_commitment_index()
+            except ValueError as e:
+                raise CheckBlockError("CheckBlock() : " + str(e))
+            commit_script = block.vtx[0].vout[index].scriptPubKey
+            if not (6 + 32 <= len(commit_script) <= 6 + 32 + 1):
+                raise CheckBlockError("CheckBlock() : invalid segwit commitment length")
+            commitment = commit_script[6:6 + 32]
+            commit = commit_script[6:6 + 32]
+            if commit != Hash(root + nonce):
+                raise CheckBlockError("CheckBlock() : invalid segwit commitment")
 
 __all__ = (
         'Hash',
@@ -885,6 +982,8 @@ def CheckBlock(block, fCheckPoW = True, fCheckMerkleRoot = True, cur_time=None):
         'CMutableTxOut',
         'CTransaction',
         'CMutableTransaction',
+        'CTxWitness',
+        'CTxInWitness',
         'CBlockHeader',
         'CBlock',
         'CoreChainParams',

bitcoin/core/serialize.py

@@ -87,23 +87,23 @@ class Serializable(object):
 
     __slots__ = []
 
-    def stream_serialize(self, f):
+    def stream_serialize(self, f, **kwargs):
         """Serialize to a stream"""
         raise NotImplementedError
 
     @classmethod
-    def stream_deserialize(cls, f):
+    def stream_deserialize(cls, f, **kwargs):
         """Deserialize from a stream"""
         raise NotImplementedError
 
-    def serialize(self):
+    def serialize(self, params={}):
         """Serialize, returning bytes"""
         f = _BytesIO()
-        self.stream_serialize(f)
+        self.stream_serialize(f, **params)
         return f.getvalue()
 
     @classmethod
-    def deserialize(cls, buf, allow_padding=False):
+    def deserialize(cls, buf, allow_padding=False, params={}):
         """Deserialize bytes, returning an instance
 
         allow_padding - Allow buf to include extra padding. (default False)
@@ -112,7 +112,7 @@ def deserialize(cls, buf, allow_padding=False):
         deserialization DeserializationExtraDataError will be raised.
         """
         fd = _BytesIO(buf)
-        r = cls.stream_deserialize(fd)
+        r = cls.stream_deserialize(fd, **params)
         if not allow_padding:
             padding = fd.read()
             if len(padding) != 0:
@@ -242,17 +242,17 @@ class VectorSerializer(Serializer):
     # and should be rethought at some point.
 
     @classmethod
-    def stream_serialize(cls, inner_cls, objs, f):
+    def stream_serialize(cls, inner_cls, objs, f, inner_params={}):
         VarIntSerializer.stream_serialize(len(objs), f)
         for obj in objs:
-            inner_cls.stream_serialize(obj, f)
+            inner_cls.stream_serialize(obj, f, **inner_params)
 
     @classmethod
-    def stream_deserialize(cls, inner_cls, f):
+    def stream_deserialize(cls, inner_cls, f, inner_params={}):
         n = VarIntSerializer.stream_deserialize(f)
         r = []
         for i in range(n):
-            r.append(inner_cls.stream_deserialize(f))
+            r.append(inner_cls.stream_deserialize(f, **inner_params))
         return r
 
 

bitcoin/tests/test_segwit.py

@@ -12,11 +12,19 @@
 from __future__ import absolute_import, division, print_function, unicode_literals
 
 import unittest
+import random
+import sys
 
+import bitcoin
 from bitcoin.core import *
 from bitcoin.core.script import *
 from bitcoin.wallet import *
 
+if sys.version > '3':
+    _bytes = bytes
+else:
+    _bytes = lambda x: bytes(bytearray(x))
+
 # Test serialization
 
 
@@ -107,4 +115,27 @@ def test_p2sh_p2wsh_signaturehash(self):
                 0, SIGHASH_SINGLE|SIGHASH_ANYONECANPAY, value, SIGVERSION_WITNESS_V0), 
             x('511e8e52ed574121fc1b654970395502128263f62662e076dc6baf05c2e6a99b'))
 
+    def test_checkblock(self):
+        # (No witness) coinbase generated by Bitcoin Core
+        str_coinbase = '01000000010000000000000000000000000000000000000000000000000000000000000000ffffffff03520101ffffffff0100f2052a01000000232102960c90bc04a631cb17922e4f5d80ac75fd590a88b8baaa5a3d5086ac85e4d788ac00000000'
+        # No witness transaction
+        str_tx = '0100000001e3d0c1d051a3abe79d5951dcab86f71d8926aff5caed5ff9bd72cb593e4ebaf5010000006b483045022100a28e1c57e160296953e1af85c5034bb1b907c12c50367da75ba547874a8a8c52022040682e888ddce7fd5c72519a9f28f22f5164c8af864d33332bbc7f65596ad4ae012102db30394fd5cc8288bed607b9382338240c014a98c9c0febbfb380db74ceb30a3ffffffff020d920000000000001976a914ad781c8ffcc18b2155433cba2da4601180a66eef88aca3170000000000001976a914bacb1c4b0725e61e385c7093b4260533953c8e1688ac00000000'
+        # SegWit transaction
+        str_w_tx = '0100000000010115e180dc28a2327e687facc33f10f2a20da717e5548406f7ae8b4c811072f8560100000000ffffffff0100b4f505000000001976a9141d7cd6c75c2e86f4cbf98eaed221b30bd9a0b92888ac02483045022100df7b7e5cda14ddf91290e02ea10786e03eb11ee36ec02dd862fe9a326bbcb7fd02203f5b4496b667e6e281cc654a2da9e4f08660c620a1051337fa8965f727eb19190121038262a6c6cec93c2d3ecd6c6072efea86d02ff8e3328bbd0242b20af3425990ac00000000'
+        witness_nonce = _bytes(random.getrandbits(8) for _ in range(32))
+
+        coinbase = CMutableTransaction.deserialize(x(str_coinbase))
+        coinbase.wit = CTxWitness([CTxInWitness(CScriptWitness([witness_nonce]))])
+
+        tx_legacy = CTransaction.deserialize(x(str_tx))
+        tx_segwit = CTransaction.deserialize(x(str_w_tx))
+
+        witness_merkle_root = CBlock.build_witness_merkle_tree_from_txs((coinbase, tx_legacy, tx_segwit))[-1]
+
+        commitment = Hash(witness_merkle_root + witness_nonce)
+        commitment_script = bitcoin.core.WITNESS_COINBASE_SCRIPTPUBKEY_MAGIC + commitment
+        coinbase.vout.append(CTxOut(0, CScript(commitment_script)))
+
+        block = CBlock(2, b'\x00'*32, b'\x00'*32, 0, 0, 0, (coinbase, tx_legacy, tx_segwit))
 
+        CheckBlock(block, False, True)