From ae90caaf039d00556e9fa21caf19824b543a3a78 Mon Sep 17 00:00:00 2001 From: Kristinn Vikar Date: Fri, 21 Jun 2024 15:03:11 +0000 Subject: [PATCH] Clustering performance improvements --- pkg/protocols/dns/cluster.go | 7 ++- pkg/protocols/http/cluster.go | 8 ++- pkg/protocols/ssl/ssl.go | 8 ++- pkg/templates/cluster.go | 113 ++++++++++++++++------------------ 4 files changed, 66 insertions(+), 70 deletions(-) diff --git a/pkg/protocols/dns/cluster.go b/pkg/protocols/dns/cluster.go index f2f5734785..f89372b65d 100644 --- a/pkg/protocols/dns/cluster.go +++ b/pkg/protocols/dns/cluster.go @@ -6,9 +6,6 @@ package dns // are similar enough to be considered one and can be checked by // just adding the matcher/extractors for the request and the correct IDs. func (request *Request) CanCluster(other *Request) bool { - if len(request.Resolvers) > 0 || request.Trace || request.ID != "" { - return false - } if request.Name != other.Name || request.class != other.class || request.Retries != other.Retries || @@ -25,3 +22,7 @@ func (request *Request) CanCluster(other *Request) bool { } return true } + +func (request *Request) IsClusterable() bool { + return (len(request.Resolvers) > 0 || request.Trace || request.ID != "") +} diff --git a/pkg/protocols/http/cluster.go b/pkg/protocols/http/cluster.go index 3839817454..1ce164ec45 100644 --- a/pkg/protocols/http/cluster.go +++ b/pkg/protocols/http/cluster.go @@ -11,9 +11,6 @@ import ( // are similar enough to be considered one and can be checked by // just adding the matcher/extractors for the request and the correct IDs. func (request *Request) CanCluster(other *Request) bool { - if len(request.Payloads) > 0 || len(request.Fuzzing) > 0 || len(request.Raw) > 0 || len(request.Body) > 0 || request.Unsafe || request.NeedsRequestCondition() || request.Name != "" { - return false - } if request.Method != other.Method || request.MaxRedirects != other.MaxRedirects || request.DisableCookie != other.DisableCookie || @@ -28,3 +25,8 @@ func (request *Request) CanCluster(other *Request) bool { } return true } + + +func (request *Request) IsClusterable() bool { + return (len(request.Payloads) > 0 || len(request.Fuzzing) > 0 || len(request.Raw) > 0 || len(request.Body) > 0 || request.Unsafe || request.NeedsRequestCondition() || request.Name != "") +} diff --git a/pkg/protocols/ssl/ssl.go b/pkg/protocols/ssl/ssl.go index 49d612b882..e81753534f 100644 --- a/pkg/protocols/ssl/ssl.go +++ b/pkg/protocols/ssl/ssl.go @@ -99,17 +99,19 @@ type Request struct { options *protocols.ExecutorOptions } + // CanCluster returns true if the request can be clustered. func (request *Request) CanCluster(other *Request) bool { - if len(request.CipherSuites) > 0 || request.MinVersion != "" || request.MaxVersion != "" { - return false - } if request.Address != other.Address || request.ScanMode != other.ScanMode { return false } return true } +func (request *Request) IsClusterable() bool { + return (len(request.CipherSuites) > 0 || request.MinVersion != "" || request.MaxVersion != "") +} + // Compile compiles the request generators preparing any requests possible. func (request *Request) Compile(options *protocols.ExecutorOptions) error { request.options = options diff --git a/pkg/templates/cluster.go b/pkg/templates/cluster.go index 603378e986..f82b60ae30 100644 --- a/pkg/templates/cluster.go +++ b/pkg/templates/cluster.go @@ -14,7 +14,6 @@ import ( "github.com/projectdiscovery/nuclei/v3/pkg/scan" "github.com/projectdiscovery/nuclei/v3/pkg/templates/types" cryptoutil "github.com/projectdiscovery/utils/crypto" - mapsutil "github.com/projectdiscovery/utils/maps" ) // Cluster clusters a list of templates into a lesser number if possible based @@ -42,90 +41,82 @@ import ( // Finally, the engine creates a single executer with a clusteredexecuter for all templates // in a cluster. func Cluster(list []*Template) [][]*Template { + http := make(map[string]*Template) + dns := make(map[string]*Template) + ssl := make(map[string]*Template) + final := [][]*Template{} - skip := mapsutil.NewSyncLockMap[string, struct{}]() + // Split up templates that might be clusterable for _, template := range list { key := template.Path - if skip.Has(key) { - continue - } - - // We only cluster http, dns and ssl requests as of now. - // Take care of requests that can't be clustered first. - if len(template.RequestsHTTP) == 0 && len(template.RequestsDNS) == 0 && len(template.RequestsSSL) == 0 { - _ = skip.Set(key, struct{}{}) - final = append(final, []*Template{template}) - continue - } - // it is not possible to cluster flow and multiprotocol due to dependent execution if template.Flow != "" || template.Options.IsMultiProtocol { - _ = skip.Set(key, struct{}{}) final = append(final, []*Template{template}) continue } - _ = skip.Set(key, struct{}{}) - - var templateType types.ProtocolType switch { case len(template.RequestsDNS) == 1: - templateType = types.DNSProtocol + if template.RequestsDNS[0].IsClusterable() { + dns[key] = template + } else { + final = append(final, []*Template{template}) + } case len(template.RequestsHTTP) == 1: - templateType = types.HTTPProtocol + if template.RequestsHTTP[0].IsClusterable() { + http[key] = template + } else { + final = append(final, []*Template{template}) + } case len(template.RequestsSSL) == 1: - templateType = types.SSLProtocol + if template.RequestsSSL[0].IsClusterable() { + ssl[key] = template + } else { + final = append(final, []*Template{template}) + } + default: + final = append(final, []*Template{template}) } + } - // Find any/all similar matching request that is identical to - // this one and cluster them together for http protocol only. - cluster := []*Template{} - for _, other := range list { - otherKey := other.Path + // Cluster together dns, http and ssl individually - if skip.Has(otherKey) { - continue + for key, template := range dns { + cluster := []*Template{template} + delete(dns, key) + for otherKey, other := range dns { + if template.RequestsDNS[0].CanCluster(other.RequestsDNS[0]) { + delete(dns, otherKey) + cluster = append(cluster, other) } + } + final = append(final, cluster) + } - // it is not possible to cluster flow and multiprotocol due to dependent execution - if other.Flow != "" || other.Options.IsMultiProtocol { - _ = skip.Set(otherKey, struct{}{}) - final = append(final, []*Template{other}) - continue + for key, template := range http { + cluster := []*Template{template} + delete(http, key) + for otherKey, other := range http { + if template.RequestsHTTP[0].CanCluster(other.RequestsHTTP[0]) { + delete(http, otherKey) + cluster = append(cluster, other) } + } + final = append(final, cluster) + } - switch templateType { - case types.DNSProtocol: - if len(other.RequestsDNS) != 1 { - continue - } else if template.RequestsDNS[0].CanCluster(other.RequestsDNS[0]) { - _ = skip.Set(otherKey, struct{}{}) - cluster = append(cluster, other) - } - case types.HTTPProtocol: - if len(other.RequestsHTTP) != 1 { - continue - } else if template.RequestsHTTP[0].CanCluster(other.RequestsHTTP[0]) { - _ = skip.Set(otherKey, struct{}{}) - cluster = append(cluster, other) - } - case types.SSLProtocol: - if len(other.RequestsSSL) != 1 { - continue - } else if template.RequestsSSL[0].CanCluster(other.RequestsSSL[0]) { - _ = skip.Set(otherKey, struct{}{}) - cluster = append(cluster, other) - } + for key, template := range ssl { + cluster := []*Template{template} + delete(ssl, key) + for otherKey, other := range ssl { + if template.RequestsSSL[0].CanCluster(other.RequestsSSL[0]) { + delete(ssl, otherKey) + cluster = append(cluster, other) } } - if len(cluster) > 0 { - cluster = append(cluster, template) - final = append(final, cluster) - } else { - final = append(final, []*Template{template}) - } + final = append(final, cluster) } return final }