Fix bug when reading many messages - size guard was triggered

2 files changed, 24 insertions, 7 deletions

diff --git a/src/ProtocolBuffers/CodedInputStream.cs b/src/ProtocolBuffers/CodedInputStream.cs
index 313bddf3..e652af0d 100644
--- a/src/ProtocolBuffers/CodedInputStream.cs
+++ b/src/ProtocolBuffers/CodedInputStream.cs
@@ -61,9 +61,9 @@ namespace Google.ProtocolBuffers {
     private readonly Stream input;
     private uint lastTag = 0;
 
-    const int DefaultRecursionLimit = 64;
-    const int DefaultSizeLimit = 64 << 20; // 64MB
-    const int BufferSize = 4096;
+    internal const int DefaultRecursionLimit = 64;
+    internal const int DefaultSizeLimit = 64 << 20; // 64MB
+    internal const int BufferSize = 4096;
     
     /// <summary>
     /// The total number of bytes read before the current buffer. The
@@ -741,7 +741,7 @@ namespace Google.ProtocolBuffers {
     /// Read one byte from the input.
     /// </summary>
     /// <exception cref="InvalidProtocolBufferException">
-    /// he end of the stream or the current limit was reached
+    /// the end of the stream or the current limit was reached
     /// </exception>
     public byte ReadRawByte() {
       if (bufferPos == bufferSize) {
diff --git a/src/ProtocolBuffers/MessageStreamIterator.cs b/src/ProtocolBuffers/MessageStreamIterator.cs
index e8cc4306..5ddfc62a 100644
--- a/src/ProtocolBuffers/MessageStreamIterator.cs
+++ b/src/ProtocolBuffers/MessageStreamIterator.cs
@@ -18,6 +18,7 @@ namespace Google.ProtocolBuffers {
 
     private readonly StreamProvider streamProvider;
     private readonly ExtensionRegistry extensionRegistry;
+    private readonly int sizeLimit;
 
     /// <summary>
     /// Delegate created via reflection trickery (once per type) to create a builder
@@ -103,17 +104,22 @@ namespace Google.ProtocolBuffers {
       TBuilder builder = builderBuilder();
       input.ReadMessage(builder, registry);
       return builder.Build();
-        }
+    }
 #pragma warning restore 0414
 
     private static readonly uint ExpectedTag = WireFormat.MakeTag(1, WireFormat.WireType.LengthDelimited);
 
-    private MessageStreamIterator(StreamProvider streamProvider, ExtensionRegistry extensionRegistry) {
+    private MessageStreamIterator(StreamProvider streamProvider, ExtensionRegistry extensionRegistry, int sizeLimit) {
       if (messageReader == null) {
         throw typeInitializationException;
       }
       this.streamProvider = streamProvider;
       this.extensionRegistry = extensionRegistry;
+      this.sizeLimit = sizeLimit;
+    }
+
+    private MessageStreamIterator(StreamProvider streamProvider, ExtensionRegistry extensionRegistry) 
+      : this (streamProvider, extensionRegistry, CodedInputStream.DefaultSizeLimit) {
     }
 
     /// <summary>
@@ -121,7 +127,16 @@ namespace Google.ProtocolBuffers {
     /// but the specified extension registry.
     /// </summary>
     public MessageStreamIterator<TMessage> WithExtensionRegistry(ExtensionRegistry newRegistry) {
-      return new MessageStreamIterator<TMessage>(streamProvider, newRegistry);
+      return new MessageStreamIterator<TMessage>(streamProvider, newRegistry, sizeLimit);
+    }
+
+    /// <summary>
+    /// Creates a new instance which uses the same stream provider and extension registry as this one,
+    /// but with the specified size limit. Note that this must be big enough for the largest message
+    /// and the tag and size preceding it.
+    /// </summary>
+    public MessageStreamIterator<TMessage> WithSizeLimit(int newSizeLimit) {
+      return new MessageStreamIterator<TMessage>(streamProvider, extensionRegistry, newSizeLimit);
     }
 
     public static MessageStreamIterator<TMessage> FromFile(string file) {
@@ -135,12 +150,14 @@ namespace Google.ProtocolBuffers {
     public IEnumerator<TMessage> GetEnumerator() {
       using (Stream stream = streamProvider()) {
         CodedInputStream input = CodedInputStream.CreateInstance(stream);
+        input.SetSizeLimit(sizeLimit);
         uint tag;
         while ((tag = input.ReadTag()) != 0) {
           if (tag != ExpectedTag) {
             throw InvalidProtocolBufferException.InvalidMessageStreamTag();
           }
           yield return messageReader(input, extensionRegistry);
+          input.ResetSizeCounter();
         }
       }
     }