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}