Allow multiport with ephemeral & unix socket support (#1078)

* Allow multiport with ephemeral & unix socket support

* Fix unix tests

Author: abhinavsingh

proxy/core/listener/base.py

@@ -8,7 +8,6 @@
     :copyright: (c) 2013-present by Abhinav Singh and contributors.
     :license: BSD, see LICENSE for more details.
 """
-import os
 import socket
 import logging
 import argparse
@@ -60,5 +59,3 @@ def setup(self) -> None:
     def shutdown(self) -> None:
         assert self._socket
         self._socket.close()
-        if self.flags.unix_socket_path:
-            os.remove(self.flags.unix_socket_path)

proxy/core/listener/tcp.py

@@ -65,7 +65,10 @@ def __init__(self, *args: Any, port: Optional[int] = None, **kwargs: Any) -> Non
         super().__init__(*args, **kwargs)
 
     def listen(self) -> socket.socket:
-        sock = socket.socket(self.flags.family, socket.SOCK_STREAM)
+        sock = socket.socket(
+            socket.AF_INET6 if self.flags.hostname.version == 6 else socket.AF_INET,
+            socket.SOCK_STREAM,
+        )
         sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
         sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
         # s.setsockopt(socket.SOL_TCP, socket.TCP_FASTOPEN, 5)

proxy/core/listener/unix.py

@@ -8,6 +8,7 @@
     :copyright: (c) 2013-present by Abhinav Singh and contributors.
     :license: BSD, see LICENSE for more details.
 """
+import os
 import socket
 import logging
 
@@ -30,9 +31,8 @@ class UnixSocketListener(BaseListener):
     """Unix socket domain listener."""
 
     def listen(self) -> socket.socket:
-        sock = socket.socket(self.flags.family, socket.SOCK_STREAM)
+        sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
         sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-        sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
         sock.bind(self.flags.unix_socket_path)
         sock.listen(self.flags.backlog)
         sock.setblocking(False)
@@ -41,3 +41,8 @@ def listen(self) -> socket.socket:
             self.flags.unix_socket_path,
         )
         return sock
+
+    def shutdown(self) -> None:
+        super().shutdown()
+        if self.flags.unix_socket_path:
+            os.remove(self.flags.unix_socket_path)

proxy/proxy.py

@@ -195,10 +195,24 @@ def setup(self) -> None:
         # Override flags.port to match the actual port
         # we are listening upon.  This is necessary to preserve
         # the server port when `--port=0` is used.
-        self.flags.port = cast(
-            'TcpSocketListener',
-            self.listeners.pool[0],
-        )._port
+        if not self.flags.unix_socket_path:
+            self.flags.port = cast(
+                'TcpSocketListener',
+                self.listeners.pool[0],
+            )._port
+        # --ports flag can also use 0 as value for ephemeral port selection.
+        # Here, we override flags.ports to reflect actual listening ports.
+        ports = []
+        offset = 1 if self.flags.unix_socket_path or self.flags.port else 0
+        for index in range(offset, offset + len(self.flags.ports)):
+            ports.append(
+                cast(
+                    'TcpSocketListener',
+                    self.listeners.pool[index],
+                )._port,
+            )
+        self.flags.ports = ports
+        # Write ports to port file
         self._write_port_file()
         # Setup EventManager
         if self.flags.enable_events:
@@ -278,7 +292,12 @@ def _delete_pid_file(self) -> None:
     def _write_port_file(self) -> None:
         if self.flags.port_file:
             with open(self.flags.port_file, 'wb') as port_file:
-                port_file.write(bytes_(self.flags.port))
+                if not self.flags.unix_socket_path:
+                    port_file.write(bytes_(self.flags.port))
+                    port_file.write(b'\n')
+                for port in self.flags.ports:
+                    port_file.write(bytes_(port))
+                    port_file.write(b'\n')
 
     def _delete_port_file(self) -> None:
         if self.flags.port_file \

tests/core/test_listener.py

@@ -69,15 +69,11 @@ def test_unix_path_listener(self, mock_socket: mock.Mock, mock_remove: mock.Mock
             socket.AF_UNIX,
             socket.SOCK_STREAM,
         )
-        self.assertEqual(sock.setsockopt.call_count, 2)
+        self.assertEqual(sock.setsockopt.call_count, 1)
         self.assertEqual(
             sock.setsockopt.call_args_list[0][0],
             (socket.SOL_SOCKET, socket.SO_REUSEADDR, 1),
         )
-        self.assertEqual(
-            sock.setsockopt.call_args_list[1][0],
-            (socket.IPPROTO_TCP, socket.TCP_NODELAY, 1),
-        )
         sock.bind.assert_called_with(sock_path)
         sock.listen.assert_called_with(flags.backlog)
         sock.setblocking.assert_called_with(False)