Added socket count option to use with massdns

internal/app/cmd/resolve.go

@@ -51,6 +51,7 @@ The <file> argument can be omitted if the domains to resolve are read from stdin
 	resolveFlags.BoolVar(&resolveOptions.SkipSanitize, "skip-sanitize", resolveOptions.SkipSanitize, "do not sanitize the list of domains to test")
 	resolveFlags.BoolVar(&resolveOptions.SkipWildcard, "skip-wildcard-filter", resolveOptions.SkipWildcard, "do not perform wildcard detection and filtering")
 	resolveFlags.BoolVar(&resolveOptions.SkipValidation, "skip-validation", resolveOptions.SkipValidation, "do not validate results with trusted resolvers")
+	resolveFlags.IntVar(&resolveOptions.SocketCount, "socket-count", resolveOptions.SocketCount, "Number of sockets that massdns will use per process")
 
 	must(cobra.MarkFlagFilename(resolveFlags, "bin"))
 	must(cobra.MarkFlagFilename(resolveFlags, "resolvers"))

internal/app/ctx/options.go

@@ -76,6 +76,8 @@ type ResolveOptions struct {
 	Domain     string
 	Wordlist   string
 	DomainFile string
+
+	SocketCount int
 }
 
 // DefaultResolveOptions creates a new ResolveOptions struct with default values.
@@ -121,6 +123,8 @@ func DefaultResolveOptions() *ResolveOptions {
 		Domain:     "",
 		Wordlist:   "",
 		DomainFile: "",
+
+		SocketCount: 1,
 	}
 }
 

internal/usecase/resolve/massresolver.go

@@ -21,7 +21,7 @@ func NewDefaultMassResolver(binPath string) *DefaultMassResolver {
 }
 
 // Resolve calls massdns to resolve the domains contained in the input file.
