[tor-commits] [stem/master] Accept string and hash digests

[atagar at torproject.org]

commit bdc96f01bf91c3f7758e12a58e5e388b88734ad9
Author: Damian Johnson <atagar at torproject.org>
Date:   Mon Feb 5 10:14:14 2018 -0800

    Accept string and hash digests
    
    The RELAY cells actually work from a truncated digest. Only the two ends know
    the full sha1 digest for the message series that has been sent. Adjusting the
    RelayCell so it takes all three forms that the digest comes in.
---
 stem/client/cell.py      | 12 ++++++++++++
 test/unit/client/cell.py |  8 ++++++++
 2 files changed, 20 insertions(+)

diff --git a/stem/client/cell.py b/stem/client/cell.py
index ed8932c4..17f27202 100644
--- a/stem/client/cell.py
+++ b/stem/client/cell.py
@@ -308,6 +308,18 @@ class RelayCell(CircuitCell):
   IS_FIXED_SIZE = True
 
   def __init__(self, circ_id, command, data, digest = 0, stream_id = 0):
+    if 'hashlib.HASH' in str(type(digest)):
+      # Unfortunately hashlib generates from a dynamic private class so
+      # isinstance() isn't such a great option.
+
+      digest = Size.LONG.unpack(digest.digest()[:4])
+    elif isinstance(digest, str):
+      digest = Size.LONG.unpack(digest[:4])
+    elif isinstance(digest, int):
+      pass
+    else:
+      raise ValueError('RELAY cell digest must be a hash, string, or int but was a %s' % type(digest).__name__)
+
     super(RelayCell, self).__init__(circ_id)
     self.command, self.command_int = RelayCommand.get(command)
     self.data = data
diff --git a/test/unit/client/cell.py b/test/unit/client/cell.py
index f53355dd..19744c6d 100644
--- a/test/unit/client/cell.py
+++ b/test/unit/client/cell.py
@@ -3,6 +3,7 @@ Unit tests for the stem.client.cell.
 """
 
 import datetime
+import hashlib
 import os
 import unittest
 
@@ -153,6 +154,13 @@ class TestCell(unittest.TestCase):
       self.assertEqual(digest, cell.digest)
       self.assertEqual(stream_id, cell.stream_id)
 
+    digest = hashlib.sha1('hi')
+    self.assertEqual(3257622417, RelayCell(5, 'RELAY_BEGIN_DIR', '', digest, 564346860).digest)
+    self.assertEqual(3257622417, RelayCell(5, 'RELAY_BEGIN_DIR', '', digest.digest(), 564346860).digest)
+    self.assertEqual(3257622417, RelayCell(5, 'RELAY_BEGIN_DIR', '', 3257622417, 564346860).digest)
+    self.assertRaisesRegexp(ValueError, 'RELAY cell digest must be a hash, string, or int but was a list', RelayCell, 5, 'RELAY_BEGIN_DIR', '', [], 564346860)
+    self.assertRaisesRegexp(ValueError, "Invalid enumeration 'NO_SUCH_COMMAND', options are RELAY_BEGIN, RELAY_DATA", RelayCell, 5, 'NO_SUCH_COMMAND', '', 5, 564346860)
+
   def test_destroy_cell(self):
     for cell_bytes, (circ_id, reason, reason_int) in DESTROY_CELLS.items():
       self.assertEqual(cell_bytes, DestroyCell(circ_id, reason).pack(5))