1// Copyright 2012 Google, Inc. All rights reserved.
2//
3// Use of this source code is governed by a BSD-style license
4// that can be found in the LICENSE file in the root of the source
5// tree.
6
7package tcpassembly
8
9import (
10	"github.com/google/gopacket"
11	"github.com/google/gopacket/layers"
12	"net"
13	"reflect"
14	"testing"
15	"time"
16)
17
18var netFlow gopacket.Flow
19
20func init() {
21	netFlow, _ = gopacket.FlowFromEndpoints(
22		layers.NewIPEndpoint(net.IP{1, 2, 3, 4}),
23		layers.NewIPEndpoint(net.IP{5, 6, 7, 8}))
24}
25
26func TestSequenceOverflow(t *testing.T) {
27	if want, got := 20, Sequence((1<<32)-10).Difference(Sequence(10)); want != got {
28		t.Errorf("overflow diff failure: got %d want %d", got, want)
29	}
30}
31
32type testSequence struct {
33	in   layers.TCP
34	want []Reassembly
35}
36
37type testFactory struct {
38	reassembly []Reassembly
39}
40
41func (t *testFactory) New(a, b gopacket.Flow) Stream {
42	return t
43}
44func (t *testFactory) Reassembled(r []Reassembly) {
45	t.reassembly = r
46	for i := 0; i < len(r); i++ {
47		t.reassembly[i].Seen = time.Time{}
48	}
49}
50func (t *testFactory) ReassemblyComplete() {
51}
52
53func test(t *testing.T, s []testSequence) {
54	fact := &testFactory{}
55	p := NewStreamPool(fact)
56	a := NewAssembler(p)
57	a.MaxBufferedPagesPerConnection = 4
58	for i, test := range s {
59		fact.reassembly = []Reassembly{}
60		a.Assemble(netFlow, &test.in)
61		if !reflect.DeepEqual(fact.reassembly, test.want) {
62			t.Fatalf("test %v:\nwant: %v\n got: %v\n", i, test.want, fact.reassembly)
63		}
64	}
65}
66
67func TestReorder(t *testing.T) {
68	test(t, []testSequence{
69		{
70			in: layers.TCP{
71				SrcPort:   1,
72				DstPort:   2,
73				Seq:       1001,
74				BaseLayer: layers.BaseLayer{Payload: []byte{1, 2, 3}},
75			},
76			want: []Reassembly{},
77		},
78		{
79			in: layers.TCP{
80				SrcPort:   1,
81				DstPort:   2,
82				Seq:       1004,
83				BaseLayer: layers.BaseLayer{Payload: []byte{2, 2, 3}},
84			},
85			want: []Reassembly{},
86		},
87		{
88			in: layers.TCP{
89				SrcPort:   1,
90				DstPort:   2,
91				Seq:       1010,
92				BaseLayer: layers.BaseLayer{Payload: []byte{4, 2, 3}},
93			},
94			want: []Reassembly{},
95		},
96		{
97			in: layers.TCP{
98				SrcPort:   1,
99				DstPort:   2,
100				Seq:       1007,
101				BaseLayer: layers.BaseLayer{Payload: []byte{3, 2, 3}},
102			},
103			want: []Reassembly{
104				Reassembly{
105					Skip:  -1,
106					Bytes: []byte{1, 2, 3},
107				},
108				Reassembly{
109					Bytes: []byte{2, 2, 3},
110				},
111				Reassembly{
112					Bytes: []byte{3, 2, 3},
113				},
114				Reassembly{
115					Bytes: []byte{4, 2, 3},
116				},
117			},
118		},
119		{
120			in: layers.TCP{
121				SrcPort:   1,
122				DstPort:   2,
123				Seq:       1016,
124				BaseLayer: layers.BaseLayer{Payload: []byte{2, 2, 3}},
125			},
126			want: []Reassembly{},
127		},
128		{
129			in: layers.TCP{
130				SrcPort:   1,
131				DstPort:   2,
132				Seq:       1019,
133				BaseLayer: layers.BaseLayer{Payload: []byte{3, 2, 3}},
134			},
135			want: []Reassembly{},
136		},
137		{
138			in: layers.TCP{
139				SrcPort:   1,
140				DstPort:   2,
141				Seq:       1013,
142				BaseLayer: layers.BaseLayer{Payload: []byte{1, 2, 3}},
143			},
144			want: []Reassembly{
145				Reassembly{
146					Bytes: []byte{1, 2, 3},
147				},
148				Reassembly{
149					Bytes: []byte{2, 2, 3},
150				},
151				Reassembly{
152					Bytes: []byte{3, 2, 3},
153				},
154			},
155		},
156	})
157}
158
159func TestMaxPerSkip(t *testing.T) {
160	test(t, []testSequence{
161		{
162			in: layers.TCP{
163				SrcPort:   1,
164				DstPort:   2,
165				Seq:       1000,
166				SYN:       true,
167				BaseLayer: layers.BaseLayer{Payload: []byte{1, 2, 3}},
168			},
169			want: []Reassembly{
170				Reassembly{
171					Start: true,
172					Bytes: []byte{1, 2, 3},
173				},
174			},
175		},
176		{
177			in: layers.TCP{
178				SrcPort:   1,
179				DstPort:   2,
180				Seq:       1007,
181				BaseLayer: layers.BaseLayer{Payload: []byte{3, 2, 3}},
182			},
183			want: []Reassembly{},
184		},
185		{
186			in: layers.TCP{
187				SrcPort:   1,
188				DstPort:   2,
189				Seq:       1010,
190				BaseLayer: layers.BaseLayer{Payload: []byte{4, 2, 3}},
191			},
192			want: []Reassembly{},
193		},
194		{
195			in: layers.TCP{
196				SrcPort:   1,
197				DstPort:   2,
198				Seq:       1013,
199				BaseLayer: layers.BaseLayer{Payload: []byte{5, 2, 3}},
200			},
201			want: []Reassembly{},
202		},
203		{
204			in: layers.TCP{
205				SrcPort:   1,
206				DstPort:   2,
207				Seq:       1016,
208				BaseLayer: layers.BaseLayer{Payload: []byte{6, 2, 3}},
209			},
210			want: []Reassembly{
211				Reassembly{
212					Skip:  3,
213					Bytes: []byte{3, 2, 3},
214				},
215				Reassembly{
216					Bytes: []byte{4, 2, 3},
217				},
218				Reassembly{
219					Bytes: []byte{5, 2, 3},
220				},
221				Reassembly{
222					Bytes: []byte{6, 2, 3},
223				},
224			},
225		},
226	})
227}
228
229func TestReorderFast(t *testing.T) {
230	test(t, []testSequence{
231		{
232			in: layers.TCP{
233				SrcPort:   1,
234				DstPort:   2,
235				SYN:       true,
236				Seq:       1000,
237				BaseLayer: layers.BaseLayer{Payload: []byte{1, 2, 3}},
238			},
239			want: []Reassembly{
240				Reassembly{
241					Start: true,
242					Bytes: []byte{1, 2, 3},
243				},
244			},
245		},
246		{
247			in: layers.TCP{
248				SrcPort:   1,
249				DstPort:   2,
250				Seq:       1007,
251				BaseLayer: layers.BaseLayer{Payload: []byte{3, 2, 3}},
252			},
253			want: []Reassembly{},
254		},
255		{
256			in: layers.TCP{
257				SrcPort:   1,
258				DstPort:   2,
259				Seq:       1004,
260				BaseLayer: layers.BaseLayer{Payload: []byte{2, 2, 3}},
261			},
262			want: []Reassembly{
263				Reassembly{
264					Bytes: []byte{2, 2, 3},
265				},
266				Reassembly{
267					Bytes: []byte{3, 2, 3},
268				},
269			},
270		},
271	})
272}
273
274func TestOverlap(t *testing.T) {
275	test(t, []testSequence{
276		{
277			in: layers.TCP{
278				SrcPort:   1,
279				DstPort:   2,
280				SYN:       true,
281				Seq:       1000,
282				BaseLayer: layers.BaseLayer{Payload: []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 0}},
283			},
284			want: []Reassembly{
285				Reassembly{
286					Start: true,
287					Bytes: []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 0},
288				},
289			},
290		},
291		{
292			in: layers.TCP{
293				SrcPort:   1,
294				DstPort:   2,
295				Seq:       1007,
296				BaseLayer: layers.BaseLayer{Payload: []byte{7, 8, 9, 0, 1, 2, 3, 4, 5}},
297			},
298			want: []Reassembly{
299				Reassembly{
300					Bytes: []byte{1, 2, 3, 4, 5},
301				},
302			},
303		},
304		{
305			in: layers.TCP{
306				SrcPort:   1,
307				DstPort:   2,
308				Seq:       1010,
309				BaseLayer: layers.BaseLayer{Payload: []byte{0, 1, 2, 3, 4, 5, 6, 7}},
310			},
311			want: []Reassembly{
312				Reassembly{
313					Bytes: []byte{6, 7},
314				},
315			},
316		},
317	})
318}
319
320func TestBufferedOverlap(t *testing.T) {
321	test(t, []testSequence{
322		{
323			in: layers.TCP{
324				SrcPort:   1,
325				DstPort:   2,
326				Seq:       1007,
327				BaseLayer: layers.BaseLayer{Payload: []byte{7, 8, 9, 0, 1, 2, 3, 4, 5}},
328			},
329			want: []Reassembly{},
330		},
331		{
332			in: layers.TCP{
333				SrcPort:   1,
334				DstPort:   2,
335				Seq:       1010,
336				BaseLayer: layers.BaseLayer{Payload: []byte{0, 1, 2, 3, 4, 5, 6, 7}},
337			},
338			want: []Reassembly{},
339		},
340		{
341			in: layers.TCP{
342				SrcPort:   1,
343				DstPort:   2,
344				SYN:       true,
345				Seq:       1000,
346				BaseLayer: layers.BaseLayer{Payload: []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 0}},
347			},
348			want: []Reassembly{
349				Reassembly{
350					Start: true,
351					Bytes: []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 0},
352				},
353				Reassembly{
354					Bytes: []byte{1, 2, 3, 4, 5},
355				},
356				Reassembly{
357					Bytes: []byte{6, 7},
358				},
359			},
360		},
361	})
362}
363
364func TestOverrun1(t *testing.T) {
365	test(t, []testSequence{
366		{
367			in: layers.TCP{
368				SrcPort:   1,
369				DstPort:   2,
370				SYN:       true,
371				Seq:       0xFFFFFFFF,
372				BaseLayer: layers.BaseLayer{Payload: []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 0}},
373			},
374			want: []Reassembly{
375				Reassembly{
376					Start: true,
377					Bytes: []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 0},
378				},
379			},
380		},
381		{
382			in: layers.TCP{
383				SrcPort:   1,
384				DstPort:   2,
385				Seq:       10,
386				BaseLayer: layers.BaseLayer{Payload: []byte{1, 2, 3, 4}},
387			},
388			want: []Reassembly{
389				Reassembly{
390					Bytes: []byte{1, 2, 3, 4},
391				},
392			},
393		},
394	})
395}
396
397func TestOverrun2(t *testing.T) {
398	test(t, []testSequence{
399		{
400			in: layers.TCP{
401				SrcPort:   1,
402				DstPort:   2,
403				Seq:       10,
404				BaseLayer: layers.BaseLayer{Payload: []byte{1, 2, 3, 4}},
405			},
406			want: []Reassembly{},
407		},
408		{
409			in: layers.TCP{
410				SrcPort:   1,
411				DstPort:   2,
412				SYN:       true,
413				Seq:       0xFFFFFFFF,
414				BaseLayer: layers.BaseLayer{Payload: []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 0}},
415			},
416			want: []Reassembly{
417				Reassembly{
418					Start: true,
419					Bytes: []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 0},
420				},
421				Reassembly{
422					Bytes: []byte{1, 2, 3, 4},
423				},
424			},
425		},
426	})
427}
428
429func TestCacheLargePacket(t *testing.T) {
430	data := make([]byte, pageBytes*3)
431	test(t, []testSequence{
432		{
433			in: layers.TCP{
434				SrcPort:   1,
435				DstPort:   2,
436				Seq:       1001,
437				BaseLayer: layers.BaseLayer{Payload: data},
438			},
439			want: []Reassembly{},
440		},
441		{
442			in: layers.TCP{
443				SrcPort:   1,
444				DstPort:   2,
445				Seq:       1000,
446				SYN:       true,
447				BaseLayer: layers.BaseLayer{Payload: []byte{}},
448			},
449			want: []Reassembly{
450				Reassembly{
451					Start: true,
452					Bytes: []byte{},
453				},
454				Reassembly{
455					Bytes: data[:pageBytes],
456				},
457				Reassembly{
458					Bytes: data[pageBytes : pageBytes*2],
459				},
460				Reassembly{
461					Bytes: data[pageBytes*2 : pageBytes*3],
462				},
463			},
464		},
465	})
466}
467
468func BenchmarkSingleStream(b *testing.B) {
469	t := layers.TCP{
470		SrcPort:   1,
471		DstPort:   2,
472		SYN:       true,
473		Seq:       1000,
474		BaseLayer: layers.BaseLayer{Payload: []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 0}},
475	}
476	a := NewAssembler(NewStreamPool(&testFactory{}))
477	for i := 0; i < b.N; i++ {
478		a.Assemble(netFlow, &t)
479		if t.SYN {
480			t.SYN = false
481			t.Seq++
482		}
483		t.Seq += 10
484	}
485}
486
487func BenchmarkSingleStreamSkips(b *testing.B) {
488	t := layers.TCP{
489		SrcPort:   1,
490		DstPort:   2,
491		SYN:       true,
492		Seq:       1000,
493		BaseLayer: layers.BaseLayer{Payload: []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 0}},
494	}
495	a := NewAssembler(NewStreamPool(&testFactory{}))
496	skipped := false
497	for i := 0; i < b.N; i++ {
498		if i%10 == 9 {
499			t.Seq += 10
500			skipped = true
501		} else if skipped {
502			t.Seq -= 20
503		}
504		a.Assemble(netFlow, &t)
505		if t.SYN {
506			t.SYN = false
507			t.Seq++
508		}
509		t.Seq += 10
510		if skipped {
511			t.Seq += 10
512			skipped = false
513		}
514	}
515}
516
517func BenchmarkSingleStreamLoss(b *testing.B) {
518	t := layers.TCP{
519		SrcPort:   1,
520		DstPort:   2,
521		SYN:       true,
522		Seq:       1000,
523		BaseLayer: layers.BaseLayer{Payload: []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 0}},
524	}
525	a := NewAssembler(NewStreamPool(&testFactory{}))
526	for i := 0; i < b.N; i++ {
527		a.Assemble(netFlow, &t)
528		t.SYN = false
529		t.Seq += 11
530	}
531}
532
533func BenchmarkMultiStreamGrow(b *testing.B) {
534	t := layers.TCP{
535		SrcPort:   1,
536		DstPort:   2,
537		Seq:       0,
538		BaseLayer: layers.BaseLayer{Payload: []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 0}},
539	}
540	a := NewAssembler(NewStreamPool(&testFactory{}))
541	for i := 0; i < b.N; i++ {
542		t.SrcPort = layers.TCPPort(i)
543		a.Assemble(netFlow, &t)
544		t.Seq += 10
545	}
546}
547
548func BenchmarkMultiStreamConn(b *testing.B) {
549	t := layers.TCP{
550		SrcPort:   1,
551		DstPort:   2,
552		Seq:       0,
553		SYN:       true,
554		BaseLayer: layers.BaseLayer{Payload: []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 0}},
555	}
556	a := NewAssembler(NewStreamPool(&testFactory{}))
557	for i := 0; i < b.N; i++ {
558		t.SrcPort = layers.TCPPort(i)
559		a.Assemble(netFlow, &t)
560		if i%65536 == 65535 {
561			if t.SYN {
562				t.SYN = false
563				t.Seq++
564			}
565			t.Seq += 10
566		}
567	}
568}
569