Core acceptor pool doc, cleanup and standalone example (#393)

* Better document acceptor module and add a TCP Echo Server example

* autopep8 formating

* Rename ThreadlessWork --> Work class

* Make initialize, is_inactive and shutdown as optional interface methods.

Also introduce Readables & Writables custom types.

* Move websocket code into its own module

* Add websocket client example

* Cleanup websocket client

Author: abhinavsingh

Makefile

@@ -25,6 +25,7 @@ devtools:
 	pushd dashboard && npm run devtools && popd
 
 autopep8:
+	autopep8 --recursive --in-place --aggressive examples
 	autopep8 --recursive --in-place --aggressive proxy
 	autopep8 --recursive --in-place --aggressive tests
 	autopep8 --recursive --in-place --aggressive setup.py
@@ -73,8 +74,8 @@ lib-clean:
 	rm -rf .hypothesis
 
 lib-lint:
-	flake8 --ignore=W504 --max-line-length=127 --max-complexity=19 proxy/ tests/ setup.py
-	mypy --strict --ignore-missing-imports proxy/ tests/ setup.py
+	flake8 --ignore=W504 --max-line-length=127 --max-complexity=19 examples/ proxy/ tests/ setup.py
+	mypy --strict --ignore-missing-imports examples/ proxy/ tests/ setup.py
 
 lib-test: lib-clean lib-version lib-lint
 	pytest -v tests/

examples/README.md

@@ -0,0 +1,5 @@
+# Proxy.py Library Examples
+
+This directory contains examples that demonstrate `proxy.py` core library capabilities.
+
+Looking for `proxy.py` plugin examples?  Check [proxy/plugin](https://github.com/abhinavsingh/proxy.py/tree/develop/proxy/plugin) directory.

examples/tcp_echo_server.py

@@ -0,0 +1,75 @@
+# -*- 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 time
+import socket
+import selectors
+
+from typing import Dict
+
+from proxy.core.acceptor import AcceptorPool, Work
+from proxy.common.flags import Flags
+from proxy.common.types import Readables, Writables
+
+
+class EchoServerHandler(Work):
+    """EchoServerHandler implements Work interface.
+
+    An instance of EchoServerHandler is created for each client
+    connection.  EchoServerHandler lifecycle is controlled by
+    Threadless core using asyncio.  Implementation must provide
+    get_events and handle_events method.  Optionally, also implement
+    intialize, is_inactive and shutdown method.
+    """
+
+    def get_events(self) -> Dict[socket.socket, int]:
+        # We always want to read from client
+        # Register for EVENT_READ events
+        events = {self.client.connection: selectors.EVENT_READ}
+        # If there is pending buffer for client
+        # also register for EVENT_WRITE events
+        if self.client.has_buffer():
+            events[self.client.connection] |= selectors.EVENT_WRITE
+        return events
+
+    def handle_events(
+            self,
+            readables: Readables,
+            writables: Writables) -> bool:
+        """Return True to shutdown work."""
+        if self.client.connection in readables:
+            data = self.client.recv()
+            if data is None:
+                # Client closed connection, signal shutdown
+                return True
+            # Queue data back to client
+            self.client.queue(data)
+
+        if self.client.connection in writables:
+            self.client.flush()
+
+        return False
+
+
+def main() -> None:
+    # This example requires `threadless=True`
+    pool = AcceptorPool(
+        flags=Flags(num_workers=1, threadless=True),
+        work_klass=EchoServerHandler)
+    try:
+        pool.setup()
+        while True:
+            time.sleep(1)
+    finally:
+        pool.shutdown()
+
+
+if __name__ == '__main__':
+    main()

examples/websocket_client.py

@@ -0,0 +1,44 @@
+# -*- 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 time
+from proxy.http.websocket import WebsocketClient, WebsocketFrame, websocketOpcodes
+
+
+# globals
+client: WebsocketClient
+last_dispatch_time: float
+static_frame = memoryview(WebsocketFrame.text(b'hello'))
+num_echos = 10
+
+
+def on_message(frame: WebsocketFrame) -> None:
+    """WebsocketClient on_message callback."""
+    global client, num_echos, last_dispatch_time
+    print('Received %r after %d millisec' % (frame.data, (time.time() - last_dispatch_time) * 1000))
+    assert(frame.data == b'hello' and frame.opcode == websocketOpcodes.TEXT_FRAME)
+    if num_echos > 0:
+        client.queue(static_frame)
+        last_dispatch_time = time.time()
+        num_echos -= 1
+    else:
+        client.close()
+
+
+if __name__ == '__main__':
+    # Constructor establishes socket connection
+    client = WebsocketClient(b'echo.websocket.org', 80, b'/', on_message=on_message)
+    # Perform handshake
+    client.handshake()
+    # Queue some data for client
+    client.queue(static_frame)
+    last_dispatch_time = time.time()
+    # Start event loop
+    client.run()

proxy/common/flags.py

@@ -21,8 +21,9 @@
 import sys
 import inspect
 
-from typing import Optional, Union, Dict, List, TypeVar, Type, cast, Any, Tuple
+from typing import Optional, Dict, List, TypeVar, Type, cast, Any, Tuple
 
+from .types import IpAddress
 from .utils import text_, bytes_
 from .constants import DEFAULT_LOG_LEVEL, DEFAULT_LOG_FILE, DEFAULT_LOG_FORMAT, DEFAULT_BACKLOG, DEFAULT_BASIC_AUTH
 from .constants import DEFAULT_TIMEOUT, DEFAULT_DEVTOOLS_WS_PATH, DEFAULT_DISABLE_HTTP_PROXY, DEFAULT_DISABLE_HEADERS
@@ -67,8 +68,7 @@ def __init__(
             ca_signing_key_file: Optional[str] = None,
             ca_file: Optional[str] = None,
             num_workers: int = 0,
-            hostname: Union[ipaddress.IPv4Address,
-                            ipaddress.IPv6Address] = DEFAULT_IPV6_HOSTNAME,
+            hostname: IpAddress = DEFAULT_IPV6_HOSTNAME,
             port: int = DEFAULT_PORT,
             backlog: int = DEFAULT_BACKLOG,
             static_server_dir: str = DEFAULT_STATIC_SERVER_DIR,
@@ -99,8 +99,7 @@ def __init__(
         self.ca_signing_key_file: Optional[str] = ca_signing_key_file
         self.ca_file = ca_file
         self.num_workers: int = num_workers if num_workers > 0 else multiprocessing.cpu_count()
-        self.hostname: Union[ipaddress.IPv4Address,
-                             ipaddress.IPv6Address] = hostname
+        self.hostname: IpAddress = hostname
         self.family: socket.AddressFamily = socket.AF_INET6 if hostname.version == 6 else socket.AF_INET
         self.port: int = port
         self.backlog: int = backlog
@@ -161,7 +160,8 @@ def initialize(
         # Setup limits
         Flags.set_open_file_limit(args.open_file_limit)
 
-        # Prepare list of plugins to load based upon --enable-* and --disable-* flags
+        # Prepare list of plugins to load based upon --enable-* and --disable-*
+        # flags
         default_plugins: List[Tuple[str, bool]] = []
         if args.enable_dashboard:
             default_plugins.append((PLUGIN_WEB_SERVER, True))
@@ -249,8 +249,7 @@ def initialize(
                 opts.get(
                     'ca_file',
                     args.ca_file)),
-            hostname=cast(Union[ipaddress.IPv4Address,
-                                ipaddress.IPv6Address],
+            hostname=cast(IpAddress,
                           opts.get('hostname', ipaddress.ip_address(args.hostname))),
             port=cast(int, opts.get('port', args.port)),
             backlog=cast(int, opts.get('backlog', args.backlog)),

proxy/common/types.py

@@ -9,8 +9,9 @@
     :license: BSD, see LICENSE for more details.
 """
 import queue
+import ipaddress
 
-from typing import TYPE_CHECKING, Dict, Any
+from typing import TYPE_CHECKING, Dict, Any, List, Union
 
 from typing_extensions import Protocol
 
@@ -23,3 +24,8 @@
 class HasFileno(Protocol):
     def fileno(self) -> int:
         ...     # pragma: no cover
+
+
+Readables = List[Union[int, HasFileno]]
+Writables = List[Union[int, HasFileno]]
+IpAddress = Union[ipaddress.IPv4Address, ipaddress.IPv6Address]

proxy/common/utils.py

@@ -101,7 +101,8 @@ def build_http_pkt(line: List[bytes],
 def build_websocket_handshake_request(
         key: bytes,
         method: bytes = b'GET',
-        url: bytes = b'/') -> bytes:
+        url: bytes = b'/',
+        host: bytes = b'localhost') -> bytes:
     """
     Build and returns a Websocket handshake request packet.
 
@@ -112,6 +113,7 @@ def build_websocket_handshake_request(
     return build_http_request(
         method, url,
         headers={
+            b'Host': host,
             b'Connection': b'upgrade',
             b'Upgrade': b'websocket',
             b'Sec-WebSocket-Key': key,

proxy/core/acceptor/__init__.py

@@ -10,8 +10,12 @@
 """
 from .acceptor import Acceptor
 from .pool import AcceptorPool
+from .work import Work
+from .threadless import Threadless
 
 __all__ = [
     'Acceptor',
     'AcceptorPool',
+    'Work',
+    'Threadless',
 ]

proxy/core/acceptor/acceptor.py

@@ -14,32 +14,36 @@
 import selectors
 import socket
 import threading
-# import time
+
 from multiprocessing import connection
 from multiprocessing.reduction import send_handle, recv_handle
 from typing import Optional, Type, Tuple
 
+from .work import Work
+from .threadless import Threadless
+
 from ..connection import TcpClientConnection
-from ..threadless import ThreadlessWork, Threadless
 from ..event import EventQueue, eventNames
 from ...common.flags import Flags
 
 logger = logging.getLogger(__name__)
 
 
 class Acceptor(multiprocessing.Process):
-    """Socket client acceptor.
+    """Socket server acceptor process.
 
-    Accepts client connection over received server socket handle and
-    starts a new work thread.
+    Accepts client connection over received server socket handle at startup.  Spawns a separate
+    thread to handle each client request.  However, when `--threadless` is enabled, Acceptor also
+    pre-spawns a `Threadless` process at startup.  Accepted client connections are passed to
+    `Threadless` process which internally uses asyncio event loop to handle client connections.
     """
 
     def __init__(
             self,
             idd: int,
             work_queue: connection.Connection,
             flags: Flags,
-            work_klass: Type[ThreadlessWork],
+            work_klass: Type[Work],
             lock: multiprocessing.synchronize.Lock,
             event_queue: Optional[EventQueue] = None) -> None:
         super().__init__()
@@ -108,11 +112,7 @@ def run_once(self) -> None:
             if len(events) == 0:
                 return
             conn, addr = self.sock.accept()
-
-        # now = time.time()
-        # fileno: int = conn.fileno()
         self.start_work(conn, addr)
-        # logger.info('Work started for fd %d in %f seconds', fileno, time.time() - now)
 
     def run(self) -> None:
         self.selector = selectors.DefaultSelector()

proxy/core/acceptor/pool.py

@@ -18,7 +18,8 @@
 from typing import List, Optional, Type
 
 from .acceptor import Acceptor
-from ..threadless import ThreadlessWork
+from .work import Work
+
 from ..event import EventQueue, EventDispatcher
 from ...common.flags import Flags
 
@@ -31,11 +32,20 @@ class AcceptorPool:
     """AcceptorPool.
 
     Pre-spawns worker processes to utilize all cores available on the system.  Server socket connection is
-    dispatched over a pipe to workers.  Each worker accepts incoming client request and spawns a
-    separate thread to handle the client request.
+    dispatched over a pipe to workers.  Each Acceptor instance accepts for new client connection.
+
+    Example usage:
+
+    pool = AcceptorPool(flags=..., work_klass=...)
+    try:
+        pool.setup()
+        while True:
+            time.sleep(1)
+    finally:
+        pool.shutdown()
     """
 
-    def __init__(self, flags: Flags, work_klass: Type[ThreadlessWork]) -> None:
+    def __init__(self, flags: Flags, work_klass: Type[Work]) -> None:
         self.flags = flags
         self.socket: Optional[socket.socket] = None
         self.acceptors: List[Acceptor] = []