Utility classes used to find open TCP/UDP ports.

socketutils.d

/**
 *  Copyright 2002-2020 the original author or authors.
 *
 *  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
 *
 *  https://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.
 */

/**
 *  Simple utility methods for working with network sockets; for example,
 *  for finding available ports on localhost.
 *
 *  This is a port of code originally written for the Java Spring Framework.
 *  See https://github.com/spring-projects/spring-framework/blob/main/spring-core/src/main/java/org/springframework/util/SocketUtils.java
 *
 *  Authors: Sam Brannen, Ben Hale, Arjen Poutsma, Gunnar Hillert, Gary Russell
 *  Authors: Vijay Nayar (converted to D)
 */
module funnel.utils.socketutils;

import std.conv : to;
import std.socket : InternetAddress, Socket, TcpSocket, UdpSocket, SocketOSException;
import std.random : uniform;
import std.container.rbtree : RedBlackTree;

/**
 * The default minimum value for port ranges used when finding an available socket port.
 */
enum ushort PORT_RANGE_MIN = 1024;

/**
 * The default maximum value for port ranges used when finding an available
 * socket port.
 */
enum ushort PORT_RANGE_MAX = 65535;

/**
 * Find an available TCP port randomly selected from the range
 * \[ [PORT_RANGE_MIN], [PORT_RANGE_MAX] \].
 * Returns: an available TCP port number
 * Throws: Exception if no available port could be found
 */
ushort findAvailableTcpPort() {
  return findAvailableTcpPort(PORT_RANGE_MIN);
}

unittest {
  foreach (ushort i; 0..10) {
    ushort port = findAvailableTcpPort();
    assert(port >= PORT_RANGE_MIN && port <= PORT_RANGE_MAX);
  }
}

/**
 * Find an available TCP port randomly selected from the range
 * \[ [minPort], [PORT_RANGE_MAX] \].
 * Params:
 *   minPort = the minimum port number
 * Returns: an available TCP port number
 * Throws: Exception if no available port could be found
 */
ushort findAvailableTcpPort(ushort minPort) {
  return findAvailableTcpPort(minPort, PORT_RANGE_MAX);
}

unittest {
  foreach (ushort i; 0..10) {
    ushort portMin = cast(ushort)(PORT_RANGE_MIN + 1000 * i);
    ushort port = findAvailableTcpPort(portMin);
    assert(port >= portMin && port <= PORT_RANGE_MAX);
  }
}

/**
 * Find an available TCP port randomly selected from the range
 * \[ [minPort], [maxPort] \].
 * Params:
 *   minPort = the minimum port number
 *   maxPort = the maximum port number
 * Returns: an available TCP port number
 * Throws: Exception if no available port could be found
 */
ushort findAvailableTcpPort(ushort minPort, ushort maxPort) {
  return findAvailablePort!TcpSocket(minPort, maxPort);
}

unittest {
  foreach (ushort i; 0..10) {
    ushort portMin = cast(ushort)(PORT_RANGE_MIN + 1000 * i);
    ushort portMax = cast(ushort)(PORT_RANGE_MAX - 1000 * i);
    ushort port = findAvailableTcpPort(portMin, portMax);
    assert(port >= portMin && port <= portMax);
  }
}

/**
 * Find the requested number of available TCP ports, each randomly selected
 * from the range \[ [PORT_RANGE_MIN], [PORT_RANGE_MAX] \].
 * Params:
 *   numRequested = the number of available ports to find
 * Returns: a sorted set of available TCP port numbers
 * Throws: Exception if the requested number of available ports could not be found
 */
RedBlackTree!ushort findAvailableTcpPorts(ushort numRequested) {
  return findAvailableTcpPorts(numRequested, PORT_RANGE_MIN, PORT_RANGE_MAX);
}

unittest {
  RedBlackTree!ushort portSet = findAvailableTcpPorts(10);
  assert(portSet.length == 10);
}


/**
 * Find the requested number of available TCP ports, each randomly selected
 * from the range \[ [minPort], [maxPort] \].
 * Params:
 *   numRequested = the number of available ports to find
 *   minPort = the minimum port number
 *   maxPort = the maximum port number
 * Returns: a sorted set of available TCP port numbers
 * Throws: Exception if the requested number of available ports could not be found
 */
