mirror of
https://github.com/projectdiscovery/nuclei.git
synced 2025-12-17 21:35:26 +00:00
85 lines
2.6 KiB
Go
85 lines
2.6 KiB
Go
package core
|
|
|
|
import (
|
|
"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()
|
|
}
|
|
|
|
func (w *WorkPool) Refresh() {
|
|
if w.Default.Size != w.config.TypeConcurrency {
|
|
w.Default.Resize(w.config.TypeConcurrency)
|
|
}
|
|
if w.Headless.Size != w.config.HeadlessTypeConcurrency {
|
|
w.Headless.Resize(w.config.HeadlessTypeConcurrency)
|
|
}
|
|
}
|