// Copyright 2019 The Fuchsia 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 <fbl/auto_call.h>
#include <fs-management/ramdisk.h>
#include <unittest/unittest.h>
#include <zircon/assert.h>
#include <zircon/device/device.h>
#include <zircon/device/ramdisk.h>

#include "libgpt-tests.h"

// generate a random number  between [1, max]
#define random_non_zero_length(max) ((random() % max) + 1)

extern bool gUseRamDisk;
extern char gDevPath[PATH_MAX];

namespace {

using gpt::GptDevice;

// TODO(auradkar): consolidate this guid_t definition with one in
// ulib/gpt/gpt.cpp.
struct guid_t {
    uint32_t data1;
    uint16_t data2;
    uint16_t data3;
    uint8_t data4[8];
};

constexpr guid_t kGuid = {0x0, 0x1, 0x2, {0x3, 0x4, 0x5, 0x6, 0x7, 0x8, 0x9, 0xa}};

constexpr uint64_t partition_size(const gpt_partition_t* p) {
    return p->last - p->first + 1;
}

bool destroy_gpt(int fd, uint64_t block_size, uint64_t offset,
                 uint64_t block_count) {
    BEGIN_HELPER;

    char zero[block_size];
    memset(zero, 0, sizeof(zero));

    ASSERT_GT(block_count, 0, "Block count should be greater than zero");
    ASSERT_GT(block_size, 0, "Block count should be greater than zero");

    uint64_t first = offset;
    uint64_t last = offset + block_count - 1;

    for (uint64_t i = first; i <= last; i++) {
        ASSERT_EQ(pwrite(fd, zero, sizeof(zero), block_size * i),
                  (ssize_t)sizeof(zero),
                  "Failed to pwrite");
    }

    ASSERT_EQ(fsync(fd), 0, "Failed to fsync");
    END_HELPER;
}

class Partitions {
public:
    Partitions(uint32_t count, uint64_t first, uint64_t last);

    // Returns partition at index index. Returns null if index is out of range.
    const gpt_partition_t* GetPartition(uint32_t index) const;

    // Returns number of parition created/removed.
    uint32_t GetCount() const { return partition_count_; }

    // Marks a partition as created in GPT.
    bool MarkCreated(uint32_t index);

    // Mark a partition as removed in GPT.
    bool ClearCreated(uint32_t index);

    // Returns true if the GPT should have the partition.
    bool IsCreated(uint32_t index) const;

    // Returns number of partition that should exist on GPT.
    uint32_t CreatedCount() const;

    // Returns true if two partitions are the same.
    bool Compare(const gpt_partition_t* in_mem_partition,
                 const gpt_partition_t* on_disk_partition) const;

    // Returns true if the partition p exists in partitions_.
    bool Find(const gpt_partition_t* p, uint32_t* out_index) const;

private:
    // List of partitions
    gpt_partition_t partitions_[gpt::kPartitionCount];

    // A variable to track whether a partition is created on GPT or not
    bool created_[gpt::kPartitionCount] = {};

