Use compiled structs and some code cleanup

electrumx/lib/script.py

@@ -32,6 +32,8 @@ from collections import namedtuple
 
 from electrumx.lib.enum import Enumeration
 from electrumx.lib.hash import hash160
+from electrumx.lib.util import unpack_le_uint16_from, unpack_le_uint32_from, \
+    pack_le_uint16, pack_le_uint32
 
 
 class ScriptError(Exception):
@@ -196,10 +198,10 @@ class Script(object):
                         dlen = script[n]
                         n += 1
                     elif op == OpCodes.OP_PUSHDATA2:
-                        dlen, = struct.unpack('<H', script[n: n + 2])
+                        dlen, = unpack_le_uint16_from(script[n: n + 2])
                         n += 2
                     else:
-                        dlen, = struct.unpack('<I', script[n: n + 4])
+                        dlen, = unpack_le_uint32_from(script[n: n + 4])
                         n += 4
                     if n + dlen > len(script):
                         raise IndexError
@@ -225,8 +227,8 @@ class Script(object):
         if n < 256:
             return bytes([OpCodes.OP_PUSHDATA1, n]) + data
         if n < 65536:
-            return bytes([OpCodes.OP_PUSHDATA2]) + struct.pack('<H', n) + data
+            return bytes([OpCodes.OP_PUSHDATA2]) + pack_le_uint16(n) + data
-        return bytes([OpCodes.OP_PUSHDATA4]) + struct.pack('<I', n) + data
+        return bytes([OpCodes.OP_PUSHDATA4]) + pack_le_uint32(n) + data
 
     @classmethod
     def opcode_name(cls, opcode):

electrumx/lib/util.py

