From 5e1632e2cca5b8d537bc66525d631ebd72f6949b Mon Sep 17 00:00:00 2001
From: Akihiko Odaki <nekomanma@pixiv.co.jp>
Date: Wed, 26 Sep 2018 19:34:19 +0900
Subject: [PATCH] Add Visual Studio detection by the COM API

The new COM API is expected to work on Visual Studio 2017 and newer
installations. It is also compatible with non-x86 runtime.
---
 .gitignore          |  1 +
 src/cv2pdb.vcxproj  | 10 ++++++++++
 src/main.cpp        |  2 ++
 src/mspdb.cpp       | 38 ++++++++++++++++++++++++++++++++++++++
 src/packages.config |  4 ++++
 5 files changed, 55 insertions(+)
 create mode 100644 src/packages.config

diff --git a/.gitignore b/.gitignore
index 8b430f3..ead68ad 100644
--- a/.gitignore
+++ b/.gitignore
@@ -5,3 +5,4 @@
 *.user
 /src/.vs
 /src/*.opensdf
+/src/packages
diff --git a/src/cv2pdb.vcxproj b/src/cv2pdb.vcxproj
index d79bdc2..af3b87b 100644
--- a/src/cv2pdb.vcxproj
+++ b/src/cv2pdb.vcxproj
@@ -208,8 +208,18 @@
       <ExcludedFromBuild Condition="'$(Configuration)|$(Platform)'=='Release|Win32'">true</ExcludedFromBuild>
     </MASM>
   </ItemGroup>
+  <ItemGroup>
+    <None Include="packages.config" />
+  </ItemGroup>
   <Import Project="$(VCTargetsPath)\Microsoft.Cpp.targets" />
   <ImportGroup Label="ExtensionTargets">
     <Import Project="$(VCTargetsPath)\BuildCustomizations\masm.targets" />
+    <Import Project="packages\Microsoft.VisualStudio.Setup.Configuration.Native.1.16.30\build\native\Microsoft.VisualStudio.Setup.Configuration.Native.targets" Condition="Exists('packages\Microsoft.VisualStudio.Setup.Configuration.Native.1.16.30\build\native\Microsoft.VisualStudio.Setup.Configuration.Native.targets')" />
   </ImportGroup>
+  <Target Name="EnsureNuGetPackageBuildImports" BeforeTargets="PrepareForBuild">
+    <PropertyGroup>
+      <ErrorText>This project references NuGet package(s) that are missing on this computer. Enable NuGet Package Restore to download them.  For more information, see http://go.microsoft.com/fwlink/?LinkID=322105. The missing file is {0}.</ErrorText>
+    </PropertyGroup>
+    <Error Condition="!Exists('packages\Microsoft.VisualStudio.Setup.Configuration.Native.1.16.30\build\native\Microsoft.VisualStudio.Setup.Configuration.Native.targets')" Text="$([System.String]::Format('$(ErrorText)', 'packages\Microsoft.VisualStudio.Setup.Configuration.Native.1.16.30\build\native\Microsoft.VisualStudio.Setup.Configuration.Native.targets'))" />
+  </Target>
 </Project>
\ No newline at end of file
diff --git a/src/main.cpp b/src/main.cpp
index 70b1563..608f197 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -122,6 +122,8 @@ int T_main(int argc, TCHAR* argv[])
 	const TCHAR* pdbref = 0;
 	bool debug = false;
 
+	CoInitialize(nullptr);
+
 	while (argc > 1 && argv[1][0] == '-')
 	{
 		argv++;
diff --git a/src/mspdb.cpp b/src/mspdb.cpp
index cb3e812..59c994a 100644
--- a/src/mspdb.cpp
+++ b/src/mspdb.cpp
@@ -6,7 +6,14 @@
 
 #include "mspdb.h"
 
+#include <atlstr.h>
+#include <comdef.h>
 #include <windows.h>
+#include <Setup.Configuration.h>
+
+_COM_SMARTPTR_TYPEDEF(ISetupConfiguration, __uuidof(ISetupConfiguration));
+_COM_SMARTPTR_TYPEDEF(ISetupInstance, __uuidof(ISetupInstance));
+_COM_SMARTPTR_TYPEDEF(IEnumSetupInstances, __uuidof(IEnumSetupInstances));
 
 #pragma comment(lib, "rpcrt4.lib")
 
@@ -81,6 +88,35 @@ bool tryLoadMsPdb(const char* version, const char* mspdb, const char* path = 0)
 	return modMsPdb != 0;
 }
 
+bool tryLoadMsPdbCom(const char* mspdb, const char* path = 0)
+{
+	ISetupConfigurationPtr query;
+	ISetupInstancePtr instance;
+	IEnumSetupInstancesPtr instances;
+	BSTR installDir;
+	unsigned long fetched;
+
+	auto result = query.CreateInstance(__uuidof(SetupConfiguration));
+	if ((FAILED(result) && result != REGDB_E_CLASSNOTREG) || FAILED(query->EnumInstances(&instances)))
+		return false;
+
+	while (!modMsPdb)
+	{
+		if (FAILED(instances->Next(1, &instance, &fetched)) || !fetched)
+			return false;
+
+		if (FAILED(instance->GetInstallationPath(&installDir)))
+			continue;
+
+		CStringA modpath = installDir;
+		modpath += "\\Common7\\IDE\\";
+		modpath += mspdb;
+		tryLoadLibrary(modpath);
+	}
+
+	return true;
+}
+
 bool tryLoadMsPdbVS2017(const char* mspdb, const char* path = 0)
 {
 	const char* key = "SOFTWARE\\Microsoft\\VisualStudio\\SxS\\VS7";
@@ -177,6 +213,8 @@ void tryLoadMsPdb140(bool throughPath)
 	{
 		if(throughPath)
 			modMsPdb = LoadLibraryA(mspdb140_dll);
+		if(!modMsPdb && !throughPath)
+			tryLoadMsPdbCom(mspdb140_dll);
 		if(!modMsPdb && !throughPath)
 			tryLoadMsPdbVS2017(mspdb140_dll);
 		if (!modMsPdb && !throughPath)
diff --git a/src/packages.config b/src/packages.config
new file mode 100644
index 0000000..8bfdb5a
--- /dev/null
+++ b/src/packages.config
@@ -0,0 +1,4 @@
+<?xml version="1.0" encoding="utf-8"?>
+<packages>
+  <package id="Microsoft.VisualStudio.Setup.Configuration.Native" version="1.16.30" targetFramework="native" developmentDependency="true" />
+</packages>
\ No newline at end of file