Mercurial > hg

--- a/mercurial/exewrapper.c	Thu Sep 02 14:08:45 2021 -0700
+++ b/mercurial/exewrapper.c	Sun Sep 19 01:23:16 2021 -0400
@@ -29,6 +29,134 @@
 #define _countof(array) (sizeof(array) / sizeof(array[0]))
 #endif

+#if PY_MAJOR_VERSION >= 3
+
+#pragma comment(lib, "Advapi32.lib")
+
+/* python.org installations */
+#define CORE_PATH L"SOFTWARE\\Python\\PythonCore"
+
+/* Microsoft Store installations */
+#define LOOKASIDE_PATH                                                         \
+	L"SOFTWARE\\Microsoft\\AppModel\\Lookaside\\user\\Software\\Python\\"  \
+	L"PythonCore"
+
+static wchar_t *_locate_python_for_key(HKEY root, LPCWSTR subkey, size_t *size)
+{
+	wchar_t installPathKey[512];
+	wchar_t *executable = NULL;
+	DWORD type;
+	DWORD sz = 0;
+	HKEY ip_key;
+	LSTATUS status;
+
+	_snwprintf_s(installPathKey, sizeof(installPathKey), _TRUNCATE,
+	             L"%ls\\%d.%d\\InstallPath", subkey, PY_MAJOR_VERSION,
+	             PY_MINOR_VERSION);
+
+	status =
+	    RegOpenKeyExW(root, installPathKey, 0, KEY_QUERY_VALUE, &ip_key);
+
+	if (status != ERROR_SUCCESS)
+		return NULL;
+
+	status = RegQueryValueExW(ip_key, L"ExecutablePath", NULL, &type,
+	                          (LPBYTE)executable, &sz);
+	if (status == ERROR_SUCCESS) {
+		/* Allocate extra space so path\to\python.exe can be converted
+		 * to path\to\python39.dll + NUL.
+		 */
+		*size = sz + sizeof(_T(HGPYTHONLIB ".dll")) + sizeof(wchar_t);
+		executable = malloc(*size);
+
+		if (executable) {
+			status =
+			    RegQueryValueExW(ip_key, L"ExecutablePath", NULL,
+			                     &type, (LPBYTE)executable, &sz);
+
+			if (status != ERROR_SUCCESS) {
+				free(executable);
+				executable = NULL;
+			} else {
+				/* Not all values are stored NUL terminated */
+				executable[sz] = L'\0';
+			}
+		}
+	}
+
+	RegCloseKey(ip_key);
+	return executable;
+}
+
+static HMODULE load_system_py3(void)
+{
+	wchar_t *subkeys[] = {CORE_PATH, LOOKASIDE_PATH};
+	HKEY roots[] = {HKEY_LOCAL_MACHINE, HKEY_CURRENT_USER};
+
+	/* Give priority to python.org installs, because MS Store installs can
+	 * break with user profile corruption, and also use some NTFS feature
+	 * that MSYS doesn't understand.
+	 */
+	for (int s = 0; s < _countof(subkeys); s++) {
+		for (int r = 0; r < _countof(roots); r++) {
+			size_t size = 0;
+			wchar_t *py =
+			    _locate_python_for_key(roots[r], subkeys[s], &size);
+			wchar_t *cut = NULL;
+			HMODULE pydll;
+
+			if (!py) {
+				continue;
+			}
+
+			/* Cut off the python executable component */
+			cut = wcsrchr(py, L'\\');
+			if (cut == NULL) {
+				free(py);
+				continue;
+			}
+			*cut = 0;
+
+			wcscat_s(py, size, _T("\\" HGPYTHONLIB ".dll"));
+
+			pydll = LoadLibrary(py);
+
+			/* Also load python3.dll, so we don't pick up a random
+			 * one on PATH. We don't search {sys.prefix}\DLLs like
+			 * python.exe because this is commented as "not been a
+			 * normal install layout for a while", and don't search
+			 * LOAD_LIBRARY_SEARCH_APPLICATION_DIR because it's not
+			 * clear what the use case is.
+			 */
+			if (pydll != NULL) {
+				HANDLE py3dll = NULL;
+
+				*cut = 0;
+				wcscat_s(py, size, L"\\python3.dll");
+
+				py3dll = LoadLibrary(py);
+				if (py3dll == NULL) {
+					const char *err;
+					err = "failed to load python3.dll "
+					      "for " HGPYTHONLIB ".dll";
+					fprintf(stderr, "abort: %s (0x%X)\n",
+					        err, GetLastError());
+					exit(255);
+				}
+			}
+
+			free(py);
+
+			if (pydll != NULL) {
+				return pydll;
+			}
+		}
+	}
+
+	return NULL;
+}
+#endif
+
 static TCHAR pyscript[MAX_PATH + 10];
 static TCHAR pyhome[MAX_PATH + 10];
 static TCHAR pydllfile[MAX_PATH + 10];
@@ -108,6 +236,12 @@
 		Py_SetPythonHome(pyhome);
 	}

+#if PY_MAJOR_VERSION >= 3
+	if (pydll == NULL) {
+		pydll = load_system_py3();
+	}
+#endif
+
 	if (pydll == NULL) {
 		pydll = LoadLibrary(_T(HGPYTHONLIB) _T(".dll"));
 		if (pydll == NULL) {