[tor-commits] [ooni-probe/master] Better error handling in tcp test template

[art at torproject.org]

commit 880ddb4fdfb95006e4792350d752925b36474dcf
Author: Arturo Filastò <art at fuffa.org>
Date:   Thu Nov 22 11:21:58 2012 +0100

    Better error handling in tcp test template
---
 ooni/templates/tcpt.py |   33 +++++++++++++++++++++++++++++++--
 1 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/ooni/templates/tcpt.py b/ooni/templates/tcpt.py
index 3312cbe..e592274 100644
--- a/ooni/templates/tcpt.py
+++ b/ooni/templates/tcpt.py
@@ -7,10 +7,37 @@ from ooni.utils import log
 
 class TCPSender(protocol.Protocol):
     report = None
+    payload_len = None
+    received_data = ''
     def dataReceived(self, data):
-        self.report['received'].append(data)
+        """
+        We receive data until the total amount of data received reaches that
+        which we have sent. At that point we append the received data to the
+        report and we fire the callback of the test template sendPayload
+        function.
+
+        This is used in pair with a TCP Echo server.
+
+        The reason why we put the data received inside of an array is that in
+        future we may want to expand this to support state and do something
+        similar to what daphne does, but without the mutation.
+
+        XXX Actually daphne will probably be refactored to be a subclass of the
+        TCP Test Template.
+        """
+        if self.payload_len:
+            self.received_data += data
+            if len(self.received_data) >= self.payload_len:
+                self.transport.loseConnection()
+                self.report['received'].append(data)
+                self.deferred.callback(self.report['received'])
 
     def sendPayload(self, payload):
+        """
+        Write the payload to the wire and set the expected size of the payload
+        we are to receive.
+        """
+        self.payload_len = len(payload)
         self.report['sent'].append(payload)
         self.transport.write(payload)
 
@@ -37,15 +64,17 @@ class TCPTest(NetTestCase):
         def closeConnection(p):
             p.transport.loseConnection()
             log.debug("Closing connection")
-            d1.callback(None)
+            d1.callback(self.report['received'])
 
         def errback(failure):
             self.report['error'] = str(failure)
             log.exception(failure)
+            d1.callback(self.report['received'])
 
         def connected(p):
             log.debug("Connected to %s:%s" % (self.address, self.port))
             p.report = self.report
+            p.deferred = d1
             p.sendPayload(payload)
             reactor.callLater(self.timeout, closeConnection, p)