mirror of
https://github.com/projectdiscovery/nuclei.git
synced 2025-12-17 19:45:28 +00:00
3e54ca54b0
* feat: fixed leak * add go leak unit test in sdk * added goleak unit tests * bugfix: add random user agents to fuzzing requests * misc * misc * fix lint + use utils pr + misc * fix ratelimit memleak in sdk * close protocolstate shared resources in nuclei sdk/lib * add missing close references * ignore read/write loop of intransit connections * close unnecessary idle conns * add ignore method * using fixed utils * dep update --------- Co-authored-by: Ice3man <nizamulrana@gmail.com> Co-authored-by: mzack <marco.rivoli.nvh@gmail.com> Co-authored-by: sandeep <8293321+ehsandeep@users.noreply.github.com>
92 lines
2.9 KiB
Go
92 lines
2.9 KiB
Go
package core
|
|
|
|
import (
|
|
"context"
|
|
|
|
"github.com/projectdiscovery/gologger"
|
|
"github.com/projectdiscovery/nuclei/v3/pkg/templates/types"
|
|
syncutil "github.com/projectdiscovery/utils/sync"
|
|
)
|
|
|
|
// WorkPool implements an execution pool for executing different
|
|
// types of task with different concurrency requirements.
|
|
//
|
|
// It also allows Configuration of such requirements. This is used
|
|
// for per-module like separate headless concurrency etc.
|
|
type WorkPool struct {
|
|
Headless *syncutil.AdaptiveWaitGroup
|
|
Default *syncutil.AdaptiveWaitGroup
|
|
config WorkPoolConfig
|
|
}
|
|
|
|
// WorkPoolConfig is the configuration for work pool
|
|
type WorkPoolConfig struct {
|
|
// InputConcurrency is the concurrency for inputs values.
|
|
InputConcurrency int
|
|
// TypeConcurrency is the concurrency for the request type templates.
|
|
TypeConcurrency int
|
|
// HeadlessInputConcurrency is the concurrency for headless inputs values.
|
|
HeadlessInputConcurrency int
|
|
// TypeConcurrency is the concurrency for the headless request type templates.
|
|
HeadlessTypeConcurrency int
|
|
}
|
|
|
|
// NewWorkPool returns a new WorkPool instance
|
|
func NewWorkPool(config WorkPoolConfig) *WorkPool {
|
|
headlessWg, _ := syncutil.New(syncutil.WithSize(config.HeadlessTypeConcurrency))
|
|
defaultWg, _ := syncutil.New(syncutil.WithSize(config.TypeConcurrency))
|
|
|
|
return &WorkPool{
|
|
config: config,
|
|
Headless: headlessWg,
|
|
Default: defaultWg,
|
|
}
|
|
}
|
|
|
|
// Wait waits for all the work pool wait groups to finish
|
|
func (w *WorkPool) Wait() {
|
|
w.Default.Wait()
|
|
w.Headless.Wait()
|
|
}
|
|
|
|
// InputPool returns a work pool for an input type
|
|
func (w *WorkPool) InputPool(templateType types.ProtocolType) *syncutil.AdaptiveWaitGroup {
|
|
var count int
|
|
if templateType == types.HeadlessProtocol {
|
|
count = w.config.HeadlessInputConcurrency
|
|
} else {
|
|
count = w.config.InputConcurrency
|
|
}
|
|
swg, _ := syncutil.New(syncutil.WithSize(count))
|
|
return swg
|
|
}
|
|
|
|
func (w *WorkPool) RefreshWithConfig(config WorkPoolConfig) {
|
|
if w.config.TypeConcurrency != config.TypeConcurrency {
|
|
w.config.TypeConcurrency = config.TypeConcurrency
|
|
}
|
|
if w.config.HeadlessTypeConcurrency != config.HeadlessTypeConcurrency {
|
|
w.config.HeadlessTypeConcurrency = config.HeadlessTypeConcurrency
|
|
}
|
|
if w.config.InputConcurrency != config.InputConcurrency {
|
|
w.config.InputConcurrency = config.InputConcurrency
|
|
}
|
|
if w.config.HeadlessInputConcurrency != config.HeadlessInputConcurrency {
|
|
w.config.HeadlessInputConcurrency = config.HeadlessInputConcurrency
|
|
}
|
|
w.Refresh(context.Background())
|
|
}
|
|
|
|
func (w *WorkPool) Refresh(ctx context.Context) {
|
|
if w.Default.Size != w.config.TypeConcurrency {
|
|
if err := w.Default.Resize(ctx, w.config.TypeConcurrency); err != nil {
|
|
gologger.Warning().Msgf("Could not resize workpool: %s\n", err)
|
|
}
|
|
}
|
|
if w.Headless.Size != w.config.HeadlessTypeConcurrency {
|
|
if err := w.Headless.Resize(ctx, w.config.HeadlessTypeConcurrency); err != nil {
|
|
gologger.Warning().Msgf("Could not resize workpool: %s\n", err)
|
|
}
|
|
}
|
|
}
|