From 4f5e5ab8044a8fb66fd14bc3ed924e5b779e6fa6 Mon Sep 17 00:00:00 2001
From: ChsBuffer <33744752+chsbuffer@users.noreply.github.com>
Date: Thu, 11 Nov 2021 13:04:38 +0800
Subject: [PATCH] Fix wait CreateUnicastIP Fix Guard close file

---
 Netch/Controllers/Guard.cs         | 35 ++++++++---------
 Netch/Controllers/TUNController.cs | 12 ++++--
 Netch/Interops/RouteHelper.cs      | 60 +++++++++++++++++++++++-------
 Netch/NativeMethods.txt            |  1 +
 4 files changed, 74 insertions(+), 34 deletions(-)

diff --git a/Netch/Controllers/Guard.cs b/Netch/Controllers/Guard.cs
index 066b1bc4..ae465b0d 100644
--- a/Netch/Controllers/Guard.cs
+++ b/Netch/Controllers/Guard.cs
@@ -61,13 +61,6 @@ namespace Netch.Controllers
 
         public Process Instance { get; }
 
-        ~Guard()
-        {
-            _logFileStream?.Dispose();
-            _logStreamWriter?.Dispose();
-            Instance.Dispose();
-        }
-
         protected async Task StartGuardAsync(string argument, ProcessPriorityClass priority = ProcessPriorityClass.Normal)
         {
             State = State.Starting;
@@ -83,8 +76,12 @@ namespace Netch.Controllers
 
             if (RedirectOutput)
             {
-                Task.Run(() => ReadOutput(Instance.StandardOutput)).Forget();
-                Task.Run(() => ReadOutput(Instance.StandardError)).Forget();
+                WaitAllReadOutputTaskAsync(new[]
+                    {
+                        ReadOutputAsync(Instance.StandardOutput),
+                        ReadOutputAsync(Instance.StandardError)
+                    })
+                    .Forget();
 
                 if (!StartedKeywords.Any())
                 {
@@ -114,12 +111,21 @@ namespace Netch.Controllers
             }
         }
 
-        private void ReadOutput(TextReader reader)
+        private async Task WaitAllReadOutputTaskAsync(Task[] tasks)
+        {
+            await Task.WhenAll(tasks);
+            _logStreamWriter?.Close();
+            _logFileStream?.Close();
+            Instance.Dispose();
+            State = State.Stopped;
+        }
+
+        private async Task ReadOutputAsync(TextReader reader)
         {
             string? line;
-            while ((line = reader.ReadLine()) != null)
+            while ((line = await reader.ReadLineAsync()) != null)
             {
-                _logStreamWriter!.WriteLine(line);
+                await _logStreamWriter!.WriteLineAsync(line);
                 OnReadNewLine(line);
 
                 if (State == State.Starting)
@@ -133,8 +139,6 @@ namespace Netch.Controllers
                     }
                 }
             }
-
-            State = State.Stopped;
         }
 
         public virtual async Task StopAsync()
@@ -144,9 +148,6 @@ namespace Netch.Controllers
 
         protected async Task StopGuardAsync()
         {
-            _logStreamWriter?.Close();
-            _logFileStream?.Close();
-
             try
             {
                 if (Instance is { HasExited: false })
diff --git a/Netch/Controllers/TUNController.cs b/Netch/Controllers/TUNController.cs
index 4d328a09..fe50c919 100644
--- a/Netch/Controllers/TUNController.cs
+++ b/Netch/Controllers/TUNController.cs
@@ -96,10 +96,14 @@ namespace Netch.Controllers
             var tunIndex = _tunNetworkInterface.GetIndex();
             _tun = NetRoute.TemplateBuilder(_tunConfig.Gateway, tunIndex);
 
-            RouteHelper.CreateUnicastIP(AddressFamily.InterNetwork,
-                _tunConfig.Address,
-                (byte)Utils.Utils.SubnetToCidr(_tunConfig.Netmask),
-                (ulong)tunIndex);
+            if (!RouteHelper.CreateUnicastIP(AddressFamily.InterNetwork,
+                    _tunConfig.Address,
+                    (byte)Utils.Utils.SubnetToCidr(_tunConfig.Netmask),
+                    (ulong)tunIndex))
+            {
+                Log.Error("Create Unicast IP failed");
+                throw new MessageException("Create Unicast IP failed");
+            }
 
             await SetupRouteTableAsync();
         }
diff --git a/Netch/Interops/RouteHelper.cs b/Netch/Interops/RouteHelper.cs
index 059a0a36..ad33f68e 100644
--- a/Netch/Interops/RouteHelper.cs
+++ b/Netch/Interops/RouteHelper.cs
@@ -1,8 +1,11 @@
+using System;
 using System.Net.Sockets;
 using System.Runtime.InteropServices;
 using System.Threading;
 using Windows.Win32.Foundation;
+using Windows.Win32.Networking.WinSock;
 using Windows.Win32.NetworkManagement.IpHelper;
+using Serilog;
 using static Windows.Win32.PInvoke;
 
 namespace Netch.Interops
@@ -40,25 +43,56 @@ namespace Netch.Interops
                 return false;
             }
 
