test/functional/test_framework/blockstore.py

   1 #!/usr/bin/env python3
   2 # Copyright (c) 2015-2017 The Bitcoin Core developers
   3 # Distributed under the MIT software license, see the accompanying
   4 # file COPYING or http://www.opensource.org/licenses/mit-license.php.
   5 """BlockStore and TxStore helper classes."""
   6
   7 from .mininode import *
   8 from io import BytesIO
   9 import dbm.dumb as dbmd
  10
  11 logger = logging.getLogger("TestFramework.blockstore")
  12
  13 class BlockStore():
  14     """BlockStore helper class.
  15
  16     BlockStore keeps a map of blocks and implements helper functions for
  17     responding to getheaders and getdata, and for constructing a getheaders
  18     message.
  19     """
  20
  21     def __init__(self, datadir):
  22         self.blockDB = dbmd.open(datadir + "/blocks", 'c')
  23         self.currentBlock = 0
  24         self.headers_map = dict()
  25
  26     def close(self):
  27         self.blockDB.close()
  28
  29     def erase(self, blockhash):
  30         del self.blockDB[repr(blockhash)]
  31
  32     # lookup an entry and return the item as raw bytes
  33     def get(self, blockhash):
  34         value = None
  35         try:
  36             value = self.blockDB[repr(blockhash)]
  37         except KeyError:
  38             return None
  39         return value
  40
  41     # lookup an entry and return it as a CBlock
  42     def get_block(self, blockhash):
  43         ret = None
  44         serialized_block = self.get(blockhash)
  45         if serialized_block is not None:
  46             f = BytesIO(serialized_block)
  47             ret = CBlock()
  48             ret.deserialize(f)
  49             ret.calc_sha256()
  50         return ret
  51
  52     def get_header(self, blockhash):
  53         try:
  54             return self.headers_map[blockhash]
  55         except KeyError:
  56             return None
  57
  58     # Note: this pulls full blocks out of the database just to retrieve
  59     # the headers -- perhaps we could keep a separate data structure
  60     # to avoid this overhead.
  61     def headers_for(self, locator, hash_stop, current_tip=None):
  62         if current_tip is None:
  63             current_tip = self.currentBlock
  64         current_block_header = self.get_header(current_tip)
  65         if current_block_header is None:
  66             return None
  67
  68         response = msg_headers()
  69         headersList = [ current_block_header ]
  70         maxheaders = 2000
  71         while (headersList[0].sha256 not in locator.vHave):
  72             prevBlockHash = headersList[0].hashPrevBlock
  73             prevBlockHeader = self.get_header(prevBlockHash)
  74             if prevBlockHeader is not None:
  75                 headersList.insert(0, prevBlockHeader)
  76             else:
  77                 break
  78         headersList = headersList[:maxheaders] # truncate if we have too many
  79         hashList = [x.sha256 for x in headersList]
  80         index = len(headersList)
  81         if (hash_stop in hashList):
  82             index = hashList.index(hash_stop)+1
  83         response.headers = headersList[:index]
  84         return response
  85
  86     def add_block(self, block):
  87         block.calc_sha256()
  88         try:
  89             self.blockDB[repr(block.sha256)] = bytes(block.serialize())
  90         except TypeError as e:
  91             logger.exception("Unexpected error")
  92         self.currentBlock = block.sha256
  93         self.headers_map[block.sha256] = CBlockHeader(block)
  94
  95     def add_header(self, header):
  96         self.headers_map[header.sha256] = header
  97
  98     # lookup the hashes in "inv", and return p2p messages for delivering
  99     # blocks found.
 100     def get_blocks(self, inv):
 101         responses = []
 102         for i in inv:
 103             if (i.type == 2 or i.type == (2 | (1 << 30))): # MSG_BLOCK or MSG_WITNESS_BLOCK
 104                 data = self.get(i.hash)
 105                 if data is not None:
 106                     # Use msg_generic to avoid re-serialization
 107                     responses.append(msg_generic(b"block", data))
 108         return responses
 109
 110     def get_locator(self, current_tip=None):
 111         if current_tip is None:
 112             current_tip = self.currentBlock
 113         r = []
 114         counter = 0
 115         step = 1
 116         lastBlock = self.get_block(current_tip)
 117         while lastBlock is not None:
 118             r.append(lastBlock.hashPrevBlock)
 119             for i in range(step):
 120                 lastBlock = self.get_block(lastBlock.hashPrevBlock)
 121                 if lastBlock is None:
 122                     break
 123             counter += 1
 124             if counter > 10:
 125                 step *= 2
 126         locator = CBlockLocator()
 127         locator.vHave = r
 128         return locator
 129
 130 class TxStore():
 131     def __init__(self, datadir):
 132         self.txDB = dbmd.open(datadir + "/transactions", 'c')
 133
 134     def close(self):
 135         self.txDB.close()
 136
 137     # lookup an entry and return the item as raw bytes
 138     def get(self, txhash):
 139         value = None
 140         try:
 141             value = self.txDB[repr(txhash)]
 142         except KeyError:
 143             return None
 144         return value
 145
 146     def add_transaction(self, tx):
 147         tx.calc_sha256()
 148         try:
 149             self.txDB[repr(tx.sha256)] = bytes(tx.serialize())
 150         except TypeError as e:
 151             logger.exception("Unexpected error")
 152
 153     def get_transactions(self, inv):
 154         responses = []
 155         for i in inv:
 156             if (i.type == 1 or i.type == (1 | (1 << 30))): # MSG_TX or MSG_WITNESS_TX
 157                 tx = self.get(i.hash)
 158                 if tx is not None:
 159                     responses.append(msg_generic(b"tx", tx))
 160         return responses