Merge pull request #579 from rthalley/proc-order

Processing order

9 files changed, 118 insertions, 0 deletions

diff --git a/dns/rdata.py b/dns/rdata.py
index acf34ae..7846205 100644
--- a/dns/rdata.py
+++ b/dns/rdata.py
@@ -23,6 +23,7 @@ import binascii
 import io
 import inspect
 import itertools
+import random
 
 import dns.wire
 import dns.exception
@@ -459,6 +460,14 @@ class Rdata:
             # against *as_value*.
             return tuple(as_value(v) for v in value)
 
+    # Processing order
+
+    @classmethod
+    def _processing_order(cls, iterable):
+        items = list(iterable)
+        random.shuffle(items)
+        return items
+
 
 class GenericRdata(Rdata):
 
diff --git a/dns/rdataset.py b/dns/rdataset.py
index 45f8bd6..e69ee23 100644
--- a/dns/rdataset.py
+++ b/dns/rdataset.py
@@ -309,6 +309,20 @@ class Rdataset(dns.set.Set):
             return True
         return False
 
+    def processing_order(self):
+        """Return rdatas in a valid processing order according to the type's
+        specification.  For example, MX records are in preference order from
+        lowest to highest preferences, with items of the same perference
+        shuffled.
+
+        For types that do not define a processing order, the rdatas are
+        simply shuffled.
+        """
+        if len(self) == 0:
+            return []
+        else:
+            return self[0]._processing_order(iter(self))
+
 
 @dns.immutable.immutable
 class ImmutableRdataset(Rdataset):
diff --git a/dns/rdtypes/ANY/URI.py b/dns/rdtypes/ANY/URI.py
index 60a43c8..ccbd2ce 100644
--- a/dns/rdtypes/ANY/URI.py
+++ b/dns/rdtypes/ANY/URI.py
@@ -67,3 +67,13 @@ class URI(dns.rdata.Rdata):
         if len(target) == 0:
             raise dns.exception.FormError('URI target may not be empty')
         return cls(rdclass, rdtype, priority, weight, target)
+
+    def _processing_priority(self):
+        return self.priority
+
+    def _processing_weight(self):
+        return self.weight
+
+    @classmethod
+    def _processing_order(cls, iterable):
+        return dns.rdtypes.util.weighted_processing_order(iterable)
diff --git a/dns/rdtypes/IN/NAPTR.py b/dns/rdtypes/IN/NAPTR.py
index a4058ad..1f072f8 100644
--- a/dns/rdtypes/IN/NAPTR.py
+++ b/dns/rdtypes/IN/NAPTR.py
@@ -89,3 +89,10 @@ class NAPTR(dns.rdata.Rdata):
         replacement = parser.get_name(origin)
         return cls(rdclass, rdtype, order, preference, strings[0], strings[1],
                    strings[2], replacement)
+
+    def _processing_priority(self):
+        return (self.order, self.preference)
+
+    @classmethod
+    def _processing_order(cls, iterable):
+        return dns.rdtypes.util.priority_processing_order(iterable)
diff --git a/dns/rdtypes/IN/PX.py b/dns/rdtypes/IN/PX.py
index 3a744f6..8abfb29 100644
--- a/dns/rdtypes/IN/PX.py
+++ b/dns/rdtypes/IN/PX.py
@@ -63,3 +63,10 @@ class PX(dns.rdata.Rdata):
         map822 = parser.get_name(origin)
         mapx400 = parser.get_name(origin)
         return cls(rdclass, rdtype, preference, map822, mapx400)
+
+    def _processing_priority(self):
+        return self.preference
+
+    @classmethod
+    def _processing_order(cls, iterable):
+        return dns.rdtypes.util.priority_processing_order(iterable)
diff --git a/dns/rdtypes/IN/SRV.py b/dns/rdtypes/IN/SRV.py
index 99ff70b..6d9b683 100644
--- a/dns/rdtypes/IN/SRV.py
+++ b/dns/rdtypes/IN/SRV.py
@@ -63,3 +63,13 @@ class SRV(dns.rdata.Rdata):
         (priority, weight, port) = parser.get_struct('!HHH')
         target = parser.get_name(origin)
         return cls(rdclass, rdtype, priority, weight, port, target)
