Ensure stream position to be reset before returning task

Neo4j.Driver/Neo4j.Driver.Tests/IO/ChunkReaderTests.cs

@@ -201,6 +201,24 @@ public void ShouldResetInternalBufferPositionsAfterOneChunkSpanningMessageIsRead
                 count.Should().Be(1);
                 logger.Verify(l => l.Trace(It.IsRegex("^\\d+ bytes left in chunk buffer.*compacting\\.$"), It.IsAny<object[]>()), Times.AtLeast(size / Constants.ChunkBufferSize));
             }
+
+            [Fact]
+            public void ShouldResetBufferStreamPosition()
+            {
+                var data = IOExtensions.GenerateBoltMessages(1000, 128 * 1024);
+
+                var logger = new Mock<ILogger>();
+                var reader = new ChunkReader(new MemoryStream(data.ToArray()), logger.Object);
+
+                var bufferStream = new MemoryStream();
+                bufferStream.Write(IOExtensions.GenerateBoltMessage(1035));
+
+                var bufferPosition = bufferStream.Position;
+
+                var count = reader.ReadNextMessages(bufferStream);
+
+                bufferStream.Position.Should().Be(bufferPosition);
+            }
 
         }
 
@@ -328,6 +346,24 @@ public async void ShouldResetInternalBufferPositionsAfterOneChunkSpanningMessage
                 logger.Verify(l => l.Trace(It.IsRegex("^\\d+ bytes left in chunk buffer.*compacting\\.$"), It.IsAny<object[]>()), Times.AtLeast(size / Constants.ChunkBufferSize));
             }
 
+            [Fact]
+            public async void ShouldResetBufferStreamPosition()
+            {
+                var data = IOExtensions.GenerateBoltMessages(1000, 128 * 1024);
+
+                var logger = new Mock<ILogger>();
+                var reader = new ChunkReader(new MemoryStream(data.ToArray()), logger.Object);
+
+                var bufferStream = new MemoryStream();
+                bufferStream.Write(IOExtensions.GenerateBoltMessage(1035));
+
+                var bufferPosition = bufferStream.Position;
+
+                var count = await reader.ReadNextMessagesAsync(bufferStream);
+
+                bufferStream.Position.Should().Be(bufferPosition);
+            }
+
         }
 
     }

Neo4j.Driver/Neo4j.Driver/Internal/IO/ChunkReader.cs

@@ -16,8 +16,10 @@
 // limitations under the License.
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.IO;
 using System.Text;
+using System.Threading;
 using System.Threading.Tasks;
 using Neo4j.Driver.V1;
 
@@ -150,18 +152,19 @@ public Task<int> ReadNextMessagesAsync(Stream messageStream)
 
             // Track and manage target stream's positions.
             var previousPosition = messageStream.Position;
-            taskCompletionSource.Task.ContinueWith(t =>
-            {
-                messageStream.Position = previousPosition;
-            }, TaskContinuationOptions.ExecuteSynchronously);
             messageStream.Position = messageStream.Length;
 
             ReadNextChunkLoopAsync(
                 messageStream,
                 taskCompletionSource,
                 TaskExtensions.GetCompletedTask());
 
-            return taskCompletionSource.Task;
+            return taskCompletionSource.Task.ContinueWith(t =>
+            {
+                messageStream.Position = previousPosition;
+
+                return t;
+            }, TaskContinuationOptions.ExecuteSynchronously).Unwrap();
         }
 
         private Task ReadNextChunkLoopAsync(Stream messageStream, TaskCompletionSource<int> taskCompletionSource,