Support writing FEC.
am: a778e5b5f6

Change-Id: I967d3d1fb684299693d521a7e4edae118fe7a30d
diff --git a/Android.mk b/Android.mk
index 062babb..4b06c27 100644
--- a/Android.mk
+++ b/Android.mk
@@ -108,12 +108,14 @@
     libbz \
     libbspatch \
     libbrotli \
+    libfec_rs \
     libpuffpatch \
     libverity_tree \
     $(ue_update_metadata_protos_exported_static_libraries)
 ue_libpayload_consumer_exported_shared_libraries := \
     libbase \
     libcrypto \
+    libfec \
     $(ue_update_metadata_protos_exported_shared_libraries)
 
 ue_libpayload_consumer_src_files := \
@@ -157,8 +159,6 @@
 ifeq ($(local_use_fec),1)
 ue_libpayload_consumer_src_files += \
     payload_consumer/fec_file_descriptor.cc
-ue_libpayload_consumer_exported_shared_libraries += \
-    libfec
 endif  # local_use_fec == 1
 
 ifeq ($(HOST_OS),linux)
diff --git a/payload_consumer/filesystem_verifier_action_unittest.cc b/payload_consumer/filesystem_verifier_action_unittest.cc
index 9eeecb1..15ec77a 100644
--- a/payload_consumer/filesystem_verifier_action_unittest.cc
+++ b/payload_consumer/filesystem_verifier_action_unittest.cc
@@ -277,7 +277,7 @@
 }
 
 #ifdef __ANDROID__
-TEST_F(FilesystemVerifierActionTest, VerityHashTreeTest) {
+TEST_F(FilesystemVerifierActionTest, WriteVerityTest) {
   FilesystemVerifierActionTestDelegate delegate;
   processor_.set_delegate(&delegate);
 
@@ -302,16 +302,21 @@
   part.hash_tree_data_size = filesystem_size;
   part.hash_tree_offset = filesystem_size;
   part.hash_tree_size = 3 * 4096;
+  part.fec_data_offset = 0;
+  part.fec_data_size = filesystem_size + part.hash_tree_size;
+  part.fec_offset = part.fec_data_size;
+  part.fec_size = 2 * 4096;
+  part.fec_roots = 2;
   // for i in {1..$((200 * 4096))}; do echo -n -e '\x1' >> part; done
   // avbtool add_hashtree_footer --image part --partition_size $((256 * 4096))
-  //     --partition_name part --do_not_generate_fec
-  //     --do_not_append_vbmeta_image --output_vbmeta_image vbmeta
+  //     --partition_name part --do_not_append_vbmeta_image
+  //     --output_vbmeta_image vbmeta
   // truncate -s $((256 * 4096)) part
   // sha256sum part | xxd -r -p | hexdump -v -e '/1 "0x%02x, "'
-  part.target_hash = {0xf0, 0x2c, 0x81, 0xf5, 0xec, 0x30, 0xa6, 0x99,
-                      0x1b, 0x41, 0x72, 0x16, 0x38, 0x48, 0xe5, 0x68,
-                      0x06, 0x7c, 0x3b, 0x88, 0xb5, 0x97, 0xa9, 0x29,
-                      0xa5, 0x7d, 0xdd, 0xa5, 0x9f, 0x5c, 0x15, 0x84};
+  part.target_hash = {0x28, 0xd4, 0x96, 0x75, 0x4c, 0xf5, 0x8a, 0x3e,
+                      0x31, 0x85, 0x08, 0x92, 0x85, 0x62, 0xf0, 0x37,
+                      0xbc, 0x8d, 0x7e, 0xa4, 0xcb, 0x24, 0x18, 0x7b,
+                      0xf3, 0xeb, 0xb5, 0x8d, 0x6f, 0xc8, 0xd8, 0x1a};
   // avbtool info_image --image vbmeta | grep Salt | cut -d':' -f 2 |
   //     xxd -r -p | hexdump -v -e '/1 "0x%02x, "'
   part.hash_tree_salt = {0x9e, 0xcb, 0xf8, 0xd5, 0x0b, 0xb4, 0x43,
diff --git a/payload_consumer/verity_writer_android.cc b/payload_consumer/verity_writer_android.cc
index b374817..06d1489 100644
--- a/payload_consumer/verity_writer_android.cc
+++ b/payload_consumer/verity_writer_android.cc
@@ -23,6 +23,10 @@
 
 #include <base/logging.h>
 #include <base/posix/eintr_wrapper.h>
+#include <fec/ecc.h>
+extern "C" {
+#include <fec.h>
+}
 
 #include "update_engine/common/utils.h"
 
@@ -78,7 +82,7 @@
         int fd = HANDLE_EINTR(open(partition_->target_path.c_str(), O_WRONLY));
         if (fd < 0) {
           PLOG(ERROR) << "Failed to open " << partition_->target_path
-                      << " to write verity data.";
+                      << " to write hash tree.";
           return false;
         }
         ScopedFdCloser fd_closer(&fd);
@@ -92,9 +96,97 @@
     }
   }
   if (partition_->fec_size != 0) {
-    // TODO(senj): Update FEC data.
+    uint64_t fec_data_end =
+        partition_->fec_data_offset + partition_->fec_data_size;
+    if (offset < fec_data_end && offset + size >= fec_data_end) {
+      LOG(INFO) << "Writing verity FEC to " << partition_->target_path;
+      TEST_AND_RETURN_FALSE(EncodeFEC(partition_->target_path,
+                                      partition_->fec_data_offset,
+                                      partition_->fec_data_size,
+                                      partition_->fec_offset,
+                                      partition_->fec_size,
+                                      partition_->fec_roots,
+                                      partition_->block_size,
+                                      false /* verify_mode */));
+    }
   }
   return true;
 }
 
