1#!/usr/bin/python3
2#
3# Copyright 2007 Google Inc. All Rights Reserved.
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#     http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16#
17"""Pure python code for finding unused ports on a host.
18
19This module provides a pick_unused_port() function.
20It can also be called via the command line for use in shell scripts.
21When called from the command line, it takes one optional argument, which,
22if given, is sent to portserver instead of portpicker's PID.
23To reserve a port for the lifetime of a bash script, use $BASHPID as this
24argument.
25
26There is a race condition between picking a port and your application code
27binding to it.  The use of a port server to prevent that is recommended on
28loaded test hosts running many tests at a time.
29
30If your code can accept a bound socket as input rather than being handed a
31port number consider using socket.bind(('localhost', 0)) to bind to an
32available port without a race condition rather than using this library.
33
34Typical usage:
35  test_port = portpicker.pick_unused_port()
36"""
37
38from __future__ import print_function
39
40import logging
41import os
42import random
43import socket
44import sys
45
46if sys.platform == 'win32':
47    import _winapi
48else:
49    _winapi = None
50
51# The legacy Bind, IsPortFree, etc. names are not exported.
52__all__ = ('bind', 'is_port_free', 'pick_unused_port', 'return_port',
53           'add_reserved_port', 'get_port_from_port_server')
54
55_PROTOS = [(socket.SOCK_STREAM, socket.IPPROTO_TCP),
56           (socket.SOCK_DGRAM, socket.IPPROTO_UDP)]
57
58
59# Ports that are currently available to be given out.
60_free_ports = set()
61
62# Ports that are reserved or from the portserver that may be returned.
63_owned_ports = set()
64
65# Ports that we chose randomly that may be returned.
66_random_ports = set()
67
68
69class NoFreePortFoundError(Exception):
70    """Exception indicating that no free port could be found."""
71
72
73def add_reserved_port(port):
74    """Add a port that was acquired by means other than the port server."""
75    _free_ports.add(port)
76
77
78def return_port(port):
79    """Return a port that is no longer being used so it can be reused."""
80    if port in _random_ports:
81        _random_ports.remove(port)
82    elif port in _owned_ports:
83        _owned_ports.remove(port)
84        _free_ports.add(port)
85    elif port in _free_ports:
86        logging.info("Returning a port that was already returned: %s", port)
87    else:
88        logging.info("Returning a port that wasn't given by portpicker: %s",
89                     port)
90
91
92def bind(port, socket_type, socket_proto):
93    """Try to bind to a socket of the specified type, protocol, and port.
94
95    This is primarily a helper function for PickUnusedPort, used to see
96    if a particular port number is available.
97
98    For the port to be considered available, the kernel must support at least
99    one of (IPv6, IPv4), and the port must be available on each supported
100    family.
101
102    Args:
103      port: The port number to bind to, or 0 to have the OS pick a free port.
104      socket_type: The type of the socket (ex: socket.SOCK_STREAM).
105      socket_proto: The protocol of the socket (ex: socket.IPPROTO_TCP).
106
107    Returns:
108      The port number on success or None on failure.
109    """
110    got_socket = False
111    for family in (socket.AF_INET6, socket.AF_INET):
112        try:
113            sock = socket.socket(family, socket_type, socket_proto)
114            got_socket = True
115        except socket.error:
116            continue
117        try:
118            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
119            sock.bind(('', port))
120            if socket_type == socket.SOCK_STREAM:
121                sock.listen(1)
122            port = sock.getsockname()[1]
123        except socket.error:
124            return None
125        finally:
126            sock.close()
127    return port if got_socket else None
128
129Bind = bind  # legacy API. pylint: disable=invalid-name
130
131
132def is_port_free(port):
133    """Check if specified port is free.
134
135    Args:
136      port: integer, port to check
137    Returns:
138      boolean, whether it is free to use for both TCP and UDP
139    """
140    return bind(port, *_PROTOS[0]) and bind(port, *_PROTOS[1])
141
142IsPortFree = is_port_free  # legacy API. pylint: disable=invalid-name
143
144
145def pick_unused_port(pid=None, portserver_address=None):
146    """A pure python implementation of PickUnusedPort.
147
148    Args:
149      pid: PID to tell the portserver to associate the reservation with. If
150        None, the current process's PID is used.
151      portserver_address: The address (path) of a unix domain socket
152        with which to connect to a portserver, a leading '@'
153        character indicates an address in the "abstract namespace".  OR
154        On systems without socket.AF_UNIX, this is an AF_INET address.
155        If None, or no port is returned by the portserver at the provided
156        address, the environment will be checked for a PORTSERVER_ADDRESS
157        variable.  If that is not set, no port server will be used.
158
159    Returns:
160      A port number that is unused on both TCP and UDP.
161
162    Raises:
163      NoFreePortFoundError: No free port could be found.
164    """
165    try:  # Instead of `if _free_ports:` to handle the race condition.
166        port = _free_ports.pop()
167    except KeyError:
168        pass
169    else:
170        _owned_ports.add(port)
171        return port
172    # Provide access to the portserver on an opt-in basis.
173    if portserver_address:
174        port = get_port_from_port_server(portserver_address, pid=pid)
175        if port:
176            return port
177    if 'PORTSERVER_ADDRESS' in os.environ:
178        port = get_port_from_port_server(os.environ['PORTSERVER_ADDRESS'],
179                                         pid=pid)
180        if port:
181            return port
182    return _pick_unused_port_without_server()
183
184PickUnusedPort = pick_unused_port  # legacy API. pylint: disable=invalid-name
185
186
187def _pick_unused_port_without_server():  # Protected. pylint: disable=invalid-name
188    """Pick an available network port without the help of a port server.
189
190    This code ensures that the port is available on both TCP and UDP.
191
192    This function is an implementation detail of PickUnusedPort(), and
193    should not be called by code outside of this module.
194
195    Returns:
196      A port number that is unused on both TCP and UDP.
197
198    Raises:
199      NoFreePortFoundError: No free port could be found.
200    """
201    # Next, try a few times to get an OS-assigned port.
202    # Ambrose discovered that on the 2.6 kernel, calling Bind() on UDP socket
203    # returns the same port over and over. So always try TCP first.
204    for _ in range(10):
205        # Ask the OS for an unused port.
206        port = bind(0, _PROTOS[0][0], _PROTOS[0][1])
207        # Check if this port is unused on the other protocol.
208        if port and bind(port, _PROTOS[1][0], _PROTOS[1][1]):
209            _random_ports.add(port)
210            return port
211
212    # Try random ports as a last resort.
213    rng = random.Random()
214    for _ in range(10):
215        port = int(rng.randrange(15000, 25000))
216        if is_port_free(port):
217            _random_ports.add(port)
218            return port
219
220    # Give up.
221    raise NoFreePortFoundError()
222
223
224def _get_linux_port_from_port_server(portserver_address, pid):
225    # An AF_UNIX address may start with a zero byte, in which case it is in the
226    # "abstract namespace", and doesn't have any filesystem representation.
227    # See 'man 7 unix' for details.
228    # The convention is to write '@' in the address to represent this zero byte.
229    if portserver_address[0] == '@':
230        portserver_address = '\0' + portserver_address[1:]
231
232    try:
233        # Create socket.
234        if hasattr(socket, 'AF_UNIX'):
235            sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) # pylint: disable=no-member
236        else:
237            # fallback to AF_INET if this is not unix
238            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
239        try:
240            # Connect to portserver.
241            sock.connect(portserver_address)
242
243            # Write request.
244            sock.sendall(('%d\n' % pid).encode('ascii'))
245
246            # Read response.
247            # 1K should be ample buffer space.
248            return sock.recv(1024)
249        finally:
250            sock.close()
251    except socket.error as error:
252        print('Socket error when connecting to portserver:', error,
253              file=sys.stderr)
254        return None
255
256
257def _get_windows_port_from_port_server(portserver_address, pid):
258    if portserver_address[0] == '@':
259        portserver_address = '\\\\.\\pipe\\' + portserver_address[1:]
260
261    try:
262        handle = _winapi.CreateFile(
263            portserver_address,
264            _winapi.GENERIC_READ | _winapi.GENERIC_WRITE,
265            0,
266            0,
267            _winapi.OPEN_EXISTING,
268            0,
269            0)
270
271        _winapi.WriteFile(handle, ('%d\n' % pid).encode('ascii'))
272        data, _ = _winapi.ReadFile(handle, 6, 0)
273        return data
274    except FileNotFoundError as error:
275        print('File error when connecting to portserver:', error,
276              file=sys.stderr)
277        return None
278
279def get_port_from_port_server(portserver_address, pid=None):
280    """Request a free a port from a system-wide portserver.
281
282    This follows a very simple portserver protocol:
283    The request consists of our pid (in ASCII) followed by a newline.
284    The response is a port number and a newline, 0 on failure.
285
286    This function is an implementation detail of pick_unused_port().
287    It should not normally be called by code outside of this module.
288
289    Args:
290      portserver_address: The address (path) of a unix domain socket
291        with which to connect to the portserver.  A leading '@'
292        character indicates an address in the "abstract namespace."
293        On systems without socket.AF_UNIX, this is an AF_INET address.
294      pid: The PID to tell the portserver to associate the reservation with.
295        If None, the current process's PID is used.
296
297    Returns:
298      The port number on success or None on failure.
299    """
300    if not portserver_address:
301        return None
302
303    if pid is None:
304        pid = os.getpid()
305
306    if _winapi:
307        buf = _get_windows_port_from_port_server(portserver_address, pid)
308    else:
309        buf = _get_linux_port_from_port_server(portserver_address, pid)
310
311    if buf is None:
312        return None
313
314    try:
315        port = int(buf.split(b'\n')[0])
316    except ValueError:
317        print('Portserver failed to find a port.', file=sys.stderr)
318        return None
319    _owned_ports.add(port)
320    return port
321
322
323GetPortFromPortServer = get_port_from_port_server  # legacy API. pylint: disable=invalid-name
324
325
326def main(argv):
327    """If passed an arg, treat it as a PID, otherwise portpicker uses getpid."""
328    port = pick_unused_port(pid=int(argv[1]) if len(argv) > 1 else None)
329    if not port:
330        sys.exit(1)
331    print(port)
332
333
334if __name__ == '__main__':
335    main(sys.argv)
336