diff --git a/input_raw.go b/input_raw.go index cc39a0c..77545d4 100644 --- a/input_raw.go +++ b/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()) -} diff --git a/tcp/tcp_message.go b/tcp/tcp_message.go index b7db474..00f3969 100644 --- a/tcp/tcp_message.go +++ b/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)) +} diff --git a/tcp/tcp_test.go b/tcp/tcp_test.go index f2cb16c..dff52dd 100644 --- a/tcp/tcp_test.go +++ b/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")