fix: Large buffer will cause FlushAsync block

HMBSbige · Sep 18, 2021 · 52124f0 · 52124f0
1 parent 4b05448
commit 52124f0
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 39 deletions.
diff --git a/Shadowsocks.Protocol/TcpClients/ShadowsocksPipeReader.cs b/Shadowsocks.Protocol/TcpClients/ShadowsocksPipeReader.cs
@@ -4,7 +4,6 @@
 using Shadowsocks.Crypto;
 using Shadowsocks.Protocol.Models;
 using System;
-using System.Buffers;
 using System.IO.Pipelines;
 using System.Threading;
 using System.Threading.Tasks;
@@ -56,10 +55,25 @@ private Task WrapAsync(IShadowsocksCrypto decryptor, CancellationToken cancellat
 
 							try
 							{
-								if (ReceiveFromRemote(ref buffer))
+								while (!buffer.IsEmpty)
 								{
-									var writerFlushResult = await Writer.FlushAsync(cancellationToken);
-									if (writerFlushResult.IsCompleted)
+									var oldLength = buffer.Length;
+
+									var memory = Writer.GetMemory(BufferSize);
+
+									var outLength = decryptor.DecryptTCP(ref buffer, memory.Span);
+
+									Writer.Advance(outLength);
+									if (outLength > 0)
+									{
+										var writerFlushResult = await Writer.FlushAsync(cancellationToken);
+										if (writerFlushResult.IsCompleted)
+										{
+											goto NoData;
+										}
+									}
+
+									if (oldLength == buffer.Length)
 									{
 										break;
 									}
@@ -75,7 +89,7 @@ private Task WrapAsync(IShadowsocksCrypto decryptor, CancellationToken cancellat
 								InternalReader.AdvanceTo(buffer.Start, buffer.End);
 							}
 						}
-
+					NoData:
 						await Writer.CompleteAsync();
 					}
 					catch (Exception ex)
@@ -89,32 +103,6 @@ private Task WrapAsync(IShadowsocksCrypto decryptor, CancellationToken cancellat
 				},
 				default
 			);
-
-			bool ReceiveFromRemote(ref ReadOnlySequence<byte> sequence)
-			{
-				var result = false;
-				while (!sequence.IsEmpty)
-				{
-					var oldLength = sequence.Length;
-
-					var span = Writer.GetSpan(BufferSize);
-
-					var outLength = decryptor.DecryptTCP(ref sequence, span);
-
-					Writer.Advance(outLength);
-					if (outLength > 0)
-					{
-						result = true;
-					}
-
-					if (oldLength == sequence.Length)
-					{
-						break;
-					}
-				}
-
-				return result;
-			}
 		}
 
 		public override void AdvanceTo(SequencePosition consumed)

diff --git a/Shadowsocks.Protocol/TcpClients/ShadowsocksPipeWriter.cs b/Shadowsocks.Protocol/TcpClients/ShadowsocksPipeWriter.cs
@@ -52,17 +52,13 @@ private Task WrapAsync(IShadowsocksCrypto encryptor, CancellationToken cancellat
 							var result = await Reader.ReadAsync(cancellationToken);
 							var buffer = result.Buffer;
 
-							if (!buffer.IsEmpty)
+							foreach (var segment in buffer)
 							{
-								foreach (var segment in buffer)
-								{
-									SendToRemote(segment.Span);
-								}
-
+								SendToRemote(segment.Span);
 								var flushResult = await InternalWriter.FlushAsync(cancellationToken);
 								if (flushResult.IsCompleted)
 								{
-									break;
+									goto NoData;
 								}
 							}
 
@@ -73,7 +69,7 @@ private Task WrapAsync(IShadowsocksCrypto encryptor, CancellationToken cancellat
 								break;
 							}
 						}
-
+					NoData:
 						await Reader.CompleteAsync();
 					}
 					catch (Exception ex)