From 82576f8de211265c3c32383cf003c65c21a7a9e6 Mon Sep 17 00:00:00 2001
From: Tim Wojtulewicz <tim@corelight.com>
Date: Wed, 1 Jun 2022 09:21:23 -0700
Subject: [PATCH] GH-2136: Clean up DNS_Mgr before shutting down

---
 src/DNS_Mgr.cc          | 24 ++++++++++++++++++++++++
 src/DNS_Mgr.h           |  7 +++++++
 src/iosource/Manager.cc |  7 ++++---
 3 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/src/DNS_Mgr.cc b/src/DNS_Mgr.cc
index c6934a43b8..33bb667afa 100644
--- a/src/DNS_Mgr.cc
+++ b/src/DNS_Mgr.cc
@@ -560,6 +560,12 @@ DNS_Mgr::~DNS_Mgr()
 	ares_library_cleanup();
 	}
 
+void DNS_Mgr::Done()
+	{
+	shutting_down = true;
+	Flush();
+	}
+
 void DNS_Mgr::RegisterSocket(int fd, bool read, bool write)
 	{
 	if ( read && socket_fds.count(fd) == 0 )
@@ -713,6 +719,9 @@ static void resolve_lookup_cb(DNS_Mgr::LookupCallback* callback, const std::stri
 
 ValPtr DNS_Mgr::Lookup(const std::string& name, int request_type)
 	{
+	if ( shutting_down )
+		return nullptr;
+
 	if ( request_type == T_A || request_type == T_AAAA )
 		return LookupHost(name);
 
@@ -761,6 +770,9 @@ ValPtr DNS_Mgr::Lookup(const std::string& name, int request_type)
 
 TableValPtr DNS_Mgr::LookupHost(const std::string& name)
 	{
+	if ( shutting_down )
+		return nullptr;
+
 	if ( mode == DNS_FAKE )
 		return fake_name_lookup_result(name);
 
@@ -809,6 +821,9 @@ TableValPtr DNS_Mgr::LookupHost(const std::string& name)
 
 StringValPtr DNS_Mgr::LookupAddr(const IPAddr& addr)
 	{
+	if ( shutting_down )
+		return nullptr;
+
 	if ( mode == DNS_FAKE )
 		return make_intrusive<StringVal>(fake_addr_lookup_result(addr));
 
@@ -853,6 +868,9 @@ StringValPtr DNS_Mgr::LookupAddr(const IPAddr& addr)
 
 void DNS_Mgr::LookupHost(const std::string& name, LookupCallback* callback)
 	{
+	if ( shutting_down )
+		return;
+
 	if ( mode == DNS_FAKE )
 		{
 		resolve_lookup_cb(callback, fake_name_lookup_result(name));
@@ -893,6 +911,9 @@ void DNS_Mgr::LookupHost(const std::string& name, LookupCallback* callback)
 
 void DNS_Mgr::LookupAddr(const IPAddr& addr, LookupCallback* callback)
 	{
+	if ( shutting_down )
+		return;
+
 	if ( mode == DNS_FAKE )
 		{
 		resolve_lookup_cb(callback, fake_addr_lookup_result(addr));
@@ -932,6 +953,9 @@ void DNS_Mgr::LookupAddr(const IPAddr& addr, LookupCallback* callback)
 
 void DNS_Mgr::Lookup(const std::string& name, int request_type, LookupCallback* callback)
 	{
+	if ( shutting_down )
+		return;
+
 	if ( mode == DNS_FAKE )
 		{
 		resolve_lookup_cb(callback, fake_lookup_result(name, request_type));
diff --git a/src/DNS_Mgr.h b/src/DNS_Mgr.h
index 2adfc3ac01..44e3a82464 100644
--- a/src/DNS_Mgr.h
+++ b/src/DNS_Mgr.h
@@ -97,6 +97,11 @@ public:
 	explicit DNS_Mgr(DNS_MgrMode mode);
 	~DNS_Mgr() override;
 
+	/**
+	 * Finalizes the source when it's being closed.
+	 */
+	void Done() override;
+
 	/**
 	 * Finalizes the manager initialization. This should be called only after all
 	 * of the scripts have been parsed at startup.
@@ -337,6 +342,8 @@ protected:
 
 	std::set<int> socket_fds;
 	std::set<int> write_socket_fds;
+
+	bool shutting_down = false;
 	};
 
 extern DNS_Mgr* dns_mgr;
diff --git a/src/iosource/Manager.cc b/src/iosource/Manager.cc
index 6c9876c398..439efda6c2 100644
--- a/src/iosource/Manager.cc
+++ b/src/iosource/Manager.cc
@@ -61,11 +61,12 @@ Manager::~Manager()
 	delete wakeup;
 	wakeup = nullptr;
 
-	for ( SourceList::iterator i = sources.begin(); i != sources.end(); ++i )
-		{
-		auto src = *i;
+	// Make sure all of the sources are done before we try to delete any of them.
+	for ( auto& src : sources )
 		src->src->Done();
 
+	for ( auto& src : sources )
+		{
 		if ( src->manage_lifetime )
 			delete src->src;