1#
2# Copyright (C) 2018 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
17"""Unittests for verity_utils.py."""
18
19import copy
20import math
21import os.path
22import random
23
24import common
25import sparse_img
26from rangelib import RangeSet
27from test_utils import (
28    get_testdata_dir, ReleaseToolsTestCase, SkipIfExternalToolsUnavailable)
29from verity_utils import (
30    CreateHashtreeInfoGenerator, CreateVerityImageBuilder, HashtreeInfo,
31    VerifiedBootVersion1HashtreeInfoGenerator)
32
33BLOCK_SIZE = common.BLOCK_SIZE
34
35
36class VerifiedBootVersion1HashtreeInfoGeneratorTest(ReleaseToolsTestCase):
37
38  def setUp(self):
39    self.testdata_dir = get_testdata_dir()
40
41    self.partition_size = 1024 * 1024
42    self.prop_dict = {
43        'verity': 'true',
44        'verity_fec': 'true',
45        'system_verity_block_device': '/dev/block/system',
46        'system_size': self.partition_size
47    }
48
49    self.hash_algorithm = "sha256"
50    self.fixed_salt = (
51        "aee087a5be3b982978c923f566a94613496b417f2af592639bc80d141e34dfe7")
52    self.expected_root_hash = (
53        "0b7c4565e87b1026e11fbab91c0bc29e185c847a5b44d40e6e86e461e8adf80d")
54
55  def _CreateSimg(self, raw_data):  # pylint: disable=no-self-use
56    output_file = common.MakeTempFile()
57    raw_image = common.MakeTempFile()
58    with open(raw_image, 'wb') as f:
59      f.write(raw_data)
60
61    cmd = ["img2simg", raw_image, output_file, '4096']
62    common.RunAndCheckOutput(cmd)
63    return output_file
64
65  def _GenerateImage(self):
66    partition_size = 1024 * 1024
67    prop_dict = {
68        'partition_size': str(partition_size),
69        'verity': 'true',
70        'verity_block_device': '/dev/block/system',
71        'verity_key': os.path.join(self.testdata_dir, 'testkey'),
72        'verity_fec': 'true',
73        'verity_signer_cmd': 'verity_signer',
74    }
75    verity_image_builder = CreateVerityImageBuilder(prop_dict)
76    self.assertIsNotNone(verity_image_builder)
77    adjusted_size = verity_image_builder.CalculateMaxImageSize()
78
79    raw_image = bytearray(adjusted_size)
80    for i in range(adjusted_size):
81      raw_image[i] = ord('0') + i % 10
82
83    output_file = self._CreateSimg(raw_image)
84
85    # Append the verity metadata.
86    verity_image_builder.Build(output_file)
87
88    return output_file
89
90  @SkipIfExternalToolsUnavailable()
91  def test_CreateHashtreeInfoGenerator(self):
92    image_file = sparse_img.SparseImage(self._GenerateImage())
93
94    generator = CreateHashtreeInfoGenerator(
95        'system', image_file, self.prop_dict)
96    self.assertEqual(
97        VerifiedBootVersion1HashtreeInfoGenerator, type(generator))
98    self.assertEqual(self.partition_size, generator.partition_size)
99    self.assertTrue(generator.fec_supported)
100
101  @SkipIfExternalToolsUnavailable()
102  def test_DecomposeSparseImage(self):
103    image_file = sparse_img.SparseImage(self._GenerateImage())
104
105    generator = VerifiedBootVersion1HashtreeInfoGenerator(
106        self.partition_size, 4096, True)
107    generator.DecomposeSparseImage(image_file)
108    self.assertEqual(991232, generator.filesystem_size)
109    self.assertEqual(12288, generator.hashtree_size)
110    self.assertEqual(32768, generator.metadata_size)
111
112  @SkipIfExternalToolsUnavailable()
113  def test_ParseHashtreeMetadata(self):
114    image_file = sparse_img.SparseImage(self._GenerateImage())
115    generator = VerifiedBootVersion1HashtreeInfoGenerator(
116        self.partition_size, 4096, True)
117    generator.DecomposeSparseImage(image_file)
118
119    # pylint: disable=protected-access
120    generator._ParseHashtreeMetadata()
121
122    self.assertEqual(
123        self.hash_algorithm, generator.hashtree_info.hash_algorithm)
124    self.assertEqual(self.fixed_salt, generator.hashtree_info.salt)
125    self.assertEqual(self.expected_root_hash, generator.hashtree_info.root_hash)
126
127  @SkipIfExternalToolsUnavailable()
128  def test_ValidateHashtree_smoke(self):
129    generator = VerifiedBootVersion1HashtreeInfoGenerator(
130        self.partition_size, 4096, True)
131    generator.image = sparse_img.SparseImage(self._GenerateImage())
132
133    generator.hashtree_info = info = HashtreeInfo()
134    info.filesystem_range = RangeSet(data=[0, 991232 // 4096])
135    info.hashtree_range = RangeSet(
136        data=[991232 // 4096, (991232 + 12288) // 4096])
137    info.hash_algorithm = self.hash_algorithm
138    info.salt = self.fixed_salt
139    info.root_hash = self.expected_root_hash
140
141    self.assertTrue(generator.ValidateHashtree())
142
143  @SkipIfExternalToolsUnavailable()
144  def test_ValidateHashtree_failure(self):
145    generator = VerifiedBootVersion1HashtreeInfoGenerator(
146        self.partition_size, 4096, True)
147    generator.image = sparse_img.SparseImage(self._GenerateImage())
148
149    generator.hashtree_info = info = HashtreeInfo()
150    info.filesystem_range = RangeSet(data=[0, 991232 // 4096])
151    info.hashtree_range = RangeSet(
152        data=[991232 // 4096, (991232 + 12288) // 4096])
153    info.hash_algorithm = self.hash_algorithm
154    info.salt = self.fixed_salt
155    info.root_hash = "a" + self.expected_root_hash[1:]
156
157    self.assertFalse(generator.ValidateHashtree())
158
159  @SkipIfExternalToolsUnavailable()
160  def test_Generate(self):
161    image_file = sparse_img.SparseImage(self._GenerateImage())
162    generator = CreateHashtreeInfoGenerator('system', 4096, self.prop_dict)
163    info = generator.Generate(image_file)
164
165    self.assertEqual(RangeSet(data=[0, 991232 // 4096]), info.filesystem_range)
166    self.assertEqual(RangeSet(data=[991232 // 4096, (991232 + 12288) // 4096]),
167                     info.hashtree_range)
168    self.assertEqual(self.hash_algorithm, info.hash_algorithm)
169    self.assertEqual(self.fixed_salt, info.salt)
170    self.assertEqual(self.expected_root_hash, info.root_hash)
171
172
173class VerifiedBootVersion1VerityImageBuilderTest(ReleaseToolsTestCase):
174
175  DEFAULT_PARTITION_SIZE = 4096 * 1024
176  DEFAULT_PROP_DICT = {
177      'partition_size': str(DEFAULT_PARTITION_SIZE),
178      'verity': 'true',
179      'verity_block_device': '/dev/block/system',
180      'verity_key': os.path.join(get_testdata_dir(), 'testkey'),
181      'verity_fec': 'true',
182      'verity_signer_cmd': 'verity_signer',
183  }
184
185  def test_init(self):
186    prop_dict = copy.deepcopy(self.DEFAULT_PROP_DICT)
187    verity_image_builder = CreateVerityImageBuilder(prop_dict)
188    self.assertIsNotNone(verity_image_builder)
189    self.assertEqual(1, verity_image_builder.version)
190
191  def test_init_MissingProps(self):
192    prop_dict = copy.deepcopy(self.DEFAULT_PROP_DICT)
193    del prop_dict['verity']
194    self.assertIsNone(CreateVerityImageBuilder(prop_dict))
195
196    prop_dict = copy.deepcopy(self.DEFAULT_PROP_DICT)
197    del prop_dict['verity_block_device']
198    self.assertIsNone(CreateVerityImageBuilder(prop_dict))
199
200  @SkipIfExternalToolsUnavailable()
201  def test_CalculateMaxImageSize(self):
202    verity_image_builder = CreateVerityImageBuilder(self.DEFAULT_PROP_DICT)
203    size = verity_image_builder.CalculateMaxImageSize()
204    self.assertLess(size, self.DEFAULT_PARTITION_SIZE)
205
206    # Same result by explicitly passing the partition size.
207    self.assertEqual(
208        verity_image_builder.CalculateMaxImageSize(),
209        verity_image_builder.CalculateMaxImageSize(
210            self.DEFAULT_PARTITION_SIZE))
211
212  @staticmethod
213  def _BuildAndVerify(prop, verify_key):
214    verity_image_builder = CreateVerityImageBuilder(prop)
215    image_size = verity_image_builder.CalculateMaxImageSize()
216
217    # Build the sparse image with verity metadata.
218    input_dir = common.MakeTempDir()
219    image = common.MakeTempFile(suffix='.img')
220    cmd = ['mkuserimg_mke2fs', input_dir, image, 'ext4', '/system',
221           str(image_size), '-j', '0', '-s']
222    common.RunAndCheckOutput(cmd)
223    verity_image_builder.Build(image)
224
225    # Verify the verity metadata.
226    cmd = ['verity_verifier', image, '-mincrypt', verify_key]
227    common.RunAndCheckOutput(cmd)
228
229  @SkipIfExternalToolsUnavailable()
230  def test_Build(self):
231    self._BuildAndVerify(
232        self.DEFAULT_PROP_DICT,
233        os.path.join(get_testdata_dir(), 'testkey_mincrypt'))
234
235  @SkipIfExternalToolsUnavailable()
236  def test_Build_ValidationCheck(self):
237    # A validity check for the test itself: the image shouldn't be verifiable
238    # with wrong key.
239    self.assertRaises(
240        common.ExternalError,
241        self._BuildAndVerify,
242        self.DEFAULT_PROP_DICT,
243        os.path.join(get_testdata_dir(), 'verity_mincrypt'))
244
245  @SkipIfExternalToolsUnavailable()
246  def test_Build_FecDisabled(self):
247    prop_dict = copy.deepcopy(self.DEFAULT_PROP_DICT)
248    del prop_dict['verity_fec']
249    self._BuildAndVerify(
250        prop_dict,
251        os.path.join(get_testdata_dir(), 'testkey_mincrypt'))
252
253  @SkipIfExternalToolsUnavailable()
254  def test_Build_SquashFs(self):
255    verity_image_builder = CreateVerityImageBuilder(self.DEFAULT_PROP_DICT)
256    verity_image_builder.CalculateMaxImageSize()
257
258    # Build the sparse image with verity metadata.
259    input_dir = common.MakeTempDir()
260    image = common.MakeTempFile(suffix='.img')
261    cmd = ['mksquashfsimage.sh', input_dir, image, '-s']
262    common.RunAndCheckOutput(cmd)
263    verity_image_builder.PadSparseImage(image)
264    verity_image_builder.Build(image)
265
266    # Verify the verity metadata.
267    cmd = ["verity_verifier", image, '-mincrypt',
268           os.path.join(get_testdata_dir(), 'testkey_mincrypt')]
269    common.RunAndCheckOutput(cmd)
270
271
272class VerifiedBootVersion2VerityImageBuilderTest(ReleaseToolsTestCase):
273
274  DEFAULT_PROP_DICT = {
275      'partition_size': str(4096 * 1024),
276      'partition_name': 'system',
277      'avb_avbtool': 'avbtool',
278      'avb_hashtree_enable': 'true',
279      'avb_add_hashtree_footer_args': '',
280  }
281
282  def test_init(self):
283    prop_dict = copy.deepcopy(self.DEFAULT_PROP_DICT)
284    verity_image_builder = CreateVerityImageBuilder(prop_dict)
285    self.assertIsNotNone(verity_image_builder)
286    self.assertEqual(2, verity_image_builder.version)
287
288  def test_init_MissingProps(self):
289    prop_dict = copy.deepcopy(self.DEFAULT_PROP_DICT)
290    del prop_dict['avb_hashtree_enable']
291    verity_image_builder = CreateVerityImageBuilder(prop_dict)
292    self.assertIsNone(verity_image_builder)
293
294  @SkipIfExternalToolsUnavailable()
295  def test_Build(self):
296    prop_dict = copy.deepcopy(self.DEFAULT_PROP_DICT)
297    verity_image_builder = CreateVerityImageBuilder(prop_dict)
298    self.assertIsNotNone(verity_image_builder)
299    self.assertEqual(2, verity_image_builder.version)
300
301    input_dir = common.MakeTempDir()
302    image_dir = common.MakeTempDir()
303    system_image = os.path.join(image_dir, 'system.img')
304    system_image_size = verity_image_builder.CalculateMaxImageSize()
305    cmd = ['mkuserimg_mke2fs', input_dir, system_image, 'ext4', '/system',
306           str(system_image_size), '-j', '0', '-s']
307    common.RunAndCheckOutput(cmd)
308    verity_image_builder.Build(system_image)
309
310    # Additionally make vbmeta image so that we can verify with avbtool.
311    vbmeta_image = os.path.join(image_dir, 'vbmeta.img')
312    cmd = ['avbtool', 'make_vbmeta_image', '--include_descriptors_from_image',
313           system_image, '--output', vbmeta_image]
314    common.RunAndCheckOutput(cmd)
315
316    # Verify the verity metadata.
317    cmd = ['avbtool', 'verify_image', '--image', vbmeta_image]
318    common.RunAndCheckOutput(cmd)
319
320  def _test_CalculateMinPartitionSize_SetUp(self):
321    # To test CalculateMinPartitionSize(), by using 200MB to 2GB image size.
322    #   -  51200 = 200MB * 1024 * 1024 / 4096
323    #   - 524288 = 2GB * 1024 * 1024 * 1024 / 4096
324    image_sizes = [BLOCK_SIZE * random.randint(51200, 524288) + offset
325                   for offset in range(BLOCK_SIZE)]
326
327    prop_dict = {
328        'partition_size': None,
329        'partition_name': 'system',
330        'avb_avbtool': 'avbtool',
331        'avb_hashtree_enable': 'true',
332        'avb_add_hashtree_footer_args': None,
333    }
334    builder = CreateVerityImageBuilder(prop_dict)
335    self.assertEqual(2, builder.version)
336    return image_sizes, builder
337
338  def test_CalculateMinPartitionSize_LinearFooterSize(self):
339    """Tests with footer size which is linear to partition size."""
340    image_sizes, builder = self._test_CalculateMinPartitionSize_SetUp()
341    for image_size in image_sizes:
342      for ratio in 0.95, 0.56, 0.22:
343        expected_size = common.RoundUpTo4K(int(math.ceil(image_size / ratio)))
344        self.assertEqual(
345            expected_size,
346            builder.CalculateMinPartitionSize(
347                image_size, lambda x, ratio=ratio: int(x * ratio)))
348
349  def test_AVBCalcMinPartitionSize_SlowerGrowthFooterSize(self):
350    """Tests with footer size which grows slower than partition size."""
351
352    def _SizeCalculator(partition_size):
353      """Footer size is the power of 0.95 of partition size."""
354      # Minus footer size to return max image size.
355      return partition_size - int(math.pow(partition_size, 0.95))
356
357    image_sizes, builder = self._test_CalculateMinPartitionSize_SetUp()
358    for image_size in image_sizes:
359      min_partition_size = builder.CalculateMinPartitionSize(
360          image_size, _SizeCalculator)
361      # Checks min_partition_size can accommodate image_size.
362      self.assertGreaterEqual(
363          _SizeCalculator(min_partition_size),
364          image_size)
365      # Checks min_partition_size (round to BLOCK_SIZE) is the minimum.
366      self.assertLess(
367          _SizeCalculator(min_partition_size - BLOCK_SIZE),
368          image_size)
369
370  def test_CalculateMinPartitionSize_FasterGrowthFooterSize(self):
371    """Tests with footer size which grows faster than partition size."""
372
373    def _SizeCalculator(partition_size):
374      """Max image size is the power of 0.95 of partition size."""
375      # Max image size grows less than partition size, which means
376      # footer size grows faster than partition size.
377      return int(math.pow(partition_size, 0.95))
378
379    image_sizes, builder = self._test_CalculateMinPartitionSize_SetUp()
380    for image_size in image_sizes:
381      min_partition_size = builder.CalculateMinPartitionSize(
382          image_size, _SizeCalculator)
383      # Checks min_partition_size can accommodate image_size.
384      self.assertGreaterEqual(
385          _SizeCalculator(min_partition_size),
386          image_size)
387      # Checks min_partition_size (round to BLOCK_SIZE) is the minimum.
388      self.assertLess(
389          _SizeCalculator(min_partition_size - BLOCK_SIZE),
390          image_size)
391