3 files changed, 26 insertions, 25 deletions

diff --git a/base/pickle.cc b/base/pickle.cc
index 0fb1d9c..356d5df 100644
--- a/base/pickle.cc
+++ b/base/pickle.cc
@@ -309,20 +309,21 @@ char* Pickle::BeginWriteData(int length) {
   return data_ptr;
 }
 
-void Pickle::TrimWriteData(int length) {
+void Pickle::TrimWriteData(int new_length) {
   DCHECK(variable_buffer_offset_ != 0);
 
-  VariableLengthBuffer *buffer = reinterpret_cast<VariableLengthBuffer*>(
+  // Fetch the the variable buffer size
+  int* cur_length = reinterpret_cast<int*>(
       reinterpret_cast<char*>(header_) + variable_buffer_offset_);
 
-  DCHECK_GE(buffer->length, length);
-
-  int old_length = buffer->length;
-  int trimmed_bytes = old_length - length;
-  if (trimmed_bytes > 0) {
-    header_->payload_size -= trimmed_bytes;
-    buffer->length = length;
+  if (new_length < 0 || new_length > *cur_length) {
+    NOTREACHED() << "Invalid length in TrimWriteData.";
+    return;
   }
+
+  // Update the payload size and variable buffer size
+  header_->payload_size -= (*cur_length - new_length);
+  *cur_length = new_length;
 }
 
 bool Pickle::Resize(size_t new_capacity) {
diff --git a/base/pickle.h b/base/pickle.h
index 7c4a007..da2a485 100644
--- a/base/pickle.h
+++ b/base/pickle.h
@@ -239,12 +239,6 @@ class Pickle {
   static const int kPayloadUnit;
 
  private:
-  // A buffer of variable length; used internally
-  struct VariableLengthBuffer {
-    int length;
-    char data;   // This is variable length.
-  };
-
   Header* header_;
   size_t header_size_;  // Supports extra data between header and payload.
   // Allocation size of payload (or -1 if allocation is const).
diff --git a/base/pickle_unittest.cc b/base/pickle_unittest.cc
index 6780e21..4c38677 100644
--- a/base/pickle_unittest.cc
+++ b/base/pickle_unittest.cc
@@ -26,14 +26,13 @@
 // THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-#include <windows.h>
-#include <string.h>
+
+#include <string>
 
 #include "base/basictypes.h"
-#include "base/check_handler.h"
 #include "base/logging.h"
-#include "base/scoped_ptr.h"
 #include "base/pickle.h"
+#include "base/scoped_ptr.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 namespace {
@@ -94,10 +93,13 @@ TEST(PickleTest, EncodeDecode) {
   EXPECT_TRUE(pickle.WriteBool(testbool2));
   EXPECT_TRUE(pickle.WriteData(testdata, testdatalen));
 
-  char* dest = pickle.BeginWriteData(testdatalen);
+  // Over allocate BeginWriteData so we can test TrimWriteData.
+  char* dest = pickle.BeginWriteData(testdatalen + 100);
   EXPECT_TRUE(dest);
   memcpy(dest, testdata, testdatalen);
 
+  pickle.TrimWriteData(testdatalen);
+
   VerifyResult(pickle);
 
   // test copy constructor
@@ -206,11 +208,15 @@ TEST(PickleTest, Resize) {
   EXPECT_EQ(cur_payload, pickle.payload_size());
 }
 
-TEST(PickleTest, HeaderPadding) {
-  struct CustomHeader : Pickle::Header {
-    int blah;
-  };
+namespace {
 
+struct CustomHeader : Pickle::Header {
+  int blah;
+};
+
+}  // namespace
+
+TEST(PickleTest, HeaderPadding) {
   const uint32 kMagic = 0x12345678;
 
   Pickle pickle(sizeof(CustomHeader));
@@ -235,4 +241,4 @@ TEST(PickleTest, EqualsOperator) {
   Pickle copy;
   copy = copy_refs_source_buffer;
   ASSERT_EQ(source.size(), copy.size());
-}
\ No newline at end of file
+}