Merge pull request #592 from projectdiscovery/bugfix-progress-logic

Progress tracking logic
dev
PD-Team 2021-03-05 03:29:38 +05:30 committed by GitHub
commit 0c5be835de
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
8 changed files with 64 additions and 67 deletions

View File

@ -70,7 +70,7 @@ func (p *Progress) Init(hostCount int64, rulesCount int, requestCount int64) {
p.stats.AddCounter("total", uint64(requestCount))
if p.active {
if err := p.stats.Start(makePrintCallback(), p.tickDuration); err != nil {
if err := p.stats.Start(printCallback, p.tickDuration); err != nil {
gologger.Warning().Msgf("Couldn't start statistics: %s", err)
}
}
@ -91,20 +91,20 @@ func (p *Progress) IncrementMatched() {
p.stats.IncrementCounter("matched", 1)
}
// DecrementRequests decrements the number of requests from total.
func (p *Progress) DecrementRequests(count int64) {
// IncrementErrorsBy increments the error counter by count.
func (p *Progress) IncrementErrorsBy(count int64) {
p.stats.IncrementCounter("errors", int(count))
}
// IncrementFailedRequestsBy increments the number of requests counter by count along with errors.
func (p *Progress) IncrementFailedRequestsBy(count int64) {
// mimic dropping by incrementing the completed requests
p.stats.IncrementCounter("requests", int(count))
p.stats.IncrementCounter("errors", int(count))
}
const bufferSize = 128
func makePrintCallback() func(stats clistats.StatisticsClient) {
func printCallback(stats clistats.StatisticsClient) {
builder := &strings.Builder{}
builder.Grow(bufferSize)
return func(stats clistats.StatisticsClient) {
builder.WriteRune('[')
startedAt, _ := stats.GetStatic("startedAt")
duration := time.Since(startedAt.(time.Time))
@ -146,8 +146,6 @@ func makePrintCallback() func(stats clistats.StatisticsClient) {
builder.WriteRune('\n')
gologger.Print().Msgf("%s", builder.String())
builder.Reset()
}
}
// getMetrics returns a map of important metrics for client
@ -194,6 +192,8 @@ func fmtDuration(d time.Duration) string {
// Stop stops the progress bar execution
func (p *Progress) Stop() {
if p.active {
// Print one final summary
printCallback(p.stats)
if err := p.stats.Stop(); err != nil {
gologger.Warning().Msgf("Couldn't stop statistics: %s", err)
}

View File

@ -25,7 +25,7 @@ func (r *Request) ExecuteWithResults(input string, metadata, previous output.Int
compiledRequest, err := r.Make(domain)
if err != nil {
r.options.Output.Request(r.options.TemplateID, domain, "dns", err)
r.options.Progress.DecrementRequests(1)
r.options.Progress.IncrementFailedRequestsBy(1)
return errors.Wrap(err, "could not build request")
}
@ -38,7 +38,7 @@ func (r *Request) ExecuteWithResults(input string, metadata, previous output.Int
resp, err := r.dnsClient.Do(compiledRequest)
if err != nil {
r.options.Output.Request(r.options.TemplateID, domain, "dns", err)
r.options.Progress.DecrementRequests(1)
r.options.Progress.IncrementFailedRequestsBy(1)
return errors.Wrap(err, "could not send dns request")
}
r.options.Progress.IncrementRequests()

View File

@ -71,7 +71,7 @@ func (r *Request) ExecuteWithResults(input string, metadata, previous output.Int
wg.Wait()
if err != nil {
r.options.Output.Request(r.options.TemplateID, input, "file", err)
r.options.Progress.DecrementRequests(1)
r.options.Progress.IncrementFailedRequestsBy(1)
return errors.Wrap(err, "could not send file request")
}
r.options.Progress.IncrementRequests()

View File

@ -18,7 +18,7 @@ func (r *Request) ExecuteWithResults(input string, metadata, previous output.Int
instance, err := r.options.Browser.NewInstance()
if err != nil {
r.options.Output.Request(r.options.TemplateID, input, "headless", err)
r.options.Progress.DecrementRequests(1)
r.options.Progress.IncrementFailedRequestsBy(1)
return errors.Wrap(err, "could get html element")
}
defer instance.Close()
@ -26,13 +26,13 @@ func (r *Request) ExecuteWithResults(input string, metadata, previous output.Int
parsed, err := url.Parse(input)
if err != nil {
r.options.Output.Request(r.options.TemplateID, input, "headless", err)
r.options.Progress.DecrementRequests(1)
r.options.Progress.IncrementFailedRequestsBy(1)
return errors.Wrap(err, "could get html element")
}
out, page, err := instance.Run(parsed, r.Steps, time.Duration(r.options.Options.PageTimeout)*time.Second)
if err != nil {
r.options.Output.Request(r.options.TemplateID, input, "headless", err)
r.options.Progress.DecrementRequests(1)
r.options.Progress.IncrementFailedRequestsBy(1)
return errors.Wrap(err, "could get html element")
}
defer page.Close()

View File

@ -68,7 +68,6 @@ func (r *Request) executeRaceRequest(reqURL string, previous output.InternalEven
mutex.Lock()
if err != nil {
requestErr = multierr.Append(requestErr, err)
r.options.Progress.DecrementRequests(1)
}
mutex.Unlock()
}(requests[i])
@ -95,7 +94,7 @@ func (r *Request) executeParallelHTTP(reqURL string, dynamicValues, previous out
break
}
if err != nil {
r.options.Progress.DecrementRequests(int64(generator.Total()))
r.options.Progress.IncrementFailedRequestsBy(int64(generator.Total()))
return err
}
swg.Add()
@ -107,7 +106,6 @@ func (r *Request) executeParallelHTTP(reqURL string, dynamicValues, previous out
mutex.Lock()
if err != nil {
requestErr = multierr.Append(requestErr, err)
r.options.Progress.DecrementRequests(1)
}
mutex.Unlock()
}(request)
@ -154,7 +152,7 @@ func (r *Request) executeTurboHTTP(reqURL string, dynamicValues, previous output
break
}
if err != nil {
r.options.Progress.DecrementRequests(int64(generator.Total()))
r.options.Progress.IncrementFailedRequestsBy(int64(generator.Total()))
return err
}
request.pipelinedClient = pipeclient
@ -167,7 +165,6 @@ func (r *Request) executeTurboHTTP(reqURL string, dynamicValues, previous output
mutex.Lock()
if err != nil {
requestErr = multierr.Append(requestErr, err)
r.options.Progress.DecrementRequests(1)
}
mutex.Unlock()
}(request)
@ -203,7 +200,7 @@ func (r *Request) ExecuteWithResults(reqURL string, dynamicValues, previous outp
break
}
if err != nil {
r.options.Progress.DecrementRequests(int64(generator.Total()))
r.options.Progress.IncrementFailedRequestsBy(int64(generator.Total()))
return err
}
@ -223,7 +220,7 @@ func (r *Request) ExecuteWithResults(reqURL string, dynamicValues, previous outp
r.options.Progress.IncrementRequests()
if request.original.options.Options.StopAtFirstMatch && gotOutput {
r.options.Progress.DecrementRequests(int64(generator.Total()))
r.options.Progress.IncrementErrorsBy(int64(generator.Total()))
break
}
}
@ -300,7 +297,7 @@ func (r *Request) executeRequest(reqURL string, request *generatedRequest, previ
resp.Body.Close()
}
r.options.Output.Request(r.options.TemplateID, reqURL, "http", err)
r.options.Progress.DecrementRequests(1)
r.options.Progress.IncrementErrorsBy(1)
return err
}
defer func() {

View File

@ -23,7 +23,7 @@ func (r *Request) ExecuteWithResults(input string, metadata, previous output.Int
address, err := getAddress(input)
if err != nil {
r.options.Output.Request(r.options.TemplateID, input, "network", err)
r.options.Progress.DecrementRequests(1)
r.options.Progress.IncrementFailedRequestsBy(1)
return errors.Wrap(err, "could not get address from url")
}
@ -50,7 +50,7 @@ func (r *Request) executeAddress(actualAddress, address, input string, shouldUse
if !strings.Contains(actualAddress, ":") {
err := errors.New("no port provided in network protocol request")
r.options.Output.Request(r.options.TemplateID, address, "network", err)
r.options.Progress.DecrementRequests(1)
r.options.Progress.IncrementFailedRequestsBy(1)
return err
}
@ -71,7 +71,7 @@ func (r *Request) executeAddress(actualAddress, address, input string, shouldUse
}
if err != nil {
r.options.Output.Request(r.options.TemplateID, address, "network", err)
r.options.Progress.DecrementRequests(1)
r.options.Progress.IncrementFailedRequestsBy(1)
return errors.Wrap(err, "could not connect to server request")
}
defer conn.Close()
@ -92,7 +92,7 @@ func (r *Request) executeAddress(actualAddress, address, input string, shouldUse
}
if err != nil {
r.options.Output.Request(r.options.TemplateID, address, "network", err)
r.options.Progress.DecrementRequests(1)
r.options.Progress.IncrementFailedRequestsBy(1)
return errors.Wrap(err, "could not write request to server")
}
reqBuilder.Grow(len(input.Data))
@ -101,7 +101,7 @@ func (r *Request) executeAddress(actualAddress, address, input string, shouldUse
_, err = conn.Write(data)
if err != nil {
r.options.Output.Request(r.options.TemplateID, address, "network", err)
r.options.Progress.DecrementRequests(1)
r.options.Progress.IncrementFailedRequestsBy(1)
return errors.Wrap(err, "could not write request to server")
}
@ -117,7 +117,7 @@ func (r *Request) executeAddress(actualAddress, address, input string, shouldUse
}
if err != nil {
r.options.Output.Request(r.options.TemplateID, address, "network", err)
r.options.Progress.DecrementRequests(1)
r.options.Progress.IncrementFailedRequestsBy(1)
return errors.Wrap(err, "could not write request to server")
}
@ -137,7 +137,7 @@ func (r *Request) executeAddress(actualAddress, address, input string, shouldUse
n, err := conn.Read(final)
if err != nil && err != io.EOF {
r.options.Output.Request(r.options.TemplateID, address, "network", err)
r.options.Progress.DecrementRequests(1)
r.options.Progress.IncrementFailedRequestsBy(1)
return errors.Wrap(err, "could not read from server")
}
responseBuilder.Write(final[:n])

View File

@ -98,7 +98,7 @@ func (r *Request) ExecuteWithResults(input string, metadata, previous output.Int
wg.Wait()
if err != nil {
r.options.Output.Request(r.options.TemplateID, input, "file", err)
r.options.Progress.DecrementRequests(1)
r.options.Progress.IncrementFailedRequestsBy(1)
return errors.Wrap(err, "could not send file request")
}
r.options.Progress.IncrementRequests()

View File

@ -55,7 +55,7 @@ type Options struct {
Retries int
// Rate-Limit is the maximum number of requests per specified target
RateLimit int
//`PageTimeout is the maximum time to wait for a page in seconds
// PageTimeout is the maximum time to wait for a page in seconds
PageTimeout int
// OfflineHTTP is a flag that specific offline processing of http response
// using same matchers/extractors from http protocol without the need