增强WinDivertRouter

xljiulang · Oct 19, 2022 · 3cb9a81 · 3cb9a81
1 parent a38115c
commit 3cb9a81
Show file tree

Hide file tree

Showing 4 changed files with 136 additions and 47 deletions.
diff --git a/WindivertDotnet.Test/WindivertDotnet.Test.csproj b/WindivertDotnet.Test/WindivertDotnet.Test.csproj
@@ -1,27 +1,27 @@
-<Project Sdk="Microsoft.NET.Sdk">
+<Project Sdk="Microsoft.NET.Sdk">
 
-  <PropertyGroup>
-    <TargetFramework>net6.0</TargetFramework>
-    <Nullable>enable</Nullable>
+	<PropertyGroup>
+		<TargetFramework>net6.0</TargetFramework>
+		<Nullable>enable</Nullable>
+		<AllowUnsafeBlocks>True</AllowUnsafeBlocks>
+		<IsPackable>false</IsPackable>
+	</PropertyGroup>
 
-    <IsPackable>false</IsPackable>
-  </PropertyGroup>
+	<ItemGroup>
+		<PackageReference Include="Microsoft.NET.Test.Sdk" Version="16.11.0" />
+		<PackageReference Include="xunit" Version="2.4.1" />
+		<PackageReference Include="xunit.runner.visualstudio" Version="2.4.3">
+			<IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
+			<PrivateAssets>all</PrivateAssets>
+		</PackageReference>
+		<PackageReference Include="coverlet.collector" Version="3.1.0">
+			<IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
+			<PrivateAssets>all</PrivateAssets>
+		</PackageReference>
+	</ItemGroup>
 
-  <ItemGroup>
-    <PackageReference Include="Microsoft.NET.Test.Sdk" Version="16.11.0" />
-    <PackageReference Include="xunit" Version="2.4.1" />
-    <PackageReference Include="xunit.runner.visualstudio" Version="2.4.3">
-      <IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
-      <PrivateAssets>all</PrivateAssets>
-    </PackageReference>
-    <PackageReference Include="coverlet.collector" Version="3.1.0">
-      <IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
-      <PrivateAssets>all</PrivateAssets>
-    </PackageReference>
-  </ItemGroup>
-
-  <ItemGroup>
-    <ProjectReference Include="..\WindivertDotnet\WindivertDotnet.csproj" />
-  </ItemGroup>
+	<ItemGroup>
+		<ProjectReference Include="..\WindivertDotnet\WindivertDotnet.csproj" />
+	</ItemGroup>
 
 </Project>
diff --git a/WindivertDotnet.Test/WindivertRouterTest.cs b/WindivertDotnet.Test/WindivertRouterTest.cs
@@ -0,0 +1,67 @@
+using System.Net;
+using Xunit;
+
+namespace WindivertDotnet.Test
+{
+    public class WindivertRouterTest
+    {
+        [Fact]
+        public void IPv4LoopbackTest()
+        {
+            var router = new WinDivertRouter(IPAddress.Loopback);
+            Assert.Equal(IPAddress.Loopback, router.SrcAddress);
+            Assert.Equal(IPAddress.Loopback, router.DstAddress);
+            Assert.True(router.IsOutbound);
+            Assert.True(router.IsLoopback);
+        }
+
+        [Fact]
+        public unsafe void IPv4LoopbackAddrTest()
+        {
+            var router = new WinDivertRouter(IPAddress.Loopback);
+            using var addr = router.CreateAddress();
+
+            Assert.Equal(addr.Network->IfIdx, router.InterfaceIndex);
+            Assert.True(addr.Flags.HasFlag(WinDivertAddressFlag.Outbound));
+            Assert.True(addr.Flags.HasFlag(WinDivertAddressFlag.Loopback));
+        }
+
+        [Fact]
+        public void IPv6LoopbackTest()
+        {
+            var router = new WinDivertRouter(IPAddress.IPv6Loopback);
+            Assert.Equal(IPAddress.IPv6Loopback, router.SrcAddress);
+            Assert.Equal(IPAddress.IPv6Loopback, router.DstAddress);
+            Assert.True(router.IsOutbound);
+            Assert.True(router.IsLoopback);
+        }
+
+        [Fact]
+        public void IPv41111Test()
+        {
+            var dstAddrr = IPAddress.Parse("1.1.1.1");
+            var router = new WinDivertRouter(dstAddrr);
+            Assert.NotEqual(IPAddress.Any, router.SrcAddress);
+            Assert.Equal(dstAddrr, router.DstAddress);
+            Assert.True(router.IsOutbound);
+            Assert.False(router.IsLoopback);
+        }
+
+        [Fact]
+        public unsafe void IPv41111AddrTest()
+        {
+            var dstAddrr = IPAddress.Parse("1.1.1.1");
+            var router = new WinDivertRouter(dstAddrr);
+
+            using var addr = new WinDivertAddress
+            {
+                Flags = WinDivertAddressFlag.Loopback
+            };
+            router.ApplyToAddress(addr);
+
+            Assert.Equal(addr.Network->IfIdx, router.InterfaceIndex);
+            Assert.True(addr.Flags.HasFlag(WinDivertAddressFlag.Outbound));
+            Assert.False(addr.Flags.HasFlag(WinDivertAddressFlag.Loopback));
+        }
+    }
+}
diff --git a/WindivertDotnet/WinDivertRouter.cs b/WindivertDotnet/WinDivertRouter.cs
@@ -34,6 +34,11 @@ public class WinDivertRouter
         /// </summary>
         public bool IsOutbound { get; }
 
