1package dgram
2
3import (
4 "bytes"
5 "errors"
6 "testing"
7)
8
9func testCodec(t *testing.T) *Codec {
10 t.Helper()
11 key := bytes.Repeat([]byte{7}, KeySize)
12 c, err := NewCodec(key)
13 if err != nil {
14 t.Fatal(err)
15 }
16 return c
17}
18
19func TestCodecOverheadMatchesSealedPacketGrowth(t *testing.T) {
20 c := testCodec(t)
21 plaintext := []byte("hello")
22 pkt := c.Seal(1, 1, plaintext, nil)
23 if got, want := len(pkt), HeaderSize+len(plaintext)+c.Overhead(); got != want {
24 t.Fatalf("sealed packet length = %d, want %d", got, want)
25 }
26}
27
28func TestSealOpen(t *testing.T) {
29 c := testCodec(t)
30 pkt := c.Seal(1, 42, []byte("hello"), []byte("aad"))
31 ctr, pt, err := c.Open(pkt, 1, []byte("aad"), NewReplayWindow())
32 if err != nil {
33 t.Fatal(err)
34 }
35 if ctr != 42 || string(pt) != "hello" {
36 t.Fatalf("ctr=%d pt=%q", ctr, pt)
37 }
38 if _, _, err := c.Open(pkt, 2, []byte("aad"), nil); !errors.Is(err, ErrDirection) {
39 t.Fatalf("wrong direction err=%v", err)
40 }
41}
42
43func TestReplayReject(t *testing.T) {
44 c := testCodec(t)
45 rw := NewReplayWindow()
46 pkt := c.Seal(1, 1, []byte("x"), nil)
47 if _, _, err := c.Open(pkt, 1, nil, rw); err != nil {
48 t.Fatal(err)
49 }
50 if _, _, err := c.Open(pkt, 1, nil, rw); !errors.Is(err, ErrReplay) {
51 t.Fatalf("err=%v", err)
52 }
53 old := c.Seal(1, 0, []byte("old"), nil)
54 if _, _, err := c.Open(old, 1, nil, rw); err != nil {
55 t.Fatalf("within window should pass: %v", err)
56 }
57}
58
59func TestReplayWindowAllowsLargeFragmentReordering(t *testing.T) {
60 c := testCodec(t)
61 rw := NewReplayWindow()
62 last := uint64(maxFragmentCount)
63 if _, _, err := c.Open(c.Seal(1, last, []byte("last"), nil), 1, nil, rw); err != nil {
64 t.Fatal(err)
65 }
66 if _, _, err := c.Open(c.Seal(1, 1, []byte("first"), nil), 1, nil, rw); err != nil {
67 t.Fatalf("fragment within max fragment window rejected: %v", err)
68 }
69 if _, _, err := c.Open(c.Seal(1, last+2, []byte("advance"), nil), 1, nil, rw); err != nil {
70 t.Fatal(err)
71 }
72 tooOld := c.Seal(1, 0, []byte("old"), nil)
73 if _, _, err := c.Open(tooOld, 1, nil, rw); !errors.Is(err, ErrReplay) {
74 t.Fatalf("err=%v, want ErrReplay", err)
75 }
76}
77
78func TestFragmentation(t *testing.T) {
79 payload := bytes.Repeat([]byte("abcdef"), 400)
80 frags, err := FragmentPayload(9, payload, 100)
81 if err != nil {
82 t.Fatal(err)
83 }
84 if len(frags) < 2 {
85 t.Fatalf("got %d fragments", len(frags))
86 }
87 for _, f := range frags {
88 b, err := MarshalFragment(f)
89 if err != nil {
90 t.Fatal(err)
91 }
92 g, err := UnmarshalFragment(b)
93 if err != nil {
94 t.Fatal(err)
95 }
96 if g.Seq != f.Seq || g.Index != f.Index || !bytes.Equal(g.Data, f.Data) {
97 t.Fatalf("bad roundtrip")
98 }
99 }
100}
101
102func TestReassemblyDisorderDuplicates(t *testing.T) {
103 payload := []byte("abcdefghijklmnopqrstuvwxyz")
104 frags, err := FragmentPayload(3, payload, 30)
105 if err != nil {
106 t.Fatal(err)
107 }
108 r := NewReassembler()
109 order := []int{1, 0, 1}
110 var got []byte
111 var done bool
112 for _, idx := range order {
113 out, ok, err := r.Add(frags[idx])
114 if err != nil {
115 t.Fatal(err)
116 }
117 if ok {
118 got, done = out, true
119 }
120 }
121 if !done || !bytes.Equal(got, payload) {
122 t.Fatalf("done=%v got=%q", done, got)
123 }
124}
125
126func TestReassemblyRejectsOversizedCountBeforeInsert(t *testing.T) {
127 r := NewReassembler()
128 if _, _, err := r.Add(Fragment{Seq: 99, Index: 0, Count: uint16(maxFragmentCount + 1), Data: []byte("x")}); !errors.Is(err, ErrFragment) {
129 t.Fatalf("err=%v, want ErrFragment", err)
130 }
131 if len(r.inflight) != 0 || len(r.order) != 0 {
132 t.Fatalf("oversized fragment was inserted: inflight=%d order=%d", len(r.inflight), len(r.order))
133 }
134}
135
136func TestFragmentPayloadRejectsOverMaxFragmentCount(t *testing.T) {
137 payload := bytes.Repeat([]byte("x"), maxFragmentCount+1)
138 if _, err := FragmentPayload(1, payload, fragmentHdr+1); !errors.Is(err, ErrFragment) {
139 t.Fatalf("err=%v, want ErrFragment", err)
140 }
141}
142
143func TestReassemblyCountMismatchRemovesStaleAssembly(t *testing.T) {
144 r := NewReassembler()
145 if _, _, err := r.Add(Fragment{Seq: 7, Index: 0, Count: 2, Data: []byte("stale")}); err != nil {
146 t.Fatal(err)
147 }
148 if _, _, err := r.Add(Fragment{Seq: 7, Index: 0, Count: 3, Data: []byte("a")}); !errors.Is(err, ErrFragment) {
149 t.Fatalf("err=%v, want ErrFragment", err)
150 }
151 for _, f := range []Fragment{
152 {Seq: 7, Index: 0, Count: 3, Data: []byte("a")},
153 {Seq: 7, Index: 1, Count: 3, Data: []byte("b")},
154 {Seq: 7, Index: 2, Count: 3, Data: []byte("c")},
155 } {
156 out, ok, err := r.Add(f)
157 if err != nil {
158 t.Fatal(err)
159 }
160 if ok && string(out) != "abc" {
161 t.Fatalf("out=%q", out)
162 }
163 }
164}
165
166func TestReassemblyInflightBounded(t *testing.T) {
167 r := NewReassembler()
168 for i := range uint64(maxReassemblyInflight + 10) {
169 if _, _, err := r.Add(Fragment{Seq: i, Index: 0, Count: 2, Data: []byte("x")}); err != nil {
170 t.Fatal(err)
171 }
172 }
173 if len(r.inflight) > maxReassemblyInflight {
174 t.Fatalf("inflight=%d, want <= %d", len(r.inflight), maxReassemblyInflight)
175 }
176}