Handle if the last data in async path.

Author: yyjdelete

src/DotNetty.Handlers/Tls/TlsHandler.cs

@@ -27,6 +27,7 @@ public sealed class TlsHandler : ByteToMessageDecoder
 
         static readonly Exception ChannelClosedException = new IOException("Channel is closed");
         static readonly Action<Task, object> HandshakeCompletionCallback = new Action<Task, object>(HandleHandshakeCompleted);
+        static readonly Action<Task<int>, object> UnwrapCompletedCallback = new Action<Task<int>, object>(UnwrapCompleted);
 
         readonly SslStream sslStream;
         readonly MediationStream mediationStream;
@@ -40,6 +41,7 @@ public sealed class TlsHandler : ByteToMessageDecoder
         bool firedChannelRead;
         volatile FlushMode flushMode = FlushMode.ForceFlush;
         IByteBuffer pendingSslStreamReadBuffer;
+        int pendingSslStreamReadLength;
         Task<int> pendingSslStreamReadFuture;
 
         public TlsHandler(TlsSettings settings)
@@ -342,10 +344,11 @@ void Unwrap(IChannelHandlerContext ctx, IByteBuffer packet, int offset, int leng
                     Contract.Assert(this.pendingSslStreamReadBuffer != null);
 
                     outputBuffer = this.pendingSslStreamReadBuffer;
-                    outputBufferLength = outputBuffer.WritableBytes;
+                    outputBufferLength = this.pendingSslStreamReadLength;
 
                     this.pendingSslStreamReadFuture = null;
                     this.pendingSslStreamReadBuffer = null;
+                    this.pendingSslStreamReadLength = 0;
                 }
                 else
                 {
@@ -358,6 +361,9 @@ void Unwrap(IChannelHandlerContext ctx, IByteBuffer packet, int offset, int leng
                     int currentPacketLength = packetLengths[packetIndex];
                     this.mediationStream.ExpandSource(currentPacketLength);
 
+                    while (true)
+                    {
+                        int totalRead = 0;
                         if (currentReadFuture != null)
                         {
                             // there was a read pending already, so we make sure we completed that first
@@ -366,10 +372,11 @@ void Unwrap(IChannelHandlerContext ctx, IByteBuffer packet, int offset, int leng
                             {
                                 // we did feed the whole current packet to SslStream yet it did not produce any result -> move to the next packet in input
 
-                            continue;
+                                break;
                             }
 
                             int read = currentReadFuture.Result;
+                            totalRead += read;
 
                             if (read == 0)
                             {
@@ -382,19 +389,19 @@ void Unwrap(IChannelHandlerContext ctx, IByteBuffer packet, int offset, int leng
 
                             currentReadFuture = null;
                             outputBuffer = null;
-                        if (this.mediationStream.SourceReadableBytes == 0)
+                            if (this.mediationStream.TotalReadableBytes == 0)
                             {
                                 // we just made a frame available for reading but there was already pending read so SslStream read it out to make further progress there
 
                                 if (read < outputBufferLength)
                                 {
                                     // SslStream returned non-full buffer and there's no more input to go through ->
                                     // typically it means SslStream is done reading current frame so we skip
-                                continue;
+                                    break;
                                 }
 
                                 // we've read out `read` bytes out of current packet to fulfil previously outstanding read
-                            outputBufferLength = currentPacketLength - read;
+                                outputBufferLength = currentPacketLength - totalRead;
                                 if (outputBufferLength <= 0)
                                 {
                                     // after feeding to SslStream current frame it read out more bytes than current packet size
@@ -417,27 +424,15 @@ void Unwrap(IChannelHandlerContext ctx, IByteBuffer packet, int offset, int leng
                         outputBuffer = ctx.Allocator.Buffer(outputBufferLength);
                         currentReadFuture = this.ReadFromSslStreamAsync(outputBuffer, outputBufferLength);
                     }
+                }
 
-                // read out the rest of SslStream's output (if any) at risk of going async
-                // using FallbackReadBufferSize - buffer size we're ok to have pinned with the SslStream until it's done reading
-                while (true)
-                {
                 if (currentReadFuture != null)
                 {
-                        if (!currentReadFuture.IsCompleted)
-                        {
-                            break;
-                        }
-                        int read = currentReadFuture.Result;
-                        AddBufferToOutput(outputBuffer, read, output);
-                    }
-                    outputBuffer = ctx.Allocator.Buffer(FallbackReadBufferSize);
-                    currentReadFuture = this.ReadFromSslStreamAsync(outputBuffer, FallbackReadBufferSize);
-                }
-
                     pending = true;
                     this.pendingSslStreamReadBuffer = outputBuffer;
                     this.pendingSslStreamReadFuture = currentReadFuture;
+                    this.pendingSslStreamReadLength = outputBufferLength;
+                }
             }
             catch (Exception ex)
             {
@@ -458,6 +453,91 @@ void Unwrap(IChannelHandlerContext ctx, IByteBuffer packet, int offset, int leng
                         outputBuffer.SafeRelease();
                     }
                 }
+
+                if (pending)
+                {
+                    //Can't use ExecuteSynchronously here for it may change the order of output if task is already completed here.
+                    this.pendingSslStreamReadFuture?.ContinueWith(UnwrapCompletedCallback, this, TaskContinuationOptions.None);
+                }
+            }
+        }
+
+        static void UnwrapCompleted(Task<int> task, object state)
+        {
+            // Mono(with legacy provider) finish ReadAsync in async, 
+            // so extra check is needed to receive data in async
+            var self = (TlsHandler)state;
+            Debug.Assert(self.capturedContext.Executor.InEventLoop);
+
+            //Ignore task completed in Unwrap
+            if (task == self.pendingSslStreamReadFuture)
+            {
+                IByteBuffer buf = self.pendingSslStreamReadBuffer;
+                int outputBufferLength = self.pendingSslStreamReadLength;
+
+                self.pendingSslStreamReadFuture = null;
+                self.pendingSslStreamReadBuffer = null;
+                self.pendingSslStreamReadLength = 0;
+
+                while (true)
+                {
+                    switch (task.Status)
+                    {
+                        case TaskStatus.RanToCompletion:
+                            {
+                                //The logic is the same as the one in Unwrap()
+                                var read = task.Result;
+                                //Stream Closed
+                                if (read == 0)
+                                    return;
+                                self.capturedContext.FireChannelRead(buf.SetWriterIndex(buf.WriterIndex + read));
+
+                                if (self.mediationStream.TotalReadableBytes == 0)
+                                {
+                                    self.capturedContext.FireChannelReadComplete();
+                                    self.mediationStream.ResetSource(self.capturedContext.Allocator);
+
+                                    if (read < outputBufferLength)
+                                    {
+                                        // SslStream returned non-full buffer and there's no more input to go through ->
+                                        // typically it means SslStream is done reading current frame so we skip
+                                        return;
+                                    }
+                                }
+
+                                outputBufferLength = self.mediationStream.TotalReadableBytes;
+                                if (outputBufferLength <= 0)
+                                    outputBufferLength = FallbackReadBufferSize;
+
+                                buf = self.capturedContext.Allocator.Buffer(outputBufferLength);
+                                task = self.ReadFromSslStreamAsync(buf, outputBufferLength);
+                                if (task.IsCompleted)
+                                {
+                                    continue;
+                                }
+
+                                self.pendingSslStreamReadFuture = task;
+                                self.pendingSslStreamReadBuffer = buf;
+                                self.pendingSslStreamReadLength = outputBufferLength;
+                                task.ContinueWith(UnwrapCompletedCallback, self, TaskContinuationOptions.ExecuteSynchronously);
+                                return;
+                            }
+
+                        case TaskStatus.Canceled:
+                        case TaskStatus.Faulted:
+                            {
+                                buf.SafeRelease();
+                                self.HandleFailure(task.Exception);
+                                return;
+                            }
+
+                        default:
+                            {
+                                buf.SafeRelease();
+                                throw new ArgumentOutOfRangeException(nameof(task), "Unexpected task status: " + task.Status);
+                            }
+                    }
+                }
             }
         }