    // Number of partitions_ that is populated with valid information
    uint32_t partition_count_;
};

Partitions::Partitions(uint32_t count, uint64_t first, uint64_t last) {
    ZX_ASSERT(count > 0);
    ZX_ASSERT(count <= gpt::kPartitionCount);
    partition_count_ = count;
    guid_t guid = kGuid;

    uint64_t part_first = first, part_last;
    uint64_t part_max_len = (last - first) / partition_count_;
    ZX_ASSERT(part_max_len > 0);

    memset(partitions_, 0, sizeof(partitions_));
    for (uint32_t i = 0; i < partition_count_; i++) {
        part_last = part_first + random_non_zero_length(part_max_len);
        guid.data1 = i;
        memcpy(partitions_[i].type, &guid, sizeof(partitions_[i].type));
        memcpy(partitions_[i].guid, &guid, sizeof(partitions_[i].type));
        partitions_[i].first = part_first;
        partitions_[i].last = part_last;
        partitions_[i].flags = 0;
        memset(partitions_[i].name, 0, sizeof(partitions_[i].name));
        snprintf(reinterpret_cast<char*>(partitions_[i].name),
                 sizeof(partitions_[i].name), "%u_part", i);

        part_first += part_max_len;
    }
}

const gpt_partition_t* Partitions::GetPartition(uint32_t index) const {
    if (index >= partition_count_) {
        return nullptr;
    }

    return &partitions_[index];
}

bool Partitions::MarkCreated(uint32_t index) {
    BEGIN_HELPER;
    ASSERT_LT(index, partition_count_, "Index out of range");

    created_[index] = true;
    END_HELPER;
}

bool Partitions::ClearCreated(uint32_t index) {
    BEGIN_HELPER;
    ASSERT_LT(index, partition_count_, "Index out of range");

    created_[index] = false;
    END_HELPER;
}

bool Partitions::IsCreated(uint32_t index) const {
    BEGIN_HELPER;
    ASSERT_LT(index, partition_count_, "Index out of range");

    return created_[index];
    END_HELPER;
}

uint32_t Partitions::CreatedCount() const {
    uint32_t created_count = 0;

    for (uint32_t i = 0; i < GetCount(); i++) {
        if (IsCreated(i)) {
            created_count++;
        }
    }
    return (created_count);
}

bool Partitions::Compare(const gpt_partition_t* in_mem_partition,
                         const gpt_partition_t* on_disk_partition) const {
    if (memcmp(in_mem_partition->type, on_disk_partition->type,
               sizeof(in_mem_partition->type)) != 0) {
        return false;
    }

    if (memcmp(in_mem_partition->guid, on_disk_partition->guid,
               sizeof(in_mem_partition->guid)) != 0) {
        return false;
    }

    if (in_mem_partition->first != on_disk_partition->first) {
        return false;
    }

    if (in_mem_partition->last != on_disk_partition->last) {
        return false;
    }

    if (in_mem_partition->flags != on_disk_partition->flags) {
        return false;
    }

    // In mem partition name is a c-string whereas on-disk partition name
    // is stored as UTF-16. We need to convert UTF-16 to c-string before we
    // compare.
    char name[GPT_NAME_LEN];
    memset(name, 0, GPT_NAME_LEN);
    utf16_to_cstring(name, (const uint16_t*)on_disk_partition->name,
                     GPT_NAME_LEN / 2);

    if (strncmp(name, reinterpret_cast<const char*>(in_mem_partition->name),
                GPT_NAME_LEN / 2) != 0) {
        return false;
    }

    return true;
}

bool Partitions::Find(const gpt_partition_t* p, uint32_t* out_index) const {
    for (uint32_t i = 0; i < partition_count_; i++) {
        if (Compare(GetPartition(i), p)) {
            *out_index = i;
            return true;
        }
    }

    return false;
}

// Defines a libgpt test function which can be passed to the TestWrapper.
typedef bool (*TestFunction)(LibGptTest* libGptTest);
#define RUN_TEST_WRAP(test_name) RUN_TEST_MEDIUM(TestWrapper<test_name>)

// A test wrapper which runs a libgpt test.
template <TestFunction TestFunc>
bool TestWrapper(void) {
    BEGIN_TEST;

    LibGptTest libGptTest(gUseRamDisk);
    ASSERT_TRUE(libGptTest.Init(), "Setting up the block device");
    // Run the test. This should pass.
    ASSERT_TRUE(TestFunc(&libGptTest));
    ASSERT_TRUE(libGptTest.Teardown(),
                "Tearing down and cleaning up the block device");

    END_TEST;
}

bool LibGptTest::Reset() {
    BEGIN_HELPER;
    fbl::unique_ptr<GptDevice> gpt;

    // explicitly close the fd, if open, before we attempt to reopen it.
    fd_.reset();

    fd_.reset(open(disk_path_, O_RDWR));

    ASSERT_TRUE(fd_.is_valid(), "Could not open block device\n");
    ASSERT_EQ(GptDevice::Create(fd_.get(), GetBlockSize(), GetBlockCount(),
                                &gpt),
              ZX_OK);
    gpt_ = std::move(gpt);
    END_HELPER;
}

bool LibGptTest::Finalize() {
    BEGIN_HELPER;
    ASSERT_FALSE(gpt_->Valid(), "Valid GPT on uninitialized disk");

    ASSERT_EQ(gpt_->Finalize(), ZX_OK, "Failed to finalize");
    ASSERT_TRUE(gpt_->Valid(), "Invalid GPT after finalize");
    END_HELPER;
}

bool LibGptTest::Sync() {
    BEGIN_HELPER;

    ASSERT_EQ(gpt_->Sync(), ZX_OK, "Failed to sync");
    ASSERT_TRUE(gpt_->Valid(), "Invalid GPT after sync");

    END_HELPER;
}

bool LibGptTest::ReadRange() {
    BEGIN_HELPER;

    gpt_->Range(&usable_start_block_, &usable_last_block_);

    // TODO(auradkar): GptDevice doesn't export api to get GPT-metadata size.
    // If it does, we can keep better track of metadata size it says it needs
    // and metadata it actually uses.
    ASSERT_LT(GetUsableStartBlock(), GetBlockCount(), "Range starts after EOD");
    ASSERT_LT(GetUsableStartBlock(), GetUsableLastBlock(), "Invalid range");
    ASSERT_LT(GetUsableLastBlock(), GetBlockCount(),
              "Range end greater than block count");
    ASSERT_GT(GetUsableBlockCount(), 0, "GPT occupied all available blocks");

    END_HELPER;
}

bool LibGptTest::PrepDisk(bool sync) {
    BEGIN_HELPER;

    if (sync) {
        Sync();
    } else {
        Finalize();
    }

    ASSERT_TRUE(ReadRange(), "Read range failed");
    END_HELPER;
}

bool LibGptTest::InitDisk(const char* disk_path) {
    BEGIN_HELPER;

    use_ramdisk_ = false;
    snprintf(disk_path_, PATH_MAX, "%s", disk_path);
    fbl::unique_fd fd(open(disk_path_, O_RDWR));
    ASSERT_TRUE(fd.is_valid(), "Could not open block device to fetch info");
    block_info_t block_info;
    ssize_t rc = ioctl_block_get_info(fd.get(), &block_info);
    ASSERT_EQ(rc, sizeof(block_info), "Could not query block device info");

    blk_size_ = block_info.block_size;
    blk_count_ = block_info.block_count;

    ASSERT_GE(GetDiskSize(), kAccptableMinimumSize,
              "Insufficient disk space for tests");
    fd_ = std::move(fd);

    END_HELPER;
}

bool LibGptTest::InitRamDisk() {
    BEGIN_HELPER;
    ASSERT_EQ(create_ramdisk(GetBlockSize(), GetBlockCount(), disk_path_),
              ZX_OK, "Could not create ramdisk");
    int fd = open(disk_path_, O_RDWR);
    if (fd < 0) {
        return false;
    }

    fd_.reset(fd);
    END_HELPER;
}

bool LibGptTest::Init() {
    BEGIN_HELPER;
    auto error = fbl::MakeAutoCall([this]() { Teardown(); });
    if (use_ramdisk_) {
        ASSERT_TRUE(InitRamDisk());
    } else {
        ASSERT_TRUE(InitDisk(gDevPath));
    }

    // TODO(auradkar): All tests assume that the disks don't have an initialized
    // disk. If tests find an GPT initialized disk at the beginning of test,
    // they fail. The tests leave disks in initalized state.
    //
    // To either uninitialize an initialized disk as a part of setup
    // test needs to know where gpt lies on the disk. As of now libgpt doesn't
    // export an api to get the location(s) of the gpt on disk. So, we assume
    // here that gpt lies in first few (GptMetadataBlocksCount()) blocks on the
    // device. We also ignore any backup copies on the device.
    // Once there exists an api in libgpt to get size and location(s) of gpt,
    // we can setup/cleanup before/after running tests in a better way.
    ASSERT_TRUE(destroy_gpt(fd_.get(), GetBlockSize(), 0,
                            GptMetadataBlocksCount()),
                "Failed to destroy gpt");

    ASSERT_TRUE(Reset());
    error.cancel();
    END_HELPER;
}

bool LibGptTest::TearDownDisk() {
    BEGIN_HELPER;
    ASSERT_FALSE(use_ramdisk_);
    END_HELPER;
}

bool LibGptTest::TearDownRamDisk() {
    BEGIN_HELPER;
    ASSERT_EQ(destroy_ramdisk(disk_path_), ZX_OK);
    END_HELPER;
}

bool LibGptTest::Teardown() {
    BEGIN_HELPER;

    if (use_ramdisk_) {
        ASSERT_TRUE(TearDownRamDisk());
    } else {
        ASSERT_TRUE(TearDownDisk());
    }

    END_HELPER;
}

} // namespace

