Merge pull request #4833 from projectdiscovery/maint-memory

Adding memguardian + various optimizations
This commit is contained in:
Mzack9999 2024-03-14 23:39:43 +01:00 committed by GitHub
commit b7f76cfd4b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
13 changed files with 221 additions and 72 deletions

12
go.mod
View File

@ -30,7 +30,7 @@ require (
github.com/remeh/sizedwaitgroup v1.0.0
github.com/rs/xid v1.5.0
github.com/segmentio/ksuid v1.0.4
github.com/shirou/gopsutil/v3 v3.23.7 // indirect
github.com/shirou/gopsutil/v3 v3.24.2 // indirect
github.com/spaolacci/murmur3 v1.1.0 // indirect
github.com/spf13/cast v1.5.1
github.com/syndtr/goleveldb v1.0.0
@ -94,7 +94,7 @@ require (
github.com/projectdiscovery/tlsx v1.1.6
github.com/projectdiscovery/uncover v1.0.7
github.com/projectdiscovery/useragent v0.0.40
github.com/projectdiscovery/utils v0.0.84-0.20240312214300-d3ba70dbb9ca
github.com/projectdiscovery/utils v0.0.84-0.20240313184656-e3ec80f4dd42
github.com/projectdiscovery/wappalyzergo v0.0.112
github.com/redis/go-redis/v9 v9.1.0
github.com/sashabaranov/go-openai v1.15.3
@ -204,6 +204,7 @@ require (
github.com/projectdiscovery/stringsutil v0.0.2 // indirect
github.com/quic-go/quic-go v0.40.1 // indirect
github.com/refraction-networking/utls v1.6.1 // indirect
github.com/shirou/gopsutil v3.21.11+incompatible // indirect
github.com/shoenig/go-m1cpu v0.1.6 // indirect
github.com/sirupsen/logrus v1.9.3 // indirect
github.com/skeema/knownhosts v1.2.1 // indirect
@ -227,6 +228,7 @@ require (
github.com/yuin/goldmark-emoji v1.0.1 // indirect
github.com/zeebo/blake3 v0.2.3 // indirect
golang.org/x/arch v0.3.0 // indirect
golang.org/x/sync v0.6.0 // indirect
gopkg.in/djherbis/times.v1 v1.3.0 // indirect
mellium.im/sasl v0.3.1 // indirect
)
@ -281,8 +283,8 @@ require (
github.com/projectdiscovery/networkpolicy v0.0.8
github.com/rivo/uniseg v0.4.6 // indirect
github.com/saintfish/chardet v0.0.0-20230101081208-5e3ef4b5456d // indirect
github.com/tklauser/go-sysconf v0.3.11 // indirect
github.com/tklauser/numcpus v0.6.0 // indirect
github.com/tklauser/go-sysconf v0.3.12 // indirect
github.com/tklauser/numcpus v0.6.1 // indirect
github.com/trivago/tgo v1.0.7
github.com/ulikunitz/xz v0.5.11 // indirect
github.com/valyala/bytebufferpool v1.0.0 // indirect
@ -290,7 +292,7 @@ require (
github.com/ysmood/goob v0.4.0 // indirect
github.com/ysmood/gson v0.7.3 // indirect
github.com/ysmood/leakless v0.8.0 // indirect
github.com/yusufpapurcu/wmi v1.2.3 // indirect
github.com/yusufpapurcu/wmi v1.2.4 // indirect
github.com/zmap/rc2 v0.0.0-20190804163417-abaa70531248 // indirect
github.com/zmap/zcrypto v0.0.0-20231219022726-a1f61fb1661c // indirect
go.etcd.io/bbolt v1.3.8 // indirect

25
go.sum
View File

@ -890,8 +890,10 @@ github.com/projectdiscovery/uncover v1.0.7 h1:ut+2lTuvmftmveqF5RTjMWAgyLj8ltPQC7
github.com/projectdiscovery/uncover v1.0.7/go.mod h1:HFXgm1sRPuoN0D4oATljPIdmbo/EEh1wVuxQqo/dwFE=
github.com/projectdiscovery/useragent v0.0.40 h1:1LUhReSGPkhqsM5n40OOC9dIoNqMGs1dyGFJcOmg2Fo=
github.com/projectdiscovery/useragent v0.0.40/go.mod h1:EvK1x3s948Gtqb/XOahXcauyejCL/rSgy5d1IAvsKT4=
github.com/projectdiscovery/utils v0.0.84-0.20240312214300-d3ba70dbb9ca h1:GY9lUYDlENXPSFPJH01Bm1BfhrUF2jpnUBR+K4VPJIs=
github.com/projectdiscovery/utils v0.0.84-0.20240312214300-d3ba70dbb9ca/go.mod h1:wzMfHBq2I9oy+DEiMfUYV86g1D7eXKaQsgWnqFpmMtI=
github.com/projectdiscovery/utils v0.0.84-0.20240311212130-16ce15974a4a h1:njYY24OsTQJ80L8O+QxcYFljiVl83xp/BWz3dsIJF30=
github.com/projectdiscovery/utils v0.0.84-0.20240311212130-16ce15974a4a/go.mod h1:bvcudEteeZ5MIZeBCXEfpcgj9h3tyB9qtnmc7zQR92w=
github.com/projectdiscovery/utils v0.0.84-0.20240313184656-e3ec80f4dd42 h1:l22rSOP8i6HXu1QfAtIot8NvmJgUmBHEn6Mih7s8Gak=
github.com/projectdiscovery/utils v0.0.84-0.20240313184656-e3ec80f4dd42/go.mod h1:VsoXXTuNAAziuodKWakLyurVXaV4tNTJU4Eo8umyr3Q=
github.com/projectdiscovery/wappalyzergo v0.0.112 h1:QPpp5jmj1lqLd5mFdFKQ9VvcYhQNqyU9Mr+IB0US2zA=
github.com/projectdiscovery/wappalyzergo v0.0.112/go.mod h1:hc/o+fgM8KtdpFesjfBTmHTwsR+yBd+4kYZW/DGy/x8=
github.com/projectdiscovery/yamldoc-go v1.0.4 h1:eZoESapnMw6WAHiVgRwNqvbJEfNHEH148uthhFbG5jE=
@ -964,8 +966,10 @@ github.com/seh-msft/burpxml v1.0.1/go.mod h1:lTViCHPtGGS0scK0B4krm6Ld1kVZLWzQccw
github.com/sergi/go-diff v1.0.0/go.mod h1:0CfEIISq7TuYL3j771MWULgwwjU+GofnZX9QAmXWZgo=
github.com/sergi/go-diff v1.2.0 h1:XU+rvMAioB0UC3q1MFrIQy4Vo5/4VsRDQQXHsEya6xQ=
github.com/sergi/go-diff v1.2.0/go.mod h1:STckp+ISIX8hZLjrqAeVduY0gWCT9IjLuqbuNXdaHfM=
github.com/shirou/gopsutil/v3 v3.23.7 h1:C+fHO8hfIppoJ1WdsVm1RoI0RwXoNdfTK7yWXV0wVj4=
github.com/shirou/gopsutil/v3 v3.23.7/go.mod h1:c4gnmoRC0hQuaLqvxnx1//VXQ0Ms/X9UnJF8pddY5z4=
github.com/shirou/gopsutil v3.21.11+incompatible h1:+1+c1VGhc88SSonWP6foOcLhvnKlUeu/erjjvaPEYiI=
github.com/shirou/gopsutil v3.21.11+incompatible/go.mod h1:5b4v6he4MtMOwMlS0TUMTu2PcXUg8+E1lC7eC3UO/RA=
github.com/shirou/gopsutil/v3 v3.24.2 h1:kcR0erMbLg5/3LcInpw0X/rrPSqq4CDPyI6A6ZRC18Y=
github.com/shirou/gopsutil/v3 v3.24.2/go.mod h1:tSg/594BcA+8UdQU2XcW803GWYgdtauFFPgJCJKZlVk=
github.com/shoenig/go-m1cpu v0.1.6 h1:nxdKQNcEB6vzgA2E2bvzKIYRuNj7XNJ4S/aRSwKzFtM=
github.com/shoenig/go-m1cpu v0.1.6/go.mod h1:1JJMcUBvfNwpq05QDQVAnx3gUHr9IYF7GNg9SUEw2VQ=
github.com/shoenig/test v0.6.4 h1:kVTaSd7WLz5WZ2IaoM0RSzRsUD+m8wRR+5qvntpn4LU=
@ -1044,10 +1048,10 @@ github.com/tidwall/tinyqueue v0.1.1 h1:SpNEvEggbpyN5DIReaJ2/1ndroY8iyEGxPYxoSaym
github.com/tidwall/tinyqueue v0.1.1/go.mod h1:O/QNHwrnjqr6IHItYrzoHAKYhBkLI67Q096fQP5zMYw=
github.com/tim-ywliu/nested-logrus-formatter v1.3.2 h1:jugNJ2/CNCI79SxOJCOhwUHeN3O7/7/bj+ZRGOFlCSw=
github.com/tim-ywliu/nested-logrus-formatter v1.3.2/go.mod h1:oGPmcxZB65j9Wo7mCnQKSrKEJtVDqyjD666SGmyStXI=
github.com/tklauser/go-sysconf v0.3.11 h1:89WgdJhk5SNwJfu+GKyYveZ4IaJ7xAkecBo+KdJV0CM=
github.com/tklauser/go-sysconf v0.3.11/go.mod h1:GqXfhXY3kiPa0nAXPDIQIWzJbMCB7AmcWpGR8lSZfqI=
github.com/tklauser/numcpus v0.6.0 h1:kebhY2Qt+3U6RNK7UqpYNA+tJ23IBEGKkB7JQBfDYms=
github.com/tklauser/numcpus v0.6.0/go.mod h1:FEZLMke0lhOUG6w2JadTzp0a+Nl8PF/GFkQ5UVIcaL4=
github.com/tklauser/go-sysconf v0.3.12 h1:0QaGUFOdQaIVdPgfITYzaTegZvdCjmYO52cSFAEVmqU=
github.com/tklauser/go-sysconf v0.3.12/go.mod h1:Ho14jnntGE1fpdOqQEEaiKRpvIavV0hSfmBq8nJbHYI=
github.com/tklauser/numcpus v0.6.1 h1:ng9scYS7az0Bk4OZLvrNXNSAO2Pxr1XXRAPyjhIx+Fk=
github.com/tklauser/numcpus v0.6.1/go.mod h1:1XfjsgE2zo8GVw7POkMbHENHzVg3GzmoZ9fESEdAacY=
github.com/tmc/grpc-websocket-proxy v0.0.0-20170815181823-89b8d40f7ca8/go.mod h1:ncp9v5uamzpCO7NfCPTXjqaC+bZgJeR0sMTm6dMHP7U=
github.com/trivago/tgo v1.0.7 h1:uaWH/XIy9aWYWpjm2CU3RpcqZXmX2ysQ9/Go+d9gyrM=
github.com/trivago/tgo v1.0.7/go.mod h1:w4dpD+3tzNIIiIfkWWa85w5/B77tlvdZckQ+6PkFnhc=
@ -1116,8 +1120,8 @@ github.com/yuin/goldmark v1.5.4 h1:2uY/xC0roWy8IBEGLgB1ywIoEJFGmRrX21YQcvGZzjU=
github.com/yuin/goldmark v1.5.4/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
github.com/yuin/goldmark-emoji v1.0.1 h1:ctuWEyzGBwiucEqxzwe0SOYDXPAucOrE9NQC18Wa1os=
github.com/yuin/goldmark-emoji v1.0.1/go.mod h1:2w1E6FEWLcDQkoTE+7HU6QF1F6SLlNGjRIBbIZQFqkQ=
github.com/yusufpapurcu/wmi v1.2.3 h1:E1ctvB7uKFMOJw3fdOW32DwGE9I7t++CRUEMKvFoFiw=
github.com/yusufpapurcu/wmi v1.2.3/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0=
github.com/yusufpapurcu/wmi v1.2.4 h1:zFUKzehAFReQwLys1b/iSMl+JQGSCSjtVqQn9bBrPo0=
github.com/yusufpapurcu/wmi v1.2.4/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0=
github.com/zeebo/assert v1.1.0 h1:hU1L1vLTHsnO8x8c9KAR5GmM5QscxHg5RNU5z5qbUWY=
github.com/zeebo/assert v1.1.0/go.mod h1:Pq9JiuJQpG8JLJdtkwrJESF0Foym2/D9XMU5ciN/wJ0=
github.com/zeebo/blake3 v0.2.3 h1:TFoLXsjeXqRNFxSbk35Dk4YtszE/MQQGK10BH4ptoTg=
@ -1387,6 +1391,7 @@ golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.7.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.10.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.11.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.15.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
golang.org/x/sys v0.17.0 h1:25cE3gD+tdBA7lp7QfhuV+rJiE9YXTcS3VG1SqssI/Y=
golang.org/x/sys v0.17.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=

View File

@ -64,6 +64,22 @@ func (operators *Operators) Compile() error {
return nil
}
func (operators *Operators) HasDSL() bool {
for _, matcher := range operators.Matchers {
if len(matcher.DSL) > 0 {
return true
}
}
for _, extractor := range operators.Extractors {
if len(extractor.DSL) > 0 {
return true
}
}
return false
}
// GetMatchersCondition returns the condition for the matchers
func (operators *Operators) GetMatchersCondition() matchers.ConditionType {
return operators.matchersCondition

View File

@ -73,6 +73,10 @@ var decolorizerRegex = regexp.MustCompile(`\x1B\[[0-9;]*[a-zA-Z]`)
// InternalEvent is an internal output generation structure for nuclei.
type InternalEvent map[string]interface{}
func (ie InternalEvent) Set(k string, v interface{}) {
ie[k] = v
}
// InternalWrappedEvent is a wrapped event with operators result added to it.
type InternalWrappedEvent struct {
// Mutex is internal field which is implicitly used

View File

@ -0,0 +1,98 @@
package protocolstate
import (
"sync"
"time"
"github.com/projectdiscovery/gologger"
"github.com/projectdiscovery/utils/env"
httputil "github.com/projectdiscovery/utils/http"
"github.com/projectdiscovery/utils/memguardian"
)
var (
MaxThreadsOnLowMemory = env.GetEnvOrDefault("MEMGUARDIAN_THREADS", 0)
MaxBytesBufferAllocOnLowMemory = env.GetEnvOrDefault("MEMGUARDIAN_ALLOC", 0)
memTimer *time.Ticker
)
func StartActiveMemGuardian() {
if memguardian.DefaultMemGuardian == nil {
return
}
memTimer := time.NewTicker(memguardian.DefaultInterval)
go func() {
for range memTimer.C {
if IsLowOnMemory() {
_ = GlobalGuardBytesBufferAlloc()
} else {
GlobalRestoreBytesBufferAlloc()
}
}
}()
}
func StopActiveMemGuardian() {
if memguardian.DefaultMemGuardian == nil {
return
}
memTimer.Stop()
}
func IsLowOnMemory() bool {
if memguardian.DefaultMemGuardian != nil && memguardian.DefaultMemGuardian.Warning.Load() {
return true
}
return false
}
// GuardThreads on caller
func GuardThreadsOrDefault(current int) int {
if MaxThreadsOnLowMemory > 0 {
return MaxThreadsOnLowMemory
}
fraction := int(current / 5)
if fraction > 0 {
return fraction
}
return 1
}
var muGlobalChange sync.Mutex
// Global setting
func GlobalGuardBytesBufferAlloc() error {
if muGlobalChange.TryLock() {
return nil
}
defer muGlobalChange.Unlock()
// if current capacity was not reduced decrease it
if MaxBytesBufferAllocOnLowMemory > 0 && httputil.DefaultBytesBufferAlloc == httputil.GetPoolSize() {
gologger.Debug().Msgf("reducing bytes.buffer pool size to: %d", MaxBytesBufferAllocOnLowMemory)
delta := httputil.GetPoolSize() - int64(MaxBytesBufferAllocOnLowMemory)
return httputil.ChangePoolSize(-delta)
}
return nil
}
// Global setting
func GlobalRestoreBytesBufferAlloc() {
if muGlobalChange.TryLock() {
return
}
defer muGlobalChange.Unlock()
if httputil.DefaultBytesBufferAlloc != httputil.GetPoolSize() {
delta := httputil.DefaultBytesBufferAlloc - httputil.GetPoolSize()
gologger.Debug().Msgf("restoring bytes.buffer pool size to: %d", httputil.DefaultBytesBufferAlloc)
_ = httputil.ChangePoolSize(delta)
}
}

View File

@ -18,13 +18,16 @@ import (
)
// Dialer is a shared fastdialer instance for host DNS resolution
var Dialer *fastdialer.Dialer
var (
Dialer *fastdialer.Dialer
)
// Init creates the Dialer instance based on user configuration
func Init(options *types.Options) error {
if Dialer != nil {
return nil
}
lfaAllowed = options.AllowLocalFileAccess
opts := fastdialer.DefaultOptions
if options.DialerTimeout > 0 {
@ -142,6 +145,8 @@ func Init(options *types.Options) error {
return Dialer.Dial(ctx, "tcp", addr)
})
StartActiveMemGuardian()
return nil
}
@ -202,4 +207,5 @@ func Close() {
if Dialer != nil {
Dialer.Close()
}
StopActiveMemGuardian()
}

View File

@ -1,8 +0,0 @@
package tostring
import "unsafe"
// UnsafeToString converts byte slice to string with zero allocations
func UnsafeToString(bs []byte) string {
return *(*string)(unsafe.Pointer(&bs))
}

View File

@ -14,6 +14,7 @@ import (
"github.com/projectdiscovery/nuclei/v3/pkg/protocols"
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/expressions"
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/generators"
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolstate"
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/http/httpclientpool"
httputil "github.com/projectdiscovery/nuclei/v3/pkg/protocols/utils/http"
"github.com/projectdiscovery/rawhttp"
@ -411,6 +412,10 @@ func (request *Request) Compile(options *protocols.ExecutorOptions) error {
}
}
if len(request.Payloads) > 0 {
// specifically for http requests high concurrency and and threads will lead to memory exausthion, hence reduce the maximum parallelism
if protocolstate.IsLowOnMemory() {
request.Threads = protocolstate.GuardThreadsOrDefault(request.Threads)
}
// if we have payloads, adjust threads if none specified
request.Threads = options.GetThreadsForNPayloadRequests(request.Requests(), request.Threads)
}

View File

@ -112,21 +112,21 @@ func (request *Request) responseToDSLMap(resp *http.Response, host, matched, raw
data[k] = v
}
for _, cookie := range resp.Cookies() {
data[strings.ToLower(cookie.Name)] = cookie.Value
request.setHashOrDefault(data, strings.ToLower(cookie.Name), cookie.Value)
}
for k, v := range resp.Header {
k = strings.ToLower(strings.ReplaceAll(strings.TrimSpace(k), "-", "_"))
data[k] = strings.Join(v, " ")
request.setHashOrDefault(data, k, strings.Join(v, " "))
}
data["host"] = host
data["type"] = request.Type().String()
data["matched"] = matched
data["request"] = rawReq
data["response"] = rawResp
request.setHashOrDefault(data, "request", rawReq)
request.setHashOrDefault(data, "response", rawResp)
data["status_code"] = resp.StatusCode
data["body"] = body
data["all_headers"] = headers
data["header"] = headers
request.setHashOrDefault(data, "body", body)
request.setHashOrDefault(data, "all_headers", headers)
request.setHashOrDefault(data, "header", headers)
data["duration"] = duration.Seconds()
data["template-id"] = request.options.TemplateID
data["template-info"] = request.options.TemplateInfo
@ -140,6 +140,15 @@ func (request *Request) responseToDSLMap(resp *http.Response, host, matched, raw
return data
}
// TODO: disabling hdd storage while testing backpressure mechanism
func (request *Request) setHashOrDefault(data output.InternalEvent, k string, v string) {
// if hash, err := request.options.Storage.SetString(v); err == nil {
// data[k] = hash
// } else {
data[k] = v
//}
}
// MakeResultEvent creates a result event from internal wrapped event
func (request *Request) MakeResultEvent(wrapped *output.InternalWrappedEvent) []*output.ResultEvent {
return protocols.MakeDefaultResultEvent(request, wrapped)

View File

@ -28,7 +28,7 @@ import (
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/helpers/eventcreator"
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/helpers/responsehighlighter"
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/interactsh"
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/tostring"
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/protocolstate"
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/http/httpclientpool"
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/http/httputils"
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/http/signer"
@ -168,6 +168,10 @@ func (request *Request) executeParallelHTTP(input *contextargs.Context, dynamicV
// Workers that keeps enqueuing new requests
maxWorkers := request.Threads
if protocolstate.IsLowOnMemory() {
maxWorkers = protocolstate.GuardThreadsOrDefault(request.Threads)
}
// Stop-at-first-match logic while executing requests
// parallely using threads
shouldStop := (request.options.Options.StopAtFirstMatch || request.StopAtFirstMatch || request.options.StopAtFirstMatch)
@ -684,7 +688,7 @@ func (request *Request) executeRequest(input *contextargs.Context, generatedRequ
// In case of interactsh markers and request times out, still send
// a callback event so in case we receive an interaction, correlation is possible.
// Also, to log failed use-cases.
outputEvent := request.responseToDSLMap(&http.Response{}, input.MetaInput.Input, formedURL, tostring.UnsafeToString(dumpedRequest), "", "", "", 0, generatedRequest.meta)
outputEvent := request.responseToDSLMap(&http.Response{}, input.MetaInput.Input, formedURL, convUtil.String(dumpedRequest), "", "", "", 0, generatedRequest.meta)
if i := strings.LastIndex(hostname, ":"); i != -1 {
hostname = hostname[:i]
}
@ -695,8 +699,8 @@ func (request *Request) executeRequest(input *contextargs.Context, generatedRequ
outputEvent["ip"] = httpclientpool.Dialer.GetDialedIP(hostname)
}
event := &output.InternalWrappedEvent{InternalEvent: outputEvent}
if request.CompiledOperators != nil {
event := &output.InternalWrappedEvent{}
if request.CompiledOperators != nil && request.CompiledOperators.HasDSL() {
event.InternalEvent = outputEvent
}
callback(event)

View File

@ -14,9 +14,9 @@ import (
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/contextargs"
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/generators"
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/helpers/eventcreator"
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/tostring"
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/utils"
templateTypes "github.com/projectdiscovery/nuclei/v3/pkg/templates/types"
"github.com/projectdiscovery/utils/conversion"
)
var _ protocols.Request = &Request{}
@ -60,7 +60,7 @@ func (request *Request) ExecuteWithResults(input *contextargs.Context, metadata
gologger.Error().Msgf("Could not read file path %s: %s\n", data, err)
return
}
dataStr := tostring.UnsafeToString(buffer)
dataStr := conversion.String(buffer)
resp, err := readResponseFromString(dataStr)
if err != nil {
@ -86,7 +86,7 @@ func (request *Request) ExecuteWithResults(input *contextargs.Context, metadata
return
}
outputEvent := request.responseToDSLMap(resp, data, data, data, tostring.UnsafeToString(dumpedResponse), tostring.UnsafeToString(body), utils.HeadersToString(resp.Header), 0, nil)
outputEvent := request.responseToDSLMap(resp, data, data, data, conversion.String(dumpedResponse), conversion.String(body), utils.HeadersToString(resp.Header), 0, nil)
// add response fields to template context and merge templatectx variables to output event
request.options.AddTemplateVars(input.MetaInput, request.Type(), request.GetID(), outputEvent)
if request.options.HasTemplateCtx(input.MetaInput) {

View File

@ -7,13 +7,13 @@ package dedupe
import (
"crypto/sha1"
"os"
"unsafe"
"github.com/syndtr/goleveldb/leveldb"
"github.com/syndtr/goleveldb/leveldb/errors"
"github.com/projectdiscovery/nuclei/v3/pkg/output"
"github.com/projectdiscovery/nuclei/v3/pkg/types"
"github.com/projectdiscovery/utils/conversion"
)
// Storage is a duplicate detecting storage for nuclei scan events.
@ -75,29 +75,29 @@ func (s *Storage) Close() {
func (s *Storage) Index(result *output.ResultEvent) (bool, error) {
hasher := sha1.New()
if result.TemplateID != "" {
_, _ = hasher.Write(unsafeToBytes(result.TemplateID))
_, _ = hasher.Write(conversion.Bytes(result.TemplateID))
}
if result.MatcherName != "" {
_, _ = hasher.Write(unsafeToBytes(result.MatcherName))
_, _ = hasher.Write(conversion.Bytes(result.MatcherName))
}
if result.ExtractorName != "" {
_, _ = hasher.Write(unsafeToBytes(result.ExtractorName))
_, _ = hasher.Write(conversion.Bytes(result.ExtractorName))
}
if result.Type != "" {
_, _ = hasher.Write(unsafeToBytes(result.Type))
_, _ = hasher.Write(conversion.Bytes(result.Type))
}
if result.Host != "" {
_, _ = hasher.Write(unsafeToBytes(result.Host))
_, _ = hasher.Write(conversion.Bytes(result.Host))
}
if result.Matched != "" {
_, _ = hasher.Write(unsafeToBytes(result.Matched))
_, _ = hasher.Write(conversion.Bytes(result.Matched))
}
for _, v := range result.ExtractedResults {
_, _ = hasher.Write(unsafeToBytes(v))
_, _ = hasher.Write(conversion.Bytes(v))
}
for k, v := range result.Metadata {
_, _ = hasher.Write(unsafeToBytes(k))
_, _ = hasher.Write(unsafeToBytes(types.ToString(v)))
_, _ = hasher.Write(conversion.Bytes(k))
_, _ = hasher.Write(conversion.Bytes(types.ToString(v)))
}
hash := hasher.Sum(nil)
@ -112,12 +112,3 @@ func (s *Storage) Index(result *output.ResultEvent) (bool, error) {
}
return false, err
}
// unsafeToBytes converts a string to byte slice and does it with
// zero allocations.
//
// Reference - https://stackoverflow.com/questions/59209493/how-to-use-unsafe-get-a-byte-slice-from-a-string-without-memory-copy
func unsafeToBytes(data string) []byte {
var buf = (*[]byte)(unsafe.Pointer(&data))
return *buf
}

View File

@ -10,19 +10,32 @@ import (
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/contextargs"
)
type ScanContextOption func(*ScanContext)
func WithEvents() ScanContextOption {
return func(sc *ScanContext) {
sc.withEvents = true
}
}
type ScanContext struct {
context.Context
// exported / configurable fields
Input *contextargs.Context
// callbacks or hooks
OnError func(error)
OnResult func(e *output.InternalWrappedEvent)
OnError func(error)
OnResult func(e *output.InternalWrappedEvent)
OnWarning func(string)
// unexported state fields
errors []error
warnings []string
events []*output.InternalWrappedEvent
results []*output.ResultEvent
// what to log
withEvents bool
// might not be required but better to sync
m sync.Mutex
@ -37,7 +50,8 @@ func NewScanContext(input *contextargs.Context) *ScanContext {
func (s *ScanContext) GenerateResult() []*output.ResultEvent {
s.m.Lock()
defer s.m.Unlock()
return aggregateResults(s.events)
return s.results
}
// LogEvent logs events to all events and triggeres any callbacks
@ -48,10 +62,16 @@ func (s *ScanContext) LogEvent(e *output.InternalWrappedEvent) {
// do not log nil events
return
}
if s.OnResult != nil {
s.OnResult(e)
}
s.events = append(s.events, e)
if s.withEvents {
s.events = append(s.events, e)
}
s.results = append(s.results, e.Results...)
}
// LogError logs error to all events and triggeres any callbacks
@ -68,10 +88,11 @@ func (s *ScanContext) LogError(err error) {
s.errors = append(s.errors, err)
errorMessage := joinErrors(s.errors)
results := aggregateResults(s.events)
for _, result := range results {
for _, result := range s.results {
result.Error = errorMessage
}
for _, e := range s.events {
e.InternalEvent["error"] = errorMessage
}
@ -82,6 +103,11 @@ func (s *ScanContext) LogWarning(format string, args ...any) {
s.m.Lock()
defer s.m.Unlock()
val := fmt.Sprintf(format, args...)
if s.OnWarning != nil {
s.OnWarning(val)
}
s.warnings = append(s.warnings, val)
for _, e := range s.events {
@ -91,15 +117,6 @@ func (s *ScanContext) LogWarning(format string, args ...any) {
}
}
// aggregateResults aggregates results from multiple events
func aggregateResults(events []*output.InternalWrappedEvent) []*output.ResultEvent {
var results []*output.ResultEvent
for _, e := range events {
results = append(results, e.Results...)
}
return results
}
// joinErrors joins multiple errors and returns a single error string
func joinErrors(errors []error) string {
var errorMessages []string