nuclei/v2/pkg/core/workpool.go

66 lines
1.8 KiB
Go

package core
import (
"github.com/remeh/sizedwaitgroup"
"github.com/projectdiscovery/nuclei/v2/pkg/templates/types"
)
// WorkPool implements an execution pool for executing different
// types of task with different concurreny requirements.
//
// It also allows Configuration of such requirements. This is used
// for per-module like separate headless concurrency etc.
type WorkPool struct {
Headless *sizedwaitgroup.SizedWaitGroup
Default *sizedwaitgroup.SizedWaitGroup
config WorkPoolConfig
}
// WorkPoolConfig is the configuration for workpool
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 := sizedwaitgroup.New(config.HeadlessTypeConcurrency)
defaultWg := sizedwaitgroup.New(config.TypeConcurrency)
return &WorkPool{
config: config,
Headless: &headlessWg,
Default: &defaultWg,
}
}
// Wait waits for all the workpool waitgroups to finish
func (w *WorkPool) Wait() {
w.Default.Wait()
w.Headless.Wait()
}
// InputWorkPool is a workpool per-input
type InputWorkPool struct {
Waitgroup *sizedwaitgroup.SizedWaitGroup
}
// InputPool returns a workpool for an input type
func (w *WorkPool) InputPool(templateType types.ProtocolType) *InputWorkPool {
var count int
if templateType == types.HeadlessProtocol {
count = w.config.HeadlessInputConcurrency
} else {
count = w.config.InputConcurrency
}
swg := sizedwaitgroup.New(count)
return &InputWorkPool{Waitgroup: &swg}
}