// Creates "partitions->GetCount()"" number of partitions on GPT.
//  The information needed to create a partitions is passed in "partitions".
bool AddPartitionHelper(LibGptTest* libGptTest, Partitions* partitions) {
    BEGIN_HELPER;

    ASSERT_GT(partitions->GetCount(), 0, "At least one partition is required");
    for (uint32_t i = 0; i < partitions->GetCount(); i++) {
        const gpt_partition_t* p = partitions->GetPartition(i);
        ASSERT_EQ(libGptTest->AddPartition(reinterpret_cast<const char*>(p->name),
                                           p->type, p->guid, p->first,
                                           partition_size(p), p->flags),
                  0, "Add partition failed");
        partitions->MarkCreated(i);
    }

    END_HELPER;
}

// Removes randomly selected "remove_count" number of partitions.
bool RemovePartitionsHelper(LibGptTest* libGptTest,
                            Partitions* partitions,
                            uint32_t remove_count) {
    BEGIN_HELPER;
    uint32_t index;

    ASSERT_LE(remove_count, partitions->GetCount(),
              "Remove count exceeds whats available");
    ASSERT_LE(remove_count, partitions->CreatedCount(),
              "Cannot remove more partitions than created");

    for (uint32_t i = 0; i < remove_count; i++) {
        while (true) {
            index = static_cast<uint32_t>(random()) % partitions->GetCount();
            if (partitions->IsCreated(index)) {
                break;
            }
        }
        ASSERT_TRUE(partitions->IsCreated(index),
                    "Partition already removed");
        const gpt_partition_t* p = partitions->GetPartition(index);
        ASSERT_EQ(libGptTest->RemovePartition(p->guid), 0,
                  "Failed to remove partition");
        partitions->ClearCreated(index);
    }
    END_HELPER;
}

