Repository: twisted/treq

Issue:
treq.collect() should close the HTTP transport when the collector function throws an exception
If the collector function throws an exception, that means there is no way to handle remaining body data. In some cases this may actually be a deliberate indication that the response is no good; for example, I implemented a length-constrained variant of `treq.content()`, and if the body exceeds that it's indication of a buggy or malicious server.

 As such, the only reasonable thing to do in this situation is to close the transport.

Comments:

Patch:
diff --git a/changelog.d/347.feature.rst b/changelog.d/347.feature.rst
new file mode 100644
index 0000000..05bff3b
--- /dev/null
+++ b/changelog.d/347.feature.rst
@@ -0,0 +1,1 @@
+When the collector passed to ``treq.collect(response, collector)`` throws an exception, that error will now be returned to the caller of ``collect()`` via the result ``Deferred``, and the underlying HTTP transport will be closed.
diff --git a/src/treq/content.py b/src/treq/content.py
index 36219f5..8982ce9 100644
--- a/src/treq/content.py
+++ b/src/treq/content.py
@@ -2,7 +2,7 @@ import cgi
 import json
 
 from twisted.internet.defer import Deferred, succeed
-
+from twisted.python.failure import Failure
 from twisted.internet.protocol import Protocol
 from twisted.web.client import ResponseDone
 from twisted.web.http import PotentialDataLoss
@@ -30,9 +30,16 @@ class _BodyCollector(Protocol):
         self.collector = collector
 
     def dataReceived(self, data):
-        self.collector(data)
+        try:
+            self.collector(data)
+        except BaseException:
+            self.transport.loseConnection()
+            self.finished.errback(Failure())
+            self.finished = None
 
     def connectionLost(self, reason):
+        if self.finished is None:
+            return
         if reason.check(ResponseDone):
             self.finished.callback(None)
         elif reason.check(PotentialDataLoss):
@@ -48,9 +55,13 @@ def collect(response, collector):
 
     This function may only be called **once** for a given response.
 
+    If the ``collector`` raises an exception, it will be set as the error value
+    on response ``Deferred`` returned from this function, and the underlying
+    HTTP transport will be closed.
+
     :param IResponse response: The HTTP response to collect the body from.
-    :param collector: A callable to be called each time data is available
-        from the response body.
+    :param collector: A callable to be called each time data is available from
+        the response body.
     :type collector: single argument callable
 
     :rtype: Deferred that fires with None when the entire body has been read.


Test patch:
diff --git a/src/treq/test/test_content.py b/src/treq/test/test_content.py
index 9acb9d7..0d83ddf 100644
--- a/src/treq/test/test_content.py
+++ b/src/treq/test/test_content.py
@@ -2,13 +2,17 @@ from unittest import mock
 
 from twisted.python.failure import Failure
 
+from twisted.internet.error import ConnectionDone
 from twisted.trial.unittest import TestCase
 from twisted.web.http_headers import Headers
 from twisted.web.client import ResponseDone, ResponseFailed
 from twisted.web.http import PotentialDataLoss
+from twisted.web.resource import Resource
+from twisted.web.server import NOT_DONE_YET
 
 from treq import collect, content, json_content, text_content
 from treq.client import _BufferedResponse
+from treq.testing import StubTreq
 
 
 class ContentTests(TestCase):
@@ -215,3 +219,54 @@ class ContentTests(TestCase):
         self.protocol.connectionLost(Failure(ResponseDone()))
 
         self.assertEqual(self.successResultOf(d), u'ᚠᚡ')
+
+
+class UnfinishedResponse(Resource):
+    """Write some data, but never finish."""
+
+    isLeaf = True
+
+    def __init__(self):
+        Resource.__init__(self)
+        # Track how requests finished.
+        self.request_finishes = []
+
+    def render(self, request):
+        request.write(b"HELLO")
+        request.notifyFinish().addBoth(self.request_finishes.append)
+        return NOT_DONE_YET
+
+
+class MoreRealisticContentTests(TestCase):
+    """Tests involving less mocking."""
+
+    def test_exception_handling(self):
+        """
+        An exception in the collector function:
+
+            1. Always gets returned in the result ``Deferred`` from
+               ``treq.collect()``.
+
+            2. Closes the transport.
+        """
+        resource = UnfinishedResponse()
+        stub = StubTreq(resource)
+        response = self.successResultOf(stub.request("GET", "http://127.0.0.1/"))
+        self.assertEqual(response.code, 200)
+
+        def error(data):
+            1 / 0
+
+        d = collect(response, error)
+
+        # Exceptions in the collector are passed on to the caller via the
+        # response Deferred:
+        self.failureResultOf(d, ZeroDivisionError)
+
+        # An exception in the protocol results in the transport for the request
+        # being closed.
+        stub.flush()
+        self.assertEqual(len(resource.request_finishes), 1)
+        self.assertIsInstance(
+            resource.request_finishes[0].value, ConnectionDone
+        )