+bool VerityWriterAndroid::EncodeFEC(const std::string& path,
+                                    uint64_t data_offset,
+                                    uint64_t data_size,
+                                    uint64_t fec_offset,
+                                    uint64_t fec_size,
+                                    uint32_t fec_roots,
+                                    uint32_t block_size,
+                                    bool verify_mode) {
+  TEST_AND_RETURN_FALSE(data_size % block_size == 0);
+  TEST_AND_RETURN_FALSE(fec_roots >= 0 && fec_roots < FEC_RSM);
+  // This is the N in RS(M, N), which is the number of bytes for each rs block.
+  size_t rs_n = FEC_RSM - fec_roots;
+  uint64_t rounds = utils::DivRoundUp(data_size / block_size, rs_n);
+  TEST_AND_RETURN_FALSE(rounds * fec_roots * block_size == fec_size);
+
+  std::unique_ptr<void, decltype(&free_rs_char)> rs_char(
+      init_rs_char(FEC_PARAMS(fec_roots)), &free_rs_char);
+  TEST_AND_RETURN_FALSE(rs_char != nullptr);
+
+  int fd = HANDLE_EINTR(open(path.c_str(), verify_mode ? O_RDONLY : O_RDWR));
+  if (fd < 0) {
+    PLOG(ERROR) << "Failed to open " << path << " to write FEC.";
+    return false;
+  }
+  ScopedFdCloser fd_closer(&fd);
+
+  for (size_t i = 0; i < rounds; i++) {
+    // Encodes |block_size| number of rs blocks each round so that we can read
+    // one block each time instead of 1 byte to increase random read
+    // performance. This uses about 1 MiB memory for 4K block size.
+    brillo::Blob rs_blocks(block_size * rs_n);
+    for (size_t j = 0; j < rs_n; j++) {
+      brillo::Blob buffer(block_size, 0);
+      uint64_t offset =
+          fec_ecc_interleave(i * rs_n * block_size + j, rs_n, rounds);
+      // Don't read past |data_size|, treat them as 0.
+      if (offset < data_size) {
+        ssize_t bytes_read = 0;
+        TEST_AND_RETURN_FALSE(utils::PReadAll(fd,
+                                              buffer.data(),
+                                              buffer.size(),
+                                              data_offset + offset,
+                                              &bytes_read));
+        TEST_AND_RETURN_FALSE(bytes_read ==
+                              static_cast<ssize_t>(buffer.size()));
+      }
+      for (size_t k = 0; k < buffer.size(); k++) {
+        rs_blocks[k * rs_n + j] = buffer[k];
+      }
+    }
+    brillo::Blob fec(block_size * fec_roots);
+    for (size_t j = 0; j < block_size; j++) {
+      // Encode [j * rs_n : (j + 1) * rs_n) in |rs_blocks| and write |fec_roots|
+      // number of parity bytes to |j * fec_roots| in |fec|.
+      encode_rs_char(rs_char.get(),
+                     rs_blocks.data() + j * rs_n,
+                     fec.data() + j * fec_roots);
+    }
+
+    if (verify_mode) {
+      brillo::Blob fec_read(fec.size());
+      ssize_t bytes_read = 0;
+      TEST_AND_RETURN_FALSE(utils::PReadAll(
+          fd, fec_read.data(), fec_read.size(), fec_offset, &bytes_read));
+      TEST_AND_RETURN_FALSE(bytes_read ==
+                            static_cast<ssize_t>(fec_read.size()));
+      TEST_AND_RETURN_FALSE(fec == fec_read);
+    } else {
+      TEST_AND_RETURN_FALSE(
+          utils::PWriteAll(fd, fec.data(), fec.size(), fec_offset));
+    }
+    fec_offset += fec.size();
+  }
+
+  return true;
+}
 }  // namespace chromeos_update_engine