-            // Create a Handle to be notified of IP address changes
+            // https://docs.microsoft.com/en-us/windows/win32/api/netioapi/nf-netioapi-createunicastipaddressentry#remarks
+
             HANDLE handle = default;
-            using var obj = new Semaphore(1, 1);
-            void Callback(void* context, MIB_UNICASTIPADDRESS_ROW* row, MIB_NOTIFICATION_TYPE type) => obj.Release(1);
+            using var obj = new Semaphore(0, 1);
 
-            // Use NotifyUnicastIpAddressChange to determine when the address is ready
-            obj.WaitOne();
-            NotifyUnicastIpAddressChange((ushort)ADDRESS_FAMILY.AF_INET, Callback, null, new BOOLEAN(), ref handle);
-
-            if (CreateUnicastIpAddressEntry(&addr) != 0)
+            void Callback(void* context, MIB_UNICASTIPADDRESS_ROW* row, MIB_NOTIFICATION_TYPE type)
             {
-                // ignored return state because i feel great
-                CancelMibChangeNotify2(handle);
-                return false;
+                if (type != MIB_NOTIFICATION_TYPE.MibInitialNotification)
+                {
+                    // TODO pass error
+                    NTSTATUS state;
+                    if ((state = GetUnicastIpAddressEntry(row)) != 0)
+                    {
+                        Log.Error("CreateUnicastIpAddressEntry failed: {0}", state);
+                        return;
+                    }
+
+                    if (row -> DadState == NL_DAD_STATE.IpDadStatePreferred)
+                    {
+                        try
+                        {
+                            obj.Release();
+                        }
+                        catch (Exception e)
+                        {
+                            // i don't trust win32 api
+                            Log.Error(e, "semaphore disposed");
+                        }
+                    }
+                }
             }
 
-            obj.WaitOne();
+            NotifyUnicastIpAddressChange((ushort)ADDRESS_FAMILY.AF_INET, Callback, null, new BOOLEAN(byte.MaxValue), ref handle);
 
-            return true;
+            try
+            {
+                NTSTATUS state;
+                if ((state = CreateUnicastIpAddressEntry(&addr)) != 0)
+                {
+                    Log.Error("CreateUnicastIpAddressEntry failed: {0}", state);
+                    return false;
+                }
+
+                obj.WaitOne();
+                return true;
+            }
+            finally
+            {
+                CancelMibChangeNotify2(handle);
+            }
         }
 
         [DllImport("RouteHelper.bin", CallingConvention = CallingConvention.Cdecl)]
diff --git a/Netch/NativeMethods.txt b/Netch/NativeMethods.txt
index b55523f7..f5ef68c3 100644
--- a/Netch/NativeMethods.txt
+++ b/Netch/NativeMethods.txt
@@ -6,6 +6,7 @@ ADDRESS_FAMILY
 MIB_UNICASTIPADDRESS_ROW
 InitializeUnicastIpAddressEntry
 CreateUnicastIpAddressEntry
+GetUnicastIpAddressEntry
 NotifyUnicastIpAddressChange
 CancelMibChangeNotify2