make use of buf.NewWithSize() in Trojan's PacketReader and PacketWriter

proxy/trojan/protocol.go

@@ -21,7 +21,6 @@ var (
 )
 
 const (
-	maxLength       = 8192
 	commandTCP byte = 1
 	commandUDP byte = 3
 )
@@ -136,15 +135,22 @@ func (w *PacketWriter) WriteTo(payload []byte, addr gonet.Addr) (int, error) {
 }
 
 func (w *PacketWriter) writePacket(payload []byte, dest net.Destination) (int, error) { // nolint: unparam
-	buffer := buf.StackNew()
-	defer buffer.Release()
-
 	length := len(payload)
 	lengthBuf := [2]byte{}
 	binary.BigEndian.PutUint16(lengthBuf[:], uint16(length))
-	if err := addrParser.WriteAddressPort(&buffer, dest.Address, dest.Port); err != nil {
+
+	addrPortBuf := buf.StackNew()
+	if err := addrParser.WriteAddressPort(&addrPortBuf, dest.Address, dest.Port); err != nil {
+		return 0, err
+	}
+
+	buffer := buf.NewWithSize(addrPortBuf.Len() + 2 + 2 + int32(length))
+	defer buffer.Release()
+
+	if _, err := buffer.Write(addrPortBuf.Bytes()); err != nil {
 		return 0, err
 	}
+	addrPortBuf.Release()
 	if _, err := buffer.Write(lengthBuf[:]); err != nil {
 		return 0, err
 	}
@@ -255,36 +261,22 @@ func (r *PacketReader) ReadMultiBufferWithMetadata() (*PacketPayload, error) {
 		return nil, newError("failed to read payload length").Base(err)
 	}
 
-	remain := int(binary.BigEndian.Uint16(lengthBuf[:]))
-	if remain > maxLength {
-		return nil, newError("oversize payload")
-	}
+	length := binary.BigEndian.Uint16(lengthBuf[:])
 
 	var crlf [2]byte
 	if _, err := io.ReadFull(r, crlf[:]); err != nil {
 		return nil, newError("failed to read crlf").Base(err)
 	}
 
 	dest := net.UDPDestination(addr, port)
-	var mb buf.MultiBuffer
-	for remain > 0 {
-		length := buf.Size
-		if remain < length {
-			length = remain
-		}
-
-		b := buf.New()
-		mb = append(mb, b)
-		n, err := b.ReadFullFrom(r, int32(length))
-		if err != nil {
-			buf.ReleaseMulti(mb)
-			return nil, newError("failed to read payload").Base(err)
-		}
 
-		remain -= int(n)
+	b := buf.NewWithSize(int32(length))
+	_, err = b.ReadFullFrom(r, int32(length))
+	if err != nil {
+		return nil, newError("failed to read payload").Base(err)
 	}
 
-	return &PacketPayload{Target: dest, Buffer: mb}, nil
+	return &PacketPayload{Target: dest, Buffer: buf.MultiBuffer{b}}, nil
 }
 
 type PacketConnectionReader struct {

proxy/trojan/protocol_test.go

@@ -57,11 +57,12 @@ func TestUDPRequest(t *testing.T) {
 			Password: "password",
 		}),
 	}
+
 	payload := []byte("test string")
-	data := buf.New()
+	data := buf.NewWithSize(int32(len(payload)))
 	common.Must2(data.Write(payload))
 
-	buffer := buf.New()
+	buffer := buf.NewWithSize(int32(len(payload)))
 	defer buffer.Release()
 
 	destination := net.Destination{Network: net.Network_UDP, Address: net.LocalHostIP, Port: 1234}