[tor-commits] [stem/master] Provide all cells circuits reply with

[atagar at torproject.org]

commit 6b5ff7cd17759fcf628237f9277503d36969d850
Author: Damian Johnson <atagar at torproject.org>
Date:   Tue Apr 24 11:44:38 2018 -0700

    Provide all cells circuits reply with
    
    Oops, callers of send() on our circuit only got the first cell sent in
    repsonse. These are fixed sized cells, so when the response is larger
    than 509 bytes (such as in the case of descriptors) multiple will be
    received.
    
    It's important that this returns a list rather than yielding a generator
    because we need to ensure we issue recv() calls to drain messages off the
    socket. Otherwise the next call will hang.
---
 stem/client/__init__.py        | 11 ++++++++---
 stem/descriptor/remote.py      |  2 +-
 test/unit/descriptor/remote.py |  9 ++++++++-
 3 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/stem/client/__init__.py b/stem/client/__init__.py
index 11a97171..aa4aa274 100644
--- a/stem/client/__init__.py
+++ b/stem/client/__init__.py
@@ -227,6 +227,8 @@ class Circuit(object):
     :param stem.client.RelayCommand command: command to be issued
     :param bytes data: message payload
     :param int stream_id: specific stream this concerns
+
+    :returns: **list** of :class:`~stem.client.cell.RelayCell` responses
     """
 
     with self.relay._orport_lock:
@@ -242,11 +244,14 @@ class Circuit(object):
         header, payload = split(cell.pack(self.relay.link_protocol), 3)
         encrypted_payload = header + self.forward_key.update(payload)
 
+        reply = []
         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(self.relay.link_protocol)[3:])
-        return stem.client.cell.RelayCell._unpack(decrypted, self.id, self.relay.link_protocol)
+        for cell in stem.client.cell.Cell.unpack(self.relay._orport.recv(), self.relay.link_protocol):
+          decrypted = self.backward_key.update(cell.pack(self.relay.link_protocol)[3:])
+          reply.append(stem.client.cell.RelayCell._unpack(decrypted, self.id, self.relay.link_protocol))
+
+        return reply
       except:
         self.forward_digest = orig_digest
         self.forward_key = orig_key
diff --git a/stem/descriptor/remote.py b/stem/descriptor/remote.py
index 5655c84c..5bd186c3 100644
--- a/stem/descriptor/remote.py
+++ b/stem/descriptor/remote.py
@@ -287,7 +287,7 @@ def _download_from_orport(endpoint, resource):
   with stem.client.Relay.connect(endpoint.address, endpoint.port, link_protocol) as relay:
     with relay.create_circuit() as circ:
       circ.send('RELAY_BEGIN_DIR', stream_id = 1)
-      lines = circ.send('RELAY_DATA', resource, stream_id = 1).data.splitlines()
+      lines = b''.join([cell.data for cell in circ.send('RELAY_DATA', resource, stream_id = 1)]).splitlines()
       first_line = lines.pop(0)
 
       if first_line != 'HTTP/1.0 200 OK':
diff --git a/test/unit/descriptor/remote.py b/test/unit/descriptor/remote.py
index 7d92c695..040d4c6c 100644
--- a/test/unit/descriptor/remote.py
+++ b/test/unit/descriptor/remote.py
@@ -130,10 +130,17 @@ ORPORT_DESCRIPTOR = 'HTTP/1.0 200 OK\n' + HEADER + '\n\n' + TEST_DESCRIPTOR
 
 
 def _orport_mock(data):
+  cells = []
+
+  for hunk in [data[i:i + 50] for i in range(0, len(data), 50)]:
+    cell = Mock()
+    cell.data = hunk
+    cells.append(cell)
+
   connect_mock = MagicMock()
   relay_mock = connect_mock().__enter__()
   circ_mock = relay_mock.create_circuit().__enter__()
-  circ_mock.send().data = data
+  circ_mock.send.return_value = cells
   return connect_mock