Add some utility functions and tests

electrumx/lib/util.py

@@ -330,10 +330,38 @@ def protocol_version(client_req, min_tuple, max_tuple):
     return result, client_min
 
 
-unpack_int32_from = Struct('<i').unpack_from
-unpack_int64_from = Struct('<q').unpack_from
-unpack_uint16_from = Struct('<H').unpack_from
-unpack_uint32_from = Struct('<I').unpack_from
-unpack_uint64_from = Struct('<Q').unpack_from
+structi = Struct('<i')
+structq = Struct('<q')
+structH = Struct('<H')
+structI = Struct('<I')
+structQ = Struct('<Q')
+structB = Struct('B')
+
+unpack_int32_from = structi.unpack_from
+unpack_int64_from = structq.unpack_from
+unpack_uint16_from = structH.unpack_from
+unpack_uint32_from = structI.unpack_from
+unpack_uint64_from = structQ.unpack_from
+
+pack_le_int32 = structi.pack
+pack_le_int64 = structq.pack
+pack_le_uint16 = structH.pack
+pack_le_uint32 = structI.pack
+pack_le_uint64 = structQ.pack
+pack_byte = structB.pack
 
 hex_to_bytes = bytes.fromhex
+
+
+def pack_varint(n):
+    if n < 253:
+        return pack_byte(n)
+    if n < 65536:
+        return pack_byte(253) + pack_le_uint16(n)
+    if n < 4294967296:
+        return pack_byte(254) + pack_le_uint32(n)
+    return pack_byte(255) + pack_le_uint64(n)
+
+
+def pack_varbytes(data):
+    return pack_varint(len(data)) + data

tests/lib/test_util.py

@@ -2,7 +2,7 @@ import os
 
 import pytest
 
-from electrumx.lib import util
+from electrumx.lib import util, tx
 
 
 def test_cachedproperty():
@@ -223,3 +223,21 @@ def test_unpackers():
 def test_hex_transforms():
     h = "AABBCCDDEEFF"
     assert util.hex_to_bytes(h) == b'\xaa\xbb\xcc\xdd\xee\xff'
+
+
+def test_pack_varint():
+    tests = list(range(0, 258))
+    tests.extend([1024, 65535, 65536, 4294967295, 4294967296, 8294967296])
+
+    for n in tests:
+        data = util.pack_varint(n)
+        deser = tx.Deserializer(data)
+        assert deser._read_varint() == n
+
+def test_pack_varbytes():
+    tests = [b'', b'1', b'2' * 253, b'3' * 254, b'4' * 256, b'5' * 65536]
+
+    for test in tests:
+        data = util.pack_varbytes(test)
+        deser = tx.Deserializer(data)
+        assert deser._read_varbytes() == test