From d3814ca945face2daae004dfda1aacd9163ff4c0 Mon Sep 17 00:00:00 2001
From: Connection Refused <cutecr@pm.me>
Date: Wed, 10 Nov 2021 14:09:38 +0800
Subject: [PATCH] [RouteHelper] Add WaitForUnicastIP

---
 RouteHelper/Based.h                     | 15 ++++++++++
 RouteHelper/README.md                   |  8 ++---
 RouteHelper/RouteHelper.cpp             | 40 +++++++++++++++----------
 RouteHelper/RouteHelper.vcxproj         |  5 ++++
 RouteHelper/RouteHelper.vcxproj.filters | 11 +++++++
 RouteHelper/WaitGroup.cpp               | 19 ++++++++++++
 RouteHelper/WaitGroup.h                 | 18 +++++++++++
 build.ps1                               |  4 +--
 8 files changed, 98 insertions(+), 22 deletions(-)
 create mode 100644 RouteHelper/Based.h
 create mode 100644 RouteHelper/WaitGroup.cpp
 create mode 100644 RouteHelper/WaitGroup.h

diff --git a/RouteHelper/Based.h b/RouteHelper/Based.h
new file mode 100644
index 00000000..c7938101
--- /dev/null
+++ b/RouteHelper/Based.h
@@ -0,0 +1,15 @@
+#pragma once
+#ifndef BASED_H
+#define BASED_H
+#include <mutex>
+#include <atomic>
+#include <shared_mutex>
+#include <condition_variable>
+
+#include <WinSock2.h>
+#include <WS2tcpip.h>
+#include <ws2ipdef.h>
+#include <iphlpapi.h>
+#include <netioapi.h>
+#include <Windows.h>
+#endif
diff --git a/RouteHelper/README.md b/RouteHelper/README.md
index 138307cb..419d6a65 100644
--- a/RouteHelper/README.md
+++ b/RouteHelper/README.md
@@ -1,9 +1,9 @@
 # RouteHelper
 ```cpp
 ULONG ConvertLuidToIndex(ULONG64 id);
+void WaitForUnicastIP();
 BOOL CreateIPv4(const char* address, const char* netmask, ULONG index);
 BOOL CreateUnicastIP(USHORT inet, const char* address, UINT8 cidr, ULONG index);
-BOOL RefreshIPTable(USHORT inet, ULONG index);
 BOOL CreateRoute(USHORT inet, const char* address, UINT8 cidr, const char* gateway, ULONG index, ULONG metric);
 BOOL DeleteRoute(USHORT inet, const char* address, UINT8 cidr, const char* gateway, ULONG index, ULONG metric);
 ```
