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

     1  package httpServeHTTP
     2  
     3  import (
     4  	"bufio"
     5  	"bytes"
     6  	"encoding/base64"
     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  	"net/http"
    12  	"reflect"
    13  	"strconv"
    14  	"strings"
    15  )
    16  
    17  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
    22  	MyServerTemp(server, w, r)
    23  	id := utils.CatGoroutineID()
    24  	go func() {
    25  		t := reflect.ValueOf(r.Body)
    26  		var headerBase string
    27  		body := ""
    28  		for k, v := range r.Header {
    29  			headerBase += k + ": " + strings.Join(v, ",") + "\n"
    30  		}
    31  		tranceID := TraceId + "." + strconv.Itoa(global.AgentId) + ".0.0.0"
    32  		headerBase += "dt-traceid:" + tranceID
    33  		if t.Kind() == reflect.Ptr {
    34  			buf := t.
    35  				Elem().
    36  				FieldByName("src").
    37  				Elem().Elem().
    38  				FieldByName("R").
    39  				Elem().Elem().
    40  				FieldByName("buf").Bytes()
    41  			buf = buf[:bytes.IndexByte(buf, 0)]
    42  			reader := bufio.NewReader(bytes.NewReader(buf))
    43  			var reqArr []string
    44  			for {
    45  				line, _, err := reader.ReadLine()
    46  				if err != nil {
    47  					break
    48  				}
    49  				reqArr = append(reqArr, string(line))
    50  			}
    51  			body = reqArr[len(reqArr)-1]
    52  		}
    53  		header := base64.StdEncoding.EncodeToString([]byte(headerBase))
    54  		scheme := "http"
    55  		if r.TLS != nil {
    56  			scheme = "https"
    57  		}
    58  		onlyKey := int(worker.GetId())
    59  
    60  		HookGroup := &request.UploadReq{
    61  			Type:     36,
    62  			InvokeId: onlyKey,
    63  			Detail: request.Detail{
    64  				AgentId: global.AgentId,
    65  				Function: request.Function{
    66  					Method:        r.Method,
    67  					Url:           scheme + "://" + r.Host + r.RequestURI,
    68  					Uri:           r.URL.Path,
    69  					Protocol:      r.Proto,
    70  					ClientIp:      r.RemoteAddr,
    71  					Language:      "GO",
    72  					ReplayRequest: false,
    73  					ReqHeader:     header,
    74  					ReqBody:       body,
    75  					QueryString:   r.URL.RawQuery,
    76  					Pool:          []request.Pool{},
    77  					TraceId:       tranceID,
    78  				},
    79  			},
    80  		}
    81  		var resBody string
    82  		var resH string
    83  		res, ok := global.ResponseMap.Load(id)
    84  		if ok {
    85  			global.ResponseMap.Delete(id)
    86  			resBody = res.(string)
    87  		}
    88  		value2, ok2 := global.ResponseHeaderMap.Load(id)
    89  		if ok2 {
    90  			global.ResponseHeaderMap.Delete(id)
    91  			resH = value2.(string)
    92  		}
    93  		for k, v := range w.Header() {
    94  			resH += k + ": " + strings.Join(v, ",") + "\n"
    95  		}
    96  		resHeader := base64.StdEncoding.EncodeToString([]byte(resH))
    97  		HookGroup.Detail.ResHeader = resHeader
    98  		HookGroup.Detail.ResBody = resBody
    99  		goroutineIDs := make(map[string]bool)
   100  		global.PoolTreeMap.Range(func(key, value interface{}) bool {
   101  			if value.(*request.PoolTree).IsThisBegin(id) {
   102  				global.PoolTreeMap.Delete(key)
   103  				value.(*request.PoolTree).FMT(&HookGroup.Detail.Function.Pool, worker, goroutineIDs, HookGroup.Detail.Function.TraceId)
   104  				return false
   105  			}
   106  			return true
   107  		})
   108  		api.ReportUpload(*HookGroup)
   109  		request.RunMapGCbYGoroutineID(goroutineIDs)
   110  	}()
   111  	return
   112  }
   113  
   114  func MyServerTemp(server *http.ServeMux, w http.ResponseWriter, r *http.Request) {
   115  	for i := 0; i < 100; i++ {
   116  
   117  	}
   118  	return
   119  }