Repository: abhinavsingh/proxy.py
Pull request #898 (2021-12-21 21:07:03)

Add `conn_close` kwarg to packet builder utilities
- passing `True` will automatically add `Connection: close` header to the packet

codecov[bot]: # [Codecov](https://codecov.io/gh/abhinavsingh/proxy.py/pull/898?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) Report
> Merging [#898](https://codecov.io/gh/abhinavsingh/proxy.py/pull/898?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) (cfae52f) into [develop](https://codecov.io/gh/abhinavsingh/proxy.py/commit/da1795d2ed78175addb4da16852903e8e9b1ccd3?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) (da1795d) will **decrease** coverage by `0.08%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/abhinavsingh/proxy.py/pull/898/graphs/tree.svg?width=650&height=150&src=pr&token=Zh9J7b4la2&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh)](https://codecov.io/gh/abhinavsingh/proxy.py/pull/898?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh)

```diff
@@             Coverage Diff             @@
##           develop     #898      +/-   ##
===========================================
- Coverage    87.00%   86.91%   -0.09%     
===========================================
  Files          130      130              
  Lines         5833     5833              
  Branches       583      582       -1     
===========================================
- Hits          5075     5070       -5     
- Misses         652      654       +2     
- Partials       106      109       +3     
```

| Flag | Coverage Δ | |
|---|---|---|
| pytest | `86.83% <100.00%> (-0.02%)` | :arrow_down: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Impacted Files](https://codecov.io/gh/abhinavsingh/proxy.py/pull/898?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) | Coverage Δ | |
|---|---|---|
| [proxy/dashboard/dashboard.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/898/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvZGFzaGJvYXJkL2Rhc2hib2FyZC5weQ==) | `43.39% <ø> (ø)` | |
| [proxy/http/exception/http\_request\_rejected.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/898/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvaHR0cC9leGNlcHRpb24vaHR0cF9yZXF1ZXN0X3JlamVjdGVkLnB5) | `100.00% <ø> (ø)` | |
| [proxy/http/exception/proxy\_auth\_failed.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/898/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvaHR0cC9leGNlcHRpb24vcHJveHlfYXV0aF9mYWlsZWQucHk=) | `100.00% <ø> (ø)` | |
| [proxy/http/exception/proxy\_conn\_failed.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/898/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvaHR0cC9leGNlcHRpb24vcHJveHlfY29ubl9mYWlsZWQucHk=) | `100.00% <ø> (ø)` | |
| [proxy/http/server/web.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/898/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvaHR0cC9zZXJ2ZXIvd2ViLnB5) | `63.35% <ø> (ø)` | |
| [proxy/plugin/filter\_by\_client\_ip.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/898/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvcGx1Z2luL2ZpbHRlcl9ieV9jbGllbnRfaXAucHk=) | `71.42% <ø> (ø)` | |
| [proxy/plugin/filter\_by\_upstream.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/898/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvcGx1Z2luL2ZpbHRlcl9ieV91cHN0cmVhbS5weQ==) | `85.71% <ø> (ø)` | |
| [proxy/plugin/filter\_by\_url\_regex.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/898/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvcGx1Z2luL2ZpbHRlcl9ieV91cmxfcmVnZXgucHk=) | `73.68% <ø> (ø)` | |
| [proxy/plugin/shortlink.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/898/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvcGx1Z2luL3Nob3J0bGluay5weQ==) | `52.38% <ø> (ø)` | |
| [...ests/http/exceptions/test\_http\_request\_rejected.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/898/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-dGVzdHMvaHR0cC9leGNlcHRpb25zL3Rlc3RfaHR0cF9yZXF1ZXN0X3JlamVjdGVkLnB5) | `100.00% <ø> (ø)` | |
| ... and [7 more](https://codecov.io/gh/abhinavsingh/proxy.py/pull/898/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/abhinavsingh/proxy.py/pull/898?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/abhinavsingh/proxy.py/pull/898?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh). Last update [da1795d...cfae52f](https://codecov.io/gh/abhinavsingh/proxy.py/pull/898?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh).


Patch:
diff --git a/README.md b/README.md
index b87bfff..e874368 100644
--- a/README.md
+++ b/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'
   ```
 
diff --git a/examples/https_connect_tunnel.py b/examples/https_connect_tunnel.py
index 135a472..a20a4b3 100644
--- a/examples/https_connect_tunnel.py
+++ b/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,
         ),
     )
 
diff --git a/proxy/common/utils.py b/proxy/common/utils.py
index 54a2cb4..3a51dbf 100644
--- a/proxy/common/utils.py
+++ b/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
diff --git a/proxy/dashboard/dashboard.py b/proxy/dashboard/dashboard.py
index 9bb2a3e..cfdb5c8 100644
--- a/proxy/dashboard/dashboard.py
+++ b/proxy/dashboard/dashboard.py
@@ -83,8 +83,8 @@ class ProxyDashboard(HttpWebServerBasePlugin):
                         headers={
                             b'Location': b'/dashboard/',
                             b'Content-Length': b'0',
-                            b'Connection': b'close',
                         },
+                        conn_close=True,
                     ),
                 ),
             )
diff --git a/proxy/http/exception/http_request_rejected.py b/proxy/http/exception/http_request_rejected.py
index a0fa810..1b3b1fd 100644
--- a/proxy/http/exception/http_request_rejected.py
+++ b/proxy/http/exception/http_request_rejected.py
@@ -41,6 +41,7 @@ class HttpRequestRejected(HttpProtocolException):
                     reason=self.reason,
                     headers=self.headers,
                     body=self.body,
+                    conn_close=True,
                 ),
             )
         return None
diff --git a/proxy/http/exception/proxy_auth_failed.py b/proxy/http/exception/proxy_auth_failed.py
index d82211d..6c96498 100644
--- a/proxy/http/exception/proxy_auth_failed.py
+++ b/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,
         ),
     )
 
diff --git a/proxy/http/exception/proxy_conn_failed.py b/proxy/http/exception/proxy_conn_failed.py
index cfba1d2..3c90557 100644
--- a/proxy/http/exception/proxy_conn_failed.py
+++ b/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,
         ),
     )
 
diff --git a/proxy/http/server/web.py b/proxy/http/server/web.py
index bf7a301..c645267 100644
--- a/proxy/http/server/web.py
+++ b/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 @@ class HttpWebServerPlugin(HttpProtocolHandlerPlugin):
             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 @@ class HttpWebServerPlugin(HttpProtocolHandlerPlugin):
                     reason=b'OK',
                     headers=headers,
                     body=gzip.compress(content) if do_compress else content,
+                    conn_close=True,
                 ),
             )
         except FileNotFoundError:
diff --git a/proxy/plugin/filter_by_client_ip.py b/proxy/plugin/filter_by_client_ip.py
index 4b3724e..fba9810 100644
--- a/proxy/plugin/filter_by_client_ip.py
+++ b/proxy/plugin/filter_by_client_ip.py
@@ -39,9 +39,7 @@ class FilterByClientIpPlugin(HttpProxyBasePlugin):
         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
diff --git a/proxy/plugin/filter_by_upstream.py b/proxy/plugin/filter_by_upstream.py
index 0970fa6..a257fdf 100644
--- a/proxy/plugin/filter_by_upstream.py
+++ b/proxy/plugin/filter_by_upstream.py
@@ -35,9 +35,7 @@ class FilterByUpstreamHostPlugin(HttpProxyBasePlugin):
     ) -> 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
diff --git a/proxy/plugin/filter_by_url_regex.py b/proxy/plugin/filter_by_url_regex.py
index 557a8cf..c9659e3 100644
--- a/proxy/plugin/filter_by_url_regex.py
+++ b/proxy/plugin/filter_by_url_regex.py
@@ -88,7 +88,6 @@ class FilterByURLRegexPlugin(HttpProxyBasePlugin):
                 # list
                 raise HttpRequestRejected(
                     status_code=httpStatusCodes.NOT_FOUND,
-                    headers={b'Connection': b'close'},
                     reason=b'Blocked',
                 )
                 # stop looping through filter list
diff --git a/proxy/plugin/shortlink.py b/proxy/plugin/shortlink.py
index 5b1ff48..43e2534 100644
--- a/proxy/plugin/shortlink.py
+++ b/proxy/plugin/shortlink.py
@@ -72,8 +72,8 @@ class ShortLinkPlugin(HttpProxyBasePlugin):
                             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 @@ class ShortLinkPlugin(HttpProxyBasePlugin):
                             httpStatusCodes.NOT_FOUND, reason=b'NOT FOUND',
                             headers={
                                 b'Content-Length': b'0',
-                                b'Connection': b'close',
                             },
+                            conn_close=True,
                         ),
                     ),
                 )


Test patch:
diff --git a/tests/http/exceptions/test_http_request_rejected.py b/tests/http/exceptions/test_http_request_rejected.py
index de8f736..69a9611 100644
--- a/tests/http/exceptions/test_http_request_rejected.py
+++ b/tests/http/exceptions/test_http_request_rejected.py
@@ -31,19 +31,23 @@ class TestHttpRequestRejected(unittest.TestCase):
         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,
             ),
         )
diff --git a/tests/http/test_http_parser.py b/tests/http/test_http_parser.py
index 7e3c974..774bf9f 100644
--- a/tests/http/test_http_parser.py
+++ b/tests/http/test_http_parser.py
@@ -694,9 +694,7 @@ class TestHttpParser(unittest.TestCase):
         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)
diff --git a/tests/http/test_web_server.py b/tests/http/test_web_server.py
index 54a1d49..3a59fa6 100644
--- a/tests/http/test_web_server.py
+++ b/tests/http/test_web_server.py
@@ -181,10 +181,13 @@ class TestWebServerPluginWithPacFilePlugin(Assertions):
         )
         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,
             ),
         )
 
diff --git a/tests/plugin/test_http_proxy_plugins.py b/tests/plugin/test_http_proxy_plugins.py
index 7102c3e..03a3252 100644
--- a/tests/plugin/test_http_proxy_plugins.py
+++ b/tests/plugin/test_http_proxy_plugins.py
@@ -242,9 +242,7 @@ class TestHttpProxyPluginExamples(Assertions):
             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 @@ class TestHttpProxyPluginExamples(Assertions):
             build_http_response(
                 status_code=httpStatusCodes.NOT_FOUND,
                 reason=b'Blocked',
-                headers={b'Connection': b'close'},
+                conn_close=True,
             ),
         )