Stricter checks in noise; more atomicity in yamux

Author: flcl42

src/libp2p/Libp2p.Protocols.Noise/NoiseProtocol.cs

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2024 Demerzel Solutions Limited
+// SPDX-FileCopyrightText: 2025 Demerzel Solutions Limited
 // SPDX-License-Identifier: MIT
 
 using System.Buffers;
@@ -239,7 +239,7 @@ private static Task ExchangeData(Transport transport, IChannel downChannel, ICha
 
                 int bytesWritten = transport.WriteMessage(dataReadResult.Data.ToArray(), buffer.AsSpan(2));
                 BinaryPrimitives.WriteUInt16BigEndian(buffer.AsSpan(), (ushort)bytesWritten);
-                IOResult writeResult = await downChannel.WriteAsync(new ReadOnlySequence<byte>(buffer));
+                IOResult writeResult = await downChannel.WriteAsync(new ReadOnlySequence<byte>(buffer, 0, 2 + bytesWritten));
                 if (writeResult != IOResult.Ok)
                 {
                     logger?.LogDebug("End sending, due to {}", writeResult);
@@ -271,11 +271,14 @@ private static Task ExchangeData(Transport transport, IChannel downChannel, ICha
 
                 int bytesRead = transport.ReadMessage(dataReadResult.Data.ToArray(), buffer);
 
-                IOResult writeResult = await upChannel.WriteAsync(new ReadOnlySequence<byte>(buffer, 0, bytesRead));
-                if (writeResult != IOResult.Ok)
+                if (bytesRead != 0)
                 {
-                    logger?.LogDebug("Receiving data failed due to {}", dataReadResult);
-                    return;
+                    IOResult writeResult = await upChannel.WriteAsync(new ReadOnlySequence<byte>(buffer, 0, bytesRead));
+                    if (writeResult != IOResult.Ok)
+                    {
+                        logger?.LogDebug("Receiving data failed due to {}", dataReadResult);
+                        return;
+                    }
                 }
             }
         });

src/libp2p/Libp2p.Protocols.Yamux/YamuxProtocol.cs

@@ -6,6 +6,7 @@
 using Nethermind.Libp2p.Core.Exceptions;
 using Nethermind.Libp2p.Protocols.Yamux;
 using System.Buffers;
+using System.Collections.Concurrent;
 using System.Diagnostics;
 using System.Runtime.CompilerServices;
 
@@ -39,7 +40,7 @@ protected override async Task ConnectAsync(IChannel channel, IConnectionContext
         TaskAwaiter downChannelAwaiter = channel.GetAwaiter();
         channel.GetAwaiter().OnCompleted(() => context.Activity?.AddEvent(new ActivityEvent("channel closed")));
 
-        Dictionary<int, ChannelState> channels = [];
+        ConcurrentDictionary<int, ChannelState> channels = [];
         INewSessionContext? session = null;
         Timer? timer = null;
 
@@ -173,13 +174,13 @@ protected override async Task ConnectAsync(IChannel channel, IConnectionContext
                     bool spent = channels[header.StreamID].LocalWindow.TrySpend((int)data.Length);
                     if (!spent)
                     {
-                        _logger?.LogDebug("Ctx({ctx}), stream {stream id}: Window spent out of budget", session.Id, header.StreamID);
+                        _logger?.LogDebug("Ctx({ctx}), stream {stream id}: Local window spent out of budget", session.Id, header.StreamID);
                         await WriteGoAwayAsync(session.Id, channel, SessionTerminationCode.InternalError);
                         return;
                     }
 
-                    _logger?.LogDebug("Ctx({ctx}), stream {stream id}: Spent window, was {available}, became {new}", session.Id,
-                           header.StreamID, available, channels[header.StreamID].LocalWindow.Available);
+                    _logger?.LogDebug("Ctx({ctx}), stream {stream id}: Local spent window, was {available}, became {new}", session.Id,
+                               header.StreamID, available, channels[header.StreamID].LocalWindow.Available);
 
                     _ = channels[header.StreamID].Channel!.WriteAsync(data).AsTask().ContinueWith((t) =>
                     {
@@ -196,7 +197,7 @@ protected override async Task ConnectAsync(IChannel channel, IConnectionContext
                 {
                     int oldSize = channels[header.StreamID].RemoteWindow.Available;
                     int newSize = channels[header.StreamID].RemoteWindow.Extend(header.Length);
-                    _logger?.LogDebug("Ctx({ctx}), stream {stream id}: Window update requested: {old} => {new}", session.Id, header.StreamID, oldSize, newSize);
+                    _logger?.LogDebug("Ctx({ctx}), stream {stream id}: Window update received: {old} => {new}", session.Id, header.StreamID, oldSize, newSize);
                 }
 
                 if ((header.Flags & YamuxHeaderFlags.Fin) == YamuxHeaderFlags.Fin)
@@ -239,7 +240,7 @@ ChannelState CreateUpchannel(string contextId, int streamId, YamuxHeaderFlags in
 
                 upChannel.GetAwaiter().OnCompleted(() =>
                 {
-                    channels.Remove(streamId);
+                    channels.TryRemove(streamId, out ChannelState? _);
                     _logger?.LogDebug("Ctx({ctx}), stream {stream id}: Closed", contextId, streamId);
                 });
 
@@ -272,13 +273,15 @@ await WriteHeaderAsync(contextId, channel,
                             {
                                 int sendingSize = await state.RemoteWindow.SpendOrWait((int)upData.Length - i, state.Channel!.CancellationToken);
 
+                                _logger?.LogDebug("Ctx({ctx}), stream {stream id}: Remote window spend {sendingSize}", contextId, streamId, sendingSize);
+
                                 await WriteHeaderAsync(contextId, channel,
                                     new YamuxHeader
                                     {
                                         Type = YamuxHeaderType.Data,
                                         Length = sendingSize,
                                         StreamID = streamId
-                                    }, upData.Slice(i, sendingSize));
+                                    }, new ReadOnlySequence<byte>(upData.Slice(i, sendingSize).ToArray()));
                                 i += sendingSize;
                             }
                         }
@@ -307,7 +310,7 @@ await WriteHeaderAsync(contextId, channel,
                               StreamID = streamId
                           });
                         _ = upChannel.CloseAsync();
-                        channels.Remove(streamId);
+                        channels.TryRemove(streamId, out ChannelState? _);
 
                         _logger?.LogDebug("Ctx({ctx}), stream {stream id}: Unexpected error, closing: {error}", contextId, streamId, e.Message);
                     }
@@ -344,16 +347,19 @@ void ExtendWindow(IChannel channel, string sessionId, int streamId, IOResult res
         {
             if (result == IOResult.Ok)
             {
-                int extendedBy = channels[streamId].LocalWindow.ExtendIfNeeded();
-                if (extendedBy is not 0)
+                if (channels.TryGetValue(streamId, out ChannelState? channelState))
                 {
-                    _ = WriteHeaderAsync(sessionId, channel,
-                        new YamuxHeader
-                        {
-                            Type = YamuxHeaderType.WindowUpdate,
-                            Length = extendedBy,
-                            StreamID = streamId
-                        });
+                    int extendedBy = channelState.LocalWindow.ExtendIfNeeded();
+                    if (extendedBy is not 0)
+                    {
+                        _ = WriteHeaderAsync(sessionId, channel,
+                            new YamuxHeader
+                            {
+                                Type = YamuxHeaderType.WindowUpdate,
+                                Length = extendedBy,
+                                StreamID = streamId
+                            });
+                    }
                 }
             }
         }