From 6218bab623bf43644c2563fc9d623721c64b122e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=99=88=E5=9B=BD=E4=BC=9F?= <366193849@qq.com>
Date: Fri, 17 Sep 2021 08:35:12 +0800
Subject: [PATCH] =?UTF-8?q?=E6=AD=A3=E7=A1=AE=E5=85=B3=E9=97=ADWinDivert?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 FastGithub.Dns/DnsInterceptor.cs | 59 ++++----------------------------
 1 file changed, 7 insertions(+), 52 deletions(-)

diff --git a/FastGithub.Dns/DnsInterceptor.cs b/FastGithub.Dns/DnsInterceptor.cs
index 2743dac..614a6d1 100644
--- a/FastGithub.Dns/DnsInterceptor.cs
+++ b/FastGithub.Dns/DnsInterceptor.cs
@@ -10,7 +10,6 @@ using System.Runtime.InteropServices;
 using System.Runtime.Versioning;
 using System.Threading;
 using WinDivertSharp;
-using WinDivertSharp.WinAPI;
 
 namespace FastGithub.Dns
 {
@@ -57,16 +56,21 @@ namespace FastGithub.Dns
                 return;
             }
 
-            DnsFlushResolverCache();
+            cancellationToken.Register(hwnd =>
+            {
+                WinDivert.WinDivertClose((IntPtr)hwnd!);
+                DnsFlushResolverCache();
+            }, handle);
 
             var packetLength = 0U;
             var packetBuffer = new byte[ushort.MaxValue];
             using var winDivertBuffer = new WinDivertBuffer(packetBuffer);
             var winDivertAddress = new WinDivertAddress();
 
+            DnsFlushResolverCache();
             while (cancellationToken.IsCancellationRequested == false)
             {
-                if (this.WinDivertRecvEx(handle, winDivertBuffer, ref winDivertAddress, ref packetLength, cancellationToken))
+                if (WinDivert.WinDivertRecv(handle, winDivertBuffer, ref winDivertAddress, ref packetLength))
                 {
                     try
                     {
@@ -81,55 +85,6 @@ namespace FastGithub.Dns
                     WinDivert.WinDivertSend(handle, winDivertBuffer, packetLength, ref winDivertAddress);
                 }
             }
-
-            WinDivert.WinDivertClose(handle);
-            DnsFlushResolverCache();
-        }
-
-        /// <summary>
-        /// 接收数据
-        /// </summary>
-        /// <param name="handle"></param>
-        /// <param name="winDivertBuffer"></param>
-        /// <param name="winDivertAddress"></param>
-        /// <param name="packetLength"></param>
-        /// <param name="cancellationToken"></param>
-        /// <returns></returns>
-        private bool WinDivertRecvEx(IntPtr handle, WinDivertBuffer winDivertBuffer, ref WinDivertAddress winDivertAddress, ref uint packetLength, CancellationToken cancellationToken)
-        {
-            using var waitHandle = new EventWaitHandle(false, EventResetMode.ManualReset, null);
-            var overlapped = new NativeOverlapped
-            {
-                EventHandle = waitHandle.SafeWaitHandle.DangerousGetHandle()
-            };
-
-            winDivertAddress.Reset();
-            if (WinDivert.WinDivertRecvEx(handle, winDivertBuffer, 0, ref winDivertAddress, ref packetLength, ref overlapped))
-            {
-                return true;
-            }
-
-            var error = Marshal.GetLastWin32Error();
-            if (error != ERROR_IO_PENDING)
-            {
-                this.logger.LogWarning($"Unknown IO error ID {error} while awaiting overlapped result.");
-                return false;
-            }
-
-            while (waitHandle.WaitOne(TimeSpan.FromSeconds(1d)) == false)
-            {
-                cancellationToken.ThrowIfCancellationRequested();
-            }
-
-            var asyncPacketLength = 0U;
-            if (Kernel32.GetOverlappedResult(handle, ref overlapped, ref asyncPacketLength, false) == false)
-            {
-                this.logger.LogWarning("Failed to get overlapped result.");
-                return false;
-            }
-
-            packetLength = asyncPacketLength;
-            return true;
         }
 
         /// <summary>