Checks added to ensure infra information is resolved

enum/enum.go

@@ -46,6 +46,7 @@ type Enumeration struct {
 	nameSrc        *enumSource
 	subTask        *subdomainTask
 	dnsTask        *dNSTask
+	store          *dataManager
 }
 
 // NewEnumeration returns an initialized Enumeration that has not been started yet.
@@ -68,6 +69,7 @@ func NewEnumeration(cfg *config.Config, sys systems.System) *Enumeration {
 
 	e.dnsTask = newDNSTask(e)
 	e.subTask = newSubdomainTask(e)
+	e.store = newDataManager(e)
 	return e
 }
 
@@ -109,7 +111,7 @@ func (e *Enumeration) Start(ctx context.Context) error {
 	stages = append(stages, pipeline.FIFO("filter", e.makeFilterTaskFunc()))
 
 	if !e.Config.Passive {
-		stages = append(stages, pipeline.DynamicPool("store", newDataManager(e), defaultPipelineTasks))
+		stages = append(stages, pipeline.DynamicPool("store", e.store, defaultPipelineTasks))
 		stages = append(stages, pipeline.FIFO("", e.subTask))
 	}
 	if e.Config.Active {
@@ -124,12 +126,18 @@ func (e *Enumeration) Start(ctx context.Context) error {
 	 * by the user and names acquired from the graph database can be brought
 	 * into the enumeration
 	 */
-	go e.submitKnownNames()
+	e.submitKnownNames()
 	e.submitProvidedNames()
 	e.submitDomainNames()
 	e.submitASNs()
 
-	return pipeline.NewPipeline(stages...).Execute(e.ctx, e.nameSrc, e.makeOutputSink())
+	err := pipeline.NewPipeline(stages...).Execute(e.ctx, e.nameSrc, e.makeOutputSink())
+	if !e.Config.Passive {
+		// Ensure all data has been stored
+		e.store.signalDone <- struct{}{}
+		<-e.store.confirmDone
+	}
+	return err
 }
 
 func (e *Enumeration) startupAndCleanup() {

enum/store.go

@@ -25,15 +25,19 @@ import (
 
 // dataManager is the stage that stores all data processed by the pipeline.
 type dataManager struct {
-	enum  *Enumeration
-	queue queue.Queue
+	enum        *Enumeration
+	queue       queue.Queue
+	signalDone  chan struct{}
+	confirmDone chan struct{}
 }
 
 // newDataManager returns a dataManager specific to the provided Enumeration.
 func newDataManager(e *Enumeration) *dataManager {
 	dm := &dataManager{
-		enum:  e,
-		queue: queue.NewQueue(),
+		enum:        e,
+		queue:       queue.NewQueue(),
+		signalDone:  make(chan struct{}, 2),
+		confirmDone: make(chan struct{}, 2),
 	}
 
 	go dm.processASNRequests()
@@ -444,68 +448,70 @@ func (dm *dataManager) addrRequest(ctx context.Context, req *requests.AddrReques
 }
 
 func (dm *dataManager) processASNRequests() {
-	graph := dm.enum.Graph
-	uuid := dm.enum.Config.UUID.String()
-loop:
 	for {
 		select {
-		case <-dm.enum.ctx.Done():
-			break loop
-		case <-dm.enum.done:
-			break loop
-		case <-dm.queue.Signal():
-			e, found := dm.queue.Next()
-			if !found {
-				continue loop
+		case <-dm.signalDone:
+			if dm.queue.Len() > 0 {
+				dm.queue.Process(func(e interface{}) {
+					if qar, ok := e.(*queuedAddrRequest); ok {
+						dm.findInfraInfo(qar.Req)
+					}
+				})
 			}
+			dm.confirmDone <- struct{}{}
+			return
+		case <-dm.queue.Signal():
+			dm.queue.Process(func(e interface{}) {
+				if qar, ok := e.(*queuedAddrRequest); ok {
+					dm.findInfraInfo(qar.Req)
+				}
+			})
+		}
+	}
+}
 
-			qar, ok := e.(*queuedAddrRequest)
-			if !ok {
-				continue loop
-			}
-			req := qar.Req
-
-			if r := dm.enum.Sys.Cache().AddrSearch(req.Address); r != nil {
-				_ = graph.UpsertInfrastructure(r.ASN, r.Description, req.Address, r.Prefix, r.Source, uuid)
-				continue loop
-			}
+func (dm *dataManager) findInfraInfo(req *requests.AddrRequest) {
+	graph := dm.enum.Graph
+	uuid := dm.enum.Config.UUID.String()
 
-			for _, src := range dm.enum.srcs {
-				src.Request(dm.enum.ctx, &requests.ASNRequest{Address: req.Address})
-			}
-			time.Sleep(10 * time.Second)
+	if r := dm.enum.Sys.Cache().AddrSearch(req.Address); r != nil {
+		_ = graph.UpsertInfrastructure(r.ASN, r.Description, req.Address, r.Prefix, r.Source, uuid)
+		return
+	}
 
-			if r := dm.enum.Sys.Cache().AddrSearch(req.Address); r != nil {
-				_ = graph.UpsertInfrastructure(r.ASN, r.Description, req.Address, r.Prefix, r.Source, uuid)
-				continue loop
-			}
+	for _, src := range dm.enum.srcs {
+		src.Request(dm.enum.ctx, &requests.ASNRequest{Address: req.Address})
+	}
 
-			asn := 0
-			desc := "Unknown"
-			prefix := fakePrefix(req.Address)
-			_ = graph.UpsertInfrastructure(asn, desc, req.Address, prefix, "RIR", uuid)
+	for i := 0; i < 20; i++ {
+		if r := dm.enum.Sys.Cache().AddrSearch(req.Address); r != nil {
+			_ = graph.UpsertInfrastructure(r.ASN, r.Description, req.Address, r.Prefix, r.Source, uuid)
+			return
+		}
+		time.Sleep(500 * time.Millisecond)
+	}
 
-			first, cidr, err := net.ParseCIDR(prefix)
-			if err != nil {
-				continue loop
-			}
-			if ones, _ := cidr.Mask.Size(); ones == 0 {
-				continue loop
-			}
+	asn := 0
+	desc := "Unknown"
+	prefix := fakePrefix(req.Address)
+	_ = graph.UpsertInfrastructure(asn, desc, req.Address, prefix, "RIR", uuid)
 
-			dm.enum.Sys.Cache().Update(&requests.ASNRequest{
-				Address:     first.String(),
-				ASN:         asn,
-				Prefix:      cidr.String(),
-				Description: desc,
-				Tag:         requests.RIR,
-				Source:      "RIR",
-			})
-		}
+	first, cidr, err := net.ParseCIDR(prefix)
+	if err != nil {
+		return
+	}
+	if ones, _ := cidr.Mask.Size(); ones == 0 {
+		return
 	}
 
-	// Empty the queue
-	dm.queue.Process(func(e interface{}) {})
+	dm.enum.Sys.Cache().Update(&requests.ASNRequest{
+		Address:     first.String(),
+		ASN:         asn,
+		Prefix:      cidr.String(),
+		Description: desc,
+		Tag:         requests.RIR,
+		Source:      "RIR",
+	})
 }
 
 func fakePrefix(addr string) string {