Skip to content

Commit

Permalink
Core acceptor pool doc, cleanup and standalone example (#393)
Browse files Browse the repository at this point in the history
* 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
  • Loading branch information
abhinavsingh committed Jul 7, 2020
1 parent 9be6c29 commit c884338
Show file tree
Hide file tree
Showing 27 changed files with 467 additions and 239 deletions.
5 changes: 3 additions & 2 deletions Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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/
Expand Down
5 changes: 5 additions & 0 deletions examples/README.md
Original file line number Diff line number Diff line change
@@ -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.
75 changes: 75 additions & 0 deletions examples/tcp_echo_server.py
Original file line number Diff line number Diff line change
@@ -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()
44 changes: 44 additions & 0 deletions examples/websocket_client.py
Original file line number Diff line number Diff line change
@@ -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()
15 changes: 7 additions & 8 deletions proxy/common/flags.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -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)),
Expand Down
8 changes: 7 additions & 1 deletion proxy/common/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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]
4 changes: 3 additions & 1 deletion proxy/common/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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,
Expand Down
4 changes: 4 additions & 0 deletions proxy/core/acceptor/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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',
]
20 changes: 10 additions & 10 deletions proxy/core/acceptor/acceptor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__()
Expand Down Expand Up @@ -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()
Expand Down
18 changes: 14 additions & 4 deletions proxy/core/acceptor/pool.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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] = []
Expand Down
Loading

0 comments on commit c884338

Please sign in to comment.