github.com/HXSecurity/DongTai-agent-go@v0.4.2/core/grpc/clientConn/replacement.go (about)

     1  package clientConn
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"strconv"
     7  	"strings"
     8  
     9  	"github.com/HXSecurity/DongTai-agent-go/global"
    10  	"github.com/HXSecurity/DongTai-agent-go/model/request"
    11  	"github.com/HXSecurity/DongTai-agent-go/utils"
    12  	"google.golang.org/grpc"
    13  	"google.golang.org/grpc/metadata"
    14  )
    15  
    16  const (
    17  	TraceId = iota
    18  	AgentId
    19  	RoutineId
    20  	NextKey
    21  	OnlyKey
    22  )
    23  
    24  func Invoke(cl *grpc.ClientConn, ctx context.Context, method string, args, reply interface{}, opts ...grpc.CallOption) error {
    25  	outmd, _ := metadata.FromIncomingContext(ctx)
    26  	worker, _ := utils.NewWorker(global.AgentId)
    27  	var tranceid string
    28  	if len(outmd.Get("dt-traceid")) > 0 {
    29  		tranceid = outmd.Get("dt-traceid")[0]
    30  	}
    31  	if tranceid == "" {
    32  		tranceid = global.TargetTraceId + "." + strconv.Itoa(global.AgentId) + ".0.1." + strconv.Itoa(int(worker.GetId()))
    33  	} else {
    34  		four := strconv.Itoa(int(worker.GetId()))
    35  		tranceids := strings.Split(tranceid, ".")
    36  		tranceids[AgentId] = strconv.Itoa(global.AgentId)
    37  		num, _ := strconv.Atoi(tranceids[NextKey])
    38  		tranceids[NextKey] = strconv.Itoa(num + 1)
    39  		tranceids[OnlyKey] = four
    40  		newId := ""
    41  		for i := 0; i < len(tranceids); i++ {
    42  			if i == OnlyKey {
    43  				newId += tranceids[i]
    44  			} else {
    45  				newId += tranceids[i] + "."
    46  			}
    47  		}
    48  		tranceid = newId
    49  	}
    50  	md := metadata.Pairs("dt-traceid", tranceid,
    51  		"protocol", "ProtoBuf",
    52  		"requestURL", cl.Target()+method,
    53  		"requestURI", method,
    54  		"headers", "traceid:"+tranceid,
    55  	)
    56  	fmt.Println(tranceid)
    57  	ctx = metadata.NewOutgoingContext(ctx, md)
    58  	err := InvokeT(cl, ctx, method, args, reply, opts...)
    59  	request.FmtHookPool(request.PoolReq{
    60  		Args:            request.Collect(args),
    61  		Reqs:            request.Collect(reply),
    62  		Source:          false,
    63  		OriginClassName: "grpc.(*ClientConn)",
    64  		MethodName:      "Invoke",
    65  		ClassName:       "grpc.(*ClientConn)",
    66  		TraceId:         tranceid,
    67  		Plugin:          "GRPC",
    68  	})
    69  	return err
    70  }
    71  func InvokeT(cl *grpc.ClientConn, ctx context.Context, method string, args, reply interface{}, opts ...grpc.CallOption) error {
    72  	return nil
    73  }