From d608ffaeb29bd6ad5557d92a84df18d763f93c0a Mon Sep 17 00:00:00 2001 From: Mzack9999 Date: Mon, 13 Feb 2023 12:16:41 +0100 Subject: [PATCH] clear after stop (#3312) * clear after stop * fixing data races * adding atomic cache * fixing lint errors * fixing imports --- v2/cmd/integration-test/http.go | 2 +- v2/go.mod | 4 +- v2/go.sum | 9 ++- .../protocols/common/interactsh/interactsh.go | 27 ++++---- .../http/httpclientpool/clientpool.go | 28 ++++++++- v2/pkg/protocols/http/request.go | 2 +- v2/pkg/utils/atomcache/atomcache.go | 62 +++++++++++++++++++ 7 files changed, 111 insertions(+), 23 deletions(-) create mode 100644 v2/pkg/utils/atomcache/atomcache.go diff --git a/v2/cmd/integration-test/http.go b/v2/cmd/integration-test/http.go index 478b69565..d9107a3a9 100644 --- a/v2/cmd/integration-test/http.go +++ b/v2/cmd/integration-test/http.go @@ -19,7 +19,7 @@ import ( ) var httpTestcases = map[string]testutils.TestCase{ - //"http/raw-unsafe-request.yaml": &httpRawUnsafeRequest{}, + // "http/raw-unsafe-request.yaml": &httpRawUnsafeRequest{}, "http/get-headers.yaml": &httpGetHeaders{}, "http/get-query-string.yaml": &httpGetQueryString{}, "http/get-redirects.yaml": &httpGetRedirects{}, diff --git a/v2/go.mod b/v2/go.mod index 40973601d..dbfe15b5a 100644 --- a/v2/go.mod +++ b/v2/go.mod @@ -88,6 +88,8 @@ require ( ) require ( + cloud.google.com/go/compute v1.18.0 // indirect + cloud.google.com/go/compute/metadata v0.2.3 // indirect github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.4.10 // indirect github.com/aws/aws-sdk-go-v2/internal/v4a v1.0.19 // indirect github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.9.11 // indirect @@ -162,7 +164,7 @@ require ( go.etcd.io/etcd/tests/v3 v3.5.0-alpha.0 // indirect go.etcd.io/etcd/v3 v3.5.0-alpha.0 // indirect go.uber.org/atomic v1.10.0 // indirect - google.golang.org/genproto v0.0.0-20220804142021-4e6b2dfa6612 // indirect + google.golang.org/genproto v0.0.0-20230124163310-31e0e69b6fc2 // indirect google.golang.org/grpc v1.51.0 // indirect gopkg.in/cheggaaa/pb.v1 v1.0.28 // indirect gopkg.in/djherbis/times.v1 v1.3.0 // indirect diff --git a/v2/go.sum b/v2/go.sum index d81b8bcc6..1110a643f 100644 --- a/v2/go.sum +++ b/v2/go.sum @@ -29,7 +29,10 @@ cloud.google.com/go/bigquery v1.4.0/go.mod h1:S8dzgnTigyfTmLBfrtrhyYhwRxG72rYxvf cloud.google.com/go/bigquery v1.5.0/go.mod h1:snEHRnqQbz117VIFhE8bmtwIDY80NLUZUMb4Nv6dBIg= cloud.google.com/go/bigquery v1.7.0/go.mod h1://okPTzCYNXSlb24MZs83e2Do+h+VXtc4gLoIoXIAPc= cloud.google.com/go/bigquery v1.8.0/go.mod h1:J5hqkt3O0uAFnINi6JXValWIb1v0goeZM77hZzJN/fQ= -cloud.google.com/go/compute/metadata v0.2.0 h1:nBbNSZyDpkNlo3DepaaLKVuO7ClyifSAmNloSCZrHnQ= +cloud.google.com/go/compute v1.18.0 h1:FEigFqoDbys2cvFkZ9Fjq4gnHBP55anJ0yQyau2f9oY= +cloud.google.com/go/compute v1.18.0/go.mod h1:1X7yHxec2Ga+Ss6jPyjxRxpu2uu7PLgsOVXvgU0yacs= +cloud.google.com/go/compute/metadata v0.2.3 h1:mg4jlk7mCAj6xXp9UJ4fjI9VUI5rubuGBW5aJ7UnBMY= +cloud.google.com/go/compute/metadata v0.2.3/go.mod h1:VAV5nSsACxMJvgaAuX6Pk2AawlZn8kiOGuCv6gTkwuA= cloud.google.com/go/datastore v1.0.0/go.mod h1:LXYbyblFSglQ5pkeyhO+Qmw7ukd3C+pD7TKLgZqpHYE= cloud.google.com/go/datastore v1.1.0/go.mod h1:umbIZjpQpHh4hmRpGhH4tLFup+FVzqBi1b3c64qFpCk= cloud.google.com/go/firestore v1.1.0/go.mod h1:ulACoGHTpvq5r8rxGJ4ddJZBZqakUQqClKRT5SZwBmk= @@ -1800,8 +1803,8 @@ google.golang.org/genproto v0.0.0-20210402141018-6c239bbf2bb1/go.mod h1:9lPAdzaE google.golang.org/genproto v0.0.0-20210413151531-c14fb6ef47c3/go.mod h1:P3QM42oQyzQSnHPnZ/vqoCdDmzH28fzWByN9asMeM8A= google.golang.org/genproto v0.0.0-20210510173355-fb37daa5cd7a/go.mod h1:P3QM42oQyzQSnHPnZ/vqoCdDmzH28fzWByN9asMeM8A= google.golang.org/genproto v0.0.0-20210624195500-8bfb893ecb84/go.mod h1:SzzZ/N+nwJDaO1kznhnlzqS8ocJICar6hYhVyhi++24= -google.golang.org/genproto v0.0.0-20220804142021-4e6b2dfa6612 h1:NX3L5YesD5qgxxrPHdKqHH38Ao0AG6poRXG+JljPsGU= -google.golang.org/genproto v0.0.0-20220804142021-4e6b2dfa6612/go.mod h1:iHe1svFLAZg9VWz891+QbRMwUv9O/1Ww+/mngYeThbc= +google.golang.org/genproto v0.0.0-20230124163310-31e0e69b6fc2 h1:O97sLx/Xmb/KIZHB/2/BzofxBs5QmmR0LcihPtllmbc= +google.golang.org/genproto v0.0.0-20230124163310-31e0e69b6fc2/go.mod h1:RGgjbofJ8xD9Sq1VVhDM1Vok1vRONV+rg+CjzG4SZKM= google.golang.org/grpc v1.8.0/go.mod h1:yo6s7OP7yaDglbqo1J04qKzAhqBH6lvTonzMVmEdcZw= google.golang.org/grpc v1.12.0/go.mod h1:yo6s7OP7yaDglbqo1J04qKzAhqBH6lvTonzMVmEdcZw= google.golang.org/grpc v1.17.0/go.mod h1:6QZJwpn2B+Zp71q/5VxRsJ6NXXVCE5NRUHRo+f3cWCs= diff --git a/v2/pkg/protocols/common/interactsh/interactsh.go b/v2/pkg/protocols/common/interactsh/interactsh.go index 225e6a7a4..81ee94993 100644 --- a/v2/pkg/protocols/common/interactsh/interactsh.go +++ b/v2/pkg/protocols/common/interactsh/interactsh.go @@ -24,6 +24,7 @@ import ( "github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/helpers/responsehighlighter" "github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/helpers/writer" "github.com/projectdiscovery/nuclei/v2/pkg/reporting" + "github.com/projectdiscovery/nuclei/v2/pkg/utils/atomcache" "github.com/projectdiscovery/retryablehttp-go" ) @@ -32,13 +33,13 @@ type Client struct { // interactsh is a client for interactsh server. interactsh *client.Client // requests is a stored cache for interactsh-url->request-event data. - requests *ccache.Cache + requests *atomcache.Cache // interactions is a stored cache for interactsh-interaction->interactsh-url data - interactions *ccache.Cache + interactions *atomcache.Cache // matchedTemplates is a stored cache to track matched templates - matchedTemplates *ccache.Cache + matchedTemplates *atomcache.Cache // interactshURLs is a stored cache to track track multiple interactsh markers - interactshURLs *ccache.Cache + interactshURLs *atomcache.Cache options *Options eviction time.Duration @@ -51,7 +52,7 @@ type Client struct { firstTimeGroup sync.Once generated uint32 // decide to wait if we have a generated url - matched bool + matched atomic.Bool } var ( @@ -108,14 +109,14 @@ const defaultMaxInteractionsCount = 5000 func New(options *Options) (*Client, error) { configure := ccache.Configure() configure = configure.MaxSize(options.CacheSize) - cache := ccache.New(configure) + cache := atomcache.NewWithCache(ccache.New(configure)) interactionsCfg := ccache.Configure() interactionsCfg = interactionsCfg.MaxSize(defaultMaxInteractionsCount) - interactionsCache := ccache.New(interactionsCfg) + interactionsCache := atomcache.NewWithCache(ccache.New(interactionsCfg)) - matchedTemplateCache := ccache.New(ccache.Configure().MaxSize(defaultMaxInteractionsCount)) - interactshURLCache := ccache.New(ccache.Configure().MaxSize(defaultMaxInteractionsCount)) + matchedTemplateCache := atomcache.NewWithCache(ccache.New(ccache.Configure().MaxSize(defaultMaxInteractionsCount))) + interactshURLCache := atomcache.NewWithCache(ccache.New(ccache.Configure().MaxSize(defaultMaxInteractionsCount))) interactClient := &Client{ eviction: options.Eviction, @@ -172,7 +173,6 @@ func (c *Client) firstTimeInitializeClient() error { interactsh.StartPolling(c.pollDuration, func(interaction *server.Interaction) { item := c.requests.Get(interaction.UniqueID) - if item == nil { // If we don't have any request for this ID, add it to temporary // lru cache, so we can correlate when we get an add request. @@ -230,7 +230,7 @@ func (c *Client) processInteractionForRequest(interaction *server.Interaction, d } if writer.WriteResult(data.Event, c.options.Output, c.options.Progress, c.options.IssuesClient) { - c.matched = true + c.matched.Store(true) if _, ok := data.Event.InternalEvent[stopAtFirstMatchAttribute]; ok || c.options.StopAtFirstMatch { c.matchedTemplates.Set(hash(data.Event.InternalEvent[templateIdAttribute].(string), data.Event.InternalEvent["host"].(string)), true, defaultInteractionDuration) } @@ -262,9 +262,8 @@ func (c *Client) Close() bool { c.interactsh.Close() } - closeCache := func(cc *ccache.Cache) { + closeCache := func(cc *atomcache.Cache) { if cc != nil { - cc.Clear() cc.Stop() } } @@ -273,7 +272,7 @@ func (c *Client) Close() bool { closeCache(c.matchedTemplates) closeCache(c.interactshURLs) - return c.matched + return c.matched.Load() } // ReplaceMarkers replaces the default {{interactsh-url}} placeholders with interactsh urls diff --git a/v2/pkg/protocols/http/httpclientpool/clientpool.go b/v2/pkg/protocols/http/httpclientpool/clientpool.go index fb882ef1b..c972616b0 100644 --- a/v2/pkg/protocols/http/httpclientpool/clientpool.go +++ b/v2/pkg/protocols/http/httpclientpool/clientpool.go @@ -59,7 +59,29 @@ func Init(options *types.Options) error { type ConnectionConfiguration struct { // DisableKeepAlive of the connection DisableKeepAlive bool - Cookiejar *cookiejar.Jar + cookiejar *cookiejar.Jar + mu sync.RWMutex +} + +func (cc *ConnectionConfiguration) SetCookieJar(cookiejar *cookiejar.Jar) { + cc.mu.Lock() + defer cc.mu.Unlock() + + cc.cookiejar = cookiejar +} + +func (cc *ConnectionConfiguration) GetCookieJar() *cookiejar.Jar { + cc.mu.RLock() + defer cc.mu.RUnlock() + + return cc.cookiejar +} + +func (cc *ConnectionConfiguration) HasCookieJar() bool { + cc.mu.RLock() + defer cc.mu.RUnlock() + + return cc.cookiejar != nil } // Configuration contains the custom configuration options for a client @@ -244,8 +266,8 @@ func wrappedGet(options *types.Options, configuration *Configuration) (*retryabl } var jar *cookiejar.Jar - if configuration.Connection != nil && configuration.Connection.Cookiejar != nil { - jar = configuration.Connection.Cookiejar + if configuration.Connection != nil && configuration.Connection.HasCookieJar() { + jar = configuration.Connection.GetCookieJar() } else if configuration.CookieReuse { if jar, err = cookiejar.New(&cookiejar.Options{PublicSuffixList: publicsuffix.List}); err != nil { return nil, errors.Wrap(err, "could not create cookiejar") diff --git a/v2/pkg/protocols/http/request.go b/v2/pkg/protocols/http/request.go index 97c493666..03e3a080d 100644 --- a/v2/pkg/protocols/http/request.go +++ b/v2/pkg/protocols/http/request.go @@ -557,7 +557,7 @@ func (request *Request) executeRequest(input *contextargs.Context, generatedRequ httpclient := request.httpClient if input.CookieJar != nil { connConfiguration := request.connConfiguration - connConfiguration.Connection.Cookiejar = input.CookieJar + connConfiguration.Connection.SetCookieJar(input.CookieJar) client, err := httpclientpool.Get(request.options.Options, connConfiguration) if err != nil { return errors.Wrap(err, "could not get http client") diff --git a/v2/pkg/utils/atomcache/atomcache.go b/v2/pkg/utils/atomcache/atomcache.go new file mode 100644 index 000000000..ee6cc84b7 --- /dev/null +++ b/v2/pkg/utils/atomcache/atomcache.go @@ -0,0 +1,62 @@ +package atomcache + +import ( + "sync" + "sync/atomic" + "time" + + "github.com/karlseguin/ccache" +) + +type Cache struct { + *ccache.Cache + Closed atomic.Bool + mu sync.RWMutex +} + +func NewWithCache(c *ccache.Cache) *Cache { + return &Cache{Cache: c} +} + +func (c *Cache) Get(key string) *ccache.Item { + if c.Closed.Load() { + return nil + } + c.mu.RLock() + defer c.mu.RUnlock() + + return c.Cache.Get(key) +} + +func (c *Cache) Set(key string, value interface{}, duration time.Duration) { + if c.Closed.Load() { + return + } + c.mu.Lock() + defer c.mu.Unlock() + + c.Cache.Set(key, value, duration) +} + +func (c *Cache) Delete(key string) bool { + if c.Closed.Load() { + return false + } + + c.mu.Lock() + defer c.mu.Unlock() + + return c.Cache.Delete(key) +} + +func (c *Cache) Stop() { + if c.Closed.Load() { + return + } + c.Closed.Store(true) + + c.mu.Lock() + defer c.mu.Unlock() + + c.Cache.Stop() +}