@@ -12,15 +12,15 @@ BOOL DeleteRoute(USHORT inet, const char* address, UINT8 cidr, const char* gatew
 [DllImport("RouteHelper.bin", CallingConvention = CallingConvention.Cdecl)]
 public static extern uint ConvertLuidToIndex(ulong id);
 
+[DllImport("RouteHelper.bin", CallingConvention = CallingConvention.Cdecl)]
+public static extern void WaitForUnicastIP();
+
 [DllImport("RouteHelper.bin", CallingConvention = CallingConvention.Cdecl)]
 public static extern bool CreateIPv4(string address, string netmask, uint index);
 
 [DllImport("RouteHelper.bin", CallingConvention = CallingConvention.Cdecl)]
 public static extern bool CreateUnicastIP(AddressFamily inet, string address, byte cidr, uint index);
 
-[DllImport("RouteHelper.bin", CallingConvention = CallingConvention.Cdecl)]
-public static extern bool RefreshIPTable(AddressFamily inet, uint index);
-
 [DllImport("RouteHelper.bin", CallingConvention = CallingConvention.Cdecl)]
 public static extern bool CreateRoute(AddressFamily inet, string address, byte cidr, string gateway, uint index, uint metric);
 
diff --git a/RouteHelper/RouteHelper.cpp b/RouteHelper/RouteHelper.cpp
index ca219e5a..262df686 100644
--- a/RouteHelper/RouteHelper.cpp
+++ b/RouteHelper/RouteHelper.cpp
@@ -1,9 +1,5 @@
-#include <WinSock2.h>
-#include <WS2tcpip.h>
-#include <ws2ipdef.h>
-#include <iphlpapi.h>
-#include <netioapi.h>
-#include <Windows.h>
+#include "Based.h"
+#include "WaitGroup.h"
 
 BOOL APIENTRY DllMain(HMODULE hModule, DWORD ul_reason_for_call, LPVOID lpReserved)
 {
@@ -14,6 +10,16 @@ BOOL APIENTRY DllMain(HMODULE hModule, DWORD ul_reason_for_call, LPVOID lpReserv
     return TRUE;
 }
 
+WaitGroup wg;
+void UnicastIPChangeCallback(PVOID ctx, PMIB_UNICASTIPADDRESS_ROW row, MIB_NOTIFICATION_TYPE type)
+{
+    UNREFERENCED_PARAMETER(ctx);
+    UNREFERENCED_PARAMETER(row);
+    UNREFERENCED_PARAMETER(type);
+
+    wg.Done();
+}
+
 bool make(PMIB_IPFORWARD_ROW2 rule, USHORT inet, const char* address, UINT8 cidr, const char* gateway, ULONG index, ULONG metric)
 {
     rule->InterfaceIndex = index;
@@ -76,6 +82,18 @@ extern "C" {
         return index;
     }
 
+    __declspec(dllexport) void __cdecl WaitForUnicastIP()
+    {
+        wg.Add(1);
+
+        HANDLE hCallback = NULL;
+        NotifyUnicastIpAddressChange(AF_INET, UnicastIPChangeCallback, NULL, FALSE, &hCallback);
+
+        wg.Wait();
+
+        CancelMibChangeNotify2(hCallback);
+    }
+
     __declspec(dllexport) BOOL __cdecl CreateIPv4(const char* address, const char* netmask, ULONG index)
     {
         ULONG addr = 0;
@@ -127,16 +145,6 @@ extern "C" {
         return (NO_ERROR == CreateUnicastIpAddressEntry(&addr)) ? TRUE : FALSE;
     }
 
-    __declspec(dllexport) BOOL __cdecl RefreshIPTable(USHORT inet, ULONG index)
-    {
-        if (NO_ERROR != FlushIpPathTable(inet))
-        {
-            return FALSE;
-        }
-
-        return (NO_ERROR == FlushIpNetTable2(inet, index)) ? TRUE : FALSE;
-    }
-
     __declspec(dllexport) BOOL __cdecl CreateRoute(USHORT inet, const char* address, UINT8 cidr, const char* gateway, ULONG index, ULONG metric = 1)
     {
         MIB_IPFORWARD_ROW2 rule;
diff --git a/RouteHelper/RouteHelper.vcxproj b/RouteHelper/RouteHelper.vcxproj
index fa7b6789..64a61463 100644
--- a/RouteHelper/RouteHelper.vcxproj
+++ b/RouteHelper/RouteHelper.vcxproj
@@ -93,6 +93,11 @@
   </ItemDefinitionGroup>
   <ItemGroup>
     <ClCompile Include="RouteHelper.cpp" />
+    <ClCompile Include="WaitGroup.cpp" />
+  </ItemGroup>
+  <ItemGroup>
+    <ClInclude Include="Based.h" />
+    <ClInclude Include="WaitGroup.h" />
   </ItemGroup>
   <Import Project="$(VCTargetsPath)\Microsoft.Cpp.targets" />
   <ImportGroup Label="ExtensionTargets">
diff --git a/RouteHelper/RouteHelper.vcxproj.filters b/RouteHelper/RouteHelper.vcxproj.filters
index e15645d2..e03fb878 100644
--- a/RouteHelper/RouteHelper.vcxproj.filters
+++ b/RouteHelper/RouteHelper.vcxproj.filters
@@ -14,5 +14,16 @@
     <ClCompile Include="RouteHelper.cpp">
       <Filter>Source</Filter>
     </ClCompile>
+    <ClCompile Include="WaitGroup.cpp">
+      <Filter>Source</Filter>
+    </ClCompile>
+  </ItemGroup>
+  <ItemGroup>
+    <ClInclude Include="WaitGroup.h">
+      <Filter>Header</Filter>
+    </ClInclude>
+    <ClInclude Include="Based.h">
+      <Filter>Header</Filter>
+    </ClInclude>
   </ItemGroup>
 </Project>
\ No newline at end of file
diff --git a/RouteHelper/WaitGroup.cpp b/RouteHelper/WaitGroup.cpp
new file mode 100644
index 00000000..e850ae16
--- /dev/null
+++ b/RouteHelper/WaitGroup.cpp
@@ -0,0 +1,19 @@
+#include "WaitGroup.h"
+
+void WaitGroup::Add(int size)
+{
+	this->counter += size;
+}
+
+void WaitGroup::Done()
+{
+	if (--this->counter <= 0)
+		this->condition.notify_all();
+}
+
+void WaitGroup::Wait()
+{
+	std::unique_lock<std::mutex> lock(this->mutex);
+
+	condition.wait(lock, [&] { return counter <= 0; });
+}
diff --git a/RouteHelper/WaitGroup.h b/RouteHelper/WaitGroup.h
new file mode 100644
index 00000000..a564b912
--- /dev/null
+++ b/RouteHelper/WaitGroup.h
@@ -0,0 +1,18 @@
+#pragma once
+#ifndef WAITGROUP_H
+#define WAITGROUP_H
+#include "Based.h"
+
+class WaitGroup {
+public:
+	void Add(int size);
+	void Done();
+	void Wait();
+
+private:
+	std::mutex mutex;
+	std::atomic_int counter;
+	std::condition_variable condition;
+};
+
+#endif
diff --git a/build.ps1 b/build.ps1
index 3e7ecba5..556c8e54 100644
--- a/build.ps1
+++ b/build.ps1
@@ -19,7 +19,7 @@ param (
 
 	[Parameter()]
 	[bool]
-	$PublishReadyToRun = $True
+	$PublishReadyToRun = $False
 )
 
 Push-Location (Split-Path $MyInvocation.MyCommand.Path -Parent)
@@ -54,7 +54,7 @@ if ( -Not ( Test-Path ".\Netch\bin\$Configuration" ) ) {
 		-r 'win-x64' `
 		-p:Platform='x64' `
 		-p:SelfContained=$SelfContained `
-		-p:PublishTrimmed=$SelfContained `
+		-p:PublishTrimmed=$PublishReadyToRun `
 		-p:PublishSingleFile=$PublishSingleFile `
 		-p:PublishReadyToRun=$PublishReadyToRun `
 		-p:PublishReadyToRunShowWarnings=$PublishReadyToRun `