src/tests/portpicker_test.py

#!/usr/bin/python3
#
# Copyright 2007 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""Unittests for portpicker."""

# pylint: disable=invalid-name,protected-access,missing-class-docstring,missing-function-docstring

from contextlib import ExitStack
import errno
import os
import socket
import subprocess
import sys
import time
import unittest
from unittest import mock

import portpicker
_winapi = portpicker._winapi

# pylint: disable=invalid-name,protected-access,missing-class-docstring,missing-function-docstring


class CommonTestMixin:
    def IsUnusedTCPPort(self, port):
        return self._bind(port, socket.SOCK_STREAM, socket.IPPROTO_TCP)

    def IsUnusedUDPPort(self, port):
        return self._bind(port, socket.SOCK_DGRAM, socket.IPPROTO_UDP)

    def setUp(self):
        super().setUp()
        # So we can Bind even if portpicker.bind is stubbed out.
        self._bind = portpicker.bind
        portpicker._owned_ports.clear()
        portpicker._free_ports.clear()
        portpicker._random_ports.clear()


@unittest.skipIf(
        ('PORTSERVER_ADDRESS' not in os.environ) and
        not hasattr(socket, 'AF_UNIX'),
        'no existing port server; test launching code requires AF_UNIX.')
class PickUnusedPortTestWithAPortServer(CommonTestMixin, unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.portserver_process = None
        if 'PORTSERVER_ADDRESS' not in os.environ:
            # Launch a portserver child process for our tests to use if we are
            # able to. Obviously not host-exclusive, but good for integration
            # testing purposes on CI without a portserver of its own.
            cls.portserver_address = '@pid%d-test-ports' % os.getpid()
            try:
                cls.portserver_process = subprocess.Popen(
                        ['portserver.py',  # Installed in PATH within the venv.
                         '--portserver_address=%s' % cls.portserver_address])
            except EnvironmentError as err:
                raise unittest.SkipTest(
                        'Unable to launch portserver.py: %s' % err)
            linux_addr = '\0' + cls.portserver_address[1:]  # The @ means 0.
            # loop for a few seconds waiting for that socket to work.
            err = '???'
            for _ in range(123):
                time.sleep(0.05)
                try:
                    ps_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
                    ps_sock.connect(linux_addr)
                except socket.error as err:  # pylint: disable=unused-variable
                    continue
                ps_sock.close()
                break
            else:
                # The socket failed or never accepted connections, assume our
                # portserver setup attempt failed and bail out.
                if cls.portserver_process.poll() is not None:
                    cls.portserver_process.kill()
                    cls.portserver_process.wait()
                cls.portserver_process = None
                raise unittest.SkipTest(
                        'Unable to connect to our own portserver.py: %s' % err)
            # Point child processes at our shiny portserver process.
            os.environ['PORTSERVER_ADDRESS'] = cls.portserver_address

    @classmethod
    def tearDownClass(cls):
        if cls.portserver_process:
            if os.environ.get('PORTSERVER_ADDRESS') == cls.portserver_address:
                del os.environ['PORTSERVER_ADDRESS']
            if cls.portserver_process.poll() is None:
                cls.portserver_process.kill()
                cls.portserver_process.wait()
            cls.portserver_process = None

    def testPickUnusedCanSuccessfullyUsePortServer(self):

        with mock.patch.object(portpicker, '_pick_unused_port_without_server'):
            portpicker._pick_unused_port_without_server.side_effect = (
                Exception('eek!')
            )

            # Since _PickUnusedPortWithoutServer() raises an exception, if we
            # can successfully obtain a port, the portserver must be working.
            port = portpicker.pick_unused_port()
            self.assertTrue(self.IsUnusedTCPPort(port))
            self.assertTrue(self.IsUnusedUDPPort(port))

    def testPickUnusedCanSuccessfullyUsePortServerAddressKwarg(self):

        with mock.patch.object(portpicker, '_pick_unused_port_without_server'):
            portpicker._pick_unused_port_without_server.side_effect = (
                Exception('eek!')
            )

            # Since _PickUnusedPortWithoutServer() raises an exception, and
            # we've temporarily removed PORTSERVER_ADDRESS from os.environ, if
            # we can successfully obtain a port, the portserver must be working.
            addr = os.environ.pop('PORTSERVER_ADDRESS')
            try:
                port = portpicker.pick_unused_port(portserver_address=addr)
                self.assertTrue(self.IsUnusedTCPPort(port))
                self.assertTrue(self.IsUnusedUDPPort(port))
            finally:
                os.environ['PORTSERVER_ADDRESS'] = addr

    def testGetPortFromPortServer(self):
        """Exercise the get_port_from_port_server() helper function."""
        for _ in range(10):
            port = portpicker.get_port_from_port_server(
                os.environ['PORTSERVER_ADDRESS'])
            self.assertTrue(self.IsUnusedTCPPort(port))
            self.assertTrue(self.IsUnusedUDPPort(port))


class PickUnusedPortTest(CommonTestMixin, unittest.TestCase):

    def testPickUnusedPortActuallyWorks(self):
        """This test can be flaky."""
        for _ in range(10):
            port = portpicker.pick_unused_port()
            self.assertTrue(self.IsUnusedTCPPort(port))
            self.assertTrue(self.IsUnusedUDPPort(port))

    def testSendsPidToPortServer(self):
        with ExitStack() as stack:
            if _winapi:
                create_file_mock = mock.Mock()
                create_file_mock.return_value = 0
                read_file_mock = mock.Mock()
                write_file_mock = mock.Mock()
                read_file_mock.return_value = (b'42768\n', 0)
                stack.enter_context(
                    mock.patch('_winapi.CreateFile', new=create_file_mock))
                stack.enter_context(
                    mock.patch('_winapi.WriteFile', new=write_file_mock))
                stack.enter_context(
                    mock.patch('_winapi.ReadFile', new=read_file_mock))
                port = portpicker.get_port_from_port_server(
                    'portserver', pid=1234)
                write_file_mock.assert_called_once_with(0, b'1234\n')
            else:
                server = mock.Mock()
                server.recv.return_value = b'42768\n'
                stack.enter_context(
                    mock.patch.object(socket, 'socket', return_value=server))
                port = portpicker.get_port_from_port_server(
                    'portserver', pid=1234)
                server.sendall.assert_called_once_with(b'1234\n')

        self.assertEqual(port, 42768)

    def testPidDefaultsToOwnPid(self):
        with ExitStack() as stack:
            stack.enter_context(
                mock.patch.object(os, 'getpid', return_value=9876))

            if _winapi:
                create_file_mock = mock.Mock()
                create_file_mock.return_value = 0
                read_file_mock = mock.Mock()
                write_file_mock = mock.Mock()
                read_file_mock.return_value = (b'52768\n', 0)
                stack.enter_context(
                    mock.patch('_winapi.CreateFile', new=create_file_mock))
                stack.enter_context(
                    mock.patch('_winapi.WriteFile', new=write_file_mock))
                stack.enter_context(
                    mock.patch('_winapi.ReadFile', new=read_file_mock))
                port = portpicker.get_port_from_port_server('portserver')
                write_file_mock.assert_called_once_with(0, b'9876\n')
            else:
                server = mock.Mock()
                server.recv.return_value = b'52768\n'
                stack.enter_context(
                    mock.patch.object(socket, 'socket', return_value=server))
                port = portpicker.get_port_from_port_server('portserver')
                server.sendall.assert_called_once_with(b'9876\n')

        self.assertEqual(port, 52768)

    @mock.patch.dict(os.environ,{'PORTSERVER_ADDRESS': 'portserver'})
    def testReusesPortServerPorts(self):
        with ExitStack() as stack:
            if _winapi:
                read_file_mock = mock.Mock()
                read_file_mock.side_effect = [
                    (b'12345\n', 0),
                    (b'23456\n', 0),
                    (b'34567\n', 0),
                ]
                stack.enter_context(mock.patch('_winapi.CreateFile'))
                stack.enter_context(mock.patch('_winapi.WriteFile'))
                stack.enter_context(
                    mock.patch('_winapi.ReadFile', new=read_file_mock))
            else:
                server = mock.Mock()
                server.recv.side_effect = [b'12345\n', b'23456\n', b'34567\n']
                stack.enter_context(
                    mock.patch.object(socket, 'socket', return_value=server))

            self.assertEqual(portpicker.pick_unused_port(), 12345)
            self.assertEqual(portpicker.pick_unused_port(), 23456)
            portpicker.return_port(12345)
            self.assertEqual(portpicker.pick_unused_port(), 12345)

    @mock.patch.dict(os.environ,{'PORTSERVER_ADDRESS': ''})
    def testDoesntReuseRandomPorts(self):
        ports = set()
        for _ in range(10):
            try:
                port = portpicker.pick_unused_port()
            except portpicker.NoFreePortFoundError:
                # This sometimes happens when not using portserver. Just
                # skip to the next attempt.
                continue
            ports.add(port)
            portpicker.return_port(port)
        self.assertGreater(len(ports), 5)  # Allow some random reuse.

    def testReturnsReservedPorts(self):
        with mock.patch.object(portpicker, '_pick_unused_port_without_server'):
            portpicker._pick_unused_port_without_server.side_effect = (
                Exception('eek!'))
            # Arbitrary port. In practice you should get this from somewhere
            # that assigns ports.
            reserved_port = 28465
            portpicker.add_reserved_port(reserved_port)
            ports = set()
            for _ in range(10):
                port = portpicker.pick_unused_port()
                ports.add(port)
                portpicker.return_port(port)
            self.assertEqual(len(ports), 1)
            self.assertEqual(ports.pop(), reserved_port)

    @mock.patch.dict(os.environ,{'PORTSERVER_ADDRESS': ''})
    def testFallsBackToRandomAfterRunningOutOfReservedPorts(self):
        # Arbitrary port. In practice you should get this from somewhere
        # that assigns ports.
        reserved_port = 23456
        portpicker.add_reserved_port(reserved_port)
        self.assertEqual(portpicker.pick_unused_port(), reserved_port)
        self.assertNotEqual(portpicker.pick_unused_port(), reserved_port)

    def testRandomlyChosenPorts(self):
        # Unless this box is under an overwhelming socket load, this test
        # will heavily exercise the "pick a port randomly" part of the
        # port picking code, but may never hit the "OS assigns a port"
        # code.
        ports = 0
        for _ in range(100):
            try:
                port = portpicker._pick_unused_port_without_server()
            except portpicker.NoFreePortFoundError:
                # Without the portserver, pick_unused_port can sometimes fail
                # to find a free port. Check that it passes most of the time.
                continue
            self.assertTrue(self.IsUnusedTCPPort(port))
            self.assertTrue(self.IsUnusedUDPPort(port))
            ports += 1
        # Getting a port shouldn't have failed very often, even on machines
        # with a heavy socket load.
        self.assertGreater(ports, 95)

    def testOSAssignedPorts(self):
        self.last_assigned_port = None

        def error_for_explicit_ports(port, socket_type, socket_proto):
            # Only successfully return a port if an OS-assigned port is
            # requested, or if we're checking that the last OS-assigned port
            # is unused on the other protocol.
            if port in (0, self.last_assigned_port):
                self.last_assigned_port = self._bind(port, socket_type,
                                                     socket_proto)
                return self.last_assigned_port
            return None

        with mock.patch.object(portpicker, 'bind', error_for_explicit_ports):
            # Without server, this can be little flaky, so check that it
            # passes most of the time.
            ports = 0
            for _ in range(100):
                try:
                    port = portpicker._pick_unused_port_without_server()
                except portpicker.NoFreePortFoundError:
                    continue
                self.assertTrue(self.IsUnusedTCPPort(port))
                self.assertTrue(self.IsUnusedUDPPort(port))
                ports += 1
            self.assertGreater(ports, 70)

    def pickUnusedPortWithoutServer(self):
        # Try a few times to pick a port, to avoid flakiness and to make sure
        # the code path we want was exercised.
        for _ in range(5):
            try:
                port = portpicker._pick_unused_port_without_server()
            except portpicker.NoFreePortFoundError:
                continue
            else:
                self.assertTrue(self.IsUnusedTCPPort(port))
                self.assertTrue(self.IsUnusedUDPPort(port))
                return
        self.fail("Failed to find a free port")

    def testPickPortsWithoutServer(self):
        # Test the first part of _pick_unused_port_without_server, which
        # tries a few random ports and checks is_port_free.
        self.pickUnusedPortWithoutServer()

        # Now test the second part, the fallback from above, which asks the
        # OS for a port.
        def mock_port_free(unused_port):
            return False

        with mock.patch.object(portpicker, 'is_port_free', mock_port_free):
            self.pickUnusedPortWithoutServer()

    def checkIsPortFree(self):
        """This might be flaky unless this test is run with a portserver."""
        # The port should be free initially.
        port = portpicker.pick_unused_port()
        self.assertTrue(portpicker.is_port_free(port))

        cases = [
            (socket.AF_INET,  socket.SOCK_STREAM, None),
            (socket.AF_INET6, socket.SOCK_STREAM, 1),
            (socket.AF_INET,  socket.SOCK_DGRAM,  None),
            (socket.AF_INET6, socket.SOCK_DGRAM,  1),
        ]

        # Using v6only=0 on Windows doesn't result in collisions
        if not _winapi:
            cases.extend([
                (socket.AF_INET6, socket.SOCK_STREAM, 0),
                (socket.AF_INET6, socket.SOCK_DGRAM,  0),
            ])

        for (sock_family, sock_type, v6only) in cases:
            # Occupy the port on a subset of possible protocols.
            try:
                sock = socket.socket(sock_family, sock_type, 0)
            except socket.error:
                print('Kernel does not support sock_family=%d' % sock_family,
                      file=sys.stderr)
                # Skip this case, since we cannot occupy a port.
                continue

            if not hasattr(socket, 'IPPROTO_IPV6'):
                v6only = None

            if v6only is not None:
                try:
                    sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY,
                                    v6only)
                except socket.error:
                    print('Kernel does not support IPV6_V6ONLY=%d' % v6only,
                          file=sys.stderr)
                    # Don't care; just proceed with the default.

            # Socket may have been taken in the mean time, so catch the
            # socket.error with errno set to EADDRINUSE and skip this
            # attempt.
            try:
                sock.bind(('', port))
            except socket.error as e:
                if e.errno == errno.EADDRINUSE:
                    raise portpicker.NoFreePortFoundError
                raise

            # The port should be busy.
            self.assertFalse(portpicker.is_port_free(port))
            sock.close()

            # Now it's free again.
            self.assertTrue(portpicker.is_port_free(port))

    def testIsPortFree(self):
        # This can be quite flaky on a busy host, so try a few times.
        for _ in range(10):
            try:
                self.checkIsPortFree()
            except portpicker.NoFreePortFoundError:
                pass
            else:
                return
        self.fail("checkPortIsFree failed every time.")

    def testIsPortFreeException(self):
        port = portpicker.pick_unused_port()
        with mock.patch.object(socket, 'socket') as mock_sock:
            mock_sock.side_effect = socket.error('fake socket error', 0)
            self.assertFalse(portpicker.is_port_free(port))

    def testThatLegacyCapWordsAPIsExist(self):
        """The original APIs were CapWords style, 1.1 added PEP8 names."""
        self.assertEqual(portpicker.bind, portpicker.Bind)
        self.assertEqual(portpicker.is_port_free, portpicker.IsPortFree)
        self.assertEqual(portpicker.pick_unused_port, portpicker.PickUnusedPort)
        self.assertEqual(portpicker.get_port_from_port_server,
                         portpicker.GetPortFromPortServer)


