1#!/usr/bin/python
2
3# Copyright 2016, The Android Open Source Project
4#
5# Permission is hereby granted, free of charge, to any person
6# obtaining a copy of this software and associated documentation
7# files (the "Software"), to deal in the Software without
8# restriction, including without limitation the rights to use, copy,
9# modify, merge, publish, distribute, sublicense, and/or sell copies
10# of the Software, and to permit persons to whom the Software is
11# furnished to do so, subject to the following conditions:
12#
13# The above copyright notice and this permission notice shall be
14# included in all copies or substantial portions of the Software.
15#
16# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
17# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
18# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
19# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
20# BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
21# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
22# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
23# SOFTWARE.
24
25
26"""Unit-test for ImageHandler."""
27
28
29import imp
30import os
31import sys
32import tempfile
33import unittest
34
35sys.dont_write_bytecode = True
36avbtool = imp.load_source('avbtool', './avbtool')
37
38# The file test_file.bin and test_file.bin.sparse are generated using
39# the following python code:
40#
41#  with open('test_file.bin', 'w+b') as f:
42#    f.write('Barfoo43'*128*12)
43#  os.system('img2simg test_file.bin test_file.bin.sparse')
44#  image = avbtool.ImageHandler('test_file.bin.sparse')
45#  image.append_dont_care(12*1024)
46#  image.append_fill('\x01\x02\x03\x04', 12*1024)
47#  image.append_raw('Foobar42'*128*12)
48#  image.append_dont_care(12*1024)
49#  del image
50#  os.system('rm -f test_file.bin')
51#  os.system('simg2img test_file.bin.sparse test_file.bin')
52#
53# and manually verified to be correct. The content of the raw and
54# sparse files are as follows (the line with "Fill with 0x04030201" is
55# a simg_dump.py bug):
56#
57# $ hexdump -C test_file.bin
58# 00000000  42 61 72 66 6f 6f 34 33  42 61 72 66 6f 6f 34 33  |Barfoo43Barfoo43|
59# *
60# 00003000  00 00 00 00 00 00 00 00  00 00 00 00 00 00 00 00  |................|
61# *
62# 00006000  01 02 03 04 01 02 03 04  01 02 03 04 01 02 03 04  |................|
63# *
64# 00009000  46 6f 6f 62 61 72 34 32  46 6f 6f 62 61 72 34 32  |Foobar42Foobar42|
65# *
66# 0000c000  00 00 00 00 00 00 00 00  00 00 00 00 00 00 00 00  |................|
67# *
68# 0000f000
69#
70# $ system/core/libsparse/simg_dump.py -v test_file.bin.sparse
71# test_file.bin.sparse: Total of 15 4096-byte output blocks in 5 input chunks.
72#             input_bytes      output_blocks
73# chunk    offset     number  offset  number
74#    1         40      12288       0       3 Raw data
75#    2      12340          0       3       3 Don't care
76#    3      12352          4       6       3 Fill with 0x04030201
77#    4      12368      12288       9       3 Raw data
78#    5      24668          0      12       3 Don't care
79#           24668                 15         End
80#
81
82
83class ImageHandler(unittest.TestCase):
84
85  TEST_FILE_SPARSE_PATH = 'test/data/test_file.bin.sparse'
86  TEST_FILE_PATH = 'test/data/test_file.bin'
87  TEST_FILE_SIZE = 61440
88  TEST_FILE_BLOCK_SIZE = 4096
89
90  def _file_contents_equal(self, path1, path2, size):
91    f1 = open(path1, 'r')
92    f2 = open(path2, 'r')
93    if f1.read(size) != f2.read(size):
94      return False
95    return True
96
97  def _file_size(self, f):
98    old_pos = f.tell()
99    f.seek(0, os.SEEK_END)
100    size = f.tell()
101    f.seek(old_pos)
102    return size
103
104  def _clone_sparse_file(self):
105    f = tempfile.NamedTemporaryFile()
106    f.write(open(self.TEST_FILE_SPARSE_PATH).read())
107    f.flush()
108    return f
109
110  def _unsparsify(self, path):
111    f = tempfile.NamedTemporaryFile()
112    os.system('simg2img {} {}'.format(path, f.name))
113    return f
114
115  def testRead(self):
116    """Checks that reading from a sparse file works as intended."""
117    ih = avbtool.ImageHandler(self.TEST_FILE_SPARSE_PATH)
118
119    # Check that we start at offset 0.
120    self.assertEqual(ih.tell(), 0)
121
122    # Check that reading advances the cursor.
123    self.assertEqual(ih.read(14), bytearray('Barfoo43Barfoo'))
124    self.assertEqual(ih.tell(), 14)
125    self.assertEqual(ih.read(2), bytearray('43'))
126    self.assertEqual(ih.tell(), 16)
127
128    # Check reading in the middle of a fill chunk gets the right data.
129    ih.seek(0x6000 + 1)
130    self.assertEqual(ih.read(4), bytearray('\x02\x03\x04\x01'))
131
132    # Check we can cross the chunk boundary correctly.
133    ih.seek(0x3000 - 10)
134    self.assertEqual(ih.read(12), bytearray('43Barfoo43\x00\x00'))
135    ih.seek(0x9000 - 3)
136    self.assertEqual(ih.read(5), bytearray('\x02\x03\x04Fo'))
137
138    # Check reading at end of file is a partial read.
139    ih.seek(0xf000 - 2)
140    self.assertEqual(ih.read(16), bytearray('\x00\x00'))
141
142  def testTruncate(self):
143    """Checks that we can truncate a sparse file correctly."""
144    # Check truncation at all possible boundaries (including start and end).
145    for size in range(0, self.TEST_FILE_SIZE + self.TEST_FILE_BLOCK_SIZE,
146                      self.TEST_FILE_BLOCK_SIZE):
147      sparse_file = self._clone_sparse_file()
148      ih = avbtool.ImageHandler(sparse_file.name)
149      ih.truncate(size)
150      unsparse_file = self._unsparsify(sparse_file.name)
151      self.assertEqual(self._file_size(unsparse_file), size)
152      self.assertTrue(self._file_contents_equal(unsparse_file.name,
153                                                self.TEST_FILE_PATH,
154                                                size))
155
156    # Check truncation to grow the file.
157    grow_size = 8192
158    sparse_file = self._clone_sparse_file()
159    ih = avbtool.ImageHandler(sparse_file.name)
160    ih.truncate(self.TEST_FILE_SIZE + grow_size)
161    unsparse_file = self._unsparsify(sparse_file.name)
162    self.assertEqual(self._file_size(unsparse_file),
163                     self.TEST_FILE_SIZE + grow_size)
164    self.assertTrue(self._file_contents_equal(unsparse_file.name,
165                                              self.TEST_FILE_PATH,
166                                              self.TEST_FILE_SIZE))
167    unsparse_file.seek(self.TEST_FILE_SIZE)
168    self.assertEqual(unsparse_file.read(), '\0'*grow_size)
169
170  def testAppendRaw(self):
171    """Checks that we can append raw data correctly."""
172    sparse_file = self._clone_sparse_file()
173    ih = avbtool.ImageHandler(sparse_file.name)
174    data = 'SomeData'*4096
175    ih.append_raw(data)
176    unsparse_file = self._unsparsify(sparse_file.name)
177    self.assertTrue(self._file_contents_equal(unsparse_file.name,
178                                              self.TEST_FILE_PATH,
179                                              self.TEST_FILE_SIZE))
180    unsparse_file.seek(self.TEST_FILE_SIZE)
181    self.assertEqual(unsparse_file.read(), data)
182
183  def testAppendFill(self):
184    """Checks that we can append fill data correctly."""
185    sparse_file = self._clone_sparse_file()
186    ih = avbtool.ImageHandler(sparse_file.name)
187    data = 'ABCD'*4096
188    ih.append_fill('ABCD', len(data))
189    unsparse_file = self._unsparsify(sparse_file.name)
190    self.assertTrue(self._file_contents_equal(unsparse_file.name,
191                                              self.TEST_FILE_PATH,
192                                              self.TEST_FILE_SIZE))
193    unsparse_file.seek(self.TEST_FILE_SIZE)
194    self.assertEqual(unsparse_file.read(), data)
195
196  def testDontCare(self):
197    """Checks that we can append DONT_CARE data correctly."""
198    sparse_file = self._clone_sparse_file()
199    ih = avbtool.ImageHandler(sparse_file.name)
200    data = '\0'*40960
201    ih.append_dont_care(len(data))
202    unsparse_file = self._unsparsify(sparse_file.name)
203    self.assertTrue(self._file_contents_equal(unsparse_file.name,
204                                              self.TEST_FILE_PATH,
205                                              self.TEST_FILE_SIZE))
206    unsparse_file.seek(self.TEST_FILE_SIZE)
207    self.assertEqual(unsparse_file.read(), data)
208
209
210if __name__ == '__main__':
211  unittest.main()
212