Source code for proxy.http.handler

# -*- 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 ssl
import time
import errno
import socket
import asyncio
import logging
import selectors
from typing import Any, List, Type, Tuple, Optional

from .parser import HttpParser, httpParserTypes, httpParserStates
from .plugin import HttpProtocolHandlerPlugin
from .exception import HttpProtocolException
from .protocols import httpProtocols
from .responses import BAD_REQUEST_RESPONSE_PKT
from ..core.base import BaseTcpServerHandler
from .connection import HttpClientConnection
from ..common.types import Readables, Writables, SelectableEvents
from ..common.constants import DEFAULT_SELECTOR_SELECT_TIMEOUT


logger = logging.getLogger(__name__)


[docs]class HttpProtocolHandler(BaseTcpServerHandler[HttpClientConnection]): """HTTP, HTTPS, HTTP2, WebSockets protocol handler. Accepts `Client` connection and delegates to HttpProtocolHandlerPlugin. """ def __init__(self, *args: Any, **kwargs: Any): super().__init__(*args, **kwargs) self.start_time: float = time.time() self.last_activity: float = self.start_time self.request: HttpParser = HttpParser( httpParserTypes.REQUEST_PARSER, enable_proxy_protocol=self.flags.enable_proxy_protocol, ) self.selector: Optional[selectors.DefaultSelector] = None if not self.flags.threadless: self.selector = selectors.DefaultSelector() self.plugin: Optional[HttpProtocolHandlerPlugin] = None ## # initialize, is_inactive, shutdown, get_events, handle_events # overrides Work class definitions. ##
[docs] @staticmethod def create(*args: Any) -> HttpClientConnection: # pragma: no cover return HttpClientConnection(*args)
[docs] def initialize(self) -> None: super().initialize() if self._encryption_enabled(): self.work = HttpClientConnection( conn=self.work.connection, addr=self.work.addr, )
[docs] def is_inactive(self) -> bool: if not self.work.has_buffer() and \ self._connection_inactive_for() > self.flags.timeout: return True return False
[docs] def shutdown(self) -> None: try: # Flush pending buffer in threaded mode only. # # For threadless mode, BaseTcpServerHandler implements # the must_flush_before_shutdown logic automagically. if self.selector and self.work.has_buffer(): self._flush() # Invoke plugin.on_client_connection_close if self.plugin: self.plugin.on_client_connection_close() logger.debug( 'Closing client connection %s has buffer %s' % (self.work.address, self.work.has_buffer()), ) conn = self.work.connection # Unwrap if wrapped before shutdown. if self._encryption_enabled() and \ isinstance(self.work.connection, ssl.SSLSocket): conn = self.work.connection.unwrap() conn.shutdown(socket.SHUT_WR) logger.debug('Client connection shutdown successful') except OSError: pass finally: # Section 4.2.2.13 of RFC 1122 tells us that a close() with any pending readable data # could lead to an immediate reset being sent. # # "A host MAY implement a 'half-duplex' TCP close sequence, so that an application # that has called CLOSE cannot continue to read data from the connection. # If such a host issues a CLOSE call while received data is still pending in TCP, # or if new data is received after CLOSE is called, its TCP SHOULD send a RST to # show that data was lost." # self.work.connection.close() logger.debug('Client connection closed') super().shutdown()
[docs] async def get_events(self) -> SelectableEvents: # Get default client events events: SelectableEvents = await super().get_events() # HttpProtocolHandlerPlugin.get_descriptors if self.plugin: plugin_read_desc, plugin_write_desc = await self.plugin.get_descriptors() for rfileno in plugin_read_desc: if rfileno not in events: events[rfileno] = selectors.EVENT_READ else: events[rfileno] |= selectors.EVENT_READ for wfileno in plugin_write_desc: if wfileno not in events: events[wfileno] = selectors.EVENT_WRITE else: events[wfileno] |= selectors.EVENT_WRITE return events
# We override super().handle_events and never call it
[docs] async def handle_events( self, readables: Readables, writables: Writables, ) -> bool: """Returns True if proxy must tear down.""" # Flush buffer for ready to write sockets teardown = await self.handle_writables(writables) if teardown: return True # Invoke plugin.write_to_descriptors if self.plugin: teardown = await self.plugin.write_to_descriptors(writables) if teardown: return True # Read from ready to read sockets teardown = await self.handle_readables(readables) if teardown: return True # Invoke plugin.read_from_descriptors if self.plugin: teardown = await self.plugin.read_from_descriptors(readables) if teardown: return True return False
[docs] def handle_data(self, data: memoryview) -> Optional[bool]: """Handles incoming data from client.""" if data is None: logger.debug('Client closed connection, tearing down...') self.work.closed = True return True try: # We don't parse incoming data any further after 1st HTTP request packet. # # Plugins can utilize on_client_data for such cases and # apply custom logic to handle request data sent after 1st # valid request. if self.request.state != httpParserStates.COMPLETE: if self._parse_first_request(data): return True # HttpProtocolHandlerPlugin.on_client_data # Can raise HttpProtocolException to tear down the connection elif self.plugin: self.plugin.on_client_data(data) except HttpProtocolException as e: logger.info('HttpProtocolException: %s' % e) response: Optional[memoryview] = e.response(self.request) if response: self.work.queue(response) return True return False
[docs] async def handle_writables(self, writables: Writables) -> bool: if self.work.connection.fileno() in writables and self.work.has_buffer(): logger.debug('Client is write ready, flushing...') self.last_activity = time.time() # TODO(abhinavsingh): This hook could just reside within server recv block # instead of invoking when flushed to client. # # Invoke plugin.on_response_chunk chunk = self.work.buffer if self.plugin: chunk = self.plugin.on_response_chunk(chunk) try: # Call super() for client flush teardown = await super().handle_writables(writables) if teardown: return True except BrokenPipeError: logger.warning( # pragma: no cover 'BrokenPipeError when flushing buffer for client', ) return True except OSError: logger.warning( # pragma: no cover 'OSError when flushing buffer to client', ) return True return False
[docs] async def handle_readables(self, readables: Readables) -> bool: if self.work.connection.fileno() in readables: logger.debug('Client is read ready, receiving...') self.last_activity = time.time() try: teardown = await super().handle_readables(readables) if teardown: return teardown except ssl.SSLWantReadError: # Try again later logger.warning( 'SSLWantReadError encountered while reading from client, will retry ...', ) return False except socket.error as e: if e.errno == errno.ECONNRESET: # Most requests for mobile devices will end up # with client closed connection. Using `debug` # here to avoid flooding the logs. logger.debug('%r' % e) else: logger.warning( 'Exception when receiving from %s connection#%d with reason %r' % (self.work.tag, self.work.connection.fileno(), e), exc_info=True, ) return True return False
## # Internal methods ##
[docs] def _initialize_plugin( self, klass: Type['HttpProtocolHandlerPlugin'], ) -> HttpProtocolHandlerPlugin: """Initializes passed HTTP protocol handler plugin class.""" return klass( self.uid, self.flags, self.work, self.request, self.event_queue, self.upstream_conn_pool, )
[docs] def _discover_plugin_klass(self, protocol: int) -> Optional[Type['HttpProtocolHandlerPlugin']]: """Discovers and return matching HTTP handler plugin matching protocol.""" if b'HttpProtocolHandlerPlugin' in self.flags.plugins: for klass in self.flags.plugins[b'HttpProtocolHandlerPlugin']: k: Type['HttpProtocolHandlerPlugin'] = klass if protocol in k.protocols(): return k return None
[docs] def _parse_first_request(self, data: memoryview) -> bool: # Parse http request try: self.request.parse(data) except HttpProtocolException as e: # noqa: WPS329 self.work.queue(BAD_REQUEST_RESPONSE_PKT) raise e except Exception as e: self.work.queue(BAD_REQUEST_RESPONSE_PKT) raise HttpProtocolException( 'Error when parsing request: %r' % data.tobytes(), ) from e if not self.request.is_complete: return False # Bail out if http protocol is unknown if self.request.http_handler_protocol == httpProtocols.UNKNOWN: self.work.queue(BAD_REQUEST_RESPONSE_PKT) return True # Discover which HTTP handler plugin is capable of # handling the current incoming request klass = self._discover_plugin_klass( self.request.http_handler_protocol, ) if klass is None: # No matching protocol class found. # Return bad request response and # close the connection. self.work.queue(BAD_REQUEST_RESPONSE_PKT) return True assert klass is not None self.plugin = self._initialize_plugin(klass) # Invoke plugin.on_request_complete output = self.plugin.on_request_complete() if isinstance(output, bool): return output assert isinstance(output, ssl.SSLSocket) logger.debug( 'Updated client conn to %s', output, ) self.work._conn = output return False
[docs] def _connection_inactive_for(self) -> float: return time.time() - self.last_activity
## # run() and _run_once() are here to maintain backward compatibility # with threaded mode. These methods are only called when running # in threaded mode. ##
[docs] def run(self) -> None: """run() method is not used when in --threadless mode. This is here just to maintain backward compatibility with threaded mode. """ loop = asyncio.new_event_loop() try: self.initialize() while True: # Tear down if client buffer is empty and connection is inactive if self.is_inactive(): logger.debug( 'Client buffer is empty and maximum inactivity has reached ' 'between client and server connection, tearing down...', ) break if loop.run_until_complete(self._run_once()): break except KeyboardInterrupt: # pragma: no cover pass except ssl.SSLError as e: logger.exception('ssl.SSLError', exc_info=e) except Exception as e: logger.exception( 'Exception while handling connection %r' % self.work.connection, exc_info=e, ) finally: self.shutdown() if self.selector: self.selector.close() loop.close()
[docs] async def _run_once(self) -> bool: events, readables, writables = await self._selected_events() try: return await self.handle_events(readables, writables) finally: assert self.selector # TODO: Like Threadless we should not unregister # work fds repeatedly. for fd in events: self.selector.unregister(fd)
# FIXME: Returning events is only necessary because we cannot use async context manager # for < Python 3.8. As a reason, this method is no longer a context manager and caller # is responsible for unregistering the descriptors.
[docs] async def _selected_events(self) -> Tuple[SelectableEvents, Readables, Writables]: assert self.selector events = await self.get_events() for fd in events: self.selector.register(fd, events[fd]) ev = self.selector.select(timeout=DEFAULT_SELECTOR_SELECT_TIMEOUT) readables = [] writables = [] for key, mask in ev: if mask & selectors.EVENT_READ: readables.append(key.fd) if mask & selectors.EVENT_WRITE: writables.append(key.fd) return (events, readables, writables)
[docs] def _flush(self) -> None: assert self.selector logger.debug('Flushing pending data') try: self.selector.register( self.work.connection, selectors.EVENT_WRITE, ) while self.work.has_buffer(): logging.debug('Waiting for client read ready') ev: List[ Tuple[selectors.SelectorKey, int] ] = self.selector.select(timeout=DEFAULT_SELECTOR_SELECT_TIMEOUT) if len(ev) == 0: continue self.work.flush(self.flags.max_sendbuf_size) except BrokenPipeError: pass finally: self.selector.unregister(self.work.connection)