Limit concurrency in file protocol

This commit is contained in:
Ice3man543 2021-01-14 22:43:08 +05:30
parent 8b93d5e1d2
commit 6a739c2d0c

View File

@ -10,13 +10,21 @@ import (
"github.com/projectdiscovery/nuclei/v2/pkg/output"
"github.com/projectdiscovery/nuclei/v2/pkg/protocols"
"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/tostring"
"github.com/remeh/sizedwaitgroup"
)
var _ protocols.Request = &Request{}
// ExecuteWithResults executes the protocol requests and returns results instead of writing them.
func (r *Request) ExecuteWithResults(input string, metadata output.InternalEvent, callback protocols.OutputEventCallback) error {
wg := sizedwaitgroup.New(r.options.Options.RateLimit)
err := r.getInputPaths(input, func(data string) {
wg.Add()
go func(data string) {
defer wg.Done()
file, err := os.Open(data)
if err != nil {
gologger.Error().Msgf("Could not open file path %s: %s\n", data, err)
@ -57,7 +65,9 @@ func (r *Request) ExecuteWithResults(input string, metadata output.InternalEvent
}
}
callback(event)
}(data)
})
wg.Wait()
if err != nil {
r.options.Output.Request(r.options.TemplateID, input, "file", err)
r.options.Progress.DecrementRequests(1)