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 <mrheinheimer@atlassian.com>

.gitignore

@@ -1,4 +1,6 @@
 .idea
+.vscode
+v2/vendor
 integration_tests/nuclei
 integration_tests/integration-test
 v2/cmd/nuclei/main

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
 

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, err)
-						w.Options.HostErrorsCache.MarkFailed(input)
-					}
 				}
 				if len(template.Executers) == 1 {
 					mainErr = err

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, err)
-					e.options.HostErrorsCache.MarkFailed(input)
-				}
 			}
 			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, err)
-					e.options.HostErrorsCache.MarkFailed(input)
-				}
 			}
 			gologger.Warning().Msgf("[%s] Could not execute request for %s: %s\n", e.options.TemplateID, input, err)
 		}

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)
 }

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")

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) {
+				if request.options.HostErrorsCache != nil {
-					request.options.HostErrorsCache.MarkFailed(reqURL)
+					request.options.HostErrorsCache.MarkFailed(reqURL, err)
 				}
 				requestErr = err
 			}

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

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) {
+	if err != nil && e.options.HostErrorsCache != nil {
-		e.options.HostErrorsCache.MarkFailed(input)
+		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) {
+	if err != nil && e.options.HostErrorsCache != nil {
-		e.options.HostErrorsCache.MarkFailed(input)
+		e.options.HostErrorsCache.MarkFailed(input, err)
 	}
 	return err
 }