From f7013462ccc10781bf1ac722d6f0f8c5b7c31b7b Mon Sep 17 00:00:00 2001
From: Christian Ullrich <chris@chrullrich.net>
Date: Tue, 6 Sep 2016 10:02:06 +0200
Subject: [PATCH] Pin any DLL as soon as seen when looking for _putenv().

---
 src/port/win32env.c | 56 ++++++++++++++++++++++++++++++++++++-----------------
 1 file changed, 38 insertions(+), 18 deletions(-)

diff --git a/src/port/win32env.c b/src/port/win32env.c
index d6391b0..48c4d82 100644
--- a/src/port/win32env.c
+++ b/src/port/win32env.c
@@ -75,26 +75,27 @@ pgwin32_putenv(const char *envval)
 		char	   *modulename;
 		HMODULE		hmodule;
 		PUTENVPROC	putenvFunc;
+		bool		pinned;
 	}			rtmodules[] =
 	{
-		{ "msvcrt",    NULL, NULL },
-		{ "msvcrtd",   NULL, NULL },	/* Visual Studio 6.0 / mingw */
-		{ "msvcr70",   NULL, NULL },
-		{ "msvcr70d",  NULL, NULL },	/* Visual Studio 2002 */
-		{ "msvcr71",   NULL, NULL },
-		{ "msvcr71d",  NULL, NULL },	/* Visual Studio 2003 */
-		{ "msvcr80",   NULL, NULL },
-		{ "msvcr80d",  NULL, NULL },	/* Visual Studio 2005 */
-		{ "msvcr90",   NULL, NULL },
-		{ "msvcr90d",  NULL, NULL },	/* Visual Studio 2008 */
-		{ "msvcr100",  NULL, NULL },
-		{ "msvcr100d", NULL, NULL },	/* Visual Studio 2010 */
-		{ "msvcr110",  NULL, NULL },
-		{ "msvcr110d", NULL, NULL },	/* Visual Studio 2012 */
-		{ "msvcr120",  NULL, NULL },
-		{ "msvcr120d", NULL, NULL },	/* Visual Studio 2013 */
-		{ "ucrtbase",  NULL, NULL },
-		{ "ucrtbased", NULL, NULL },	/* Visual Studio 2015 and later */
+		{ "msvcrt",    NULL, NULL, false },
+		{ "msvcrtd",   NULL, NULL, false },	/* Visual Studio 6.0 / mingw */
+		{ "msvcr70",   NULL, NULL, false },
+		{ "msvcr70d",  NULL, NULL, false },	/* Visual Studio 2002 */
+		{ "msvcr71",   NULL, NULL, false },
+		{ "msvcr71d",  NULL, NULL, false },	/* Visual Studio 2003 */
+		{ "msvcr80",   NULL, NULL, false },
+		{ "msvcr80d",  NULL, NULL, false },	/* Visual Studio 2005 */
+		{ "msvcr90",   NULL, NULL, false },
+		{ "msvcr90d",  NULL, NULL, false },	/* Visual Studio 2008 */
+		{ "msvcr100",  NULL, NULL, false },
+		{ "msvcr100d", NULL, NULL, false },	/* Visual Studio 2010 */
+		{ "msvcr110",  NULL, NULL, false },
+		{ "msvcr110d", NULL, NULL, false },	/* Visual Studio 2012 */
+		{ "msvcr120",  NULL, NULL, false },
+		{ "msvcr120d", NULL, NULL, false },	/* Visual Studio 2013 */
+		{ "ucrtbase",  NULL, NULL, false },
+		{ "ucrtbased", NULL, NULL, false },	/* Visual Studio 2015 and later */
 		{ NULL, NULL, NULL }
 	};
 	int			i;
@@ -113,6 +114,25 @@ pgwin32_putenv(const char *envval)
 				}
 				else
 				{
+					/*
+					 * Pin this DLL handle as soon as possible to avoid it
+					 * to be unloaded until the process terminates.
+					 */
+					if (!rtmodules[i].pinned)
+					{
+						/*
+						 * This leaks the new reference, but that does
+						 * not matter because we pin the DLL anyway.
+						 */
+						HMODULE tmp;
+						BOOL res = GetModuleHandleEx(
+								GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS
+									| GET_MODULE_HANDLE_EX_FLAG_PIN,
+								(LPCTSTR)rtmodules[i].hmodule,
+								&tmp);
+						rtmodules[i].pinned = res != 0;
+					}
+
 					rtmodules[i].putenvFunc = (PUTENVPROC) GetProcAddress(rtmodules[i].hmodule, "_putenv");
 					if (rtmodules[i].putenvFunc == NULL)
 					{
-- 
2.10.2.windows.1

