summaryrefslogtreecommitdiff
path: root/split_patch_writer.cc
blob: f863484645c2feddfcc0f079b6291ad4b74af391 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
// Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "bsdiff/split_patch_writer.h"

#include <algorithm>

#include "bsdiff/logging.h"

namespace bsdiff {

bool SplitPatchWriter::Init(size_t new_size) {
  new_size_ = new_size;
  // Fail gracefully if re-initialized.
  if (current_patch_ || patches_.empty())
    return false;

  size_t expected_patches = (new_size_ + new_chunk_size_ - 1) / new_chunk_size_;
  if (expected_patches == 0)
    expected_patches = 1;
  if (expected_patches != patches_.size()) {
    LOG(ERROR) << "Expected " << expected_patches << " for a new file of size "
               << new_size_ << " split in chunks of " << new_chunk_size_
               << " but got " << patches_.size() << " instead.";
    return false;
  }

  return patches_[0]->Init(
      std::min(static_cast<uint64_t>(new_size_), new_chunk_size_));
}

bool SplitPatchWriter::WriteDiffStream(const uint8_t* data, size_t size) {
  return WriteToStream(&PatchWriterInterface::WriteDiffStream, &diff_sizes_,
                       data, size);
}

bool SplitPatchWriter::WriteExtraStream(const uint8_t* data, size_t size) {
  return WriteToStream(&PatchWriterInterface::WriteExtraStream, &extra_sizes_,
                       data, size);
}

bool SplitPatchWriter::AddControlEntry(const ControlEntry& entry) {
  ControlEntry remaining(entry);
  while (written_output_ + remaining.diff_size + remaining.extra_size >=
         (current_patch_ + 1) * new_chunk_size_) {
    // We need to write some of the current ControlEntry to the current patch
    // and move on to the next patch if there are more bytes to write.
    uint64_t remaining_bytes =
        (current_patch_ + 1) * new_chunk_size_ - written_output_;
    // The offset_increment is always 0 in this case since we don't plan to read
    // for the old file in the current_patch anymore.
    ControlEntry current_patch_entry(0, 0, 0);

    current_patch_entry.diff_size =
        std::min(remaining.diff_size, remaining_bytes);
    remaining_bytes -= current_patch_entry.diff_size;
    remaining.diff_size -= current_patch_entry.diff_size;

    // This will be positive only if we used all the diff_size bytes.
    current_patch_entry.extra_size =
        std::min(remaining.extra_size, remaining_bytes);
    remaining_bytes -= current_patch_entry.extra_size;
    remaining.extra_size -= current_patch_entry.extra_size;

    AddControlEntryToCurrentPatch(current_patch_entry);

    if (remaining.diff_size + remaining.extra_size > 0) {
      current_patch_++;
      if (current_patch_ >= patches_.size()) {
        LOG(ERROR) << "Writing past the last patch";
        return false;
      }
      if (!patches_[current_patch_]->Init(std::min(
              new_size_ - current_patch_ * new_chunk_size_, new_chunk_size_))) {
        LOG(ERROR) << "Failed to initialize patch " << current_patch_;
        return false;
      }
      if (!remaining.diff_size) {
        // When no diff need to be sent to the output, we can just push the
        // existing old_pos_ as part of the current triplet, since the extra
        // stream doesn't use the old_pos_;
        remaining.offset_increment += old_pos_;
        old_pos_ = 0;
      }
      // Need to add a dummy control entry at the beginning of the patch to
      // offset the old_pos in the new patch, which would start at 0.
      if (old_pos_ != 0) {
        if (!patches_[current_patch_]->AddControlEntry(
                ControlEntry(0, 0, old_pos_)))
          return false;
      }
    } else {
      // There was no need to write more bytes past the current patch, so just
      // update the old_pos_ we are tracking for the next patch, if any.
      old_pos_ += remaining.offset_increment;
      return true;
    }
  }

  // Trivial entries will be ignored.
  return AddControlEntryToCurrentPatch(remaining);
}

bool SplitPatchWriter::Close() {
  uint64_t missing_bytes = 0;
  for (auto size : diff_sizes_)
    missing_bytes += size;
  for (auto size : extra_sizes_)
    missing_bytes += size;
  if (missing_bytes > 0) {
    LOG(ERROR) << "Close() called but there are " << missing_bytes
               << " bytes missing from Write*Stream() calls";
    return false;
  }

  // |current_patch_| holds the last patch that was Init()'ed. If there are more
  // patches in the list those have not been initialized/closed, which is a
  // programming error.
  if (current_patch_ + 1 != patches_.size()) {
    LOG(ERROR)
        << "Close() called but no bytes habe been written to the last patch";
    return false;
  }

  // Close all the remaining streams.
  for (; closed_patches_ < patches_.size(); closed_patches_++) {
    if (!patches_[closed_patches_]->Close())
      return false;
  }
  return true;
}

bool SplitPatchWriter::AddControlEntryToCurrentPatch(
    const ControlEntry& entry) {
  // Ignore trivial control entries that don't modify the state.
  if (!entry.diff_size && !entry.extra_size && !entry.offset_increment)
    return true;

  if (current_patch_ >= patches_.size()) {
    LOG(ERROR) << "Writing past the last patch";
    return false;
  }
  old_pos_ += entry.diff_size + entry.offset_increment;
  written_output_ += entry.diff_size + entry.extra_size;
  // Register the diff/extra sizes as required bytes for the current patch.
  diff_sizes_[current_patch_] += entry.diff_size;
  extra_sizes_[current_patch_] += entry.extra_size;
  return patches_[current_patch_]->AddControlEntry(entry);
}

bool SplitPatchWriter::WriteToStream(WriteStreamMethod method,
                                     std::vector<size_t>* sizes_vector,
                                     const uint8_t* data,
                                     size_t size) {
  size_t written = 0;
  for (size_t i = closed_patches_; i <= current_patch_ && written < size; i++) {
    if ((*sizes_vector)[i]) {
      size_t flush_size = std::min(size - written, (*sizes_vector)[i]);
      if (!(patches_[i]->*method)(data + written, flush_size))
        return false;
      written += flush_size;
      (*sizes_vector)[i] -= flush_size;
    }

    if (i < current_patch_ && !diff_sizes_[i] && !extra_sizes_[i]) {
      // All bytes expected for the patch i are already sent.
      if (!patches_[i]->Close())
        return false;
      closed_patches_++;
    }
  }
  if (written < size) {
    LOG(ERROR) << "Calling Write*Stream() before the corresponding "
                  "AddControlEntry() is not supported.";
    return false;
  }
  return true;
}

}  // namespace bsdiff