1#
2# Copyright (C) 2016 The Android Open Source Project
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#      http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15#
16
17import os
18from hashlib import sha1
19
20import common
21from blockimgdiff import BlockImageDiff, HeapItem, ImgdiffStats, Transfer
22from images import DataImage, EmptyImage, FileImage
23from rangelib import RangeSet
24from test_utils import ReleaseToolsTestCase
25
26
27class HealpItemTest(ReleaseToolsTestCase):
28
29  class Item(object):
30    def __init__(self, score):
31      self.score = score
32
33  def test_init(self):
34    item1 = HeapItem(self.Item(15))
35    item2 = HeapItem(self.Item(20))
36    item3 = HeapItem(self.Item(15))
37    self.assertTrue(item1)
38    self.assertTrue(item2)
39    self.assertTrue(item3)
40
41    self.assertNotEqual(item1, item2)
42    self.assertEqual(item1, item3)
43    # HeapItem uses negated scores.
44    self.assertGreater(item1, item2)
45    self.assertLessEqual(item1, item3)
46    self.assertTrue(item1 <= item3)
47    self.assertFalse(item2 >= item1)
48
49  def test_clear(self):
50    item = HeapItem(self.Item(15))
51    self.assertTrue(item)
52
53    item.clear()
54    self.assertFalse(item)
55
56
57class BlockImageDiffTest(ReleaseToolsTestCase):
58
59  def test_GenerateDigraphOrder(self):
60    """Make sure GenerateDigraph preserves the order.
61
62    t0: <0-5> => <...>
63    t1: <0-7> => <...>
64    t2: <0-4> => <...>
65    t3: <...> => <0-10>
66
67    t0, t1 and t2 must go before t3, i.e. t3.goes_after =
68    { t0:..., t1:..., t2:... }. But the order of t0-t2 must be preserved.
69    """
70
71    src = EmptyImage()
72    tgt = EmptyImage()
73    block_image_diff = BlockImageDiff(tgt, src)
74
75    transfers = block_image_diff.transfers
76    t0 = Transfer("t1", "t1", RangeSet("10-15"), RangeSet("0-5"), "t1hash",
77                  "t1hash", "move", transfers)
78    t1 = Transfer("t2", "t2", RangeSet("20-25"), RangeSet("0-7"), "t2hash",
79                  "t2hash", "move", transfers)
80    t2 = Transfer("t3", "t3", RangeSet("30-35"), RangeSet("0-4"), "t3hash",
81                  "t3hash", "move", transfers)
82    t3 = Transfer("t4", "t4", RangeSet("0-10"), RangeSet("40-50"), "t4hash",
83                  "t4hash", "move", transfers)
84
85    block_image_diff.GenerateDigraph()
86    t3_goes_after_copy = t3.goes_after.copy()
87
88    # Elements in the set must be in the transfer evaluation order.
89    elements = list(t3_goes_after_copy)
90    self.assertEqual(t0, elements[0])
91    self.assertEqual(t1, elements[1])
92    self.assertEqual(t2, elements[2])
93
94    # Now switch the order of t0, t1 and t2.
95    transfers[0], transfers[1], transfers[2] = (
96        transfers[2], transfers[0], transfers[1])
97    t3.goes_after.clear()
98    t3.goes_before.clear()
99    block_image_diff.GenerateDigraph()
100
101    # The goes_after must be different from last run.
102    self.assertNotEqual(t3_goes_after_copy, t3.goes_after)
103
104    # Assert that each element must agree with the transfer order.
105    elements = list(t3.goes_after)
106    self.assertEqual(t2, elements[0])
107    self.assertEqual(t0, elements[1])
108    self.assertEqual(t1, elements[2])
109
110  def test_ReviseStashSize(self):
111    """ReviseStashSize should convert transfers to 'new' commands as needed.
112
113    t1: diff <20-29> => <11-15>
114    t2: diff <11-15> => <20-29>
115    """
116
117    src = EmptyImage()
118    tgt = EmptyImage()
119    block_image_diff = BlockImageDiff(tgt, src, version=3)
120
121    transfers = block_image_diff.transfers
122    Transfer("t1", "t1", RangeSet("11-15"), RangeSet("20-29"), "t1hash",
123             "t1hash", "diff", transfers)
124    Transfer("t2", "t2", RangeSet("20-29"), RangeSet("11-15"), "t2hash",
125             "t2hash", "diff", transfers)
126
127    block_image_diff.GenerateDigraph()
128    block_image_diff.FindVertexSequence()
129    block_image_diff.ReverseBackwardEdges()
130
131    # Sufficient cache to stash 5 blocks (size * 0.8 >= 5).
132    common.OPTIONS.cache_size = 7 * 4096
133    self.assertEqual((0, 5), block_image_diff.ReviseStashSize())
134
135    # Insufficient cache to stash 5 blocks (size * 0.8 < 5).
136    common.OPTIONS.cache_size = 6 * 4096
137    self.assertEqual((10, 0), block_image_diff.ReviseStashSize())
138
139  def test_ReviseStashSize_bug_33687949(self):
140    """ReviseStashSize() should "free" the used stash _after_ the command.
141
142    t1: diff <1-5> => <11-15>
143    t2: diff <11-15> => <21-25>
144    t3: diff <11-15 30-39> => <1-5 30-39>
145
146    For transfer t3, the used stash "11-15" should not be freed until the
147    command finishes. Assume the allowed cache size is 12-block, it should
148    convert the command to 'new' due to insufficient cache (12 < 5 + 10).
149    """
150
151    src = EmptyImage()
152    tgt = EmptyImage()
153    block_image_diff = BlockImageDiff(tgt, src, version=3)
154
155    transfers = block_image_diff.transfers
156    t1 = Transfer("t1", "t1", RangeSet("11-15"), RangeSet("1-5"), "t1hash",
157                  "t1hash", "diff", transfers)
158    t2 = Transfer("t2", "t2", RangeSet("21-25"), RangeSet("11-15"), "t2hash",
159                  "t2hash", "diff", transfers)
160    t3 = Transfer("t3", "t3", RangeSet("1-5 30-39"), RangeSet("11-15 30-39"),
161                  "t3hash", "t3hash", "diff", transfers)
162
163    block_image_diff.GenerateDigraph()
164
165    # Instead of calling FindVertexSequence() and ReverseBackwardEdges(), we
166    # just set up the stash_before and use_stash manually. Otherwise it will
167    # reorder the transfer, which makes testing ReviseStashSize() harder.
168    t1.stash_before.append((0, RangeSet("11-15")))
169    t2.use_stash.append((0, RangeSet("11-15")))
170    t1.stash_before.append((1, RangeSet("11-15")))
171    t3.use_stash.append((1, RangeSet("11-15")))
172
173    # Insufficient cache to stash 15 blocks (size * 0.8 < 15).
174    common.OPTIONS.cache_size = 15 * 4096
175    self.assertEqual((15, 5), block_image_diff.ReviseStashSize())
176
177  def test_FileTypeSupportedByImgdiff(self):
178    self.assertTrue(
179        BlockImageDiff.FileTypeSupportedByImgdiff(
180            "/system/priv-app/Settings/Settings.apk"))
181    self.assertTrue(
182        BlockImageDiff.FileTypeSupportedByImgdiff(
183            "/system/framework/am.jar"))
184    self.assertTrue(
185        BlockImageDiff.FileTypeSupportedByImgdiff(
186            "/system/etc/security/otacerts.zip"))
187
188    self.assertFalse(
189        BlockImageDiff.FileTypeSupportedByImgdiff(
190            "/system/framework/arm/boot.oat"))
191    self.assertFalse(
192        BlockImageDiff.FileTypeSupportedByImgdiff(
193            "/system/priv-app/notanapk"))
194
195  def test_CanUseImgdiff(self):
196    block_image_diff = BlockImageDiff(EmptyImage(), EmptyImage())
197    self.assertTrue(
198        block_image_diff.CanUseImgdiff(
199            "/system/app/app1.apk", RangeSet("10-15"), RangeSet("0-5")))
200    self.assertTrue(
201        block_image_diff.CanUseImgdiff(
202            "/vendor/app/app2.apk", RangeSet("20 25"), RangeSet("30-31"), True))
203
204    self.assertDictEqual(
205        {
206            ImgdiffStats.USED_IMGDIFF: {"/system/app/app1.apk"},
207            ImgdiffStats.USED_IMGDIFF_LARGE_APK: {"/vendor/app/app2.apk"},
208        },
209        block_image_diff.imgdiff_stats.stats)
210
211
212  def test_CanUseImgdiff_ineligible(self):
213    # Disabled by caller.
214    block_image_diff = BlockImageDiff(EmptyImage(), EmptyImage(),
215                                      disable_imgdiff=True)
216    self.assertFalse(
217        block_image_diff.CanUseImgdiff(
218            "/system/app/app1.apk", RangeSet("10-15"), RangeSet("0-5")))
219
220    # Unsupported file type.
221    block_image_diff = BlockImageDiff(EmptyImage(), EmptyImage())
222    self.assertFalse(
223        block_image_diff.CanUseImgdiff(
224            "/system/bin/gzip", RangeSet("10-15"), RangeSet("0-5")))
225
226    # At least one of the ranges is in non-monotonic order.
227    self.assertFalse(
228        block_image_diff.CanUseImgdiff(
229            "/system/app/app2.apk", RangeSet("10-15"),
230            RangeSet("15-20 30 10-14")))
231
232    # At least one of the ranges is incomplete.
233    src_ranges = RangeSet("0-5")
234    src_ranges.extra['incomplete'] = True
235    self.assertFalse(
236        block_image_diff.CanUseImgdiff(
237            "/vendor/app/app4.apk", RangeSet("10-15"), src_ranges))
238
239    # The stats are correctly logged.
240    self.assertDictEqual(
241        {
242            ImgdiffStats.SKIPPED_NONMONOTONIC: {'/system/app/app2.apk'},
243            ImgdiffStats.SKIPPED_INCOMPLETE: {'/vendor/app/app4.apk'},
244        },
245        block_image_diff.imgdiff_stats.stats)
246
247
248class ImgdiffStatsTest(ReleaseToolsTestCase):
249
250  def test_Log(self):
251    imgdiff_stats = ImgdiffStats()
252    imgdiff_stats.Log("/system/app/app2.apk", ImgdiffStats.USED_IMGDIFF)
253    self.assertDictEqual(
254        {
255            ImgdiffStats.USED_IMGDIFF: {'/system/app/app2.apk'},
256        },
257        imgdiff_stats.stats)
258
259  def test_Log_invalidInputs(self):
260    imgdiff_stats = ImgdiffStats()
261
262    self.assertRaises(AssertionError, imgdiff_stats.Log, "/system/bin/gzip",
263                      ImgdiffStats.USED_IMGDIFF)
264
265    self.assertRaises(AssertionError, imgdiff_stats.Log, "/system/app/app1.apk",
266                      "invalid reason")
267
268
269class DataImageTest(ReleaseToolsTestCase):
270
271  def test_read_range_set(self):
272    data = "file" + ('\0' * 4092)
273    image = DataImage(data)
274    self.assertEqual(data, "".join(image.ReadRangeSet(image.care_map)))
275
276
277class FileImageTest(ReleaseToolsTestCase):
278
279  def setUp(self):
280    self.file_path = common.MakeTempFile()
281    self.data = os.urandom(4096 * 4)
282    with open(self.file_path, 'wb') as f:
283      f.write(self.data)
284    self.file = FileImage(self.file_path)
285
286  def test_totalsha1(self):
287    self.assertEqual(sha1(self.data).hexdigest(), self.file.TotalSha1())
288
289  def test_ranges(self):
290    blocksize = self.file.blocksize
291    for s in range(4):
292      for e in range(s, 4):
293        expected_data = self.data[s * blocksize : e * blocksize]
294
295        rs = RangeSet([s, e])
296        data = b''.join(self.file.ReadRangeSet(rs))
297        self.assertEqual(expected_data, data)
298
299        sha1sum = self.file.RangeSha1(rs)
300        self.assertEqual(sha1(expected_data).hexdigest(), sha1sum)
301
302        tmpfile = common.MakeTempFile()
303        with open(tmpfile, 'wb') as f:
304          self.file.WriteRangeDataToFd(rs, f)
305        with open(tmpfile, 'rb') as f:
306          self.assertEqual(expected_data, f.read())
307
308  def test_read_all(self):
309    data = b''.join(self.file.ReadRangeSet(self.file.care_map))
310    self.assertEqual(self.data, data)
311