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

     1  package newServer
     2  
     3  import (
     4  	"context"
     5  	"encoding/base64"
     6  	"fmt"
     7  	"github.com/HXSecurity/DongTai-agent-go/api"
     8  	"github.com/HXSecurity/DongTai-agent-go/global"
     9  	"github.com/HXSecurity/DongTai-agent-go/model/request"
    10  	"github.com/HXSecurity/DongTai-agent-go/utils"
    11  	"google.golang.org/grpc"
    12  	"google.golang.org/grpc/metadata"
    13  	"strconv"
    14  )
    15  
    16  const (
    17  	TraceId = iota
    18  	AgentId
    19  	RoutineId
    20  	NextKey
    21  	OnlyKey
    22  )
    23  
    24  func NewServer(opt ...grpc.ServerOption) *grpc.Server {
    25  	opt = append(opt, grpc.UnaryInterceptor(interceptor))
    26  	return NewServerT(opt...)
    27  }
    28  
    29  // interceptor 一元拦截器
    30  func interceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
    31  	md, _ := metadata.FromIncomingContext(ctx)
    32  	worker, _ := utils.NewWorker(global.AgentId)
    33  	dt := md.Get("dt-traceid")
    34  	var Traceid = global.TargetTraceId + "." + strconv.Itoa(global.AgentId) + ".0.0." + strconv.Itoa(int(worker.GetId()))
    35  	if len(dt) != 0 {
    36  		Traceid = dt[0]
    37  	}
    38  
    39  	id := utils.CatGoroutineID()
    40  	request.FmtHookPool(request.PoolReq{
    41  		Reqs:            request.Collect(req),
    42  		Args:            request.Collect(req),
    43  		Source:          true,
    44  		OriginClassName: "grpc",
    45  		MethodName:      "NewServer",
    46  		ClassName:       "grpc",
    47  		TraceId:         Traceid,
    48  	})
    49  	// 获取metadata
    50  	res, err := handler(ctx, req)
    51  	go func() {
    52  		worker, _ := utils.NewWorker(global.AgentId)
    53  		onlyKey := int(worker.GetId())
    54  		header := base64.StdEncoding.EncodeToString([]byte("dt-traceid:" + Traceid))
    55  		HookGroup := &request.UploadReq{
    56  			Type:     36,
    57  			InvokeId: onlyKey,
    58  			Detail: request.Detail{
    59  				AgentId: global.AgentId,
    60  				Function: request.Function{
    61  					Method:        "RPC",
    62  					Url:           info.FullMethod,
    63  					Uri:           info.FullMethod,
    64  					Protocol:      "ProtoBuf",
    65  					ClientIp:      "",
    66  					Language:      "GO",
    67  					Scheme:        "GRPC",
    68  					ReplayRequest: false,
    69  					ReqHeader:     header,
    70  					ReqBody:       "",
    71  					QueryString:   "",
    72  					Pool:          []request.Pool{},
    73  					TraceId:       Traceid,
    74  				},
    75  			},
    76  		}
    77  
    78  		goroutineIDs := make(map[string]bool)
    79  		global.PoolTreeMap.Range(func(key, value interface{}) bool {
    80  			if value.(*request.PoolTree).IsThisBegin(id) {
    81  				global.PoolTreeMap.Delete(key)
    82  				value.(*request.PoolTree).FMT(&HookGroup.Detail.Function.Pool, worker, goroutineIDs, HookGroup.Detail.Function.TraceId)
    83  				return false
    84  			}
    85  			return true
    86  		})
    87  
    88  		fmt.Println(HookGroup.Detail.Function.Url)
    89  		api.ReportUpload(*HookGroup)
    90  		request.RunMapGCbYGoroutineID(goroutineIDs)
    91  	}()
    92  	return res, err
    93  }
    94  
    95  func NewServerT(opt ...grpc.ServerOption) *grpc.Server {
    96  	return nil
    97  }