Error correction: Append codes to verified partitions

Append error-correcting codes to verified partitions provided that
PRODUCT_SUPPORTS_VERITY_FEC is true.

This moves verity metadata to be after the hash tree, and requires
matching changes from
  Ide48f581bbba77aed6132f77b309db71630d81ed

Bug: 21893453
Change-Id: I6945cbab99e214566a1f9d3702333f2dbbc35816
diff --git a/tools/releasetools/build_image.py b/tools/releasetools/build_image.py
index 470a108..cd750e8 100755
--- a/tools/releasetools/build_image.py
+++ b/tools/releasetools/build_image.py
@@ -33,6 +33,7 @@
 OPTIONS = common.OPTIONS
 
 FIXED_SALT = "aee087a5be3b982978c923f566a94613496b417f2af592639bc80d141e34dfe7"
+BLOCK_SIZE = 4096
 
 def RunCommand(cmd):
   """Echo and run the given command.
@@ -48,6 +49,14 @@
   print "%s" % (output.rstrip(),)
   return (output, p.returncode)
 
+def GetVerityFECSize(partition_size):
+  cmd = "fec -s %d" % partition_size
+  status, output = commands.getstatusoutput(cmd)
+  if status:
+    print output
+    return False, 0
+  return True, int(output)
+
 def GetVerityTreeSize(partition_size):
   cmd = "build_verity_tree -s %d"
   cmd %= partition_size
@@ -67,7 +76,22 @@
     return False, 0
   return True, int(output)
 
-def AdjustPartitionSizeForVerity(partition_size):
+def GetVeritySize(partition_size, fec_supported):
+  success, verity_tree_size = GetVerityTreeSize(partition_size)
+  if not success:
+    return 0
+  success, verity_metadata_size = GetVerityMetadataSize(partition_size)
+  if not success:
+    return 0
+  verity_size = verity_tree_size + verity_metadata_size
+  if fec_supported:
+    success, fec_size = GetVerityFECSize(partition_size + verity_size)
+    if not success:
+      return 0
+    return verity_size + fec_size
+  return verity_size
+
+def AdjustPartitionSizeForVerity(partition_size, fec_supported):
   """Modifies the provided partition size to account for the verity metadata.
 
   This information is used to size the created image appropriately.
@@ -76,13 +100,43 @@
   Returns:
     The size of the partition adjusted for verity metadata.
   """
-  success, verity_tree_size = GetVerityTreeSize(partition_size)
-  if not success:
-    return 0
-  success, verity_metadata_size = GetVerityMetadataSize(partition_size)
-  if not success:
-    return 0
-  return partition_size - verity_tree_size - verity_metadata_size
+  key = "%d %d" % (partition_size, fec_supported)
+  if key in AdjustPartitionSizeForVerity.results:
+    return AdjustPartitionSizeForVerity.results[key]
+
+  hi = partition_size
+  if hi % BLOCK_SIZE != 0:
+    hi = (hi // BLOCK_SIZE) * BLOCK_SIZE
+
+  # verity tree and fec sizes depend on the partition size, which
+  # means this estimate is always going to be unnecessarily small
+  lo = partition_size - GetVeritySize(hi, fec_supported)
+  result = lo
+
+  # do a binary search for the optimal size
+  while lo < hi:
+    i = ((lo + hi) // (2 * BLOCK_SIZE)) * BLOCK_SIZE
+    size = i + GetVeritySize(i, fec_supported)
+    if size <= partition_size:
+      if result < i:
+        result = i
+      lo = i + BLOCK_SIZE
+    else:
+      hi = i
+
+  AdjustPartitionSizeForVerity.results[key] = result
+  return result
+
+AdjustPartitionSizeForVerity.results = {}
+
+def BuildVerityFEC(sparse_image_path, verity_fec_path, prop_dict):
+  cmd = "fec -e %s %s" % (sparse_image_path, verity_fec_path)
+  print cmd
+  status, output = commands.getstatusoutput(cmd)
+  if status:
+    print "Could not build FEC data! Error: %s" % output
+    return False
+  return True
 
 def BuildVerityTree(sparse_image_path, verity_image_path, prop_dict):
   cmd = "build_verity_tree -A %s %s %s" % (
@@ -130,12 +184,12 @@
 
 def BuildVerifiedImage(data_image_path, verity_image_path,
                        verity_metadata_path):
-  if not Append2Simg(data_image_path, verity_metadata_path,
-                     "Could not append verity metadata!"):
-    return False
   if not Append2Simg(data_image_path, verity_image_path,
                      "Could not append verity tree!"):
     return False
+  if not Append2Simg(data_image_path, verity_metadata_path,
+                     "Could not append verity metadata!"):
+    return False
   return True
 
 def UnsparseImage(sparse_image_path, replace=True):
@@ -154,7 +208,7 @@
     return False, None
   return True, unsparse_image_path
 
-def MakeVerityEnabledImage(out_file, prop_dict):
+def MakeVerityEnabledImage(out_file, fec_supported, prop_dict):
   """Creates an image that is verifiable using dm-verity.
 
   Args:
