From 68658c0f4fe5420226df5849b642f98fb7f5d984 Mon Sep 17 00:00:00 2001
From: Tao Bao <tbao@google.com>
Date: Mon, 1 Jun 2015 13:40:49 -0700
Subject: [PATCH] Add post-install verification for BBOTAs

Similar to the assertations in file-based OTA, we perform verification
for block-based OTAs (BBOTAs) after updating a partition, for both of
the incremental and full OTAs. It increases the update time (~20s on
Nexus 6), but will capture unnoticed errors right away.

Bug: 21500869
Change-Id: I89ac8fe623b855721b7affd07cf9426a23433ab0
---
 tools/releasetools/blockimgdiff.py          | 15 ++++++++++-----
 tools/releasetools/common.py                | 16 ++++++++++++++++
 tools/releasetools/ota_from_target_files.py |  1 +
 tools/releasetools/sparse_img.py            | 13 +++++++++----
 4 files changed, 36 insertions(+), 9 deletions(-)

diff --git a/tools/releasetools/blockimgdiff.py b/tools/releasetools/blockimgdiff.py
index 0a387ec2c6..3402572495 100644
--- a/tools/releasetools/blockimgdiff.py
+++ b/tools/releasetools/blockimgdiff.py
@@ -74,7 +74,7 @@ class Image(object):
   def ReadRangeSet(self, ranges):
     raise NotImplementedError
 
-  def TotalSha1(self):
+  def TotalSha1(self, include_clobbered_blocks=False):
     raise NotImplementedError
 
 
@@ -87,7 +87,10 @@ class EmptyImage(Image):
   file_map = {}
   def ReadRangeSet(self, ranges):
     return ()
-  def TotalSha1(self):
+  def TotalSha1(self, include_clobbered_blocks=False):
+    # EmptyImage always carries empty clobbered_blocks, so
+    # include_clobbered_blocks can be ignored.
+    assert self.clobbered_blocks.size() == 0
     return sha1().hexdigest()
 
 
@@ -136,8 +139,9 @@ class DataImage(Image):
   def ReadRangeSet(self, ranges):
     return [self.data[s*self.blocksize:e*self.blocksize] for (s, e) in ranges]
 
-  def TotalSha1(self):
-    # DataImage always carries empty clobbered_blocks.
+  def TotalSha1(self, include_clobbered_blocks=False):
+    # DataImage always carries empty clobbered_blocks, so
+    # include_clobbered_blocks can be ignored.
     assert self.clobbered_blocks.size() == 0
     return sha1(self.data).hexdigest()
 
@@ -201,7 +205,8 @@ class Transfer(object):
 #
 #    TotalSha1(): a function that returns (as a hex string) the SHA-1
 #      hash of all the data in the image (ie, all the blocks in the
-#      care_map minus clobbered_blocks).
+#      care_map minus clobbered_blocks, or including the clobbered
+#      blocks if include_clobbered_blocks is True).
 #
 # When creating a BlockImageDiff, the src image may be None, in which
 # case the list of transfers produced will never read from the
diff --git a/tools/releasetools/common.py b/tools/releasetools/common.py
index f54c8bd18a..f4e25d5a97 100644
--- a/tools/releasetools/common.py
+++ b/tools/releasetools/common.py
@@ -1182,6 +1182,7 @@ class BlockDifference(object):
     if progress:
       script.ShowProgress(progress, 0)
     self._WriteUpdate(script, output_zip)
+    self._WritePostInstallVerifyScript(script)
 
   def WriteVerifyScript(self, script):
     partition = self.partition
@@ -1220,6 +1221,21 @@ class BlockDifference(object):
       script.AppendExtra(('abort("%s partition has unexpected contents");\n'
                           'endif;') % (partition,))
 
+  def _WritePostInstallVerifyScript(self, script):
+    partition = self.partition
+    script.Print('Verifying the updated %s image...' % (partition,))
+    # Unlike pre-install verification, clobbered_blocks should not be ignored.
+    ranges = self.tgt.care_map
+    ranges_str = ranges.to_string_raw()
+    script.AppendExtra('if range_sha1("%s", "%s") == "%s" then' % (
+                       self.device, ranges_str,
+                       self.tgt.TotalSha1(include_clobbered_blocks=True)))
+    script.Print('Verified the updated %s image.' % (partition,))
+    script.AppendExtra(
+        'else\n'
+        '  abort("%s partition has unexpected contents after OTA update");\n'
+        'endif;' % (partition,))
+
   def _WriteUpdate(self, script, output_zip):
     ZipWrite(output_zip,
              '{}.transfer.list'.format(self.path),
diff --git a/tools/releasetools/ota_from_target_files.py b/tools/releasetools/ota_from_target_files.py
index c4d0c1b74d..b0913e457b 100755
--- a/tools/releasetools/ota_from_target_files.py
+++ b/tools/releasetools/ota_from_target_files.py
@@ -917,6 +917,7 @@ else
 
   system_diff.WriteScript(script, output_zip,
                           progress=0.8 if vendor_diff else 0.9)
+
   if vendor_diff:
     vendor_diff.WriteScript(script, output_zip, progress=0.1)
 
diff --git a/tools/releasetools/sparse_img.py b/tools/releasetools/sparse_img.py
index 2ac97ac965..51a1643be8 100644
--- a/tools/releasetools/sparse_img.py
+++ b/tools/releasetools/sparse_img.py
@@ -118,11 +118,16 @@ class SparseImage(object):
   def ReadRangeSet(self, ranges):
     return [d for d in self._GetRangeData(ranges)]
 
-  def TotalSha1(self):
-    """Return the SHA-1 hash of all data in the 'care' regions but not in
-    clobbered_blocks of this image."""
+  def TotalSha1(self, include_clobbered_blocks=False):
+    """Return the SHA-1 hash of all data in the 'care' regions.
+
+    If include_clobbered_blocks is True, it returns the hash including the
+    clobbered_blocks."""
+    ranges = self.care_map
+    if not include_clobbered_blocks:
+      ranges.subtract(self.clobbered_blocks)
     h = sha1()
-    for d in self._GetRangeData(self.care_map.subtract(self.clobbered_blocks)):
+    for d in self._GetRangeData(ranges):
       h.update(d)
     return h.hexdigest()