def get_connection_protocol_version(self,
input_buffer: InputBuffer) -> int:
if (input_buffer.length <
(constants.STARTING_SEQUENCE_BYTES_LEN +
constants.BX_HDR_COMMON_OFF + constants.VERSION_NUM_LEN)):
return self.CURRENT_PROTOCOL_VERSION
header_buf = input_buffer.peek_message(VersionMessage.HEADER_LENGTH)
if header_buf[:constants.
STARTING_SEQUENCE_BYTES_LEN] == constants.STARTING_SEQUENCE_BYTES:
command, payload_len = VersionMessage.unpack(header_buf)
header_len = VersionMessage.HEADER_LENGTH
else:
command = bytearray(header_buf[:constants.MSG_TYPE_LEN])
payload_len = constants.MSG_TYPE_LEN
header_len = constants.BX_HDR_COMMON_OFF
if command != self.version_message_command:
if constants.HTTP_MESSAGE in command:
raise NonVersionMessageError(
msg=
"Instead of a version hello message, we received an HTTP request: {} with payload length: {} "
"Ignoring and closing connection. ".format(
header_buf, payload_len),
is_known=True)
if command in constants.BITCOIN_MESSAGES:
raise NonVersionMessageError(
msg="Received some kind of bitcoin peering message: {}. "
"Ignoring and closing connection.".format(command),
is_known=True)
logger.debug(
"Received message of type {} instead of hello message. Use current version of protocol {}.",
command, self.CURRENT_PROTOCOL_VERSION)
return self.CURRENT_PROTOCOL_VERSION
if payload_len < self.VERSION_MESSAGE_MAIN_LENGTH:
return 1
version_buf = input_buffer.get_slice(header_len,
header_len + VERSION_NUM_LEN)
version, = struct.unpack_from("<L", version_buf, 0)
return version
class TestInputBuffer(unittest.TestCase):
def setUp(self):
self.in_buf = InputBuffer()
self.length_to_add = 20
self.data1 = bytearray([i for i in range(1, 21)])
self.data2 = bytearray([i for i in range(21, 41)])
self.data3 = bytearray([i for i in range(41, 61)])
def test_endswith(self):
self.assertFalse(self.in_buf.endswith(5))
self.in_buf.add_bytes(bytearray([1, 2, 3, 4, 5, 6]))
self.assertTrue(self.in_buf.endswith(bytearray([5, 6])))
self.assertFalse(self.in_buf.endswith(bytearray([4, 6])))
with self.assertRaises(ValueError):
self.in_buf.endswith([5, 6])
self.assertEqual(
self.in_buf.endswith(bytearray([0, 1, 2, 3, 4, 5, 6])), False)
self.assertEqual(True, self.in_buf.endswith(bytearray(0)))
def test_add_bytes(self):
self.in_buf.add_bytes(self.data1)
self.in_buf.add_bytes(self.data2)
self.assertEqual(2 * self.length_to_add, self.in_buf.length)
self.assertEqual(deque([self.data1, self.data2]),
self.in_buf.input_list)
def test_remove_bytes(self):
with self.assertRaises(AssertionError):
self.in_buf.remove_bytes(5)
self.in_buf.remove_bytes(0)
self.assertEqual(0, self.in_buf.length, 0)
self.make_input_buffer()
self.assertEqual(bytearray([i for i in range(1, 6)]),
self.in_buf.remove_bytes(5))
self.assertEqual(55, self.in_buf.length, 55)
self.assertEqual(bytearray([i for i in range(6, 26)]),
self.in_buf.remove_bytes(20))
self.assertEqual(35, self.in_buf.length)
self.assertEqual(bytearray([i for i in range(26, 61)]),
self.in_buf.remove_bytes(35))
self.assertEqual(0, self.in_buf.length)
def test_peek_message(self):
self.make_input_buffer()
# Edge Case: peek_message returns all bytes when it peeks a number greater than the message length.
self.assertEqual(bytearray([i for i in range(1, 61)]),
self.in_buf.peek_message(70))
self.assertIn(bytearray([i for i in range(1, 5)]),
self.in_buf.peek_message(5))
self.assertIn(bytearray([i for i in range(1, 31)]),
self.in_buf.peek_message(30))
self.assertIn(bytearray([i for i in range(1, 61)]),
self.in_buf.peek_message(60))
def test_get_slice(self):
self.make_input_buffer()
with self.assertRaises(ValueError):
self.in_buf.get_slice(5, None)
with self.assertRaises(ValueError):
self.in_buf.get_slice(None, 6)
with self.assertRaises(ValueError):
self.in_buf.get_slice(100, 60)
self.assertEqual(bytearray([i for i in range(1, 6)]),
self.in_buf.get_slice(0, 5))
self.assertEqual(self.in_buf[:5], self.in_buf.get_slice(0, 5))
self.assertEqual(bytearray([i for i in range(11, 32)]),
self.in_buf.get_slice(10, 31))
self.assertEqual(self.in_buf[10:31], self.in_buf.get_slice(10, 31))
self.assertEqual(bytearray([i for i in range(35, 61)]),
self.in_buf.get_slice(34, 60))
self.assertEqual(self.in_buf[34:], self.in_buf.get_slice(34, 60))
def make_input_buffer(self):
self.in_buf.add_bytes(self.data1)
self.in_buf.add_bytes(self.data2)
self.in_buf.add_bytes(self.data3)