Skip to content

Commit

Permalink
clear after stop (#3312)
Browse files Browse the repository at this point in the history
* clear after stop

* fixing data races

* adding atomic cache

* fixing lint errors

* fixing imports
  • Loading branch information
Mzack9999 committed Feb 13, 2023
1 parent 20da19a commit d608ffa
Show file tree
Hide file tree
Showing 7 changed files with 111 additions and 23 deletions.
2 changes: 1 addition & 1 deletion v2/cmd/integration-test/http.go
Original file line number Diff line number Diff line change
Expand Up @@ -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{},
Expand Down
4 changes: 3 additions & 1 deletion v2/go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
9 changes: 6 additions & 3 deletions v2/go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down Expand Up @@ -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=
Expand Down
27 changes: 13 additions & 14 deletions v2/pkg/protocols/common/interactsh/interactsh.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)

Expand All @@ -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
Expand All @@ -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 (
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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)
}
Expand Down Expand Up @@ -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()
}
}
Expand All @@ -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
Expand Down
28 changes: 25 additions & 3 deletions v2/pkg/protocols/http/httpclientpool/clientpool.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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")
Expand Down
2 changes: 1 addition & 1 deletion v2/pkg/protocols/http/request.go
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
62 changes: 62 additions & 0 deletions v2/pkg/utils/atomcache/atomcache.go
Original file line number Diff line number Diff line change
@@ -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()
}

0 comments on commit d608ffa

Please sign in to comment.