@@ -168,20 +168,6 @@ def int_to_bytes(value):
     return value.to_bytes((value.bit_length() + 7) // 8, 'big')
 
 
-def int_to_varint(value):
-    '''Converts an integer to a Bitcoin-like varint bytes'''
-    if value < 0:
-        raise ValueError("attempt to write size < 0")
-    elif value < 253:
-        return pack('<B', value)
-    elif value < 2**16:
-        return b'\xfd' + pack('<H', value)
-    elif value < 2**32:
-        return b'\xfe' + pack('<I', value)
-    elif value < 2**64:
-        return b'\xff' + pack('<Q', value)
-
-
 def increment_byte_string(bs):
     '''Return the lexicographically next byte string of the same length.
 

electrumx/server/daemon.py

@@ -17,8 +17,8 @@ from time import strptime
 
 import aiohttp
 
-from electrumx.lib.util import int_to_varint, hex_to_bytes, class_logger, \
+from electrumx.lib.util import hex_to_bytes, class_logger,\
-    unpack_le_uint16_from
+    unpack_le_uint16_from, pack_varint
 from electrumx.lib.hash import hex_str_to_hash, hash_to_hex_str
 from electrumx.lib.tx import DeserializerDecred
 from aiorpcx import JSONRPC
@@ -356,7 +356,7 @@ class LegacyRPCDaemon(Daemon):
         raw_block = header
         num_txs = len(transactions)
         if num_txs > 0:
-            raw_block += int_to_varint(num_txs)
+            raw_block += pack_varint(num_txs)
             raw_block += b''.join(transactions)
         else:
             raw_block += b'\x00'

electrumx/server/history.py

@@ -14,9 +14,9 @@ import bisect
 import time
 from collections import defaultdict
 from functools import partial
-from struct import pack, unpack
 
 import electrumx.lib.util as util
+from electrumx.lib.util import pack_be_uint16, unpack_be_uint16_from
 from electrumx.lib.hash import hash_to_hex_str, HASHX_LEN
 
 
@@ -81,7 +81,7 @@ class History(object):
 
         keys = []
         for key, hist in self.db.iterator(prefix=b''):
-            flush_id, = unpack('>H', key[-2:])
+            flush_id, = unpack_be_uint16_from(key[-2:])
             if flush_id > utxo_flush_count:
                 keys.append(key)
 
@@ -126,7 +126,7 @@ class History(object):
     def flush(self):
         start_time = time.time()
         self.flush_count += 1
-        flush_id = pack('>H', self.flush_count)
+        flush_id = pack_be_uint16(self.flush_count)
         unflushed = self.unflushed
 
         with self.db.write_batch() as batch:
@@ -250,7 +250,7 @@ class History(object):
         write_size = 0
         keys_to_delete.update(hist_map)
         for n, chunk in enumerate(util.chunks(full_hist, max_row_size)):
-            key = hashX + pack('>H', n)
+            key = hashX + pack_be_uint16(n)
             if hist_map.get(key) == chunk:
                 keys_to_delete.remove(key)
             else:
@@ -302,7 +302,7 @@ class History(object):
         # Loop over 2-byte prefixes
         cursor = self.comp_cursor
         while write_size < limit and cursor < 65536:
-            prefix = pack('>H', cursor)
+            prefix = pack_be_uint16(cursor)
             write_size += self._compact_prefix(prefix, write_items,
                                                keys_to_delete)
             cursor += 1

electrumx/wallet/bip32.py

@@ -15,7 +15,8 @@ import ecdsa.numbertheory as NT
 
 from electrumx.lib.coins import Coin
 from electrumx.lib.hash import Base58, hmac_sha512, hash160
-from electrumx.lib.util import cachedproperty, bytes_to_int, int_to_bytes
+from electrumx.lib.util import cachedproperty, bytes_to_int, int_to_bytes, \
+    pack_be_uint32, unpack_be_uint32_from
 
 
 class DerivationError(Exception):
@@ -65,7 +66,7 @@ class _KeyBase(object):
             raise ValueError('raw_serkey must have length 33')
 
         return (ver_bytes + bytes([self.depth])
-                + self.parent_fingerprint() + struct.pack('>I', self.n)
+                + self.parent_fingerprint() + pack_be_uint32(self.n)
                 + self.chain_code + raw_serkey)
 
     def fingerprint(self):
@@ -142,7 +143,7 @@ class PubKey(_KeyBase):
         if not 0 <= n < (1 << 31):
             raise ValueError('invalid BIP32 public key child number')
 
-        msg = self.pubkey_bytes + struct.pack('>I', n)
+        msg = self.pubkey_bytes + pack_be_uint32(n)
         L, R = self._hmac_sha512(msg)
 
         curve = self.CURVE
@@ -244,7 +245,7 @@ class PrivKey(_KeyBase):
         else:
             serkey = self.public_key.pubkey_bytes
 
-        msg = serkey + struct.pack('>I', n)
+        msg = serkey + pack_be_uint32(n)
         L, R = self._hmac_sha512(msg)
 
         curve = self.CURVE
@@ -282,7 +283,7 @@ def _from_extended_key(ekey):
     is_public, coin = Coin.lookup_xverbytes(ekey[:4])
     depth = ekey[4]
     fingerprint = ekey[5:9]   # Not used
-    n, = struct.unpack('>I', ekey[9:13])
+    n, = unpack_be_uint32_from(ekey[9:13])
     chain_code = ekey[13:45]
 
     if is_public:

tests/lib/test_util.py

@@ -71,25 +71,14 @@ def test_increment_byte_string():
     assert util.increment_byte_string(b'\x01\x01') == b'\x01\x02'
     assert util.increment_byte_string(b'\xff\xff') is None
 
+
 def test_bytes_to_int():
     assert util.bytes_to_int(b'\x07[\xcd\x15') == 123456789
 
+
 def test_int_to_bytes():
     assert util.int_to_bytes(456789) == b'\x06\xf8U'
 
-def test_int_to_varint():
-    with pytest.raises(ValueError):
-        util.int_to_varint(-1)
-    assert util.int_to_varint(0) == b'\0'
-    assert util.int_to_varint(5) == b'\5'
-    assert util.int_to_varint(252) == b'\xfc'
-    assert util.int_to_varint(253) == b'\xfd\xfd\0'
-    assert util.int_to_varint(65535) == b'\xfd\xff\xff'
-    assert util.int_to_varint(65536) == b'\xfe\0\0\1\0'
-    assert util.int_to_varint(2**32-1) == b'\xfe\xff\xff\xff\xff'
-    assert util.int_to_varint(2**32) == b'\xff\0\0\0\0\1\0\0\0'
-    assert util.int_to_varint(2**64-1) \
-        == b'\xff\xff\xff\xff\xff\xff\xff\xff\xff'
 
 def test_LogicalFile(tmpdir):
     prefix = os.path.join(tmpdir, 'log')
@@ -234,6 +223,20 @@ def test_pack_varint():
         deser = tx.Deserializer(data)
         assert deser._read_varint() == n
 
+    import struct
+    with pytest.raises(struct.error):
+        util.pack_varint(-1)
+    assert util.pack_varint(0) == b'\0'
+    assert util.pack_varint(5) == b'\5'
+    assert util.pack_varint(252) == b'\xfc'
+    assert util.pack_varint(253) == b'\xfd\xfd\0'
+    assert util.pack_varint(65535) == b'\xfd\xff\xff'
+    assert util.pack_varint(65536) == b'\xfe\0\0\1\0'
+    assert util.pack_varint(2**32-1) == b'\xfe\xff\xff\xff\xff'
+    assert util.pack_varint(2**32) == b'\xff\0\0\0\0\1\0\0\0'
+    assert util.pack_varint(2**64-1) \
+           == b'\xff\xff\xff\xff\xff\xff\xff\xff\xff'
+
 def test_pack_varbytes():
     tests = [b'', b'1', b'2' * 253, b'3' * 254, b'4' * 256, b'5' * 65536]
 

tests/server/test_compaction.py

@@ -2,12 +2,11 @@
 
 import array
 import asyncio
-from collections import defaultdict
 from os import environ, urandom
-from struct import pack
 import random
 
 from electrumx.lib.hash import HASHX_LEN
+from electrumx.lib.util import pack_be_uint16
 from electrumx.server.env import Env
 from electrumx.server.db import DB
 
@@ -49,7 +48,7 @@ def check_hashX_compaction(history):
     hist_list = []
     hist_map = {}
     for flush_count, count in pairs:
-        key = hashX + pack('>H', flush_count)
+        key = hashX + pack_be_uint16(flush_count)
         hist = full_hist[cum * 4: (cum+count) * 4]
         hist_map[key] = hist
         hist_list.append(hist)
@@ -65,10 +64,10 @@ def check_hashX_compaction(history):
     assert len(keys_to_delete) == 3
     assert len(hist_map) == len(pairs)
     for n, item in enumerate(write_items):
-        assert item == (hashX + pack('>H', n),
+        assert item == (hashX + pack_be_uint16(n),
                         full_hist[n * row_size: (n + 1) * row_size])
     for flush_count, count in pairs:
-        assert hashX + pack('>H', flush_count) in keys_to_delete
+        assert hashX + pack_be_uint16(flush_count) in keys_to_delete
 
     # Check re-compaction is null
     hist_map = {key: value for key, value in write_items}
@@ -87,7 +86,7 @@ def check_hashX_compaction(history):
     write_size = history._compact_hashX(hashX, hist_map, hist_list,
                                         write_items, keys_to_delete)
     assert write_size == len(hist_list[-1])
-    assert write_items == [(hashX + pack('>H', 2), hist_list[-1])]
+    assert write_items == [(hashX + pack_be_uint16(2), hist_list[-1])]
     assert len(keys_to_delete) == 1
     assert write_items[0][0] in keys_to_delete
     assert len(hist_map) == len(pairs)