Merge branch 'main' of github.com:HXSecurity/DongTai-agent-go

Author: piexlmax

core/grpc/clientConn/replacement.go

@@ -3,13 +3,14 @@ package clientConn
 import (
 	"context"
 	"fmt"
+	"strconv"
+	"strings"
+
 	"github.com/HXSecurity/DongTai-agent-go/global"
 	"github.com/HXSecurity/DongTai-agent-go/model/request"
 	"github.com/HXSecurity/DongTai-agent-go/utils"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/metadata"
-	"strconv"
-	"strings"
 )
 
 const (
@@ -63,6 +64,7 @@ func Invoke(cl *grpc.ClientConn, ctx context.Context, method string, args, reply
 		MethodName:      "Invoke",
 		ClassName:       "grpc.(*ClientConn)",
 		TraceId:         tranceid,
+		Plugin:          "GRPC",
 	})
 	return err
 }

core/kafkaGo/kafkaWriter/install.go

@@ -0,0 +1,31 @@
+package kafkaWriter
+
+import (
+	"fmt"
+
+	"github.com/HXSecurity/DongTai-agent-go/model"
+	"github.com/brahma-adshonor/gohook"
+	"github.com/segmentio/kafka-go"
+)
+
+func init() {
+	model.HookMap["kafkaGoWriter"] = new(KafkaWriter)
+}
+
+type KafkaWriter struct {
+}
+
+func (h *KafkaWriter) Hook() {
+	w := &kafka.Writer{}
+	err := gohook.HookMethod(w, "WriteMessages", WriteMessages, WriteMessagesT)
+	if err != nil {
+		fmt.Println(err, "kafkaGoWriter")
+	} else {
+		fmt.Println("kafkaGoWriter")
+	}
+}
+
+func (h *KafkaWriter) UnHook() {
+	cl := &kafka.Writer{}
+	gohook.UnHookMethod(cl, "WriteMessages")
+}

core/kafkaGo/kafkaWriter/replacement.go

@@ -0,0 +1,83 @@
+package kafkaWriter
+
+import (
+	"context"
+	"fmt"
+	"reflect"
+	"strconv"
+	"strings"
+
+	"github.com/HXSecurity/DongTai-agent-go/global"
+	"github.com/HXSecurity/DongTai-agent-go/model/request"
+	"github.com/HXSecurity/DongTai-agent-go/utils"
+	"github.com/segmentio/kafka-go"
+	"google.golang.org/grpc/metadata"
+)
+
+const (
+	TraceId = iota
+	AgentId
+	RoutineId
+	NextKey
+	OnlyKey
+)
+
+func WriteMessages(w *kafka.Writer, ctx context.Context, msgs ...kafka.Message) error {
+	for idx, msg := range msgs {
+		traceId := getTraceId(ctx)
+		msgs[idx].Headers = append(msg.Headers, kafka.Header{
+			Key:   "dt-traceid",
+			Value: []byte(traceId),
+		})
+
+		fmt.Println("traceId:", traceId, "msg:", msg.Value)
+		v := reflect.ValueOf(msg.Value)
+		request.FmtHookPool(request.PoolReq{
+			Args:            request.Collect(msg.Value),
+			NeedHook:        request.Collect(v.Pointer()),
+			Source:          false,
+			OriginClassName: "kafka.(*Writer)",
+			MethodName:      "WriteMessages",
+			ClassName:       "kafka.(*Writer)",
+			TraceId:         traceId,
+			Plugin:          "KAFKA",
+		})
+	}
+
+	err := WriteMessagesT(w, ctx, msgs...)
+	return err
+}
+
+func getTraceId(ctx context.Context) string {
+	outmd, _ := metadata.FromIncomingContext(ctx)
+	worker, _ := utils.NewWorker(global.AgentId)
+	var tranceid string
+	if len(outmd.Get("dt-traceid")) > 0 {
+		tranceid = outmd.Get("dt-traceid")[0]
+	}
+	if tranceid == "" {
+		tranceid = global.TargetTraceId + "." + strconv.Itoa(global.AgentId) + ".0.1." + strconv.Itoa(int(worker.GetId()))
+	} else {
+		four := strconv.Itoa(int(worker.GetId()))
+		tranceids := strings.Split(tranceid, ".")
+		tranceids[AgentId] = strconv.Itoa(global.AgentId)
+		num, _ := strconv.Atoi(tranceids[NextKey])
+		tranceids[NextKey] = strconv.Itoa(num + 1)
+		tranceids[OnlyKey] = four
+		newId := ""
+		for i := 0; i < len(tranceids); i++ {
+			if i == OnlyKey {
+				newId += tranceids[i]
+			} else {
+				newId += tranceids[i] + "."
+			}
+		}
+		tranceid = newId
+	}
+
+	return tranceid
+}
+
+func WriteMessagesT(w *kafka.Writer, ctx context.Context, msgs ...kafka.Message) error {
+	return nil
+}

