Use ipaddress type and avoid explicit Connection: close header to upstream

Makefile

@@ -36,10 +36,10 @@ coverage:
 	open htmlcov/index.html
 
 lint:
-	flake8 --ignore=E501,W504 --builtins="unicode" proxy.py
-	flake8 --ignore=E501,W504 tests.py
 	autopep8 --recursive --in-place --aggressive --aggressive proxy.py
 	autopep8 --recursive --in-place --aggressive --aggressive tests.py
+	flake8 --ignore=E501,W504 --builtins="unicode" proxy.py
+	flake8 --ignore=E501,W504 tests.py
 
 container:
 	docker build -t $(LATEST_TAG) -t $(IMAGE_TAG) .

README.md

@@ -3,21 +3,23 @@
 [![License](https://img.shields.io/github/license/abhinavsingh/proxy.py.svg)](https://opensource.org/licenses/BSD-3-Clause)
 [![PyPi Downloads](https://img.shields.io/pypi/dm/proxy.py.svg)](https://pypi.org/project/proxy.py/)
 [![Build Status](https://travis-ci.org/abhinavsingh/proxy.py.svg?branch=develop)](https://travis-ci.org/abhinavsingh/proxy.py/)
+[![No Dependencies](https://david-dm.org/dwyl/esta.svg)](https://github.com/abhinavsingh/proxy.py)
 [![Coverage](https://coveralls.io/repos/github/abhinavsingh/proxy.py/badge.svg?branch=develop)](https://coveralls.io/github/abhinavsingh/proxy.py?branch=develop)
 
 [![Python 3.5](https://img.shields.io/badge/python-3.5-blue.svg)](https://www.python.org/downloads/release/python-350/)
 [![Python 3.6](https://img.shields.io/badge/python-3.6-blue.svg)](https://www.python.org/downloads/release/python-360/)
 [![Python 3.7](https://img.shields.io/badge/python-3.7-blue.svg)](https://www.python.org/downloads/release/python-370/)
 
 [![Maintenance](https://img.shields.io/badge/Maintained%3F-yes-green.svg)](https://gitHub.com/abhinavsingh/proxy.py/graphs/commit-activity)
+[![contributions welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg?style=flat)](https://github.com/abhinavsingh/proxy.py/issues)
 [![Ask Me Anything](https://img.shields.io/badge/Ask%20me-anything-1abc9c.svg)](https://twitter.com/imoracle)
 
 Features
 --------
 
 - Lightweight
     - Distributed as a single file module `~50KB`
-    - Uses only `RAM <20MB` for general use cases
+    - Uses only `~5-20MB` RAM
     - No external dependency other than standard Python library
 - Programmable
     - Optionally enable builtin Web Server

proxy.py

@@ -50,8 +50,8 @@
 DEFAULT_CLIENT_RECVBUF_SIZE = DEFAULT_BUFFER_SIZE
 DEFAULT_SERVER_RECVBUF_SIZE = DEFAULT_BUFFER_SIZE
 DEFAULT_DISABLE_HEADERS: List[str] = []
-DEFAULT_IPV4_HOSTNAME = '127.0.0.1'
-DEFAULT_IPV6_HOSTNAME = '::'
+DEFAULT_IPV4_HOSTNAME = ipaddress.IPv4Address('127.0.0.1')
+DEFAULT_IPV6_HOSTNAME = ipaddress.IPv6Address('::1')
 DEFAULT_PORT = 8899
 DEFAULT_IPV4 = False
 DEFAULT_DISABLE_HTTP_PROXY = False
@@ -209,8 +209,9 @@ def __init__(
         self.socket: Optional[socket.socket] = None
         self.running: bool = False
         self.family = socket.AF_INET if self.ipv4 else socket.AF_INET6
-        self.hostname: str = hostname if hostname not in [DEFAULT_IPV4_HOSTNAME,
-                                                          DEFAULT_IPV6_HOSTNAME] \
+        self.hostname: Union[ipaddress.IPv4Address, ipaddress.IPv6Address] = \
+            hostname if hostname not in [DEFAULT_IPV4_HOSTNAME,
+                                         DEFAULT_IPV6_HOSTNAME] \
             else DEFAULT_IPV4_HOSTNAME if self.ipv4 else DEFAULT_IPV6_HOSTNAME
 
     @abstractmethod
@@ -234,7 +235,7 @@ def run(self):
         try:
             self.socket = socket.socket(self.family, socket.SOCK_STREAM)
             self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-            self.socket.bind((self.hostname, self.port))
+            self.socket.bind((str(self.hostname), self.port))
             self.socket.listen(self.backlog)
             logger.info('Started server on %s:%d' % (self.hostname, self.port))
             while self.running:
@@ -485,7 +486,7 @@ def parse(self, raw) -> None:
                     self.state = HttpParser.states.RCVING_BODY
                     self.body += raw
                     if len(
-                            self.body) >= int(
+                        self.body) >= int(
                             self.headers[b'content-length'][1]):
                         self.state = HttpParser.states.COMPLETE
                 elif self.is_chunked_encoded_response():
@@ -562,7 +563,7 @@ def process_header(self, raw):
         parts = raw.split(COLON)
         key = parts[0].strip()
         value = COLON.join(parts[1:]).strip()
-        self.headers[key.lower()] = (key, value)
+        self.add_headers([(key, value)])
 
     def build_url(self):
         if not self.url:
@@ -623,19 +624,19 @@ def has_upstream_server(self):
         return True if self.host is not None else False
 
     def add_header(self, key: bytes, value: bytes) -> None:
-        self.headers[key] = (key, value)
+        self.headers[key.lower()] = (key, value)
 
     def add_headers(self, headers: List[Tuple[bytes, bytes]]) -> None:
         for (key, value) in headers:
             self.add_header(key, value)
 
     def del_header(self, header: bytes) -> None:
-        if header in self.headers:
-            del self.headers[header]
+        if header.lower() in self.headers:
+            del self.headers[header.lower()]
 
     def del_headers(self, headers: List[bytes]) -> None:
         for key in headers:
-            self.del_header(key)
+            self.del_header(key.lower())
 
 
 class HttpProtocolException(Exception):
@@ -968,11 +969,20 @@ def on_request_complete(self):
         # for general http requests, re-build request packet
         # and queue for the server with appropriate headers
         else:
-            # remove args.disable_headers before dispatching to upstream
-            self.request.add_headers(
-                [(b'Via', b'1.1 proxy.py v%s' % version), (b'Connection', b'Close')])
+            # - proxy-connection header is a mistake, it doesn't seem to be
+            #   officially documented in any specification, drop it.
+            # - proxy-authorization is of no use for upstream, remove it.
             self.request.del_headers(
-                [b'proxy-authorization', b'proxy-connection', b'connection', b'keep-alive'])
+                [b'proxy-authorization', b'proxy-connection'])
+            # - For HTTP/1.0, connection header defaults to close
+            # - For HTTP/1.1, connection header defaults to keep-alive
+            # Respect headers sent by client instead of manipulating
+            # Connection or Keep-Alive header.  However, note that per
+            # 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)])
+            # Disable args.disable_headers before dispatching to upstream
             self.server.queue(
                 self.request.build(
                     disable_headers=self.config.disable_headers))
@@ -988,10 +998,10 @@ def access_log(self):
                 (self.client.addr[0],
                  self.client.addr[1],
                  text_(
-                    self.request.method),
-                    text_(host),
-                    text_(port),
-                    self.response.total_size))
+                     self.request.method),
+                 text_(host),
+                 text_(port),
+                 self.response.total_size))
         elif self.request.method:
             logger.info(
                 '%s:%s - %s %s:%s%s - %s %s - %s bytes' %
@@ -1015,7 +1025,7 @@ def connect_upstream(self, host, port):
             logger.debug('Connected to upstream %s:%s' % (host, port))
         except Exception as e:  # TimeoutError, socket.gaierror
             self.server.closed = True
-            raise ProxyConnectionFailed(host, port, repr(e))
+            raise ProxyConnectionFailed(host, port, repr(e)) from e
 
 
 class HttpWebServerPlugin(HttpProtocolBasePlugin):
@@ -1316,28 +1326,30 @@ def init_parser() -> argparse.ArgumentParser:
         type=str,
         default=DEFAULT_BASIC_AUTH,
         help='Default: No authentication. Specify colon separated user:password '
-        'to enable basic authentication.')
+             'to enable basic authentication.')
     parser.add_argument(
         '--client-recvbuf-size',
         type=int,
         default=DEFAULT_CLIENT_RECVBUF_SIZE,
         help='Default: 1 MB. Maximum amount of data received from the '
-        'client in a single recv() operation. Bump this '
-        'value for faster uploads at the expense of '
-        'increased RAM.')
+             'client in a single recv() operation. Bump this '
+             'value for faster uploads at the expense of '
+             'increased RAM.')
     parser.add_argument(
         '--disable-headers',
         type=str,
         default=COMMA.join(DEFAULT_DISABLE_HEADERS),
         help='Default: None.  Comma separated list of headers to remove before '
-        'dispatching client request to upstream server.')
+             'dispatching client request to upstream server.')
     parser.add_argument(
         '--disable-http-proxy',
         action='store_true',
         default=DEFAULT_DISABLE_HTTP_PROXY,
         help='Default: False.  Whether to disable proxy.HttpProxyPlugin.')
-    parser.add_argument('--hostname', type=str, default=DEFAULT_IPV4_HOSTNAME,
-                        help='Default: 127.0.0.1. Server IP address.')
+    parser.add_argument('--hostname',
+                        type=str,
+                        default=str(DEFAULT_IPV6_HOSTNAME),
+                        help='Default: ::1. Server IP address.')
     parser.add_argument('--ipv4', action='store_true', default=DEFAULT_IPV4,
                         help='Whether to listen on IPv4 address. '
                              'By default server only listens on IPv6.')
@@ -1351,8 +1363,8 @@ def init_parser() -> argparse.ArgumentParser:
         type=str,
         default=DEFAULT_LOG_LEVEL,
         help='Valid options: DEBUG, INFO (default), WARNING, ERROR, CRITICAL. '
-        'Both upper and lowercase values are allowed. '
-        'You may also simply use the leading character e.g. --log-level d')
+             'Both upper and lowercase values are allowed. '
+             'You may also simply use the leading character e.g. --log-level d')
     parser.add_argument('--log-file', type=str, default=DEFAULT_LOG_FILE,
                         help='Default: sys.stdout. Log file destination.')
     parser.add_argument('--log-format', type=str, default=DEFAULT_LOG_FORMAT,
@@ -1364,20 +1376,20 @@ def init_parser() -> argparse.ArgumentParser:
         type=int,
         default=DEFAULT_OPEN_FILE_LIMIT,
         help='Default: 1024. Maximum number of files (TCP connections) '
-        'that proxy.py can open concurrently.')
+             'that proxy.py can open concurrently.')
     parser.add_argument(
         '--pac-file',
         type=str,
         default=DEFAULT_PAC_FILE,
         help='A file (Proxy Auto Configuration) or string to serve when '
-        'the server receives a direct file request. '
-        'Using this option enables proxy.HttpWebServerPlugin.')
+             'the server receives a direct file request. '
+             'Using this option enables proxy.HttpWebServerPlugin.')
     parser.add_argument(
         '--pac-file-url-path',
         type=str,
         default=DEFAULT_PAC_FILE_URL_PATH,
         help='Default: %s. Web server path to serve the PAC file.' %
-        text_(DEFAULT_PAC_FILE_URL_PATH))
+             text_(DEFAULT_PAC_FILE_URL_PATH))
     parser.add_argument(
         '--pid-file',
         type=str,
@@ -1395,9 +1407,9 @@ def init_parser() -> argparse.ArgumentParser:
         type=int,
         default=DEFAULT_SERVER_RECVBUF_SIZE,
         help='Default: 1 MB. Maximum amount of data received from the '
-        'server in a single recv() operation. Bump this '
-        'value for faster downloads at the expense of '
-        'increased RAM.')
+             'server in a single recv() operation. Bump this '
+             'value for faster downloads at the expense of '
+             'increased RAM.')
     parser.add_argument(
         '--version',
         '-v',
@@ -1451,12 +1463,14 @@ def main(args) -> None:
             default_plugins += 'proxy.HttpWebServerPlugin,'
         config.plugins = load_plugins('%s%s' % (default_plugins, args.plugins))
 
-        server = MultiCoreRequestDispatcher(hostname=args.hostname,
-                                            port=args.port,
-                                            backlog=args.backlog,
-                                            ipv4=args.ipv4,
-                                            num_workers=args.num_workers,
-                                            config=config)
+        server = MultiCoreRequestDispatcher(
+            hostname=ipaddress.ip_address(
+                args.hostname),
+            port=args.port,
+            backlog=args.backlog,
+            ipv4=args.ipv4,
+            num_workers=args.num_workers,
+            config=config)
         if args.pid_file:
             with open(args.pid_file, 'wb') as pid_file:
                 pid_file.write(bytes_(str(os.getpid())))

tests.py

@@ -145,8 +145,9 @@ def baseTestCase(self, ipv4=True):
                     socket.SOCK_STREAM,
                     0)
                 sock.connect(
-                    (proxy.DEFAULT_IPV4_HOSTNAME if ipv4 else proxy.DEFAULT_IPV6_HOSTNAME,
-                     self.ipv4_port if ipv4 else self.ipv6_port))
+                    (str(
+                        proxy.DEFAULT_IPV4_HOSTNAME if ipv4 else proxy.DEFAULT_IPV6_HOSTNAME),
+                        self.ipv4_port if ipv4 else self.ipv6_port))
                 sock.sendall(b'HELLO')
                 data = sock.recv(proxy.DEFAULT_BUFFER_SIZE)
                 self.assertEqual(data, b'WORLD')
@@ -209,7 +210,7 @@ def testHttpProxyConnection(self, _mock_tcp_proxy):
                 with socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0) as sock:
                     try:
                         sock.connect(
-                            (proxy.DEFAULT_IPV4_HOSTNAME, self.tcp_port))
+                            (str(proxy.DEFAULT_IPV4_HOSTNAME), self.tcp_port))
                         sock.send(proxy.CRLF.join([
                             b'GET http://httpbin.org/get HTTP/1.1',
                             b'Host: httpbin.org',
@@ -975,7 +976,6 @@ def assert_data_queued(self, mock_server_connection, server):
             b'Host: localhost:%d' % self.http_server_port,
             b'Accept: */*',
             b'Via: %s' % b'1.1 proxy.py v%s' % proxy.version,
-            b'Connection: Close',
             proxy.CRLF
         ]))
 
@@ -1107,7 +1107,7 @@ def test_main(
         proxy.main(['--basic-auth', 'user:pass'])
         self.assertTrue(mock_set_open_file_limit.called)
         mock_multicore_dispatcher.assert_called_with(
-            hostname=proxy.DEFAULT_IPV4_HOSTNAME,
+            hostname=proxy.DEFAULT_IPV6_HOSTNAME,
             port=proxy.DEFAULT_PORT,
             ipv4=proxy.DEFAULT_IPV4,
             backlog=proxy.DEFAULT_BACKLOG,