EDNS option support

5 files changed, 146 insertions, 6 deletions

diff --git a/ChangeLog b/ChangeLog
index f0e92bc..ffd8444 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -1,3 +1,7 @@
+2009-03-30  Bob Halley  <halley@dnspython.org>
+
+	* Added EDNS option support.
+
 2008-10-16  Bob Halley  <halley@dnspython.org>
 
 	* dns/rdtypes/ANY/DS.py: The from_text() parser for DS RRs did not
diff --git a/dns/__init__.py b/dns/__init__.py
index 2d4cfd7..0be068f 100644
--- a/dns/__init__.py
+++ b/dns/__init__.py
@@ -1,4 +1,4 @@
-# Copyright (C) 2003-2007 Nominum, Inc.
+# Copyright (C) 2003-2009 Nominum, Inc.
 #
 # Permission to use, copy, modify, and distribute this software and its
 # documentation for any purpose with or without fee is hereby granted,
@@ -18,6 +18,7 @@
 __all__ = [
     'dnssec',
     'e164',
+    'edns',
     'entropy',
     'exception',
     'flags',
diff --git a/dns/edns.py b/dns/edns.py
new file mode 100644
index 0000000..093e36a
--- /dev/null
+++ b/dns/edns.py
@@ -0,0 +1,95 @@
+# Copyright (C) 2009 Nominum, Inc.
+#
+# Permission to use, copy, modify, and distribute this software and its
+# documentation for any purpose with or without fee is hereby granted,
+# provided that the above copyright notice and this permission notice
+# appear in all copies.
+#
+# THE SOFTWARE IS PROVIDED "AS IS" AND NOMINUM DISCLAIMS ALL WARRANTIES
+# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL NOMINUM BE LIABLE FOR
+# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
+# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT
+# OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+
+"""EDNS Options"""
+
+NSID = 3
+
+class Option(object):
+    """Base class for all EDNS option types.
+    """
+
+    def __init__(self, otype):
+        """Initialize an option.
+        @param rdtype: The rdata type
+        @type rdtype: int
+        """
+        self.otype = otype
+
+    def to_wire(self, file):
+        """Convert an option to wire format.
+        """
+        raise NotImplementedError
+
+    def from_wire(cls, otype, wire, current, olen):
+        """Build an EDNS option object from wire format
+
+        @param otype: The option type
+        @type otype: int
+        @param wire: The wire-format message
+        @type wire: string
+        @param current: The offet in wire of the beginning of the rdata.
+        @type current: int
+        @param olen: The length of the wire-format option data
+        @type olen: int
+        @rtype: dns.ends.Option instance"""
+        raise NotImplementedError
+
+    from_wire = classmethod(from_wire)
+
+class GenericOption(Option):
+    """Generate Rdata Class
+
+    This class is used for EDNS option types for which we have no better
+    implementation.
+    """
+
+    def __init__(self, otype, data):
+        super(GenericOption, self).__init__(otype)
+        self.data = data
+
+    def to_wire(self, file):
+        file.write(self.data)
+
+    def from_wire(cls, otype, wire, current, olen):
+        return cls(otype, wire[current : current + olen])
+
+    from_wire = classmethod(from_wire)
+
+
+_type_to_class = {
+}
+
+def get_option_class(otype):
+    cls = _type_to_class.get(otype)
+    if cls is None:
+        cls = GenericOption
+    return cls
+
+def option_from_wire(otype, wire, current, olen):
+    """Build an EDNS option object from wire format
+
+    @param otype: The option type
+    @type otype: int
+    @param wire: The wire-format message
+    @type wire: string
+    @param current: The offet in wire of the beginning of the rdata.
+    @type current: int
+    @param olen: The length of the wire-format option data
+    @type olen: int
+    @rtype: dns.ends.Option instance"""
+
+    cls = get_option_class(otype)
+    return cls.from_wire(otype, wire, current, olen)
diff --git a/dns/message.py b/dns/message.py
index ac8554c..4786bed 100644
--- a/dns/message.py
+++ b/dns/message.py
@@ -1,4 +1,4 @@
-# Copyright (C) 2001-2007 Nominum, Inc.
+# Copyright (C) 2001-2009 Nominum, Inc.
 #
 # Permission to use, copy, modify, and distribute this software and its
 # documentation for any purpose with or without fee is hereby granted,
@@ -84,6 +84,8 @@ class Message(object):
     @type ednsflags: long
     @ivar payload: The EDNS payload size.  The default is 0.
     @type payload: int
+    @ivar options: The EDNS options
+    @type options: list of dns.edns.Option objects
     @ivar request_payload: The associated request's EDNS payload size.
     @type request_payload: int
     @ivar keyring: The TSIG keyring to use.  The default is None.
@@ -143,6 +145,7 @@ class Message(object):
         self.edns = -1
         self.ednsflags = 0
         self.payload = 0
+        self.options = []
         self.request_payload = 0
         self.keyring = None
         self.keyname = None
@@ -400,7 +403,7 @@ class Message(object):
         for rrset in self.authority:
             r.add_rrset(dns.renderer.AUTHORITY, rrset, **kw)
         if self.edns >= 0:
-            r.add_edns(self.edns, self.ednsflags, self.payload)
+            r.add_edns(self.edns, self.ednsflags, self.payload, self.options)
         for rrset in self.additional:
             r.add_rrset(dns.renderer.ADDITIONAL, rrset, **kw)
         r.write_header()
@@ -450,7 +453,7 @@ class Message(object):
         self.tsig_error = tsig_error
         self.other_data = other_data
 
-    def use_edns(self, edns=0, ednsflags=0, payload=1280, request_payload=None):
+    def use_edns(self, edns=0, ednsflags=0, payload=1280, request_payload=None, options=None):
         """Configure EDNS behavior.
         @param edns: The EDNS level to use.  Specifying None, False, or -1
         means 'do not use EDNS', and in this case the other parameters are
@@ -477,13 +480,17 @@ class Message(object):
             ednsflags = 0
             payload = 0
             request_payload = 0
+            options = []
         else:
             # make sure the EDNS version in ednsflags agrees with edns
             ednsflags &= 0xFF00FFFFL
             ednsflags |= (edns << 16)
+            if options is None:
+                options = []
         self.edns = edns
         self.ednsflags = ednsflags
         self.payload = payload
+        self.options = options
         self.request_payload = request_payload
 
     def want_dnssec(self, wanted=True):
@@ -613,6 +620,18 @@ class _WireReader(object):
                 self.message.payload = rdclass
                 self.message.ednsflags = ttl
                 self.message.edns = (ttl & 0xff0000) >> 16
+                self.message.options = []
+                current = self.current
+                optslen = rdlen
+                while optslen > 0:
+                    (otype, olen) = \
+                            struct.unpack('!HH',
+                                          self.wire[current:current + 4])
+                    current = current + 4
+                    opt = dns.edns.option_from_wire(otype, self.wire, current, olen)
+                    self.message.options.append(opt)
+                    current = current + olen
+                    optslen = optslen - 4 - olen
                 seen_opt = True
             elif rdtype == dns.rdatatype.TSIG:
                 if not (section is self.message.additional and
diff --git a/dns/renderer.py b/dns/renderer.py
index 07bce98..ef53899 100644
--- a/dns/renderer.py
+++ b/dns/renderer.py
@@ -1,4 +1,4 @@
-# Copyright (C) 2001-2007 Nominum, Inc.
+# Copyright (C) 2001-2009 Nominum, Inc.
 #
 # Permission to use, copy, modify, and distribute this software and its
 # documentation for any purpose with or without fee is hereby granted,
@@ -202,7 +202,7 @@ class Renderer(object):
             raise dns.exception.TooBig
         self.counts[section] += n
 
-    def add_edns(self, edns, ednsflags, payload):
+    def add_edns(self, edns, ednsflags, payload, options=None):
         """Add an EDNS OPT record to the message.
 
         @param edns: The EDNS level to use.
@@ -212,6 +212,8 @@ class Renderer(object):
         @param payload: The EDNS sender's payload field, which is the maximum
         size of UDP datagram the sender can handle.
         @type payload: int
+        @param options: The EDNS options list
+        @type options: list of dns.edns.Option instances
         @see: RFC 2671
         """
 
@@ -222,6 +224,25 @@ class Renderer(object):
         before = self.output.tell()
         self.output.write(struct.pack('!BHHIH', 0, dns.rdatatype.OPT, payload,
                                       ednsflags, 0))
+        if not options is None:
+            lstart = self.output.tell()
+            for opt in options:
+                stuff = struct.pack("!HH", opt.otype, 0)
+                self.output.write(stuff)
+                start = self.output.tell()
+                opt.to_wire(self.output)
+                end = self.output.tell()
+                assert end - start < 65536
+                self.output.seek(start - 2)
+                stuff = struct.pack("!H", end - start)
+                self.output.write(stuff)
+                self.output.seek(0, 2)
+            lend = self.output.tell()
+            assert lend - lstart < 65536
+            self.output.seek(lstart - 2)
+            stuff = struct.pack("!H", lend - lstart)
+            self.output.write(stuff)
+            self.output.seek(0, 2)
         after = self.output.tell()
         if after >= self.max_size:
             self._rollback(before)