go.mod

@@ -20,6 +20,7 @@ require (
 	github.com/jinzhu/now v1.1.3 // indirect
 	github.com/json-iterator/go v1.1.11 // indirect
 	github.com/modern-go/reflect2 v1.0.1 // indirect
+	github.com/segmentio/kafka-go v0.4.31
 	github.com/smartystreets/goconvey v1.7.2 // indirect
 	github.com/tklauser/go-sysconf v0.3.9 // indirect
 	golang.org/x/crypto v0.0.0-20210817164053-32db794688a5 // indirect

go.sum

@@ -87,6 +87,8 @@ github.com/jtolds/gls v4.20.0+incompatible h1:xdiiI2gbIgH/gLH7ADydsJ1uDOEzR8yvV7
 github.com/jtolds/gls v4.20.0+incompatible/go.mod h1:QJZ7F/aHp+rZTRtaJ1ow/lLfFfVYBRgL+9YlvaHOwJU=
 github.com/julienschmidt/httprouter v1.3.0 h1:U0609e9tgbseu3rBINet9P48AI/D3oJs4dN7jwJOQ1U=
 github.com/julienschmidt/httprouter v1.3.0/go.mod h1:JR6WtHb+2LUe8TCKY3cZOxFyyO8IZAc4RVcycCCAKdM=
+github.com/klauspost/compress v1.14.2 h1:S0OHlFk/Gbon/yauFJ4FfJJF5V0fc5HbBTJazi28pRw=
+github.com/klauspost/compress v1.14.2/go.mod h1:/3/Vjq9QcHkK5uEr5lBEmyoZ1iFhe47etQ6QUkpK6sk=
 github.com/leodido/go-urn v1.2.0 h1:hpXL4XnriNwQ/ABnpepYM/1vCLWNDfUNts8dX3xTG6Y=
 github.com/leodido/go-urn v1.2.0/go.mod h1:+8+nEpDfqqsY+g338gtMEUOtuK+4dEMhiQEgxpxOKII=
 github.com/mattn/go-isatty v0.0.12 h1:wuysRhFDzyxgEmMf5xjvJ2M9dZoWAXNNr5LSBS7uHXY=
@@ -98,13 +100,17 @@ github.com/modern-go/reflect2 v1.0.1 h1:9f412s+6RmYXLWZSEzVVgPGK7C2PphHj5RJrvfx9
 github.com/modern-go/reflect2 v1.0.1/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3RllmbCylyMrvgv0=
 github.com/parnurzeal/gorequest v0.2.16 h1:T/5x+/4BT+nj+3eSknXmCTnEVGSzFzPGdpqmUVVZXHQ=
 github.com/parnurzeal/gorequest v0.2.16/go.mod h1:3Kh2QUMJoqw3icWAecsyzkpY7UzRfDhbRdTjtNwNiUE=
+github.com/pierrec/lz4/v4 v4.1.14 h1:+fL8AQEZtz/ijeNnpduH0bROTu0O3NZAlPjQxGn8LwE=
+github.com/pierrec/lz4/v4 v4.1.14/go.mod h1:gZWDp/Ze/IJXGXf23ltt2EXimqmTUXEy0GFuRQyBid4=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
 github.com/rogpeppe/fastuuid v1.2.0/go.mod h1:jVj6XXZzXRy/MSR5jhDC/2q6DgLz+nrA6LYCDYWNEvQ=
 github.com/rogpeppe/go-charset v0.0.0-20180617210344-2471d30d28b4/go.mod h1:qgYeAmZ5ZIpBWTGllZSQnw97Dj+woV0toclVaRGI8pc=
+github.com/segmentio/kafka-go v0.4.31 h1:+ImsrkJRju9j1D9U44rvRGRlpsI9GnwD8s9WTFagNLQ=
+github.com/segmentio/kafka-go v0.4.31/go.mod h1:m1lXeqJtIFYZayv0shM/tjrAFljvWLTprxBHd+3PnaU=
 github.com/shirou/gopsutil v3.21.10+incompatible h1:AL2kpVykjkqeN+MFe1WcwSBVUjGjvdU8/ubvCuXAjrU=
 github.com/shirou/gopsutil v3.21.10+incompatible/go.mod h1:5b4v6he4MtMOwMlS0TUMTu2PcXUg8+E1lC7eC3UO/RA=
 github.com/smartystreets/assertions v1.2.0 h1:42S6lae5dvLc7BrLu/0ugRtcFVjoJNMC/N3yZFZkDFs=
@@ -115,6 +121,7 @@ github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
 github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
 github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA=
+github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/tklauser/go-sysconf v0.3.9 h1:JeUVdAOWhhxVcU6Eqr/ATFHgXk/mmiItdKeJPev3vTo=
@@ -125,10 +132,15 @@ github.com/ugorji/go v1.1.7 h1:/68gy2h+1mWMrwZFeD1kQialdSzAb432dtpeJ42ovdo=
 github.com/ugorji/go v1.1.7/go.mod h1:kZn38zHttfInRq0xu/PH0az30d+z6vm202qpg1oXVMw=
 github.com/ugorji/go/codec v1.1.7 h1:2SvQaVZ1ouYrrKKwoSk2pzd4A9evlKJb9oTL+OaLUSs=
 github.com/ugorji/go/codec v1.1.7/go.mod h1:Ax+UKWsSmolVDwsd+7N3ZtXu+yMGCf907BLYF3GoBXY=
+github.com/xdg/scram v0.0.0-20180814205039-7eeb5667e42c h1:u40Z8hqBAAQyv+vATcGgV0YCnDjqSL7/q/JyPhhJSPk=
+github.com/xdg/scram v0.0.0-20180814205039-7eeb5667e42c/go.mod h1:lB8K/P019DLNhemzwFU4jHLhdvlE6uDZjXFejJXr49I=
+github.com/xdg/stringprep v1.0.0 h1:d9X0esnoa3dFsV0FG35rAT0RIhYFlPq7MiP+DW89La0=
+github.com/xdg/stringprep v1.0.0/go.mod h1:Jhud4/sHMO4oL310DaZAKk9ZaJ08SJfe+sJh0HrGL1Y=
 go.opentelemetry.io/proto/otlp v0.7.0/go.mod h1:PqfVotwruBrMGOCsRd/89rSnXhoiJIqeYNgFYFoEGnI=
 golang.org/x/arch v0.0.0-20190312162104-788fe5ffcd8c h1:Rx/HTKi09myZ25t1SOlDHmHOy/mKxNAcu0hP1oPX9qM=
 golang.org/x/arch v0.0.0-20190312162104-788fe5ffcd8c/go.mod h1:flIaEI6LNU6xOCD5PaJvn9wGP0agmIOqjrtsKGRguv4=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
+golang.org/x/crypto v0.0.0-20190506204251-e1dfcc566284/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
 golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
 golang.org/x/crypto v0.0.0-20210817164053-32db794688a5 h1:HWj/xjIHfjYU5nVXpTM0s39J9CbLn7Cc5a7IC5rwsMQ=
 golang.org/x/crypto v0.0.0-20210817164053-32db794688a5/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=

hook/kafkaGo.go

@@ -0,0 +1,18 @@
+package hook
+
+type KafkaGo struct {
+}
+
+func (g *KafkaGo) GetHook() []string {
+	return []string{
+		"kafkaGoWriter",
+	}
+}
+
+func (g *KafkaGo) HookAll() {
+	Hook(g.GetHook())
+}
+
+func (g *KafkaGo) UnHookAll() {
+	UnHook(g.GetHook())
+}

model/request/engine.go

@@ -2,10 +2,11 @@ package request
 
 import (
 	"fmt"
-	"github.com/HXSecurity/DongTai-agent-go/global"
-	"github.com/HXSecurity/DongTai-agent-go/utils"
 	"reflect"
 	"strconv"
+
+	"github.com/HXSecurity/DongTai-agent-go/global"
+	"github.com/HXSecurity/DongTai-agent-go/utils"
 )
 
 type AgentRegisterReq struct {
@@ -109,6 +110,7 @@ type Pool struct {
 	SourceHash       []string      `json:"sourceHash"`
 	RetClassName     string        `json:"retClassName"`
 	TraceId          string        `json:"traceId"`
+	Plugin           string        `json:"plugin"`
 }
 
 type PoolReq struct {
@@ -122,6 +124,7 @@ type PoolReq struct {
 	NeedCatch       []interface{}
 	ArgsStr         string
 	TraceId         string
+	Plugin          string
 }
 
 type Log struct {
@@ -162,9 +165,9 @@ func (p *PoolTree) IsThisBegin(GoroutineID string) bool {
 
 func (p *PoolTree) FMT(pools *[]Pool, w *utils.Worker, goroutineIDs map[string]bool, TraceId string) {
 	p.Pool.InvokeId = int(w.GetId())
-	//if(p.Pool.ClassName == "grpc.(*ClientConn)"){
+	// if(p.Pool.ClassName == "grpc.(*ClientConn)"){
 	//	p.Pool.TraceId = TraceId
-	//}
+	// }
 	*pools = append(*pools, *p.Pool)
 	goroutineIDs[p.GoroutineID] = true
 	fmt.Println(p.Pool.ClassName, p.Pool.MethodName)
@@ -251,6 +254,7 @@ func FmtHookPool(p PoolReq) Pool {
 		RetClassName:     RetClassNames,
 		Args:             ArgsStr,
 		TraceId:          p.TraceId,
+		Plugin:           p.Plugin,
 	}
 
 	poolTree := PoolTree{

run/kafkaGo/kafkaGo.go

@@ -0,0 +1,13 @@
+package kafkaGo
+
+import (
+	_ "github.com/HXSecurity/DongTai-agent-go/core/kafkaGo/kafkaWriter"
+	"github.com/HXSecurity/DongTai-agent-go/global"
+	"github.com/HXSecurity/DongTai-agent-go/hook"
+)
+
+func init() {
+	g := new(hook.KafkaGo)
+	global.AllHooks = append(global.AllHooks, g)
+	g.HookAll()
+}

service/engine.go

@@ -5,23 +5,25 @@ import (
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
-	"github.com/HXSecurity/DongTai-agent-go/api"
-	"github.com/HXSecurity/DongTai-agent-go/global"
-	"github.com/HXSecurity/DongTai-agent-go/hook"
-	"github.com/HXSecurity/DongTai-agent-go/utils"
-	"github.com/pkg/errors"
 	"os/exec"
 	"path/filepath"
 	"regexp"
 	"strconv"
 	"strings"
 	"time"
 
-	//"github.com/HXSecurity/DongTai-agent-go/api"
-	"github.com/HXSecurity/DongTai-agent-go/model/request"
+	"github.com/HXSecurity/DongTai-agent-go/api"
+	"github.com/HXSecurity/DongTai-agent-go/global"
+	"github.com/HXSecurity/DongTai-agent-go/hook"
+	"github.com/HXSecurity/DongTai-agent-go/utils"
+	"github.com/pkg/errors"
+
 	"net"
 	"os"
 	"runtime"
+
+	// "github.com/HXSecurity/DongTai-agent-go/api"
+	"github.com/HXSecurity/DongTai-agent-go/model/request"
 )
 
 var live bool