Update RestrictDns in WinFw

5 files changed, 26 insertions, 23 deletions

diff --git a/windows/winfw/src/winfw/fwcontext.cpp b/windows/winfw/src/winfw/fwcontext.cpp
index 1c9c78c526..610936de0b 100644
--- a/windows/winfw/src/winfw/fwcontext.cpp
+++ b/windows/winfw/src/winfw/fwcontext.cpp
@@ -140,9 +140,9 @@ bool FwContext::applyPolicyConnected
 (
 	const WinFwSettings &settings,
 	const WinFwRelay &relay,
-	const wchar_t *tunnelInterfaceAlias,
-	const wchar_t *v4DnsHost,
-	const wchar_t *v6DnsHost
+	const std::wstring &tunnelInterfaceAlias,
+	const wfp::IpAddress &v4DnsHost,
+	const std::optional<wfp::IpAddress> &v6DnsHost
 )
 {
 	Ruleset ruleset;
@@ -166,10 +166,9 @@ bool FwContext::applyPolicyConnected
 
 	ruleset.emplace_back(std::make_unique<rules::RestrictDns>(
 		tunnelInterfaceAlias,
-		wfp::IpAddress(std::wstring(v4DnsHost)),
-		nullptr != v6DnsHost ? std::make_optional<wfp::IpAddress>(std::wstring(v6DnsHost)) : std::nullopt,
-		std::wstring(relay.ip),
-		relay.port
+		v4DnsHost,
+		v6DnsHost,
+		53 == relay.port ? std::make_optional(wfp::IpAddress(relay.ip)) : std::nullopt
 	));
 
 	return applyRuleset(ruleset);
diff --git a/windows/winfw/src/winfw/fwcontext.h b/windows/winfw/src/winfw/fwcontext.h
index 9d5b34c51b..552b075869 100644
--- a/windows/winfw/src/winfw/fwcontext.h
+++ b/windows/winfw/src/winfw/fwcontext.h
@@ -35,9 +35,9 @@ public:
 	(
 		const WinFwSettings &settings,
 		const WinFwRelay &relay,
-		const wchar_t *tunnelInterfaceAlias,
-		const wchar_t *v4DnsHost,
-		const wchar_t *v6DnsHost
+		const std::wstring &tunnelInterfaceAlias,
+		const wfp::IpAddress &v4DnsHost,
+		const std::optional<wfp::IpAddress> &v6DnsHost
 	);
 	bool applyPolicyBlocked(const WinFwSettings &settings);
 
diff --git a/windows/winfw/src/winfw/rules/restrictdns.cpp b/windows/winfw/src/winfw/rules/restrictdns.cpp
index efa4c8421b..e1c81bfd53 100644
--- a/windows/winfw/src/winfw/rules/restrictdns.cpp
+++ b/windows/winfw/src/winfw/rules/restrictdns.cpp
@@ -12,17 +12,16 @@ using namespace wfp::conditions;
 namespace rules
 {
 
-RestrictDns::RestrictDns(const std::wstring &tunnelInterfaceAlias,
+RestrictDns::RestrictDns(
+	const std::wstring &tunnelInterfaceAlias,
 	const wfp::IpAddress v4DnsHost,
 	std::optional<wfp::IpAddress> v6DnsHost,
-	wfp::IpAddress relay,
-	uint16_t relayPort)
+	std::optional<wfp::IpAddress> allowHost
+)
 	: m_tunnelInterfaceAlias(tunnelInterfaceAlias)
 	, m_v4DnsHost(v4DnsHost)
 	, m_v6DnsHost(v6DnsHost)
-	, m_relayHost(relay)
-	, m_relayPort(relayPort)
-
+	, m_allowHost(allowHost)
 {
 }
 
@@ -72,12 +71,12 @@ bool RestrictDns::apply(IObjectInstaller &objectInstaller)
 		wfp::ConditionBuilder conditionBuilder(FWPM_LAYER_ALE_AUTH_CONNECT_V4);
 		conditionBuilder.add_condition(ConditionPort::Remote(53));
 
-		if (53 == m_relayPort)
+		if (m_allowHost.has_value())
 		{
 			//
-			// Allow relay traffic over port 53
+			// Allow DNS traffic over select host
 			//
-			conditionBuilder.add_condition(ConditionIp::Remote(m_relayHost, CompareNeq()));
+			conditionBuilder.add_condition(ConditionIp::Remote(*m_allowHost, CompareNeq()));
 		}
 
 		if (!objectInstaller.addFilter(filterBuilder, conditionBuilder))
diff --git a/windows/winfw/src/winfw/rules/restrictdns.h b/windows/winfw/src/winfw/rules/restrictdns.h
index cdfe1f4697..0d730e4a32 100644
--- a/windows/winfw/src/winfw/rules/restrictdns.h
+++ b/windows/winfw/src/winfw/rules/restrictdns.h
@@ -13,7 +13,7 @@ class RestrictDns : public IFirewallRule
 {
 public:
 
-	RestrictDns(const std::wstring &tunnelInterfaceAlias, const wfp::IpAddress v4DnsHost, std::optional<wfp::IpAddress> v6DnsHost, wfp::IpAddress relay, uint16_t relayPort);
+	RestrictDns(const std::wstring &tunnelInterfaceAlias, const wfp::IpAddress v4DnsHost, std::optional<wfp::IpAddress> v6DnsHost, std::optional<wfp::IpAddress> allowHost);
 
 	bool apply(IObjectInstaller &objectInstaller) override;
 
@@ -22,8 +22,7 @@ private:
 	const std::wstring m_tunnelInterfaceAlias;
 	const wfp::IpAddress m_v4DnsHost;
 	const std::optional<wfp::IpAddress> m_v6DnsHost;
-	const uint16_t m_relayPort;
-	const wfp::IpAddress m_relayHost;
+	const std::optional<wfp::IpAddress> m_allowHost;
 
 };
 
diff --git a/windows/winfw/src/winfw/winfw.cpp b/windows/winfw/src/winfw/winfw.cpp
index bb79c2eff6..c053087355 100644
--- a/windows/winfw/src/winfw/winfw.cpp
+++ b/windows/winfw/src/winfw/winfw.cpp
@@ -205,7 +205,13 @@ WinFw_ApplyPolicyConnected(
 
 	try
 	{
-		return g_fwContext->applyPolicyConnected(settings, relay, tunnelInterfaceAlias, v4DnsHost, v6DnsHost);
+		return g_fwContext->applyPolicyConnected(
+			settings,
+			relay,
+			tunnelInterfaceAlias,
+			wfp::IpAddress(v4DnsHost),
+			nullptr != v6DnsHost ? std::make_optional(wfp::IpAddress(v6DnsHost)) : std::nullopt
+		);
 	}
 	catch (std::exception &err)
 	{