Add FillWithMatchedRules method

5 files changed, 21 insertions, 2 deletions

diff --git a/connection_handler.go b/connection_handler.go
index 4fc2a95..de7f508 100644
--- a/connection_handler.go
+++ b/connection_handler.go
@@ -35,6 +35,7 @@ type ConnectionHandler interface {
 	Complete(handler *StreamHandler)
 	Storage() Storage
 	PatternsDatabase() hyperscan.StreamDatabase
+	PatternsDatabaseSize() int
 }
 
 type connectionHandlerImpl struct {
@@ -203,6 +204,8 @@ func (ch *connectionHandlerImpl) Complete(handler *StreamHandler) {
 		ServerDocuments: len(server.documentsIDs),
 		ProcessedAt:     time.Now(),
 	}
+	ch.factory.rulesManager.FillWithMatchedRules(&connection, client.patternMatches, server.patternMatches)
+
 	_, err := ch.Storage().Insert(Connections).One(connection)
 	if err != nil {
 		log.WithError(err).WithField("connection", connection).Error("failed to insert a connection")
@@ -228,6 +231,10 @@ func (ch *connectionHandlerImpl) PatternsDatabase() hyperscan.StreamDatabase {
 	return ch.factory.rulesDatabase.database
 }
 
+func (ch *connectionHandlerImpl) PatternsDatabaseSize() int {
+	return ch.factory.rulesDatabase.databaseSize
+}
+
 func (sf StreamFlow) Hash() uint64 {
 	return sf[0].FastHash() ^ sf[1].FastHash() ^ sf[2].FastHash() ^ sf[3].FastHash()
 }
diff --git a/connections.go b/connections.go
index e3adaf2..3c22a9a 100644
--- a/connections.go
+++ b/connections.go
@@ -17,4 +17,5 @@ type Connection struct {
 	ClientDocuments int       `json:"client_documents" bson:"client_documents"`
 	ServerDocuments int       `json:"server_documents" bson:"server_documents"`
 	ProcessedAt     time.Time `json:"processed_at" bson:"processed_at"`
+	MatchedRules    []RowID   `json:"matched_rules" bson:"matched_rules"`
 }
diff --git a/rules_manager.go b/rules_manager.go
index 69439e0..e358fed 100644
--- a/rules_manager.go
+++ b/rules_manager.go
@@ -54,6 +54,7 @@ type Rule struct {
 
 type RulesDatabase struct {
 	database hyperscan.StreamDatabase
+	databaseSize int
 	version  RowID
 }
 
@@ -160,11 +161,17 @@ func (rm RulesManager) generateDatabase(version RowID) error {
 
 	rm.databaseUpdated <- RulesDatabase{
 		database: database,
+		databaseSize: len(patterns),
 		version:  version,
 	}
 	return nil
 }
 
+func (rm RulesManager) FillWithMatchedRules(connection *Connection, clientMatches map[uint][]PatternSlice,
+	serverMatches map[uint][]PatternSlice) {
+
+}
+
 func (p Pattern) BuildPattern() error {
 	if p.compiledPattern != nil {
 		return nil
diff --git a/stream_handler.go b/stream_handler.go
index 2d80f60..f2309ad 100644
--- a/stream_handler.go
+++ b/stream_handler.go
@@ -44,7 +44,7 @@ func NewStreamHandler(connection ConnectionHandler, streamFlow StreamFlow, scann
 		timestamps:     make([]time.Time, 0, InitialBlockCount),
 		lossBlocks:     make([]bool, 0, InitialBlockCount),
 		documentsIDs:   make([]RowID, 0, 1),               // most of the time the stream fit in one document
-		patternMatches: make(map[uint][]PatternSlice, 10), // TODO: change with exactly value
+		patternMatches: make(map[uint][]PatternSlice, connection.PatternsDatabaseSize()),
 		scanner:        scanner,
 	}
 
@@ -127,7 +127,7 @@ func (sh *StreamHandler) resetCurrentDocument() {
 	}
 }
 
-func (sh *StreamHandler) onMatch(id uint, from uint64, to uint64, flags uint, context interface{}) error {
+func (sh *StreamHandler) onMatch(id uint, from uint64, to uint64, _ uint, _ interface{}) error {
 	patternSlices, isPresent := sh.patternMatches[id]
 	if isPresent {
 		if len(patternSlices) > 0 {
diff --git a/stream_handler_test.go b/stream_handler_test.go
index 962048f..254f8c3 100644
--- a/stream_handler_test.go
+++ b/stream_handler_test.go
@@ -332,6 +332,10 @@ func (tch *testConnectionHandler) PatternsDatabase() hyperscan.StreamDatabase {
 	return tch.patterns
 }
 
+func (tch *testConnectionHandler) PatternsDatabaseSize() int {
+	return 8
+}
+
 func (tch *testConnectionHandler) Complete(handler *StreamHandler) {
 	tch.onComplete(handler)
 }