Down-integrate from internal branch

1 files changed, 61 insertions, 1 deletions

diff --git a/src/google/protobuf/io/coded_stream_unittest.cc b/src/google/protobuf/io/coded_stream_unittest.cc
index ff268ab9..2daab194 100644
--- a/src/google/protobuf/io/coded_stream_unittest.cc
+++ b/src/google/protobuf/io/coded_stream_unittest.cc
@@ -44,7 +44,6 @@
 #include <google/protobuf/testing/googletest.h>
 #include <gtest/gtest.h>
 #include <google/protobuf/io/zero_copy_stream_impl.h>
-#include <google/protobuf/stubs/strutil.h>
 
 
 // This declares an unsigned long long integer literal in a portable way.
@@ -125,6 +124,13 @@ namespace {
 
 class CodedStreamTest : public testing::Test {
  protected:
+  // Helper method used by tests for bytes warning. See implementation comment
+  // for further information.
+  static void SetupTotalBytesLimitWarningTest(
+      int total_bytes_limit, int warning_threshold,
+      vector<string>* out_errors, vector<string>* out_warnings);
+
+  // Buffer used during most of the tests. This assumes tests run sequentially.
   static const int kBufferSize = 1024 * 64;
   static uint8 buffer_[kBufferSize];
 };
@@ -1022,6 +1028,59 @@ TEST_F(CodedStreamTest, TotalBytesLimitNotValidMessageEnd) {
   EXPECT_FALSE(coded_input.ConsumedEntireMessage());
 }
 
+// This method is used by the tests below.
+// It constructs a CodedInputStream with the given limits and tries to read 2KiB
+// of data from it. Then it returns the logged errors and warnings in the given
+// vectors.
+void CodedStreamTest::SetupTotalBytesLimitWarningTest(
+    int total_bytes_limit, int warning_threshold,
+    vector<string>* out_errors, vector<string>* out_warnings) {
+  ArrayInputStream raw_input(buffer_, sizeof(buffer_), 128);
+
+  ScopedMemoryLog scoped_log;
+  {
+    CodedInputStream input(&raw_input);
+    input.SetTotalBytesLimit(total_bytes_limit, warning_threshold);
+    string str;
+    EXPECT_TRUE(input.ReadString(&str, 2048));
+  }
+
+  *out_errors = scoped_log.GetMessages(ERROR);
+  *out_warnings = scoped_log.GetMessages(WARNING);
+}
+
+TEST_F(CodedStreamTest, TotalBytesLimitWarning) {
+  vector<string> errors;
+  vector<string> warnings;
+  SetupTotalBytesLimitWarningTest(10240, 1024, &errors, &warnings);
+
+  EXPECT_EQ(0, errors.size());
+
+  ASSERT_EQ(2, warnings.size());
+  EXPECT_PRED_FORMAT2(testing::IsSubstring,
+    "Reading dangerously large protocol message.  If the message turns out to "
+    "be larger than 10240 bytes, parsing will be halted for security reasons.",
+    warnings[0]);
+  EXPECT_PRED_FORMAT2(testing::IsSubstring,
+    "The total number of bytes read was 2048",
+    warnings[1]);
+}
+
+TEST_F(CodedStreamTest, TotalBytesLimitWarningDisabled) {
+  vector<string> errors;
+  vector<string> warnings;
+
+  // Test with -1
+  SetupTotalBytesLimitWarningTest(10240, -1, &errors, &warnings);
+  EXPECT_EQ(0, errors.size());
+  EXPECT_EQ(0, warnings.size());
+
+  // Test again with -2, expecting the same result
+  SetupTotalBytesLimitWarningTest(10240, -2, &errors, &warnings);
+  EXPECT_EQ(0, errors.size());
+  EXPECT_EQ(0, warnings.size());
+}
+
 
 TEST_F(CodedStreamTest, RecursionLimit) {
   ArrayInputStream input(buffer_, sizeof(buffer_));
@@ -1060,6 +1119,7 @@ TEST_F(CodedStreamTest, RecursionLimit) {
   EXPECT_FALSE(coded_input.IncrementRecursionDepth());     // 7
 }
 
+
 class ReallyBigInputStream : public ZeroCopyInputStream {
  public:
   ReallyBigInputStream() : backup_amount_(0), buffer_count_(0) {}