Cleanup DNS requests when modules are unloaded, fixes unloading m_dnsbl during the middle of queries

5 files changed, 28 insertions, 15 deletions

diff --git a/include/dns.h b/include/dns.h
index a8ef82d44..97cb0a677 100644
--- a/include/dns.h
+++ b/include/dns.h
@@ -48,6 +48,7 @@ enum DNSError
 
 class DNSRequestTimeout; // Forward declarations
 struct DNSRecord;
+class Module;
 
 /** The request
  */
@@ -57,12 +58,14 @@ class DNSRequest
 	DNSRequestTimeout *timeout;
 
  public:
+	Module *creator;
+
 	/* Address we're looking up */
 	Anope::string address;
 	/* QueryType, A, AAAA, PTR etc */
 	QueryType QT;
 
-	DNSRequest(const Anope::string &addr, QueryType qt, bool cache = false);
+	DNSRequest(const Anope::string &addr, QueryType qt, bool cache = false, Module *c = NULL);
 
 	virtual ~DNSRequest();
 
@@ -158,6 +161,8 @@ class DNSManager : public Timer
 	void AddCache(DNSRecord *rr);
 	bool CheckCache(DNSRequest *request);
 	void Tick(time_t now);
+
+	void Cleanup(Module *mod);
 };
 
 /** A DNS timeout, one is made for every DNS request to detect timeouts
diff --git a/include/modules.h b/include/modules.h
index 60a16a978..ac76a4534 100644
--- a/include/modules.h
+++ b/include/modules.h
@@ -1294,17 +1294,6 @@ class service_reference : public dynamic_reference<T>
 
 	inline T *operator->()
 	{
-		if (this->invalid)
-		{
-			this->invalid = false;
-			this->ref = NULL;
-		}
-		if (!this->ref)
-		{
-			this->ref = static_cast<T *>(ModuleManager::GetService(this->name));
-			if (this->ref)
-				this->ref->AddReference(this);
-		}
 		return this->ref;
 	}
 };
diff --git a/modules/extra/m_dnsbl.cpp b/modules/extra/m_dnsbl.cpp
index 42e505c31..ce82b68a4 100644
--- a/modules/extra/m_dnsbl.cpp
+++ b/modules/extra/m_dnsbl.cpp
@@ -23,7 +23,7 @@ class DNSBLResolver : public DNSRequest
 	bool add_to_akill;
 
  public:
-	DNSBLResolver(User *u, const Blacklist &b, const Anope::string &host, bool add_akill) : DNSRequest(host, DNS_QUERY_A, true), user(u), blacklist(b), add_to_akill(add_akill) { }
+	DNSBLResolver(Module *c, User *u, const Blacklist &b, const Anope::string &host, bool add_akill) : DNSRequest(host, DNS_QUERY_A, true, c), user(u), blacklist(b), add_to_akill(add_akill) { }
 
 	void OnLookupComplete(const DNSRecord *)
 	{
@@ -120,7 +120,7 @@ class ModuleDNSBL : public Module
 			try
 			{
 				Anope::string dnsbl_host = user_ip.addr() + "." + b.name;
-				new DNSBLResolver(u, b, dnsbl_host, this->add_to_akill);
+				new DNSBLResolver(this, u, b, dnsbl_host, this->add_to_akill);
 			}
 			catch (const SocketException &ex)
 			{
diff --git a/src/dns.cpp b/src/dns.cpp
index 329c03c86..7b96cf660 100644
--- a/src/dns.cpp
+++ b/src/dns.cpp
@@ -7,7 +7,7 @@ static inline unsigned short GetRandomID()
 	return random();
 }
 
-DNSRequest::DNSRequest(const Anope::string &addr, QueryType qt, bool cache) : address(addr), QT(qt)
+DNSRequest::DNSRequest(const Anope::string &addr, QueryType qt, bool cache, Module *c) : address(addr), QT(qt), creator(c)
 {
 	if (!DNSEngine)
 		DNSEngine = new DNSManager();
@@ -531,6 +531,22 @@ void DNSManager::Tick(time_t now)
 	}
 }
 
+void DNSManager::Cleanup(Module *mod)
+{
+	for (std::map<short, DNSRequest *>::iterator it = this->requests.begin(), it_end = this->requests.end(); it != it_end;)
+	{
+		short id = it->first;
+		DNSRequest *req = it->second;
+		++it;
+
+		if (req->creator && req->creator == mod)
+		{
+			delete req;
+			this->requests.erase(id);
+		}
+	}
+}
+
 DNSRequestTimeout::DNSRequestTimeout(DNSRequest *r, time_t timeout) : Timer(timeout), request(r)
 {
 	this->done = false;
diff --git a/src/modulemanager.cpp b/src/modulemanager.cpp
index bacf6c8d1..dc1258340 100644
--- a/src/modulemanager.cpp
+++ b/src/modulemanager.cpp
@@ -240,6 +240,9 @@ int ModuleManager::UnloadModule(Module *m, User *u)
 
 	FOREACH_MOD(I_OnModuleUnload, OnModuleUnload(u, m));
 
+	if (DNSEngine)
+		DNSEngine->Cleanup(m);
+
 	DeleteModule(m);
 	return MOD_ERR_OK;
 }