Martino Ferrari
106 lines
2.3 KiB
Go
106 lines
2.3 KiB
Go
package main
|
||
|
||
import (
|
||
"sync"
|
||
"time"
|
||
)
|
||
|
||
// fragmentSet holds the received fragments for one (counter, type) pair.
|
||
type fragmentSet struct {
|
||
total int
|
||
received int
|
||
fragments [][]byte // indexed by fragmentIdx
|
||
lastSeen time.Time
|
||
}
|
||
|
||
// Reassembler reassembles fragmented UDP payloads.
|
||
// Key: uint64(counter)<<8 | uint64(pktType)
|
||
type Reassembler struct {
|
||
mu sync.Mutex
|
||
sets map[uint64]*fragmentSet
|
||
expiry time.Duration // how long to keep incomplete sets
|
||
}
|
||
|
||
// NewReassembler creates a Reassembler that expires stale fragment sets after ttl.
|
||
func NewReassembler(ttl time.Duration) *Reassembler {
|
||
r := &Reassembler{
|
||
sets: make(map[uint64]*fragmentSet),
|
||
expiry: ttl,
|
||
}
|
||
go r.gcLoop()
|
||
return r
|
||
}
|
||
|
||
func reassemblyKey(counter uint32, pktType uint8) uint64 {
|
||
return uint64(counter)<<8 | uint64(pktType)
|
||
}
|
||
|
||
// AddFragment registers one fragment. Returns the reassembled payload and true
|
||
// when all fragments for this (counter, type) have been received, otherwise
|
||
// returns nil, false.
|
||
func (r *Reassembler) AddFragment(h PacketHeader, payload []byte) ([]byte, bool) {
|
||
key := reassemblyKey(h.Counter, h.Type)
|
||
total := int(h.TotalFragments)
|
||
idx := int(h.FragmentIdx)
|
||
|
||
// Fast path: single-fragment packet.
|
||
if total == 1 && idx == 0 {
|
||
return payload, true
|
||
}
|
||
|
||
r.mu.Lock()
|
||
defer r.mu.Unlock()
|
||
|
||
fs, ok := r.sets[key]
|
||
if !ok {
|
||
fs = &fragmentSet{
|
||
total: total,
|
||
fragments: make([][]byte, total),
|
||
}
|
||
r.sets[key] = fs
|
||
}
|
||
|
||
if idx >= len(fs.fragments) {
|
||
// Stale or corrupt – discard.
|
||
return nil, false
|
||
}
|
||
if fs.fragments[idx] == nil {
|
||
fs.fragments[idx] = payload
|
||
fs.received++
|
||
}
|
||
fs.lastSeen = time.Now()
|
||
|
||
if fs.received < fs.total {
|
||
return nil, false
|
||
}
|
||
|
||
// All fragments received – concatenate in order.
|
||
delete(r.sets, key)
|
||
|
||
total_len := 0
|
||
for _, f := range fs.fragments {
|
||
total_len += len(f)
|
||
}
|
||
out := make([]byte, 0, total_len)
|
||
for _, f := range fs.fragments {
|
||
out = append(out, f...)
|
||
}
|
||
return out, true
|
||
}
|
||
|
||
// gcLoop periodically removes fragment sets that have been incomplete for too long.
|
||
func (r *Reassembler) gcLoop() {
|
||
ticker := time.NewTicker(r.expiry / 2)
|
||
defer ticker.Stop()
|
||
for range ticker.C {
|
||
r.mu.Lock()
|
||
now := time.Now()
|
||
for k, fs := range r.sets {
|
||
if now.Sub(fs.lastSeen) > r.expiry {
|
||
delete(r.sets, k)
|
||
}
|
||
}
|
||
r.mu.Unlock()
|
||
}
|
||
}
|