From 178d984996c518848e8c1133b6ce52ffaa621579 Mon Sep 17 00:00:00 2001
From: David Adrian <davidcadrian@gmail.com>
Date: Tue, 18 Feb 2020 14:54:03 -0500
Subject: [PATCH] Fix synchronization for Monitor

The goroutine running the monitor isn't actually closed. This PR updates
the API to allow that Goroutine to properly block program exit. This can
be leveraged as we continue to make the configuration non-global.
---
 bin/bin.go |  6 +++++-
 monitor.go | 13 ++++++++++++-
 2 files changed, 17 insertions(+), 2 deletions(-)

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)