/*
 * 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 <buildboxrun_chrootmanager.h>

#include <buildboxcommon_casclient.h>
#include <buildboxcommon_exception.h>
#include <buildboxcommon_fileutils.h>
#include <buildboxcommon_grpcclient.h>
#include <buildboxcommon_grpcerror.h>
#include <buildboxcommon_logging.h>
#include <buildboxcommon_mergeutil.h>
#include <buildboxcommon_permissions.h>
#include <buildboxcommon_runner.h>
#include <buildboxcommon_scopeguard.h>
#include <buildboxcommon_stringutils.h>

#include <array>
#include <filesystem>
#include <optional>
#include <stdlib.h>
#include <string.h>
#include <string>
#include <unistd.h>

namespace buildboxcommon {
namespace buildboxrun {
namespace userchroot {

ActionResult UserChrootRunner::execute(const Command &command,
                                       const Digest &inputRootDigest,
                                       const Platform &platform)
{
    // Check the Platform properties to see if we need to merge
    Digest digest(inputRootDigest);
    std::optional<Digest> chrootTreeDigest = hasChrootTreeDigest(platform);
    std::optional<Digest> chrootRootDigest = hasChrootRootDigest(platform);

    ActionResult result;

    auto *result_metadata = result.mutable_execution_metadata();
    Runner::metadata_mark_input_download_start(result_metadata);
    if (chrootTreeDigest) {
        // fetchChrootIntoLocalCas() has been omitted here as it doesn't work
        // with tree digests and will be removed from the chrootRootDigest case
        // above in a future commit because pre-loading the chroot into local
        // CAS no longer provides the desired performance optimisation from
        // commit 60126ba9 as the tree cache was expanded to also cover
        // subdirectories in commit 3dca1f5d.

        Digest mergedDigest;
        mergeDigests(inputRootDigest, chrootTreeDigest.value(), &mergedDigest,
                     Command::TREE_ONLY);
        BUILDBOX_RUNNER_LOG(DEBUG, "inputRootDigest = "
                                       << inputRootDigest
                                       << ", chrootTreeDigest = "
                                       << chrootTreeDigest.value()
                                       << ", mergedDigest = " << mergedDigest);
        digest.Swap(&mergedDigest);
    }
    else if (chrootRootDigest) {
        if (d_use_localcas_protocol) {
            // Making sure that the chroot is pre-loaded in local CAS if
            // connected to one.

            fetchChrootIntoLocalCas(chrootRootDigest.value());
        }

        Digest mergedDigest;
        mergeDigests(inputRootDigest, chrootRootDigest.value(), &mergedDigest,
                     Command::DIRECTORY_ONLY);
        BUILDBOX_RUNNER_LOG(DEBUG, "inputRootDigest = "
                                       << inputRootDigest
                                       << ", chrootRootDigest = "
                                       << chrootRootDigest.value()
                                       << ", mergedDigest = " << mergedDigest);
        digest.Swap(&mergedDigest);
    }
    else {
        // The root directory must have `unix_mode = 0755` to signal casd that
        // it requires to be staged with those permissions.
        Directory rootDirectory =
            d_casClient->fetchMessage<Directory>(inputRootDigest);

        const auto rootDirectoryUnixMode =
            rootDirectory.node_properties().unix_mode().value();
        if (rootDirectoryUnixMode != s_rootDirectoryRequiredUnixMode) {
            BUILDBOX_RUNNER_LOG(
                DEBUG,
                "inputRootDigest="
                    << inputRootDigest
                    << " does not have node_properties.unix_mode == "
                    << s_rootDirectoryRequiredUnixMode
                    << ". Setting it and uploading the new Directory to CAS.");

            rootDirectory.mutable_node_properties()
                ->mutable_unix_mode()
                ->set_value(s_rootDirectoryRequiredUnixMode);

            const Digest newRootDirectoryDigest =
                d_casClient->uploadMessage(rootDirectory);

            BUILDBOX_RUNNER_LOG(DEBUG, "inputRootDigest="
                                           << inputRootDigest
                                           << " is now rooted at "
                                           << newRootDirectoryDigest << ".");

            digest = newRootDirectoryDigest;
        }
    }

    const auto stagedDir = this->stageDirectory(
        digest, [this]() { return this->getProcessCredentials(); });
    Runner::metadata_mark_input_download_end(result_metadata);

    // Injecting files
    for (const auto &[srcPath, chrootPath] : d_mappedPaths) {
        injectPath(srcPath, chrootPath);
    }
    // Injecting files indexed by platform properties
    for (const auto &property : platform.properties()) {
        if (property.name() == "inject") {
            const auto &content = property.value();
            if (d_mappedPathsByProperties.find(content) ==
                d_mappedPathsByProperties.end()) {
                BUILDBOX_LOG_WARNING(
                    "Unrecognized content by property to inject: " << content);
                continue;
            }

            for (const auto &[srcPath, chrootPath] :
                 d_mappedPathsByProperties.at(content)) {
                BUILDBOX_LOG_DEBUG("Injecting content by property: "
                                   << content << ":" << srcPath << ":"
                                   << chrootPath);
                injectPath(srcPath, chrootPath);
            }
        }
    }

    std::ostringstream workingDir;
    workingDir << stagedDir->getPath() << "/" << command.working_directory();

    // Verify that the required directories (such as "/dev") exist in the
    // staged directory.
    // If that is not the case, throw and cause the runner to return that error
    // in the `ActionResult`:
    assertRequiredDirectoriesExist();

    // Making sure that the permissions of the stage directory are set to at
    // least 0770 when this function is done so that it can later be cleaned
    // up.
    const std::string stagedDirPath(stagedDir->getPath());
    const buildboxcommon::ScopeGuard chmodGuard([this, stagedDirPath]() {
        // Need sudo subprocess
        if (!d_user.empty()) {
            const auto preUnstageCommand =
                createPreUnstageCommand(stagedDirPath);
            BUILDBOX_RUNNER_LOG(DEBUG, "Executing pre-unstage command"
                                           << logging::printableCommandLine(
                                                  preUnstageCommand));
            int preUnstageRc =
                SystemUtils::executeCommandAndWait(preUnstageCommand, true);
            if (preUnstageRc != 0) {
                BUILDBOX_LOG_WARNING(
                    "Failed to run pre-unstage command, rc: " << preUnstageRc);
            }
        }
        // No need to sudo
        else {
            buildboxcommon::Runner::recursively_chmod_directories(
                stagedDirPath.c_str(), PERMISSION_RWXALL);
        }
    });

    {
        createOutputDirectories(command, workingDir.str());

        // userchroot doesn't allow symlinks in the root path
        char resolved_path[PATH_MAX];
        if (realpath(stagedDir->getPath(), resolved_path) == NULL) {
            BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
                std::system_error, errno, std::system_category,
                "error in realpath() failed for path \""
                    << stagedDir->getPath() << "\"");
        }

        ChrootManager userChrootMgr(d_userchroot_bin, resolved_path, d_sudoBin,
                                    d_user);
        // environment variables set in file, which is sourced in command
        const std::vector<std::string> commandLine =
            userChrootMgr.generateCommandLine(command);

        BUILDBOX_RUNNER_LOG(
            DEBUG, "Executing " << logging::printableCommandLine(commandLine));
        executeAndStore(commandLine, &result);
    }

    BUILDBOX_RUNNER_LOG(DEBUG, "Capturing command outputs...");
    Runner::metadata_mark_output_upload_start(result_metadata);
    stagedDir->captureAllOutputs(command, &result);
    Runner::metadata_mark_output_upload_end(result_metadata);
    BUILDBOX_RUNNER_LOG(DEBUG, "Finished capturing command outputs");

    return result;
}

void UserChrootRunner::printSpecialUsage()
{
    std::clog << "    --userchroot-bin=PATH       Path to userchroot "
                 "executable. Will default to looking in PATH if not set.\n";
    std::clog
        << "    --local-only-cas-instance-name        CAS instance name used "
           "to upload blobs of merged chroot trees.\n"
           "                                    This allows switching from "
           "a proxy to a server instance to prevent those blobs from\n"
           "                                    being forwarded to a remote "
           "CAS (they are only needed in the local CAS for staging).\n";
    std::clog
        << "    --inject-path=[CONTENT:]HOST_PATH:CHROOT_PATH   Map a "
           "path from HOST_PATH on the host filesystem to CHROOT_PATH "
           "inside the constructed chroot. The contents of the path will be "
           "copied into the chroot. If CONTENT is set, the injection only "
           "happens if the action has a property `inject=CONTENT`\n";
    std::clog << "    --sudo-bin=PATH       Path to sudo executable. Will "
                 "default to looking in PATH if not set.\n";
    std::clog << "    --find-bin=PATH       Path to find executable. Will "
                 "default to looking in PATH if not set. Required to be valid "
                 "if --user is specified.\n";
    std::clog << "    --chmod-bin=PATH       Path to chmod executable. Will "
                 "default to looking in PATH if not set. Required to be valid "
                 "if --user is specified.\n";
    std::clog << "    --user=USER       Run userchroot as another user.\n";
    std::clog
        << "    --group=GROUP     Advise localcas this process is in GROUP.\n";
}

bool UserChrootRunner::parseArg(std::string_view arg_view)
{
    if (this->d_casRemote.parseArg(arg_view)) {
        // The optional local-only CAS uses the same connection options as the
        // regular CAS connection. `buildboxcommon::Runner` doesn't expose
        // these connection options, so parse them here as well but return
        // false such that `Runner` can still use them for the regular CAS
        // connection.
        return false;
    }

    bool argumentParsed = false;
    if (arg_view.size() >= 2 && arg_view.substr(0, 2) == "--") {
        arg_view.remove_prefix(2);
        size_t assign_pos = arg_view.find('=');
        if (assign_pos != std::string_view::npos) {
            std::string_view key = arg_view.substr(0, assign_pos);
            std::string value(arg_view.substr(assign_pos + 1));
            if (key == "userchroot-bin") {
                this->d_userchroot_bin = value;
                argumentParsed = true;
            }
            else if (key == "local-only-cas-instance-name") {
                this->d_localOnlyCasInstanceName = value;
                argumentParsed = true;
            }
            else if (key == "inject-path") {
                const auto parts = StringUtils::split(value, ":");
                if (parts.size() == 2) {
                    std::filesystem::path src{parts[0]}, target{parts[1]};
                    this->d_mappedPaths.emplace_back(src, target);
                    argumentParsed = true;
                }
                else if (parts.size() == 3) {
                    const auto &content = parts[0];
                    std::filesystem::path src{parts[1]}, target{parts[2]};
                    this->d_mappedPathsByProperties[std::string(content)]
                        .emplace_back(src, target);
                    argumentParsed = true;
                }
                else {
                    BUILDBOX_LOG_WARNING(
                        "Invalid --inject-path option: " << value);
                }
            }
            else if (key == "sudo-bin") {
                this->d_sudoBin = value;
                argumentParsed = true;
            }
            else if (key == "find-bin") {
                this->d_findBin = value;
                argumentParsed = true;
            }
            else if (key == "chmod-bin") {
                this->d_chmodBin = value;
                argumentParsed = true;
            }
            else if (key == "user") {
                this->d_user = value;
                argumentParsed = true;
            }
            else if (key == "group") {
                this->d_group = value;
                argumentParsed = true;
            }
        }
    }
    return argumentParsed;
}

void UserChrootRunner::fetchChrootIntoLocalCas(const Digest &digest) const
{
    BUILDBOX_RUNNER_LOG(INFO, "Calling LocalCAS.FetchTree(chrootDigest="
                                  << digest
                                  << ") to make sure that the chroot is "
                                     "available from Local CAS");

    try {
        const bool fetchFiles = true;
        d_casClient->fetchTree(digest, fetchFiles);
    }
    catch (const buildboxcommon::GrpcError &e) {
        BUILDBOX_RUNNER_LOG(ERROR, "LocalCAS.FetchTree(chrootDigest="
                                       << digest
                                       << ") failed: " << e.status.error_code()
                                       << ": " << e.status.error_message());
        throw e;
    }
}

void UserChrootRunner::uploadMissingBlobs(const digest_string_map &data) const
{
    // A different instance name could have been specified for these requests.
    // This special case allows switching casd instances from a proxy to a
    // server so that the blobs in the merged tree do not end up being
    // forwarded to a remote CAS and only kept locally.
    auto casClient = d_casClient;
    std::vector<Digest> digests;
    digests.reserve(data.size());
    for (const auto &digestBlob : data) {
        digests.push_back(digestBlob.first);
    }

    if (!d_localOnlyCasInstanceName.empty()) {
        BUILDBOX_LOG_DEBUG("Uploading blobs to local-only CAS instance \""
                           << d_localOnlyCasInstanceName << "\"");

        // Use the same connection options as the regular CAS connection with
        // the exception of the instance name.
        ConnectionOptions localOnlyCasRemote = d_casRemote;
        localOnlyCasRemote.d_instanceName = d_localOnlyCasInstanceName;

        auto client = std::make_shared<GrpcClient>();
        client->init(localOnlyCasRemote);
        casClient = std::make_shared<CASClient>(std::move(client));
        casClient->init();
    }

    const std::vector<Digest> missingDigests =
        casClient->findMissingBlobs(digests);

    if (missingDigests.empty()) {
        return;
    }

    BUILDBOX_RUNNER_LOG(DEBUG, "Uploading " << missingDigests.size()
                                            << " missing blob(s) out of "
                                            << digests.size() << " total");

    std::vector<CASClient::UploadRequest> uploadRequests;
    uploadRequests.reserve(missingDigests.size());
    for (const auto &digest : missingDigests) {
        uploadRequests.emplace_back(digest, data.at(digest));
    }

    const std::vector<CASClient::UploadResult> results =
        casClient->uploadBlobs(uploadRequests);

    // verify that all blobs have been successfully uploaded
    // throw on any failures
    bool uploadSuccess = true;
    std::ostringstream oss;
    for (const auto &result : results) {
        if (!result.status.ok()) {
            oss << "Failed to upload a merged digest(" << result.digest
                << "), status = [" << result.status.error_code() << ": \""
                << result.status.error_message() << "\"]\n";
            uploadSuccess = false;
        }
    }

    if (!uploadSuccess) {
        BUILDBOXCOMMON_THROW_EXCEPTION(std::runtime_error, oss.str());
    }
}

Digest UserChrootRunner::mergeTreesAndSetRootDirectoryProperty(
    const DirectoryTree &inputTree, const DirectoryTree &chrootTree,
    digest_string_map *mergedDirectoryBlobs)
{
    digest_string_map blobs;

    const Digest rootDigest = mergeTrees(inputTree, chrootTree, &blobs);

    if (rootDigest == Digest()) {
        return Digest();
    }

    // The merge succeeded. We will now modify the root `Directory` of the
    // result to set its `unix_mode` property:
    Directory newRootDirectory;
    newRootDirectory.ParseFromString(blobs.at(rootDigest));
    newRootDirectory.mutable_node_properties()->mutable_unix_mode()->set_value(
        s_rootDirectoryRequiredUnixMode);

    // Serializing and recalculating its hash:
    const std::string newRootDirectoryBlob =
        newRootDirectory.SerializeAsString();
    Digest newRootDirectoryDigest =
        DigestGenerator::hash(newRootDirectoryBlob);

    // And swapping it for the old root:
    if (rootDigest != newRootDirectoryDigest) {
        blobs.erase(rootDigest);
        blobs.emplace(newRootDirectoryDigest, newRootDirectoryBlob);
    }

    // Everything succeeded, we can write to the output arguments:
    *mergedDirectoryBlobs = std::move(blobs);

    return newRootDirectoryDigest;
}

void UserChrootRunner::mergeDigests(
    const Digest &inputDigest, const Digest &chrootDigest,
    Digest *mergedRootDigest,
    const Command_OutputDirectoryFormat &chrootDirectoryFormat)
{
    // acquire the input trees of both digests
    // getTree and fetchMessage will throw on error
    const DirectoryTree inputTree = this->d_casClient->getTree(inputDigest);

    DirectoryTree chrootTree;
    // If the chroot is avalble as a Tree Message download that
    if (chrootDirectoryFormat == Command::TREE_ONLY) {
        chrootTree = this->d_casClient->getTreeFromTreeMessage(chrootDigest);
    }
    // If not use getTree - Internally this will recursively download starting
    // from the root directory to build up the tree.
    else {
        chrootTree = this->d_casClient->getTree(chrootDigest);
    }

    digest_string_map mergedDirectoryBlobs;
    const Digest rootDigest = mergeTreesAndSetRootDirectoryProperty(
        inputTree, chrootTree, &mergedDirectoryBlobs);

    if (rootDigest == Digest()) {
        std::ostringstream errMsg;
        errMsg << "Error merging inputDigest(" << inputDigest
               << ") with chrootDigest(" << chrootDigest << ")";
        const grpc::Status grpcStatus(grpc::StatusCode::FAILED_PRECONDITION,
                                      errMsg.str());
        // Written to errorStatusFile in Runner::main()
        GrpcError::throwGrpcError(grpcStatus);
    }

    // Upload any missing digests to CAS so they can be staged
    uploadMissingBlobs(mergedDirectoryBlobs);

    *mergedRootDigest = rootDigest;
}

void UserChrootRunner::assertRequiredDirectoriesExist() const
{
    static const std::array<std::string, 2> expected_directories = {"/dev",
                                                                    "/tmp"};

    for (const auto &directory : expected_directories) {
        const std::string expected_path = this->d_stage_path + directory;

        if (!FileUtils::isDirectory(expected_path.c_str())) {
            std::ostringstream errMsg;
            errMsg << "Staged chroot does not contain expected directory \""
                   << directory << "\"";
            // Written to errorStatusFile in Runner::main()
            const grpc::Status grpcStatus(
                grpc::StatusCode::FAILED_PRECONDITION, errMsg.str());
            GrpcError::throwGrpcError(grpcStatus);
        }
    }
}

ProcessCredentials UserChrootRunner::getProcessCredentials() const
{
    ProcessCredentials creds{};
    creds.uid =
        d_user.empty() ? geteuid() : SystemUtils::getUidFromName(d_user);
    creds.gid =
        d_group.empty() ? getegid() : SystemUtils::getGidFromName(d_group);
    return creds;
}

std::vector<std::string> UserChrootRunner::createPreUnstageCommand(
    const std::string &stagedDirPath) const
{
    std::unordered_map<std::string, std::string> binPaths;
    binPaths["sudo"] = SystemUtils::getPathToCommand(d_sudoBin);
    binPaths["find"] = SystemUtils::getPathToCommand(d_findBin);
    binPaths["chmod"] = SystemUtils::getPathToCommand(d_chmodBin);
    for (const auto &binPathPair : binPaths) {
        if (binPathPair.second.empty()) {
            BUILDBOXCOMMON_THROW_EXCEPTION(std::runtime_error,
                                           "Cannot find path to command "
                                               << binPathPair.first);
        }
    }

    return {
        binPaths["sudo"],
        "-u",
        d_user,
        binPaths["find"],
        stagedDirPath,
        "-xdev",
        "-user",
        d_user,
        "-type",
        "d",
        "-execdir",
        binPaths["chmod"],
        "777",
        "{}",
        "+",
    };
}

void UserChrootRunner::injectPath(const std::filesystem::path &srcPath,
                                  const std::filesystem::path &chrootPath)
{
    // Are we copying a directory or a single file?
    //
    // A directory has it's contents copied recursively into the directory
    // specified by the chroot part of the path mapping entry. If the path
    // exists already and is not a directory, the entry is skipped.
    //
    // A file is copied directly to the specified path. If the path already
    // exists and is not a file, the entry is skipped.
    if (buildboxcommon::FileUtils::isDirectory(srcPath.c_str())) {
        auto staged_path_fd = buildboxcommon::FileUtils::openInRoot(
            d_stage_path, chrootPath.c_str(),
            O_DIRECTORY | O_RDONLY | O_CLOEXEC, true);
        if (staged_path_fd < 0) {
            BUILDBOX_RUNNER_LOG(
                WARNING, "Skipping path mapping `"
                             << srcPath << ":" << chrootPath
                             << "`. Failed to open destination directory.");
            return;
        }
        buildboxcommon::FileDescriptor destfd(staged_path_fd);
        if (!buildboxcommon::FileUtils::isDirectory(destfd.get())) {
            BUILDBOX_RUNNER_LOG(
                WARNING, "Skipping path mapping `"
                             << srcPath << ":" << chrootPath
                             << "`. Source is a directory but the destination "
                                "already exists and is not a directory.");
            return;
        }

        const buildboxcommon::FileDescriptor srcdirfd(
            open(srcPath.c_str(), O_DIRECTORY | O_RDONLY | O_CLOEXEC));
        buildboxcommon::FileUtils::copyRecursively(srcdirfd.get(),
                                                   destfd.get());
    }
    else if (buildboxcommon::FileUtils::isRegularFile(srcPath.c_str())) {
        const std::filesystem::path chrootParent =
            std::filesystem::path(chrootPath).parent_path();
        const buildboxcommon::FileDescriptor chrootdirfd(
            buildboxcommon::FileUtils::openInRoot(
                d_stage_path, chrootParent.c_str(),
                O_DIRECTORY | O_RDONLY | O_CLOEXEC, true));

        if (buildboxcommon::FileUtils::existsInDir(
                chrootdirfd.get(), chrootPath.filename().c_str()) &&
            !buildboxcommon::FileUtils::isRegularFileNoFollow(
                chrootdirfd.get(), chrootPath.filename().c_str())) {
            BUILDBOX_RUNNER_LOG(
                WARNING,
                "Skipping path mapping `"
                    << srcPath << ":" << chrootPath
                    << "`. Source is a regular file but the destination "
                       "already exists and is not a regular file.");
            return;
        }

        buildboxcommon::FileUtils::copyFile(AT_FDCWD, srcPath.c_str(),
                                            chrootdirfd.get(),
                                            chrootPath.filename().c_str());
    }
    else {
        BUILDBOX_RUNNER_LOG(WARNING,
                            "Skipping path mapping `"
                                << srcPath << ":" << chrootPath
                                << "`. Source is not a file or directory.");
        return;
    }
}

} // namespace userchroot
} // namespace buildboxrun
} // namespace buildboxcommon
