Repository: abhinavsingh/proxy.py
Pull request #220 (2019-12-03 03:02:02)

Response parser now reaches COMPLETE even when no body is expected
Previously it was callee's responsibility to mark state complete on stream close.

codecov[bot]: # [Codecov](https://codecov.io/gh/abhinavsingh/proxy.py/pull/220?src=pr&el=h1) Report
> Merging [#220](https://codecov.io/gh/abhinavsingh/proxy.py/pull/220?src=pr&el=desc) into [develop](https://codecov.io/gh/abhinavsingh/proxy.py/commit/269484df2e89bc659124177d339d4fc59f280cba?src=pr&el=desc) will **decrease** coverage by `0.11%`.
> The diff coverage is `78.94%`.

[![Impacted file tree graph](https://codecov.io/gh/abhinavsingh/proxy.py/pull/220/graphs/tree.svg?width=650&token=Zh9J7b4la2&height=150&src=pr)](https://codecov.io/gh/abhinavsingh/proxy.py/pull/220?src=pr&el=tree)

```diff
@@             Coverage Diff             @@
##           develop     #220      +/-   ##
===========================================
- Coverage    73.45%   73.33%   -0.12%     
===========================================
  Files           64       69       +5     
  Lines         2727     2753      +26     
===========================================
+ Hits          2003     2019      +16     
- Misses         724      734      +10
```


| [Impacted Files](https://codecov.io/gh/abhinavsingh/proxy.py/pull/220?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [proxy/core/ssh/\_\_init\_\_.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/220/diff?src=pr&el=tree#diff-cHJveHkvY29yZS9zc2gvX19pbml0X18ucHk=) | `0% <0%> (ø)` | |
| [proxy/core/ssh/tunnel.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/220/diff?src=pr&el=tree#diff-cHJveHkvY29yZS9zc2gvdHVubmVsLnB5) | `0% <0%> (ø)` | |
| [proxy/core/ssh/client.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/220/diff?src=pr&el=tree#diff-cHJveHkvY29yZS9zc2gvY2xpZW50LnB5) | `0% <0%> (ø)` | |
| [proxy/core/connection/connection.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/220/diff?src=pr&el=tree#diff-cHJveHkvY29yZS9jb25uZWN0aW9uL2Nvbm5lY3Rpb24ucHk=) | `97.87% <100%> (ø)` | |
| [proxy/core/connection/\_\_init\_\_.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/220/diff?src=pr&el=tree#diff-cHJveHkvY29yZS9jb25uZWN0aW9uL19faW5pdF9fLnB5) | `100% <100%> (ø)` | |
| [proxy/http/parser.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/220/diff?src=pr&el=tree#diff-cHJveHkvaHR0cC9wYXJzZXIucHk=) | `95.45% <100%> (-0.6%)` | :arrow_down: |
| [proxy/core/connection/client.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/220/diff?src=pr&el=tree#diff-cHJveHkvY29yZS9jb25uZWN0aW9uL2NsaWVudC5weQ==) | `100% <100%> (ø)` | |
| [proxy/core/connection/server.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/220/diff?src=pr&el=tree#diff-cHJveHkvY29yZS9jb25uZWN0aW9uL3NlcnZlci5weQ==) | `100% <100%> (ø)` | |
| ... and [4 more](https://codecov.io/gh/abhinavsingh/proxy.py/pull/220/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/abhinavsingh/proxy.py/pull/220?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/abhinavsingh/proxy.py/pull/220?src=pr&el=footer). Last update [269484d...ca60849](https://codecov.io/gh/abhinavsingh/proxy.py/pull/220?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/proxy/core/connection/__init__.py b/proxy/core/connection/__init__.py
new file mode 100644
index 0000000..ee44bc1
--- /dev/null
+++ b/proxy/core/connection/__init__.py
@@ -0,0 +1,21 @@
+# -*- coding: utf-8 -*-
+"""
+    proxy.py
+    ~~~~~~~~
+    ⚡⚡⚡ Fast, Lightweight, Pluggable, TLS interception capable proxy server focused on
+    Network monitoring, controls & Application development, testing, debugging.
+
+    :copyright: (c) 2013-present by Abhinav Singh and contributors.
+    :license: BSD, see LICENSE for more details.
+"""
+from .connection import TcpConnection, TcpConnectionUninitializedException, tcpConnectionTypes
+from .client import TcpClientConnection
+from .server import TcpServerConnection
+
+__all__ = [
+    'TcpConnection',
+    'TcpConnectionUninitializedException',
+    'TcpServerConnection',
+    'TcpClientConnection',
+    'tcpConnectionTypes',
+]
diff --git a/proxy/core/connection/client.py b/proxy/core/connection/client.py
new file mode 100644
index 0000000..28995a5
--- /dev/null
+++ b/proxy/core/connection/client.py
@@ -0,0 +1,32 @@
+# -*- coding: utf-8 -*-
+"""
+    proxy.py
+    ~~~~~~~~
+    ⚡⚡⚡ Fast, Lightweight, Pluggable, TLS interception capable proxy server focused on
+    Network monitoring, controls & Application development, testing, debugging.
+
+    :copyright: (c) 2013-present by Abhinav Singh and contributors.
+    :license: BSD, see LICENSE for more details.
+"""
+import socket
+import ssl
+from typing import Union, Tuple, Optional
+
+from .connection import TcpConnection, tcpConnectionTypes, TcpConnectionUninitializedException
+
+
+class TcpClientConnection(TcpConnection):
+    """An accepted client connection request."""
+
+    def __init__(self,
+                 conn: Union[ssl.SSLSocket, socket.socket],
+                 addr: Tuple[str, int]):
+        super().__init__(tcpConnectionTypes.CLIENT)
+        self._conn: Optional[Union[ssl.SSLSocket, socket.socket]] = conn
+        self.addr: Tuple[str, int] = addr
+
+    @property
+    def connection(self) -> Union[ssl.SSLSocket, socket.socket]:
+        if self._conn is None:
+            raise TcpConnectionUninitializedException()
+        return self._conn
diff --git a/proxy/core/connection.py b/proxy/core/connection/connection.py
similarity index 65%
rename from proxy/core/connection.py
rename to proxy/core/connection/connection.py
index 7b5fe6c..29c15c1 100644
--- a/proxy/core/connection.py
+++ b/proxy/core/connection/connection.py
@@ -12,10 +12,9 @@ import socket
 import ssl
 import logging
 from abc import ABC, abstractmethod
-from typing import NamedTuple, Optional, Union, Tuple, List
+from typing import NamedTuple, Optional, Union, List
 
-from ..common.constants import DEFAULT_BUFFER_SIZE
-from ..common.utils import new_socket_connection
+from ...common.constants import DEFAULT_BUFFER_SIZE
 
 logger = logging.getLogger(__name__)
 
@@ -91,40 +90,3 @@ class TcpConnection(ABC):
             self.buffer[0] = memoryview(mv.tobytes()[sent:])
         logger.debug('flushed %d bytes to %s' % (sent, self.tag))
         return sent
-
-
-class TcpServerConnection(TcpConnection):
-    """Establishes connection to upstream server."""
-
-    def __init__(self, host: str, port: int):
-        super().__init__(tcpConnectionTypes.SERVER)
-        self._conn: Optional[Union[ssl.SSLSocket, socket.socket]] = None
-        self.addr: Tuple[str, int] = (host, int(port))
-
-    @property
-    def connection(self) -> Union[ssl.SSLSocket, socket.socket]:
-        if self._conn is None:
-            raise TcpConnectionUninitializedException()
-        return self._conn
-
-    def connect(self) -> None:
-        if self._conn is not None:
-            return
-        self._conn = new_socket_connection(self.addr)
-
-
-class TcpClientConnection(TcpConnection):
-    """An accepted client connection request."""
-
-    def __init__(self,
-                 conn: Union[ssl.SSLSocket, socket.socket],
-                 addr: Tuple[str, int]):
-        super().__init__(tcpConnectionTypes.CLIENT)
-        self._conn: Optional[Union[ssl.SSLSocket, socket.socket]] = conn
-        self.addr: Tuple[str, int] = addr
-
-    @property
-    def connection(self) -> Union[ssl.SSLSocket, socket.socket]:
-        if self._conn is None:
-            raise TcpConnectionUninitializedException()
-        return self._conn
diff --git a/proxy/core/connection/server.py b/proxy/core/connection/server.py
new file mode 100644
index 0000000..cbb9806
--- /dev/null
+++ b/proxy/core/connection/server.py
@@ -0,0 +1,36 @@
+# -*- coding: utf-8 -*-
+"""
+    proxy.py
+    ~~~~~~~~
+    ⚡⚡⚡ Fast, Lightweight, Pluggable, TLS interception capable proxy server focused on
+    Network monitoring, controls & Application development, testing, debugging.
+
+    :copyright: (c) 2013-present by Abhinav Singh and contributors.
+    :license: BSD, see LICENSE for more details.
+"""
+import socket
+import ssl
+from typing import Optional, Union, Tuple
+
+from .connection import TcpConnection, tcpConnectionTypes, TcpConnectionUninitializedException
+from ...common.utils import new_socket_connection
+
+
+class TcpServerConnection(TcpConnection):
+    """Establishes connection to upstream server."""
+
+    def __init__(self, host: str, port: int):
+        super().__init__(tcpConnectionTypes.SERVER)
+        self._conn: Optional[Union[ssl.SSLSocket, socket.socket]] = None
+        self.addr: Tuple[str, int] = (host, int(port))
+
+    @property
+    def connection(self) -> Union[ssl.SSLSocket, socket.socket]:
+        if self._conn is None:
+            raise TcpConnectionUninitializedException()
+        return self._conn
+
+    def connect(self) -> None:
+        if self._conn is not None:
+            return
+        self._conn = new_socket_connection(self.addr)
diff --git a/proxy/core/ssh/__init__.py b/proxy/core/ssh/__init__.py
new file mode 100644
index 0000000..232621f
--- /dev/null
+++ b/proxy/core/ssh/__init__.py
@@ -0,0 +1,10 @@
+# -*- coding: utf-8 -*-
+"""
+    proxy.py
+    ~~~~~~~~
+    ⚡⚡⚡ Fast, Lightweight, Pluggable, TLS interception capable proxy server focused on
+    Network monitoring, controls & Application development, testing, debugging.
+
+    :copyright: (c) 2013-present by Abhinav Singh and contributors.
+    :license: BSD, see LICENSE for more details.
+"""
diff --git a/proxy/core/ssh/client.py b/proxy/core/ssh/client.py
new file mode 100644
index 0000000..650d894
--- /dev/null
+++ b/proxy/core/ssh/client.py
@@ -0,0 +1,28 @@
+# -*- coding: utf-8 -*-
+"""
+    proxy.py
+    ~~~~~~~~
+    ⚡⚡⚡ Fast, Lightweight, Pluggable, TLS interception capable proxy server focused on
+    Network monitoring, controls & Application development, testing, debugging.
+
+    :copyright: (c) 2013-present by Abhinav Singh and contributors.
+    :license: BSD, see LICENSE for more details.
+"""
+import socket
+import ssl
+from typing import Union
+
+from ..connection import TcpClientConnection
+
+
+class SshClient(TcpClientConnection):
+    """Overrides TcpClientConnection.
+
+    This is necessary because paramiko fileno() can be used for polling
+    but not for send / recv.
+    """
+
+    @property
+    def connection(self) -> Union[ssl.SSLSocket, socket.socket]:
+        # Dummy return to comply with
+        return socket.socket()
diff --git a/proxy/core/tunnel.py b/proxy/core/ssh/tunnel.py
similarity index 89%
rename from proxy/core/tunnel.py
rename to proxy/core/ssh/tunnel.py
index 612fb52..e3a61b5 100644
--- a/proxy/core/tunnel.py
+++ b/proxy/core/ssh/tunnel.py
@@ -8,7 +8,6 @@
     :copyright: (c) 2013-present by Abhinav Singh and contributors.
     :license: BSD, see LICENSE for more details.
 """
-import socket
 from typing import Tuple, Callable
 
 import paramiko
@@ -26,7 +25,7 @@ class Tunnel:
             remote_addr: Tuple[str, int],
             private_pem_key: str,
             remote_proxy_port: int,
-            conn_handler: Callable[[socket.socket], None]) -> None:
+            conn_handler: Callable[[paramiko.channel.Channel], None]) -> None:
         self.remote_addr = remote_addr
         self.ssh_username = ssh_username
         self.private_pem_key = private_pem_key
@@ -45,11 +44,11 @@ class Tunnel:
                 key_filename=self.private_pem_key
             )
             print('SSH connection established...')
-            transport = ssh.get_transport()
+            transport: paramiko.transport.Transport = ssh.get_transport()
             transport.request_port_forward('', self.remote_proxy_port)
             print('Tunnel port forward setup successful...')
             while True:
-                conn: socket.socket = transport.accept(timeout=1)
+                conn: paramiko.channel.Channel = transport.accept(timeout=1)
                 e = transport.get_exception()
                 if e:
                     raise e
diff --git a/proxy/http/parser.py b/proxy/http/parser.py
index 80945b5..ece1ce0 100644
--- a/proxy/http/parser.py
+++ b/proxy/http/parser.py
@@ -126,6 +126,11 @@ class HttpParser:
         return b'transfer-encoding' in self.headers and \
                self.headers[b'transfer-encoding'][1].lower() == b'chunked'
 
+    def body_expected(self) -> bool:
+        return (b'content-length' in self.headers and
+                int(self.header(b'content-length')) > 0) or \
+            self.is_chunked_encoded()
+
     def parse(self, raw: bytes) -> None:
         """Parses Http request out of raw bytes.
 
@@ -134,7 +139,6 @@ class HttpParser:
         raw = self.buffer + raw
         self.buffer = b''
 
-        # TODO(abhinavsingh): Someday clean this up.
         more = True if len(raw) > 0 else False
         while more and self.state != httpParserStates.COMPLETE:
             if self.state in (
@@ -159,6 +163,8 @@ class HttpParser:
                         self.body = self.chunk_parser.body
                         self.state = httpParserStates.COMPLETE
                     more = False
+                else:
+                    raise NotImplementedError('Parser shouldn\'t have reached here')
             else:
                 more, raw = self.process(raw)
         self.buffer = raw
@@ -181,33 +187,14 @@ class HttpParser:
             else:
                 self.process_header(line)
 
-        # TODO(abhinavsingh): Can these be generalized instead of per-case handling?
         # When server sends a response line without any header or body e.g.
         # HTTP/1.1 200 Connection established\r\n\r\n
         if self.state == httpParserStates.LINE_RCVD and \
                 self.type == httpParserTypes.RESPONSE_PARSER and \
                 raw == CRLF:
             self.state = httpParserStates.COMPLETE
-        # When connect request is received without a following host header
-        # See
-        # `TestHttpParser.test_connect_request_without_host_header_request_parse`
-        # for details
-        #
-        # When raw request has ended with \r\n\r\n and no more http headers are expected
-        # See `TestHttpParser.test_request_parse_without_content_length` and
-        # `TestHttpParser.test_response_parse_without_content_length` for details
-        elif self.state == httpParserStates.HEADERS_COMPLETE and \
-                self.type == httpParserTypes.REQUEST_PARSER and \
-                self.method != httpMethods.POST and \
-                raw == b'':
-            self.state = httpParserStates.COMPLETE
         elif self.state == httpParserStates.HEADERS_COMPLETE and \
-                self.type == httpParserTypes.REQUEST_PARSER and \
-                self.method == httpMethods.POST and \
-                not self.is_chunked_encoded() and \
-                (b'content-length' not in self.headers or
-                 (b'content-length' in self.headers and
-                  int(self.headers[b'content-length'][1]) == 0)) and \
+                not self.body_expected() and \
                 raw == b'':
             self.state = httpParserStates.COMPLETE
 


Test patch:
diff --git a/tests/core/test_connection.py b/tests/core/test_connection.py
index 6808a46..3cd63ad 100644
--- a/tests/core/test_connection.py
+++ b/tests/core/test_connection.py
@@ -73,7 +73,7 @@ class TestTcpConnection(unittest.TestCase):
         mock_socket.return_value.connect.assert_called_with(
             (str(DEFAULT_IPV6_HOSTNAME), DEFAULT_PORT, 0, 0))
 
-    @mock.patch('proxy.core.connection.new_socket_connection')
+    @mock.patch('proxy.core.connection.server.new_socket_connection')
     def testTcpServerIgnoresDoubleConnectSilently(
             self,
             mock_new_socket_connection: mock.Mock) -> None:
@@ -93,7 +93,7 @@ class TestTcpConnection(unittest.TestCase):
         mock_socket.return_value.connect.assert_called_with(
             (str(DEFAULT_IPV4_HOSTNAME), DEFAULT_PORT))
 
-    @mock.patch('proxy.core.connection.new_socket_connection')
+    @mock.patch('proxy.core.connection.server.new_socket_connection')
     def testTcpServerConnectionProperty(
             self,
             mock_new_socket_connection: mock.Mock) -> None:
diff --git a/tests/http/test_http_parser.py b/tests/http/test_http_parser.py
index 623af08..ee13621 100644
--- a/tests/http/test_http_parser.py
+++ b/tests/http/test_http_parser.py
@@ -329,7 +329,12 @@ class TestHttpParser(unittest.TestCase):
         and it is responsibility of callee to change state to httpParserStates.COMPLETE
         when server stream closes.
 
-        See https://github.com/abhinavsingh/py/issues/20 for details.
+        See https://github.com/abhinavsingh/proxy.py/issues/20 for details.
+
+        Post commit https://github.com/abhinavsingh/proxy.py/commit/269484df2e89bc659124177d339d4fc59f280cba
+        HttpParser would reach state COMPLETE also for RESPONSE_PARSER types and no longer
+        it is callee responsibility to change state on stream close.  This was important because
+        pipelined responses not trigger stream close but may receive multiple responses.
         """
         self.parser.type = httpParserTypes.RESPONSE_PARSER
         self.parser.parse(b'HTTP/1.0 200 OK' + CRLF)
@@ -343,7 +348,7 @@ class TestHttpParser(unittest.TestCase):
         ]))
         self.assertEqual(
             self.parser.state,
-            httpParserStates.HEADERS_COMPLETE)
+            httpParserStates.COMPLETE)
 
     def test_response_parse(self) -> None:
         self.parser.type = httpParserTypes.RESPONSE_PARSER
@@ -513,3 +518,12 @@ class TestHttpParser(unittest.TestCase):
             httpMethods.GET, b'/', protocol_version=b'HTTP/1.0',
         ))
         self.assertFalse(self.parser.is_http_1_1_keep_alive())
+
+    def test_paramiko_doc(self) -> None:
+        response = b'HTTP/1.1 304 Not Modified\r\nDate: Tue, 03 Dec 2019 02:31:55 GMT\r\nConnection: keep-alive' \
+                   b'\r\nLast-Modified: Sun, 23 Jun 2019 22:58:21 GMT\r\nETag: "5d10040d-1af2c"' \
+                   b'\r\nX-Cname-TryFiles: True\r\nX-Served: Nginx\r\nX-Deity: web02\r\nCF-Cache-Status: DYNAMIC' \
+                   b'\r\nServer: cloudflare\r\nCF-RAY: 53f2208c6fef6c38-SJC\r\n\r\n'
+        self.parser = HttpParser(httpParserTypes.RESPONSE_PARSER)
+        self.parser.parse(response)
+        self.assertEqual(self.parser.state, httpParserStates.COMPLETE)