Recreate EDNS if options or payload are changed

dnspython >= 2.0.0 introduces an OPT class, instances of which include
the edns, ednsflags, and payload attributes.  However, the payload
attribute is immutable because it is actually the TTL value of the OPT
record.  This commit changes the code to use use_edns() instead of
trying to set payload directly.  Additionally, the EDNS options in
dnspython >= 2.0.0 are implemented as a tuple, so use_edns() is to
modify the option as well. Fixes #74.

1 files changed, 18 insertions, 16 deletions

diff --git a/dnsviz/query.py b/dnsviz/query.py
index 19afdcd..1b8cdf6 100644
--- a/dnsviz/query.py
+++ b/dnsviz/query.py
@@ -375,7 +375,8 @@ class ReduceUDPMaxPayloadOnTimeoutHandler(DNSResponseHandler):
     def handle(self, response_wire, response, response_time):
         timeouts = self._get_num_timeouts(response)
         if not self._params['tcp'] and timeouts >= self._timeouts and self._request.payload > self._reduced_payload:
-            self._request.payload = self._reduced_payload
+            self._request.use_edns(self._request.edns, self._request.ednsflags,
+                    self._reduced_payload, options=self._request.options)
             return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_TIMEOUT, None, RETRY_ACTION_CHANGE_UDP_MAX_PAYLOAD, self._reduced_payload)
 
 class ClearEDNSFlagOnTimeoutHandler(DNSResponseHandler):
@@ -412,12 +413,10 @@ class RemoveEDNSOptionOnTimeoutHandler(DNSResponseHandler):
 
     def handle(self, response_wire, response, response_time):
         timeouts = self._get_num_timeouts(response)
-        try:
+        if not self._params['tcp'] and timeouts >= self._timeouts and self._request.options:
             opt = self._request.options[0]
-        except IndexError:
-            opt = None
-        if not self._params['tcp'] and timeouts >= self._timeouts and opt is not None:
-            self._request.options.remove(opt)
+            self._request.use_edns(self._request.edns, self._request.ednsflags,
+                    self._request.payload, options=self._request.options[1:])
             return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_TIMEOUT, None, RETRY_ACTION_REMOVE_EDNS_OPTION, opt.otype)
 
 class DisableEDNSOnTimeoutHandler(DNSResponseHandler):
@@ -464,12 +463,10 @@ class RemoveEDNSOptionOnRcodeHandler(DNSResponseHandler):
         self._rcode = rcode
 
     def handle(self, response_wire, response, response_time):
-        try:
+        if isinstance(response, dns.message.Message) and response.rcode() == self._rcode and self._request.options:
             opt = self._request.options[0]
-        except IndexError:
-            opt = None
-        if isinstance(response, dns.message.Message) and response.rcode() == self._rcode and opt is not None:
-            self._request.options.remove(opt)
+            self._request.use_edns(self._request.edns, self._request.ednsflags,
+                    self._request.payload, options=self._request.options[1:])
             return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_RCODE, response.rcode(), RETRY_ACTION_REMOVE_EDNS_OPTION, opt.otype)
 
 class AddServerCookieOnBADCOOKIE(DNSResponseHandler):
@@ -608,7 +605,8 @@ class PMTUBoundingHandler(DNSResponseHandler):
             if timeouts >= self._initial_timeouts:
                 self._lower_bound = self._reduced_payload
                 self._upper_bound = self._request.payload - 1
-                self._request.payload = self._reduced_payload
+                self._request.use_edns(self._request.edns, self._request.ednsflags,
+                        self._reduced_payload, options=self._request.options)
                 self._state = self.REDUCED_PAYLOAD
                 return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_TIMEOUT, None, RETRY_ACTION_CHANGE_UDP_MAX_PAYLOAD, self._reduced_payload)
 
@@ -633,7 +631,8 @@ class PMTUBoundingHandler(DNSResponseHandler):
                 #XXX this is cheating because we're not reporting the change to UDP
                 self._params['tcp'] = False
                 payload = response_len - 1
-                self._request.payload = payload
+                self._request.use_edns(self._request.edns, self._request.ednsflags,
+                        payload, options=self._request.options)
                 self._state = self.TCP_MINUS_ONE
                 return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_DIAGNOSTIC, response_len, RETRY_ACTION_CHANGE_UDP_MAX_PAYLOAD, payload)
 
@@ -641,7 +640,8 @@ class PMTUBoundingHandler(DNSResponseHandler):
             if is_timeout:
                 self._upper_bound = self._request.payload - 1
                 payload = self._lower_bound + (self._upper_bound + 1 - self._lower_bound)//2
-                self._request.payload = payload
+                self._request.use_edns(self._request.edns, self._request.ednsflags,
+                        payload, options=self._request.options)
                 self._state = self.PICKLE
                 return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_TIMEOUT, None, RETRY_ACTION_CHANGE_UDP_MAX_PAYLOAD, payload)
             # if the size of the message is less than the watermark, then perhaps we were rate limited
@@ -675,7 +675,8 @@ class PMTUBoundingHandler(DNSResponseHandler):
             elif is_timeout:
                 self._upper_bound = self._request.payload - 1
                 payload = self._lower_bound + (self._upper_bound + 1 - self._lower_bound)//2
-                self._request.payload = payload
+                self._request.use_edns(self._request.edns, self._request.ednsflags,
+                        payload, options=self._request.options)
                 return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_TIMEOUT, None, RETRY_ACTION_CHANGE_UDP_MAX_PAYLOAD, payload)
             # if the size of the message is less than the watermark, then perhaps we were rate limited
             elif response_len < self._water_mark:
@@ -691,7 +692,8 @@ class PMTUBoundingHandler(DNSResponseHandler):
             elif is_valid:
                 self._lower_bound = self._request.payload
                 payload = self._lower_bound + (self._upper_bound + 1 - self._lower_bound)//2
-                self._request.payload = payload
+                self._request.use_edns(self._request.edns, self._request.ednsflags,
+                        payload, options=self._request.options)
                 return DNSQueryRetryAttempt(response_time, RETRY_CAUSE_DIAGNOSTIC, response_len, RETRY_ACTION_CHANGE_UDP_MAX_PAYLOAD, payload)
 
         elif self._state == self.TCP_FINAL: