// Copyright 2016 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "imageloader_impl.h"

#include <stdint.h>

#include <list>
#include <string>
#include <vector>

#include "component.h"
#include "mock_helper_process.h"
#include "test_utilities.h"
#include "verity_mounter.h"

#include <base/files/file_path.h>
#include <base/files/file_util.h>
#include <base/files/scoped_temp_dir.h>
#include <base/memory/ptr_util.h>
#include <gmock/gmock.h>
#include <gtest/gtest.h>

namespace imageloader {

using testing::_;

class ImageLoaderTest : public testing::Test {
 public:
  ImageLoaderTest() {
    CHECK(scoped_temp_dir_.CreateUniqueTempDir());
    temp_dir_ = scoped_temp_dir_.path();
    CHECK(base::SetPosixFilePermissions(temp_dir_, kComponentDirPerms));
  }

  ImageLoaderConfig GetConfig(const char* path) {
    std::vector<uint8_t> key(std::begin(kDevPublicKey),
                             std::end(kDevPublicKey));
    ImageLoaderConfig config(key, path, "/foo");
    return config;
  }

  base::ScopedTempDir scoped_temp_dir_;
  base::FilePath temp_dir_;
};

// Test the RegisterComponent public interface.
TEST_F(ImageLoaderTest, RegisterComponentAndGetVersion) {
  ImageLoaderImpl loader(GetConfig(temp_dir_.value().c_str()));
  ASSERT_TRUE(loader.RegisterComponent(kTestComponentName, kTestDataVersion,
                                       GetTestComponentPath().value()));

  base::FilePath comp_dir = temp_dir_.Append(kTestComponentName);
  ASSERT_TRUE(base::DirectoryExists(comp_dir));

  base::FilePath hint_file = comp_dir.Append("latest-version");
  ASSERT_TRUE(base::PathExists(hint_file));

  std::string hint_file_contents;
  ASSERT_TRUE(
      base::ReadFileToStringWithMaxSize(hint_file, &hint_file_contents, 4096));
  EXPECT_EQ(kTestDataVersion, hint_file_contents);

  base::FilePath version_dir = comp_dir.Append(kTestDataVersion);
  ASSERT_TRUE(base::DirectoryExists(version_dir));

  // Make sure it actually checks the reported version against the real version.
  EXPECT_FALSE(loader.RegisterComponent(kTestComponentName, kTestUpdatedVersion,
                                        GetTestComponentPath().value()));

  // Now copy a new version into place.
  EXPECT_TRUE(
      loader.RegisterComponent(kTestComponentName, kTestUpdatedVersion,
                               GetTestComponentPath(kTestUpdatedVersion).value()));

  std::string hint_file_contents2;
  ASSERT_TRUE(
      base::ReadFileToStringWithMaxSize(hint_file, &hint_file_contents2, 4096));
  EXPECT_EQ(kTestUpdatedVersion, hint_file_contents2);

  base::FilePath version_dir2 = comp_dir.Append(kTestUpdatedVersion);
  ASSERT_TRUE(base::DirectoryExists(version_dir2));

  EXPECT_EQ(kTestUpdatedVersion,
            loader.GetComponentVersion(kTestComponentName));

  // Reject rollback to an older version.
  EXPECT_FALSE(loader.RegisterComponent(kTestComponentName, kTestDataVersion,
                                        GetTestComponentPath().value()));

  EXPECT_EQ(kTestUpdatedVersion,
            loader.GetComponentVersion(kTestComponentName));
}

// Pretend ImageLoader crashed, by creating an incomplete installation, and then
// attempt registration with ImageLoader.
TEST_F(ImageLoaderTest, RegisterComponentAfterCrash) {
  // Now create the junk there.
  const std::string junk_contents ="Bad file contents";
  const base::FilePath junk_path =
      temp_dir_.Append(kTestComponentName).Append(kTestDataVersion);
  ASSERT_TRUE(base::CreateDirectory(junk_path));
  ASSERT_EQ(static_cast<int>(junk_contents.size()),
            base::WriteFile(junk_path.Append("junkfile"), junk_contents.data(),
                            junk_contents.size()));
  ImageLoaderImpl loader(GetConfig(temp_dir_.value().c_str()));
  ASSERT_TRUE(loader.RegisterComponent(kTestComponentName, kTestDataVersion,
                                       GetTestComponentPath().value()));
}

TEST_F(ImageLoaderTest, MountValidImage) {
  std::vector<uint8_t> key(std::begin(kDevPublicKey), std::end(kDevPublicKey));
  auto helper_mock = base::MakeUnique<MockHelperProcess>();
  EXPECT_CALL(*helper_mock, SendMountCommand(_, _, _)).Times(2);
  ON_CALL(*helper_mock, SendMountCommand(_, _, _))
      .WillByDefault(testing::Return(true));

  base::ScopedTempDir scoped_mount_dir;
  ASSERT_TRUE(scoped_mount_dir.CreateUniqueTempDir());

  ImageLoaderConfig config(key, temp_dir_.value().c_str(),
                           scoped_mount_dir.path().value().c_str());
  ImageLoaderImpl loader(std::move(config));

  // We previously tested RegisterComponent, so assume this works if it reports
  // true.
  ASSERT_TRUE(loader.RegisterComponent(kTestComponentName, kTestDataVersion,
                                       GetTestComponentPath().value()));

  const std::string expected_path =
      scoped_mount_dir.path().value() + "/PepperFlashPlayer/22.0.0.158";
  EXPECT_EQ(expected_path,
            loader.LoadComponent(kTestComponentName, helper_mock.get()));

  // Let's also test mounting the component at a fixed point.
  const std::string expected_path2 =
      scoped_mount_dir.path().value() + "/FixedMountPoint";
  EXPECT_TRUE(loader.LoadComponent(kTestComponentName, expected_path2,
                                   helper_mock.get()));
}

TEST_F(ImageLoaderTest, MountInvalidImage) {
  std::vector<uint8_t> key(std::begin(kDevPublicKey), std::end(kDevPublicKey));
  auto helper_mock = base::MakeUnique<MockHelperProcess>();
  EXPECT_CALL(*helper_mock, SendMountCommand(_, _, _)).Times(0);
  ON_CALL(*helper_mock, SendMountCommand(_, _, _))
      .WillByDefault(testing::Return(true));

  base::ScopedTempDir scoped_mount_dir;
  ASSERT_TRUE(scoped_mount_dir.CreateUniqueTempDir());

  ImageLoaderConfig config(key, temp_dir_.value().c_str(),
                           scoped_mount_dir.path().value().c_str());
  ImageLoaderImpl loader(std::move(config));

  // We previously tested RegisterComponent, so assume this works if it reports
  // true.
  ASSERT_TRUE(loader.RegisterComponent(kTestComponentName, kTestDataVersion,
                                       GetTestComponentPath().value()));

  base::FilePath table = temp_dir_.Append(kTestComponentName)
                              .Append(kTestDataVersion)
                              .Append("table");
  std::string contents = "corrupt";
  ASSERT_EQ(static_cast<int>(contents.size()),
            base::WriteFile(table, contents.data(), contents.size()));
  ASSERT_EQ("", loader.LoadComponent(kTestComponentName, helper_mock.get()));
}

TEST_F(ImageLoaderTest, SetupTable) {
  std::string base_table = "0 40 verity payload=ROOT_DEV hashtree=HASH_DEV "
      "hashstart=40 alg=sha256 root_hexdigest="
      "34663b9920632778d38a0943a5472cae196bd4bf1d7dfa191506e7a8e7ec84d2 "
      "salt=fcfc9b5a329e44be73a323188ae75ca644122d920161f672f6935623831d07e2";

  // Make sure excess newlines are rejected.
  std::string bad_table = base_table + "\n\n";
  EXPECT_FALSE(VerityMounter::SetupTable(&bad_table, "/dev/loop6"));

  // Make sure it does the right replacements on a simple base table.
  std::string good_table = base_table;
  EXPECT_TRUE(VerityMounter::SetupTable(&good_table, "/dev/loop6"));

  std::string known_good_table =
      "0 40 verity payload=/dev/loop6 hashtree=/dev/loop6 "
      "hashstart=40 alg=sha256 root_hexdigest="
      "34663b9920632778d38a0943a5472cae196bd4bf1d7dfa191506e7a8e7ec84d2 "
      "salt=fcfc9b5a329e44be73a323188ae75ca644122d920161f672f6935623831d07e2 "
      "error_behavior=eio";
  EXPECT_EQ(known_good_table, good_table);

  // Make sure the newline is stripped.
  std::string good_table_newline = base_table + "\n";
  EXPECT_TRUE(VerityMounter::SetupTable(&good_table_newline, "/dev/loop6"));
  EXPECT_EQ(known_good_table, good_table_newline);

  // Make sure error_behavior isn't appended twice.
  std::string good_table_error = base_table + " error_behavior=eio\n";
  EXPECT_TRUE(VerityMounter::SetupTable(&good_table_error, "/dev/loop6"));
  EXPECT_EQ(known_good_table, good_table_error);
}

}   // namespace imageloader