From 9efba05e0ca02dcda33d206e2c61c71bf4a9c511 Mon Sep 17 00:00:00 2001 From: Mike Rheinheimer Date: Mon, 18 Jul 2022 15:35:53 -0500 Subject: [PATCH] expose hosterrorscache.Cache as an interface (#2291) * expose hosterrorscache as an interface, change signature to capture the error reason * use the hosterrorscache.CacheInterface as struct field so users of Nuclei embedded can provide their own cache implementation Co-authored-by: Mike Rheinheimer --- .gitignore | 2 ++ v2/internal/runner/runner.go | 7 ++++--- v2/pkg/core/workflow_execute.go | 4 +--- v2/pkg/protocols/common/executer/executer.go | 8 ++----- .../common/hosterrorscache/hosterrorscache.go | 21 ++++++++++++++----- .../hosterrorscache/hosterrorscache_test.go | 9 ++++---- v2/pkg/protocols/http/request.go | 4 ++-- v2/pkg/protocols/protocols.go | 2 +- v2/pkg/templates/cluster.go | 8 +++---- 9 files changed, 37 insertions(+), 28 deletions(-) diff --git a/.gitignore b/.gitignore index 9a9771c71..b1640ce4a 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,6 @@ .idea +.vscode +v2/vendor integration_tests/nuclei integration_tests/integration-test v2/cmd/nuclei/main diff --git a/v2/internal/runner/runner.go b/v2/internal/runner/runner.go index 9091668cf..6531a8255 100644 --- a/v2/internal/runner/runner.go +++ b/v2/internal/runner/runner.go @@ -65,7 +65,7 @@ type Runner struct { hmapInputProvider *hybrid.Input browser *engine.Browser ratelimiter ratelimit.Limiter - hostErrors *hosterrorscache.Cache + hostErrors hosterrorscache.CacheInterface resumeCfg *types.ResumeCfg pprofServer *http.Server } @@ -345,7 +345,8 @@ func (r *Runner) RunEnumeration() error { } var cache *hosterrorscache.Cache if r.options.MaxHostError > 0 { - cache = hosterrorscache.New(r.options.MaxHostError, hosterrorscache.DefaultMaxHostsCount).SetVerbose(r.options.Verbose) + cache = hosterrorscache.New(r.options.MaxHostError, hosterrorscache.DefaultMaxHostsCount) + cache.SetVerbose(r.options.Verbose) } r.hostErrors = cache @@ -581,7 +582,7 @@ func (r *Runner) readNewTemplatesFile() ([]string, error) { file, err := os.Open(additionsFile) if err != nil { return nil, err - } + } defer file.Close() templatesList := []string{} diff --git a/v2/pkg/core/workflow_execute.go b/v2/pkg/core/workflow_execute.go index aae9955ad..015a74829 100644 --- a/v2/pkg/core/workflow_execute.go +++ b/v2/pkg/core/workflow_execute.go @@ -59,9 +59,7 @@ func (e *Engine) runWorkflowStep(template *workflows.WorkflowTemplate, input str } if err != nil { if w.Options.HostErrorsCache != nil { - if w.Options.HostErrorsCache.CheckError(err) { - w.Options.HostErrorsCache.MarkFailed(input) - } + w.Options.HostErrorsCache.MarkFailed(input, err) } if len(template.Executers) == 1 { mainErr = err diff --git a/v2/pkg/protocols/common/executer/executer.go b/v2/pkg/protocols/common/executer/executer.go index a8386a44f..6543b6cfd 100644 --- a/v2/pkg/protocols/common/executer/executer.go +++ b/v2/pkg/protocols/common/executer/executer.go @@ -96,9 +96,7 @@ func (e *Executer) Execute(input string) (bool, error) { }) if err != nil { if e.options.HostErrorsCache != nil { - if e.options.HostErrorsCache.CheckError(err) { - e.options.HostErrorsCache.MarkFailed(input) - } + e.options.HostErrorsCache.MarkFailed(input, err) } gologger.Warning().Msgf("[%s] Could not execute request for %s: %s\n", e.options.TemplateID, input, err) } @@ -139,9 +137,7 @@ func (e *Executer) ExecuteWithResults(input string, callback protocols.OutputEve }) if err != nil { if e.options.HostErrorsCache != nil { - if e.options.HostErrorsCache.CheckError(err) { - e.options.HostErrorsCache.MarkFailed(input) - } + e.options.HostErrorsCache.MarkFailed(input, err) } gologger.Warning().Msgf("[%s] Could not execute request for %s: %s\n", e.options.TemplateID, input, err) } diff --git a/v2/pkg/protocols/common/hosterrorscache/hosterrorscache.go b/v2/pkg/protocols/common/hosterrorscache/hosterrorscache.go index c509ff847..5067df906 100644 --- a/v2/pkg/protocols/common/hosterrorscache/hosterrorscache.go +++ b/v2/pkg/protocols/common/hosterrorscache/hosterrorscache.go @@ -11,6 +11,15 @@ import ( "github.com/projectdiscovery/gologger" ) +// CacheInterface defines the signature of the hosterrorscache so that +// users of Nuclei as embedded lib may implement their own cache +type CacheInterface interface { + SetVerbose(verbose bool) // log verbosely + Close() // close the cache + Check(value string) bool // return true if the host should be skipped + MarkFailed(value string, err error) // record a failure (and cause) for the host +} + // Cache is a cache for host based errors. It allows skipping // certain hosts based on an error threshold. // @@ -33,9 +42,8 @@ func New(maxHostError, maxHostsCount int) *Cache { } // SetVerbose sets the cache to log at verbose level -func (c *Cache) SetVerbose(verbose bool) *Cache { +func (c *Cache) SetVerbose(verbose bool) { c.verbose = verbose - return c } // Close closes the host errors cache @@ -99,7 +107,10 @@ func (c *Cache) Check(value string) bool { } // MarkFailed marks a host as failed previously -func (c *Cache) MarkFailed(value string) { +func (c *Cache) MarkFailed(value string, err error) { + if !c.checkError(err) { + return + } finalValue := c.normalizeCacheValue(value) if !c.failedTargets.Has(finalValue) { _ = c.failedTargets.Set(finalValue, 1) @@ -118,9 +129,9 @@ func (c *Cache) MarkFailed(value string) { var checkErrorRegexp = regexp.MustCompile(`(no address found for host|Client\.Timeout exceeded while awaiting headers|could not resolve host)`) -// CheckError checks if an error represents a type that should be +// checkError checks if an error represents a type that should be // added to the host skipping table. -func (c *Cache) CheckError(err error) bool { +func (c *Cache) checkError(err error) bool { errString := err.Error() return checkErrorRegexp.MatchString(errString) } diff --git a/v2/pkg/protocols/common/hosterrorscache/hosterrorscache_test.go b/v2/pkg/protocols/common/hosterrorscache/hosterrorscache_test.go index a366a65bd..9586eb99d 100644 --- a/v2/pkg/protocols/common/hosterrorscache/hosterrorscache_test.go +++ b/v2/pkg/protocols/common/hosterrorscache/hosterrorscache_test.go @@ -1,6 +1,7 @@ package hosterrorscache import ( + "fmt" "testing" "github.com/stretchr/testify/require" @@ -9,20 +10,20 @@ import ( func TestCacheCheckMarkFailed(t *testing.T) { cache := New(3, DefaultMaxHostsCount) - cache.MarkFailed("http://example.com:80") + cache.MarkFailed("http://example.com:80", fmt.Errorf("no address found for host")) if value, err := cache.failedTargets.Get("http://example.com:80"); err == nil && value != nil { require.Equal(t, 1, value, "could not get correct number of marked failed hosts") } - cache.MarkFailed("example.com:80") + cache.MarkFailed("example.com:80", fmt.Errorf("Client.Timeout exceeded while awaiting headers")) if value, err := cache.failedTargets.Get("example.com:80"); err == nil && value != nil { require.Equal(t, 2, value, "could not get correct number of marked failed hosts") } - cache.MarkFailed("example.com") + cache.MarkFailed("example.com", fmt.Errorf("could not resolve host")) if value, err := cache.failedTargets.Get("example.com"); err == nil && value != nil { require.Equal(t, 1, value, "could not get correct number of marked failed hosts") } for i := 0; i < 3; i++ { - cache.MarkFailed("test") + cache.MarkFailed("test", fmt.Errorf("could not resolve host")) } value := cache.Check("test") diff --git a/v2/pkg/protocols/http/request.go b/v2/pkg/protocols/http/request.go index 38abb7d62..1c8e764c8 100644 --- a/v2/pkg/protocols/http/request.go +++ b/v2/pkg/protocols/http/request.go @@ -299,8 +299,8 @@ func (request *Request) ExecuteWithResults(reqURL string, dynamicValues, previou return true, nil } if err != nil { - if request.options.HostErrorsCache != nil && request.options.HostErrorsCache.CheckError(err) { - request.options.HostErrorsCache.MarkFailed(reqURL) + if request.options.HostErrorsCache != nil { + request.options.HostErrorsCache.MarkFailed(reqURL, err) } requestErr = err } diff --git a/v2/pkg/protocols/protocols.go b/v2/pkg/protocols/protocols.go index 19cb05f05..57d1ecbb9 100644 --- a/v2/pkg/protocols/protocols.go +++ b/v2/pkg/protocols/protocols.go @@ -62,7 +62,7 @@ type ExecuterOptions struct { // Interactsh is a client for interactsh oob polling server Interactsh *interactsh.Client // HostErrorsCache is an optional cache for handling host errors - HostErrorsCache *hosterrorscache.Cache + HostErrorsCache hosterrorscache.CacheInterface // Stop execution once first match is found StopAtFirstMatch bool // Variables is a list of variables from template diff --git a/v2/pkg/templates/cluster.go b/v2/pkg/templates/cluster.go index 82319e0a6..78fad59b8 100644 --- a/v2/pkg/templates/cluster.go +++ b/v2/pkg/templates/cluster.go @@ -206,8 +206,8 @@ func (e *ClusterExecuter) Execute(input string) (bool, error) { } } }) - if err != nil && e.options.HostErrorsCache != nil && e.options.HostErrorsCache.CheckError(err) { - e.options.HostErrorsCache.MarkFailed(input) + if err != nil && e.options.HostErrorsCache != nil { + e.options.HostErrorsCache.MarkFailed(input, err) } return results, err } @@ -228,8 +228,8 @@ func (e *ClusterExecuter) ExecuteWithResults(input string, callback protocols.Ou } } }) - if err != nil && e.options.HostErrorsCache != nil && e.options.HostErrorsCache.CheckError(err) { - e.options.HostErrorsCache.MarkFailed(input) + if err != nil && e.options.HostErrorsCache != nil { + e.options.HostErrorsCache.MarkFailed(input, err) } return err }