Propagate errors from callbacks correctly; Fix upgrade (#3)

Author: 1st1

httptools/parser/parser.pyx

@@ -34,6 +34,8 @@ cdef class HttpParser:
         _proto_on_message_complete, _proto_on_chunk_header, \
         _proto_on_chunk_complete, _proto_on_message_begin
 
+        object _last_error
+
         Py_buffer py_buf
 
     def __cinit__(self):
@@ -90,6 +92,8 @@ cdef class HttpParser:
             protocol, 'on_chunk_complete', None)
         self._csettings.on_chunk_complete = cb_on_chunk_complete
 
+        self._last_error = None
+
     cdef _maybe_call_on_header(self):
         if self._current_header_name is not None:
             current_header_name = self._current_header_name
@@ -157,14 +161,18 @@ cdef class HttpParser:
 
         PyBuffer_Release(&self.py_buf)
 
+        if self._cparser.http_errno != cparser.HPE_OK:
+            ex =  parser_error_from_errno(
+                <cparser.http_errno> self._cparser.http_errno)
+            if isinstance(ex, HttpParserCallbackError):
+                if self._last_error is not None:
+                    ex.__context__ = self._last_error
+                    self._last_error = None
+            raise ex
+
         if self._cparser.upgrade:
             raise HttpParserUpgrade(nb)
 
-        # TODO: Handle parser->upgrade
-
-        if self._cparser.http_errno != cparser.HPE_OK:
-            raise parser_error_from_errno(
-                <cparser.http_errno> self._cparser.http_errno)
         if nb != data_len:
             raise HttpParserError('not all of the data was parsed')
 
@@ -199,57 +207,120 @@ cdef class HttpResponseParser(HttpParser):
 
 cdef int cb_on_message_begin(cparser.http_parser* parser) except -1:
     cdef HttpParser pyparser = <HttpParser>parser.data
-    pyparser._proto_on_message_begin()
+    try:
+        pyparser._proto_on_message_begin()
+    except BaseException as ex:
+        pyparser._last_error = ex
+        return -1
+    else:
+        return 0
 
 
 cdef int cb_on_url(cparser.http_parser* parser,
                    const char *at, size_t length) except -1:
     cdef HttpParser pyparser = <HttpParser>parser.data
-    pyparser._proto_on_url(at[:length])
+    try:
+        pyparser._proto_on_url(at[:length])
+    except BaseException as ex:
+        pyparser._last_error = ex
+        return -1
+    else:
+        return 0
 
 
 cdef int cb_on_status(cparser.http_parser* parser,
                       const char *at, size_t length) except -1:
     cdef HttpParser pyparser = <HttpParser>parser.data
-    pyparser._proto_on_status(at[:length])
+    try:
+        pyparser._proto_on_status(at[:length])
+    except BaseException as ex:
+        pyparser._last_error = ex
+        return -1
+    else:
+        return 0
 
 
 cdef int cb_on_header_field(cparser.http_parser* parser,
                             const char *at, size_t length) except -1:
     cdef HttpParser pyparser = <HttpParser>parser.data
-    pyparser._on_header_field(at[:length])
+    try:
+        pyparser._on_header_field(at[:length])
+    except BaseException as ex:
+        pyparser._last_error = ex
+        return -1
+    else:
+        return 0
 
 
 cdef int cb_on_header_value(cparser.http_parser* parser,
                             const char *at, size_t length) except -1:
     cdef HttpParser pyparser = <HttpParser>parser.data
-    pyparser._on_header_value(at[:length])
+    try:
+        pyparser._on_header_value(at[:length])
+    except BaseException as ex:
+        pyparser._last_error = ex
+        return -1
+    else:
+        return 0
 
 
 cdef int cb_on_headers_complete(cparser.http_parser* parser) except -1:
     cdef HttpParser pyparser = <HttpParser>parser.data
-    pyparser._on_headers_complete()
+    try:
+        pyparser._on_headers_complete()
+    except BaseException as ex:
+        pyparser._last_error = ex
+        return -1
+    else:
+        if pyparser._cparser.upgrade:
+            return 1
+        else:
+            return 0
 
 
 cdef int cb_on_body(cparser.http_parser* parser,
                     const char *at, size_t length) except -1:
     cdef HttpParser pyparser = <HttpParser>parser.data
-    pyparser._proto_on_body(at[:length])
+    try:
+        pyparser._proto_on_body(at[:length])
+    except BaseException as ex:
+        pyparser._last_error = ex
+        return -1
+    else:
+        return 0
 
 
 cdef int cb_on_message_complete(cparser.http_parser* parser) except -1:
     cdef HttpParser pyparser = <HttpParser>parser.data
