core/fpdfapi/edit/cpdf_pagecontentmanager.cpp - pdfium - Git at Google

 // Copyright 2018 The PDFium Authors
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.

 #include "core/fpdfapi/edit/cpdf_pagecontentmanager.h"

 #include <stdint.h>

 #include <map>
 #include <numeric>
 #include <set>
 #include <sstream>
 #include <utility>
 #include <vector>

 #include "core/fpdfapi/page/cpdf_pageobject.h"
 #include "core/fpdfapi/page/cpdf_pageobjectholder.h"
 #include "core/fpdfapi/parser/cpdf_array.h"
 #include "core/fpdfapi/parser/cpdf_dictionary.h"
 #include "core/fpdfapi/parser/cpdf_document.h"
 #include "core/fpdfapi/parser/cpdf_reference.h"
 #include "core/fpdfapi/parser/cpdf_stream.h"
 #include "core/fpdfapi/parser/object_tree_traversal_util.h"
 #include "third_party/abseil-cpp/absl/types/variant.h"
 #include "third_party/base/check.h"
 #include "third_party/base/containers/adapters.h"
 #include "third_party/base/containers/contains.h"
 #include "third_party/base/numerics/safe_conversions.h"

 CPDF_PageContentManager::CPDF_PageContentManager(
     CPDF_PageObjectHolder* page_obj_holder,
     CPDF_Document* document)
     : page_obj_holder_(page_obj_holder),
       document_(document),
       objects_with_multi_refs_(GetObjectsWithMultipleReferences(document_)) {
   RetainPtr<CPDF_Dictionary> page_dict = page_obj_holder_->GetMutableDict();
   RetainPtr<CPDF_Object> contents_obj =
       page_dict->GetMutableObjectFor("Contents");
   RetainPtr<CPDF_Array> contents_array = ToArray(contents_obj);
   if (contents_array) {
     CHECK(contents_array->IsInline());
     contents_ = std::move(contents_array);
     return;
   }

   RetainPtr<CPDF_Reference> contents_reference = ToReference(contents_obj);
   if (contents_reference) {
     RetainPtr<CPDF_Object> indirect_obj =
         contents_reference->GetMutableDirect();
     if (!indirect_obj)
       return;

     contents_array.Reset(indirect_obj->AsMutableArray());
     if (contents_array) {
       if (pdfium::Contains(objects_with_multi_refs_,
                            contents_array->GetObjNum())) {
         RetainPtr<CPDF_Array> cloned_contents_array =
             pdfium::WrapRetain(contents_array->Clone()->AsMutableArray());
         page_dict->SetFor("Contents", cloned_contents_array);
         contents_ = std::move(cloned_contents_array);
       } else {
         contents_ = std::move(contents_array);
       }
     } else if (indirect_obj->IsStream()) {
       contents_ = pdfium::WrapRetain(indirect_obj->AsMutableStream());
     }
   }
 }

 CPDF_PageContentManager::~CPDF_PageContentManager() {
   ExecuteScheduledRemovals();
 }

 RetainPtr<CPDF_Stream> CPDF_PageContentManager::GetStreamByIndex(
     size_t stream_index) {
   RetainPtr<CPDF_Stream> contents_stream = GetContentsStream();
   if (contents_stream) {
     return stream_index == 0 ? contents_stream : nullptr;
   }

   RetainPtr<CPDF_Array> contents_array = GetContentsArray();
   if (!contents_array) {
     return nullptr;
   }

   RetainPtr<CPDF_Reference> stream_reference =
       ToReference(contents_array->GetMutableObjectAt(stream_index));
   if (!stream_reference)
     return nullptr;

   return ToStream(stream_reference->GetMutableDirect());
 }

 size_t CPDF_PageContentManager::AddStream(fxcrt::ostringstream* buf) {
   auto new_stream = document_->NewIndirect<CPDF_Stream>();
   new_stream->SetDataFromStringstream(buf);

   // If there is one Content stream (not in an array), now there will be two, so
   // create an array with the old and the new one. The new one's index is 1.
   RetainPtr<CPDF_Stream> contents_stream = GetContentsStream();
   if (contents_stream) {
     auto new_contents_array = document_->NewIndirect<CPDF_Array>();
     new_contents_array->AppendNew<CPDF_Reference>(document_,
                                                   contents_stream->GetObjNum());
     new_contents_array->AppendNew<CPDF_Reference>(document_,
                                                   new_stream->GetObjNum());

     RetainPtr<CPDF_Dictionary> page_dict = page_obj_holder_->GetMutableDict();
     page_dict->SetNewFor<CPDF_Reference>("Contents", document_,
                                          new_contents_array->GetObjNum());
     contents_ = std::move(new_contents_array);
     return 1;
   }

   // If there is an array, just add the new stream to it, at the last position.
   RetainPtr<CPDF_Array> contents_array = GetContentsArray();
   if (contents_array) {
     contents_array->AppendNew<CPDF_Reference>(document_,
                                               new_stream->GetObjNum());
     return contents_array->size() - 1;
   }

   // There were no Contents, so add the new stream as the single Content stream.
   // Its index is 0.
   RetainPtr<CPDF_Dictionary> page_dict = page_obj_holder_->GetMutableDict();
   page_dict->SetNewFor<CPDF_Reference>("Contents", document_,
                                        new_stream->GetObjNum());
   contents_ = std::move(new_stream);
   return 0;
 }

 void CPDF_PageContentManager::UpdateStream(size_t stream_index,
                                            fxcrt::ostringstream* buf) {
   // If `buf` is now empty, remove the stream instead of setting the data.
   if (buf->tellp() <= 0) {
     ScheduleRemoveStreamByIndex(stream_index);
     return;
   }

   RetainPtr<CPDF_Stream> existing_stream = GetStreamByIndex(stream_index);
   CHECK(existing_stream);
   if (!pdfium::Contains(objects_with_multi_refs_,
                         existing_stream->GetObjNum())) {
     existing_stream->SetDataFromStringstreamAndRemoveFilter(buf);
     return;
   }

   if (GetContentsStream()) {
     auto new_stream = document_->NewIndirect<CPDF_Stream>();
     new_stream->SetDataFromStringstream(buf);
     RetainPtr<CPDF_Dictionary> page_dict = page_obj_holder_->GetMutableDict();
     page_dict->SetNewFor<CPDF_Reference>("Contents", document_,
                                          new_stream->GetObjNum());
   }

   RetainPtr<CPDF_Array> contents_array = GetContentsArray();
   if (!contents_array) {
     return;
   }

   RetainPtr<CPDF_Reference> stream_reference =
       ToReference(contents_array->GetMutableObjectAt(stream_index));
   if (!stream_reference) {
     return;
   }

   auto new_stream = document_->NewIndirect<CPDF_Stream>();
   new_stream->SetDataFromStringstream(buf);
   stream_reference->SetRef(document_, new_stream->GetObjNum());
 }

 void CPDF_PageContentManager::ScheduleRemoveStreamByIndex(size_t stream_index) {
   streams_to_remove_.insert(stream_index);
 }

 void CPDF_PageContentManager::ExecuteScheduledRemovals() {
   // This method assumes there are no dirty streams in the
   // CPDF_PageObjectHolder. If there were any, their indexes would need to be
   // updated.
   // Since CPDF_PageContentManager is only instantiated in
   // CPDF_PageContentGenerator::GenerateContent(), which cleans up the dirty
   // streams first, this should always be true.
   DCHECK(!page_obj_holder_->HasDirtyStreams());

   if (streams_to_remove_.empty()) {
     return;
   }

   RetainPtr<CPDF_Stream> contents_stream = GetContentsStream();
   if (contents_stream) {
     // Only stream that can be removed is 0.
     if (streams_to_remove_.find(0) != streams_to_remove_.end()) {
       RetainPtr<CPDF_Dictionary> page_dict = page_obj_holder_->GetMutableDict();
       page_dict->RemoveFor("Contents");
     }
     return;
   }

   RetainPtr<CPDF_Array> contents_array = GetContentsArray();
   if (!contents_array) {
     return;
   }

   // Initialize a vector with the old stream indexes. This will be used to build
   // a map from the old to the new indexes.
   std::vector<size_t> streams_left(contents_array->size());
   std::iota(streams_left.begin(), streams_left.end(), 0);

   // In reverse order so as to not change the indexes in the middle of the loop,
   // remove the streams.
   for (size_t stream_index : pdfium::base::Reversed(streams_to_remove_)) {
     contents_array->RemoveAt(stream_index);
     streams_left.erase(streams_left.begin() + stream_index);
   }

   // Create a mapping from the old to the new stream indexes, shifted due to the
   // deletion of the |streams_to_remove_|.
   std::map<size_t, size_t> stream_index_mapping;
   for (size_t i = 0; i < streams_left.size(); ++i) {
     stream_index_mapping[streams_left[i]] = i;
   }

   // Update the page objects' content stream indexes.
   for (const auto& obj : *page_obj_holder_) {
     int32_t old_stream_index = obj->GetContentStream();
     int32_t new_stream_index = pdfium::base::checked_cast<int32_t>(
         stream_index_mapping[old_stream_index]);
     obj->SetContentStream(new_stream_index);
   }

   // Even if there is a single content stream now, keep the array with a single
   // element. It's valid, a second stream might be added in the near future, and
   // the complexity of removing it is not worth it.
 }

 RetainPtr<CPDF_Stream> CPDF_PageContentManager::GetContentsStream() {
   if (absl::holds_alternative<RetainPtr<CPDF_Stream>>(contents_)) {
     return absl::get<RetainPtr<CPDF_Stream>>(contents_);
   }
   return nullptr;
 }

 RetainPtr<CPDF_Array> CPDF_PageContentManager::GetContentsArray() {
   if (absl::holds_alternative<RetainPtr<CPDF_Array>>(contents_)) {
     return absl::get<RetainPtr<CPDF_Array>>(contents_);
   }
   return nullptr;
 }
	// Copyright 2018 The PDFium Authors
	// Use of this source code is governed by a BSD-style license that can be
	// found in the LICENSE file.

	#include "core/fpdfapi/edit/cpdf_pagecontentmanager.h"

	#include <stdint.h>

	#include <map>
	#include <numeric>
	#include <set>
	#include <sstream>
	#include <utility>
	#include <vector>

	#include "core/fpdfapi/page/cpdf_pageobject.h"
	#include "core/fpdfapi/page/cpdf_pageobjectholder.h"
	#include "core/fpdfapi/parser/cpdf_array.h"
	#include "core/fpdfapi/parser/cpdf_dictionary.h"
	#include "core/fpdfapi/parser/cpdf_document.h"
	#include "core/fpdfapi/parser/cpdf_reference.h"
	#include "core/fpdfapi/parser/cpdf_stream.h"
	#include "core/fpdfapi/parser/object_tree_traversal_util.h"
	#include "third_party/abseil-cpp/absl/types/variant.h"
	#include "third_party/base/check.h"
	#include "third_party/base/containers/adapters.h"
	#include "third_party/base/containers/contains.h"
	#include "third_party/base/numerics/safe_conversions.h"

	CPDF_PageContentManager::CPDF_PageContentManager(
	CPDF_PageObjectHolder* page_obj_holder,
	CPDF_Document* document)
	: page_obj_holder_(page_obj_holder),
	document_(document),
	objects_with_multi_refs_(GetObjectsWithMultipleReferences(document_)) {
	RetainPtr<CPDF_Dictionary> page_dict = page_obj_holder_->GetMutableDict();
	RetainPtr<CPDF_Object> contents_obj =
	page_dict->GetMutableObjectFor("Contents");
	RetainPtr<CPDF_Array> contents_array = ToArray(contents_obj);
	if (contents_array) {
	CHECK(contents_array->IsInline());
	contents_ = std::move(contents_array);
	return;
	}

	RetainPtr<CPDF_Reference> contents_reference = ToReference(contents_obj);
	if (contents_reference) {
	RetainPtr<CPDF_Object> indirect_obj =
	contents_reference->GetMutableDirect();
	if (!indirect_obj)
	return;

	contents_array.Reset(indirect_obj->AsMutableArray());
	if (contents_array) {
	if (pdfium::Contains(objects_with_multi_refs_,
	contents_array->GetObjNum())) {
	RetainPtr<CPDF_Array> cloned_contents_array =
	pdfium::WrapRetain(contents_array->Clone()->AsMutableArray());
	page_dict->SetFor("Contents", cloned_contents_array);
	contents_ = std::move(cloned_contents_array);
	} else {
	contents_ = std::move(contents_array);
	}
	} else if (indirect_obj->IsStream()) {
	contents_ = pdfium::WrapRetain(indirect_obj->AsMutableStream());
	}
	}
	}

	CPDF_PageContentManager::~CPDF_PageContentManager() {
	ExecuteScheduledRemovals();
	}

	RetainPtr<CPDF_Stream> CPDF_PageContentManager::GetStreamByIndex(
	size_t stream_index) {
	RetainPtr<CPDF_Stream> contents_stream = GetContentsStream();
	if (contents_stream) {
	return stream_index == 0 ? contents_stream : nullptr;
	}

	RetainPtr<CPDF_Array> contents_array = GetContentsArray();
	if (!contents_array) {
	return nullptr;
	}

	RetainPtr<CPDF_Reference> stream_reference =
	ToReference(contents_array->GetMutableObjectAt(stream_index));
	if (!stream_reference)
	return nullptr;

	return ToStream(stream_reference->GetMutableDirect());
	}

	size_t CPDF_PageContentManager::AddStream(fxcrt::ostringstream* buf) {
	auto new_stream = document_->NewIndirect<CPDF_Stream>();
	new_stream->SetDataFromStringstream(buf);

	// If there is one Content stream (not in an array), now there will be two, so
	// create an array with the old and the new one. The new one's index is 1.
	RetainPtr<CPDF_Stream> contents_stream = GetContentsStream();
	if (contents_stream) {
	auto new_contents_array = document_->NewIndirect<CPDF_Array>();
	new_contents_array->AppendNew<CPDF_Reference>(document_,
	contents_stream->GetObjNum());
	new_contents_array->AppendNew<CPDF_Reference>(document_,
	new_stream->GetObjNum());

	RetainPtr<CPDF_Dictionary> page_dict = page_obj_holder_->GetMutableDict();
	page_dict->SetNewFor<CPDF_Reference>("Contents", document_,
	new_contents_array->GetObjNum());
	contents_ = std::move(new_contents_array);
	return 1;
	}

	// If there is an array, just add the new stream to it, at the last position.
	RetainPtr<CPDF_Array> contents_array = GetContentsArray();
	if (contents_array) {
	contents_array->AppendNew<CPDF_Reference>(document_,
	new_stream->GetObjNum());
	return contents_array->size() - 1;
	}

	// There were no Contents, so add the new stream as the single Content stream.
	// Its index is 0.
	RetainPtr<CPDF_Dictionary> page_dict = page_obj_holder_->GetMutableDict();
	page_dict->SetNewFor<CPDF_Reference>("Contents", document_,
	new_stream->GetObjNum());
	contents_ = std::move(new_stream);
	return 0;
	}

	void CPDF_PageContentManager::UpdateStream(size_t stream_index,
	fxcrt::ostringstream* buf) {
	// If `buf` is now empty, remove the stream instead of setting the data.
	if (buf->tellp() <= 0) {
	ScheduleRemoveStreamByIndex(stream_index);
	return;
	}

	RetainPtr<CPDF_Stream> existing_stream = GetStreamByIndex(stream_index);
	CHECK(existing_stream);
	if (!pdfium::Contains(objects_with_multi_refs_,
	existing_stream->GetObjNum())) {
	existing_stream->SetDataFromStringstreamAndRemoveFilter(buf);
	return;
	}

	if (GetContentsStream()) {
	auto new_stream = document_->NewIndirect<CPDF_Stream>();
	new_stream->SetDataFromStringstream(buf);
	RetainPtr<CPDF_Dictionary> page_dict = page_obj_holder_->GetMutableDict();
	page_dict->SetNewFor<CPDF_Reference>("Contents", document_,
	new_stream->GetObjNum());
	}

	RetainPtr<CPDF_Array> contents_array = GetContentsArray();
	if (!contents_array) {
	return;
	}

	RetainPtr<CPDF_Reference> stream_reference =
	ToReference(contents_array->GetMutableObjectAt(stream_index));
	if (!stream_reference) {
	return;
	}

	auto new_stream = document_->NewIndirect<CPDF_Stream>();
	new_stream->SetDataFromStringstream(buf);
	stream_reference->SetRef(document_, new_stream->GetObjNum());
	}

	void CPDF_PageContentManager::ScheduleRemoveStreamByIndex(size_t stream_index) {
	streams_to_remove_.insert(stream_index);
	}

	void CPDF_PageContentManager::ExecuteScheduledRemovals() {
	// This method assumes there are no dirty streams in the
	// CPDF_PageObjectHolder. If there were any, their indexes would need to be
	// updated.
	// Since CPDF_PageContentManager is only instantiated in
	// CPDF_PageContentGenerator::GenerateContent(), which cleans up the dirty
	// streams first, this should always be true.
	DCHECK(!page_obj_holder_->HasDirtyStreams());

	if (streams_to_remove_.empty()) {
	return;
	}

	RetainPtr<CPDF_Stream> contents_stream = GetContentsStream();
	if (contents_stream) {
	// Only stream that can be removed is 0.
	if (streams_to_remove_.find(0) != streams_to_remove_.end()) {
	RetainPtr<CPDF_Dictionary> page_dict = page_obj_holder_->GetMutableDict();
	page_dict->RemoveFor("Contents");
	}
	return;
	}

	RetainPtr<CPDF_Array> contents_array = GetContentsArray();
	if (!contents_array) {
	return;
	}

	// Initialize a vector with the old stream indexes. This will be used to build
	// a map from the old to the new indexes.
	std::vector<size_t> streams_left(contents_array->size());
	std::iota(streams_left.begin(), streams_left.end(), 0);

	// In reverse order so as to not change the indexes in the middle of the loop,
	// remove the streams.
	for (size_t stream_index : pdfium::base::Reversed(streams_to_remove_)) {
	contents_array->RemoveAt(stream_index);
	streams_left.erase(streams_left.begin() + stream_index);
	}

	// Create a mapping from the old to the new stream indexes, shifted due to the
	// deletion of the \|streams_to_remove_\|.
	std::map<size_t, size_t> stream_index_mapping;
	for (size_t i = 0; i < streams_left.size(); ++i) {
	stream_index_mapping[streams_left[i]] = i;
	}

	// Update the page objects' content stream indexes.
	for (const auto& obj : *page_obj_holder_) {
	int32_t old_stream_index = obj->GetContentStream();
	int32_t new_stream_index = pdfium::base::checked_cast<int32_t>(
	stream_index_mapping[old_stream_index]);
	obj->SetContentStream(new_stream_index);
	}

	// Even if there is a single content stream now, keep the array with a single
	// element. It's valid, a second stream might be added in the near future, and
	// the complexity of removing it is not worth it.
	}

	RetainPtr<CPDF_Stream> CPDF_PageContentManager::GetContentsStream() {
	if (absl::holds_alternative<RetainPtr<CPDF_Stream>>(contents_)) {
	return absl::get<RetainPtr<CPDF_Stream>>(contents_);
	}
	return nullptr;
	}

	RetainPtr<CPDF_Array> CPDF_PageContentManager::GetContentsArray() {
	if (absl::holds_alternative<RetainPtr<CPDF_Array>>(contents_)) {
	return absl::get<RetainPtr<CPDF_Array>>(contents_);
	}
	return nullptr;
	}