hengyoush · mannkafai · Jan 16, 2025 · Jan 17, 2025 · Jan 18, 2025 · Jan 18, 2025
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -237,6 +237,22 @@ jobs:
             fi
             popd
 
+      - name: Test NATS
+        uses: cilium/little-vm-helper@97c89f004bd0ab4caeacfe92ebc956e13e362e6b # v0.0.19
+        with:
+          provision: 'false'
+          cmd: |
+            set -ex
+            uname -a
+            cat /etc/issue
+            pushd /host
+            if [ -f "/var/lib/kyanos/btf/current.btf" ]; then  
+                bash /host/testdata/test_nats.sh 'sudo /host/kyanos/kyanos $kyanos_log_option --btf /var/lib/kyanos/btf/current.btf'
+            else  
+                bash /host/testdata/test_nats.sh 'sudo /host/kyanos/kyanos $kyanos_log_option'
+            fi
+            popd
+
       - name: Test filter by comm
         uses: cilium/little-vm-helper@97c89f004bd0ab4caeacfe92ebc956e13e362e6b # v0.0.19
         with:

diff --git a/agent/protocol/common.go b/agent/protocol/common.go
@@ -4,7 +4,7 @@ import (
 	"kyanos/agent/buffer"
 )
 
-func matchByTimestamp(reqStream *ParsedMessageQueue, respStream *ParsedMessageQueue) []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
@@ -36,7 +36,7 @@ func (h *HTTPStreamParser) Match(reqStreams map[StreamId]*ParsedMessageQueue, re
 	if !ok1 || !ok2 {
 		return []Record{}
 	}
-	return matchByTimestamp(reqStream, respStream)
+	return MatchByTimestamp(reqStream, respStream)
 }
 
 func (h *HTTPStreamParser) FindBoundary(streamBuffer *buffer.StreamBuffer, messageType MessageType, startPos int) int {

diff --git a/agent/protocol/nats/filter.go b/agent/protocol/nats/filter.go
@@ -0,0 +1,42 @@
+package nats
+
+import (
+	"kyanos/agent/protocol"
+	"kyanos/bpf"
+	"kyanos/common"
+	"slices"
-	"kyanos/agent/protocol"
-	"kyanos/bpf"
-	"kyanos/common"
-	"slices"
+	"slices"
+
+	"kyanos/agent/protocol"
+	"kyanos/bpf"
+	"kyanos/common"
-	"kyanos/agent/protocol"
-	"kyanos/bpf"
-	"kyanos/common"
-	"slices"
+	"slices"
+
+	"kyanos/agent/protocol"
+	"kyanos/bpf"
+	"kyanos/common"
+)
+
+var _ protocol.ProtocolFilter = NatsFilter{}
+
+type NatsFilter struct {
+	Protocols []string
+	Subjects  []string
+}
+
+func (filter NatsFilter) FilterByProtocol(protocol bpf.AgentTrafficProtocolT) bool {
+	return protocol == bpf.AgentTrafficProtocolTKProtocolNATS
+}
+
+func (filter NatsFilter) FilterByRequest() bool {
+	return len(filter.Protocols) > 0 || len(filter.Subjects) > 0
+}
+
+func (filter NatsFilter) FilterByResponse() bool {
+	return false
+}
+
+func (filter NatsFilter) Filter(req protocol.ParsedMessage, rsp protocol.ParsedMessage) bool {
+	natsReq, ok := req.(*NatsMessage)
+	if !ok {
+		common.ProtocolParserLog.Warnf("[NATSFilter] cast to NatsMessage failed: %v\n", req)
+		return false
+	}
+	if len(filter.Protocols) > 0 && !slices.Contains(filter.Protocols, natsReq.ProtocolCode.String()) {
+		return false
+	}
+	if len(filter.Subjects) > 0 && len(natsReq.Subject) > 0 && !slices.Contains(filter.Subjects, natsReq.Subject) {
+		return false
+	}
+	return true
+}