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/endsley_patch_writer.h"
6 
7 #include <string.h>
8 
9 #include <algorithm>
10 
11 #include "bsdiff/brotli_compressor.h"
12 #include "bsdiff/bz2_compressor.h"
13 #include "bsdiff/logging.h"
14 
15 namespace {
16 
17 constexpr uint8_t kEndsleyMagicHeader[] = "ENDSLEY/BSDIFF43";
18 
19 void EncodeInt64(int64_t x, uint8_t* buf) {
20   uint64_t y = x < 0 ? (1ULL << 63ULL) - x : x;
21   for (int i = 0; i < 8; ++i) {
22     buf[i] = y & 0xff;
23     y /= 256;
24   }
25 }
26 
27 // The minimum size that we would consider flushing out.
28 constexpr size_t kMinimumFlushSize = 1024 * 1024;  // 1 MiB
29 
30 }  // namespace
31 
32 namespace bsdiff {
33 
34 bool EndsleyPatchWriter::Init(size_t new_size) {
35   switch (compressor_type_) {
36     case CompressorType::kNoCompression:
37       // The patch is uncompressed and it will need exactly:
38       //   new_size + 24 * len(control_entries) + sizeof(header)
39       // We don't know the length of the control entries yet, but we can reserve
40       // enough space to hold at least |new_size|.
41       patch_->clear();
42       patch_->reserve(new_size);
43       break;
44     case CompressorType::kBrotli:
45       compressor_.reset(new BrotliCompressor(brotli_quality_));
46       if (!compressor_) {
47         LOG(ERROR) << "Error creating brotli compressor.";
48         return false;
49       }
50       break;
51     case CompressorType::kBZ2:
52       compressor_.reset(new BZ2Compressor());
53       if (!compressor_) {
54         LOG(ERROR) << "Error creating BZ2 compressor.";
55         return false;
56       }
57       break;
58   }
59 
60   // Header is the magic followed by the new length.
61   uint8_t header[24];
62   memcpy(header, kEndsleyMagicHeader, 16);
63   EncodeInt64(new_size, header + 16);
64   EmitBuffer(header, sizeof(header));
65   return true;
66 }
67 
68 bool EndsleyPatchWriter::WriteDiffStream(const uint8_t* data, size_t size) {
69   if (!size)
70     return true;
71   // Speed-up the common case where the diff stream data is added right after
72   // the control entry that refers to it.
73   if (control_.empty() && pending_diff_ >= size) {
74     pending_diff_ -= size;
75     EmitBuffer(data, size);
76     return true;
77   }
78 
79   diff_data_.insert(diff_data_.end(), data, data + size);
80   return true;
81 }
82 
83 bool EndsleyPatchWriter::WriteExtraStream(const uint8_t* data, size_t size) {
84   if (!size)
85     return true;
86   // Speed-up the common case where the extra stream data is added right after
87   // the diff stream data and the control entry that refers to it. Note that
88   // the diff data comes first so we need to make sure it is all out.
89   if (control_.empty() && !pending_diff_ && pending_extra_ >= size) {
90     pending_extra_ -= size;
91     EmitBuffer(data, size);
92     return true;
93   }
94 
95   extra_data_.insert(extra_data_.end(), data, data + size);
96   return true;
97 }
98 
99 bool EndsleyPatchWriter::AddControlEntry(const ControlEntry& entry) {
100   // Speed-up the common case where the control entry is added when there's
101   // nothing else pending.
102   if (control_.empty() && diff_data_.empty() && extra_data_.empty() &&
103       !pending_diff_ && !pending_extra_) {
104     pending_diff_ = entry.diff_size;
105     pending_extra_ = entry.extra_size;
106     EmitControlEntry(entry);
107     return true;
108   }
109 
110   control_.push_back(entry);
111   pending_control_data_ += entry.diff_size + entry.extra_size;
112 
113   // Check whether it is worth Flushing the entries now that the we have more
114   // control entries. We need control entries to write enough output data, and
115   // we need that output data to be at least 50% of the available diff and extra
116   // data. This last requirement is to reduce the overhead of removing the
117   // flushed data.
118   if (pending_control_data_ > kMinimumFlushSize &&
119       (diff_data_.size() + extra_data_.size()) / 2 <= pending_control_data_) {
120     Flush();
121   }
122 
123   return true;
124 }
125 
126 bool EndsleyPatchWriter::Close() {
127   // Flush any pending data.
128   Flush();
129 
130   if (pending_diff_ || pending_extra_ || !control_.empty()) {
131     LOG(ERROR) << "Insufficient data sent to diff/extra streams";
132     return false;
133   }
134 
135   if (!diff_data_.empty() || !extra_data_.empty()) {
136     LOG(ERROR) << "Pending data to diff/extra not flushed out on Close()";
137     return false;
138   }
139 
140   if (compressor_) {
141     if (!compressor_->Finish())
142       return false;
143     *patch_ = compressor_->GetCompressedData();
144   }
145 
146   return true;
147 }
148 
149 void EndsleyPatchWriter::EmitControlEntry(const ControlEntry& entry) {
150   // Generate the 24 byte control entry.
151   uint8_t buf[24];
152   EncodeInt64(entry.diff_size, buf);
153   EncodeInt64(entry.extra_size, buf + 8);
154   EncodeInt64(entry.offset_increment, buf + 16);
155   EmitBuffer(buf, sizeof(buf));
156 }
157 
158 void EndsleyPatchWriter::EmitBuffer(const uint8_t* data, size_t size) {
159   if (compressor_) {
160     compressor_->Write(data, size);
161   } else {
162     patch_->insert(patch_->end(), data, data + size);
163   }
164 }
165 
166 void EndsleyPatchWriter::Flush() {
167   size_t used_diff = 0;
168   size_t used_extra = 0;
169   size_t used_control = 0;
170   do {
171     if (!pending_diff_ && !pending_extra_ && used_control < control_.size()) {
172       // We can emit a control entry in these conditions.
173       const ControlEntry& entry = control_[used_control];
174       used_control++;
175 
176       pending_diff_ = entry.diff_size;
177       pending_extra_ = entry.extra_size;
178       pending_control_data_ -= entry.extra_size + entry.diff_size;
179       EmitControlEntry(entry);
180     }
181 
182     if (pending_diff_) {
183       size_t diff_size = std::min(diff_data_.size() - used_diff, pending_diff_);
184       EmitBuffer(diff_data_.data() + used_diff, diff_size);
185       pending_diff_ -= diff_size;
186       used_diff += diff_size;
187     }
188 
189     if (!pending_diff_ && pending_extra_) {
190       size_t extra_size =
191           std::min(extra_data_.size() - used_extra, pending_extra_);
192       EmitBuffer(extra_data_.data() + used_extra, extra_size);
193       pending_extra_ -= extra_size;
194       used_extra += extra_size;
195     }
196   } while (!pending_diff_ && !pending_extra_ && used_control < control_.size());
197 
198   if (used_diff)
199     diff_data_.erase(diff_data_.begin(), diff_data_.begin() + used_diff);
200   if (used_extra)
201     extra_data_.erase(extra_data_.begin(), extra_data_.begin() + used_extra);
202   if (used_control)
203     control_.erase(control_.begin(), control_.begin() + used_control);
204 }
205 
206 }  // namespace bsdiff
207