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

#include <buildboxcommon_exception.h>
#include <buildboxcommon_fileutils.h>
#include <buildboxcommon_futuregroup.h>
#include <buildboxcommon_logging.h>
#include <buildboxcommon_metricnames.h>
#include <buildboxcommon_permissions.h>
#include <buildboxcommon_protos.h>
#include <buildboxcommon_timeutils.h>

#include <buildboxcommonmetrics_countingmetric.h>
#include <buildboxcommonmetrics_countingmetricutil.h>

#include <cerrno>
#include <dirent.h>
#include <exception>
#include <fcntl.h>
#include <future>
#include <limits.h>
#include <optional>
#include <stdlib.h>
#include <string>
#include <sys/stat.h>
#include <sys/types.h>
#include <system_error>
#include <unistd.h>
#include <vector>

using namespace buildboxcommon;
using namespace build::bazel::remote::execution::v2;

namespace {

// Get the mode of staged directory based on "unix_mode" and "SubtreeReadOnly"
// properties
mode_t directoryMode(const std::optional<mode_t> &propertyMode,
                     bool readOnly = false)
{
    mode_t mode = propertyMode.value_or(PERMISSION_RWXALL);
    if (readOnly) {
        mode &= ~(S_IWGRP | S_IWOTH);
    }
    return mode;
}

void createDirectory(const std::string &path,
                     const mode_t mode = PERMISSION_RWXUSR_RXALL)
{
    const auto mkdir_status = mkdir(path.c_str(), mode);
    if (mkdir_status != 0) {
        BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
            std::system_error, errno, std::system_category,
            "Could not create directory [" << path << "]");
    }
}

} // namespace

HardLinkStager::HardLinkStager(LocalCas *cas_storage,
                               const std::shared_ptr<ThreadPool> &threadPool)
    : FileStager(cas_storage), d_threadPool(threadPool), d_euid(geteuid())
{
    google::protobuf::util::TimeUtil::FromString("2011-11-11T11:11:11Z",
                                                 &d_deterministicTime);
}

HardLinkStager::HardLinkStagedDirectory::HardLinkStagedDirectory(
    const std::string &path)
    : d_staged_path(path)
{
}

HardLinkStager::HardLinkStagedDirectory::~HardLinkStagedDirectory()
{
    unstage(d_staged_path);
}

std::unique_ptr<FileStager::StagedDirectory> HardLinkStager::stage(
    const Digest &root_digest, const std::string &path,
    const std::optional<ProcessCredentials> &accessCredentials,
    const FileStagedCallback &callback)
{
    const bool path_is_existing_directory =
        buildboxcommon::FileUtils::isDirectory(path.c_str());

    if (path_is_existing_directory &&
        !buildboxcommon::FileUtils::directoryIsEmpty(path.c_str())) {
        BUILDBOXCOMMON_THROW_EXCEPTION(std::invalid_argument,
                                       "Could not stage "
                                           << toString(root_digest) << " as ["
                                           << path << "] is not empty");
    }

    bool access_by_owner = false;
    std::optional<gid_t> accessByGroup;

    if (accessCredentials.has_value()) {
        if (accessCredentials->uid == d_euid) {
            // If the process accessing the staged files runs as the same user
            // as buildbox-casd, it can modify the files and thus, it's not
            // safe to hardlink files without risking cache corruption.
            access_by_owner = true;
        }
        if (SystemUtils::getGids().count(accessCredentials->gid) > 0) {
            // only set the gid if casd process user is within the same group
            accessByGroup = accessCredentials->gid;
        }
    }

    if (!path_is_existing_directory) {
        createDirectory(path);
    }
    else {
        FileDescriptor fd(
            open(path.c_str(), O_RDONLY | O_NOFOLLOW | O_CLOEXEC));
        if (fd.get() == -1) {
            BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(std::system_error, errno,
                                                  std::system_category,
                                                  "Could not open " << path);
        }
        struct stat st{};
        if (fstat(fd.get(), &st) != 0) {
            BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(std::system_error, errno,
                                                  std::system_category,
                                                  "Could not stat " << path);
        }
        if ((st.st_mode & PERMISSION_RWXALL) != PERMISSION_RWXALL) {
            auto rc = fchmod(fd.get(), st.st_mode | PERMISSION_RWXALL);
            if (rc < 0) {
                BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
                    std::system_error, errno, std::system_category,
                    "Could not chmod " << path);
            }
        }
    }

    try {
        stageDirectory(root_digest, path, access_by_owner, accessByGroup,
                       callback);
    }
    catch (const std::exception &e) {
        BUILDBOX_LOG_DEBUG("Error staging root digest " << root_digest
                                                        << " in [" << path
                                                        << "]: " << e.what());

        // If the staging fails, we roll back the stage directory to the
        // its status before the call:
        unstage(path);

        throw;
    }

    // This object will keep the directory staged until it's destruction:
    return std::make_unique<HardLinkStagedDirectory>(path);
}

