Packet multiplexing: async nvarchar(max) read returns corrupted value, throws ArgumentException or spins when the value starts in a packet with header ID 0

<Project Sdk="Microsoft.NET.Sdk">
  <PropertyGroup>
    <OutputType>Exe</OutputType>
    <TargetFramework>net10.0</TargetFramework>
    <ImplicitUsings>enable</ImplicitUsings>
  </PropertyGroup>
  <ItemGroup>
    <PackageReference Include="Microsoft.Data.SqlClient" Version="7.1.1" />
  </ItemGroup>
</Project>

// Repro: async reads of nvarchar(max) with packet multiplexing enabled return corrupted strings,
// throw ArgumentException from TryReadPlpUnicodeChars, or spin forever, when a value starts in a
// TDS packet whose header packet ID is 0.
//
// Usage: dotnet run -c Release -- [--server 127.0.0.1,1433] [--packet-ids asis|nozero] [--compat]
// Auth: Integrated Security, or set SQL_AUTH="User ID=...;Password=...".
using System.Collections.Concurrent;
using System.Diagnostics;
using System.Net;
using System.Net.Sockets;
using System.Text;
using Microsoft.Data.SqlClient;

bool compat = args.Contains("--compat");
AppContext.SetSwitch("Switch.Microsoft.Data.SqlClient.UseCompatibilityAsyncBehaviour", compat);
AppContext.SetSwitch("Switch.Microsoft.Data.SqlClient.UseCompatibilityProcessSni", compat);

string server = Arg("--server", "127.0.0.1,1433");
string packetIds = Arg("--packet-ids", "asis");
int iterations = int.Parse(Arg("--iterations", "40"));
string auth = Environment.GetEnvironmentVariable("SQL_AUTH") ?? "Integrated Security=true";
const string Db = "plp_repro";
const int Rows = 3000;

int[] lengths = Lengths(Rows);
await Setup($"Server={server};{auth};Encrypt=False;Database=master;Pooling=false");

using var proxy = new ChoppyProxy(server, packetIds);
string cs = $"Server=127.0.0.1,{proxy.Port};{auth};Encrypt=False;Database={Db};Pooling=false;Command Timeout=600";
Console.WriteLine($"SqlClient {FileVersionInfo.GetVersionInfo(typeof(SqlConnection).Assembly.Location).ProductVersion}, compat={compat}, packet IDs {packetIds}");

int corrupt = 0, hung = 0, shown = 0, shownErrors = 0;
var errors = new ConcurrentDictionary<string, int>();
await Parallel.ForAsync(0, iterations, new ParallelOptions { MaxDegreeOfParallelism = 4 }, async (_, ct) =>
{
    var read = Task.Run(ReadAll, ct);
    // A stuck read spins inside SqlClient and ignores cancellation, so it is abandoned.
    if (await Task.WhenAny(read, Task.Delay(TimeSpan.FromSeconds(30), ct)) != read)
    {
        Interlocked.Increment(ref hung);
    }
});

Console.WriteLine($"Reads: {iterations}, corrupt values: {corrupt}, hung reads: {hung}, server packets with ID 0: {proxy.PacketsWithIdZero}");
foreach (var (error, count) in errors)
{
    Console.WriteLine($"  {count} x {error}");
}

Console.Out.Flush();
Environment.Exit(corrupt + hung + errors.Count == 0 ? 0 : 1);

async Task ReadAll()
{
    try
    {
        await using var conn = new SqlConnection(cs);
        await conn.OpenAsync();
        await using var cmd = new SqlCommand("SELECT Id, Body FROM dbo.items ORDER BY Id", conn);
        await using var reader = await cmd.ExecuteReaderAsync();
        while (await reader.ReadAsync())
        {
            int id = reader.GetInt32(0);
            string body = reader.GetString(1);
            if (body != Expected(id, lengths[id]))
            {
                Interlocked.Increment(ref corrupt);
                if (Interlocked.Increment(ref shown) <= 5)
                {
                    Console.WriteLine($"Row {id}: expected {lengths[id]} chars, got {body.Length}");
                }
            }
        }
    }
    catch (Exception ex)
    {
        errors.AddOrUpdate($"{ex.GetType().Name}: {ex.Message}", 1, static (_, n) => n + 1);
        if (Interlocked.Increment(ref shownErrors) == 1)
        {
            Console.WriteLine(ex);
        }
    }
}

string Arg(string name, string fallback)
{
    int i = Array.IndexOf(args, name);
    return i >= 0 && i + 1 < args.Length ? args[i + 1] : fallback;
}

static int[] Lengths(int rows)
{
    var rnd = new Random(1);
    return Enumerable.Range(0, rows).Select(_ =>
    {
        double p = rnd.NextDouble();
        return p < 0.60 ? rnd.Next(0, 400) : p < 0.90 ? rnd.Next(400, 4000) : p < 0.98 ? rnd.Next(4000, 20000) : rnd.Next(20000, 120000);
    }).ToArray();
}

static string Expected(int id, int length)
{
    const string pattern = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789";
    var sb = new StringBuilder(length);
    for (int k = 0; k < length; k++)
    {
        sb.Append(pattern[(id + k) % pattern.Length]);
    }

    return sb.ToString();
}

