From 6d33a0d51c3ce6b3a8fd59f086dddcc85d187360 Mon Sep 17 00:00:00 2001
From: Toby <tobyxdd@gmail.com>
Date: Sun, 11 Feb 2024 13:05:05 -0800
Subject: [PATCH] fix: incorrect verdict handling that caused packets to pass
 through even after they had been blocked (#52)

---
 engine/tcp.go | 17 +++++++++++++----
 engine/udp.go | 17 ++++++++++++-----
 2 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/engine/tcp.go b/engine/tcp.go
index f153f15..a4744dd 100644
--- a/engine/tcp.go
+++ b/engine/tcp.go
@@ -65,7 +65,7 @@ func (f *tcpStreamFactory) New(ipFlow, tcpFlow gopacket.Flow, tcp *layers.TCP, a
 		ctx.Verdict = tcpVerdictAcceptStream
 		f.Logger.TCPStreamAction(info, ruleset.ActionAllow, true)
 		// a tcpStream with no activeEntries is a no-op
-		return &tcpStream{}
+		return &tcpStream{finalVerdict: tcpVerdictAcceptStream}
 	}
 	// Create entries for each analyzer
 	entries := make([]*tcpStreamEntry, 0, len(ans))
@@ -109,6 +109,7 @@ type tcpStream struct {
 	ruleset       ruleset.Ruleset
 	activeEntries []*tcpStreamEntry
 	doneEntries   []*tcpStreamEntry
+	finalVerdict  tcpVerdict
 }
 
 type tcpStreamEntry struct {
@@ -119,8 +120,13 @@ type tcpStreamEntry struct {
 }
 
 func (s *tcpStream) Accept(tcp *layers.TCP, ci gopacket.CaptureInfo, dir reassembly.TCPFlowDirection, nextSeq reassembly.Sequence, start *bool, ac reassembly.AssemblerContext) bool {
-	// Only accept packets if we still have active entries
-	return len(s.activeEntries) > 0
+	if len(s.activeEntries) > 0 {
+		return true
+	} else {
+		ctx := ac.(*tcpContext)
+		ctx.Verdict = s.finalVerdict
+		return false
+	}
 }
 
 func (s *tcpStream) ReassembledSG(sg reassembly.ScatterGather, ac reassembly.AssemblerContext) {
@@ -152,7 +158,9 @@ func (s *tcpStream) ReassembledSG(sg reassembly.ScatterGather, ac reassembly.Ass
 		}
 		action := result.Action
 		if action != ruleset.ActionMaybe && action != ruleset.ActionModify {
-			ctx.Verdict = actionToTCPVerdict(action)
+			verdict := actionToTCPVerdict(action)
+			s.finalVerdict = verdict
+			ctx.Verdict = verdict
 			s.logger.TCPStreamAction(s.info, action, false)
 			// Verdict issued, no need to process any more packets
 			s.closeActiveEntries()
@@ -160,6 +168,7 @@ func (s *tcpStream) ReassembledSG(sg reassembly.ScatterGather, ac reassembly.Ass
 	}
 	if len(s.activeEntries) == 0 && ctx.Verdict == tcpVerdictAccept {
 		// All entries are done but no verdict issued, accept stream
+		s.finalVerdict = tcpVerdictAcceptStream
 		ctx.Verdict = tcpVerdictAcceptStream
 		s.logger.TCPStreamAction(s.info, ruleset.ActionAllow, true)
 	}
diff --git a/engine/udp.go b/engine/udp.go
index 227413b..e97c3ac 100644
--- a/engine/udp.go
+++ b/engine/udp.go
@@ -65,7 +65,7 @@ func (f *udpStreamFactory) New(ipFlow, udpFlow gopacket.Flow, udp *layers.UDP, u
 		uc.Verdict = udpVerdictAcceptStream
 		f.Logger.UDPStreamAction(info, ruleset.ActionAllow, true)
 		// a udpStream with no activeEntries is a no-op
-		return &udpStream{}
+		return &udpStream{finalVerdict: udpVerdictAcceptStream}
 	}
 	// Create entries for each analyzer
 	entries := make([]*udpStreamEntry, 0, len(ans))
@@ -167,6 +167,7 @@ type udpStream struct {
 	ruleset       ruleset.Ruleset
 	activeEntries []*udpStreamEntry
 	doneEntries   []*udpStreamEntry
+	finalVerdict  udpVerdict
 }
 
 type udpStreamEntry struct {
@@ -177,8 +178,12 @@ type udpStreamEntry struct {
 }
 
 func (s *udpStream) Accept(udp *layers.UDP, rev bool, uc *udpContext) bool {
-	// Only accept packets if we still have active entries
-	return len(s.activeEntries) > 0
+	if len(s.activeEntries) > 0 {
+		return true
+	} else {
+		uc.Verdict = s.finalVerdict
+		return false
+	}
 }
 
 func (s *udpStream) Feed(udp *layers.UDP, rev bool, uc *udpContext) {
@@ -221,16 +226,18 @@ func (s *udpStream) Feed(udp *layers.UDP, rev bool, uc *udpContext) {
 			}
 		}
 		if action != ruleset.ActionMaybe {
-			var final bool
-			uc.Verdict, final = actionToUDPVerdict(action)
+			verdict, final := actionToUDPVerdict(action)
+			uc.Verdict = verdict
 			s.logger.UDPStreamAction(s.info, action, false)
 			if final {
+				s.finalVerdict = verdict
 				s.closeActiveEntries()
 			}
 		}
 	}
 	if len(s.activeEntries) == 0 && uc.Verdict == udpVerdictAccept {
 		// All entries are done but no verdict issued, accept stream
+		s.finalVerdict = udpVerdictAcceptStream
 		uc.Verdict = udpVerdictAcceptStream
 		s.logger.UDPStreamAction(s.info, ruleset.ActionAllow, true)
 	}