More code coverage

src/Microsoft.Extensions.ServiceDiscovery.Dns/Resolver/DnsResolver.cs

@@ -196,7 +196,6 @@ public ValueTask<AddressResult[]> ResolveIPAddressesAsync(string name, AddressFa
 
             foreach (var answer in response.Answers)
             {
-
                 switch (answer.Type)
                 {
                     case QueryType.CNAME:

tests/Microsoft.Extensions.ServiceDiscovery.Dns.Tests/Resolver/LoopbackDnsServer.cs

@@ -36,7 +36,7 @@ public void DisableTcpFallback()
         _tcpSocket.Close();
     }
 
-    private static async Task<int> ProcessRequestCore(ArraySegment<byte> message, Func<LoopbackDnsResponseBuilder, Task> action, Memory<byte> responseBuffer)
+    private static async Task<int> ProcessRequestCore(IPEndPoint remoteEndPoint, ArraySegment<byte> message, Func<LoopbackDnsResponseBuilder, IPEndPoint, Task> action, Memory<byte> responseBuffer)
     {
         DnsDataReader reader = new DnsDataReader(message);
 
@@ -51,66 +51,20 @@ private static async Task<int> ProcessRequestCore(ArraySegment<byte> message, Fu
         responseBuilder.Flags = header.QueryFlags | QueryFlags.HasResponse;
         responseBuilder.ResponseCode = QueryResponseCode.NoError;
 
-        await action(responseBuilder);
+        await action(responseBuilder, remoteEndPoint);
 
-        DnsDataWriter writer = new(responseBuffer);
-        if (!writer.TryWriteHeader(new DnsMessageHeader
-        {
-            TransactionId = responseBuilder.TransactionId,
-            QueryFlags = responseBuilder.Flags | (QueryFlags)responseBuilder.ResponseCode,
-            QueryCount = (ushort)responseBuilder.Questions.Count,
-            AnswerCount = (ushort)responseBuilder.Answers.Count,
-            AuthorityCount = (ushort)responseBuilder.Authorities.Count,
-            AdditionalRecordCount = (ushort)responseBuilder.Additionals.Count
-        }))
-        {
-            throw new InvalidOperationException("Failed to write header");
-        }
-
-        foreach (var (questionName, questionType, questionClass) in responseBuilder.Questions)
-        {
-            if (!writer.TryWriteQuestion(questionName, questionType, questionClass))
-            {
-                throw new InvalidOperationException("Failed to write question");
-            }
-        }
-
-        foreach (var answer in responseBuilder.Answers)
-        {
-            if (!writer.TryWriteResourceRecord(answer))
-            {
-                throw new InvalidOperationException("Failed to write answer");
-            }
-        }
-
-        foreach (var authority in responseBuilder.Authorities)
-        {
-            if (!writer.TryWriteResourceRecord(authority))
-            {
-                throw new InvalidOperationException("Failed to write authority");
-            }
-        }
-
-        foreach (var additional in responseBuilder.Additionals)
-        {
-            if (!writer.TryWriteResourceRecord(additional))
-            {
-                throw new InvalidOperationException("Failed to write additional records");
-            }
-        }
-
-        return writer.Position;
+        return responseBuilder.Write(responseBuffer);
     }
 
-    public async Task ProcessUdpRequest(Func<LoopbackDnsResponseBuilder, Task> action)
+    public async Task ProcessUdpRequest(Func<LoopbackDnsResponseBuilder, IPEndPoint, Task> action)
     {
         byte[] buffer = ArrayPool<byte>.Shared.Rent(512);
         try
         {
             EndPoint remoteEndPoint = new IPEndPoint(IPAddress.Any, 0);
             SocketReceiveFromResult result = await _dnsSocket.ReceiveFromAsync(buffer, remoteEndPoint);
 
-            int bytesWritten = await ProcessRequestCore(new ArraySegment<byte>(buffer, 0, result.ReceivedBytes), action, buffer.AsMemory(0, 512));
+            int bytesWritten = await ProcessRequestCore((IPEndPoint)result.RemoteEndPoint, new ArraySegment<byte>(buffer, 0, result.ReceivedBytes), action, buffer.AsMemory(0, 512));
 
             await _dnsSocket.SendToAsync(buffer.AsMemory(0, bytesWritten), SocketFlags.None, result.RemoteEndPoint);
         }
@@ -120,15 +74,18 @@ public async Task ProcessUdpRequest(Func<LoopbackDnsResponseBuilder, Task> actio
         }
     }
 
-    public async Task ProcessTcpRequest(Func<LoopbackDnsResponseBuilder, Task> action)
+    public Task ProcessUdpRequest(Func<LoopbackDnsResponseBuilder, Task> action)
+    {
+        return ProcessUdpRequest((builder, _) => action(builder));
+    }
+
+    public async Task ProcessTcpRequest(Func<LoopbackDnsResponseBuilder, IPEndPoint, Task> action)
     {
         using Socket tcpClient = await _tcpSocket.AcceptAsync();
 
         byte[] buffer = ArrayPool<byte>.Shared.Rent(8 * 1024);
         try
         {
-            EndPoint remoteEndPoint = new IPEndPoint(IPAddress.Any, 0);
-
             int bytesRead = 0;
             int length = -1;
             while (length < 0 || bytesRead < length + 2)
@@ -143,7 +100,7 @@ public async Task ProcessTcpRequest(Func<LoopbackDnsResponseBuilder, Task> actio
                 }
             }
 
-            int bytesWritten = await ProcessRequestCore(new ArraySegment<byte>(buffer, 2, length), action, buffer.AsMemory(2));
+            int bytesWritten = await ProcessRequestCore((IPEndPoint)tcpClient.RemoteEndPoint!, new ArraySegment<byte>(buffer, 2, length), action, buffer.AsMemory(2));
             BinaryPrimitives.WriteUInt16BigEndian(buffer.AsSpan(0, 2), (ushort)bytesWritten);
             await tcpClient.SendAsync(buffer.AsMemory(0, bytesWritten + 2), SocketFlags.None);
         }
@@ -152,6 +109,11 @@ public async Task ProcessTcpRequest(Func<LoopbackDnsResponseBuilder, Task> actio
             ArrayPool<byte>.Shared.Return(buffer);
         }
     }
+
+    public Task ProcessTcpRequest(Func<LoopbackDnsResponseBuilder, Task> action)
+    {
+        return ProcessTcpRequest((builder, _) => action(builder));
+    }
 }
 
 internal sealed class LoopbackDnsResponseBuilder
@@ -176,6 +138,71 @@ public LoopbackDnsResponseBuilder(string name, QueryType type, QueryClass @class
     public List<DnsResourceRecord> Answers { get; } = new List<DnsResourceRecord>();
     public List<DnsResourceRecord> Authorities { get; } = new List<DnsResourceRecord>();
     public List<DnsResourceRecord> Additionals { get; } = new List<DnsResourceRecord>();
+
+    public int Write(Memory<byte> responseBuffer)
+    {
+        DnsDataWriter writer = new(responseBuffer);
+        if (!writer.TryWriteHeader(new DnsMessageHeader
+        {
+            TransactionId = TransactionId,
+            QueryFlags = Flags | (QueryFlags)ResponseCode,
+            QueryCount = (ushort)Questions.Count,
+            AnswerCount = (ushort)Answers.Count,
+            AuthorityCount = (ushort)Authorities.Count,
+            AdditionalRecordCount = (ushort)Additionals.Count
+        }))
+        {
+            throw new InvalidOperationException("Failed to write header");
+        }
+
+        foreach (var (questionName, questionType, questionClass) in Questions)
+        {
+            if (!writer.TryWriteQuestion(questionName, questionType, questionClass))
+            {
+                throw new InvalidOperationException("Failed to write question");
+            }
+        }
+
+        foreach (var answer in Answers)
+        {
+            if (!writer.TryWriteResourceRecord(answer))
+            {
+                throw new InvalidOperationException("Failed to write answer");
+            }
+        }
+
+        foreach (var authority in Authorities)
+        {
+            if (!writer.TryWriteResourceRecord(authority))
+            {
+                throw new InvalidOperationException("Failed to write authority");
+            }
+        }
+
+        foreach (var additional in Additionals)
+        {
+            if (!writer.TryWriteResourceRecord(additional))
+            {
+                throw new InvalidOperationException("Failed to write additional records");
+            }
+        }
+
+        return writer.Position;
+    }
+
+    public byte[] GetMessageBytes()
+    {
+        byte[] buffer = ArrayPool<byte>.Shared.Rent(512);
+        try
+        {
+            int bytesWritten = Write(buffer.AsMemory(0, 512));
+            return buffer.AsSpan(0, bytesWritten).ToArray();
+        }
+        finally
+        {
+            ArrayPool<byte>.Shared.Return(buffer);
+        }
+    }
 }
 
 internal static class LoopbackDnsServerExtensions

tests/Microsoft.Extensions.ServiceDiscovery.Dns.Tests/Resolver/ResolveAddressesTests.cs

@@ -109,7 +109,6 @@ public async Task ResolveIPv4_Aliases_OutOfOrder_Success()
     [Fact]
     public async Task ResolveIPv4_Aliases_Loop_ReturnsEmpty()
     {
-        IPAddress address = IPAddress.Parse("172.213.245.111");
         _ = DnsServer.ProcessUdpRequest(builder =>
         {
             builder.Answers.AddCname("www.example1.com", 3600, "www.example2.com");
@@ -123,6 +122,58 @@ public async Task ResolveIPv4_Aliases_Loop_ReturnsEmpty()
         Assert.Empty(results);
     }
 
+    [Fact]
+    public async Task ResolveIPv4_Aliases_Loop_Reverse_ReturnsEmpty()
+    {
+        _ = DnsServer.ProcessUdpRequest(builder =>
+        {
+            builder.Answers.AddCname("www.example3.com", 3600, "www.example1.com");
+            builder.Answers.AddCname("www.example2.com", 3600, "www.example3.com");
+            builder.Answers.AddCname("www.example1.com", 3600, "www.example2.com");
+            return Task.CompletedTask;
+        });
+
+        AddressResult[] results = await Resolver.ResolveIPAddressesAsync("www.example1.com", AddressFamily.InterNetwork);
+
+        Assert.Empty(results);
+    }
+
+    [Fact]
+    public async Task ResolveIPv4_Alias_And_Address()
+    {
+        IPAddress address = IPAddress.Parse("172.213.245.111");
+        _ = DnsServer.ProcessUdpRequest(builder =>
+        {
+            builder.Answers.AddCname("www.example1.com", 3600, "www.example2.com");
+            builder.Answers.AddCname("www.example2.com", 3600, "www.example3.com");
+            builder.Answers.AddAddress("www.example2.com", 3600, address);
+            return Task.CompletedTask;
+        });
+
+        AddressResult[] results = await Resolver.ResolveIPAddressesAsync("www.example1.com", AddressFamily.InterNetwork);
+
+        Assert.Empty(results);
+    }
+
+    [Fact]
+    public async Task ResolveIPv4_DuplicateAlias()
+    {
+        IPAddress address = IPAddress.Parse("172.213.245.111");
+        _ = DnsServer.ProcessUdpRequest(builder =>
+        {
+            builder.Answers.AddCname("www.example1.com", 3600, "www.example2.com");
+            builder.Answers.AddCname("www.example2.com", 3600, "www.example3.com");
+            builder.Answers.AddCname("www.example2.com", 3600, "www.example4.com");
+            builder.Answers.AddAddress("www.example2.com", 3600, address);
+            builder.Answers.AddAddress("www.example4.com", 3600, address);
+            return Task.CompletedTask;
+        });
+
+        AddressResult[] results = await Resolver.ResolveIPAddressesAsync("www.example1.com", AddressFamily.InterNetwork);
+
+        Assert.Empty(results);
+    }
+
     [Fact]
     public async Task ResolveIPv4_Aliases_NotFound_Success()
     {
@@ -167,4 +218,70 @@ public async Task Resolve_Timeout_ReturnsEmpty()
         AddressResult[] result = await Resolver.ResolveIPAddressesAsync("example.com", AddressFamily.InterNetwork);
         Assert.Empty(result);
     }
-}
+
+    [Theory]
+    [InlineData("not-example.com", (int)QueryType.A, (int)QueryClass.Internet)]
+    [InlineData("example.com", (int)QueryType.AAAA, (int)QueryClass.Internet)]
+    [InlineData("example.com", (int)QueryType.A, 0)]
+    public async Task Resolve_QuestionMismatch_ReturnsEmpty(string name, int type, int @class)
+    {
+        Options.Timeout = TimeSpan.FromSeconds(1);
+
+        IPAddress address = IPAddress.Parse("172.213.245.111");
+        _ = DnsServer.ProcessUdpRequest(builder =>
+        {
+            builder.Questions[0] = (name, (QueryType)type, (QueryClass)@class);
+            builder.Answers.AddAddress("www.example4.com", 3600, address);
+            return Task.CompletedTask;
+        });
+
+        AddressResult[] result = await Resolver.ResolveIPAddressesAsync("example.com", AddressFamily.InterNetwork);
+        Assert.Empty(result);
+    }
+
+    [Fact]
+    public async Task Resolve_HeaderMismatch_Ignores()
+    {
+        string name = "example.com";
+        Options.Timeout = TimeSpan.FromSeconds(5);
+
+        SemaphoreSlim responseSemaphore = new SemaphoreSlim(0, 1);
+        SemaphoreSlim requestSemaphore = new SemaphoreSlim(0, 1);
+
+        IPEndPoint clientAddress = null!;
+
+        IPAddress address = IPAddress.Parse("172.213.245.111");
+        ushort transactionId = 0x1234;
+        _ = DnsServer.ProcessUdpRequest((builder, clientAddr) =>
+        {
+            clientAddress = clientAddr;
+            transactionId = (ushort)(builder.TransactionId + 1);
+
+            builder.Answers.AddAddress(name, 3600, address);
+            requestSemaphore.Release();
+            return responseSemaphore.WaitAsync();
+        });
+
+        ValueTask<AddressResult[]> task = Resolver.ResolveIPAddressesAsync(name, AddressFamily.InterNetwork);
+
+        await requestSemaphore.WaitAsync().WaitAsync(Options.Timeout);
+
+        using Socket socket = new Socket(clientAddress.AddressFamily, SocketType.Dgram, ProtocolType.Udp);
+        LoopbackDnsResponseBuilder responseBuilder = new LoopbackDnsResponseBuilder(name, QueryType.A, QueryClass.Internet)
+        {
+            TransactionId = transactionId,
+            ResponseCode = QueryResponseCode.NoError
+        };
+
+        responseBuilder.Questions.Add((name, QueryType.A, QueryClass.Internet));
+        responseBuilder.Answers.AddAddress(name, 3600, IPAddress.Loopback);
+        socket.SendTo(responseBuilder.GetMessageBytes(), clientAddress);
+
+        responseSemaphore.Release();
+
+        AddressResult[] results = await task;
+        AddressResult result = Assert.Single(results);
+
+        Assert.Equal(address, result.Address);
+    }
+}

tests/Microsoft.Extensions.ServiceDiscovery.Dns.Tests/Resolver/TcpFailoverTests.cs

@@ -78,4 +78,53 @@ public async Task TcpFailover_TcpNotAvailable_EmptyResult()
         AddressResult[] results = await Resolver.ResolveIPAddressesAsync("www.example.com", AddressFamily.InterNetwork);
         Assert.Empty(results);
     }
