Response parser now reaches COMPLETE even when no body is expected

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',
+]

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

proxy/core/connection.py → proxy/core/connection/connection.py

@@ -12,10 +12,9 @@
 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 @@ def flush(self) -> int:
             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

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)

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.
+"""

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()

proxy/core/tunnel.py → 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 @@ def __init__(
             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 @@ def run(self) -> None:
                 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

proxy/http/parser.py

@@ -126,6 +126,11 @@ def is_chunked_encoded(self) -> bool:
         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 @@ def parse(self, raw: bytes) -> None:
         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 @@ def parse(self, raw: bytes) -> None:
                         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 @@ def process(self, raw: bytes) -> Tuple[bool, bytes]:
             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
 

tests/core/test_connection.py

@@ -73,7 +73,7 @@ def testTcpServerEstablishesIPv6Connection(
         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 @@ def testTcpServerEstablishesIPv4Connection(
         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:

tests/http/test_http_parser.py

@@ -329,7 +329,12 @@ def test_response_parse_without_content_length(self) -> None:
         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 @@ def test_response_parse_without_content_length(self) -> None:
         ]))
         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 @@ def test_is_not_http_1_1_keep_alive_for_http_1_0(self) -> None:
             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)