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

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