Repository: abhinavsingh/proxy.py
Pull request #833 (2021-11-30 23:39:40)

[ProxyPoolPlugin] Avoid remote proxy of private IP requests
Additionally, as `on_client_data` is only invoked after initial client request has been received, refactor to avoid unnecessary invocations.

codecov[bot]: # [Codecov](https://codecov.io/gh/abhinavsingh/proxy.py/pull/833?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) Report
> Merging [#833](https://codecov.io/gh/abhinavsingh/proxy.py/pull/833?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) (fc5a0b6) into [develop](https://codecov.io/gh/abhinavsingh/proxy.py/commit/921f2b56e0370373357a56cf129c872f1feefd15?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) (921f2b5) will **decrease** coverage by `0.16%`.
> The diff coverage is `50.00%`.

[![Impacted file tree graph](https://codecov.io/gh/abhinavsingh/proxy.py/pull/833/graphs/tree.svg?width=650&height=150&src=pr&token=Zh9J7b4la2&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh)](https://codecov.io/gh/abhinavsingh/proxy.py/pull/833?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh)

```diff
@@             Coverage Diff             @@
##           develop     #833      +/-   ##
===========================================
- Coverage    87.05%   86.89%   -0.17%     
===========================================
  Files          128      128              
  Lines         5732     5738       +6     
  Branches       473      474       +1     
===========================================
- Hits          4990     4986       -4     
- Misses         637      646       +9     
- Partials       105      106       +1     
```

| Flag | Coverage Δ | |
|---|---|---|
| pytest | `86.77% <50.00%> (-0.17%)` | :arrow_down: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Impacted Files](https://codecov.io/gh/abhinavsingh/proxy.py/pull/833?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) | Coverage Δ | |
|---|---|---|
| [proxy/http/plugin.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/833/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvaHR0cC9wbHVnaW4ucHk=) | `100.00% <ø> (ø)` | |
| [proxy/plugin/proxy\_pool.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/833/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvcGx1Z2luL3Byb3h5X3Bvb2wucHk=) | `31.76% <28.57%> (-0.29%)` | :arrow_down: |
| [proxy/http/handler.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/833/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvaHR0cC9oYW5kbGVyLnB5) | `72.39% <50.00%> (-1.36%)` | :arrow_down: |
| [proxy/http/parser/parser.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/833/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvaHR0cC9wYXJzZXIvcGFyc2VyLnB5) | `94.27% <100.00%> (-0.03%)` | :arrow_down: |
| [tests/http/test\_http\_proxy\_tls\_interception.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/833/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-dGVzdHMvaHR0cC90ZXN0X2h0dHBfcHJveHlfdGxzX2ludGVyY2VwdGlvbi5weQ==) | `97.97% <100.00%> (-1.02%)` | :arrow_down: |
| [proxy/http/server/web.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/833/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvaHR0cC9zZXJ2ZXIvd2ViLnB5) | `63.35% <0.00%> (-0.63%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/abhinavsingh/proxy.py/pull/833?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/abhinavsingh/proxy.py/pull/833?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh). Last update [921f2b5...fc5a0b6](https://codecov.io/gh/abhinavsingh/proxy.py/pull/833?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh).


Patch:
diff --git a/proxy/http/handler.py b/proxy/http/handler.py
index 1256def..b6cd6f2 100644
--- a/proxy/http/handler.py
+++ b/proxy/http/handler.py
@@ -189,19 +189,12 @@ class HttpProtocolHandler(BaseTcpServerHandler):
         return False
 
     def handle_data(self, data: memoryview) -> Optional[bool]:
+        """Handles incoming data from client."""
         if data is None:
             logger.debug('Client closed connection, tearing down...')
             self.work.closed = True
             return True
-
         try:
-            # HttpProtocolHandlerPlugin.on_client_data
-            # Can raise HttpProtocolException to tear down the connection
-            for plugin in self.plugins.values():
-                optional_data = plugin.on_client_data(data)
-                if optional_data is None:
-                    break
-                data = optional_data
             # Don't parse incoming data any further after 1st request has completed.
             #
             # This specially does happen for pipeline requests.
@@ -209,7 +202,7 @@ class HttpProtocolHandler(BaseTcpServerHandler):
             # Plugins can utilize on_client_data for such cases and
             # apply custom logic to handle request data sent after 1st
             # valid request.
-            if data and self.request.state != httpParserStates.COMPLETE:
+            if self.request.state != httpParserStates.COMPLETE:
                 # Parse http request
                 #
                 # TODO(abhinavsingh): Remove .tobytes after parser is
@@ -229,6 +222,14 @@ class HttpProtocolHandler(BaseTcpServerHandler):
                                     plugin_.client._conn = upgraded_sock
                         elif isinstance(upgraded_sock, bool) and upgraded_sock is True:
                             return True
+            else:
+                # HttpProtocolHandlerPlugin.on_client_data
+                # Can raise HttpProtocolException to tear down the connection
+                for plugin in self.plugins.values():
+                    optional_data = plugin.on_client_data(data)
+                    if optional_data is None:
+                        break
+                    data = optional_data
         except HttpProtocolException as e:
             logger.debug('HttpProtocolException raised')
             response: Optional[memoryview] = e.response(self.request)
diff --git a/proxy/http/parser/parser.py b/proxy/http/parser/parser.py
index 9df644f..2785b24 100644
--- a/proxy/http/parser/parser.py
+++ b/proxy/http/parser/parser.py
@@ -15,7 +15,7 @@
 from typing import TypeVar, Optional, Dict, Type, Tuple, List
 
 from ...common.constants import DEFAULT_DISABLE_HEADERS, COLON, DEFAULT_ENABLE_PROXY_PROTOCOL
-from ...common.constants import HTTP_1_1, HTTP_1_0, SLASH, CRLF
+from ...common.constants import HTTP_1_1, SLASH, CRLF
 from ...common.constants import WHITESPACE, DEFAULT_HTTP_PORT
 from ...common.utils import build_http_request, build_http_response, find_http_line, text_
 from ...common.flag import flags
@@ -271,7 +271,7 @@ class HttpParser:
                 self.body = self.chunk.body
                 self.state = httpParserStates.COMPLETE
             more = False
-        elif b'content-length' in self.headers:
+        elif self.content_expected:
             self.state = httpParserStates.RCVING_BODY
             if self.body is None:
                 self.body = b''
@@ -283,13 +283,15 @@ class HttpParser:
                 self.state = httpParserStates.COMPLETE
             more, raw = len(raw) > 0, raw[total_size - received_size:]
         else:
-            # HTTP/1.0 scenario only
-            assert self.version == HTTP_1_0
             self.state = httpParserStates.RCVING_BODY
             # Received a packet without content-length header
             # and no transfer-encoding specified.
             #
+            # This can happen for both HTTP/1.0 and HTTP/1.1 scenarios.
+            # Currently, we consume the remaining buffer as body.
+            #
             # Ref https://github.com/abhinavsingh/proxy.py/issues/398
+            #
             # See TestHttpParser.test_issue_398 scenario
             self.body = raw
             more, raw = False, b''
diff --git a/proxy/http/plugin.py b/proxy/http/plugin.py
index d5510b5..a1e169c 100644
--- a/proxy/http/plugin.py
+++ b/proxy/http/plugin.py
@@ -94,6 +94,7 @@ class HttpProtocolHandlerPlugin(ABC):
 
     @abstractmethod
     def on_client_data(self, raw: memoryview) -> Optional[memoryview]:
+        """Called only after original request has been completely received."""
         return raw  # pragma: no cover
 
     @abstractmethod
diff --git a/proxy/http/proxy/plugin.py b/proxy/http/proxy/plugin.py
index 81392d1..a0f7763 100644
--- a/proxy/http/proxy/plugin.py
+++ b/proxy/http/proxy/plugin.py
@@ -121,6 +121,8 @@ class HttpProxyBasePlugin(ABC):
         Essentially, if you return None from within before_upstream_connection,
         be prepared to handle_client_data and not handle_client_request.
 
+        Only called after initial request from client has been received.
+
         Raise HttpRequestRejected to tear down the connection
         Return None to drop the connection
         """
diff --git a/proxy/http/proxy/server.py b/proxy/http/proxy/server.py
index 2dafc43..7a7622b 100644
--- a/proxy/http/proxy/server.py
+++ b/proxy/http/proxy/server.py
@@ -903,7 +903,6 @@ class HttpProxyPlugin(HttpProtocolHandlerPlugin):
     def emit_request_complete(self) -> None:
         if not self.flags.enable_events:
             return
-
         assert self.request.port
         self.event_queue.publish(
             request_id=self.uid.hex,
@@ -924,7 +923,6 @@ class HttpProxyPlugin(HttpProtocolHandlerPlugin):
     def emit_response_events(self, chunk_size: int) -> None:
         if not self.flags.enable_events:
             return
-
         if self.response.state == httpParserStates.COMPLETE:
             self.emit_response_complete()
         elif self.response.state == httpParserStates.RCVING_BODY:
@@ -935,7 +933,6 @@ class HttpProxyPlugin(HttpProtocolHandlerPlugin):
     def emit_response_headers_complete(self) -> None:
         if not self.flags.enable_events:
             return
-
         self.event_queue.publish(
             request_id=self.uid.hex,
             event_name=eventNames.RESPONSE_HEADERS_COMPLETE,
@@ -948,7 +945,6 @@ class HttpProxyPlugin(HttpProtocolHandlerPlugin):
     def emit_response_chunk_received(self, chunk_size: int) -> None:
         if not self.flags.enable_events:
             return
-
         self.event_queue.publish(
             request_id=self.uid.hex,
             event_name=eventNames.RESPONSE_CHUNK_RECEIVED,
@@ -962,7 +958,6 @@ class HttpProxyPlugin(HttpProtocolHandlerPlugin):
     def emit_response_complete(self) -> None:
         if not self.flags.enable_events:
             return
-
         self.event_queue.publish(
             request_id=self.uid.hex,
             event_name=eventNames.RESPONSE_COMPLETE,
diff --git a/proxy/plugin/proxy_pool.py b/proxy/plugin/proxy_pool.py
index 1005aca..c39c6af 100644
--- a/proxy/plugin/proxy_pool.py
+++ b/proxy/plugin/proxy_pool.py
@@ -10,10 +10,12 @@
 """
 import random
 import logging
+import ipaddress
 
 from typing import Dict, List, Optional, Any
 
 from ..common.flag import flags
+from ..common.utils import text_
 
 from ..http import Url, httpMethods
 from ..http.parser import HttpParser
@@ -78,15 +80,22 @@ class ProxyPoolPlugin(TcpUpstreamConnectionHandler, HttpProxyBasePlugin):
     ) -> Optional[HttpParser]:
         """Avoids establishing the default connection to upstream server
         by returning None.
+
+        TODO(abhinavsingh): Ideally connection to upstream proxy endpoints
+        must be bootstrapped within it's own re-usable and garbage collected pool,
+        to avoid establishing a new upstream proxy connection for each client request.
+
+        See :class:`~proxy.core.connection.pool.ConnectionPool` which is a work
+        in progress for SSL cache handling.
         """
-        # TODO(abhinavsingh): Ideally connection to upstream proxy endpoints
-        # must be bootstrapped within it's own re-usable and gc'd pool, to avoid establishing
-        # a fresh upstream proxy connection for each client request.
-        #
-        # See :class:`~proxy.core.connection.pool.ConnectionPool` which is a work
-        # in progress for SSL cache handling.
-        #
-        # Implement your own logic here e.g. round-robin, least connection etc.
+        # We don't want to send private IP requests to remote proxies
+        try:
+            if ipaddress.ip_address(text_(request.host)).is_private:
+                return request
+        except ValueError:
+            pass
+        # Choose a random proxy from the pool
+        # TODO: Implement your own logic here e.g. round-robin, least connection etc.
         endpoint = random.choice(self.flags.proxy_pool)[0].split(':', 1)
         if endpoint[0] == 'localhost' and endpoint[1] == '8899':
             return request


Test patch:
diff --git a/tests/http/test_http_proxy_tls_interception.py b/tests/http/test_http_proxy_tls_interception.py
index 1fcfa71..e10ec0f 100644
--- a/tests/http/test_http_proxy_tls_interception.py
+++ b/tests/http/test_http_proxy_tls_interception.py
@@ -145,9 +145,8 @@ class TestHttpProxyTlsInterception(Assertions):
         # Assert our mocked plugins invocations
         self.plugin.return_value.get_descriptors.assert_called()
         self.plugin.return_value.write_to_descriptors.assert_called_with([])
-        self.plugin.return_value.on_client_data.assert_called_with(
-            connect_request,
-        )
+        # on_client_data is only called after initial request has completed
+        self.plugin.return_value.on_client_data.assert_not_called()
         self.plugin.return_value.on_request_complete.assert_called()
         self.plugin.return_value.read_from_descriptors.assert_called_with([
             self._conn.fileno(),