// Verifies all the partitions that exists on GPT are the ones that are created
// by the test and vice-versa.
bool PartitionVerify(LibGptTest* libGptTest, const Partitions* partitions) {
    BEGIN_HELPER;
    bool found[gpt::kPartitionCount] = {};
    uint32_t found_index;

    // Check what's found on disk is created by us
    // iteratre over all partition that are present on disk and make sure
    // that we intended to create them.
    // Note: The index of an entry/partition need not match with the index of
    // the partition in "Partition* partition".
    for (uint32_t i = 0; i < gpt::kPartitionCount; i++) {
        gpt_partition_t* p = libGptTest->GetPartition(i);

        if (p == NULL) {
            continue;
        }

        ASSERT_TRUE(partitions->Find(p, &found_index),
                    "Found an entry on GPT that we did not create");

        ASSERT_TRUE(partitions->IsCreated(found_index),
                    "Removed entry reincarnated");
        found[found_index] = true;
    }

    // Check what's created is found on disk
    for (uint32_t i = 0; i < partitions->GetCount(); i++) {
        if (partitions->IsCreated(i)) {
            ASSERT_TRUE(found[i], "Created partition is missing on disk");
        }
    }

    END_TEST;
}

// Creates partitions and verifies them.
bool AddPartitions(LibGptTest* libGptTest, Partitions* partitions, bool sync) {
    BEGIN_HELPER;

    ASSERT_TRUE(AddPartitionHelper(libGptTest, partitions),
                "AddPartitionHelper failed");

    if (sync) {
        ASSERT_TRUE(libGptTest->Sync(), "Sync failed");
    }

    ASSERT_TRUE(PartitionVerify(libGptTest, partitions),
                "Partition verify failed");
    ASSERT_EQ(partitions->GetCount(), partitions->CreatedCount(),
              "Not as many created as we wanted to");

    END_HELPER;
}

