1 // Copyright 2017 The Chromium OS Authors. All rights reserved. 2 // Use of this source code is governed by a BSD-style license that can be 3 // found in the LICENSE file. 4 5 #include "bsdiff/split_patch_writer.h" 6 7 #include <algorithm> 8 9 #include "bsdiff/logging.h" 10 11 namespace bsdiff { 12 13 bool SplitPatchWriter::Init(size_t new_size) { 14 new_size_ = new_size; 15 // Fail gracefully if re-initialized. 16 if (current_patch_ || patches_.empty()) 17 return false; 18 19 size_t expected_patches = (new_size_ + new_chunk_size_ - 1) / new_chunk_size_; 20 if (expected_patches == 0) 21 expected_patches = 1; 22 if (expected_patches != patches_.size()) { 23 LOG(ERROR) << "Expected " << expected_patches << " for a new file of size " 24 << new_size_ << " split in chunks of " << new_chunk_size_ 25 << " but got " << patches_.size() << " instead."; 26 return false; 27 } 28 29 return patches_[0]->Init( 30 std::min(static_cast<uint64_t>(new_size_), new_chunk_size_)); 31 } 32 33 bool SplitPatchWriter::WriteDiffStream(const uint8_t* data, size_t size) { 34 return WriteToStream(&PatchWriterInterface::WriteDiffStream, &diff_sizes_, 35 data, size); 36 } 37 38 bool SplitPatchWriter::WriteExtraStream(const uint8_t* data, size_t size) { 39 return WriteToStream(&PatchWriterInterface::WriteExtraStream, &extra_sizes_, 40 data, size); 41 } 42 43 bool SplitPatchWriter::AddControlEntry(const ControlEntry& entry) { 44 ControlEntry remaining(entry); 45 while (written_output_ + remaining.diff_size + remaining.extra_size >= 46 (current_patch_ + 1) * new_chunk_size_) { 47 // We need to write some of the current ControlEntry to the current patch 48 // and move on to the next patch if there are more bytes to write. 49 uint64_t remaining_bytes = 50 (current_patch_ + 1) * new_chunk_size_ - written_output_; 51 // The offset_increment is always 0 in this case since we don't plan to read 52 // for the old file in the current_patch anymore. 53 ControlEntry current_patch_entry(0, 0, 0); 54 55 current_patch_entry.diff_size = 56 std::min(remaining.diff_size, remaining_bytes); 57 remaining_bytes -= current_patch_entry.diff_size; 58 remaining.diff_size -= current_patch_entry.diff_size; 59 60 // This will be positive only if we used all the diff_size bytes. 61 current_patch_entry.extra_size = 62 std::min(remaining.extra_size, remaining_bytes); 63 remaining_bytes -= current_patch_entry.extra_size; 64 remaining.extra_size -= current_patch_entry.extra_size; 65 66 AddControlEntryToCurrentPatch(current_patch_entry); 67 68 if (remaining.diff_size + remaining.extra_size > 0) { 69 current_patch_++; 70 if (current_patch_ >= patches_.size()) { 71 LOG(ERROR) << "Writing past the last patch"; 72 return false; 73 } 74 if (!patches_[current_patch_]->Init(std::min( 75 new_size_ - current_patch_ * new_chunk_size_, new_chunk_size_))) { 76 LOG(ERROR) << "Failed to initialize patch " << current_patch_; 77 return false; 78 } 79 if (!remaining.diff_size) { 80 // When no diff need to be sent to the output, we can just push the 81 // existing old_pos_ as part of the current triplet, since the extra 82 // stream doesn't use the old_pos_; 83 remaining.offset_increment += old_pos_; 84 old_pos_ = 0; 85 } 86 // Need to add a dummy control entry at the beginning of the patch to 87 // offset the old_pos in the new patch, which would start at 0. 88 if (old_pos_ != 0) { 89 if (!patches_[current_patch_]->AddControlEntry( 90 ControlEntry(0, 0, old_pos_))) 91 return false; 92 } 93 } else { 94 // There was no need to write more bytes past the current patch, so just 95 // update the old_pos_ we are tracking for the next patch, if any. 96 old_pos_ += remaining.offset_increment; 97 return true; 98 } 99 } 100 101 // Trivial entries will be ignored. 102 return AddControlEntryToCurrentPatch(remaining); 103 } 104 105 bool SplitPatchWriter::Close() { 106 uint64_t missing_bytes = 0; 107 for (auto size : diff_sizes_) 108 missing_bytes += size; 109 for (auto size : extra_sizes_) 110 missing_bytes += size; 111 if (missing_bytes > 0) { 112 LOG(ERROR) << "Close() called but there are " << missing_bytes 113 << " bytes missing from Write*Stream() calls"; 114 return false; 115 } 116 117 // |current_patch_| holds the last patch that was Init()'ed. If there are more 118 // patches in the list those have not been initialized/closed, which is a 119 // programming error. 120 if (current_patch_ + 1 != patches_.size()) { 121 LOG(ERROR) 122 << "Close() called but no bytes habe been written to the last patch"; 123 return false; 124 } 125 126 // Close all the remaining streams. 127 for (; closed_patches_ < patches_.size(); closed_patches_++) { 128 if (!patches_[closed_patches_]->Close()) 129 return false; 130 } 131 return true; 132 } 133 134 bool SplitPatchWriter::AddControlEntryToCurrentPatch( 135 const ControlEntry& entry) { 136 // Ignore trivial control entries that don't modify the state. 137 if (!entry.diff_size && !entry.extra_size && !entry.offset_increment) 138 return true; 139 140 if (current_patch_ >= patches_.size()) { 141 LOG(ERROR) << "Writing past the last patch"; 142 return false; 143 } 144 old_pos_ += entry.diff_size + entry.offset_increment; 145 written_output_ += entry.diff_size + entry.extra_size; 146 // Register the diff/extra sizes as required bytes for the current patch. 147 diff_sizes_[current_patch_] += entry.diff_size; 148 extra_sizes_[current_patch_] += entry.extra_size; 149 return patches_[current_patch_]->AddControlEntry(entry); 150 } 151 152 bool SplitPatchWriter::WriteToStream(WriteStreamMethod method, 153 std::vector<size_t>* sizes_vector, 154 const uint8_t* data, 155 size_t size) { 156 size_t written = 0; 157 for (size_t i = closed_patches_; i <= current_patch_ && written < size; i++) { 158 if ((*sizes_vector)[i]) { 159 size_t flush_size = std::min(size - written, (*sizes_vector)[i]); 160 if (!(patches_[i]->*method)(data + written, flush_size)) 161 return false; 162 written += flush_size; 163 (*sizes_vector)[i] -= flush_size; 164 } 165 166 if (i < current_patch_ && !diff_sizes_[i] && !extra_sizes_[i]) { 167 // All bytes expected for the patch i are already sent. 168 if (!patches_[i]->Close()) 169 return false; 170 closed_patches_++; 171 } 172 } 173 if (written < size) { 174 LOG(ERROR) << "Calling Write*Stream() before the corresponding " 175 "AddControlEntry() is not supported."; 176 return false; 177 } 178 return true; 179 } 180 181 } // namespace bsdiff 182