+
+    [Fact]
+    public async Task TcpFailover_HeaderMismatch_ReturnsEmpty()
+    {
+        Options.Timeout = TimeSpan.FromSeconds(1);
+        IPAddress address = IPAddress.Parse("172.213.245.111");
+
+        _ = DnsServer.ProcessUdpRequest(builder =>
+        {
+            builder.Flags |= QueryFlags.ResultTruncated;
+            return Task.CompletedTask;
+        });
+
+        _ = DnsServer.ProcessTcpRequest(builder =>
+        {
+            builder.TransactionId++;
+            builder.Answers.AddAddress("www.example.com", 3600, address);
+            return Task.CompletedTask;
+        });
+
+        AddressResult[] result = await Resolver.ResolveIPAddressesAsync("example.com", AddressFamily.InterNetwork);
+        Assert.Empty(result);
+    }
+
+    [Theory]
+    [InlineData("not-example.com", (int)QueryType.A, (int)QueryClass.Internet)]
+    [InlineData("example.com", (int)QueryType.AAAA, (int)QueryClass.Internet)]
+    [InlineData("example.com", (int)QueryType.A, 0)]
+    public async Task TcpFailover_QuestionMismatch_ReturnsEmpty(string name, int type, int @class)
+    {
+        Options.Timeout = TimeSpan.FromSeconds(1);
+        IPAddress address = IPAddress.Parse("172.213.245.111");
+
+        _ = DnsServer.ProcessUdpRequest(builder =>
+        {
+            builder.Flags |= QueryFlags.ResultTruncated;
+            return Task.CompletedTask;
+        });
+
+        _ = DnsServer.ProcessTcpRequest(builder =>
+        {
+            builder.Questions[0] = (name, (QueryType)type, (QueryClass)@class);
+            builder.Answers.AddAddress("www.example.com", 3600, address);
+            return Task.CompletedTask;
+        });
+
+        AddressResult[] result = await Resolver.ResolveIPAddressesAsync("example.com", AddressFamily.InterNetwork);
+        Assert.Empty(result);
+    }
 }