// Removes partitions and verifies them.
bool RemovePartitions(LibGptTest* libGptTest, Partitions* partitions,
                      uint32_t remove_count, bool sync) {
    BEGIN_HELPER;

    ASSERT_TRUE(RemovePartitionsHelper(libGptTest, partitions, remove_count),
                "RemovePartitionsHelper failed");
    if (sync) {
        ASSERT_TRUE(libGptTest->Sync(), "Sync failed");
    }

    ASSERT_TRUE(PartitionVerify(libGptTest, partitions),
                "Partition verify failed");
    ASSERT_EQ(partitions->GetCount() - partitions->CreatedCount(), remove_count,
              "Not as many removed as we wanted to");

    END_HELPER;
}

// Removes all partitions and verifies them.
bool RemoveAllPartitions(LibGptTest* libGptTest, Partitions* partitions, bool sync) {
    BEGIN_HELPER;

    ASSERT_LE(partitions->GetCount(), partitions->CreatedCount(),
              "Not all partitions populated");
    ASSERT_EQ(libGptTest->RemoveAllPartitions(), 0,
              "Failed to remove all partition");

    for (uint32_t i = 0; i < partitions->GetCount(); i++) {
        partitions->ClearCreated(i);
    }

    ASSERT_TRUE(PartitionVerify(libGptTest, partitions),
                "Partition verify failed");
    ASSERT_EQ(partitions->CreatedCount(), 0,
              "Not as many removed as we wanted to");
    END_HELPER;
}

// Tests if we can create a GptDevice.
bool CreateTest(LibGptTest* libGptTest) {
    BEGIN_TEST;

    ASSERT_FALSE(libGptTest->IsGptValid(),
                 "Valid GPT on uninitialized disk");
    ASSERT_TRUE(libGptTest->Reset(), "Failed to reset Test");
    ASSERT_FALSE(libGptTest->IsGptValid(), "Valid GPT after reset");
    END_TEST;
}

// Tests Finalize initializes GPT in-memory only and doesn't commit to disk.
bool FinalizeTest(LibGptTest* libGptTest) {
    BEGIN_TEST;

    ASSERT_TRUE(libGptTest->Finalize(), "Finalize failed");

    // Finalize initializes GPT but doesn't write changes to disk.
    // Resetting the Test should bring invalid gpt back.
    ASSERT_TRUE(libGptTest->Reset(), "Failed to reset Test");
    ASSERT_FALSE(libGptTest->IsGptValid(),
                 "Valid GPT after finalize and reset");
    END_TEST;
}

// Tests Finalize initializes GPT and writes it to disk.
bool SyncTest(LibGptTest* libGptTest) {
    BEGIN_TEST;

    ASSERT_FALSE(libGptTest->IsGptValid(),
                 "Valid GPT on uninitialized disk");

    // Sync should write changes to disk. Resetting should bring valid gpt back.
    ASSERT_TRUE(libGptTest->Sync(), "Sync failed");
    ASSERT_TRUE(libGptTest->Reset(), "Failed to reset Test");
    ASSERT_TRUE(libGptTest->IsGptValid(), "Invalid GPT after sync and reset");
    END_TEST;
}

// Tests the range the GPT blocks falls within disk.
bool RangeTest(LibGptTest* libGptTest) {
    BEGIN_TEST;

    ASSERT_TRUE(libGptTest->Finalize(), "Finalize failed");
    ASSERT_TRUE(libGptTest->ReadRange(), "Failed to read range");

    END_TEST;
}

// Test adding total_partitions partitions to GPT w/o writing to disk
template <uint32_t total_partitions, bool sync>
bool AddPartitionTest(LibGptTest* libGptTest) {
    BEGIN_TEST;
    ASSERT_TRUE(libGptTest->PrepDisk(sync), "Failed to setup disk");

    Partitions partitions(total_partitions, libGptTest->GetUsableStartBlock(),
                          libGptTest->GetUsableLastBlock());

    ASSERT_TRUE(AddPartitions(libGptTest, &partitions, sync),
                "AddPartitions failed");
    END_TEST;
}

