/*
 * Copyright 2019 Bloomberg Finance LP
 *
 * 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 <buildboxrun_userchroot.h>

#include <buildboxcommon_digestgenerator.h>
#include <buildboxcommon_fileutils.h>
#include <buildboxcommon_temporarydirectory.h>

#include <gtest/gtest.h>

using namespace buildboxcommon;
using namespace buildboxrun;
using namespace userchroot;

namespace {
const auto digestFunctionInitializer = []() {
    buildboxcommon::DigestGenerator::init();
    return 0;
}();
} // namespace

TEST(PrefixArgumentTests, TestCustomParser)
{
    UserChrootRunner runner;
    EXPECT_FALSE(runner.parseArg("--userchroot-bin"));
    EXPECT_TRUE(runner.parseArg("--userchroot-bin=/bin/userchroot"));
    EXPECT_FALSE(runner.parseArg("--some-other-argument=/bin/userchroot"));
    EXPECT_FALSE(runner.parseArg("--verbose"));
    EXPECT_TRUE(
        runner.parseArg("--local-only-cas-instance-name=local-instance"));
    EXPECT_FALSE(runner.parseArg("--sudo-bin"));
    EXPECT_TRUE(runner.parseArg("--sudo-bin=/bin/sudo"));
    EXPECT_TRUE(runner.parseArg("--find-bin=/my/path/bin/find"));
    EXPECT_TRUE(runner.parseArg("--chmod-bin=/my/path/bin/chmod"));
    EXPECT_TRUE(runner.parseArg("--user=foo"));
    EXPECT_TRUE(runner.parseArg("--group=bar"));
    EXPECT_TRUE(runner.parseArg("--inject-path=local/foo:remote/foo"));
    EXPECT_TRUE(runner.parseArg("--inject-path=recc:local/recc:remote/recc"));
    EXPECT_TRUE(runner.parseArg(
        "--inject-path=recc:local/recc_conf:remote/recc_conf"));
    EXPECT_FALSE(
        runner.parseArg("--inject-path=recc:recc:local/recc:remote/recc"));

    EXPECT_EQ(runner.user(), "foo");
    EXPECT_EQ(runner.group(), "bar");
    EXPECT_EQ(runner.userchrootBin(), "/bin/userchroot");
    EXPECT_EQ(2, runner.mappedPathsByProperties().at("recc").size());
}

class UserChrootRunnerValidateChrootFixture : public UserChrootRunner,
                                              public ::testing::Test {

  protected:
    UserChrootRunnerValidateChrootFixture()
    {
        this->d_stage_path = stage_directory.name();
        this->setOutputPath(this->d_stage_path + "/");
    }

    buildboxcommon::TemporaryDirectory stage_directory;
};

TEST_F(UserChrootRunnerValidateChrootFixture, TestValidChroot)
{
    // A valid chroot must contain `/tmp` and `/dev`:
    buildboxcommon::FileUtils::createDirectory(
        (this->d_stage_path + "/tmp").c_str());

    buildboxcommon::FileUtils::createDirectory(
        (this->d_stage_path + "/dev").c_str());

    ASSERT_NO_THROW(this->assertRequiredDirectoriesExist());
}

TEST_F(UserChrootRunnerValidateChrootFixture, TestChrootMissingDirectories)
{
    buildboxcommon::FileUtils::createDirectory(
        (this->d_stage_path + "/dir1").c_str());

    buildboxcommon::FileUtils::createDirectory(
        (this->d_stage_path + "/tmp").c_str());

    ASSERT_THROW(
        {
            try {
                this->assertRequiredDirectoriesExist();
            }
            catch (const GrpcError &e) {
                std::ostringstream errMsg;
                errMsg
                    << "Staged chroot does not contain expected directory \""
                    << "/dev"
                    << "\"";
                const grpc::Status grpcStatus(
                    grpc::StatusCode::FAILED_PRECONDITION, errMsg.str());
                ASSERT_EQ(e.status.error_code(), grpcStatus.error_code());
                ASSERT_EQ(e.status.error_message(),
                          grpcStatus.error_message());
                throw;
            }
        },
        GrpcError);
}

TEST_F(UserChrootRunnerValidateChrootFixture, TestErrorReturnsEmptyDigest)
{
    DirectoryTree inputTree;
    DirectoryTree chrootTree;

    digest_string_map blobs;
    ASSERT_EQ(this->mergeTreesAndSetRootDirectoryProperty(inputTree,
                                                          chrootTree, &blobs),
              Digest());
}

void prepareChrootTree(DirectoryTree *tree)
{
    /* Creates the following directory structure:
     *
     * ./
     *   include/
     *           time.h
     *           sys/
     *               stat.h
     *   local/
     *         lib/
     *             libc.so
     *   var/
     */

    // ./include/sys
    Directory sys_directory;
    FileNode *sysFileNodes = sys_directory.add_files();
    sysFileNodes->set_name("stat.h");
    sysFileNodes->set_is_executable(false);
    sysFileNodes->mutable_digest()->CopyFrom(
        DigestGenerator::hash("stat_h_contents"));
    const auto sys_directory_digest = DigestGenerator::hash(sys_directory);

    // ./include
    Directory include_directory;
    FileNode *includeFileNodes = include_directory.add_files();
    includeFileNodes->set_name("time.h");
    includeFileNodes->set_is_executable(false);
    includeFileNodes->mutable_digest()->CopyFrom(
        DigestGenerator::hash("time_h_contents"));
    DirectoryNode *sysNode = include_directory.add_directories();
    sysNode->set_name("sys");
    sysNode->mutable_digest()->CopyFrom(sys_directory_digest);
    const auto include_directory_digest =
        DigestGenerator::hash(include_directory);

    // ./local/lib
    Directory lib_directory;
    FileNode *libFileNode = lib_directory.add_files();
    libFileNode->set_name("libc.so");
    libFileNode->set_is_executable(false);
    libFileNode->mutable_digest()->CopyFrom(
        DigestGenerator::hash("libc_so_contents"));
    const auto lib_directory_digest = DigestGenerator::hash(lib_directory);

    // ./local
    Directory local_directory;
    DirectoryNode *libNode = local_directory.add_directories();
    libNode->set_name("lib");
    libNode->mutable_digest()->CopyFrom(lib_directory_digest);
    const auto local_directory_digest = DigestGenerator::hash(local_directory);

    // ./var
    Directory var_directory;
    const auto var_directory_digest = DigestGenerator::hash(var_directory);

    // .
    Directory root_directory;
    // add include to root
    DirectoryNode *includeNode = root_directory.add_directories();
    includeNode->set_name("include");
    includeNode->mutable_digest()->CopyFrom(include_directory_digest);

    // add local to root
    DirectoryNode *localNode = root_directory.add_directories();
    localNode->set_name("local");
    localNode->mutable_digest()->CopyFrom(local_directory_digest);

    // add var to root
    DirectoryNode *varNode = root_directory.add_directories();
    varNode->set_name("var");
    varNode->mutable_digest()->CopyFrom(var_directory_digest);

    // create the tree
    tree->root = DigestGenerator::hash(root_directory);
    tree->add(root_directory);
    tree->add(include_directory);
    tree->add(sys_directory);
    tree->add(local_directory);
    tree->add(lib_directory);
    tree->add(var_directory);
}

