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

     1  package kafkaWriter
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"reflect"
     7  	"strconv"
     8  	"strings"
     9  
    10  	"github.com/HXSecurity/DongTai-agent-go/global"
    11  	"github.com/HXSecurity/DongTai-agent-go/model/request"
    12  	"github.com/HXSecurity/DongTai-agent-go/utils"
    13  	"github.com/segmentio/kafka-go"
    14  	"google.golang.org/grpc/metadata"
    15  )
    16  
    17  const (
    18  	TraceId = iota
    19  	AgentId
    20  	RoutineId
    21  	NextKey
    22  	OnlyKey
    23  )
    24  
    25  func WriteMessages(w *kafka.Writer, ctx context.Context, msgs ...kafka.Message) error {
    26  	for idx, msg := range msgs {
    27  		traceId := getTraceId(ctx)
    28  		msgs[idx].Headers = append(msg.Headers, kafka.Header{
    29  			Key:   "dt-traceid",
    30  			Value: []byte(traceId),
    31  		})
    32  
    33  		fmt.Println("traceId:", traceId, "msg:", msg.Value)
    34  		v := reflect.ValueOf(msg.Value)
    35  		request.FmtHookPool(request.PoolReq{
    36  			Args:            request.Collect(msg.Value),
    37  			NeedHook:        request.Collect(v.Pointer()),
    38  			Source:          false,
    39  			OriginClassName: "kafka.(*Writer)",
    40  			MethodName:      "WriteMessages",
    41  			ClassName:       "kafka.(*Writer)",
    42  			TraceId:         traceId,
    43  			Plugin:          "KAFKA",
    44  		})
    45  	}
    46  
    47  	err := WriteMessagesT(w, ctx, msgs...)
    48  	return err
    49  }
    50  
    51  func getTraceId(ctx context.Context) string {
    52  	outmd, _ := metadata.FromIncomingContext(ctx)
    53  	worker, _ := utils.NewWorker(global.AgentId)
    54  	var tranceid string
    55  	if len(outmd.Get("dt-traceid")) > 0 {
    56  		tranceid = outmd.Get("dt-traceid")[0]
    57  	}
    58  	if tranceid == "" {
    59  		tranceid = global.TargetTraceId + "." + strconv.Itoa(global.AgentId) + ".0.1." + strconv.Itoa(int(worker.GetId()))
    60  	} else {
    61  		four := strconv.Itoa(int(worker.GetId()))
    62  		tranceids := strings.Split(tranceid, ".")
    63  		tranceids[AgentId] = strconv.Itoa(global.AgentId)
    64  		num, _ := strconv.Atoi(tranceids[NextKey])
    65  		tranceids[NextKey] = strconv.Itoa(num + 1)
    66  		tranceids[OnlyKey] = four
    67  		newId := ""
    68  		for i := 0; i < len(tranceids); i++ {
    69  			if i == OnlyKey {
    70  				newId += tranceids[i]
    71  			} else {
    72  				newId += tranceids[i] + "."
    73  			}
    74  		}
    75  		tranceid = newId
    76  	}
    77  
    78  	return tranceid
    79  }
    80  
    81  func WriteMessagesT(w *kafka.Writer, ctx context.Context, msgs ...kafka.Message) error {
    82  	return nil
    83  }