From 337216f9a77420e4401c8d32fa485fa3c0797440 Mon Sep 17 00:00:00 2001
From: Tom Thorogood <me+github@tomthorogood.co.uk>
Date: Mon, 11 Mar 2019 00:16:14 +1030
Subject: [PATCH] Use net.Buffers for writing TCP message (#934)

---
 client.go | 21 ++++++++-------------
 server.go | 15 +++++----------
 2 files changed, 13 insertions(+), 23 deletions(-)

diff --git a/client.go b/client.go
index 599f14c7..86894ec2 100644
--- a/client.go
+++ b/client.go
@@ -3,7 +3,6 @@ package dns
 // A client implementation.
 
 import (
-	"bytes"
 	"context"
 	"crypto/tls"
 	"encoding/binary"
@@ -353,23 +352,19 @@ func (co *Conn) WriteMsg(m *Msg) (err error) {
 
 // Write implements the net.Conn Write method.
 func (co *Conn) Write(p []byte) (n int, err error) {
-	switch t := co.Conn.(type) {
+	switch co.Conn.(type) {
 	case *net.TCPConn, *tls.Conn:
-		w := t.(io.Writer)
-
-		lp := len(p)
-		if lp < 2 {
-			return 0, io.ErrShortBuffer
-		}
-		if lp > MaxMsgSize {
+		if len(p) > MaxMsgSize {
 			return 0, &Error{err: "message too large"}
 		}
-		l := make([]byte, 2, lp+2)
-		binary.BigEndian.PutUint16(l, uint16(lp))
-		p = append(l, p...)
-		n, err := io.Copy(w, bytes.NewReader(p))
+
+		l := make([]byte, 2)
+		binary.BigEndian.PutUint16(l, uint16(len(p)))
+
+		n, err := (&net.Buffers{l, p}).WriteTo(co.Conn)
 		return int(n), err
 	}
+
 	return co.Conn.Write(p)
 }
 
diff --git a/server.go b/server.go
index 42ca7324..4bd4674d 100644
--- a/server.go
+++ b/server.go
@@ -3,7 +3,6 @@
 package dns
 
 import (
-	"bytes"
 	"context"
 	"crypto/tls"
 	"encoding/binary"
@@ -701,18 +700,14 @@ func (w *response) Write(m []byte) (int, error) {
 	case w.udp != nil:
 		return WriteToSessionUDP(w.udp, m, w.udpSession)
 	case w.tcp != nil:
-		lm := len(m)
-		if lm < 2 {
-			return 0, io.ErrShortBuffer
-		}
-		if lm > MaxMsgSize {
+		if len(m) > MaxMsgSize {
 			return 0, &Error{err: "message too large"}
 		}
-		l := make([]byte, 2, 2+lm)
-		binary.BigEndian.PutUint16(l, uint16(lm))
-		m = append(l, m...)
 
-		n, err := io.Copy(w.tcp, bytes.NewReader(m))
+		l := make([]byte, 2)
+		binary.BigEndian.PutUint16(l, uint16(len(m)))
+
+		n, err := (&net.Buffers{l, m}).WriteTo(w.tcp)
 		return int(n), err
 	default:
 		panic("dns: internal error: udp and tcp both nil")