From 41d6d7ec572f9b8388a5dd573b5638d47dc827e2 Mon Sep 17 00:00:00 2001
From: Adam Demasi <adam@hbang.com.au>
Date: Thu, 19 Sep 2024 23:51:28 +0930
Subject: [PATCH] Check if the ActiveX control is registered on launch

---
 launcher/LaunchUpdateSite.c | 36 +++++++++++++++++++++++++++++-------
 launcher/RegisterServer.c   |  8 ++++++--
 launcher/RegisterServer.h   |  5 +++++
 3 files changed, 40 insertions(+), 9 deletions(-)
 create mode 100644 launcher/RegisterServer.h

diff --git a/launcher/LaunchUpdateSite.c b/launcher/LaunchUpdateSite.c
index f64f4a3..c140978 100644
--- a/launcher/LaunchUpdateSite.c
+++ b/launcher/LaunchUpdateSite.c
@@ -7,15 +7,19 @@
 #include <wchar.h>
 #include "Exec.h"
 #include "HResult.h"
+#include "MsgBox.h"
 #include "Registry.h"
+#include "RegisterServer.h"
 #include "User.h"
 #include "VersionInfo.h"
-#include "MsgBox.h"
 
 const LPTSTR UpdateSiteURLHttp      = L"http://legacyupdate.net/windowsupdate/v6/";
 const LPTSTR UpdateSiteURLHttps     = L"https://legacyupdate.net/windowsupdate/v6/";
 const LPTSTR UpdateSiteFirstRunFlag = L"?firstrun=true";
 
+DEFINE_GUID(IID_ILegacyUpdateCtrl,  0xC33085BB, 0xC3E1, 0x4D27, 0xA2, 0x14, 0xAF, 0x01, 0x95, 0x3D, 0xF5, 0xE5);
+DEFINE_GUID(CLSID_LegacyUpdateCtrl, 0xAD28E0DF, 0x5F5A, 0x40B5, 0x94, 0x32, 0x85, 0xEF, 0xD9, 0x7D, 0x1F, 0x9F);
+
 static BOOL CanUseSSLConnection() {
 	// Get the Windows Update website URL set by Legacy Update setup
 	LPWSTR data;
@@ -54,15 +58,33 @@ void LaunchUpdateSite(int argc, LPWSTR *argv, int nCmdShow) {
 	// If running on 2k/XP, make sure we're elevated. If not, show Run As prompt.
 	if (!IsOSVersionOrLater(6, 0) && !IsUserAdmin()) {
 		LPWSTR filename;
-		DWORD filenameSize;
-		GetOwnFileName(&filename, &filenameSize);
-
-		INT_PTR execResult = (INT_PTR)ShellExecute(NULL, L"runas", filename, GetCommandLineW(), NULL, nCmdShow);
+		GetOwnFileName(&filename);
+		hr = Exec(L"runas", filename, GetCommandLineW(), NULL, nCmdShow, FALSE, NULL);
+		LocalFree(filename);
 
 		// Access denied happens when the user clicks No/Cancel.
-		if (execResult <= 32 && execResult != SE_ERR_ACCESSDENIED) {
-			hr = HRESULT_FROM_WIN32(GetLastError());
+		if (hr == HRESULT_FROM_WIN32(ERROR_CANCELLED)) {
+			hr = S_OK;
 		}
+
+		goto end;
+	}
+
+	// Can we instantiate our own ActiveX control? If not, try to register it.
+	hr = CoCreateInstance(&CLSID_LegacyUpdateCtrl, NULL, CLSCTX_LOCAL_SERVER, &IID_ILegacyUpdateCtrl, (void **)&browser);
+	if (hr == REGDB_E_CLASSNOTREG) {
+		hr = RegisterServer(TRUE, TRUE);
+		if (!SUCCEEDED(hr)) {
+			goto end;
+		}
+
+		hr = CoCreateInstance(&CLSID_LegacyUpdateCtrl, NULL, CLSCTX_LOCAL_SERVER, &IID_ILegacyUpdateCtrl, (void **)&browser);
+		if (!SUCCEEDED(hr)) {
+			goto end;
+		}
+
+		IUnknown_Release(browser);
+	} else if (!SUCCEEDED(hr)) {
 		goto end;
 	}
 
diff --git a/launcher/RegisterServer.c b/launcher/RegisterServer.c
index 98290c8..696f556 100644
--- a/launcher/RegisterServer.c
+++ b/launcher/RegisterServer.c
@@ -37,7 +37,7 @@ HRESULT RegisterDll(LPWSTR path, BOOL state) {
 	return status == 0 ? S_OK : E_FAIL;
 }
 
-void RegisterServer(BOOL state) {
+HRESULT RegisterServer(BOOL state, BOOL forLaunch) {
 	// Ensure elevation
 	HRESULT hr = S_OK;
 	LPWSTR installPath;
@@ -96,5 +96,9 @@ end:
 		MsgBox(NULL, title, /*GetMessageForHresult(hr)*/ text, MB_OK | MB_ICONERROR);
 	}
 
-	PostQuitMessage(hr);
+	if (!forLaunch) {
+		PostQuitMessage(hr);
+	}
+
+	return hr;
 }
diff --git a/launcher/RegisterServer.h b/launcher/RegisterServer.h
new file mode 100644
index 0000000..549d586
--- /dev/null
+++ b/launcher/RegisterServer.h
@@ -0,0 +1,5 @@
+#pragma once
+
+#include <windows.h>
+
+HRESULT RegisterServer(BOOL state, BOOL forLaunch);