adding shared fastdialer + missing cleanup

This commit is contained in:
mzack 2021-04-18 11:57:43 +02:00
parent 27f15a9c93
commit 54540f3a22
6 changed files with 47 additions and 35 deletions

View File

@ -18,6 +18,7 @@ import (
"github.com/projectdiscovery/nuclei/v2/pkg/projectfile"
"github.com/projectdiscovery/nuclei/v2/pkg/protocols"
"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/clusterer"
"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/protocolinit"
"github.com/projectdiscovery/nuclei/v2/pkg/protocols/headless/engine"
"github.com/projectdiscovery/nuclei/v2/pkg/reporting"
"github.com/projectdiscovery/nuclei/v2/pkg/reporting/exporters/disk"
@ -218,6 +219,7 @@ func (r *Runner) Close() {
if r.projectFile != nil {
r.projectFile.Close()
}
protocolinit.Close()
}
// RunEnumeration sets up the input layer for giving input nuclei.

View File

@ -2,6 +2,7 @@ package protocolinit
import (
"github.com/corpix/uarand"
"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/protocolstate"
"github.com/projectdiscovery/nuclei/v2/pkg/protocols/dns/dnsclientpool"
"github.com/projectdiscovery/nuclei/v2/pkg/protocols/http/httpclientpool"
"github.com/projectdiscovery/nuclei/v2/pkg/protocols/network/networkclientpool"
@ -12,6 +13,9 @@ import (
func Init(options *types.Options) error {
uarand.Default = uarand.NewWithCustomList(userAgents)
if err := protocolstate.Init(options); err != nil {
return err
}
if err := dnsclientpool.Init(options); err != nil {
return err
}
@ -58,3 +62,7 @@ var userAgents = []string{
"Mozilla/5.0 (Windows NT 5.1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/34.0.1866.237 Safari/537.36",
"Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/34.0.1847.137 Safari/4E423F",
}
func Close() {
protocolstate.Dialer.Close()
}

View File

@ -0,0 +1,31 @@
package protocolstate
import (
"github.com/pkg/errors"
"github.com/projectdiscovery/fastdialer/fastdialer"
"github.com/projectdiscovery/nuclei/v2/pkg/types"
)
var Dialer *fastdialer.Dialer
func Init(options *types.Options) error {
opts := fastdialer.DefaultOptions
if options.SystemResolvers {
opts.EnableFallback = true
}
if options.ResolversFile != "" {
opts.BaseResolvers = options.InternalResolversList
}
dialer, err := fastdialer.NewDialer(opts)
if err != nil {
errors.Wrap(err, "could not create dialer")
}
Dialer = dialer
return nil
}
func Close() {
if Dialer != nil {
Dialer.Close()
}
}

View File

@ -5,25 +5,13 @@ import (
"net/http"
"time"
"github.com/pkg/errors"
"github.com/projectdiscovery/fastdialer/fastdialer"
"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/protocolstate"
"github.com/projectdiscovery/nuclei/v2/pkg/types"
)
// newhttpClient creates a new http client for headless communication with a timeout
func newhttpClient(options *types.Options) (*http.Client, error) {
opts := fastdialer.DefaultOptions
if options.SystemResolvers {
opts.EnableFallback = true
}
if options.ResolversFile != "" {
opts.BaseResolvers = options.InternalResolversList
}
dialer, err := fastdialer.NewDialer(opts)
if err != nil {
return nil, errors.Wrap(err, "could not create dialer")
}
dialer := protocolstate.Dialer
transport := &http.Transport{
DialContext: dialer.Dial,
MaxIdleConns: 500,

View File

@ -15,6 +15,7 @@ import (
"github.com/pkg/errors"
"github.com/projectdiscovery/fastdialer/fastdialer"
"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/protocolstate"
"github.com/projectdiscovery/nuclei/v2/pkg/types"
"github.com/projectdiscovery/rawhttp"
"github.com/projectdiscovery/retryablehttp-go"
@ -97,14 +98,7 @@ func wrappedGet(options *types.Options, configuration *Configuration) (*retryabl
var err error
if Dialer == nil {
opts := fastdialer.DefaultOptions
if options.SystemResolvers {
opts.EnableFallback = true
}
if options.ResolversFile != "" {
opts.BaseResolvers = options.InternalResolversList
}
Dialer, err = fastdialer.NewDialer(opts)
Dialer = protocolstate.Dialer
}
if err != nil {
return nil, errors.Wrap(err, "could not create dialer")

View File

@ -1,8 +1,8 @@
package networkclientpool
import (
"github.com/pkg/errors"
"github.com/projectdiscovery/fastdialer/fastdialer"
"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/protocolstate"
"github.com/projectdiscovery/nuclei/v2/pkg/types"
)
@ -16,18 +16,7 @@ func Init(options *types.Options) error {
if normalClient != nil {
return nil
}
opts := fastdialer.DefaultOptions
if options.SystemResolvers {
opts.EnableFallback = true
}
if options.ResolversFile != "" {
opts.BaseResolvers = options.InternalResolversList
}
dialer, err := fastdialer.NewDialer(opts)
if err != nil {
return errors.Wrap(err, "could not create dialer")
}
normalClient = dialer
normalClient = protocolstate.Dialer
return nil
}