Skip to content

Commit f57e0cc

Browse files
zan8inzan8in
committed
fixed the memory leak bug
1 parent ce881ae commit f57e0cc

File tree

4 files changed

+259
-102
lines changed

4 files changed

+259
-102
lines changed

‎pkg/core/checker.go

+67-95
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ func (c *Checker) Check(target string, pocItem poc.Poc) (err error) {
3232

3333
c.FastClient.MaxRedirect = c.Options.Config.ConfigHttp.MaxRedirect
3434
c.FastClient.DialTimeout = c.Options.Config.ConfigHttp.DialTimeout
35+
c.FastClient.UserAgent = c.Options.Config.ConfigHttp.UserAgent
3536

3637
pocHandler := ""
3738
if strings.Contains(pocItem.Expression, "&&") && !strings.Contains(pocItem.Expression, "||") {
@@ -41,109 +42,88 @@ func (c *Checker) Check(target string, pocItem poc.Poc) (err error) {
4142
pocHandler = poc.ALLOR
4243
}
4344

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-
}
51-
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-
}
60-
61-
tempRequest, err := http2.ParseRequest(c.OriginalRequest)
62-
c.VariableMap["request"] = tempRequest
45+
if !strings.HasPrefix(target, "http://") && !strings.HasPrefix(target, "https://") {
46+
target = "http://" + target
47+
}
6348

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-
}
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+
}
7057

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
7764
}
65+
c.VariableMap["request"] = tempRequest
7866

79-
// update set cel and variablemap
8067
if len(pocItem.Set) > 0 {
8168
c.UpdateVariableMap(pocItem.Set)
8269
}
8370

84-
// update payloads cel and variablemap
8571
if len(pocItem.Payloads.Payloads) > 0 {
8672
c.UpdateVariableMap(pocItem.Payloads.Payloads)
8773
}
8874

89-
// rule
9075
for _, ruleMap := range pocItem.Rules {
9176
k := ruleMap.Key
9277
rule := ruleMap.Value
9378

94-
// translate : http
95-
if pocItem.Transport != "tcp" && pocItem.Transport != "udp" {
96-
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-
}
106-
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-
}
114-
115-
// set result function eg: r1() r2()
116-
c.CustomLib.WriteRuleFunctionsROptions(k, isVul.Value().(bool))
117-
118-
// update output cel and variablemap
119-
if len(rule.Output) > 0 {
120-
c.UpdateVariableMap(rule.Output)
121-
}
122-
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)})
124-
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-
}
130-
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-
}
136-
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-
}
79+
utils.RandSleep(500) // firewall just test.
80+
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+
}
87+
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+
}
95+
96+
// set result function eg: r1() r2()
97+
c.CustomLib.WriteRuleFunctionsROptions(k, isVul.Value().(bool))
98+
99+
// update output cel and variableMap
100+
if len(rule.Output) > 0 {
101+
c.UpdateVariableMap(rule.Output)
102+
}
103+
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)})
105+
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+
}
111+
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+
}
117+
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
147127
}
148128
}
149129

@@ -163,14 +143,6 @@ func (c *Checker) Check(target string, pocItem poc.Poc) (err error) {
163143
return err
164144
}
165145

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.
174146
func (c *Checker) UpdateVariableMap(args yaml.MapSlice) {
175147
for _, item := range args {
176148
key := item.Key.(string)

‎pkg/core/excute.go

+1-2
Original file line numberDiff line numberDiff line change
@@ -47,14 +47,13 @@ func (e *Engine) executeTargets(poc1 poc.Poc) {
4747
}
4848
}()
4949

50-
wg := e.workPool.NewPool(e.workPool.config.TargetConcurrencyType)
51-
5250
allTargets := e.options.Targets
5351
if len(allTargets) == 0 {
5452
log.Log().Error("executeTargets failed, no targets")
5553
return
5654
}
5755

56+
wg := e.workPool.NewPool(e.workPool.config.TargetConcurrencyType)
5857
for _, target := range allTargets {
5958
wg.WaitGroup.Add()
6059
go func(target string, poc1 poc.Poc) {

0 commit comments

Comments
 (0)