formats/pdf/writer.cc

// Copyright 2026 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

#include "formats/pdf/writer.h"

#include <cstdint>
#include <string>
#include <utility>

#include "absl/base/nullability.h"
#include "absl/status/status.h"
#include "absl/status/status_macros.h"
#include "absl/status/statusor.h"
#include "absl/strings/cord.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "crypto/algorithms.h"
#include "formats/hasher_data.h"
#include "formats/pdf/embedder.h"
#include "proto/generator.pb.h"
#include "riegeli/bytes/cord_reader.h"
#include "riegeli/bytes/cord_writer.h"

namespace credentio {
namespace {

absl::StatusOr<int64_t> RewriteToCordForEmbedding(
    riegeli::Reader& input, int64_t estimated_manifest_store_size,
    absl::Cord* result) {
  if (!input.Seek(0)) {
    return input.StatusOrAnnotate(
        absl::DataLossError("failed to seek to start of input"));
  }
  riegeli::CordWriter output(result);
  PdfEmbedder embedder;
  ABSL_ASSIGN_OR_RETURN(int64_t manifest_store_offset,
                        embedder.RewriteForEmbedding(
                            input, estimated_manifest_store_size, output));
  if (!output.Close()) {
    return output.status();
  }
  return manifest_store_offset;
}
}  // namespace

absl::Status PdfWriter::Write(GenerateManifestStoreRequest request,
                              riegeli::Reader* absl_nonnull source,
                              riegeli::Writer* absl_nonnull destination) {
  if (request.has_hard_binding()) {
    return absl::InvalidArgumentError(
        "request must not include `hard_binding` (binding is computed by "
        "PdfWriter)");
  }

  // The two-phase approach is required for PDF: The sections to be hashed can't
  // be finalized until the manifest store size is known, because the manifest
  // size affects the offsets of objects in the cross-reference table.

  // Phase 1: Rewrite the asset to enable embedding, compute preliminary hard
  // binding, produce a manifest store template, and embed a zero-filled
  // placeholder in the asset.
  HashAlgorithm alg = HashAlgorithm::kSha256;
  ABSL_ASSIGN_OR_RETURN(*request.mutable_hard_binding()->mutable_data_hash(),
                        DataHashTemplate(alg));
  ABSL_ASSIGN_OR_RETURN(ManifestStoreTemplate manifest_store_template,
                        generator_.GenerateManifestStoreTemplate(request));
  constexpr int64_t kMaxManifestStoreSizeBytes = 10 * 1024 * 1024;  // 10 MiB
  if (manifest_store_template.manifest_store_size() <= 0 ||
      manifest_store_template.manifest_store_size() >
          kMaxManifestStoreSizeBytes) {
    return absl::InvalidArgumentError(
        absl::StrCat("invalid manifest store template size: ",
                     manifest_store_template.manifest_store_size()));
  }
  absl::Cord phase1_result;
  ABSL_ASSIGN_OR_RETURN(
      int64_t manifest_store_offset,
      RewriteToCordForEmbedding(*source,
                                /*estimated_manifest_store_size=*/
                                manifest_store_template.manifest_store_size(),
                                &phase1_result),
      _ << "embedding manifest store template");

  // Phase 2: Compute final hard binding based on the updated content, finalize
  // the manifest store, and embed the final manifest store in the asset.
  riegeli::CordReader phase2_reader(&phase1_result);
  FinalizeManifestStoreRequest finalize_request;
  *finalize_request.mutable_manifest_store_template() =
      std::move(manifest_store_template);
  ABSL_ASSIGN_OR_RETURN(
      *finalize_request.mutable_hard_binding()->mutable_data_hash(),
      DigestDataHash(
          phase2_reader, HashAlgorithm::kSha256,
          /*manifest_store_offset=*/manifest_store_offset,
          /*manifest_store_length=*/
          finalize_request.manifest_store_template().manifest_store_size()));
  ABSL_ASSIGN_OR_RETURN(auto response,
                        generator_.FinalizeManifestStore(finalize_request));
  if (response.manifest_store().size() !=
      finalize_request.manifest_store_template().manifest_store_size()) {
    return absl::InternalError("manifest store size does not match template");
  }
  ABSL_RETURN_IF_ERROR(PdfEmbedder().Embed(
      response.manifest_store(), manifest_store_offset, phase1_result));
  if (!destination->Write(phase1_result)) {
    return absl::InternalError("failed to write result");
  }
  return absl::OkStatus();
}
}  // namespace credentio