async Task Setup(string master)
{
    await using (var conn = new SqlConnection(master))
    {
        await conn.OpenAsync();
        await Exec(conn, $"IF DB_ID('{Db}') IS NOT NULL BEGIN ALTER DATABASE [{Db}] SET SINGLE_USER WITH ROLLBACK IMMEDIATE; DROP DATABASE [{Db}]; END; CREATE DATABASE [{Db}];");
    }

    await using var c = new SqlConnection(master.Replace("Database=master", $"Database={Db}"));
    await c.OpenAsync();
    await Exec(c, "CREATE TABLE dbo.items (Id int NOT NULL PRIMARY KEY, Len int NOT NULL, Body nvarchar(max) NOT NULL DEFAULT N'');");
    for (int start = 0; start < Rows; start += 1000)
    {
        var values = Enumerable.Range(start, Math.Min(1000, Rows - start)).Select(i => $"({i},{lengths[i]})");
        await Exec(c, "INSERT INTO dbo.items (Id, Len) VALUES " + string.Join(",", values));
    }

    // Same content as Expected(): the pattern rotated by Id, cut to Len.
    await Exec(c, "UPDATE dbo.items SET Body = SUBSTRING(REPLICATE(CAST(N'ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789' AS nvarchar(max)), Len / 36 + 2), Id % 36 + 1, Len);");

    static async Task Exec(SqlConnection c, string sql)
    {
        await using var cmd = new SqlCommand(sql, c) { CommandTimeout = 600 };
        await cmd.ExecuteNonQueryAsync();
    }
}

/// <summary>
/// TCP proxy that hands the server's bytes to the client in random small pieces with short pauses,
/// so async reads run out of data in the middle of a value, as on a real network.
/// With "nozero" it renumbers the TDS header packet IDs 1..255 so that no packet has ID 0.
/// </summary>
sealed class ChoppyProxy : IDisposable
{
    private readonly TcpListener _listener = new(IPAddress.Loopback, 0);
    private readonly string _server;
    private readonly string _packetIds;
    private long _packetsWithIdZero;
    private int _connections;

    public ChoppyProxy(string server, string packetIds)
    {
        _server = server;
        _packetIds = packetIds;
        _listener.Start();
        _ = Task.Run(async () =>
        {
            while (true)
            {
                var client = await _listener.AcceptTcpClientAsync();
                _ = Task.Run(() => Handle(client, Interlocked.Increment(ref _connections)));
            }
        });
    }

    public int Port => ((IPEndPoint)_listener.LocalEndpoint).Port;
    public long PacketsWithIdZero => Interlocked.Read(ref _packetsWithIdZero);

    private async Task Handle(TcpClient client, int seed)
    {
        using var _ = client;
        using var server = new TcpClient();
        string[] hostPort = _server.Split(',');
        await server.ConnectAsync(hostPort[0], hostPort.Length > 1 ? int.Parse(hostPort[1]) : 1433);
        client.NoDelay = server.NoDelay = true;
        var up = client.GetStream().CopyToAsync(server.GetStream());
        var down = Down(server.GetStream(), client.GetStream(), new Random(seed));
        await Task.WhenAny(up, down);
    }

    private async Task Down(NetworkStream from, NetworkStream to, Random rnd)
    {
        var buffer = new byte[65536];
        int headerPos = 0, packetLength = 0, payloadLeft = 0, nextId = 0;
        try
        {
            int n;
            while ((n = await from.ReadAsync(buffer)) > 0)
            {
                // Walk the 8-byte TDS packet headers: bytes 2-3 are the packet length, byte 6 is the packet ID.
                for (int i = 0; i < n; i++)
                {
                    if (payloadLeft > 0)
                    {
                        int skip = Math.Min(payloadLeft, n - i);
                        payloadLeft -= skip;
                        i += skip - 1;
                        continue;
                    }

                    switch (headerPos)
                    {
                        case 2:
                            packetLength = buffer[i] << 8;
                            break;
                        case 3:
                            packetLength |= buffer[i];
                            break;
                        case 6:
                            if (_packetIds == "nozero")
                            {
                                buffer[i] = (byte)(nextId = (nextId % 255) + 1);
                            }

                            if (buffer[i] == 0)
                            {
                                Interlocked.Increment(ref _packetsWithIdZero);
                            }

                            break;
                    }

                    if (++headerPos == 8)
                    {
                        headerPos = 0;
                        payloadLeft = packetLength - 8;
                    }
                }

                for (int offset = 0, chunk; offset < n; offset += chunk)
                {
                    chunk = Math.Min(n - offset, rnd.Next(1, 3001));
                    await to.WriteAsync(buffer.AsMemory(offset, chunk));
                    if (rnd.NextDouble() < 0.3)
                    {
                        // ~200 µs busy pause: the reader drains what it has and has to wait for more.
                        long until = Stopwatch.GetTimestamp() + (Stopwatch.Frequency / 5000);
                        while (Stopwatch.GetTimestamp() < until)
                        {
                        }
                    }
                }
            }
        }
        catch (IOException)
        {
        }
    }

    public void Dispose() => _listener.Stop();
}