Test plugin examples

Makefile

@@ -39,7 +39,7 @@ release: package
 	twine upload dist/*
 
 coverage:
-	coverage3 run --source=proxy tests.py
+	coverage3 run --source=proxy,plugin_examples tests.py
 	coverage3 html
 	open htmlcov/index.html
 

plugin_examples.py

@@ -49,7 +49,12 @@ class ProposedRestApiPlugin(proxy.HttpProxyBasePlugin):
     Used to test and develop client side applications
     without need of an actual upstream REST API server.
 
-    Returns proposed REST API mock responses to the client."""
+    Returns proposed REST API mock responses to the client
+    without establishing upstream connection.
+
+    Note: This plugin won't work if your client is making
+    HTTPS connection to api.example.com.
+    """
 
     API_SERVER = b'api.example.com'
 
@@ -76,6 +81,11 @@ class ProposedRestApiPlugin(proxy.HttpProxyBasePlugin):
     }
 
     def before_upstream_connection(self, request: proxy.HttpParser) -> Optional[proxy.HttpParser]:
+        # Return None to disable establishing connection to upstream
+        # Most likely our api.example.com won't even exist under development scenario
+        return None
+
+    def handle_client_request(self, request: proxy.HttpParser) -> Optional[proxy.HttpParser]:
         if request.host != self.API_SERVER:
             return request
         assert request.path
@@ -94,9 +104,6 @@ def before_upstream_connection(self, request: proxy.HttpParser) -> Optional[prox
             ))
         return None
 
-    def handle_client_request(self, request: proxy.HttpParser) -> Optional[proxy.HttpParser]:
-        return request
-
     def handle_upstream_chunk(self, chunk: bytes) -> bytes:
         return chunk
 
@@ -107,15 +114,16 @@ def on_upstream_connection_close(self) -> None:
 class RedirectToCustomServerPlugin(proxy.HttpProxyBasePlugin):
     """Modifies client request to redirect all incoming requests to a fixed server address."""
 
-    UPSTREAM_SERVER = b'http://localhost:8899'
+    UPSTREAM_SERVER = b'http://localhost:8899/'
 
     def before_upstream_connection(self, request: proxy.HttpParser) -> Optional[proxy.HttpParser]:
         # Redirect all non-https requests to inbuilt WebServer.
         if request.method != proxy.httpMethods.CONNECT:
-            request.url = urlparse.urlsplit(self.UPSTREAM_SERVER)
-            # This command will re-parse modified url and
-            # update host, port, path fields
-            request.set_line_attributes()
+            request.set_url(self.UPSTREAM_SERVER)
+            # Update Host header too, otherwise upstream can reject our request
+            if request.has_header(b'Host'):
+                request.del_header(b'Host')
+            request.add_header(b'Host', urlparse.urlsplit(self.UPSTREAM_SERVER).netloc)
         return request
 
     def handle_client_request(self, request: proxy.HttpParser) -> Optional[proxy.HttpParser]:

proxy.py

@@ -598,6 +598,10 @@ def del_headers(self, headers: List[bytes]) -> None:
         for key in headers:
             self.del_header(key.lower())
 
+    def set_url(self, url: bytes) -> None:
+        self.url = urlparse.urlsplit(url)
+        self.set_line_attributes()
+
     def set_line_attributes(self) -> None:
         if self.type == httpParserTypes.REQUEST_PARSER:
             if self.method == httpMethods.CONNECT and self.url:
@@ -700,13 +704,12 @@ def process_line(self, raw: bytes) -> None:
         line = raw.split(WHITESPACE)
         if self.type == httpParserTypes.REQUEST_PARSER:
             self.method = line[0].upper()
-            self.url = urlparse.urlsplit(line[1])
+            self.set_url(line[1])
             self.version = line[2]
         else:
             self.version = line[0]
             self.code = line[1]
             self.reason = WHITESPACE.join(line[2:])
-        self.set_line_attributes()
 
     def process_header(self, raw: bytes) -> None:
         parts = raw.split(COLON)
@@ -1226,7 +1229,7 @@ def get_descriptors(
     def write_to_descriptors(self, w: List[Union[int, _HasFileno]]) -> bool:
         if self.request.has_upstream_server() and \
                 self.server and not self.server.closed and \
-                self.server.buffer_size() > 0 and \
+                self.server.has_buffer() and \
                 self.server.connection in w:
             logger.debug('Server is write ready, flushing buffer')
             try:
@@ -1378,16 +1381,20 @@ def on_request_complete(self) -> Union[socket.socket, bool]:
         if not self.request.has_upstream_server():
             return False
 
+        self.authenticate()
+
         # Note: can raise HttpRequestRejected exception
         # Invoke plugin.before_upstream_connection
+        do_connect = True
         for plugin in self.plugins.values():
             r = plugin.before_upstream_connection(self.request)
             if r is None:
-                return False
+                do_connect = False
+                break
             self.request = r
 
-        self.authenticate()
-        self.connect_upstream()
+        if do_connect:
+            self.connect_upstream()
 
         for plugin in self.plugins.values():
             assert self.request is not None
@@ -1444,7 +1451,7 @@ def on_request_complete(self) -> Union[socket.socket, bool]:
             # https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Connection
             # connection headers are meant for communication between client and
             # first intercepting proxy.
-            self.request.add_headers([(b'Via', b'1.1 proxy.py v%s' % version)])
+            self.request.add_headers([(b'Via', b'1.1 %s' % PROXY_AGENT_HEADER_VALUE)])
             # Disable args.disable_headers before dispatching to upstream
             self.server.queue(
                 self.request.build(
@@ -2095,7 +2102,8 @@ 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.setblocking(False)
-        self.client = TcpClientConnection(conn=conn, addr=self.addr)
+        if self.config.encryption_enabled():
+            self.client = TcpClientConnection(conn=conn, addr=self.addr)
         if b'ProtocolHandlerPlugin' in self.config.plugins:
             for klass in self.config.plugins[b'ProtocolHandlerPlugin']:
                 instance = klass(self.config, self.client, self.request)