From c252b205b14e40813e75ea7cebc83965669e7b86 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=99=88=E5=9B=BD=E4=BC=9F?= <366193849@qq.com>
Date: Wed, 15 Sep 2021 13:28:57 +0800
Subject: [PATCH] =?UTF-8?q?=E5=BC=82=E6=AD=A5=E6=8E=A5=E6=94=B6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 FastGithub.Dns/DnsInterceptor.cs | 56 +++++++++++++++++++++++++-------
 1 file changed, 44 insertions(+), 12 deletions(-)

diff --git a/FastGithub.Dns/DnsInterceptor.cs b/FastGithub.Dns/DnsInterceptor.cs
index c970c55..7305f19 100644
--- a/FastGithub.Dns/DnsInterceptor.cs
+++ b/FastGithub.Dns/DnsInterceptor.cs
@@ -10,6 +10,7 @@ using System.Runtime.InteropServices;
 using System.Runtime.Versioning;
 using System.Threading;
 using WinDivertSharp;
+using WinDivertSharp.WinAPI;
 
 namespace FastGithub.Dns
 {
@@ -19,7 +20,8 @@ namespace FastGithub.Dns
     [SupportedOSPlatform("windows")]
     sealed class DnsInterceptor
     {
-        const string DNS_FILTER = "udp.DstPort == 53";
+        private const string DNS_FILTER = "udp.DstPort == 53";
+        private const int ERROR_IO_PENDING = 997;
         private readonly FastGithubConfig fastGithubConfig;
         private readonly ILogger<DnsInterceptor> logger;
         private readonly TimeSpan ttl = TimeSpan.FromSeconds(10d);
@@ -47,7 +49,7 @@ namespace FastGithub.Dns
         /// DNS拦截
         /// </summary>
         /// <param name="cancellationToken"></param>
-        public void Intercept(CancellationToken cancellationToken)
+        unsafe public void Intercept(CancellationToken cancellationToken)
         {
             var handle = WinDivert.WinDivertOpen(DNS_FILTER, WinDivertLayer.Network, 0, WinDivertOpenFlags.None);
             if (handle == IntPtr.Zero)
@@ -55,28 +57,58 @@ namespace FastGithub.Dns
                 return;
             }
 
+            DnsFlushResolverCache();
+
             var packetLength = 0U;
             var packetBuffer = new byte[ushort.MaxValue];
             using var winDivertBuffer = new WinDivertBuffer(packetBuffer);
             var winDivertAddress = new WinDivertAddress();
+            using var ioEvent = new AutoResetEvent(false);
 
-            DnsFlushResolverCache();
             while (cancellationToken.IsCancellationRequested == false)
             {
-                if (WinDivert.WinDivertRecv(handle, winDivertBuffer, ref winDivertAddress, ref packetLength))
+                winDivertAddress.Reset();
+                var overlapped = new NativeOverlapped
                 {
-                    try
+                    EventHandle = ioEvent.SafeWaitHandle.DangerousGetHandle()
+                };
+
+                // false表示异步完成
+                if (WinDivert.WinDivertRecvEx(handle, winDivertBuffer, 0, ref winDivertAddress, ref packetLength, ref overlapped) == false)
+                {
+                    var error = Marshal.GetLastWin32Error();
+                    if (error != ERROR_IO_PENDING)
                     {
-                        this.ProcessDnsPacket(packetBuffer, ref packetLength);
-                    }
-                    catch (Exception ex)
-                    {
-                        this.logger.LogWarning(ex.Message);
+                        this.logger.LogWarning($"Unknown IO error ID {error} while awaiting overlapped result.");
+                        continue;
                     }
 
-                    WinDivert.WinDivertHelperCalcChecksums(winDivertBuffer, packetLength, ref winDivertAddress, WinDivertChecksumHelperParam.All);
-                    WinDivert.WinDivertSend(handle, winDivertBuffer, packetLength, ref winDivertAddress);
+                    while (ioEvent.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.");
+                        continue;
+                    }
+
+                    packetLength = asyncPacketLength;
                 }
+
+                try
+                {
+                    this.ProcessDnsPacket(packetBuffer, ref packetLength);
+                }
+                catch (Exception ex)
+                {
+                    this.logger.LogWarning(ex.Message);
+                }
+
+                WinDivert.WinDivertHelperCalcChecksums(winDivertBuffer, packetLength, ref winDivertAddress, WinDivertChecksumHelperParam.All);
+                WinDivert.WinDivertSend(handle, winDivertBuffer, packetLength, ref winDivertAddress);
             }
 
             WinDivert.WinDivertClose(handle);