mirror of
https://github.com/fumiama/WireGold.git
synced 2026-06-04 23:40:26 +08:00
143 lines
3.5 KiB
Go
143 lines
3.5 KiB
Go
package head
|
||
|
||
import (
|
||
"encoding/binary"
|
||
"errors"
|
||
"net"
|
||
"unsafe"
|
||
|
||
blake2b "github.com/minio/blake2b-simd"
|
||
"github.com/sirupsen/logrus"
|
||
)
|
||
|
||
// Packet 是发送和接收的最小单位
|
||
type Packet struct {
|
||
// Ver 协议版本
|
||
Ver uint16
|
||
// DataSZ len(Data)
|
||
// 不得超过 65507-head 字节
|
||
DataSZ uint16
|
||
// Proto 详见 head
|
||
Proto uint8
|
||
// TTL is time to live
|
||
TTL uint8
|
||
// SrcPort 源端口
|
||
SrcPort uint16
|
||
// DstPort 目的端口
|
||
DstPort uint16
|
||
// Flags 高3位为标志(xDM),低13位为分片偏移
|
||
Flags uint16
|
||
// Src 源 ip (ipv4)
|
||
Src net.IP
|
||
// Dst 目的 ip (ipv4)
|
||
Dst net.IP
|
||
// Hash 使用 BLAKE2 生成加密前 Packet 的摘要
|
||
// 生成时 Hash 全 0
|
||
// https://github.com/minio/blake2b-simd
|
||
Hash [32]byte
|
||
// Data 承载的数据
|
||
Data []byte
|
||
// 记录还有多少字节未到达
|
||
rembytes uint16
|
||
}
|
||
|
||
// NewPacket 生成一个新包
|
||
func NewPacket(proto uint8, srcPort uint16, dst net.IP, dstPort uint16, data []byte) *Packet {
|
||
logrus.Debugln("[packet] new: [proto:", proto, ", srcport:", srcPort, ", dstport:", dstPort, ", dst:", dst, ", data:", data)
|
||
return &Packet{
|
||
Ver: 1,
|
||
Proto: proto,
|
||
TTL: 16,
|
||
SrcPort: srcPort,
|
||
DstPort: dstPort,
|
||
Dst: dst,
|
||
Data: data,
|
||
}
|
||
}
|
||
|
||
// Unmarshal 将 data 的数据解码到自身
|
||
func (p *Packet) Unmarshal(data []byte) (complete bool, err error) {
|
||
if len(data) < 12 {
|
||
err = errors.New("data len < 12")
|
||
return
|
||
}
|
||
if p.DataSZ == 0 && len(p.Data) == 0 {
|
||
p.Ver = binary.LittleEndian.Uint16(data[:2])
|
||
if p.Ver != 1 {
|
||
err = errors.New("unknown protocol version")
|
||
return
|
||
}
|
||
p.DataSZ = binary.LittleEndian.Uint16(data[2:4])
|
||
p.Data = make([]byte, p.DataSZ)
|
||
pt := binary.LittleEndian.Uint16(data[4:6])
|
||
p.Proto = uint8(pt)
|
||
p.TTL = uint8(pt >> 8)
|
||
p.SrcPort = binary.LittleEndian.Uint16(data[6:8])
|
||
p.DstPort = binary.LittleEndian.Uint16(data[8:10])
|
||
p.rembytes = p.DataSZ
|
||
}
|
||
|
||
p.Flags = binary.LittleEndian.Uint16(data[10:12])
|
||
|
||
p.Src = make(net.IP, 4)
|
||
copy(p.Src, data[12:16])
|
||
p.Dst = make(net.IP, 4)
|
||
copy(p.Dst, data[16:20])
|
||
copy(p.Hash[:], data[20:52])
|
||
p.rembytes -= uint16(copy(p.Data[p.Flags<<3:], data[52:]))
|
||
|
||
complete = p.rembytes == 0
|
||
|
||
return
|
||
}
|
||
|
||
// Marshal 将自身数据编码为 []byte
|
||
// offset 必须为 8 的倍数,表示偏移的 8 位
|
||
func (p *Packet) Marshal(src net.IP, offset uint16, dontfrag, hasmore bool) []byte {
|
||
p.TTL--
|
||
if p.TTL == 0 {
|
||
return nil
|
||
}
|
||
|
||
p.DataSZ = uint16(len(p.Data))
|
||
if src != nil {
|
||
p.Src = src
|
||
offset >>= 3
|
||
if dontfrag {
|
||
offset |= 0x4000
|
||
}
|
||
if hasmore {
|
||
offset |= 0x2000
|
||
}
|
||
p.Flags = offset
|
||
}
|
||
|
||
packet := make([]byte, 52+len(p.Data))
|
||
binary.LittleEndian.PutUint16(packet[:2], p.Ver)
|
||
binary.LittleEndian.PutUint16(packet[2:4], p.DataSZ)
|
||
binary.LittleEndian.PutUint16(packet[4:6], (uint16(p.TTL)<<8)|uint16(p.Proto))
|
||
binary.LittleEndian.PutUint16(packet[6:8], p.SrcPort)
|
||
binary.LittleEndian.PutUint16(packet[8:10], p.DstPort)
|
||
binary.LittleEndian.PutUint16(packet[10:12], p.Flags)
|
||
copy(packet[12:16], p.Src.To4())
|
||
copy(packet[16:20], p.Dst.To4())
|
||
copy(packet[20:52], p.Hash[:])
|
||
copy(packet[52:], p.Data)
|
||
|
||
// logrus.Debugln("[packet] marshaled packet:", hex.EncodeToString(packet))
|
||
|
||
return packet
|
||
}
|
||
|
||
// FillHash 生成 p.Data 的 Hash
|
||
func (p *Packet) FillHash() {
|
||
sum := blake2b.New256().Sum(p.Data)
|
||
p.Hash = *(*[32]byte)(*(*unsafe.Pointer)(unsafe.Pointer(&sum)))
|
||
}
|
||
|
||
// IsVaildHash 验证 packet 合法性
|
||
func (p *Packet) IsVaildHash() bool {
|
||
sum := blake2b.New256().Sum(p.Data)
|
||
return *(*[32]byte)(*(*unsafe.Pointer)(unsafe.Pointer(&sum))) == p.Hash
|
||
}
|