// Test adding total_partitions partitions to GPT and test removing remove_count
// partitions later w/o writing to disk.
template <uint32_t total_partitions, uint32_t remove_count, bool sync>
bool RemovePartitionTest(LibGptTest* libGptTest) {
    BEGIN_TEST;
    ASSERT_TRUE(libGptTest->PrepDisk(sync), "Failed to setup disk");

    Partitions partitions(total_partitions, libGptTest->GetUsableStartBlock(),
                          libGptTest->GetUsableLastBlock());

    ASSERT_TRUE(AddPartitions(libGptTest, &partitions, sync),
                "AddPartitions failed");
    ASSERT_TRUE(RemovePartitions(libGptTest, &partitions, remove_count, sync),
                "RemovePartitions failed");

    END_TEST;
}

// Test removing all total_partititions from GPT w/o syncing.
template <uint32_t total_partitions, bool sync>
bool RemovePartitionAllTest(LibGptTest* libGptTest) {
    BEGIN_TEST;

    ASSERT_TRUE(libGptTest->PrepDisk(sync), "Failed to setup disk");

    Partitions partitions(total_partitions, libGptTest->GetUsableStartBlock(),
                          libGptTest->GetUsableLastBlock());

    ASSERT_TRUE(AddPartitions(libGptTest, &partitions, sync),
                "AddPartitions failed");
    ASSERT_TRUE(RemoveAllPartitions(libGptTest, &partitions, sync),
                "RemoveAllPartitions failed");

    END_TEST;
}

// Test if Diffs after adding partitions reflect all the changes.
template <uint32_t total_partitions>
bool DiffsTest(LibGptTest* libGptTest) {
    BEGIN_TEST;
    uint32_t diffs;

    ASSERT_NE(libGptTest->GetDiffs(0, &diffs), 0,
              "GetDiffs should fail before PrepDisk");
    ASSERT_TRUE(libGptTest->PrepDisk(false), "Failed to setup disk");
    ASSERT_NE(libGptTest->GetDiffs(0, &diffs), 0,
              "GetDiffs for non-existing partition should fail");

    Partitions partitions(total_partitions, libGptTest->GetUsableStartBlock(),
                          libGptTest->GetUsableLastBlock());
    ASSERT_TRUE(AddPartitions(libGptTest, &partitions, false),
                "AddPartitions failed");
    ASSERT_EQ(libGptTest->GetDiffs(0, &diffs), 0,
              "Diffs zero after adding partition");

    ASSERT_EQ(diffs, GPT_DIFF_TYPE | GPT_DIFF_GUID | GPT_DIFF_FIRST | GPT_DIFF_LAST | GPT_DIFF_NAME,
              "Unexpected diff after creating partition");
    ASSERT_TRUE(libGptTest->Sync(), "Failed to sync");
    ASSERT_EQ(libGptTest->GetDiffs(0, &diffs), 0, "GetDiffs failed");
    ASSERT_EQ(diffs, 0, "Diffs not zero after syncing partition");

    END_TEST;
}

BEGIN_TEST_CASE(libgpt_tests)
RUN_TEST_WRAP(CreateTest)
RUN_TEST_WRAP(FinalizeTest)
RUN_TEST_WRAP(SyncTest)
RUN_TEST_WRAP(RangeTest)
RUN_TEST_WRAP((AddPartitionTest<3, false>))
RUN_TEST_WRAP((AddPartitionTest<20, true>))
RUN_TEST_WRAP((RemovePartitionTest<12, 4, false>))
RUN_TEST_WRAP((RemovePartitionTest<3, 2, true>))
RUN_TEST_WRAP((RemovePartitionTest<11, 11, false>))
RUN_TEST_WRAP((RemovePartitionAllTest<12, true>))
RUN_TEST_WRAP((RemovePartitionAllTest<15, false>))
RUN_TEST_WRAP((DiffsTest<9>))
END_TEST_CASE(libgpt_tests);