RedBlackTree!ushort findAvailableTcpPorts(ushort numRequested, ushort minPort, ushort maxPort) {
  return findAvailablePorts!TcpSocket(numRequested, minPort, maxPort);
}

unittest {
  ushort portMin = 20_000;
  ushort portMax = 30_000;
  RedBlackTree!ushort portSet = findAvailableTcpPorts(10, portMin, portMax);
  assert(portSet.length == 10);
  foreach (ushort port; portSet) {
    assert(port >= portMin && port <= portMax);
  }
}


/**
 * Find an available UDP port randomly selected from the range
 * \[ [PORT_RANGE_MIN], [PORT_RANGE_MAX] \].
 * Returns: an available UDP port number
 * Throws: Exception if no available port could be found
 */
ushort findAvailableUdpPort() {
  return findAvailableUdpPort(PORT_RANGE_MIN);
}

unittest {
  foreach (ushort i; 0..10) {
    ushort port = findAvailableUdpPort();
    assert(port >= PORT_RANGE_MIN && port <= PORT_RANGE_MAX);
  }
}

/**
 * Find an available UDP port randomly selected from the range
 * [{@code minPort}, {@value #PORT_RANGE_MAX}].
 * @param minPort the minimum port number
 * @return an available UDP port number
 * @throws Exception if no available port could be found
 */
ushort findAvailableUdpPort(ushort minPort) {
  return findAvailableUdpPort(minPort, PORT_RANGE_MAX);
}

unittest {
  foreach (ushort i; 0..10) {
    ushort portMin = cast(ushort)(PORT_RANGE_MIN + 1000 * i);
    ushort port = findAvailableUdpPort(portMin);
    assert(port >= portMin && port <= PORT_RANGE_MAX);
  }
}

/**
 * Find an available UDP port randomly selected from the range
 * \[ [minPort], [maxPort] \].
 * Params:
 *   minPort = the minimum port number
 *   maxPort = the maximum port number
 * Returns: an available UDP port number
 * Throws: Exception if no available port could be found
 */
ushort findAvailableUdpPort(ushort minPort, ushort maxPort) {
  return findAvailablePort!UdpSocket(minPort, maxPort);
}

unittest {
  foreach (ushort i; 0..10) {
    ushort portMin = cast(ushort)(PORT_RANGE_MIN + 1000 * i);
    ushort portMax = cast(ushort)(PORT_RANGE_MAX - 1000 * i);
    ushort port = findAvailableUdpPort(portMin, portMax);
    assert(port >= portMin && port <= portMax);
  }
}

/**
 * Find the requested number of available UDP ports, each randomly selected
 * from the range \[ [PORT_RANGE_MIN], [PORT_RANGE_MAX] \].
 * Params:
 *   numRequested = the number of available ports to find
 * Returns: a sorted set of available UDP port numbers
 * Throws: Exception if the requested number of available ports could not be found
 */
RedBlackTree!ushort findAvailableUdpPorts(ushort numRequested) {
  return findAvailableUdpPorts(numRequested, PORT_RANGE_MIN, PORT_RANGE_MAX);
}

unittest {
  RedBlackTree!ushort portSet = findAvailableUdpPorts(10);
  assert(portSet.length == 10);
}

/**
 * Find the requested number of available UDP ports, each randomly selected
 * from the range [{@code minPort}, {@code maxPort}].
 * Params:
 *   numRequested = the number of available ports to find
 *   minPort = the minimum port number
 *   maxPort = the maximum port number
 * Returns: a sorted set of available UDP port numbers
 * Throws: Exception if the requested number of available ports could not be found
 */
RedBlackTree!ushort findAvailableUdpPorts(ushort numRequested, ushort minPort, ushort maxPort) {
  return findAvailablePorts!UdpSocket(numRequested, minPort, maxPort);
}

unittest {
  ushort portMin = 20_000;
  ushort portMax = 30_000;
  RedBlackTree!ushort portSet = findAvailableUdpPorts(10, portMin, portMax);
  assert(portSet.length == 10);
  foreach (ushort port; portSet) {
    assert(port >= portMin && port <= portMax);
  }
}

/**
 * Determine if the specified port for this {@code SocketType} is
 * currently available on {@code localhost}.
 */
