feat: support rocketMQ

agent/protocol/rocketmq/filter.go

@@ -0,0 +1,27 @@
+package rocketmq
+
+import (
+	"kyanos/agent/protocol"
+	"kyanos/bpf"
+)
+
+type Filter struct {
+}
+
+func (m Filter) Filter(req protocol.ParsedMessage, resp protocol.ParsedMessage) bool {
+	return true
+}
+
+func (m Filter) FilterByProtocol(p bpf.AgentTrafficProtocolT) bool {
+	return p == bpf.AgentTrafficProtocolTKProtocolRocketMQ
+}
+
+func (m Filter) FilterByRequest() bool {
+	return false
+}
+
+func (m Filter) FilterByResponse() bool {
+	return false
+}
+
+var _ protocol.ProtocolFilter = Filter{}

agent/protocol/rocketmq/rocketmq.go

@@ -0,0 +1,171 @@
+package rocketmq
+
+import (
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"kyanos/agent/buffer"
+	"kyanos/agent/protocol"
+	"kyanos/bpf"
+	"kyanos/common"
+)
+
+func init() {
+	protocol.ParsersMap[bpf.AgentTrafficProtocolTKProtocolRocketMQ] = func() protocol.ProtocolStreamParser {
+		return &RocketMQStreamParser{}
+	}
+}
+
+func (r *RocketMQMessage) FormatToString() string {
+	return fmt.Sprintf("base=[%s] command=[%s] payload=[%s]", r.FrameBase.String(), "todo", r.Body)
+}
+
+func (r *RocketMQMessage) IsReq() bool {
+	return r.isReq
+}
+
+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))
+	common.ProtocolParserLog.Debugln("==============", buffer)
+
+	if len(buffer) < 8 {
+		common.ProtocolParserLog.Warn("Buffer too small for header, needs more data.")
+		return protocol.ParseResult{
+			ParseState: protocol.NeedsMoreData,
+		}
+	}
+
+	frameSize := int(binary.BigEndian.Uint32(buffer[:4]))
+	if frameSize <= 0 {
+		common.ProtocolParserLog.Warnf("Invalid frame size: %d", frameSize)
+		return protocol.ParseResult{ParseState: protocol.Invalid, ReadBytes: 4}
+	}
+
+	if frameSize > len(buffer) {
+		common.ProtocolParserLog.Debugf("Frame size %d exceeds buffer length %d, needs more data.", frameSize, len(buffer))
+		return protocol.ParseResult{ParseState: protocol.NeedsMoreData}
+	}
+
+	headerLength := binary.BigEndian.Uint32(buffer[4:8])
+	headerDataLen := headerLength & 0xFFFFFF
+	serializedType := byte((headerLength >> 24) & 0xFF)
+
+	if 8+int(headerDataLen) > frameSize || len(buffer) < 8+int(headerDataLen) {
+		common.ProtocolParserLog.Warnf("Incomplete header detected: headerDataLen=%d, frameSize=%d.", headerDataLen, frameSize)
+		return protocol.ParseResult{ParseState: protocol.NeedsMoreData}
+	}
+
+	headerBody := buffer[8 : 8+headerDataLen]
+	body := buffer[8+headerDataLen : frameSize]
+
+	message, err := r.parseHeader(headerBody, serializedType)
+	if err != nil {
+		common.ProtocolParserLog.Errorf("Failed to parse header: %v", err)
+		return protocol.ParseResult{ParseState: protocol.Invalid, ReadBytes: int(frameSize)}
+	}
+
+	message.Body = body
+	message.isReq = messageType == protocol.Request
+	fb, ok := protocol.CreateFrameBase(streamBuffer, frameSize)
+
+	if !ok {
+		common.ProtocolParserLog.Warnf("Failed to create FrameBase for frameSize=%d", frameSize)
+		return protocol.ParseResult{
+			ParseState: protocol.Ignore,
+			ReadBytes:  frameSize,
+		}
+	}
+
+	common.ProtocolParserLog.Debugf("Successfully parsed message: %+v", message)
+	message.FrameBase = fb
+	return protocol.ParseResult{
+		ParseState:     protocol.Success,
+		ReadBytes:      frameSize,
+		ParsedMessages: []protocol.ParsedMessage{message},
+	}
+
+}
+
+func (parser *RocketMQStreamParser) parseHeader(headerBody []byte, serializedType byte) (*RocketMQMessage, error) {
+	fmt.Println(serializedType)
+	message := &RocketMQMessage{}
+	if serializedType == 0 {
+		if len(headerBody) < 18 {
+			return nil, errors.New("invalid header size")
+		}
+
+		message.RequestCode = int16(binary.BigEndian.Uint16(headerBody[:2]))
+		message.LanguageFlag = headerBody[2]
+		message.VersionFlag = int16(binary.BigEndian.Uint16(headerBody[3:5]))
+		message.Opaque = int32(binary.BigEndian.Uint32(headerBody[5:9]))
+		message.RequestFlag = int32(binary.BigEndian.Uint32(headerBody[9:13]))
+		message.RemarkLength = int32(binary.BigEndian.Uint32(headerBody[13:17]))
+
+		if int(message.RemarkLength) > len(headerBody[17:]) {
+			return nil, errors.New("invalid remark length")
+		}
+		message.Remark = headerBody[17 : 17+message.RemarkLength]
+		propertiesStart := 17 + message.RemarkLength
+		if len(headerBody[propertiesStart:]) < 4 {
+			return nil, errors.New("invalid properties length")
+		}
+		message.PropertiesLen = int32(binary.BigEndian.Uint32(headerBody[propertiesStart:]))
+		message.Properties = headerBody[propertiesStart+4 : propertiesStart+4+message.PropertiesLen]
+	} else {
+		return nil, errors.New("unsupported serialization type")
+	}
+	return message, nil
+}
+
+func (r *RocketMQStreamParser) FindBoundary(streamBuffer *buffer.StreamBuffer, messageType protocol.MessageType, startPos int) int {
+	buffer := streamBuffer.Head().Buffer()
+	common.ProtocolParserLog.Debugf("FindBoundary starting at position: %d, buffer length: %d", startPos, len(buffer))
+
+	for i := startPos; i <= len(buffer)-8; i++ {
+		frameSize := int(binary.BigEndian.Uint32(buffer[i : i+4]))
+
+		if frameSize <= 0 || frameSize > len(buffer)-i {
+			common.ProtocolParserLog.Warnf("Skipping invalid frameSize=%d at position=%d", frameSize, i)
+			continue
+		}
+
+		if i+frameSize <= len(buffer) {
+			common.ProtocolParserLog.Debugf("Found boundary at position=%d with frameSize=%d", i, frameSize)
+			return i
+		}
+	}
+
+	common.ProtocolParserLog.Warn("No valid boundary found, returning -1.")
+	return -1
+}
+
+func (r *RocketMQStreamParser) Match(reqStream *[]protocol.ParsedMessage, respStream *[]protocol.ParsedMessage) []protocol.Record {
+	common.ProtocolParserLog.Debugf("Matching %d requests with %d responses.", len(*reqStream), len(*respStream))
+	records := []protocol.Record{}
+
+	reqMap := make(map[int32]*RocketMQMessage)
+	for _, msg := range *reqStream {
+		req := msg.(*RocketMQMessage)
+		reqMap[req.Opaque] = req
+	}
+
+	for _, msg := range *respStream {
+		resp := msg.(*RocketMQMessage)
+		if req, ok := reqMap[resp.Opaque]; ok {
+			records = append(records, protocol.Record{
+				Req:  req,
+				Resp: resp,
+			})
+			delete(reqMap, resp.Opaque)
+		} else {
+			common.ProtocolParserLog.Warnf("No matching request found for response Opaque=%d", resp.Opaque)
+		}
+	}
+
+	if len(reqMap) > 0 {
+		common.ProtocolParserLog.Warnf("Unmatched requests remain: %d", len(reqMap))
+	}
+
+	return records
+}

