add log and fixed trace bug

piexlmax · piexlmax · commit de94df15f363 · 2022-05-13T10:53:43.000+08:00
diff --git a/api/base_config.go b/api/base_config.go
@@ -5,12 +5,14 @@ import (
 	"github.com/HXSecurity/DongTai-agent-go/global"
 	"github.com/HXSecurity/DongTai-agent-go/utils"
 	"github.com/parnurzeal/gorequest"
+	"time"
 )
 
 // NewRequest
 // 空白的请求工具 预留做后续的统一参数处理
 func NewRequest() *gorequest.SuperAgent {
 	request := gorequest.New()
+	request.Timeout(30 * time.Second)
 	return request
 }
 
diff --git a/api/engine.go b/api/engine.go
@@ -13,7 +13,9 @@ import (
 	req: api.AgentRegisterReq
 */
 func AgentRegister(req request.AgentRegisterReq) (AgentId int, err error) {
+	fmt.Println("探针注册开始等待返回")
 	resp, body, errs := POST("/api/v1/agent/register", req).End()
+	fmt.Println("探针注册返回:", string(body))
 	if len(errs) > 0 {
 		for _, v := range errs {
 			fmt.Println(v)
@@ -31,7 +33,7 @@ func AgentRegister(req request.AgentRegisterReq) (AgentId int, err error) {
 			AgentId = res.Data.Id
 			fmt.Println("注册成功，探针ID为", res.Data.Id)
 			return AgentId, nil
-		} 
+		}
 		return 0, errors.New("注册失败，失败原因" + res.Msg)
 
 	}
diff --git a/core/base/execCmdRun/install.go b/core/base/execCmdRun/install.go
diff --git a/core/base/execCmdStart/install.go b/core/base/execCmdStart/install.go
@@ -0,0 +1,29 @@
+package execCmdStart
+
+import (
+	"fmt"
+	"github.com/HXSecurity/DongTai-agent-go/model"
+	"github.com/brahma-adshonor/gohook"
+	"os/exec"
+)
+
+func init() {
+	model.HookMap["execCmdStart"] = new(ExecCmdStart)
+}
+
+type ExecCmdStart struct {
+}
+
+func (h *ExecCmdStart) Hook() {
+	var cmd *exec.Cmd
+	err := gohook.HookMethod(cmd, "Start", Start, StartT)
+	if err != nil {
+		fmt.Println(err, "execCmdStart")
+	} else {
+		fmt.Println("execCmdStart")
+	}
+}
+
+func (h *ExecCmdStart) UnHook() {
+	gohook.UnHook(exec.Command)
+}
diff --git a/core/base/execCmdStart/replacement.go b/core/base/execCmdStart/replacement.go
@@ -1,13 +1,13 @@
-package execCmdRun
+package execCmdStart
 
 import (
 	"github.com/HXSecurity/DongTai-agent-go/model/request"
 	"os/exec"
 	"reflect"
 )
 
-func Run(cmd *exec.Cmd) error {
-	e := RunT(cmd)
+func Start(cmd *exec.Cmd) error {
+	e := StartT(cmd)
 	//
 	var u uintptr
 	value := reflect.ValueOf(cmd)
@@ -18,12 +18,12 @@ func Run(cmd *exec.Cmd) error {
 		NeedHook:        request.Collect(u),
 		Source:          false,
 		OriginClassName: "exec.(*Cmd)",
-		MethodName:      "Run",
+		MethodName:      "Start",
 		ClassName:       "exec.(*Cmd)",
 	})
 	return e
 }
 
-func RunT(cmd *exec.Cmd) error {
+func StartT(cmd *exec.Cmd) error {
 	return nil
 }
diff --git a/core/grpc/clientConn/replacement.go b/core/grpc/clientConn/replacement.go
@@ -62,6 +62,7 @@ func Invoke(cl *grpc.ClientConn, ctx context.Context, method string, args, reply
 		OriginClassName: "grpc.(*ClientConn)",
 		MethodName:      "Invoke",
 		ClassName:       "grpc.(*ClientConn)",
+		TraceId:         tranceid,
 	})
 	return err
 }
diff --git a/core/grpc/newServer/replacement.go b/core/grpc/newServer/replacement.go
@@ -11,7 +11,6 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/metadata"
 	"strconv"
-	"strings"
 )
 
 const (
@@ -36,17 +35,6 @@ func interceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInf
 	if len(dt) != 0 {
 		Traceid = dt[0]
 	}
-	four := strconv.Itoa(int(worker.GetId()))
-	tranceids := strings.Split(Traceid, ".")
-	tranceids[OnlyKey] = four
-	newId := ""
-	for i := 0; i < len(tranceids); i++ {
-		if i == OnlyKey {
-			newId += tranceids[i]
-		} else {
-			newId += tranceids[i] + "."
-		}
-	}
 
 	id := utils.CatGoroutineID()
 	request.FmtHookPool(request.PoolReq{
@@ -56,13 +44,14 @@ func interceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInf
 		OriginClassName: "grpc",
 		MethodName:      "NewServer",
 		ClassName:       "grpc",
+		TraceId:         Traceid,
 	})
 	// 获取metadata
 	res, err := handler(ctx, req)
 	go func() {
 		worker, _ := utils.NewWorker(global.AgentId)
 		onlyKey := int(worker.GetId())
-		header := base64.StdEncoding.EncodeToString([]byte("dt-traceid:" + newId))
+		header := base64.StdEncoding.EncodeToString([]byte("dt-traceid:" + Traceid))
 		HookGroup := &request.UploadReq{
 			Type:     36,
 			InvokeId: onlyKey,
@@ -81,7 +70,7 @@ func interceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInf
 					ReqBody:       "",
 					QueryString:   "",
 					Pool:          []request.Pool{},
-					TraceId:       newId,
+					TraceId:       Traceid,
 				},
 			},
 		}
