vendor/github.com/google/gopacket/tcpassembly/assembly_test.go

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