// Copyright 2018 The Fuchsia Authors
//
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT

#include <object/suspend_token_dispatcher.h>

#include <err.h>

#include <fbl/alloc_checker.h>
#include <fbl/unique_ptr.h>
#include <lib/counters.h>
#include <kernel/auto_lock.h>
#include <object/process_dispatcher.h>
#include <object/thread_dispatcher.h>
#include <zircon/rights.h>

KCOUNTER(dispatcher_suspend_token_create_count, "dispatcher.suspend_token.create")
KCOUNTER(dispatcher_suspend_token_destroy_count, "dispatcher.suspend_token.destroy")

namespace {

// Suspends a process or thread.
// TODO(ZX-858): Add support for jobs.
zx_status_t SuspendTask(fbl::RefPtr<Dispatcher> task) {
    if (auto thread = DownCastDispatcher<ThreadDispatcher>(&task)) {
        if (thread.get() == ThreadDispatcher::GetCurrent())
            return ZX_ERR_NOT_SUPPORTED;
        return thread->Suspend();
    }

    if (auto process = DownCastDispatcher<ProcessDispatcher>(&task)) {
        if (process.get() == ProcessDispatcher::GetCurrent())
            return ZX_ERR_NOT_SUPPORTED;
        return process->Suspend();
    }

    return ZX_ERR_WRONG_TYPE;
}

// Resumes a process or thread.
// TODO(ZX-858): Add support for jobs.
void ResumeTask(fbl::RefPtr<Dispatcher> task) {
    if (auto thread = DownCastDispatcher<ThreadDispatcher>(&task)) {
        thread->Resume();
        return;
    }

    if (auto process = DownCastDispatcher<ProcessDispatcher>(&task)) {
        process->Resume();
        return;
    }

    __UNREACHABLE;
}

} // namespace

zx_status_t SuspendTokenDispatcher::Create(fbl::RefPtr<Dispatcher> task,
                                           KernelHandle<SuspendTokenDispatcher>* handle,
                                           zx_rights_t* rights) {
    fbl::AllocChecker ac;
    KernelHandle new_handle(fbl::AdoptRef(new (&ac) SuspendTokenDispatcher()));
    if (!ac.check())
        return ZX_ERR_NO_MEMORY;

    zx_status_t status = SuspendTask(task);
    if (status != ZX_OK)
        return status;

    // Save the task after suspending so that on_zero_handles() resumes it.
    new_handle.dispatcher()->task_ = ktl::move(task);

    *rights = default_rights();
    *handle = ktl::move(new_handle);
    return ZX_OK;
}

SuspendTokenDispatcher::SuspendTokenDispatcher() {
    kcounter_add(dispatcher_suspend_token_create_count, 1);
}

SuspendTokenDispatcher::~SuspendTokenDispatcher() {
    kcounter_add(dispatcher_suspend_token_destroy_count, 1);
}

void SuspendTokenDispatcher::on_zero_handles() {
    // This is only called once and we're done with |task_| afterwards so we can move it out.
    if (task_)
        ResumeTask(ktl::move(task_));
}