Improve: pool buffer alloc

component/simple-obfs/http.go

@@ -34,15 +34,15 @@ func (ho *HTTPObfs) Read(b []byte) (int, error) {
 	}
 
 	if ho.firstResponse {
-		buf := pool.BufPool.Get().([]byte)
+		buf := pool.Get(pool.RelayBufferSize)
 		n, err := ho.Conn.Read(buf)
 		if err != nil {
-			pool.BufPool.Put(buf[:cap(buf)])
+			pool.Put(buf)
 			return 0, err
 		}
 		idx := bytes.Index(buf[:n], []byte("\r\n\r\n"))
 		if idx == -1 {
-			pool.BufPool.Put(buf[:cap(buf)])
+			pool.Put(buf)
 			return 0, io.EOF
 		}
 		ho.firstResponse = false
@@ -52,7 +52,7 @@ func (ho *HTTPObfs) Read(b []byte) (int, error) {
 			ho.buf = buf[:idx+4+length]
 			ho.offset = idx + 4 + n
 		} else {
-			pool.BufPool.Put(buf[:cap(buf)])
+			pool.Put(buf)
 		}
 		return n, nil
 	}

component/simple-obfs/tls.go

@@ -29,12 +29,12 @@ type TLSObfs struct {
 }
 
 func (to *TLSObfs) read(b []byte, discardN int) (int, error) {
-	buf := pool.BufPool.Get().([]byte)
-	_, err := io.ReadFull(to.Conn, buf[:discardN])
+	buf := pool.Get(discardN)
+	_, err := io.ReadFull(to.Conn, buf)
 	if err != nil {
 		return 0, err
 	}
-	pool.BufPool.Put(buf[:cap(buf)])
+	pool.Put(buf)
 
 	sizeBuf := make([]byte, 2)
 	_, err = io.ReadFull(to.Conn, sizeBuf)
@@ -102,15 +102,11 @@ func (to *TLSObfs) write(b []byte) (int, error) {
 		return len(b), err
 	}
 
-	size := pool.BufPool.Get().([]byte)
-	binary.BigEndian.PutUint16(size[:2], uint16(len(b)))
-
 	buf := &bytes.Buffer{}
 	buf.Write([]byte{0x17, 0x03, 0x03})
-	buf.Write(size[:2])
+	binary.Write(buf, binary.BigEndian, uint16(len(b)))
 	buf.Write(b)
 	_, err := to.Conn.Write(buf.Bytes())
-	pool.BufPool.Put(size[:cap(size)])
 	return len(b), err
 }
 

component/vmess/aead.go

@@ -22,8 +22,8 @@ func newAEADWriter(w io.Writer, aead cipher.AEAD, iv []byte) *aeadWriter {
 }
 
 func (w *aeadWriter) Write(b []byte) (n int, err error) {
-	buf := pool.BufPool.Get().([]byte)
-	defer pool.BufPool.Put(buf[:cap(buf)])
+	buf := pool.Get(pool.RelayBufferSize)
+	defer pool.Put(buf)
 	length := len(b)
 	for {
 		if length == 0 {
@@ -73,7 +73,7 @@ func (r *aeadReader) Read(b []byte) (int, error) {
 		n := copy(b, r.buf[r.offset:])
 		r.offset += n
 		if r.offset == len(r.buf) {
-			pool.BufPool.Put(r.buf[:cap(r.buf)])
+			pool.Put(r.buf)
 			r.buf = nil
 		}
 		return n, nil
@@ -89,10 +89,10 @@ func (r *aeadReader) Read(b []byte) (int, error) {
 		return 0, errors.New("Buffer is larger than standard")
 	}
 
-	buf := pool.BufPool.Get().([]byte)
+	buf := pool.Get(size)
 	_, err = io.ReadFull(r.Reader, buf[:size])
 	if err != nil {
-		pool.BufPool.Put(buf[:cap(buf)])
+		pool.Put(buf)
 		return 0, err
 	}
 
@@ -107,7 +107,7 @@ func (r *aeadReader) Read(b []byte) (int, error) {
 	realLen := size - r.Overhead()
 	n := copy(b, buf[:realLen])
 	if len(b) >= realLen {
-		pool.BufPool.Put(buf[:cap(buf)])
+		pool.Put(buf)
 		return n, nil
 	}
 

component/vmess/chunk.go

@@ -34,7 +34,7 @@ func (cr *chunkReader) Read(b []byte) (int, error) {
 		n := copy(b, cr.buf[cr.offset:])
 		cr.offset += n
 		if cr.offset == len(cr.buf) {
-			pool.BufPool.Put(cr.buf[:cap(cr.buf)])
+			pool.Put(cr.buf)
 			cr.buf = nil
 		}
 		return n, nil
@@ -59,15 +59,15 @@ func (cr *chunkReader) Read(b []byte) (int, error) {
 		return size, nil
 	}
 
-	buf := pool.BufPool.Get().([]byte)
-	_, err = io.ReadFull(cr.Reader, buf[:size])
+	buf := pool.Get(size)
+	_, err = io.ReadFull(cr.Reader, buf)
 	if err != nil {
-		pool.BufPool.Put(buf[:cap(buf)])
+		pool.Put(buf)
 		return 0, err
 	}
-	n := copy(b, cr.buf[:])
+	n := copy(b, buf)
 	cr.offset = n
-	cr.buf = buf[:size]
+	cr.buf = buf
 	return n, nil
 }
 
@@ -76,8 +76,8 @@ type chunkWriter struct {
 }
 
 func (cw *chunkWriter) Write(b []byte) (n int, err error) {
-	buf := pool.BufPool.Get().([]byte)
-	defer pool.BufPool.Put(buf[:cap(buf)])
+	buf := pool.Get(pool.RelayBufferSize)
+	defer pool.Put(buf)
 	length := len(b)
 	for {
 		if length == 0 {