[tor-commits] [stem/master] Unpack only a single cell

[atagar at torproject.org]

commit e548a715bc50c4c1b8a773d49b1c6d00442a1ea7
Author: Damian Johnson <atagar at torproject.org>
Date:   Tue Jan 16 09:46:44 2018 -0800

    Unpack only a single cell
    
    When first establishing a link the first cell we receive (VERSIONS) determines
    the link protocol version used to unpack all future cells. As such we
    definitely want an 'unpack just one cell' function like endosome.
---
 stem/client/cell.py      | 33 ++++++++++++++-------------------
 test/unit/client/cell.py | 16 ++++++++++------
 2 files changed, 24 insertions(+), 25 deletions(-)

diff --git a/stem/client/cell.py b/stem/client/cell.py
index 98ab198c..8a6480c7 100644
--- a/stem/client/cell.py
+++ b/stem/client/cell.py
@@ -95,39 +95,34 @@ class Cell(object):
   @staticmethod
   def unpack(content, link_version):
     """
-    Unpacks encoded bytes into a series of cells.
+    Unpacks the first cell.
 
     :param bytes content: payload to decode
     :param int link_version: link protocol version
 
-    :returns: **list** of :class:`~stem.client.cell.Cell` subclasses
+    :returns: (:class:`~stem.client.cell.Cell`, remainder) tuple
 
     :raises:
       * ValueError if content is malformed
       * NotImplementedError if unable to unpack this cell type
     """
 
-    cells = []
+    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)
 
-    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)))
+    if cls.IS_FIXED_SIZE:
+      payload_len = FIXED_PAYLOAD_LEN
+    else:
+      payload_len, content = Size.SHORT.pop(content)
 
-      payload = content[:payload_len]
-      content = content[payload_len:]
+    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)))
 
-      cells.append(cls._unpack(payload, link_version, circ_id))
+    payload = content[:payload_len]
+    content = content[payload_len:]
 
-    return cells
+    return cls._unpack(payload, link_version, circ_id), content
 
   @classmethod
   def _pack(cls, link_version, payload, circ_id = 0):
diff --git a/test/unit/client/cell.py b/test/unit/client/cell.py
index cf7eb684..a0eab2e5 100644
--- a/test/unit/client/cell.py
+++ b/test/unit/client/cell.py
@@ -86,11 +86,12 @@ class TestCell(unittest.TestCase):
       (7, '\x1a\xa5\xb3\xbd\x88\xb1C'),
     )
 
-    link_cells = Cell.unpack(test_data('new_link_cells'), 2)
-    self.assertEqual(4, len(link_cells))
-    self.assertEqual(VersionsCell([3, 4, 5]), link_cells[0])
+    content = test_data('new_link_cells')
 
-    certs_cell = link_cells[1]
+    version_cell, content = Cell.unpack(content, 2)
+    self.assertEqual(VersionsCell([3, 4, 5]), version_cell)
+
+    certs_cell, content = Cell.unpack(content, 2)
     self.assertEqual(CertsCell, type(certs_cell))
     self.assertEqual(len(expected_certs), len(certs_cell.certificates))
 
@@ -98,14 +99,17 @@ class TestCell(unittest.TestCase):
       self.assertEqual(cert_type, certs_cell.certificates[i].type)
       self.assertTrue(certs_cell.certificates[i].value.startswith(cert_prefix))
 
-    self.assertEqual(AuthChallengeCell('\x89Y\t\x99\xb2\x1e\xd9*V\xb6\x1bn\n\x05\xd8/\xe3QH\x85\x13Z\x17\xfc\x1c\x00{\xa9\xae\x83^K', [1, 3]), link_cells[2])
+    auth_challenge_cell, content = Cell.unpack(content, 2)
+    self.assertEqual(AuthChallengeCell('\x89Y\t\x99\xb2\x1e\xd9*V\xb6\x1bn\n\x05\xd8/\xe3QH\x85\x13Z\x17\xfc\x1c\x00{\xa9\xae\x83^K', [1, 3]), auth_challenge_cell)
 
-    netinfo_cell = link_cells[3]
+    netinfo_cell, content = Cell.unpack(content, 2)
     self.assertEqual(NetinfoCell, type(netinfo_cell))
     self.assertEqual(datetime.datetime(2018, 1, 14, 1, 46, 56), netinfo_cell.timestamp)
     self.assertEqual(Address(type='IPv4', type_int=4, value='127.0.0.1', value_bin='\x7f\x00\x00\x01'), netinfo_cell.receiver_address)
     self.assertEqual([Address(type='IPv4', type_int=4, value='97.113.15.2', value_bin='aq\x0f\x02')], netinfo_cell.sender_addresses)
 
+    self.assertEqual('', content)  # check that we've consumed all of the bytes
+
   def test_padding_packing(self):
     for cell_bytes, payload in PADDING_CELLS.items():
       self.assertEqual(cell_bytes, PaddingCell.pack(2, payload))