Organize portpicker to build properly am: fd0833193b am: a00b5f9f36 am: a18902a24e

Original change: https://android-review.googlesource.com/c/platform/external/python/portpicker/+/2115373

Change-Id: I2edf0c3cd9ba7790337a07a11a6506e9acced718
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>

3 files changed, 370 insertions, 16 deletions

diff --git a/Android.bp b/Android.bp
index 96f8529..eb55539 100644
--- a/Android.bp
+++ b/Android.bp
@@ -26,19 +26,3 @@ license {
     ],
 }
 
-python_library {
-    name: "py-portpicker",
-    host_supported: true,
-    srcs: [
-        "src/portpicker.py",
-        "src/portserver.py",
-    ],
-    version: {
-        py2: {
-            enabled: false,
-        },
-        py3: {
-            enabled: true,
-        },
-    },
-}
diff --git a/src/Android.bp b/src/Android.bp
new file mode 100644
index 0000000..0d5b8ac
--- /dev/null
+++ b/src/Android.bp
@@ -0,0 +1,35 @@
+// Copyright 2022 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.
+package {
+    default_applicable_licenses: ["external_python_portpicker_license"],
+}
+
+python_library {
+    name: "py-portpicker",
+    host_supported: true,
+    srcs: [
+        "__init__.py",
+        "portpicker.py",
+        "portserver.py",
+    ],
+    version: {
+        py2: {
+            enabled: false,
+        },
+        py3: {
+            enabled: true,
+        },
+    },
+    pkg_path: "portpicker",
+}
diff --git a/src/__init__.py b/src/__init__.py
new file mode 100644
index 0000000..fc2825b
--- /dev/null
+++ b/src/__init__.py
@@ -0,0 +1,335 @@
+#!/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()
+"""
+
+from __future__ import print_function
+
+import logging
+import os
+import random
+import socket
+import sys
+
+if sys.platform == 'win32':
+    import _winapi
+else:
+    _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.
+    """
+    got_socket = False
+    for family in (socket.AF_INET6, socket.AF_INET):
+        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:
+            sock.close()
+    return port if got_socket else None
+
+Bind = bind  # legacy API. pylint: disable=invalid-name
+
+
+def is_port_free(port):
+    """Check if specified port is free.
+
+    Args:
+      port: integer, port to check
+    Returns:
+      boolean, whether it is free to use for both TCP and UDP
+    """
+    return bind(port, *_PROTOS[0]) and bind(port, *_PROTOS[1])
+
+IsPortFree = is_port_free  # legacy API. pylint: disable=invalid-name
+
+
+def pick_unused_port(pid=None, portserver_address=None):
+    """A pure python implementation of PickUnusedPort.
+
+    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.
+
+    Returns:
+      A port number that is unused on both TCP and UDP.
+
+    Raises:
+      NoFreePortFoundError: No free port could be found.
+    """
+    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()
+
+PickUnusedPort = pick_unused_port  # legacy API. pylint: disable=invalid-name
+
+
+def _pick_unused_port_without_server():  # Protected. pylint: disable=invalid-name
+    """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.
+
+    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.
+    for _ in range(10):
+        # Ask the OS for an unused port.
+        port = bind(0, _PROTOS[0][0], _PROTOS[0][1])
+        # Check if this port is unused on the other protocol.
+        if port and bind(port, _PROTOS[1][0], _PROTOS[1][1]):
+            _random_ports.add(port)
+            return port
+
+    # Try random ports as a last resort.
+    rng = random.Random()
+    for _ in range(10):
+        port = int(rng.randrange(15000, 25000))
+        if is_port_free(port):
+            _random_ports.add(port)
+            return port
+
+    # Give up.
+    raise NoFreePortFoundError()
+
+
+def _get_linux_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) # pylint: disable=no-member
+        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 _get_windows_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 = _get_windows_port_from_port_server(portserver_address, pid)
+    else:
+        buf = _get_linux_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
+
+
+GetPortFromPortServer = get_port_from_port_server  # legacy API. pylint: disable=invalid-name
+
+
+def main(argv):
+    """If passed an arg, treat it as a PID, otherwise portpicker uses getpid."""
+    port = pick_unused_port(pid=int(argv[1]) if len(argv) > 1 else None)
+    if not port:
+        sys.exit(1)
+    print(port)
+
+
+if __name__ == '__main__':
+    main(sys.argv)