From 450eec89fb95512462303b3c56e3831e455cbf6a Mon Sep 17 00:00:00 2001
From: Neil Booth <kyuupichan@gmail.com>
Date: Fri, 10 Aug 2018 17:39:47 +0900
Subject: [PATCH] Add some utility functions and tests

---
 electrumx/lib/util.py  | 38 +++++++++++++++++++++++++++++++++-----
 tests/lib/test_util.py | 20 +++++++++++++++++++-
 2 files changed, 52 insertions(+), 6 deletions(-)

diff --git a/electrumx/lib/util.py b/electrumx/lib/util.py
index 8f49e2d..991cd88 100644
--- a/electrumx/lib/util.py
+++ b/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
diff --git a/tests/lib/test_util.py b/tests/lib/test_util.py
index 0351471..6186da6 100644
--- a/tests/lib/test_util.py
+++ b/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