From 81c6ebd583dc7701d09963f103d4edd7367fd239 Mon Sep 17 00:00:00 2001
From: xljiulang <366193849@qq.com>
Date: Sun, 16 Oct 2022 18:47:26 +0800
Subject: [PATCH] =?UTF-8?q?=E5=BA=94=E7=94=A8cancellationToken?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../Dns/DnsInterceptor.cs                     | 14 +++++--------
 .../FastGithub.PacketIntercept.csproj         |  2 +-
 .../Tcp/TcpInterceptor.cs                     | 21 ++++++-------------
 3 files changed, 12 insertions(+), 25 deletions(-)

diff --git a/FastGithub.PacketIntercept/Dns/DnsInterceptor.cs b/FastGithub.PacketIntercept/Dns/DnsInterceptor.cs
index 7ebb1aa..0e47414 100644
--- a/FastGithub.PacketIntercept/Dns/DnsInterceptor.cs
+++ b/FastGithub.PacketIntercept/Dns/DnsInterceptor.cs
@@ -61,19 +61,15 @@ namespace FastGithub.PacketIntercept.Dns
         public async Task InterceptAsync(CancellationToken cancellationToken)
         {
             using var divert = new WinDivert(filter, WinDivertLayer.Network);
-            cancellationToken.Register(d =>
-            {
-                ((WinDivert)d!).Dispose();
-                DnsFlushResolverCache();
-            }, divert);
-
             using var packet = new WinDivertPacket();
             using var addr = new WinDivertAddress();
 
             DnsFlushResolverCache();
+            cancellationToken.Register(DnsFlushResolverCache);
+
             while (cancellationToken.IsCancellationRequested == false)
             {
-                await divert.RecvAsync(packet, addr);
+                await divert.RecvAsync(packet, addr, cancellationToken);
                 try
                 {
                     this.ModifyDnsPacket(packet, addr);
@@ -84,7 +80,7 @@ namespace FastGithub.PacketIntercept.Dns
                 }
                 finally
                 {
-                    await divert.SendAsync(packet, addr);
+                    await divert.SendAsync(packet, addr, cancellationToken);
                 }
             }
         }
@@ -159,7 +155,7 @@ namespace FastGithub.PacketIntercept.Dns
             }
             else
             {
-                addr.Flags ^= WinDivertAddressFlag.Outbound;
+                addr.Flags &= ~WinDivertAddressFlag.Outbound;
             }
 
             packet.CalcChecksums(addr);
diff --git a/FastGithub.PacketIntercept/FastGithub.PacketIntercept.csproj b/FastGithub.PacketIntercept/FastGithub.PacketIntercept.csproj
index a2e9348..890c1ee 100644
--- a/FastGithub.PacketIntercept/FastGithub.PacketIntercept.csproj
+++ b/FastGithub.PacketIntercept/FastGithub.PacketIntercept.csproj
@@ -7,7 +7,7 @@
 	<ItemGroup>
 		<FrameworkReference Include="Microsoft.AspNetCore.App" />
 		<PackageReference Include="DNS" Version="7.0.0" />
-		<PackageReference Include="WindivertDotnet" Version="1.0.2" />
+		<PackageReference Include="WindivertDotnet" Version="1.0.4" />
 	</ItemGroup>
 
 	<ItemGroup>
diff --git a/FastGithub.PacketIntercept/Tcp/TcpInterceptor.cs b/FastGithub.PacketIntercept/Tcp/TcpInterceptor.cs
index ed0331a..2d38bb1 100644
--- a/FastGithub.PacketIntercept/Tcp/TcpInterceptor.cs
+++ b/FastGithub.PacketIntercept/Tcp/TcpInterceptor.cs
@@ -50,7 +50,10 @@ namespace FastGithub.PacketIntercept.Tcp
                 return;
             }
 
-            using var divert = new WinDivert(this.filter, WinDivertLayer.Network, 0, WinDivertFlag.None);
+            using var divert = new WinDivert(this.filter, WinDivertLayer.Network);
+            using var packet = new WinDivertPacket();
+            using var addr = new WinDivertAddress();
+
             if (Socket.OSSupportsIPv4)
             {
                 this.logger.LogInformation($"{IPAddress.Loopback}:{this.oldServerPort} <=> {IPAddress.Loopback}:{this.newServerPort}");
@@ -59,13 +62,10 @@ namespace FastGithub.PacketIntercept.Tcp
             {
                 this.logger.LogInformation($"{IPAddress.IPv6Loopback}:{this.oldServerPort} <=> {IPAddress.IPv6Loopback}:{this.newServerPort}");
             }
-            cancellationToken.Register(d => ((WinDivert)d!).Dispose(), divert);
 
-            using var packet = new WinDivertPacket();
-            using var addr = new WinDivertAddress();
             while (cancellationToken.IsCancellationRequested == false)
             {
-                await divert.RecvAsync(packet, addr);
+                await divert.RecvAsync(packet, addr, cancellationToken);
                 try
                 {
                     this.ModifyTcpPacket(packet, addr);
@@ -76,7 +76,7 @@ namespace FastGithub.PacketIntercept.Tcp
                 }
                 finally
                 {
-                    await divert.SendAsync(packet, addr);
+                    await divert.SendAsync(packet, addr, cancellationToken);
                 }
             }
         }
@@ -89,15 +89,6 @@ namespace FastGithub.PacketIntercept.Tcp
         unsafe private void ModifyTcpPacket(WinDivertPacket packet, WinDivertAddress addr)
         {
             var result = packet.GetParseResult();
-            if (result.IPV4Header != null && result.IPV4Header->SrcAddr.Equals(IPAddress.Loopback) == false)
-            {
-                return;
-            }
-            if (result.IPV6Header != null && result.IPV6Header->SrcAddr.Equals(IPAddress.IPv6Loopback) == false)
-            {
-                return;
-            }
-
             if (result.TcpHeader->DstPort == oldServerPort)
             {
                 result.TcpHeader->DstPort = this.newServerPort;