Merge pull request #590 from projectdiscovery/bugfix-race-condition-sync

Improvement in race condition
dev
Ice3man 2021-03-01 12:30:09 +05:30 committed by GitHub
commit b532a5b3e0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 35 additions and 13 deletions

View File

@ -27,31 +27,53 @@ const defaultMaxWorkers = 150
// executeRaceRequest executes race condition request for a URL
func (r *Request) executeRaceRequest(reqURL string, previous output.InternalEvent, callback protocols.OutputEventCallback) error {
var requests []*generatedRequest
// Requests within race condition should be dumped once and the output prefilled to allow DSL language to work
// This will introduce a delay and will populate in hacky way the field "request" of outputEvent
generator := r.newGenerator()
requestForDump, err := generator.Make(reqURL, nil)
if err != nil {
return err
}
r.setCustomHeaders(requestForDump)
dumpedRequest, err := dump(requestForDump, reqURL)
if err != nil {
return err
}
if r.options.Options.Debug || r.options.Options.DebugRequests {
gologger.Info().Msgf("[%s] Dumped HTTP request for %s\n\n", r.options.TemplateID, reqURL)
gologger.Print().Msgf("%s", string(dumpedRequest))
}
previous["request"] = string(dumpedRequest)
maxWorkers := r.RaceNumberRequests
swg := sizedwaitgroup.New(maxWorkers)
var requestErr error
mutex := &sync.Mutex{}
// Pre-Generate requests
for i := 0; i < r.RaceNumberRequests; i++ {
generator := r.newGenerator()
request, err := generator.Make(reqURL, nil)
if err != nil {
return err
}
requests = append(requests, request)
}
wg := sync.WaitGroup{}
var requestErr error
mutex := &sync.Mutex{}
for i := 0; i < r.RaceNumberRequests; i++ {
swg.Add()
wg.Add(1)
go func(httpRequest *generatedRequest) {
defer wg.Done()
err := r.executeRequest(reqURL, httpRequest, previous, callback)
mutex.Lock()
if err != nil {
requestErr = multierr.Append(requestErr, err)
}
mutex.Unlock()
swg.Done()
}(request)
}(requests[i])
}
swg.Wait()
wg.Wait()
return requestErr
}
@ -217,7 +239,7 @@ func (r *Request) executeRequest(reqURL string, request *generatedRequest, previ
err error
)
// For race conditions we can't dump the request body at this point as it's already waiting the open-gate event
// For race conditions we can't dump the request body at this point as it's already waiting the open-gate event, already handled with a similar code within the race function
if !request.original.Race {
dumpedRequest, err = dump(request, reqURL)
if err != nil {