-func (m *DefaultMassResolver) Resolve(r io.Reader, output string, total int, resolversFilename string, qps int) error {
+func (m *DefaultMassResolver) Resolve(r io.Reader, output string, total int, resolversFilename string, qps int, socketcount int) error {
 	var template string
 
 	if total == 0 {
@@ -38,7 +38,7 @@ func (m *DefaultMassResolver) Resolve(r io.Reader, output string, total int, res
 	)
 
 	bar.Start()
-	err := m.massdns.Resolve(r, output, resolversFilename, qps)
+	err := m.massdns.Resolve(r, output, resolversFilename, qps, socketcount)
 	bar.Stop()
 
 	return err

internal/usecase/resolve/massresolver_test.go

@@ -15,13 +15,13 @@ func TestMassResolverNew(t *testing.T) {
 func TestMassResolverResolve_OK(t *testing.T) {
 	r := NewDefaultMassResolver("")
 
-	err := r.Resolve(strings.NewReader("example.com"), "", 0, "", 10)
+	err := r.Resolve(strings.NewReader("example.com"), "", 0, "", 10, 1)
 	assert.EqualError(t, err, "exec: no command", "should not call massdns because of invalid path")
 }
 
 func TestMassResolverResolve_WithTotal(t *testing.T) {
 	r := NewDefaultMassResolver("")
 
-	err := r.Resolve(strings.NewReader("example.com"), "", 100, "", 10)
+	err := r.Resolve(strings.NewReader("example.com"), "", 100, "", 10, 1)
 	assert.EqualError(t, err, "exec: no command")
 }

internal/usecase/resolve/resolve.go

@@ -45,7 +45,7 @@ type WorkfileCreator interface {
 // MassResolver resolves the domains contained in the input file using the resolvers present in the resolvers file
 // and saves the results. Queries per second can be limited by setting the qps argument (0 for unlimited).
 type MassResolver interface {
-	Resolve(reader io.Reader, output string, total int, resolversFilename string, qps int) error
+	Resolve(reader io.Reader, output string, total int, resolversFilename string, qps int, socketcount int) error
 }
 
 // ResultSaver saves the results as direction by the options specified.
@@ -250,11 +250,13 @@ func (s *Service) resolvePublic(reader *DomainReader) error {
 	resolvers := s.workfiles.PublicResolvers
 	ratelimit := s.Options.RateLimit
 	resolverString := "public"
+	socketcount := s.Options.SocketCount
 
 	if s.Options.TrustedOnly {
 		resolvers = s.workfiles.TrustedResolvers
 		ratelimit = s.Options.RateLimitTrusted
 		resolverString = "trusted"
+		socketcount = 1
 	}
 
 	console.Printf("%sResolving domains with %s resolvers%s\n", console.ColorBrightWhite, resolverString, console.ColorReset)
@@ -265,6 +267,7 @@ func (s *Service) resolvePublic(reader *DomainReader) error {
 		s.domainCount,
 		resolvers,
 		ratelimit,
+		socketcount,
 	)
 
 	if err != nil {
@@ -382,6 +385,7 @@ func (s *Service) resolveTrusted() error {
 		s.domainCount,
 		s.workfiles.TrustedResolvers,
 		s.Options.RateLimitTrusted,
+		1,
 	)
 
 	if err != nil {

pkg/massdns/resolver.go

@@ -7,7 +7,7 @@ import (
 
 // Runner is an interface that runs the commands required to execute the massdns binary.
 type Runner interface {
-	Run(reader io.Reader, output string, resolversFile string, qps int) error
+	Run(reader io.Reader, output string, resolversFile string, qps int, socketcount int) error
 }
 
 // Resolver uses massdns to resolve a batch of domain names.
@@ -32,10 +32,10 @@ func NewResolver(binPath string) *Resolver {
 
 // Resolve reads domain names from the reader and saves the answers to a file.
 // It uses the resolvers and queries per second limit specified.
-func (r *Resolver) Resolve(reader io.Reader, output string, resolversFile string, qps int) error {
+func (r *Resolver) Resolve(reader io.Reader, output string, resolversFile string, qps int, socketcount int) error {
 	r.reader = NewLineReader(reader, qps)
 
-	if err := r.runner.Run(r.reader, output, resolversFile, qps); err != nil {
+	if err := r.runner.Run(r.reader, output, resolversFile, qps, socketcount); err != nil {
 		return err
 	}
 

pkg/massdns/resolver_test.go

@@ -41,7 +41,7 @@ func TestResolve(t *testing.T) {
 			resolver := NewResolver("massdns")
 			resolver.runner = stubRunner{returns: test.haveRunnerError}
 
-			gotErr := resolver.Resolve(strings.NewReader(""), "", "", 10)
+			gotErr := resolver.Resolve(strings.NewReader(""), "", "", 10, 1)
 
 			assert.Equal(t, test.wantErr, gotErr != nil, gotErr)
 		})
@@ -55,7 +55,7 @@ func TestCurrent(t *testing.T) {
 	gotCurrent := resolver.Current()
 	assert.Equal(t, 0, gotCurrent)
 
-	resolver.Resolve(strings.NewReader("example.com\n"), "", "", 0)
+	resolver.Resolve(strings.NewReader("example.com\n"), "", "", 0, 1)
 	gotCurrent = resolver.Current()
 	assert.Equal(t, 1, gotCurrent)
 }

pkg/massdns/runner.go

@@ -19,9 +19,9 @@ func newDefaultRunner(binPath string) *defaultRunner {
 }
 
 // Run executes massdns on the specified domains files and saves the results to the output file.
-func (runner *defaultRunner) Run(r io.Reader, output string, resolvers string, qps int) error {
+func (runner *defaultRunner) Run(r io.Reader, output string, resolvers string, qps int, socketcount int) error {
 	// Create massdns program arguments
-	massdnsArgs := runner.createMassdnsArgs(output, resolvers, qps)
+	massdnsArgs := runner.createMassdnsArgs(output, resolvers, qps, socketcount)
 
 	// Create a new exec.Cmd and set Stdin and Stdout to our custom handlers to avoid file operations
 	massdns := runner.execCommand(runner.binPath, massdnsArgs...)
@@ -36,9 +36,9 @@ func (runner *defaultRunner) Run(r io.Reader, output string, resolvers string, q
 }
 
 // createMassdnsArgs creates the command line arguments for massdns.
-func (runner *defaultRunner) createMassdnsArgs(output string, resolvers string, qps int) []string {
+func (runner *defaultRunner) createMassdnsArgs(output string, resolvers string, qps int, socketcount int) []string {
 	// Default command line
-	args := []string{"-q", "-r", resolvers, "-o", "Snl", "-t", "A", "--root", "--retry", "REFUSED", "--retry", "SERVFAIL", "-w", output}
+	args := []string{"-q", "-r", resolvers, "-o", "Snl", "-t", "A", "--root", "--retry", "REFUSED", "--retry", "SERVFAIL", "--socket-count", strconv.Itoa(socketcount), "-w", output}
 
 	// Set the massdns hashmap size manually to prevent it from accumulating DNS query on start
 	if qps > 0 {

pkg/massdns/runner_test.go

@@ -67,12 +67,12 @@ func TestDefaultRunnerRun(t *testing.T) {
 
 func TestCreateMassdnsArgs_DefaultRateLimit(t *testing.T) {
 	runner := defaultRunner{}
-	gotArgs := runner.createMassdnsArgs("output.txt", "resolvers.txt", 0)
-	assert.ElementsMatch(t, []string{"-q", "-r", "resolvers.txt", "-o", "Snl", "-t", "A", "--root", "--retry", "REFUSED", "--retry", "SERVFAIL", "-w", "output.txt"}, gotArgs)
+	gotArgs := runner.createMassdnsArgs("output.txt", "resolvers.txt", 0, 1)
+	assert.ElementsMatch(t, []string{"-q", "-r", "resolvers.txt", "-o", "Snl", "-t", "A", "--root", "--retry", "REFUSED", "--retry", "SERVFAIL", "--socket-count", "1", "-w", "output.txt"}, gotArgs)
 }
 
 func TestCreateMassdnsArgs_CustomRateLimit(t *testing.T) {
 	runner := defaultRunner{}
-	gotArgs := runner.createMassdnsArgs("output.txt", "resolvers.txt", 100)
-	assert.ElementsMatch(t, []string{"-q", "-r", "resolvers.txt", "-o", "Snl", "-t", "A", "--root", "--retry", "REFUSED", "--retry", "SERVFAIL", "-w", "output.txt", "-s", "100"}, gotArgs)
+	gotArgs := runner.createMassdnsArgs("output.txt", "resolvers.txt", 100, 1)
+	assert.ElementsMatch(t, []string{"-q", "-r", "resolvers.txt", "-o", "Snl", "-t", "A", "--root", "--retry", "REFUSED", "--retry", "SERVFAIL", "--socket-count", "1", "-w", "output.txt", "-s", "100"}, gotArgs)
 }