void prepareInputTree(DirectoryTree *tree, const bool isExecutable = true)
{
    /* Creates the following directory structure:
     *
     * ./
     *   include/
     *          time.h
     *   src/
     *       build.sh*
     *       headers/
     *               file1.h
     *               file2.h
     *               file3.h
     *       cpp/
     *           file1.cpp
     *           file2.cpp
     *           file3.cpp
     *           symlink: file4.cpp --> file3.cpp
     */

    // ./include/time.h

    Directory include_directory;
    FileNode *includeFileNodes = include_directory.add_files();
    includeFileNodes->set_name("time.h");
    includeFileNodes->set_is_executable(false);
    includeFileNodes->mutable_digest()->CopyFrom(
        DigestGenerator::hash("time_h_contents_input_root"));

    const auto include_directory_digest =
        DigestGenerator::hash(include_directory);

    // ./src/headers
    Directory headers_directory;
    std::vector<std::string> headerFiles = {"file1.h", "file2.h", "file3.h"};
    for (const auto &file : headerFiles) {
        FileNode *fileNode = headers_directory.add_files();
        fileNode->set_name(file);
        fileNode->set_is_executable(false);
        fileNode->mutable_digest()->CopyFrom(
            DigestGenerator::hash(file + "_contents"));
    }
    const auto headers_directory_digest =
        DigestGenerator::hash(headers_directory);

    // ./src/cpp
    Directory cpp_directory;
    std::vector<std::string> cppFiles = {"file1.cpp", "file2.cpp",
                                         "file3.cpp"};
    for (const auto &file : cppFiles) {
        FileNode *fileNode = cpp_directory.add_files();
        fileNode->set_name(file);
        fileNode->set_is_executable(false);
        fileNode->mutable_digest()->CopyFrom(
            DigestGenerator::hash(file + "_contents"));
    }
    SymlinkNode *symNode = cpp_directory.add_symlinks();
    symNode->set_name("file4.cpp");
    symNode->set_target("file3.cpp");
    const auto cpp_directory_digest = DigestGenerator::hash(cpp_directory);

    // ./src
    Directory src_directory;
    DirectoryNode *headersNode = src_directory.add_directories();
    headersNode->set_name("headers");
    headersNode->mutable_digest()->CopyFrom(headers_directory_digest);
    DirectoryNode *cppNode = src_directory.add_directories();
    cppNode->set_name("cpp");
    cppNode->mutable_digest()->CopyFrom(cpp_directory_digest);
    FileNode *fileNode = src_directory.add_files();
    fileNode->set_name("build.sh");
    fileNode->set_is_executable(isExecutable);
    fileNode->mutable_digest()->CopyFrom(
        DigestGenerator::hash("build.sh_contents"));
    const auto src_directory_digest = DigestGenerator::hash(src_directory);

    // .
    Directory root_directory;
    DirectoryNode *srcNode = root_directory.add_directories();
    srcNode->set_name("src");
    srcNode->mutable_digest()->CopyFrom(src_directory_digest);

    // ./include
    DirectoryNode *includeNode = root_directory.add_directories();
    includeNode->set_name("include");
    includeNode->mutable_digest()->CopyFrom(include_directory_digest);

    // create the tree
    tree->root = DigestGenerator::hash(root_directory);
    tree->add(root_directory);
    tree->add(include_directory);
    tree->add(src_directory);
    tree->add(headers_directory);
    tree->add(cpp_directory);
}

