[WebM] Fix output truncated if using the same file name for I/O

Open output file in DoFinalize() to ensure that the input file has been
closed already.

Fixes #210

Change-Id: I935941b31c667e49be030c8da9f953d8387c7a9d

packager/media/formats/webm/multi_segment_segmenter.cc

@@ -31,7 +31,11 @@ bool MultiSegmentSegmenter::GetIndexRangeStartAndEnd(uint64_t* start,
   return false;
 }
 
-Status MultiSegmentSegmenter::DoInitialize(std::unique_ptr<MkvWriter> writer) {
+Status MultiSegmentSegmenter::DoInitialize() {
+  std::unique_ptr<MkvWriter> writer(new MkvWriter);
+  Status status = writer->Open(options().output_file_name);
+  if (!status.ok())
+    return status;
   writer_ = std::move(writer);
   return WriteSegmentHeader(0, writer_.get());
 }

packager/media/formats/webm/multi_segment_segmenter.h

@@ -34,7 +34,7 @@ class MultiSegmentSegmenter : public Segmenter {
 
  protected:
   // Segmenter implementation overrides.
-  Status DoInitialize(std::unique_ptr<MkvWriter> writer) override;
+  Status DoInitialize() override;
   Status DoFinalize() override;
 
  private:

packager/media/formats/webm/segmenter.cc

@@ -48,8 +48,7 @@ Segmenter::Segmenter(const MuxerOptions& options)
 
 Segmenter::~Segmenter() {}
 
-Status Segmenter::Initialize(std::unique_ptr<MkvWriter> writer,
+Status Segmenter::Initialize(StreamInfo* info,
-                             StreamInfo* info,
                              ProgressListener* progress_listener,
                              MuxerListener* muxer_listener,
                              KeySource* encryption_key_source,
@@ -107,7 +106,7 @@ Status Segmenter::Initialize(std::unique_ptr<MkvWriter> writer,
   if (!status.ok())
     return status;
 
-  return DoInitialize(std::move(writer));
+  return DoInitialize();
 }
 
 Status Segmenter::Finalize() {

packager/media/formats/webm/segmenter.h

@@ -39,7 +39,6 @@ class Segmenter {
   /// Initialize the segmenter.
   /// Calling other public methods of this class without this method returning
   /// Status::OK results in an undefined behavior.
-  /// @param writer contains the output file (or init file in multi-segment).
   /// @param info The stream info for the stream being segmented.
   /// @param muxer_listener receives muxer events. Can be NULL.
   /// @param encryption_key_source points to the key source which contains
@@ -58,8 +57,7 @@ class Segmenter {
   ///        it is UHD1. Otherwise it is UHD2.
   /// @param clear_time specifies clear lead duration in seconds.
   /// @return OK on success, an error status otherwise.
-  Status Initialize(std::unique_ptr<MkvWriter> writer,
+  Status Initialize(StreamInfo* info,
-                    StreamInfo* info,
                     ProgressListener* progress_listener,
                     MuxerListener* muxer_listener,
                     KeySource* encryption_key_source,
@@ -118,7 +116,7 @@ class Segmenter {
     return cluster_length_in_time_scale_;
   }
 
-  virtual Status DoInitialize(std::unique_ptr<MkvWriter> writer) = 0;
+  virtual Status DoInitialize() = 0;
   virtual Status DoFinalize() = 0;
 
  private:

packager/media/formats/webm/segmenter_test_base.h

@@ -51,13 +51,10 @@ class SegmentTestBase : public ::testing::Test {
       std::unique_ptr<webm::Segmenter>* result) const {
     std::unique_ptr<S> segmenter(new S(options));
 
-    std::unique_ptr<MkvWriter> writer(new MkvWriter());
-    ASSERT_OK(writer->Open(options.output_file_name));
     ASSERT_OK(segmenter->Initialize(
-        std::move(writer), info, NULL /* progress_listener */,
+        info, NULL /* progress_listener */, NULL /* muxer_listener */,
-        NULL /* muxer_listener */, key_source, 0 /* max_sd_pixels */,
+        key_source, 0 /* max_sd_pixels */, 0 /* max_hd_pixels */,
-        0 /* max_hd_pixels */, 0 /* max_uhd1_pixels */,
+        0 /* max_uhd1_pixels */, 1 /* clear_lead_in_seconds */));
-        1 /* clear_lead_in_seconds */));
     *result = std::move(segmenter);
   }
 

packager/media/formats/webm/single_segment_segmenter.cc

@@ -18,8 +18,15 @@ SingleSegmentSegmenter::SingleSegmentSegmenter(const MuxerOptions& options)
 
 SingleSegmentSegmenter::~SingleSegmentSegmenter() {}
 
-Status SingleSegmentSegmenter::DoInitialize(std::unique_ptr<MkvWriter> writer) {
+Status SingleSegmentSegmenter::DoInitialize() {
-  writer_ = std::move(writer);
+  if (!writer_) {
+    std::unique_ptr<MkvWriter> writer(new MkvWriter);
+    Status status = writer->Open(options().output_file_name);
+    if (!status.ok())
+      return status;
+    writer_ = std::move(writer);
+  }
+
   Status ret = WriteSegmentHeader(0, writer_.get());
   init_end_ = writer_->Position() - 1;
   seek_head()->set_cluster_pos(init_end_ + 1 - segment_payload_pos());

packager/media/formats/webm/single_segment_segmenter.h

@@ -45,7 +45,7 @@ class SingleSegmentSegmenter : public Segmenter {
   }
 
   // Segmenter implementation overrides.
-  Status DoInitialize(std::unique_ptr<MkvWriter> writer) override;
+  Status DoInitialize() override;
   Status DoFinalize() override;
 
  private:

packager/media/formats/webm/two_pass_single_segment_segmenter.cc

@@ -69,22 +69,20 @@ TwoPassSingleSegmentSegmenter::TwoPassSingleSegmentSegmenter(
 
 TwoPassSingleSegmentSegmenter::~TwoPassSingleSegmentSegmenter() {}
 
-Status TwoPassSingleSegmentSegmenter::DoInitialize(
+Status TwoPassSingleSegmentSegmenter::DoInitialize() {
-    std::unique_ptr<MkvWriter> writer) {
   // Assume the amount of time to copy the temp file as the same amount
   // of time as to make it.
   set_progress_target(info()->duration() * 2);
 
-  real_writer_ = std::move(writer);
-
   if (!TempFilePath(options().temp_dir, &temp_file_name_))
     return Status(error::FILE_FAILURE, "Unable to create temporary file.");
   std::unique_ptr<MkvWriter> temp(new MkvWriter);
   Status status = temp->Open(temp_file_name_);
   if (!status.ok())
     return status;
+  set_writer(std::move(temp));
 
-  return SingleSegmentSegmenter::DoInitialize(std::move(temp));
+  return SingleSegmentSegmenter::DoInitialize();
 }
 
 Status TwoPassSingleSegmentSegmenter::DoFinalize() {
@@ -98,18 +96,23 @@ Status TwoPassSingleSegmentSegmenter::DoFinalize() {
   seek_head()->set_cluster_pos(cues_pos + cues_size);
 
   // Write the header to the real output file.
+  std::unique_ptr<MkvWriter> real_writer(new MkvWriter);
+  Status status = real_writer->Open(options().output_file_name);
+  if (!status.ok())
+    return status;
+
   const uint64_t file_size = writer()->Position() + cues_size;
-  Status temp = WriteSegmentHeader(file_size, real_writer_.get());
+  Status temp = WriteSegmentHeader(file_size, real_writer.get());
   if (!temp.ok())
     return temp;
-  DCHECK_EQ(real_writer_->Position(), static_cast<int64_t>(header_size));
+  DCHECK_EQ(real_writer->Position(), static_cast<int64_t>(header_size));
 
   // Write the cues to the real output file.
-  set_index_start(real_writer_->Position());
+  set_index_start(real_writer->Position());
-  if (!cues()->Write(real_writer_.get()))
+  if (!cues()->Write(real_writer.get()))
     return Status(error::FILE_FAILURE, "Error writing Cues data.");
-  set_index_end(real_writer_->Position() - 1);
+  set_index_end(real_writer->Position() - 1);
-  DCHECK_EQ(real_writer_->Position(),
+  DCHECK_EQ(real_writer->Position(),
             static_cast<int64_t>(segment_payload_pos() + cues_pos + cues_size));
 
   // Close the temp file and open it for reading.
@@ -124,7 +127,7 @@ Status TwoPassSingleSegmentSegmenter::DoFinalize() {
     return Status(error::FILE_FAILURE, "Error reading temp file.");
 
   // Copy the rest of the data over.
-  if (!CopyFileWithClusterRewrite(temp_reader.get(), real_writer_.get(),
+  if (!CopyFileWithClusterRewrite(temp_reader.get(), real_writer.get(),
                                   cluster()->Size())) {
     return Status(error::FILE_FAILURE, "Error copying temp file.");
   }
@@ -135,7 +138,7 @@ Status TwoPassSingleSegmentSegmenter::DoFinalize() {
     LOG(WARNING) << "Unable to delete temporary file " << temp_file_name_;
   }
 
-  return real_writer_->Close();
+  return real_writer->Close();
 }
 
 bool TwoPassSingleSegmentSegmenter::CopyFileWithClusterRewrite(

packager/media/formats/webm/two_pass_single_segment_segmenter.h

@@ -29,7 +29,7 @@ class TwoPassSingleSegmentSegmenter : public SingleSegmentSegmenter {
   ~TwoPassSingleSegmentSegmenter() override;
 
   // Segmenter implementation overrides.
-  Status DoInitialize(std::unique_ptr<MkvWriter> writer) override;
+  Status DoInitialize() override;
   Status DoFinalize() override;
 
  private:
@@ -41,7 +41,6 @@ class TwoPassSingleSegmentSegmenter : public SingleSegmentSegmenter {
                                   MkvWriter* dest,
                                   uint64_t last_size);
 
-  std::unique_ptr<MkvWriter> real_writer_;
   std::string temp_file_name_;
 
   DISALLOW_COPY_AND_ASSIGN(TwoPassSingleSegmentSegmenter);

packager/media/formats/webm/webm_muxer.cc

@@ -38,11 +38,6 @@ Status WebMMuxer::Initialize() {
                   "WebM does not support protection scheme other than 'cenc'.");
   }
 
-  std::unique_ptr<MkvWriter> writer(new MkvWriter);
-  Status status = writer->Open(options().output_file_name);
-  if (!status.ok())
-    return status;
-
   if (!options().segment_template.empty()) {
     segmenter_.reset(new MultiSegmentSegmenter(options()));
   } else {
@@ -50,9 +45,9 @@ Status WebMMuxer::Initialize() {
   }
 
   Status initialized = segmenter_->Initialize(
-      std::move(writer), streams()[0]->info().get(), progress_listener(),
+      streams()[0]->info().get(), progress_listener(), muxer_listener(),
-      muxer_listener(), encryption_key_source(), max_sd_pixels(),
+      encryption_key_source(), max_sd_pixels(), max_hd_pixels(),
-      max_hd_pixels(), max_uhd1_pixels(), clear_lead_in_seconds());
+      max_uhd1_pixels(), clear_lead_in_seconds());
 
   if (!initialized.ok())
     return initialized;