Add support for WMI event notification

6 files changed, 305 insertions, 0 deletions

diff --git a/windns/src/windns/netconfigeventsink.cpp b/windns/src/windns/netconfigeventsink.cpp
new file mode 100644
index 0000000000..6a506e3f3f
--- /dev/null
+++ b/windns/src/windns/netconfigeventsink.cpp
@@ -0,0 +1,31 @@
+#include "stdafx.h"
+#include "netconfigeventsink.h"
+#include "windns/netconfighelpers.h"
+
+NetConfigEventSink::NetConfigEventSink(std::shared_ptr<wmi::IConnection> connection, std::shared_ptr<ConfigManager> configManager)
+	: m_connection(connection)
+	, m_configManager(configManager)
+{
+}
+
+void NetConfigEventSink::update(CComPtr<IWbemClassObject> instance)
+{
+	DnsConfig config(instance);
+
+	ConfigManager::Mutex mutex(*m_configManager);
+
+	//
+	// This is OK because the config manager will reject updates
+	// that set our DNS servers.
+	//
+	auto updated = m_configManager->updateConfig(std::move(config));
+
+	if (updated)
+	{
+		//
+		// Override current settings to use our DNS servers.
+		//
+		auto servers = m_configManager->getServers();
+		nchelpers::SetDnsServers(*m_connection, instance, &servers);
+	}
+}
diff --git a/windns/src/windns/netconfigeventsink.h b/windns/src/windns/netconfigeventsink.h
new file mode 100644
index 0000000000..35ab829a08
--- /dev/null
+++ b/windns/src/windns/netconfigeventsink.h
@@ -0,0 +1,20 @@
+#pragma once
+
+#include "windns/wmi/eventsink.h"
+#include "windns/wmi/iconnection.h"
+#include "windns/configmanager.h"
+#include <memory>
+
+class NetConfigEventSink : public wmi::IEventSink
+{
+public:
+
+	explicit NetConfigEventSink(std::shared_ptr<wmi::IConnection> connection, std::shared_ptr<ConfigManager> configManager);
+
+	void update(CComPtr<IWbemClassObject> instance) override;
+
+private:
+
+	std::shared_ptr<wmi::IConnection> m_connection;
+	std::shared_ptr<ConfigManager> m_configManager;
+};
diff --git a/windns/src/windns/wmi/eventsink.cpp b/windns/src/windns/wmi/eventsink.cpp
new file mode 100644
index 0000000000..516660ff0a
--- /dev/null
+++ b/windns/src/windns/wmi/eventsink.cpp
@@ -0,0 +1,93 @@
+#include "stdafx.h"
+#include "eventsink.h"
+#include "windns/comhelpers.h"
+
+namespace wmi
+{
+
+EventSink::EventSink(std::shared_ptr<IEventSink> eventSink)
+	: m_references(0)
+	, m_callbacks(0)
+	, m_eventSink(eventSink)
+{
+}
+
+EventSink::~EventSink()
+{
+}
+
+bool EventSink::processing() const
+{
+	//
+	// Cancelling the notification registration WILL NOT wait for the completion of
+	// callbacks currently in progress :-(
+	//
+	// (Observed on Win10)
+	//
+
+	return 0 != InterlockedAdd(&m_callbacks, 0);
+}
+
+ULONG STDMETHODCALLTYPE EventSink::AddRef()
+{
+	return InterlockedIncrement(&m_references);
+}
+
+ULONG STDMETHODCALLTYPE EventSink::Release()
+{
+	auto refs = InterlockedDecrement(&m_references);
+
+	if (refs == 0)
+	{
+		delete this;
+	}
+
+	return refs;
+}
+
+HRESULT STDMETHODCALLTYPE EventSink::QueryInterface(REFIID riid, void **ppv)
+{
+	if (IID_IUnknown == riid || IID_IWbemObjectSink == riid)
+	{
+		*ppv = (IWbemObjectSink *)this;
+		AddRef();
+
+		return WBEM_S_NO_ERROR;
+	}
+
+	return E_NOINTERFACE;
+}
+
+HRESULT STDMETHODCALLTYPE EventSink::Indicate
+(
+	LONG numObjects,
+	IWbemClassObject __RPC_FAR *__RPC_FAR *objects
+)
+{
+	InterlockedIncrement(&m_callbacks);
+
+	for (LONG i = 0; i < numObjects; ++i)
+	{
+		CComPtr<IWbemClassObject> eventRecord(objects[i]);
+
+		auto rawTarget = ComGetPropertyAlways(eventRecord, L"TargetInstance");
+
+		CComQIPtr<IWbemClassObject> target(V_UNKNOWN(&rawTarget));
+
+		m_eventSink->update(target);
+	}
+
+	InterlockedDecrement(&m_callbacks);
+
+	return WBEM_S_NO_ERROR;
+}
+
+HRESULT STDMETHODCALLTYPE EventSink::SetStatus
+(
+	LONG, HRESULT, BSTR, IWbemClassObject __RPC_FAR *
+)
+{
+	return WBEM_S_NO_ERROR;
+}
+
+}
diff --git a/windns/src/windns/wmi/eventsink.h b/windns/src/windns/wmi/eventsink.h
new file mode 100644
index 0000000000..2f0783d81f
--- /dev/null
+++ b/windns/src/windns/wmi/eventsink.h
@@ -0,0 +1,55 @@
+#pragma once
+
+#include <memory>
+#include <atlbase.h>
+#include <wbemidl.h>
+
+namespace wmi
+{
+
+struct IEventSink
+{
+	virtual ~IEventSink() = 0
+	{
+	}
+
+	virtual void update(CComPtr<IWbemClassObject> instance) = 0;
+};
+
+class EventSink : public IWbemObjectSink
+{
+public:
+
+	EventSink(std::shared_ptr<IEventSink> eventSink);
+	~EventSink();
+
+	bool processing() const;
+
+	ULONG STDMETHODCALLTYPE AddRef() override;
+	ULONG STDMETHODCALLTYPE Release() override;
+	HRESULT STDMETHODCALLTYPE QueryInterface(REFIID riid, void **ppv) override;
+
+	HRESULT STDMETHODCALLTYPE Indicate
+	(
+		LONG numObjects,
+		IWbemClassObject __RPC_FAR *__RPC_FAR *objects
+	)
+	override;
+
+	HRESULT STDMETHODCALLTYPE SetStatus
+	(
+		LONG flags,
+		HRESULT result,
+		BSTR param,
+		IWbemClassObject __RPC_FAR *object
+	)
+	override;
+
+private:
+
+	LONG m_references;
+	mutable LONG m_callbacks;
+	std::shared_ptr<IEventSink> m_eventSink;
+};
+
+}
diff --git a/windns/src/windns/wmi/notification.cpp b/windns/src/windns/wmi/notification.cpp
new file mode 100644
index 0000000000..a9427f6520
--- /dev/null
+++ b/windns/src/windns/wmi/notification.cpp
@@ -0,0 +1,72 @@
+#include "stdafx.h"
+#include "notification.h"
+#include "windns/comhelpers.h"
+
+namespace wmi
+{
+
+Notification::Notification(std::shared_ptr<IConnection> connection, CComPtr<EventSink> eventSink)
+	: m_connection(connection)
+	, m_eventSink(eventSink)
+{
+}
+
+Notification::~Notification()
+{
+	//
+	// TODO: Revise to avoid exceptions in dtor.
+	//
+	deactivate();
+}
+
+void Notification::activate(const std::wstring &query)
+{
+	CComPtr<IUnsecuredApartment> apartment;
+
+	auto status = CoCreateInstance(CLSID_UnsecuredApartment, nullptr, CLSCTX_LOCAL_SERVER, IID_IUnsecuredApartment, (void**)&apartment);
+	VALIDATE_COM(status, "Create unsecured COM apartment");
+
+	CComPtr<IUnknown> unknownEventSink;
+
+	status = m_eventSink->QueryInterface(IID_IUnknown, (void**)&unknownEventSink);
+	VALIDATE_COM(status, "Retrieve IUnkown interface for event sink");
+
+	CComPtr<IUnknown> unknownForwarder;
+
+	status = apartment->CreateObjectStub(unknownEventSink, &unknownForwarder);
+	VALIDATE_COM(status, "Create forwarder for event sink");
+
+	status = unknownForwarder->QueryInterface(IID_IWbemObjectSink, (void**)&m_forwarder);
+	VALIDATE_COM(status, "Retrieve sink interface on event sink forwarder");
+
+	status = m_connection->services()->ExecNotificationQueryAsync(_bstr_t("WQL"), _bstr_t(query.c_str()), 0, nullptr, m_forwarder);
+	VALIDATE_COM(status, "Register notification query with WMI");
+}
+
+void Notification::deactivate()
+{
+	if (nullptr == m_forwarder)
+	{
+		return;
+	}
+
+	auto status = m_connection->services()->CancelAsyncCall(m_forwarder);
+	VALIDATE_COM(status, "Cancel notification query");
+
+	m_forwarder.Release();
+
+	//
+	// This is a hack-solution for a corner case issue.
+	//
+	// Since cancelling the notification registration does not wait for in-progress callbacks
+	// to complete, we have to implement the corresponding logic ourselves.
+	//
+	// Using a Sleep() here is preferable to introducing a critical section in the callback.
+	//
+	while (m_eventSink->processing())
+	{
+		Sleep(100);
+	}
+}
+
+}
diff --git a/windns/src/windns/wmi/notification.h b/windns/src/windns/wmi/notification.h
new file mode 100644
index 0000000000..40a57f1e57
--- /dev/null
+++ b/windns/src/windns/wmi/notification.h
@@ -0,0 +1,34 @@
+#pragma once
+
+#include "iconnection.h"
+#include "eventsink.h"
+#include <atlbase.h>
+#include <memory>
+
+namespace wmi
+{
+
+class Notification
+{
+public:
+
+	Notification(std::shared_ptr<IConnection> connection, CComPtr<EventSink> eventSink);
+	~Notification();
+
+	Notification(const Notification &) = delete;
+	Notification &operator=(const Notification &) = delete;
+	Notification(Notification &&) = delete;
+	Notification &operator=(Notification &&) = delete;
+
+	void activate(const std::wstring &query);
+	void deactivate();
+
+private:
+
+	std::shared_ptr<IConnection> m_connection;
+	CComPtr<EventSink> m_eventSink;
+
+	CComPtr<IWbemObjectSink> m_forwarder;
+};
+
+}