From 46474eeb425ee571e4510f6d63e9431a46291727 Mon Sep 17 00:00:00 2001
From: webhe <webhe@tencent.com>
Date: Sun, 28 Jul 2024 12:02:19 +0800
Subject: [PATCH] Use the PSH bit in the tcp protocol to solve the problem of
 capturing the full request/response body when there is no content-length or
 chunked field in the request/response header

---
 internal/capture/capture.go |  4 ++++
 internal/tcp/tcp_message.go | 16 +++++++++++-----
 internal/tcp/tcp_packet.go  | 27 ++++++++++++++-------------
 3 files changed, 29 insertions(+), 18 deletions(-)

diff --git a/internal/capture/capture.go b/internal/capture/capture.go
index b2a31f1d8..451b01f12 100644
--- a/internal/capture/capture.go
+++ b/internal/capture/capture.go
@@ -513,6 +513,10 @@ func http1StartHint(pckt *tcp.Packet) (isRequest, isResponse bool) {
 }
 
 func http1EndHint(m *tcp.Message) bool {
+	if len(m.Packets()) == 0 {
+		return false
+	}
+
 	if m.MissingChunk() {
 		return false
 	}
diff --git a/internal/tcp/tcp_message.go b/internal/tcp/tcp_message.go
index e2a2188aa..9c43f4da4 100644
--- a/internal/tcp/tcp_message.go
+++ b/internal/tcp/tcp_message.go
@@ -279,7 +279,7 @@ func (parser *MessageParser) wait() {
 }
 
 func (parser *MessageParser) parsePacket(pcapPkt *PcapPacket) *Packet {
-	pckt, err := ParsePacket(pcapPkt.Data, pcapPkt.LType, pcapPkt.LTypeLen, pcapPkt.Ci, false)
+	pckt, err := ParsePacket(pcapPkt.Data, pcapPkt.LType, pcapPkt.LTypeLen, pcapPkt.Ci, true)
 	if err != nil {
 		if _, empty := err.(EmptyPacket); !empty {
 			stats.Add("packet_error", 1)
@@ -313,7 +313,7 @@ func containsOrEmpty(element net.IP, ipList []net.IP) bool {
 }
 
 func (parser *MessageParser) processPacket(pckt *Packet) {
-	if pckt == nil {
+	if pckt == nil || (len(pckt.Payload) == 0 && !pckt.FIN) {
 		return
 	}
 
@@ -356,14 +356,14 @@ func (parser *MessageParser) processPacket(pckt *Packet) {
 }
 
 func (parser *MessageParser) addPacket(m *Message, pckt *Packet) bool {
-	if !m.add(pckt) {
+	if !pckt.FIN && len(pckt.Payload) == 0 || (len(pckt.Payload) != 0 && !m.add(pckt)) {
 		return false
 	}
 
 	// If we are using protocol parsing, like HTTP, depend on its parsing func.
 	// For the binary procols wait for message to expire
 	if parser.End != nil {
-		if parser.End(m) {
+		if parser.End(m) || pckt.FIN || (pckt.Direction == DirIncoming && pckt.PSH && pckt.ACK) {
 			parser.Emit(m)
 			return true
 		}
@@ -405,6 +405,10 @@ func (parser *MessageParser) Read() *Message {
 func (parser *MessageParser) Emit(m *Message) {
 	stats.Add("message_count", 1)
 
+	if len(m.packets) == 0 {
+		return
+	}
+
 	delete(parser.m, m.packets[0].MessageID())
 
 	parser.messages <- m
@@ -430,7 +434,9 @@ func (parser *MessageParser) timer(now time.Time) {
 			if parser.End == nil || parser.allowIncompete {
 				parser.Emit(m)
 			}
-
+			if len(m.packets) == 0 {
+				continue
+			}
 			delete(parser.m, m.packets[0].MessageID())
 		}
 	}
diff --git a/internal/tcp/tcp_packet.go b/internal/tcp/tcp_packet.go
index 43fb99425..e1037341e 100644
--- a/internal/tcp/tcp_packet.go
+++ b/internal/tcp/tcp_packet.go
@@ -57,19 +57,19 @@ calllers must make sure that ParsePacket has'nt returned any error before callin
 function.
 */
 type Packet struct {
-	Direction          Dir
-	messageID          uint64
-	SrcIP, DstIP       net.IP
-	Version            uint8
-	SrcPort, DstPort   uint16
-	Ack, Seq           uint32
-	ACK, SYN, FIN, RST bool
-	Lost               uint32
-	Retry              int
-	CaptureLength      int
-	Timestamp          time.Time
-	Payload            []byte
-	buf                []byte
+	Direction               Dir
+	messageID               uint64
+	SrcIP, DstIP            net.IP
+	Version                 uint8
+	SrcPort, DstPort        uint16
+	Ack, Seq                uint32
+	ACK, SYN, FIN, RST, PSH bool
+	Lost                    uint32
+	Retry                   int
+	CaptureLength           int
+	Timestamp               time.Time
+	Payload                 []byte
+	buf                     []byte
 
 	created time.Time
 	gc      bool
@@ -204,6 +204,7 @@ func (pckt *Packet) parse(data []byte, lType, lTypeLen int, cp *gopacket.Capture
 	pckt.FIN = transLayer[13]&0x01 != 0
 	pckt.SYN = transLayer[13]&0x02 != 0
 	pckt.RST = transLayer[13]&0x04 != 0
+	pckt.PSH = transLayer[13]&0x08 != 0
 	pckt.ACK = transLayer[13]&0x10 != 0
 	pckt.Lost = uint32(cp.Length - cp.CaptureLength)