TEST_F(UserChrootRunnerValidateChrootFixture, TestUnixModePropertyIsSet)
{
    DirectoryTree inputTree;
    prepareInputTree(&inputTree);

    DirectoryTree chrootTree;
    prepareChrootTree(&chrootTree);

    digest_string_map blobs;

    const Digest mergedRootDigest =
        this->mergeTreesAndSetRootDirectoryProperty(inputTree, chrootTree,
                                                    &blobs);
    ASSERT_NE(mergedRootDigest, Digest());

    Directory mergedRootDirectory;
    ASSERT_TRUE(
        mergedRootDirectory.ParseFromString(blobs.at(mergedRootDigest)));

    ASSERT_EQ(mergedRootDirectory.node_properties().unix_mode().value(), 0755);
}

TEST_F(UserChrootRunnerValidateChrootFixture, TestMergedTreeIsUnchanged)
{
    DirectoryTree inputTree;
    prepareInputTree(&inputTree);

    DirectoryTree chrootTree;
    prepareChrootTree(&chrootTree);

    // Merging the tree with the `UserChrootRunner` helper:
    digest_string_map mergedTreeBlobs;
    const Digest mergedRootDigest =
        this->mergeTreesAndSetRootDirectoryProperty(inputTree, chrootTree,
                                                    &mergedTreeBlobs);

    ASSERT_NE(mergedRootDigest, Digest());

    // And using `MergeUtil` for comparison.
    // This will produce the same tree only that the root won't have its
    // `unix_mode` field set.
    Digest referenceMergedRootDigest;
    digest_string_map referenceBlobs;
    const std::vector<DirectoryTree> treesToMerge = {chrootTree, inputTree};
    ASSERT_TRUE(MergeUtil::createMergedLayersDigest(
        treesToMerge, &referenceMergedRootDigest, &referenceBlobs));

    ASSERT_EQ(mergedTreeBlobs.size(), referenceBlobs.size());

    // The root `Directory`'s are equal, except for their
    // `node_properties.unix_mode`. Therefore, after removing that field their
    // digests match.
    Directory mergedRootDirectory;
    ASSERT_TRUE(mergedRootDirectory.ParseFromString(
        mergedTreeBlobs.at(mergedRootDigest)));

    ASSERT_EQ(mergedRootDirectory.node_properties().unix_mode().value(), 0755);
    mergedRootDirectory.clear_node_properties();

    ASSERT_EQ(DigestGenerator::hash(mergedRootDirectory.SerializeAsString()),
              referenceMergedRootDigest);

    // The rest of the tree was untouched, so it matches the one produced by
    // `MergeUtil`:
    for (const auto &entry : mergedTreeBlobs) {
        const Digest digest = entry.first;
        if (digest != mergedRootDigest) {
            ASSERT_EQ(referenceBlobs.count(digest), 1);
        }
    }
}

TEST_F(UserChrootRunnerValidateChrootFixture, TestInputRootOverwrite)
{
    DirectoryTree inputTree;
    prepareInputTree(&inputTree);

    DirectoryTree chrootTree;
    prepareChrootTree(&chrootTree);

    // Merging the tree with the `UserChrootRunner` helper:
    digest_string_map mergedTreeBlobs;
    const Digest mergedRootDigest =
        this->mergeTreesAndSetRootDirectoryProperty(inputTree, chrootTree,
                                                    &mergedTreeBlobs);

    Directory mergedRootDirectory;
    ASSERT_TRUE(mergedRootDirectory.ParseFromString(
        mergedTreeBlobs.at(mergedRootDigest)));

    Digest includeDirectoryDigest = Digest();
    Directory includeDirectory;

    for (auto &directoryNode : mergedRootDirectory.directories()) {
        if (directoryNode.name() == "include") {
            includeDirectoryDigest = directoryNode.digest();
            ASSERT_TRUE(includeDirectory.ParseFromString(
                mergedTreeBlobs.at(includeDirectoryDigest)));
        }
    }
    ASSERT_TRUE(includeDirectoryDigest != Digest());

    Digest timeHFileDigest = Digest();

    for (auto &fileNode : includeDirectory.files()) {
        if (fileNode.name() == "time.h") {
            //            ASSERT_EQ("", fileNode.SerializeAsString());
            timeHFileDigest = fileNode.digest();
        }
    }
    ASSERT_EQ(timeHFileDigest,
              DigestGenerator::hash("time_h_contents_input_root"));
}
