/*
 * Copyright (C) 2018 The Android Open Source Project
 *
 * 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 <modprobe/modprobe.h>
#include <modprobe/utils.h>

#include <fnmatch.h>
#include <grp.h>
#include <pwd.h>
#include <sys/stat.h>
#include <sys/wait.h>

#include <algorithm>
#include <map>
#include <set>
#include <string>
#include <thread>
#include <vector>

#include <android-base/file.h>
#include <android-base/logging.h>
#include <android-base/parseint.h>
#include <android-base/strings.h>
#include <android-base/unique_fd.h>

using android::modprobe::CanonicalizeModulePath;

Modprobe::Modprobe(const std::vector<std::string>& base_paths, const std::string load_file,  bool use_blocklist, ExternalPortState external_port_state)
    : Modprobe(ModuleConfig::Parse(base_paths, external_port_state, load_file), use_blocklist) {}

Modprobe::Modprobe(ModuleConfig config, bool use_blocklist)
    : module_aliases_(std::move(config.module_aliases)),
      module_deps_(std::move(config.module_deps)),
      module_pre_softdep_(std::move(config.module_pre_softdep)),
      module_post_softdep_(std::move(config.module_post_softdep)),
      module_load_(std::move(config.module_load)),
      module_options_(std::move(config.module_options)),
      module_blocklist_(std::move(config.module_blocklist)),
      blocklist_enabled(use_blocklist) {}

std::vector<std::string> Modprobe::GetDependencies(const std::string& module) {
    auto it = module_deps_.find(module);
    if (it == module_deps_.end()) {
        return {};
    }
    return it->second;
}

bool Modprobe::InsmodWithDeps(const std::string& module_name, const std::string& parameters) {
    if (module_name.empty()) {
        LOG(ERROR) << "Need valid module name, given: " << module_name;
        return false;
    }

    auto dependencies = GetDependencies(module_name);
    if (dependencies.empty()) {
        LOG(ERROR) << "Module " << module_name << " not in dependency file";
        return false;
    }

    // load module dependencies in reverse order
    for (auto dep = dependencies.rbegin(); dep != dependencies.rend() - 1; ++dep) {
        LOG(VERBOSE) << "Loading hard dep for '" << module_name << "': " << *dep;
        if (!LoadWithAliases(*dep, true)) {
            return false;
        }
    }

    // try to load soft pre-dependencies
    for (const auto& [module, softdep] : module_pre_softdep_) {
        if (module_name == module) {
            LOG(VERBOSE) << "Loading soft pre-dep for '" << module << "': " << softdep;
            LoadWithAliases(softdep, false);
        }
    }

    // load target module itself with args
    if (!Insmod(dependencies[0], parameters)) {
        return false;
    }

    // try to load soft post-dependencies
    for (const auto& [module, softdep] : module_post_softdep_) {
        if (module_name == module) {
            LOG(VERBOSE) << "Loading soft post-dep for '" << module << "': " << softdep;
            LoadWithAliases(softdep, false);
        }
    }

    return true;
}

bool Modprobe::LoadWithAliases(const std::string& module_name, bool strict,
                               const std::string& parameters) {
    std::set<std::string> modules_to_load;
    bool module_loaded = false;
    {
        std::lock_guard guard(module_loaded_lock_);

        auto canonical_name = CanonicalizeModulePath(module_name);
        if (module_loaded_.count(canonical_name)) {
            return true;
        }
        modules_to_load.insert(std::move(canonical_name));
        // use aliases to expand list of modules to load (multiple modules
        // may alias themselves to the requested name)
        for (const auto& [alias, aliased_module] : module_aliases_) {
            if (fnmatch(alias.c_str(), module_name.c_str(), 0) != 0) continue;
            LOG(VERBOSE) << "Found alias for '" << module_name << "': '" << aliased_module;
            if (module_loaded_.count(CanonicalizeModulePath(aliased_module))) continue;
            modules_to_load.emplace(aliased_module);
        }
    }

    // attempt to load all modules aliased to this name
    for (const auto& module : modules_to_load) {
        if (!ModuleExists(module)) continue;
        if (InsmodWithDeps(module, parameters)) module_loaded = true;
    }

    if (strict && !module_loaded) {
        LOG(ERROR) << "LoadWithAliases was unable to load " << module_name
                   << ", tried: " << android::base::Join(modules_to_load, ", ");
        return false;
    }
    return true;
}

bool Modprobe::IsBlocklisted(const std::string& module_name) {
    if (!blocklist_enabled) return false;

    auto canonical_name = CanonicalizeModulePath(module_name);
    auto dependencies = GetDependencies(canonical_name);
    for (auto dep = dependencies.begin(); dep != dependencies.end(); ++dep) {
        if (module_blocklist_.count(CanonicalizeModulePath(*dep))) return true;
    }

    return module_blocklist_.count(canonical_name) > 0;
}

// Another option to load kernel modules. load independent modules dependencies
// in parallel and then update dependency list of other remaining modules,
// repeat these steps until all modules are loaded.
// Discard all blocklist.
// Softdeps are taken care in InsmodWithDeps().
bool Modprobe::LoadModulesParallel(int num_threads) {
    bool ret = true;
    std::map<std::string, std::vector<std::string>> mod_with_deps;

    // Get dependencies
    for (const auto& module : module_load_) {
        // Skip blocklist modules
        if (IsBlocklisted(module)) {
            LOG(VERBOSE) << "LMP: Blocklist: Module " << module << " skipping...";
            continue;
        }
        auto dependencies = GetDependencies(CanonicalizeModulePath(module));
        if (dependencies.empty()) {
            LOG(ERROR) << "LMP: Hard-dep: Module " << module
                       << " not in .dep file";
            return false;
        }
        mod_with_deps[CanonicalizeModulePath(module)] = dependencies;
    }

    while (!mod_with_deps.empty()) {
        std::vector<std::thread> threads;
        std::vector<std::string> mods_path_to_load;
        std::mutex vector_lock;

        // Find independent modules
        for (const auto& [it_mod, it_dep] : mod_with_deps) {
            auto itd_last = it_dep.rbegin();
            if (itd_last == it_dep.rend())
                continue;

            auto cnd_last = CanonicalizeModulePath(*itd_last);
            // Hard-dependencies cannot be blocklisted
            if (IsBlocklisted(cnd_last)) {
                LOG(ERROR) << "LMP: Blocklist: Module-dep " << cnd_last
                           << " : failed to load module " << it_mod;
                return false;
            }

            std::string str = "load_sequential=1";
            auto it = module_options_[cnd_last].find(str);
            if (it != std::string::npos) {
                module_options_[cnd_last].erase(it, it + str.size());

                if (!LoadWithAliases(cnd_last, true)) {
                    return false;
                }
            } else {
                if (std::find(mods_path_to_load.begin(), mods_path_to_load.end(),
                            cnd_last) == mods_path_to_load.end()) {
                    mods_path_to_load.emplace_back(cnd_last);
                }
            }
        }

        // Load independent modules in parallel
        auto thread_function = [&] {
            std::unique_lock lk(vector_lock);
            while (!mods_path_to_load.empty()) {
                auto ret_load = true;
                auto mod_to_load = std::move(mods_path_to_load.back());
                mods_path_to_load.pop_back();

                lk.unlock();
                ret_load &= LoadWithAliases(mod_to_load, true);
                lk.lock();
                if (!ret_load) {
                    ret &= ret_load;
                }
            }
        };

        std::generate_n(std::back_inserter(threads), num_threads,
                        [&] { return std::thread(thread_function); });

        // Wait for the threads.
        for (auto& thread : threads) {
            thread.join();
        }

        if (!ret) return ret;

        std::lock_guard guard(module_loaded_lock_);
        // Remove loaded module form mod_with_deps and soft dependencies of other modules
        for (const auto& module_loaded : module_loaded_)
            mod_with_deps.erase(module_loaded);

        // Remove loaded module form dependencies of other modules which are not loaded yet
        for (const auto& module_loaded_path : module_loaded_paths_) {
            for (auto& [mod, deps] : mod_with_deps) {
                auto it = std::find(deps.begin(), deps.end(), module_loaded_path);
                if (it != deps.end()) {
                    deps.erase(it);
                }
            }
        }
    }

    return ret;
}

bool Modprobe::LoadListedModules(bool strict) {
    auto ret = true;
    for (const auto& module : module_load_) {
        if (!LoadWithAliases(module, true)) {
            if (IsBlocklisted(module)) continue;
            ret = false;
            if (strict) break;
        }
    }
    return ret;
}

bool Modprobe::Remove(const std::string& module_name) {
    auto dependencies = GetDependencies(CanonicalizeModulePath(module_name));
    for (auto dep = dependencies.begin(); dep != dependencies.end(); ++dep) {
        Rmmod(*dep);
    }
    Rmmod(module_name);
    return true;
}

std::vector<std::string> Modprobe::ListModules(const std::string& pattern) {
    std::vector<std::string> rv;
    for (const auto& [module, deps] : module_deps_) {
        // Attempt to match both the canonical module name and the module filename.
        if (!fnmatch(pattern.c_str(), module.c_str(), 0)) {
            rv.emplace_back(module);
        } else if (!fnmatch(pattern.c_str(), android::base::Basename(deps[0]).c_str(), 0)) {
            rv.emplace_back(deps[0]);
        }
    }
    return rv;
}

bool Modprobe::GetAllDependencies(const std::string& module,
                                  std::vector<std::string>* pre_dependencies,
                                  std::vector<std::string>* dependencies,
                                  std::vector<std::string>* post_dependencies) {
    std::string canonical_name = CanonicalizeModulePath(module);
    if (pre_dependencies) {
        pre_dependencies->clear();
        for (const auto& [it_module, it_softdep] : module_pre_softdep_) {
            if (canonical_name == it_module) {
                pre_dependencies->emplace_back(it_softdep);
            }
        }
    }
    if (dependencies) {
        dependencies->clear();
        auto hard_deps = GetDependencies(canonical_name);
        if (hard_deps.empty()) {
            return false;
        }
        for (auto dep = hard_deps.rbegin(); dep != hard_deps.rend(); dep++) {
            dependencies->emplace_back(*dep);
        }
    }
    if (post_dependencies) {
        for (const auto& [it_module, it_softdep] : module_post_softdep_) {
            if (canonical_name == it_module) {
                post_dependencies->emplace_back(it_softdep);
            }
        }
    }
    return true;
}