github.com/projectdiscovery/nuclei/v2@v2.9.15/pkg/protocols/dns/dnsclientpool/clientpool.go (about)

     1  package dnsclientpool
     2  
     3  import (
     4  	"strconv"
     5  	"strings"
     6  	"sync"
     7  
     8  	"github.com/pkg/errors"
     9  	"github.com/projectdiscovery/nuclei/v2/pkg/types"
    10  	"github.com/projectdiscovery/retryabledns"
    11  )
    12  
    13  var (
    14  	poolMutex    *sync.RWMutex
    15  	normalClient *retryabledns.Client
    16  	clientPool   map[string]*retryabledns.Client
    17  )
    18  
    19  // defaultResolvers contains the list of resolvers known to be trusted.
    20  var defaultResolvers = []string{
    21  	"1.1.1.1:53", // Cloudflare
    22  	"1.0.0.1:53", // Cloudflare
    23  	"8.8.8.8:53", // Google
    24  	"8.8.4.4:53", // Google
    25  }
    26  
    27  // Init initializes the client pool implementation
    28  func Init(options *types.Options) error {
    29  	// Don't create clients if already created in the past.
    30  	if normalClient != nil {
    31  		return nil
    32  	}
    33  	poolMutex = &sync.RWMutex{}
    34  	clientPool = make(map[string]*retryabledns.Client)
    35  
    36  	resolvers := defaultResolvers
    37  	if options.ResolversFile != "" {
    38  		resolvers = options.InternalResolversList
    39  	}
    40  	var err error
    41  	normalClient, err = retryabledns.New(resolvers, 1)
    42  	if err != nil {
    43  		return errors.Wrap(err, "could not create dns client")
    44  	}
    45  	return nil
    46  }
    47  
    48  // Configuration contains the custom configuration options for a client
    49  type Configuration struct {
    50  	// Retries contains the retries for the dns client
    51  	Retries int
    52  	// Resolvers contains the specific per request resolvers
    53  	Resolvers []string
    54  }
    55  
    56  // Hash returns the hash of the configuration to allow client pooling
    57  func (c *Configuration) Hash() string {
    58  	builder := &strings.Builder{}
    59  	builder.WriteString("r")
    60  	builder.WriteString(strconv.Itoa(c.Retries))
    61  	builder.WriteString("l")
    62  	builder.WriteString(strings.Join(c.Resolvers, ""))
    63  	hash := builder.String()
    64  	return hash
    65  }
    66  
    67  // Get creates or gets a client for the protocol based on custom configuration
    68  func Get(options *types.Options, configuration *Configuration) (*retryabledns.Client, error) {
    69  	if !(configuration.Retries > 1) && len(configuration.Resolvers) == 0 {
    70  		return normalClient, nil
    71  	}
    72  	hash := configuration.Hash()
    73  	poolMutex.RLock()
    74  	if client, ok := clientPool[hash]; ok {
    75  		poolMutex.RUnlock()
    76  		return client, nil
    77  	}
    78  	poolMutex.RUnlock()
    79  
    80  	resolvers := defaultResolvers
    81  	if options.ResolversFile != "" {
    82  		resolvers = options.InternalResolversList
    83  	} else if len(configuration.Resolvers) > 0 {
    84  		resolvers = configuration.Resolvers
    85  	}
    86  	client, err := retryabledns.New(resolvers, configuration.Retries)
    87  	if err != nil {
    88  		return nil, errors.Wrap(err, "could not create dns client")
    89  	}
    90  
    91  	poolMutex.Lock()
    92  	clientPool[hash] = client
    93  	poolMutex.Unlock()
    94  	return client, nil
    95  }