+
+    def _processing_priority(self):
+        return self.priority
+
+    def _processing_weight(self):
+        return self.weight
+
+    @classmethod
+    def _processing_order(cls, iterable):
+        return dns.rdtypes.util.weighted_processing_order(iterable)
diff --git a/dns/rdtypes/mxbase.py b/dns/rdtypes/mxbase.py
index 7bacbd8..5641823 100644
--- a/dns/rdtypes/mxbase.py
+++ b/dns/rdtypes/mxbase.py
@@ -23,6 +23,7 @@ import dns.exception
 import dns.immutable
 import dns.rdata
 import dns.name
+import dns.rdtypes.util
 
 
 @dns.immutable.immutable
@@ -59,6 +60,13 @@ class MXBase(dns.rdata.Rdata):
         exchange = parser.get_name(origin)
         return cls(rdclass, rdtype, preference, exchange)
 
+    def _processing_priority(self):
+        return self.preference
+
+    @classmethod
+    def _processing_order(cls, iterable):
+        return dns.rdtypes.util.priority_processing_order(iterable)
+
 
 @dns.immutable.immutable
 class UncompressedMX(MXBase):
diff --git a/dns/rdtypes/svcbbase.py b/dns/rdtypes/svcbbase.py
index 4a3a217..585c05e 100644
--- a/dns/rdtypes/svcbbase.py
+++ b/dns/rdtypes/svcbbase.py
@@ -534,3 +534,10 @@ class SVCBBase(dns.rdata.Rdata):
                 value = pcls.from_wire_parser(parser, origin)
             params[key] = value
         return cls(rdclass, rdtype, priority, target, params)
+
+    def _processing_priority(self):
+        return self.priority
+
+    @classmethod
+    def _processing_order(cls, iterable):
+        return dns.rdtypes.util.priority_processing_order(iterable)
diff --git a/dns/rdtypes/util.py b/dns/rdtypes/util.py
index 30be37d..7fc08cd 100644
--- a/dns/rdtypes/util.py
+++ b/dns/rdtypes/util.py
@@ -15,6 +15,8 @@
 # ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT
 # OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 
+import collections
+import random
 import struct
 
 import dns.exception
@@ -183,3 +185,47 @@ class Bitmap:
             bitmap = parser.get_counted_bytes()
             windows.append((window, bitmap))
         return cls(windows)
+
+
+def _priority_table(items):
+    by_priority = collections.defaultdict(list)
+    for rdata in items:
+        by_priority[rdata._processing_priority()].append(rdata)
+    return by_priority
+
+def priority_processing_order(iterable):
+    items = list(iterable)
+    if len(items) == 1:
+        return items
+    by_priority = _priority_table(items)
+    ordered = []
+    for k in sorted(by_priority.keys()):
+        rdatas = by_priority[k]
+        random.shuffle(rdatas)
+        ordered.extend(rdatas)
+    return ordered
+
+_no_weight = 0.1
+
+def weighted_processing_order(iterable):
+    items = list(iterable)
+    if len(items) == 1:
+        return items
+    by_priority = _priority_table(items)
+    ordered = []
+    for k in sorted(by_priority.keys()):
+        rdatas = by_priority[k]
+        total = sum(rdata._processing_weight() or _no_weight
+                    for rdata in rdatas)
+        while len(rdatas) > 1:
+            r = random.uniform(0, total)
+            for (n, rdata) in enumerate(rdatas):
+                weight = rdata._processing_weight() or _no_weight
+                if weight > r:
+                    break
+                r -= weight
+            total -= weight
+            ordered.append(rdata)
+            del rdatas[n]
+        ordered.append(rdatas[0])
+    return ordered