1 files changed, 83 insertions, 44 deletions

diff --git a/includes/libs/DnsSrvDiscoverer.php b/includes/libs/DnsSrvDiscoverer.php
index 1b91a932ba83..d5f1829bd968 100644
--- a/includes/libs/DnsSrvDiscoverer.php
+++ b/includes/libs/DnsSrvDiscoverer.php
@@ -27,28 +27,76 @@ class DnsSrvDiscoverer {
 	/**
 	 * @var string
 	 */
+	private $service;
+
+	/**
+	 * @var string
+	 */
+	private $protocol;
+
+	/**
+	 * @var string
+	 */
 	private $domain;
 
 	/**
-	 * @param string $domain
+	 * @var callable
 	 */
-	public function __construct( $domain ) {
+	private $resolver;
+
+	/**
+	 * Construct a new discoverer for the given domain, service, and protocol.
+	 *
+	 * @param string $service Name of the service to discover.
+	 * @param string $protocol Service protocol. Defaults to 'tcp'
+	 * @param ?string $domain The hostname/domain on which to perform discovery
+	 *  of the given service and protocol. Defaults to null which effectively
+	 *  performs a query relative to the host's configured search domain.
+	 * @param ?callable $resolver Resolver function. Defaults to using
+	 *  dns_get_record. Primarily useful in testing.
+	 */
+	public function __construct(
+		string $service,
+		string $protocol = 'tcp',
+		?string $domain = null,
+		?callable $resolver = null
+	) {
+		$this->service = $service;
+		$this->protocol = $protocol;
 		$this->domain = $domain;
+
+		$this->resolver = $resolver ?? static function ( $srv ) {
+			return dns_get_record( $srv, DNS_SRV );
+		};
 	}
 
 	/**
-	 * Fetch the servers with a DNS SRV request
+	 * Queries the resolver for an SRV resource record matching the service,
+	 * protocol, and domain and returns all target/port/priority/weight
+	 * records.
 	 *
 	 * @return array
 	 */
-	public function getServers() {
+	public function getRecords() {
 		$result = [];
-		foreach ( $this->getDnsRecords() as $record ) {
+
+		$records = ( $this->resolver )( $this->getSrvName() );
+
+		// Respect RFC 2782 with regard to a single '.' entry denoting a valid
+		// empty response
+		if (
+			!$records
+			|| ( count( $records ) === 1 && $records[0]['target'] === '.' )
+		) {
+			return $result;
+		}
+
+		foreach ( $records as $record ) {
 			$result[] = [
 				'target' => $record['target'],
-				'port' => $record['port'],
-				'pri' => $record['pri'],
-				'weight' => $record['weight'],
+				'port' => (int)$record['port'],
+				'pri' => (int)$record['pri'],
+				'weight' => (int)$record['weight'],
 			];
 		}
 
@@ -56,53 +104,44 @@ class DnsSrvDiscoverer {
 	}
 
 	/**
-	 * Pick a server according to the priority fields.
-	 * Note that weight is currently ignored.
+	 * Performs discovery for the domain, service, and protocol, and returns a
+	 * list of resolved server name/ip and port number pairs sorted by each
+	 * record's priority, with servers of the same priority randomly shuffled.
 	 *
-	 * @param array $servers from getServers
-	 * @return array|bool
+	 * @return array
 	 */
-	public function pickServer( array $servers ) {
-		if ( !$servers ) {
-			return false;
-		}
+	public function getServers() {
+		$records = $this->getRecords();
 
-		$srvsByPrio = [];
-		foreach ( $servers as $server ) {
-			$srvsByPrio[$server['pri']][] = $server;
-		}
+		usort( $records, static function ( $a, $b ) {
+			if ( $a['pri'] === $b['pri'] ) {
+				return mt_rand( 0, 1 ) ? 1 : -1;
+			}
 
-		$min = min( array_keys( $srvsByPrio ) );
-		if ( count( $srvsByPrio[$min] ) == 1 ) {
-			return $srvsByPrio[$min][0];
-		} else {
-			// Choose randomly
-			$rand = mt_rand( 0, count( $srvsByPrio[$min] ) - 1 );
+			return $a['pri'] - $b['pri'];
+		} );
 
-			return $srvsByPrio[$min][$rand];
-		}
-	}
+		$serversAndPorts = [];
 
-	/**
-	 * @param array $server
-	 * @param array[] $servers
-	 * @return array[]
-	 */
-	public function removeServer( $server, array $servers ) {
-		foreach ( $servers as $i => $srv ) {
-			if ( $srv['target'] === $server['target'] && $srv['port'] === $server['port'] ) {
-				unset( $servers[$i] );
-				break;
-			}
+		foreach ( $records as $record ) {
+			$serversAndPorts[] = [ $record['target'], $record['port'] ];
 		}
 
-		return array_values( $servers );
+		return $serversAndPorts;
 	}
 
 	/**
-	 * @return array[]
+	 * Returns the SRV resource record name.
+	 *
+	 * @return string
 	 */
-	protected function getDnsRecords() {
-		return dns_get_record( $this->domain, DNS_SRV );
+	public function getSrvName(): string {
+		$srv = "_{$this->service}._{$this->protocol}";
+
+		if ( $this->domain === null || $this->domain === '' ) {
+			return $srv;
+		}
+
+		return "$srv.{$this->domain}";
 	}
 }