Add conn_close kwarg to packet builder utilities

README.md

@@ -1512,8 +1512,7 @@ As a decorator:
 - Generate HTTP GET request with headers
 
   ```python
-  >>> build_http_request(b'GET', b'/',
-          headers={b'Connection': b'close'})
+  >>> build_http_request(b'GET', b'/', conn_close=True)
   b'GET / HTTP/1.1\r\nConnection: close\r\n\r\n'
   ```
 

examples/https_connect_tunnel.py

@@ -31,8 +31,8 @@ class HttpsConnectTunnelHandler(BaseTcpTunnelHandler):
     PROXY_TUNNEL_UNSUPPORTED_SCHEME = memoryview(
         build_http_response(
             httpStatusCodes.BAD_REQUEST,
-            headers={b'Connection': b'close'},
             reason=b'Unsupported protocol scheme',
+            conn_close=True,
         ),
     )
 

proxy/common/utils.py

@@ -75,12 +75,14 @@ def build_http_request(
     protocol_version: bytes = HTTP_1_1,
     headers: Optional[Dict[bytes, bytes]] = None,
     body: Optional[bytes] = None,
+    conn_close: bool = False,
 ) -> bytes:
     """Build and returns a HTTP request packet."""
-    if headers is None:
-        headers = {}
     return build_http_pkt(
-        [method, url, protocol_version], headers, body,
+        [method, url, protocol_version],
+        headers or {},
+        body,
+        conn_close,
     )
 
 
@@ -90,6 +92,7 @@ def build_http_response(
     reason: Optional[bytes] = None,
     headers: Optional[Dict[bytes, bytes]] = None,
     body: Optional[bytes] = None,
+    conn_close: bool = False,
 ) -> bytes:
     """Build and returns a HTTP response packet."""
     line = [protocol_version, bytes_(status_code)]
@@ -108,7 +111,7 @@ def build_http_response(
             not has_transfer_encoding and \
             not has_content_length:
         headers[b'Content-Length'] = bytes_(len(body))
-    return build_http_pkt(line, headers, body)
+    return build_http_pkt(line, headers, body, conn_close)
 
 
 def build_http_header(k: bytes, v: bytes) -> bytes:
@@ -120,12 +123,15 @@ def build_http_pkt(
     line: List[bytes],
     headers: Optional[Dict[bytes, bytes]] = None,
     body: Optional[bytes] = None,
+    conn_close: bool = False,
 ) -> bytes:
     """Build and returns a HTTP request or response packet."""
     pkt = WHITESPACE.join(line) + CRLF
-    if headers is not None:
-        for k, v in headers.items():
-            pkt += build_http_header(k, v) + CRLF
+    headers = headers or {}
+    if conn_close:
+        headers[b'Connection'] = b'close'
+    for k, v in headers.items():
+        pkt += build_http_header(k, v) + CRLF
     pkt += CRLF
     if body:
         pkt += body

proxy/dashboard/dashboard.py

@@ -83,8 +83,8 @@ def handle_request(self, request: HttpParser) -> None:
                         headers={
                             b'Location': b'/dashboard/',
                             b'Content-Length': b'0',
-                            b'Connection': b'close',
                         },
+                        conn_close=True,
                     ),
                 ),
             )

proxy/http/exception/http_request_rejected.py

@@ -41,6 +41,7 @@ def response(self, _request: HttpParser) -> Optional[memoryview]:
                     reason=self.reason,
                     headers=self.headers,
                     body=self.body,
+                    conn_close=True,
                 ),
             )
         return None

proxy/http/exception/proxy_auth_failed.py

@@ -33,9 +33,9 @@ class ProxyAuthenticationFailed(HttpProtocolException):
             headers={
                 PROXY_AGENT_HEADER_KEY: PROXY_AGENT_HEADER_VALUE,
                 b'Proxy-Authenticate': b'Basic',
-                b'Connection': b'close',
             },
             body=b'Proxy Authentication Required',
+            conn_close=True,
         ),
     )
 

proxy/http/exception/proxy_conn_failed.py

@@ -30,9 +30,9 @@ class ProxyConnectionFailed(HttpProtocolException):
             reason=b'Bad Gateway',
             headers={
                 PROXY_AGENT_HEADER_KEY: PROXY_AGENT_HEADER_VALUE,
-                b'Connection': b'close',
             },
             body=b'Bad Gateway',
+            conn_close=True,
         ),
     )
 

proxy/http/server/web.py

@@ -81,8 +81,8 @@ class HttpWebServerPlugin(HttpProtocolHandlerPlugin):
             headers={
                 b'Server': PROXY_AGENT_HEADER_VALUE,
                 b'Content-Length': b'0',
-                b'Connection': b'close',
             },
+            conn_close=True,
         ),
     )
 
@@ -93,8 +93,8 @@ class HttpWebServerPlugin(HttpProtocolHandlerPlugin):
             headers={
                 b'Server': PROXY_AGENT_HEADER_VALUE,
                 b'Content-Length': b'0',
-                b'Connection': b'close',
             },
+            conn_close=True,
         ),
     )
 
