Merge pull request #30 from piexlmax/main

fixed:grpc

Author: 奇淼（piexlmax

core/gin/ginEngineServerHTTP/replacement.go

@@ -11,6 +11,7 @@ import (
 	"github.com/gin-gonic/gin"
 	"net/http"
 	"reflect"
+	"strconv"
 	"strings"
 )
 
@@ -24,6 +25,8 @@ func MyServer(server *gin.Engine, w http.ResponseWriter, r *http.Request) {
 		for k, v := range r.Header {
 			headerBase += k + ": " + strings.Join(v, ",") + "\n"
 		}
+		tranceID := global.TraceId + "." + strconv.Itoa(global.AgentId) + ".0.0.0"
+		headerBase += "dt-traceid:" + tranceID
 		if t.Kind() == reflect.Ptr {
 			buf := t.
 				Elem().

core/grpc/clientConn/replacement.go

@@ -20,20 +20,21 @@ func Invoke(cl *grpc.ClientConn, ctx context.Context, method string, args, reply
 	}
 
 	if tranceid == "" {
-		tranceid = global.TraceId + "." + strconv.Itoa(global.AgentId) + ".0." + strconv.Itoa(int(worker.GetId()))
+		tranceid = global.TraceId + "." + strconv.Itoa(global.AgentId) + ".0.0." + strconv.Itoa(int(worker.GetId()))
 	} else {
 		four := strconv.Itoa(int(worker.GetId()))
 		tranceids := strings.Split(tranceid, ".")
 		tranceids[1] = strconv.Itoa(global.AgentId)
-		int, _ := strconv.Atoi(tranceids[2])
-		tranceids[2] = strconv.Itoa(int + 1)
-		tranceids[3] = four
+		num, _ := strconv.Atoi(tranceids[3])
+		tranceids[3] = strconv.Itoa(num + 1)
+		tranceids[4] = four
 		newId := ""
-		for i := range tranceids {
-			if i == 3 {
-				newId = tranceids[i]
+		for i := 0; i < len(tranceids); i++ {
+			if i == 4 {
+				newId += tranceids[i]
+			} else {
+				newId += tranceids[i] + "."
 			}
-			newId = tranceids[i] + "."
 		}
 		tranceid = newId
 	}

core/grpc/newServer/replacement.go

@@ -10,6 +10,8 @@ import (
 	"github.com/HXSecurity/DongTai-agent-go/utils"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/metadata"
+	"strconv"
+	"strings"
 )
 
 func NewServer(opt ...grpc.ServerOption) *grpc.Server {
@@ -20,6 +22,24 @@ func NewServer(opt ...grpc.ServerOption) *grpc.Server {
 // interceptor 一元拦截器
 func interceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
 	md, _ := metadata.FromIncomingContext(ctx)
+	Traceid := md.Get("dt-traceid")[0]
+	worker, _ := utils.NewWorker(global.AgentId)
+	four := strconv.Itoa(int(worker.GetId()))
+	tranceids := strings.Split(Traceid, ".")
+	tranceids[1] = strconv.Itoa(global.AgentId)
+	num, _ := strconv.Atoi(tranceids[3])
+	tranceids[3] = strconv.Itoa(num + 1)
+	tranceids[4] = four
+	newId := ""
+	for i := 0; i < len(tranceids); i++ {
+		if i == 4 {
+			newId += tranceids[i]
+		} else {
+			newId += tranceids[i] + "."
+		}
+	}
+	global.TraceId = tranceids[0]
+
 	id := utils.CatGoroutineID()
 	request.FmtHookPool(request.PoolReq{
 		Reqs:            request.Collect(req),
@@ -33,17 +53,17 @@ func interceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInf
 	go func() {
 		worker, _ := utils.NewWorker(global.AgentId)
 		onlyKey := int(worker.GetId())
-		header := base64.StdEncoding.EncodeToString([]byte(md.Get("headers")[0]))
+		header := base64.StdEncoding.EncodeToString([]byte("dt-traceid:" + newId))
 		HookGroup := &request.UploadReq{
 			Type:     36,
 			InvokeId: onlyKey,
 			Detail: request.Detail{
 				AgentId: global.AgentId,
 				Function: request.Function{
 					Method:        "RPC",
-					Url:           md.Get("requestURL")[0],
-					Uri:           md.Get("requestURI")[0],
-					Protocol:      md.Get("protocol")[0],
+					Url:           info.FullMethod,
+					Uri:           info.FullMethod,
+					Protocol:      "ProtoBuf",
 					ClientIp:      "",
 					Language:      "GO",
 					Scheme:        "GRPC",
@@ -52,7 +72,7 @@ func interceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInf
 					ReqBody:       "",
 					QueryString:   "",
 					Pool:          []request.Pool{},
-					TraceId:       md.Get("dt-traceid")[0],
+					TraceId:       newId,
 				},
 			},
 		}

core/http/httpServeHTTP/replacement.go

@@ -10,6 +10,7 @@ import (
 	"github.com/HXSecurity/DongTai-agent-go/utils"
 	"net/http"
 	"reflect"
+	"strconv"
 	"strings"
 )
 
@@ -23,6 +24,8 @@ func MyServer(server *http.ServeMux, w http.ResponseWriter, r *http.Request) {
 		for k, v := range r.Header {
 			headerBase += k + ": " + strings.Join(v, ",") + "\n"
 		}
+		tranceID := global.TraceId + "." + strconv.Itoa(global.AgentId) + ".0.0.0"
+		headerBase += "dt-traceid:" + tranceID
 		if t.Kind() == reflect.Ptr {
 			buf := t.
 				Elem().

core/httpRouter/httpRouter/replacement.go

@@ -12,6 +12,7 @@ import (
 	"github.com/julienschmidt/httprouter"
 	"net/http"
 	"reflect"
+	"strconv"
 	"strings"
 )
 
@@ -35,6 +36,8 @@ func MyHttpRouterServer(server *httprouter.Router, w http.ResponseWriter, r *htt
 		for k, v := range r.Header {
 			headerBase += k + ": " + strings.Join(v, ",") + "\n"
 		}
+		tranceID := global.TraceId + "." + strconv.Itoa(global.AgentId) + ".0.0.0"
+		headerBase += "dt-traceid:" + tranceID
 		if t.Kind() == reflect.Ptr {
 			buf := t.
 				Elem().

run/base/base.go

@@ -22,7 +22,8 @@ import (
 	_ "github.com/HXSecurity/DongTai-agent-go/core/base/urlUrlString"
 	"github.com/HXSecurity/DongTai-agent-go/hook"
 	"github.com/HXSecurity/DongTai-agent-go/service"
-	"github.com/google/uuid"
+	"github.com/HXSecurity/DongTai-agent-go/utils"
+	"strconv"
 
 	"github.com/HXSecurity/DongTai-agent-go/global"
 )
@@ -33,5 +34,7 @@ func init() {
 	global.InitViper()
 	_ = service.AgentRegister()
 	b.HookAll()
-	global.TraceId = uuid.New().String() + "-" + uuid.New().String()
+	worker, _ := utils.NewWorker(global.AgentId)
+	strconv.Itoa(int(worker.GetId()))
+	global.TraceId = strconv.Itoa(int(worker.GetId())) + "-" + strconv.Itoa(int(worker.GetId()))
 }