dns

commit 2f4ad6437420a56d016b5d1ac5eb61c2a7234254
parent c245affd0ab63714b15dc1e6bd67307059e8da4f
Author: Oliver Lowe <o@olowe.co>
Date:   Sat, 20 Nov 2021 11:19:50 +1100

Use the message length field to only read that much

Using io.ReadAll doesn't work properly. The socket has a queue that
may be entirely drained by this call, which means we could
accidentally read a bit of the next message if there is one in the
queue.

Diffstat:
M
dns.go
 | 
15
++++++++++++---

1 file changed, 12 insertions(+), 3 deletions(-)
diff --git a/dns.go b/dns.go
@@ -87,9 +87,18 @@ func dnsStreamExchange(b []byte, conn net.Conn) ([]byte, error) {
 	if _, err := conn.Write(m); err != nil {
 		return nil, err
 	}
-	buf, err := io.ReadAll(conn)
+
+	b = make([]byte, 1280)
+	if _, err := io.ReadFull(conn, b[:2]); err != nil {
+		return nil, fmt.Errorf("read length: %w", err)
+	}
+	l := int(b[0])<<8 | int(b[1])
+	if l > len(b) {
+		b = make([]byte, l)
+	}
+	n, err := io.ReadFull(conn, b[:l])
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("read after length: %w", err)
 	}
-	return buf[2:], nil
+	return b[:n], nil
 }