dns: marshal to io.writer instead of byte slice

dns/client.go

@@ -1,6 +1,7 @@
 package dns
 
 import (
+	"bytes"
 	"encoding/binary"
 	"errors"
 	"io"
@@ -9,7 +10,6 @@ import (
 	"time"
 
 	"github.com/nadoo/glider/common/log"
-	"github.com/nadoo/glider/common/pool"
 	"github.com/nadoo/glider/proxy"
 )
 
@@ -257,10 +257,8 @@ func (c *Client) AddRecord(record string) error {
 
 	b, _ := m.Marshal()
 
-	buf := pool.GetWriteBuffer()
+	var buf bytes.Buffer
-	defer pool.PutWriteBuffer(buf)
+	binary.Write(&buf, binary.BigEndian, uint16(len(b)))
-
-	binary.Write(buf, binary.BigEndian, uint16(len(b)))
 	buf.Write(b)
 
 	c.cache.Put(getKey(m.Question), buf.Bytes(), LongTTL)

dns/message.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"encoding/binary"
 	"errors"
+	"io"
 	"math/rand"
 	"net"
 	"strings"
@@ -89,29 +90,17 @@ func (m *Message) AddAnswer(rr *RR) error {
 
 // Marshal marshals message struct to []byte.
 func (m *Message) Marshal() ([]byte, error) {
-	var buf bytes.Buffer
+	buf := &bytes.Buffer{}
 
 	m.Header.SetQdcount(1)
 	m.Header.SetAncount(len(m.Answers))
 
-	b, err := m.Header.Marshal()
+	// no error when write to bytes.Buffer
-	if err != nil {
+	m.Header.MarshalTo(buf)
-		return nil, err
+	m.Question.MarshalTo(buf)
-	}
-	buf.Write(b)
-
-	b, err = m.Question.Marshal()
-	if err != nil {
-		return nil, err
-	}
-	buf.Write(b)
 
 	for _, answer := range m.Answers {
-		b, err := answer.Marshal()
+		answer.MarshalTo(buf)
-		if err != nil {
-			return nil, err
-		}
-		buf.Write(b)
 	}
 
 	return buf.Bytes(), nil
@@ -208,11 +197,9 @@ func (h *Header) setFlag(QR uint16, Opcode uint16, AA uint16,
 	h.Bits = QR<<15 + Opcode<<11 + AA<<10 + TC<<9 + RD<<8 + RA<<7 + RCODE
 }
 
-// Marshal marshals header struct to []byte.
+// MarshalTo marshals header struct to []byte and write to w.
-func (h *Header) Marshal() ([]byte, error) {
+func (h *Header) MarshalTo(w io.Writer) (int, error) {
-	var buf bytes.Buffer
+	return HeaderLen, binary.Write(w, binary.BigEndian, h)
-	err := binary.Write(&buf, binary.BigEndian, h)
-	return buf.Bytes(), err
 }
 
 // UnmarshalHeader unmarshals []bytes to Header.
@@ -267,15 +254,15 @@ func NewQuestion(qtype uint16, domain string) *Question {
 	}
 }
 
-// Marshal marshals Question struct to []byte.
+// MarshalTo marshals Question struct to []byte and write to w.
-func (q *Question) Marshal() ([]byte, error) {
+func (q *Question) MarshalTo(w io.Writer) (int, error) {
-	var buf bytes.Buffer
+	n, _ := MarshalDomainTo(w, q.QNAME)
 
-	buf.Write(MarshalDomain(q.QNAME))
+	binary.Write(w, binary.BigEndian, q.QTYPE)
-	binary.Write(&buf, binary.BigEndian, q.QTYPE)
+	binary.Write(w, binary.BigEndian, q.QCLASS)
-	binary.Write(&buf, binary.BigEndian, q.QCLASS)
+	n += 4
 
-	return buf.Bytes(), nil
+	return n, nil
 }
 
 // UnmarshalQuestion unmarshals []bytes to Question.
@@ -345,18 +332,20 @@ func NewRR() *RR {
 	return rr
 }
 
-// Marshal marshals RR struct to []byte.
+// MarshalTo marshals RR struct to []byte and write to w.
-func (rr *RR) Marshal() ([]byte, error) {
+func (rr *RR) MarshalTo(w io.Writer) (int, error) {
-	var buf bytes.Buffer
+	n, _ := MarshalDomainTo(w, rr.NAME)
 
-	buf.Write(MarshalDomain(rr.NAME))
+	binary.Write(w, binary.BigEndian, rr.TYPE)
-	binary.Write(&buf, binary.BigEndian, rr.TYPE)
+	binary.Write(w, binary.BigEndian, rr.CLASS)
-	binary.Write(&buf, binary.BigEndian, rr.CLASS)
+	binary.Write(w, binary.BigEndian, rr.TTL)
-	binary.Write(&buf, binary.BigEndian, rr.TTL)
+	binary.Write(w, binary.BigEndian, rr.RDLENGTH)
-	binary.Write(&buf, binary.BigEndian, rr.RDLENGTH)
+	n += 10
-	buf.Write(rr.RDATA)
 
-	return buf.Bytes(), nil
+	w.Write(rr.RDATA)
+	n += len(rr.RDATA)
+
+	return n, nil
 }
 
 // UnmarshalRR unmarshals []bytes to RR.
@@ -399,17 +388,17 @@ func (m *Message) UnmarshalRR(start int, rr *RR) (n int, err error) {
 	return n, nil
 }
 
-// MarshalDomain marshals domain string struct to []byte.
+// MarshalDomainTo marshals domain string struct to []byte and write to w.
-func MarshalDomain(domain string) []byte {
+func MarshalDomainTo(w io.Writer, domain string) (int, error) {
-	var buf bytes.Buffer
+	n := 1
-
 	for _, seg := range strings.Split(domain, ".") {
-		binary.Write(&buf, binary.BigEndian, byte(len(seg)))
+		binary.Write(w, binary.BigEndian, byte(len(seg)))
-		binary.Write(&buf, binary.BigEndian, []byte(seg))
+		binary.Write(w, binary.BigEndian, []byte(seg))
+		n += 1 + len(seg)
 	}
-	binary.Write(&buf, binary.BigEndian, byte(0x00))
+	binary.Write(w, binary.BigEndian, byte(0x00))
 
-	return buf.Bytes()
+	return n, nil
 }
 
 // UnmarshalDomain gets domain from bytes.

proxy/http/request.go

@@ -109,16 +109,12 @@ func parseRequest(r *bufio.Reader) (*request, error) {
 	return req, nil
 }
 
-func (r *request) Marshal() []byte {
+func (r *request) WriteBuf(buf *bytes.Buffer) {
-	var buf bytes.Buffer
+	writeStartLine(buf, r.method, r.ruri, r.proto)
-	writeStartLine(&buf, r.method, r.ruri, r.proto)
+	writeHeaders(buf, r.header)
-	writeHeaders(&buf, r.header)
-	return buf.Bytes()
 }
 
-func (r *request) MarshalAbs() []byte {
+func (r *request) WriteAbsBuf(buf *bytes.Buffer) {
-	var buf bytes.Buffer
+	writeStartLine(buf, r.method, r.absuri, r.proto)
-	writeStartLine(&buf, r.method, r.absuri, r.proto)
+	writeHeaders(buf, r.header)
-	writeHeaders(&buf, r.header)
-	return buf.Bytes()
 }

proxy/http/server.go

@@ -120,8 +120,12 @@ func (s *HTTP) servHTTP(req *request, c *conn.Conn) {
 	}
 	defer rc.Close()
 
+	buf := pool.GetWriteBuffer()
+	defer pool.PutWriteBuffer(buf)
+
 	// send request to remote server
-	_, err = rc.Write(req.Marshal())
+	req.WriteBuf(buf)
+	_, err = rc.Write(buf.Bytes())
 	if err != nil {
 		return
 	}
@@ -159,13 +163,12 @@ func (s *HTTP) servHTTP(req *request, c *conn.Conn) {
 	header.Set("Proxy-Connection", "close")
 	header.Set("Connection", "close")
 
-	buf := pool.GetWriteBuffer()
+	buf.Reset()
 	writeStartLine(buf, proto, code, status)
 	writeHeaders(buf, header)
 
 	log.F("[http] %s <-> %s", c.RemoteAddr(), req.target)
 	c.Write(buf.Bytes())
-	pool.PutWriteBuffer(buf)
 
 	b := pool.GetBuffer(conn.TCPBufSize)
 	io.CopyBuffer(c, r, b)

proxy/obfs/tls.go

@@ -64,20 +64,22 @@ func (c *TLSObfsConn) Write(b []byte) (int, error) {
 		return c.handshake(b)
 	}
 
+	buf := pool.GetWriteBuffer()
+	defer pool.PutWriteBuffer(buf)
+
 	n := len(b)
 	for i := 0; i < n; i += chunkSize {
+		buf.Reset()
 		end := i + chunkSize
 		if end > n {
 			end = n
 		}
 
-		buf := pool.GetWriteBuffer()
 		buf.Write([]byte{0x17, 0x03, 0x03})
 		binary.Write(buf, binary.BigEndian, uint16(len(b[i:end])))
 		buf.Write(b[i:end])
 
 		_, err := c.Conn.Write(buf.Bytes())
-		pool.PutWriteBuffer(buf)
 		if err != nil {
 			return 0, err
 		}