From 149824cc20ab94303f3f25567da2acb3258a1293 Mon Sep 17 00:00:00 2001 From: nadoo <287492+nadoo@users.noreply.github.com> Date: Mon, 20 Apr 2020 12:41:53 +0800 Subject: [PATCH] dns: marshal to io.writer instead of byte slice --- dns/client.go | 8 ++--- dns/message.go | 83 +++++++++++++++++++------------------------ proxy/http/request.go | 16 ++++----- proxy/http/server.go | 9 +++-- proxy/obfs/tls.go | 6 ++-- 5 files changed, 55 insertions(+), 67 deletions(-) diff --git a/dns/client.go b/dns/client.go index c2bcefe..0ef42c1 100644 --- a/dns/client.go +++ b/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() - defer pool.PutWriteBuffer(buf) - - binary.Write(buf, binary.BigEndian, uint16(len(b))) + var buf bytes.Buffer + binary.Write(&buf, binary.BigEndian, uint16(len(b))) buf.Write(b) c.cache.Put(getKey(m.Question), buf.Bytes(), LongTTL) diff --git a/dns/message.go b/dns/message.go index efc0f23..f1d1920 100644 --- a/dns/message.go +++ b/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() - if err != nil { - return nil, err - } - buf.Write(b) - - b, err = m.Question.Marshal() - if err != nil { - return nil, err - } - buf.Write(b) + // no error when write to bytes.Buffer + m.Header.MarshalTo(buf) + m.Question.MarshalTo(buf) for _, answer := range m.Answers { - b, err := answer.Marshal() - if err != nil { - return nil, err - } - buf.Write(b) + answer.MarshalTo(buf) } 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. -func (h *Header) Marshal() ([]byte, error) { - var buf bytes.Buffer - err := binary.Write(&buf, binary.BigEndian, h) - return buf.Bytes(), err +// MarshalTo marshals header struct to []byte and write to w. +func (h *Header) MarshalTo(w io.Writer) (int, error) { + return HeaderLen, binary.Write(w, binary.BigEndian, h) } // UnmarshalHeader unmarshals []bytes to Header. @@ -267,15 +254,15 @@ func NewQuestion(qtype uint16, domain string) *Question { } } -// Marshal marshals Question struct to []byte. -func (q *Question) Marshal() ([]byte, error) { - var buf bytes.Buffer +// MarshalTo marshals Question struct to []byte and write to w. +func (q *Question) MarshalTo(w io.Writer) (int, error) { + n, _ := MarshalDomainTo(w, q.QNAME) - buf.Write(MarshalDomain(q.QNAME)) - binary.Write(&buf, binary.BigEndian, q.QTYPE) - binary.Write(&buf, binary.BigEndian, q.QCLASS) + binary.Write(w, binary.BigEndian, q.QTYPE) + binary.Write(w, 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. -func (rr *RR) Marshal() ([]byte, error) { - var buf bytes.Buffer +// MarshalTo marshals RR struct to []byte and write to w. +func (rr *RR) MarshalTo(w io.Writer) (int, error) { + n, _ := MarshalDomainTo(w, rr.NAME) - buf.Write(MarshalDomain(rr.NAME)) - binary.Write(&buf, binary.BigEndian, rr.TYPE) - binary.Write(&buf, binary.BigEndian, rr.CLASS) - binary.Write(&buf, binary.BigEndian, rr.TTL) - binary.Write(&buf, binary.BigEndian, rr.RDLENGTH) - buf.Write(rr.RDATA) + binary.Write(w, binary.BigEndian, rr.TYPE) + binary.Write(w, binary.BigEndian, rr.CLASS) + binary.Write(w, binary.BigEndian, rr.TTL) + binary.Write(w, binary.BigEndian, rr.RDLENGTH) + n += 10 - 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. -func MarshalDomain(domain string) []byte { - var buf bytes.Buffer - +// MarshalDomainTo marshals domain string struct to []byte and write to w. +func MarshalDomainTo(w io.Writer, domain string) (int, error) { + n := 1 for _, seg := range strings.Split(domain, ".") { - binary.Write(&buf, binary.BigEndian, byte(len(seg))) - binary.Write(&buf, binary.BigEndian, []byte(seg)) + binary.Write(w, binary.BigEndian, byte(len(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. diff --git a/proxy/http/request.go b/proxy/http/request.go index f22a103..26b4090 100644 --- a/proxy/http/request.go +++ b/proxy/http/request.go @@ -109,16 +109,12 @@ func parseRequest(r *bufio.Reader) (*request, error) { return req, nil } -func (r *request) Marshal() []byte { - var buf bytes.Buffer - writeStartLine(&buf, r.method, r.ruri, r.proto) - writeHeaders(&buf, r.header) - return buf.Bytes() +func (r *request) WriteBuf(buf *bytes.Buffer) { + writeStartLine(buf, r.method, r.ruri, r.proto) + writeHeaders(buf, r.header) } -func (r *request) MarshalAbs() []byte { - var buf bytes.Buffer - writeStartLine(&buf, r.method, r.absuri, r.proto) - writeHeaders(&buf, r.header) - return buf.Bytes() +func (r *request) WriteAbsBuf(buf *bytes.Buffer) { + writeStartLine(buf, r.method, r.absuri, r.proto) + writeHeaders(buf, r.header) } diff --git a/proxy/http/server.go b/proxy/http/server.go index 1e1d77f..21e922d 100644 --- a/proxy/http/server.go +++ b/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) diff --git a/proxy/obfs/tls.go b/proxy/obfs/tls.go index 224d413..fed65fc 100644 --- a/proxy/obfs/tls.go +++ 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 }