void HardLinkStager::unstage(const std::string &path)
{
    try {
        buildboxcommon::FileUtils::clearDirectory(path.c_str());
    }
    catch (const std::system_error &e) {
        BUILDBOX_LOG_ERROR("Error unstaging " << path << ": " << e.what());
    }
}

void HardLinkStager::stageDirectory(const Digest &digest,
                                    const std::string &path,
                                    bool access_by_owner,
                                    const std::optional<gid_t> &accessByGroup,
                                    const FileStagedCallback &callback)
{
    const bool is_root_directory = true;
    stageDirectoriesRecursively(digest, is_root_directory, path,
                                access_by_owner, accessByGroup, callback);
}

void HardLinkStager::stageDirectoriesRecursively(
    const Digest &digest, const bool is_root_directory,
    const std::string &path, bool access_by_owner,
    const std::optional<gid_t> &accessByGroup,
    const FileStagedCallback &callback, bool readonly)
{
    Directory directory = getDirectory(digest);

    for (const NodeProperty &property :
         directory.node_properties().properties()) {
        if (property.name() == "SubtreeReadOnly") {
            readonly = property.value() == "true";
        }
    }

    // PRE: The top level directory already exists at this point.
    if (!is_root_directory) {
        createDirectory(path);
    }

    const FileDescriptor dirfd(
        open(path.c_str(), O_DIRECTORY | O_RDONLY | O_CLOEXEC));
    if (dirfd.get() < 0) {
        BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
            std::system_error, errno, std::system_category,
            "Could not open directory " << path);
    }

    const auto compare_files = [](const FileNode &f1, const FileNode &f2) {
        return f1.digest().hash() < f2.digest().hash();
    };

    // Sort files so that files from the same objects subdirectory in the
    // local cache will be staged together. This will improve hit rate in
    // the kernel directory entry cache.
    std::sort(directory.mutable_files()->begin(),
              directory.mutable_files()->end(), compare_files);

    {
        FutureGroup<void> futureGroup(d_threadPool.get());
        std::vector<std::shared_future<void>> futures;
        for (const FileNode &file : directory.files()) {
            auto stageFileLambda = [&,
                                    stage_path = path + "/" + file.name()]() {
                stageFile(file, dirfd.get(), stage_path, access_by_owner,
                          accessByGroup, callback);
            };
            futures.emplace_back(futureGroup.add(stageFileLambda));
        }
        for (auto &f : futures) {
            f.get();
        }
    }

    for (const DirectoryNode &subdirectory : directory.directories()) {
        const std::string subdirectory_path = path + "/" + subdirectory.name();
        stageDirectoriesRecursively(subdirectory.digest(), false,
                                    subdirectory_path, access_by_owner,
                                    accessByGroup, callback, readonly);
    }

    for (const SymlinkNode &symlink : directory.symlinks()) {
        const std::string symlink_path = path + "/" + symlink.name();
        stageSymLink(symlink.target(), symlink_path);
    }

    buildboxcommon::buildboxcommonmetrics::CountingMetricUtil::
        recordCounterMetric(
            CommonMetricNames::COUNTER_NUM_HARDLINK_STAGER_STAGED_FILES,
            directory.files_size());

    buildboxcommon::buildboxcommonmetrics::CountingMetricUtil::
        recordCounterMetric(
            CommonMetricNames::COUNTER_NUM_HARDLINK_STAGER_STAGED_DIRECTORIES,
            directory.directories_size());

    buildboxcommon::buildboxcommonmetrics::CountingMetricUtil::
        recordCounterMetric(
            CommonMetricNames::COUNTER_NUM_HARDLINK_STAGER_STAGED_SYMLINKS,
            directory.symlinks_size());

    // Update mode and timestamps of directory after staging sub-entries.
    // When is_root_directory, allow chmod and setFileTimes to fail due to
    // permission error to accommodate the case if root dir is pre-created by
    // another user and runner has option `--workspace-path`. This went through
    // some refactoring and error check had been missing.
    std::optional<mode_t> propertyMode;
    if (directory.node_properties().has_unix_mode()) {
        propertyMode = directory.node_properties().unix_mode().value();
    }
    const mode_t dirMode = directoryMode(propertyMode, readonly);
    if (chmod(path.c_str(), dirMode) != 0 &&
        !(is_root_directory && errno == EPERM)) {
        BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
            std::system_error, errno, std::system_category,
            "Could not chmod directory " << path << " to " << dirMode);
    }
    // Update gid if the directory has g+w bit and gid is provided
    if ((dirMode & S_IWGRP) && accessByGroup.has_value()) {
        if (chown(path.c_str(), -1, *accessByGroup) != 0 &&
            !(is_root_directory && errno == EPERM)) {
            BUILDBOXCOMMON_THROW_EXCEPTION(
                std::runtime_error,
                "`chown` failed, cannot change gid of directory "
                    << path << " to " << *accessByGroup << ", "
                    << std::strerror(errno));
        }
    }

    const google::protobuf::Timestamp mtime =
        directory.node_properties().has_mtime()
            ? directory.node_properties().mtime()
            : d_deterministicTime;
    const auto mtimePoint = buildboxcommon::TimeUtils::parse_timestamp(mtime);
    try {
        FileUtils::setFileTimes(dirfd.get(), std::nullopt, mtimePoint);
    }
    catch (std::system_error &e) {
        if (!(is_root_directory && e.code().value() == EPERM)) {
            throw;
        }
    }
}

