perf(writer): drop some lambda expressions (#224)

* perf(writer): drop some lambda expressions

* perf(writer): drop more lambda expressions

* fix: resolve conflicts

* perf(writer): drop more lambda expressions

* fix: modify function name

* perf(writer): drop more lambda expressions

* fix: some error

.gitignore

@@ -1,2 +1,3 @@
 .idea
 vendor/
+.DS_Store

binary/pool.go

@@ -26,7 +26,7 @@ func SelectWriter() *Writer {
 func PutWriter(w *Writer) {
 	// See https://golang.org/issue/23199
 	const maxSize = 1 << 16
-	if w.Cap() < maxSize { // 对于大Buffer直接丢弃
+	if (*bytes.Buffer)(w).Cap() < maxSize { // 对于大Buffer直接丢弃
 		w.Reset()
 		bufferPool.Put(w)
 	}

binary/writer.go

@@ -4,8 +4,6 @@ import (
 	"bytes"
 	"encoding/binary"
 	"encoding/hex"
-
-	"github.com/Mrs4s/MiraiGo/utils"
 )
 
 // Writer 写入
@@ -15,15 +13,48 @@ func NewWriterF(f func(writer *Writer)) []byte {
 	w := SelectWriter()
 	f(w)
 	b := append([]byte(nil), w.Bytes()...)
-	PutWriter(w)
+	w.put()
 	return b
 }
 
 // OpenWriterF must call func cl to close
 func OpenWriterF(f func(*Writer)) (b []byte, cl func()) {
 	w := SelectWriter()
 	f(w)
-	return w.Bytes(), func() { PutWriter(w) }
+	return w.Bytes(), w.put
+}
+
+func (w *Writer) AllocUInt16Head() (pos int) {
+	pos = (*bytes.Buffer)(w).Len()
+	(*bytes.Buffer)(w).Write([]byte{0, 0})
+	return
+}
+
+/*
+func (w *Writer) WriteUInt16HeadAt(pos int) {
+	newdata := (*bytes.Buffer)(w).Bytes()[pos:]
+	binary.BigEndian.PutUint16(newdata, uint16(len(newdata)))
+}
+*/
+
+func (w *Writer) WriteUInt16HeadUsingTotalBufferLenAt(pos int) {
+	binary.BigEndian.PutUint16((*bytes.Buffer)(w).Bytes()[pos:], uint16((*bytes.Buffer)(w).Len()))
+}
+
+func (w *Writer) WriteUInt16HeadExcludeSelfAt(pos int) {
+	newdata := (*bytes.Buffer)(w).Bytes()[pos:]
+	binary.BigEndian.PutUint16(newdata, uint16(len(newdata)-2))
+}
+
+func (w *Writer) AllocUInt32Head() (pos int) {
+	pos = (*bytes.Buffer)(w).Len()
+	(*bytes.Buffer)(w).Write([]byte{0, 0, 0, 0})
+	return
+}
+
+func (w *Writer) WriteUInt32HeadAt(pos int) {
+	newdata := (*bytes.Buffer)(w).Bytes()[pos:]
+	binary.BigEndian.PutUint32(newdata, uint32(len(newdata)))
 }
 
 func (w *Writer) Write(b []byte) {
@@ -35,8 +66,8 @@ func (w *Writer) WriteHex(h string) {
 	w.Write(b)
 }
 
-func (w *Writer) WriteByte(b byte) {
-	(*bytes.Buffer)(w).WriteByte(b)
+func (w *Writer) WriteByte(b byte) error {
+	return (*bytes.Buffer)(w).WriteByte(b)
 }
 
 func (w *Writer) WriteUInt16(v uint16) {
@@ -58,13 +89,13 @@ func (w *Writer) WriteUInt64(v uint64) {
 }
 
 func (w *Writer) WriteString(v string) {
-	payload := utils.S2B(v)
-	w.WriteUInt32(uint32(len(payload) + 4))
-	w.Write(payload)
+	w.WriteUInt32(uint32(len(v) + 4))
+	(*bytes.Buffer)(w).WriteString(v)
 }
 
 func (w *Writer) WriteStringShort(v string) {
-	w.WriteBytesShort(utils.S2B(v))
+	w.WriteUInt16(uint16(len(v)))
+	(*bytes.Buffer)(w).WriteString(v)
 }
 
 func (w *Writer) WriteBool(b bool) {
@@ -86,27 +117,6 @@ func (w *Writer) WriteIntLvPacket(offset int, f func(*Writer)) {
 	cl()
 }
 
-func (w *Writer) WriteUniPacket(commandName string, sessionId, extraData, body []byte) {
-	w1 := SelectWriter()
-	{ // WriteIntLvPacket
-		w1.WriteString(commandName)
-		w1.WriteUInt32(8)
-		w1.Write(sessionId)
-		if len(extraData) == 0 {
-			w1.WriteUInt32(0x04)
-		} else {
-			w1.WriteUInt32(uint32(len(extraData) + 4))
-			w1.Write(extraData)
-		}
-	}
-	data := w1.Bytes()
-	w.WriteUInt32(uint32(len(data) + 4))
-	w.Write(data)
-	PutWriter(w1)
-	w.WriteUInt32(uint32(len(body) + 4)) // WriteIntLvPacket
-	w.Write(body)
-}
-
 func (w *Writer) WriteBytesShort(data []byte) {
 	w.WriteUInt16(uint16(len(data)))
 	w.Write(data)
@@ -124,14 +134,14 @@ func (w *Writer) Bytes() []byte {
 	return (*bytes.Buffer)(w).Bytes()
 }
 
-func (w *Writer) Cap() int {
-	return (*bytes.Buffer)(w).Cap()
-}
-
 func (w *Writer) Reset() {
 	(*bytes.Buffer)(w).Reset()
 }
 
 func (w *Writer) Grow(n int) {
 	(*bytes.Buffer)(w).Grow(n)
 }
+
+func (w *Writer) put() {
+	PutWriter(w)
+}

client/builders.go

@@ -1097,12 +1097,7 @@ func (c *QQClient) buildGroupAdminSetPacket(groupCode, member int64, flag bool)
 	b, cl := binary.OpenWriterF(func(w *binary.Writer) {
 		w.WriteUInt32(uint32(groupCode))
 		w.WriteUInt32(uint32(member))
-		w.WriteByte(func() byte {
-			if flag {
-				return 1
-			}
-			return 0
-		}())
+		w.WriteBool(flag)
 	})
 	payload := c.packOIDBPackage(1372, 1, b)
 	cl()

client/internal/network/transport.go

@@ -19,9 +19,7 @@ type Transport struct {
 	// conn *TCPListener
 }
 
-func (t *Transport) packBody(req *Request) []byte {
-	w := binary.SelectWriter()
-	defer binary.PutWriter(w)
+func (t *Transport) packBody(req *Request, w *binary.Writer) {
 	w.WriteIntLvPacket(4, func(writer *binary.Writer) {
 		if req.Type == RequestTypeLogin {
 			writer.WriteUInt32(uint32(req.SequenceID))
@@ -58,7 +56,6 @@ func (t *Transport) packBody(req *Request) []byte {
 	})
 	// w.WriteUInt32(uint32(len(req.Body) + 4))
 	// w.Write(req.Body)
-	return append([]byte(nil), w.Bytes()...)
 }
 
 // PackPacket packs a packet.
@@ -67,16 +64,10 @@ func (t *Transport) PackPacket(req *Request) []byte {
 	if len(t.Sig.D2) == 0 {
 		req.EncryptType = EncryptTypeEmptyKey
 	}
-	body := t.packBody(req)
-	// encrypt body
-	switch req.EncryptType {
-	case EncryptTypeD2Key:
-		body = binary.NewTeaCipher(t.Sig.D2Key).Encrypt(body)
-	case EncryptTypeEmptyKey:
-		body = binary.NewTeaCipher(emptyKey).Encrypt(body)
-	}
 
-	head := binary.NewWriterF(func(w *binary.Writer) {
+	return binary.NewWriterF(func(w *binary.Writer) {
+		pos := w.AllocUInt32Head()
+		// vvv w.Write(head) vvv
 		w.WriteUInt32(uint32(req.Type))
 		w.WriteByte(byte(req.EncryptType))
 		switch req.Type {
@@ -93,14 +84,24 @@ func (t *Transport) PackPacket(req *Request) []byte {
 		}
 		w.WriteByte(0x00)
 		w.WriteString(strconv.FormatInt(req.Uin, 10))
+		// ^^^ w.Write(head) ^^^
+		w.Write(binary.NewWriterF(func(w *binary.Writer) {
+			// encrypt body
+			switch req.EncryptType {
+			case EncryptTypeD2Key:
+				wt, cl := binary.OpenWriterF(func(w *binary.Writer) { t.packBody(req, w) })
+				w.EncryptAndWrite(t.Sig.D2Key, wt)
+				cl()
+			case EncryptTypeEmptyKey:
+				wt, cl := binary.OpenWriterF(func(w *binary.Writer) { t.packBody(req, w) })
+				w.EncryptAndWrite(emptyKey, wt)
+				cl()
+			default:
+				t.packBody(req, w)
+			}
+		}))
+		w.WriteUInt32HeadAt(pos)
 	})
-
-	w := binary.SelectWriter()
-	defer binary.PutWriter(w)
-	w.WriteUInt32(uint32(len(head)+len(body)) + 4)
-	w.Write(head)
-	w.Write(body)
-	return append([]byte(nil), w.Bytes()...) // copy
 }
 
 func (t *Transport) parse(head []byte) *Request {

client/tlv_decoders.go

@@ -111,7 +111,13 @@ func (c *QQClient) decodeT119(data, ek []byte) {
 	s.PsKeyMap = psKeyMap
 	s.Pt4TokenMap = pt4TokenMap
 	if len(c.PasswordMd5[:]) > 0 {
-		key := md5.Sum(append(append(c.PasswordMd5[:], []byte{0x00, 0x00, 0x00, 0x00}...), binary.NewWriterF(func(w *binary.Writer) { w.WriteUInt32(uint32(c.Uin)) })...))
+		data, cl := binary.OpenWriterF(func(w *binary.Writer) {
+			w.Write(c.PasswordMd5[:])
+			w.WriteUInt32(0) // []byte{0x00, 0x00, 0x00, 0x00}...
+			w.WriteUInt32(uint32(c.Uin))
+		})
+		key := md5.Sum(data)
+		cl()
 		decrypted := binary.NewTeaCipher(key[:]).Decrypt(c.sig.EncryptedA1)
 		if len(decrypted) > 51+16 {
 			dr := binary.NewReader(decrypted)

internal/packets/global.go

@@ -31,18 +31,18 @@ type IncomingPacket struct {
 
 func BuildCode2DRequestPacket(seq uint32, j uint64, cmd uint16, bodyFunc func(writer *binary.Writer)) []byte {
 	return binary.NewWriterF(func(w *binary.Writer) {
-		body := binary.NewWriterF(bodyFunc)
 		w.WriteByte(2)
-		w.WriteUInt16(uint16(43 + len(body) + 1))
+		pos := w.AllocUInt16Head()
 		w.WriteUInt16(cmd)
 		w.Write(make([]byte, 21))
 		w.WriteByte(3)
 		w.WriteUInt16(0)
 		w.WriteUInt16(50) // version
 		w.WriteUInt32(seq)
 		w.WriteUInt64(j)
-		w.Write(body)
+		bodyFunc(w)
 		w.WriteByte(3)
+		w.WriteUInt16HeadUsingTotalBufferLenAt(pos)
 	})
 }
 

internal/tlv/t1.go

@@ -13,13 +13,13 @@ func T1(uin uint32, ip []byte) []byte {
 	}
 	return binary.NewWriterF(func(w *binary.Writer) {
 		w.WriteUInt16(0x01)
-		w.WriteBytesShort(binary.NewWriterF(func(w *binary.Writer) {
-			w.WriteUInt16(1)
-			w.WriteUInt32(rand.Uint32())
-			w.WriteUInt32(uin)
-			w.WriteUInt32(uint32(time.Now().UnixNano() / 1e6))
-			w.Write(ip)
-			w.WriteUInt16(0)
-		}))
+		pos := w.AllocUInt16Head()
+		w.WriteUInt16(1)
+		w.WriteUInt32(rand.Uint32())
+		w.WriteUInt32(uin)
+		w.WriteUInt32(uint32(time.Now().UnixNano() / 1e6))
+		w.Write(ip)
+		w.WriteUInt16(0)
+		w.WriteUInt16HeadExcludeSelfAt(pos)
 	})
 }

internal/tlv/t100.go

@@ -7,13 +7,13 @@ import (
 func T100(ssoVersion, protocol, mainSigMap uint32) []byte {
 	return binary.NewWriterF(func(w *binary.Writer) {
 		w.WriteUInt16(0x100)
-		w.WriteBytesShort(binary.NewWriterF(func(w *binary.Writer) {
-			w.WriteUInt16(1)
-			w.WriteUInt32(ssoVersion)
-			w.WriteUInt32(16)
-			w.WriteUInt32(protocol)
-			w.WriteUInt32(0)          // App client version
-			w.WriteUInt32(mainSigMap) // 34869472
-		}))
+		pos := w.AllocUInt16Head()
+		w.WriteUInt16(1)
+		w.WriteUInt32(ssoVersion)
+		w.WriteUInt32(16)
+		w.WriteUInt32(protocol)
+		w.WriteUInt32(0)          // App client version
+		w.WriteUInt32(mainSigMap) // 34869472
+		w.WriteUInt16HeadExcludeSelfAt(pos)
 	})
 }

internal/tlv/t106.go

@@ -2,7 +2,6 @@ package tlv
 
 import (
 	"crypto/md5"
-	binary2 "encoding/binary"
 	"math/rand"
 	"strconv"
 	"time"
@@ -13,7 +12,19 @@ import (
 func T106(uin, salt, appId, ssoVer uint32, passwordMd5 [16]byte, guidAvailable bool, guid, tgtgtKey []byte, wtf uint32) []byte {
 	return binary.NewWriterF(func(w *binary.Writer) {
 		w.WriteUInt16(0x106)
-		body := binary.NewWriterF(func(w *binary.Writer) {
+		pos := w.AllocUInt16Head()
+		keydata, kcl := binary.OpenWriterF(func(w *binary.Writer) {
+			w.Write(passwordMd5[:])
+			w.WriteUInt32(0) // []byte{0x00, 0x00, 0x00, 0x00}...
+			if salt != 0 {
+				w.WriteUInt32(salt)
+			} else {
+				w.WriteUInt32(uin)
+			}
+		})
+		key := md5.Sum(keydata)
+		kcl()
+		body, cl := binary.OpenWriterF(func(w *binary.Writer) {
 			w.WriteUInt16(4)
 			w.WriteUInt32(rand.Uint32())
 			w.WriteUInt32(ssoVer)
@@ -25,33 +36,27 @@ func T106(uin, salt, appId, ssoVer uint32, passwordMd5 [16]byte, guidAvailable b
 				w.WriteUInt64(uint64(uin))
 			}
 			w.WriteUInt32(uint32(time.Now().UnixNano() / 1e6))
-			w.Write([]byte{0x00, 0x00, 0x00, 0x00}) // fake ip
+			w.WriteUInt32(0) // fake ip w.Write([]byte{0x00, 0x00, 0x00, 0x00})
 			w.WriteByte(0x01)
 			w.Write(passwordMd5[:])
 			w.Write(tgtgtKey)
 			w.WriteUInt32(wtf)
 			w.WriteBool(guidAvailable)
 			if len(guid) == 0 {
-				for i := 0; i < 4; i++ {
-					w.WriteUInt32(rand.Uint32())
-				}
+				w.WriteUInt32(rand.Uint32())
+				w.WriteUInt32(rand.Uint32())
+				w.WriteUInt32(rand.Uint32())
+				w.WriteUInt32(rand.Uint32())
 			} else {
 				w.Write(guid)
 			}
 			w.WriteUInt32(appId)
 			w.WriteUInt32(1) // password login
-			w.WriteBytesShort([]byte(strconv.FormatInt(int64(uin), 10)))
+			w.WriteStringShort(strconv.FormatInt(int64(uin), 10))
 			w.WriteUInt16(0)
 		})
-		w.WriteBytesShort(binary.NewWriterF(func(w *binary.Writer) {
-			b := make([]byte, 4)
-			if salt != 0 {
-				binary2.BigEndian.PutUint32(b, salt)
-			} else {
-				binary2.BigEndian.PutUint32(b, uin)
-			}
-			key := md5.Sum(append(append(passwordMd5[:], []byte{0x00, 0x00, 0x00, 0x00}...), b...))
-			w.EncryptAndWrite(key[:], body)
-		}))
+		w.EncryptAndWrite(key[:], body)
+		w.WriteUInt16HeadExcludeSelfAt(pos)
+		cl()
 	})
 }