diff --git a/README.md b/README.md index b87bfff2f6..e8743686c5 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 135a472b6b..a20a4b3433 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 54a2cb4162..3a51dbf0f4 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 9bb2a3eef3..cfdb5c8c65 100644 --- a/proxy/dashboard/dashboard.py +++ b/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, ), ), ) diff --git a/proxy/http/exception/http_request_rejected.py b/proxy/http/exception/http_request_rejected.py index a0fa810fc1..1b3b1fd742 100644 --- a/proxy/http/exception/http_request_rejected.py +++ b/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 diff --git a/proxy/http/exception/proxy_auth_failed.py b/proxy/http/exception/proxy_auth_failed.py index d82211d092..6c9649880b 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 cfba1d26dd..3c90557e8e 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 bf7a301c63..c645267f57 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 @@ 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: diff --git a/proxy/plugin/filter_by_client_ip.py b/proxy/plugin/filter_by_client_ip.py index 4b3724e96d..fba981012d 100644 --- a/proxy/plugin/filter_by_client_ip.py +++ b/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 diff --git a/proxy/plugin/filter_by_upstream.py b/proxy/plugin/filter_by_upstream.py index 0970fa6eac..a257fdf286 100644 --- a/proxy/plugin/filter_by_upstream.py +++ b/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 diff --git a/proxy/plugin/filter_by_url_regex.py b/proxy/plugin/filter_by_url_regex.py index 557a8cff3c..c9659e3732 100644 --- a/proxy/plugin/filter_by_url_regex.py +++ b/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 diff --git a/proxy/plugin/shortlink.py b/proxy/plugin/shortlink.py index 5b1ff48ffa..43e2534ea6 100644 --- a/proxy/plugin/shortlink.py +++ b/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, ), ), ) diff --git a/tests/http/exceptions/test_http_request_rejected.py b/tests/http/exceptions/test_http_request_rejected.py index de8f736bc4..69a9611339 100644 --- a/tests/http/exceptions/test_http_request_rejected.py +++ b/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, ), ) diff --git a/tests/http/test_http_parser.py b/tests/http/test_http_parser.py index 7e3c974377..774bf9f3d2 100644 --- a/tests/http/test_http_parser.py +++ b/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) diff --git a/tests/http/test_web_server.py b/tests/http/test_web_server.py index 54a1d49060..3a59fa63af 100644 --- a/tests/http/test_web_server.py +++ b/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, ), ) diff --git a/tests/plugin/test_http_proxy_plugins.py b/tests/plugin/test_http_proxy_plugins.py index 7102c3ef7a..03a32527ed 100644 --- a/tests/plugin/test_http_proxy_plugins.py +++ b/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, ), )