@@ -180,6 +234,7 @@
   # get partial image paths
   verity_image_path = os.path.join(tempdir_name, "verity.img")
   verity_metadata_path = os.path.join(tempdir_name, "verity_metadata.img")
+  verity_fec_path = os.path.join(tempdir_name, "verity_fec.img")
 
   # build the verity tree and get the root hash and salt
   if not BuildVerityTree(out_file, verity_image_path, prop_dict):
@@ -201,6 +256,16 @@
     shutil.rmtree(tempdir_name, ignore_errors=True)
     return False
 
+  if fec_supported:
+    # build FEC for the entire partition, including metadata
+    if not BuildVerityFEC(out_file, verity_fec_path, prop_dict):
+      shutil.rmtree(tempdir_name, ignore_errors=True)
+      return False
+
+    if not Append2Simg(out_file, verity_fec_path, "Could not append FEC!"):
+      shutil.rmtree(tempdir_name, ignore_errors=True)
+      return False
+
   shutil.rmtree(tempdir_name, ignore_errors=True)
   return True
 
@@ -248,12 +313,14 @@
 
   is_verity_partition = "verity_block_device" in prop_dict
   verity_supported = prop_dict.get("verity") == "true"
+  verity_fec_supported = prop_dict.get("verity_fec") == "true"
+
   # Adjust the partition size to make room for the hashes if this is to be
   # verified.
   if verity_supported and is_verity_partition and fs_spans_partition:
     partition_size = int(prop_dict.get("partition_size"))
-
-    adjusted_size = AdjustPartitionSizeForVerity(partition_size)
+    adjusted_size = AdjustPartitionSizeForVerity(partition_size,
+                                                 verity_fec_supported)
     if not adjusted_size:
       return False
     prop_dict["partition_size"] = str(adjusted_size)
@@ -366,7 +433,7 @@
             "%d" % (mount_point, image_size, partition_size))
       return False
     if verity_supported and is_verity_partition:
-      if 2 * image_size - AdjustPartitionSizeForVerity(image_size) > partition_size:
+      if 2 * image_size - AdjustPartitionSizeForVerity(image_size, verity_fec_supported) > partition_size:
         print "Error: No more room on %s to fit verity data" % mount_point
         return False
     prop_dict["original_partition_size"] = prop_dict["partition_size"]
@@ -374,7 +441,7 @@
 
   # create the verified image if this is to be verified
   if verity_supported and is_verity_partition:
-    if not MakeVerityEnabledImage(out_file, prop_dict):
+    if not MakeVerityEnabledImage(out_file, verity_fec_supported, prop_dict):
       return False
 
   if run_fsck and prop_dict.get("skip_fsck") != "true":
@@ -416,7 +483,8 @@
       "skip_fsck",
       "verity",
       "verity_key",
-      "verity_signer_cmd"
+      "verity_signer_cmd",
+      "verity_fec"
       )
   for p in common_props:
     copy_prop(p, p)