def get_open_listen_tcp_ports():
    netstat = subprocess.run(['netstat', '-lnt'], capture_output=True,
                             encoding='utf-8')
    if netstat.returncode != 0:
        raise unittest.SkipTest('Unable to run netstat -lnt to list binds.')
    rows = (line.split() for line in netstat.stdout.splitlines())
    listen_addrs = (row[3] for row in rows if row[0].startswith('tcp'))
    listen_ports = [int(addr.split(':')[-1]) for addr in listen_addrs]
    return listen_ports


@unittest.skipUnless((sys.executable and os.access(sys.executable, os.X_OK))
                     or (os.environ.get('TEST_PORTPICKER_CLI') and
                         os.access(os.environ['TEST_PORTPICKER_CLI'], os.X_OK)),
                     'sys.executable portpicker.__file__ not launchable and '
                     ' no TEST_PORTPICKER_CLI supplied.')
class PortpickerCommandLineTests(unittest.TestCase):
    def setUp(self):
        self.main_py = portpicker.__file__

    def _run_portpicker(self, pp_args, env_override=None):
        env = dict(os.environ)
        if env_override:
            env.update(env_override)
        if os.environ.get('TEST_PORTPICKER_CLI'):
            pp_command = [os.environ['TEST_PORTPICKER_CLI']]
        else:
            pp_command = [sys.executable, '-m', 'portpicker']
        return subprocess.run(pp_command + pp_args,
                              capture_output=True,
                              env=env,
                              encoding='utf-8',
                              check=False)

    def test_command_line_help(self):
        cmd = self._run_portpicker(['-h'])
        self.assertNotEqual(0, cmd.returncode)
        self.assertIn('usage', cmd.stdout)
        self.assertIn('passed an arg', cmd.stdout)
        cmd = self._run_portpicker(['--help'])
        self.assertNotEqual(0, cmd.returncode)
        self.assertIn('usage', cmd.stdout)
        self.assertIn('passed an arg', cmd.stdout)

    def test_command_line_help_text_dedented(self):
        cmd = self._run_portpicker(['-h'])
        self.assertNotEqual(0, cmd.returncode)
        self.assertIn('\nIf passed an arg', cmd.stdout)
        self.assertIn('\n  #!/bin/bash', cmd.stdout)
        self.assertIn('\nOlder versions ', cmd.stdout)

    def test_command_line_interface(self):
        cmd = self._run_portpicker([str(os.getpid())])
        cmd.check_returncode()
        port = int(cmd.stdout)
        self.assertNotEqual(0, port, msg=cmd)
        listen_ports = sorted(get_open_listen_tcp_ports())
        self.assertNotIn(port, listen_ports, msg='expected nothing to be bound to port.')

    def test_command_line_interface_no_portserver(self):
        cmd = self._run_portpicker([str(os.getpid())],
                                   env_override={'PORTSERVER_ADDRESS': ''})
        cmd.check_returncode()
        port = int(cmd.stdout)
        self.assertNotEqual(0, port, msg=cmd)
        listen_ports = sorted(get_open_listen_tcp_ports())
        self.assertNotIn(port, listen_ports, msg='expected nothing to be bound to port.')

    def test_command_line_interface_no_portserver_bind_timeout(self):
        # This test is timing sensitive and leaves that bind process hanging
        # around consuming resources until it dies on its own unless the test
        # runner kills the process group upon exit.
        timeout = 9.5
        before = time.monotonic()
        cmd = self._run_portpicker([str(os.getpid()), str(timeout)],
                                   env_override={'PORTSERVER_ADDRESS': ''})
        self.assertEqual(0, cmd.returncode, msg=(cmd.stdout, cmd.stderr))
        port = int(cmd.stdout)
        self.assertNotEqual(0, port, msg=cmd)
        if 'WARNING' in cmd.stderr:
            raise unittest.SkipTest('bind timeout not supported on this platform.')
        listen_ports = sorted(get_open_listen_tcp_ports())
        self.assertIn(port, listen_ports, msg='expected port to be bound. '
                      '%f seconds elapsed of %f bind timeout.' %
                      (time.monotonic() - before, timeout))


if __name__ == '__main__':
    unittest.main()