More code review feedback

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

@@ -3,6 +3,7 @@
 
 using System.Buffers;
 using System.Buffers.Binary;
+using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.Runtime.InteropServices;
 
@@ -23,7 +24,9 @@ public DnsDataReader(ReadOnlyMemory<byte> buffer, byte[]? returnToPool = null)
 
     public bool TryReadHeader(out DnsMessageHeader header)
     {
-        if (_buffer.Length - _position < DnsMessageHeader.HeaderLength)
+        Debug.Assert(_position == 0);
+
+        if (_buffer.Length < DnsMessageHeader.HeaderLength)
         {
             header = default;
             return false;

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

@@ -2,21 +2,28 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Buffers.Binary;
+using System.Runtime.InteropServices;
 using System.Security.Cryptography;
 
 namespace Microsoft.Extensions.ServiceDiscovery.Dns.Resolver;
 
 // RFC 1035 4.1.1. Header section format
+[StructLayout(LayoutKind.Explicit, Size = HeaderLength)]
 internal struct DnsMessageHeader
 {
     internal const int HeaderLength = 12;
 
+    [FieldOffset(0)]
     private ushort _transactionId;
+    [FieldOffset(2)]
     private ushort _flags;
-
+    [FieldOffset(4)]
     private ushort _queryCount;
+    [FieldOffset(6)]
     private ushort _answerCount;
+    [FieldOffset(8)]
     private ushort _authorityCount;
+    [FieldOffset(10)]
     private ushort _additionalRecordCount;
 
     internal ushort QueryCount

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

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Buffers.Binary;
+using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.Text;
 
@@ -25,7 +26,11 @@ internal static bool TryWriteQName(Span<byte> destination, string name, out int
         //     padding is used.
         //
 
-        if (!Encoding.ASCII.TryGetBytes(name, destination.IsEmpty ? destination : destination.Slice(1), out int length) || destination.Length < length + 2)
+        // The is assumed to be already validated and puny-encoded if needed
+        Debug.Assert(name.Length <= MaxDomainNameLength);
+        Debug.Assert(Ascii.IsValid(name));
+
+        if (destination.IsEmpty || !Encoding.ASCII.TryGetBytes(name, destination.Slice(1), out int length) || destination.Length < length + 2)
         {
             // buffer too small
             written = 0;
@@ -38,13 +43,14 @@ internal static bool TryWriteQName(Span<byte> destination, string name, out int
         while (true)
         {
             // figure out the next label and prepend the length
-            int index = nameBuffer.Slice(1).IndexOf<byte>((byte)'.');
+            int index = nameBuffer.Slice(1).IndexOf((byte)'.');
             int labelLen = index == -1 ? nameBuffer.Length - 1 : index;
 
             // https://www.rfc-editor.org/rfc/rfc1035#section-2.3.4
             // labels          63 octets or less
             if (labelLen > 63)
             {
+                // this should never happen, as we validate the name before calling this method
                 throw new ArgumentException("Label is too long");
             }
 
@@ -197,41 +203,6 @@ internal static bool TryReadService(ReadOnlySpan<byte> buffer, out ushort priori
         return true;
     }
 
-    internal static bool TryWriteService(Span<byte> buffer, ushort priority, ushort weight, ushort port, string target, out int bytesWritten)
-    {
-        // https://www.rfc-editor.org/rfc/rfc2782
-        if (!BinaryPrimitives.TryWriteUInt16BigEndian(buffer, priority) ||
-            !BinaryPrimitives.TryWriteUInt16BigEndian(buffer.Slice(2), weight) ||
-            !BinaryPrimitives.TryWriteUInt16BigEndian(buffer.Slice(4), port) ||
-            !TryWriteQName(buffer.Slice(6), target, out bytesWritten))
-        {
-            bytesWritten = 0;
-            return false;
-        }
-
-        bytesWritten += 6;
-        return true;
-    }
-
-    internal static bool TryWriteSoa(Span<byte> buffer, string primaryNameServer, string responsibleMailAddress, uint serial, uint refresh, uint retry, uint expire, uint minimum, out int bytesWritten)
-    {
-        // https://www.rfc-editor.org/rfc/rfc1035#section-3.3.13
-        if (!TryWriteQName(buffer, primaryNameServer, out int w1) ||
-            !TryWriteQName(buffer.Slice(w1), responsibleMailAddress, out int w2) ||
-            !BinaryPrimitives.TryWriteUInt32BigEndian(buffer.Slice(w1 + w2), serial) ||
-            !BinaryPrimitives.TryWriteUInt32BigEndian(buffer.Slice(w1 + w2 + 4), refresh) ||
-            !BinaryPrimitives.TryWriteUInt32BigEndian(buffer.Slice(w1 + w2 + 8), retry) ||
-            !BinaryPrimitives.TryWriteUInt32BigEndian(buffer.Slice(w1 + w2 + 12), expire) ||
-            !BinaryPrimitives.TryWriteUInt32BigEndian(buffer.Slice(w1 + w2 + 16), minimum))
-        {
-            bytesWritten = 0;
-            return false;
-        }
-
-        bytesWritten = w1 + w2 + 20;
-        return true;
-    }
-
     internal static bool TryReadSoa(ReadOnlySpan<byte> buffer, [NotNullWhen(true)] out string? primaryNameServer, [NotNullWhen(true)] out string? responsibleMailAddress, out uint serial, out uint refresh, out uint retry, out uint expire, out uint minimum, out int bytesRead)
     {
         // https://www.rfc-editor.org/rfc/rfc1035#section-3.3.13

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

@@ -552,6 +552,7 @@ internal static SendQueryError ValidateResponse(in DnsResponse response)
                 throw new InvalidOperationException("Invalid response: Header mismatch");
             }
 
+            // transfer ownership of buffer to the caller
             buffer = null!;
             return (responseReader, header);
         }

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

@@ -61,7 +61,8 @@ private static async Task<int> ProcessRequestCore(ReadOnlyMemory<byte> message,
         }))
         {
             throw new InvalidOperationException("Failed to write header");
-        };
+        }
+        ;
 
         foreach (var (questionName, questionType, questionClass) in responseBuilder.Questions)
         {
@@ -198,23 +199,40 @@ public static List<DnsResourceRecord> AddCname(this List<DnsResourceRecord> reco
     public static List<DnsResourceRecord> AddService(this List<DnsResourceRecord> records, string name, int ttl, ushort priority, ushort weight, ushort port, string target)
     {
         byte[] buff = new byte[256];
-        if (!DnsPrimitives.TryWriteService(buff, priority, weight, port, target, out int length))
+
+        // https://www.rfc-editor.org/rfc/rfc2782
+        if (!BinaryPrimitives.TryWriteUInt16BigEndian(buff, priority) ||
+            !BinaryPrimitives.TryWriteUInt16BigEndian(buff.AsSpan(2), weight) ||
+            !BinaryPrimitives.TryWriteUInt16BigEndian(buff.AsSpan(4), port) ||
+            !DnsPrimitives.TryWriteQName(buff.AsSpan(6), target, out int length))
         {
             throw new InvalidOperationException("Failed to encode SRV record");
         }
 
+        length += 6;
+
         records.Add(new DnsResourceRecord(name, QueryType.SRV, QueryClass.Internet, ttl, buff.AsMemory(0, length)));
         return records;
     }
 
     public static List<DnsResourceRecord> AddStartOfAuthority(this List<DnsResourceRecord> records, string name, int ttl, string mname, string rname, uint serial, uint refresh, uint retry, uint expire, uint minimum)
     {
         byte[] buff = new byte[256];
-        if (!DnsPrimitives.TryWriteSoa(buff, mname, rname, serial, refresh, retry, expire, minimum, out int length))
+
+        // https://www.rfc-editor.org/rfc/rfc1035#section-3.3.13
+        if (!DnsPrimitives.TryWriteQName(buff, mname, out int w1) ||
+            !DnsPrimitives.TryWriteQName(buff.AsSpan(w1), rname, out int w2) ||
+            !BinaryPrimitives.TryWriteUInt32BigEndian(buff.AsSpan(w1 + w2), serial) ||
+            !BinaryPrimitives.TryWriteUInt32BigEndian(buff.AsSpan(w1 + w2 + 4), refresh) ||
+            !BinaryPrimitives.TryWriteUInt32BigEndian(buff.AsSpan(w1 + w2 + 8), retry) ||
+            !BinaryPrimitives.TryWriteUInt32BigEndian(buff.AsSpan(w1 + w2 + 12), expire) ||
+            !BinaryPrimitives.TryWriteUInt32BigEndian(buff.AsSpan(w1 + w2 + 16), minimum))
         {
             throw new InvalidOperationException("Failed to encode SOA record");
         }
 
+        int length = w1 + w2 + 20;
+
         records.Add(new DnsResourceRecord(name, QueryType.SOA, QueryClass.Internet, ttl, buff.AsMemory(0, length)));
         return records;
     }