Replace text parsing and WMI calls in 'driverlogic'

1 files changed, 70 insertions, 171 deletions

diff --git a/windows/nsis-plugins/src/driverlogic/context.cpp b/windows/nsis-plugins/src/driverlogic/context.cpp
index a01fd29bac..d0ad9dfb67 100644
--- a/windows/nsis-plugins/src/driverlogic/context.cpp
+++ b/windows/nsis-plugins/src/driverlogic/context.cpp
@@ -1,243 +1,142 @@
 #include "stdafx.h"
 #include "context.h"
 #include <libcommon/string.h>
-#include <libcommon/wmi/connection.h>
-#include <libcommon/wmi/resultset.h>
-#include <libcommon/wmi/wmi.h>
+#include <libcommon/error.h>
 #include <log/log.h>
+#include <winsock2.h>
+#include <ws2ipdef.h>
+#include <iphlpapi.h>
+#include <windows.h>
 #include <vector>
+#include <list>
 #include <stdexcept>
-#include <algorithm>
-#include <memory>
-#include <sstream>
-
-using namespace common;
 
 namespace
 {
 
-std::vector<std::wstring> BlockToRows(const std::wstring &textBlock)
+std::set<Context::NetworkAdapter> GetAllAdapters()
 {
-	//
-	// This is such a hack :-(
-	//
-	// It only works because the tokenizer is greedy and because we don't care about
-	// empty lines for this usage.
-	//
-	return common::string::Tokenize(textBlock, L"\r\n");
-}
+	ULONG bufferSize = 0;
 
-void LogAllAdapters(wmi::Connection &connection)
-{
-	auto resultset = connection.query(L"SELECT * from Win32_NetworkAdapter");
-
-	struct NetworkAdapter
-	{
-		size_t interfaceIndex;
-		std::wstring manufacturer;
-		std::wstring name;
-		std::wstring pnpDeviceId;
-		std::wstring alias;
-	};
+	const ULONG flags = GAA_FLAG_SKIP_UNICAST | GAA_FLAG_SKIP_ANYCAST | GAA_FLAG_SKIP_MULTICAST | GAA_FLAG_SKIP_DNS_SERVER;
 
-	std::vector<NetworkAdapter> adapters;
+	auto status = GetAdaptersAddresses(AF_INET, flags, nullptr, nullptr, &bufferSize);
 
-	//
-	// Find all adapters and extract the most important data.
-	//
+	THROW_UNLESS(ERROR_BUFFER_OVERFLOW, status, "Probe for adapter listing buffer size");
 
-	auto StringOrNa = [](const _variant_t &variant)
-	{
-		if (VT_BSTR == V_VT(&variant))
-		{
-			return std::wstring(V_BSTR(&variant));
-		}
+	// Memory is cheap, this avoids a looping construct.
+	bufferSize *= 2;
 
-		return std::wstring(L"n/a");
-	};
+	std::vector<uint8_t> buffer(bufferSize);
 
-	while(resultset.advance())
-	{
-		auto interfaceIndex = wmi::WmiGetPropertyAlways(resultset.result(), L"InterfaceIndex");
-		auto manufacturer = wmi::WmiGetProperty(resultset.result(), L"Manufacturer");
-		auto name = wmi::WmiGetProperty(resultset.result(), L"Name");
-		auto pnpDeviceId = wmi::WmiGetProperty(resultset.result(), L"PNPDeviceID");
-		auto alias = wmi::WmiGetProperty(resultset.result(), L"NetConnectionID");
+	status = GetAdaptersAddresses(AF_INET, flags, nullptr,
+		reinterpret_cast<PIP_ADAPTER_ADDRESSES>(&buffer[0]), &bufferSize);
 
-		NetworkAdapter adapter;
+	THROW_UNLESS(ERROR_SUCCESS, status, "Retrieve adapter listing");
 
-		adapter.interfaceIndex = static_cast<size_t>(V_UI8(&interfaceIndex));
-		adapter.manufacturer = StringOrNa(manufacturer);
-		adapter.name = StringOrNa(name);
-		adapter.pnpDeviceId = StringOrNa(pnpDeviceId);
-		adapter.alias = StringOrNa(alias);
+	std::set<Context::NetworkAdapter> adapters;
 
-		adapters.emplace_back(adapter);
+	for (auto it = (PIP_ADAPTER_ADDRESSES)&buffer[0]; nullptr != it; it = it->Next)
+	{
+		adapters.emplace(Context::NetworkAdapter(common::string::ToWide(it->AdapterName),
+			it->Description, it->FriendlyName));
 	}
 
-	//
-	// Flatten the adapter information so we can log it more easily.
-	//
+	return adapters;
+}
 
-	std::vector<std::wstring> details;
+std::set<Context::NetworkAdapter> GetTapAdapters(const std::set<Context::NetworkAdapter> &adapters)
+{
+	std::set<Context::NetworkAdapter> tapAdapters;
 
 	for (const auto &adapter : adapters)
 	{
-		details.emplace_back(L"Adapter");
-
-		{
-			std::wstringstream ss;
+		static const wchar_t name[] = L"TAP-Windows Adapter V9";
 
-			ss << L"    InterfaceIndex: " << adapter.interfaceIndex;
+		//
+		// Compare partial name, because once you start having more TAP adapters
+		// they're named "TAP-Windows Adapter V9 #2" and so on.
+		//
 
-			details.emplace_back(ss.str());
+		if (0 == adapter.name.compare(0, _countof(name) - 1, name))
+		{
+			tapAdapters.insert(adapter);
 		}
-
-		details.emplace_back(std::wstring(L"    Manufacturer: ").append(adapter.manufacturer));
-		details.emplace_back(std::wstring(L"    Name: ").append(adapter.name));
-		details.emplace_back(std::wstring(L"    PnpDeviceId: ").append(adapter.pnpDeviceId));
-		details.emplace_back(std::wstring(L"    Alias: ").append(adapter.alias));
 	}
 
-	PluginLogWithDetails(L"Adapters known to WMI", details);
+	return tapAdapters;
 }
 
-std::wstring DoubleBackslashes(const std::wstring &str)
+void LogAdapters(const std::wstring &description, const std::set<Context::NetworkAdapter> &adapters)
 {
-	auto result(str);
+	//
+	// Flatten the information so we can log it more easily.
+	//
 
-	size_t offset = 0;
+	std::vector<std::wstring> details;
 
-	for (size_t index = 0; index < str.size(); ++index)
+	for (const auto &adapter : adapters)
 	{
-		if (L'\\' == str[index])
-		{
-			result.insert(index + offset, 1, L'\\');
-			++offset;
-		}
+		details.emplace_back(L"Adapter");
+
+		details.emplace_back(std::wstring(L"    Guid: ").append(adapter.guid));
+		details.emplace_back(std::wstring(L"    Name: ").append(adapter.name));
+		details.emplace_back(std::wstring(L"    Alias: ").append(adapter.alias));
 	}
 
-	return result;
+	PluginLogWithDetails(description, details);
 }
 
 } // anonymous namespace
 
-Context::Context()
-	: m_connection(wmi::Connection::Namespace::Cimv2)
+Context::BaselineStatus Context::establishBaseline()
 {
-}
+	m_baseline = GetAllAdapters();
 
-Context::BaselineStatus Context::establishBaseline(const std::wstring &textBlock)
-{
-	m_baseline = ParseVirtualNics(textBlock);
+	auto tapAdapters = GetTapAdapters(m_baseline);
 
-	if (m_baseline.empty())
+	if (tapAdapters.empty())
 	{
-		return BaselineStatus::NO_INTERFACES_PRESENT;
+		return BaselineStatus::NO_TAP_ADAPTERS_PRESENT;
 	}
 
-	for (const auto &nic : m_baseline)
+	for (const auto &adapter : tapAdapters)
 	{
-		if (0 == _wcsicmp(nic.alias.c_str(), L"mullvad"))
+		if (0 == _wcsicmp(adapter.alias.c_str(), L"mullvad"))
 		{
-			return BaselineStatus::MULLVAD_INTERFACE_PRESENT;
+			return BaselineStatus::MULLVAD_ADAPTER_PRESENT;
 		}
 	}
 
-	return BaselineStatus::SOME_INTERFACES_PRESENT;
+	return BaselineStatus::SOME_TAP_ADAPTERS_PRESENT;
 }
 
-void Context::recordCurrentState(const std::wstring &textBlock)
+void Context::recordCurrentState()
 {
-	m_currentState = ParseVirtualNics(textBlock);
-}
-
-Context::VirtualNic Context::getNewAdapter()
-{
-	std::vector<VirtualNic> added;
-
-	for (const auto &nic : m_currentState)
-	{
-		if (m_baseline.end() == m_baseline.find(nic))
-		{
-			added.push_back(nic);
-		}
-	}
-
-	if (added.size() != 1)
-	{
-		throw std::runtime_error("Unable to identify newly added virtual adapter");
-	}
-
-	return added[0];
+	m_currentState = GetAllAdapters();
 }
 
-std::set<Context::VirtualNic> Context::ParseVirtualNics(const std::wstring &textBlock)
+Context::NetworkAdapter Context::getNewAdapter()
 {
-	// ROOT\NET\0000
-	//     Name: TAP - Windows Adapter V9
-	//     Hardware IDs :
-	//         tap0901
-	// 1 matching device(s) found.
-
-	std::set<VirtualNic> nics;
-
-	auto text = BlockToRows(textBlock);
+	std::list<NetworkAdapter> added;
 
-	size_t line = 0;
+	const auto baselineTaps = GetTapAdapters(m_baseline);
+	const auto currentTaps = GetTapAdapters(m_currentState);
 
-	while (nullptr != wcschr(text.at(line).c_str(), L'\\'))
+	for (const auto &adapter : currentTaps)
 	{
-		auto nameDelimiter = wcschr(text.at(line + 1).c_str(), L':');
-
-		if (nullptr == nameDelimiter)
+		if (baselineTaps.end() == baselineTaps.find(adapter))
 		{
-			throw std::runtime_error("Unexpected formatting in input data");
+			added.push_back(adapter);
 		}
-
-		VirtualNic nic;
-
-		nic.node = text.at(line);
-		nic.name = std::wstring(nameDelimiter + 2);
-		nic.alias = GetNicAlias(nic.node, nic.name);
-
-		nics.emplace(std::move(nic));
-		line += 4;
 	}
 
-	return nics;
-}
-
-std::wstring Context::GetNicAlias(const std::wstring &node, const std::wstring &name)
-{
-	//
-	// The name cannot be used when querying WMI, because WMI sometimes normalizes the
-	// names in its dataset, thereby destroying their uniqueness.
-	//
-	// E.g. if a network interface has a name of "TAP-Windows Adapter V9 #2" it will
-	// sometimes be reported by WMI as "TAP-Windows Adapter V9".
-	//
-	// Also, the node string cannot be used as-is. We have to double the backslashes in it
-	// or the string will be rejected by WMI.
-	//
-
-	const auto formattedNode = DoubleBackslashes(node);
-
-	std::wstringstream ss;
-
-	ss << L"SELECT * FROM Win32_NetworkAdapter WHERE PNPDeviceID = \"" << formattedNode << L"\"";
-
-	auto resultset = m_connection.query(ss.str().c_str());
-
-	if (false == resultset.advance())
+	if (added.size() != 1)
 	{
-		PluginLog(std::wstring(L"WMI query failed for adapter: ").append(name));
-		LogAllAdapters(m_connection);
+		LogAdapters(L"Enumerable network adapters", m_currentState);
 
-		throw std::runtime_error("Unable to look up virtual adapter using WMI");
+		throw std::runtime_error("Unable to identify recently added TAP adapter");
 	}
 
-	auto alias = wmi::WmiGetPropertyAlways(resultset.result(), L"NetConnectionID");
-
-	return V_BSTR(&alias);
+	return *added.begin();
 }