agent/protocol/rocketmq/types.go

@@ -0,0 +1,27 @@
+package rocketmq
+
+import (
+	"kyanos/agent/protocol"
+)
+
+var _ protocol.ParsedMessage = &RocketMQMessage{}
+
+type RocketMQMessage struct {
+	protocol.FrameBase
+	RequestCode   int16
+	LanguageFlag  byte
+	VersionFlag   int16
+	Opaque        int32
+	RequestFlag   int32
+	RemarkLength  int32
+	Remark        []byte
+	PropertiesLen int32
+	Properties    []byte
+	Body          []byte
+	isReq         bool
+}
+
+var _ protocol.ProtocolStreamParser = &RocketMQStreamParser{}
+
+type RocketMQStreamParser struct {
+}

bpf/pktlatency.h

@@ -51,6 +51,7 @@ enum traffic_protocol_t {
   kProtocolKafka,
   kProtocolMux,
   kProtocolAMQP,
+  kProtocolRocketMQ,
   kNumProtocols
 };
 

bpf/protocol_inference.h

@@ -127,6 +127,28 @@ static __always_inline enum message_type_t is_http_protocol(const char *old_buf,
   return kUnknown;
 }
 
+static __always_inline enum message_type_t is_rocketmq_protocol(const char *old_buf, size_t count) {
+  if (count < 8) {
+    return 0;
+  }
+
+  int32_t frame_size = 0;
+  bpf_probe_read_user(&frame_size, sizeof(int32_t), old_buf);
+
+  if (frame_size <= 0 || frame_size > 64 * 1024 * 1024) {
+    return kUnknown;
+  }
+
+  char serialized_type = 0;
+  bpf_probe_read_user(&serialized_type, 1, old_buf + 4);
+
+  if (serialized_type != 0x0 && serialized_type != 0x1) {
+    return kUnknown;
+  }
+
+  return kRequest;
+}
+
 static __always_inline struct protocol_message_t infer_protocol(const char *buf, size_t count, struct conn_info_t *conn_info) {
   struct protocol_message_t protocol_message;
   protocol_message.protocol = kProtocolUnknown;
@@ -137,6 +159,8 @@ static __always_inline struct protocol_message_t infer_protocol(const char *buf,
     protocol_message.protocol = kProtocolMySQL;
   } else if (is_redis_protocol(buf, count)) {
     protocol_message.protocol = kProtocolRedis;
+  } else if (is_rocketmq_protocol(buf,count)) {
+    protocol_message.protocol = kProtocolRocketMQ;
   }
   conn_info->prev_count = count;
   if (count == 4) {

cmd/rocketmq.go

@@ -0,0 +1,26 @@
+package cmd
+
+import (
+	"kyanos/agent/protocol/rocketmq"
+
+	"github.com/spf13/cobra"
+)
+
+var rocketmqCmd *cobra.Command = &cobra.Command{
+	Use:   "rocketmq",
+	Short: "watch RocketMQ message",
+	Run: func(cmd *cobra.Command, args []string) {
+		options.MessageFilter = rocketmq.Filter{}
+		options.LatencyFilter = initLatencyFilter(cmd)
+		options.SizeFilter = initSizeFilter(cmd)
+		startAgent()
+	},
+}
+
+func init() {
+	rocketmqCmd.PersistentFlags().SortFlags = false
+	copy := *rocketmqCmd
+	watchCmd.AddCommand(&copy)
+	copy2 := *rocketmqCmd
+	statCmd.AddCommand(&copy2)
+}

cmd/watch.go

@@ -9,12 +9,13 @@ import (
 var maxRecords int
 var supportedProtocols = []string{"http", "redis", "mysql"}
 var watchCmd = &cobra.Command{
-	Use: "watch [http|redis|mysql] [flags]",
+	Use: "watch [http|redis|mysql|rocketmq] [flags]",
 	Example: `
 sudo kyanos watch
 sudo kyanos watch http --side server --pid 1234 --path /foo/bar --host ubuntu.com
 sudo kyanos watch redis --comands GET,SET --keys foo,bar --key-prefix app1:
 sudo kyanos watch mysql --latency 100 --req-size 1024 --resp-size 2048
+sudo kyanos watch rocketmq
 	`,
 	Short:            "Capture the request/response recrods",
 	PersistentPreRun: func(cmd *cobra.Command, args []string) { Mode = WatchMode },