From 4c9cdf0bb9120a2df0c601268360d7d0540f8886 Mon Sep 17 00:00:00 2001 From: "rui.zheng" Date: Thu, 4 Feb 2016 13:45:26 +0800 Subject: [PATCH] add buffer pools --- README.md | 2 +- socks5.go | 190 +++++++++++++++++++++++++++++++++++------------------- 2 files changed, 124 insertions(+), 68 deletions(-) diff --git a/README.md b/README.md index 2e5d38e..9192c60 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ gosocks5 ======== -golang and socks 5 +golang and SOCKSV5 diff --git a/socks5.go b/socks5.go index 8a32b59..02151da 100644 --- a/socks5.go +++ b/socks5.go @@ -4,7 +4,7 @@ package gosocks5 import ( - "bytes" + //"bytes" "encoding/binary" "errors" "fmt" @@ -12,6 +12,7 @@ import ( //"log" "net" "strconv" + "sync" ) const ( @@ -61,16 +62,33 @@ var ( ErrAuthFailure = errors.New("Auth failure") ) +// buffer pools +var ( + sPool = sync.Pool{ + New: func() interface{} { + return make([]byte, 576) + }, + } // small buff pool + lPool = sync.Pool{ + New: func() interface{} { + return make([]byte, 64*1024+262) + }, + } // large buff pool for udp +) + /* Method selection -+----+----------+----------+ -|VER | NMETHODS | METHODS | -+----+----------+----------+ -| 1 | 1 | 1 to 255 | -+----+----------+----------+ + +----+----------+----------+ + |VER | NMETHODS | METHODS | + +----+----------+----------+ + | 1 | 1 | 1 to 255 | + +----+----------+----------+ */ func ReadMethods(r io.Reader) ([]uint8, error) { - b := make([]byte, 257) + //b := make([]byte, 257) + b := sPool.Get().([]byte) + defer sPool.Put(b) + n, err := io.ReadAtLeast(r, b, 2) if err != nil { return nil, err @@ -91,7 +109,10 @@ func ReadMethods(r io.Reader) ([]uint8, error) { } } - return b[2:length], nil + methods := make([]byte, int(b[1])) + copy(methods, b[2:length]) + + return methods, nil } func WriteMethod(method uint8, w io.Writer) error { @@ -101,11 +122,11 @@ func WriteMethod(method uint8, w io.Writer) error { /* Username/Password authentication request - +----+------+----------+------+----------+ - |VER | ULEN | UNAME | PLEN | PASSWD | - +----+------+----------+------+----------+ - | 1 | 1 | 1 to 255 | 1 | 1 to 255 | - +----+------+----------+------+----------+ + +----+------+----------+------+----------+ + |VER | ULEN | UNAME | PLEN | PASSWD | + +----+------+----------+------+----------+ + | 1 | 1 | 1 to 255 | 1 | 1 to 255 | + +----+------+----------+------+----------+ */ type UserPassRequest struct { Version byte @@ -122,7 +143,10 @@ func NewUserPassRequest(ver byte, u, p string) *UserPassRequest { } func ReadUserPassRequest(r io.Reader) (*UserPassRequest, error) { - b := make([]byte, 513) + // b := make([]byte, 513) + b := sPool.Get().([]byte) + defer sPool.Put(b) + n, err := io.ReadAtLeast(r, b, 2) if err != nil { return nil, err @@ -159,7 +183,10 @@ func ReadUserPassRequest(r io.Reader) (*UserPassRequest, error) { } func (req *UserPassRequest) Write(w io.Writer) error { - b := make([]byte, 513) + // b := make([]byte, 513) + b := sPool.Get().([]byte) + defer sPool.Put(b) + b[0] = req.Version ulen := len(req.Username) b[1] = byte(ulen) @@ -183,11 +210,11 @@ func (req *UserPassRequest) String() string { /* Username/Password authentication response - +----+--------+ - |VER | STATUS | - +----+--------+ - | 1 | 1 | - +----+--------+ + +----+--------+ + |VER | STATUS | + +----+--------+ + | 1 | 1 | + +----+--------+ */ type UserPassResponse struct { Version byte @@ -202,7 +229,10 @@ func NewUserPassResponse(ver, status byte) *UserPassResponse { } func ReadUserPassResponse(r io.Reader) (*UserPassResponse, error) { - b := make([]byte, 2) + // b := make([]byte, 2) + b := sPool.Get().([]byte) + defer sPool.Put(b) + if _, err := io.ReadFull(r, b); err != nil { return nil, err } @@ -229,6 +259,14 @@ func (res *UserPassResponse) String() string { res.Version, res.Status) } +/* +Address + +------+----------+----------+ + | ATYP | ADDR | PORT | + +------+----------+----------+ + | 1 | Variable | 2 | + +------+----------+----------+ +*/ type Addr struct { Type uint8 Host string @@ -287,11 +325,11 @@ func (addr *Addr) String() string { /* The SOCKSv5 request -+----+-----+-------+------+----------+----------+ -|VER | CMD | RSV | ATYP | DST.ADDR | DST.PORT | -+----+-----+-------+------+----------+----------+ -| 1 | 1 | X'00' | 1 | Variable | 2 | -+----+-----+-------+------+----------+----------+ + +----+-----+-------+------+----------+----------+ + |VER | CMD | RSV | ATYP | DST.ADDR | DST.PORT | + +----+-----+-------+------+----------+----------+ + | 1 | 1 | X'00' | 1 | Variable | 2 | + +----+-----+-------+------+----------+----------+ */ type Request struct { Cmd uint8 @@ -306,7 +344,10 @@ func NewRequest(cmd uint8, addr *Addr) *Request { } func ReadRequest(r io.Reader) (*Request, error) { - b := make([]byte, 262) + // b := make([]byte, 262) + b := sPool.Get().([]byte) + defer sPool.Put(b) + n, err := io.ReadAtLeast(r, b, 5) if err != nil { return nil, err @@ -348,11 +389,13 @@ func ReadRequest(r io.Reader) (*Request, error) { } func (r *Request) Write(w io.Writer) (err error) { - b := make([]byte, 262) + //b := make([]byte, 262) + b := sPool.Get().([]byte) + defer sPool.Put(b) b[0] = Ver5 b[1] = r.Cmd - // b[2] = 0 //rsv + b[2] = 0 //rsv b[3] = AddrIPv4 // default length := 10 @@ -375,11 +418,11 @@ func (r *Request) String() string { /* The SOCKSv5 reply -+----+-----+-------+------+----------+----------+ -|VER | REP | RSV | ATYP | BND.ADDR | BND.PORT | -+----+-----+-------+------+----------+----------+ -| 1 | 1 | X'00' | 1 | Variable | 2 | -+----+-----+-------+------+----------+----------+ + +----+-----+-------+------+----------+----------+ + |VER | REP | RSV | ATYP | BND.ADDR | BND.PORT | + +----+-----+-------+------+----------+----------+ + | 1 | 1 | X'00' | 1 | Variable | 2 | + +----+-----+-------+------+----------+----------+ */ type Reply struct { Rep uint8 @@ -394,7 +437,10 @@ func NewReply(rep uint8, addr *Addr) *Reply { } func ReadReply(r io.Reader) (*Reply, error) { - b := make([]byte, 262) + // b := make([]byte, 262) + b := sPool.Get().([]byte) + defer sPool.Put(b) + n, err := io.ReadAtLeast(r, b, 5) if err != nil { return nil, err @@ -437,11 +483,13 @@ func ReadReply(r io.Reader) (*Reply, error) { } func (r *Reply) Write(w io.Writer) (err error) { - b := make([]byte, 262) + // b := make([]byte, 262) + b := sPool.Get().([]byte) + defer sPool.Put(b) b[0] = Ver5 b[1] = r.Rep - // b[2] = 0 //rsv + b[2] = 0 //rsv b[3] = AddrIPv4 // default length := 10 @@ -465,11 +513,11 @@ func (r *Reply) String() string { /* UDP request -+----+------+------+----------+----------+----------+ -|RSV | FRAG | ATYP | DST.ADDR | DST.PORT | DATA | -+----+------+------+----------+----------+----------+ -| 2 | 1 | 1 | Variable | 2 | Variable | -+----+------+------+----------+----------+----------+ + +----+------+------+----------+----------+----------+ + |RSV | FRAG | ATYP | DST.ADDR | DST.PORT | DATA | + +----+------+------+----------+----------+----------+ + | 2 | 1 | 1 | Variable | 2 | Variable | + +----+------+------+----------+----------+----------+ */ type UDPHeader struct { Rsv uint16 @@ -485,6 +533,23 @@ func NewUDPHeader(rsv uint16, frag uint8, addr *Addr) *UDPHeader { } } +func (h *UDPHeader) Write(w io.Writer) error { + b := sPool.Get().([]byte) + defer sPool.Put(b) + + binary.BigEndian.PutUint16(b[:2], h.Rsv) + b[2] = h.Frag + + addr := h.Addr + if addr == nil { + addr = &Addr{} + } + length, _ := addr.Encode(b[3:]) + + _, err := w.Write(b[:3+length]) + return err +} + func (h *UDPHeader) String() string { return fmt.Sprintf("%d %d %d %s", h.Rsv, h.Frag, h.Addr.Type, h.Addr.String()) @@ -503,7 +568,10 @@ func NewUDPDatagram(header *UDPHeader, data []byte) *UDPDatagram { } func ReadUDPDatagram(r io.Reader) (*UDPDatagram, error) { - b := make([]byte, 65797) + // b := make([]byte, 65797) + b := lPool.Get().([]byte) + defer lPool.Put(b) + n, err := io.ReadAtLeast(r, b, 5) if err != nil { return nil, err @@ -526,7 +594,7 @@ func ReadUDPDatagram(r io.Reader) (*UDPDatagram, error) { default: return nil, ErrBadAddrType } - // extended feature, for udp over tcp + // extended feature, for udp over tcp, using reserved field for data length dlen := int(header.Rsv) if n < hlen+dlen { if _, err := io.ReadFull(r, b[n:hlen+dlen]); err != nil { @@ -540,38 +608,26 @@ func ReadUDPDatagram(r io.Reader) (*UDPDatagram, error) { return nil, err } + data := make([]byte, dlen) + copy(data, b[hlen:n]) + d := &UDPDatagram{ Header: header, - Data: b[hlen:n], + Data: data, } return d, nil } func (d *UDPDatagram) Write(w io.Writer) error { - buffer := &bytes.Buffer{} - - b := make([]byte, 259) - if d.Header != nil { - binary.BigEndian.PutUint16(b[:2], d.Header.Rsv) - buffer.Write(b[:2]) - buffer.WriteByte(d.Header.Frag) - - b[0] = AddrIPv4 - b[1] = 0 - length := 7 - - if d.Header.Addr != nil { - length, _ = d.Header.Addr.Encode(b) - } - buffer.Write(b[:length]) - } else { - b[3] = AddrIPv4 - buffer.Write(b[:10]) + h := d.Header + if h == nil { + h = &UDPHeader{} } - - buffer.Write(d.Data) - _, err := w.Write(buffer.Bytes()) + if err := h.Write(w); err != nil { + return err + } + _, err := w.Write(d.Data) return err }