Added missing tests for the verify_vbmeta_image methods & fixed bugs. am: 443bf320de am: 20c425e04b

Change-Id: I2bcb3c9ab5006137b9b50ab67c35905fe3426918

9 files changed, 164 insertions, 19 deletions

diff --git a/Android.bp b/Android.bp
index f66848a..dfb1cc1 100644
--- a/Android.bp
+++ b/Android.bp
@@ -148,7 +148,7 @@ python_test_host {
         "aftl_proto",
     ],
     data: [
-        "test/data/testkey_rsa4096.pem",
+        "test/data/**/*.*",
     ],
     test_suites: ["general-tests"],
     version: {
diff --git a/aftltool b/aftltool
index 26db6de..258f55f 100755
--- a/aftltool
+++ b/aftltool
@@ -1014,6 +1014,9 @@ class AftlDescriptor(object):
     Returns:
       True if the calculated signature matches AftlIcpEntry's. False otherwise.
     """
+    if not self.icp_entries:
+      return False
+
     icp_verified = 0
     for icp_entry in self.icp_entries:
       verified = False
@@ -1184,8 +1187,13 @@ class Aftl(avbtool.Avb):
     except ValueError as e:
       sys.stderr.write('The image does not contain a valid VBMeta structure: '
                        '{}.\n'.format(e))
-      return None
-    (footer, header, _, _) = self._parse_image(image)
+      return None, None
+
+    try:
+      (footer, header, _, _) = self._parse_image(image)
+    except avbtool.AvbError as e:
+      sys.stderr.write('The image cannot be parsed: {}.\n'.format(e))
+      return None, None
 
     # Seeks for the start of the vbmeta image and calculates its size.
     offset = 0
@@ -1200,7 +1208,7 @@ class Aftl(avbtool.Avb):
       image.seek(offset)
     except RuntimeError as e:
       sys.stderr.write('Given vbmeta image offset is invalid: {}.\n'.format(e))
-      return None
+      return None, None
     return image.read(vbmeta_image_size), footer
 
   def get_aftl_descriptor(self, image_filename):
@@ -1214,6 +1222,8 @@ class Aftl(avbtool.Avb):
     """
     # Reads the vbmeta image bytes.
     vbmeta_image, _ = self.get_vbmeta_image(image_filename)
+    if not vbmeta_image:
+      return None
 
     try:
       image = avbtool.ImageHandler(image_filename)
@@ -1232,6 +1242,10 @@ class Aftl(avbtool.Avb):
 
     # Parses the header for the AftlDescriptor size.
     tmp_header_bytes = image.read(AftlIcpHeader.SIZE)
+    if not tmp_header_bytes or len(tmp_header_bytes) != AftlIcpHeader.SIZE:
+      sys.stderr.write('This image does not contain a AftlDescriptor.\n')
+      return None
+
     try:
       tmp_header = AftlIcpHeader(tmp_header_bytes)
     except AftlError as e:
diff --git a/aftltool_test.py b/aftltool_test.py
index 1e29c1e..0b32514 100755
--- a/aftltool_test.py
+++ b/aftltool_test.py
@@ -44,7 +44,7 @@ import proto.trillian_pb2
 # Workaround for b/149307145 in order to pick up the test data from the right
 # location independent where the script is called from.
 # TODO(b/149307145): Remove workaround once the referenced bug is fixed.
-TESTDATA_PATH = os.path.dirname(os.path.realpath(__file__))
+TEST_EXEC_PATH = os.path.dirname(os.path.realpath(__file__))
 
 
 class AftltoolTestCase(unittest.TestCase):
@@ -248,6 +248,20 @@ class AftltoolTestCase(unittest.TestCase):
 
     super(AftltoolTestCase, self).tearDown()
 
+  def get_testdata_path(self, relative_path):
+    """Retrieves the absolute path for testdata given the relative path.
+
+    Arguments:
+      relative_path: The relative path to the testdata in the testdata
+        directory.
+
+    Returns:
+      The absolute path to the testdata.
+    """
+    rel_path_parts = ['test', 'data']
+    rel_path_parts.extend(relative_path.split(os.path.sep))
+    return os.path.join(TEST_EXEC_PATH, *rel_path_parts)
+
 
 class AftltoolTest(AftltoolTestCase):
 
@@ -371,6 +385,33 @@ class AftlDescriptorTest(AftltoolTestCase):
     self.assertEqual(len(d.icp_entries), 2)
     self.assertTrue(d.is_valid())
 
+  def test_verify_vbmeta_image(self):
+    """Tests the verify_vbmeta_image method."""
+    # Valid vbmeta image without footer with 1 AftlDescriptor.
+    tool = aftltool.Aftl()
+    vbmeta_image, _ = tool.get_vbmeta_image(
+        self.get_testdata_path('aftltool/aftl_output_vbmeta_with_1_icp.img'))
+    desc = tool.get_aftl_descriptor(
+        self.get_testdata_path('aftltool/aftl_output_vbmeta_with_1_icp.img'))
+    self.assertTrue(desc.verify_vbmeta_image(
+        vbmeta_image, [self.get_testdata_path('aftltool/aftl_pubkey_1.pub')]))
+
+    # Valid vbmeta image without footer with 2 AftlDescriptor.
+    vbmeta_image, _ = tool.get_vbmeta_image(
+        self.get_testdata_path('aftltool/aftl_output_vbmeta_with_2_icp.img'))
+    desc = tool.get_aftl_descriptor(
+        self.get_testdata_path('aftltool/aftl_output_vbmeta_with_2_icp.img'))
+    self.assertTrue(desc.verify_vbmeta_image(
+        vbmeta_image, [self.get_testdata_path('aftltool/aftl_pubkey_1.pub')]))
+
+    # Valid vbmeta image but checked with the public key of another log.
+    self.assertFalse(desc.verify_vbmeta_image(
+        vbmeta_image, [self.get_testdata_path('aftltool/aftl_pubkey_2.pub')]))
+
+    # Valid vbmeta image but checked with invalid public key file.
+    self.assertFalse(desc.verify_vbmeta_image(
+        vbmeta_image, [self.get_testdata_path('large_blob.bin')]))
+
   def test_save(self):
     """Tests save method."""
     buf = io.BytesIO()
@@ -634,6 +675,28 @@ class AftlIcpEntryTest(AftltoolTestCase):
 
     os.remove(key_file)
 
+  def test_verify_vbmeta_image(self):
+    """Tests the verify_vbmeta_image method."""
+    # Valid vbmeta image without footer with 1 AftlDescriptor.
+    tool = aftltool.Aftl()
+    vbmeta_image, _ = tool.get_vbmeta_image(
+        self.get_testdata_path('aftltool/aftl_output_vbmeta_with_1_icp.img'))
+    desc = tool.get_aftl_descriptor(
+        self.get_testdata_path('aftltool/aftl_output_vbmeta_with_1_icp.img'))
+
+    self.assertEqual(desc.icp_header.icp_count, 1)
+    entry = desc.icp_entries[0]
+    self.assertTrue(entry.verify_vbmeta_image(
+        vbmeta_image, self.get_testdata_path('aftltool/aftl_pubkey_1.pub')))
+
+    # Valid vbmeta image but checked with the public key of another log.
+    self.assertFalse(entry.verify_vbmeta_image(
+        vbmeta_image, self.get_testdata_path('aftltool/aftl_pubkey_2.pub')))
+
+    # Valid vbmeta image but checked with invalid public key file.
+    self.assertFalse(entry.verify_vbmeta_image(
+        vbmeta_image, self.get_testdata_path('large_blob.bin')))
+
   def test_print_desc(self):
     """Tests print_desc method."""
     buf = io.BytesIO()
@@ -950,18 +1013,57 @@ class AftlTest(AftltoolTestCase):
     super(AftlTest, self).setUp()
     self.mock_aftl_host = 'test.foo.bar:9000'
 
+  def test_get_vbmeta_image(self):
+    """Tests the get_vbmeta_image method."""
+    tool = aftltool.Aftl()
+
+    # Valid vbmeta image without footer and AftlDescriptor.
+    image, footer = tool.get_vbmeta_image(
+        self.get_testdata_path('aftltool/aftl_input_vbmeta.img'))
+    self.assertIsNotNone(image)
+    self.assertEqual(len(image), 4352)
+    self.assertIsNone(footer)
+
+    # Valid vbmeta image without footer but with AftlDescriptor.
+    image, footer = tool.get_vbmeta_image(
+        self.get_testdata_path('aftltool/aftl_output_vbmeta_with_1_icp.img'))
+    self.assertIsNotNone(image)
+    self.assertEqual(len(image), 4352)
+    self.assertIsNone(footer)
+
+    # Invalid vbmeta image.
+    image, footer = tool.get_vbmeta_image(
+        self.get_testdata_path('large_blob.bin'))
+    self.assertIsNone(image)
+    self.assertIsNone(footer)
+
+  def test_get_aftl_descriptor(self):
+    """Tests the get_aftl_descriptor method."""
+    tool = aftltool.Aftl()
+
+    # Valid vbmeta image without footer with AftlDescriptor.
+    desc = tool.get_aftl_descriptor(
+        self.get_testdata_path('aftltool/aftl_output_vbmeta_with_1_icp.img'))
+    self.assertIsInstance(desc, aftltool.AftlDescriptor)
+
+    # Valid vbmeta image without footer and AftlDescriptor.
+    desc = tool.get_aftl_descriptor(
+        self.get_testdata_path('aftltool/aftl_input_vbmeta.img'))
+    self.assertIsNone(desc)
+
+    # Invalid vbmeta image.
+    desc = tool.get_aftl_descriptor(self.get_testdata_path('large_blob.bin'))
+    self.assertIsNone(desc)
+
   # pylint: disable=no-member
   def test_request_inclusion_proof(self):
     """Tests the request_inclusion_proof method."""
     aftl_comms = AftlMockCommunication(self.mock_aftl_host, self.test_afi_resp)
     aftl = aftltool.Aftl()
-    icp = aftl.request_inclusion_proof(self.mock_aftl_host,
-                                       'a'*1024, '1',
-                                       os.path.join(TESTDATA_PATH, 'test',
-                                                    'data',
-                                                    'testkey_rsa4096.pem'),
-                                       None, None, None,
-                                       aftl_comms=aftl_comms)
+    icp = aftl.request_inclusion_proof(
+        self.mock_aftl_host, 'a' * 1024, '1',
+        self.get_testdata_path('testkey_rsa4096.pem'), None, None, None,
+        aftl_comms=aftl_comms)
     self.assertEqual(icp.leaf_index,
                      self.test_afi_resp.fw_info_proof.proof.leaf_index)
     self.assertEqual(icp.proof_hash_count,
@@ -989,11 +1091,10 @@ class AftlTest(AftltoolTestCase):
                                        aftltool.AftlError('Comms error'))
     aftl = aftltool.Aftl()
     with self.assertRaises(aftltool.AftlError):
-      aftl.request_inclusion_proof(self.mock_aftl_host,
-                                   'a'*1024, 'version_inc',
-                                   'test/data/testkey_rsa4096.pem',
-                                   None, None, None,
-                                   aftl_comms=aftl_comms)
+      aftl.request_inclusion_proof(
+          self.mock_aftl_host, 'a' * 1024, 'version_inc',
+          self.get_testdata_path('testkey_rsa4096.pem'), None, None, None,
+          aftl_comms=aftl_comms)
 
 if __name__ == '__main__':
   unittest.main(verbosity=2)
diff --git a/avbtool b/avbtool
index 2523449..c20ff0c 100755
--- a/avbtool
+++ b/avbtool
@@ -2577,6 +2577,9 @@ class Avb(object):
       AvbVBMetaHeader, the third argument is a list of
       AvbDescriptor-derived instances, and the fourth argument is the
       size of |image|.
+
+    Raises:
+      AvbError: In case the image cannot be parsed.
     """
     assert isinstance(image, ImageHandler)
     footer = None
@@ -2797,7 +2800,6 @@ class Avb(object):
       padding_needed = padded_size - len(vbmeta_blob)
       output.write('\0' * padding_needed)
 
-
   def _generate_vbmeta_blob(self, algorithm_name, key_path,
                             public_key_metadata_path, descriptors,
                             chain_partitions,
@@ -4095,7 +4097,6 @@ class AvbTool(object):
     self._add_common_args(sub_parser)
     sub_parser.set_defaults(func=self.make_vbmeta_image)
 
-
     sub_parser = subparsers.add_parser('add_hash_footer',
                                        help='Add hashes and footer to image.')
     sub_parser.add_argument('--image',
diff --git a/test/data/aftltool/aftl_input_vbmeta.img b/test/data/aftltool/aftl_input_vbmeta.img
new file mode 100644
index 0000000..4660701
--- /dev/null
+++ b/test/data/aftltool/aftl_input_vbmeta.img
diff --git a/test/data/aftltool/aftl_output_vbmeta_with_1_icp.img b/test/data/aftltool/aftl_output_vbmeta_with_1_icp.img
new file mode 100644
index 0000000..25e0869
--- /dev/null
+++ b/test/data/aftltool/aftl_output_vbmeta_with_1_icp.img
diff --git a/test/data/aftltool/aftl_output_vbmeta_with_2_icp.img b/test/data/aftltool/aftl_output_vbmeta_with_2_icp.img
new file mode 100644
index 0000000..e53178c
--- /dev/null
+++ b/test/data/aftltool/aftl_output_vbmeta_with_2_icp.img
diff --git a/test/data/aftltool/aftl_pubkey_1.pub b/test/data/aftltool/aftl_pubkey_1.pub
new file mode 100644
index 0000000..8bfd816
--- /dev/null
+++ b/test/data/aftltool/aftl_pubkey_1.pub
@@ -0,0 +1,15 @@
+-----BEGIN PUBLIC KEY-----
+MIICIjANBgkqhkiG9w0BAQEFAAOCAg8AMIICCgKCAgEA4ilqCNsenNA013iCdwgD
+YPxZ853nbHG9lMBp9boXiwRcqT/8bUKHIL7YX5z7s+QoRYVY3rkMKppRabclXzyx
+H59YnPMaU4uv7NqwWzjgaZo7E+vo7IF+KBjV3cJulId5Av0yIYUCsrwd7MpGtWdC
+Q3S+7Vd4zwzCKEhcvliNIhnNlp1U3wNkPCxOyCAsMEn6k8O5ar12ke5TvxDv15db
+rPDeHh8G2OYWoCkWL+lSN35L2kOJqKqVbLKWrrOd96RCYrrtbPCi580OADJRcUlG
+lgcjwmNwmypBWvQMZ6ITj0P0ksHnl1zZz1DE2rXe1goLI1doghb5KxLaezlR8c2C
+E3w/uo9KJgNmNgUVzzqZZ6FE0moyIDNOpP7KtZAL0DvEZj6jqLbB0ccPQElrg52m
+Dv2/A3nYSr0mYBKeskT4+Bg7PGgoC8p7WyLSxMyzJEDYdtrj9OFx6eZaA23oqTQx
+k3Qq5H8RfNBeeSUEeKF7pKH/7gyqZ2bNzBFMA2EBZgBozwRfaeN/HCv3qbaCnwvu
+6caacmAsK+RxiYxSL1QsJqyhCWWGxVyenmxdc1KG/u5ypi7OIioztyzR3t2tAzD3
+Nb+2t8lgHBRxbV24yiPlnvPmB1ZYEctXnlRR9Evpl1o9xA9NnybPHKr9rozN39CZ
+V/USB8K6ao1y5xPZxa8CZksCAwEAAQ==
+-----END PUBLIC KEY-----
+
diff --git a/test/data/aftltool/aftl_pubkey_2.pub b/test/data/aftltool/aftl_pubkey_2.pub
new file mode 100644
index 0000000..470fefa
--- /dev/null
+++ b/test/data/aftltool/aftl_pubkey_2.pub
@@ -0,0 +1,14 @@
+-----BEGIN PUBLIC KEY-----
+MIICIjANBgkqhkiG9w0BAQEFAAOCAg8AMIICCgKCAgEAylNLlTCf4FzlYDZIOWGf
+YQw3CusPo4bbcExfXP+0A5G94LXz02haTzlxHi3rIfiCqcUiJH1YrPWUFsD4dju+
+MKmb0lAPYSMjA5VMO3XINn5z0C+CUMq0/6rIw+n4iH/bfC3wqh0C/qqxDWM+pm5f
+duC0sk30jQ9+SMQqo/kxNY6Tzv3C1rLxDZ5lxsZ4+8IbAj7gtQFoibgwa9w2hbkZ
+dmXbyF/G8wjFPvQREC5VQik4RVPKnAB9r8ZDR9D3Myjfs/84KgMnpK7YVvsDqEcO
+4jB79t3AAVU/d6vl3hCGjTCXDZxaBzbB0PGJzDYX9cWn0tclu1WlpC7YY2YLZ8ai
+EWkMp7dXJxFGtQ7gn/RuULE+Li/H8jJkwQzVu2tb9geGDodXwNJYba4xs0CgLDjr
+L5LuiFsKhh+GBTAE/6JbQkgb97tI3ClmqhBNjvPki9sKMVL6hcUTzPHj4zxKTY8V
+yB+/1WEETcVyq1SDy0VrWPDgxXbvqkeMcMxr/8JfL6alsTCClPSlisLlhkXGrp+t
+FlkthdsD4M8BtccQHRuru3fAL8Kk5XOE7qeOcs0cDgzzmkZY1Pg4g4TpL3yiBvfi
+a3wHB2u7gF1XdLeMPNJa3v5pMwxHZzokQ8q01pC+gf/wNldY7oGWnTi5JHu/EyFV
+vWW/HX+Hd7cWppyAQFRqWwsCAwEAAQ==
+-----END PUBLIC KEY-----