void HardLinkStager::stageFile(const FileNode &file, int dirfd,
                               const std::string &path, bool access_by_owner,
                               const std::optional<gid_t> &accessByGroup,
                               const FileStagedCallback &callback)
{
    // If the file will be accessed by the owner or we require properties,
    // we must copy instead of hardlinking.
    bool can_hardlink = !access_by_owner &&
                        !file.node_properties().has_unix_mode() &&
                        !file.node_properties().has_mtime();

    if (can_hardlink) {
        const auto cas_path = d_cas_storage->pathAtUnchecked(file.digest());

        hardLinkAt(cas_path.d_directory_fd, cas_path.d_file_name, dirfd,
                   file.name());

        if (file.is_executable()) {
            // This is problematic as it affects all hardlinks of the file.
            try {
                buildboxcommon::FileUtils::makeExecutable(path.c_str());
            }
            catch (const std::runtime_error &e) {
                BUILDBOX_LOG_ERROR("Could not give executable permission to "
                                   << path << ": " << e.what());
                throw e;
            }
        }
    }
    else {
        mode_t mode = 0;

        if (!access_by_owner) {
            // Falling back to file copies to be able to apply the mtime.
            // Match the file mode of CAS objects to minimize differences
            // between hardlinking and copy fallback.
            mode = S_IRUSR | S_IRGRP | S_IROTH;
            if (d_cas_storage->objectFilesWritable()) {
                mode |= S_IWUSR;
            }
        }
        else {
            // All files will be copied, there is no mix between hardlinks and
            // file copies. Use a more common default mode where the owner
            // can modify the file.
            mode = S_IWUSR | S_IRUSR | S_IRGRP | S_IROTH;
        }

        if (file.is_executable()) {
            mode |= S_IXUSR | S_IXGRP | S_IXOTH;
        }

        if (file.node_properties().has_unix_mode()) {
            // Override the default mode
            mode = static_cast<mode_t>(
                file.node_properties().unix_mode().value());
            if (mode & ~PERMISSION_RWXALL_SPECIAL) {
                BUILDBOXCOMMON_THROW_EXCEPTION(std::runtime_error,
                                               "Invalid `unix_mode` 0"
                                                   << std::oct << mode);
            }

            if (!access_by_owner && (mode & ~PERMISSION_RWXALL)) {
                // The effect of the file mode depends on the file owner.
                // Reject custom modes if the file will be accessed by a
                // different user to prevent unexpected and possibly dangerous
                // (suid) permissions.
                BUILDBOXCOMMON_THROW_EXCEPTION(
                    std::runtime_error, "Custom `unix_mode` 0"
                                            << std::oct << mode
                                            << " not allowed for other users");
            }
        }

        if (!d_cas_storage->copyBlob(file.digest(), dirfd, file.name().c_str(),
                                     mode)) {
            BUILDBOXCOMMON_THROW_EXCEPTION(
                std::runtime_error,
                "Failed to stage " << file.digest() << ": Blob not available");
        }

        // Update gid if the file has g+w bit and gid is provided
        if ((mode & S_IWGRP) && accessByGroup.has_value()) {
            if (fchownat(dirfd, file.name().c_str(), -1, *accessByGroup, 0) ==
                -1) {
                BUILDBOXCOMMON_THROW_EXCEPTION(
                    std::runtime_error,
                    "`fchownat` failed, cannot change gid of file to "
                        << *accessByGroup << ", " << std::strerror(errno));
            }
        }

        // Use a different mtime only if set as property
        google::protobuf::Timestamp mtime =
            file.node_properties().has_mtime() ? file.node_properties().mtime()
                                               : d_deterministicTime;
        const auto mtimePoint =
            buildboxcommon::TimeUtils::parse_timestamp(mtime);
        buildboxcommon::FileUtils::setFileTimes(dirfd, file.name().c_str(),
                                                std::nullopt, mtimePoint);

        const auto &properties = file.node_properties().properties();
        if (!properties.empty()) {
            const NodeProperty &property = *properties.begin();
            BUILDBOXCOMMON_THROW_EXCEPTION(
                std::runtime_error, "Unsupported FileNode property "
                                        << property.name() << " with value "
                                        << property.value());
        }

        if (callback) {
            struct stat st{};
            if (fstatat(dirfd, file.name().c_str(), &st,
                        AT_SYMLINK_NOFOLLOW) != 0) {
                BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
                    std::system_error, errno, std::system_category,
                    "Failed to get file status for " << file.name());
            }
            (*callback)(st, file.digest());
        }
    }
}

