[tor-commits] [stem/master] Add check_digest instance method to BaseRelayCell

[atagar at torproject.org]

commit 69f293222f550978e603906a4b3e34cb371447b0
Author: Dave Rolek <dmr-x at riseup.net>
Date:   Sat Aug 18 02:40:01 2018 +0000

    Add check_digest instance method to BaseRelayCell
    
    Similar to how apply_digest() is used by a RELAY cell sender, this
    collects the logic needed for a receiver to confirm that a cell is fully
    decrypted and that the cell's integrity is intact.
    
    This is new functionality. Prior to this change, stem.client had no
    facility to check the digest of a RELAY cell according to the spec.
    
    Not yet used.
---
 stem/client/cell.py | 38 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/stem/client/cell.py b/stem/client/cell.py
index 43dc076c..5235331e 100644
--- a/stem/client/cell.py
+++ b/stem/client/cell.py
@@ -354,6 +354,44 @@ class BaseRelayCell(CircuitCell):
     # unlike everywhere else, we actually want to use the subclass type, NOT *this* class
     return cls(circ_id, content)
 
+  def check_digest(self, digest):
+    """
+    Calculates the running digest of the cell payload per the spec, returning
+    whether the cell's unpacked digest matched, along with the updated digest
+    if so.
+
+    :param HASH digest: running digest held with the relay
+
+    :returns: (digest_matches, digest) tuple of object copies updated as follows:
+      * digest_matches: **bool** indicating whether the digest matches
+      * digest: updated via digest.update(payload), if the digest matches;
+        otherwise a copy of the original
+
+    :raises: **ValueError** if payload is the wrong size
+    """
+
+    command, recognized, stream_id, digest_from_cell, data_len, data, unused = RelayCell._unpack_payload(self.payload)
+
+    # running digest is calculated using a zero'd digest field in the payload
+    prepared_payload = RelayCell._pack_payload(command, recognized, stream_id, 0, data_len, data, unused, pad_remainder = False)
+
+    if len(prepared_payload) != FIXED_PAYLOAD_LEN:
+      # this should never fail
+      # if it did, it indicates a programming error either within stem.client.cell or a consumer
+      raise ValueError('Payload should be %i bytes, but was %i' % (FIXED_PAYLOAD_LEN, len(prepared_payload)))
+
+    new_digest = digest.copy()
+    new_digest.update(prepared_payload)
+
+    digest_matches = (RelayCell._coerce_digest(new_digest) == digest_from_cell)
+
+    # only return the new_digest if the digest check passed
+    # even if not, return a copy of the original
+    # this allows a consumer to always assume the returned digest is a different object
+    digest_to_return = new_digest if digest_matches else digest.copy()
+
+    return digest_matches, digest_to_return
+
   def __hash__(self):
     return stem.util._hash_attr(self, 'circ_id', 'payload', cache = True)