Add stream_handler tests

2 files changed, 196 insertions, 19 deletions

diff --git a/stream_handler.go b/stream_handler.go
index 80d91d6..ce580fc 100644
--- a/stream_handler.go
+++ b/stream_handler.go
@@ -135,7 +135,7 @@ func (sh *StreamHandler) onMatch(id uint, from uint64, to uint64, flags uint, co
 		// new from == new match
 		sh.patternMatches[id] = append(patternSlices, PatternSlice{from, to})
 	} else {
-		patternSlices = make([]PatternSlice, InitialPatternSliceSize)
+		patternSlices = make([]PatternSlice, 1, InitialPatternSliceSize)
 		patternSlices[0] = PatternSlice{from, to}
 		sh.patternMatches[id] = patternSlices
 	}
@@ -169,7 +169,7 @@ func (sh *StreamHandler) generateDocumentKey() string {
 	endpointsHash := sh.streamKey[0].FastHash() ^ sh.streamKey[1].FastHash() ^
 		sh.streamKey[2].FastHash() ^ sh.streamKey[3].FastHash()
 	binary.BigEndian.PutUint64(hash, endpointsHash)
-	binary.BigEndian.PutUint64(hash[8:], uint64(sh.timestamps[0].UnixNano()))
+	binary.BigEndian.PutUint64(hash[8:], uint64(sh.firstPacketSeen.UnixNano()))
 	binary.BigEndian.PutUint16(hash[8:], uint16(len(sh.documentsKeys)))
 
 	return fmt.Sprintf("%x", hash)
diff --git a/stream_handler_test.go b/stream_handler_test.go
index a1004dc..425c1b7 100644
--- a/stream_handler_test.go
+++ b/stream_handler_test.go
@@ -3,6 +3,7 @@ package main
 import (
 	"context"
 	"errors"
+	"fmt"
 	"github.com/flier/gohs/hyperscan"
 	"github.com/google/gopacket/layers"
 	"github.com/google/gopacket/tcpassembly"
@@ -23,9 +24,16 @@ const dstPort = 8080
 func TestReassemblingEmptyStream(t *testing.T) {
 	patterns, err := hyperscan.NewStreamDatabase(hyperscan.NewPattern("/nope/", 0))
 	require.Nil(t, err)
-	streamHandler := createTestStreamHandler(t, testStorage{}, patterns)
+	scratch, err := hyperscan.NewScratch(patterns)
+	require.Nil(t, err)
+	streamHandler := createTestStreamHandler(testStorage{}, patterns, scratch)
 
-	streamHandler.Reassembled([]tcpassembly.Reassembly{})
+	streamHandler.Reassembled([]tcpassembly.Reassembly{{
+		Bytes: []byte{},
+		Skip:  0,
+		Start: true,
+		End:   true,
+	}})
 	assert.Len(t, streamHandler.indexes, 0, "indexes")
 	assert.Len(t, streamHandler.timestamps, 0, "timestamps")
 	assert.Len(t, streamHandler.lossBlocks, 0)
@@ -42,18 +50,26 @@ func TestReassemblingEmptyStream(t *testing.T) {
 	}
 	streamHandler.ReassemblyComplete()
 	assert.Equal(t, 42, expected)
+
+	err = scratch.Free()
+	require.Nil(t, err, "free scratch")
+	err = patterns.Close()
+	require.Nil(t, err, "close stream database")
 }
 
-func TestReassemblingSingleDocumentStream(t *testing.T) {
+
+func TestReassemblingSingleDocument(t *testing.T) {
 	patterns, err := hyperscan.NewStreamDatabase(hyperscan.NewPattern("/impossible_to_match/", 0))
 	require.Nil(t, err)
+	scratch, err := hyperscan.NewScratch(patterns)
+	require.Nil(t, err)
 	storage := &testStorage{}
-	streamHandler := createTestStreamHandler(t, storage, patterns)
+	streamHandler := createTestStreamHandler(storage, patterns, scratch)
 
 	payloadLen := 256
-	firstTime := time.Unix(1000000000, 0)
-	middleTime := time.Unix(1000000010, 0)
-	lastTime := time.Unix(1000000020, 0)
+	firstTime := time.Unix(0, 0)
+	middleTime := time.Unix(10, 0)
+	lastTime := time.Unix(20, 0)
 	data := make([]byte, MaxDocumentSize)
 	rand.Read(data)
 	reassembles := make([]tcpassembly.Reassembly, MaxDocumentSize / payloadLen)
@@ -85,7 +101,7 @@ func TestReassemblingSingleDocumentStream(t *testing.T) {
 	storage.insertFunc = func(ctx context.Context, collectionName string, document interface{}) (i interface{}, err error) {
 		od := document.(OrderedDocument)
 		assert.Equal(t, "connection_streams", collectionName)
-		assert.Equal(t, "bb41a60281cfae830000b6b3a7640000", od[0].Value)
+		assert.Equal(t, "bb41a60281cfae830000000000000000", od[0].Value)
 		assert.Equal(t, nil, od[1].Value)
 		assert.Equal(t, 0, od[2].Value)
 		assert.Equal(t, data, od[3].Value)
@@ -102,37 +118,198 @@ func TestReassemblingSingleDocumentStream(t *testing.T) {
 		inserted = false
 	}
 
-	assert.Equal(t, data, streamHandler.buffer.Bytes(), "buffer should contains the same bytes of reassembles")
-	assert.Equal(t, indexes, streamHandler.indexes, "indexes")
-	assert.Equal(t, timestamps, streamHandler.timestamps, "timestamps")
-	assert.Equal(t, lossBlocks, streamHandler.lossBlocks, "lossBlocks")
+	completed := false
+	streamHandler.connection.(*testConnectionHandler).onComplete = func(handler *StreamHandler) {
+		completed = true
+	}
+	streamHandler.ReassemblyComplete()
+
 	assert.Equal(t, len(data), streamHandler.currentIndex)
 	assert.Equal(t, firstTime, streamHandler.firstPacketSeen)
 	assert.Equal(t, lastTime, streamHandler.lastPacketSeen)
-	assert.Len(t, streamHandler.documentsKeys, 0)
+	assert.Len(t, streamHandler.documentsKeys, 1)
 	assert.Equal(t, len(data), streamHandler.streamLength)
 	assert.Len(t, streamHandler.patternMatches, 0)
 
+	assert.Equal(t, true, inserted, "inserted")
+	assert.Equal(t, true, completed, "completed")
+
+	err = scratch.Free()
+	require.Nil(t, err, "free scratch")
+	err = patterns.Close()
+	require.Nil(t, err, "close stream database")
+}
+
+
+func TestReassemblingMultipleDocuments(t *testing.T) {
+	patterns, err := hyperscan.NewStreamDatabase(hyperscan.NewPattern("/impossible_to_match/", 0))
+	require.Nil(t, err)
+	scratch, err := hyperscan.NewScratch(patterns)
+	require.Nil(t, err)
+	storage := &testStorage{}
+	streamHandler := createTestStreamHandler(storage, patterns, scratch)
+
+	payloadLen := 256
+	firstTime := time.Unix(0, 0)
+	middleTime := time.Unix(10, 0)
+	lastTime := time.Unix(20, 0)
+	dataSize := MaxDocumentSize*2
+	data := make([]byte, dataSize)
+	rand.Read(data)
+	reassembles := make([]tcpassembly.Reassembly, dataSize / payloadLen)
+	indexes := make([]int, dataSize / payloadLen)
+	timestamps := make([]time.Time, dataSize / payloadLen)
+	lossBlocks := make([]bool, dataSize / payloadLen)
+	for i := 0; i < len(reassembles); i++ {
+		var seen time.Time
+		if i == 0 {
+			seen = firstTime
+		} else if i == len(reassembles)-1 {
+			seen = lastTime
+		} else {
+			seen = middleTime
+		}
+
+		reassembles[i] = tcpassembly.Reassembly{
+			Bytes: data[i*payloadLen:(i+1)*payloadLen],
+			Skip:  0,
+			Start: i == 0,
+			End:   i == len(reassembles)-1,
+			Seen:  seen,
+		}
+		indexes[i] = i*payloadLen % MaxDocumentSize
+		timestamps[i] = seen
+	}
+
+	inserted := 0
+	storage.insertFunc = func(ctx context.Context, collectionName string, document interface{}) (i interface{}, err error) {
+		od := document.(OrderedDocument)
+		blockLen := MaxDocumentSize / payloadLen
+		assert.Equal(t, "connection_streams", collectionName)
+		assert.Equal(t, fmt.Sprintf("bb41a60281cfae83000%v000000000000", inserted), od[0].Value)
+		assert.Equal(t, nil, od[1].Value)
+		assert.Equal(t, inserted, od[2].Value)
+		assert.Equal(t, data[MaxDocumentSize*inserted:MaxDocumentSize*(inserted+1)], od[3].Value)
+		assert.Equal(t, indexes[blockLen*inserted:blockLen*(inserted+1)], od[4].Value)
+		assert.Equal(t, timestamps[blockLen*inserted:blockLen*(inserted+1)], od[5].Value)
+		assert.Equal(t, lossBlocks[blockLen*inserted:blockLen*(inserted+1)], od[6].Value)
+		assert.Len(t, od[7].Value, 0)
+		inserted += 1
+
+		return nil, nil
+	}
+
+	streamHandler.Reassembled(reassembles)
+	if !assert.Equal(t, 1, inserted) {
+		inserted = 1
+	}
+
 	completed := false
 	streamHandler.connection.(*testConnectionHandler).onComplete = func(handler *StreamHandler) {
 		completed = true
 	}
 	streamHandler.ReassemblyComplete()
+
+	assert.Equal(t, MaxDocumentSize, streamHandler.currentIndex)
+	assert.Equal(t, firstTime, streamHandler.firstPacketSeen)
+	assert.Equal(t, lastTime, streamHandler.lastPacketSeen)
+	assert.Len(t, streamHandler.documentsKeys, 2)
+	assert.Equal(t, len(data), streamHandler.streamLength)
+	assert.Len(t, streamHandler.patternMatches, 0)
+
+	assert.Equal(t, 2, inserted, "inserted")
+	assert.Equal(t, true, completed, "completed")
+
+	err = scratch.Free()
+	require.Nil(t, err, "free scratch")
+	err = patterns.Close()
+	require.Nil(t, err, "close stream database")
+}
+
+func TestReassemblingPatternMatching(t *testing.T) {
+	a, err := hyperscan.ParsePattern("/a{8}/i")
+	require.Nil(t, err)
+	a.Id = 0
+	a.Flags |= hyperscan.SomLeftMost
+	b, err := hyperscan.ParsePattern("/b[c]+b/i")
+	require.Nil(t, err)
+	b.Id = 1
+	b.Flags |= hyperscan.SomLeftMost
+	d, err := hyperscan.ParsePattern("/[d]+e[d]+/i")
+	require.Nil(t, err)
+	d.Id = 2
+	d.Flags |= hyperscan.SomLeftMost
+
+	payload := "aaaaaaaa0aaaaaaaaaa0bbbcccbbb0dddeddddedddd"
+	expected := map[uint][]PatternSlice{
+		0: {{0, 8}, {9, 17}, {10, 18}, {11, 19}},
+		1: {{22, 27}},
+		2: {{30, 38}, {34, 43}},
+	}
+
+	patterns, err := hyperscan.NewStreamDatabase(a, b, d)
+	require.Nil(t, err)
+	scratch, err := hyperscan.NewScratch(patterns)
+	require.Nil(t, err)
+	storage := &testStorage{}
+	streamHandler := createTestStreamHandler(storage, patterns, scratch)
+
+	seen := time.Unix(0, 0)
+	inserted := false
+	storage.insertFunc = func(ctx context.Context, collectionName string, document interface{}) (i interface{}, err error) {
+		od := document.(OrderedDocument)
+		assert.Equal(t, "connection_streams", collectionName)
+		assert.Equal(t, "bb41a60281cfae830000000000000000", od[0].Value)
+		assert.Equal(t, nil, od[1].Value)
+		assert.Equal(t, 0, od[2].Value)
+		assert.Equal(t, []byte(payload), od[3].Value)
+		assert.Equal(t, []int{0}, od[4].Value)
+		assert.Equal(t, []time.Time{seen}, od[5].Value)
+		assert.Equal(t, []bool{false}, od[6].Value)
+		assert.Equal(t, expected, od[7].Value)
+		inserted = true
+
+		return nil, nil
+	}
+
+	streamHandler.Reassembled([]tcpassembly.Reassembly{{
+		Bytes: []byte(payload),
+		Skip:  0,
+		Start: true,
+		End:   true,
+		Seen:  seen,
+	}})
+	assert.Equal(t, false, inserted)
+
+	completed := false
+	streamHandler.connection.(*testConnectionHandler).onComplete = func(handler *StreamHandler) {
+		completed = true
+	}
+	streamHandler.ReassemblyComplete()
+
+	assert.Equal(t, len(payload), streamHandler.currentIndex)
+	assert.Equal(t, seen, streamHandler.firstPacketSeen)
+	assert.Equal(t, seen, streamHandler.lastPacketSeen)
+	assert.Len(t, streamHandler.documentsKeys, 1)
+	assert.Equal(t, len(payload), streamHandler.streamLength)
+
 	assert.Equal(t, true, inserted, "inserted")
 	assert.Equal(t, true, completed, "completed")
+
+	err = scratch.Free()
+	require.Nil(t, err, "free scratch")
+	err = patterns.Close()
+	require.Nil(t, err, "close stream database")
 }
 
 
-func createTestStreamHandler(t *testing.T, storage Storage, patterns hyperscan.StreamDatabase) StreamHandler {
+func createTestStreamHandler(storage Storage, patterns hyperscan.StreamDatabase, scratch *hyperscan.Scratch) StreamHandler {
 	testConnectionHandler := &testConnectionHandler{
 		storage: storage,
 		context: context.Background(),
 		patterns: patterns,
 	}
 
-	scratch, err := hyperscan.NewScratch(patterns)
-	require.Nil(t, err)
-
 	srcIp := layers.NewIPEndpoint(net.ParseIP(testSrcIp))
 	dstIp := layers.NewIPEndpoint(net.ParseIP(testDstIp))
 	srcPort := layers.NewTCPPortEndpoint(srcPort)