[tor-commits] [stem/master] Move cell decryption into circuit

[atagar at torproject.org]

commit 981760fd7996c849945a1b41e6c304d191267d26
Author: Damian Johnson <atagar at torproject.org>
Date:   Mon Feb 5 18:37:00 2018 -0800

    Move cell decryption into circuit
---
 stem/client/__init__.py |  4 ++++
 stem/client/cell.py     | 21 ++++++---------------
 2 files changed, 10 insertions(+), 15 deletions(-)

diff --git a/stem/client/__init__.py b/stem/client/__init__.py
index 4f9ed17c..8d34b626 100644
--- a/stem/client/__init__.py
+++ b/stem/client/__init__.py
@@ -234,6 +234,10 @@ class Circuit(object):
       encrypted_payload = header + self.forward_key.update(payload)
 
       self.relay._orport.send(encrypted_payload)
+      reply = next(stem.client.cell.Cell.unpack(self.relay._orport.recv(), self.relay.link_protocol))
+
+      decrypted = self.backward_key.update(reply.pack(3)[3:])
+      return stem.client.cell.RelayCell._unpack(decrypted, self.id, 3)
     except:
       self.forward_digest = orig_digest
       self.forward_key = orig_key
diff --git a/stem/client/cell.py b/stem/client/cell.py
index c4a19940..041e1c39 100644
--- a/stem/client/cell.py
+++ b/stem/client/cell.py
@@ -292,6 +292,7 @@ class RelayCell(CircuitCell):
   :var stem.client.RelayCommand command: command to be issued
   :var int command_int: integer value of our command
   :var bytes data: payload of the cell
+  :var int recognized: zero if endpoint is this hop, non-zero otherwise
   :var int digest: running digest held with the relay
   :var int stream_id: specific stream this concerns
   """
@@ -307,7 +308,7 @@ class RelayCell(CircuitCell):
   VALUE = 3
   IS_FIXED_SIZE = True
 
-  def __init__(self, circ_id, command, data, digest = 0, stream_id = 0, raw_content = None):
+  def __init__(self, circ_id, command, data, digest = 0, stream_id = 0, recognized = 0):
     if 'hashlib.HASH' in str(type(digest)):
       # Unfortunately hashlib generates from a dynamic private class so
       # isinstance() isn't such a great option.
@@ -323,9 +324,9 @@ class RelayCell(CircuitCell):
     super(RelayCell, self).__init__(circ_id)
     self.command, self.command_int = RelayCommand.get(command)
     self.data = data
+    self.recognized = recognized
     self.digest = digest
     self.stream_id = stream_id
-    self._raw_content = raw_content
 
     if not stream_id and self.command in STREAM_ID_REQUIRED:
       raise ValueError('%s relay cells require a stream id' % self.command)
@@ -335,7 +336,7 @@ class RelayCell(CircuitCell):
   def pack(self, link_protocol):
     payload = io.BytesIO()
     payload.write(Size.CHAR.pack(self.command_int))
-    payload.write(Size.SHORT.pack(0))  # 'recognized' field
+    payload.write(Size.SHORT.pack(self.recognized))
     payload.write(Size.SHORT.pack(self.stream_id))
     payload.write(Size.LONG.pack(self.digest))
     payload.write(Size.SHORT.pack(len(self.data)))
@@ -343,28 +344,18 @@ class RelayCell(CircuitCell):
 
     return RelayCell._pack(link_protocol, payload.getvalue(), self.circ_id)
 
-  def decrypt(self, circ):
-    # TODO: clearly funky, just a spot to start...
-
-    if not self._raw_content:
-      raise ValueError('Only received cells can be decrypted')
-
-    decrypted = circ.backward_key.update(self._raw_content)
-    return RelayCell._unpack(decrypted, self.circ_id, 3)
-
-
   @classmethod
   def _unpack(cls, content, circ_id, link_protocol):
     orig_content = content
 
     command, content = Size.CHAR.pop(content)
-    _, content = Size.SHORT.pop(content)  # 'recognized' field
+    recognized, content = Size.SHORT.pop(content)  # 'recognized' field
     stream_id, content = Size.SHORT.pop(content)
     digest, content = Size.LONG.pop(content)
     data_len, content = Size.SHORT.pop(content)
     data, content = split(content, data_len)
 
-    return RelayCell(circ_id, command, data, digest, stream_id, orig_content)
+    return RelayCell(circ_id, command, data, digest, stream_id, recognized)
 
   def __hash__(self):
     return _hash_attr(self, 'command_int', 'stream_id', 'digest', 'data')