From fc0361dcdcbc41541ba6378f63197e37d8a38939 Mon Sep 17 00:00:00 2001
From: Leonid Bugaev <leonsbox@gmail.com>
Date: Wed, 31 Aug 2016 16:09:39 +0300
Subject: [PATCH] Fix handling of connection: close for POST requests

---
 raw_socket_listener/listener.go    | 27 ++++++++++++++----
 raw_socket_listener/tcp_message.go | 45 ++++++++++++++++--------------
 2 files changed, 45 insertions(+), 27 deletions(-)

diff --git a/raw_socket_listener/listener.go b/raw_socket_listener/listener.go
index f85597b3..8e035e2f 100644
--- a/raw_socket_listener/listener.go
+++ b/raw_socket_listener/listener.go
@@ -115,8 +115,6 @@ func NewListener(addr string, port string, engine int, trackResponse bool, expir
 	// Special case for testing
 	if l.port != 0 {
 		switch engine {
-		case EngineRawSocket:
-			go l.readRAWSocket()
 		case EnginePcap:
 			go l.readPcap()
 		case EnginePcapFile:
@@ -559,7 +557,14 @@ func (t *Listener) readPcapFile() {
 			if tcpLayer := packet.Layer(layers.LayerTypeTCP); tcpLayer != nil {
 				tcp, _ := tcpLayer.(*layers.TCP)
 				data = append(tcp.LayerContents(), tcp.LayerPayload()...)
-				copy(data[2:4], []byte{0, 1})
+
+				if tcp.SrcPort >= 32768 && tcp.SrcPort <= 61000 {
+					copy(data[0:2], []byte{0, 0})
+					copy(data[2:4], []byte{0, 1})
+				} else {
+					copy(data[0:2], []byte{0, 1})
+					copy(data[2:4], []byte{0, 0})
+				}
 			} else {
 				continue
 			}
@@ -576,10 +581,11 @@ func (t *Listener) readPcapFile() {
 			}
 
 			dataOffset := (data[12] & 0xF0) >> 4
+			isFIN := data[13]&0x01 != 0
 
 			// We need only packets with data inside
 			// Check that the buffer is larger than the size of the TCP header
-			if len(data) <= int(dataOffset*4) {
+			if len(data) <= int(dataOffset*4) && !isFIN {
 				continue
 			}
 
@@ -663,8 +669,6 @@ func (t *Listener) processTCPPacket(packet *TCPPacket) {
 		}
 	}()
 
-	// log.Println("Processing packet:", packet.Ack, packet.Seq, packet.ID)
-
 	var message *TCPMessage
 
 	isIncoming := packet.DestPort == t.port
@@ -693,6 +697,14 @@ func (t *Listener) processTCPPacket(packet *TCPPacket) {
 		packet.UpdateAck(parentAck)
 	}
 
+	if isIncoming && packet.IsFIN {
+		if ma, ok := t.respAliases[packet.Seq]; ok {
+			if ma.packets[0].SrcPort == packet.SrcPort {
+				packet.UpdateAck(ma.Ack)
+			}
+		}
+	}
+
 	if alias, ok := t.ackAliases[packet.Ack]; ok {
 		packet.UpdateAck(alias)
 	}
@@ -764,8 +776,11 @@ func (t *Listener) processTCPPacket(packet *TCPPacket) {
 
 	// If message contains only single packet immediately dispatch it
 	if message.complete {
+		// log.Println("COMPLETE!", isIncoming, message)
 		if isIncoming {
 			if t.trackResponse {
+				// log.Println("Found response!", message.ResponseID, t.messages)
+
 				if resp, ok := t.messages[message.ResponseID]; ok {
 					if resp.complete {
 						t.dispatchMessage(resp)
diff --git a/raw_socket_listener/tcp_message.go b/raw_socket_listener/tcp_message.go
index 7a5185cc..5cb566e0 100644
--- a/raw_socket_listener/tcp_message.go
+++ b/raw_socket_listener/tcp_message.go
@@ -352,6 +352,26 @@ func (t *TCPMessage) updateBodyType() {
 		return
 	}
 
+	var lengthB, encB, connB []byte
+
+	proto.ParseHeaders(t.packetsData(), func(header, value []byte)bool{
+		if proto.HeadersEqual(header, []byte("Content-Length")) {
+			lengthB = value
+			return false
+		}
+
+		if proto.HeadersEqual(header, []byte("Transfer-Encoding")) {
+			encB = value
+			return false
+		}
+
+		if proto.HeadersEqual(header, []byte("Connection")) {
+			connB = value
+		}
+
+		return true
+	})
+
 	switch t.methodType {
 	case httpMethodNotFound:
 		return
@@ -359,27 +379,6 @@ func (t *TCPMessage) updateBodyType() {
 		t.bodyType = httpBodyEmpty
 		return
 	case httpMethodWithBody:
-		var lengthB, encB, connB []byte
-
-		proto.ParseHeaders(t.packetsData(), func(header, value []byte)bool{
-			if proto.HeadersEqual(header, []byte("Content-Length")) {
-				lengthB = value
-				return false
-			}
-
-			if proto.HeadersEqual(header, []byte("Transfer-Encoding")) {
-				encB = value
-				return false
-			}
-
-			if proto.HeadersEqual(header, []byte("Connection")) {
-				connB = value
-				return false
-			}
-
-			return true
-		})
-
 		if len(lengthB) > 0 {
 			t.contentLength, _ = strconv.Atoi(string(lengthB))
 
@@ -461,6 +460,10 @@ func (t *TCPMessage) setAssocMessage(m *TCPMessage) {
 // UpdateResponseAck should be called after packet is added
 func (t *TCPMessage) UpdateResponseAck() uint32 {
 	lastPacket := t.packets[len(t.packets)-1]
+	if lastPacket.IsFIN && len(t.packets) > 1 {
+		lastPacket = t.packets[len(t.packets)-2]
+	}
+
 	respAck := lastPacket.Seq + uint32(len(lastPacket.Data))
 
 	if t.ResponseAck != respAck {