Call ResetConsumedMessageSize from ReadMessageEndAsync

9 files changed, 26 insertions, 3 deletions

diff --git a/lib/netstd/Thrift/Protocol/TBinaryProtocol.cs b/lib/netstd/Thrift/Protocol/TBinaryProtocol.cs
index ba2a7abbc..7feb698b4 100644
--- a/lib/netstd/Thrift/Protocol/TBinaryProtocol.cs
+++ b/lib/netstd/Thrift/Protocol/TBinaryProtocol.cs
@@ -253,6 +253,7 @@ namespace Thrift.Protocol
         public override Task ReadMessageEndAsync(CancellationToken cancellationToken)
         {
             cancellationToken.ThrowIfCancellationRequested();
+            Transport.ResetConsumedMessageSize();
             return Task.CompletedTask;
         }
 
diff --git a/lib/netstd/Thrift/Protocol/TCompactProtocol.cs b/lib/netstd/Thrift/Protocol/TCompactProtocol.cs
index b899d3dbc..dd38f4870 100644
--- a/lib/netstd/Thrift/Protocol/TCompactProtocol.cs
+++ b/lib/netstd/Thrift/Protocol/TCompactProtocol.cs
@@ -459,6 +459,7 @@ namespace Thrift.Protocol
         public override Task ReadMessageEndAsync(CancellationToken cancellationToken)
         {
             cancellationToken.ThrowIfCancellationRequested();
+            Transport.ResetConsumedMessageSize();
             return Task.CompletedTask;
         }
 
diff --git a/lib/netstd/Thrift/Protocol/TJSONProtocol.cs b/lib/netstd/Thrift/Protocol/TJSONProtocol.cs
index c100d8627..82e758b1e 100644
--- a/lib/netstd/Thrift/Protocol/TJSONProtocol.cs
+++ b/lib/netstd/Thrift/Protocol/TJSONProtocol.cs
@@ -694,7 +694,9 @@ namespace Thrift.Protocol
 
         public override async Task ReadMessageEndAsync(CancellationToken cancellationToken)
         {
+            cancellationToken.ThrowIfCancellationRequested();
             await ReadJsonArrayEndAsync(cancellationToken);
+            Transport.ResetConsumedMessageSize();
         }
 
         public override async ValueTask<TStruct> ReadStructBeginAsync(CancellationToken cancellationToken)
diff --git a/lib/netstd/Thrift/Protocol/TProtocolDecorator.cs b/lib/netstd/Thrift/Protocol/TProtocolDecorator.cs
index 1ea9fb9ae..c75cc6319 100644
--- a/lib/netstd/Thrift/Protocol/TProtocolDecorator.cs
+++ b/lib/netstd/Thrift/Protocol/TProtocolDecorator.cs
@@ -156,7 +156,9 @@ namespace Thrift.Protocol
 
         public override async Task ReadMessageEndAsync(CancellationToken cancellationToken)
         {
+            cancellationToken.ThrowIfCancellationRequested();
             await _wrappedProtocol.ReadMessageEndAsync(cancellationToken);
+            Transport.ResetConsumedMessageSize();
         }
 
         public override async ValueTask<TStruct> ReadStructBeginAsync(CancellationToken cancellationToken)
diff --git a/lib/netstd/Thrift/Transport/Layered/TBufferedTransport.cs b/lib/netstd/Thrift/Transport/Layered/TBufferedTransport.cs
index 271a9946c..7474b7f93 100644
--- a/lib/netstd/Thrift/Transport/Layered/TBufferedTransport.cs
+++ b/lib/netstd/Thrift/Transport/Layered/TBufferedTransport.cs
@@ -179,6 +179,11 @@ namespace Thrift.Transport
             }
         }
 
+        public override void ResetConsumedMessageSize(long newSize = -1)
+        {
+            base.ResetConsumedMessageSize(newSize);
+            ReadBuffer.ResetConsumedMessageSize(newSize);
+        }
 
         private void CheckNotDisposed()
         {
diff --git a/lib/netstd/Thrift/Transport/Layered/TFramedTransport.cs b/lib/netstd/Thrift/Transport/Layered/TFramedTransport.cs
index 722df7672..694ea0f40 100644
--- a/lib/netstd/Thrift/Transport/Layered/TFramedTransport.cs
+++ b/lib/netstd/Thrift/Transport/Layered/TFramedTransport.cs
@@ -186,5 +186,11 @@ namespace Thrift.Transport
             }
             IsDisposed = true;
         }
+
+        public override void ResetConsumedMessageSize(long newSize = -1)
+        {
+            base.ResetConsumedMessageSize(newSize);
+            ReadBuffer.ResetConsumedMessageSize(newSize);
+        }
     }
 }
diff --git a/lib/netstd/Thrift/Transport/Layered/TLayeredTransport.cs b/lib/netstd/Thrift/Transport/Layered/TLayeredTransport.cs
index 4d39b392e..98fbd227d 100644
--- a/lib/netstd/Thrift/Transport/Layered/TLayeredTransport.cs
+++ b/lib/netstd/Thrift/Transport/Layered/TLayeredTransport.cs
@@ -1,4 +1,4 @@
-// Licensed to the Apache Software Foundation(ASF) under one
+// Licensed to the Apache Software Foundation(ASF) under one
 // or more contributor license agreements.See the NOTICE file
 // distributed with this work for additional information
 // regarding copyright ownership.The ASF licenses this file
@@ -41,5 +41,10 @@ namespace Thrift.Transport
         {
             InnerTransport.CheckReadBytesAvailable(numBytes);
         }
+
+        public override void ResetConsumedMessageSize(long newSize = -1)
+        {
+            InnerTransport.ResetConsumedMessageSize(newSize);
+        }
     }
 }
diff --git a/lib/netstd/Thrift/Transport/TEndpointTransport.cs b/lib/netstd/Thrift/Transport/TEndpointTransport.cs
index 51a2a17cb..6c7810105 100644
--- a/lib/netstd/Thrift/Transport/TEndpointTransport.cs
+++ b/lib/netstd/Thrift/Transport/TEndpointTransport.cs
@@ -43,7 +43,7 @@ namespace Thrift.Transport
         /// <summary>
         /// Resets RemainingMessageSize to the configured maximum 
         /// </summary>
-        protected void ResetConsumedMessageSize(long newSize = -1)
+        public override void ResetConsumedMessageSize(long newSize = -1)
         {
             // full reset 
             if (newSize < 0)
diff --git a/lib/netstd/Thrift/Transport/TTransport.cs b/lib/netstd/Thrift/Transport/TTransport.cs
index 3f4245aa2..2f87ca0be 100644
--- a/lib/netstd/Thrift/Transport/TTransport.cs
+++ b/lib/netstd/Thrift/Transport/TTransport.cs
@@ -1,4 +1,4 @@
-// Licensed to the Apache Software Foundation(ASF) under one
+// Licensed to the Apache Software Foundation(ASF) under one
 // or more contributor license agreements.See the NOTICE file
 // distributed with this work for additional information
 // regarding copyright ownership.The ASF licenses this file
@@ -35,6 +35,7 @@ namespace Thrift.Transport
         public abstract TConfiguration Configuration { get; }
         public abstract void UpdateKnownMessageSize(long size);
         public abstract void CheckReadBytesAvailable(long numBytes);
+        public abstract void ResetConsumedMessageSize(long newSize = -1);
         public void Dispose()
         {
             Dispose(true);