Implement full context support in ZDNS and better SIGTERM/SIGINT handling (#538)

* added context and compiles. Killing now kills everything much faster

* add ctx handling for the SIGINT and SIGTERM os signals

* fix up tests with context.Background

* have status handler print a Scan Aborted msg if a user SIGINT's or SIGTERM's a scan rather than Scan Complete

* comment and ensure chan was written to and not closed

---------

Co-authored-by: Zakir Durumeric <zakird@gmail.com>

Author: phillip-stephens

examples/all_nameservers_lookup/main.go

@@ -14,6 +14,7 @@
 package main
 
 import (
+	"context"
 	"net"
 
 	"github.com/miekg/dns"
@@ -33,15 +34,15 @@ func main() {
 	}
 	defer resolver.Close()
 	// LookupAllNameserversIterative will query all root nameservers, and then all TLD nameservers, and then all authoritative nameservers for the domain.
-	result, _, status, err := resolver.LookupAllNameserversIterative(dnsQuestion, nil)
+	result, _, status, err := resolver.LookupAllNameserversIterative(context.Background(), dnsQuestion, nil)
 	if err != nil {
 		log.Fatal("Error looking up domain: ", err)
 	}
 	log.Warnf("Result: %v", result)
 	log.Warnf("Status: %v", status)
 	log.Info("We can also specify which root nameservers to use by setting the argument.")
 
-	result, _, status, err = resolver.LookupAllNameserversIterative(dnsQuestion, []zdns.NameServer{{IP: net.ParseIP("198.41.0.4"), Port: 53}}) // a.root-servers.net
+	result, _, status, err = resolver.LookupAllNameserversIterative(context.Background(), dnsQuestion, []zdns.NameServer{{IP: net.ParseIP("198.41.0.4"), Port: 53}}) // a.root-servers.net
 	if err != nil {
 		log.Fatal("Error looking up domain: ", err)
 	}
@@ -50,7 +51,7 @@ func main() {
 
 	log.Info("You can query multiple recursive resolvers as well")
 
-	externalResult, _, status, err := resolver.LookupAllNameserversExternal(dnsQuestion, []zdns.NameServer{{IP: net.ParseIP("1.1.1.1"), Port: 53}, {IP: net.ParseIP("8.8.8.8"), Port: 53}}) // Cloudflare and Google recursive resolvers, respectively
+	externalResult, _, status, err := resolver.LookupAllNameserversExternal(context.Background(), dnsQuestion, []zdns.NameServer{{IP: net.ParseIP("1.1.1.1"), Port: 53}, {IP: net.ParseIP("8.8.8.8"), Port: 53}}) // Cloudflare and Google recursive resolvers, respectively
 	if err != nil {
 		log.Fatal("Error looking up domain: ", err)
 	}

src/cli/cli.go

@@ -15,6 +15,7 @@
 package cli
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"net"
@@ -32,13 +33,13 @@ import (
 var parser *flags.Parser
 
 type InputHandler interface {
-	FeedChannel(in chan<- string, wg *sync.WaitGroup) error
+	FeedChannel(ctx context.Context, in chan<- string, wg *sync.WaitGroup) error
 }
 type OutputHandler interface {
 	WriteResults(results <-chan string, wg *sync.WaitGroup) error
 }
 type StatusHandler interface {
-	LogPeriodicUpdates(statusChan <-chan zdns.Status, wg *sync.WaitGroup) error
+	LogPeriodicUpdates(statusChan <-chan zdns.Status, statusAbortChan <-chan struct{}, wg *sync.WaitGroup) error
 }
 
 // GeneralOptions core options for all ZDNS modules

src/cli/iohandlers/file_handlers.go

@@ -16,6 +16,7 @@ package iohandlers
 
 import (
 	"bufio"
+	"context"
 	"os"
 	"sync"
 
@@ -36,7 +37,7 @@ func NewFileInputHandler(filepath string) *FileInputHandler {
 	}
 }
 
-func (h *FileInputHandler) FeedChannel(in chan<- string, wg *sync.WaitGroup) error {
+func (h *FileInputHandler) FeedChannel(ctx context.Context, in chan<- string, wg *sync.WaitGroup) error {
 	defer close(in)
 	defer (*wg).Done()
 
@@ -52,6 +53,9 @@ func (h *FileInputHandler) FeedChannel(in chan<- string, wg *sync.WaitGroup) err
 	}
 	s := bufio.NewScanner(f)
 	for s.Scan() {
+		if util.HasCtxExpired(ctx) {
+			return nil // context expired, exiting
+		}
 		in <- s.Text()
 	}
 	if err := s.Err(); err != nil {

src/cli/iohandlers/status_handler.go

@@ -47,7 +47,7 @@ func NewStatusHandler(filePath string) *StatusHandler {
 }
 
 // LogPeriodicUpdates prints a per-second update to the user scan progress and per-status statistics
-func (h *StatusHandler) LogPeriodicUpdates(statusChan <-chan zdns.Status, wg *sync.WaitGroup) error {
+func (h *StatusHandler) LogPeriodicUpdates(statusChan <-chan zdns.Status, statusAbortChan <-chan struct{}, wg *sync.WaitGroup) error {
 	defer wg.Done()
 	// open file for writing
 	var f *os.File
@@ -66,20 +66,25 @@ func (h *StatusHandler) LogPeriodicUpdates(statusChan <-chan zdns.Status, wg *sy
 			}
 		}(f)
 	}
-	if err := h.statusLoop(statusChan, f); err != nil {
+	if err := h.statusLoop(statusChan, statusAbortChan, f); err != nil {
 		return errors.Wrap(err, "error encountered in status loop")
 	}
 	return nil
 }
 
 // statusLoop will print a per-second summary of the scan progress and per-status statistics
-func (h *StatusHandler) statusLoop(statusChan <-chan zdns.Status, statusFile *os.File) error {
+// statusChan is a channel that will receive the statuses from each lookup and updates it's internal state stats
+// statusAbortChan is used for the main thread to notify the status loop that it has aborted and we should notify the
+// user appropriately
+// statusFile is where to write the status updates to
+func (h *StatusHandler) statusLoop(statusChan <-chan zdns.Status, statusAbortChan <-chan struct{}, statusFile *os.File) error {
 	// initialize stats
 	stats := scanStats{
 		statusOccurance: make(map[zdns.Status]int),
 		scanStartTime:   time.Now(),
 	}
 	ticker := time.NewTicker(time.Second)
+	scanAborted := false
 statusLoop:
 	for {
 		select {
@@ -111,13 +116,26 @@ statusLoop:
 				stats.statusOccurance[status] = 0
 			}
 			stats.statusOccurance[status] += 1
+		case _, ok := <-statusAbortChan:
+			if ok {
+				// We can't exit the loop here because the lookup threads will block on writing to the status channel.
+				// Instead, we'll set a flag and when we break out of this loop later (because the status channel is closed)
+				// we'll print an abort message for the user.
+				scanAborted = true
+			}
 		}
 	}
+	scanStateString := "Scan Complete"
+	if scanAborted {
+		scanStateString = "Scan Aborted"
+	}
+
 	timeSinceStart := time.Since(stats.scanStartTime)
-	s := fmt.Sprintf("%02dh:%02dm:%02ds; Scan Complete; %d names scanned; %.02f names/sec; %.01f%% success rate; %s\n",
+	s := fmt.Sprintf("%02dh:%02dm:%02ds; %s; %d names scanned; %.02f names/sec; %.01f%% success rate; %s\n",
 		int(timeSinceStart.Hours()),
 		int(timeSinceStart.Minutes())%60,
 		int(timeSinceStart.Seconds())%60,
+		scanStateString,
 		stats.domainsScanned,
 		float64(stats.domainsScanned)/time.Since(stats.scanStartTime).Seconds(),
 		float64(stats.domainsSuccess*100)/float64(stats.domainsScanned),

src/cli/iohandlers/stream_handlers.go

@@ -16,12 +16,15 @@ package iohandlers
 
 import (
 	"bufio"
+	"context"
 	"io"
 	"sync"
 
 	"github.com/pkg/errors"
 
 	log "github.com/sirupsen/logrus"
+
+	"github.com/zmap/zdns/v2/src/internal/util"
 )
 
 type StreamInputHandler struct {
@@ -34,12 +37,15 @@ func NewStreamInputHandler(r io.Reader) *StreamInputHandler {
 	}
 }
 
-func (h *StreamInputHandler) FeedChannel(in chan<- string, wg *sync.WaitGroup) error {
+func (h *StreamInputHandler) FeedChannel(ctx context.Context, in chan<- string, wg *sync.WaitGroup) error {
 	defer close(in)
 	defer (*wg).Done()
 
 	s := bufio.NewScanner(h.reader)
 	for s.Scan() {
+		if util.HasCtxExpired(ctx) {
+			return nil // context expired, exiting
+		}
 		in <- s.Text()
 	}
 	if err := s.Err(); err != nil {

src/cli/iohandlers/string_slice_input_handler.go

@@ -15,9 +15,12 @@
 package iohandlers
 
 import (
+	"context"
 	"sync"
 
 	log "github.com/sirupsen/logrus"
+
+	"github.com/zmap/zdns/v2/src/internal/util"
 )
 
 // StringSliceInputHandler Feeds a channel with the strings in the slice.
@@ -32,10 +35,13 @@ func NewStringSliceInputHandler(domains []string) *StringSliceInputHandler {
 	return &StringSliceInputHandler{Names: domains}
 }
 
-func (h *StringSliceInputHandler) FeedChannel(in chan<- string, wg *sync.WaitGroup) error {
+func (h *StringSliceInputHandler) FeedChannel(ctx context.Context, in chan<- string, wg *sync.WaitGroup) error {
 	defer close(in)
 	defer wg.Done()
 	for _, name := range h.Names {
+		if util.HasCtxExpired(ctx) {
+			return nil
+		}
 		in <- name
 	}
 	return nil

src/cli/modules.go

@@ -26,7 +26,7 @@ import (
 
 type LookupModule interface {
 	CLIInit(gc *CLIConf, rc *zdns.ResolverConfig) error
-	Lookup(resolver *zdns.Resolver, lookupName string, nameServer *zdns.NameServer) (interface{}, zdns.Trace, zdns.Status, error)
+	Lookup(ctx context.Context, resolver *zdns.Resolver, lookupName string, nameServer *zdns.NameServer) (interface{}, zdns.Trace, zdns.Status, error)
 	Help() string                 // needed to satisfy the ZCommander interface in ZFlags.
 	GetDescription() string       // needed to add a command to the parser, printed to the user. Printed to the user when they run the help command for a given module
 	Validate(args []string) error // needed to satisfy the ZCommander interface in ZFlags
@@ -183,17 +183,17 @@ func (lm *BasicLookupModule) NewFlags() interface{} {
 // non-Iterative query -> we'll send a query to the nameserver provided. If none provided, a random nameserver from the resolver's external nameservers will be used
 // iterative + all-Nameservers query -> we'll send a query to each root NS and query all nameservers down the chain.
 // iterative query -> we'll send a query to a random root NS and query all nameservers down the chain.
-func (lm *BasicLookupModule) Lookup(resolver *zdns.Resolver, lookupName string, nameServer *zdns.NameServer) (interface{}, zdns.Trace, zdns.Status, error) {
+func (lm *BasicLookupModule) Lookup(ctx context.Context, resolver *zdns.Resolver, lookupName string, nameServer *zdns.NameServer) (interface{}, zdns.Trace, zdns.Status, error) {
 	if lm.LookupAllNameServers && lm.IsIterative {
-		return resolver.LookupAllNameserversIterative(&zdns.Question{Name: lookupName, Type: lm.DNSType, Class: lm.DNSClass}, nil)
+		return resolver.LookupAllNameserversIterative(ctx, &zdns.Question{Name: lookupName, Type: lm.DNSType, Class: lm.DNSClass}, nil)
 	}
 	if lm.LookupAllNameServers {
-		return resolver.LookupAllNameserversExternal(&zdns.Question{Name: lookupName, Type: lm.DNSType, Class: lm.DNSClass}, nil)
+		return resolver.LookupAllNameserversExternal(ctx, &zdns.Question{Name: lookupName, Type: lm.DNSType, Class: lm.DNSClass}, nil)
 	}
 	if lm.IsIterative {
-		return resolver.IterativeLookup(context.Background(), &zdns.Question{Name: lookupName, Type: lm.DNSType, Class: lm.DNSClass})
+		return resolver.IterativeLookup(ctx, &zdns.Question{Name: lookupName, Type: lm.DNSType, Class: lm.DNSClass})
 	}
-	return resolver.ExternalLookup(context.Background(), &zdns.Question{Type: lm.DNSType, Class: lm.DNSClass, Name: lookupName}, nameServer)
+	return resolver.ExternalLookup(ctx, &zdns.Question{Type: lm.DNSType, Class: lm.DNSClass, Name: lookupName}, nameServer)
 }
 
 func GetLookupModule(name string) (LookupModule, error) {

src/cli/worker_manager.go

@@ -14,17 +14,20 @@
 package cli
 
 import (
+	"context"
 	"encoding/csv"
 	"encoding/json"
 	"fmt"
 	"io"
 	"math/rand"
 	"net"
 	"os"
+	"os/signal"
 	"runtime"
 	"strconv"
 	"strings"
 	"sync"
+	"syscall"
 	"time"
 
 	"github.com/zmap/zcrypto/x509"
@@ -482,6 +485,10 @@ func populateLocalAddresses(gc *CLIConf, config *zdns.ResolverConfig) (*zdns.Res
 }
 
 func Run(gc CLIConf) {
+	// Create a context that is cancelled on SIGINT or SIGTERM.
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	gc = *populateCLIConfig(&gc)
 	resolverConfig := populateResolverConfig(&gc)
 	// Log any information about the resolver configuration, according to log level
@@ -502,10 +509,21 @@ func Run(gc CLIConf) {
 	//	- process until inChan closes, then wg.done()
 	// Once we processing threads have all finished, wait until the
 	// output and metadata threads have completed
-	inChan := make(chan string)
-	outChan := make(chan string)
-	metaChan := make(chan routineMetadata, gc.Threads)
-	statusChan := make(chan zdns.Status)
+	inChan := make(chan string)                        // input handler feeds inChan with input
+	outChan := make(chan string)                       // lookup workers write to outChan, output handler reads from outChan
+	metaChan := make(chan routineMetadata, gc.Threads) // lookup workers write to metaChan, metadata is collected at end of scan
+	statusChan := make(chan zdns.Status)               // lookup workers write to status chan after each lookup, status handler reads from statusChan
+	statusAbortChan := make(chan struct{})             // used by this thread to signal to the status handler that we've aborted
+
+	sigs := make(chan os.Signal, 1)
+	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
+
+	go func() {
+		<-sigs                        // SIGINT or SIGTERM received
+		cancel()                      // signal to all goroutines to clean up
+		statusAbortChan <- struct{}{} // signal to status handler that we've aborted so it will print a suitable message
+		// to the user vs. the 'Scan Complete' it normally prints
+	}()
 	var routineWG sync.WaitGroup
 
 	inHandler := gc.InputHandler
@@ -525,7 +543,7 @@ func Run(gc CLIConf) {
 
 	// Use handlers to populate the input and output/results channel
 	go func() {
-		if inErr := inHandler.FeedChannel(inChan, &routineWG); inErr != nil {
+		if inErr := inHandler.FeedChannel(ctx, inChan, &routineWG); inErr != nil {
 			log.Fatal(fmt.Sprintf("could not feed input channel: %v", inErr))
 		}
 	}()
@@ -539,7 +557,7 @@ func Run(gc CLIConf) {
 
 	if !gc.QuietStatusUpdates {
 		go func() {
-			if statusErr := statusHandler.LogPeriodicUpdates(statusChan, &routineWG); statusErr != nil {
+			if statusErr := statusHandler.LogPeriodicUpdates(statusChan, statusAbortChan, &routineWG); statusErr != nil {
 				log.Fatal(fmt.Sprintf("could not log periodic status updates: %v", statusErr))
 			}
 		}()
@@ -554,7 +572,7 @@ func Run(gc CLIConf) {
 	for i := 0; i < gc.Threads; i++ {
 		i := i
 		go func(threadID int) {
-			initWorkerErr := doLookupWorker(&gc, resolverConfig, inChan, outChan, metaChan, statusChan, &lookupWG)
+			initWorkerErr := doLookupWorker(ctx, &gc, resolverConfig, inChan, outChan, metaChan, statusChan, &lookupWG)
 			if initWorkerErr != nil {
 				log.Fatalf("could not start lookup worker #%d: %v", i, initWorkerErr)
 			}
@@ -564,6 +582,7 @@ func Run(gc CLIConf) {
 	close(outChan)
 	close(metaChan)
 	close(statusChan)
+	close(statusAbortChan)
 	routineWG.Wait()
 	if gc.MetadataFilePath != "" {
 		// we're done processing data. aggregate all the data from individual routines
@@ -610,25 +629,27 @@ func Run(gc CLIConf) {
 }
 
 // doLookupWorker is a single worker thread that processes lookups from the input channel. It calls wg.Done when it is finished.
-func doLookupWorker(gc *CLIConf, rc *zdns.ResolverConfig, inputChan <-chan string, outputChan chan<- string, metaChan chan<- routineMetadata, statusChan chan<- zdns.Status, wg *sync.WaitGroup) error {
+func doLookupWorker(ctx context.Context, gc *CLIConf, rc *zdns.ResolverConfig, inputChan <-chan string, outputChan chan<- string, metaChan chan<- routineMetadata, statusChan chan<- zdns.Status, wg *sync.WaitGroup) error {
 	defer wg.Done()
 	resolver, err := zdns.InitResolver(rc)
 	if err != nil {
 		return fmt.Errorf("could not init resolver: %w", err)
 	}
+	defer resolver.Close() // close the resolver, freeing up resources
 	var metadata routineMetadata
 	metadata.Status = make(map[zdns.Status]int)
 
 	for line := range inputChan {
-		handleWorkerInput(gc, rc, line, resolver, &metadata, outputChan, statusChan)
+		if util.HasCtxExpired(ctx) {
+			break
+		}
+		handleWorkerInput(ctx, gc, rc, line, resolver, &metadata, outputChan, statusChan)
 	}
-	// close the resolver, freeing up resources
-	resolver.Close()
 	metaChan <- metadata
 	return nil
 }
 
-func handleWorkerInput(gc *CLIConf, rc *zdns.ResolverConfig, line string, resolver *zdns.Resolver, metadata *routineMetadata, outputChan chan<- string, statusChan chan<- zdns.Status) {
+func handleWorkerInput(ctx context.Context, gc *CLIConf, rc *zdns.ResolverConfig, line string, resolver *zdns.Resolver, metadata *routineMetadata, outputChan chan<- string, statusChan chan<- zdns.Status) {
 	// we'll process each module sequentially, parallelism is per-domain
 	res := zdns.Result{Results: make(map[string]zdns.SingleModuleResult, len(gc.ActiveModules))}
 	// get the fields that won't change for each lookup module
@@ -685,7 +706,7 @@ func handleWorkerInput(gc *CLIConf, rc *zdns.ResolverConfig, line string, resolv
 		res.Class = dns.Class(gc.Class).String()
 
 		startTime := time.Now()
-		innerRes, trace, status, err = module.Lookup(resolver, lookupName, nameServer)
+		innerRes, trace, status, err = module.Lookup(ctx, resolver, lookupName, nameServer)
 
 		lookupRes := zdns.SingleModuleResult{
 			Timestamp: time.Now().Format(gc.TimeFormat),

src/modules/alookup/a_lookup.go

@@ -15,6 +15,8 @@
 package alookup
 
 import (
+	"context"
+
 	"github.com/pkg/errors"
 
 	"github.com/zmap/zdns/v2/src/cli"
@@ -51,8 +53,8 @@ func (aMod *ALookupModule) Init(ipv4Lookup bool, ipv6Lookup bool) {
 	aMod.IPv6Lookup = ipv6Lookup
 }
 
-func (aMod *ALookupModule) Lookup(r *zdns.Resolver, lookupName string, nameServer *zdns.NameServer) (interface{}, zdns.Trace, zdns.Status, error) {
-	ipResult, trace, status, err := r.DoTargetedLookup(lookupName, nameServer, aMod.baseModule.IsIterative, aMod.IPv4Lookup, aMod.IPv6Lookup)
+func (aMod *ALookupModule) Lookup(ctx context.Context, r *zdns.Resolver, lookupName string, nameServer *zdns.NameServer) (interface{}, zdns.Trace, zdns.Status, error) {
+	ipResult, trace, status, err := r.DoTargetedLookup(ctx, lookupName, nameServer, aMod.baseModule.IsIterative, aMod.IPv4Lookup, aMod.IPv6Lookup)
 	return ipResult, trace, status, err
 }