diff --git a/hook/base.go b/hook/base.go
@@ -12,7 +12,7 @@ func (b *Base) GetHook() []string {
 		"jsonNewDecoder",
 		"runtimeConcatstrings",
 		"execCommand",
-		"execCmdRun",
+		"execCmdStart",
 		"bufioWriterWrite",
 		"bufioWriterWriteString",
 		"runtimesSringtoslicebyte",
diff --git a/model/request/engine.go b/model/request/engine.go
@@ -121,6 +121,7 @@ type PoolReq struct {
 	NeedHook        []interface{}
 	NeedCatch       []interface{}
 	ArgsStr         string
+	TraceId         string
 }
 
 type Log struct {
@@ -161,9 +162,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)"){
-		p.Pool.TraceId = TraceId
-	}
+	//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)
@@ -249,6 +250,7 @@ func FmtHookPool(p PoolReq) Pool {
 		CallerMethod:     callerMethod,
 		RetClassName:     RetClassNames,
 		Args:             ArgsStr,
+		TraceId:          p.TraceId,
 	}
 
 	poolTree := PoolTree{
diff --git a/run/base/base.go b/run/base/base.go
@@ -3,7 +3,7 @@ package base
 import (
 	_ "github.com/HXSecurity/DongTai-agent-go/core/base/bufioWriterWrite"
 	_ "github.com/HXSecurity/DongTai-agent-go/core/base/bufioWriterWriteString"
-	_ "github.com/HXSecurity/DongTai-agent-go/core/base/execCmdRun"
+	_ "github.com/HXSecurity/DongTai-agent-go/core/base/execCmdStart"
 	_ "github.com/HXSecurity/DongTai-agent-go/core/base/execCommand"
 	_ "github.com/HXSecurity/DongTai-agent-go/core/base/fmtSprintf"
 	_ "github.com/HXSecurity/DongTai-agent-go/core/base/htmlTemplateExecuteTemplate"
diff --git a/service/engine.go b/service/engine.go
@@ -73,6 +73,7 @@ func RunAgent() {
 }
 
 func AgentRegister() (err error) {
+	max := 6
 	breaker := CreateCircuitBreaker()
 	live = true
 	OS := runtime.GOOS
@@ -146,7 +147,8 @@ func AgentRegister() (err error) {
 	}
 	go func() {
 		for {
-			fmt.Println("等待当前程序http启动完成")
+			max = max - 1
+			fmt.Printf("等待当前程序当前程序启动完成，剩余%d次", max)
 			time.Sleep(1 * time.Second)
 			ip, err := utils.ExternalIP()
 			if err != nil {
@@ -164,7 +166,9 @@ func AgentRegister() (err error) {
 				cmd.Stderr = &strErr
 				cmd.Stdout = &out
 			}
+			fmt.Println("等待当前程序端口返回")
 			err = cmd.Run()
+			fmt.Println("等待当前程序端口返回成功")
 			output := out.String()
 			if err != nil {
 				return
@@ -187,10 +191,15 @@ func AgentRegister() (err error) {
 				r := regexp.MustCompile(`:::\s*(.*?)\s* `)
 				matches = r.FindAllStringSubmatch(output, -1)
 			}
-			if matches != nil {
-				if matches[0] != nil {
-					if matches[0][1] != "" {
-						req.ServerPort = matches[0][1]
+			if matches != nil || max == 0 {
+				if matches[0] != nil || max == 0 {
+					if matches[0][1] != "" || max == 0 {
+						fmt.Println("当前程序启动完成")
+						if max == 0 {
+							req.ServerPort = ""
+						} else {
+							req.ServerPort = matches[0][1]
+						}
 						req.ServerAddr = ip.String()
 						agentId, err := api.AgentRegister(req)
 						if err != nil {

Original file line number	Diff line number	Diff line change
`@@ -5,12 +5,14 @@ import (`
`5`	`5`	`"github.com/HXSecurity/DongTai-agent-go/global"`
`6`	`6`	`"github.com/HXSecurity/DongTai-agent-go/utils"`
`7`	`7`	`"github.com/parnurzeal/gorequest"`
	`8`	`+ "time"`
`8`	`9`	`)`
`9`	`10`
`10`	`11`	`// NewRequest`
`11`	`12`	`// 空白的请求工具预留做后续的统一参数处理`
`12`	`13`	`func NewRequest() *gorequest.SuperAgent {`
`13`	`14`	`request := gorequest.New()`
	`15`	`+ request.Timeout(30 * time.Second)`
`14`	`16`	`return request`
`15`	`17`	`}`
`16`	`18`
Original file line number	Diff line number	Diff line change
`@@ -62,6 +62,7 @@ func Invoke(cl *grpc.ClientConn, ctx context.Context, method string, args, reply`
`62`	`62`	`OriginClassName: "grpc.(*ClientConn)",`
`63`	`63`	`MethodName: "Invoke",`
`64`	`64`	`ClassName: "grpc.(*ClientConn)",`
	`65`	`+ TraceId: tranceid,`
`65`	`66`	`})`
`66`	`67`	`return err`
`67`	`68`	`}`