-    pyparser._proto_on_message_complete()
+    try:
+        pyparser._proto_on_message_complete()
+    except BaseException as ex:
+        pyparser._last_error = ex
+        return -1
+    else:
+        return 0
 
 
 cdef int cb_on_chunk_header(cparser.http_parser* parser) except -1:
     cdef HttpParser pyparser = <HttpParser>parser.data
-    pyparser._on_chunk_header()
+    try:
+        pyparser._on_chunk_header()
+    except BaseException as ex:
+        pyparser._last_error = ex
+        return -1
+    else:
+        return 0
 
 
 cdef int cb_on_chunk_complete(cparser.http_parser* parser) except -1:
     cdef HttpParser pyparser = <HttpParser>parser.data
-    pyparser._on_chunk_complete()
+    try:
+        pyparser._on_chunk_complete()
+    except BaseException as ex:
+        pyparser._last_error = ex
+        return -1
+    else:
+        return 0
 
 
 cdef parser_error_from_errno(cparser.http_errno errno):

tests/test_parser.py

@@ -57,6 +57,16 @@
 
 Hot diggity dogg'''
 
+UPGRADE_RESPONSE1 = b'''HTTP/1.1 101 Switching Protocols
+UPGRADE: websocket
+SEC-WEBSOCKET-ACCEPT: rVg+XakFNFOxk3ZH0lzrZBmg0aU=
+TRANSFER-ENCODING: chunked
+CONNECTION: upgrade
+DATE: Sat, 07 May 2016 23:44:32 GMT
+SERVER: Python/3.4 aiohttp/1.0.3
+
+data'''.replace(b'\n', b'\r\n')
+
 
 class TestResponseParser(unittest.TestCase):
 
@@ -67,7 +77,7 @@ def test_parser_response_1(self):
         m.on_header.side_effect = headers.__setitem__
 
         p = httptools.HttpResponseParser(m)
-        p.feed_data(RESPONSE1_HEAD)
+        p.feed_data(memoryview(RESPONSE1_HEAD))
 
         self.assertEqual(p.get_http_version(), '1.1')
         self.assertEqual(p.get_status_code(), 200)
@@ -92,9 +102,8 @@ def test_parser_response_1(self):
         self.assertFalse(m.on_chunk_complete.called)
 
         with self.assertRaisesRegex(
-            httptools.HttpParserError,
-            'data received after completed connection'):
-
+                httptools.HttpParserError,
+                'data received after completed connection'):
             p.feed_data(b'12123123')
 
     def test_parser_response_2(self):
@@ -136,6 +145,78 @@ def test_parser_response_5(self):
 
         m.on_message_complete.assert_called_once_with()
 
+    def test_parser_response_cb_on_status_1(self):
+        class Error(Exception):
+            pass
+
+        m = mock.Mock()
+        m.on_status.side_effect = Error()
+
+        p = httptools.HttpResponseParser(m)
+        try:
+            p.feed_data(RESPONSE1_HEAD + RESPONSE1_BODY)
+        except httptools.HttpParserCallbackError as ex:
+            self.assertIsInstance(ex.__context__, Error)
+        else:
+            self.fail('HttpParserCallbackError was not raised')
+
+    def test_parser_response_cb_on_body_1(self):
+        class Error(Exception):
+            pass
+
+        m = mock.Mock()
+        m.on_body.side_effect = Error()
+
+        p = httptools.HttpResponseParser(m)
+        try:
+            p.feed_data(RESPONSE1_HEAD + RESPONSE1_BODY)
+        except httptools.HttpParserCallbackError as ex:
+            self.assertIsInstance(ex.__context__, Error)
+        else:
+            self.fail('HttpParserCallbackError was not raised')
+
+    def test_parser_response_cb_on_message_complete_1(self):
+        class Error(Exception):
+            pass
+
+        m = mock.Mock()
+        m.on_message_complete.side_effect = Error()
+
+        p = httptools.HttpResponseParser(m)
+        try:
+            p.feed_data(RESPONSE1_HEAD + RESPONSE1_BODY)
+        except httptools.HttpParserCallbackError as ex:
+            self.assertIsInstance(ex.__context__, Error)
+        else:
+            self.fail('HttpParserCallbackError was not raised')
+
+    def test_parser_upgrade_response_1(self):
+        m = mock.Mock()
+
+        headers = {}
+        m.on_header.side_effect = headers.__setitem__
+
+        p = httptools.HttpResponseParser(m)
+        try:
+            p.feed_data(UPGRADE_RESPONSE1)
+        except httptools.HttpParserUpgrade as ex:
+            offset = ex.args[0]
+        else:
+            self.fail('HttpParserUpgrade was not raised')
+
+        self.assertEqual(UPGRADE_RESPONSE1[offset:], b'data')
+
+        self.assertEqual(p.get_http_version(), '1.1')
+        self.assertEqual(p.get_status_code(), 101)
+
+        m.on_status.assert_called_once_with(b'Switching Protocols')
+
+        m.on_headers_complete.assert_called_once_with()
+        self.assertEqual(m.on_header.call_count, 6)
+        self.assertEqual(len(headers), 6)
+
+        m.on_message_complete.assert_called_once_with()
+
 
 class TestRequestParser(unittest.TestCase):
 
@@ -194,6 +275,36 @@ def test_parser_request_chunked_2(self):
              b'Host': b'bar',
              b'Vary': b'*'})
 
+    def test_parser_request_chunked_cb_error_1(self):
+        class Error(Exception):
+            pass
+
+        m = mock.Mock()
+        m.on_chunk_header.side_effect = Error()
+
+        p = httptools.HttpRequestParser(m)
+        try:
+            p.feed_data(CHUNKED_REQUEST1_1)
+        except httptools.HttpParserCallbackError as ex:
+            self.assertIsInstance(ex.__context__, Error)
+        else:
+            self.fail('HttpParserCallbackError was not raised')
+
+    def test_parser_request_chunked_cb_error_2(self):
+        class Error(Exception):
+            pass
+
+        m = mock.Mock()
+        m.on_chunk_complete.side_effect = Error()
+
+        p = httptools.HttpRequestParser(m)
+        try:
+            p.feed_data(CHUNKED_REQUEST1_1)
+        except httptools.HttpParserCallbackError as ex:
+            self.assertIsInstance(ex.__context__, Error)
+        else:
+            self.fail('HttpParserCallbackError was not raised')
+
     def test_parser_request_chunked_3(self):
         m = mock.Mock()
         p = httptools.HttpRequestParser(m)
@@ -237,6 +348,62 @@ def test_parser_request_upgrade_1(self):
             b'Host': b'example.com',
             b'Upgrade': b'WebSocket'})
 
+    def test_parser_request_error_in_on_header(self):
+        class Error(Exception):
+            pass
+        m = mock.Mock()
+        m.on_header.side_effect = Error()
+        p = httptools.HttpRequestParser(m)
+
+        try:
+            p.feed_data(UPGRADE_REQUEST1)
+        except httptools.HttpParserCallbackError as ex:
+            self.assertIsInstance(ex.__context__, Error)
+        else:
+            self.fail('HttpParserCallbackError was not raised')
+
+    def test_parser_request_error_in_on_message_begin(self):
+        class Error(Exception):
+            pass
+        m = mock.Mock()
+        m.on_message_begin.side_effect = Error()
+        p = httptools.HttpRequestParser(m)
+
+        try:
+            p.feed_data(UPGRADE_REQUEST1)
+        except httptools.HttpParserCallbackError as ex:
+            self.assertIsInstance(ex.__context__, Error)
+        else:
+            self.fail('HttpParserCallbackError was not raised')
+
+    def test_parser_request_error_in_cb_on_url(self):
+        class Error(Exception):
+            pass
+        m = mock.Mock()
+        m.on_url.side_effect = Error()
+        p = httptools.HttpRequestParser(m)
+
+        try:
+            p.feed_data(UPGRADE_REQUEST1)
+        except httptools.HttpParserCallbackError as ex:
+            self.assertIsInstance(ex.__context__, Error)
+        else:
+            self.fail('HttpParserCallbackError was not raised')
+
+    def test_parser_request_error_in_cb_on_headers_complete(self):
+        class Error(Exception):
+            pass
+        m = mock.Mock()
+        m.on_headers_complete.side_effect = Error()
+        p = httptools.HttpRequestParser(m)
+
+        try:
+            p.feed_data(UPGRADE_REQUEST1)
+        except httptools.HttpParserCallbackError as ex:
+            self.assertIsInstance(ex.__context__, Error)
+        else:
+            self.fail('HttpParserCallbackError was not raised')
+
     def test_parser_request_2(self):
         p = httptools.HttpRequestParser(None)
         with self.assertRaises(httptools.HttpParserInvalidMethodError):