Use buffer in socks proxy

common/alloc/buffer.go

@@ -17,26 +17,39 @@ func (b *Buffer) Release() {
 	b.pool = nil
 }
 
-func (b *Buffer) Clear() {
+func (b *Buffer) Clear() *Buffer {
 	b.Value = b.head[:0]
+	return b
 }
 
-func (b *Buffer) Append(data []byte) {
+func (b *Buffer) AppendBytes(bytes ...byte) *Buffer {
+	b.Value = append(b.Value, bytes...)
+	return b
+}
+
+func (b *Buffer) Append(data []byte) *Buffer {
 	b.Value = append(b.Value, data...)
+	return b
 }
 
-func (b *Buffer) Slice(from, to int) {
+func (b *Buffer) Slice(from, to int) *Buffer {
 	b.Value = b.Value[from:to]
+	return b
 }
 
-func (b *Buffer) SliceFrom(from int) {
+func (b *Buffer) SliceFrom(from int) *Buffer {
 	b.Value = b.Value[from:]
+	return b
 }
 
 func (b *Buffer) Len() int {
 	return len(b.Value)
 }
 
+func (b *Buffer) IsFull() bool {
+	return len(b.Value) == cap(b.Value)
+}
+
 type bufferPool struct {
 	chain        chan []byte
 	bufferSize   int
@@ -90,8 +103,18 @@ func (p *bufferPool) cleanup(tick <-chan time.Time) {
 	}
 }
 
-var smallPool = newBufferPool(8*1024, 256, 2048)
+var smallPool = newBufferPool(1024, 16, 64)
+var mediumPool = newBufferPool(8*1024, 256, 2048)
+var largePool = newBufferPool(64*1024, 16, 64)
 
-func NewBuffer() *Buffer {
+func NewSmallBuffer() *Buffer {
 	return smallPool.allocate()
 }
+
+func NewBuffer() *Buffer {
+	return mediumPool.allocate()
+}
+
+func NewLargeBuffer() *Buffer {
+	return largePool.allocate()
+}

proxy/socks/protocol/socks.go

@@ -4,6 +4,7 @@ import (
 	"encoding/binary"
 	"io"
 
+	"github.com/v2ray/v2ray-core/common/alloc"
 	"github.com/v2ray/v2ray-core/common/errors"
 	"github.com/v2ray/v2ray-core/common/log"
 	v2net "github.com/v2ray/v2ray-core/common/net"
@@ -39,8 +40,10 @@ func (request *Socks5AuthenticationRequest) HasAuthMethod(method byte) bool {
 }
 
 func ReadAuthentication(reader io.Reader) (auth Socks5AuthenticationRequest, auth4 Socks4AuthenticationRequest, err error) {
-	buffer := make([]byte, 256)
-	nBytes, err := reader.Read(buffer)
+	buffer := alloc.NewSmallBuffer()
+	defer buffer.Release()
+
+	nBytes, err := reader.Read(buffer.Value)
 	if err != nil {
 		return
 	}
@@ -50,22 +53,22 @@ func ReadAuthentication(reader io.Reader) (auth Socks5AuthenticationRequest, aut
 		return
 	}
 
-	if buffer[0] == socks4Version {
-		auth4.Version = buffer[0]
-		auth4.Command = buffer[1]
-		auth4.Port = binary.BigEndian.Uint16(buffer[2:4])
-		copy(auth4.IP[:], buffer[4:8])
+	if buffer.Value[0] == socks4Version {
+		auth4.Version = buffer.Value[0]
+		auth4.Command = buffer.Value[1]
+		auth4.Port = binary.BigEndian.Uint16(buffer.Value[2:4])
+		copy(auth4.IP[:], buffer.Value[4:8])
 		err = NewSocksVersion4Error()
 		return
 	}
 
-	auth.version = buffer[0]
+	auth.version = buffer.Value[0]
 	if auth.version != socksVersion {
 		err = errors.NewProtocolVersionError(int(auth.version))
 		return
 	}
 
-	auth.nMethods = buffer[1]
+	auth.nMethods = buffer.Value[1]
 	if auth.nMethods <= 0 {
 		log.Info("Zero length of authentication methods")
 		err = errors.NewCorruptedPacketError()
@@ -77,7 +80,7 @@ func ReadAuthentication(reader io.Reader) (auth Socks5AuthenticationRequest, aut
 		err = errors.NewCorruptedPacketError()
 		return
 	}
-	copy(auth.authMethods[:], buffer[2:nBytes])
+	copy(auth.authMethods[:], buffer.Value[2:nBytes])
 	return
 }
 
@@ -113,29 +116,31 @@ func (request Socks5UserPassRequest) AuthDetail() string {
 }
 
 func ReadUserPassRequest(reader io.Reader) (request Socks5UserPassRequest, err error) {
-	buffer := make([]byte, 256)
-	_, err = reader.Read(buffer[0:2])
+	buffer := alloc.NewSmallBuffer()
+	defer buffer.Release()
+
+	_, err = reader.Read(buffer.Value[0:2])
 	if err != nil {
 		return
 	}
-	request.version = buffer[0]
-	nUsername := buffer[1]
-	nBytes, err := reader.Read(buffer[:nUsername])
+	request.version = buffer.Value[0]
+	nUsername := buffer.Value[1]
+	nBytes, err := reader.Read(buffer.Value[:nUsername])
 	if err != nil {
 		return
 	}
-	request.username = string(buffer[:nBytes])
+	request.username = string(buffer.Value[:nBytes])
 
-	_, err = reader.Read(buffer[0:1])
+	_, err = reader.Read(buffer.Value[0:1])
 	if err != nil {
 		return
 	}
-	nPassword := buffer[0]
-	nBytes, err = reader.Read(buffer[:nPassword])
+	nPassword := buffer.Value[0]
+	nBytes, err = reader.Read(buffer.Value[:nPassword])
 	if err != nil {
 		return
 	}
-	request.password = string(buffer[:nBytes])
+	request.password = string(buffer.Value[:nBytes])
 	return
 }
 
@@ -177,8 +182,10 @@ type Socks5Request struct {
 }
 
 func ReadRequest(reader io.Reader) (request *Socks5Request, err error) {
-	buffer := make([]byte, 256)
-	nBytes, err := reader.Read(buffer[:4])
+	buffer := alloc.NewSmallBuffer()
+	defer buffer.Release()
+
+	nBytes, err := reader.Read(buffer.Value[:4])
 	if err != nil {
 		return
 	}
@@ -187,10 +194,10 @@ func ReadRequest(reader io.Reader) (request *Socks5Request, err error) {
 		return
 	}
 	request = &Socks5Request{
-		Version: buffer[0],
-		Command: buffer[1],
+		Version: buffer.Value[0],
+		Command: buffer.Value[1],
 		// buffer[2] is a reserved field
-		AddrType: buffer[3],
+		AddrType: buffer.Value[3],
 	}
 	switch request.AddrType {
 	case AddrTypeIPv4:
@@ -203,12 +210,12 @@ func ReadRequest(reader io.Reader) (request *Socks5Request, err error) {
 			return
 		}
 	case AddrTypeDomain:
-		nBytes, err = reader.Read(buffer[0:1])
+		nBytes, err = reader.Read(buffer.Value[0:1])
 		if err != nil {
 			return
 		}
-		domainLength := buffer[0]
-		nBytes, err = reader.Read(buffer[:domainLength])
+		domainLength := buffer.Value[0]
+		nBytes, err = reader.Read(buffer.Value[:domainLength])
 		if err != nil {
 			return
 		}
@@ -218,7 +225,7 @@ func ReadRequest(reader io.Reader) (request *Socks5Request, err error) {
 			err = errors.NewCorruptedPacketError()
 			return
 		}
-		request.Domain = string(buffer[:domainLength])
+		request.Domain = string(buffer.Value[:domainLength])
 	case AddrTypeIPv6:
 		nBytes, err = reader.Read(request.IPv6[:])
 		if err != nil {
@@ -234,7 +241,7 @@ func ReadRequest(reader io.Reader) (request *Socks5Request, err error) {
 		return
 	}
 
-	nBytes, err = reader.Read(buffer[:2])
+	nBytes, err = reader.Read(buffer.Value[:2])
 	if err != nil {
 		return
 	}
@@ -243,7 +250,7 @@ func ReadRequest(reader io.Reader) (request *Socks5Request, err error) {
 		return
 	}
 
-	request.Port = binary.BigEndian.Uint16(buffer)
+	request.Port = binary.BigEndian.Uint16(buffer.Value[:2])
 	return
 }
 
@@ -305,26 +312,16 @@ func (r *Socks5Response) SetDomain(domain string) {
 	r.Domain = domain
 }
 
-func (r *Socks5Response) toBytes() []byte {
-	buffer := make([]byte, 0, 300)
-	buffer = append(buffer, r.Version)
-	buffer = append(buffer, r.Error)
-	buffer = append(buffer, 0x00) // reserved
-	buffer = append(buffer, r.AddrType)
+func (r *Socks5Response) Write(buffer *alloc.Buffer) {
+	buffer.AppendBytes(r.Version, r.Error, 0x00 /* reserved */, r.AddrType)
 	switch r.AddrType {
 	case 0x01:
-		buffer = append(buffer, r.IPv4[:]...)
+		buffer.Append(r.IPv4[:])
 	case 0x03:
-		buffer = append(buffer, byte(len(r.Domain)))
-		buffer = append(buffer, []byte(r.Domain)...)
+		buffer.AppendBytes(byte(len(r.Domain)))
+		buffer.Append([]byte(r.Domain))
 	case 0x04:
-		buffer = append(buffer, r.IPv6[:]...)
+		buffer.Append(r.IPv6[:])
 	}
-	buffer = append(buffer, byte(r.Port>>8), byte(r.Port))
-	return buffer
-}
-
-func WriteResponse(writer io.Writer, response *Socks5Response) error {
-	_, err := writer.Write(response.toBytes())
-	return err
+	buffer.AppendBytes(byte(r.Port>>8), byte(r.Port))
 }

proxy/socks/protocol/socks4.go

@@ -1,6 +1,7 @@
 package protocol
 
 import (
+	"github.com/v2ray/v2ray-core/common/alloc"
 	"github.com/v2ray/v2ray-core/common/errors"
 )
 
@@ -39,13 +40,8 @@ func NewSocks4AuthenticationResponse(result byte, port uint16, ip []byte) *Socks
 	}
 }
 
-func (r *Socks4AuthenticationResponse) ToBytes(buffer []byte) []byte {
-	if buffer == nil {
-		buffer = make([]byte, 8)
-	}
-	buffer[1] = r.result
-	buffer[2] = byte(r.port >> 8)
-	buffer[3] = byte(r.port)
-	copy(buffer[4:], r.ip)
-	return buffer
+func (r *Socks4AuthenticationResponse) Write(buffer *alloc.Buffer) {
+	buffer.AppendBytes(
+		byte(0x00), r.result, byte(r.port>>8), byte(r.port),
+		r.ip[0], r.ip[1], r.ip[2], r.ip[3])
 }

proxy/socks/protocol/socks4_test.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"testing"
 
+	"github.com/v2ray/v2ray-core/common/alloc"
 	"github.com/v2ray/v2ray-core/testing/unit"
 )
 
@@ -32,6 +33,10 @@ func TestSocks4AuthenticationResponseToBytes(t *testing.T) {
 		port:   443,
 		ip:     []byte{1, 2, 3, 4},
 	}
-	responseBytes := response.ToBytes(nil)
-	assert.Bytes(responseBytes).Equals([]byte{0x00, 0x10, 0x01, 0xBB, 0x01, 0x02, 0x03, 0x04})
+
+	buffer := alloc.NewSmallBuffer().Clear()
+	defer buffer.Release()
+
+	response.Write(buffer)
+	assert.Bytes(buffer.Value).Equals([]byte{0x00, 0x10, 0x01, 0xBB, 0x01, 0x02, 0x03, 0x04})
 }

proxy/socks/protocol/socks_test.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"testing"
 
+	"github.com/v2ray/v2ray-core/common/alloc"
 	"github.com/v2ray/v2ray-core/testing/unit"
 )
 
@@ -67,7 +68,7 @@ func TestRequestRead(t *testing.T) {
 	assert.Uint16(request.Port).Named("Port").Equals(53)
 }
 
-func TestResponseToBytes(t *testing.T) {
+func TestResponseWrite(t *testing.T) {
 	assert := unit.Assert(t)
 
 	response := Socks5Response{
@@ -79,7 +80,10 @@ func TestResponseToBytes(t *testing.T) {
 		[16]byte{},
 		uint16(53),
 	}
-	rawResponse := response.toBytes()
+	buffer := alloc.NewSmallBuffer().Clear()
+	defer buffer.Release()
+
+	response.Write(buffer)
 	expectedBytes := []byte{
 		socksVersion,
 		ErrorSuccess,
@@ -88,5 +92,5 @@ func TestResponseToBytes(t *testing.T) {
 		0x72, 0x72, 0x72, 0x72,
 		byte(0x00), byte(0x035),
 	}
-	assert.Bytes(rawResponse).Named("raw response").Equals(expectedBytes)
+	assert.Bytes(buffer.Value).Named("raw response").Equals(expectedBytes)
 }

proxy/socks/protocol/udp.go

@@ -75,9 +75,7 @@ func ReadUDPRequest(packet []byte) (request Socks5UDPRequest, err error) {
 		return
 	}
 
-	request.Data = alloc.NewBuffer()
-	request.Data.Clear()
-	request.Data.Append(packet[dataBegin:])
+	request.Data = alloc.NewBuffer().Clear().Append(packet[dataBegin:])
 
 	return
 }

proxy/socks/socks.go

@@ -133,7 +133,10 @@ func (server *SocksServer) handleSocks5(reader *v2net.TimeOutReader, writer io.W
 	if request.Command == protocol.CmdBind || request.Command == protocol.CmdUdpAssociate {
 		response := protocol.NewSocks5Response()
 		response.Error = protocol.ErrorCommandNotSupported
-		err = protocol.WriteResponse(writer, response)
+
+		responseBuffer := alloc.NewSmallBuffer().Clear()
+		response.Write(responseBuffer)
+		_, err = writer.Write(responseBuffer.Value)
 		if err != nil {
 			log.Error("Socks failed to write response: %v", err)
 			return err
@@ -152,7 +155,9 @@ func (server *SocksServer) handleSocks5(reader *v2net.TimeOutReader, writer io.W
 	response.IPv4[2] = 0
 	response.IPv4[3] = 0
 
-	err = protocol.WriteResponse(writer, response)
+	responseBuffer := alloc.NewSmallBuffer().Clear()
+	response.Write(responseBuffer)
+	_, err = writer.Write(responseBuffer.Value)
 	if err != nil {
 		log.Error("Socks failed to write response: %v", err)
 		return err
@@ -187,7 +192,12 @@ func (server *SocksServer) handleUDP(reader *v2net.TimeOutReader, writer io.Writ
 		response.AddrType = protocol.AddrTypeDomain
 		response.Domain = udpAddr.Domain()
 	}
-	err := protocol.WriteResponse(writer, response)
+
+	responseBuffer := alloc.NewSmallBuffer()
+	response.Write(responseBuffer)
+	_, err := writer.Write(responseBuffer.Value)
+	responseBuffer.Release()
+
 	if err != nil {
 		log.Error("Socks failed to write response: %v", err)
 		return err
@@ -209,7 +219,11 @@ func (server *SocksServer) handleSocks4(reader io.Reader, writer io.Writer, auth
 		result = protocol.Socks4RequestRejected
 	}
 	socks4Response := protocol.NewSocks4AuthenticationResponse(result, auth.Port, auth.IP[:])
-	writer.Write(socks4Response.ToBytes(nil))
+
+	responseBuffer := alloc.NewSmallBuffer().Clear()
+	socks4Response.Write(responseBuffer)
+	writer.Write(responseBuffer.Value)
+	responseBuffer.Release()
 
 	if result == protocol.Socks4RequestRejected {
 		return errors.NewInvalidOperationError("Socks4 command " + strconv.Itoa(int(auth.Command)))

proxy/vmess/vmessin.go

@@ -91,7 +91,7 @@ func (handler *VMessInboundHandler) HandleConnection(connection *net.TCPConn) er
 	}
 
 	// Optimize for small response packet
-	buffer := alloc.NewBuffer()
+	buffer := alloc.NewLargeBuffer()
 	buffer.Clear()
 	buffer.Append(request.ResponseHeader)