Document FilterByClientIpPlugin & ModifyChunkResponsePlugin (#387)

Author: abhinavsingh

README.md

@@ -57,6 +57,8 @@ Table of Contents
         * [Cache Responses Plugin](#cacheresponsesplugin)
         * [Man-In-The-Middle Plugin](#maninthemiddleplugin)
         * [Proxy Pool Plugin](#proxypoolplugin)
+        * [FilterByClientIpPlugin](#filterbyclientipplugin)
+        * [ModifyChunkResponsePlugin](#modifychunkresponseplugin)
     * [HTTP Web Server Plugins](#http-web-server-plugins)
         * [Reverse Proxy](#reverse-proxy)
         * [Web Server Route](#web-server-route)
@@ -669,6 +671,57 @@ Make a curl request via `8899` proxy:
 Verify that `8899` proxy forwards requests to upstream proxies
 by checking respective logs.
 
+### FilterByClientIpPlugin
+
+Reject traffic from specific IP addresses.  By default this
+plugin blocks traffic from `127.0.0.1` and `::1`.
+
+Start `proxy.py` as:
+
+```bash
+❯ proxy \
+    --plugins proxy.plugin.FilterByClientIpPlugin
+```
+
+Send a request using `curl -v -x localhost:8899 http://google.com`:
+
+```bash
+... [redacted] ...
+> Proxy-Connection: Keep-Alive
+>
+< HTTP/1.1 418 I'm a tea pot
+< Connection: close
+<
+* Closing connection 0
+```
+
+Modify plugin to your taste e.g. Allow specific IP addresses only.
+
+### ModifyChunkResponsePlugin
+
+This plugin demonstrate how to modify chunked encoded responses.  In able to do so, this plugin uses `proxy.py` core to parse the chunked encoded response.  Then we reconstruct the response using custom hardcoded chunks, ignoring original chunks received from upstream server.
+
+Start `proxy.py` as:
+
+```bash
+❯ proxy \
+    --plugins proxy.plugin.ModifyChunkResponsePlugin
+```
+
+Verify using `curl -v -x localhost:8899 http://httpbin.org/stream/5`:
+
+```bash
+... [redacted] ...
+modify
+chunk
+response
+plugin
+* Connection #0 to host localhost left intact
+* Closing connection 0
+```
+
+Modify `ModifyChunkResponsePlugin` to your taste. Example, instead of sending hardcoded chunks, parse and modify the original `JSON` chunks received from the upstream server.
+
 ## HTTP Web Server Plugins
 
 ### Reverse Proxy

proxy/http/parser.py

@@ -15,7 +15,7 @@
 from .chunk_parser import ChunkParser, chunkParserStates
 
 from ..common.constants import DEFAULT_DISABLE_HEADERS, COLON, CRLF, WHITESPACE, HTTP_1_1, DEFAULT_HTTP_PORT
-from ..common.utils import build_http_request, find_http_line, text_
+from ..common.utils import build_http_request, build_http_response, find_http_line, text_
 
 
 HttpParserStates = NamedTuple('HttpParserStates', [
@@ -237,7 +237,8 @@ def build_path(self) -> bytes:
         return url
 
     def build(self, disable_headers: Optional[List[bytes]] = None) -> bytes:
-        assert self.method and self.version and self.path
+        """Rebuild the request object."""
+        assert self.method and self.version and self.path and self.type == httpParserTypes.REQUEST_PARSER
         if disable_headers is None:
             disable_headers = DEFAULT_DISABLE_HEADERS
         body: Optional[bytes] = ChunkParser.to_chunks(self.body) \
@@ -250,6 +251,16 @@ def build(self, disable_headers: Optional[List[bytes]] = None) -> bytes:
             body=body
         )
 
+    def build_response(self) -> bytes:
+        """Rebuild the response object."""
+        assert self.code and self.version and self.body and self.type == httpParserTypes.RESPONSE_PARSER
+        return build_http_response(
+            status_code=int(self.code),
+            protocol_version=self.version,
+            reason=self.reason,
+            headers={} if not self.headers else {self.headers[k][0]: self.headers[k][1] for k in self.headers},
+            body=self.body if not self.is_chunked_encoded() else ChunkParser.to_chunks(self.body))
+
     def has_upstream_server(self) -> bool:
         """Host field SHOULD be None for incoming local WebServer requests."""
         return True if self.host is not None else False

proxy/plugin/__init__.py

@@ -19,6 +19,7 @@
 from .reverse_proxy import ReverseProxyPlugin
 from .proxy_pool import ProxyPoolPlugin
 from .filter_by_client_ip import FilterByClientIpPlugin
+from .modify_chunk_response_plugin import ModifyChunkResponsePlugin
 
 __all__ = [
     'CacheResponsesPlugin',
@@ -33,4 +34,5 @@
     'ReverseProxyPlugin',
     'ProxyPoolPlugin',
     'FilterByClientIpPlugin',
+    'ModifyChunkResponsePlugin',
 ]

proxy/plugin/modify_chunk_response_plugin.py

@@ -0,0 +1,51 @@
+# -*- 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 typing import Optional, Any
+
+from ..http.parser import HttpParser, httpParserTypes, httpParserStates
+from ..http.proxy import HttpProxyBasePlugin
+
+
+class ModifyChunkResponsePlugin(HttpProxyBasePlugin):
+    """Accumulate & modify chunk responses as received from upstream."""
+
+    DEFAULT_CHUNKS = [
+        b'modify',
+        b'chunk',
+        b'response',
+        b'plugin',
+    ]
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        # Create a new http protocol parser for response payloads
+        self.response = HttpParser(httpParserTypes.RESPONSE_PARSER)
+
+    def before_upstream_connection(
+            self, request: HttpParser) -> Optional[HttpParser]:
+        return request
+
+    def handle_client_request(
+            self, request: HttpParser) -> Optional[HttpParser]:
+        return request
+
+    def handle_upstream_chunk(self, chunk: memoryview) -> memoryview:
+        # Parse the response.
+        # Note that these chunks also include headers
+        self.response.parse(chunk.tobytes())
+        # If response is complete, modify and dispatch to client
+        if self.response.state == httpParserStates.COMPLETE:
+            self.response.body = b'\n'.join(self.DEFAULT_CHUNKS) + b'\n'
+            self.client.queue(memoryview(self.response.build_response()))
+        return memoryview(b'')
+
+    def on_upstream_connection_close(self) -> None:
+        pass