Skip to content

Commit 8237d44

Browse files
author
piexlmax
committed
fixed:grpc
1 parent 5458479 commit 8237d44

File tree

11 files changed

+56
-27
lines changed

11 files changed

+56
-27
lines changed

‎core/gin/ginEngineServerHTTP/replacement.go

+7-3
Original file line numberDiff line numberDiff line change
@@ -18,14 +18,18 @@ import (
1818
func MyServer(server *gin.Engine, w http.ResponseWriter, r *http.Request) {
1919
MyServerTemp(server, w, r)
2020
id := utils.CatGoroutineID()
21+
worker, _ := utils.NewWorker(global.AgentId)
2122
go func() {
2223
t := reflect.ValueOf(r.Body)
2324
var headerBase string
2425
body := ""
2526
for k, v := range r.Header {
2627
headerBase += k + ": " + strings.Join(v, ",") + "\n"
2728
}
28-
tranceID := global.TraceId + "." + strconv.Itoa(global.AgentId) + ".0.0.0"
29+
30+
TraceId := global.TraceId + "-" + strconv.Itoa(int(worker.GetId()))
31+
global.TargetTraceId = TraceId
32+
tranceID := TraceId + "." + strconv.Itoa(global.AgentId) + ".0.0.0"
2933
headerBase += "dt-traceid:" + tranceID
3034
if t.Kind() == reflect.Ptr {
3135
buf := t.
@@ -52,7 +56,6 @@ func MyServer(server *gin.Engine, w http.ResponseWriter, r *http.Request) {
5256
if r.TLS != nil {
5357
scheme = "https"
5458
}
55-
worker, _ := utils.NewWorker(global.AgentId)
5659
onlyKey := int(worker.GetId())
5760
HookGroup := &request.UploadReq{
5861
Type: 36,
@@ -71,6 +74,7 @@ func MyServer(server *gin.Engine, w http.ResponseWriter, r *http.Request) {
7174
ReqBody: body,
7275
QueryString: r.URL.RawQuery,
7376
Pool: []request.Pool{},
77+
TraceId: tranceID,
7478
},
7579
},
7680
}
@@ -97,7 +101,7 @@ func MyServer(server *gin.Engine, w http.ResponseWriter, r *http.Request) {
97101
global.PoolTreeMap.Range(func(key, value interface{}) bool {
98102
if value.(*request.PoolTree).IsThisBegin(id) {
99103
global.PoolTreeMap.Delete(key)
100-
value.(*request.PoolTree).FMT(&HookGroup.Detail.Function.Pool, worker, goroutineIDs, "")
104+
value.(*request.PoolTree).FMT(&HookGroup.Detail.Function.Pool, worker, goroutineIDs, HookGroup.Detail.Function.TraceId)
101105
return false
102106
}
103107
return true

‎core/goChi/httpRouter/replacement.go

+2-1
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,7 @@ func ChiRouterServer(server *chi.Mux, w http.ResponseWriter, r *http.Request) {
8181
ReqBody: body,
8282
QueryString: r.URL.RawQuery,
8383
Pool: []request.Pool{},
84+
//TraceId: tranceID,
8485
},
8586
},
8687
}
@@ -106,7 +107,7 @@ func ChiRouterServer(server *chi.Mux, w http.ResponseWriter, r *http.Request) {
106107
global.PoolTreeMap.Range(func(key, value interface{}) bool {
107108
if value.(*request.PoolTree).IsThisBegin(id) {
108109
global.PoolTreeMap.Delete(key)
109-
value.(*request.PoolTree).FMT(&HookGroup.Detail.Function.Pool, worker, goroutineIDs, "")
110+
value.(*request.PoolTree).FMT(&HookGroup.Detail.Function.Pool, worker, goroutineIDs, HookGroup.Detail.Function.TraceId)
110111
return true
111112
}
112113
return true

‎core/gorilla/gorillaRpcServerHTTP/replacement.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,7 @@ func MyServer(server *rpc.Server, w http.ResponseWriter, r *http.Request) {
9999
if value.(*request.PoolTree).IsThisBegin(id) {
100100
onlyKey += 1
101101
global.PoolTreeMap.Delete(key)
102-
value.(*request.PoolTree).FMT(&HookGroup.Detail.Function.Pool, worker, goroutineIDs, "")
102+
value.(*request.PoolTree).FMT(&HookGroup.Detail.Function.Pool, worker, goroutineIDs, HookGroup.Detail.Function.TraceId)
103103
return false
104104
}
105105
return true

‎core/grpc/clientConn/replacement.go

+16-7
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package clientConn
22

33
import (
44
"context"
5+
"fmt"
56
"github.com/HXSecurity/DongTai-agent-go/global"
67
"github.com/HXSecurity/DongTai-agent-go/model/request"
78
"github.com/HXSecurity/DongTai-agent-go/utils"
@@ -11,26 +12,33 @@ import (
1112
"strings"
1213
)
1314

15+
const (
16+
TraceId = iota
17+
AgentId
18+
RoutineId
19+
NextKey
20+
OnlyKey
21+
)
22+
1423
func Invoke(cl *grpc.ClientConn, ctx context.Context, method string, args, reply interface{}, opts ...grpc.CallOption) error {
1524
outmd, _ := metadata.FromIncomingContext(ctx)
1625
worker, _ := utils.NewWorker(global.AgentId)
1726
var tranceid string
1827
if len(outmd.Get("dt-traceid")) > 0 {
1928
tranceid = outmd.Get("dt-traceid")[0]
2029
}
21-
2230
if tranceid == "" {
23-
tranceid = global.TraceId + "." + strconv.Itoa(global.AgentId) + ".0.0." + strconv.Itoa(int(worker.GetId()))
31+
tranceid = global.TargetTraceId + "." + strconv.Itoa(global.AgentId) + ".0.1." + strconv.Itoa(int(worker.GetId()))
2432
} else {
2533
four := strconv.Itoa(int(worker.GetId()))
2634
tranceids := strings.Split(tranceid, ".")
27-
tranceids[1] = strconv.Itoa(global.AgentId)
28-
num, _ := strconv.Atoi(tranceids[3])
29-
tranceids[3] = strconv.Itoa(num + 1)
30-
tranceids[4] = four
35+
tranceids[AgentId] = strconv.Itoa(global.AgentId)
36+
num, _ := strconv.Atoi(tranceids[NextKey])
37+
tranceids[NextKey] = strconv.Itoa(num + 1)
38+
tranceids[OnlyKey] = four
3139
newId := ""
3240
for i := 0; i < len(tranceids); i++ {
33-
if i == 4 {
41+
if i == OnlyKey {
3442
newId += tranceids[i]
3543
} else {
3644
newId += tranceids[i] + "."
@@ -44,6 +52,7 @@ func Invoke(cl *grpc.ClientConn, ctx context.Context, method string, args, reply
4452
"requestURI", method,
4553
"headers", "traceid:"+tranceid,
4654
)
55+
fmt.Println(tranceid)
4756
ctx = metadata.NewOutgoingContext(ctx, md)
4857
err := InvokeT(cl, ctx, method, args, reply, opts...)
4958
request.FmtHookPool(request.PoolReq{

‎core/grpc/newServer/replacement.go

+11-6
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,14 @@ import (
1414
"strings"
1515
)
1616

17+
const (
18+
TraceId = iota
19+
AgentId
20+
RoutineId
21+
NextKey
22+
OnlyKey
23+
)
24+
1725
func NewServer(opt ...grpc.ServerOption) *grpc.Server {
1826
opt = append(opt, grpc.UnaryInterceptor(interceptor))
1927
return NewServerT(opt...)
@@ -26,19 +34,16 @@ func interceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInf
2634
worker, _ := utils.NewWorker(global.AgentId)
2735
four := strconv.Itoa(int(worker.GetId()))
2836
tranceids := strings.Split(Traceid, ".")
29-
tranceids[1] = strconv.Itoa(global.AgentId)
30-
num, _ := strconv.Atoi(tranceids[3])
31-
tranceids[3] = strconv.Itoa(num + 1)
32-
tranceids[4] = four
37+
tranceids[AgentId] = strconv.Itoa(global.AgentId)
38+
tranceids[OnlyKey] = four
3339
newId := ""
3440
for i := 0; i < len(tranceids); i++ {
35-
if i == 4 {
41+
if i == OnlyKey {
3642
newId += tranceids[i]
3743
} else {
3844
newId += tranceids[i] + "."
3945
}
4046
}
41-
global.TraceId = tranceids[0]
4247

4348
id := utils.CatGoroutineID()
4449
request.FmtHookPool(request.PoolReq{

‎core/http/httpServeHTTP/replacement.go

+7-3
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,10 @@ import (
1515
)
1616

1717
func MyServer(server *http.ServeMux, w http.ResponseWriter, r *http.Request) {
18+
worker, _ := utils.NewWorker(global.AgentId)
19+
20+
TraceId := global.TraceId + "-" + strconv.Itoa(int(worker.GetId()))
21+
global.TargetTraceId = TraceId
1822
MyServerTemp(server, w, r)
1923
id := utils.CatGoroutineID()
2024
go func() {
@@ -24,7 +28,7 @@ func MyServer(server *http.ServeMux, w http.ResponseWriter, r *http.Request) {
2428
for k, v := range r.Header {
2529
headerBase += k + ": " + strings.Join(v, ",") + "\n"
2630
}
27-
tranceID := global.TraceId + "." + strconv.Itoa(global.AgentId) + ".0.0.0"
31+
tranceID := TraceId + "." + strconv.Itoa(global.AgentId) + ".0.0.0"
2832
headerBase += "dt-traceid:" + tranceID
2933
if t.Kind() == reflect.Ptr {
3034
buf := t.
@@ -51,7 +55,6 @@ func MyServer(server *http.ServeMux, w http.ResponseWriter, r *http.Request) {
5155
if r.TLS != nil {
5256
scheme = "https"
5357
}
54-
worker, _ := utils.NewWorker(global.AgentId)
5558
onlyKey := int(worker.GetId())
5659

5760
HookGroup := &request.UploadReq{
@@ -71,6 +74,7 @@ func MyServer(server *http.ServeMux, w http.ResponseWriter, r *http.Request) {
7174
ReqBody: body,
7275
QueryString: r.URL.RawQuery,
7376
Pool: []request.Pool{},
77+
TraceId: tranceID,
7478
},
7579
},
7680
}
@@ -96,7 +100,7 @@ func MyServer(server *http.ServeMux, w http.ResponseWriter, r *http.Request) {
96100
global.PoolTreeMap.Range(func(key, value interface{}) bool {
97101
if value.(*request.PoolTree).IsThisBegin(id) {
98102
global.PoolTreeMap.Delete(key)
99-
value.(*request.PoolTree).FMT(&HookGroup.Detail.Function.Pool, worker, goroutineIDs, "")
103+
value.(*request.PoolTree).FMT(&HookGroup.Detail.Function.Pool, worker, goroutineIDs, HookGroup.Detail.Function.TraceId)
100104
return false
101105
}
102106
return true

‎core/httpRouter/httpRouter/replacement.go

+6-4
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ import (
1717
)
1818

1919
func MyHttpRouterServer(server *httprouter.Router, w http.ResponseWriter, r *http.Request) {
20-
20+
worker, err := utils.NewWorker(global.AgentId)
2121
request.FmtHookPool(request.PoolReq{
2222
Args: request.Collect(r.Host),
2323
Reqs: request.Collect(r.Host),
@@ -36,7 +36,9 @@ func MyHttpRouterServer(server *httprouter.Router, w http.ResponseWriter, r *htt
3636
for k, v := range r.Header {
3737
headerBase += k + ": " + strings.Join(v, ",") + "\n"
3838
}
39-
tranceID := global.TraceId + "." + strconv.Itoa(global.AgentId) + ".0.0.0"
39+
TraceId := global.TraceId + "-" + strconv.Itoa(int(worker.GetId()))
40+
global.TargetTraceId = TraceId
41+
tranceID := TraceId + "." + strconv.Itoa(global.AgentId) + ".0.0.0"
4042
headerBase += "dt-traceid:" + tranceID
4143
if t.Kind() == reflect.Ptr {
4244
buf := t.
@@ -63,7 +65,6 @@ func MyHttpRouterServer(server *httprouter.Router, w http.ResponseWriter, r *htt
6365
if r.TLS != nil {
6466
scheme = "https"
6567
}
66-
worker, err := utils.NewWorker(global.AgentId)
6768
if err != nil {
6869
fmt.Println(err)
6970
}
@@ -85,6 +86,7 @@ func MyHttpRouterServer(server *httprouter.Router, w http.ResponseWriter, r *htt
8586
ReqBody: body,
8687
QueryString: r.URL.RawQuery,
8788
Pool: []request.Pool{},
89+
TraceId: tranceID,
8890
},
8991
},
9092
}
@@ -110,7 +112,7 @@ func MyHttpRouterServer(server *httprouter.Router, w http.ResponseWriter, r *htt
110112
global.PoolTreeMap.Range(func(key, value interface{}) bool {
111113
if value.(*request.PoolTree).IsThisBegin(id) {
112114
global.PoolTreeMap.Delete(key)
113-
value.(*request.PoolTree).FMT(&HookGroup.Detail.Function.Pool, worker, goroutineIDs, "")
115+
value.(*request.PoolTree).FMT(&HookGroup.Detail.Function.Pool, worker, goroutineIDs, HookGroup.Detail.Function.TraceId)
114116
return true
115117
}
116118
return true

‎global/status.go

+1
Original file line numberDiff line numberDiff line change
@@ -25,4 +25,5 @@ var (
2525
ResponseMap = sync.Map{}
2626
ResponseHeaderMap = sync.Map{}
2727
TraceId string
28+
TargetTraceId string
2829
)

‎hook/http.go

+3
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,9 @@ type Http struct {
66
func (h *Http) GetHook() []string {
77
return []string{
88
"httpServeHTTP",
9+
"httpRequestFormValue",
10+
"httpServeHTTP",
11+
"urlURLQuery",
912
}
1013
}
1114

‎run/base/base.go

+1-2
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,5 @@ func init() {
3535
_ = service.AgentRegister()
3636
b.HookAll()
3737
worker, _ := utils.NewWorker(global.AgentId)
38-
strconv.Itoa(int(worker.GetId()))
39-
global.TraceId = strconv.Itoa(int(worker.GetId())) + "-" + strconv.Itoa(int(worker.GetId()))
38+
global.TraceId = strconv.Itoa(int(worker.GetId()))
4039
}

‎run/http/hookHttp.go

+1
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import (
44
_ "github.com/HXSecurity/DongTai-agent-go/core/http/httpRequestCookie"
55
_ "github.com/HXSecurity/DongTai-agent-go/core/http/httpRequestFormValue"
66
_ "github.com/HXSecurity/DongTai-agent-go/core/http/httpServeHTTP"
7+
_ "github.com/HXSecurity/DongTai-agent-go/core/http/urlURLQuery"
78
"github.com/HXSecurity/DongTai-agent-go/global"
89
"github.com/HXSecurity/DongTai-agent-go/hook"
910
)

0 commit comments

Comments
 (0)