3 files changed, 49 insertions, 4 deletions
diff --git a/base/pickle.cc b/base/pickle.cc
index 06a3be1..3f376e3 100644
--- a/base/pickle.cc
+++ b/base/pickle.cc
@@ -41,11 +41,21 @@ Pickle::Pickle(int header_size)
 
 Pickle::Pickle(const char* data, int data_len)
     : header_(reinterpret_cast<Header*>(const_cast<char*>(data))),
-      header_size_(data_len - header_->payload_size),
+      header_size_(0),
       capacity_(kCapacityReadOnly),
       variable_buffer_offset_(0) {
-  DCHECK(header_size_ >= sizeof(Header));
-  DCHECK(header_size_ == AlignInt(header_size_, sizeof(uint32)));
+  if (data_len >= static_cast<int>(sizeof(Header)))
+    header_size_ = data_len - header_->payload_size;
+
+  if (header_size_ > static_cast<unsigned int>(data_len))
+    header_size_ = 0;
+
+  if (header_size_ != AlignInt(header_size_, sizeof(uint32)))
+    header_size_ = 0;
+
+  // If there is anything wrong with the data, we're not going to use it.
+  if (!header_size_)
+    header_ = NULL;
 }
 
 Pickle::Pickle(const Pickle& other)
diff --git a/base/pickle.h b/base/pickle.h
index c7aee67..6006e62 100644
--- a/base/pickle.h
+++ b/base/pickle.h
@@ -177,10 +177,12 @@ class Pickle {
   // Returns the address of the byte immediately following the currently valid
   // header + payload.
   char* end_of_payload() {
+    // We must have a valid header_.
     return payload() + payload_size();
   }
   const char* end_of_payload() const {
-    return payload() + payload_size();
+    // This object may be invalid.
+    return header_ ? payload() + payload_size() : NULL;
   }
 
   size_t capacity() const {
diff --git a/base/pickle_unittest.cc b/base/pickle_unittest.cc
index aea3830..fdc0664 100644
--- a/base/pickle_unittest.cc
+++ b/base/pickle_unittest.cc
@@ -87,6 +87,39 @@ TEST(PickleTest, EncodeDecode) {
   VerifyResult(pickle3);
 }
 
+// Tests that we can handle really small buffers.
+TEST(PickleTest, SmallBuffer) {
+  scoped_array<char> buffer(new char[1]);
+
+  // We should not touch the buffer.
+  Pickle pickle(buffer.get(), 1);
+
+  void* iter = NULL;
+  int data;
+  EXPECT_FALSE(pickle.ReadInt(&iter, &data));
+}
+
+// Tests that we can handle improper headers.
+TEST(PickleTest, BigSize) {
+  int buffer[] = { 0x56035200, 25, 40, 50 };
+
+  Pickle pickle(reinterpret_cast<char*>(buffer), sizeof(buffer));
+
+  void* iter = NULL;
+  int data;
+  EXPECT_FALSE(pickle.ReadInt(&iter, &data));
+}
+
+TEST(PickleTest, UnalignedSize) {
+  int buffer[] = { 10, 25, 40, 50 };
+
+  Pickle pickle(reinterpret_cast<char*>(buffer), sizeof(buffer));
+
+  void* iter = NULL;
+  int data;
+  EXPECT_FALSE(pickle.ReadInt(&iter, &data));
+}
+
 TEST(PickleTest, ZeroLenStr) {
   Pickle pickle;
   EXPECT_TRUE(pickle.WriteString(""));