void HardLinkStager::stageSymLink(const std::string &target,
                                  const std::string &link_path)
{
    const int link_status = symlink(target.c_str(), link_path.c_str());

    if (link_status != 0) {
        BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
            std::system_error, errno, std::system_category,
            "Could not create symlink " << link_path << " -> " << target);
    }
}

void HardLinkStager::hardLinkAt(const int target_dir_fd,
                                const std::string &target,
                                const int link_dir_fd,
                                const std::string &link_path)
{
    // linkat is used instead of link, so that we may leverage
    // directory file descriptors which are cached. This should
    // result in a performance improvement.
    const int link_status = linkat(target_dir_fd, target.c_str(), link_dir_fd,
                                   link_path.c_str(), 0);

    if (link_status != 0) {
        BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
            std::system_error, errno, std::system_category,
            "Could not create hard link " << link_path << " -> " << target);
    }
}

Directory HardLinkStager::getDirectory(const Digest &digest)
{
    const auto blob_data = d_cas_storage->readBlob(digest);
    if (blob_data == nullptr) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::invalid_argument,
            "Could not read blob with digest: " << digest);
    }

    Directory directory;
    if (directory.ParseFromString(*blob_data)) {
        return directory;
    }

    BUILDBOXCOMMON_THROW_EXCEPTION(
        std::runtime_error,
        "Could not parse Directory from blob with digest: " << digest);
}