diff --git a/payload_consumer/verity_writer_android.h b/payload_consumer/verity_writer_android.h
index c26663b..05a5856 100644
--- a/payload_consumer/verity_writer_android.h
+++ b/payload_consumer/verity_writer_android.h
@@ -18,6 +18,7 @@
 #define UPDATE_ENGINE_PAYLOAD_CONSUMER_VERITY_WRITER_ANDROID_H_
 
 #include <memory>
+#include <string>
 
 #include <verity/hash_tree_builder.h>
 
@@ -33,6 +34,21 @@
   bool Init(const InstallPlan::Partition& partition) override;
   bool Update(uint64_t offset, const uint8_t* buffer, size_t size) override;
 
+  // Read [data_offset : data_offset + data_size) from |path| and encode FEC
+  // data, if |verify_mode|, then compare the encoded FEC with the one in
+  // |path|, otherwise write the encoded FEC to |path|. We can't encode as we go
+  // in each Update() like hash tree, because for every rs block, its data are
+  // spreaded across entire |data_size|, unless we can cache all data in
+  // memory, we have to re-read them from disk.
+  static bool EncodeFEC(const std::string& path,
+                        uint64_t data_offset,
+                        uint64_t data_size,
+                        uint64_t fec_offset,
+                        uint64_t fec_size,
+                        uint32_t fec_roots,
+                        uint32_t block_size,
+                        bool verify_mode);
+
  private:
   const InstallPlan::Partition* partition_ = nullptr;
 
diff --git a/payload_consumer/verity_writer_android_unittest.cc b/payload_consumer/verity_writer_android_unittest.cc
index 56335c1..f943ce8 100644
--- a/payload_consumer/verity_writer_android_unittest.cc
+++ b/payload_consumer/verity_writer_android_unittest.cc
@@ -34,6 +34,7 @@
     partition_.hash_tree_offset = 4096;
     partition_.hash_tree_size = 4096;
     partition_.hash_tree_algorithm = "sha1";
+    partition_.fec_roots = 2;
   }
 
   VerityWriterAndroid verity_writer_;
@@ -97,4 +98,23 @@
   EXPECT_EQ(part_data, actual_part);
 }
 
+TEST_F(VerityWriterAndroidTest, FECTest) {
+  partition_.fec_data_offset = 0;
+  partition_.fec_data_size = 4096;
+  partition_.fec_offset = 4096;
+  partition_.fec_size = 2 * 4096;
+  brillo::Blob part_data(3 * 4096, 0x1);
+  test_utils::WriteFileVector(partition_.target_path, part_data);
+  ASSERT_TRUE(verity_writer_.Init(partition_));
+  EXPECT_TRUE(verity_writer_.Update(0, part_data.data(), part_data.size()));
+  brillo::Blob actual_part;
+  utils::ReadFile(partition_.target_path, &actual_part);
+  // Write FEC data.
+  for (size_t i = 4096; i < part_data.size(); i += 2) {
+    part_data[i] = 0x8e;
+    part_data[i + 1] = 0x8f;
+  }
+  EXPECT_EQ(part_data, actual_part);
+}
+
 }  // namespace chromeos_update_engine
diff --git a/payload_generator/delta_diff_utils.cc b/payload_generator/delta_diff_utils.cc
index 9d14474..41de623 100644
--- a/payload_generator/delta_diff_utils.cc
+++ b/payload_generator/delta_diff_utils.cc
@@ -306,8 +306,9 @@
     LOG(INFO) << "Skipping verity hash tree blocks: "
               << ExtentsToString({new_part.verity.hash_tree_extent});
     new_visited_blocks.AddExtent(new_part.verity.hash_tree_extent);
