Add prepend_length to Message.to_wire(). (#1001)

If a caller passes prepend_length=True, the wire format will include the
2 byte encoded message length before the message itself.  This is useful
for callers planning to send the message over TCP, DoT, and DoQ.

diff --git a/dns/asyncquery.py b/dns/asyncquery.py
index ecf9c1a5faace190cc7bbc900ef685aa06c9db24..13d317d148a2ff0e47c909eabb330502d0330d19 100644 (file)
--- a/dns/asyncquery.py
+++ b/dns/asyncquery.py
@@ -292,14 +292,12 @@ async def send_tcp(
     """
 
     if isinstance(what, dns.message.Message):
-        wire = what.to_wire()
+        tcpmsg = what.to_wire(prepend_length=True)
     else:
-        wire = what
-    l = len(wire)
-    # copying the wire into tcpmsg is inefficient, but lets us
-    # avoid writev() or doing a short write that would get pushed
-    # onto the net
-    tcpmsg = struct.pack("!H", l) + wire
+        # copying the wire into tcpmsg is inefficient, but lets us
+        # avoid writev() or doing a short write that would get pushed
+        # onto the net
+        tcpmsg = len(what).to_bytes(2, 'big') + what
     sent_time = time.time()
     await sock.sendall(tcpmsg, _timeout(expiration, sent_time))
     return (len(tcpmsg), sent_time)

diff --git a/dns/message.py b/dns/message.py
index 70049efd3320cd73b40859bfb2ae85c40038dc9e..54611d6fc40a3279b541bb87a073b29f52130236 100644 (file)
--- a/dns/message.py
+++ b/dns/message.py
@@ -527,6 +527,7 @@ class Message:
         max_size: int = 0,
         multi: bool = False,
         tsig_ctx: Optional[Any] = None,
+        prepend_length : bool = False,
         **kw: Dict[str, Any],
     ) -> bytes:
         """Return a string containing the message in DNS compressed wire
@@ -549,6 +550,10 @@ class Message:
         *tsig_ctx*, a ``dns.tsig.HMACTSig`` or ``dns.tsig.GSSTSig`` object, the
         ongoing TSIG context, used when signing zone transfers.
 
+        *prepend_length", a ``bool``, should be set to ``True`` if the caller
+        wants the message length prepended to the message itself.  This is
+        useful for messages sent over TCP, TLS (DoT), or QUIC (DoQ).
+
         Raises ``dns.exception.TooBig`` if *max_size* was exceeded.
 
         Returns a ``bytes``.
@@ -598,7 +603,10 @@ class Message:
             r.write_header()
             if multi:
                 self.tsig_ctx = ctx
-        return r.get_wire()
+        wire = r.get_wire()
+        if prepend_length:
+            wire = len(wire).to_bytes(2, 'big') + wire
+        return wire
 
     @staticmethod
     def _make_tsig(

diff --git a/dns/query.py b/dns/query.py
index 0d71125156603e7f60f5710b2a4e8c88048eecc5..a5fc6023dd36c183bfa2c1d2d2a50f34c061c96b 100644 (file)
--- a/dns/query.py
+++ b/dns/query.py
@@ -864,14 +864,12 @@ def send_tcp(
     """
 
     if isinstance(what, dns.message.Message):
-        wire = what.to_wire()
+        tcpmsg = what.to_wire(prepend_length=True)
     else:
-        wire = what
-    l = len(wire)
-    # copying the wire into tcpmsg is inefficient, but lets us
-    # avoid writev() or doing a short write that would get pushed
-    # onto the net
-    tcpmsg = struct.pack("!H", l) + wire
+        # copying the wire into tcpmsg is inefficient, but lets us
+        # avoid writev() or doing a short write that would get pushed
+        # onto the net
+        tcpmsg = len(what).to_bytes(2, 'big') + what
     sent_time = time.time()
     _net_write(sock, tcpmsg, expiration)
     return (len(tcpmsg), sent_time)