2 files changed, 73 insertions, 33 deletions

diff --git a/windows/winfw/src/winfw/objectpurger.cpp b/windows/winfw/src/winfw/objectpurger.cpp
index dce36c99c8..7baf0c4a22 100644
--- a/windows/winfw/src/winfw/objectpurger.cpp
+++ b/windows/winfw/src/winfw/objectpurger.cpp
@@ -5,6 +5,8 @@
 #include "libwfp/filterengine.h"
 #include "libwfp/objectdeleter.h"
 #include "libwfp/transaction.h"
+#include "libwfp/objectenumerator.h"
+#include <set>
 #include <algorithm>
 
 namespace
@@ -12,45 +14,59 @@ namespace
 
 using ObjectDeleter = std::function<void(wfp::FilterEngine &, const GUID &)>;
 
-template<typename TRange>
-void RemoveRange(wfp::FilterEngine &engine, ObjectDeleter deleter, TRange range)
+template<typename T>
+bool HasMullvadProvider(T obj)
 {
-	std::for_each(range.first, range.second, [&](const auto &record)
-	{
-		const GUID &objectId = record.second;
-		deleter(engine, objectId);
-	});
+	return nullptr != obj.providerKey && *obj.providerKey == MullvadGuids::Provider();
 }
 
-} // anonymous namespace
-
-//static
-ObjectPurger::RemovalFunctor ObjectPurger::GetRemoveFiltersFunctor()
+template<typename T>
+bool HasPersistentMullvadProvider(const T &obj)
 {
-	return [](wfp::FilterEngine &engine)
-	{
-		const auto registry = MullvadGuids::DetailedRegistry(MullvadGuids::IdentityQualifier::IncludeAll);
-
-		// Resolve correct overload.
-		void (*deleter)(wfp::FilterEngine &, const GUID &) = wfp::ObjectDeleter::DeleteFilter;
-
-		RemoveRange(engine, deleter, registry.equal_range(WfpObjectType::Filter));
-	};
+	return nullptr != obj.providerKey && *obj.providerKey == MullvadGuids::ProviderPersistent();
 }
 
+} // anonymous namespace
+
 //static
 ObjectPurger::RemovalFunctor ObjectPurger::GetRemoveAllFunctor()
 {
 	return [](wfp::FilterEngine &engine)
 	{
-		const auto registry = MullvadGuids::DetailedRegistry(MullvadGuids::IdentityQualifier::IncludeAll);
+		std::unordered_set<GUID> filtersToRemove;
+		wfp::ObjectEnumerator::Filters(engine, [&](const auto &filter) -> bool
+		{
+			// Delete both non-persistent and persistent filters
+			if (HasMullvadProvider(filter) || HasPersistentMullvadProvider(filter))
+			{
+				filtersToRemove.insert(filter.filterKey);
+			}
+			return true;
+		});
+
+		std::unordered_set<GUID> sublayersToRemove;
+		wfp::ObjectEnumerator::Sublayers(engine, [&](const auto &sublayer) -> bool
+		{
+			// Delete both non-persistent and persistent sublayers
+			if (HasMullvadProvider(sublayer) || HasPersistentMullvadProvider(sublayer))
+			{
+				sublayersToRemove.insert(sublayer.subLayerKey);
+			}
+			return true;
+		});
 
-		// Resolve correct overload.
-		void(*deleter)(wfp::FilterEngine &, const GUID &) = wfp::ObjectDeleter::DeleteFilter;
+		for (const auto &filter : filtersToRemove)
+		{
+			wfp::ObjectDeleter::DeleteFilter(engine, filter);
+		}
 
-		RemoveRange(engine, deleter, registry.equal_range(WfpObjectType::Filter));
-		RemoveRange(engine, wfp::ObjectDeleter::DeleteSublayer, registry.equal_range(WfpObjectType::Sublayer));
-		RemoveRange(engine, wfp::ObjectDeleter::DeleteProvider, registry.equal_range(WfpObjectType::Provider));
+		for (const auto &sublayer : sublayersToRemove)
+		{
+			wfp::ObjectDeleter::DeleteSublayer(engine, sublayer);
+		}
+
+		wfp::ObjectDeleter::DeleteProvider(engine, MullvadGuids::Provider());
+		wfp::ObjectDeleter::DeleteProvider(engine, MullvadGuids::ProviderPersistent());
 	};
 }
 
@@ -59,14 +75,39 @@ ObjectPurger::RemovalFunctor ObjectPurger::GetRemoveNonPersistentFunctor()
 {
 	return [](wfp::FilterEngine &engine)
 	{
-		const auto registry = MullvadGuids::DetailedRegistry(MullvadGuids::IdentityQualifier::IncludeDeprecated);
+		std::unordered_set<GUID> filtersToRemove;
+		wfp::ObjectEnumerator::Filters(engine, [&](const auto &filter) -> bool
+		{
+			// Delete only non-persistent filters
+			if (HasMullvadProvider(filter))
+			{
+				filtersToRemove.insert(filter.filterKey);
+			}
+			return true;
+		});
+
+		std::unordered_set<GUID> sublayersToRemove;
+		wfp::ObjectEnumerator::Sublayers(engine, [&](const auto &sublayer) -> bool
+		{
+			// Delete only non-persistent sublayers
+			if (HasMullvadProvider(sublayer))
+			{
+				sublayersToRemove.insert(sublayer.subLayerKey);
+			}
+			return true;
+		});
+
+		for (const auto &filter : filtersToRemove)
+		{
+			wfp::ObjectDeleter::DeleteFilter(engine, filter);
+		}
 
-		// Resolve correct overload.
-		void(*deleter)(wfp::FilterEngine &, const GUID &) = wfp::ObjectDeleter::DeleteFilter;
+		for (const auto &sublayer : sublayersToRemove)
+		{
+			wfp::ObjectDeleter::DeleteSublayer(engine, sublayer);
+		}
 
-		RemoveRange(engine, deleter, registry.equal_range(WfpObjectType::Filter));
-		RemoveRange(engine, wfp::ObjectDeleter::DeleteSublayer, registry.equal_range(WfpObjectType::Sublayer));
-		RemoveRange(engine, wfp::ObjectDeleter::DeleteProvider, registry.equal_range(WfpObjectType::Provider));
+		wfp::ObjectDeleter::DeleteProvider(engine, MullvadGuids::Provider());
 	};
 }
 
diff --git a/windows/winfw/src/winfw/objectpurger.h b/windows/winfw/src/winfw/objectpurger.h
index 7728aac694..91ed61ddaf 100644
--- a/windows/winfw/src/winfw/objectpurger.h
+++ b/windows/winfw/src/winfw/objectpurger.h
@@ -13,7 +13,6 @@ public:
 
 	using RemovalFunctor = std::function<void(wfp::FilterEngine &engine)>;
 
-	static RemovalFunctor GetRemoveFiltersFunctor();
 	static RemovalFunctor GetRemoveAllFunctor();
 	static RemovalFunctor GetRemoveNonPersistentFunctor();