| #!/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. |
| # |
| """Pure python code for finding unused ports on a host. |
| |
| This module provides a pick_unused_port() function. |
| It can also be called via the command line for use in shell scripts. |
| When called from the command line, it takes one optional argument, which, |
| if given, is sent to portserver instead of portpicker's PID. |
| To reserve a port for the lifetime of a bash script, use $BASHPID as this |
| argument. |
| |
| There is a race condition between picking a port and your application code |
| binding to it. The use of a port server to prevent that is recommended on |
| loaded test hosts running many tests at a time. |
| |
| If your code can accept a bound socket as input rather than being handed a |
| port number consider using socket.bind(('localhost', 0)) to bind to an |
| available port without a race condition rather than using this library. |
| |
| Typical usage: |
| test_port = portpicker.pick_unused_port() |
| """ |
| |
| # pylint: disable=consider-using-f-string |
| # Some people still use this on old Pythons despite our test matrix and |
| # supported versions. Be kind for now, until it gets in our way. |
| from __future__ import print_function |
| |
| import logging |
| import os |
| import random |
| import socket |
| import sys |
| import time |
| |
| _winapi = None # pylint: disable=invalid-name |
| if sys.platform == 'win32': |
| try: |
| import _winapi |
| except ImportError: |
| _winapi = None |
| |
| # The legacy Bind, IsPortFree, etc. names are not exported. |
| __all__ = ('bind', 'is_port_free', 'pick_unused_port', 'return_port', |
| 'add_reserved_port', 'get_port_from_port_server') |
| |
| _PROTOS = [(socket.SOCK_STREAM, socket.IPPROTO_TCP), |
| (socket.SOCK_DGRAM, socket.IPPROTO_UDP)] |
| |
| |
| # Ports that are currently available to be given out. |
| _free_ports = set() |
| |
| # Ports that are reserved or from the portserver that may be returned. |
| _owned_ports = set() |
| |
| # Ports that we chose randomly that may be returned. |
| _random_ports = set() |
| |
| |
| class NoFreePortFoundError(Exception): |
| """Exception indicating that no free port could be found.""" |
| |
| |
| def add_reserved_port(port): |
| """Add a port that was acquired by means other than the port server.""" |
| _free_ports.add(port) |
| |
| |
| def return_port(port): |
| """Return a port that is no longer being used so it can be reused.""" |
| if port in _random_ports: |
| _random_ports.remove(port) |
| elif port in _owned_ports: |
| _owned_ports.remove(port) |
| _free_ports.add(port) |
| elif port in _free_ports: |
| logging.info("Returning a port that was already returned: %s", port) |
| else: |
| logging.info("Returning a port that wasn't given by portpicker: %s", |
| port) |
| |
| |
| def bind(port, socket_type, socket_proto): |
| """Try to bind to a socket of the specified type, protocol, and port. |
| |
| This is primarily a helper function for PickUnusedPort, used to see |
| if a particular port number is available. |
| |
| For the port to be considered available, the kernel must support at least |
| one of (IPv6, IPv4), and the port must be available on each supported |
| family. |
| |
| Args: |
| port: The port number to bind to, or 0 to have the OS pick a free port. |
| socket_type: The type of the socket (ex: socket.SOCK_STREAM). |
| socket_proto: The protocol of the socket (ex: socket.IPPROTO_TCP). |
| |
| Returns: |
| The port number on success or None on failure. |
| """ |
| return _bind(port, socket_type, socket_proto) |
| |
| |
| def _bind(port, socket_type, socket_proto, return_socket=None, |
| return_family=socket.AF_INET6): |
| """Internal implementation of bind. |
| |
| Args: |
| port, socket_type, socket_proto: see bind(). |
| return_socket: If supplied, a list that we will append an open bound |
| reuseaddr socket on the port in question to. |
| return_family: The socket family to return in return_socket. |
| |
| Returns: |
| The port number on success or None on failure. |
| """ |
| # Our return family must come last when returning a bound socket |
| # as we cannot keep it bound while testing a bind on the other |
| # family with many network stack configurations. |
| if return_socket is None or return_family == socket.AF_INET: |
| socket_families = (socket.AF_INET6, socket.AF_INET) |
| elif return_family == socket.AF_INET6: |
| socket_families = (socket.AF_INET, socket.AF_INET6) |
| else: |
| raise ValueError('unknown return_family %s' % return_family) |
| got_socket = False |
| for family in socket_families: |
| try: |
| sock = socket.socket(family, socket_type, socket_proto) |
| got_socket = True |
| except socket.error: |
| continue |
| try: |
| sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) |
| sock.bind(('', port)) |
| if socket_type == socket.SOCK_STREAM: |
| sock.listen(1) |
| port = sock.getsockname()[1] |
| except socket.error: |
| return None |
| finally: |
| if return_socket is None or family != return_family: |
| sock.close() |
| if return_socket is not None and family == return_family: |
| return_socket.append(sock) |
| break # Final iteration due to pre-loop logic; don't close. |
| return port if got_socket else None |
| |
| |
| def is_port_free(port): |
| """Check if specified port is free. |
| |
| Args: |
| port: integer, port to check |
| |
| Returns: |
| bool, whether port is free to use for both TCP and UDP. |
| """ |
| return _is_port_free(port) |
| |
| |
| def _is_port_free(port, return_sockets=None): |
| """Internal implementation of is_port_free. |
| |
| Args: |
| port: integer, port to check |
| return_sockets: If supplied, a list that we will append open bound |
| sockets on the port in question to rather than closing them. |
| |
| Returns: |
| bool, whether port is free to use for both TCP and UDP. |
| """ |
| return (_bind(port, *_PROTOS[0], return_socket=return_sockets) and |
| _bind(port, *_PROTOS[1], return_socket=return_sockets)) |
| |
| |
| def pick_unused_port(pid=None, portserver_address=None): |
| """Picks an unused port and reserves it for use by a given process id. |
| |
| Args: |
| pid: PID to tell the portserver to associate the reservation with. If |
| None, the current process's PID is used. |
| portserver_address: The address (path) of a unix domain socket |
| with which to connect to a portserver, a leading '@' |
| character indicates an address in the "abstract namespace". OR |
| On systems without socket.AF_UNIX, this is an AF_INET address. |
| If None, or no port is returned by the portserver at the provided |
| address, the environment will be checked for a PORTSERVER_ADDRESS |
| variable. If that is not set, no port server will be used. |
| |
| If no portserver is used, no pid based reservation is managed by any |
| central authority. Race conditions and duplicate assignments may occur. |
| |
| Returns: |
| A port number that is unused on both TCP and UDP. |
| |
| Raises: |
| NoFreePortFoundError: No free port could be found. |
| """ |
| return _pick_unused_port(pid, portserver_address) |
| |
| |
| def _pick_unused_port(pid=None, portserver_address=None, |
| noserver_bind_timeout=0): |
| """Internal implementation of pick_unused_port. |
| |
| Args: |
| pid, portserver_address: See pick_unused_port(). |
| noserver_bind_timeout: If no portserver was used, this is the number of |
| seconds we will attempt to keep a child process around with the ports |
| returned open and bound SO_REUSEADDR style to help avoid race condition |
| port reuse. A non-zero value attempts os.fork(). Do not use it in a |
| multithreaded process. |
| """ |
| try: # Instead of `if _free_ports:` to handle the race condition. |
| port = _free_ports.pop() |
| except KeyError: |
| pass |
| else: |
| _owned_ports.add(port) |
| return port |
| # Provide access to the portserver on an opt-in basis. |
| if portserver_address: |
| port = get_port_from_port_server(portserver_address, pid=pid) |
| if port: |
| return port |
| if 'PORTSERVER_ADDRESS' in os.environ: |
| port = get_port_from_port_server(os.environ['PORTSERVER_ADDRESS'], |
| pid=pid) |
| if port: |
| return port |
| return _pick_unused_port_without_server(bind_timeout=noserver_bind_timeout) |
| |
| |
| def _spawn_bound_port_holding_daemon(port, bound_sockets, timeout): |
| """If possible, fork()s a daemon process to hold bound_sockets open. |
| |
| Emits a warning to stderr if it cannot. |
| |
| Args: |
| port: The port number the sockets are bound to (informational). |
| bound_sockets: The list of bound sockets our child process will hold |
| open. If the list is empty, no action is taken. |
| timeout: A positive number of seconds the child should sleep for before |
| closing the sockets and exiting. |
| """ |
| if bound_sockets and timeout > 0: |
| try: |
| fork_pid = os.fork() # This concept only works on POSIX. |
| except Exception as err: # pylint: disable=broad-except |
| print('WARNING: Cannot timeout unbinding close of port', port, |
| ' closing on exit. -', err, file=sys.stderr) |
| else: |
| if fork_pid == 0: |
| # This child process inherits and holds bound_sockets open |
| # for bind_timeout seconds. |
| try: |
| # Close the stdio fds as may be connected to |
| # a pipe that will cause a grandparent process |
| # to wait on before returning. (cl/427587550) |
| os.close(sys.stdin.fileno()) |
| os.close(sys.stdout.fileno()) |
| os.close(sys.stderr.fileno()) |
| time.sleep(timeout) |
| for held_socket in bound_sockets: |
| held_socket.close() |
| finally: |
| sys.exit(0) |
| |
| |
| def _pick_unused_port_without_server(bind_timeout=0): |
| """Pick an available network port without the help of a port server. |
| |
| This code ensures that the port is available on both TCP and UDP. |
| |
| This function is an implementation detail of PickUnusedPort(), and |
| should not be called by code outside of this module. |
| |
| Args: |
| bind_timeout: number of seconds to attempt to keep a child process |
| process around bound SO_REUSEADDR style to the port. If we cannot |
| do that we emit a warning to stderr. |
| |
| Returns: |
| A port number that is unused on both TCP and UDP. |
| |
| Raises: |
| NoFreePortFoundError: No free port could be found. |
| """ |
| # Next, try a few times to get an OS-assigned port. |
| # Ambrose discovered that on the 2.6 kernel, calling Bind() on UDP socket |
| # returns the same port over and over. So always try TCP first. |
| port = None |
| bound_sockets = [] if bind_timeout > 0 else None |
| for _ in range(10): |
| # Ask the OS for an unused port. |
| port = _bind(0, socket.SOCK_STREAM, socket.IPPROTO_TCP, bound_sockets) |
| # Check if this port is unused on the other protocol. |
| if (port and port not in _random_ports and |
| _bind(port, socket.SOCK_DGRAM, socket.IPPROTO_UDP, bound_sockets)): |
| _random_ports.add(port) |
| _spawn_bound_port_holding_daemon(port, bound_sockets, bind_timeout) |
| return port |
| if bound_sockets: |
| for held_socket in bound_sockets: |
| held_socket.close() |
| del bound_sockets[:] |
| |
| # Try random ports as a last resort. |
| rng = random.Random() |
| for _ in range(10): |
| port = int(rng.randrange(15000, 25000)) |
| if port not in _random_ports: |
| if _is_port_free(port, bound_sockets): |
| _random_ports.add(port) |
| _spawn_bound_port_holding_daemon( |
| port, bound_sockets, bind_timeout) |
| return port |
| if bound_sockets: |
| for held_socket in bound_sockets: |
| held_socket.close() |
| del bound_sockets[:] |
| |
| # Give up. |
| raise NoFreePortFoundError() |
| |
| |
| def _posix_get_port_from_port_server(portserver_address, pid): |
| # An AF_UNIX address may start with a zero byte, in which case it is in the |
| # "abstract namespace", and doesn't have any filesystem representation. |
| # See 'man 7 unix' for details. |
| # The convention is to write '@' in the address to represent this zero byte. |
| if portserver_address[0] == '@': |
| portserver_address = '\0' + portserver_address[1:] |
| |
| try: |
| # Create socket. |
| if hasattr(socket, 'AF_UNIX'): |
| sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) |
| else: |
| # fallback to AF_INET if this is not unix |
| sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) |
| try: |
| # Connect to portserver. |
| sock.connect(portserver_address) |
| |
| # Write request. |
| sock.sendall(('%d\n' % pid).encode('ascii')) |
| |
| # Read response. |
| # 1K should be ample buffer space. |
| return sock.recv(1024) |
| finally: |
| sock.close() |
| except socket.error as error: |
| print('Socket error when connecting to portserver:', error, |
| file=sys.stderr) |
| return None |
| |
| |
| def _windows_get_port_from_port_server(portserver_address, pid): |
| if portserver_address[0] == '@': |
| portserver_address = '\\\\.\\pipe\\' + portserver_address[1:] |
| |
| try: |
| handle = _winapi.CreateFile( |
| portserver_address, |
| _winapi.GENERIC_READ | _winapi.GENERIC_WRITE, |
| 0, |
| 0, |
| _winapi.OPEN_EXISTING, |
| 0, |
| 0) |
| |
| _winapi.WriteFile(handle, ('%d\n' % pid).encode('ascii')) |
| data, _ = _winapi.ReadFile(handle, 6, 0) |
| return data |
| except FileNotFoundError as error: |
| print('File error when connecting to portserver:', error, |
| file=sys.stderr) |
| return None |
| |
| |
| def get_port_from_port_server(portserver_address, pid=None): |
| """Request a free a port from a system-wide portserver. |
| |
| This follows a very simple portserver protocol: |
| The request consists of our pid (in ASCII) followed by a newline. |
| The response is a port number and a newline, 0 on failure. |
| |
| This function is an implementation detail of pick_unused_port(). |
| It should not normally be called by code outside of this module. |
| |
| Args: |
| portserver_address: The address (path) of a unix domain socket |
| with which to connect to the portserver. A leading '@' |
| character indicates an address in the "abstract namespace." |
| On systems without socket.AF_UNIX, this is an AF_INET address. |
| pid: The PID to tell the portserver to associate the reservation with. |
| If None, the current process's PID is used. |
| |
| Returns: |
| The port number on success or None on failure. |
| """ |
| if not portserver_address: |
| return None |
| |
| if pid is None: |
| pid = os.getpid() |
| |
| if _winapi: |
| buf = _windows_get_port_from_port_server(portserver_address, pid) |
| else: |
| buf = _posix_get_port_from_port_server(portserver_address, pid) |
| |
| if buf is None: |
| return None |
| |
| try: |
| port = int(buf.split(b'\n')[0]) |
| except ValueError: |
| print('Portserver failed to find a port.', file=sys.stderr) |
| return None |
| _owned_ports.add(port) |
| return port |
| |
| |
| # Legacy APIs. |
| # pylint: disable=invalid-name |
| Bind = bind |
| GetPortFromPortServer = get_port_from_port_server |
| IsPortFree = is_port_free |
| PickUnusedPort = pick_unused_port |
| # pylint: enable=invalid-name |
| |
| |
| def main(argv): |
| """If passed an arg, treat it as a PID, otherwise we use getppid(). |
| |
| A second optional argument can be a bind timeout in seconds that will be |
| used ONLY if no portserver is found. We attempt to leave a process around |
| holding the port open and bound with SO_REUSEADDR set for timeout seconds. |
| If the timeout bind was not possible, a warning is emitted to stderr. |
| |
| #!/bin/bash |
| port="$(python -m portpicker $$ 1.23)" |
| test_my_server "$port" |
| |
| This will pick a port for your script's PID and assign it to $port, if no |
| portserver was used, it attempts to keep a socket bound to $port for 1.23 |
| seconds after the portpicker process has exited. This is a convenient hack |
| to attempt to prevent port reallocation during scripts outside of |
| portserver managed environments. |
| |
| Older versions of the portpicker CLI ignore everything beyond the first arg. |
| Older versions also used getpid() instead of getppid(), so script users are |
| strongly encouraged to be explicit and pass $$ or your languages equivalent |
| to associate the port with the PID of the controlling process. |
| """ |
| # Our command line is trivial so I avoid an argparse import. If we ever |
| # grow more than 1-2 args, switch to a using argparse. |
| if '-h' in argv or '--help' in argv: |
| print(argv[0], 'usage:\n') |
| import inspect |
| print(inspect.getdoc(main)) |
| sys.exit(1) |
| pid=int(argv[1]) if len(argv) > 1 else os.getppid() |
| bind_timeout=float(argv[2]) if len(argv) > 2 else 0 |
| port = _pick_unused_port(pid=pid, noserver_bind_timeout=bind_timeout) |
| if not port: |
| sys.exit(1) |
| print(port) |
| |
| |
| if __name__ == '__main__': |
| main(sys.argv) |