[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
This commit is contained in:
Kongqun Yang 2017-03-16 13:56:11 -07:00 committed by KongQun Yang
parent 15fd745fa7
commit 0c3fb49eeb
10 changed files with 43 additions and 41 deletions

View File

@ -52,7 +52,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());
}

View File

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

View File

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

View File

@ -38,7 +38,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
@ -57,8 +56,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,
StreamInfo* info,
Status Initialize(StreamInfo* info,
ProgressListener* progress_listener,
MuxerListener* muxer_listener,
KeySource* encryption_key_source,
@ -119,7 +117,7 @@ class Segmenter {
int track_id() const { return track_id_; }
uint64_t segment_payload_pos() const { return segment_payload_pos_; }
virtual Status DoInitialize(std::unique_ptr<MkvWriter> writer) = 0;
virtual Status DoInitialize() = 0;
virtual Status DoFinalize() = 0;
private:

View File

@ -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 */,
NULL /* muxer_listener */, key_source, 0 /* max_sd_pixels */,
0 /* max_hd_pixels */, 0 /* max_uhd1_pixels */,
1 /* clear_lead_in_seconds */));
info, NULL /* progress_listener */, NULL /* muxer_listener */,
key_source, 0 /* max_sd_pixels */, 0 /* max_hd_pixels */,
0 /* max_uhd1_pixels */, 1 /* clear_lead_in_seconds */));
*result = std::move(segmenter);
}

View File

@ -52,8 +52,15 @@ bool SingleSegmentSegmenter::GetIndexRangeStartAndEnd(uint64_t* start,
return true;
}
Status SingleSegmentSegmenter::DoInitialize(std::unique_ptr<MkvWriter> writer) {
Status SingleSegmentSegmenter::DoInitialize() {
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());

View File

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

View File

@ -68,22 +68,20 @@ TwoPassSingleSegmentSegmenter::TwoPassSingleSegmentSegmenter(
TwoPassSingleSegmentSegmenter::~TwoPassSingleSegmentSegmenter() {}
Status TwoPassSingleSegmentSegmenter::DoInitialize(
std::unique_ptr<MkvWriter> writer) {
Status TwoPassSingleSegmentSegmenter::DoInitialize() {
// 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() {
@ -94,18 +92,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());
if (!cues()->Write(real_writer_.get()))
set_index_start(real_writer->Position());
if (!cues()->Write(real_writer.get()))
return Status(error::FILE_FAILURE, "Error writing Cues data.");
set_index_end(real_writer_->Position() - 1);
DCHECK_EQ(real_writer_->Position(),
set_index_end(real_writer->Position() - 1);
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.
@ -120,7 +123,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.");
}
@ -131,7 +134,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(

View File

@ -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);

View File

@ -37,11 +37,6 @@ Status WebMMuxer::InitializeMuxer() {
"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 {
@ -49,9 +44,9 @@ Status WebMMuxer::InitializeMuxer() {
}
Status initialized = segmenter_->Initialize(
std::move(writer), streams()[0].get(), progress_listener(),
muxer_listener(), encryption_key_source(), max_sd_pixels(),
max_hd_pixels(), max_uhd1_pixels(), clear_lead_in_seconds());
streams()[0].get(), progress_listener(), muxer_listener(),
encryption_key_source(), max_sd_pixels(), max_hd_pixels(),
max_uhd1_pixels(), clear_lead_in_seconds());
if (!initialized.ok())
return initialized;