//
// Copyright (C) 2018 The Android Open Source Project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

#include "update_engine/payload_generator/payload_generation_config.h"

#include <brillo/secure_blob.h>
#include <gtest/gtest.h>

#include "update_engine/common/test_utils.h"
#include "update_engine/payload_generator/extent_ranges.h"
#include "update_engine/payload_generator/extent_utils.h"

namespace chromeos_update_engine {

namespace {
// dd if=/dev/zero of=part bs=4096 count=2
// avbtool add_hashtree_footer --image part --partition_size $((24 * 4096))
//     --partition_name system
constexpr uint64_t kImageSize = 24 * 4096;

// hexdump -s $((2 * 4096)) -n 64 -v -e '/1 "0x%02x, "' part
constexpr uint64_t kHashTreeOffset = 2 * 4096;
const uint8_t kHashTree[] = {
    0x62, 0x4b, 0x5a, 0x4d, 0xa2, 0x97, 0xa0, 0xc8, 0x08, 0x03, 0xa6,
    0x95, 0x4c, 0x4c, 0x7a, 0x2d, 0xac, 0x50, 0xde, 0x74, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x62,
    0x4b, 0x5a, 0x4d, 0xa2, 0x97, 0xa0, 0xc8, 0x08, 0x03, 0xa6, 0x95,
    0x4c, 0x4c, 0x7a, 0x2d, 0xac, 0x50, 0xde, 0x74, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00};

// hexdump -s $((3 * 4096)) -n 128 -v -e '/1 "0x%02x, "' part
constexpr uint64_t kFECOffset = 3 * 4096;
const uint8_t kFEC[] = {
    0xec, 0x8e, 0x93, 0xd8, 0xf9, 0xa3, 0xd6, 0x9b, 0xa4, 0x06, 0x5f, 0xc8,
    0x6c, 0xcc, 0x4f, 0x87, 0x07, 0x0f, 0xac, 0xaf, 0x29, 0x8f, 0x97, 0x02,
    0xb2, 0xfe, 0xb2, 0xfe, 0xe5, 0x9f, 0xf2, 0xdf, 0xe6, 0x4a, 0x36, 0x66,
    0x04, 0xda, 0xa7, 0xd3, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0xec, 0x8e, 0x93, 0xd8, 0xf9, 0xa3, 0xd6, 0x9b,
    0xa4, 0x06, 0x5f, 0xc8, 0x6c, 0xcc, 0x4f, 0x87, 0x07, 0x0f, 0xac, 0xaf,
    0x29, 0x8f, 0x97, 0x02, 0xb2, 0xfe, 0xb2, 0xfe, 0xe5, 0x9f, 0xf2, 0xdf,
    0xe6, 0x4a, 0x36, 0x66, 0x04, 0xda, 0xa7, 0xd3, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00};

// hexdump -s $((5 * 4096)) -n 512 -v -e '/1 "0x%02x, "' part
constexpr uint64_t kVBMetaImageOffset = 5 * 4096;
const uint8_t kVBMetaImage[] = {
    0x41, 0x56, 0x42, 0x30, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xe8,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0xe8, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0xe8, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x61, 0x76, 0x62, 0x74,
    0x6f, 0x6f, 0x6c, 0x20, 0x31, 0x2e, 0x31, 0x2e, 0x30, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xd8, 0x00, 0x00, 0x00, 0x01,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x20, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x20, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x10, 0x00,
    0x00, 0x00, 0x10, 0x00, 0x00, 0x00, 0x10, 0x00, 0x00, 0x00, 0x00, 0x02,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x30, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x20, 0x00, 0x73, 0x68, 0x61, 0x31, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x06, 0x00, 0x00, 0x00, 0x14, 0x00, 0x00, 0x00, 0x14,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x73, 0x79, 0x73, 0x74, 0x65, 0x6d, 0x1f, 0xab,
    0x7a, 0x6b, 0xf6, 0xb1, 0x3a, 0x1f, 0xdb, 0x34, 0xa3, 0xfc, 0xc8, 0x73,
    0x0b, 0x23, 0x61, 0xb3, 0x04, 0xe2, 0x4f, 0x6c, 0xd0, 0x1e, 0x39, 0x9d,
    0xaa, 0x73, 0x35, 0x53, 0xa7, 0x74, 0x1f, 0x81, 0xd0, 0xa6, 0xa9, 0x5f,
    0x19, 0x9f, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00};

// hexdump -s $((24 * 4096 - 64)) -n 64 -v -e '/1 "0x%02x, "' part
constexpr uint64_t kAVBFooterOffset = 24 * 4096 - 64;
const uint8_t kAVBFooter[] = {
    0x41, 0x56, 0x42, 0x66, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x20, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x50, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00};

// avbtool info_image --image part | grep Salt | cut -d':' -f 2 | xxd -r -p |
//     hexdump -v -e '/1 "0x%02x, "'
const uint8_t kHashTreeSalt[] = {0x1f, 0xab, 0x7a, 0x6b, 0xf6, 0xb1, 0x3a,
                                 0x1f, 0xdb, 0x34, 0xa3, 0xfc, 0xc8, 0x73,
                                 0x0b, 0x23, 0x61, 0xb3, 0x04, 0xe2};

brillo::Blob GetAVBPartition() {
  brillo::Blob part(kImageSize);
  memcpy(part.data() + kHashTreeOffset, kHashTree, sizeof(kHashTree));
  memcpy(part.data() + kFECOffset, kFEC, sizeof(kFEC));
  memcpy(part.data() + kVBMetaImageOffset, kVBMetaImage, sizeof(kVBMetaImage));
  memcpy(part.data() + kAVBFooterOffset, kAVBFooter, sizeof(kAVBFooter));
  return part;
}

}  // namespace

class PayloadGenerationConfigAndroidTest : public ::testing::Test {
 protected:
  void SetUp() override {
    image_config_.partitions.emplace_back("system");
    image_config_.partitions[0].path = temp_file_.path();
  }

