Fix bug when reading many messages - size guard was triggered
diff --git a/src/ProtocolBuffers.Test/MessageStreamIteratorTest.cs b/src/ProtocolBuffers.Test/MessageStreamIteratorTest.cs
index 7eebde3..c7f31c0 100644
--- a/src/ProtocolBuffers.Test/MessageStreamIteratorTest.cs
+++ b/src/ProtocolBuffers.Test/MessageStreamIteratorTest.cs
@@ -3,6 +3,7 @@
using System.IO;
using NUnit.Framework;
using NestedMessage = Google.ProtocolBuffers.TestProtos.TestAllTypes.Types.NestedMessage;
+using Google.ProtocolBuffers.TestProtos;
namespace Google.ProtocolBuffers {
[TestFixture]
@@ -19,5 +20,31 @@
Assert.AreEqual(1500, messages[1].Bb);
Assert.IsFalse(messages[2].HasBb);
}
+
+ [Test]
+ public void ManyMessagesShouldNotTriggerSizeAlert() {
+ int messageSize = TestUtil.GetAllSet().SerializedSize;
+ // Enough messages to trigger the alert unless we've reset the size
+ // Note that currently we need to make this big enough to copy two whole buffers,
+ // as otherwise when we refill the buffer the second type, the alert triggers instantly.
+ int correctCount = (CodedInputStream.BufferSize * 2) / messageSize + 1;
+ using (MemoryStream stream = new MemoryStream()) {
+ MessageStreamWriter<TestAllTypes> writer = new MessageStreamWriter<TestAllTypes>(stream);
+ for (int i = 0; i < correctCount; i++) {
+ writer.Write(TestUtil.GetAllSet());
+ }
+ writer.Flush();
+
+ stream.Position = 0;
+
+ int count = 0;
+ foreach (var message in MessageStreamIterator<TestAllTypes>.FromStreamProvider(() => stream)
+ .WithSizeLimit(CodedInputStream.BufferSize * 2)) {
+ count++;
+ TestUtil.AssertAllFieldsSet(message);
+ }
+ Assert.AreEqual(correctCount, count);
+ }
+ }
}
}
diff --git a/src/ProtocolBuffers/CodedInputStream.cs b/src/ProtocolBuffers/CodedInputStream.cs
index 313bddf..e652af0 100644
--- a/src/ProtocolBuffers/CodedInputStream.cs
+++ b/src/ProtocolBuffers/CodedInputStream.cs
@@ -61,9 +61,9 @@
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 @@
/// 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 e8cc430..5ddfc62 100644
--- a/src/ProtocolBuffers/MessageStreamIterator.cs
+++ b/src/ProtocolBuffers/MessageStreamIterator.cs
@@ -18,6 +18,7 @@
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 @@
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 @@
/// 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 @@
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();
}
}
}