[tor-commits] [stem/master] Have Cell.unpack() provide a list of cells

[atagar at torproject.org]

commit 3db7fab9c3ea84a68c154a0a33829bbcf5da8cdf
Author: Damian Johnson <atagar at torproject.org>
Date:   Sat Jan 13 12:07:18 2018 -0800

    Have Cell.unpack() provide a list of cells
    
    On reflection ORPort socket responses contain a series of cells, not
    just one.   Changing our unpack function to parse a series. We'll need
    CERTS, AUTH_CHALLENGE, and NETINFO cell support before we can open links.
---
 stem/client/cell.py                  |  33 ++++++++++++++++++++++-----------
 test/unit/client/__init__.py         |  17 +++++++++++++++++
 test/unit/client/cell.py             |  11 ++++++++---
 test/unit/client/data/new_link_cells | Bin 0 -> 2043 bytes
 4 files changed, 47 insertions(+), 14 deletions(-)

diff --git a/stem/client/cell.py b/stem/client/cell.py
index e29e404f..4d4667d0 100644
--- a/stem/client/cell.py
+++ b/stem/client/cell.py
@@ -42,7 +42,7 @@ import sys
 from stem import UNDEFINED
 from stem.client import ZERO, Size
 
-MAX_FIXED_PAYLOAD_LEN = 509
+FIXED_PAYLOAD_LEN = 509
 
 
 class Cell(object):
@@ -96,28 +96,39 @@ class Cell(object):
   @staticmethod
   def unpack(content, link_version):
     """
-    Unpacks encoded bytes into a Cell subclass.
+    Unpacks encoded bytes into a series of cells.
 
     :param bytes content: payload to decode
     :param int link_version: link protocol version
 
-    :returns: :class:`~stem.client.cell.Cell` subclass
+    :returns: **list** of :class:`~stem.client.cell.Cell` subclasses
 
     :raises:
       * ValueError if content is malformed
       * NotImplementedError if unable to unpack this cell type
     """
 
-    circ_id, content = Size.LONG.pop(content) if link_version > 3 else Size.SHORT.pop(content)
-    command, content = Size.CHAR.pop(content)
-    cls = Cell.by_value(command)
+    cells = []
 
-    if cls.IS_FIXED_SIZE:
-      payload_len = MAX_FIXED_PAYLOAD_LEN
-    else:
-      payload_len, content = Size.SHORT.pop(content)
+    while content:
+      circ_id, content = Size.SHORT.pop(content) if link_version < 4 else Size.LONG.pop(content)
+      command, content = Size.CHAR.pop(content)
+      cls = Cell.by_value(command)
+
+      if cls.IS_FIXED_SIZE:
+        payload_len = FIXED_PAYLOAD_LEN
+      else:
+        payload_len, content = Size.SHORT.pop(content)
+
+      if len(content) < payload_len:
+        raise ValueError('%s cell should have a payload of %i bytes, but only had %i' % (cls.NAME, payload_len, len(content)))
+
+      payload = content[:payload_len].rstrip(ZERO)  # strip padding
+      content = content[payload_len:]
+
+      cells.append(cls._unpack(payload, link_version, circ_id))
 
-    return cls._unpack(content, link_version, circ_id)
+    return cells
 
   @classmethod
   def _pack(cls, link_version, payload, circ_id = 0):
diff --git a/test/unit/client/__init__.py b/test/unit/client/__init__.py
index e35416a8..f5b82e7e 100644
--- a/test/unit/client/__init__.py
+++ b/test/unit/client/__init__.py
@@ -5,3 +5,20 @@ Unit tests for stem.client.* contents.
 __all__ = [
   'cell',
 ]
+
+import os
+
+TEST_DATA = os.path.join(os.path.dirname(__file__), 'data')
+
+
+def test_data(filename):
+  """
+  Provides test data in the given file.
+
+  :param str filename: test data to provide
+
+  :returns: **bytes** with the data
+  """
+
+  with open(os.path.join(TEST_DATA, filename), 'rb') as data_file:
+    return data_file.read()
diff --git a/test/unit/client/cell.py b/test/unit/client/cell.py
index 9502975e..3a2a55f2 100644
--- a/test/unit/client/cell.py
+++ b/test/unit/client/cell.py
@@ -5,6 +5,7 @@ Unit tests for the stem.client.cell.
 import unittest
 
 from stem.client.cell import Cell, VersionsCell
+from test.unit.client import test_data
 
 
 class TestCell(unittest.TestCase):
@@ -31,12 +32,16 @@ class TestCell(unittest.TestCase):
   def test_unpack_not_implemented(self):
     self.assertRaisesRegexp(NotImplementedError, 'Unpacking not yet implemented for AUTHORIZE cells', Cell.unpack, '\x00\x00\x84\x00\x06\x00\x01\x00\x02\x00\x03', 2)
 
+  def test_unpack_for_new_link(self):
+    # TODO: we need to support more cell types before we can test this
+    self.assertRaisesRegexp(NotImplementedError, 'Unpacking not yet implemented for CERTS cells', Cell.unpack, test_data('new_link_cells'), 2)
+
   def test_versions_pack(self):
     self.assertEqual('\x00\x00\x07\x00\x00', VersionsCell.pack([]))
     self.assertEqual('\x00\x00\x07\x00\x02\x00\x01', VersionsCell.pack([1]))
     self.assertEqual('\x00\x00\x07\x00\x06\x00\x01\x00\x02\x00\x03', VersionsCell.pack([1, 2, 3]))
 
   def test_versions_unpack(self):
-    self.assertEqual([], Cell.unpack('\x00\x00\x07\x00\x00', 2).versions)
-    self.assertEqual([1], Cell.unpack('\x00\x00\x07\x00\x02\x00\x01', 2).versions)
-    self.assertEqual([1, 2, 3], Cell.unpack('\x00\x00\x07\x00\x06\x00\x01\x00\x02\x00\x03', 2).versions)
+    self.assertEqual([], Cell.unpack('\x00\x00\x07\x00\x00', 2)[0].versions)
+    self.assertEqual([1], Cell.unpack('\x00\x00\x07\x00\x02\x00\x01', 2)[0].versions)
+    self.assertEqual([1, 2, 3], Cell.unpack('\x00\x00\x07\x00\x06\x00\x01\x00\x02\x00\x03', 2)[0].versions)
diff --git a/test/unit/client/data/new_link_cells b/test/unit/client/data/new_link_cells
new file mode 100644
index 00000000..38e68613
Binary files /dev/null and b/test/unit/client/data/new_link_cells differ