/*
 * Copyright (C) 2020 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 <android-base/logging.h>
#include <android-base/properties.h>
#include <android-base/strings.h>
#include <gflags/gflags.h>
#include <libsnapshot/capabilities.h>
#include <snapuserd/snapuserd_client.h>
#include <storage_literals/storage_literals.h>
#include "android-base/file.h"

#include "snapuserd_daemon.h"

using namespace android::storage_literals;

DEFINE_string(socket, android::snapshot::kSnapuserdSocket, "Named socket or socket path.");
DEFINE_bool(no_socket, false,
            "If true, no socket is used. Each additional argument is an INIT message.");
DEFINE_bool(socket_handoff, false,
            "If true, perform a socket hand-off with an existing snapuserd instance, then exit.");
DEFINE_bool(user_snapshot, false, "If true, user-space snapshots are used");
DEFINE_bool(io_uring, false, "If true, io_uring feature is enabled");
DEFINE_bool(o_direct, false, "If true, enable direct reads on source device");
DEFINE_bool(skip_verification, false, "If true, skip verification of partitions");
DEFINE_int32(cow_op_merge_size, 0, "number of operations to be processed at once");
DEFINE_int32(worker_count, 4, "number of worker threads used to serve I/O requests to dm-user");
DEFINE_int32(verify_block_size, 1_MiB, "block sized used during verification of snapshots");
DEFINE_int32(num_verify_threads, 3, "number of threads used during verification phase");
DEFINE_bool(ublk, false, "If true, use ublk instead of dm-user");
DEFINE_int32(fd_num, -1, "fd used to communicate with Init for udev events");

namespace android {
namespace snapshot {

bool Daemon::IsUserspaceSnapshotsEnabled() {
    const std::string UNKNOWN = "unknown";
    const std::string vendor_release =
            android::base::GetProperty("ro.vendor.build.version.release_or_codename", UNKNOWN);

    // If the vendor is on Android S, install process will forcefully take the
    // userspace snapshots path.
    //
    // We will not reach here post OTA reboot as the binary will be from vendor
    // ramdisk which is on Android S.
    if (vendor_release.find("12") != std::string::npos) {
        LOG(INFO) << "Userspace snapshots enabled as vendor partition is on Android: "
                  << vendor_release;
        return true;
    }

    return android::base::GetBoolProperty("ro.virtual_ab.userspace.snapshots.enabled", false);
}

bool Daemon::StartDaemon(int argc, char** argv) {
    int arg_start = gflags::ParseCommandLineFlags(&argc, &argv, true);

    // Daemon launched from first stage init and during selinux transition
    // will have the command line "-user_snapshot" flag set if the user-space
    // snapshots are enabled.
    //
    // Daemon launched as a init service during "socket-handoff" and when OTA
    // is applied will check for the property. This is ok as the system
    // properties are valid at this point. We can't do this during first
    // stage init and hence use the command line flags to get the information.
    bool user_snapshots = FLAGS_user_snapshot;
    if (!user_snapshots) {
        user_snapshots = IsUserspaceSnapshotsEnabled();
    }
    if (user_snapshots) {
        LOG(INFO) << "Starting daemon for user-space snapshots.....";
        return StartServerForUserspaceSnapshots(arg_start, argc, argv);
    } else {
        LOG(ERROR) << "Userspace snapshots not enabled. No support for legacy snapshots";
    }
    return false;
}

bool Daemon::SendUeventRequest(int socket, const std::string& device_path) {
    ssize_t ret =
            TEMP_FAILURE_RETRY(send(socket, device_path.c_str(), device_path.size(), MSG_NOSIGNAL));

    if (ret < 0) {
        PLOG(ERROR) << "Failed to send uevent request for: " << device_path;
        return false;
    }
    // check and treat partial send as an error
    if (static_cast<size_t>(ret) != device_path.size()) {
        LOG(ERROR) << "Partial send for uevent request for: " << device_path << ". Sent " << ret
                   << " of " << device_path.size() << " bytes.";
        return false;
    }

    LOG(INFO) << "Sent uevent request for: " << device_path;
    return true;
}

bool Daemon::SendDoneNotification(int socket_fd) {
    const char* done_message = "DONE";
    const size_t message_length = strlen(done_message) + 1;

    ssize_t bytes_sent =
            TEMP_FAILURE_RETRY(send(socket_fd, done_message, message_length, MSG_NOSIGNAL));

    if (bytes_sent < 0) {
        PLOG(ERROR) << "Failed to send '" << done_message << "' message to init via socketpair fd "
                    << socket_fd;
        return false;
    }
    if (static_cast<size_t>(bytes_sent) != message_length) {
        LOG(ERROR) << "Partial send for '" << done_message << "' message to init: sent "
                   << bytes_sent << " of " << message_length << " bytes on fd " << socket_fd;
        return false;
    }

    LOG(INFO) << "Successfully sent '" << done_message << "' message to init via socketpair fd "
              << socket_fd;
    return true;
}

bool Daemon::StartServerForUserspaceSnapshots(int arg_start, int argc, char** argv) {
    int socket_ = -1;

    sigfillset(&signal_mask_);
    sigdelset(&signal_mask_, SIGINT);
    sigdelset(&signal_mask_, SIGTERM);
    sigdelset(&signal_mask_, SIGUSR1);

    // Masking signals here ensure that after this point, we won't handle INT/TERM
    // until after we call into ppoll()
    signal(SIGINT, Daemon::SignalHandler);
    signal(SIGTERM, Daemon::SignalHandler);
    signal(SIGPIPE, Daemon::SignalHandler);
    signal(SIGUSR1, Daemon::SignalHandler);

    MaskAllSignalsExceptIntAndTerm();

    gflags::CommandLineFlagInfo ublk_flag_info;
    bool ublk_flag_is_default = true;
    if (gflags::GetCommandLineFlagInfo("ublk", &ublk_flag_info)) {
        ublk_flag_is_default = ublk_flag_info.is_default;
    }

    if (ublk_flag_is_default) {
        LOG(INFO) << "No explicit mode flag provided, checking for hint file...";
        std::string mode_from_hint;
        if (android::base::ReadFileToString(kSnapuserdModeHintFile, &mode_from_hint)) {
            mode_from_hint = android::base::Trim(mode_from_hint);
            if (mode_from_hint == kSnapuserdModeUblk) {
                LOG(INFO) << "Using ublk mode based on hint file.";
                FLAGS_ublk = true;
                ublk_flag_is_default = false;  // Hint file has priority
            } else if (mode_from_hint == kSnapuserdModeDmUser) {
                LOG(INFO) << "Using dm-user mode based on hint file.";
                FLAGS_ublk = false;
                ublk_flag_is_default = false;  // Hint file has priority
            }
        }
    } else {
        LOG(INFO) << "Mode explicitly set via command-line flag: "
                  << (FLAGS_ublk ? "ublk" : "dm-user");
    }

    if (ublk_flag_is_default) {
        LOG(INFO) << "No hint file found or hint was invalid, proceeding with auto-detection.";
        // If -ublk or -noublk was not explicitly passed on the command-line,
        // fall back to auto-detection. Otherwise, the flag passed by the
        // caller (libsnapshot) is respected.
        // Check if test wants to force dm_user mode.
        if (android::base::GetProperty("snapuserd.test.force_dm_user", "") == "true") {
            LOG(INFO) << "Forcing dm-user mode via test property.";
            FLAGS_ublk = false;
        } else {
            // Otherwise, perform normal auto-detection.
            FLAGS_ublk = IsUblkEnabled();
            LOG(INFO) << "Auto-detection selected mode: " << (FLAGS_ublk ? "ublk" : "dm-user");
        }
    }
    // Set block_server_opener_
    user_server_.Initialize(FLAGS_ublk);

    user_server_.SetServerRunning();

    if (FLAGS_socket_handoff) {
        return user_server_.RunForSocketHandoff();
    }
    if (!FLAGS_no_socket) {
        if (!user_server_.Start(FLAGS_socket)) {
            return false;
        }
        return user_server_.Run();
    }

    UeventHelperCallback uevent_helper = [&](const std::string& device_path) {
        if (socket_ == -1) {
            LOG(INFO) << "socket_ not initialized, dropping req for device:" << device_path;
            return false;
        }
        return SendUeventRequest(socket_, device_path);
    };
    if (FLAGS_fd_num != -1) {
        LOG(INFO) << "using fd from socketpair() " << FLAGS_fd_num;
        socket_ = FLAGS_fd_num;
        user_server_.SetUeventHelper(uevent_helper);
    }
    for (int i = arg_start; i < argc; i++) {
        auto parts = android::base::Split(argv[i], ",");
        if (parts.size() != 4) {
            LOG(ERROR) << "Malformed message, expected four sub-arguments.";
            return false;
        }
        HandlerOptions options = {
                .num_worker_threads = FLAGS_worker_count,
                .use_iouring = FLAGS_io_uring,
                .o_direct = FLAGS_o_direct,
                .skip_verification = FLAGS_skip_verification,
                .cow_op_merge_size = static_cast<uint32_t>(FLAGS_cow_op_merge_size),
                .verify_block_size = static_cast<uint32_t>(FLAGS_verify_block_size),
                .num_verification_threads = static_cast<uint32_t>(FLAGS_num_verify_threads),
        };
        auto handler = user_server_.AddHandler(parts[0], parts[1], parts[2], parts[3], options);
        if (!handler || !user_server_.StartHandler(parts[0])) {
            return false;
        }
    }
    if (FLAGS_ublk && FLAGS_fd_num != -1) {
        SendDoneNotification(socket_);
        close(socket_);
        socket_ = -1;
    }
    LOG(INFO) << "DONE Starting daemon, notifying init";

    // We reach this point only during selinux transition during device boot.
    // At this point, all threads are spin up and are ready to serve the I/O
    // requests for dm-user. Lets inform init.
    auto client = std::make_unique<SnapuserdClient>();
    client->NotifyTransitionDaemonIsReady();

    // Skip the accept() call to avoid spurious log spam. The server will still
    // run until all handlers have completed.
    return user_server_.WaitForSocket();
}

void Daemon::MaskAllSignalsExceptIntAndTerm() {
    sigset_t signal_mask;
    sigfillset(&signal_mask);
    sigdelset(&signal_mask, SIGINT);
    sigdelset(&signal_mask, SIGTERM);
    sigdelset(&signal_mask, SIGPIPE);
    sigdelset(&signal_mask, SIGUSR1);
    if (sigprocmask(SIG_SETMASK, &signal_mask, NULL) != 0) {
        PLOG(ERROR) << "Failed to set sigprocmask";
    }
}

void Daemon::MaskAllSignals() {
    sigset_t signal_mask;
    sigfillset(&signal_mask);
    if (sigprocmask(SIG_SETMASK, &signal_mask, NULL) != 0) {
        PLOG(ERROR) << "Couldn't mask all signals";
    }
}

void Daemon::Interrupt() {
    // TODO: We cannot access system property during first stage init.
    // Until we remove the dm-snapshot code, we will have this check
    // and verify it through a temp variable.
    if (user_server_.IsServerRunning()) {
        user_server_.Interrupt();
    }
}

void Daemon::ReceivedSocketSignal() {
    if (user_server_.IsServerRunning()) {
        user_server_.ReceivedSocketSignal();
    }
}

void Daemon::SignalHandler(int signal) {
    LOG(DEBUG) << "Snapuserd received signal: " << signal;
    switch (signal) {
        case SIGINT:
        case SIGTERM: {
            Daemon::Instance().Interrupt();
            break;
        }
        case SIGPIPE: {
            LOG(ERROR) << "Received SIGPIPE signal";
            break;
        }
        case SIGUSR1: {
            LOG(INFO) << "Received SIGUSR1, attaching to proxy socket";
            Daemon::Instance().ReceivedSocketSignal();
            break;
        }
        default:
            LOG(ERROR) << "Received unknown signal " << signal;
            break;
    }
}

}  // namespace snapshot
}  // namespace android

int main(int argc, char** argv) {
    android::base::InitLogging(argv, &android::base::KernelLogger);

    android::snapshot::Daemon& daemon = android::snapshot::Daemon::Instance();

    if (!daemon.StartDaemon(argc, argv)) {
        LOG(ERROR) << "Snapuserd daemon failed to start";
        exit(EXIT_FAILURE);
    }

    return 0;
}
