From d233c90d5083b6aabb6053b603932179a0c69d29 Mon Sep 17 00:00:00 2001 From: wdvxdr Date: Tue, 23 Mar 2021 22:37:36 +0800 Subject: [PATCH] =?UTF-8?q?change=20writer=20without=20buffer=20=E4=BD=BF?= =?UTF-8?q?=E7=94=A8=E4=BA=86sync.Pool,=E4=BB=8E=E6=B1=A0=E4=B8=AD?= =?UTF-8?q?=E5=8F=96=E5=87=BA=E6=9D=A5=E7=9A=84=E8=87=AA=E5=B8=A6=E4=B8=80?= =?UTF-8?q?=E5=AE=9A=E5=AE=B9=E9=87=8F=EF=BC=8C=E6=89=80=E4=BB=A5=E6=97=A0?= =?UTF-8?q?=E9=9C=80=E4=BD=BF=E7=94=A8Buffer?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- binary/writer.go | 39 ++++++--------- binary/writer_test.go | 93 ------------------------------------ protocol/packets/builders.go | 48 +++++++++---------- protocol/packets/global.go | 87 ++++++++++++++++----------------- 4 files changed, 81 insertions(+), 186 deletions(-) diff --git a/binary/writer.go b/binary/writer.go index 86742eba..43dca05d 100644 --- a/binary/writer.go +++ b/binary/writer.go @@ -1,53 +1,44 @@ package binary import ( - "bytes" "encoding/binary" "encoding/hex" "sync" ) -type Writer bytes.Buffer +type Writer struct { + b []byte +} var bufferPool = sync.Pool{ New: func() interface{} { - return new(bytes.Buffer) + return new(Writer) }, } -func NewBuffer() *bytes.Buffer { - return bufferPool.Get().(*bytes.Buffer) +func NewWriter() *Writer { + return bufferPool.Get().(*Writer) } -func PutBuffer(buf *bytes.Buffer) { +func PutBuffer(w *Writer) { // See https://golang.org/issue/23199 const maxSize = 1 << 16 - if buf.Cap() < maxSize { // 对于大Buffer直接丢弃 - buf.Reset() - bufferPool.Put(buf) + if cap(w.b) < maxSize { // 对于大Buffer直接丢弃 + w.b = w.b[:0] + bufferPool.Put(w) } } -func NewWriter() *Writer { - return new(Writer) -} - -/* -func PutWriter(w *Writer) { - PutBuffer((*bytes.Buffer)(w)) -} -*/ - func NewWriterF(f func(writer *Writer)) []byte { - w := NewBuffer() - f((*Writer)(w)) + w := NewWriter() + f(w) b := append([]byte(nil), w.Bytes()...) PutBuffer(w) return b } func (w *Writer) Write(b []byte) { - (*bytes.Buffer)(w).Write(b) + w.b = append(w.b, b...) } func (w *Writer) WriteHex(h string) { @@ -56,7 +47,7 @@ func (w *Writer) WriteHex(h string) { } func (w *Writer) WriteByte(b byte) { - (*bytes.Buffer)(w).WriteByte(b) + w.b = append(w.b, b) } func (w *Writer) WriteUInt16(v uint16) { @@ -138,5 +129,5 @@ func (w *Writer) WriteTlvLimitedSize(data []byte, limit int) { } func (w *Writer) Bytes() []byte { - return (*bytes.Buffer)(w).Bytes() + return w.b } diff --git a/binary/writer_test.go b/binary/writer_test.go index 5c42d7bf..578095c2 100644 --- a/binary/writer_test.go +++ b/binary/writer_test.go @@ -1,71 +1,10 @@ package binary import ( - "bytes" - "encoding/hex" - "fmt" "math/rand" - "sync" "testing" ) -func NewWriterFOld(f func(writer *Writer)) []byte { - w := new(bytes.Buffer) - f((*Writer)(w)) - return w.Bytes() -} - -func TestNewWriterF(t *testing.T) { - wg := sync.WaitGroup{} - for i := 0; i < 10000; i++ { - wg.Add(1) - go func() { - defer wg.Done() - test := make([]byte, 1024) - rand.Read(test) - b1 := NewWriterFOld(func(writer *Writer) { - writer.Write(test) - writer.Write(NewWriterFOld(func(writer *Writer) { - writer.Write(test) - writer.Write(NewWriterFOld(func(writer *Writer) { - writer.Write(test) - })) - })) - }) - - b2 := NewWriterF(func(writer *Writer) { - writer.Write(test) - writer.Write(NewWriterF(func(writer *Writer) { - writer.Write(test) - writer.Write(NewWriterF(func(writer *Writer) { - writer.Write(test) - })) - })) - }) - - if !bytes.Equal(b1, b2) { - fmt.Println(hex.EncodeToString(b1)) - fmt.Println(hex.EncodeToString(b2)) - t.Error("Not equal!!!") - } - }() - } - wg.Wait() -} - -func BenchmarkNewWriterFOld128(b *testing.B) { - test := make([]byte, 128) - rand.Read(test) - b.StartTimer() - b.RunParallel(func(pb *testing.PB) { - for pb.Next() { - NewWriterFOld(func(w *Writer) { - w.Write(test) - }) - } - }) -} - func BenchmarkNewWriterF128(b *testing.B) { test := make([]byte, 128) rand.Read(test) @@ -79,21 +18,6 @@ func BenchmarkNewWriterF128(b *testing.B) { }) } -func BenchmarkNewWriterFOld128_3(b *testing.B) { - test := make([]byte, 128) - rand.Read(test) - b.StartTimer() - b.RunParallel(func(pb *testing.PB) { - for pb.Next() { - NewWriterFOld(func(w *Writer) { - w.Write(test) - w.Write(test) - w.Write(test) - }) - } - }) -} - func BenchmarkNewWriterF128_3(b *testing.B) { test := make([]byte, 128) rand.Read(test) @@ -109,23 +33,6 @@ func BenchmarkNewWriterF128_3(b *testing.B) { }) } -func BenchmarkNewWriterFOld128_5(b *testing.B) { - test := make([]byte, 128) - rand.Read(test) - b.StartTimer() - b.RunParallel(func(pb *testing.PB) { - for pb.Next() { - NewWriterFOld(func(w *Writer) { - w.Write(test) - w.Write(test) - w.Write(test) - w.Write(test) - w.Write(test) - }) - } - }) -} - func BenchmarkNewWriterF128_5(b *testing.B) { test := make([]byte, 128) rand.Read(test) diff --git a/protocol/packets/builders.go b/protocol/packets/builders.go index 84a87af1..1b4789e9 100644 --- a/protocol/packets/builders.go +++ b/protocol/packets/builders.go @@ -7,35 +7,35 @@ import ( ) func BuildLoginPacket(uin int64, bodyType byte, key, body, extraData []byte) []byte { - w := binary.NewWriter() - w.WriteIntLvPacket(4, func(w *binary.Writer) { - w.WriteUInt32(0x00_00_00_0A) - w.WriteByte(bodyType) + return binary.NewWriterF(func(w *binary.Writer) { w.WriteIntLvPacket(4, func(w *binary.Writer) { - w.Write(extraData) + w.WriteUInt32(0x00_00_00_0A) + w.WriteByte(bodyType) + w.WriteIntLvPacket(4, func(w *binary.Writer) { + w.Write(extraData) + }) + w.WriteByte(0x00) + w.WriteString(strconv.FormatInt(uin, 10)) + if len(key) == 0 { + w.Write(body) + } else { + w.EncryptAndWrite(key, body) + } }) - w.WriteByte(0x00) - w.WriteString(strconv.FormatInt(uin, 10)) - if len(key) == 0 { - w.Write(body) - } else { - w.EncryptAndWrite(key, body) - } }) - return w.Bytes() } func BuildUniPacket(uin int64, seq uint16, commandName string, encryptType byte, sessionID, extraData, key, body []byte) []byte { - w := binary.NewWriter() - w.WriteIntLvPacket(4, func(w *binary.Writer) { - w.WriteUInt32(0x0B) - w.WriteByte(encryptType) - w.WriteUInt32(uint32(seq)) - w.WriteByte(0) - w.WriteString(strconv.FormatInt(uin, 10)) - w.EncryptAndWrite(key, binary.NewWriterF(func(w *binary.Writer) { - w.WriteUniPacket(commandName, sessionID, extraData, body) - })) + return binary.NewWriterF(func(w *binary.Writer) { + w.WriteIntLvPacket(4, func(w *binary.Writer) { + w.WriteUInt32(0x0B) + w.WriteByte(encryptType) + w.WriteUInt32(uint32(seq)) + w.WriteByte(0) + w.WriteString(strconv.FormatInt(uin, 10)) + w.EncryptAndWrite(key, binary.NewWriterF(func(w *binary.Writer) { + w.WriteUniPacket(commandName, sessionID, extraData, body) + })) + }) }) - return w.Bytes() } diff --git a/protocol/packets/global.go b/protocol/packets/global.go index e4b05f10..2da69db9 100644 --- a/protocol/packets/global.go +++ b/protocol/packets/global.go @@ -33,26 +33,23 @@ type IEncryptMethod interface { } func BuildOicqRequestPacket(uin int64, commandId uint16, encrypt IEncryptMethod, key []byte, bodyFunc func(writer *binary.Writer)) []byte { - b := binary.NewWriter() - bodyFunc(b) - - body := encrypt.DoEncrypt(b.Bytes(), key) - p := binary.NewWriter() - p.WriteByte(0x02) - p.WriteUInt16(27 + 2 + uint16(len(body))) - p.WriteUInt16(8001) - p.WriteUInt16(commandId) - p.WriteUInt16(1) - p.WriteUInt32(uint32(uin)) - p.WriteByte(3) - p.WriteByte(encrypt.Id()) - p.WriteByte(0) - p.WriteUInt32(2) - p.WriteUInt32(0) - p.WriteUInt32(0) - p.Write(body) - p.WriteByte(0x03) - return p.Bytes() + body := encrypt.DoEncrypt(binary.NewWriterF(bodyFunc), key) + return binary.NewWriterF(func(p *binary.Writer) { + p.WriteByte(0x02) + p.WriteUInt16(27 + 2 + uint16(len(body))) + p.WriteUInt16(8001) + p.WriteUInt16(commandId) + p.WriteUInt16(1) + p.WriteUInt32(uint32(uin)) + p.WriteByte(3) + p.WriteByte(encrypt.Id()) + p.WriteByte(0) + p.WriteUInt32(2) + p.WriteUInt32(0) + p.WriteUInt32(0) + p.Write(body) + p.WriteByte(0x03) + }) } func BuildCode2DRequestPacket(seq uint32, j uint64, cmd uint16, bodyFunc func(writer *binary.Writer)) []byte { @@ -73,35 +70,35 @@ func BuildCode2DRequestPacket(seq uint32, j uint64, cmd uint16, bodyFunc func(wr } func BuildSsoPacket(seq uint16, appID uint32, commandName, imei string, extData, outPacketSessionId, body, ksid []byte) []byte { - p := binary.NewWriter() - p.WriteIntLvPacket(4, func(writer *binary.Writer) { - writer.WriteUInt32(uint32(seq)) - writer.WriteUInt32(appID) - writer.WriteUInt32(appID) - writer.Write([]byte{0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00}) - if len(extData) == 0 || len(extData) == 4 { + return binary.NewWriterF(func(p *binary.Writer) { + p.WriteIntLvPacket(4, func(writer *binary.Writer) { + writer.WriteUInt32(uint32(seq)) + writer.WriteUInt32(appID) + writer.WriteUInt32(appID) + writer.Write([]byte{0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00}) + if len(extData) == 0 || len(extData) == 4 { + writer.WriteUInt32(0x04) + } else { + writer.WriteUInt32(uint32(len(extData) + 4)) + writer.Write(extData) + } + writer.WriteString(commandName) + writer.WriteIntLvPacket(4, func(w *binary.Writer) { + w.Write(outPacketSessionId) + }) + writer.WriteString(imei) + writer.WriteUInt32(0x04) + { + writer.WriteUInt16(uint16(len(ksid)) + 2) + writer.Write(ksid) + } writer.WriteUInt32(0x04) - } else { - writer.WriteUInt32(uint32(len(extData) + 4)) - writer.Write(extData) - } - writer.WriteString(commandName) - writer.WriteIntLvPacket(4, func(w *binary.Writer) { - w.Write(outPacketSessionId) }) - writer.WriteString(imei) - writer.WriteUInt32(0x04) - { - writer.WriteUInt16(uint16(len(ksid)) + 2) - writer.Write(ksid) - } - writer.WriteUInt32(0x04) - }) - p.WriteIntLvPacket(4, func(writer *binary.Writer) { - writer.Write(body) + p.WriteIntLvPacket(4, func(writer *binary.Writer) { + writer.Write(body) + }) }) - return p.Bytes() } func ParseIncomingPacket(payload, d2key []byte) (*IncomingPacket, error) {