diff --git a/common/httputilz/httputilz_test.go b/common/httputilz/httputilz_test.go index ffc4f1645..8431c4503 100644 --- a/common/httputilz/httputilz_test.go +++ b/common/httputilz/httputilz_test.go @@ -91,3 +91,21 @@ func TestParseRequestMalformed(t *testing.T) { _, _, _, _, err := ParseRequest("GET\r\n\r\n", false) require.Error(t, err) } + +func TestParseRequestPreservesQueryParameters(t *testing.T) { + raw := strings.Join([]string{ + "GET /api/v1/users?page=2&limit=50 HTTP/1.1", + "Host: example.com", + "Authorization: Bearer token123", + "", + "", + }, "\r\n") + + method, path, headers, _, err := ParseRequest(raw, false) + require.NoError(t, err) + require.Equal(t, "GET", method) + require.Equal(t, "/api/v1/users?page=2&limit=50", path) + require.Equal(t, []string{"example.com"}, headers["Host"]) + require.Equal(t, []string{"Bearer token123"}, headers["Authorization"]) +} + diff --git a/runner/ports_optimization_test.go b/runner/ports_optimization_test.go index 232db95ed..018f2da1c 100644 --- a/runner/ports_optimization_test.go +++ b/runner/ports_optimization_test.go @@ -91,8 +91,21 @@ func TestSwitchPortForFallback(t *testing.T) { inputProtocol: httpx.HTTP, expectedNewPort: "8080", // non-default port stays the same }, + { + name: "HTTP:3000 fallback should keep port 3000", + inputPort: "3000", + inputProtocol: httpx.HTTP, + expectedNewPort: "3000", + }, + { + name: "HTTPS:5000 fallback should keep port 5000", + inputPort: "5000", + inputProtocol: httpx.HTTPS, + expectedNewPort: "5000", + }, } + for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { newPort := getPortForFallback(tc.inputPort, tc.inputProtocol) diff --git a/runner/resume.go b/runner/resume.go index ce2d3752e..cc757b68b 100644 --- a/runner/resume.go +++ b/runner/resume.go @@ -1,8 +1,119 @@ package runner +import ( + "fmt" + "os" + "path/filepath" + "sync" + + "github.com/projectdiscovery/goconfig" +) + +type resumeSaveState struct { + ResumeFrom string `json:"resumeFrom,omitempty"` + Index int `json:"index,omitempty"` +} + type ResumeCfg struct { - ResumeFrom string - Index int - current string - currentIndex int + sync.RWMutex `json:"-"` + ResumeFrom string `json:"resumeFrom,omitempty"` + Index int `json:"index,omitempty"` + resumeBaseline int + current string + currentIndex int + completed map[int]string + completedIdx int + completedTarget string +} + +func (r *ResumeCfg) init() { + if r.completed == nil { + r.completed = make(map[int]string) + r.completedIdx = r.Index + r.completedTarget = r.ResumeFrom + r.resumeBaseline = r.Index + } +} + +// NextIndex increments the dispatched index and returns whether the item should be skipped. +func (r *ResumeCfg) NextIndex(target string) (int, bool) { + r.Lock() + defer r.Unlock() + r.init() + + r.currentIndex++ + r.current = target + + if r.currentIndex <= r.resumeBaseline { + return r.currentIndex, true + } + return r.currentIndex, false +} + +// MarkCompleted records that an item at the given index has fully finished processing. +func (r *ResumeCfg) MarkCompleted(index int, target string) { + r.Lock() + defer r.Unlock() + r.init() + + r.completed[index] = target + + for { + nextIdx := r.completedIdx + 1 + if tgt, exists := r.completed[nextIdx]; exists { + r.completedIdx = nextIdx + r.completedTarget = tgt + delete(r.completed, nextIdx) + } else { + break + } + } + + r.Index = r.completedIdx + r.ResumeFrom = r.completedTarget +} + +// CurrentCompleted returns the current contiguous completed index and target. +func (r *ResumeCfg) CurrentCompleted() (int, string) { + r.RLock() + defer r.RUnlock() + return r.Index, r.ResumeFrom +} + +// Save atomically writes the ResumeCfg to the specified file path. +func (r *ResumeCfg) Save(filePath string) error { + r.RLock() + state := resumeSaveState{ + ResumeFrom: r.ResumeFrom, + Index: r.Index, + } + r.RUnlock() + + dir := filepath.Dir(filePath) + if dir == "" { + dir = "." + } + tempFile, err := os.CreateTemp(dir, fmt.Sprintf(".%s-*.tmp", filepath.Base(filePath))) + if err != nil { + return err + } + tempPath := tempFile.Name() + _ = tempFile.Close() + + if err := goconfig.Save(state, tempPath); err != nil { + _ = os.Remove(tempPath) + return err + } + + if f, err := os.OpenFile(tempPath, os.O_RDWR, 0600); err == nil { + _ = f.Sync() + _ = f.Close() + } + + if err := os.Rename(tempPath, filePath); err != nil { + _ = os.Remove(tempPath) + return err + } + + return nil } diff --git a/runner/resume_test.go b/runner/resume_test.go new file mode 100644 index 000000000..212879e5d --- /dev/null +++ b/runner/resume_test.go @@ -0,0 +1,204 @@ +package runner + +import ( + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "path/filepath" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/projectdiscovery/goconfig" + fileutil "github.com/projectdiscovery/utils/file" + "github.com/stretchr/testify/require" +) + +func TestResumeCfg_AtomicSave(t *testing.T) { + tempDir := t.TempDir() + resumePath := filepath.Join(tempDir, "test_resume.cfg") + + cfg := &ResumeCfg{ + Index: 42, + ResumeFrom: "https://example.com", + } + + err := cfg.Save(resumePath) + require.NoError(t, err, "Save should succeed") + require.True(t, fileutil.FileExists(resumePath), "Resume file should exist") + + var loadedCfg ResumeCfg + err = goconfig.Load(&loadedCfg, resumePath) + require.NoError(t, err, "Loading saved config should succeed") + require.Equal(t, 42, loadedCfg.Index, "Loaded index should match") + require.Equal(t, "https://example.com", loadedCfg.ResumeFrom, "Loaded target should match") +} + +func TestResumeCfg_ContiguousCompletionTracking(t *testing.T) { + cfg := &ResumeCfg{} + + // Dispatch 5 items + targets := []string{"t1", "t2", "t3", "t4", "t5"} + for _, target := range targets { + _, skip := cfg.NextIndex(target) + require.False(t, skip) + } + + // Complete item 1 -> index should be 1 + cfg.MarkCompleted(1, "t1") + idx, tgt := cfg.CurrentCompleted() + require.Equal(t, 1, idx) + require.Equal(t, "t1", tgt) + + // Complete item 3 out-of-order -> index should still remain 1 because item 2 is in-flight + cfg.MarkCompleted(3, "t3") + idx, tgt = cfg.CurrentCompleted() + require.Equal(t, 1, idx, "Index must not advance past incomplete in-flight item 2") + require.Equal(t, "t1", tgt) + + // Complete item 5 out-of-order -> index should still be 1 + cfg.MarkCompleted(5, "t5") + idx, tgt = cfg.CurrentCompleted() + require.Equal(t, 1, idx) + require.Equal(t, "t1", tgt) + + // Complete item 2 -> index should jump to 3 (since 1, 2, 3 are now all done, but 4 is still in-flight) + cfg.MarkCompleted(2, "t2") + idx, tgt = cfg.CurrentCompleted() + require.Equal(t, 3, idx, "Index should advance to 3 after missing item 2 completes") + require.Equal(t, "t3", tgt) + + // Complete item 4 -> index should jump to 5 (since 4 and 5 are now complete) + cfg.MarkCompleted(4, "t4") + idx, tgt = cfg.CurrentCompleted() + require.Equal(t, 5, idx, "Index should advance to 5 once all items complete") + require.Equal(t, "t5", tgt) +} + +func TestRunner_MultiThreadedInterruptAndResume(t *testing.T) { + if testing.Short() { + t.Skip("skipping interrupt-and-resume integration test in short mode") + } + + // Set up mock HTTP server + var serverRequests int32 + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt32(&serverRequests, 1) + // Small delay to simulate in-flight concurrency + time.Sleep(10 * time.Millisecond) + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("ok")) + })) + defer ts.Close() + + u, err := url.Parse(ts.URL) + require.NoError(t, err) + + tempDir := t.TempDir() + resumeFile := filepath.Join(tempDir, "resume.cfg") + + // Generate target list with query params to create unique URLs pointing to local test server + const totalTargets = 30 + var targets []string + for i := 1; i <= totalTargets; i++ { + targets = append(targets, fmt.Sprintf("%s:%s?id=%d", u.Hostname(), u.Port(), i)) + } + + var firstRunProcessed sync.Map + var firstRunCount int32 + const interruptThreshold = 10 + + opts1 := &Options{ + InputTargetHost: targets, + Threads: 4, + Delay: 0, + NoColor: true, + resumeCfg: &ResumeCfg{}, + OnResult: func(r Result) { + if r.Err == nil { + firstRunProcessed.Store(r.URL, true) + atomic.AddInt32(&firstRunCount, 1) + } + }, + } + + r1, err := New(opts1) + require.NoError(t, err) + + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + r1.RunEnumeration() + }() + + // Monitor progress and interrupt when threshold reached (with 60s timeout guard) + deadline := time.After(60 * time.Second) + ticker := time.NewTicker(5 * time.Millisecond) + defer ticker.Stop() + + interrupted := false + for !interrupted { + select { + case <-deadline: + t.Fatal("timed out waiting for interrupt threshold") + case <-ticker.C: + if atomic.LoadInt32(&firstRunCount) >= interruptThreshold { + r1.Interrupt() + interrupted = true + } + } + } + + wg.Wait() + + require.Less(t, atomic.LoadInt32(&firstRunCount), int32(totalTargets), + "the interrupted run must leave targets for the resumed run") + + // Save resume config atomically to the temp resume file + err = r1.options.resumeCfg.Save(resumeFile) + require.NoError(t, err) + require.True(t, fileutil.FileExists(resumeFile), "Resume file must exist") + + var savedCfg ResumeCfg + err = goconfig.Load(&savedCfg, resumeFile) + require.NoError(t, err) + require.True(t, savedCfg.Index > 0, "Saved index must be greater than 0") + require.NotEmpty(t, savedCfg.ResumeFrom, "Saved ResumeFrom must not be empty") + + // --- Resumed Scan --- + var secondRunProcessed sync.Map + opts2 := &Options{ + InputTargetHost: targets, + Threads: 4, + Delay: 0, + NoColor: true, + Resume: true, + resumeCfg: &ResumeCfg{Index: savedCfg.Index, ResumeFrom: savedCfg.ResumeFrom}, + OnResult: func(r Result) { + if r.Err == nil { + secondRunProcessed.Store(r.URL, true) + } + }, + } + + r2, err := New(opts2) + require.NoError(t, err) + + r2.RunEnumeration() + + // Assert that across run 1 + run 2, 100% of targets were processed + allProcessed := make(map[string]bool) + firstRunProcessed.Range(func(key, value any) bool { + allProcessed[key.(string)] = true + return true + }) + secondRunProcessed.Range(func(key, value any) bool { + allProcessed[key.(string)] = true + return true + }) + + require.Equal(t, totalTargets, len(allProcessed), "100% of targets must be processed with no targets dropped") +} diff --git a/runner/runner.go b/runner/runner.go index eed5d886b..973b1a870 100644 --- a/runner/runner.go +++ b/runner/runner.go @@ -49,7 +49,6 @@ import ( "github.com/pkg/errors" "github.com/projectdiscovery/clistats" - "github.com/projectdiscovery/goconfig" "github.com/projectdiscovery/httpx/common/hashes" "github.com/projectdiscovery/retryablehttp-go" sliceutil "github.com/projectdiscovery/utils/slice" @@ -90,7 +89,7 @@ type Runner struct { hm *hybrid.HybridMap excludeCdn bool stats clistats.StatisticsClient - ratelimiter ratelimit.Limiter + ratelimiter *ratelimit.Limiter HostErrorsCache gcache.Cache[string, int] browser *Browser ditClassifier *dit.Classifier @@ -416,11 +415,11 @@ func New(options *Options) (*Runner, error) { runner.hm = hm if options.RateLimitMinute > 0 { - runner.ratelimiter = *ratelimit.New(context.Background(), uint(options.RateLimitMinute), time.Minute) + runner.ratelimiter = ratelimit.New(context.Background(), uint(options.RateLimitMinute), time.Minute) } else if options.RateLimit > 0 { - runner.ratelimiter = *ratelimit.New(context.Background(), uint(options.RateLimit), time.Second) + runner.ratelimiter = ratelimit.New(context.Background(), uint(options.RateLimit), time.Second) } else { - runner.ratelimiter = *ratelimit.NewUnlimited(context.Background()) + runner.ratelimiter = ratelimit.NewUnlimited(context.Background()) } if options.HostMaxErrors >= 0 { @@ -917,7 +916,9 @@ func (r *Runner) Close() { // nolint:errcheck // ignore r.hm.Close() r.hp.Dialer.Close() - r.ratelimiter.Stop() + if r.ratelimiter != nil { + r.ratelimiter.Stop() + } if r.options.HostMaxErrors >= 0 { r.HostErrorsCache.Purge() @@ -1504,6 +1505,7 @@ func (r *Runner) RunEnumeration() { }(nextStep) wg, _ := syncutil.New(syncutil.WithSize(r.options.Threads)) + var completionWG sync.WaitGroup processItem := func(k string) error { select { @@ -1512,10 +1514,11 @@ func (r *Runner) RunEnumeration() { default: } + var itemIndex int if r.options.resumeCfg != nil { - r.options.resumeCfg.current = k - r.options.resumeCfg.currentIndex++ - if r.options.resumeCfg.currentIndex <= r.options.resumeCfg.Index { + var skip bool + itemIndex, skip = r.options.resumeCfg.NextIndex(k) + if skip { return nil } } @@ -1528,16 +1531,19 @@ func (r *Runner) RunEnumeration() { } } + itemWG := &sync.WaitGroup{} + itemWG.Add(1) + runProcess := func(times int) { for i := 0; i < times; i++ { if len(r.options.requestURIs) > 0 { for _, p := range r.options.requestURIs { scanopts := r.scanopts.Clone() scanopts.RequestURI = p - r.process(k, wg, r.hp, protocol, scanopts, output) + r.process(k, wg, r.hp, protocol, scanopts, output, itemWG) } } else { - r.process(k, wg, r.hp, protocol, &r.scanopts, output) + r.process(k, wg, r.hp, protocol, &r.scanopts, output, itemWG) } } } @@ -1552,6 +1558,17 @@ func (r *Runner) RunEnumeration() { runProcess(cnt) } + itemWG.Done() + + if r.options.resumeCfg != nil { + completionWG.Add(1) + go func(idx int, target string, iwg *sync.WaitGroup) { + defer completionWG.Done() + iwg.Wait() + r.options.resumeCfg.MarkCompleted(idx, target) + }(itemIndex, k, itemWG) + } + return nil } @@ -1569,6 +1586,7 @@ func (r *Runner) RunEnumeration() { } wg.Wait() + completionWG.Wait() close(output) @@ -1665,7 +1683,11 @@ func (r *Runner) Process(t string, wg *syncutil.AdaptiveWaitGroup, protocol stri r.process(t, wg, r.hp, protocol, scanopts, output) } -func (r *Runner) process(t string, wg *syncutil.AdaptiveWaitGroup, hp *httpx.HTTPX, protocol string, scanopts *ScanOptions, output chan Result) { +func (r *Runner) process(t string, wg *syncutil.AdaptiveWaitGroup, hp *httpx.HTTPX, protocol string, scanopts *ScanOptions, output chan Result, itemWGs ...*sync.WaitGroup) { + var itemWG *sync.WaitGroup + if len(itemWGs) > 0 { + itemWG = itemWGs[0] + } // attempts to set the workpool size to the number of threads if r.options.Threads > 0 && wg.Size != r.options.Threads { if err := wg.Resize(context.Background(), r.options.Threads); err != nil { @@ -1685,9 +1707,17 @@ func (r *Runner) process(t string, wg *syncutil.AdaptiveWaitGroup, hp *httpx.HTT for _, prot := range protocols { // sleep for delay time time.Sleep(r.options.Delay) + if itemWG != nil { + itemWG.Add(1) + } wg.Add() go func(target httpx.Target, method, protocol string) { - defer wg.Done() + defer func() { + if itemWG != nil { + itemWG.Done() + } + wg.Done() + }() result := r.analyze(hp, protocol, target, method, t, scanopts) output <- result if scanopts.TLSProbe && result.TLSData != nil { @@ -1695,10 +1725,10 @@ func (r *Runner) process(t string, wg *syncutil.AdaptiveWaitGroup, hp *httpx.HTT if !r.testAndSet(tt) { continue } - r.process(tt, wg, hp, protocol, scanopts, output) + r.process(tt, wg, hp, protocol, scanopts, output, itemWG) } if r.testAndSet(result.TLSData.SubjectCN) { - r.process(result.TLSData.SubjectCN, wg, hp, protocol, scanopts, output) + r.process(result.TLSData.SubjectCN, wg, hp, protocol, scanopts, output, itemWG) } } if scanopts.CSPProbe && result.CSPData != nil { @@ -1709,7 +1739,7 @@ func (r *Runner) process(t string, wg *syncutil.AdaptiveWaitGroup, hp *httpx.HTT if !r.testAndSet(tt) { continue } - r.process(tt, wg, hp, protocol, scanopts, output) + r.process(tt, wg, hp, protocol, scanopts, output, itemWG) } } }(target, method, prot) @@ -1733,9 +1763,17 @@ func (r *Runner) process(t string, wg *syncutil.AdaptiveWaitGroup, hp *httpx.HTT for _, method := range scanopts.Methods { // sleep for delay time time.Sleep(r.options.Delay) + if itemWG != nil { + itemWG.Add(1) + } wg.Add() go func(port int, target httpx.Target, method, protocol string) { - defer wg.Done() + defer func() { + if itemWG != nil { + itemWG.Done() + } + wg.Done() + }() if urlx, err := r.parseURL(target.Host); err != nil { gologger.Warning().Msgf("failed to update port of %v got %v", target.Host, err) } else { @@ -1749,10 +1787,10 @@ func (r *Runner) process(t string, wg *syncutil.AdaptiveWaitGroup, hp *httpx.HTT if !r.testAndSet(tt) { continue } - r.process(tt, wg, hp, protocol, scanopts, output) + r.process(tt, wg, hp, protocol, scanopts, output, itemWG) } if r.testAndSet(result.TLSData.SubjectCN) { - r.process(result.TLSData.SubjectCN, wg, hp, protocol, scanopts, output) + r.process(result.TLSData.SubjectCN, wg, hp, protocol, scanopts, output, itemWG) } } }(port, target, method, wantedProtocol) @@ -2911,10 +2949,10 @@ func extractPotentialFavIconsURLs(resp []byte) (candidates []string, baseHref st // SaveResumeConfig to file func (r *Runner) SaveResumeConfig() error { - var resumeCfg ResumeCfg - resumeCfg.Index = r.options.resumeCfg.currentIndex - resumeCfg.ResumeFrom = r.options.resumeCfg.current - return goconfig.Save(resumeCfg, DefaultResumeFile) + if r.options.resumeCfg == nil { + return nil + } + return r.options.resumeCfg.Save(DefaultResumeFile) } // JSON the result