@@ -147,7 +147,6 @@ def read_and_build_static_file_response(path: str, min_compression_limit: int) -
             headers = {
                 b'Content-Type': bytes_(content_type),
                 b'Cache-Control': b'max-age=86400',
-                b'Connection': b'close',
             }
             do_compress = len(content) > min_compression_limit
             if do_compress:
@@ -160,6 +159,7 @@ def read_and_build_static_file_response(path: str, min_compression_limit: int) -
                     reason=b'OK',
                     headers=headers,
                     body=gzip.compress(content) if do_compress else content,
+                    conn_close=True,
                 ),
             )
         except FileNotFoundError:

proxy/plugin/filter_by_client_ip.py

@@ -39,9 +39,7 @@ def before_upstream_connection(
         assert not self.flags.unix_socket_path and self.client.addr
         if self.client.addr[0] in self.flags.filtered_client_ips.split(','):
             raise HttpRequestRejected(
-                status_code=httpStatusCodes.I_AM_A_TEAPOT, reason=b'I\'m a tea pot',
-                headers={
-                    b'Connection': b'close',
-                },
+                status_code=httpStatusCodes.I_AM_A_TEAPOT,
+                reason=b'I\'m a tea pot',
             )
         return request

proxy/plugin/filter_by_upstream.py

@@ -35,9 +35,7 @@ def before_upstream_connection(
     ) -> Optional[HttpParser]:
         if text_(request.host) in self.flags.filtered_upstream_hosts.split(','):
             raise HttpRequestRejected(
-                status_code=httpStatusCodes.I_AM_A_TEAPOT, reason=b'I\'m a tea pot',
-                headers={
-                    b'Connection': b'close',
-                },
+                status_code=httpStatusCodes.I_AM_A_TEAPOT,
+                reason=b'I\'m a tea pot',
             )
         return request

proxy/plugin/filter_by_url_regex.py

@@ -88,7 +88,6 @@ def handle_client_request(
                 # list
                 raise HttpRequestRejected(
                     status_code=httpStatusCodes.NOT_FOUND,
-                    headers={b'Connection': b'close'},
                     reason=b'Blocked',
                 )
                 # stop looping through filter list

proxy/plugin/shortlink.py

@@ -72,8 +72,8 @@ def handle_client_request(
                             headers={
                                 b'Location': b'http://' + self.SHORT_LINKS[request.host] + path,
                                 b'Content-Length': b'0',
-                                b'Connection': b'close',
                             },
+                            conn_close=True,
                         ),
                     ),
                 )
@@ -84,8 +84,8 @@ def handle_client_request(
                             httpStatusCodes.NOT_FOUND, reason=b'NOT FOUND',
                             headers={
                                 b'Content-Length': b'0',
-                                b'Connection': b'close',
                             },
+                            conn_close=True,
                         ),
                     ),
                 )

tests/http/exceptions/test_http_request_rejected.py

@@ -31,19 +31,23 @@ def test_status_code_response(self) -> None:
         self.assertEqual(
             e.response(self.request), CRLF.join([
                 b'HTTP/1.1 200 OK',
+                b'Connection: close',
                 CRLF,
             ]),
         )
 
     def test_body_response(self) -> None:
         e = HttpRequestRejected(
-            status_code=httpStatusCodes.NOT_FOUND, reason=b'NOT FOUND',
+            status_code=httpStatusCodes.NOT_FOUND,
+            reason=b'NOT FOUND',
             body=b'Nothing here',
         )
         self.assertEqual(
             e.response(self.request),
             build_http_response(
                 httpStatusCodes.NOT_FOUND,
-                reason=b'NOT FOUND', body=b'Nothing here',
+                reason=b'NOT FOUND',
+                body=b'Nothing here',
+                conn_close=True,
             ),
         )

tests/http/test_http_parser.py

@@ -694,9 +694,7 @@ def test_is_not_http_1_1_keep_alive_with_close_header(self) -> None:
         self.parser.parse(
             build_http_request(
                 httpMethods.GET, b'/',
-                headers={
-                    b'Connection': b'close',
-                },
+                conn_close=True,
             ),
         )
         self.assertFalse(self.parser.is_http_1_1_keep_alive)

tests/http/test_web_server.py

@@ -181,10 +181,13 @@ async def test_pac_file_served_from_disk(self) -> None:
         )
         self._conn.send.called_once_with(
             build_http_response(
-                200, reason=b'OK', headers={
+                200,
+                reason=b'OK',
+                headers={
                     b'Content-Type': b'application/x-ns-proxy-autoconfig',
-                    b'Connection': b'close',
-                }, body=self.expected_response,
+                },
+                body=self.expected_response,
+                conn_close=True,
             ),
         )
 

tests/plugin/test_http_proxy_plugins.py

@@ -242,9 +242,7 @@ async def test_filter_by_upstream_host_plugin(self) -> None:
             build_http_response(
                 status_code=httpStatusCodes.I_AM_A_TEAPOT,
                 reason=b'I\'m a tea pot',
-                headers={
-                    b'Connection': b'close',
-                },
+                conn_close=True,
             ),
         )
 
@@ -376,6 +374,6 @@ async def test_filter_by_url_regex_plugin(self) -> None:
             build_http_response(
                 status_code=httpStatusCodes.NOT_FOUND,
                 reason=b'Blocked',
-                headers={b'Connection': b'close'},
+                conn_close=True,
             ),
         )