diff --git a/runner/atomic_resume.go b/runner/atomic_resume.go new file mode 100644 index 00000000..e79cafd6 --- /dev/null +++ b/runner/atomic_resume.go @@ -0,0 +1,72 @@ +package runner + +import ( + "os" + "path/filepath" + "runtime" + "sync" + + "gopkg.in/yaml.v3" +) + +var atomicResumeMutex sync.Mutex + +// syncDir attempts to fsync the directory containing targetPath on POSIX platforms. +func syncDir(dirPath string) error { + if runtime.GOOS == "windows" { + return nil + } + d, err := os.Open(dirPath) + if err != nil { + return err + } + defer d.Close() + return d.Sync() +} + +// SaveAtomic writes data to a temporary file in targetPath's directory, syncs to disk, and atomically renames it over targetPath. +func SaveAtomic(targetPath string, data []byte) error { + atomicResumeMutex.Lock() + defer atomicResumeMutex.Unlock() + + dir := filepath.Dir(targetPath) + tmpFile, err := os.CreateTemp(dir, "httpx-resume-*.tmp") + if err != nil { + return err + } + tmpName := tmpFile.Name() + defer os.Remove(tmpName) + + if _, err := tmpFile.Write(data); err != nil { + tmpFile.Close() + return err + } + if err := tmpFile.Sync(); err != nil { + tmpFile.Close() + return err + } + if err := tmpFile.Close(); err != nil { + return err + } + + if err := os.Rename(tmpName, targetPath); err != nil { + return err + } + + return syncDir(dir) +} + +// SaveResumeConfigAtomic serializes the resume config and writes it using SaveAtomic. +func (r *Runner) SaveResumeConfigAtomic() error { + if r.options == nil || r.options.resumeCfg == nil { + return nil + } + var resumeCfg ResumeCfg + resumeCfg.Index = r.options.resumeCfg.currentIndex + resumeCfg.ResumeFrom = r.options.resumeCfg.current + data, err := yaml.Marshal(resumeCfg) + if err != nil { + return err + } + return SaveAtomic(DefaultResumeFile, data) +} diff --git a/runner/atomic_resume_test.go b/runner/atomic_resume_test.go new file mode 100644 index 00000000..257b5808 --- /dev/null +++ b/runner/atomic_resume_test.go @@ -0,0 +1,49 @@ +package runner + +import ( + "os" + "path/filepath" + "sync" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestSaveAtomic(t *testing.T) { + tempDir := t.TempDir() + targetPath := filepath.Join(tempDir, "test_resume.cfg") + + // 1. Basic atomic write + data := []byte("resume_index: 42\nresume_from: example.com\n") + err := SaveAtomic(targetPath, data) + require.NoError(t, err) + + readData, err := os.ReadFile(targetPath) + require.NoError(t, err) + require.Equal(t, data, readData) + + // 2. Overwrite atomically + newData := []byte("resume_index: 100\nresume_from: target.org\n") + err = SaveAtomic(targetPath, newData) + require.NoError(t, err) + + readData, err = os.ReadFile(targetPath) + require.NoError(t, err) + require.Equal(t, newData, readData) + + // 3. Concurrent saves + var wg sync.WaitGroup + for i := 0; i < 10; i++ { + wg.Add(1) + go func(idx int) { + defer wg.Done() + payload := []byte("concurrent_data") + _ = SaveAtomic(targetPath, payload) + }(i) + } + wg.Wait() + + finalData, err := os.ReadFile(targetPath) + require.NoError(t, err) + require.Equal(t, []byte("concurrent_data"), finalData) +} diff --git a/runner/runner.go b/runner/runner.go index ea0cac3f..9bbff456 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" @@ -80,25 +79,25 @@ import ( // Runner is a client for running the enumeration process. type Runner struct { - seenMux sync.Mutex - options *Options - hp *httpx.HTTPX - wappalyzer *wappalyzer.Wappalyze - cpeDetector *CPEDetector - wpDetector *WordPressDetector - scanopts ScanOptions - hm *hybrid.HybridMap - excludeCdn bool - stats clistats.StatisticsClient - ratelimiter ratelimit.Limiter - HostErrorsCache gcache.Cache[string, int] - browser *Browser - ditClassifier *dit.Classifier - pHashClusters []pHashCluster - simHashes gcache.Cache[uint64, []string] - httpApiEndpoint *Server - authProvider authprovider.AuthProvider - interruptCh chan struct{} + seenMux sync.Mutex + options *Options + hp *httpx.HTTPX + wappalyzer *wappalyzer.Wappalyze + cpeDetector *CPEDetector + wpDetector *WordPressDetector + scanopts ScanOptions + hm *hybrid.HybridMap + excludeCdn bool + stats clistats.StatisticsClient + ratelimiter ratelimit.Limiter + HostErrorsCache gcache.Cache[string, int] + browser *Browser + ditClassifier *dit.Classifier + pHashClusters []pHashCluster + simHashes gcache.Cache[uint64, []string] + httpApiEndpoint *Server + authProvider authprovider.AuthProvider + interruptCh chan struct{} } func (r *Runner) HTTPX() *httpx.HTTPX { @@ -2667,60 +2666,60 @@ retry: } result := Result{ - Timestamp: time.Now(), - Request: request, - LinkRequest: linkRequest, - ResponseHeaders: responseHeaders, - RawHeaders: rawResponseHeaders, - Scheme: parsed.Scheme, - Port: finalPort, - Path: finalPath, - Raw: resp.Raw, - URL: fullURL, - Input: origInput, - ContentLength: resp.ContentLength, - ChainStatusCodes: chainStatusCodes, - Chain: chainItems, - StatusCode: resp.StatusCode, - Location: resp.GetHeaderPart("Location", ";"), - ContentType: resp.GetHeaderPart("Content-Type", ";"), - Title: title, - str: builder.String(), - VHost: isvhost, - WebServer: serverHeader, - ResponseBody: serverResponseRaw, - BodyPreview: bodyPreview, - WebSocket: isWebSocket, - TLSData: resp.TLSData, - CSPData: resp.CSPData, - Pipeline: pipeline, - HTTP2: http2, - Method: method, - Host: parsed.Hostname(), - HostIP: ip, - A: ips4, - AAAA: ips6, - CNAMEs: cnames, - CDN: isCDN, - CDNName: cdnName, - CDNType: cdnType, - ResponseTime: resp.Duration.String(), - Technologies: technologies, - FinalURL: finalURL, - FavIconMMH3: faviconMMH3, - FavIconMD5: faviconMD5, - FaviconPath: faviconPath, - FaviconURL: faviconURL, - Hashes: hashesMap, - Extracts: extractResult, - JarmHash: jarmhash, - Lines: resp.Lines, - Words: resp.Words, - ASN: asnResponse, - ExtractRegex: extractRegex, - ScreenshotBytes: screenshotBytes, - HeadlessBody: headlessBody, - KnowledgeBase: r.classifyPage(headlessBody, respData, pHash), + Timestamp: time.Now(), + Request: request, + LinkRequest: linkRequest, + ResponseHeaders: responseHeaders, + RawHeaders: rawResponseHeaders, + Scheme: parsed.Scheme, + Port: finalPort, + Path: finalPath, + Raw: resp.Raw, + URL: fullURL, + Input: origInput, + ContentLength: resp.ContentLength, + ChainStatusCodes: chainStatusCodes, + Chain: chainItems, + StatusCode: resp.StatusCode, + Location: resp.GetHeaderPart("Location", ";"), + ContentType: resp.GetHeaderPart("Content-Type", ";"), + Title: title, + str: builder.String(), + VHost: isvhost, + WebServer: serverHeader, + ResponseBody: serverResponseRaw, + BodyPreview: bodyPreview, + WebSocket: isWebSocket, + TLSData: resp.TLSData, + CSPData: resp.CSPData, + Pipeline: pipeline, + HTTP2: http2, + Method: method, + Host: parsed.Hostname(), + HostIP: ip, + A: ips4, + AAAA: ips6, + CNAMEs: cnames, + CDN: isCDN, + CDNName: cdnName, + CDNType: cdnType, + ResponseTime: resp.Duration.String(), + Technologies: technologies, + FinalURL: finalURL, + FavIconMMH3: faviconMMH3, + FavIconMD5: faviconMD5, + FaviconPath: faviconPath, + FaviconURL: faviconURL, + Hashes: hashesMap, + Extracts: extractResult, + JarmHash: jarmhash, + Lines: resp.Lines, + Words: resp.Words, + ASN: asnResponse, + ExtractRegex: extractRegex, + ScreenshotBytes: screenshotBytes, + HeadlessBody: headlessBody, + KnowledgeBase: r.classifyPage(headlessBody, respData, pHash), TechnologyDetails: technologyDetails, Resolvers: resolvers, RequestRaw: requestDump, @@ -2937,12 +2936,9 @@ func extractPotentialFavIconsURLs(resp []byte) (candidates []string, baseHref st return candidates, baseHref, nil } -// SaveResumeConfig to file +// SaveResumeConfig saves the current resume configuration state to file atomically. 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) + return r.SaveResumeConfigAtomic() } // JSON the result