From 11d61dcb4d518f7ba33ae6f3817a279cee36398d Mon Sep 17 00:00:00 2001
From: Leonid Bugaev <leonsbox@gmail.com>
Date: Fri, 6 Aug 2021 21:24:18 +0300
Subject: [PATCH] Make packet proccessing multi threaded

If packet processign is not fast enough, packet capture may miss packets
Now it use 10 goroutines, which distribute work based on the ephemeral port number
---
 tcp/tcp_message.go | 80 +++++++++++++++++++++++++++++++++-------------
 1 file changed, 57 insertions(+), 23 deletions(-)

diff --git a/tcp/tcp_message.go b/tcp/tcp_message.go
index 09d1503..4dab30c 100644
--- a/tcp/tcp_message.go
+++ b/tcp/tcp_message.go
@@ -7,6 +7,7 @@ import (
 	"net"
 	"reflect"
 	"sort"
+	"sync"
 	"time"
 	"unsafe"
 )
@@ -64,6 +65,7 @@ type Message struct {
 	packets  []*Packet
 	parser   *MessageParser
 	feedback interface{}
+	Idx      uint16
 	Stats
 }
 
@@ -198,7 +200,8 @@ type HintStart func(*Packet) (IsRequest, IsOutgoing bool)
 // MessageParser holds data of all tcp messages in progress(still receiving/sending packets).
 // message is identified by its source port and dst port, and last 4bytes of src IP.
 type MessageParser struct {
-	m map[uint64]*Message
+	m  []map[uint64]*Message
+	mL []sync.RWMutex
 
 	messageExpire  time.Duration // the maximum time to wait for the final packet, minimum is 100ms
 	allowIncompete bool
@@ -226,18 +229,24 @@ func NewMessageParser(messages chan *Message, ports []uint16, ips []net.IP, mess
 	parser.packets = make(chan *PcapPacket, 10000)
 
 	if messages == nil {
-		messages = make(chan *Message, 1000)
+		messages = make(chan *Message, 100)
 	}
 	parser.messages = messages
-
-	parser.m = make(map[uint64]*Message)
 	parser.ticker = time.NewTicker(time.Millisecond * 100)
 	parser.close = make(chan struct{}, 1)
 
 	parser.ports = ports
 	parser.ips = ips
 
-	go parser.wait()
+	for i := 0; i < 10; i++ {
+		parser.m = append(parser.m, make(map[uint64]*Message))
+		parser.mL = append(parser.mL, sync.RWMutex{})
+	}
+
+	for i := 0; i < 10; i++ {
+		go parser.wait(i)
+	}
+
 	return parser
 }
 
@@ -249,7 +258,7 @@ func (parser *MessageParser) PacketHandler(packet *PcapPacket) {
 	parser.packets <- packet
 }
 
-func (parser *MessageParser) wait() {
+func (parser *MessageParser) wait(index int) {
 	var (
 		now time.Time
 	)
@@ -258,7 +267,7 @@ func (parser *MessageParser) wait() {
 		case pckt := <-parser.packets:
 			parser.processPacket(parser.parsePacket(pckt))
 		case now = <-parser.ticker.C:
-			parser.timer(now)
+			parser.timer(now, index)
 		case <-parser.close:
 			parser.ticker.Stop()
 			// parser.Close should wait for this function to return
@@ -298,19 +307,28 @@ func (parser *MessageParser) processPacket(pckt *Packet) {
 
 	// Trying to build unique hash, but there is small chance of collision
 	// No matter if it is request or response, all packets in the same message have same
-	m, ok := parser.m[pckt.MessageID()]
+	mID := pckt.MessageID()
+	mIDX := pckt.SrcPort % 10
+
+	parser.mL[mIDX].Lock()
+	m, ok := parser.m[mIDX][mID]
+	if !ok {
+		parser.mL[mIDX].Unlock()
+
+		mIDX = pckt.DstPort % 10
+		parser.mL[mIDX].Lock()
+		m, ok = parser.m[mIDX][mID]
+
+		if !ok {
+			parser.mL[mIDX].Unlock()
+		}
+	}
+
 	switch {
 	case ok:
-		if m.Direction == DirUnknown {
-			if in, out := parser.Start(pckt); in || out {
-				if in {
-					m.Direction = DirIncoming
-				} else {
-					m.Direction = DirOutcoming
-				}
-			}
-		}
 		parser.addPacket(m, pckt)
+
+		parser.mL[mIDX].Unlock()
 		return
 	case pckt.Direction == DirUnknown && parser.Start != nil:
 		if in, out := parser.Start(pckt); in || out {
@@ -322,12 +340,25 @@ func (parser *MessageParser) processPacket(pckt *Packet) {
 		}
 	}
 
+	if pckt.Direction == DirIncoming {
+		mIDX = pckt.SrcPort % 10
+	} else {
+		mIDX = pckt.DstPort % 10
+	}
+
+	parser.mL[mIDX].Lock()
+
 	m = new(Message)
 	m.Direction = pckt.Direction
-	parser.m[pckt.MessageID()] = m
+
+	parser.m[mIDX][mID] = m
+
+	m.Idx = mIDX
 	m.Start = pckt.Timestamp
 	m.parser = parser
 	parser.addPacket(m, pckt)
+
+	parser.mL[mIDX].Unlock()
 }
 
 func (parser *MessageParser) addPacket(m *Message, pckt *Packet) bool {
@@ -354,7 +385,7 @@ func (parser *MessageParser) Read() *Message {
 func (parser *MessageParser) Emit(m *Message) {
 	stats.Add("message_count", 1)
 
-	delete(parser.m, m.packets[0].MessageID())
+	delete(parser.m[m.Idx], m.packets[0].MessageID())
 
 	parser.messages <- m
 }
@@ -365,13 +396,14 @@ func GetUnexportedField(field reflect.Value) interface{} {
 
 var failMsg int
 
-func (parser *MessageParser) timer(now time.Time) {
+func (parser *MessageParser) timer(now time.Time, index int) {
 	packetLen = 0
+	parser.mL[index].Lock()
 
 	packetQueueLen.Set(int64(len(parser.packets)))
-	messageQueueLen.Set(int64(len(parser.m)))
+	messageQueueLen.Set(int64(len(parser.m[index])))
 
-	for _, m := range parser.m {
+	for _, m := range parser.m[index] {
 		if now.Sub(m.End) > parser.messageExpire {
 			m.TimedOut = true
 			stats.Add("message_timeout_count", 1)
@@ -380,9 +412,11 @@ func (parser *MessageParser) timer(now time.Time) {
 				parser.Emit(m)
 			}
 
-			delete(parser.m, m.packets[0].MessageID())
+			delete(parser.m[index], m.packets[0].MessageID())
 		}
 	}
+
+	parser.mL[index].Unlock()
 }
 
 func (parser *MessageParser) Close() error {