11package binary
22
33import (
4- "bytes"
54 "encoding/binary"
65 "io"
76 "math"
@@ -10,93 +9,63 @@ import (
109 "unsafe"
1110
1211 ftea "github.com/fumiama/gofastTEA"
12+ "github.com/fumiama/orbyte"
13+ "github.com/fumiama/orbyte/pbuf"
1314
1415 "github.com/LagrangeDev/LagrangeGo/utils"
1516)
1617
17- type Builder struct {
18- buffer bytes.Buffer
18+ type teacfg struct {
1919 key ftea.TEA
2020 usetea bool
21- hasput bool
22- io.Writer
23- io.ReaderFrom
2421}
2522
26- // NewWriterF from https://github.com/Mrs4s/MiraiGo/blob/master/binary/writer.go
27- func NewWriterF (f func (writer * Builder )) []byte {
28- w := NewBuilder ()
29- f (w )
30- b := make ([]byte , w .Len ())
31- copy (b , w .ToBytes ())
32- return b
33- }
34-
35- // ToBytes from https://github.com/Mrs4s/MiraiGo/blob/master/binary/writer.go
36- func ToBytes (i any ) []byte {
37- return NewWriterF (func (w * Builder ) {
38- // TODO: more types
39- switch t := i .(type ) {
40- case int16 :
41- w .WriteU16 (uint16 (t ))
42- case int32 :
43- w .WriteU32 (uint32 (t ))
44- }
45- })
46- }
47-
48- func (b * Builder ) init (key []byte ) * Builder {
49- b .key = ftea .NewTeaCipher (key )
50- b .usetea = len (key ) == 16
51- b .hasput = false
52- return b
53- }
54-
55- func (b * Builder ) Len () int {
56- return b .buffer .Len ()
23+ func (tc * teacfg ) init (key []byte ) {
24+ tc .key = ftea .NewTeaCipher (key )
25+ tc .usetea = len (key ) == 16
5726}
5827
59- // ToReader GC 不安全, 确保在 Builder 被回收前使用
60- func (b * Builder ) ToReader () io.Reader {
61- return & b .buffer
62- }
63-
64- // Buffer GC 不安全, 确保在 Builder 被回收前使用
65- func (b * Builder ) Buffer () * bytes.Buffer {
66- return & b .buffer
28+ func (b * Builder ) p (f func (* pbuf.UserBuffer [teacfg ])) {
29+ (* orbyte.Item [pbuf.UserBuffer [teacfg ]])(b ).P (f )
6730}
6831
6932// ToBytes return data with tea encryption if key is set
7033//
7134// GC 安全, 返回的数据在 Builder 被销毁之后仍能被正确读取,
7235// 但是只能调用一次, 调用后 Builder 即失效
73- func (b * Builder ) ToBytes () []byte {
74- if b .usetea {
75- return b .key .Encrypt (b .buffer .Bytes ())
76- }
77- buf := make ([]byte , b .Len ())
78- copy (buf , b .buffer .Bytes ())
79- return buf
36+ func (b * Builder ) ToBytes () (data []byte ) {
37+ b .p (func (ub * pbuf.UserBuffer [teacfg ]) {
38+ if ub .DAT .usetea {
39+ data = ub .DAT .key .Encrypt (ub .Bytes ())
40+ return
41+ }
42+ data = make ([]byte , ub .Len ())
43+ copy (data , ub .Bytes ())
44+ })
45+ return
8046}
8147
8248// Pack TLV with tea encryption if key is set
8349//
8450// GC 安全, 返回的数据在 Builder 被销毁之后仍能被正确读取,
85- // 但是只能调用一次, 调用后 Builder 即失效
86- func (b * Builder ) Pack (typ uint16 ) []byte {
87- buf := make ([]byte , b .Len ()+ 2 + 2 + 16 )
88-
89- n := 0
90- if b .usetea {
91- n = b .key .EncryptTo (b .buffer .Bytes (), buf [2 + 2 :])
92- } else {
93- n = copy (buf [2 + 2 :], b .buffer .Bytes ())
94- }
51+ // 调用后 Builder 仍有效
52+ func (b * Builder ) Pack (typ uint16 ) (data []byte ) {
53+ b .p (func (buf * pbuf.UserBuffer [teacfg ]) {
54+ data = make ([]byte , buf .Len ()+ 2 + 2 + 16 )
55+
56+ n := 0
57+ if buf .DAT .usetea {
58+ n = buf .DAT .key .EncryptTo (buf .Bytes (), data [2 + 2 :])
59+ } else {
60+ n = copy (data [2 + 2 :], buf .Bytes ())
61+ }
9562
96- binary .BigEndian .PutUint16 (buf [0 :2 ], typ ) // type
97- binary .BigEndian .PutUint16 (buf [2 :2 + 2 ], uint16 (n )) // length
63+ binary .BigEndian .PutUint16 (data [0 :2 ], typ ) // type
64+ binary .BigEndian .PutUint16 (data [2 :2 + 2 ], uint16 (n )) // length
9865
99- return buf [:n + 2 + 2 ]
66+ data = data [:n + 2 + 2 ]
67+ })
68+ return
10069}
10170
10271func (b * Builder ) WriteBool (v bool ) * Builder {
@@ -140,7 +109,10 @@ func (b *Builder) WritePacketString(s, prefix string, withPrefix bool) *Builder
140109
141110// Write for impl. io.Writer
142111func (b * Builder ) Write (p []byte ) (n int , err error ) {
143- return b .buffer .Write (p )
112+ b .p (func (ub * pbuf.UserBuffer [teacfg ]) {
113+ n , err = ub .Write (p )
114+ })
115+ return
144116}
145117
146118func (b * Builder ) EncryptAndWrite (key []byte , data []byte ) * Builder {
@@ -150,7 +122,10 @@ func (b *Builder) EncryptAndWrite(key []byte, data []byte) *Builder {
150122
151123// ReadFrom for impl. io.ReaderFrom
152124func (b * Builder ) ReadFrom (r io.Reader ) (n int64 , err error ) {
153- return io .Copy (& b .buffer , r )
125+ b .p (func (ub * pbuf.UserBuffer [teacfg ]) {
126+ n , err = io .Copy (ub , r )
127+ })
128+ return
154129}
155130
156131func (b * Builder ) WriteLenBytes (v []byte ) * Builder {
@@ -169,21 +144,25 @@ func (b *Builder) WriteLenString(v string) *Builder {
169144}
170145
171146func (b * Builder ) WriteStruct (data ... any ) * Builder {
172- for _ , data := range data {
173- _ = binary .Write (& b .buffer , binary .BigEndian , data )
174- }
147+ b .p (func (ub * pbuf.UserBuffer [teacfg ]) {
148+ for _ , data := range data {
149+ _ = binary .Write (ub , binary .BigEndian , data )
150+ }
151+ })
175152 return b
176153}
177154
178155func (b * Builder ) WriteU8 (v uint8 ) * Builder {
179- b .buffer .WriteByte (v )
156+ b .p (func (ub * pbuf.UserBuffer [teacfg ]) {
157+ ub .WriteByte (v )
158+ })
180159 return b
181160}
182161
183162func writeint [T ~ uint16 | ~ uint32 | ~ uint64 ](b * Builder , v T ) * Builder {
184163 buf := make ([]byte , 8 )
185164 binary .BigEndian .PutUint64 (buf , uint64 (v ))
186- b .buffer . Write (buf [8 - unsafe .Sizeof (v ):])
165+ b .Write (buf [8 - unsafe .Sizeof (v ):])
187166 return b
188167}
189168
0 commit comments