Chore: use protobytes replace most of bytes.Buffer

transport/vmess/conn.go

@@ -1,7 +1,6 @@
 package vmess
 
 import (
-	"bytes"
 	"crypto/aes"
 	"crypto/cipher"
 	"crypto/hmac"
@@ -16,6 +15,7 @@ import (
 	"net"
 	"time"
 
+	"github.com/Dreamacro/protobytes"
 	"golang.org/x/crypto/chacha20poly1305"
 )
 
@@ -57,48 +57,46 @@ func (vc *Conn) Read(b []byte) (int, error) {
 func (vc *Conn) sendRequest() error {
 	timestamp := time.Now()
 
-	mbuf := &bytes.Buffer{}
+	mbuf := protobytes.BytesWriter{}
 
 	if !vc.isAead {
 		h := hmac.New(md5.New, vc.id.UUID.Bytes())
 		binary.Write(h, binary.BigEndian, uint64(timestamp.Unix()))
-		mbuf.Write(h.Sum(nil))
+		mbuf.PutSlice(h.Sum(nil))
 	}
 
-	buf := &bytes.Buffer{}
+	buf := protobytes.BytesWriter{}
 
 	// Ver IV Key V Opt
-	buf.WriteByte(Version)
-	buf.Write(vc.reqBodyIV[:])
-	buf.Write(vc.reqBodyKey[:])
-	buf.WriteByte(vc.respV)
-	buf.WriteByte(vc.option)
+	buf.PutUint8(Version)
+	buf.PutSlice(vc.reqBodyIV[:])
+	buf.PutSlice(vc.reqBodyKey[:])
+	buf.PutUint8(vc.respV)
+	buf.PutUint8(vc.option)
 
 	p := mathRand.Intn(16)
 	// P Sec Reserve Cmd
-	buf.WriteByte(byte(p<<4) | vc.security)
-	buf.WriteByte(0)
+	buf.PutUint8(byte(p<<4) | vc.security)
+	buf.PutUint8(0)
 	if vc.dst.UDP {
-		buf.WriteByte(CommandUDP)
+		buf.PutUint8(CommandUDP)
 	} else {
-		buf.WriteByte(CommandTCP)
+		buf.PutUint8(CommandTCP)
 	}
 
 	// Port AddrType Addr
-	binary.Write(buf, binary.BigEndian, uint16(vc.dst.Port))
-	buf.WriteByte(vc.dst.AddrType)
-	buf.Write(vc.dst.Addr)
+	buf.PutUint16be(uint16(vc.dst.Port))
+	buf.PutUint8(vc.dst.AddrType)
+	buf.PutSlice(vc.dst.Addr)
 
 	// padding
 	if p > 0 {
-		padding := make([]byte, p)
-		rand.Read(padding)
-		buf.Write(padding)
+		buf.ReadFull(rand.Reader, p)
 	}
 
 	fnv1a := fnv.New32a()
 	fnv1a.Write(buf.Bytes())
-	buf.Write(fnv1a.Sum(nil))
+	buf.PutSlice(fnv1a.Sum(nil))
 
 	if !vc.isAead {
 		block, err := aes.NewCipher(vc.id.CmdKey)
@@ -108,7 +106,7 @@ func (vc *Conn) sendRequest() error {
 
 		stream := cipher.NewCFBEncrypter(block, hashTimestamp(timestamp))
 		stream.XORKeyStream(buf.Bytes(), buf.Bytes())
-		mbuf.Write(buf.Bytes())
+		mbuf.PutSlice(buf.Bytes())
 		_, err = vc.Conn.Write(mbuf.Bytes())
 		return err
 	}

transport/vmess/header.go

@@ -1,7 +1,6 @@
 package vmess
 
 import (
-	"bytes"
 	"crypto/aes"
 	"crypto/cipher"
 	"crypto/hmac"
@@ -11,6 +10,8 @@ import (
 	"hash"
 	"hash/crc32"
 	"time"
+
+	"github.com/Dreamacro/protobytes"
 )
 
 const (
@@ -49,14 +50,11 @@ func (h *hMacCreator) Create() hash.Hash {
 }
 
 func createAuthID(cmdKey []byte, time int64) [16]byte {
-	buf := &bytes.Buffer{}
-	binary.Write(buf, binary.BigEndian, time)
-
-	random := make([]byte, 4)
-	rand.Read(random)
-	buf.Write(random)
+	buf := protobytes.BytesWriter{}
+	buf.PutUint64be(uint64(time))
+	buf.ReadFull(rand.Reader, 4)
 	zero := crc32.ChecksumIEEE(buf.Bytes())
-	binary.Write(buf, binary.BigEndian, zero)
+	buf.PutUint32be(zero)
 
 	aesBlock, _ := aes.NewCipher(kdf(cmdKey[:], kdfSaltConstAuthIDEncryptionKey)[:16])
 	var result [16]byte
@@ -92,12 +90,12 @@ func sealVMessAEADHeader(key [16]byte, data []byte, t time.Time) []byte {
 		payloadHeaderAEADEncrypted = payloadHeaderAEAD.Seal(nil, payloadHeaderAEADNonce, data, generatedAuthID[:])
 	}
 
-	outputBuffer := &bytes.Buffer{}
+	outputBuffer := protobytes.BytesWriter{}
 
-	outputBuffer.Write(generatedAuthID[:])
-	outputBuffer.Write(payloadHeaderLengthAEADEncrypted)
-	outputBuffer.Write(connectionNonce)
-	outputBuffer.Write(payloadHeaderAEADEncrypted)
+	outputBuffer.PutSlice(generatedAuthID[:])
+	outputBuffer.PutSlice(payloadHeaderLengthAEADEncrypted)
+	outputBuffer.PutSlice(connectionNonce)
+	outputBuffer.PutSlice(payloadHeaderAEADEncrypted)
 
 	return outputBuffer.Bytes()
 }