142 lines
3.1 KiB
Go
142 lines
3.1 KiB
Go
package protocol
|
|
|
|
import (
|
|
"encoding/binary"
|
|
"errors"
|
|
"io"
|
|
|
|
"ananta-ps/pkg/crypto/chacha8"
|
|
)
|
|
|
|
type Frame struct {
|
|
Mode byte
|
|
Payload []byte
|
|
}
|
|
|
|
func ReadFrame(r io.Reader, cryptor *chacha8.UxChaCha8) (*Frame, error) {
|
|
head := make([]byte, 5)
|
|
if _, err := io.ReadFull(r, head); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if cryptor != nil {
|
|
cryptor.XorInPlace(head)
|
|
}
|
|
|
|
payloadLength := int(binary.LittleEndian.Uint32(head[0:4]))
|
|
if payloadLength < 0 || payloadLength > 8*1024*1024 {
|
|
return nil, errors.New("invalid frame payload length after decrypt")
|
|
}
|
|
|
|
payload := make([]byte, payloadLength)
|
|
if payloadLength > 0 {
|
|
if _, err := io.ReadFull(r, payload); err != nil {
|
|
return nil, err
|
|
}
|
|
if cryptor != nil {
|
|
cryptor.XorInPlace(payload)
|
|
}
|
|
}
|
|
|
|
return &Frame{
|
|
Mode: head[4],
|
|
Payload: payload,
|
|
}, nil
|
|
}
|
|
|
|
func WriteFrame(w io.Writer, mode byte, payload []byte, cryptor *chacha8.UxChaCha8) error {
|
|
buf := make([]byte, 5+len(payload))
|
|
binary.LittleEndian.PutUint32(buf[0:4], uint32(len(payload)))
|
|
buf[4] = mode
|
|
copy(buf[5:], payload)
|
|
|
|
if cryptor != nil {
|
|
cryptor.XorInPlace(buf)
|
|
}
|
|
|
|
for len(buf) > 0 {
|
|
n, err := w.Write(buf)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if n <= 0 || n > len(buf) {
|
|
return io.ErrShortWrite
|
|
}
|
|
buf = buf[n:]
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func ParseRpcPacket(payload []byte) RpcPacket {
|
|
if len(payload) < 5 {
|
|
return RpcPacket{Kind: PacketKindUnknown, Body: payload}
|
|
}
|
|
|
|
tag := payload[0]
|
|
methodID := binary.LittleEndian.Uint32(payload[1:5])
|
|
|
|
switch tag {
|
|
case 0x03: // Notify
|
|
return RpcPacket{
|
|
Kind: PacketKindNotify,
|
|
MethodID: methodID,
|
|
Body: payload[5:],
|
|
}
|
|
case 0x01: // Invoke
|
|
if len(payload) >= 9 {
|
|
invokeID := int32(binary.LittleEndian.Uint32(payload[5:9]))
|
|
return RpcPacket{
|
|
Kind: PacketKindInvoke,
|
|
MethodID: methodID,
|
|
InvokeID: invokeID,
|
|
Body: payload[9:],
|
|
}
|
|
}
|
|
case 0x02: // Return
|
|
if len(payload) >= 13 {
|
|
invokeID := int32(binary.LittleEndian.Uint32(payload[5:9]))
|
|
errCode := int32(binary.LittleEndian.Uint32(payload[9:13]))
|
|
return RpcPacket{
|
|
Kind: PacketKindReturn,
|
|
MethodID: methodID,
|
|
InvokeID: invokeID,
|
|
Error: errCode,
|
|
Body: payload[13:],
|
|
}
|
|
}
|
|
}
|
|
|
|
return RpcPacket{
|
|
Kind: PacketKindUnknown,
|
|
MethodID: methodID,
|
|
Body: payload,
|
|
}
|
|
}
|
|
|
|
func BuildNotify(methodID uint32, body []byte) []byte {
|
|
buf := make([]byte, 5+len(body))
|
|
buf[0] = 0x03
|
|
binary.LittleEndian.PutUint32(buf[1:5], methodID)
|
|
copy(buf[5:], body)
|
|
return buf
|
|
}
|
|
|
|
func BuildInvoke(methodID uint32, invokeID int32, body []byte) []byte {
|
|
buf := make([]byte, 9+len(body))
|
|
buf[0] = 0x01
|
|
binary.LittleEndian.PutUint32(buf[1:5], methodID)
|
|
binary.LittleEndian.PutUint32(buf[5:9], uint32(invokeID))
|
|
copy(buf[9:], body)
|
|
return buf
|
|
}
|
|
|
|
func BuildReturn(methodID uint32, invokeID int32, errCode int32, body []byte) []byte {
|
|
buf := make([]byte, 13+len(body))
|
|
buf[0] = 0x02
|
|
binary.LittleEndian.PutUint32(buf[1:5], methodID)
|
|
binary.LittleEndian.PutUint32(buf[5:9], uint32(invokeID))
|
|
binary.LittleEndian.PutUint32(buf[9:13], uint32(errCode))
|
|
copy(buf[13:], body)
|
|
return buf
|
|
}
|