Initial release

Common/Client/go/udpsprotocol/reassembler.go

@@ -0,0 +1,105 @@
+package udpsprotocol
+
+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()
+	}
+}