bughz/glider
1
0
Fork 0
mirror of https://github.com/nadoo/glider.git synced 2025-02-23 01:15:41 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

dns: marshal to io.writer instead of byte slice

Browse Source
This commit is contained in:
nadoo 2020-04-20 12:41:53 +08:00
parent a0542a028e
commit 149824cc20
5 changed files with 55 additions and 67 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
dns/client.go
View File

@ -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)

83
dns/message.go
View File

@ -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.

16
proxy/http/request.go
View File

@ -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)
}

9
proxy/http/server.go
View File

@ -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)

6
proxy/obfs/tls.go
View File

@ -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
}