Top-level notion of work not client (#695)

* Top-level notion of work not client

* Update ssl echo server example

Author: abhinavsingh

examples/https_connect_tunnel.py

@@ -53,7 +53,7 @@ def handle_data(self, data: memoryview) -> Optional[bool]:
 
         # Drop the request if not a CONNECT request
         if self.request.method != httpMethods.CONNECT:
-            self.client.queue(
+            self.work.queue(
                 HttpsConnectTunnelHandler.PROXY_TUNNEL_UNSUPPORTED_SCHEME,
             )
             return True
@@ -66,7 +66,7 @@ def handle_data(self, data: memoryview) -> Optional[bool]:
         self.connect_upstream()
 
         # Queue tunnel established response to client
-        self.client.queue(
+        self.work.queue(
             HttpsConnectTunnelHandler.PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT,
         )
 

examples/ssl_echo_server.py

@@ -27,19 +27,19 @@ def initialize(self) -> None:
         # here using wrap_socket() utility.
         assert self.flags.keyfile is not None and self.flags.certfile is not None
         conn = wrap_socket(
-            self.client.connection,
+            self.work.connection,
             self.flags.keyfile,
             self.flags.certfile,
         )
         conn.setblocking(False)
         # Upgrade plain TcpClientConnection to SSL connection object
-        self.client = TcpClientConnection(
-            conn=conn, addr=self.client.addr,
+        self.work = TcpClientConnection(
+            conn=conn, addr=self.work.addr,
         )
 
     def handle_data(self, data: memoryview) -> Optional[bool]:
         # echo back to client
-        self.client.queue(data)
+        self.work.queue(data)
         return None
 
 

examples/tcp_echo_server.py

@@ -20,11 +20,11 @@ class EchoServerHandler(BaseTcpServerHandler):
     """Sets client socket to non-blocking during initialization."""
 
     def initialize(self) -> None:
-        self.client.connection.setblocking(False)
+        self.work.connection.setblocking(False)
 
     def handle_data(self, data: memoryview) -> Optional[bool]:
         # echo back to client
-        self.client.queue(data)
+        self.work.queue(data)
         return None
 
 

proxy/core/acceptor/work.py

@@ -25,15 +25,18 @@ class Work(ABC):
 
     def __init__(
             self,
-            client: TcpClientConnection,
+            work: TcpClientConnection,
             flags: argparse.Namespace,
             event_queue: Optional[EventQueue] = None,
             uid: Optional[UUID] = None,
     ) -> None:
-        self.client = client
+        # Work uuid
+        self.uid: UUID = uid if uid is not None else uuid4()
         self.flags = flags
+        # Eventing core queue
         self.event_queue = event_queue
-        self.uid: UUID = uid if uid is not None else uuid4()
+        # Accept work
+        self.work = work
 
     @abstractmethod
     def get_events(self) -> Dict[socket.socket, int]:

proxy/core/base/tcp_server.py

@@ -45,7 +45,7 @@ class BaseTcpServerHandler(Work):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
         self.must_flush_before_shutdown = False
-        logger.debug('Connection accepted from {0}'.format(self.client.addr))
+        logger.debug('Connection accepted from {0}'.format(self.work.addr))
 
     @abstractmethod
     def handle_data(self, data: memoryview) -> Optional[bool]:
@@ -57,14 +57,14 @@ def get_events(self) -> Dict[socket.socket, int]:
         # We always want to read from client
         # Register for EVENT_READ events
         if self.must_flush_before_shutdown is False:
-            events[self.client.connection] = selectors.EVENT_READ
+            events[self.work.connection] = selectors.EVENT_READ
         # If there is pending buffer for client
         # also register for EVENT_WRITE events
-        if self.client.has_buffer():
-            if self.client.connection in events:
-                events[self.client.connection] |= selectors.EVENT_WRITE
+        if self.work.has_buffer():
+            if self.work.connection in events:
+                events[self.work.connection] |= selectors.EVENT_WRITE
             else:
-                events[self.client.connection] = selectors.EVENT_WRITE
+                events[self.work.connection] = selectors.EVENT_WRITE
         return events
 
     def handle_events(
@@ -79,32 +79,32 @@ def handle_events(
         if teardown:
             logger.debug(
                 'Shutting down client {0} connection'.format(
-                    self.client.addr,
+                    self.work.addr,
                 ),
             )
         return teardown
 
     def handle_writables(self, writables: Writables) -> bool:
         teardown = False
-        if self.client.connection in writables and self.client.has_buffer():
+        if self.work.connection in writables and self.work.has_buffer():
             logger.debug(
-                'Flushing buffer to client {0}'.format(self.client.addr),
+                'Flushing buffer to client {0}'.format(self.work.addr),
             )
-            self.client.flush()
+            self.work.flush()
             if self.must_flush_before_shutdown is True:
-                if not self.client.has_buffer():
+                if not self.work.has_buffer():
                     teardown = True
                     self.must_flush_before_shutdown = False
         return teardown
 
     def handle_readables(self, readables: Readables) -> bool:
         teardown = False
-        if self.client.connection in readables:
-            data = self.client.recv(self.flags.client_recvbuf_size)
+        if self.work.connection in readables:
+            data = self.work.recv(self.flags.client_recvbuf_size)
             if data is None:
                 logger.debug(
                     'Connection closed by client {0}'.format(
-                        self.client.addr,
+                        self.work.addr,
                     ),
                 )
                 teardown = True
@@ -113,13 +113,13 @@ def handle_readables(self, readables: Readables) -> bool:
                 if isinstance(r, bool) and r is True:
                     logger.debug(
                         'Implementation signaled shutdown for client {0}'.format(
-                            self.client.addr,
+                            self.work.addr,
                         ),
                     )
-                    if self.client.has_buffer():
+                    if self.work.has_buffer():
                         logger.debug(
                             'Client {0} has pending buffer, will be flushed before shutting down'.format(
-                                self.client.addr,
+                                self.work.addr,
                             ),
                         )
                         self.must_flush_before_shutdown = True

proxy/core/base/tcp_tunnel.py

@@ -43,7 +43,7 @@ def handle_data(self, data: memoryview) -> Optional[bool]:
         pass    # pragma: no cover
 
     def initialize(self) -> None:
-        self.client.connection.setblocking(False)
+        self.work.connection.setblocking(False)
 
     def shutdown(self) -> None:
         if self.upstream:
@@ -87,7 +87,7 @@ def handle_events(
                 print('Connection closed by server')
                 return True
             # tunnel data to client
-            self.client.queue(data)
+            self.work.queue(data)
         if self.upstream and self.upstream.connection in writables:
             self.upstream.flush()
         return False

proxy/http/handler.py

@@ -89,25 +89,25 @@ def __init__(self, *args: Any, **kwargs: Any):
 
     def initialize(self) -> None:
         """Optionally upgrades connection to HTTPS, set conn in non-blocking mode and initializes plugins."""
-        conn = self._optionally_wrap_socket(self.client.connection)
+        conn = self._optionally_wrap_socket(self.work.connection)
         conn.setblocking(False)
         # Update client connection reference if connection was wrapped
         if self._encryption_enabled():
-            self.client = TcpClientConnection(conn=conn, addr=self.client.addr)
+            self.work = TcpClientConnection(conn=conn, addr=self.work.addr)
         if b'HttpProtocolHandlerPlugin' in self.flags.plugins:
             for klass in self.flags.plugins[b'HttpProtocolHandlerPlugin']:
                 instance: HttpProtocolHandlerPlugin = klass(
                     self.uid,
                     self.flags,
-                    self.client,
+                    self.work,
                     self.request,
                     self.event_queue,
                 )
                 self.plugins[instance.name()] = instance
-        logger.debug('Handling connection %r' % self.client.connection)
+        logger.debug('Handling connection %r' % self.work.connection)
 
     def is_inactive(self) -> bool:
-        if not self.client.has_buffer() and \
+        if not self.work.has_buffer() and \
                 self._connection_inactive_for() > self.flags.timeout:
             return True
         return False
@@ -127,20 +127,20 @@ def shutdown(self) -> None:
             logger.debug(
                 'Closing client connection %r '
                 'at address %r has buffer %s' %
-                (self.client.connection, self.client.addr, self.client.has_buffer()),
+                (self.work.connection, self.work.addr, self.work.has_buffer()),
             )
 
-            conn = self.client.connection
+            conn = self.work.connection
             # Unwrap if wrapped before shutdown.
             if self._encryption_enabled() and \
-                    isinstance(self.client.connection, ssl.SSLSocket):
-                conn = self.client.connection.unwrap()
+                    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:
-            self.client.connection.close()
+            self.work.connection.close()
             logger.debug('Client connection closed')
             super().shutdown()
 
@@ -196,7 +196,7 @@ def handle_events(
     def handle_data(self, data: memoryview) -> Optional[bool]:
         if data is None:
             logger.debug('Client closed connection, tearing down...')
-            self.client.closed = True
+            self.work.closed = True
             return True
 
         try:
@@ -227,7 +227,7 @@ def handle_data(self, data: memoryview) -> Optional[bool]:
                             logger.debug(
                                 'Updated client conn to %s', upgraded_sock,
                             )
-                            self.client._conn = upgraded_sock
+                            self.work._conn = upgraded_sock
                             for plugin_ in self.plugins.values():
                                 if plugin_ != plugin:
                                     plugin_.client._conn = upgraded_sock
@@ -237,20 +237,20 @@ def handle_data(self, data: memoryview) -> Optional[bool]:
             logger.debug('HttpProtocolException raised')
             response: Optional[memoryview] = e.response(self.request)
             if response:
-                self.client.queue(response)
+                self.work.queue(response)
             return True
         return False
 
     def handle_writables(self, writables: Writables) -> bool:
-        if self.client.connection in writables and self.client.has_buffer():
+        if self.work.connection in writables and self.work.has_buffer():
             logger.debug('Client is ready for writes, flushing buffer')
             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.client.buffer
+            chunk = self.work.buffer
             for plugin in self.plugins.values():
                 chunk = plugin.on_response_chunk(chunk)
                 if chunk is None:
@@ -272,7 +272,7 @@ def handle_writables(self, writables: Writables) -> bool:
         return False
 
     def handle_readables(self, readables: Readables) -> bool:
-        if self.client.connection in readables:
+        if self.work.connection in readables:
             logger.debug('Client is ready for reads, reading')
             self.last_activity = time.time()
             try:
@@ -290,7 +290,7 @@ def handle_readables(self, readables: Readables) -> bool:
                 else:
                     logger.exception(
                         'Exception while receiving from %s connection %r with reason %r' %
-                        (self.client.tag, self.client.connection, e),
+                        (self.work.tag, self.work.connection, e),
                     )
                 return True
         return False
@@ -324,7 +324,7 @@ def run(self) -> None:
         except Exception as e:
             logger.exception(
                 'Exception while handling connection %r' %
-                self.client.connection, exc_info=e,
+                self.work.connection, exc_info=e,
             )
         finally:
             self.shutdown()
@@ -377,24 +377,24 @@ def _run_once(self) -> bool:
 
     def _flush(self) -> None:
         assert self.selector
-        if not self.client.has_buffer():
+        if not self.work.has_buffer():
             return
         try:
             self.selector.register(
-                self.client.connection,
+                self.work.connection,
                 selectors.EVENT_WRITE,
             )
-            while self.client.has_buffer():
+            while self.work.has_buffer():
                 ev: List[
                     Tuple[selectors.SelectorKey, int]
                 ] = self.selector.select(timeout=1)
                 if len(ev) == 0:
                     continue
-                self.client.flush()
+                self.work.flush()
         except BrokenPipeError:
             pass
         finally:
-            self.selector.unregister(self.client.connection)
+            self.selector.unregister(self.work.connection)
 
     def _connection_inactive_for(self) -> float:
         return time.time() - self.last_activity

tests/http/exceptions/test_http_proxy_auth_failed.py

@@ -63,9 +63,9 @@ def test_proxy_auth_fails_without_cred(self, mock_server_conn: mock.Mock) -> Non
 
         self.protocol_handler._run_once()
         mock_server_conn.assert_not_called()
-        self.assertEqual(self.protocol_handler.client.has_buffer(), True)
+        self.assertEqual(self.protocol_handler.work.has_buffer(), True)
         self.assertEqual(
-            self.protocol_handler.client.buffer[0], ProxyAuthenticationFailed.RESPONSE_PKT,
+            self.protocol_handler.work.buffer[0], ProxyAuthenticationFailed.RESPONSE_PKT,
         )
         self._conn.send.assert_not_called()
 
@@ -92,9 +92,9 @@ def test_proxy_auth_fails_with_invalid_cred(self, mock_server_conn: mock.Mock) -
 
         self.protocol_handler._run_once()
         mock_server_conn.assert_not_called()
-        self.assertEqual(self.protocol_handler.client.has_buffer(), True)
+        self.assertEqual(self.protocol_handler.work.has_buffer(), True)
         self.assertEqual(
-            self.protocol_handler.client.buffer[0], ProxyAuthenticationFailed.RESPONSE_PKT,
+            self.protocol_handler.work.buffer[0], ProxyAuthenticationFailed.RESPONSE_PKT,
         )
         self._conn.send.assert_not_called()
 
@@ -121,7 +121,7 @@ def test_proxy_auth_works_with_valid_cred(self, mock_server_conn: mock.Mock) ->
 
         self.protocol_handler._run_once()
         mock_server_conn.assert_called_once()
-        self.assertEqual(self.protocol_handler.client.has_buffer(), False)
+        self.assertEqual(self.protocol_handler.work.has_buffer(), False)
 
     @mock.patch('proxy.http.proxy.server.TcpServerConnection')
     def test_proxy_auth_works_with_mixed_case_basic_string(self, mock_server_conn: mock.Mock) -> None:
@@ -146,4 +146,4 @@ def test_proxy_auth_works_with_mixed_case_basic_string(self, mock_server_conn: m
 
         self.protocol_handler._run_once()
         mock_server_conn.assert_called_once()
-        self.assertEqual(self.protocol_handler.client.has_buffer(), False)
+        self.assertEqual(self.protocol_handler.work.has_buffer(), False)

tests/http/test_http_proxy_tls_interception.py

@@ -201,7 +201,7 @@ def mock_connection() -> Any:
         )
         self.assertEqual(self._conn.setblocking.call_count, 2)
         self.assertEqual(
-            self.protocol_handler.client.connection,
+            self.protocol_handler.work.connection,
             self.mock_ssl_wrap.return_value,
         )