4 files changed, 48 insertions, 16 deletions

diff --git a/core/fpdfapi/parser/cpdf_object_unittest.cpp b/core/fpdfapi/parser/cpdf_object_unittest.cpp
index dd170ba90e..169e0f1e56 100644
--- a/core/fpdfapi/parser/cpdf_object_unittest.cpp
+++ b/core/fpdfapi/parser/cpdf_object_unittest.cpp
@@ -828,6 +828,30 @@ TEST(PDFStreamTest, SetDataAndRemoveFilter) {
   EXPECT_FALSE(stream->GetDict()->KeyExist("DecodeParms"));
 }
 
+TEST(PDFStreamTest, LengthInDictionaryOnCreate) {
+  static constexpr uint32_t kBufSize = 100;
+  // The length field should be created on stream create.
+  {
+    std::unique_ptr<uint8_t, FxFreeDeleter> data;
+    data.reset(FX_Alloc(uint8_t, kBufSize));
+    auto stream = pdfium::MakeUnique<CPDF_Stream>(
+        std::move(data), kBufSize, pdfium::MakeUnique<CPDF_Dictionary>());
+    EXPECT_EQ(static_cast<int>(kBufSize),
+              stream->GetDict()->GetIntegerFor("Length"));
+  }
+  // The length field should be corrected on stream create.
+  {
+    std::unique_ptr<uint8_t, FxFreeDeleter> data;
+    data.reset(FX_Alloc(uint8_t, kBufSize));
+    auto dict = pdfium::MakeUnique<CPDF_Dictionary>();
+    dict->SetNewFor<CPDF_Number>("Length", 30000);
+    auto stream = pdfium::MakeUnique<CPDF_Stream>(std::move(data), kBufSize,
+                                                  std::move(dict));
+    EXPECT_EQ(static_cast<int>(kBufSize),
+              stream->GetDict()->GetIntegerFor("Length"));
+  }
+}
+
 TEST(PDFDictionaryTest, CloneDirectObject) {
   CPDF_IndirectObjectHolder objects_holder;
   auto dict = pdfium::MakeUnique<CPDF_Dictionary>();
diff --git a/core/fpdfapi/parser/cpdf_object_walker_unittest.cpp b/core/fpdfapi/parser/cpdf_object_walker_unittest.cpp
index 66c559d3ca..c5ffe844ad 100644
--- a/core/fpdfapi/parser/cpdf_object_walker_unittest.cpp
+++ b/core/fpdfapi/parser/cpdf_object_walker_unittest.cpp
@@ -38,6 +38,8 @@ std::string Walk(CPDF_Object* object) {
       result << " Stream";
     else if (obj->IsReference())
       result << " Ref";
+    else if (obj->IsNumber())
+      result << " Num";
     else if (obj->IsNull())
       result << " Null";
     else
@@ -72,13 +74,14 @@ TEST(CPDF_ObjectWalkerTest, CombinedObject) {
   array->Add(pdfium::MakeUnique<CPDF_Stream>(
       nullptr, 0, pdfium::MakeUnique<CPDF_Dictionary>()));
   dict->SetFor("3", std::move(array));
-  EXPECT_EQ(Walk(dict.get()), "Dict Str Bool Arr Ref Null Stream Dict");
+  // The last number for stream length.
+  EXPECT_EQ(Walk(dict.get()), "Dict Str Bool Arr Ref Null Stream Dict Num");
 }
 
 TEST(CPDF_ObjectWalkerTest, GetParent) {
-  auto level_4 = pdfium::MakeUnique<CPDF_Null>();
+  auto level_4 = pdfium::MakeUnique<CPDF_Number>(0);
   auto level_3 = pdfium::MakeUnique<CPDF_Dictionary>();
-  level_3->SetFor("AnyObj", std::move(level_4));
+  level_3->SetFor("Length", std::move(level_4));
   auto level_2 =
       pdfium::MakeUnique<CPDF_Stream>(nullptr, 0, std::move(level_3));
   auto level_1 = pdfium::MakeUnique<CPDF_Array>();
@@ -86,7 +89,7 @@ TEST(CPDF_ObjectWalkerTest, GetParent) {
   auto level_0 = pdfium::MakeUnique<CPDF_Dictionary>();
   level_0->SetFor("Array", std::move(level_1));
 
-  // We have <</Array [ stream( << /AnyObj null >>) ]>>
+  // We have <</Array [ stream( << /Length 0 >>) ]>>
   // In this case each step will increase depth.
   // And on each step the prev object should be parent for current.
   const CPDF_Object* cur_parent = nullptr;
diff --git a/core/fpdfapi/parser/cpdf_stream.cpp b/core/fpdfapi/parser/cpdf_stream.cpp
index 1c8660efb0..fadaec1b7e 100644
--- a/core/fpdfapi/parser/cpdf_stream.cpp
+++ b/core/fpdfapi/parser/cpdf_stream.cpp
@@ -22,7 +22,9 @@ CPDF_Stream::CPDF_Stream() {}
 CPDF_Stream::CPDF_Stream(std::unique_ptr<uint8_t, FxFreeDeleter> pData,
                          uint32_t size,
                          std::unique_ptr<CPDF_Dictionary> pDict)
-    : m_dwSize(size), m_pDict(std::move(pDict)), m_pDataBuf(std::move(pData)) {}
+    : m_pDict(std::move(pDict)) {
+  SetData(std::move(pData), size);
+}
 
 CPDF_Stream::~CPDF_Stream() {
   m_ObjNum = kInvalidObjNum;
@@ -54,14 +56,7 @@ void CPDF_Stream::InitStream(const uint8_t* pData,
                              uint32_t size,
                              std::unique_ptr<CPDF_Dictionary> pDict) {
   m_pDict = std::move(pDict);
-  m_bMemoryBased = true;
-  m_pFile = nullptr;
-  m_pDataBuf.reset(FX_Alloc(uint8_t, size));
-  if (pData)
-    memcpy(m_pDataBuf.get(), pData, size);
-  m_dwSize = size;
-  if (m_pDict)
-    m_pDict->SetNewFor<CPDF_Number>("Length", static_cast<int>(m_dwSize));
+  SetData(pData, size);
 }
 
 void CPDF_Stream::InitStreamFromFile(
@@ -110,10 +105,19 @@ void CPDF_Stream::SetDataAndRemoveFilter(std::ostringstream* stream) {
 }
 
 void CPDF_Stream::SetData(const uint8_t* pData, uint32_t size) {
+  std::unique_ptr<uint8_t, FxFreeDeleter> data_copy;
+  if (pData) {
+    data_copy.reset(FX_Alloc(uint8_t, size));
+    memcpy(data_copy.get(), pData, size);
+  }
+  SetData(std::move(data_copy), size);
+}
+
+void CPDF_Stream::SetData(std::unique_ptr<uint8_t, FxFreeDeleter> pData,
+                          uint32_t size) {
   m_bMemoryBased = true;
-  m_pDataBuf.reset(FX_Alloc(uint8_t, size));
-  if (pData)
-    memcpy(m_pDataBuf.get(), pData, size);
+  m_pFile = nullptr;
+  m_pDataBuf = std::move(pData);
   m_dwSize = size;
   if (!m_pDict)
     m_pDict = pdfium::MakeUnique<CPDF_Dictionary>();
diff --git a/core/fpdfapi/parser/cpdf_stream.h b/core/fpdfapi/parser/cpdf_stream.h
index 1dbbd0343e..3c3b674a48 100644
--- a/core/fpdfapi/parser/cpdf_stream.h
+++ b/core/fpdfapi/parser/cpdf_stream.h
@@ -41,6 +41,7 @@ class CPDF_Stream : public CPDF_Object {
 
   // Does not takes ownership of |pData|, copies into internally-owned buffer.
   void SetData(const uint8_t* pData, uint32_t size);
+  void SetData(std::unique_ptr<uint8_t, FxFreeDeleter> pData, uint32_t size);
   void SetData(std::ostringstream* stream);
   // Set data and remove "Filter" and "DecodeParms" fields from stream
   // dictionary.