Fix BodySize for multi-packet headers

.gitignore

@@ -15,3 +15,5 @@
 gor
 
 *.mprof
+
+*.pcap

Makefile

@@ -76,6 +76,9 @@ run-arg:
 file-server:
 	go run $(SOURCE) file-server $(FADDR)
 
+readpcap:
+	go run $(SOURCE) --input-raw $(FILE) --input-raw-engine pcap_file --output-stdout
+
 record:
 	$(RUN) go run $(SOURCE) --input-dummy=0 --output-file=requests.gor --verbose --debug
 

input_raw.go

@@ -24,6 +24,7 @@ type RAWInput struct {
 const (
 	EngineRawSocket = 1 << iota
 	EnginePcap
+	EnginePcapFile
 )
 
 // NewRAWInput constructor for RAWInput. Accepts address with port as argument.
@@ -69,6 +70,12 @@ func (i *RAWInput) listen(address string) {
 
 	host, port, err := net.SplitHostPort(address)
 
+	if i.engine == EnginePcapFile {
+		host = address
+		port = "1"
+		err = nil
+	}
+
 	if err != nil {
 		log.Fatal("input-raw: error while parsing address", err)
 	}

plugins.go

@@ -98,6 +98,8 @@ func InitPlugins() {
 	engine := EnginePcap
 	if Settings.inputRAWEngine == "raw_socket" {
 		engine = EngineRawSocket
+	} else if Settings.inputRAWEngine == "pcap_file" {
+		engine = EnginePcapFile
 	}
 
 	for _, options := range Settings.inputRAW {

raw_socket_listener/listener.go

@@ -81,6 +81,7 @@ type request struct {
 const (
 	EngineRawSocket = 1 << iota
 	EnginePcap
+	EnginePcapFile
 )
 
 // NewListener creates and initializes new Listener object
@@ -118,6 +119,8 @@ func NewListener(addr string, port string, engine int, trackResponse bool, expir
 			go l.readRAWSocket()
 		case EnginePcap:
 			go l.readPcap()
+		case EnginePcapFile:
+			go l.readPcapFile()
 		default:
 			log.Fatal("Unknown traffic interception engine:", engine)
 		}
@@ -171,6 +174,8 @@ func (t *Listener) dispatchMessage(message *TCPMessage) {
 		return
 	}
 
+	log.Println("MESSAGE:", message, message.BodySize(), message.contentLength, message.methodType, message.bodyType)
+
 	t.deleteMessage(message)
 
 	if message.methodType == httpMethodNotFound {
@@ -466,6 +471,53 @@ func (t *Listener) readPcap() {
 	t.readyCh <- true
 }
 
+func (t *Listener) readPcapFile() {
+	if handle, err := pcap.OpenOffline(t.addr); err != nil {
+		log.Fatal(err)
+	} else {
+		t.readyCh <- true
+		packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
+
+		for {
+			packet, err := packetSource.NextPacket()
+			if err == io.EOF {
+				break
+			} else if err != nil {
+				log.Println("Error:", err)
+				continue
+			}
+
+			var addr, data []byte
+
+			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})
+			} else {
+				log.Println("Can't find TCP layer", packet)
+				continue
+			}
+
+			if ipLayer := packet.Layer(layers.LayerTypeIPv4); ipLayer != nil {
+			  ip, _ := ipLayer.(*layers.IPv4)
+			  addr = ip.SrcIP
+			} else if ipLayer = packet.Layer(layers.LayerTypeIPv6); ipLayer != nil {
+				ip, _ := ipLayer.(*layers.IPv6)
+				addr = ip.SrcIP
+			} else {
+				log.Println("Can't find IP layer", packet)
+				continue
+			}
+
+			newBuf := make([]byte, len(data)+16)
+			copy(newBuf[:16], addr)
+			copy(newBuf[16:], data)
+
+			t.packetsChan <- newBuf
+		}
+	}
+}
+
 func (t *Listener) readRAWSocket() {
 	conn, e := net.ListenPacket("ip:tcp", t.addr)
 	t.conn = conn

raw_socket_listener/tcp_message.go

@@ -66,13 +66,13 @@ func (t *TCPMessage) Bytes() (output []byte) {
 
 // Size returns total body size
 func (t *TCPMessage) BodySize() (size int) {
-	if len(t.packets) == 0 {
+	if len(t.packets) == 0 || t.headerPacket == -1 {
 		return 0
 	}
 
-	size += len(proto.Body(t.packets[0].Data))
+	size += len(proto.Body(t.packets[t.headerPacket].Data))
 
-	for _, p := range t.packets[1:] {
+	for _, p := range t.packets[t.headerPacket + 1:] {
 		size += len(p.Data)
 	}
 
@@ -266,6 +266,7 @@ func (t *TCPMessage) updateMethodType() {
 
 	if t.IsIncoming {
 		var method []byte
+
 		if mIdx := bytes.IndexByte(d[:8], ' '); mIdx != -1 {
 			method = d[:mIdx]
 

raw_socket_listener/tcp_message_test.go

@@ -215,3 +215,32 @@ func TestTCPMessageBodyType(t *testing.T) {
 		}
 	}
 }
+
+
+func TestTCPMessageBodySize(t *testing.T) {
+	testCases := []struct {
+		direction        bool
+		payloads         []string
+		expectedSize     int
+	}{
+		{true, []string{"GET / HTTP/1.1\r\n\r\n"}, 0},
+		{true, []string{"POST / HTTP/1.1\r\nContent-Length: 2\r\n\r\nab"}, 2},
+		{true, []string{"GET / HTTP/1.1\r\n", "Content-Length: 2\r\n\r\nab"}, 2},
+		{true, []string{"GET / HTTP/1.1\r\n", "Content-Length: 2\r\n\r\n", "ab"}, 2},
+	}
+
+	for _, tc := range testCases {
+		msg := buildMessage(buildPacket(tc.direction, 1, 1, []byte(tc.payloads[0])))
+
+		if len(tc.payloads) > 1 {
+			for _, p := range tc.payloads[1:] {
+				seq := uint32(1 + msg.Size())
+				msg.AddPacket(buildPacket(tc.direction, 1, seq, []byte(p)))
+			}
+		}
+
+		if msg.BodySize() != tc.expectedSize {
+			t.Errorf("Expected %d, got %d", tc.expectedSize, msg.BodySize())
+		}
+	}
+}