/
codec_encoder.go
205 lines (172 loc) · 6.16 KB
/
codec_encoder.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
package codec
import (
"bytes"
"errors"
"fmt"
"git.golaxy.org/core"
"git.golaxy.org/plugins/gtp"
"git.golaxy.org/plugins/util/binaryutil"
"io"
)
// IEncoder 消息包编码器接口
type IEncoder interface {
io.Reader
io.WriterTo
// Reset 重置缓存
Reset()
// Encode 编码消息包,写入缓存
Encode(flags gtp.Flags, msg gtp.MsgReader) error
// EncodeWriter 编码消息包,写入指定writer
EncodeWriter(writer io.Writer, flags gtp.Flags, msg gtp.MsgReader) error
// EncodeBuff 编码消息包,写入指定buffer
EncodeBuff(buff *bytes.Buffer, flags gtp.Flags, msg gtp.MsgReader) error
// EncodeBytes 编码消息包,返回可回收bytes
EncodeBytes(flags gtp.Flags, msg gtp.MsgReader) (binaryutil.RecycleBytes, error)
}
// Encoder 消息包编码器
type Encoder struct {
EncryptionModule IEncryptionModule // 加密模块
MACModule IMACModule // MAC模块
CompressionModule ICompressionModule // 压缩模块
Encryption bool // 开启加密
PatchMAC bool // 开启MAC
CompressedSize int // 启用压缩阀值(字节),<=0表示不开启
buffer bytes.Buffer // buffer
}
// Read implements io.Reader
func (e *Encoder) Read(p []byte) (int, error) {
return e.buffer.Read(p)
}
// WriteTo implements io.WriterTo
func (e *Encoder) WriteTo(w io.Writer) (int64, error) {
if w == nil {
return 0, fmt.Errorf("gtp: %w: w is nil", core.ErrArgs)
}
return e.buffer.WriteTo(w)
}
// Reset 重置缓存
func (e *Encoder) Reset() {
e.buffer.Reset()
}
// Encode 编码消息包,写入缓存
func (e *Encoder) Encode(flags gtp.Flags, msg gtp.MsgReader) error {
return e.EncodeWriter(&e.buffer, flags, msg)
}
// EncodeWriter 编码消息包,写入指定writer
func (e *Encoder) EncodeWriter(writer io.Writer, flags gtp.Flags, msg gtp.MsgReader) error {
if writer == nil {
return fmt.Errorf("gtp: %w: writer is nil", core.ErrArgs)
}
mpBuf, err := e.encode(flags, msg)
if err != nil {
return err
}
defer mpBuf.Release()
_, err = writer.Write(mpBuf.Data())
if err != nil {
return fmt.Errorf("gtp: write msg-packet failed, %w", err)
}
return nil
}
// EncodeBuff 编码消息包,写入指定buffer
func (e *Encoder) EncodeBuff(buff *bytes.Buffer, flags gtp.Flags, msg gtp.MsgReader) error {
if buff == nil {
return fmt.Errorf("gtp: %w: buff is nil", core.ErrArgs)
}
return e.EncodeWriter(buff, flags, msg)
}
// EncodeBytes 编码消息包,返回可回收bytes
func (e *Encoder) EncodeBytes(flags gtp.Flags, msg gtp.MsgReader) (binaryutil.RecycleBytes, error) {
return e.encode(flags, msg)
}
// encode 编码消息包
func (e *Encoder) encode(flags gtp.Flags, msg gtp.MsgReader) (ret binaryutil.RecycleBytes, err error) {
if msg == nil {
return binaryutil.MakeNonRecycleBytes(nil), fmt.Errorf("gtp: %w: msg is nil", core.ErrArgs)
}
head := gtp.MsgHead{}
head.MsgId = msg.MsgId()
head.Flags = flags.Setd(gtp.Flag_Encrypted, false).
Setd(gtp.Flag_MAC, false).
Setd(gtp.Flag_Compressed, false)
// 预估追加的数据大小,因为后续数据可能会被压缩,所以此为评估值,只要保证不会内存溢出即可
msgAddition := 0
if e.Encryption {
if e.EncryptionModule == nil {
return binaryutil.MakeNonRecycleBytes(nil), errors.New("gtp: setting EncryptionModule is nil, msg can't be encrypted")
}
encAddition, err := e.EncryptionModule.SizeOfAddition(msg.Size())
if err != nil {
return binaryutil.MakeNonRecycleBytes(nil), fmt.Errorf("gtp: encrypt SizeOfAddition failed, %w", err)
}
msgAddition += encAddition
if e.PatchMAC {
if e.MACModule == nil {
return binaryutil.MakeNonRecycleBytes(nil), errors.New("gtp: setting MACModule is nil, msg can't be patch MAC")
}
msgAddition += e.MACModule.SizeofMAC(msg.Size() + encAddition)
}
}
mpBuf := binaryutil.MakeRecycleBytes(binaryutil.BytesPool.Get(head.Size() + msg.Size() + msgAddition))
defer func() {
if err != nil {
mpBuf.Release()
}
}()
// 写入消息
mn, err := msg.Read(mpBuf.Data()[head.Size():])
if err != nil {
return binaryutil.MakeNonRecycleBytes(nil), fmt.Errorf("gtp: write msg failed, %w", err)
}
end := head.Size() + mn
// 消息长度达到阀值,需要压缩消息
if e.CompressedSize > 0 && msg.Size() >= e.CompressedSize {
if e.CompressionModule == nil {
return binaryutil.MakeNonRecycleBytes(nil), errors.New("gtp: setting CompressionModule is nil, msg can't be compress")
}
compressedBuf, compressed, err := e.CompressionModule.Compress(mpBuf.Data()[head.Size():end])
if err != nil {
return binaryutil.MakeNonRecycleBytes(nil), fmt.Errorf("gtp: compress msg failed, %w", err)
}
defer compressedBuf.Release()
if compressed {
head.Flags.Set(gtp.Flag_Compressed, true)
copy(mpBuf.Data()[head.Size():], compressedBuf.Data())
end = head.Size() + len(compressedBuf.Data())
}
}
// 加密消息
if e.Encryption {
head.Flags.Set(gtp.Flag_Encrypted, true)
// 补充MAC
if e.PatchMAC {
head.Flags.Set(gtp.Flag_MAC, true)
if _, err = head.Read(mpBuf.Data()); err != nil {
return binaryutil.MakeNonRecycleBytes(nil), fmt.Errorf("gtp: failed to write msg-packet-head for patch msg-mac, %w", err)
}
macBuf, err := e.MACModule.PatchMAC(head.MsgId, head.Flags, mpBuf.Data()[head.Size():end])
if err != nil {
return binaryutil.MakeNonRecycleBytes(nil), fmt.Errorf("gtp: patch msg-mac failed, %w", err)
}
defer macBuf.Release()
copy(mpBuf.Data()[head.Size():], macBuf.Data())
end = head.Size() + len(macBuf.Data())
}
// 加密消息体
encryptBuf, err := e.EncryptionModule.Transforming(mpBuf.Data()[head.Size():end], mpBuf.Data()[head.Size():end])
if err != nil {
return binaryutil.MakeNonRecycleBytes(nil), fmt.Errorf("gtp: encrypt msg failed, %w", err)
}
defer encryptBuf.Release()
copy(mpBuf.Data()[head.Size():], encryptBuf.Data())
end = head.Size() + len(encryptBuf.Data())
}
// 调整消息大小
mpBuf = binaryutil.SliceRecycleBytes(mpBuf, 0, end)
// 写入消息头
head.Len = uint32(len(mpBuf.Data()))
if _, err = head.Read(mpBuf.Data()); err != nil {
return binaryutil.MakeNonRecycleBytes(nil), fmt.Errorf("gtp: write msg-packet-head failed, %w", err)
}
return mpBuf, nil
}