From a977e61aedea3d6d72c715a55bb5db097496f607 Mon Sep 17 00:00:00 2001
From: Bob Halley <halley@dnspython.org>
Date: Tue, 20 Feb 2024 14:01:51 -0800
Subject: [PATCH] Message from_wire() now stores the original wire format,
 (#1057)

and to_wire() records its rendered output (minus any length prefix).
---
 dns/message.py        |  3 +++
 doc/message-class.rst | 10 ++++++++--
 tests/test_message.py | 14 ++++++++++++++
 3 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/dns/message.py b/dns/message.py
index 8513db95..74a451fc 100644
--- a/dns/message.py
+++ b/dns/message.py
@@ -161,6 +161,7 @@ class Message:
         self.index: IndexType = {}
         self.errors: List[MessageError] = []
         self.time = 0.0
+        self.wire: Optional[bytes] = None
 
     @property
     def question(self) -> List[dns.rrset.RRset]:
@@ -645,6 +646,7 @@ class Message:
             if multi:
                 self.tsig_ctx = ctx
         wire = r.get_wire()
+        self.wire = wire
         if prepend_length:
             wire = len(wire).to_bytes(2, "big") + wire
         return wire
@@ -1259,6 +1261,7 @@ class _WireReader:
         factory = _message_factory_from_opcode(dns.opcode.from_flags(flags))
         self.message = factory(id=id)
         self.message.flags = dns.flags.Flag(flags)
+        self.message.wire = self.parser.wire
         self.initialize_message(self.message)
         self.one_rr_per_rrset = self.message._get_one_rr_per_rrset(
             self.one_rr_per_rrset
diff --git a/doc/message-class.rst b/doc/message-class.rst
index 08d99586..5a833691 100644
--- a/doc/message-class.rst
+++ b/doc/message-class.rst
@@ -28,7 +28,7 @@ DNS opcodes that do not have a more specific class.
    .. attribute:: ednsflags
 
       An ``int``, the EDNS flags.
-      
+
    .. attribute:: payload
 
       An ``int``, the EDNS payload size.  The default is 0.
@@ -71,7 +71,7 @@ DNS opcodes that do not have a more specific class.
    .. attribute:: original_id
 
       An ``int``, the TSIG original id; defaults to the message's id.
-      
+
    .. attribute:: tsig_error
 
       An ``int``, the TSIG error code.  The default is 0.
@@ -127,6 +127,12 @@ DNS opcodes that do not have a more specific class.
       is ``{}``.  Indexing improves the performance of finding RRsets.
       Indexing can be disabled by setting the index to ``None``.
 
+   .. attribute:: wire
+
+      A ``bytes`` or ``None``, the encoded wire format data used to create this
+      message with ``dns.message.from_wire()`` or the output most recently generated by
+      ``to_wire()``.
+
 The following constants may be used to specify sections in the
 ``find_rrset()`` and ``get_rrset()`` methods:
 
diff --git a/tests/test_message.py b/tests/test_message.py
index bbd45718..89db217e 100644
--- a/tests/test_message.py
+++ b/tests/test_message.py
@@ -170,6 +170,20 @@ class MessageTestCase(unittest.TestCase):
         self.assertEqual(len(m2.options), 1)
         self.assertEqual(m2.options[0], opt)
 
+    def test_keeping_wire(self):
+        m = dns.message.from_wire(goodwire)
+        self.assertEqual(m.wire, goodwire)
+        m = dns.message.make_query("example", "A")
+        self.assertEqual(m.wire, None)
+
+    def test_recording_wire(self):
+        m = dns.message.from_wire(goodwire)
+        self.assertEqual(m.wire, goodwire)
+        m.wire = None
+        wire = m.to_wire()
+        self.assertEqual(wire, goodwire)
+        self.assertEqual(m.wire, goodwire)
+
     def test_TooBig(self):
         def bad():
             q = dns.message.from_text(query_text)
-- 
2.47.3