80 lines
1.7 KiB
Go
80 lines
1.7 KiB
Go
package v5
|
|
|
|
import (
|
|
"bytes"
|
|
|
|
"github.com/metacubex/quic-go"
|
|
)
|
|
|
|
func fragWriteNative(quicConn quic.Connection, packet Packet, buf *bytes.Buffer, fragSize int) (err error) {
|
|
fullPayload := packet.DATA
|
|
off := 0
|
|
fragID := uint8(0)
|
|
fragCount := uint8((len(fullPayload) + fragSize - 1) / fragSize) // round up
|
|
packet.FRAG_TOTAL = fragCount
|
|
for off < len(fullPayload) {
|
|
payloadSize := len(fullPayload) - off
|
|
if payloadSize > fragSize {
|
|
payloadSize = fragSize
|
|
}
|
|
frag := packet
|
|
frag.FRAG_ID = fragID
|
|
frag.SIZE = uint16(payloadSize)
|
|
frag.DATA = fullPayload[off : off+payloadSize]
|
|
off += payloadSize
|
|
fragID++
|
|
buf.Reset()
|
|
err = frag.WriteTo(buf)
|
|
if err != nil {
|
|
return
|
|
}
|
|
data := buf.Bytes()
|
|
err = quicConn.SendMessage(data)
|
|
if err != nil {
|
|
return
|
|
}
|
|
packet.ADDR.TYPE = AtypNone // avoid "fragment 2/2: address in non-first fragment"
|
|
}
|
|
return
|
|
}
|
|
|
|
type deFragger struct {
|
|
pkgID uint16
|
|
frags []*Packet
|
|
count uint8
|
|
}
|
|
|
|
func (d *deFragger) Feed(m Packet) *Packet {
|
|
if m.FRAG_TOTAL <= 1 {
|
|
return &m
|
|
}
|
|
if m.FRAG_ID >= m.FRAG_TOTAL {
|
|
// wtf is this?
|
|
return nil
|
|
}
|
|
if d.count == 0 || m.PKT_ID != d.pkgID {
|
|
// new message, clear previous state
|
|
d.pkgID = m.PKT_ID
|
|
d.frags = make([]*Packet, m.FRAG_TOTAL)
|
|
d.count = 1
|
|
d.frags[m.FRAG_ID] = &m
|
|
} else if d.frags[m.FRAG_ID] == nil {
|
|
d.frags[m.FRAG_ID] = &m
|
|
d.count++
|
|
if int(d.count) == len(d.frags) {
|
|
// all fragments received, assemble
|
|
var data []byte
|
|
for _, frag := range d.frags {
|
|
data = append(data, frag.DATA...)
|
|
}
|
|
p := d.frags[0] // recover from first fragment
|
|
p.SIZE = uint16(len(data))
|
|
p.DATA = data
|
|
p.FRAG_ID = 0
|
|
p.FRAG_TOTAL = 1
|
|
d.count = 0
|
|
return p
|
|
}
|
|
}
|
|
return nil
|
|
}
|