private bool isPortAvailable(SocketT : Socket)(ushort port) {
  /**/import std.stdio;
  try {
    SocketT socket = new SocketT();
    socket.bind(new InternetAddress("localhost", port));
    socket.close();
    return true;
  } catch (SocketOSException ex) {
    return false;
  }
}

unittest {
  try {
    ushort tcpPort = findAvailableTcpPort();
    TcpSocket tcpSocket = new TcpSocket();
    tcpSocket.bind(new InternetAddress("localhost", tcpPort));
    assert(isPortAvailable!TcpSocket(tcpPort) == false);
    tcpSocket.close();
  } catch (Exception ex) {
    assert(false);
  }
}

unittest {
  try {
    ushort udpPort = findAvailableUdpPort();
    UdpSocket udpSocket = new UdpSocket();
    udpSocket.bind(new InternetAddress("localhost", udpPort));
    assert(isPortAvailable!UdpSocket(udpPort) == false);
    udpSocket.close();
  } catch (Exception ex) {
    assert(false);
  }
}

/**
 * Find a pseudo-random port number within the range
 * \[ [minPort], [maxPort] \].
 * Params:
 *   minPort = the minimum port number
 *   maxPort = the maximum port number
 * Returns: a random port number within the specified range
 */
private ushort findRandomPort(ushort minPort, ushort maxPort)
in (minPort <= maxPort, "'minPort' must be greater than 'maxPort'.")
{
  return uniform!"[]"(minPort, maxPort);
}

/**
 * Find an available port for the given socket type, randomly selected
 * from the range \[ [minPort], [maxPort] \].
 * Params:
 *   minPort = the minimum port number
 *   maxPort = the maximum port number
 * Returns: an available port number for this socket type
 * Throws: Exception if no available port could be found
 */
ushort findAvailablePort(SocketT : Socket)(ushort minPort, ushort maxPort)
in (minPort > 0, "'minPort' must be greater than 0")
in (maxPort >= minPort, "'maxPort' must be greater than or equal to 'minPort'")
in (maxPort <= PORT_RANGE_MAX,
    "'maxPort' must be less than or equal to " ~ PORT_RANGE_MAX.to!string)
{
  import std.format;
  ushort portRange = cast(ushort)(maxPort - minPort);
  ushort candidatePort;
  int searchCounter = 0;
  do {
    if (searchCounter > portRange) {
      throw new Exception(format(
          "Could not find an available %s port in the range [%d, %d] after %d attempts",
          typeid(SocketT).toString(), minPort, maxPort, searchCounter));
    }
    candidatePort = findRandomPort(minPort, maxPort);
    searchCounter++;
  }
  while (!isPortAvailable!SocketT(candidatePort));

  return candidatePort;
}

/**
 * Find the requested number of available ports for this {@code SocketType},
 * each randomly selected from the range [{@code minPort}, {@code maxPort}].
 * Params:
 *   numRequested = the number of available ports to find
 *   minPort = the minimum port number
 *   maxPort = the maximum port number
 * Returns: a sorted set of available port numbers for this socket type
 * Throws: Exception if the requested number of available ports could not be found
 */
RedBlackTree!ushort findAvailablePorts(SocketT : Socket)(
    ushort numRequested, ushort minPort, ushort maxPort)
in (minPort > 0, "'minPort' must be greater than 0")
in (maxPort > minPort, "'maxPort' must be greater than 'minPort'")
in (maxPort <= PORT_RANGE_MAX,
    "'maxPort' must be less than or equal to " ~ PORT_RANGE_MAX.to!string)
in(numRequested > 0, "'numRequested' must be greater than 0")
in(maxPort - minPort >= numRequested,
    "'numRequested' must not be greater than 'maxPort' - 'minPort'")
{
  import std.format;
  RedBlackTree!ushort availablePorts = new RedBlackTree!ushort();
  ushort attemptCount = 0;
  while ((++attemptCount <= numRequested + 100) && availablePorts.length < numRequested) {
    availablePorts.stableInsert(findAvailablePort!SocketT(minPort, maxPort));
  }

  if (availablePorts.length != numRequested) {
    throw new Exception(format(
        "Could not find %d available %s ports in the range [%d, %d]",
        numRequested, typeid(SocketT).toString(), minPort, maxPort));
  }

  return availablePorts;
}