From: Bob Halley <halley@dnspython.org>
Date: Wed, 9 Sep 2020 12:35:36 +0000 (-0700)
Subject: correct and simplify weighted_processing_order()
X-Git-Tag: v2.1.0rc1~17^2~1
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=ba56a756392a996282f6b62fbecb73757a1619a0;p=thirdparty%2Fdnspython.git

correct and simplify weighted_processing_order()
---

diff --git a/dns/rdtypes/ANY/URI.py b/dns/rdtypes/ANY/URI.py
index 4e7ad01c..ccbd2ce4 100644
--- a/dns/rdtypes/ANY/URI.py
+++ b/dns/rdtypes/ANY/URI.py
@@ -76,4 +76,4 @@ class URI(dns.rdata.Rdata):
 
     @classmethod
     def _processing_order(cls, iterable):
-        return dns.rdtypes.util.weighted_processing_order(iterable, False)
+        return dns.rdtypes.util.weighted_processing_order(iterable)
diff --git a/dns/rdtypes/IN/SRV.py b/dns/rdtypes/IN/SRV.py
index 3eb227f9..6d9b683a 100644
--- a/dns/rdtypes/IN/SRV.py
+++ b/dns/rdtypes/IN/SRV.py
@@ -72,4 +72,4 @@ class SRV(dns.rdata.Rdata):
 
     @classmethod
     def _processing_order(cls, iterable):
-        return dns.rdtypes.util.weighted_processing_order(iterable, True)
+        return dns.rdtypes.util.weighted_processing_order(iterable)
diff --git a/dns/rdtypes/util.py b/dns/rdtypes/util.py
index 695754df..7fc08cde 100644
--- a/dns/rdtypes/util.py
+++ b/dns/rdtypes/util.py
@@ -190,7 +190,6 @@ class Bitmap:
 def _priority_table(items):
     by_priority = collections.defaultdict(list)
     for rdata in items:
-        key = rdata._processing_priority()
         by_priority[rdata._processing_priority()].append(rdata)
     return by_priority
 
@@ -206,43 +205,27 @@ def priority_processing_order(iterable):
         ordered.extend(rdatas)
     return ordered
 
-def _processing_weight(rdata, adjust_zero_weight):
-    weight = rdata._processing_weight()
-    if weight == 0 and adjust_zero_weight:
-        return 0.1
-    else:
-        return weight
+_no_weight = 0.1
 
-def weighted_processing_order(iterable, adjust_zero_weight=False):
+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()):
-        weights_vary = False
-        weights = []
         rdatas = by_priority[k]
-        for rdata in rdatas:
-            weight = _processing_weight(rdata, adjust_zero_weight)
-            if len(weights) > 0 and weight != weights[-1]:
-                weights_vary = True
-            weights.append(weight)
-        if weights_vary:
-            while len(rdatas) > 1:
-                items = random.choices(rdatas, weights)
-                rdata = items[0]
-                ordered.append(rdata)
-                rdatas.remove(rdata)
-                weight = _processing_weight(rdata, adjust_zero_weight)
-                weights.remove(weight)
-            ordered.append(rdatas[0])
-        elif weights[0] == 0:
-            # All the weights are 0!  (This can't happen with SRV, but
-            # can with URI.  It's not clear from the URI RFC what you do here
-            # as it doesn't discuss weight.
-            return []
-        else:
-            random.shuffle(rdatas)
-            ordered.extend(rdatas)
+        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
diff --git a/tests/test_processing_order.py b/tests/test_processing_order.py
index cc33fd4e..2fa1b271 100644
--- a/tests/test_processing_order.py
+++ b/tests/test_processing_order.py
@@ -102,12 +102,3 @@ def test_processing_all_zero_weight_srv():
             assert rds[j] in po
         seen.add(tuple(po))
     assert len(seen) == 6
-
-
-def test_processing_all_zero_weight_uri():
-    rds = dns.rdataset.from_text('in', 'uri', 300,
-                                 '10 0 "ftp://ftp1.example.com/public"',
-                                 '10 0 "ftp://ftp2.example.com/public"',
-                                 '10 0 "ftp://ftp3.example.com/public"')
-    po = rds.processing_order()
-    assert len(po) == 0