[tor-commits] [stem/master] Cell.pack() function

[atagar at torproject.org]

commit d34b517c9661081d7ac69e464c233e61c5b42ed3
Author: Damian Johnson <atagar at 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)