  ImageConfig image_config_;
  test_utils::ScopedTempFile temp_file_{
      "PayloadGenerationConfigAndroidTest.XXXXXX"};
};

TEST_F(PayloadGenerationConfigAndroidTest, LoadVerityConfigSimpleTest) {
  brillo::Blob part = GetAVBPartition();
  test_utils::WriteFileVector(temp_file_.path(), part);
  EXPECT_TRUE(image_config_.LoadImageSize());
  EXPECT_TRUE(image_config_.partitions[0].OpenFilesystem());
  EXPECT_TRUE(image_config_.LoadVerityConfig());
  const VerityConfig& verity = image_config_.partitions[0].verity;
  EXPECT_FALSE(verity.IsEmpty());
  EXPECT_EQ(ExtentForRange(0, 2), verity.hash_tree_data_extent);
  EXPECT_EQ(ExtentForRange(2, 1), verity.hash_tree_extent);
  EXPECT_EQ("sha1", verity.hash_tree_algorithm);
  brillo::Blob salt(kHashTreeSalt, std::end(kHashTreeSalt));
  EXPECT_EQ(salt, verity.hash_tree_salt);
  EXPECT_EQ(ExtentForRange(0, 3), verity.fec_data_extent);
  EXPECT_EQ(ExtentForRange(3, 2), verity.fec_extent);
  EXPECT_EQ(2u, verity.fec_roots);
}

TEST_F(PayloadGenerationConfigAndroidTest, LoadVerityConfigDisableFecTest) {
  brillo::Blob part = GetAVBPartition();
  test_utils::WriteFileVector(temp_file_.path(), part);
  image_config_.partitions[0].disable_fec_computation = true;
  EXPECT_TRUE(image_config_.LoadImageSize());
  EXPECT_TRUE(image_config_.partitions[0].OpenFilesystem());
  EXPECT_TRUE(image_config_.LoadVerityConfig());
  const VerityConfig& verity = image_config_.partitions[0].verity;
  EXPECT_FALSE(verity.IsEmpty());
  EXPECT_EQ(ExtentForRange(0, 2), verity.hash_tree_data_extent);
  EXPECT_EQ(ExtentForRange(2, 1), verity.hash_tree_extent);
  EXPECT_EQ("sha1", verity.hash_tree_algorithm);
  brillo::Blob salt(kHashTreeSalt, std::end(kHashTreeSalt));
  EXPECT_EQ(salt, verity.hash_tree_salt);
  EXPECT_EQ(0u, verity.fec_data_extent.num_blocks());
  EXPECT_EQ(0u, verity.fec_extent.num_blocks());
}

TEST_F(PayloadGenerationConfigAndroidTest,
       LoadVerityConfigInvalidHashTreeTest) {
  brillo::Blob part = GetAVBPartition();
  part[kHashTreeOffset] ^= 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, 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);
  EXPECT_TRUE(image_config_.LoadImageSize());
  EXPECT_TRUE(image_config_.LoadVerityConfig());
  EXPECT_TRUE(image_config_.partitions[0].verity.IsEmpty());
}

TEST_F(PayloadGenerationConfigAndroidTest, LoadVerityConfigTinyImageTest) {
  test_utils::WriteFileString(temp_file_.path(), "tiny");
  EXPECT_TRUE(image_config_.LoadImageSize());
  EXPECT_TRUE(image_config_.LoadVerityConfig());
  EXPECT_TRUE(image_config_.partitions[0].verity.IsEmpty());
}

}  // namespace chromeos_update_engine