commit d34b517c9661081d7ac69e464c233e61c5b42ed3 Author: Damian Johnson atagar@torproject.org Date: Thu Jan 4 11:34:30 2018 -0800

Cell.pack() function

Counterpart for endosome's pack_cell() function. This deserves more tests once we have an unpack method (assertions for packed values aren't terribly grokable). --- stem/client.py | 57 ++++++++++++++++++++++++++++++++++++++++++++++++++--- test/unit/client.py | 33 ++++++++++++++++++++++++------- 2 files changed, 80 insertions(+), 10 deletions(-)

diff --git a/stem/client.py b/stem/client.py index 8fe2878d..2a598747 100644 --- a/stem/client.py +++ b/stem/client.py @@ -36,12 +36,14 @@ import struct

from stem.util import enum

+ZERO = '\x00' +

class Cell(collections.namedtuple('Cell', ['name', 'value', 'fixed_size', 'for_circuit'])): """ Metadata for ORPort cells.

- :var str name: name of the cell type + :var str name: command of the cell :var int value: integer value of the command on the wire :var bool fixed_size: **True** if cells have a fixed length, **False** if variable @@ -54,7 +56,7 @@ class Cell(collections.namedtuple('Cell', ['name', 'value', 'fixed_size', 'for_c """ Provides cell attributes by its name.

- :parm str name: name of the cell type to fetch + :parm str name: cell command to fetch

:raise: **ValueError** if cell type is invalid """ @@ -70,7 +72,7 @@ class Cell(collections.namedtuple('Cell', ['name', 'value', 'fixed_size', 'for_c """ Provides cell attributes by its value.

- :parm int value: value of the cell type to fetch + :parm int value: cell value to fetch

:raise: **ValueError** if cell type is invalid """ @@ -81,6 +83,55 @@ class Cell(collections.namedtuple('Cell', ['name', 'value', 'fixed_size', 'for_c

raise ValueError("'%s' isn't a valid cell value" % value)

+ @staticmethod + def pack(name, link_version, payload, circ_id = None): + """ + Provides bytes that can be used on the wire for these cell attributes. + + :param str name: cell command + :param int link_version: link protocol version + :param bytes payload: cell payload + :param int circ_id: circuit id, if for a circuit + + :raise: **ValueError** if... + * cell type or circuit id is invalid + * payload is too large + """ + + attr = Cell.by_name(name) + circ_id_len = Pack.LONG if link_version > 3 else Pack.SHORT + + if attr.for_circuit and circ_id is None: + if name.startswith('CREATE'): + # Since we're initiating the circuit we pick any value from a range + # that's determined by our link version. + + circ_id = 0x80000000 if link_version > 3 else 0x01 + else: + raise ValueError('%s cells require a circ_id' % name) + elif not attr.for_circuit: + if circ_id: + raise ValueError("%s cells don't concern circuits, circ_id is unused" % name) + + circ_id = 0 # field is still mandatory for all cells + + packed_circ_id = struct.pack(circ_id_len, circ_id) + packed_command = struct.pack(Pack.CHAR, attr.value) + packed_size = b'' if attr.fixed_size else struct.pack(Pack.SHORT, len(payload)) + cell = b''.join((packed_circ_id, packed_command, packed_size, payload)) + + # pad fixed sized cells to the required length + + if attr.fixed_size: + fixed_cell_len = 514 if link_version > 3 else 512 + + if len(cell) > fixed_cell_len: + raise ValueError('Payload of %s is too large (%i bytes), must be less than %i' % (name, len(cell), fixed_cell_len)) + + cell += ZERO * (fixed_cell_len - len(cell)) + + return cell +

class Relay(object): """ diff --git a/test/unit/client.py b/test/unit/client.py index 1d2df417..6faa3e8e 100644 --- a/test/unit/client.py +++ b/test/unit/client.py @@ -2,24 +2,43 @@ Unit tests for the stem.client. """

+import struct import unittest

-from stem.client import Cell +from stem.client import Pack, Cell

class TestClient(unittest.TestCase): - def test_cell_fetching(self): + def test_cell_by_name(self): cell = Cell.by_name('NETINFO') - self.assertEqual('NETINFO', cell.name) self.assertEqual(8, cell.value) self.assertEqual(True, cell.fixed_size) self.assertEqual(False, cell.for_circuit)

- self.assertEqual(10, Cell.by_name('CREATE2').value) - self.assertEqual('CREATE2', Cell.by_value(10).name) - self.assertRaises(ValueError, Cell.by_name, 'NOPE') - self.assertRaises(ValueError, Cell.by_value, 'NOPE') self.assertRaises(ValueError, Cell.by_name, 85) self.assertRaises(ValueError, Cell.by_name, None) + + def test_cell_by_value(self): + cell = Cell.by_value(8) + self.assertEqual('NETINFO', cell.name) + self.assertEqual(8, cell.value) + self.assertEqual(True, cell.fixed_size) + self.assertEqual(False, cell.for_circuit) + + self.assertRaises(ValueError, Cell.by_value, 'NOPE') + self.assertRaises(ValueError, Cell.by_value, 85) + self.assertRaises(ValueError, Cell.by_value, None) + + def test_cell_pack(self): + version_payload = struct.pack(Pack.SHORT, 2) + + # basic link v2 and v4 VERSIONS cell + + self.assertEqual('\x00\x00\x07\x00\x02\x00\x02', Cell.pack('VERSIONS', 2, version_payload)) + self.assertEqual('\x00\x00\x00\x00\x07\x00\x02\x00\x02', Cell.pack('VERSIONS', 4, version_payload)) + + self.assertRaisesRegexp(ValueError, "'NOPE' isn't a valid cell type", Cell.pack, 'NOPE', 2, version_payload) + self.assertRaisesRegexp(ValueError, "VERSIONS cells don't concern circuits, circ_id is unused", Cell.pack, 'VERSIONS', 2, version_payload, circ_id = 5) + self.assertRaisesRegexp(ValueError, 'RELAY_EARLY cells require a circ_id', Cell.pack, 'RELAY_EARLY', 2, version_payload)