Home | History | Annotate | Download | only in bsdiff 
      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