blockimgdiff.py revision d522bdc9edbf64d15a59c6924853b2e2c8c39e90
1# 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 15from __future__ import print_function 16 17from collections import deque, OrderedDict 18from hashlib import sha1 19import array 20import common 21import functools 22import heapq 23import itertools 24import multiprocessing 25import os 26import re 27import subprocess 28import threading 29import time 30import tempfile 31 32from rangelib import RangeSet 33 34 35__all__ = ["EmptyImage", "DataImage", "BlockImageDiff"] 36 37 38def compute_patch(src, tgt, imgdiff=False): 39 srcfd, srcfile = tempfile.mkstemp(prefix="src-") 40 tgtfd, tgtfile = tempfile.mkstemp(prefix="tgt-") 41 patchfd, patchfile = tempfile.mkstemp(prefix="patch-") 42 os.close(patchfd) 43 44 try: 45 with os.fdopen(srcfd, "wb") as f_src: 46 for p in src: 47 f_src.write(p) 48 49 with os.fdopen(tgtfd, "wb") as f_tgt: 50 for p in tgt: 51 f_tgt.write(p) 52 try: 53 os.unlink(patchfile) 54 except OSError: 55 pass 56 if imgdiff: 57 p = subprocess.call(["imgdiff", "-z", srcfile, tgtfile, patchfile], 58 stdout=open("/dev/null", "a"), 59 stderr=subprocess.STDOUT) 60 else: 61 p = subprocess.call(["bsdiff", srcfile, tgtfile, patchfile]) 62 63 if p: 64 raise ValueError("diff failed: " + str(p)) 65 66 with open(patchfile, "rb") as f: 67 return f.read() 68 finally: 69 try: 70 os.unlink(srcfile) 71 os.unlink(tgtfile) 72 os.unlink(patchfile) 73 except OSError: 74 pass 75 76 77class Image(object): 78 def ReadRangeSet(self, ranges): 79 raise NotImplementedError 80 81 def TotalSha1(self, include_clobbered_blocks=False): 82 raise NotImplementedError 83 84 85class EmptyImage(Image): 86 """A zero-length image.""" 87 blocksize = 4096 88 care_map = RangeSet() 89 clobbered_blocks = RangeSet() 90 extended = RangeSet() 91 total_blocks = 0 92 file_map = {} 93 def ReadRangeSet(self, ranges): 94 return () 95 def TotalSha1(self, include_clobbered_blocks=False): 96 # EmptyImage always carries empty clobbered_blocks, so 97 # include_clobbered_blocks can be ignored. 98 assert self.clobbered_blocks.size() == 0 99 return sha1().hexdigest() 100 101 102class DataImage(Image): 103 """An image wrapped around a single string of data.""" 104 105 def __init__(self, data, trim=False, pad=False): 106 self.data = data 107 self.blocksize = 4096 108 109 assert not (trim and pad) 110 111 partial = len(self.data) % self.blocksize 112 padded = False 113 if partial > 0: 114 if trim: 115 self.data = self.data[:-partial] 116 elif pad: 117 self.data += '\0' * (self.blocksize - partial) 118 padded = True 119 else: 120 raise ValueError(("data for DataImage must be multiple of %d bytes " 121 "unless trim or pad is specified") % 122 (self.blocksize,)) 123 124 assert len(self.data) % self.blocksize == 0 125 126 self.total_blocks = len(self.data) / self.blocksize 127 self.care_map = RangeSet(data=(0, self.total_blocks)) 128 # When the last block is padded, we always write the whole block even for 129 # incremental OTAs. Because otherwise the last block may get skipped if 130 # unchanged for an incremental, but would fail the post-install 131 # verification if it has non-zero contents in the padding bytes. 132 # Bug: 23828506 133 if padded: 134 clobbered_blocks = [self.total_blocks-1, self.total_blocks] 135 else: 136 clobbered_blocks = [] 137 self.clobbered_blocks = clobbered_blocks 138 self.extended = RangeSet() 139 140 zero_blocks = [] 141 nonzero_blocks = [] 142 reference = '\0' * self.blocksize 143 144 for i in range(self.total_blocks-1 if padded else self.total_blocks): 145 d = self.data[i*self.blocksize : (i+1)*self.blocksize] 146 if d == reference: 147 zero_blocks.append(i) 148 zero_blocks.append(i+1) 149 else: 150 nonzero_blocks.append(i) 151 nonzero_blocks.append(i+1) 152 153 assert zero_blocks or nonzero_blocks or clobbered_blocks 154 155 self.file_map = dict() 156 if zero_blocks: 157 self.file_map["__ZERO"] = RangeSet(data=zero_blocks) 158 if nonzero_blocks: 159 self.file_map["__NONZERO"] = RangeSet(data=nonzero_blocks) 160 if clobbered_blocks: 161 self.file_map["__COPY"] = RangeSet(data=clobbered_blocks) 162 163 def ReadRangeSet(self, ranges): 164 return [self.data[s*self.blocksize:e*self.blocksize] for (s, e) in ranges] 165 166 def TotalSha1(self, include_clobbered_blocks=False): 167 if not include_clobbered_blocks: 168 ranges = self.care_map.subtract(self.clobbered_blocks) 169 return sha1(self.ReadRangeSet(ranges)).hexdigest() 170 else: 171 return sha1(self.data).hexdigest() 172 173 174class Transfer(object): 175 def __init__(self, tgt_name, src_name, tgt_ranges, src_ranges, style, by_id): 176 self.tgt_name = tgt_name 177 self.src_name = src_name 178 self.tgt_ranges = tgt_ranges 179 self.src_ranges = src_ranges 180 self.style = style 181 self.intact = (getattr(tgt_ranges, "monotonic", False) and 182 getattr(src_ranges, "monotonic", False)) 183 184 # We use OrderedDict rather than dict so that the output is repeatable; 185 # otherwise it would depend on the hash values of the Transfer objects. 186 self.goes_before = OrderedDict() 187 self.goes_after = OrderedDict() 188 189 self.stash_before = [] 190 self.use_stash = [] 191 192 self.id = len(by_id) 193 by_id.append(self) 194 195 def NetStashChange(self): 196 return (sum(sr.size() for (_, sr) in self.stash_before) - 197 sum(sr.size() for (_, sr) in self.use_stash)) 198 199 def ConvertToNew(self): 200 assert self.style != "new" 201 self.use_stash = [] 202 self.style = "new" 203 self.src_ranges = RangeSet() 204 205 def __str__(self): 206 return (str(self.id) + ": <" + str(self.src_ranges) + " " + self.style + 207 " to " + str(self.tgt_ranges) + ">") 208 209 210@functools.total_ordering 211class HeapItem(object): 212 def __init__(self, item): 213 self.item = item 214 # Negate the score since python's heap is a min-heap and we want 215 # the maximum score. 216 self.score = -item.score 217 def clear(self): 218 self.item = None 219 def __bool__(self): 220 return self.item is None 221 def __eq__(self, other): 222 return self.score == other.score 223 def __le__(self, other): 224 return self.score <= other.score 225 226 227# BlockImageDiff works on two image objects. An image object is 228# anything that provides the following attributes: 229# 230# blocksize: the size in bytes of a block, currently must be 4096. 231# 232# total_blocks: the total size of the partition/image, in blocks. 233# 234# care_map: a RangeSet containing which blocks (in the range [0, 235# total_blocks) we actually care about; i.e. which blocks contain 236# data. 237# 238# file_map: a dict that partitions the blocks contained in care_map 239# into smaller domains that are useful for doing diffs on. 240# (Typically a domain is a file, and the key in file_map is the 241# pathname.) 242# 243# clobbered_blocks: a RangeSet containing which blocks contain data 244# but may be altered by the FS. They need to be excluded when 245# verifying the partition integrity. 246# 247# ReadRangeSet(): a function that takes a RangeSet and returns the 248# data contained in the image blocks of that RangeSet. The data 249# is returned as a list or tuple of strings; concatenating the 250# elements together should produce the requested data. 251# Implementations are free to break up the data into list/tuple 252# elements in any way that is convenient. 253# 254# TotalSha1(): a function that returns (as a hex string) the SHA-1 255# hash of all the data in the image (ie, all the blocks in the 256# care_map minus clobbered_blocks, or including the clobbered 257# blocks if include_clobbered_blocks is True). 258# 259# When creating a BlockImageDiff, the src image may be None, in which 260# case the list of transfers produced will never read from the 261# original image. 262 263class BlockImageDiff(object): 264 def __init__(self, tgt, src=None, threads=None, version=4): 265 if threads is None: 266 threads = multiprocessing.cpu_count() // 2 267 if threads == 0: 268 threads = 1 269 self.threads = threads 270 self.version = version 271 self.transfers = [] 272 self.src_basenames = {} 273 self.src_numpatterns = {} 274 self._max_stashed_size = 0 275 self.touched_src_ranges = RangeSet() 276 self.touched_src_sha1 = None 277 278 assert version in (1, 2, 3, 4) 279 280 self.tgt = tgt 281 if src is None: 282 src = EmptyImage() 283 self.src = src 284 285 # The updater code that installs the patch always uses 4k blocks. 286 assert tgt.blocksize == 4096 287 assert src.blocksize == 4096 288 289 # The range sets in each filemap should comprise a partition of 290 # the care map. 291 self.AssertPartition(src.care_map, src.file_map.values()) 292 self.AssertPartition(tgt.care_map, tgt.file_map.values()) 293 294 @property 295 def max_stashed_size(self): 296 return self._max_stashed_size 297 298 def Compute(self, prefix): 299 # When looking for a source file to use as the diff input for a 300 # target file, we try: 301 # 1) an exact path match if available, otherwise 302 # 2) a exact basename match if available, otherwise 303 # 3) a basename match after all runs of digits are replaced by 304 # "#" if available, otherwise 305 # 4) we have no source for this target. 306 self.AbbreviateSourceNames() 307 self.FindTransfers() 308 309 # Find the ordering dependencies among transfers (this is O(n^2) 310 # in the number of transfers). 311 self.GenerateDigraph() 312 # Find a sequence of transfers that satisfies as many ordering 313 # dependencies as possible (heuristically). 314 self.FindVertexSequence() 315 # Fix up the ordering dependencies that the sequence didn't 316 # satisfy. 317 if self.version == 1: 318 self.RemoveBackwardEdges() 319 else: 320 self.ReverseBackwardEdges() 321 self.ImproveVertexSequence() 322 323 # Ensure the runtime stash size is under the limit. 324 if self.version >= 2 and common.OPTIONS.cache_size is not None: 325 self.ReviseStashSize() 326 327 # Double-check our work. 328 self.AssertSequenceGood() 329 330 self.ComputePatches(prefix) 331 self.WriteTransfers(prefix) 332 333 def HashBlocks(self, source, ranges): # pylint: disable=no-self-use 334 data = source.ReadRangeSet(ranges) 335 ctx = sha1() 336 337 for p in data: 338 ctx.update(p) 339 340 return ctx.hexdigest() 341 342 def WriteTransfers(self, prefix): 343 out = [] 344 345 total = 0 346 347 stashes = {} 348 stashed_blocks = 0 349 max_stashed_blocks = 0 350 351 free_stash_ids = [] 352 next_stash_id = 0 353 354 for xf in self.transfers: 355 356 if self.version < 2: 357 assert not xf.stash_before 358 assert not xf.use_stash 359 360 for s, sr in xf.stash_before: 361 assert s not in stashes 362 if free_stash_ids: 363 sid = heapq.heappop(free_stash_ids) 364 else: 365 sid = next_stash_id 366 next_stash_id += 1 367 stashes[s] = sid 368 if self.version == 2: 369 stashed_blocks += sr.size() 370 out.append("stash %d %s\n" % (sid, sr.to_string_raw())) 371 else: 372 sh = self.HashBlocks(self.src, sr) 373 if sh in stashes: 374 stashes[sh] += 1 375 else: 376 stashes[sh] = 1 377 stashed_blocks += sr.size() 378 self.touched_src_ranges = self.touched_src_ranges.union(sr) 379 out.append("stash %s %s\n" % (sh, sr.to_string_raw())) 380 381 if stashed_blocks > max_stashed_blocks: 382 max_stashed_blocks = stashed_blocks 383 384 free_string = [] 385 free_size = 0 386 387 if self.version == 1: 388 src_str = xf.src_ranges.to_string_raw() if xf.src_ranges else "" 389 elif self.version >= 2: 390 391 # <# blocks> <src ranges> 392 # OR 393 # <# blocks> <src ranges> <src locs> <stash refs...> 394 # OR 395 # <# blocks> - <stash refs...> 396 397 size = xf.src_ranges.size() 398 src_str = [str(size)] 399 400 unstashed_src_ranges = xf.src_ranges 401 mapped_stashes = [] 402 for s, sr in xf.use_stash: 403 sid = stashes.pop(s) 404 unstashed_src_ranges = unstashed_src_ranges.subtract(sr) 405 sh = self.HashBlocks(self.src, sr) 406 sr = xf.src_ranges.map_within(sr) 407 mapped_stashes.append(sr) 408 if self.version == 2: 409 src_str.append("%d:%s" % (sid, sr.to_string_raw())) 410 # A stash will be used only once. We need to free the stash 411 # immediately after the use, instead of waiting for the automatic 412 # clean-up at the end. Because otherwise it may take up extra space 413 # and lead to OTA failures. 414 # Bug: 23119955 415 free_string.append("free %d\n" % (sid,)) 416 free_size += sr.size() 417 else: 418 assert sh in stashes 419 src_str.append("%s:%s" % (sh, sr.to_string_raw())) 420 stashes[sh] -= 1 421 if stashes[sh] == 0: 422 free_size += sr.size() 423 free_string.append("free %s\n" % (sh)) 424 stashes.pop(sh) 425 heapq.heappush(free_stash_ids, sid) 426 427 if unstashed_src_ranges: 428 src_str.insert(1, unstashed_src_ranges.to_string_raw()) 429 if xf.use_stash: 430 mapped_unstashed = xf.src_ranges.map_within(unstashed_src_ranges) 431 src_str.insert(2, mapped_unstashed.to_string_raw()) 432 mapped_stashes.append(mapped_unstashed) 433 self.AssertPartition(RangeSet(data=(0, size)), mapped_stashes) 434 else: 435 src_str.insert(1, "-") 436 self.AssertPartition(RangeSet(data=(0, size)), mapped_stashes) 437 438 src_str = " ".join(src_str) 439 440 # all versions: 441 # zero <rangeset> 442 # new <rangeset> 443 # erase <rangeset> 444 # 445 # version 1: 446 # bsdiff patchstart patchlen <src rangeset> <tgt rangeset> 447 # imgdiff patchstart patchlen <src rangeset> <tgt rangeset> 448 # move <src rangeset> <tgt rangeset> 449 # 450 # version 2: 451 # bsdiff patchstart patchlen <tgt rangeset> <src_str> 452 # imgdiff patchstart patchlen <tgt rangeset> <src_str> 453 # move <tgt rangeset> <src_str> 454 # 455 # version 3: 456 # bsdiff patchstart patchlen srchash tgthash <tgt rangeset> <src_str> 457 # imgdiff patchstart patchlen srchash tgthash <tgt rangeset> <src_str> 458 # move hash <tgt rangeset> <src_str> 459 460 tgt_size = xf.tgt_ranges.size() 461 462 if xf.style == "new": 463 assert xf.tgt_ranges 464 out.append("%s %s\n" % (xf.style, xf.tgt_ranges.to_string_raw())) 465 total += tgt_size 466 elif xf.style == "move": 467 assert xf.tgt_ranges 468 assert xf.src_ranges.size() == tgt_size 469 if xf.src_ranges != xf.tgt_ranges: 470 if self.version == 1: 471 out.append("%s %s %s\n" % ( 472 xf.style, 473 xf.src_ranges.to_string_raw(), xf.tgt_ranges.to_string_raw())) 474 elif self.version == 2: 475 out.append("%s %s %s\n" % ( 476 xf.style, 477 xf.tgt_ranges.to_string_raw(), src_str)) 478 elif self.version >= 3: 479 # take into account automatic stashing of overlapping blocks 480 if xf.src_ranges.overlaps(xf.tgt_ranges): 481 temp_stash_usage = stashed_blocks + xf.src_ranges.size() 482 if temp_stash_usage > max_stashed_blocks: 483 max_stashed_blocks = temp_stash_usage 484 485 self.touched_src_ranges = self.touched_src_ranges.union( 486 xf.src_ranges) 487 488 out.append("%s %s %s %s\n" % ( 489 xf.style, 490 self.HashBlocks(self.tgt, xf.tgt_ranges), 491 xf.tgt_ranges.to_string_raw(), src_str)) 492 total += tgt_size 493 elif xf.style in ("bsdiff", "imgdiff"): 494 assert xf.tgt_ranges 495 assert xf.src_ranges 496 if self.version == 1: 497 out.append("%s %d %d %s %s\n" % ( 498 xf.style, xf.patch_start, xf.patch_len, 499 xf.src_ranges.to_string_raw(), xf.tgt_ranges.to_string_raw())) 500 elif self.version == 2: 501 out.append("%s %d %d %s %s\n" % ( 502 xf.style, xf.patch_start, xf.patch_len, 503 xf.tgt_ranges.to_string_raw(), src_str)) 504 elif self.version >= 3: 505 # take into account automatic stashing of overlapping blocks 506 if xf.src_ranges.overlaps(xf.tgt_ranges): 507 temp_stash_usage = stashed_blocks + xf.src_ranges.size() 508 if temp_stash_usage > max_stashed_blocks: 509 max_stashed_blocks = temp_stash_usage 510 511 self.touched_src_ranges = self.touched_src_ranges.union( 512 xf.src_ranges) 513 514 out.append("%s %d %d %s %s %s %s\n" % ( 515 xf.style, 516 xf.patch_start, xf.patch_len, 517 self.HashBlocks(self.src, xf.src_ranges), 518 self.HashBlocks(self.tgt, xf.tgt_ranges), 519 xf.tgt_ranges.to_string_raw(), src_str)) 520 total += tgt_size 521 elif xf.style == "zero": 522 assert xf.tgt_ranges 523 to_zero = xf.tgt_ranges.subtract(xf.src_ranges) 524 if to_zero: 525 out.append("%s %s\n" % (xf.style, to_zero.to_string_raw())) 526 total += to_zero.size() 527 else: 528 raise ValueError("unknown transfer style '%s'\n" % xf.style) 529 530 if free_string: 531 out.append("".join(free_string)) 532 stashed_blocks -= free_size 533 534 if self.version >= 2 and common.OPTIONS.cache_size is not None: 535 # Sanity check: abort if we're going to need more stash space than 536 # the allowed size (cache_size * threshold). There are two purposes 537 # of having a threshold here. a) Part of the cache may have been 538 # occupied by some recovery logs. b) It will buy us some time to deal 539 # with the oversize issue. 540 cache_size = common.OPTIONS.cache_size 541 stash_threshold = common.OPTIONS.stash_threshold 542 max_allowed = cache_size * stash_threshold 543 assert max_stashed_blocks * self.tgt.blocksize < max_allowed, \ 544 'Stash size %d (%d * %d) exceeds the limit %d (%d * %.2f)' % ( 545 max_stashed_blocks * self.tgt.blocksize, max_stashed_blocks, 546 self.tgt.blocksize, max_allowed, cache_size, 547 stash_threshold) 548 549 if self.version >= 3: 550 self.touched_src_sha1 = self.HashBlocks( 551 self.src, self.touched_src_ranges) 552 553 # Zero out extended blocks as a workaround for bug 20881595. 554 if self.tgt.extended: 555 out.append("zero %s\n" % (self.tgt.extended.to_string_raw(),)) 556 total += self.tgt.extended.size() 557 558 # We erase all the blocks on the partition that a) don't contain useful 559 # data in the new image and b) will not be touched by dm-verity. 560 all_tgt = RangeSet(data=(0, self.tgt.total_blocks)) 561 all_tgt_minus_extended = all_tgt.subtract(self.tgt.extended) 562 new_dontcare = all_tgt_minus_extended.subtract(self.tgt.care_map) 563 if new_dontcare: 564 out.append("erase %s\n" % (new_dontcare.to_string_raw(),)) 565 566 out.insert(0, "%d\n" % (self.version,)) # format version number 567 out.insert(1, "%d\n" % (total,)) 568 if self.version >= 2: 569 # version 2 only: after the total block count, we give the number 570 # of stash slots needed, and the maximum size needed (in blocks) 571 out.insert(2, str(next_stash_id) + "\n") 572 out.insert(3, str(max_stashed_blocks) + "\n") 573 574 with open(prefix + ".transfer.list", "wb") as f: 575 for i in out: 576 f.write(i) 577 578 if self.version >= 2: 579 self._max_stashed_size = max_stashed_blocks * self.tgt.blocksize 580 OPTIONS = common.OPTIONS 581 if OPTIONS.cache_size is not None: 582 max_allowed = OPTIONS.cache_size * OPTIONS.stash_threshold 583 print("max stashed blocks: %d (%d bytes), " 584 "limit: %d bytes (%.2f%%)\n" % ( 585 max_stashed_blocks, self._max_stashed_size, max_allowed, 586 self._max_stashed_size * 100.0 / max_allowed)) 587 else: 588 print("max stashed blocks: %d (%d bytes), limit: <unknown>\n" % ( 589 max_stashed_blocks, self._max_stashed_size)) 590 591 def ReviseStashSize(self): 592 print("Revising stash size...") 593 stashes = {} 594 595 # Create the map between a stash and its def/use points. For example, for a 596 # given stash of (idx, sr), stashes[idx] = (sr, def_cmd, use_cmd). 597 for xf in self.transfers: 598 # Command xf defines (stores) all the stashes in stash_before. 599 for idx, sr in xf.stash_before: 600 stashes[idx] = (sr, xf) 601 602 # Record all the stashes command xf uses. 603 for idx, _ in xf.use_stash: 604 stashes[idx] += (xf,) 605 606 # Compute the maximum blocks available for stash based on /cache size and 607 # the threshold. 608 cache_size = common.OPTIONS.cache_size 609 stash_threshold = common.OPTIONS.stash_threshold 610 max_allowed = cache_size * stash_threshold / self.tgt.blocksize 611 612 stashed_blocks = 0 613 new_blocks = 0 614 615 # Now go through all the commands. Compute the required stash size on the 616 # fly. If a command requires excess stash than available, it deletes the 617 # stash by replacing the command that uses the stash with a "new" command 618 # instead. 619 for xf in self.transfers: 620 replaced_cmds = [] 621 622 # xf.stash_before generates explicit stash commands. 623 for idx, sr in xf.stash_before: 624 if stashed_blocks + sr.size() > max_allowed: 625 # We cannot stash this one for a later command. Find out the command 626 # that will use this stash and replace the command with "new". 627 use_cmd = stashes[idx][2] 628 replaced_cmds.append(use_cmd) 629 print("%10d %9s %s" % (sr.size(), "explicit", use_cmd)) 630 else: 631 stashed_blocks += sr.size() 632 633 # xf.use_stash generates free commands. 634 for _, sr in xf.use_stash: 635 stashed_blocks -= sr.size() 636 637 # "move" and "diff" may introduce implicit stashes in BBOTA v3. Prior to 638 # ComputePatches(), they both have the style of "diff". 639 if xf.style == "diff" and self.version >= 3: 640 assert xf.tgt_ranges and xf.src_ranges 641 if xf.src_ranges.overlaps(xf.tgt_ranges): 642 if stashed_blocks + xf.src_ranges.size() > max_allowed: 643 replaced_cmds.append(xf) 644 print("%10d %9s %s" % (xf.src_ranges.size(), "implicit", xf)) 645 646 # Replace the commands in replaced_cmds with "new"s. 647 for cmd in replaced_cmds: 648 # It no longer uses any commands in "use_stash". Remove the def points 649 # for all those stashes. 650 for idx, sr in cmd.use_stash: 651 def_cmd = stashes[idx][1] 652 assert (idx, sr) in def_cmd.stash_before 653 def_cmd.stash_before.remove((idx, sr)) 654 655 # Add up blocks that violates space limit and print total number to 656 # screen later. 657 new_blocks += cmd.tgt_ranges.size() 658 cmd.ConvertToNew() 659 660 num_of_bytes = new_blocks * self.tgt.blocksize 661 print(" Total %d blocks (%d bytes) are packed as new blocks due to " 662 "insufficient cache size." % (new_blocks, num_of_bytes)) 663 664 def ComputePatches(self, prefix): 665 print("Reticulating splines...") 666 diff_q = [] 667 patch_num = 0 668 with open(prefix + ".new.dat", "wb") as new_f: 669 for xf in self.transfers: 670 if xf.style == "zero": 671 pass 672 elif xf.style == "new": 673 for piece in self.tgt.ReadRangeSet(xf.tgt_ranges): 674 new_f.write(piece) 675 elif xf.style == "diff": 676 src = self.src.ReadRangeSet(xf.src_ranges) 677 tgt = self.tgt.ReadRangeSet(xf.tgt_ranges) 678 679 # We can't compare src and tgt directly because they may have 680 # the same content but be broken up into blocks differently, eg: 681 # 682 # ["he", "llo"] vs ["h", "ello"] 683 # 684 # We want those to compare equal, ideally without having to 685 # actually concatenate the strings (these may be tens of 686 # megabytes). 687 688 src_sha1 = sha1() 689 for p in src: 690 src_sha1.update(p) 691 tgt_sha1 = sha1() 692 tgt_size = 0 693 for p in tgt: 694 tgt_sha1.update(p) 695 tgt_size += len(p) 696 697 if src_sha1.digest() == tgt_sha1.digest(): 698 # These are identical; we don't need to generate a patch, 699 # just issue copy commands on the device. 700 xf.style = "move" 701 else: 702 # For files in zip format (eg, APKs, JARs, etc.) we would 703 # like to use imgdiff -z if possible (because it usually 704 # produces significantly smaller patches than bsdiff). 705 # This is permissible if: 706 # 707 # - the source and target files are monotonic (ie, the 708 # data is stored with blocks in increasing order), and 709 # - we haven't removed any blocks from the source set. 710 # 711 # If these conditions are satisfied then appending all the 712 # blocks in the set together in order will produce a valid 713 # zip file (plus possibly extra zeros in the last block), 714 # which is what imgdiff needs to operate. (imgdiff is 715 # fine with extra zeros at the end of the file.) 716 imgdiff = (xf.intact and 717 xf.tgt_name.split(".")[-1].lower() 718 in ("apk", "jar", "zip")) 719 xf.style = "imgdiff" if imgdiff else "bsdiff" 720 diff_q.append((tgt_size, src, tgt, xf, patch_num)) 721 patch_num += 1 722 723 else: 724 assert False, "unknown style " + xf.style 725 726 if diff_q: 727 if self.threads > 1: 728 print("Computing patches (using %d threads)..." % (self.threads,)) 729 else: 730 print("Computing patches...") 731 diff_q.sort() 732 733 patches = [None] * patch_num 734 735 # TODO: Rewrite with multiprocessing.ThreadPool? 736 lock = threading.Lock() 737 def diff_worker(): 738 while True: 739 with lock: 740 if not diff_q: 741 return 742 tgt_size, src, tgt, xf, patchnum = diff_q.pop() 743 patch = compute_patch(src, tgt, imgdiff=(xf.style == "imgdiff")) 744 size = len(patch) 745 with lock: 746 patches[patchnum] = (patch, xf) 747 print("%10d %10d (%6.2f%%) %7s %s" % ( 748 size, tgt_size, size * 100.0 / tgt_size, xf.style, 749 xf.tgt_name if xf.tgt_name == xf.src_name else ( 750 xf.tgt_name + " (from " + xf.src_name + ")"))) 751 752 threads = [threading.Thread(target=diff_worker) 753 for _ in range(self.threads)] 754 for th in threads: 755 th.start() 756 while threads: 757 threads.pop().join() 758 else: 759 patches = [] 760 761 p = 0 762 with open(prefix + ".patch.dat", "wb") as patch_f: 763 for patch, xf in patches: 764 xf.patch_start = p 765 xf.patch_len = len(patch) 766 patch_f.write(patch) 767 p += len(patch) 768 769 def AssertSequenceGood(self): 770 # Simulate the sequences of transfers we will output, and check that: 771 # - we never read a block after writing it, and 772 # - we write every block we care about exactly once. 773 774 # Start with no blocks having been touched yet. 775 touched = array.array("B", "\0" * self.tgt.total_blocks) 776 777 # Imagine processing the transfers in order. 778 for xf in self.transfers: 779 # Check that the input blocks for this transfer haven't yet been touched. 780 781 x = xf.src_ranges 782 if self.version >= 2: 783 for _, sr in xf.use_stash: 784 x = x.subtract(sr) 785 786 for s, e in x: 787 # Source image could be larger. Don't check the blocks that are in the 788 # source image only. Since they are not in 'touched', and won't ever 789 # be touched. 790 for i in range(s, min(e, self.tgt.total_blocks)): 791 assert touched[i] == 0 792 793 # Check that the output blocks for this transfer haven't yet 794 # been touched, and touch all the blocks written by this 795 # transfer. 796 for s, e in xf.tgt_ranges: 797 for i in range(s, e): 798 assert touched[i] == 0 799 touched[i] = 1 800 801 # Check that we've written every target block. 802 for s, e in self.tgt.care_map: 803 for i in range(s, e): 804 assert touched[i] == 1 805 806 def ImproveVertexSequence(self): 807 print("Improving vertex order...") 808 809 # At this point our digraph is acyclic; we reversed any edges that 810 # were backwards in the heuristically-generated sequence. The 811 # previously-generated order is still acceptable, but we hope to 812 # find a better order that needs less memory for stashed data. 813 # Now we do a topological sort to generate a new vertex order, 814 # using a greedy algorithm to choose which vertex goes next 815 # whenever we have a choice. 816 817 # Make a copy of the edge set; this copy will get destroyed by the 818 # algorithm. 819 for xf in self.transfers: 820 xf.incoming = xf.goes_after.copy() 821 xf.outgoing = xf.goes_before.copy() 822 823 L = [] # the new vertex order 824 825 # S is the set of sources in the remaining graph; we always choose 826 # the one that leaves the least amount of stashed data after it's 827 # executed. 828 S = [(u.NetStashChange(), u.order, u) for u in self.transfers 829 if not u.incoming] 830 heapq.heapify(S) 831 832 while S: 833 _, _, xf = heapq.heappop(S) 834 L.append(xf) 835 for u in xf.outgoing: 836 del u.incoming[xf] 837 if not u.incoming: 838 heapq.heappush(S, (u.NetStashChange(), u.order, u)) 839 840 # if this fails then our graph had a cycle. 841 assert len(L) == len(self.transfers) 842 843 self.transfers = L 844 for i, xf in enumerate(L): 845 xf.order = i 846 847 def RemoveBackwardEdges(self): 848 print("Removing backward edges...") 849 in_order = 0 850 out_of_order = 0 851 lost_source = 0 852 853 for xf in self.transfers: 854 lost = 0 855 size = xf.src_ranges.size() 856 for u in xf.goes_before: 857 # xf should go before u 858 if xf.order < u.order: 859 # it does, hurray! 860 in_order += 1 861 else: 862 # it doesn't, boo. trim the blocks that u writes from xf's 863 # source, so that xf can go after u. 864 out_of_order += 1 865 assert xf.src_ranges.overlaps(u.tgt_ranges) 866 xf.src_ranges = xf.src_ranges.subtract(u.tgt_ranges) 867 xf.intact = False 868 869 if xf.style == "diff" and not xf.src_ranges: 870 # nothing left to diff from; treat as new data 871 xf.style = "new" 872 873 lost = size - xf.src_ranges.size() 874 lost_source += lost 875 876 print((" %d/%d dependencies (%.2f%%) were violated; " 877 "%d source blocks removed.") % 878 (out_of_order, in_order + out_of_order, 879 (out_of_order * 100.0 / (in_order + out_of_order)) 880 if (in_order + out_of_order) else 0.0, 881 lost_source)) 882 883 def ReverseBackwardEdges(self): 884 print("Reversing backward edges...") 885 in_order = 0 886 out_of_order = 0 887 stashes = 0 888 stash_size = 0 889 890 for xf in self.transfers: 891 for u in xf.goes_before.copy(): 892 # xf should go before u 893 if xf.order < u.order: 894 # it does, hurray! 895 in_order += 1 896 else: 897 # it doesn't, boo. modify u to stash the blocks that it 898 # writes that xf wants to read, and then require u to go 899 # before xf. 900 out_of_order += 1 901 902 overlap = xf.src_ranges.intersect(u.tgt_ranges) 903 assert overlap 904 905 u.stash_before.append((stashes, overlap)) 906 xf.use_stash.append((stashes, overlap)) 907 stashes += 1 908 stash_size += overlap.size() 909 910 # reverse the edge direction; now xf must go after u 911 del xf.goes_before[u] 912 del u.goes_after[xf] 913 xf.goes_after[u] = None # value doesn't matter 914 u.goes_before[xf] = None 915 916 print((" %d/%d dependencies (%.2f%%) were violated; " 917 "%d source blocks stashed.") % 918 (out_of_order, in_order + out_of_order, 919 (out_of_order * 100.0 / (in_order + out_of_order)) 920 if (in_order + out_of_order) else 0.0, 921 stash_size)) 922 923 def FindVertexSequence(self): 924 print("Finding vertex sequence...") 925 926 # This is based on "A Fast & Effective Heuristic for the Feedback 927 # Arc Set Problem" by P. Eades, X. Lin, and W.F. Smyth. Think of 928 # it as starting with the digraph G and moving all the vertices to 929 # be on a horizontal line in some order, trying to minimize the 930 # number of edges that end up pointing to the left. Left-pointing 931 # edges will get removed to turn the digraph into a DAG. In this 932 # case each edge has a weight which is the number of source blocks 933 # we'll lose if that edge is removed; we try to minimize the total 934 # weight rather than just the number of edges. 935 936 # Make a copy of the edge set; this copy will get destroyed by the 937 # algorithm. 938 for xf in self.transfers: 939 xf.incoming = xf.goes_after.copy() 940 xf.outgoing = xf.goes_before.copy() 941 xf.score = sum(xf.outgoing.values()) - sum(xf.incoming.values()) 942 943 # We use an OrderedDict instead of just a set so that the output 944 # is repeatable; otherwise it would depend on the hash values of 945 # the transfer objects. 946 G = OrderedDict() 947 for xf in self.transfers: 948 G[xf] = None 949 s1 = deque() # the left side of the sequence, built from left to right 950 s2 = deque() # the right side of the sequence, built from right to left 951 952 heap = [] 953 for xf in self.transfers: 954 xf.heap_item = HeapItem(xf) 955 heap.append(xf.heap_item) 956 heapq.heapify(heap) 957 958 sinks = set(u for u in G if not u.outgoing) 959 sources = set(u for u in G if not u.incoming) 960 961 def adjust_score(iu, delta): 962 iu.score += delta 963 iu.heap_item.clear() 964 iu.heap_item = HeapItem(iu) 965 heapq.heappush(heap, iu.heap_item) 966 967 while G: 968 # Put all sinks at the end of the sequence. 969 while sinks: 970 new_sinks = set() 971 for u in sinks: 972 if u not in G: continue 973 s2.appendleft(u) 974 del G[u] 975 for iu in u.incoming: 976 adjust_score(iu, -iu.outgoing.pop(u)) 977 if not iu.outgoing: new_sinks.add(iu) 978 sinks = new_sinks 979 980 # Put all the sources at the beginning of the sequence. 981 while sources: 982 new_sources = set() 983 for u in sources: 984 if u not in G: continue 985 s1.append(u) 986 del G[u] 987 for iu in u.outgoing: 988 adjust_score(iu, +iu.incoming.pop(u)) 989 if not iu.incoming: new_sources.add(iu) 990 sources = new_sources 991 992 if not G: break 993 994 # Find the "best" vertex to put next. "Best" is the one that 995 # maximizes the net difference in source blocks saved we get by 996 # pretending it's a source rather than a sink. 997 998 while True: 999 u = heapq.heappop(heap) 1000 if u and u.item in G: 1001 u = u.item 1002 break 1003 1004 s1.append(u) 1005 del G[u] 1006 for iu in u.outgoing: 1007 adjust_score(iu, +iu.incoming.pop(u)) 1008 if not iu.incoming: sources.add(iu) 1009 1010 for iu in u.incoming: 1011 adjust_score(iu, -iu.outgoing.pop(u)) 1012 if not iu.outgoing: sinks.add(iu) 1013 1014 # Now record the sequence in the 'order' field of each transfer, 1015 # and by rearranging self.transfers to be in the chosen sequence. 1016 1017 new_transfers = [] 1018 for x in itertools.chain(s1, s2): 1019 x.order = len(new_transfers) 1020 new_transfers.append(x) 1021 del x.incoming 1022 del x.outgoing 1023 1024 self.transfers = new_transfers 1025 1026 def GenerateDigraph(self): 1027 print("Generating digraph...") 1028 1029 # Each item of source_ranges will be: 1030 # - None, if that block is not used as a source, 1031 # - a transfer, if one transfer uses it as a source, or 1032 # - a set of transfers. 1033 source_ranges = [] 1034 for b in self.transfers: 1035 for s, e in b.src_ranges: 1036 if e > len(source_ranges): 1037 source_ranges.extend([None] * (e-len(source_ranges))) 1038 for i in range(s, e): 1039 if source_ranges[i] is None: 1040 source_ranges[i] = b 1041 else: 1042 if not isinstance(source_ranges[i], set): 1043 source_ranges[i] = set([source_ranges[i]]) 1044 source_ranges[i].add(b) 1045 1046 for a in self.transfers: 1047 intersections = set() 1048 for s, e in a.tgt_ranges: 1049 for i in range(s, e): 1050 if i >= len(source_ranges): break 1051 b = source_ranges[i] 1052 if b is not None: 1053 if isinstance(b, set): 1054 intersections.update(b) 1055 else: 1056 intersections.add(b) 1057 1058 for b in intersections: 1059 if a is b: continue 1060 1061 # If the blocks written by A are read by B, then B needs to go before A. 1062 i = a.tgt_ranges.intersect(b.src_ranges) 1063 if i: 1064 if b.src_name == "__ZERO": 1065 # the cost of removing source blocks for the __ZERO domain 1066 # is (nearly) zero. 1067 size = 0 1068 else: 1069 size = i.size() 1070 b.goes_before[a] = size 1071 a.goes_after[b] = size 1072 1073 def FindTransfers(self): 1074 """Parse the file_map to generate all the transfers.""" 1075 1076 def AddTransfer(tgt_name, src_name, tgt_ranges, src_ranges, style, by_id, 1077 split=False): 1078 """Wrapper function for adding a Transfer(). 1079 1080 For BBOTA v3, we need to stash source blocks for resumable feature. 1081 However, with the growth of file size and the shrink of the cache 1082 partition source blocks are too large to be stashed. If a file occupies 1083 too many blocks (greater than MAX_BLOCKS_PER_DIFF_TRANSFER), we split it 1084 into smaller pieces by getting multiple Transfer()s. 1085 1086 The downside is that after splitting, we may increase the package size 1087 since the split pieces don't align well. According to our experiments, 1088 1/8 of the cache size as the per-piece limit appears to be optimal. 1089 Compared to the fixed 1024-block limit, it reduces the overall package 1090 size by 30% volantis, and 20% for angler and bullhead.""" 1091 1092 # We care about diff transfers only. 1093 if style != "diff" or not split: 1094 Transfer(tgt_name, src_name, tgt_ranges, src_ranges, style, by_id) 1095 return 1096 1097 pieces = 0 1098 cache_size = common.OPTIONS.cache_size 1099 split_threshold = 0.125 1100 max_blocks_per_transfer = int(cache_size * split_threshold / 1101 self.tgt.blocksize) 1102 1103 # Change nothing for small files. 1104 if (tgt_ranges.size() <= max_blocks_per_transfer and 1105 src_ranges.size() <= max_blocks_per_transfer): 1106 Transfer(tgt_name, src_name, tgt_ranges, src_ranges, style, by_id) 1107 return 1108 1109 while (tgt_ranges.size() > max_blocks_per_transfer and 1110 src_ranges.size() > max_blocks_per_transfer): 1111 tgt_split_name = "%s-%d" % (tgt_name, pieces) 1112 src_split_name = "%s-%d" % (src_name, pieces) 1113 tgt_first = tgt_ranges.first(max_blocks_per_transfer) 1114 src_first = src_ranges.first(max_blocks_per_transfer) 1115 1116 Transfer(tgt_split_name, src_split_name, tgt_first, src_first, style, 1117 by_id) 1118 1119 tgt_ranges = tgt_ranges.subtract(tgt_first) 1120 src_ranges = src_ranges.subtract(src_first) 1121 pieces += 1 1122 1123 # Handle remaining blocks. 1124 if tgt_ranges.size() or src_ranges.size(): 1125 # Must be both non-empty. 1126 assert tgt_ranges.size() and src_ranges.size() 1127 tgt_split_name = "%s-%d" % (tgt_name, pieces) 1128 src_split_name = "%s-%d" % (src_name, pieces) 1129 Transfer(tgt_split_name, src_split_name, tgt_ranges, src_ranges, style, 1130 by_id) 1131 1132 empty = RangeSet() 1133 for tgt_fn, tgt_ranges in self.tgt.file_map.items(): 1134 if tgt_fn == "__ZERO": 1135 # the special "__ZERO" domain is all the blocks not contained 1136 # in any file and that are filled with zeros. We have a 1137 # special transfer style for zero blocks. 1138 src_ranges = self.src.file_map.get("__ZERO", empty) 1139 AddTransfer(tgt_fn, "__ZERO", tgt_ranges, src_ranges, 1140 "zero", self.transfers) 1141 continue 1142 1143 elif tgt_fn == "__COPY": 1144 # "__COPY" domain includes all the blocks not contained in any 1145 # file and that need to be copied unconditionally to the target. 1146 AddTransfer(tgt_fn, None, tgt_ranges, empty, "new", self.transfers) 1147 continue 1148 1149 elif tgt_fn in self.src.file_map: 1150 # Look for an exact pathname match in the source. 1151 AddTransfer(tgt_fn, tgt_fn, tgt_ranges, self.src.file_map[tgt_fn], 1152 "diff", self.transfers, self.version >= 3) 1153 continue 1154 1155 b = os.path.basename(tgt_fn) 1156 if b in self.src_basenames: 1157 # Look for an exact basename match in the source. 1158 src_fn = self.src_basenames[b] 1159 AddTransfer(tgt_fn, src_fn, tgt_ranges, self.src.file_map[src_fn], 1160 "diff", self.transfers, self.version >= 3) 1161 continue 1162 1163 b = re.sub("[0-9]+", "#", b) 1164 if b in self.src_numpatterns: 1165 # Look for a 'number pattern' match (a basename match after 1166 # all runs of digits are replaced by "#"). (This is useful 1167 # for .so files that contain version numbers in the filename 1168 # that get bumped.) 1169 src_fn = self.src_numpatterns[b] 1170 AddTransfer(tgt_fn, src_fn, tgt_ranges, self.src.file_map[src_fn], 1171 "diff", self.transfers, self.version >= 3) 1172 continue 1173 1174 AddTransfer(tgt_fn, None, tgt_ranges, empty, "new", self.transfers) 1175 1176 def AbbreviateSourceNames(self): 1177 for k in self.src.file_map.keys(): 1178 b = os.path.basename(k) 1179 self.src_basenames[b] = k 1180 b = re.sub("[0-9]+", "#", b) 1181 self.src_numpatterns[b] = k 1182 1183 @staticmethod 1184 def AssertPartition(total, seq): 1185 """Assert that all the RangeSets in 'seq' form a partition of the 1186 'total' RangeSet (ie, they are nonintersecting and their union 1187 equals 'total').""" 1188 1189 so_far = RangeSet() 1190 for i in seq: 1191 assert not so_far.overlaps(i) 1192 so_far = so_far.union(i) 1193 assert so_far == total 1194