tcp.message uuid fix

input_raw.go

@@ -149,8 +149,10 @@ func (i *RAWInput) listen(address string) {
 		log.Fatal(err)
 	}
 	pool := tcp.NewMessagePool(i.CopyBufferSize, i.Expire, Debug, i.handler)
-	pool.End = endHint
-	pool.Start = startHint
+	pool.MatchUUID(i.TrackResponse)
+	if i.Protocol == ProtocolHTTP {
+		pool.Start = http1StartHint
+	}
 	var ctx context.Context
 	ctx, i.cancelListener = context.WithCancel(context.Background())
 	errCh := i.listener.ListenBackground(ctx, pool.Handler)
@@ -196,14 +198,10 @@ func (i *RAWInput) addStats(mStats tcp.Stats) {
 	i.Unlock()
 }
 
-func startHint(pckt *tcp.Packet) (isIncoming, isOutgoing bool) {
+func http1StartHint(pckt *tcp.Packet) (isIncoming, isOutgoing bool) {
 	isIncoming = proto.HasRequestTitle(pckt.Payload)
 	if isIncoming {
 		return
 	}
 	return false, proto.HasResponseTitle(pckt.Payload)
 }
-
-func endHint(m *tcp.Message) bool {
-	return proto.HasFullPayload(m.Data())
-}

tcp/tcp_message.go

@@ -1,13 +1,13 @@
 package tcp
 
 import (
-	"bytes"
-	"crypto/sha1"
+	"encoding/binary"
 	"encoding/hex"
 	"fmt"
 	"sort"
 	"sync"
 	"time"
+	"unsafe"
 
 	"github.com/buger/goreplay/size"
 	"github.com/google/gopacket"
@@ -31,7 +31,7 @@ type Stats struct {
 type Message struct {
 	packets []*Packet
 	done    chan bool
-	buf     bytes.Buffer
+	pool    *MessagePool
 	Stats
 }
 
@@ -45,35 +45,48 @@ func NewMessage(srcAddr, dstAddr string, ipVersion uint8) (m *Message) {
 	return
 }
 
-// UUID the unique id of a TCP session it is not granted to be unique!
+// UUID returns the UUID of a TCP request and its response.
 func (m *Message) UUID() []byte {
-	var src, dst string
-	if m.IsIncoming {
-		src = m.SrcAddr
-		dst = m.DstAddr
-	} else {
-		src = m.DstAddr
-		dst = m.SrcAddr
+	var key uint64
+	pckt := m.packets[0]
+
+	// check if response or request have generated the ID before.
+	if m.pool.uuids != nil {
+		lst := len(pckt.SrcIP) - 4
+		if m.IsIncoming {
+			key = uint64(pckt.SrcPort)<<48 | uint64(pckt.DstPort)<<32 |
+				uint64(_uint32(&pckt.SrcIP[lst]))
+		} else {
+			key = uint64(pckt.DstPort)<<48 | uint64(pckt.SrcPort)<<32 |
+				uint64(_uint32(&pckt.DstIP[lst]))
+		}
+		if uuidHex, ok := m.pool.uuids[key]; ok {
+			delete(m.pool.uuids, key)
+			return uuidHex
+		}
 	}
 
-	length := len(src) + len(dst)
-	uuid := make([]byte, length)
-	copy(uuid, src)
-	copy(uuid[len(src):], dst)
-	sha := sha1.Sum(uuid)
-	uuid = make([]byte, 40)
-	hex.Encode(uuid, sha[:])
-
-	return uuid
+	id := make([]byte, 12, 12)
+	binary.BigEndian.PutUint32(id, pckt.Seq)
+	tStamp := m.End.UnixNano()
+	for i, v := range _8byte(&tStamp) {
+		id[i+4] = v
+	}
+	uuidHex := make([]byte, 24, 24)
+	hex.Encode(uuidHex[:], id[:])
+	if m.pool.uuids != nil {
+		if len(m.pool.uuids) >= 1000 {
+			m.pool.cleanUUIDs()
+		}
+		m.pool.uuids[key] = uuidHex
+	}
+	return uuidHex
 }
 
 func (m *Message) add(pckt *Packet) {
 	m.Length += len(pckt.Payload)
 	m.LostData += int(pckt.Lost)
 	m.packets = append(m.packets, pckt)
-	if len(pckt.Payload) > 0 {
-		m.buf.Write(pckt.Payload)
-	}
 	m.End = pckt.Timestamp
 }
 
@@ -84,7 +97,16 @@ func (m *Message) Packets() []*Packet {
 
 // Data returns data in this message
 func (m *Message) Data() []byte {
-	return m.buf.Bytes()
+	buf := make([]byte, m.Length, m.Length)
+	var n, nn int
+	for _, p := range m.packets {
+		if len(p.Payload) < 1 {
+			continue
+		}
+		nn = copy(buf[n:], p.Payload)
+		n += nn
+	}
+	return buf
 }
 
 // Sort a helper to sort packets
@@ -104,7 +126,7 @@ type Debugger func(int, ...interface{})
 type HintEnd func(*Message) bool
 
 // HintStart hints the pool to start the reassembling the message, see MessagePool.Start
-// when set, it will be used instead of checking SYN flag
+// when set, it will be called after checking SYN flag
 type HintStart func(*Packet) (IsIncoming, IsOutgoing bool)
 
 // MessagePool holds data of all tcp messages in progress(still receiving/sending packets).
@@ -114,6 +136,7 @@ type MessagePool struct {
 	debug         Debugger
 	maxSize       size.Size // maximum message size, default 5mb
 	pool          map[uint64]*Message
+	uuids         map[uint64][]byte
 	handler       Handler
 	messageExpire time.Duration // the maximum time to wait for the final packet, minimum is 100ms
 	End           HintEnd
@@ -141,19 +164,18 @@ func NewMessagePool(maxSize size.Size, messageExpire time.Duration, debugger Deb
 func (pool *MessagePool) Handler(packet gopacket.Packet) {
 	var in, out bool
 	pckt, err := ParsePacket(packet)
-	if err != nil || pckt == nil {
+	if err != nil {
 		go pool.say(4, fmt.Sprintf("error decoding packet(%dBytes):%s\n", packet.Metadata().CaptureLength, err))
 		return
 	}
+	if pckt == nil {
+		return
+	}
 	pool.Lock()
 	defer pool.Unlock()
-	lst := 3
-	if pckt.Version == 6 {
-		lst = 15
-	}
+	lst := len(pckt.SrcIP) - 4
 	key := uint64(pckt.SrcPort)<<48 | uint64(pckt.DstPort)<<32 |
-		uint64(pckt.SrcIP[lst])<<24 | uint64(pckt.SrcIP[lst-1])<<16 |
-		uint64(pckt.SrcIP[lst-2])<<8 | uint64(pckt.SrcIP[lst-3])
+		uint64(_uint32(&pckt.SrcIP[lst]))
 	m, ok := pool.pool[key]
 	if pckt.RST {
 		if ok {
@@ -161,8 +183,7 @@ func (pool *MessagePool) Handler(packet gopacket.Packet) {
 			<-m.done
 		}
 		key = uint64(pckt.DstPort)<<48 | uint64(pckt.SrcPort)<<32 |
-			uint64(pckt.DstIP[lst])<<24 | uint64(pckt.DstIP[lst-1])<<16 |
-			uint64(pckt.DstIP[lst-2])<<8 | uint64(pckt.DstIP[lst-3])
+			uint64(_uint32(&pckt.DstIP[lst]))
 		m, ok = pool.pool[key]
 		if ok {
 			m.done <- true
@@ -188,10 +209,34 @@ func (pool *MessagePool) Handler(packet gopacket.Packet) {
 	m.IsIncoming = in
 	pool.pool[key] = m
 	m.Start = pckt.Timestamp
+	m.pool = pool
 	go pool.dispatch(key, m)
 	pool.addPacket(m, pckt)
 }
 
+// MatchUUID instructs the pool to use same UUID for request and responses
+// this function should be called at initial stage of the pool
+func (pool *MessagePool) MatchUUID(match bool) {
+	if match {
+		pool.uuids = make(map[uint64][]byte)
+		return
+	}
+	pool.uuids = nil
+}
+
+// run GC on UUID map
+func (pool *MessagePool) cleanUUIDs() {
+	var tStamp int64
+	now := time.Now().UnixNano()
+	for k, v := range pool.uuids {
+		// there is a timestamp wrapped in every ID
+		tStamp = _int64(&v[4])
+		if time.Duration(now-tStamp) >= pool.messageExpire {
+			delete(pool.uuids, k)
+		}
+	}
+}
+
 func (pool *MessagePool) dispatch(key uint64, m *Message) {
 	select {
 	case <-m.done:
@@ -217,10 +262,20 @@ func (pool *MessagePool) addPacket(m *Message, pckt *Packet) {
 	}
 	m.add(pckt)
 	switch {
+
+	// if one of this cases matches, we dispatch the message
 	case trunc >= 0:
 	case pckt.FIN:
-	case pool.End != nil && pool.End(m):
 	default:
+		if pool.End != nil {
+			if pool.End(m) {
+				break
+			}
+		} else if m.Length > 0 && len(pckt.Payload) < 1 {
+			// generally, data are chunked and sent sequentially in packets, if one packets is empty we assume
+			// that it's just an acknowledgement packet after the final chunk.
+			break
+		}
 		return
 	}
 	m.done <- true
@@ -233,3 +288,17 @@ func (pool *MessagePool) say(level int, args ...interface{}) {
 		pool.debug(level, args...)
 	}
 }
+
+type void = unsafe.Pointer
+
+func _uint32(b *byte) uint32 {
+	return *(*uint32)(void(b))
+}
+
+func _int64(b *byte) int64 {
+	return *(*int64)(void(b))
+}
+
+func _8byte(i *int64) [8]byte {
+	return *(*[8]byte)(void(i))
+}

tcp/tcp_test.go

@@ -51,9 +51,6 @@ func TestMessageParserWithHint(t *testing.T) {
 	pool.Start = func(pckt *Packet) (bool, bool) {
 		return proto.HasRequestTitle(pckt.Payload), proto.HasResponseTitle(pckt.Payload)
 	}
-	pool.End = func(m *Message) bool {
-		return proto.HasFullPayload(m.Data())
-	}
 	packets := GetPackets(1, 30, nil)
 	packets[0].Data()[14:][20:][13] = 2  // SYN flag
 	packets[10].Data()[14:][20:][13] = 2 // SYN flag
@@ -75,8 +72,8 @@ func TestMessageParserWithHint(t *testing.T) {
 		return
 	case m = <-mssg:
 	}
-	if len(m.packets) != 7 {
-		t.Errorf("expected to have 7 packets got %d", len(m.packets))
+	if len(m.packets) != 8 {
+		t.Errorf("expected to have 8 packets got %d", len(m.packets))
 	}
 	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"))
@@ -88,8 +85,8 @@ func TestMessageParserWithHint(t *testing.T) {
 		return
 	case m = <-mssg:
 	}
-	if len(m.packets) != 7 {
-		t.Errorf("expected to have 7 packets got %d", len(m.packets))
+	if len(m.packets) != 8 {
+		t.Errorf("expected to have 8 packets got %d", len(m.packets))
 	}
 	if !bytes.HasSuffix(m.Data(), []byte("Network")) {
 		t.Errorf("expected to %q to have suffix %q", m.Data(), []byte("Network"))
@@ -101,8 +98,8 @@ func TestMessageParserWithHint(t *testing.T) {
 		return
 	case m = <-mssg:
 	}
-	if len(m.packets) != 6 {
-		t.Errorf("expected to have 6 packets got %d", len(m.packets))
+	if len(m.packets) != 2 {
+		t.Errorf("expected to have 2 packets got %d", len(m.packets))
 	}
 	if !bytes.HasSuffix(m.Data(), []byte("Content-Length: 0\r\n\r")) {
 		t.Errorf("expected to %q to have suffix %q", m.Data(), []byte("Content-Length: 0\r\n\r"))
@@ -191,16 +188,47 @@ func TestMessageTimeoutReached(t *testing.T) {
 }
 
 func TestMessageUUID(t *testing.T) {
-	m1 := &Message{}
-	m1.IsIncoming = true
-	m1.SrcAddr = "src"
-	m1.DstAddr = "dst"
-	m2 := &Message{}
-	m2.SrcAddr = "dst"
-	m2.DstAddr = "src"
-	if string(m1.UUID()) != string(m2.UUID()) {
-		t.Errorf("expected %s, to equal %s", m1.UUID(), m2.UUID())
+	packets := GetPackets(1, 10, nil)
+	packets[0].Data()[14:][20:][13] = 2 // SYN flag
+	packets[4].Data()[14:][20:][13] = 1 // FIN flag
+	packets[5].Data()[14:][20:][13] = 2 // SYN flag
+	packets[9].Data()[14:][20:][13] = 1 // FIN flag
+	var uuid, uuid1 []byte
+	pool := NewMessagePool(0, 0, nil, func(msg *Message) {
+		if len(uuid) == 0 {
+			uuid = msg.UUID()
+			return
+		}
+		uuid1 = msg.UUID()
+	})
+	pool.MatchUUID(true)
+	for _, p := range packets {
+		pool.Handler(p)
 	}
+
+	if string(uuid) != string(uuid1) {
+		t.Errorf("expected %s, to equal %s", uuid, uuid1)
+	}
+}
+
+func BenchmarkMessageUUID(b *testing.B) {
+	packets := GetPackets(1, 5, nil)
+	packets[0].Data()[14:][20:][13] = 2 // SYN flag
+	packets[4].Data()[14:][20:][13] = 1 // FIN flag
+	var uuid []byte
+	var msg *Message
+	pool := NewMessagePool(0, 0, nil, func(m *Message) {
+		msg = m
+	})
+	pool.MatchUUID(true)
+	for _, p := range packets {
+		pool.Handler(p)
+	}
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		uuid = msg.UUID()
+	}
+	_ = uuid
 }
 
 func BenchmarkPacketParseAndSort(b *testing.B) {
@@ -235,7 +263,7 @@ func BenchmarkMessageParserWithoutHint(b *testing.B) {
 }
 
 func BenchmarkMessageParserWithHint(b *testing.B) {
-	var buf [1002][]byte
+	var buf [1003][]byte
 	var chunk = []byte("1e\r\n111111111111111111111111111111\r\n")
 	buf[0] = []byte("HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nTransfer-Encoding: chunked\r\n\r\n")
 	for i := 1; i < 1000; i++ {
@@ -246,14 +274,12 @@ func BenchmarkMessageParserWithHint(b *testing.B) {
 	for i := 0; i < len(buf); i++ {
 		packets[i] = GetPackets(uint32(i+10), 1, buf[i])[0]
 	}
+	packets[1002] = GetPackets(1020, 1, nil)[0]
 	var mssg = make(chan *Message, 1)
 	pool := NewMessagePool(1<<30, time.Second*10, nil, func(m *Message) { mssg <- m })
 	pool.Start = func(pckt *Packet) (bool, bool) {
 		return false, proto.HasResponseTitle(pckt.Payload)
 	}
-	pool.End = func(m *Message) bool {
-		return proto.HasFullPayload(m.Data())
-	}
 	b.ResetTimer()
 	b.ReportMetric(float64(len(packets)), "packets/op")
 	b.ReportMetric(float64(1000), "chunks/op")