commit 240b125b01bc7ab2e517852d2e7f5947a799d345 Author: Dave Rolek dmr-x@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):