+        /// <summary>
+        /// 获取是否为回环
+        /// </summary>
+        public bool IsLoopback { get; }
+
         /// <summary>
         /// WinDivert路由
         /// </summary> 
@@ -138,27 +143,9 @@ private unsafe WinDivertRouter(IPAddress dstAddr, IPAddress? srcAddr, int? inter
             this.DstAddress = dstAddr;
             this.InterfaceIndex = interfaceIndex.Value;
             this.IsOutbound = errorCode == 0;
+            this.IsLoopback = IPAddress.IsLoopback(dstAddr) && dstAddr.Equals(this.SrcAddress);
         }
 
-        /// <summary>
-        /// 获取网络接口索引
-        /// </summary>
-        /// <param name="dstAddr">目标地址</param>
-        /// <returns></returns>
-        /// <exception cref="ArgumentException"></exception>
-        /// <exception cref="NetworkInformationException"></exception>
-        public static int GetInterfaceIndex(IPAddress dstAddr)
-        {
-            if (IsIPAddressAny(dstAddr))
-            {
-                throw new ArgumentException($"值不能为{dstAddr}", nameof(dstAddr));
-            }
-
-            var dstSockAddr = new SockAddress { IPAddress = dstAddr };
-            return GetInterfaceIndex(ref dstSockAddr);
-        }
-
-
         /// <summary>
         /// 是否为any的ip
         /// </summary>
@@ -181,26 +168,61 @@ private static int GetInterfaceIndex(ref SockAddress dstSockAddr)
             return errorCode == 0 ? ifIdx : throw new NetworkInformationException(errorCode);
         }
 
+        /// <summary>
+        /// 获取网络接口索引
+        /// </summary>
+        /// <param name="dstAddr">目标地址</param>
+        /// <returns></returns>
+        /// <exception cref="ArgumentException"></exception>
+        /// <exception cref="NetworkInformationException"></exception>
+        public static int GetInterfaceIndex(IPAddress dstAddr)
+        {
+            if (IsIPAddressAny(dstAddr))
+            {
+                throw new ArgumentException($"值不能为{dstAddr}", nameof(dstAddr));
+            }
+
+            var dstSockAddr = new SockAddress { IPAddress = dstAddr };
+            return GetInterfaceIndex(ref dstSockAddr);
+        }
+
         /// <summary>
         /// 使用路由信息创建WinDivertAddress对象
         /// </summary>
         /// <returns></returns>
-        public unsafe WinDivertAddress CreateAddress()
+        public WinDivertAddress CreateAddress()
         {
             var addr = new WinDivertAddress();
+            this.ApplyToAddress(addr);
+            return addr;
+        }
+
+        /// <summary>
+        /// 应用路由信息到指定的WinDivertAddress
+        /// 将更改Network->IfIdx、Flags.Outbound和Flags.Loopback
+        /// </summary>
+        /// <param name="addr"></param>
+        public unsafe void ApplyToAddress(WinDivertAddress addr)
+        {
             addr.Network->IfIdx = this.InterfaceIndex;
 
-            if (this.IsOutbound == true)
+            if (this.IsOutbound)
             {
                 addr.Flags |= WinDivertAddressFlag.Outbound;
             }
+            else
+            {
+                addr.Flags &= ~WinDivertAddressFlag.Outbound;
+            }
 
-            if (IPAddress.IsLoopback(this.SrcAddress) && this.SrcAddress.Equals(this.DstAddress))
+            if (this.IsLoopback)
             {
                 addr.Flags |= WinDivertAddressFlag.Loopback;
             }
-
-            return addr;
+            else
+            {
+                addr.Flags &= ~WinDivertAddressFlag.Loopback;
+            }
         }
     }
 }
diff --git a/WindivertDotnet/WindivertDotnet.csproj b/WindivertDotnet/WindivertDotnet.csproj
@@ -1,7 +1,7 @@
 <Project Sdk="Microsoft.NET.Sdk">
 
 	<PropertyGroup>
-		<Version>1.0.6</Version>
+		<Version>1.0.7</Version>
 		<TargetFrameworks>netcoreapp3.1;net6.0</TargetFrameworks>
 		<Nullable>enable</Nullable>
 		<NoWarn>IDE0079</NoWarn>