Properly handle requests with headers consisting of multiple packets

raw_socket_listener/tcp_message.go

@@ -122,6 +122,40 @@ func (t *TCPMessage) AddPacket(packet *TCPPacket) {
 	}
 }
 
+// Check if there is missing packet
+func (t *TCPMessage) isSeqMissing() bool {
+	if len(t.packets) == 1 {
+		return false
+	}
+
+	for i, p := range t.packets {
+		// If final packet
+		if len(t.packets) == i + 1 {
+			return false
+		}
+		np := t.packets[i + 1]
+
+		if np.Seq != p.Seq + uint32(len(p.Data)) {
+			return true
+		}
+	}
+
+	return false
+}
+
+var EmptyLine = []byte("\r\n\r\n")
+var ChunkEnd = []byte("0\r\n\r\n")
+
+func (t *TCPMessage) isHeadersReceived() bool {
+	for _, p := range t.packets {
+		if bytes.LastIndex(p.Data, EmptyLine) != -1 {
+			return true
+		}
+	}
+
+	return false
+}
+
 // isMultipart returns true if message contains from multiple tcp packets
 func (t *TCPMessage) IsFinished() bool {
 	payload := t.packets[0].Data
@@ -135,16 +169,23 @@ func (t *TCPMessage) IsFinished() bool {
 	if t.IsIncoming {
 		// If one GET, OPTIONS, or HEAD request
 		if bytes.Equal(m, []byte("GET ")) || bytes.Equal(m, []byte("OPTI")) || bytes.Equal(m, []byte("HEAD")) {
-			return true
+			if !t.isSeqMissing() && t.isHeadersReceived() {
+				return true
+			} else {
+				return false
+			}
 		} else {
 			// Sometimes header comes after the body :(
 			if bytes.Equal(m, []byte("POST")) || bytes.Equal(m, []byte("PUT ")) || bytes.Equal(m, []byte("PATC")) {
-				if length := proto.Header(payload, []byte("Content-Length")); len(length) > 0 {
-					l, _ := strconv.Atoi(string(length))
 
-					// If content-length equal current body length
-					if l > 0 && l == t.BodySize() {
-						return true
+				if t.isHeadersReceived() {
+					if length := proto.Header(payload, []byte("Content-Length")); len(length) > 0 {
+						l, _ := strconv.Atoi(string(length))
+
+						// If content-length equal current body length
+						if l > 0 && l == t.BodySize() {
+							return true
+						}
 					}
 				}
 			}
@@ -156,6 +197,10 @@ func (t *TCPMessage) IsFinished() bool {
 			return false
 		}
 
+		if !bytes.Equal(m, []byte("HTTP")) {
+			return false
+		}
+
 		if length := proto.Header(payload, []byte("Content-Length")); len(length) > 0 {
 			if length[0] == '0' {
 				return true
@@ -170,6 +215,10 @@ func (t *TCPMessage) IsFinished() bool {
 		} else {
 			if enc := proto.Header(payload, []byte("Transfer-Encoding")); len(enc) == 0 {
 				return true
+			} else {
+				if len(t.packets) > 1 && bytes.LastIndex(t.packets[len(t.packets) - 1].Data, ChunkEnd) != -1 {
+					return true
+				}
 			}
 		}
 	}

raw_socket_listener/tcp_message_test.go

@@ -84,7 +84,7 @@ func TestTCPMessageIsFinished(t *testing.T) {
 	methodsWithoutBodies := []string{"GET", "OPTIONS", "HEAD"}
 
 	for _, m := range methodsWithoutBodies {
-		msg := buildMessage(buildPacket(true, 1, 1, []byte(m+" / HTTP/1.1")))
+		msg := buildMessage(buildPacket(true, 1, 1, []byte(m+" / HTTP/1.1\r\n\r\n")))
 
 		if !msg.IsFinished() {
 			t.Error(m, " request should be finished")
@@ -153,3 +153,39 @@ func TestTCPMessageIsFinished(t *testing.T) {
 		t.Error("Should not mark not valid Content-Length respones as finished")
 	}
 }
+
+func TestTCPMessageIsSeqMissing(t *testing.T) {
+	p1 := buildPacket(false, 1, 1, []byte("HTTP/1.1 200 OK\r\n"))
+	p2 := buildPacket(false, 1, p1.Seq + uint32(len(p1.Data)), []byte("Content-Length: 10\r\n\r\n"))
+	p3 := buildPacket(false, 1, p2.Seq + uint32(len(p2.Data)), []byte("a"))
+
+	msg := buildMessage(p1)
+	if msg.isSeqMissing() {
+		t.Error("Should be complete if have only 1 packet")
+	}
+
+	msg.AddPacket(p3)
+	if !msg.isSeqMissing() {
+		t.Error("Should be incomplete because missing middle component")
+	}
+
+	msg.AddPacket(p2)
+	if msg.isSeqMissing() {
+		t.Error("Should be complete once missing packet added")
+	}
+}
+
+func TestTCPMessageIsHeadersReceived(t *testing.T) {
+	p1 := buildPacket(false, 1, 1, []byte("HTTP/1.1 200 OK\r\n"))
+	p2 := buildPacket(false, 1, p1.Seq + uint32(len(p1.Data)), []byte("Content-Length: 10\r\n\r\n"))
+
+	msg := buildMessage(p1)
+	if msg.isHeadersReceived() {
+		t.Error("Should be complete if have only 1 packet")
+	}
+
+	msg.AddPacket(p2)
+	if !msg.isHeadersReceived() {
+		t.Error("Should found double new line: headers received")
+	}
+}