116 lines
3.7 KiB
Python
116 lines
3.7 KiB
Python
# -*- 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.
|
|
|
|
.. spelling::
|
|
|
|
reusability
|
|
"""
|
|
import logging
|
|
|
|
from typing import Set, Dict, Tuple
|
|
|
|
from ...common.flag import flags
|
|
|
|
from .server import TcpServerConnection
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
flags.add_argument(
|
|
'--enable-conn-pool',
|
|
action='store_true',
|
|
default=False,
|
|
help='Default: False. (WIP) Enable upstream connection pooling.',
|
|
)
|
|
|
|
|
|
class ConnectionPool:
|
|
"""Manages connection pool to upstream servers.
|
|
|
|
`ConnectionPool` avoids need to reconnect with the upstream
|
|
servers repeatedly when a reusable connection is available
|
|
in the pool.
|
|
|
|
A separate pool is maintained for each upstream server.
|
|
So internally, it's a pool of pools.
|
|
|
|
TODO: Listen for read events from the connections
|
|
to remove them from the pool when peer closes the
|
|
connection. This can also be achieved lazily by
|
|
the pool users. Example, if acquired connection
|
|
is stale, reacquire.
|
|
|
|
TODO: Ideally, ConnectionPool must be shared across
|
|
all cores to make SSL session cache to also work
|
|
without additional out-of-bound synchronizations.
|
|
|
|
TODO: ConnectionPool currently WON'T work for
|
|
HTTPS connection. This is because of missing support for
|
|
session cache, session ticket, abbr TLS handshake
|
|
and other necessary features to make it work.
|
|
|
|
NOTE: However, for all HTTP only connections, ConnectionPool
|
|
can be used to save upon connection setup time and
|
|
speed-up performance of requests.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
# Pools of connection per upstream server
|
|
self.pools: Dict[Tuple[str, int], Set[TcpServerConnection]] = {}
|
|
|
|
def acquire(self, host: str, port: int) -> Tuple[bool, TcpServerConnection]:
|
|
"""Returns a connection for use with the server."""
|
|
addr = (host, port)
|
|
# Return a reusable connection if available
|
|
if addr in self.pools:
|
|
for old_conn in self.pools[addr]:
|
|
if old_conn.is_reusable():
|
|
old_conn.mark_inuse()
|
|
logger.debug(
|
|
'Reusing connection#{2} for upstream {0}:{1}'.format(
|
|
host, port, id(old_conn),
|
|
),
|
|
)
|
|
return False, old_conn
|
|
# Create new connection
|
|
new_conn = TcpServerConnection(*addr)
|
|
if addr not in self.pools:
|
|
self.pools[addr] = set()
|
|
self.pools[addr].add(new_conn)
|
|
logger.debug(
|
|
'Created new connection#{2} for upstream {0}:{1}'.format(
|
|
host, port, id(new_conn),
|
|
),
|
|
)
|
|
return True, new_conn
|
|
|
|
def release(self, conn: TcpServerConnection) -> None:
|
|
"""Release the connection.
|
|
|
|
If the connection has not been closed,
|
|
then it will be retained in the pool for reusability.
|
|
"""
|
|
if conn.closed:
|
|
logger.debug(
|
|
'Removing connection#{2} from pool from upstream {0}:{1}'.format(
|
|
conn.addr[0], conn.addr[1], id(conn),
|
|
),
|
|
)
|
|
self.pools[conn.addr].remove(conn)
|
|
else:
|
|
logger.debug(
|
|
'Retaining connection#{2} to upstream {0}:{1}'.format(
|
|
conn.addr[0], conn.addr[1], id(conn),
|
|
),
|
|
)
|
|
assert not conn.is_reusable()
|
|
# Reset for reusability
|
|
conn.reset()
|