• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env python
2#
3# Copyright (C) 2014 The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#      http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16
17from __future__ import print_function
18
19import argparse
20import bisect
21import logging
22import os
23import struct
24import threading
25from hashlib import sha1
26
27import rangelib
28
29logger = logging.getLogger(__name__)
30
31
32class SparseImage(object):
33  """Wraps a sparse image file into an image object.
34
35  Wraps a sparse image file (and optional file map and clobbered_blocks) into
36  an image object suitable for passing to BlockImageDiff. file_map contains
37  the mapping between files and their blocks. clobbered_blocks contains the set
38  of blocks that should be always written to the target regardless of the old
39  contents (i.e. copying instead of patching). clobbered_blocks should be in
40  the form of a string like "0" or "0 1-5 8".
41  """
42
43  def __init__(self, simg_fn, file_map_fn=None, clobbered_blocks=None,
44               mode="rb", build_map=True, allow_shared_blocks=False,
45               hashtree_info_generator=None):
46    self.simg_f = f = open(simg_fn, mode)
47
48    header_bin = f.read(28)
49    header = struct.unpack("<I4H4I", header_bin)
50
51    magic = header[0]
52    major_version = header[1]
53    minor_version = header[2]
54    file_hdr_sz = header[3]
55    chunk_hdr_sz = header[4]
56    self.blocksize = blk_sz = header[5]
57    self.total_blocks = total_blks = header[6]
58    self.total_chunks = total_chunks = header[7]
59
60    if magic != 0xED26FF3A:
61      raise ValueError("Magic should be 0xED26FF3A but is 0x%08X" % (magic,))
62    if major_version != 1 or minor_version != 0:
63      raise ValueError("I know about version 1.0, but this is version %u.%u" %
64                       (major_version, minor_version))
65    if file_hdr_sz != 28:
66      raise ValueError("File header size was expected to be 28, but is %u." %
67                       (file_hdr_sz,))
68    if chunk_hdr_sz != 12:
69      raise ValueError("Chunk header size was expected to be 12, but is %u." %
70                       (chunk_hdr_sz,))
71
72    logger.info(
73        "Total of %u %u-byte output blocks in %u input chunks.", total_blks,
74        blk_sz, total_chunks)
75
76    if not build_map:
77      assert not hashtree_info_generator, \
78        "Cannot generate the hashtree info without building the offset map."
79      return
80
81    pos = 0   # in blocks
82    care_data = []
83    self.offset_map = offset_map = []
84    self.clobbered_blocks = rangelib.RangeSet(data=clobbered_blocks)
85
86    for i in range(total_chunks):
87      header_bin = f.read(12)
88      header = struct.unpack("<2H2I", header_bin)
89      chunk_type = header[0]
90      chunk_sz = header[2]
91      total_sz = header[3]
92      data_sz = total_sz - 12
93
94      if chunk_type == 0xCAC1:
95        if data_sz != (chunk_sz * blk_sz):
96          raise ValueError(
97              "Raw chunk input size (%u) does not match output size (%u)" %
98              (data_sz, chunk_sz * blk_sz))
99        else:
100          care_data.append(pos)
101          care_data.append(pos + chunk_sz)
102          offset_map.append((pos, chunk_sz, f.tell(), None))
103          pos += chunk_sz
104          f.seek(data_sz, os.SEEK_CUR)
105
106      elif chunk_type == 0xCAC2:
107        fill_data = f.read(4)
108        care_data.append(pos)
109        care_data.append(pos + chunk_sz)
110        offset_map.append((pos, chunk_sz, None, fill_data))
111        pos += chunk_sz
112
113      elif chunk_type == 0xCAC3:
114        if data_sz != 0:
115          raise ValueError("Don't care chunk input size is non-zero (%u)" %
116                           (data_sz))
117        # Fills the don't care data ranges with zeros.
118        # TODO(xunchang) pass the care_map to hashtree info generator.
119        if hashtree_info_generator:
120          fill_data = '\x00' * 4
121          # In order to compute verity hashtree on device, we need to write
122          # zeros explicitly to the don't care ranges. Because these ranges may
123          # contain non-zero data from the previous build.
124          care_data.append(pos)
125          care_data.append(pos + chunk_sz)
126          offset_map.append((pos, chunk_sz, None, fill_data))
127
128        pos += chunk_sz
129
130      elif chunk_type == 0xCAC4:
131        raise ValueError("CRC32 chunks are not supported")
132
133      else:
134        raise ValueError("Unknown chunk type 0x%04X not supported" %
135                         (chunk_type,))
136
137    self.generator_lock = threading.Lock()
138
139    self.care_map = rangelib.RangeSet(care_data)
140    self.offset_index = [i[0] for i in offset_map]
141
142    # Bug: 20881595
143    # Introduce extended blocks as a workaround for the bug. dm-verity may
144    # touch blocks that are not in the care_map due to block device
145    # read-ahead. It will fail if such blocks contain non-zeroes. We zero out
146    # the extended blocks explicitly to avoid dm-verity failures. 512 blocks
147    # are the maximum read-ahead we configure for dm-verity block devices.
148    extended = self.care_map.extend(512)
149    all_blocks = rangelib.RangeSet(data=(0, self.total_blocks))
150    extended = extended.intersect(all_blocks).subtract(self.care_map)
151    self.extended = extended
152
153    self.hashtree_info = None
154    if hashtree_info_generator:
155      self.hashtree_info = hashtree_info_generator.Generate(self)
156
157    if file_map_fn:
158      self.LoadFileBlockMap(file_map_fn, self.clobbered_blocks,
159                            allow_shared_blocks)
160    else:
161      self.file_map = {"__DATA": self.care_map}
162
163  def AppendFillChunk(self, data, blocks):
164    f = self.simg_f
165
166    # Append a fill chunk
167    f.seek(0, os.SEEK_END)
168    f.write(struct.pack("<2H3I", 0xCAC2, 0, blocks, 16, data))
169
170    # Update the sparse header
171    self.total_blocks += blocks
172    self.total_chunks += 1
173
174    f.seek(16, os.SEEK_SET)
175    f.write(struct.pack("<2I", self.total_blocks, self.total_chunks))
176
177  def RangeSha1(self, ranges):
178    h = sha1()
179    for data in self._GetRangeData(ranges):
180      h.update(data)
181    return h.hexdigest()
182
183  def ReadRangeSet(self, ranges):
184    return [d for d in self._GetRangeData(ranges)]
185
186  def TotalSha1(self, include_clobbered_blocks=False):
187    """Return the SHA-1 hash of all data in the 'care' regions.
188
189    If include_clobbered_blocks is True, it returns the hash including the
190    clobbered_blocks."""
191    ranges = self.care_map
192    if not include_clobbered_blocks:
193      ranges = ranges.subtract(self.clobbered_blocks)
194    return self.RangeSha1(ranges)
195
196  def WriteRangeDataToFd(self, ranges, fd):
197    for data in self._GetRangeData(ranges):
198      fd.write(data)
199
200  def _GetRangeData(self, ranges):
201    """Generator that produces all the image data in 'ranges'.  The
202    number of individual pieces returned is arbitrary (and in
203    particular is not necessarily equal to the number of ranges in
204    'ranges'.
205
206    Use a lock to protect the generator so that we will not run two
207    instances of this generator on the same object simultaneously."""
208
209    f = self.simg_f
210    with self.generator_lock:
211      for s, e in ranges:
212        to_read = e-s
213        idx = bisect.bisect_right(self.offset_index, s) - 1
214        chunk_start, chunk_len, filepos, fill_data = self.offset_map[idx]
215
216        # for the first chunk we may be starting partway through it.
217        remain = chunk_len - (s - chunk_start)
218        this_read = min(remain, to_read)
219        if filepos is not None:
220          p = filepos + ((s - chunk_start) * self.blocksize)
221          f.seek(p, os.SEEK_SET)
222          yield f.read(this_read * self.blocksize)
223        else:
224          yield fill_data * (this_read * (self.blocksize >> 2))
225        to_read -= this_read
226
227        while to_read > 0:
228          # continue with following chunks if this range spans multiple chunks.
229          idx += 1
230          chunk_start, chunk_len, filepos, fill_data = self.offset_map[idx]
231          this_read = min(chunk_len, to_read)
232          if filepos is not None:
233            f.seek(filepos, os.SEEK_SET)
234            yield f.read(this_read * self.blocksize)
235          else:
236            yield fill_data * (this_read * (self.blocksize >> 2))
237          to_read -= this_read
238
239  def LoadFileBlockMap(self, fn, clobbered_blocks, allow_shared_blocks):
240    """Loads the given block map file.
241
242    Args:
243      fn: The filename of the block map file.
244      clobbered_blocks: A RangeSet instance for the clobbered blocks.
245      allow_shared_blocks: Whether having shared blocks is allowed.
246    """
247    remaining = self.care_map
248    self.file_map = out = {}
249
250    with open(fn) as f:
251      for line in f:
252        fn, ranges = line.split(None, 1)
253        ranges = rangelib.RangeSet.parse(ranges)
254
255        if allow_shared_blocks:
256          # Find the shared blocks that have been claimed by others. If so, tag
257          # the entry so that we can skip applying imgdiff on this file.
258          shared_blocks = ranges.subtract(remaining)
259          if shared_blocks:
260            non_shared = ranges.subtract(shared_blocks)
261            if not non_shared:
262              continue
263
264            # There shouldn't anything in the extra dict yet.
265            assert not ranges.extra, "Non-empty RangeSet.extra"
266
267            # Put the non-shared RangeSet as the value in the block map, which
268            # has a copy of the original RangeSet.
269            non_shared.extra['uses_shared_blocks'] = ranges
270            ranges = non_shared
271
272        out[fn] = ranges
273        assert ranges.size() == ranges.intersect(remaining).size()
274
275        # Currently we assume that blocks in clobbered_blocks are not part of
276        # any file.
277        assert not clobbered_blocks.overlaps(ranges)
278        remaining = remaining.subtract(ranges)
279
280    remaining = remaining.subtract(clobbered_blocks)
281    if self.hashtree_info:
282      remaining = remaining.subtract(self.hashtree_info.hashtree_range)
283
284    # For all the remaining blocks in the care_map (ie, those that
285    # aren't part of the data for any file nor part of the clobbered_blocks),
286    # divide them into blocks that are all zero and blocks that aren't.
287    # (Zero blocks are handled specially because (1) there are usually
288    # a lot of them and (2) bsdiff handles files with long sequences of
289    # repeated bytes especially poorly.)
290
291    zero_blocks = []
292    nonzero_blocks = []
293    reference = '\0' * self.blocksize
294
295    # Workaround for bug 23227672. For squashfs, we don't have a system.map. So
296    # the whole system image will be treated as a single file. But for some
297    # unknown bug, the updater will be killed due to OOM when writing back the
298    # patched image to flash (observed on lenok-userdebug MEA49). Prior to
299    # getting a real fix, we evenly divide the non-zero blocks into smaller
300    # groups (currently 1024 blocks or 4MB per group).
301    # Bug: 23227672
302    MAX_BLOCKS_PER_GROUP = 1024
303    nonzero_groups = []
304
305    f = self.simg_f
306    for s, e in remaining:
307      for b in range(s, e):
308        idx = bisect.bisect_right(self.offset_index, b) - 1
309        chunk_start, _, filepos, fill_data = self.offset_map[idx]
310        if filepos is not None:
311          filepos += (b-chunk_start) * self.blocksize
312          f.seek(filepos, os.SEEK_SET)
313          data = f.read(self.blocksize)
314        else:
315          if fill_data == reference[:4]:   # fill with all zeros
316            data = reference
317          else:
318            data = None
319
320        if data == reference:
321          zero_blocks.append(b)
322          zero_blocks.append(b+1)
323        else:
324          nonzero_blocks.append(b)
325          nonzero_blocks.append(b+1)
326
327          if len(nonzero_blocks) >= MAX_BLOCKS_PER_GROUP:
328            nonzero_groups.append(nonzero_blocks)
329            # Clear the list.
330            nonzero_blocks = []
331
332    if nonzero_blocks:
333      nonzero_groups.append(nonzero_blocks)
334      nonzero_blocks = []
335
336    assert zero_blocks or nonzero_groups or clobbered_blocks
337
338    if zero_blocks:
339      out["__ZERO"] = rangelib.RangeSet(data=zero_blocks)
340    if nonzero_groups:
341      for i, blocks in enumerate(nonzero_groups):
342        out["__NONZERO-%d" % i] = rangelib.RangeSet(data=blocks)
343    if clobbered_blocks:
344      out["__COPY"] = clobbered_blocks
345    if self.hashtree_info:
346      out["__HASHTREE"] = self.hashtree_info.hashtree_range
347
348  def ResetFileMap(self):
349    """Throw away the file map and treat the entire image as
350    undifferentiated data."""
351    self.file_map = {"__DATA": self.care_map}
352
353
354def GetImagePartitionSize(img):
355  try:
356    simg = SparseImage(img, build_map=False)
357    return simg.blocksize * simg.total_blocks
358  except ValueError:
359    return os.path.getsize(img)
360
361
362if __name__ == '__main__':
363  parser = argparse.ArgumentParser()
364  parser.add_argument('image')
365  parser.add_argument('--get_partition_size', action='store_true',
366                      help='Return partition size of the image')
367  args = parser.parse_args()
368  if args.get_partition_size:
369    print(GetImagePartitionSize(args.image))
370