/*
 * 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 <buildboxcommon_systemutils.h>
#include <buildboxcommon_temporarydirectory.h>
#include <buildboxrun_bubblewrap.h>

#include <gtest/gtest.h>

using namespace buildboxcommon;
using namespace buildboxrun;
using namespace bubblewrap;

bool contains(std::vector<std::string> vec, std::vector<std::string> subvec)
{
    auto val = vec.begin();
    auto match = subvec.begin();

    while (val != vec.end() && match != subvec.end()) {
        if (*val == *match) {
            match++;
        }
        else {
            match = subvec.begin();
        }

        val++;
    }

    return match == subvec.end();
}

BubbleWrapRunner testRunner()
{
    DigestGenerator::resetState();
    return BubbleWrapRunner("/dummy/bwrap");
}

// We test making a really simple command line; It's rather hard
// to actually test invoking `bwrap`, since this involves creating
// a RootDigest.
//
// Here we assert that what we've asked to be configured is actually
// configured; details of what the sandbox actually sandboxes aren't
// tested.
TEST(PrefixArgumentTests, TestCustomParser)
{
    std::string working_directory = SystemUtils::getCurrentWorkingDirectory();

    Command command;
    *command.add_arguments() = "./echo";
    command.set_working_directory("project");
    Command_EnvironmentVariable env;
    env.set_name("PATH");
    env.set_value("/test");
    *command.add_environment_variables() = env;

    Platform platform;
    BubbleWrapRunner runner = testRunner();

    std::vector<std::string> command_line =
        runner.generateCommandLine(command, working_directory, platform);

    EXPECT_TRUE(contains(command_line, {"--bind", working_directory, "/"}));
    EXPECT_TRUE(contains(command_line, {"--chdir", "project"}));
    EXPECT_TRUE(contains(command_line, {"--setenv", "PATH", "/test"}));
    if (BubbleWrapRunner::userNamespaceAvailable()) {
        EXPECT_TRUE(contains(command_line, {"--uid", "0"}));
        EXPECT_TRUE(contains(command_line, {"--gid", "0"}));
    }
    EXPECT_TRUE(contains(command_line, {"--unshare-net"}));
    EXPECT_TRUE(contains(command_line, {"./echo"}));
    EXPECT_TRUE(contains(command_line, {"--tmpfs", "/tmp"}));
}

TEST(PrefixArgumentTests, TestBindMount)
{
    Command command;
    *command.add_arguments() = "./echo";

    Platform platform;
    BubbleWrapRunner runner = testRunner();

    runner.parseArg("--bind-mount=/foo:/bar");

    std::vector<std::string> command_line =
        runner.generateCommandLine(command, "/", platform);

    EXPECT_TRUE(contains(command_line, {"--dev-bind", "/foo", "/bar"}));
}

TEST(PrefixArgumentTests, TestInputInTmp)
{
    TemporaryDirectory inputRoot;

    Command command;
    *command.add_arguments() = "./echo";
    command.set_working_directory("tmp");

    Platform platform;
    BubbleWrapRunner runner = testRunner();

    std::vector<std::string> command_line =
        runner.generateCommandLine(command, inputRoot.strname(), platform);

    EXPECT_FALSE(contains(command_line, {"--dir", "/tmp"}));
    EXPECT_TRUE(contains(command_line, {"--tmpfs", "/tmp"}));

    FileUtils::createDirectory((inputRoot.strname() + "/tmp").c_str());
    FileUtils::writeFileAtomically(inputRoot.strname() + "/tmp/inputfile.txt",
                                   "hello, world");

    std::vector<std::string> command_line2 =
        runner.generateCommandLine(command, inputRoot.strname(), platform);

    EXPECT_TRUE(contains(command_line2, {"--dir", "/tmp"}));
    EXPECT_FALSE(contains(command_line2, {"--tmpfs", "/tmp"}));
}

TEST(PrefixArgumentTests, TestEnableOutputsInTmpNoOutputs)
{
    TemporaryDirectory inputRoot;

    Command command;
    *command.add_arguments() = "./echo";

    Platform platform;
    BubbleWrapRunner runner = testRunner();

    runner.parseArg("--enable-tmp-outputs");

    std::vector<std::string> command_line =
        runner.generateCommandLine(command, inputRoot.strname(), platform);

    EXPECT_FALSE(contains(command_line, {"--dir", "/tmp"}));
    EXPECT_TRUE(contains(command_line, {"--tmpfs", "/tmp"}));
}

TEST(PrefixArgumentTests, TestNoTmpfsIfOutputInTmp)
{
    TemporaryDirectory inputRoot;

    Command command;
    *command.add_arguments() = "./echo";
    command.set_working_directory("tmp");
    *command.add_output_paths() = "foo.o";

    Platform platform;
    BubbleWrapRunner runner = testRunner();
    runner.parseArg("--enable-tmp-outputs");

    std::vector<std::string> command_line =
        runner.generateCommandLine(command, inputRoot.strname(), platform);

    EXPECT_TRUE(contains(command_line, {"--dir", "/tmp"}));
    EXPECT_FALSE(contains(command_line, {"--tmpfs", "/tmp"}));

    Command command2;
    *command2.add_arguments() = "./echo";
    *command2.add_output_paths() = "tmp/bar";

    std::vector<std::string> command_line2 =
        runner.generateCommandLine(command2, inputRoot.strname(), platform);

    EXPECT_TRUE(contains(command_line2, {"--dir", "/tmp"}));
    EXPECT_FALSE(contains(command_line2, {"--tmpfs", "/tmp"}));

    Command command3;
    command.set_working_directory("bar");
    *command3.add_arguments() = "./echo";
    *command3.add_output_paths() = "../tmp/foo";

    std::vector<std::string> command_line3 =
        runner.generateCommandLine(command3, inputRoot.strname(), platform);

    EXPECT_TRUE(contains(command_line3, {"--dir", "/tmp"}));
    EXPECT_FALSE(contains(command_line3, {"--tmpfs", "/tmp"}));
}

TEST(PrefixArgumentTests, TestTmpfsIfOutputInTmpWithoutEnableOption)
{
    TemporaryDirectory inputRoot;

    Command command;
    *command.add_arguments() = "./echo";
    command.set_working_directory("tmp");
    *command.add_output_paths() = "foo.o";

    Platform platform;
    BubbleWrapRunner runner = testRunner();

    std::vector<std::string> command_line =
        runner.generateCommandLine(command, inputRoot.strname(), platform);

    EXPECT_FALSE(contains(command_line, {"--dir", "/tmp"}));
    EXPECT_TRUE(contains(command_line, {"--tmpfs", "/tmp"}));

    Command command2;
    *command2.add_arguments() = "./echo";
    *command2.add_output_paths() = "tmp/bar";

    std::vector<std::string> command_line2 =
        runner.generateCommandLine(command2, inputRoot.strname(), platform);

    EXPECT_FALSE(contains(command_line2, {"--dir", "/tmp"}));
    EXPECT_TRUE(contains(command_line2, {"--tmpfs", "/tmp"}));
}

TEST(PrefixArgumentTests, TestUidGid)
{
    if (BubbleWrapRunner::userNamespaceAvailable()) {
        Command command;
        *command.add_arguments() = "./echo";

        auto platform = Platform();
        auto prop = platform.add_properties();
        prop->set_name("unixUID");
        prop->set_value("123");
        prop = platform.add_properties();
        prop->set_name("unixGID");
        prop->set_value("456");

        BubbleWrapRunner runner = testRunner();

        std::vector<std::string> command_line =
            runner.generateCommandLine(command, "/", platform);
        EXPECT_TRUE(contains(command_line, {"--uid", "123"}));
        EXPECT_TRUE(contains(command_line, {"--gid", "456"}));
    }
}

TEST(PrefixArgumentTests, TestInject)
{
    Command command;
    *command.add_arguments() = "./echo";

    auto platform = Platform();
    auto prop = platform.add_properties();
    prop->set_name("inject");
    prop->set_value("recc");

    BubbleWrapRunner runner = testRunner();
    EXPECT_TRUE(
        runner.parseArg("--ro-bind-mount=/localstore/bin/foo:/bin/bar"));
    EXPECT_TRUE(runner.parseArg(
        "--ro-bind-mount=recc:/localstore/bin/recc:/usr/bin/recc"));

    std::vector<std::string> command_line =
        runner.generateCommandLine(command, "/", platform);

    EXPECT_TRUE(contains(command_line,
                         {"--ro-bind", "/localstore/bin/foo", "/bin/bar"}));
    EXPECT_TRUE(contains(
        command_line, {"--ro-bind", "/localstore/bin/recc", "/usr/bin/recc"}));
}

TEST(PrefixArgumentTests, TestNetworkOff)
{
    Command command;
    *command.add_arguments() = "./echo";

    auto platform = Platform();
    auto prop = platform.add_properties();
    prop->set_name("network");
    prop->set_value("off");

    BubbleWrapRunner runner = testRunner();

    std::vector<std::string> command_line =
        runner.generateCommandLine(command, "/", platform);

    // Verify network sandboxing is enabled
    EXPECT_TRUE(contains(command_line, {"--unshare-net"}));
    EXPECT_TRUE(contains(command_line, {"--unshare-ipc"}));
    EXPECT_TRUE(contains(command_line, {"--unshare-uts"}));
}

TEST(PrefixArgumentTests, TestNetworkOn)
{
    Command command;
    *command.add_arguments() = "./echo";

    auto platform = Platform();
    auto prop = platform.add_properties();
    prop->set_name("network");
    prop->set_value("on");

    BubbleWrapRunner runner = testRunner();

    std::vector<std::string> command_line =
        runner.generateCommandLine(command, "/", platform);

    // Verify network sandboxing is disabled
    EXPECT_FALSE(contains(command_line, {"--unshare-net"}));
    EXPECT_FALSE(contains(command_line, {"--unshare-ipc"}));
    EXPECT_FALSE(contains(command_line, {"--unshare-uts"}));
}

TEST(PrefixArgumentTests, TestCapabilitiesWhitelist)
{
    if (BubbleWrapRunner::userNamespaceAvailable()) {
        Command command;
        *command.add_arguments() = "./echo";

        auto platform = Platform();
        auto prop = platform.add_properties();
        prop->set_name("capability");
        prop->set_value("CAP_DAC_OVERRIDE");
        prop = platform.add_properties();
        prop->set_name("capability");
        prop->set_value("CAP_FOWNER");

        BubbleWrapRunner runner = testRunner();

        std::vector<std::string> command_line =
            runner.generateCommandLine(command, "/", platform);
        EXPECT_TRUE(contains(command_line, {"--cap-add", "CAP_DAC_OVERRIDE"}));
        EXPECT_TRUE(contains(command_line, {"--cap-add", "CAP_FOWNER"}));

        // Ensure the first --cap-add option come after the --cap-drop ALL
        // option
        auto cap_drop_pos =
            std::find(command_line.begin(), command_line.end(), "--cap-drop");
        auto cap_add_pos =
            std::find(command_line.begin(), command_line.end(), "--cap-add");
        ASSERT_TRUE(cap_add_pos > cap_drop_pos);
    }
}

TEST(ArgumentParsingTests, TestCustomParser)
{
    BubbleWrapRunner runner = testRunner();
    EXPECT_TRUE(runner.parseArg("--enable-tmp-outputs"));
    EXPECT_TRUE(runner.parseArg("--bind-mount=foo:bar"));
    EXPECT_FALSE(runner.parseArg("--bind-mount=foo"));
    EXPECT_FALSE(runner.parseArg("--unknown-option"));
}

TEST(ExceptionTests, TestUnsupportedOSFamily)
{
    Command command;
    *command.add_arguments() = "./echo";

    auto platform = Platform();
    auto prop = platform.add_properties();
    prop->set_name("OSFamily");
    prop->set_value("UnsupportedOS");

    BubbleWrapRunner runner = testRunner();

    EXPECT_THROW(
        {
            try {
                runner.generateCommandLine(command, "/", platform);
            }
            catch (const std::runtime_error &e) {
                EXPECT_NE(std::string(e.what()).find(
                              "Unsupported OSFamily \"UnsupportedOS\""),
                          std::string::npos);
                throw;
            }
        },
        std::runtime_error);
}

TEST(ExceptionTests, TestUnsupportedISA)
{
    Command command;
    *command.add_arguments() = "./echo";

    auto platform = Platform();
    auto prop = platform.add_properties();
    prop->set_name("ISA");
    prop->set_value("UnsupportedISA");

    BubbleWrapRunner runner = testRunner();

    EXPECT_THROW(
        {
            try {
                runner.generateCommandLine(command, "/", platform);
            }
            catch (const std::runtime_error &e) {
                EXPECT_NE(std::string(e.what()).find(
                              "Unsupported ISA \"UnsupportedISA\""),
                          std::string::npos);
                throw;
            }
        },
        std::runtime_error);
}

TEST(ExceptionTests, TestUnsupportedCapability)
{
    if (BubbleWrapRunner::userNamespaceAvailable()) {
        Command command;
        *command.add_arguments() = "./echo";

        auto platform = Platform();
        auto prop = platform.add_properties();
        prop->set_name("capability");
        prop->set_value("CAP_SYS_ADMIN");

        BubbleWrapRunner runner = testRunner();

        EXPECT_THROW(
            {
                try {
                    runner.generateCommandLine(command, "/", platform);
                }
                catch (const std::runtime_error &e) {
                    EXPECT_NE(std::string(e.what()).find(
                                  "Unsupported capability \"CAP_SYS_ADMIN\""),
                              std::string::npos);
                    throw;
                }
            },
            std::runtime_error);
    }
}

TEST(ExceptionTests, TestUnixUIDWithoutUserNamespace)
{
    // This test can only run when user namespaces are NOT available
    if (!BubbleWrapRunner::userNamespaceAvailable()) {
        Command command;
        *command.add_arguments() = "./echo";

        auto platform = Platform();
        auto prop = platform.add_properties();
        prop->set_name("unixUID");
        prop->set_value("123");

        BubbleWrapRunner runner = testRunner();

        EXPECT_THROW(
            {
                try {
                    runner.generateCommandLine(command, "/", platform);
                }
                catch (const std::runtime_error &e) {
                    EXPECT_NE(std::string(e.what()).find(
                                  "User namespaces are not available. "
                                  "Cannot support "
                                  "unixUID."),
                              std::string::npos);
                    throw;
                }
            },
            std::runtime_error);
    }
}

TEST(ExceptionTests, TestUnixGIDWithoutUserNamespace)
{
    // This test can only run when user namespaces are NOT available
    if (!BubbleWrapRunner::userNamespaceAvailable()) {
        Command command;
        *command.add_arguments() = "./echo";

        auto platform = Platform();
        auto prop = platform.add_properties();
        prop->set_name("unixGID");
        prop->set_value("456");

        BubbleWrapRunner runner = testRunner();

        EXPECT_THROW(
            {
                try {
                    runner.generateCommandLine(command, "/", platform);
                }
                catch (const std::runtime_error &e) {
                    EXPECT_NE(
                        std::string(e.what()).find("User namespaces are not "
                                                   "available. Cannot support "
                                                   "unixGID."),
                        std::string::npos);
                    throw;
                }
            },
            std::runtime_error);
    }
}

TEST(ExceptionTests, TestCapabilityWithoutUserNamespace)
{
    // This test can only run when user namespaces are NOT available
    if (!BubbleWrapRunner::userNamespaceAvailable()) {
        Command command;
        *command.add_arguments() = "./echo";

        auto platform = Platform();
        auto prop = platform.add_properties();
        prop->set_name("capability");
        prop->set_value("CAP_NET_RAW");

        BubbleWrapRunner runner = testRunner();

        EXPECT_THROW(
            {
                try {
                    runner.generateCommandLine(command, "/", platform);
                }
                catch (const std::runtime_error &e) {
                    EXPECT_NE(std::string(e.what()).find(
                                  "User namespaces are not available. Cannot "
                                  "support adding capabilities."),
                              std::string::npos);
                    throw;
                }
            },
            std::runtime_error);
    }
}
