From 56fb63868006268ce32ccd9a0f64b1dcbe65ec70 Mon Sep 17 00:00:00 2001
From: Emiliano Ciavatta
Date: Wed, 16 Dec 2020 15:28:25 +0100
Subject: Fix crashes in stream handler

---
 docker-compose.yml |  2 +-
 pcap_importer.go   |  3 ++-
 stream_handler.go  | 10 ++++++----
 3 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/docker-compose.yml b/docker-compose.yml
index de200a0..b43ec5b 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -15,7 +15,7 @@ services:
       - mongo
     networks:
       - acheronet
-    command: ./caronte -mongo-host mongo -mongo-port 27017
+    command: ./caronte -mongo-host mongo -mongo-port 27017 -assembly_memuse_log
 
 networks:
   acheronet:
diff --git a/pcap_importer.go b/pcap_importer.go
index 1c80b3f..476d8d7 100644
--- a/pcap_importer.go
+++ b/pcap_importer.go
@@ -204,7 +204,8 @@ func (pi *PcapImporter) parsePcap(session ImportingSession, fileName string, flu
 		case packet := <-packets:
 			if packet == nil { // completed
 				if flushAll {
-					assembler.FlushAll()
+					connectionsClosed := assembler.FlushAll()
+					log.Debugf("connections closed after flush: %v", connectionsClosed)
 				}
 				handle.Close()
 				pi.releaseAssembler(assembler)
diff --git a/stream_handler.go b/stream_handler.go
index f08bd70..8042bcb 100644
--- a/stream_handler.go
+++ b/stream_handler.go
@@ -80,20 +80,22 @@ func NewStreamHandler(connection ConnectionHandler, streamFlow StreamFlow, scann
 func (sh *StreamHandler) Reassembled(reassembly []tcpassembly.Reassembly) {
 	for _, r := range reassembly {
 		skip := r.Skip
+		isLoss := skip != 0
+
 		if r.Start {
 			sh.firstPacketSeen = r.Seen
 		}
 		if r.End {
 			sh.lastPacketSeen = r.Seen
 		}
-		if skip < 0 { // start or flush ~ workaround
-			skip = 0
-		}
 
 		reassemblyLen := len(r.Bytes)
 		if reassemblyLen == 0 {
 			continue
 		}
+		if skip < 0 || skip >= reassemblyLen { // start or flush ~ workaround
+			skip = 0
+		}
 
 		if sh.buffer.Len()+len(r.Bytes)-skip > MaxDocumentSize {
 			sh.storageCurrentDocument()
@@ -106,7 +108,7 @@ func (sh *StreamHandler) Reassembled(reassembly []tcpassembly.Reassembly) {
 		}
 		sh.indexes = append(sh.indexes, sh.currentIndex)
 		sh.timestamps = append(sh.timestamps, r.Seen)
-		sh.lossBlocks = append(sh.lossBlocks, skip != 0)
+		sh.lossBlocks = append(sh.lossBlocks, isLoss)
 		sh.currentIndex += n
 		sh.streamLength += n
 
-- 
cgit v1.2.3-70-g09d2