com.mysql.cj.util.DnsSrv Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of mysql-connector-java
JDBC Type 4 driver for MySQL
There is a newer version: 8.0.33
/*
 * Copyright (c) 2019, 2020, Oracle and/or its affiliates.
 *
 * This program is free software; you can redistribute it and/or modify it under
 * the terms of the GNU General Public License, version 2.0, as published by the
 * Free Software Foundation.
 *
 * This program is also distributed with certain software (including but not
 * limited to OpenSSL) that is licensed under separate terms, as designated in a
 * particular file or component or in included license documentation. The
 * authors of MySQL hereby grant you an additional permission to link the
 * program and your derivative works with the separately licensed software that
 * they have included with MySQL.
 *
 * Without limiting anything contained in the foregoing, this file, which is
 * part of MySQL Connector/J, is also subject to the Universal FOSS Exception,
 * version 1.0, a copy of which can be found at
 * http://oss.oracle.com/licenses/universal-foss-exception.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE. See the GNU General Public License, version 2.0,
 * for more details.
 *
 * You should have received a copy of the GNU General Public License along with
 * this program; if not, write to the Free Software Foundation, Inc.,
 * 51 Franklin St, Fifth Floor, Boston, MA 02110-1301  USA
 */

package com.mysql.cj.util;

import java.util.ArrayList;
import java.util.List;
import java.util.Properties;
import java.util.Random;
import java.util.stream.Collectors;

import javax.naming.Context;
import javax.naming.NamingException;
import javax.naming.directory.Attribute;
import javax.naming.directory.Attributes;
import javax.naming.directory.DirContext;
import javax.naming.directory.InitialDirContext;

public class DnsSrv {
    public static class SrvRecord implements Comparable {
        private final int priority;
        private final int weight;
        private final int port;
        private final String target;

        public static SrvRecord buildFrom(String srvLine) {
            String[] srvParts = srvLine.split("\\s+");
            if (srvParts.length == 4) {
                int priority = Integer.parseInt(srvParts[0]);
                int weight = Integer.parseInt(srvParts[1]);
                int port = Integer.parseInt(srvParts[2]);
                String target = srvParts[3].replaceFirst("\\.$", "");
                return new SrvRecord(priority, weight, port, target);
            }
            return null;
        }

        public SrvRecord(int priority, int weight, int port, String target) {
            this.priority = priority;
            this.weight = weight;
            this.port = port;
            this.target = target;
        }

        public int getPriority() {
            return this.priority;
        }

        public int getWeight() {
            return this.weight;
        }

        public int getPort() {
            return this.port;
        }

        public String getTarget() {
            return this.target;
        }

        @Override
        public String toString() {
            return String.format("{\"Priority\": %d, \"Weight\": %d, \"Port\": %d, \"Target\": \"%s\"}", getPriority(), getWeight(), getPort(), getTarget());
        }

        @Override
        public int compareTo(SrvRecord o) {
            int priorityDiff = getPriority() - o.getPriority();
            return priorityDiff == 0 ? getWeight() - o.getWeight() : priorityDiff;
        }
    }

    public static List lookupSrvRecords(String serviceName) throws NamingException {
        List srvRecords = new ArrayList<>();

        Properties environment = new Properties();
        environment.put(Context.INITIAL_CONTEXT_FACTORY, "com.sun.jndi.dns.DnsContextFactory");
        DirContext context = new InitialDirContext(environment);
        Attributes dnsEntries = context.getAttributes(serviceName, new String[] { "SRV" });
        if (dnsEntries != null) {
            Attribute hosts = dnsEntries.get("SRV");
            if (hosts != null) {
                for (int i = 0; i < hosts.size(); i++) {
                    srvRecords.add(SrvRecord.buildFrom((String) hosts.get(i)));
                }
            }
        }

        return sortSrvRecords(srvRecords);
    }

    /**
     * Sort a list of DNS SRV records according to the sorting rules described in rfc2782.
     * 
     * @param srvRecords
     *            the list of {@link SrvRecord}s to sort.
     * @return
     *         a new list of sorted {@link SrvRecord}s.
     */
    public static List sortSrvRecords(List srvRecords) {
        // Sort srv records by their natural order, i.e., first by priority then by weight.
        List srvRecordsSortedNatural = srvRecords.stream().sorted().collect(Collectors.toList());

        // Sort records with equal priority by weight using a categorical distribution.
        Random random = new Random(System.nanoTime());
        List srvRecordsSortedRfc2782 = new ArrayList<>();

        List priorities = srvRecordsSortedNatural.stream().map(SrvRecord::getPriority).distinct().collect(Collectors.toList());
        for (Integer priority : priorities) {
            List srvRecordsSamePriority = srvRecordsSortedNatural.stream().filter(s -> s.getPriority() == priority).collect(Collectors.toList());
            while (srvRecordsSamePriority.size() > 1) {
                int recCount = srvRecordsSamePriority.size();
                int sumOfWeights = 0;
                int[] weights = new int[recCount];
                for (int i = 0; i < recCount; i++) {
                    sumOfWeights += srvRecordsSamePriority.get(i).getWeight();
                    weights[i] = sumOfWeights;
                }
                int selection = random.nextInt(sumOfWeights + 1);
                int pos = 0;
                for (; pos < recCount && weights[pos] < selection; pos++) {
                }
                srvRecordsSortedRfc2782.add(srvRecordsSamePriority.remove(pos));
            }
            srvRecordsSortedRfc2782.add(srvRecordsSamePriority.get(0));
        }

        return srvRecordsSortedRfc2782;
    }
}