fix(abcip): budget expected reply size in RequestScheduler.TakeBatch (#582); private lock + struct-handle check across read fragments (#567)

Claude-Session: https://claude.ai/code/session_01Ef2ozxS7QPanGT8uoATbTe

src/Drivers/ZB.MOM.WW.OtOpcUa.Driver.Ab.Transport/Cip/RequestScheduler.cs

@@ -19,8 +19,21 @@ public sealed class RequestScheduler : IAsyncDisposable
     private const byte MultipleServiceCode = 0x0A;
     private const byte EmbeddedServiceError = 0x1E;
 
-    private sealed class Pending(byte[] request, bool exclusive)
+    /// <summary>Smallest possible CIP reply: service, reserved, general status, additional-status size.</summary>
+    public const int ReplyHeaderBytes = 4;
+
+    /// <summary>The <c>expectedReplyBytes</c> default: the reply carries no data beyond its header (a
+    /// write, a status-only service). Every caller whose reply carries data must pass its real bound.</summary>
+    public const int StatusOnlyReply = ReplyHeaderBytes;
+
+    /// <summary>Pass as <c>expectedReplyBytes</c> when the reply size cannot be bounded: the request is
+    /// then sent alone, never packed, because its reply could fill the whole connection.</summary>
+    public const int UnknownReply = -1;
+
+    private sealed class Pending(byte[] request, bool exclusive, int expectedReplyBytes)
     {
+        // Private lock object: lock(this) would let any holder of the Pending reference contend.
+        private readonly object _gate = new();
         private bool _completed;
         private CancellationTokenRegistration _registration;
         private bool _hasRegistration;
@@ -29,13 +42,16 @@ public sealed class RequestScheduler : IAsyncDisposable
 
         /// <summary>True when this request must never share a Multiple Service Packet with another.</summary>
         public bool Exclusive { get; } = exclusive;
+
+        /// <summary>Expected whole-reply size (header included), or <see cref="UnknownReply"/>.</summary>
+        public int ExpectedReplyBytes { get; } = expectedReplyBytes;
         public TaskCompletionSource<CipReply> Tcs { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously);
         public LinkedListNode<Pending>? Node { get; set; }
 
         /// <summary>Stores the waiter's registration; disposes it at once if the waiter already completed.</summary>
         public void SetRegistration(CancellationTokenRegistration registration, RequestScheduler owner)
         {
-            lock (this)
+            lock (_gate)
             {
                 if (!_completed)
                 {
@@ -54,7 +70,7 @@ public sealed class RequestScheduler : IAsyncDisposable
         {
             CancellationTokenRegistration registration;
             bool hasRegistration;
-            lock (this)
+            lock (_gate)
             {
                 if (_completed) return false;
                 _completed = true;
@@ -125,13 +141,19 @@ public sealed class RequestScheduler : IAsyncDisposable
     /// <param name="request">The CIP request.</param>
     /// <param name="ct">Cancels this waiter only.</param>
     /// <param name="exclusive">When true the request is sent alone, never packed into a Multiple Service
-    /// Packet: for requests whose reply can fill the whole connection (the packer budgets request bytes only).</param>
-    public Task<CipReply> SendAsync(byte[] request, CancellationToken ct, bool exclusive = false)
+    /// Packet: for requests whose reply can fill the whole connection.</param>
+    /// <param name="expectedReplyBytes">Upper bound on this request's whole reply (header included) as
+    /// it would sit inside a Multiple Service Packet reply. The packer budgets this against the
+    /// connection size as well as the request bytes (#582), so a batch of small requests with large
+    /// replies (STRINGs, arrays) is split instead of overflowing. Defaults to <see cref="StatusOnlyReply"/>;
+    /// <see cref="UnknownReply"/> sends the request alone.</param>
+    public Task<CipReply> SendAsync(byte[] request, CancellationToken ct, bool exclusive = false,
+        int expectedReplyBytes = StatusOnlyReply)
     {
         ArgumentNullException.ThrowIfNull(request);
         if (ct.IsCancellationRequested) return Task.FromCanceled<CipReply>(ct);
 
-        var p = new Pending(request, exclusive);
+        var p = new Pending(request, exclusive || expectedReplyBytes < 0, Math.Max(expectedReplyBytes, UnknownReply));
         lock (_lock)
         {
             ObjectDisposedException.ThrowIf(_disposed, this);
@@ -168,19 +190,27 @@ public sealed class RequestScheduler : IAsyncDisposable
         lock (_lock)
         {
             var sizes = new List<int>();
+            var replySizes = new List<int>();
+            var replyBudget = _channel.MaxRequestBytes; // connections are symmetric: one size bounds both directions
             while (_queue.First is { } n)
             {
                 if (batch.Count > 0)
                 {
                     if (!_allowPacking || n.Value.Exclusive) break;
                     sizes.Add(n.Value.Request.Length);
-                    if (MultipleServicePacket.BuildSize(sizes) > _channel.MaxRequestBytes) break;
+                    replySizes.Add(n.Value.ExpectedReplyBytes);
+                    if (MultipleServicePacket.BuildSize(sizes) > _channel.MaxRequestBytes
+                        || ReplyEnvelopeSize(replySizes) > replyBudget)
+                    {
+                        break;
+                    }
                 }
                 else
                 {
                     // The first request always goes, even when it alone exceeds the limit: it is sent
                     // bare, and callers that must fit (LogixTagClient) fragment before they get here.
                     sizes.Add(n.Value.Request.Length);
+                    replySizes.Add(n.Value.ExpectedReplyBytes);
                 }
                 batch.Add(n.Value);
                 _queue.RemoveFirst();
@@ -189,6 +219,15 @@ public sealed class RequestScheduler : IAsyncDisposable
         }
     }
 
+    /// <summary>Expected size of the MSP reply for the given embedded reply sizes: reply header, count,
+    /// one offset per reply, then the replies themselves.</summary>
+    internal static int ReplyEnvelopeSize(IReadOnlyList<int> replySizes)
+    {
+        var total = ReplyHeaderBytes + 2 + 2 * replySizes.Count;
+        foreach (var r in replySizes) total += r;
+        return total;
+    }
+
     private async Task PumpAsync()
     {
         while (true)

src/Drivers/ZB.MOM.WW.OtOpcUa.Driver.Ab.Transport/Logix/LogixTagClient.cs

@@ -38,11 +38,18 @@ public sealed class LogixTagClient(RequestScheduler scheduler, ICipChannel chann
     /// </summary>
     /// <remarks>Packing bounds only the request size, so a packed batch of large reads can come back
     /// with 0x06 per service; that is expected, and this loop recovers it.</remarks>
+    /// <param name="path">The encoded symbolic path.</param>
+    /// <param name="elements">Element count to read.</param>
+    /// <param name="ct">Cancels this read.</param>
+    /// <param name="elementBytes">Upper bound on one element's data bytes (<see cref="LogixTypeCode.StringElementBytes"/>
+    /// for a STRING), or 0 when unknown. The scheduler budgets the expected reply against the connection
+    /// size when packing (#582); unknown sends the read alone.</param>
     public async Task<(ushort TypeCode, ushort? StructHandle, byte[] Data)> ReadAsync(
-        byte[] path, ushort elements, CancellationToken ct)
+        byte[] path, ushort elements, CancellationToken ct, int elementBytes = 0)
     {
         ArgumentNullException.ThrowIfNull(path);
-        var reply = await _scheduler.SendAsync(LogixServices.ReadTagFragmented(path, elements, 0), ct).ConfigureAwait(false);
+        var reply = await _scheduler.SendAsync(LogixServices.ReadTagFragmented(path, elements, 0), ct,
+            expectedReplyBytes: ExpectedReadReplyBytes(elements, elementBytes, 0)).ConfigureAwait(false);
         var (type, handle, first) = LogixServices.ParseReadReply(reply); // throws on any status but 0x00/0x06
         var data = new List<byte>(first.Length);
         data.AddRange(first.Span);
@@ -53,11 +60,17 @@ public sealed class LogixTagClient(RequestScheduler scheduler, ICipChannel chann
             if (lastChunk == 0)
                 throw new AbTransportException(AbStatus.BadReply, "Partial read reply carried no data; the fragment loop cannot progress.");
             reply = await _scheduler.SendAsync(
-                LogixServices.ReadTagFragmented(path, elements, (uint)data.Count), ct).ConfigureAwait(false);
-            var (fragmentType, _, chunk) = LogixServices.ParseReadReply(reply);
+                LogixServices.ReadTagFragmented(path, elements, (uint)data.Count), ct,
+                expectedReplyBytes: ExpectedReadReplyBytes(elements, elementBytes, data.Count)).ConfigureAwait(false);
+            var (fragmentType, fragmentHandle, chunk) = LogixServices.ParseReadReply(reply);
             if (fragmentType != type)
                 throw new AbTransportException(AbStatus.BadReply,
                     $"Read fragment type 0x{fragmentType:X4} differs from the first reply's 0x{type:X4}.");
+            // A mid-read change of the structure handle means the tag's type changed under us (a
+            // download / online edit); stitching the halves would yield a value of neither type.
+            if (fragmentHandle != handle)
+                throw new AbTransportException(AbStatus.BadReply,
+                    $"Read fragment structure handle {FormatHandle(fragmentHandle)} differs from the first reply's {FormatHandle(handle)}; the tag changed mid-read.");
             data.AddRange(chunk.Span);
             lastChunk = chunk.Length;
             if (data.Count > MaxReadBytes)
@@ -68,6 +81,21 @@ public sealed class LogixTagClient(RequestScheduler scheduler, ICipChannel chann
         return (type, handle, data.ToArray());
     }
 
+    private static string FormatHandle(ushort? handle) => handle is { } h ? $"0x{h:X4}" : "(none)";
+
+    /// <summary>Whole-reply bound for a Read Tag Fragmented at <paramref name="alreadyRead"/> bytes in:
+    /// reply header, type (+ struct handle), then the remaining data. Unknown element size, or an
+    /// unrepresentable total, is <see cref="RequestScheduler.UnknownReply"/>.</summary>
+    internal static int ExpectedReadReplyBytes(ushort elements, int elementBytes, int alreadyRead)
+    {
+        if (elementBytes <= 0) return RequestScheduler.UnknownReply;
+        var remaining = (long)elements * elementBytes - alreadyRead;
+        if (remaining < 0) remaining = 0;
+        // 4 header + 2 type + 2 struct handle (counted for every type: a safe over-estimate).
+        var total = RequestScheduler.ReplyHeaderBytes + 4 + remaining;
+        return total > int.MaxValue ? RequestScheduler.UnknownReply : (int)total;
+    }
+
     /// <summary>
     /// Write <paramref name="data"/> to the tag at <paramref name="path"/>. One Write Tag when the
     /// request fits <see cref="ICipChannel.MaxRequestBytes"/>; otherwise Write Tag Fragmented chunks of

src/Drivers/ZB.MOM.WW.OtOpcUa.Driver.Ab.Transport/Logix/LogixTypeCode.cs

@@ -5,6 +5,9 @@ public static class LogixTypeCode
     public const ushort Bool = 0xC1, SInt = 0xC2, Int = 0xC3, DInt = 0xC4, LInt = 0xC5, USInt = 0xC6,
         UInt = 0xC7, UDInt = 0xC8, ULInt = 0xC9, Real = 0xCA, LReal = 0xCB, BoolArray = 0xD3, Struct = 0x02A0;
 
+    /// <summary>Bytes of one Logix STRING element on the wire: 4-byte length + 82 characters + 2 pad.</summary>
+    public const int StringElementBytes = 88;
+
     /// <summary>Byte width of an atomic type, or -1 when the code is not atomic (e.g. <see cref="Struct"/>).</summary>
     public static int ByteWidth(ushort code) => code switch
     {

src/Drivers/ZB.MOM.WW.OtOpcUa.Driver.AbCip/NativeAbCipTagRuntime.cs

@@ -122,7 +122,7 @@ internal sealed class NativeAbCipTagRuntime : IAbCipTagRuntime
         var path = Path();
         try
         {
-            var (type, handle, data) = await client.ReadAsync(path,_elements, cancellationToken).ConfigureAwait(false);
+            var (type, handle, data) = await client.ReadAsync(path, _elements, cancellationToken, ExpectedElementBytes(_p.AuthoredType)).ConfigureAwait(false);
             _data = data;
             _typeCode = type;
             _structHandle = handle;
@@ -136,6 +136,21 @@ internal sealed class NativeAbCipTagRuntime : IAbCipTagRuntime
         }
     }
 
+    /// <summary>
+    ///     Upper bound on one element's wire bytes for the scheduler's reply budget (#582); 0 when the
+    ///     authored type is unknown or a structure, which sends the read alone rather than guess.
+    ///     BOOL counts 4 (a packed BOOL array reads back as DWORDs).
+    /// </summary>
+    internal static int ExpectedElementBytes(AbCipDataType? type) => type switch
+    {
+        AbCipDataType.Bool or AbCipDataType.DInt or AbCipDataType.UDInt or AbCipDataType.Real or AbCipDataType.Dt => 4,
+        AbCipDataType.SInt or AbCipDataType.USInt => 1,
+        AbCipDataType.Int or AbCipDataType.UInt => 2,
+        AbCipDataType.LInt or AbCipDataType.ULInt or AbCipDataType.LReal => 8,
+        AbCipDataType.String => LogixTypeCode.StringElementBytes,
+        _ => 0,
+    };
+
     /// <inheritdoc />
     public async Task WriteAsync(CancellationToken cancellationToken)
     {

src/Drivers/ZB.MOM.WW.OtOpcUa.Driver.AbLegacy/NativeAbLegacyTagRuntime.cs

@@ -175,11 +175,15 @@ internal sealed class NativeAbLegacyTagRuntime : IAbLegacyTagRuntime
         return (Math.Max(1, PcccServices.MaxDataBytes / _elementSize), address);
     }
 
-    private async Task<ReadOnlyMemory<byte>> ExchangeAsync(Func<ushort, byte[]> pccc, CancellationToken ct)
+    /// <summary>Bytes of an Execute PCCC reply beyond its data: CIP header, requestor id, PCCC command / status / TNS
+    /// (rounded up) — the scheduler budgets <c>PcccReplyOverhead + data</c> against the connection (#582).</summary>
+    private const int PcccReplyOverhead = 32;
+
+    private async Task<ReadOnlyMemory<byte>> ExchangeAsync(Func<ushort, byte[]> pccc, CancellationToken ct, int replyDataBytes = 0)
     {
         var tns = unchecked((ushort)Interlocked.Increment(ref s_tns));
         var request = PcccServices.ExecutePccc(_vendor, _serial, pccc(tns));
-        var reply = await _scheduler.SendAsync(request, ct).ConfigureAwait(false);
+        var reply = await _scheduler.SendAsync(request, ct, expectedReplyBytes: PcccReplyOverhead + replyDataBytes).ConfigureAwait(false);
         return PcccServices.ParseReply(reply, tns, _requestorId);
     }
 
@@ -200,7 +204,7 @@ internal sealed class NativeAbLegacyTagRuntime : IAbLegacyTagRuntime
                 tns => plc5
                     ? PcccServices.Plc5RangeRead(tns, at, (ushort)(bytes / 2), bytes)
                     : PcccServices.SlcTypedRead(tns, at, bytes),
-                ct).ConfigureAwait(false);
+                ct, bytes).ConfigureAwait(false);
             if (data.Length < bytes)
                 throw new AbTransportException(AbStatus.BadReply, $"PCCC read returned {data.Length} of {bytes} bytes.");
             if (data.Length > bytes)