feat(protocol): introduce the concept of streams to prepare for futur…

…e support of HTTP2 and Mongo (#258) Signed-off-by: 烈香 <[email protected]>
hengyoush · Jan 9, 2025 · 05b2c40 · 05b2c40
1 parent 6d507da
commit 05b2c40
Show file tree

Hide file tree

Showing 10 changed files with 88 additions and 30 deletions.
diff --git a/agent/conn/conntrack.go b/agent/conn/conntrack.go
@@ -47,10 +47,10 @@ type Connection4 struct {
 
 	reqStreamBuffer          *buffer.StreamBuffer
 	respStreamBuffer         *buffer.StreamBuffer
-	ReqQueue                 []protocol.ParsedMessage
+	ReqQueue                 map[protocol.StreamId]*protocol.ParsedMessageQueue
+	RespQueue                map[protocol.StreamId]*protocol.ParsedMessageQueue
 	lastReqMadeProgressTime  int64
 	lastRespMadeProgressTime int64
-	RespQueue                []protocol.ParsedMessage
 	StreamEvents             *KernEventStream
 	protocolParsers          map[bpf.AgentTrafficProtocolT]protocol.ProtocolStreamParser
 
@@ -83,8 +83,8 @@ func NewConnFromEvent(event *bpf.AgentConnEvtT, p *Processor) *Connection4 {
 
 		reqStreamBuffer:  buffer.New(1024 * 1024),
 		respStreamBuffer: buffer.New(1024 * 1024),
-		ReqQueue:         make([]protocol.ParsedMessage, 0),
-		RespQueue:        make([]protocol.ParsedMessage, 0),
+		ReqQueue:         make(map[protocol.StreamId]*protocol.ParsedMessageQueue),
+		RespQueue:        make(map[protocol.StreamId]*protocol.ParsedMessageQueue),
 
 		prevConn: []*Connection4{},
 
@@ -450,8 +450,8 @@ func (c *Connection4) addDataToBufferAndTryParse(data []byte, ke *bpf.AgentKernE
 	if c.Role == bpf.AgentEndpointRoleTKRoleUnknown {
 		respSteamMessageType = protocol.Unknown
 	}
-	c.parseStreamBuffer(c.reqStreamBuffer, reqSteamMessageType, &c.ReqQueue, ke)
-	c.parseStreamBuffer(c.respStreamBuffer, respSteamMessageType, &c.RespQueue, ke)
+	c.parseStreamBuffer(c.reqStreamBuffer, reqSteamMessageType, c.ReqQueue, ke)
+	c.parseStreamBuffer(c.respStreamBuffer, respSteamMessageType, c.RespQueue, ke)
 	return true
 }
 func (c *Connection4) OnSslDataEvent(data []byte, event *bpf.SslData, recordChannel chan RecordWithConn) {
@@ -467,7 +467,7 @@ func (c *Connection4) OnSslDataEvent(data []byte, event *bpf.SslData, recordChan
 		return
 	}
 
-	records := parser.Match(&c.ReqQueue, &c.RespQueue)
+	records := parser.Match(c.ReqQueue, c.RespQueue)
 	if len(records) != 0 {
 		for _, record := range records {
 			recordChannel <- RecordWithConn{record, c}
@@ -501,7 +501,7 @@ func (c *Connection4) OnSyscallEvent(data []byte, event *bpf.SyscallEventData, r
 		panic("no protocol parser!")
 	}
 
-	records := parser.Match(&c.ReqQueue, &c.RespQueue)
+	records := parser.Match(c.ReqQueue, c.RespQueue)
 	if len(records) != 0 {
 		for _, record := range records {
 			recordChannel <- RecordWithConn{record, c}
@@ -510,7 +510,7 @@ func (c *Connection4) OnSyscallEvent(data []byte, event *bpf.SyscallEventData, r
 	return true
 }
 
-func (c *Connection4) parseStreamBuffer(streamBuffer *buffer.StreamBuffer, messageType protocol.MessageType, resultQueue *[]protocol.ParsedMessage, ke *bpf.AgentKernEvt) {
+func (c *Connection4) parseStreamBuffer(streamBuffer *buffer.StreamBuffer, messageType protocol.MessageType, resultQueue map[protocol.StreamId]*protocol.ParsedMessageQueue, ke *bpf.AgentKernEvt) {
 	parser := c.GetProtocolParser(c.Protocol)
 	if parser == nil {
 		return
@@ -551,7 +551,14 @@ func (c *Connection4) parseStreamBuffer(streamBuffer *buffer.StreamBuffer, messa
 				if len(parseResult.ParsedMessages) > 0 && parseResult.ParsedMessages[0].IsReq() != (messageType == protocol.Request) {
 					streamBuffer.RemovePrefix(parseResult.ReadBytes)
 				} else {
-					*resultQueue = append(*resultQueue, parseResult.ParsedMessages...)
+					for _, parsedMessage := range parseResult.ParsedMessages {
+						streamId := parsedMessage.StreamId()
+						if resultQueue[streamId] == nil {
+							queue := protocol.ParsedMessageQueue(make([]protocol.ParsedMessage, 0))
+							resultQueue[streamId] = &queue
+						}
+						*resultQueue[streamId] = append(*resultQueue[streamId], parsedMessage)
+					}
 					streamBuffer.RemovePrefix(parseResult.ReadBytes)
 				}
 			}
@@ -748,6 +755,6 @@ func (c *Connection4) GetProtocolParser(p bpf.AgentTrafficProtocolT) protocol.Pr
 func (c *Connection4) resetParseProgress() {
 	c.reqStreamBuffer.Clear()
 	c.respStreamBuffer.Clear()
-	c.ReqQueue = c.ReqQueue[:]
-	c.RespQueue = c.RespQueue[:]
+	c.ReqQueue = make(map[protocol.StreamId]*protocol.ParsedMessageQueue)
+	c.RespQueue = make(map[protocol.StreamId]*protocol.ParsedMessageQueue)
 }
diff --git a/agent/protocol/common.go b/agent/protocol/common.go
@@ -4,7 +4,7 @@ import (
 	"kyanos/agent/buffer"
 )
 
-func matchByTimestamp(reqStream *[]ParsedMessage, respStream *[]ParsedMessage) []Record {
+func matchByTimestamp(reqStream *ParsedMessageQueue, respStream *ParsedMessageQueue) []Record {
 	if len(*reqStream) == 0 || len(*respStream) == 0 {
 		return nil
 	}

diff --git a/agent/protocol/http.go b/agent/protocol/http.go
@@ -30,7 +30,12 @@ var HTTP_BOUNDARY_MARKER = "\r\n\r\n"
 type HTTPStreamParser struct {
 }
 
-func (h *HTTPStreamParser) Match(reqStream *[]ParsedMessage, respStream *[]ParsedMessage) []Record {
+func (h *HTTPStreamParser) Match(reqStreams map[StreamId]*ParsedMessageQueue, respStreams map[StreamId]*ParsedMessageQueue) []Record {
+	reqStream, ok1 := reqStreams[0]
+	respStream, ok2 := respStreams[0]
+	if !ok1 || !ok2 {
+		return []Record{}
+	}
 	return matchByTimestamp(reqStream, respStream)
 }
 
@@ -214,6 +219,10 @@ func (req *ParsedHttpRequest) IsReq() bool {
 	return true
 }
 
+func (req *ParsedHttpRequest) StreamId() StreamId {
+	return 0
+}
+
 type ParsedHttpResponse struct {
 	FrameBase
 	buf []byte
@@ -233,6 +242,10 @@ func (resp *ParsedHttpResponse) IsReq() bool {
 	return false
 }
 
+func (req *ParsedHttpResponse) StreamId() StreamId {
+	return 0
+}
+
 var _ ProtocolFilter = HttpFilter{}
 
 type HttpFilter struct {

diff --git a/agent/protocol/mysql/mysql.go b/agent/protocol/mysql/mysql.go
@@ -57,7 +57,12 @@ func (m *MysqlParser) FindBoundary(streamBuffer *buffer.StreamBuffer, messageTyp
 	return -1
 }
 
-func (m *MysqlParser) Match(reqStream *[]ParsedMessage, respStream *[]ParsedMessage) []Record {
+func (m *MysqlParser) Match(reqStreams map[StreamId]*ParsedMessageQueue, respStreams map[StreamId]*ParsedMessageQueue) []Record {
+	reqStream, ok1 := reqStreams[0]
+	respStream, ok2 := respStreams[0]
+	if !ok1 || !ok2 {
+		return []Record{}
+	}
 	records := make([]Record, 0)
 	for len(*reqStream) != 0 {
 		reqPacket := (*reqStream)[0].(*MysqlPacket)
@@ -112,13 +117,13 @@ func (m *MysqlParser) Match(reqStream *[]ParsedMessage, respStream *[]ParsedMess
 	return records
 }
 
-func syncRespQueue(reqPacket *MysqlPacket, respStream *[]ParsedMessage) {
+func syncRespQueue(reqPacket *MysqlPacket, respStream *ParsedMessageQueue) {
 	for len(*respStream) != 0 && (*respStream)[0].TimestampNs() < reqPacket.TimestampNs() {
 		*respStream = (*respStream)[1:]
 	}
 }
 
-func getRespView(reqStream *[]ParsedMessage, respStream *[]ParsedMessage) []ParsedMessage {
+func getRespView(reqStream *ParsedMessageQueue, respStream *ParsedMessageQueue) []ParsedMessage {
 	count := 0
 	for _, resp := range *respStream {
 		if len(*reqStream) > 1 && resp.TimestampNs() > (*reqStream)[1].TimestampNs() {

diff --git a/agent/protocol/mysql/types.go b/agent/protocol/mysql/types.go
@@ -135,6 +135,10 @@ func (m *MysqlResponse) IsReq() bool {
 	return false
 }
 
+func (m *MysqlResponse) StreamId() StreamId {
+	return 0
+}
+
 const (
 	Unknwon RespStatus = iota
 	None
@@ -298,6 +302,10 @@ func (m *MysqlPacket) IsReq() bool {
 	return m.isReq
 }
 
+func (req *MysqlPacket) StreamId() StreamId {
+	return 0
+}
+
 type MysqlRequestPacket struct {
 	MysqlPacket
 	cmd byte

diff --git a/agent/protocol/protocol.go b/agent/protocol/protocol.go
@@ -9,9 +9,12 @@ import (
 
 type ProtocolCreator func() ProtocolStreamParser
 
+type StreamId int64
+
+type ParsedMessageQueue []ParsedMessage
+
 var ParsersMap map[bpf.AgentTrafficProtocolT]ProtocolCreator = make(map[bpf.AgentTrafficProtocolT]ProtocolCreator)
 
-// TODO 修改未每一个processor有自己的parser
 func GetParserByProtocol(protocol bpf.AgentTrafficProtocolT) ProtocolStreamParser {
 	parserCreator, ok := ParsersMap[protocol]
 	if ok {
@@ -60,7 +63,7 @@ func (r *Record) String(opt RecordToStringOptions) string {
 type ProtocolStreamParser interface {
 	ParseStream(streamBuffer *buffer.StreamBuffer, messageType MessageType) ParseResult
 	FindBoundary(streamBuffer *buffer.StreamBuffer, messageType MessageType, startPos int) int
-	Match(reqStream *[]ParsedMessage, respStream *[]ParsedMessage) []Record
+	Match(reqStreams map[StreamId]*ParsedMessageQueue, respStreams map[StreamId]*ParsedMessageQueue) []Record
 }
 
 type ParsedMessage interface {
@@ -69,6 +72,7 @@ type ParsedMessage interface {
 	ByteSize() int
 	IsReq() bool
 	Seq() uint64
+	StreamId() StreamId
 }
 
 type ParseState int

diff --git a/agent/protocol/redis..go b/agent/protocol/redis..go
@@ -360,6 +360,10 @@ func (m *RedisMessage) FormatToString() string {
 	return fmt.Sprintf("base=[%s] command=[%s] payload=[%s]", m.FrameBase.String(), m.command, m.payload)
 }
 
+func (req *RedisMessage) StreamId() StreamId {
+	return 0
+}
+
 func (r *RedisStreamParser) FindBoundary(streamBuffer *buffer.StreamBuffer, messageType MessageType, startPos int) int {
 	head := streamBuffer.Head().Buffer()
 	for ; startPos < len(head); startPos++ {
@@ -374,7 +378,12 @@ func (r *RedisStreamParser) FindBoundary(streamBuffer *buffer.StreamBuffer, mess
 	return -1
 }
 
-func (r *RedisStreamParser) Match(reqStream *[]ParsedMessage, respStream *[]ParsedMessage) []Record {
+func (r *RedisStreamParser) Match(reqStreams map[StreamId]*ParsedMessageQueue, respStreams map[StreamId]*ParsedMessageQueue) []Record {
+	reqStream, ok1 := reqStreams[0]
+	respStream, ok2 := respStreams[0]
+	if !ok1 || !ok2 {
+		return []Record{}
+	}
 	return matchByTimestamp(reqStream, respStream)
 }
 func ParseSize(decoder *BinaryDecoder) (int, error) {

diff --git a/agent/protocol/rocketmq/rocketmq.go b/agent/protocol/rocketmq/rocketmq.go
@@ -77,6 +77,10 @@ func (r *RocketMQMessage) IsReq() bool {
 	return r.isReq
 }
 
+func (r *RocketMQMessage) StreamId() protocol.StreamId {
+	return 0
+}
+
 func (r *RocketMQStreamParser) ParseStream(streamBuffer *buffer.StreamBuffer, messageType protocol.MessageType) protocol.ParseResult {
 	buffer := streamBuffer.Head().Buffer()
 	common.ProtocolParserLog.Debugf("ParseStream received buffer length: %d", len(buffer))
@@ -279,7 +283,12 @@ func (r *RocketMQStreamParser) FindBoundary(streamBuffer *buffer.StreamBuffer, m
 	return -1
 }
 
-func (r *RocketMQStreamParser) Match(reqStream *[]protocol.ParsedMessage, respStream *[]protocol.ParsedMessage) []protocol.Record {
+func (r *RocketMQStreamParser) Match(reqStreams map[protocol.StreamId]*protocol.ParsedMessageQueue, respStreams map[protocol.StreamId]*protocol.ParsedMessageQueue) []protocol.Record {
+	reqStream, ok1 := reqStreams[0]
+	respStream, ok2 := respStreams[0]
+	if !ok1 || !ok2 {
+		return []protocol.Record{}
+	}
 	common.ProtocolParserLog.Debugf("Matching %d requests with %d responses.", len(*reqStream), len(*respStream))
 	records := []protocol.Record{}
 

diff --git a/docs/cn/how-to-add-a-new-protocol.md b/docs/cn/how-to-add-a-new-protocol.md
@@ -82,6 +82,7 @@ type ParsedMessage interface {
 	ByteSize() int
 	IsReq() bool
 	Seq() uint64
+	StreamId() StreamId
 }
 ```
 
@@ -92,6 +93,7 @@ type ParsedMessage interface {
 | `ByteSize()`       | 返回消息的字节大小。                                                       |
 | `IsReq()`          | 判断消息是否为请求。                                                       |
 | `Seq()`            | 返回消息的字节流序列号, 可以从 `streamBuffer.Head().LeftBoundary()` 获取。 |
+| `StreamId()`       | 返回消息的 StreamId, 一般的协议可以直接返回 0，用于 HTTP2 这种多路复用的协议。 |
 
 HTTP 的例子：
 
@@ -118,7 +120,7 @@ type ParsedHttpRequest struct {
 type ProtocolStreamParser interface {
 	ParseStream(streamBuffer *buffer.StreamBuffer, messageType MessageType) ParseResult
 	FindBoundary(streamBuffer *buffer.StreamBuffer, messageType MessageType, startPos int) int
-	Match(reqStream *[]ParsedMessage, respStream *[]ParsedMessage) []Record
+	Match(reqStream *ParsedMessageQueue, respStream *ParsedMessageQueue) []Record
 }
 ```
 

diff --git a/docs/how-to-add-a-new-protocol.md b/docs/how-to-add-a-new-protocol.md
@@ -106,13 +106,14 @@ type ParsedMessage interface {
 }
 ```
 
-| Method Name        | Function                                                                                              |
-| ------------------ | ----------------------------------------------------------------------------------------------------- |
-| `FormatToString()` | Formats the message into a string representation.                                                     |
-| `TimestampNs()`    | Returns the timestamp of the message (in nanoseconds).                                                |
-| `ByteSize()`       | Returns the byte size of the message.                                                                 |
-| `IsReq()`          | Determines if the message is a request.                                                               |
-| `Seq()`            | Returns the sequence number of the byte stream.(Obtain Seq from `streamBuffer.Head().LeftBoundary()`) |
+| Method Name        | Function                                                                                                                                |
+| ------------------ | --------------------------------------------------------------------------------------------------------------------------------------- |
+| `FormatToString()` | Formats the message into a string representation.                                                                                       |
+| `TimestampNs()`    | Returns the timestamp of the message (in nanoseconds).                                                                                  |
+| `ByteSize()`       | Returns the byte size of the message.                                                                                                   |
+| `IsReq()`          | Determines if the message is a request.                                                                                                 |
+| `Seq()`            | Returns the sequence number of the byte stream.(Obtain Seq from `streamBuffer.Head().LeftBoundary()`)                                   |
+| `StreamId()`       | Return the StreamId of the message. For most protocols, you can directly return 0. This is used for multiplexed protocols like HTTP2.。 |
 
 Example for HTTP:
 
@@ -142,7 +143,7 @@ interface:
 type ProtocolStreamParser interface {
 	ParseStream(streamBuffer *buffer.StreamBuffer, messageType MessageType) ParseResult
 	FindBoundary(streamBuffer *buffer.StreamBuffer, messageType MessageType, startPos int) int
-	Match(reqStream *[]ParsedMessage, respStream *[]ParsedMessage) []Record
+	Match(reqStream *ParsedMessageQueue, respStream *ParsedMessageQueue) []Record
 }
 ```