-    // TODO(senj): Enable this when we have FEC support.
-    // new_visited_blocks.AddExtent(new_part.verity.fec_extent);
+    LOG(INFO) << "Skipping verity FEC blocks: "
+              << ExtentsToString({new_part.verity.fec_extent});
+    new_visited_blocks.AddExtent(new_part.verity.fec_extent);
   }
 
   TEST_AND_RETURN_FALSE(DeltaMovedAndZeroBlocks(
diff --git a/payload_generator/delta_diff_utils_unittest.cc b/payload_generator/delta_diff_utils_unittest.cc
index 07aa3a8..e32dde2 100644
--- a/payload_generator/delta_diff_utils_unittest.cc
+++ b/payload_generator/delta_diff_utils_unittest.cc
@@ -163,6 +163,7 @@
 
 TEST_F(DeltaDiffUtilsTest, SkipVerityExtentsTest) {
   new_part_.verity.hash_tree_extent = ExtentForRange(20, 30);
+  new_part_.verity.fec_extent = ExtentForRange(40, 50);
 
   BlobFileWriter blob_file(blob_fd_, &blob_size_);
   EXPECT_TRUE(diff_utils::DeltaReadPartition(
@@ -180,6 +181,8 @@
   for (const auto& extent : new_visited_blocks_.extent_set()) {
     EXPECT_FALSE(ExtentRanges::ExtentsOverlap(
         extent, new_part_.verity.hash_tree_extent));
+    EXPECT_FALSE(
+        ExtentRanges::ExtentsOverlap(extent, new_part_.verity.fec_extent));
   }
 }
 
diff --git a/payload_generator/payload_generation_config_android.cc b/payload_generator/payload_generation_config_android.cc
index 7d9988f..e1dee58 100644
--- a/payload_generator/payload_generation_config_android.cc
+++ b/payload_generator/payload_generation_config_android.cc
@@ -22,6 +22,7 @@
 #include <verity/hash_tree_builder.h>
 
 #include "update_engine/common/utils.h"
+#include "update_engine/payload_consumer/verity_writer_android.h"
 #include "update_engine/payload_generator/extent_ranges.h"
 
 namespace chromeos_update_engine {
@@ -87,7 +88,15 @@
   part->verity.hash_tree_extent = ExtentForBytes(
       hashtree.hash_block_size, hashtree.tree_offset, hashtree.tree_size);
 
-  // TODO(senj): Verify FEC data.
+  TEST_AND_RETURN_FALSE(VerityWriterAndroid::EncodeFEC(part->path,
+                                                       0 /* data_offset */,
+                                                       hashtree.fec_offset,
+                                                       hashtree.fec_offset,
+                                                       hashtree.fec_size,
+                                                       hashtree.fec_num_roots,
+                                                       hashtree.data_block_size,
+                                                       true /* verify_mode */));
+
   part->verity.fec_data_extent =
       ExtentForBytes(hashtree.data_block_size, 0, hashtree.fec_offset);
   part->verity.fec_extent = ExtentForBytes(
diff --git a/payload_generator/payload_generation_config_android_unittest.cc b/payload_generator/payload_generation_config_android_unittest.cc
index 0e6ebe0..53378c2 100644
--- a/payload_generator/payload_generation_config_android_unittest.cc
+++ b/payload_generator/payload_generation_config_android_unittest.cc
@@ -170,6 +170,15 @@
   EXPECT_FALSE(image_config_.LoadVerityConfig());
 }
 
+TEST_F(PayloadGenerationConfigAndroidTest, LoadVerityConfigInvalidFECTest) {
+  brillo::Blob part = GetAVBPartition();
+  part[kFECOffset] ^= 1;  // flip one bit
+  test_utils::WriteFileVector(temp_file_.path(), part);
+  EXPECT_TRUE(image_config_.LoadImageSize());
+  EXPECT_TRUE(image_config_.partitions[0].OpenFilesystem());
+  EXPECT_FALSE(image_config_.LoadVerityConfig());
+}
+
 TEST_F(PayloadGenerationConfigAndroidTest, LoadVerityConfigEmptyImageTest) {
   brillo::Blob part(kImageSize);
   test_utils::WriteFileVector(temp_file_.path(), part);