1// Package dgram provides stdlib-only helpers for authenticated datagram
2// framing, replay protection, and MTU-sized fragmentation.
3package dgram
4
5import (
6 "crypto/aes"
7 "crypto/cipher"
8 "encoding/binary"
9 "errors"
10 "fmt"
11)
12
13const (
14 KeySize = 32
15 NonceSize = 12
16 HeaderSize = 4 + 8
17 DefaultMTU = 1200
18 fragmentHdr = 8 + 2 + 2 + 2 + 2 // seq, idx, count, total, reserved
19)
20
21var (
22 ErrReplay = errors.New("dgram: replayed packet")
23 ErrDirection = errors.New("dgram: unexpected direction")
24 ErrFragment = errors.New("dgram: invalid fragment")
25)
26
27// Erase overwrites b with zeroes. This is best effort: copies made by the
28// compiler, runtime, or called code may retain the original data.
29func Erase(b []byte) {
30 clear(b)
31}
32
33type Codec struct {
34 gcm cipher.AEAD
35}
36
37func NewCodec(key []byte) (*Codec, error) {
38 if len(key) != KeySize {
39 return nil, fmt.Errorf("dgram: key must be %d bytes", KeySize)
40 }
41 block, err := aes.NewCipher(key)
42 if err != nil {
43 return nil, err
44 }
45 gcm, err := cipher.NewGCM(block)
46 if err != nil {
47 return nil, err
48 }
49 return &Codec{gcm: gcm}, nil
50}
51
52func (c *Codec) Overhead() int { return c.gcm.Overhead() }
53
54func (c *Codec) Seal(direction uint32, counter uint64, plaintext, aad []byte) []byte {
55 var nonce [NonceSize]byte
56 binary.BigEndian.PutUint32(nonce[0:4], direction)
57 binary.BigEndian.PutUint64(nonce[4:12], counter)
58 out := make([]byte, HeaderSize, HeaderSize+len(plaintext)+c.gcm.Overhead())
59 binary.BigEndian.PutUint32(out[0:4], direction)
60 binary.BigEndian.PutUint64(out[4:12], counter)
61 return c.gcm.Seal(out, nonce[:], plaintext, aad)
62}
63
64func (c *Codec) Open(packet []byte, wantDirection uint32, aad []byte, replay *ReplayWindow) (uint64, []byte, error) {
65 if len(packet) < HeaderSize+c.gcm.Overhead() {
66 return 0, nil, errors.New("dgram: packet too short")
67 }
68 dir := binary.BigEndian.Uint32(packet[0:4])
69 if dir != wantDirection {
70 return 0, nil, ErrDirection
71 }
72 ctr := binary.BigEndian.Uint64(packet[4:12])
73 if replay != nil && !replay.Check(ctr) {
74 return 0, nil, ErrReplay
75 }
76 var nonce [NonceSize]byte
77 copy(nonce[:], packet[:HeaderSize])
78 pt, err := c.gcm.Open(nil, nonce[:], packet[HeaderSize:], aad)
79 if err != nil {
80 return 0, nil, err
81 }
82 if replay != nil {
83 replay.Accept(ctr)
84 }
85 return ctr, pt, nil
86}
87
88const replayWindowSize = maxFragmentCount
89
90type ReplayWindow struct {
91 max uint64
92 seen map[uint64]struct{}
93 init bool
94}
95
96func NewReplayWindow() *ReplayWindow { return &ReplayWindow{} }
97
98func (w *ReplayWindow) Check(counter uint64) bool {
99 if !w.init {
100 return true
101 }
102 if counter > w.max {
103 return true
104 }
105 if w.max-counter >= uint64(replayWindowSize) {
106 return false
107 }
108 _, ok := w.seen[counter]
109 return !ok
110}
111
112func (w *ReplayWindow) Accept(counter uint64) {
113 if w.seen == nil {
114 w.seen = make(map[uint64]struct{})
115 }
116 if !w.init {
117 w.max, w.init = counter, true
118 }
119 if counter > w.max {
120 w.max = counter
121 }
122 w.seen[counter] = struct{}{}
123 for seen := range w.seen {
124 if w.max-seen >= uint64(replayWindowSize) {
125 delete(w.seen, seen)
126 }
127 }
128}
129
130type Fragment struct {
131 Seq uint64
132 Index uint16
133 Count uint16
134 Data []byte
135}
136
137func FragmentPayload(seq uint64, payload []byte, mtu int) ([]Fragment, error) {
138 if mtu <= fragmentHdr {
139 return nil, ErrFragment
140 }
141 maxData := mtu - fragmentHdr
142 count := (len(payload) + maxData - 1) / maxData
143 if count == 0 {
144 count = 1
145 }
146 if count > maxFragmentCount {
147 return nil, ErrFragment
148 }
149 out := make([]Fragment, count)
150 for i := 0; i < count; i++ {
151 start := i * maxData
152 end := min(len(payload), start+maxData)
153 out[i] = Fragment{Seq: seq, Index: uint16(i), Count: uint16(count), Data: append([]byte(nil), payload[start:end]...)}
154 }
155 return out, nil
156}
157
158func MarshalFragment(f Fragment) ([]byte, error) {
159 if f.Count == 0 || f.Index >= f.Count {
160 return nil, ErrFragment
161 }
162 b := make([]byte, fragmentHdr+len(f.Data))
163 binary.BigEndian.PutUint64(b[0:8], f.Seq)
164 binary.BigEndian.PutUint16(b[8:10], f.Index)
165 binary.BigEndian.PutUint16(b[10:12], f.Count)
166 binary.BigEndian.PutUint16(b[12:14], uint16(len(f.Data)))
167 copy(b[fragmentHdr:], f.Data)
168 return b, nil
169}
170
171func UnmarshalFragment(b []byte) (Fragment, error) {
172 if len(b) < fragmentHdr {
173 return Fragment{}, ErrFragment
174 }
175 f := Fragment{Seq: binary.BigEndian.Uint64(b[0:8]), Index: binary.BigEndian.Uint16(b[8:10]), Count: binary.BigEndian.Uint16(b[10:12])}
176 n := int(binary.BigEndian.Uint16(b[12:14]))
177 if f.Count == 0 || f.Index >= f.Count || len(b[fragmentHdr:]) != n {
178 return Fragment{}, ErrFragment
179 }
180 f.Data = append([]byte(nil), b[fragmentHdr:]...)
181 return f, nil
182}
183
184const (
185 maxFragmentCount = 1024
186 maxReassemblyInflight = 1024
187)
188
189type Reassembler struct {
190 inflight map[uint64]*assembly
191 order []uint64
192}
193type assembly struct {
194 count uint16
195 got int
196 parts [][]byte
197}
198
199func NewReassembler() *Reassembler { return &Reassembler{inflight: make(map[uint64]*assembly)} }
200
201// Inflight returns the number of incomplete records currently being reassembled.
202func (r *Reassembler) Inflight() int { return len(r.inflight) }
203
204func (r *Reassembler) Add(f Fragment) ([]byte, bool, error) {
205 if f.Count == 0 || int(f.Count) > maxFragmentCount || f.Index >= f.Count {
206 return nil, false, ErrFragment
207 }
208 a := r.inflight[f.Seq]
209 if a == nil {
210 r.evictOldest()
211 a = &assembly{count: f.Count, parts: make([][]byte, f.Count)}
212 r.inflight[f.Seq] = a
213 r.order = append(r.order, f.Seq)
214 }
215 if a.count != f.Count {
216 r.remove(f.Seq)
217 return nil, false, ErrFragment
218 }
219 if a.parts[f.Index] != nil {
220 return nil, false, nil
221 }
222 a.parts[f.Index] = append([]byte(nil), f.Data...)
223 a.got++
224 if a.got != int(a.count) {
225 return nil, false, nil
226 }
227 var total int
228 for _, p := range a.parts {
229 total += len(p)
230 }
231 out := make([]byte, 0, total)
232 for _, p := range a.parts {
233 out = append(out, p...)
234 }
235 r.remove(f.Seq)
236 return out, true, nil
237}
238
239func (r *Reassembler) evictOldest() {
240 for len(r.order) >= maxReassemblyInflight {
241 seq := r.order[0]
242 r.order = r.order[1:]
243 if _, ok := r.inflight[seq]; ok {
244 delete(r.inflight, seq)
245 return
246 }
247 }
248}
249
250func (r *Reassembler) remove(seq uint64) {
251 delete(r.inflight, seq)
252 for i, s := range r.order {
253 if s == seq {
254 copy(r.order[i:], r.order[i+1:])
255 r.order = r.order[:len(r.order)-1]
256 return
257 }
258 }
259}