From 0cf549278c9457e528c28adc5c2d61dc8dc25202 Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Fri, 3 Jan 2014 15:19:35 -0800
Subject: [PATCH] Simplify the Write() for TCP based connections

Simplify the code path by using io.Copy to handle partial writes. Allocate `l` large enough to avoid a re-allocation. Potential short write fix.
---
 server.go | 23 +++++++----------------
 1 file changed, 7 insertions(+), 16 deletions(-)

diff --git a/server.go b/server.go
index 158d63cc..e21e25b2 100644
--- a/server.go
+++ b/server.go
@@ -9,6 +9,8 @@
 package dns
 
 import (
+	"bytes"
+	"io"
 	"net"
 	"sync"
 	"time"
@@ -412,26 +414,15 @@ func (w *response) Write(m []byte) (int, error) {
 		return n, err
 	case w.tcp != nil:
 		lm := len(m)
-		if len(m) > MaxMsgSize {
+		if lm > MaxMsgSize {
 			return 0, &Error{err: "message too large"}
 		}
-		l := make([]byte, 2)
+		l := make([]byte, 2, 2+lm)
 		l[0], l[1] = packUint16(uint16(lm))
 		m = append(l, m...)
-		n, err := w.tcp.Write(m)
-		if err != nil {
-			return n, err
-		}
-		i := n
-		if i < lm {
-			j, err := w.tcp.Write(m[i:lm])
-			if err != nil {
-				return i, err
-			}
-			i += j
-		}
-		n = i
-		return i, nil
+		
+		n, err := io.Copy(w.tcp, bytes.NewReader(m))
+		return int(n), err
 	}
 	panic("not reached")
 }