feat: added support for context cancellation to engine (#5096)

* feat: added support for context cancellation to engine

* misc

* feat: added contexts everywhere

* misc

* misc

* use granular http timeouts and increase http timeout to 30s using multiplier

* track response header timeout in mhe

* update responseHeaderTimeout to 5sec

* skip failing windows test

---------

Co-authored-by: Tarun Koyalwar <tarun@projectdiscovery.io>
This commit is contained in:
Ice3man
2024-04-25 15:37:56 +05:30
committed by GitHub
parent 3dfcec0a36
commit 0b82e8b7aa
40 changed files with 279 additions and 113 deletions

View File

@@ -154,6 +154,7 @@ func (request *Request) Compile(options *protocols.ExecutorOptions) error {
opts := &compiler.ExecuteOptions{
Timeout: request.Timeout,
Source: &request.Init,
Context: context.Background(),
}
// register 'export' function to export variables from init code
// these are saved in args and are available in pre-condition and request code
@@ -343,7 +344,7 @@ func (request *Request) ExecuteWithResults(target *contextargs.Context, dynamicV
argsCopy.TemplateCtx = templateCtx.GetAll()
result, err := request.options.JsCompiler.ExecuteWithOptions(request.preConditionCompiled, argsCopy,
&compiler.ExecuteOptions{Timeout: request.Timeout, Source: &request.PreCondition})
&compiler.ExecuteOptions{Timeout: request.Timeout, Source: &request.PreCondition, Context: target.Context()})
if err != nil {
return errorutil.NewWithTag(request.TemplateID, "could not execute pre-condition: %s", err)
}
@@ -373,6 +374,12 @@ func (request *Request) ExecuteWithResults(target *contextargs.Context, dynamicV
return nil
}
select {
case <-input.Context().Done():
return input.Context().Err()
default:
}
if err := request.executeRequestWithPayloads(hostPort, input, hostname, value, payloadValues, func(result *output.InternalWrappedEvent) {
if result.OperatorsResult != nil && result.OperatorsResult.Matched {
gotMatches = true
@@ -419,6 +426,12 @@ func (request *Request) executeRequestParallel(ctxParent context.Context, hostPo
break
}
select {
case <-input.Context().Done():
return
default:
}
// resize check point - nop if there are no changes
if shouldFollowGlobal && sg.Size != request.options.Options.PayloadConcurrency {
sg.Resize(request.options.Options.PayloadConcurrency)
@@ -486,7 +499,7 @@ func (request *Request) executeRequestWithPayloads(hostPort string, input *conte
}
results, err := request.options.JsCompiler.ExecuteWithOptions(request.scriptCompiled, argsCopy,
&compiler.ExecuteOptions{Timeout: request.Timeout, Source: &request.Code})
&compiler.ExecuteOptions{Timeout: request.Timeout, Source: &request.Code, Context: input.Context()})
if err != nil {
// shouldn't fail even if it returned error instead create a failure event
results = compiler.ExecuteResult{"success": false, "error": err.Error()}