[tor-commits] [stem/master] Add abstract BaseRelayCell class

[atagar at torproject.org]

commit 240b125b01bc7ab2e517852d2e7f5947a799d345
Author: Dave Rolek <dmr-x at riseup.net>
Date:   Fri Aug 3 14:56:31 2018 +0000

    Add abstract BaseRelayCell class
    
    It's intended to defined common Relay cell functionality, for encrypted
    and decrypted cells.
    
    By unpacking content no further than the payload, this abstraction
    should allow a lot of flexibility in handling Relay cells prior to
    decryption.
    
    Note that unpacking is not yet possible - subclasses must be defined
    with a VALUE. (Namely: RELAY and RELAY_EARLY)
---
 stem/client/cell.py      | 36 ++++++++++++++++++++++++++++++++++++
 test/unit/client/cell.py | 31 +++++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+)

diff --git a/stem/client/cell.py b/stem/client/cell.py
index 12fa994c..a4fb1a67 100644
--- a/stem/client/cell.py
+++ b/stem/client/cell.py
@@ -14,6 +14,7 @@ Messages communicated over a Tor relay's ORPort.
     |- CircuitCell - Circuit management.
     |  |- CreateCell - Create a circuit.              (section 5.1)
     |  |- CreatedCell - Acknowledge create.           (section 5.1)
+    |  |- BaseRelayCell - End-to-end data; abstract.  (section 6.1)
     |  |- RelayCell - End-to-end data.                (section 6.1)
     |  |- DestroyCell - Stop using a circuit.         (section 5.4)
     |  |- CreateFastCell - Create a circuit, no PK.   (section 5.1)
@@ -83,6 +84,7 @@ class Cell(object):
   The following cell types explicitly don't have *unused* content:
     * PaddingCell (we consider all content part of payload)
     * VersionsCell (all content is unpacked and treated as a version specification)
+    * BaseRelayCell (we don't parse cell beyond header/body)
     * VPaddingCell (we consider all content part of payload)
 
   :var bytes unused: unused filler that padded the cell to the expected size
@@ -320,6 +322,40 @@ class CreatedCell(CircuitCell):
     super(CreatedCell, self).__init__()  # TODO: implement
 
 
+class BaseRelayCell(CircuitCell):
+  """
+  Cell whose subclasses are relayed over circuits.
+
+  :var bytes payload: raw payload, quite possibly encrypted
+  """
+
+  NAME = 'INTERNAL_BASE_RELAY'  # defined for error/other strings
+  IS_FIXED_SIZE = True  # all relay cells are fixed-size
+
+  # other attributes are deferred to subclasses, since this class cannot be directly unpacked
+
+  def __init__(self, circ_id, payload):
+    if not payload:
+      raise ValueError('Relay cells require a payload')
+    if len(payload) != FIXED_PAYLOAD_LEN:
+      raise ValueError('Payload should be %i bytes, but was %i' % (FIXED_PAYLOAD_LEN, len(payload)))
+
+    super(BaseRelayCell, self).__init__(circ_id, unused = b'')
+    self.payload = payload
+
+  def pack(self, link_protocol):
+    # unlike everywhere else, we actually want to use the subclass type, NOT *this* class
+    return type(self)._pack(link_protocol, self.payload, circ_id = self.circ_id)
+
+  @classmethod
+  def _unpack(cls, content, circ_id, link_protocol):
+    # unlike everywhere else, we actually want to use the subclass type, NOT *this* class
+    return cls(circ_id, content)
+
+  def __hash__(self):
+    return stem.util._hash_attr(self, 'circ_id', 'payload', cache = True)
+
+
 class RelayCell(CircuitCell):
   """
   Command concerning a relay circuit.
diff --git a/test/unit/client/cell.py b/test/unit/client/cell.py
index ce492638..278e0b4c 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 datetime
 import hashlib
 import os
+import struct
 import unittest
 
 from stem.client.datatype import ZERO, CertType, CloseReason, Address, Certificate
@@ -14,6 +15,7 @@ from stem.client.cell import (
   FIXED_PAYLOAD_LEN,
   Cell,
   PaddingCell,
+  BaseRelayCell,
   RelayCell,
   DestroyCell,
   CreateFastCell,
@@ -188,6 +190,35 @@ class TestCell(unittest.TestCase):
       self.assertEqual(b'', cell.unused)  # always empty
       self.assertEqual(cell_bytes, cell.pack(link_protocol))
 
+  def test_base_relay_cell(self):
+    arbitrary_circ_id = 123
+    even_more_arbitrary_link_protocol = 1234
+
+    cell = BaseRelayCell(arbitrary_circ_id, RANDOM_PAYLOAD)
+    self.assertEqual(RANDOM_PAYLOAD, cell.payload)
+    self.assertEqual(arbitrary_circ_id, cell.circ_id)
+    self.assertEqual(True, cell.IS_FIXED_SIZE)
+
+    # Cell.unpack not reachable - won't be tested
+    # but we can at least directly test _unpack, although it's a pretty simple method
+    cell_2 = BaseRelayCell._unpack(RANDOM_PAYLOAD, arbitrary_circ_id, even_more_arbitrary_link_protocol)
+    self.assertEqual(cell, cell_2)
+
+    # pack not possible, but easily callable
+    self.assertRaises(struct.error, cell.pack, even_more_arbitrary_link_protocol)
+
+    # check other values and inequality
+    for (circ_id, payload) in ((arbitrary_circ_id, ZERO * FIXED_PAYLOAD_LEN), (arbitrary_circ_id + 1, RANDOM_PAYLOAD)):
+      unequal_cell = BaseRelayCell(circ_id, payload)
+      self.assertEqual(payload, unequal_cell.payload)
+      self.assertNotEqual(cell, unequal_cell)
+
+    # invalid constructions
+    self.assertRaisesWith(ValueError, 'Relay cells require a payload', BaseRelayCell, arbitrary_circ_id, None)
+    expected_message_format = 'Payload should be %i bytes, but was ' % FIXED_PAYLOAD_LEN + '%i'
+    for payload_len in (FIXED_PAYLOAD_LEN - 1, FIXED_PAYLOAD_LEN + 1):
+      self.assertRaisesWith(ValueError, expected_message_format % payload_len, BaseRelayCell, arbitrary_circ_id, ZERO * payload_len)
+
   def test_relay_cell(self):
     for cell_bytes, (command, command_int, circ_id, stream_id, data, digest, unused, link_protocol) in RELAY_CELLS.items():
       if not unused.strip(ZERO):