v2.4.0rc9 (#1068)

* [Core] Invoke selector.close on shutdown (#1055) [Core] Invoke `selector.close` on shutdown * [CacheResponsesPlugin] Add ability to cache request packets (#1056) * Start of post encryption tests * Assertion on post encryption callback * Add `--cache-requests` flag * Clean up `on_client_data` API as we no longer have a chain within core http protocol handler * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix flake8 warnings * Fix `inconsistent-return-statements` Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * [Task] A generic payload based work abstraction (#1057) * Refactor into an internal task submodule of work * As context managers * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add missing license Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * `jupyter` notebook based tutorial (#1059) * `jupyter` notebook based tutorial * Move within `tutorial` directory * Fix spell * Add `as_non_blocking` option during wrap * `as_non_blocking=False` by defaut * `--max-sendbuf-size` flag to speed up large file upload/download (#1060) * Add `--max-sendbuf-size` flag which now defaults to 64Kb * Use `server_recvbuf_size` flag with base tunnel implementation * isort * Add to readme * [Flags] `min_compression_length` consistency (#1061) * `min_compression_length` consistency, it was used as `min_compression_limit` at a few places * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * revert back web server route * Move `serve_static_file` as a staticmethod within web plugin base Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * [Fix] Using `okResponse()` without content hangs the connection (#1062) * It hangs because of no content-length or connection close header * Fix tests * [Jupyter] Add a response notebook (#1064) * Add a response generation jupyter notebook * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Make codespell happy * precommit codespell exclude Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * [Optimize] Avoid using `tobytes` for zero-copies (#1066) * Avoid using `tobytes` where possible * `send` accepts `Union[memoryview, bytes]` now * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * [HttpParser] Memory view compliant, Zero copies (#1067) * Remove usage of `tobytes` * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix chunk parser * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Lint fixes Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * [Jupyter] Request creation notebook (#1065) * Add plugin tests, responses notebook and enhancements to `build_http_packet` * Add js code snip for ws example * Fix tests * ignore all ipynb from codespell * ignore all ipynb from codespell * Fix tests and doc spell Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
abhinavsingh · Jan 26, 2022 · 558a430 · 558a430
2 parents 3858f3a + 3fd608e
commit 558a430
Show file tree

Hide file tree

Showing 52 changed files with 2,212 additions and 542 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -125,8 +125,12 @@ repos:
   rev: v2.1.0
   hooks:
   - id: codespell
-    exclude: >-
-      ^.+\.min\.js$
+    exclude: >
+      (?x)^(
+        ^.+\.ipynb$|
+        tests/http/test_responses\.py|
+        ^.+\.min\.js$
+      )$
 
 - repo: https://github.com/adrienverge/yamllint.git
   rev: v1.26.2

diff --git a/Makefile b/Makefile
@@ -94,7 +94,7 @@ lib-clean:
 	rm -rf .hypothesis
 	# Doc RST files are cached and can cause issues
 	# See https://github.com/abhinavsingh/proxy.py/issues/642#issuecomment-1003444578
-	rm docs/pkg/*.rst
+	rm -f docs/pkg/*.rst
 
 lib-dep:
 	pip install --upgrade pip && \
@@ -134,7 +134,7 @@ lib-doc:
 	python -m tox -e build-docs && \
 	$(OPEN) .tox/build-docs/docs_out/index.html || true
 
-lib-coverage:
+lib-coverage: lib-clean
 	pytest --cov=proxy --cov=tests --cov-report=html tests/ && \
 	$(OPEN) htmlcov/index.html || true
 

diff --git a/README.md b/README.md
@@ -250,7 +250,7 @@
   - See `--enable-static-server` and `--static-server-dir` flags
 
 - Optimized for large file uploads and downloads
-  - See `--client-recvbuf-size` and `--server-recvbuf-size` flag
+  - See `--client-recvbuf-size`, `--server-recvbuf-size`, `--max-sendbuf-size` flags
 
 - `IPv4` and `IPv6` support
   - See `--hostname` flag
@@ -710,6 +710,8 @@ Start `proxy.py` as:
     --plugins proxy.plugin.CacheResponsesPlugin
 ```
 
+You may also use the `--cache-requests` flag to enable request packet caching for inspection.
+
 Verify using `curl -v -x localhost:8899 http://httpbin.org/get`:
 
 ```console
@@ -2278,13 +2280,14 @@ usage: -m [-h] [--tunnel-hostname TUNNEL_HOSTNAME] [--tunnel-port TUNNEL_PORT]
           [--work-klass WORK_KLASS] [--pid-file PID_FILE]
           [--enable-proxy-protocol] [--enable-conn-pool] [--key-file KEY_FILE]
           [--cert-file CERT_FILE] [--client-recvbuf-size CLIENT_RECVBUF_SIZE]
-          [--server-recvbuf-size SERVER_RECVBUF_SIZE] [--timeout TIMEOUT]
+          [--server-recvbuf-size SERVER_RECVBUF_SIZE]
+          [--max-sendbuf-size MAX_SENDBUF_SIZE] [--timeout TIMEOUT]
           [--disable-http-proxy] [--disable-headers DISABLE_HEADERS]
           [--ca-key-file CA_KEY_FILE] [--ca-cert-dir CA_CERT_DIR]
           [--ca-cert-file CA_CERT_FILE] [--ca-file CA_FILE]
           [--ca-signing-key-file CA_SIGNING_KEY_FILE]
           [--auth-plugin AUTH_PLUGIN] [--cache-dir CACHE_DIR]
-          [--proxy-pool PROXY_POOL] [--enable-web-server]
+          [--cache-requests] [--proxy-pool PROXY_POOL] [--enable-web-server]
           [--enable-static-server] [--static-server-dir STATIC_SERVER_DIR]
           [--min-compression-length MIN_COMPRESSION_LENGTH]
           [--enable-reverse-proxy] [--pac-file PAC_FILE]
@@ -2294,7 +2297,7 @@ usage: -m [-h] [--tunnel-hostname TUNNEL_HOSTNAME] [--tunnel-port TUNNEL_PORT]
           [--filtered-client-ips FILTERED_CLIENT_IPS]
           [--filtered-url-regex-config FILTERED_URL_REGEX_CONFIG]
 
-proxy.py v2.4.0rc8.dev17+g59a4335.d20220123
+proxy.py v2.4.0rc9.dev8+gea0253d.d20220126
 
 options:
   -h, --help            show this help message and exit
@@ -2389,6 +2392,9 @@ options:
   --server-recvbuf-size SERVER_RECVBUF_SIZE
                         Default: 128 KB. Maximum amount of data received from
                         the server in a single recv() operation.
+  --max-sendbuf-size MAX_SENDBUF_SIZE
+                        Default: 64 KB. Maximum amount of data to dispatch in
+                        a single send() operation.
   --timeout TIMEOUT     Default: 10.0. Number of seconds after which an
                         inactive connection must be dropped. Inactivity is
                         defined by no data sent or received by the client.
@@ -2425,6 +2431,7 @@ options:
                         Default: /Users/abhinavsingh/.proxy/cache. Flag only
                         applicable when cache plugin is used with on-disk
                         storage.
+  --cache-requests      Default: False. Whether to also cache request packets.
   --proxy-pool PROXY_POOL
                         List of upstream proxies to use in the pool
   --enable-web-server   Default: False. Whether to enable

diff --git a/examples/task.py b/examples/task.py
@@ -8,116 +8,64 @@
     :copyright: (c) 2013-present by Abhinav Singh and contributors.
     :license: BSD, see LICENSE for more details.
 """
-import time
+import sys
 import argparse
-import threading
-import multiprocessing
-from typing import Any
 
-from proxy.core.work import (
-    Work, ThreadlessPool, BaseLocalExecutor, BaseRemoteExecutor,
-)
+from proxy.core.work import ThreadlessPool
 from proxy.common.flag import FlagParser
-from proxy.common.backports import NonBlockingQueue
-
-
-class Task:
-    """This will be our work object."""
-
-    def __init__(self, payload: bytes) -> None:
-        self.payload = payload
-        print(payload)
-
-
-class TaskWork(Work[Task]):
-    """This will be our handler class, created for each received work."""
-
-    @staticmethod
-    def create(*args: Any) -> Task:
-        """Work core doesn't know how to create work objects for us, so
-        we must provide an implementation of create method here."""
-        return Task(*args)
-
-
-class LocalTaskExecutor(BaseLocalExecutor):
-    """We'll define a local executor which is capable of receiving
-    log lines over a non blocking queue."""
-
-    def work(self, *args: Any) -> None:
-        task_id = int(time.time())
-        uid = '%s-%s' % (self.iid, task_id)
-        self.works[task_id] = self.create(uid, *args)
-
-
-class RemoteTaskExecutor(BaseRemoteExecutor):
-
-    def work(self, *args: Any) -> None:
-        task_id = int(time.time())
-        uid = '%s-%s' % (self.iid, task_id)
-        self.works[task_id] = self.create(uid, *args)
-
-
-def start_local(flags: argparse.Namespace) -> None:
-    work_queue = NonBlockingQueue()
-    executor = LocalTaskExecutor(iid=1, work_queue=work_queue, flags=flags)
+from proxy.core.work.task import (
+    RemoteTaskExecutor, ThreadedTaskExecutor, SingleProcessTaskExecutor,
+)
 
-    t = threading.Thread(target=executor.run)
-    t.daemon = True
-    t.start()
 
-    try:
+def start_local_thread(flags: argparse.Namespace) -> None:
+    with ThreadedTaskExecutor(flags=flags) as thread:
         i = 0
         while True:
-            work_queue.put(('%d' % i).encode('utf-8'))
+            thread.executor.work_queue.put(('%d' % i).encode('utf-8'))
             i += 1
-    except KeyboardInterrupt:
-        pass
-    finally:
-        executor.running.set()
-        t.join()
 
 
-def start_remote(flags: argparse.Namespace) -> None:
-    pipe = multiprocessing.Pipe()
-    work_queue = pipe[0]
-    executor = RemoteTaskExecutor(iid=1, work_queue=pipe[1], flags=flags)
+def start_remote_process(flags: argparse.Namespace) -> None:
+    with SingleProcessTaskExecutor(flags=flags) as process:
+        i = 0
+        while True:
+            process.work_queue.send(('%d' % i).encode('utf-8'))
+            i += 1
 
-    p = multiprocessing.Process(target=executor.run)
-    p.daemon = True
-    p.start()
 
-    try:
+def start_remote_pool(flags: argparse.Namespace) -> None:
+    with ThreadlessPool(flags=flags, executor_klass=RemoteTaskExecutor) as pool:
         i = 0
         while True:
+            work_queue = pool.work_queues[i % flags.num_workers]
             work_queue.send(('%d' % i).encode('utf-8'))
             i += 1
-    except KeyboardInterrupt:
-        pass
-    finally:
-        executor.running.set()
-        p.join()
 
 
-def start_remote_pool(flags: argparse.Namespace) -> None:
-    with ThreadlessPool(flags=flags, executor_klass=RemoteTaskExecutor) as pool:
-        try:
-            i = 0
-            while True:
-                work_queue = pool.work_queues[i % flags.num_workers]
-                work_queue.send(('%d' % i).encode('utf-8'))
-                i += 1
-        except KeyboardInterrupt:
-            pass
+def main() -> None:
+    try:
+        flags = FlagParser.initialize(
+            sys.argv[2:] + ['--disable-http-proxy'],
+            work_klass='proxy.core.work.task.TaskHandler',
+        )
+        globals()['start_%s' % sys.argv[1]](flags)
+    except KeyboardInterrupt:
+        pass
 
 
 # TODO: TaskWork, LocalTaskExecutor, RemoteTaskExecutor
 # should not be needed, abstract those pieces out in the core
 # for stateless tasks.
 if __name__ == '__main__':
-    flags = FlagParser.initialize(
-        ['--disable-http-proxy'],
-        work_klass=TaskWork,
-    )
-    start_remote_pool(flags)
-    # start_remote(flags)
-    # start_local(flags)
+    if len(sys.argv) < 2:
+        print(
+            '\n'.join([
+                'Usage:',
+                '  %s <execution-mode>' % sys.argv[0],
+                '    execution-mode can be one of the following:',
+                '    "remote_pool", "remote_process", "local_thread"',
+            ]),
+        )
+        sys.exit(1)
+    main()
diff --git a/proxy/common/constants.py b/proxy/common/constants.py
@@ -74,6 +74,7 @@ def _env_threadless_compliant() -> bool:
 # Defaults
 DEFAULT_BACKLOG = 100
 DEFAULT_BASIC_AUTH = None
+DEFAULT_MAX_SEND_SIZE = 64 * 1024
 DEFAULT_BUFFER_SIZE = 128 * 1024
 DEFAULT_CA_CERT_DIR = None
 DEFAULT_CA_CERT_FILE = None
@@ -124,14 +125,13 @@ def _env_threadless_compliant() -> bool:
 DEFAULT_PORT = 8899
 DEFAULT_SERVER_RECVBUF_SIZE = DEFAULT_BUFFER_SIZE
 DEFAULT_STATIC_SERVER_DIR = os.path.join(PROXY_PY_DIR, "public")
-DEFAULT_MIN_COMPRESSION_LIMIT = 20  # In bytes
+DEFAULT_MIN_COMPRESSION_LENGTH = 20  # In bytes
 DEFAULT_THREADLESS = _env_threadless_compliant()
 DEFAULT_LOCAL_EXECUTOR = True
 DEFAULT_TIMEOUT = 10.0
 DEFAULT_VERSION = False
 DEFAULT_HTTP_PORT = 80
 DEFAULT_HTTPS_PORT = 443
-DEFAULT_MAX_SEND_SIZE = 16 * 1024
 DEFAULT_WORK_KLASS = 'proxy.http.HttpProtocolHandler'
 DEFAULT_ENABLE_PROXY_PROTOCOL = False
 # 25 milliseconds to keep the loops hot
@@ -148,6 +148,7 @@ def _env_threadless_compliant() -> bool:
 DEFAULT_CACHE_DIRECTORY_PATH = os.path.join(
     DEFAULT_DATA_DIRECTORY_PATH, 'cache',
 )
+DEFAULT_CACHE_REQUESTS = False
 
 # Cor plugins enabled by default or via flags
 DEFAULT_ABC_PLUGINS = [

diff --git a/proxy/common/flag.py b/proxy/common/flag.py
@@ -30,7 +30,7 @@
     PLUGIN_REVERSE_PROXY, DEFAULT_NUM_ACCEPTORS, PLUGIN_INSPECT_TRAFFIC,
     DEFAULT_DISABLE_HEADERS, PY2_DEPRECATION_MESSAGE, DEFAULT_DEVTOOLS_WS_PATH,
     PLUGIN_DEVTOOLS_PROTOCOL, PLUGIN_WEBSOCKET_TRANSPORT,
-    DEFAULT_DATA_DIRECTORY_PATH, DEFAULT_MIN_COMPRESSION_LIMIT,
+    DEFAULT_DATA_DIRECTORY_PATH, DEFAULT_MIN_COMPRESSION_LENGTH,
 )
 
 
@@ -335,13 +335,13 @@ def initialize(
                 args.static_server_dir,
             ),
         )
-        args.min_compression_limit = cast(
+        args.min_compression_length = cast(
             bool,
             opts.get(
-                'min_compression_limit',
+                'min_compression_length',
                 getattr(
-                    args, 'min_compression_limit',
-                    DEFAULT_MIN_COMPRESSION_LIMIT,
+                    args, 'min_compression_length',
+                    DEFAULT_MIN_COMPRESSION_LENGTH,
                 ),
             ),
         )

diff --git a/proxy/common/utils.py b/proxy/common/utils.py
@@ -26,7 +26,7 @@
 from .types import HostPort
 from .constants import (
     CRLF, COLON, HTTP_1_1, IS_WINDOWS, WHITESPACE, DEFAULT_TIMEOUT,
-    DEFAULT_THREADLESS,
+    DEFAULT_THREADLESS, PROXY_AGENT_HEADER_VALUE,
 )
 
 
@@ -84,14 +84,30 @@ def bytes_(s: Any, encoding: str = 'utf-8', errors: str = 'strict') -> Any:
 def build_http_request(
     method: bytes, url: bytes,
     protocol_version: bytes = HTTP_1_1,
+    content_type: Optional[bytes] = None,
     headers: Optional[Dict[bytes, bytes]] = None,
     body: Optional[bytes] = None,
     conn_close: bool = False,
+    no_ua: bool = False,
 ) -> bytes:
     """Build and returns a HTTP request packet."""
+    headers = headers or {}
+    if content_type is not None:
+        headers[b'Content-Type'] = content_type
+    has_transfer_encoding = False
+    has_user_agent = False
+    for k, _ in headers.items():
+        if k.lower() == b'transfer-encoding':
+            has_transfer_encoding = True
+        elif k.lower() == b'user-agent':
+            has_user_agent = True
+    if body and not has_transfer_encoding:
+        headers[b'Content-Length'] = bytes_(len(body))
+    if not has_user_agent and not no_ua:
+        headers[b'User-Agent'] = PROXY_AGENT_HEADER_VALUE
     return build_http_pkt(
         [method, url, protocol_version],
-        headers or {},
+        headers,
         body,
         conn_close,
     )
@@ -109,19 +125,14 @@ def build_http_response(
     line = [protocol_version, bytes_(status_code)]
     if reason:
         line.append(reason)
-    if headers is None:
-        headers = {}
-    has_content_length = False
+    headers = headers or {}
     has_transfer_encoding = False
     for k, _ in headers.items():
-        if k.lower() == b'content-length':
-            has_content_length = True
         if k.lower() == b'transfer-encoding':
             has_transfer_encoding = True
-    if body is not None and \
-            not has_transfer_encoding and \
-            not has_content_length:
-        headers[b'Content-Length'] = bytes_(len(body))
+            break
+    if not has_transfer_encoding:
+        headers[b'Content-Length'] = bytes_(len(body)) if body else b'0'
     return build_http_pkt(line, headers, body, conn_close)
 
 

diff --git a/proxy/core/acceptor/acceptor.py b/proxy/core/acceptor/acceptor.py
@@ -178,6 +178,7 @@ def run(self) -> None:
             for fileno in self.socks:
                 self.socks[fileno].close()
             self.socks.clear()
+            self.selector.close()
             logger.debug('Acceptor#%d shutdown', self.idd)
 
     def _recv_and_setup_socks(self) -> None:
@@ -207,7 +208,8 @@ def _start_local(self) -> None:
         self._lthread.start()
 
     def _stop_local(self) -> None:
-        if self._lthread is not None and self._local_work_queue is not None:
+        if self._lthread is not None and \
+                self._local_work_queue is not None:
             self._local_work_queue.put(False)
             self._lthread.join()