From e310fc8369899f153eef844140c56bfa94ec658b Mon Sep 17 00:00:00 2001
From: Leonid Bugaev <leonsbox@gmail.com>
Date: Sun, 3 Sep 2017 22:11:40 +0500
Subject: [PATCH] Fix zero-length 100-continue packets

Should fix https://github.com/buger/goreplay/issues/496
---
 Makefile                             |  2 +-
 raw_socket_listener/listener.go      |  9 ++-
 raw_socket_listener/listener_test.go | 89 ++++++++++++++++++++++++++--
 3 files changed, 93 insertions(+), 7 deletions(-)

diff --git a/Makefile b/Makefile
index 251d2ce..20c9538 100644
--- a/Makefile
+++ b/Makefile
@@ -83,7 +83,7 @@ file-server:
 	go run $(SOURCE) file-server $(FADDR)
 
 readpcap:
-	go run $(SOURCE) --input-raw $(FILE) --input-raw-engine pcap_file --output-null
+	go run $(SOURCE) --input-raw $(FILE) --input-raw-track-response --input-raw-engine pcap_file --output-stdout
 
 record:
 	$(RUN) go run $(SOURCE) --input-dummy=0 --output-file=requests.gor --verbose --debug
diff --git a/raw_socket_listener/listener.go b/raw_socket_listener/listener.go
index f294278..d8a1478 100644
--- a/raw_socket_listener/listener.go
+++ b/raw_socket_listener/listener.go
@@ -699,12 +699,18 @@ func (t *Listener) processTCPPacket(packet *TCPPacket) {
 
 	// Seek for 100-expect chunks
 	if parentAck, ok := t.seqWithData[packet.Seq]; ok {
+		// Skip zero-length chunks https://github.com/buger/goreplay/issues/496
+		if len(packet.Data) == 0 {
+			return
+		}
+
 		// In case if non-first data chunks comes first
 		for _, m := range t.messages {
 			if m.Ack == packet.Ack && bytes.Equal(m.packets[0].Addr, packet.Addr) {
 				t.deleteMessage(m)
 
 				if m.AssocMessage != nil {
+					m.AssocMessage.setAssocMessage(nil)
 					m.setAssocMessage(nil)
 				}
 
@@ -771,8 +777,9 @@ func (t *Listener) processTCPPacket(packet *TCPPacket) {
 				t.deleteMessage(m)
 				if m.AssocMessage != nil {
 					message.setAssocMessage(m.AssocMessage)
+					m.AssocMessage.setAssocMessage(nil)
 				}
-				// log.Println("2: Adding ack alias:", m.Ack, packet.Ack)
+
 				t.ackAliases[m.Ack] = packet.Ack
 
 				for _, pkt := range m.packets {
diff --git a/raw_socket_listener/listener_test.go b/raw_socket_listener/listener_test.go
index fc0408a..dbacdcd 100644
--- a/raw_socket_listener/listener_test.go
+++ b/raw_socket_listener/listener_test.go
@@ -46,6 +46,90 @@ func TestRawListenerInput(t *testing.T) {
 	}
 }
 
+func TestSingleAck100Continue(t *testing.T) {
+	listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, "")
+	defer listener.Close()
+
+	reqPacket1 := buildPacket(
+		true,
+		1,
+		1,
+		[]byte("POST / HTTP/1.1\r\nExpect: 100-continue\r\nContent-Length: 4\r\n\r\n"), time.Now())
+
+	respPacket1 := buildPacket(false,
+		uint32(len(reqPacket1.Data)) + reqPacket1.Seq,
+		1,
+		[]byte(""), time.Now())
+
+	respPacket2 := buildPacket( false,
+		uint32(len(reqPacket1.Data)) + reqPacket1.Seq,
+		1,
+		[]byte("HTTP/1.1 100 Continue\r\n"), time.Now())
+
+	reqPacket3 := buildPacket(true,
+		uint32(len(reqPacket1.Data)) + respPacket1.Seq,
+		reqPacket1.Seq+uint32(len(reqPacket1.Data)),
+		[]byte("DATA"), time.Now())
+
+	respPacket3 := buildPacket(false,
+		uint32(len(reqPacket3.Data)) + reqPacket3.Seq,
+		respPacket1.Seq+uint32(len(respPacket1.Data)), []byte("HTTP/1.1 200 OK\r\n\r\n"), time.Now())
+
+	result := []byte("POST / HTTP/1.1\r\nContent-Length: 4\r\n\r\nDATA")
+
+	testRawListener100Continue(t, listener, result,
+		reqPacket1,
+		respPacket1, respPacket2,
+		reqPacket3,
+		respPacket3 )
+}
+
+
+func TestDoubleAck100Continue(t *testing.T) {
+	listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, "")
+	defer listener.Close()
+
+	reqPacket1 := buildPacket(
+		true,
+		1,
+		1,
+		[]byte("POST / HTTP/1.1\r\nExpect: 100-continue\r\nContent-Length: 4\r\n\r\n"), time.Now())
+
+	respPacket1 := buildPacket(false,
+		uint32(len(reqPacket1.Data)) + reqPacket1.Seq,
+		1,
+		[]byte(""), time.Now())
+
+	respPacket2 := buildPacket( false,
+		uint32(len(reqPacket1.Data)) + reqPacket1.Seq,
+		1,
+		[]byte("HTTP/1.1 100 Continue\r\n"), time.Now())
+
+	reqPacket2 := buildPacket(true,
+		uint32(len(reqPacket1.Data)) + respPacket1.Seq,
+		reqPacket1.Seq+uint32(len(reqPacket1.Data)),
+		[]byte(""), time.Now())
+
+	reqPacket3 := buildPacket(true,
+		uint32(len(reqPacket1.Data)) + respPacket1.Seq,
+		reqPacket1.Seq+uint32(len(reqPacket1.Data)),
+		[]byte("DATA"), time.Now())
+
+	respPacket3 := buildPacket(false,
+		uint32(len(reqPacket3.Data)) + reqPacket3.Seq,
+		respPacket1.Seq+uint32(len(respPacket1.Data)),
+		[]byte("HTTP/1.1 200 OK\r\n\r\n"), time.Now())
+
+	result := []byte("POST / HTTP/1.1\r\nContent-Length: 4\r\n\r\nDATA")
+
+	testRawListener100Continue(t, listener, result,
+		reqPacket1,
+		respPacket1, respPacket2,
+		reqPacket2, reqPacket3,
+		respPacket3 )
+}
+
+
 func TestRawListenerInputResponseByClose(t *testing.T) {
 	var req, resp *TCPMessage
 
@@ -365,11 +449,6 @@ func TestRawListenerChunkedWrongOrder(t *testing.T) {
 
 	// Should re-construct message from all possible combinations
 	for i := 0; i < 6*5*4*3*2*1; i++ {
-
-		// if i < 54 || i > 57 {
-		// 	continue
-		// }
-
 		packets := permutation(i, []*TCPPacket{reqPacket1, reqPacket2, reqPacket3, reqPacket4, respPacket1, respPacket2})
 
 		t.Log("permutation:", i)