blob: e042e03cb86caae0c3a61d3774f1f6b41a80abad [file] [log] [blame]
Doug Zongker424296a2014-09-02 08:53:09 -07001# Copyright (C) 2014 The Android Open Source Project
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
Doug Zongkerfc44a512014-08-26 13:10:25 -070015from __future__ import print_function
16
17from collections import deque, OrderedDict
18from hashlib import sha1
Tao Bao8dcf7382015-05-21 14:09:49 -070019import common
Doug Zongker62338182014-09-08 08:29:55 -070020import heapq
Doug Zongkerfc44a512014-08-26 13:10:25 -070021import itertools
22import multiprocessing
23import os
Doug Zongkerfc44a512014-08-26 13:10:25 -070024import re
25import subprocess
Doug Zongkerfc44a512014-08-26 13:10:25 -070026import threading
27import tempfile
28
Dan Albert8b72aef2015-03-23 19:13:21 -070029from rangelib import RangeSet
30
Doug Zongkerfc44a512014-08-26 13:10:25 -070031
Doug Zongkerab7ca1d2014-08-26 10:40:28 -070032__all__ = ["EmptyImage", "DataImage", "BlockImageDiff"]
33
Dan Albert8b72aef2015-03-23 19:13:21 -070034
Doug Zongkerfc44a512014-08-26 13:10:25 -070035def compute_patch(src, tgt, imgdiff=False):
36 srcfd, srcfile = tempfile.mkstemp(prefix="src-")
37 tgtfd, tgtfile = tempfile.mkstemp(prefix="tgt-")
38 patchfd, patchfile = tempfile.mkstemp(prefix="patch-")
39 os.close(patchfd)
40
41 try:
42 with os.fdopen(srcfd, "wb") as f_src:
43 for p in src:
44 f_src.write(p)
45
46 with os.fdopen(tgtfd, "wb") as f_tgt:
47 for p in tgt:
48 f_tgt.write(p)
49 try:
50 os.unlink(patchfile)
51 except OSError:
52 pass
53 if imgdiff:
54 p = subprocess.call(["imgdiff", "-z", srcfile, tgtfile, patchfile],
55 stdout=open("/dev/null", "a"),
56 stderr=subprocess.STDOUT)
57 else:
58 p = subprocess.call(["bsdiff", srcfile, tgtfile, patchfile])
59
60 if p:
61 raise ValueError("diff failed: " + str(p))
62
63 with open(patchfile, "rb") as f:
64 return f.read()
65 finally:
66 try:
67 os.unlink(srcfile)
68 os.unlink(tgtfile)
69 os.unlink(patchfile)
70 except OSError:
71 pass
72
Dan Albert8b72aef2015-03-23 19:13:21 -070073
74class Image(object):
75 def ReadRangeSet(self, ranges):
76 raise NotImplementedError
77
Tao Bao68658c02015-06-01 13:40:49 -070078 def TotalSha1(self, include_clobbered_blocks=False):
Dan Albert8b72aef2015-03-23 19:13:21 -070079 raise NotImplementedError
80
81
82class EmptyImage(Image):
Doug Zongkerfc44a512014-08-26 13:10:25 -070083 """A zero-length image."""
84 blocksize = 4096
85 care_map = RangeSet()
Tao Baoff777812015-05-12 11:42:31 -070086 clobbered_blocks = RangeSet()
Tao Baoe9b61912015-07-09 17:37:49 -070087 extended = RangeSet()
Doug Zongkerfc44a512014-08-26 13:10:25 -070088 total_blocks = 0
89 file_map = {}
90 def ReadRangeSet(self, ranges):
91 return ()
Tao Bao68658c02015-06-01 13:40:49 -070092 def TotalSha1(self, include_clobbered_blocks=False):
93 # EmptyImage always carries empty clobbered_blocks, so
94 # include_clobbered_blocks can be ignored.
95 assert self.clobbered_blocks.size() == 0
Doug Zongkerab7ca1d2014-08-26 10:40:28 -070096 return sha1().hexdigest()
97
98
Dan Albert8b72aef2015-03-23 19:13:21 -070099class DataImage(Image):
Doug Zongkerab7ca1d2014-08-26 10:40:28 -0700100 """An image wrapped around a single string of data."""
101
102 def __init__(self, data, trim=False, pad=False):
103 self.data = data
104 self.blocksize = 4096
105
106 assert not (trim and pad)
107
108 partial = len(self.data) % self.blocksize
109 if partial > 0:
110 if trim:
111 self.data = self.data[:-partial]
112 elif pad:
113 self.data += '\0' * (self.blocksize - partial)
114 else:
115 raise ValueError(("data for DataImage must be multiple of %d bytes "
116 "unless trim or pad is specified") %
117 (self.blocksize,))
118
119 assert len(self.data) % self.blocksize == 0
120
121 self.total_blocks = len(self.data) / self.blocksize
122 self.care_map = RangeSet(data=(0, self.total_blocks))
Tao Baoff777812015-05-12 11:42:31 -0700123 self.clobbered_blocks = RangeSet()
Tao Baoe9b61912015-07-09 17:37:49 -0700124 self.extended = RangeSet()
Doug Zongkerab7ca1d2014-08-26 10:40:28 -0700125
126 zero_blocks = []
127 nonzero_blocks = []
128 reference = '\0' * self.blocksize
129
130 for i in range(self.total_blocks):
131 d = self.data[i*self.blocksize : (i+1)*self.blocksize]
132 if d == reference:
133 zero_blocks.append(i)
134 zero_blocks.append(i+1)
135 else:
136 nonzero_blocks.append(i)
137 nonzero_blocks.append(i+1)
138
139 self.file_map = {"__ZERO": RangeSet(zero_blocks),
140 "__NONZERO": RangeSet(nonzero_blocks)}
141
142 def ReadRangeSet(self, ranges):
143 return [self.data[s*self.blocksize:e*self.blocksize] for (s, e) in ranges]
144
Tao Bao68658c02015-06-01 13:40:49 -0700145 def TotalSha1(self, include_clobbered_blocks=False):
146 # DataImage always carries empty clobbered_blocks, so
147 # include_clobbered_blocks can be ignored.
Tao Baoff777812015-05-12 11:42:31 -0700148 assert self.clobbered_blocks.size() == 0
Dan Albert8b72aef2015-03-23 19:13:21 -0700149 return sha1(self.data).hexdigest()
Doug Zongkerab7ca1d2014-08-26 10:40:28 -0700150
Doug Zongkerfc44a512014-08-26 13:10:25 -0700151
152class Transfer(object):
153 def __init__(self, tgt_name, src_name, tgt_ranges, src_ranges, style, by_id):
154 self.tgt_name = tgt_name
155 self.src_name = src_name
156 self.tgt_ranges = tgt_ranges
157 self.src_ranges = src_ranges
158 self.style = style
159 self.intact = (getattr(tgt_ranges, "monotonic", False) and
160 getattr(src_ranges, "monotonic", False))
Tao Baob8c87172015-03-19 19:42:12 -0700161
162 # We use OrderedDict rather than dict so that the output is repeatable;
163 # otherwise it would depend on the hash values of the Transfer objects.
164 self.goes_before = OrderedDict()
165 self.goes_after = OrderedDict()
Doug Zongkerfc44a512014-08-26 13:10:25 -0700166
Doug Zongker62338182014-09-08 08:29:55 -0700167 self.stash_before = []
168 self.use_stash = []
169
Doug Zongkerfc44a512014-08-26 13:10:25 -0700170 self.id = len(by_id)
171 by_id.append(self)
172
Doug Zongker62338182014-09-08 08:29:55 -0700173 def NetStashChange(self):
174 return (sum(sr.size() for (_, sr) in self.stash_before) -
175 sum(sr.size() for (_, sr) in self.use_stash))
176
Doug Zongkerfc44a512014-08-26 13:10:25 -0700177 def __str__(self):
178 return (str(self.id) + ": <" + str(self.src_ranges) + " " + self.style +
179 " to " + str(self.tgt_ranges) + ">")
180
181
182# BlockImageDiff works on two image objects. An image object is
183# anything that provides the following attributes:
184#
185# blocksize: the size in bytes of a block, currently must be 4096.
186#
187# total_blocks: the total size of the partition/image, in blocks.
188#
189# care_map: a RangeSet containing which blocks (in the range [0,
190# total_blocks) we actually care about; i.e. which blocks contain
191# data.
192#
193# file_map: a dict that partitions the blocks contained in care_map
194# into smaller domains that are useful for doing diffs on.
195# (Typically a domain is a file, and the key in file_map is the
196# pathname.)
197#
Tao Baoff777812015-05-12 11:42:31 -0700198# clobbered_blocks: a RangeSet containing which blocks contain data
199# but may be altered by the FS. They need to be excluded when
200# verifying the partition integrity.
201#
Doug Zongkerfc44a512014-08-26 13:10:25 -0700202# ReadRangeSet(): a function that takes a RangeSet and returns the
203# data contained in the image blocks of that RangeSet. The data
204# is returned as a list or tuple of strings; concatenating the
205# elements together should produce the requested data.
206# Implementations are free to break up the data into list/tuple
207# elements in any way that is convenient.
208#
Doug Zongkerab7ca1d2014-08-26 10:40:28 -0700209# TotalSha1(): a function that returns (as a hex string) the SHA-1
210# hash of all the data in the image (ie, all the blocks in the
Tao Bao68658c02015-06-01 13:40:49 -0700211# care_map minus clobbered_blocks, or including the clobbered
212# blocks if include_clobbered_blocks is True).
Doug Zongkerab7ca1d2014-08-26 10:40:28 -0700213#
Doug Zongkerfc44a512014-08-26 13:10:25 -0700214# When creating a BlockImageDiff, the src image may be None, in which
215# case the list of transfers produced will never read from the
216# original image.
217
218class BlockImageDiff(object):
Sami Tolvanendd67a292014-12-09 16:40:34 +0000219 def __init__(self, tgt, src=None, threads=None, version=3):
Doug Zongkerfc44a512014-08-26 13:10:25 -0700220 if threads is None:
221 threads = multiprocessing.cpu_count() // 2
Dan Albert8b72aef2015-03-23 19:13:21 -0700222 if threads == 0:
223 threads = 1
Doug Zongkerfc44a512014-08-26 13:10:25 -0700224 self.threads = threads
Doug Zongker62338182014-09-08 08:29:55 -0700225 self.version = version
Dan Albert8b72aef2015-03-23 19:13:21 -0700226 self.transfers = []
227 self.src_basenames = {}
228 self.src_numpatterns = {}
Doug Zongker62338182014-09-08 08:29:55 -0700229
Sami Tolvanendd67a292014-12-09 16:40:34 +0000230 assert version in (1, 2, 3)
Doug Zongkerfc44a512014-08-26 13:10:25 -0700231
232 self.tgt = tgt
233 if src is None:
234 src = EmptyImage()
235 self.src = src
236
237 # The updater code that installs the patch always uses 4k blocks.
238 assert tgt.blocksize == 4096
239 assert src.blocksize == 4096
240
241 # The range sets in each filemap should comprise a partition of
242 # the care map.
243 self.AssertPartition(src.care_map, src.file_map.values())
244 self.AssertPartition(tgt.care_map, tgt.file_map.values())
245
246 def Compute(self, prefix):
247 # When looking for a source file to use as the diff input for a
248 # target file, we try:
249 # 1) an exact path match if available, otherwise
250 # 2) a exact basename match if available, otherwise
251 # 3) a basename match after all runs of digits are replaced by
252 # "#" if available, otherwise
253 # 4) we have no source for this target.
254 self.AbbreviateSourceNames()
255 self.FindTransfers()
256
257 # Find the ordering dependencies among transfers (this is O(n^2)
258 # in the number of transfers).
259 self.GenerateDigraph()
260 # Find a sequence of transfers that satisfies as many ordering
261 # dependencies as possible (heuristically).
262 self.FindVertexSequence()
263 # Fix up the ordering dependencies that the sequence didn't
264 # satisfy.
Doug Zongker62338182014-09-08 08:29:55 -0700265 if self.version == 1:
266 self.RemoveBackwardEdges()
267 else:
268 self.ReverseBackwardEdges()
269 self.ImproveVertexSequence()
270
Doug Zongkerfc44a512014-08-26 13:10:25 -0700271 # Double-check our work.
272 self.AssertSequenceGood()
273
274 self.ComputePatches(prefix)
275 self.WriteTransfers(prefix)
276
Dan Albert8b72aef2015-03-23 19:13:21 -0700277 def HashBlocks(self, source, ranges): # pylint: disable=no-self-use
Sami Tolvanendd67a292014-12-09 16:40:34 +0000278 data = source.ReadRangeSet(ranges)
279 ctx = sha1()
280
281 for p in data:
282 ctx.update(p)
283
284 return ctx.hexdigest()
285
Doug Zongkerfc44a512014-08-26 13:10:25 -0700286 def WriteTransfers(self, prefix):
287 out = []
288
Doug Zongkerfc44a512014-08-26 13:10:25 -0700289 total = 0
290 performs_read = False
291
Doug Zongker62338182014-09-08 08:29:55 -0700292 stashes = {}
293 stashed_blocks = 0
294 max_stashed_blocks = 0
295
296 free_stash_ids = []
297 next_stash_id = 0
298
Doug Zongkerfc44a512014-08-26 13:10:25 -0700299 for xf in self.transfers:
300
Doug Zongker62338182014-09-08 08:29:55 -0700301 if self.version < 2:
302 assert not xf.stash_before
303 assert not xf.use_stash
304
305 for s, sr in xf.stash_before:
306 assert s not in stashes
307 if free_stash_ids:
308 sid = heapq.heappop(free_stash_ids)
309 else:
310 sid = next_stash_id
311 next_stash_id += 1
312 stashes[s] = sid
313 stashed_blocks += sr.size()
Sami Tolvanendd67a292014-12-09 16:40:34 +0000314 if self.version == 2:
315 out.append("stash %d %s\n" % (sid, sr.to_string_raw()))
316 else:
317 sh = self.HashBlocks(self.src, sr)
318 if sh in stashes:
319 stashes[sh] += 1
320 else:
321 stashes[sh] = 1
322 out.append("stash %s %s\n" % (sh, sr.to_string_raw()))
Doug Zongker62338182014-09-08 08:29:55 -0700323
324 if stashed_blocks > max_stashed_blocks:
325 max_stashed_blocks = stashed_blocks
326
Jesse Zhao7b985f62015-03-02 16:53:08 -0800327 free_string = []
328
Doug Zongker62338182014-09-08 08:29:55 -0700329 if self.version == 1:
Dan Albert8b72aef2015-03-23 19:13:21 -0700330 src_str = xf.src_ranges.to_string_raw()
Sami Tolvanendd67a292014-12-09 16:40:34 +0000331 elif self.version >= 2:
Doug Zongker62338182014-09-08 08:29:55 -0700332
333 # <# blocks> <src ranges>
334 # OR
335 # <# blocks> <src ranges> <src locs> <stash refs...>
336 # OR
337 # <# blocks> - <stash refs...>
338
339 size = xf.src_ranges.size()
Dan Albert8b72aef2015-03-23 19:13:21 -0700340 src_str = [str(size)]
Doug Zongker62338182014-09-08 08:29:55 -0700341
342 unstashed_src_ranges = xf.src_ranges
343 mapped_stashes = []
344 for s, sr in xf.use_stash:
345 sid = stashes.pop(s)
346 stashed_blocks -= sr.size()
347 unstashed_src_ranges = unstashed_src_ranges.subtract(sr)
Sami Tolvanendd67a292014-12-09 16:40:34 +0000348 sh = self.HashBlocks(self.src, sr)
Doug Zongker62338182014-09-08 08:29:55 -0700349 sr = xf.src_ranges.map_within(sr)
350 mapped_stashes.append(sr)
Sami Tolvanendd67a292014-12-09 16:40:34 +0000351 if self.version == 2:
Dan Albert8b72aef2015-03-23 19:13:21 -0700352 src_str.append("%d:%s" % (sid, sr.to_string_raw()))
Tao Baobb625d22015-08-13 14:44:15 -0700353 # A stash will be used only once. We need to free the stash
354 # immediately after the use, instead of waiting for the automatic
355 # clean-up at the end. Because otherwise it may take up extra space
356 # and lead to OTA failures.
357 # Bug: 23119955
358 free_string.append("free %d\n" % (sid,))
Sami Tolvanendd67a292014-12-09 16:40:34 +0000359 else:
360 assert sh in stashes
Dan Albert8b72aef2015-03-23 19:13:21 -0700361 src_str.append("%s:%s" % (sh, sr.to_string_raw()))
Sami Tolvanendd67a292014-12-09 16:40:34 +0000362 stashes[sh] -= 1
363 if stashes[sh] == 0:
364 free_string.append("free %s\n" % (sh))
365 stashes.pop(sh)
Doug Zongker62338182014-09-08 08:29:55 -0700366 heapq.heappush(free_stash_ids, sid)
367
368 if unstashed_src_ranges:
Dan Albert8b72aef2015-03-23 19:13:21 -0700369 src_str.insert(1, unstashed_src_ranges.to_string_raw())
Doug Zongker62338182014-09-08 08:29:55 -0700370 if xf.use_stash:
371 mapped_unstashed = xf.src_ranges.map_within(unstashed_src_ranges)
Dan Albert8b72aef2015-03-23 19:13:21 -0700372 src_str.insert(2, mapped_unstashed.to_string_raw())
Doug Zongker62338182014-09-08 08:29:55 -0700373 mapped_stashes.append(mapped_unstashed)
374 self.AssertPartition(RangeSet(data=(0, size)), mapped_stashes)
375 else:
Dan Albert8b72aef2015-03-23 19:13:21 -0700376 src_str.insert(1, "-")
Doug Zongker62338182014-09-08 08:29:55 -0700377 self.AssertPartition(RangeSet(data=(0, size)), mapped_stashes)
378
Dan Albert8b72aef2015-03-23 19:13:21 -0700379 src_str = " ".join(src_str)
Doug Zongker62338182014-09-08 08:29:55 -0700380
Sami Tolvanendd67a292014-12-09 16:40:34 +0000381 # all versions:
Doug Zongker62338182014-09-08 08:29:55 -0700382 # zero <rangeset>
383 # new <rangeset>
384 # erase <rangeset>
385 #
386 # version 1:
387 # bsdiff patchstart patchlen <src rangeset> <tgt rangeset>
388 # imgdiff patchstart patchlen <src rangeset> <tgt rangeset>
389 # move <src rangeset> <tgt rangeset>
390 #
391 # version 2:
Dan Albert8b72aef2015-03-23 19:13:21 -0700392 # bsdiff patchstart patchlen <tgt rangeset> <src_str>
393 # imgdiff patchstart patchlen <tgt rangeset> <src_str>
394 # move <tgt rangeset> <src_str>
Sami Tolvanendd67a292014-12-09 16:40:34 +0000395 #
396 # version 3:
Dan Albert8b72aef2015-03-23 19:13:21 -0700397 # bsdiff patchstart patchlen srchash tgthash <tgt rangeset> <src_str>
398 # imgdiff patchstart patchlen srchash tgthash <tgt rangeset> <src_str>
399 # move hash <tgt rangeset> <src_str>
Doug Zongkerfc44a512014-08-26 13:10:25 -0700400
401 tgt_size = xf.tgt_ranges.size()
402
403 if xf.style == "new":
404 assert xf.tgt_ranges
405 out.append("%s %s\n" % (xf.style, xf.tgt_ranges.to_string_raw()))
406 total += tgt_size
407 elif xf.style == "move":
408 performs_read = True
409 assert xf.tgt_ranges
410 assert xf.src_ranges.size() == tgt_size
411 if xf.src_ranges != xf.tgt_ranges:
Doug Zongker62338182014-09-08 08:29:55 -0700412 if self.version == 1:
413 out.append("%s %s %s\n" % (
414 xf.style,
415 xf.src_ranges.to_string_raw(), xf.tgt_ranges.to_string_raw()))
416 elif self.version == 2:
417 out.append("%s %s %s\n" % (
418 xf.style,
Dan Albert8b72aef2015-03-23 19:13:21 -0700419 xf.tgt_ranges.to_string_raw(), src_str))
Sami Tolvanendd67a292014-12-09 16:40:34 +0000420 elif self.version >= 3:
Sami Tolvanen29f529f2015-04-17 16:28:08 +0100421 # take into account automatic stashing of overlapping blocks
422 if xf.src_ranges.overlaps(xf.tgt_ranges):
Tao Baoe9b61912015-07-09 17:37:49 -0700423 temp_stash_usage = stashed_blocks + xf.src_ranges.size()
Sami Tolvanen29f529f2015-04-17 16:28:08 +0100424 if temp_stash_usage > max_stashed_blocks:
425 max_stashed_blocks = temp_stash_usage
426
Sami Tolvanendd67a292014-12-09 16:40:34 +0000427 out.append("%s %s %s %s\n" % (
428 xf.style,
429 self.HashBlocks(self.tgt, xf.tgt_ranges),
Dan Albert8b72aef2015-03-23 19:13:21 -0700430 xf.tgt_ranges.to_string_raw(), src_str))
Doug Zongkerfc44a512014-08-26 13:10:25 -0700431 total += tgt_size
432 elif xf.style in ("bsdiff", "imgdiff"):
433 performs_read = True
434 assert xf.tgt_ranges
435 assert xf.src_ranges
Doug Zongker62338182014-09-08 08:29:55 -0700436 if self.version == 1:
437 out.append("%s %d %d %s %s\n" % (
438 xf.style, xf.patch_start, xf.patch_len,
439 xf.src_ranges.to_string_raw(), xf.tgt_ranges.to_string_raw()))
440 elif self.version == 2:
441 out.append("%s %d %d %s %s\n" % (
442 xf.style, xf.patch_start, xf.patch_len,
Dan Albert8b72aef2015-03-23 19:13:21 -0700443 xf.tgt_ranges.to_string_raw(), src_str))
Sami Tolvanendd67a292014-12-09 16:40:34 +0000444 elif self.version >= 3:
Sami Tolvanen29f529f2015-04-17 16:28:08 +0100445 # take into account automatic stashing of overlapping blocks
446 if xf.src_ranges.overlaps(xf.tgt_ranges):
Tao Baoe9b61912015-07-09 17:37:49 -0700447 temp_stash_usage = stashed_blocks + xf.src_ranges.size()
Sami Tolvanen29f529f2015-04-17 16:28:08 +0100448 if temp_stash_usage > max_stashed_blocks:
449 max_stashed_blocks = temp_stash_usage
450
Sami Tolvanendd67a292014-12-09 16:40:34 +0000451 out.append("%s %d %d %s %s %s %s\n" % (
452 xf.style,
453 xf.patch_start, xf.patch_len,
454 self.HashBlocks(self.src, xf.src_ranges),
455 self.HashBlocks(self.tgt, xf.tgt_ranges),
Dan Albert8b72aef2015-03-23 19:13:21 -0700456 xf.tgt_ranges.to_string_raw(), src_str))
Doug Zongkerfc44a512014-08-26 13:10:25 -0700457 total += tgt_size
458 elif xf.style == "zero":
459 assert xf.tgt_ranges
460 to_zero = xf.tgt_ranges.subtract(xf.src_ranges)
461 if to_zero:
462 out.append("%s %s\n" % (xf.style, to_zero.to_string_raw()))
463 total += to_zero.size()
464 else:
Dan Albert8b72aef2015-03-23 19:13:21 -0700465 raise ValueError("unknown transfer style '%s'\n" % xf.style)
Doug Zongkerfc44a512014-08-26 13:10:25 -0700466
Sami Tolvanendd67a292014-12-09 16:40:34 +0000467 if free_string:
468 out.append("".join(free_string))
469
Tao Bao575d68a2015-08-07 19:49:45 -0700470 if self.version >= 2 and common.OPTIONS.cache_size is not None:
Tao Bao8dcf7382015-05-21 14:09:49 -0700471 # Sanity check: abort if we're going to need more stash space than
472 # the allowed size (cache_size * threshold). There are two purposes
473 # of having a threshold here. a) Part of the cache may have been
474 # occupied by some recovery logs. b) It will buy us some time to deal
475 # with the oversize issue.
476 cache_size = common.OPTIONS.cache_size
477 stash_threshold = common.OPTIONS.stash_threshold
478 max_allowed = cache_size * stash_threshold
479 assert max_stashed_blocks * self.tgt.blocksize < max_allowed, \
480 'Stash size %d (%d * %d) exceeds the limit %d (%d * %.2f)' % (
481 max_stashed_blocks * self.tgt.blocksize, max_stashed_blocks,
482 self.tgt.blocksize, max_allowed, cache_size,
483 stash_threshold)
Doug Zongkerfc44a512014-08-26 13:10:25 -0700484
Tao Baoe9b61912015-07-09 17:37:49 -0700485 # Zero out extended blocks as a workaround for bug 20881595.
486 if self.tgt.extended:
487 out.append("zero %s\n" % (self.tgt.extended.to_string_raw(),))
488
489 # We erase all the blocks on the partition that a) don't contain useful
490 # data in the new image and b) will not be touched by dm-verity.
Doug Zongkerfc44a512014-08-26 13:10:25 -0700491 all_tgt = RangeSet(data=(0, self.tgt.total_blocks))
Tao Baoe9b61912015-07-09 17:37:49 -0700492 all_tgt_minus_extended = all_tgt.subtract(self.tgt.extended)
493 new_dontcare = all_tgt_minus_extended.subtract(self.tgt.care_map)
494 if new_dontcare:
495 out.append("erase %s\n" % (new_dontcare.to_string_raw(),))
Doug Zongkere985f6f2014-09-09 12:38:47 -0700496
497 out.insert(0, "%d\n" % (self.version,)) # format version number
498 out.insert(1, str(total) + "\n")
499 if self.version >= 2:
500 # version 2 only: after the total block count, we give the number
501 # of stash slots needed, and the maximum size needed (in blocks)
502 out.insert(2, str(next_stash_id) + "\n")
503 out.insert(3, str(max_stashed_blocks) + "\n")
Doug Zongkerfc44a512014-08-26 13:10:25 -0700504
505 with open(prefix + ".transfer.list", "wb") as f:
506 for i in out:
507 f.write(i)
508
Doug Zongker62338182014-09-08 08:29:55 -0700509 if self.version >= 2:
Tao Bao8dcf7382015-05-21 14:09:49 -0700510 max_stashed_size = max_stashed_blocks * self.tgt.blocksize
Tao Bao575d68a2015-08-07 19:49:45 -0700511 OPTIONS = common.OPTIONS
512 if OPTIONS.cache_size is not None:
513 max_allowed = OPTIONS.cache_size * OPTIONS.stash_threshold
514 print("max stashed blocks: %d (%d bytes), "
515 "limit: %d bytes (%.2f%%)\n" % (
516 max_stashed_blocks, max_stashed_size, max_allowed,
517 max_stashed_size * 100.0 / max_allowed))
518 else:
519 print("max stashed blocks: %d (%d bytes), limit: <unknown>\n" % (
520 max_stashed_blocks, max_stashed_size))
Doug Zongker62338182014-09-08 08:29:55 -0700521
Doug Zongkerfc44a512014-08-26 13:10:25 -0700522 def ComputePatches(self, prefix):
523 print("Reticulating splines...")
524 diff_q = []
525 patch_num = 0
526 with open(prefix + ".new.dat", "wb") as new_f:
527 for xf in self.transfers:
528 if xf.style == "zero":
529 pass
530 elif xf.style == "new":
531 for piece in self.tgt.ReadRangeSet(xf.tgt_ranges):
532 new_f.write(piece)
533 elif xf.style == "diff":
534 src = self.src.ReadRangeSet(xf.src_ranges)
535 tgt = self.tgt.ReadRangeSet(xf.tgt_ranges)
536
537 # We can't compare src and tgt directly because they may have
538 # the same content but be broken up into blocks differently, eg:
539 #
540 # ["he", "llo"] vs ["h", "ello"]
541 #
542 # We want those to compare equal, ideally without having to
543 # actually concatenate the strings (these may be tens of
544 # megabytes).
545
546 src_sha1 = sha1()
547 for p in src:
548 src_sha1.update(p)
549 tgt_sha1 = sha1()
550 tgt_size = 0
551 for p in tgt:
552 tgt_sha1.update(p)
553 tgt_size += len(p)
554
555 if src_sha1.digest() == tgt_sha1.digest():
556 # These are identical; we don't need to generate a patch,
557 # just issue copy commands on the device.
558 xf.style = "move"
559 else:
560 # For files in zip format (eg, APKs, JARs, etc.) we would
561 # like to use imgdiff -z if possible (because it usually
562 # produces significantly smaller patches than bsdiff).
563 # This is permissible if:
564 #
565 # - the source and target files are monotonic (ie, the
566 # data is stored with blocks in increasing order), and
567 # - we haven't removed any blocks from the source set.
568 #
569 # If these conditions are satisfied then appending all the
570 # blocks in the set together in order will produce a valid
571 # zip file (plus possibly extra zeros in the last block),
572 # which is what imgdiff needs to operate. (imgdiff is
573 # fine with extra zeros at the end of the file.)
574 imgdiff = (xf.intact and
575 xf.tgt_name.split(".")[-1].lower()
576 in ("apk", "jar", "zip"))
577 xf.style = "imgdiff" if imgdiff else "bsdiff"
578 diff_q.append((tgt_size, src, tgt, xf, patch_num))
579 patch_num += 1
580
581 else:
582 assert False, "unknown style " + xf.style
583
584 if diff_q:
585 if self.threads > 1:
586 print("Computing patches (using %d threads)..." % (self.threads,))
587 else:
588 print("Computing patches...")
589 diff_q.sort()
590
591 patches = [None] * patch_num
592
Dan Albert8b72aef2015-03-23 19:13:21 -0700593 # TODO: Rewrite with multiprocessing.ThreadPool?
Doug Zongkerfc44a512014-08-26 13:10:25 -0700594 lock = threading.Lock()
595 def diff_worker():
596 while True:
597 with lock:
Dan Albert8b72aef2015-03-23 19:13:21 -0700598 if not diff_q:
599 return
Doug Zongkerfc44a512014-08-26 13:10:25 -0700600 tgt_size, src, tgt, xf, patchnum = diff_q.pop()
601 patch = compute_patch(src, tgt, imgdiff=(xf.style == "imgdiff"))
602 size = len(patch)
603 with lock:
604 patches[patchnum] = (patch, xf)
605 print("%10d %10d (%6.2f%%) %7s %s" % (
606 size, tgt_size, size * 100.0 / tgt_size, xf.style,
607 xf.tgt_name if xf.tgt_name == xf.src_name else (
608 xf.tgt_name + " (from " + xf.src_name + ")")))
609
610 threads = [threading.Thread(target=diff_worker)
Dan Albert8b72aef2015-03-23 19:13:21 -0700611 for _ in range(self.threads)]
Doug Zongkerfc44a512014-08-26 13:10:25 -0700612 for th in threads:
613 th.start()
614 while threads:
615 threads.pop().join()
616 else:
617 patches = []
618
619 p = 0
620 with open(prefix + ".patch.dat", "wb") as patch_f:
621 for patch, xf in patches:
622 xf.patch_start = p
623 xf.patch_len = len(patch)
624 patch_f.write(patch)
625 p += len(patch)
626
627 def AssertSequenceGood(self):
628 # Simulate the sequences of transfers we will output, and check that:
629 # - we never read a block after writing it, and
630 # - we write every block we care about exactly once.
631
632 # Start with no blocks having been touched yet.
633 touched = RangeSet()
634
635 # Imagine processing the transfers in order.
636 for xf in self.transfers:
637 # Check that the input blocks for this transfer haven't yet been touched.
Doug Zongker62338182014-09-08 08:29:55 -0700638
639 x = xf.src_ranges
640 if self.version >= 2:
641 for _, sr in xf.use_stash:
642 x = x.subtract(sr)
643
644 assert not touched.overlaps(x)
Doug Zongkerfc44a512014-08-26 13:10:25 -0700645 # Check that the output blocks for this transfer haven't yet been touched.
646 assert not touched.overlaps(xf.tgt_ranges)
647 # Touch all the blocks written by this transfer.
648 touched = touched.union(xf.tgt_ranges)
649
650 # Check that we've written every target block.
651 assert touched == self.tgt.care_map
652
Doug Zongker62338182014-09-08 08:29:55 -0700653 def ImproveVertexSequence(self):
654 print("Improving vertex order...")
655
656 # At this point our digraph is acyclic; we reversed any edges that
657 # were backwards in the heuristically-generated sequence. The
658 # previously-generated order is still acceptable, but we hope to
659 # find a better order that needs less memory for stashed data.
660 # Now we do a topological sort to generate a new vertex order,
661 # using a greedy algorithm to choose which vertex goes next
662 # whenever we have a choice.
663
664 # Make a copy of the edge set; this copy will get destroyed by the
665 # algorithm.
666 for xf in self.transfers:
667 xf.incoming = xf.goes_after.copy()
668 xf.outgoing = xf.goes_before.copy()
669
670 L = [] # the new vertex order
671
672 # S is the set of sources in the remaining graph; we always choose
673 # the one that leaves the least amount of stashed data after it's
674 # executed.
675 S = [(u.NetStashChange(), u.order, u) for u in self.transfers
676 if not u.incoming]
677 heapq.heapify(S)
678
679 while S:
680 _, _, xf = heapq.heappop(S)
681 L.append(xf)
682 for u in xf.outgoing:
683 del u.incoming[xf]
684 if not u.incoming:
685 heapq.heappush(S, (u.NetStashChange(), u.order, u))
686
687 # if this fails then our graph had a cycle.
688 assert len(L) == len(self.transfers)
689
690 self.transfers = L
691 for i, xf in enumerate(L):
692 xf.order = i
693
Doug Zongkerfc44a512014-08-26 13:10:25 -0700694 def RemoveBackwardEdges(self):
695 print("Removing backward edges...")
696 in_order = 0
697 out_of_order = 0
698 lost_source = 0
699
700 for xf in self.transfers:
Doug Zongkerfc44a512014-08-26 13:10:25 -0700701 lost = 0
702 size = xf.src_ranges.size()
703 for u in xf.goes_before:
704 # xf should go before u
705 if xf.order < u.order:
706 # it does, hurray!
Doug Zongker62338182014-09-08 08:29:55 -0700707 in_order += 1
Doug Zongkerfc44a512014-08-26 13:10:25 -0700708 else:
709 # it doesn't, boo. trim the blocks that u writes from xf's
710 # source, so that xf can go after u.
Doug Zongker62338182014-09-08 08:29:55 -0700711 out_of_order += 1
Doug Zongkerfc44a512014-08-26 13:10:25 -0700712 assert xf.src_ranges.overlaps(u.tgt_ranges)
713 xf.src_ranges = xf.src_ranges.subtract(u.tgt_ranges)
714 xf.intact = False
715
716 if xf.style == "diff" and not xf.src_ranges:
717 # nothing left to diff from; treat as new data
718 xf.style = "new"
719
720 lost = size - xf.src_ranges.size()
721 lost_source += lost
Doug Zongkerfc44a512014-08-26 13:10:25 -0700722
723 print((" %d/%d dependencies (%.2f%%) were violated; "
724 "%d source blocks removed.") %
725 (out_of_order, in_order + out_of_order,
726 (out_of_order * 100.0 / (in_order + out_of_order))
727 if (in_order + out_of_order) else 0.0,
728 lost_source))
729
Doug Zongker62338182014-09-08 08:29:55 -0700730 def ReverseBackwardEdges(self):
731 print("Reversing backward edges...")
732 in_order = 0
733 out_of_order = 0
734 stashes = 0
735 stash_size = 0
736
737 for xf in self.transfers:
Doug Zongker62338182014-09-08 08:29:55 -0700738 for u in xf.goes_before.copy():
739 # xf should go before u
740 if xf.order < u.order:
741 # it does, hurray!
742 in_order += 1
743 else:
744 # it doesn't, boo. modify u to stash the blocks that it
745 # writes that xf wants to read, and then require u to go
746 # before xf.
747 out_of_order += 1
748
749 overlap = xf.src_ranges.intersect(u.tgt_ranges)
750 assert overlap
751
752 u.stash_before.append((stashes, overlap))
753 xf.use_stash.append((stashes, overlap))
754 stashes += 1
755 stash_size += overlap.size()
756
757 # reverse the edge direction; now xf must go after u
758 del xf.goes_before[u]
759 del u.goes_after[xf]
760 xf.goes_after[u] = None # value doesn't matter
761 u.goes_before[xf] = None
762
763 print((" %d/%d dependencies (%.2f%%) were violated; "
764 "%d source blocks stashed.") %
765 (out_of_order, in_order + out_of_order,
766 (out_of_order * 100.0 / (in_order + out_of_order))
767 if (in_order + out_of_order) else 0.0,
768 stash_size))
769
Doug Zongkerfc44a512014-08-26 13:10:25 -0700770 def FindVertexSequence(self):
771 print("Finding vertex sequence...")
772
773 # This is based on "A Fast & Effective Heuristic for the Feedback
774 # Arc Set Problem" by P. Eades, X. Lin, and W.F. Smyth. Think of
775 # it as starting with the digraph G and moving all the vertices to
776 # be on a horizontal line in some order, trying to minimize the
777 # number of edges that end up pointing to the left. Left-pointing
778 # edges will get removed to turn the digraph into a DAG. In this
779 # case each edge has a weight which is the number of source blocks
780 # we'll lose if that edge is removed; we try to minimize the total
781 # weight rather than just the number of edges.
782
783 # Make a copy of the edge set; this copy will get destroyed by the
784 # algorithm.
785 for xf in self.transfers:
786 xf.incoming = xf.goes_after.copy()
787 xf.outgoing = xf.goes_before.copy()
788
789 # We use an OrderedDict instead of just a set so that the output
790 # is repeatable; otherwise it would depend on the hash values of
791 # the transfer objects.
792 G = OrderedDict()
793 for xf in self.transfers:
794 G[xf] = None
795 s1 = deque() # the left side of the sequence, built from left to right
796 s2 = deque() # the right side of the sequence, built from right to left
797
798 while G:
799
800 # Put all sinks at the end of the sequence.
801 while True:
802 sinks = [u for u in G if not u.outgoing]
Dan Albert8b72aef2015-03-23 19:13:21 -0700803 if not sinks:
804 break
Doug Zongkerfc44a512014-08-26 13:10:25 -0700805 for u in sinks:
806 s2.appendleft(u)
807 del G[u]
808 for iu in u.incoming:
809 del iu.outgoing[u]
810
811 # Put all the sources at the beginning of the sequence.
812 while True:
813 sources = [u for u in G if not u.incoming]
Dan Albert8b72aef2015-03-23 19:13:21 -0700814 if not sources:
815 break
Doug Zongkerfc44a512014-08-26 13:10:25 -0700816 for u in sources:
817 s1.append(u)
818 del G[u]
819 for iu in u.outgoing:
820 del iu.incoming[u]
821
Dan Albert8b72aef2015-03-23 19:13:21 -0700822 if not G:
823 break
Doug Zongkerfc44a512014-08-26 13:10:25 -0700824
825 # Find the "best" vertex to put next. "Best" is the one that
826 # maximizes the net difference in source blocks saved we get by
827 # pretending it's a source rather than a sink.
828
829 max_d = None
830 best_u = None
831 for u in G:
832 d = sum(u.outgoing.values()) - sum(u.incoming.values())
833 if best_u is None or d > max_d:
834 max_d = d
835 best_u = u
836
837 u = best_u
838 s1.append(u)
839 del G[u]
840 for iu in u.outgoing:
841 del iu.incoming[u]
842 for iu in u.incoming:
843 del iu.outgoing[u]
844
845 # Now record the sequence in the 'order' field of each transfer,
846 # and by rearranging self.transfers to be in the chosen sequence.
847
848 new_transfers = []
849 for x in itertools.chain(s1, s2):
850 x.order = len(new_transfers)
851 new_transfers.append(x)
852 del x.incoming
853 del x.outgoing
854
855 self.transfers = new_transfers
856
857 def GenerateDigraph(self):
858 print("Generating digraph...")
859 for a in self.transfers:
860 for b in self.transfers:
Dan Albert8b72aef2015-03-23 19:13:21 -0700861 if a is b:
862 continue
Doug Zongkerfc44a512014-08-26 13:10:25 -0700863
864 # If the blocks written by A are read by B, then B needs to go before A.
865 i = a.tgt_ranges.intersect(b.src_ranges)
866 if i:
Doug Zongkerab7ca1d2014-08-26 10:40:28 -0700867 if b.src_name == "__ZERO":
868 # the cost of removing source blocks for the __ZERO domain
869 # is (nearly) zero.
870 size = 0
871 else:
872 size = i.size()
Doug Zongkerfc44a512014-08-26 13:10:25 -0700873 b.goes_before[a] = size
874 a.goes_after[b] = size
875
876 def FindTransfers(self):
Doug Zongkerfc44a512014-08-26 13:10:25 -0700877 empty = RangeSet()
878 for tgt_fn, tgt_ranges in self.tgt.file_map.items():
879 if tgt_fn == "__ZERO":
880 # the special "__ZERO" domain is all the blocks not contained
881 # in any file and that are filled with zeros. We have a
882 # special transfer style for zero blocks.
883 src_ranges = self.src.file_map.get("__ZERO", empty)
Doug Zongkerab7ca1d2014-08-26 10:40:28 -0700884 Transfer(tgt_fn, "__ZERO", tgt_ranges, src_ranges,
885 "zero", self.transfers)
Doug Zongkerfc44a512014-08-26 13:10:25 -0700886 continue
887
Tao Baoff777812015-05-12 11:42:31 -0700888 elif tgt_fn == "__COPY":
889 # "__COPY" domain includes all the blocks not contained in any
890 # file and that need to be copied unconditionally to the target.
891 Transfer(tgt_fn, None, tgt_ranges, empty, "new", self.transfers)
892 continue
893
Doug Zongkerfc44a512014-08-26 13:10:25 -0700894 elif tgt_fn in self.src.file_map:
895 # Look for an exact pathname match in the source.
896 Transfer(tgt_fn, tgt_fn, tgt_ranges, self.src.file_map[tgt_fn],
897 "diff", self.transfers)
898 continue
899
900 b = os.path.basename(tgt_fn)
901 if b in self.src_basenames:
902 # Look for an exact basename match in the source.
903 src_fn = self.src_basenames[b]
904 Transfer(tgt_fn, src_fn, tgt_ranges, self.src.file_map[src_fn],
905 "diff", self.transfers)
906 continue
907
908 b = re.sub("[0-9]+", "#", b)
909 if b in self.src_numpatterns:
910 # Look for a 'number pattern' match (a basename match after
911 # all runs of digits are replaced by "#"). (This is useful
912 # for .so files that contain version numbers in the filename
913 # that get bumped.)
914 src_fn = self.src_numpatterns[b]
915 Transfer(tgt_fn, src_fn, tgt_ranges, self.src.file_map[src_fn],
916 "diff", self.transfers)
917 continue
918
919 Transfer(tgt_fn, None, tgt_ranges, empty, "new", self.transfers)
920
921 def AbbreviateSourceNames(self):
Doug Zongkerfc44a512014-08-26 13:10:25 -0700922 for k in self.src.file_map.keys():
923 b = os.path.basename(k)
924 self.src_basenames[b] = k
925 b = re.sub("[0-9]+", "#", b)
926 self.src_numpatterns[b] = k
927
928 @staticmethod
929 def AssertPartition(total, seq):
930 """Assert that all the RangeSets in 'seq' form a partition of the
931 'total' RangeSet (ie, they are nonintersecting and their union
932 equals 'total')."""
933 so_far = RangeSet()
934 for i in seq:
935 assert not so_far.overlaps(i)
936 so_far = so_far.union(i)
937 assert so_far == total