From 47eb7d60e5f09cc40867b0ed0737c324defecfec Mon Sep 17 00:00:00 2001
From: peace-maker <peace-maker@wcfan.de>
Date: Mon, 3 Oct 2016 09:31:17 -0600
Subject: [PATCH] Fix use-after-free crash in SQLite extension (#481)

When the server crashed and the process got terminated, the SqDriver
instance was killed first (e.g. by atexit). SqDatabase tries to access
SqDriver in its destructor.
This patch tells SqDatabase to not use anything from SqDriver anymore
after SqDriver got destroyed.

Next to that, the clientprefs extension relied on the IDatabase pointer
being valid to get the driver pointer. Cache the pointer, so the dbi
system still knows the IDBThreadOperation belonged to the now gone
driver, even after the database object is gone.
---
 extensions/clientprefs/query.cpp      |  5 ++++-
 extensions/clientprefs/query.h        |  1 +
 extensions/sqlite/driver/SqDatabase.h |  4 ++++
 extensions/sqlite/driver/SqDriver.cpp | 27 +++++++++++++++++++++++++++
 extensions/sqlite/driver/SqDriver.h   |  2 ++
 5 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/extensions/clientprefs/query.cpp b/extensions/clientprefs/query.cpp
index 3da7f302..f1883ee3 100644
--- a/extensions/clientprefs/query.cpp
+++ b/extensions/clientprefs/query.cpp
@@ -92,7 +92,7 @@ void TQueryOp::RunThreadPart()
 
 IDBDriver *TQueryOp::GetDriver()
 {
-	return m_database->GetDriver();
+	return m_driver;
 }
 
 IdentityToken_t *TQueryOp::GetOwner()
@@ -114,6 +114,7 @@ TQueryOp::TQueryOp(enum querytype type, int serial)
 	m_type = type;
 	m_serial = serial;
 	m_database = NULL;
+	m_driver = NULL;
 	m_insertId = -1;
 	m_pResult = NULL;
 }
@@ -123,6 +124,7 @@ TQueryOp::TQueryOp(enum querytype type, Cookie *cookie)
 	m_type = type;
 	m_pCookie = cookie;
 	m_database = NULL;
+	m_driver = NULL;
 	m_insertId = -1;
 	m_pResult = NULL;
 	m_serial = 0;
@@ -131,6 +133,7 @@ TQueryOp::TQueryOp(enum querytype type, Cookie *cookie)
 void TQueryOp::SetDatabase(IDatabase *db)
 {
 	m_database = db;
+	m_driver = m_database->GetDriver();
 }
 
 bool TQueryOp::BindParamsAndRun()
diff --git a/extensions/clientprefs/query.h b/extensions/clientprefs/query.h
index 73febd3b..31cdf55d 100644
--- a/extensions/clientprefs/query.h
+++ b/extensions/clientprefs/query.h
@@ -100,6 +100,7 @@ public:
 
 private:
 	IDatabase *m_database;
+	IDBDriver *m_driver;
 	IQuery *m_pResult;
 
 	/* Query type */
diff --git a/extensions/sqlite/driver/SqDatabase.h b/extensions/sqlite/driver/SqDatabase.h
index 7fb9231f..900fa52b 100644
--- a/extensions/sqlite/driver/SqDatabase.h
+++ b/extensions/sqlite/driver/SqDatabase.h
@@ -64,6 +64,10 @@ public:
 	bool SetCharacterSet(const char *characterset);
 public:
 	sqlite3 *GetDb();
+	void PrepareForForcedShutdown()
+	{
+		m_Persistent = false;
+	}
 private:
 	sqlite3 *m_sq3;
 	ke::AutoPtr<ke::Mutex> m_FullLock;
diff --git a/extensions/sqlite/driver/SqDriver.cpp b/extensions/sqlite/driver/SqDriver.cpp
index 27026183..22769951 100644
--- a/extensions/sqlite/driver/SqDriver.cpp
+++ b/extensions/sqlite/driver/SqDriver.cpp
@@ -67,6 +67,32 @@ SqDriver::SqDriver()
 {
 	m_Handle = BAD_HANDLE;
 	m_bThreadSafe = false;
+	m_bShutdown = false;
+}
+
+// The extension got unloaded. Remove any open database now to avoid a use-after-free
+// of g_SqDriver in SqDatabase's destructor.
+SqDriver::~SqDriver()
+{
+	ke::AutoLock lock(&m_OpenLock);
+
+	List<SqDbInfo>::iterator iter;
+	SqDatabase *sqdb;
+	for (iter = m_Cache.begin(); iter != m_Cache.end(); iter++)
+	{
+		// Don't let SqDatabase try to remove itself from m_Cache
+		// now that we're gone.
+		sqdb = (SqDatabase *)(*iter).db;
+		sqdb->PrepareForForcedShutdown();
+
+		iter = m_Cache.erase(iter);
+	}
+	
+	if (!m_bShutdown)
+	{
+		dbi->RemoveDriver(&g_SqDriver);
+		Shutdown();
+	}
 }
 
 void SqDriver::Initialize()
@@ -76,6 +102,7 @@ void SqDriver::Initialize()
 
 void SqDriver::Shutdown()
 {
+	m_bShutdown = true;
 	if (m_bThreadSafe)
 	{
 		sqlite3_enable_shared_cache(0);
diff --git a/extensions/sqlite/driver/SqDriver.h b/extensions/sqlite/driver/SqDriver.h
index 77bb3a69..ee35235f 100644
--- a/extensions/sqlite/driver/SqDriver.h
+++ b/extensions/sqlite/driver/SqDriver.h
@@ -56,6 +56,7 @@ class SqDriver : public IDBDriver
 {
 public:
 	SqDriver();
+	~SqDriver();
 	void Initialize();
 	void Shutdown();
 public:
@@ -74,6 +75,7 @@ private:
 	ke::Mutex m_OpenLock;
 	List<SqDbInfo> m_Cache;
 	bool m_bThreadSafe;
+	bool m_bShutdown;
 };
 
 extern SqDriver g_SqDriver;