diff --git a/bin/bin.go b/bin/bin.go index eab13bd..dc49046 100644 --- a/bin/bin.go +++ b/bin/bin.go @@ -4,6 +4,7 @@ import ( "encoding/json" "os" "runtime/pprof" + "sync" "time" "fmt" @@ -132,7 +133,8 @@ func ZGrab2Main() { s.Init(flag) zgrab2.RegisterScan(moduleType, s) } - monitor := zgrab2.MakeMonitor(1) + wg := sync.WaitGroup{} + monitor := zgrab2.MakeMonitor(1, &wg) monitor.Callback = func(_ string) { dumpHeapProfile() } @@ -141,6 +143,8 @@ func ZGrab2Main() { zgrab2.Process(monitor) end := time.Now() log.Infof("finished grab at %s", end.Format(time.RFC3339)) + monitor.Stop() + wg.Wait() s := Summary{ StatusesPerModule: monitor.GetStatuses(), StartTime: start.Format(time.RFC3339), diff --git a/monitor.go b/monitor.go index 554e9f0..654a0fa 100644 --- a/monitor.go +++ b/monitor.go @@ -1,5 +1,7 @@ package zgrab2 +import "sync" + // Monitor is a collection of states per scans and a channel to communicate // those scans to the monitor type Monitor struct { @@ -34,13 +36,22 @@ func (m *Monitor) GetStatuses() map[string]*State { return m.states } +// Stop indicates the monitor is done and the internal channel should be closed. +// This function does not block, but will allow a call to Wait() on the +// WaitGroup passed to MakeMonitor to return. +func (m *Monitor) Stop() { + close(m.statusesChan) +} + // MakeMonitor returns a Monitor object that can be used to collect and send // the status of a running scan -func MakeMonitor(statusChanSize int) *Monitor { +func MakeMonitor(statusChanSize int, wg *sync.WaitGroup) *Monitor { m := new(Monitor) m.statusesChan = make(chan moduleStatus, statusChanSize) m.states = make(map[string]*State, 10) + wg.Add(1) go func() { + defer wg.Done() for s := range m.statusesChan { if m.states[s.name] == nil { m.states[s.name] = new(State)