feat: add round-robin retries via Retry Rounds/Retry Delay

Author: jjhwan-h

runner/options.go

@@ -765,6 +765,10 @@ func (options *Options) ValidateOptions() error {
 		return errors.New(fmt.Sprintf("invalid retry-delay: must be >0 when retry-rounds=%d (got %d)", options.RetryRounds, options.RetryDelay))
 	}
 
+	if options.RetryRounds > 0 && options.RetryDelay <= 0 {
+		return errors.New(fmt.Sprintf("invalid retry-delay: must be >0 when retry-rounds=%d (got %d)", options.RetryRounds, options.RetryDelay))
+	}
+
 	return nil
 }
 

runner/runner.go

@@ -23,6 +23,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"golang.org/x/exp/maps"
@@ -1259,7 +1260,7 @@ func (r *Runner) RunEnumeration() {
 	wg, _ := syncutil.New(syncutil.WithSize(r.options.Threads))
 	retryCh := make(chan retryJob)
 
-	retryCancel, retryWait := r.retryLoop(context.Background(), retryCh, output, r.analyze)
+	_, drainedCh := r.retryLoop(context.Background(), retryCh, output, r.analyze)
 
 	processItem := func(k string) error {
 		if r.options.resumeCfg != nil {
@@ -1302,11 +1303,9 @@ func (r *Runner) RunEnumeration() {
 	}
 
 	wg.Wait()
-
-	retryWait()
-	retryCancel()
-	close(retryCh)
-
+	if r.options.RetryRounds > 0 {
+		<-drainedCh
+	}
 	close(output)
 	wgoutput.Wait()
 
@@ -1333,24 +1332,30 @@ type analyzeFunc func(*httpx.HTTPX, string, httpx.Target, string, string, *ScanO
 
 func (r *Runner) retryLoop(
 	parent context.Context,
-	ch chan retryJob,
+	retryCh chan retryJob,
 	output chan<- Result,
 	analyze analyzeFunc,
-) (func(), func()) {
+) (stop func(), drained <-chan struct{}) {
+	var remaining atomic.Int64
 	ctx, cancel := context.WithCancel(parent)
-	var jobWG sync.WaitGroup
+	drainedCh := make(chan struct{})
 
 	go func() {
+		defer close(retryCh)
+
 		for {
 			select {
 			case <-ctx.Done():
 				return
-			case job := <-ch:
-				jobWG.Add(1)
+			case job, ok := <-retryCh:
+				if !ok {
+					return
+				}
+				if job.attempt == 1 {
+					remaining.Add(1)
+				}
 
 				go func(j retryJob) {
-					defer jobWG.Done()
-
 					if wait := time.Until(j.when); wait > 0 {
 						timer := time.NewTimer(wait)
 						select {
@@ -1364,28 +1369,27 @@ func (r *Runner) retryLoop(
 					res := analyze(j.hp, j.protocol, j.target, j.method, j.origInput, j.scanopts)
 					output <- res
 
-					if res.StatusCode == http.StatusTooManyRequests {
-						if j.attempt >= r.options.RetryRounds {
-							return
-						}
-
+					if res.StatusCode == http.StatusTooManyRequests && j.attempt < r.options.RetryRounds {
 						j.attempt++
 						j.when = time.Now().Add(time.Duration(r.options.RetryDelay) * time.Millisecond)
 
 						select {
 						case <-ctx.Done():
 							return
-						case ch <- j:
+						case retryCh <- j:
+							return
 						}
 					}
+
+					if remaining.Add(-1) == 0 {
+						close(drainedCh)
+					}
 				}(job)
 			}
 		}
 	}()
 
-	stop := func() { cancel() }
-	wait := func() { jobWG.Wait() }
-	return stop, wait
+	return func() { cancel() }, drainedCh
 }
 
 func logFilteredErrorPage(fileName, url string) {

runner/runner_test.go

@@ -3,7 +3,6 @@ package runner
 import (
 	"context"
 	"fmt"
-	"log"
 	"net/http"
 	"net/http/httptest"
 	"os"
@@ -324,41 +323,34 @@ func TestRunner_Process_And_RetryLoop(t *testing.T) {
 	var hits1, hits2 int32
 	srv1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if atomic.AddInt32(&hits1, 1) != 4 {
-			log.Println("serv1 429")
 			w.WriteHeader(http.StatusTooManyRequests)
 			return
 		}
-		log.Println("serv1 200")
 		w.WriteHeader(http.StatusOK)
 	}))
 	defer srv1.Close()
 
 	srv2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if atomic.AddInt32(&hits2, 1) != 3 {
-			log.Println("serv2 429")
 			w.WriteHeader(http.StatusTooManyRequests)
 			return
 		}
-		log.Println("serv2 200")
 		w.WriteHeader(http.StatusOK)
 	}))
 	defer srv2.Close()
 
 	r, err := New(&Options{
 		Threads:     1,
-		Delay:       0,
-		RetryRounds: 3,
-		RetryDelay:  200, // Duration 권장
-		Timeout:     2,
+		RetryRounds: 2,
+		RetryDelay:  5,
+		Timeout:     3,
 	})
 	require.NoError(t, err)
 
 	output := make(chan Result)
 	retryCh := make(chan retryJob)
 
-	// ctx, timeout := context.WithTimeout(context.Background(), time.Duration(r.options.Timeout))
-	// defer timeout()
-	cancel, wait := r.retryLoop(context.Background(), retryCh, output, r.analyze)
+	_, drainedCh := r.retryLoop(context.Background(), retryCh, output, r.analyze)
 
 	wg, _ := syncutil.New(syncutil.WithSize(r.options.Threads))
 	so := r.scanopts.Clone()
@@ -399,15 +391,12 @@ func TestRunner_Process_And_RetryLoop(t *testing.T) {
 	}
 
 	wg.Wait()
-	wait()
-	cancel()
-
-	close(retryCh)
+	<-drainedCh
 	close(output)
 	drainWG.Wait()
 
 	require.Equal(t, 3, s1n429)
-	require.Equal(t, 1, s1n200)
+	require.Equal(t, 0, s1n200)
 	require.Equal(t, 2, s2n429)
 	require.Equal(t, 1, s2n200)
 }