Add handling for out of order packets

2024-04-21 12:32:02 +00:00 · 2021-05-03 21:41:11 +03:00
parent c7d04738c1
commit 32c005a63d
4 changed files with 139 additions and 36 deletions
@@ -18,6 +18,7 @@ type Packet struct {
 	LinkType int
 	Info     *gopacket.CaptureInfo
 	Err      error
+	Retry    int
 }

 // NewPacket parses packet and stop at the first error encountered
@@ -19,6 +19,7 @@ package proto
 import (
 	"bufio"
 	"bytes"
+	_ "fmt"
 	"net/http"
 	"net/textproto"
 	"strings"
@@ -471,12 +472,14 @@ func HasFullPayload(data []byte, m ProtocolStateSetter) bool {
 	if state.body < 1 {
 		state.body = MIMEHeadersEndPos(data)
 		if state.body < 0 {
+			// fmt.Println("SKIPPING BODY!")
 			return false
 		}
 	}
 	if !state.headerParsed {
 		chunked := Header(data, []byte("Transfer-Encoding"))
 		if len(chunked) > 0 && bytes.Index(data, []byte("chunked")) > 0 {
+			// fmt.Println("CHUNKED DETECTED!" + string(data))
 			state.isChunked = true
 			// trailers are generally not allowed in non-chunks body
 			state.hasTrailer = len(Header(data, []byte("Trailer"))) > 0
@@ -491,6 +494,7 @@ func HasFullPayload(data []byte, m ProtocolStateSetter) bool {
 		body = data[state.body:]
 	}
 	if state.isChunked {
+		// fmt.Println("CHUNKED!!!", string(body))
 		// check chunks
 		if len(body) < 1 {
 			return false
@@ -507,10 +511,15 @@ func HasFullPayload(data []byte, m ProtocolStateSetter) bool {
 		if !state.hasTrailer {
 			return true
 		}
+
+		// fmt.Printf("CHUNKED: bodyLen: %d, actualLen: %d\n, MIME: %v", state.bodyLen, len(body), MIMEHeadersEndPos(data[state.body:]))
+
 		// trailer headers(whether chunked or plain) should end with empty line
 		return len(data) > state.body && MIMEHeadersEndPos(data[state.body:]) != -1
 	}

+	// fmt.Printf("bodyLen: %d, actualLen: %d\n", state.bodyLen, len(body))
+
 	// check for content-length header
 	return state.bodyLen == len(body)
 }
@@ -1,7 +1,6 @@
 package tcp

 import (
-	"bytes"
 	"encoding/binary"
 	"encoding/hex"
 	"fmt"
@@ -30,7 +29,6 @@ type Stats struct {
 type Message struct {
 	packets  []*Packet
 	parser   *MessageParser
-	buf      *bytes.Buffer
 	feedback interface{}
 	Stats
 }
@@ -41,7 +39,6 @@ func NewMessage(srcAddr, dstAddr string, ipVersion uint8) (m *Message) {
 	m.DstAddr = dstAddr
 	m.SrcAddr = srcAddr
 	m.IPversion = ipVersion
-	m.buf = &bytes.Buffer{}
 	return
 }

@@ -81,12 +78,36 @@ func (m *Message) UUID() []byte {
 	return uuidHex
 }

-func (m *Message) add(pckt *Packet) {
-	m.Length += len(pckt.Payload)
-	m.LostData += int(pckt.Lost)
-	m.packets = append(m.packets, pckt)
-	m.End = pckt.Timestamp
-	m.buf.Write(pckt.Payload)
+func (m *Message) add(packet *Packet) {
+	// fmt.Println("SEQ:", packet.Seq, " - ", len(packet.Payload))
+
+	// Skip duplicates
+	for _, p := range m.packets {
+		if p.Seq == packet.Seq {
+			return
+		}
+	}
+
+	// Packets not always captured in same Seq order, and sometimes we need to prepend
+	if len(m.packets) == 0 || packet.Seq > m.packets[len(m.packets)-1].Seq {
+		m.packets = append(m.packets, packet)
+	} else if packet.Seq < m.packets[0].Seq {
+		m.packets = append([]*Packet{packet}, m.packets...)
+	} else { // insert somewhere in the middle...
+		for i, p := range m.packets {
+			if packet.Seq < p.Seq {
+				m.packets = append(m.packets[:i], append([]*Packet{packet}, m.packets[i:]...)...)
+				break
+			}
+		}
+	}
+
+	m.Length += len(packet.Payload)
+	m.LostData += int(packet.Lost)
+
+	if packet.Timestamp.After(m.End) || m.End.IsZero() {
+		m.End = packet.Timestamp
+	}
 }

 // Packets returns packets of the message
@@ -96,7 +117,20 @@ func (m *Message) Packets() []*Packet {

 // Data returns data in this message
 func (m *Message) Data() []byte {
-	return m.buf.Bytes()
+	var totalLen int
+	for _, p := range m.packets {
+		totalLen += len(p.Payload)
+	}
+	tmp := make([]byte, totalLen)
+
+	// fmt.Println("Total len:", totalLen, " - ", len(m.packets), " - ", m.Length)
+
+	var i int
+	for _, p := range m.packets {
+		i += copy(tmp[i:], p.Payload)
+	}
+
+	return tmp
 }

 // SetProtocolState set feedback/data that can be used later, e.g with End or Start hint
@@ -210,6 +244,8 @@ func (parser *MessageParser) parsePacket(packet *capture.Packet) {
 	key := uint64(pckt.SrcPort)<<48 | uint64(pckt.DstPort)<<32 |
 		uint64(_uint32(pckt.SrcIP[lst:]))
 	m, ok := parser.m[key]
+
+	// If connection is closed, finish both request and response
 	if pckt.RST {
 		if ok {
 			m.doDone(key)
@@ -227,10 +263,14 @@ func (parser *MessageParser) parsePacket(packet *capture.Packet) {
 	case ok:
 		parser.addPacket(key, m, pckt)
 		return
-	case pckt.SYN:
-		in = !pckt.ACK
 	case parser.Start != nil:
 		if in, out = parser.Start(pckt); !(in || out) {
+			// Packet can be received out of order, so give it another chance
+			if packet.Retry < 3 && len(packet.Payload) > 0 {
+				// Requeue not known packets
+				packet.Retry++
+				parser.packets <- packet
+			}
 			return
 		}
 	default:
@@ -13,7 +13,7 @@ import (
 	"github.com/google/gopacket/layers"
 )

-func generateHeader(seq uint32, length uint16) []byte {
+func generateHeader(request bool, seq uint32, length uint16) []byte {
 	hdr := make([]byte, 4+24+24)
 	binary.BigEndian.PutUint32(hdr, uint32(layers.ProtocolFamilyIPv4))

@@ -27,16 +27,22 @@ func generateHeader(seq uint32, length uint16) []byte {
 	// set tcp header
 	tcp := ip[24:]
 	tcp[12] = 6 << 4
-	binary.BigEndian.PutUint16(tcp, 5535)
-	binary.BigEndian.PutUint16(tcp[2:], 8000)
+
+	if request {
+		binary.BigEndian.PutUint16(tcp, 5535)
+		binary.BigEndian.PutUint16(tcp[2:], 8000)
+	} else {
+		binary.BigEndian.PutUint16(tcp, 8000)
+		binary.BigEndian.PutUint16(tcp[2:], 5535)
+	}
 	binary.BigEndian.PutUint32(tcp[4:], seq)
 	return hdr
 }

-func GetPackets(start uint32, _len int, payload []byte) []*capture.Packet {
+func GetPackets(request bool, start uint32, _len int, payload []byte) []*capture.Packet {
 	var packets = make([]*capture.Packet, _len)
 	for i := start; i < start+uint32(_len); i++ {
-		d := append(generateHeader(i, uint16(len(payload))), payload...)
+		d := append(generateHeader(request, i, uint16(len(payload))), payload...)
 		ci := &gopacket.CaptureInfo{Length: len(d), CaptureLength: len(d), Timestamp: time.Now()}
 		packets[i-start] = capture.NewPacket(d, int(layers.LinkTypeLoop), 4, ci)
 	}
@@ -52,16 +58,16 @@ func TestMessageParserWithHint(t *testing.T) {
 	parser.End = func(m *Message) bool {
 		return proto.HasFullPayload(m.Data(), m)
 	}
-	packets := GetPackets(1, 30, nil)
+	packets := GetPackets(true, 1, 30, nil)
 	packets[0].TransLayer[13] = 2 // SYN flag
-	packets[4] = GetPackets(1, 1, []byte("HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nTransfer-Encoding: chunked\r\n\r\n7"))[0]
-	packets[5] = GetPackets(1, 1, []byte("\r\nMozilla\r\n9\r\nDeveloper\r"))[0]
-	packets[6] = GetPackets(1, 1, []byte("\n7\r\nNetwork\r\n0\r\n\r\n"))[0]
+	packets[4] = GetPackets(false, 4, 1, []byte("HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nTransfer-Encoding: chunked\r\n\r\n7"))[0]
+	packets[5] = GetPackets(false, 5, 1, []byte("\r\nMozilla\r\n9\r\nDeveloper\r"))[0]
+	packets[6] = GetPackets(false, 6, 1, []byte("\n7\r\nNetwork\r\n0\r\n\r\n"))[0]
 	packets[10].TransLayer[13] = 2 // SYN flag
-	packets[14] = GetPackets(1, 1, []byte("POST / HTTP/1.1\r\nContent-Type: text/plain\r\nContent-Length: 23\r\n\r\n"))[0]
-	packets[15] = GetPackets(1, 1, []byte("MozillaDeveloper"))[0]
-	packets[16] = GetPackets(1, 1, []byte("Network"))[0]
-	packets[24] = GetPackets(1, 1, []byte("HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nContent-Length: 0\r\n\r"))[0]
+	packets[14] = GetPackets(true, 14, 1, []byte("POST / HTTP/1.1\r\nContent-Type: text/plain\r\nContent-Length: 23\r\n\r\n"))[0]
+	packets[15] = GetPackets(true, 15, 1, []byte("MozillaDeveloper"))[0]
+	packets[16] = GetPackets(true, 16, 1, []byte("Network"))[0]
+	packets[24] = GetPackets(true, 24, 1, []byte("HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nContent-Length: 0\r\n\r"))[0]
 	packets[29].TransLayer[13] = 1 // FIN flag

 	for i := 0; i < 30; i++ {
@@ -100,10 +106,57 @@ func TestMessageParserWithHint(t *testing.T) {

 }

+func TestMessageParserWrongOrder(t *testing.T) {
+	var mssg = make(chan *Message, 3)
+	parser := NewMessageParser(1<<20, time.Second, nil, func(m *Message) { mssg <- m })
+	parser.Start = func(pckt *Packet) (bool, bool) {
+		return proto.HasRequestTitle(pckt.Payload), proto.HasResponseTitle(pckt.Payload)
+	}
+	parser.End = func(m *Message) bool {
+		return proto.HasFullPayload(m.Data(), m)
+	}
+	packets := GetPackets(true, 1, 30, nil)
+	packets[0].TransLayer[13] = 2 // SYN flag
+	packets[6] = GetPackets(false, 4, 1, []byte("HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nTransfer-Encoding: chunked\r\n\r\n7"))[0]
+	packets[5] = GetPackets(false, 5, 1, []byte("\r\nMozilla\r\n9\r\nDeveloper\r"))[0]
+	packets[4] = GetPackets(false, 6, 1, []byte("\n7\r\nNetwork\r\n0\r\n\r\n"))[0]
+	// Duplicate with same seq
+	packets[7] = GetPackets(false, 6, 1, []byte("\n7\r\nNetwork\r\n0\r\n\r\n"))[0]
+
+	packets[10].TransLayer[13] = 2 // SYN flag
+	packets[16] = GetPackets(true, 14, 1, []byte("POST / HTTP/1.1\r\nContent-Type: text/plain\r\nContent-Length: 23\r\n\r\n"))[0]
+	packets[15] = GetPackets(true, 15, 1, []byte("MozillaDeveloper"))[0]
+	packets[14] = GetPackets(true, 16, 1, []byte("Network"))[0]
+	// packets[17].TransLayer[13] = 1 // FIN flag
+	for i := 0; i < 30; i++ {
+		parser.PacketHandler(packets[i])
+	}
+	var m *Message
+	select {
+	case <-time.After(time.Second):
+		t.Errorf("can't parse packets fast enough")
+		return
+	case m = <-mssg:
+	}
+	if !bytes.HasSuffix(m.Data(), []byte("\n7\r\nNetwork\r\n0\r\n\r\n")) {
+		t.Errorf("expected to %q to have suffix %q", m.Data(), []byte("\n7\r\nNetwork\r\n0\r\n\r\n"))
+	}
+
+	select {
+	case <-time.After(time.Second):
+		t.Errorf("can't parse packets fast enough")
+		return
+	case m = <-mssg:
+	}
+	if !bytes.HasSuffix(m.Data(), []byte("Network")) {
+		t.Errorf("expected to %q to have suffix %q", m.Data(), []byte("Network"))
+	}
+}
+
 func TestMessageParserWithoutHint(t *testing.T) {
 	var mssg = make(chan *Message, 1)
 	var data [63 << 10]byte
-	packets := GetPackets(1, 10, data[:])
+	packets := GetPackets(true, 1, 10, data[:])
 	packets[0].TransLayer[13] = 2 // SYN flag
 	packets[9].TransLayer[13] = 1 // FIN flag
 	p := NewMessageParser(63<<10*10, time.Second, nil, func(m *Message) { mssg <- m })
@@ -125,8 +178,8 @@ func TestMessageParserWithoutHint(t *testing.T) {
 func TestMessageMaxSizeReached(t *testing.T) {
 	var mssg = make(chan *Message, 2)
 	var data [63 << 10]byte
-	packets := GetPackets(1, 2, data[:])
-	packets = append(packets, GetPackets(1, 1, make([]byte, 63<<10+10))...)
+	packets := GetPackets(true, 1, 2, data[:])
+	packets = append(packets, GetPackets(true, 1, 1, make([]byte, 63<<10+10))...)
 	packets[0].TransLayer[13] = 2 // SYN flag
 	packets[2].TransLayer[13] = 2 // SYN flag
 	packets[2].NetLayer[15] = 3   // changing address
@@ -165,7 +218,7 @@ func TestMessageMaxSizeReached(t *testing.T) {
 func TestMessageTimeoutReached(t *testing.T) {
 	var mssg = make(chan *Message, 2)
 	var data [63 << 10]byte
-	packets := GetPackets(1, 2, data[:])
+	packets := GetPackets(true, 1, 2, data[:])
 	packets[0].TransLayer[13] = 2 // SYN flag
 	p := NewMessageParser(1<<20, 0, nil, func(m *Message) { mssg <- m })
 	p.PacketHandler(packets[0])
@@ -181,7 +234,7 @@ func TestMessageTimeoutReached(t *testing.T) {
 }

 func TestMessageUUID(t *testing.T) {
-	packets := GetPackets(1, 10, nil)
+	packets := GetPackets(true, 1, 10, nil)
 	packets[0].TransLayer[13] = 2 // SYN flag
 	packets[4].TransLayer[13] = 1 // FIN flag
 	packets[5].TransLayer[13] = 2 // SYN flag
@@ -205,7 +258,7 @@ func TestMessageUUID(t *testing.T) {
 }

 func BenchmarkMessageUUID(b *testing.B) {
-	packets := GetPackets(1, 5, nil)
+	packets := GetPackets(true, 1, 5, nil)
 	packets[0].TransLayer[13] = 2 // SYN flag
 	packets[4].TransLayer[13] = 1 // FIN flag
 	var uuid []byte
@@ -227,7 +280,7 @@ func BenchmarkMessageUUID(b *testing.B) {
 func BenchmarkPacketParseAndSort(b *testing.B) {
 	m := new(Message)
 	m.packets = make([]*Packet, 100)
-	for i, v := range GetPackets(1, 100, nil) {
+	for i, v := range GetPackets(true, 1, 100, nil) {
 		m.packets[i], _ = ParsePacket(v)
 	}
 	b.ResetTimer()
@@ -240,7 +293,7 @@ func BenchmarkMessageParserWithoutHint(b *testing.B) {
 	// runtime.GOMAXPROCS(8)
 	var mssg = make(chan *Message, 1)
 	var chunk = []byte("111111111111111111111111111111")
-	packets := GetPackets(1, 1000, chunk)
+	packets := GetPackets(true, 1, 1000, chunk)
 	packets[0].TransLayer[13] = 2      // SYN flag
 	packets[1000-1].TransLayer[13] = 1 // FIN flag
 	p := NewMessageParser(1<<20, time.Second*2, nil, func(m *Message) {
@@ -266,7 +319,7 @@ func BenchmarkMessageParserWithHint(b *testing.B) {
 	buf[1001] = []byte("0\r\n\r\n")
 	packets := make([]*capture.Packet, len(buf))
 	for i := 0; i < len(buf); i++ {
-		packets[i] = GetPackets(1, 1, buf[i])[0]
+		packets[i] = GetPackets(false, 1, 1, buf[i])[0]
 	}
 	var mssg = make(chan *Message, 1)
 	parser := NewMessageParser(1<<30, time.Second*10, nil, func(m *Message) { mssg <- m })
@@ -288,7 +341,7 @@ func BenchmarkMessageParserWithHint(b *testing.B) {
 }

 func BenchmarkNewAndParsePacket(b *testing.B) {
-	data := append(generateHeader(1024, 10), make([]byte, 10)...)
+	data := append(generateHeader(true, 1024, 10), make([]byte, 10)...)
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
 		ParsePacket(capture.NewPacket(data, int(layers.LinkTypeLoop), 4, &gopacket.CaptureInfo{}))
@@ -297,7 +350,7 @@ func BenchmarkNewAndParsePacket(b *testing.B) {

 func BenchmarkNewPacket(b *testing.B) {
 	packet := capture.NewPacket(
-		append(generateHeader(1024, 10), make([]byte, 10)...),
+		append(generateHeader(true, 1024, 10), make([]byte, 10)...),
 		int(layers.LinkTypeLoop),
 		4,
 		&gopacket.CaptureInfo{},