STRLCPY/afrog

fixed the memory leak bug
zan8in committed 2 years ago

f57e0cc4

1 parent ce881aef

■ ■ ■ ■ ■ ■

pkg/core/checker.go

		skipped 31 lines
32	32
33	33		c.FastClient.MaxRedirect = c.Options.Config.ConfigHttp.MaxRedirect
34	34		c.FastClient.DialTimeout = c.Options.Config.ConfigHttp.DialTimeout
	35	+	c.FastClient.UserAgent = c.Options.Config.ConfigHttp.UserAgent
35	36
36	37		pocHandler := ""
37	38		if strings.Contains(pocItem.Expression, "&&") && !strings.Contains(pocItem.Expression, "\|\|") {
		skipped 3 lines
41	42		pocHandler = poc.ALLOR
42	43		}
43	44
44		-	// update request variablemap
45		-	// tempRequest := http2.AcquireProtoRequestPool()
46		-	// defer http2.ReleaseProtoRequestPool(tempRequest)
47		-	if pocItem.Transport != "tcp" && pocItem.Transport != "udp" {
48		-	if !strings.HasPrefix(target, "http://") && !strings.HasPrefix(target, "https://") {
49		-	target = "http://" + target
50		-	}
	45	+	if !strings.HasPrefix(target, "http://") && !strings.HasPrefix(target, "https://") {
	46	+	target = "http://" + target
	47	+	}
51	48
52		-	// original request
53		-	c.OriginalRequest, err = http.NewRequest("GET", target, nil)
54		-	if err != nil {
55		-	log.Log().Error(fmt.Sprintf("rule map originalRequest err, %s", err.Error()))
56		-	c.Result.IsVul = false
57		-	c.Options.ApiCallBack(c.Result)
58		-	return err
59		-	}
	49	+	// original request
	50	+	c.OriginalRequest, err = http.NewRequest("GET", target, nil)
	51	+	if err != nil {
	52	+	log.Log().Error(fmt.Sprintf("rule map originalRequest err, %s", err.Error()))
	53	+	c.Result.IsVul = false
	54	+	c.Options.ApiCallBack(c.Result)
	55	+	return err
	56	+	}
60	57
61		-	tempRequest, err := http2.ParseRequest(c.OriginalRequest)
62		-	c.VariableMap["request"] = tempRequest
63		-
64		-	if err != nil {
65		-	log.Log().Error(fmt.Sprintf("ParseRequest err, %s", err.Error()))
66		-	c.Result.IsVul = false
67		-	c.Options.ApiCallBack(c.Result)
68		-	return err
69		-	}
70		-
71		-	// set User-Agent
72		-	if len(c.Options.Config.ConfigHttp.UserAgent) > 0 {
73		-	c.OriginalRequest.Header.Set("User-Agent", c.Options.Config.ConfigHttp.UserAgent)
74		-	} else {
75		-	c.OriginalRequest.Header.Set("User-Agent", utils.RandomUA())
76		-	}
	58	+	tempRequest, err := http2.ParseRequest(c.OriginalRequest)
	59	+	if err != nil {
	60	+	log.Log().Error(fmt.Sprintf("ParseRequest err, %s", err.Error()))
	61	+	c.Result.IsVul = false
	62	+	c.Options.ApiCallBack(c.Result)
	63	+	return err
77	64		}
	65	+	c.VariableMap["request"] = tempRequest
78	66
79		-	// update set cel and variablemap
80	67		if len(pocItem.Set) > 0 {
81	68		c.UpdateVariableMap(pocItem.Set)
82	69		}
83	70
84		-	// update payloads cel and variablemap
85	71		if len(pocItem.Payloads.Payloads) > 0 {
86	72		c.UpdateVariableMap(pocItem.Payloads.Payloads)
87	73		}
88	74
89		-	// rule
90	75		for _, ruleMap := range pocItem.Rules {
91	76		k := ruleMap.Key
92	77		rule := ruleMap.Value
93	78
94		-	// translate : http
95		-	if pocItem.Transport != "tcp" && pocItem.Transport != "udp" {
	79	+	utils.RandSleep(500) // firewall just test.
96	80
97		-	// run fasthttp client
98		-	utils.RandSleep(500) // firewall just test.
99		-
100		-	err = c.FastClient.HTTPRequest(c.OriginalRequest, rule, c.VariableMap)
101		-	if err != nil {
102		-	log.Log().Error(fmt.Sprintf("rule map fasthttp.HTTPRequest err, %s", err.Error()))
103		-	c.CustomLib.WriteRuleFunctionsROptions(k, false)
104		-	continue // not return, becuase may be need test next pocitem. ？？？
105		-	}
	81	+	err = c.FastClient.HTTPRequest(c.OriginalRequest, rule, c.VariableMap)
	82	+	if err != nil {
	83	+	log.Log().Error(fmt.Sprintf("rule map fasthttp.HTTPRequest err, %s", err.Error()))
	84	+	c.CustomLib.WriteRuleFunctionsROptions(k, false)
	85	+	continue
	86	+	}
106	87
107		-	// run cel expression
108		-	isVul, err := c.CustomLib.RunEval(rule.Expression, c.VariableMap)
109		-	if err != nil {
110		-	log.Log().Error(fmt.Sprintf("rule map RunEval err, %s", err.Error()))
111		-	c.CustomLib.WriteRuleFunctionsROptions(k, false)
112		-	continue // not return, becuase may be need test next pocitem. ？？？
113		-	}
	88	+	// run cel expression
	89	+	isVul, err := c.CustomLib.RunEval(rule.Expression, c.VariableMap)
	90	+	if err != nil {
	91	+	log.Log().Error(fmt.Sprintf("rule map RunEval err, %s", err.Error()))
	92	+	c.CustomLib.WriteRuleFunctionsROptions(k, false)
	93	+	continue // not return, because may be need test next pocItem. ？？？
	94	+	}
114	95
115		-	// set result function eg: r1() r2()
116		-	c.CustomLib.WriteRuleFunctionsROptions(k, isVul.Value().(bool))
	96	+	// set result function eg: r1() r2()
	97	+	c.CustomLib.WriteRuleFunctionsROptions(k, isVul.Value().(bool))
117	98
118		-	// update output cel and variablemap
119		-	if len(rule.Output) > 0 {
120		-	c.UpdateVariableMap(rule.Output)
121		-	}
	99	+	// update output cel and variableMap
	100	+	if len(rule.Output) > 0 {
	101	+	c.UpdateVariableMap(rule.Output)
	102	+	}
122	103
123		-	c.Result.AllPocResult = append(c.Result.AllPocResult, &PocResult{IsVul: isVul.Value().(bool), ResultRequest: c.VariableMap["request"].(proto.Request), ResultResponse: c.VariableMap["response"].(proto.Response)})
	104	+	c.Result.AllPocResult = append(c.Result.AllPocResult, &PocResult{IsVul: isVul.Value().(bool), ResultRequest: c.VariableMap["request"].(proto.Request), ResultResponse: c.VariableMap["response"].(proto.Response)})
124	105
125		-	if rule.Request.Todo == poc.TODO_FAILURE_NOT_CONTINUE && !isVul.Value().(bool) {
126		-	c.Result.IsVul = false
127		-	c.Options.ApiCallBack(c.Result)
128		-	return err
129		-	}
	106	+	if rule.Request.Todo == poc.TODO_FAILURE_NOT_CONTINUE && !isVul.Value().(bool) {
	107	+	c.Result.IsVul = false
	108	+	c.Options.ApiCallBack(c.Result)
	109	+	return err
	110	+	}
130	111
131		-	if rule.Request.Todo == poc.TODO_SUCCESS_NOT_CONTINUE && isVul.Value().(bool) {
132		-	c.Result.IsVul = true
133		-	c.Options.ApiCallBack(c.Result)
134		-	return err
135		-	}
	112	+	if rule.Request.Todo == poc.TODO_SUCCESS_NOT_CONTINUE && isVul.Value().(bool) {
	113	+	c.Result.IsVul = true
	114	+	c.Options.ApiCallBack(c.Result)
	115	+	return err
	116	+	}
136	117
137		-	if pocHandler == poc.ALLOR && isVul.Value().(bool) {
138		-	c.Result.IsVul = true
139		-	c.Options.ApiCallBack(c.Result)
140		-	return err
141		-	}
142		-	if pocHandler == poc.ALLAND && !isVul.Value().(bool) {
143		-	c.Result.IsVul = false
144		-	c.Options.ApiCallBack(c.Result)
145		-	return err
146		-	}
	118	+	if pocHandler == poc.ALLOR && isVul.Value().(bool) {
	119	+	c.Result.IsVul = true
	120	+	c.Options.ApiCallBack(c.Result)
	121	+	return err
	122	+	}
	123	+	if pocHandler == poc.ALLAND && !isVul.Value().(bool) {
	124	+	c.Result.IsVul = false
	125	+	c.Options.ApiCallBack(c.Result)
	126	+	return err
147	127		}
148	128		}
149	129
		skipped 13 lines
163	143		return err
164	144		}
165	145
166		-	// print result info for debug
167		-	func (c Checker) PrintTraceInfo(result Result) {
168		-	for i, v := range result.AllPocResult {
169		-	log.Log().Info(fmt.Sprintf("\r\n%s（%d）\r\n%s\r\n\r\n%s（%d）\r\n%s\r\n", "Request:", i, v.ReadFullResultRequestInfo(), "Response:", i, v.ReadFullResultResponseInfo()))
170		-	}
171		-	}
172		-
173		-	// update set、payload、output variableMap etc.
174	146		func (c *Checker) UpdateVariableMap(args yaml.MapSlice) {
175	147		for _, item := range args {
176	148		key := item.Key.(string)
		skipped 43 lines

■ ■ ■ ■ ■ ■

pkg/core/excute.go

		skipped 46 lines
47	47		}
48	48		}()
49	49
50		-	wg := e.workPool.NewPool(e.workPool.config.TargetConcurrencyType)
51		-
52	50		allTargets := e.options.Targets
53	51		if len(allTargets) == 0 {
54	52		log.Log().Error("executeTargets failed, no targets")
55	53		return
56	54		}
57	55
	56	+	wg := e.workPool.NewPool(e.workPool.config.TargetConcurrencyType)
58	57		for _, target := range allTargets {
59	58		wg.WaitGroup.Add()
60	59		go func(target string, poc1 poc.Poc) {
		skipped 22 lines

■ ■ ■ ■ ■ ■ ■

pkg/protocols/http/http.go

		skipped 4 lines
5	5		"crypto/tls"
6	6		"errors"
7	7		"fmt"
	8	+	"github.com/zan8in/afrog/pkg/utils"
8	9		"io/ioutil"
9	10		"net/http"
10	11		"net/url"
		skipped 18 lines
29	30		type FastClient struct {
30	31		MaxRedirect int32
31	32		DialTimeout int32
	33	+	UserAgent string
32	34		}
33	35
34	36		func Init(options *config.Options) {
		skipped 24 lines
59	61
60	62		func (fc FastClient) HTTPRequest(httpRequest http.Request, rule poc.Rule, variableMap map[string]interface{}) error {
61	63		var err error
	64	+
	65	+	protoRequest, err := ParseRequest(httpRequest)
	66	+	if err != nil {
	67	+	return err
	68	+	}
	69	+
	70	+	variableMap["request"] = nil
	71	+	variableMap["response"] = nil
	72	+
	73	+	newRuleHeader := make(map[string]string)
	74	+	for k, v := range rule.Request.Headers {
	75	+	newRuleHeader[k] = fc.AssignVariableMap(v, variableMap)
	76	+	}
	77	+	rule.Request.Headers = newRuleHeader
	78	+	rule.Request.Path = fc.AssignVariableMap(strings.TrimSpace(rule.Request.Path), variableMap)
	79	+	//rule.Request.Body = fc.AssignVariableMap(strings.TrimSpace(rule.Request.Body), variableMap)
	80	+	if strings.HasPrefix(strings.ToLower(rule.Request.Headers["Content-Type"]), "multipart/form-Data") && strings.Contains(rule.Request.Body, "\n\n") {
	81	+	multipartBody, err := DealMultipart(rule.Request.Headers["Content-Type"], rule.Request.Body)
	82	+	if err != nil {
	83	+	return err
	84	+	}
	85	+	rule.Request.Body = fc.AssignVariableMap(strings.TrimSpace(multipartBody), variableMap)
	86	+	} else {
	87	+	rule.Request.Body = fc.AssignVariableMap(strings.TrimSpace(rule.Request.Body), variableMap)
	88	+	}
	89	+
	90	+	if strings.HasPrefix(rule.Request.Path, "/") {
	91	+	protoRequest.Url.Path = strings.TrimRight(httpRequest.URL.Path, "/") + rule.Request.Path
	92	+	} else if strings.HasPrefix(rule.Request.Path, "^") {
	93	+	protoRequest.Url.Path = "/" + rule.Request.Path[1:]
	94	+	}
	95	+
	96	+	protoRequest.Url.Path = strings.ReplaceAll(protoRequest.Url.Path, " ", "%20")
	97	+	protoRequest.Url.Path = strings.ReplaceAll(protoRequest.Url.Path, "+", "%20")
	98	+
	99	+	if rule.Request.Method == "POST" && len(rule.Request.Headers["Content-Type"]) == 0 {
	100	+	if rule.Request.Headers == nil {
	101	+	rule.Request.Headers = make(map[string]string)
	102	+	}
	103	+	rule.Request.Headers["Content-Type"] = "application/x-www-form-urlencoded"
	104	+	}
	105	+
	106	+	finalRequest, err := http.NewRequest(rule.Request.Method, fmt.Sprintf("%s://%s%s", protoRequest.Url.Scheme, protoRequest.Url.Host, protoRequest.Url.Path), strings.NewReader(rule.Request.Body))
	107	+	if err != nil {
	108	+	return err
	109	+	}
	110	+
	111	+	fastReq := fasthttp.AcquireRequest()
	112	+	defer fasthttp.ReleaseRequest(fastReq)
	113	+
	114	+	var rawHeader strings.Builder
	115	+	newheader := make(map[string]string)
	116	+	for k, v := range rule.Request.Headers {
	117	+	fastReq.Header.Set(k, v)
	118	+
	119	+	newheader[k] = v
	120	+
	121	+	rawHeader.WriteString(k)
	122	+	rawHeader.WriteString(": ")
	123	+	rawHeader.WriteString(v)
	124	+	rawHeader.WriteString("\n")
	125	+	}
	126	+
	127	+	if len(fc.UserAgent) > 0 {
	128	+	fastReq.Header.Set("User-Agent", fc.UserAgent)
	129	+	} else {
	130	+	fastReq.Header.Set("User-Agent", utils.RandomUA())
	131	+	}
	132	+
	133	+	fastReq.Header.SetMethod(rule.Request.Method)
	134	+
	135	+	if len(rule.Request.Headers["Host"]) == 0 {
	136	+	fastReq.Header.SetHost(finalRequest.Host)
	137	+	}
	138	+	if len(rule.Request.Headers["Content-Length"]) == 0 {
	139	+	fastReq.Header.SetContentLength(int(finalRequest.ContentLength))
	140	+	}
	141	+	fastReq.SetBody([]byte(rule.Request.Body))
	142	+	fastReq.URI().Update(finalRequest.URL.String())
	143	+	fastReq.SetRequestURI(finalRequest.URL.String())
	144	+
	145	+	fastResp := fasthttp.AcquireResponse()
	146	+	defer fasthttp.ReleaseResponse(fastResp)
	147	+
	148	+	if rule.Request.FollowRedirects {
	149	+	maxrd := 5 // follow redirects default 5
	150	+	if fc.MaxRedirect > 0 {
	151	+	maxrd = int(fc.MaxRedirect)
	152	+	}
	153	+	err = F.DoRedirects(fastReq, fastResp, maxrd)
	154	+	} else {
	155	+	dialtimeout := 6
	156	+	if fc.DialTimeout > 0 {
	157	+	dialtimeout = int(fc.DialTimeout)
	158	+	}
	159	+	err = F.DoTimeout(fastReq, fastResp, time.Second*time.Duration(dialtimeout))
	160	+	}
	161	+	if err != nil {
	162	+	errName, known := httpConnError(err)
	163	+	if known {
	164	+	log.Log().Error(fmt.Sprintf("WARN conn error: %s\n", errName))
	165	+	} else {
	166	+	log.Log().Error(fmt.Sprintf("ERR conn failure: %s %s\n", errName, err))
	167	+	}
	168	+	}
	169	+
	170	+	// set fastResp body
	171	+	var respBody []byte
	172	+	contentEncoding := strings.ToLower(string(fastResp.Header.Peek("Content-Encoding")))
	173	+	switch contentEncoding {
	174	+	case "", "none", "identity":
	175	+	respBody = fastResp.Body()
	176	+	case "gzip":
	177	+	respBody, err = fastResp.BodyGunzip()
	178	+	case "deflate":
	179	+	respBody, err = fastResp.BodyInflate()
	180	+	default:
	181	+	respBody = []byte{}
	182	+	}
	183	+	if err != nil {
	184	+	return err
	185	+	}
	186	+	fastResp.SetBody(respBody)
	187	+
	188	+	// fc.VariableMap["response"] variable assignment
	189	+	tempResultResponse := &proto.Response{}
	190	+	tempResultResponse.Status = int32(fastResp.StatusCode())
	191	+	u, err := url.Parse(fastReq.URI().String())
	192	+	if err != nil {
	193	+	return err
	194	+	}
	195	+	urlType := &proto.UrlType{
	196	+	Scheme: u.Scheme,
	197	+	Domain: u.Hostname(),
	198	+	Host: u.Host,
	199	+	Port: u.Port(),
	200	+	Path: u.Path,
	201	+	Query: u.RawQuery,
	202	+	Fragment: u.Fragment,
	203	+	}
	204	+	tempResultResponse.Url = urlType
	205	+	newheader2 := make(map[string]string)
	206	+	respHeaderSlice := strings.Split(fastResp.Header.String(), "\r\n")
	207	+	for _, h := range respHeaderSlice {
	208	+	hslice := strings.SplitN(h, ":", 2)
	209	+	if len(hslice) != 2 {
	210	+	continue
	211	+	}
	212	+	k := strings.ToLower(hslice[0])
	213	+	v := strings.TrimLeft(hslice[1], " ")
	214	+	if newheader2[k] != "" {
	215	+	newheader2[k] += v
	216	+	} else {
	217	+	newheader2[k] = v
	218	+	}
	219	+	}
	220	+	tempResultResponse.Headers = newheader2
	221	+	tempResultResponse.ContentType = string(fastResp.Header.ContentType())
	222	+	tempResultResponse.Body = fastResp.Body()
	223	+	tempResultResponse.Raw = []byte(fastResp.String())
	224	+	tempResultResponse.RawHeader = fastResp.Header.Header()
	225	+	variableMap["response"] = tempResultResponse
	226	+
	227	+	protoRequest.Method = rule.Request.Method
	228	+	protoRequest.Url = urlType
	229	+	protoRequest.RawHeader = []byte(strings.Trim(rawHeader.String(), "\n"))
	230	+	protoRequest.Raw = []byte(string(fastReq.Header.String()) + "\n" + string(fastReq.Body()))
	231	+	protoRequest.Headers = newheader
	232	+	protoRequest.ContentType = newheader["content-type"]
	233	+	protoRequest.Body = []byte(rule.Request.Body)
	234	+	variableMap["request"] = protoRequest
	235	+
	236	+	return err
	237	+	}
	238	+
	239	+	func (fc FastClient) HTTPRequest2(httpRequest http.Request, rule poc.Rule, variableMap map[string]interface{}) error {
	240	+	var err error
	241	+
	242	+	variableMap["request"] = nil
	243	+	variableMap["response"] = nil
62	244
63	245		fastReq := fasthttp.AcquireRequest()
64	246		defer fasthttp.ReleaseRequest(fastReq)
		skipped 301 lines
366	548		}
367	549
368	550		func CopyRequest(req http.Request, dstRequest fasthttp.Request, data []byte) {
369		-	dstRequest.SetRequestURI(req.URL.String())
	551	+	//dstRequest.SetRequestURI(req.URL.String())
	552	+	dstRequest.URI().Update(req.URL.String())
370	553		dstRequest.Header.SetMethod(req.Method)
371	554		for name, values := range req.Header {
372	555		// Loop over all values for the name.
		skipped 5 lines
378	561		}
379	562		}
380	563		}
	564	+	dstRequest.SetBody(data)
381	565		dstRequest.SetBodyRaw(data)
382	566		}
383	567
		skipped 72 lines

■ ■ ■ ■ ■ ■

pocs/poc.yaml

		skipped 19 lines
20	20		r1:
21	21		request:
22	22		method: POST
23		-	path: "/test1.php"
	23	+	path: "/test1.php?a=xxhh#abc"
24	24		headers:
25	25		Referer: '{{referer}}'
26	26		Domain: '{{host}}'
27	27		Reverse: '{{reverseURL}}'
28		-	body: "{{abc}}.ASPX"
	28	+	Content-Length: 100
	29	+	Host: 1.1.1.1
	30	+	body: "b={{content}}.ASPX"
29	31		expression: \|
30	32		reverse.wait(3)
31	33		output:
		skipped 5 lines
37	39		method: POST
38	40		path: "/test1.php"
39	41		body: \|
40		-	id=';echo(md5(123));//&csrftoken={{home}}
	42	+	b={{home}}
41	43		expression: \|
42		-	response.status == 200 && response.body.bcontains(b'202cb962ac59075b964b07152d234b70')
	44	+	response.status == 200
43	45		expression: r1() \|\| r2()

fixed the memory leak bug