001/*
002 * Copyright 2015-2020 the original author or authors
003 *
004 * This software is licensed under the Apache License, Version 2.0,
005 * the GNU Lesser General Public License version 2 or later ("LGPL")
006 * and the WTFPL.
007 * You may choose either license to govern your use of this software only
008 * upon the condition that you accept all of the terms of either
009 * the Apache License 2.0, the LGPL 2.1+ or the WTFPL.
010 */
011package org.minidns.util;
012
013import java.util.ArrayList;
014import java.util.Collection;
015import java.util.Collections;
016import java.util.LinkedList;
017import java.util.List;
018import java.util.SortedMap;
019import java.util.TreeMap;
020
021import org.minidns.dnsname.DnsName;
022import org.minidns.record.SRV;
023
024public class SrvUtil {
025
026    /**
027     * Sort the given collection of {@link SRV} resource records by their priority and weight.
028     * <p>
029     * Sorting by priority is easy. Sorting the buckets of SRV records with the same priority by weight requires to choose those records
030     * randomly but taking the weight into account.
031     * </p>
032     *
033     * @param srvRecords
034     *            a collection of SRV records.
035     * @return a sorted list of the given records.
036     */
037    public static List<SRV> sortSrvRecords(Collection<SRV> srvRecords) {
038        // RFC 2782, Usage rules: "If there is precisely one SRV RR, and its Target is "."
039        // (the root domain), abort."
040        if (srvRecords.size() == 1 && srvRecords.iterator().next().target.equals(DnsName.ROOT)) {
041            return Collections.emptyList();
042        }
043
044        // Create the priority buckets.
045        SortedMap<Integer, List<SRV>> buckets = new TreeMap<>();
046        for (SRV srvRecord : srvRecords) {
047            Integer priority = srvRecord.priority;
048            List<SRV> bucket = buckets.get(priority);
049            if (bucket == null) {
050                bucket = new LinkedList<>();
051                buckets.put(priority, bucket);
052            }
053            bucket.add(srvRecord);
054        }
055
056        List<SRV> sortedSrvRecords = new ArrayList<>(srvRecords.size());
057
058        for (List<SRV> bucket : buckets.values()) {
059            // The list of buckets will be sorted by priority, thanks to SortedMap. We now have determine the order of
060            // the SRV records with the same priority, i.e., within the same bucket, by their weight. This is done by
061            // creating an array 'totals' which reflects the percentage of the SRV RRs weight by the total weight of all
062            // SRV RRs in the bucket. For every entry in the bucket, we choose one using a random number and the sum of
063            // all weights left in the bucket. We then select RRs position based on the according index of the selected
064            // value in the 'total' array. This ensures that its weight is taken into account.
065            int bucketSize;
066            while ((bucketSize = bucket.size()) > 0) {
067                int[] totals = new int[bucketSize];
068
069                int zeroWeight = 1;
070                for (SRV srv : bucket) {
071                    if (srv.weight > 0) {
072                        zeroWeight = 0;
073                        break;
074                    }
075                }
076
077                int bucketWeightSum = 0, count = 0;
078                for (SRV srv : bucket) {
079                    bucketWeightSum += srv.weight + zeroWeight;
080                    totals[count++] = bucketWeightSum;
081                }
082
083                int selectedPosition;
084                if (bucketWeightSum == 0) {
085                    // If total priority is 0, then the sum of all weights in this priority bucket is 0. So we simply
086                    // select one of the weights randomly as the other algorithm performed in the else block is unable
087                    // to handle this case.
088                    selectedPosition = (int) (Math.random() * bucketSize);
089                } else {
090                    double rnd = Math.random() * bucketWeightSum;
091                    selectedPosition = bisect(totals, rnd);
092                }
093
094                SRV choosenSrvRecord = bucket.remove(selectedPosition);
095                sortedSrvRecords.add(choosenSrvRecord);
096            }
097        }
098
099        return sortedSrvRecords;
100    }
101
102    // TODO This is not yet really bisection just a stupid linear search.
103    private static int bisect(int[] array, double value) {
104        int pos = 0;
105        for (int element : array) {
106            if (value < element)
107                break;
108            pos++;
109        }
110        return pos;
111    }
112
113}