let unit tests cover proxy protocol v1

kerryjiang · Aug 18, 2024 · b6a6487 · b6a6487
1 parent c046e44
commit b6a6487
Show file tree

Hide file tree

Showing 4 changed files with 109 additions and 62 deletions.
diff --git a/test/SuperSocket.Tests/ProxyProtocolHostConfigurator.cs b/test/SuperSocket.Tests/ProxyProtocolHostConfigurator.cs
@@ -14,65 +14,19 @@
 
 namespace SuperSocket.Tests
 {
-    public class ProxyProtocolHostConfigurator : IHostConfigurator
+    public abstract class ProxyProtocolHostConfigurator : IHostConfigurator
     {
         private IHostConfigurator _innerHostConfigurator;
 
-        private static readonly byte[] _proxyProtocolV2_SIGNATURE = new byte[]
-            {
-                // Signature
-                0x0D, 0x0A, 0x0D, 0x0A,
-                0x00, 0x0D, 0x0A, 0x51,
-                0x55, 0x49, 0x54, 0x0A
-            };
+        protected IHostConfigurator InnerHostConfigurator
+        {
+            get { return _innerHostConfigurator; }
+        }
 
         private IPEndPoint _sourceIPEndPoint;
         private IPEndPoint _destinationIPEndPoint;
 
-        private byte[] CreateProxyProtocolData(IPEndPoint sourceIPEndPoint, IPEndPoint destinationIPEndPoint)
-        {
-            var isIpV4 = sourceIPEndPoint.Address.AddressFamily == AddressFamily.InterNetwork;
-            var ipAddressLength = isIpV4 ? 4 : 16;
-
-            var addressLength = isIpV4
-                ? (ipAddressLength * 2 + 4)
-                : (ipAddressLength * 2 + 4);
-
-            var data = new byte[_proxyProtocolV2_SIGNATURE.Length  + 4 + addressLength];
-
-            _proxyProtocolV2_SIGNATURE.CopyTo(data, 0);
-
-            data[_proxyProtocolV2_SIGNATURE.Length] = 0x21;
-            data[_proxyProtocolV2_SIGNATURE.Length + 1] = (byte)((_innerHostConfigurator is UdpHostConfigurator ? 0x02 : 0x01) | (isIpV4 ? 0x10 : 0x20));
-
-            var span = data.AsSpan().Slice(_proxyProtocolV2_SIGNATURE.Length);
-
-            BinaryPrimitives.WriteUInt16BigEndian(span.Slice(2, 2), (ushort)addressLength);
-
-            var spanToWrite = span.Slice(4);
-
-            var addressSpan = spanToWrite.Slice(0, ipAddressLength);
-
-            var written = 0;
-
-            sourceIPEndPoint.Address.TryWriteBytes(addressSpan, out written);
-
-            Assert.Equal(ipAddressLength, written);            
-
-            spanToWrite = spanToWrite.Slice(ipAddressLength);
-
-            addressSpan = spanToWrite.Slice(0, ipAddressLength);
-            destinationIPEndPoint.Address.TryWriteBytes(addressSpan, out written);
-
-            Assert.Equal(ipAddressLength, written);
-
-            spanToWrite = spanToWrite.Slice(ipAddressLength);
-
-            BinaryPrimitives.WriteUInt16BigEndian(spanToWrite.Slice(0, 2), (ushort)sourceIPEndPoint.Port);
-            BinaryPrimitives.WriteUInt16BigEndian(spanToWrite.Slice(2, 2), (ushort)destinationIPEndPoint.Port);
-
-            return data;
-        }
+        protected abstract byte[] CreateProxyProtocolData(IPEndPoint sourceIPEndPoint, IPEndPoint destinationIPEndPoint);
 
         public ProxyProtocolHostConfigurator(IHostConfigurator hostConfigurator, IPEndPoint sourceIPEndPoint, IPEndPoint destinationIPEndPoint)
         {
@@ -105,7 +59,7 @@ public Socket CreateClient()
         public async ValueTask<Stream> GetClientStream(Socket socket)
         {
             var stream = await _innerHostConfigurator.GetClientStream(socket);
-            
+
             stream.Write(CreateProxyProtocolData(_sourceIPEndPoint, _destinationIPEndPoint));
             await stream.FlushAsync();
 

diff --git a/test/SuperSocket.Tests/ProxyProtocolTest.cs b/test/SuperSocket.Tests/ProxyProtocolTest.cs
@@ -1,10 +1,7 @@
 using System;
 using System.Buffers;
 using System.Collections.Generic;
-using System.Text;
 using Microsoft.Extensions.Configuration;
-using Microsoft.Extensions.Hosting;
-using SuperSocket;
 using SuperSocket.ProtoBase;
 using SuperSocket.Server.Host;
 using SuperSocket.Server.Abstractions;
@@ -14,7 +11,6 @@
 using System.Linq;
 using System.Threading.Tasks;
 using SuperSocket.Server.Abstractions.Session;
-using SuperSocket.Server;
 
 namespace SuperSocket.Tests
 {
@@ -64,7 +60,7 @@ protected override IHostConfigurator CreateHostConfigurator(Type hostConfigurato
             var sourceIPEndPoint = new IPEndPoint(addressPool[_rd.Next(0, addressPool.Length)], _rd.Next(100, 9999));
             var destinationIPEndPoint = new IPEndPoint(addressPool[_rd.Next(0, addressPool.Length)], _rd.Next(100, 9999));
 
-            return new ProxyProtocolHostConfigurator(base.CreateHostConfigurator(hostConfiguratorType), sourceIPEndPoint, destinationIPEndPoint);
+            return new ProxyProtocolV2HostConfigurator(base.CreateHostConfigurator(hostConfiguratorType), sourceIPEndPoint, destinationIPEndPoint);
         }
 
         protected override Dictionary<string, string> LoadMemoryConfig(Dictionary<string, string> configSettings)
@@ -75,16 +71,20 @@ protected override Dictionary<string, string> LoadMemoryConfig(Dictionary<string
         }
 
         [Theory]
-        [InlineData(typeof(RegularHostConfigurator), "123.193.169.24", 5444, "228.111.89.87", 199)]
-        [InlineData(typeof(RegularHostConfigurator), "529f:bc3e:2e56:8365:dc06:5772:87a5:e658", 5444, "48d4:c5d6:b3e8:9859:dc0f:a8d0:e085:3518", 199)]
-        public async Task TestProxyEndPoints(Type hostConfiguratorType, string sourceIP, int sourcePort, string destinationIP, int destinationPort)
+        [InlineData(typeof(RegularHostConfigurator), 1, "123.193.169.24", 5444, "228.111.89.87", 199)]
+        [InlineData(typeof(RegularHostConfigurator), 1, "529f:bc3e:2e56:8365:dc06:5772:87a5:e658", 5444, "48d4:c5d6:b3e8:9859:dc0f:a8d0:e085:3518", 199)]
+        [InlineData(typeof(RegularHostConfigurator), 2, "123.193.169.24", 5444, "228.111.89.87", 199)]
+        [InlineData(typeof(RegularHostConfigurator), 2, "529f:bc3e:2e56:8365:dc06:5772:87a5:e658", 5444, "48d4:c5d6:b3e8:9859:dc0f:a8d0:e085:3518", 199)]
+        public async Task TestProxyEndPoints(Type hostConfiguratorType, int proxyProtocolVersion, string sourceIP, int sourcePort, string destinationIP, int destinationPort)
         {
             var hostConfigurator = base.CreateHostConfigurator(hostConfiguratorType);
 
             var sourceIPAddress = IPAddress.Parse(sourceIP);
             var destinationIPAddress = IPAddress.Parse(destinationIP);
 
-            hostConfigurator = new ProxyProtocolHostConfigurator(hostConfigurator, new IPEndPoint(sourceIPAddress, sourcePort), new IPEndPoint(destinationIPAddress, destinationPort));
+            hostConfigurator = proxyProtocolVersion == 1
+                ? new ProxyProtocolV1HostConfigurator(hostConfigurator, new IPEndPoint(sourceIPAddress, sourcePort), new IPEndPoint(destinationIPAddress, destinationPort))
+                : new ProxyProtocolV2HostConfigurator(hostConfigurator, new IPEndPoint(sourceIPAddress, sourcePort), new IPEndPoint(destinationIPAddress, destinationPort));
 
             var taskCompletionSource = new TaskCompletionSource<IAppSession>();
 

diff --git a/test/SuperSocket.Tests/ProxyProtocolV1HostConfigurator.cs b/test/SuperSocket.Tests/ProxyProtocolV1HostConfigurator.cs
@@ -0,0 +1,24 @@
+using System.Net;
+using System.Net.Sockets;
+using System.Text;
+
+namespace SuperSocket.Tests
+{
+    public class ProxyProtocolV1HostConfigurator : ProxyProtocolHostConfigurator
+    {
+        protected override byte[] CreateProxyProtocolData(IPEndPoint sourceIPEndPoint, IPEndPoint destinationIPEndPoint)
+        {
+            var protocol = InnerHostConfigurator is UdpHostConfigurator ? "UDP" : "TCP";
+            var addressVersion = sourceIPEndPoint.Address.AddressFamily == AddressFamily.InterNetwork ? 4 : 6;
+
+            var line = $"PROXY {protocol}{addressVersion} {sourceIPEndPoint.Address} {destinationIPEndPoint.Address} {sourceIPEndPoint.Port} {destinationIPEndPoint.Port}\r\n";
+
+            return Encoding.ASCII.GetBytes(line);
+        }
+
+        public ProxyProtocolV1HostConfigurator(IHostConfigurator hostConfigurator, IPEndPoint sourceIPEndPoint, IPEndPoint destinationIPEndPoint)
+            : base(hostConfigurator, sourceIPEndPoint, destinationIPEndPoint)
+        {
+        }
+    }
+}
diff --git a/test/SuperSocket.Tests/ProxyProtocolV2HostConfigurator.cs b/test/SuperSocket.Tests/ProxyProtocolV2HostConfigurator.cs
@@ -0,0 +1,69 @@
+using System;
+using System.Buffers.Binary;
+using System.Net;
+using System.Net.Sockets;
+using Xunit;
+
+namespace SuperSocket.Tests
+{
+    public class ProxyProtocolV2HostConfigurator : ProxyProtocolHostConfigurator
+    {
+        private static readonly byte[] _proxyProtocolV2_SIGNATURE = new byte[]
+            {
+                // Signature
+                0x0D, 0x0A, 0x0D, 0x0A,
+                0x00, 0x0D, 0x0A, 0x51,
+                0x55, 0x49, 0x54, 0x0A
+            };
+
+        protected override byte[] CreateProxyProtocolData(IPEndPoint sourceIPEndPoint, IPEndPoint destinationIPEndPoint)
+        {
+            var isIpV4 = sourceIPEndPoint.Address.AddressFamily == AddressFamily.InterNetwork;
+            var ipAddressLength = isIpV4 ? 4 : 16;
+
+            var addressLength = isIpV4
+                ? (ipAddressLength * 2 + 4)
+                : (ipAddressLength * 2 + 4);
+
+            var data = new byte[_proxyProtocolV2_SIGNATURE.Length  + 4 + addressLength];
+
+            _proxyProtocolV2_SIGNATURE.CopyTo(data, 0);
+
+            data[_proxyProtocolV2_SIGNATURE.Length] = 0x21;
+            data[_proxyProtocolV2_SIGNATURE.Length + 1] = (byte)((InnerHostConfigurator is UdpHostConfigurator ? 0x02 : 0x01) | (isIpV4 ? 0x10 : 0x20));
+
+            var span = data.AsSpan().Slice(_proxyProtocolV2_SIGNATURE.Length);
+
+            BinaryPrimitives.WriteUInt16BigEndian(span.Slice(2, 2), (ushort)addressLength);
+
+            var spanToWrite = span.Slice(4);
+
+            var addressSpan = spanToWrite.Slice(0, ipAddressLength);
+
+            var written = 0;
+
+            sourceIPEndPoint.Address.TryWriteBytes(addressSpan, out written);
+
+            Assert.Equal(ipAddressLength, written);            
+
+            spanToWrite = spanToWrite.Slice(ipAddressLength);
+
+            addressSpan = spanToWrite.Slice(0, ipAddressLength);
+            destinationIPEndPoint.Address.TryWriteBytes(addressSpan, out written);
+
+            Assert.Equal(ipAddressLength, written);
+
+            spanToWrite = spanToWrite.Slice(ipAddressLength);
+
+            BinaryPrimitives.WriteUInt16BigEndian(spanToWrite.Slice(0, 2), (ushort)sourceIPEndPoint.Port);
+            BinaryPrimitives.WriteUInt16BigEndian(spanToWrite.Slice(2, 2), (ushort)destinationIPEndPoint.Port);
+
+            return data;
+        }
+
+        public ProxyProtocolV2HostConfigurator(IHostConfigurator hostConfigurator, IPEndPoint sourceIPEndPoint, IPEndPoint destinationIPEndPoint)
+            : base(hostConfigurator, sourceIPEndPoint, destinationIPEndPoint)
+        {
+        }
+    }
+}